diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 9e45a5f4e..5e76572f9 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -1168,6 +1168,38 @@ def _generate_torch_header(name: str, ops: list[Op]) -> str: op_calls = "\n\n".join( f" void operator()({_format_signature(op)}) const override;" for op in ops ) + pool_members = [] + seen_pool_members = set() + + for op in ops: + for param in op.visible_params: + api_name = param.api_name + + if param.is_tensor or param.is_optional_tensor: + pool_kind = "single" + elif param.is_tensor_list: + pool_kind = "list" + else: + continue + + pool_key = (api_name, pool_kind) + + if pool_key in seen_pool_members: + continue + + seen_pool_members.add(pool_key) + + if pool_kind == "single": + pool_members.append( + " mutable TargetTensorPool> " + f"{api_name}_pool_;" + ) + else: + pool_members.append( + " mutable std::vector<" + "TargetTensorPool>> " + f"{api_name}_pools_;" + ) return _TORCH_HEADER_TEMPLATE.format( name_uc=name.upper(), @@ -1175,6 +1207,7 @@ def _generate_torch_header(name: str, ops: list[Op]) -> str: pascal=pascal, op_type=op_type, op_calls=op_calls, + pool_members="\n\n".join(pool_members), slot=_PYTORCH_SLOT, ) @@ -1182,8 +1215,6 @@ def _generate_torch_header(name: str, ops: list[Op]) -> str: def _generate_torch_method_source(name: str, op: Op) -> str: op_type = _op_cpp_type(name) conversion_lines = [] - out_device_index = f"{op.out_params[0].api_name}.device().index()" - conversion_lines.append(f" const auto device_index = {out_device_index};") def _optional_aten_type(param: Param) -> str: return _NULLOPT_BY_TYPE[param.aten_type].removesuffix("{}") @@ -1191,15 +1222,6 @@ def _optional_aten_type(param: Param) -> str: def _optional_aten_value(schema_param: Param, api_param: Param) -> str: api_name = api_param.api_name - if schema_param.aten_type == "Tensor?": - data_expr = f"const_cast({api_name}->data())" - - return ( - f"ToAtenTensor({data_expr}, {api_name}->shape(), " - f"{api_name}->strides(), {api_name}->dtype(), " - f"device_index)" - ) - if schema_param.aten_type == "Scalar?": return f"at::Scalar(*{api_name})" @@ -1225,6 +1247,25 @@ def _optional_aten_value(schema_param: Param, api_param: Param) -> str: def _append_optional_conversion(schema_param: Param, api_param: Param) -> None: api_name = api_param.api_name optional_type = _optional_aten_type(schema_param) + + if schema_param.aten_type == "Tensor?": + conversion_lines.append( + " std::optional>::Lease> " + f"{schema_param.name}_lease;" + ) + conversion_lines.append(f" {optional_type} at_{schema_param.name};") + conversion_lines.append(f" if ({api_name}.has_value()) {{") + conversion_lines.append( + f" {schema_param.name}_lease.emplace(" + f"{api_name}_pool_.Acquire(*{api_name}));" + ) + conversion_lines.append( + f" at_{schema_param.name} = {schema_param.name}_lease->Native();" + ) + conversion_lines.append(" }") + return + conversion_lines.append(f" {optional_type} at_{schema_param.name};") conversion_lines.append(f" if ({api_name}.has_value()) {{") conversion_lines.append( @@ -1245,27 +1286,39 @@ def _append_optional_conversion(schema_param: Param, api_param: Param) -> None: api_name = api_param.api_name if param.is_tensor_list: + conversion_lines.append( + f" if ({api_name}_pools_.size() < {api_name}.size()) {{" + ) + conversion_lines.append(f" {api_name}_pools_.resize({api_name}.size());") + conversion_lines.append(" }") + conversion_lines.append( + " std::vector>::Lease> " + f"{param.name}_leases;" + ) + conversion_lines.append( + f" {param.name}_leases.reserve({api_name}.size());" + ) conversion_lines.append(f" std::vector at_{param.name};") conversion_lines.append(f" at_{param.name}.reserve({api_name}.size());") - conversion_lines.append(f" for (const auto& tensor : {api_name}) {{") conversion_lines.append( - " at_" - f"{param.name}.push_back(ToAtenTensor(" - "const_cast(tensor.data()), tensor.shape(), tensor.strides(), " - "tensor.dtype(), device_index));" + f" for (std::size_t i = 0; i < {api_name}.size(); ++i) {{" + ) + conversion_lines.append( + f" {param.name}_leases.push_back(" + f"{api_name}_pools_[i].Acquire({api_name}[i]));" + ) + conversion_lines.append( + f" at_{param.name}.push_back({param.name}_leases.back().Native());" ) conversion_lines.append(" }") continue - data_expr = ( - f"{api_name}.data()" - if param.is_mutable_tensor - else f"const_cast({api_name}.data())" + conversion_lines.append( + f" auto {param.name}_lease = {api_name}_pool_.Acquire({api_name});" ) conversion_lines.append( - f" auto at_{param.name} = ToAtenTensor(\n" - f" {data_expr}, {api_name}_shape_, {api_name}_strides_,\n" - f" {api_name}_type_, device_index);" + f" auto& at_{param.name} = {param.name}_lease.Native();" ) for schema_index, param in enumerate(op.params): @@ -1325,7 +1378,6 @@ def _render_arg(schema_index, p): def _generate_torch_source(name: str, ops: list[Op]) -> str: op_type = _op_cpp_type(name) - methods = "\n\n".join(_generate_torch_method_source(name, op) for op in ops) # Guard each explicit instantiation by the matching `WITH_` macro # so a build that only enables a subset of devices does not pay the # ATen template-instantiation cost (and memory pressure) for the @@ -1338,6 +1390,8 @@ def _generate_torch_source(name: str, ops: list[Op]) -> str: for dev in _DEVICE_TYPES ) + methods = "\n\n".join(_generate_torch_method_source(name, op) for op in ops) + return _TORCH_SOURCE_TEMPLATE.format( name=name, methods=methods, @@ -1374,6 +1428,8 @@ class {pascal} : public Operator<{pascal}> {{ #define INFINI_OPS_TORCH_{name_uc}_H_ #include "base/{name}.h" +#include "target_tensor_pool.h" +#include "torch/tensor_.h" namespace infini::ops {{ @@ -1383,6 +1439,9 @@ class Operator<{op_type}, kDev, {slot}> : public {op_type} {{ using {op_type}::{pascal}; {op_calls} + + private: +{pool_members} }}; }} // namespace infini::ops diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 138312707..d2fa53c54 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -556,6 +556,13 @@ if(WITH_TORCH) ".*/flash_attn_varlen_func/flash_attn_varlen_func\\.cc$") endif() + # This source belongs to the Python module, not the core ATen backend. + set(TORCH_PYBIND_BRIDGE_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/torch/pybind11_.cc") + set(TORCH_PYBIND_BRIDGE_HEADER + "${CMAKE_CURRENT_SOURCE_DIR}/torch/pybind11_.h") + list(REMOVE_ITEM TORCH_SOURCES "${TORCH_PYBIND_BRIDGE_SOURCE}") + set(INFINI_OPS_TORCH_UNITY_BATCH_SIZE "8" CACHE STRING "Number of torch sources to include in each generated unity translation unit; set to 1 to disable") set(TORCH_COMPILE_SOURCES ${TORCH_SOURCES}) @@ -1019,7 +1026,11 @@ if(GENERATE_PYTHON_BINDINGS) list(APPEND PYBIND11_COMPILE_SOURCES ${ILUVATAR_DISPATCH_OBJECTS}) endif() - find_package(Python COMPONENTS Interpreter Development) + if(WITH_TORCH) + find_package(Python COMPONENTS Interpreter Development REQUIRED) + else() + find_package(Python COMPONENTS Interpreter Development) + endif() if(NOT pybind11_DIR) execute_process( @@ -1046,6 +1057,70 @@ if(GENERATE_PYTHON_BINDINGS) "$<$:-Wno-deprecated-declarations>" ) + if(WITH_TORCH) + find_library(TORCH_PYTHON_LIB torch_python + HINTS ${_torch_lib_dirs} REQUIRED) + + if(WITH_METAX OR WITH_MOORE) + set(_torch_bridge_include_flags "") + foreach(_dir ${TORCH_INCLUDE_DIRS} ${Python_INCLUDE_DIRS}) + list(APPEND _torch_bridge_include_flags "-isystem" "${_dir}") + endforeach() + + set(_torch_bridge_object_dir + "${CMAKE_CURRENT_BINARY_DIR}/binding_objs") + file(MAKE_DIRECTORY "${_torch_bridge_object_dir}") + set(_torch_bridge_obj + "${_torch_bridge_object_dir}/torch_pybind11_.o") + set(_torch_bridge_dep "${_torch_bridge_obj}.d") + set(_torch_bridge_depfile_arg) + if(CMAKE_GENERATOR MATCHES "Ninja") + set(_torch_bridge_depfile_arg DEPFILE "${_torch_bridge_dep}") + endif() + add_custom_command( + OUTPUT "${_torch_bridge_obj}" + COMMAND ${SYSTEM_CXX} + -std=c++17 -fPIC -O2 + "-I${CMAKE_CURRENT_SOURCE_DIR}" + "-I${PROJECT_SOURCE_DIR}/generated" + ${INFINI_RT_INCLUDE_FLAGS} + ${_torch_vendor_include_flags} + ${_torch_bridge_include_flags} + ${_torch_extra_flags} + -MMD -MF "${_torch_bridge_dep}" + -c "${TORCH_PYBIND_BRIDGE_SOURCE}" + -o "${_torch_bridge_obj}" + DEPENDS + "${TORCH_PYBIND_BRIDGE_SOURCE}" + "${TORCH_PYBIND_BRIDGE_HEADER}" + ${_torch_bridge_depfile_arg} + COMMENT "Compiling torch/pybind11_.cc with system C++ compiler" + VERBATIM + ) + set_source_files_properties("${_torch_bridge_obj}" + PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE) + target_sources(ops PRIVATE "${_torch_bridge_obj}") + else() + add_library(infini_ops_torch_bridge_obj OBJECT + ${TORCH_PYBIND_BRIDGE_SOURCE}) + set_target_properties(infini_ops_torch_bridge_obj + PROPERTIES POSITION_INDEPENDENT_CODE ON) + target_include_directories(infini_ops_torch_bridge_obj PRIVATE + $ + ${INFINI_RT_INCLUDE_DIRS} + ${TORCH_INCLUDE_DIRS} + ${Python_INCLUDE_DIRS} + ${PROJECT_SOURCE_DIR}/generated) + target_compile_definitions(infini_ops_torch_bridge_obj PRIVATE + $) + target_compile_options(infini_ops_torch_bridge_obj PRIVATE + $ + -O2) + target_sources(ops PRIVATE + $) + endif() + endif() + if(WITH_TORCH AND CMAKE_GENERATOR MATCHES "Ninja") set(INFINI_OPS_BINDING_COMPILE_JOBS "2" CACHE STRING "Maximum concurrent generated pybind11 binding compilations") @@ -1078,6 +1153,9 @@ if(GENERATE_PYTHON_BINDINGS) ${INFINI_OPS_NINETOOTHED_INCLUDE_DIRS}) endif() target_link_libraries(ops PRIVATE infiniops) + if(WITH_TORCH) + target_link_libraries(ops PRIVATE ${TORCH_PYTHON_LIB}) + endif() # Cambricon generated dispatch is compiled into the Python extension and # directly references host launch stubs emitted from `.mlu` sources. Link diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index c60748228..06e60cfcd 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -8,6 +8,9 @@ #include "tensor.h" #include "torch/device_.h" +#ifdef WITH_TORCH +#include "torch/pybind11_.h" +#endif namespace py = pybind11; @@ -131,6 +134,18 @@ inline Device DeviceFromPybind11Handle(py::handle obj) { } inline Tensor TensorFromPybind11Handle(py::handle obj) { +#ifdef WITH_TORCH + auto metadata{TryAtenTensorMetadataFromPyObject(obj.ptr())}; + + if (metadata.has_value()) { + Device device{DeviceTypeFromString(metadata->device_type), + metadata->device_index}; + + return Tensor{metadata->data, std::move(metadata->shape), metadata->dtype, + device, std::move(metadata->strides)}; + } +#endif + auto data{ reinterpret_cast(obj.attr("data_ptr")().cast())}; diff --git a/src/target_tensor_pool.h b/src/target_tensor_pool.h new file mode 100644 index 000000000..eb25d9ca6 --- /dev/null +++ b/src/target_tensor_pool.h @@ -0,0 +1,115 @@ +#ifndef INFINI_OPS_TARGET_TENSOR_POOL_H_ +#define INFINI_OPS_TARGET_TENSOR_POOL_H_ + +#include +#include +#include +#include + +namespace infini::ops { + +template +class TargetTensorPool { + struct Entry { + explicit Entry(std::unique_ptr state) + : state{std::move(state)}, busy{true} {} + + std::unique_ptr state; + + bool busy; + }; + + public: + class Lease { + public: + Lease() = default; + + Lease(const Lease&) = delete; + + Lease& operator=(const Lease&) = delete; + + Lease(Lease&& other) noexcept + : entry_{std::exchange(other.entry_, nullptr)} {} + + Lease& operator=(Lease&& other) noexcept { + if (this != &other) { + Release(); + entry_ = std::exchange(other.entry_, nullptr); + } + + return *this; + } + + ~Lease() noexcept { Release(); } + + decltype(auto) Native() const { return Adapter::Native(*entry_->state); } + + private: + explicit Lease(Entry* entry) : entry_{entry} {} + + void Release() noexcept { + if (entry_ != nullptr) { + entry_->busy = false; + entry_ = nullptr; + } + } + + Entry* entry_ = nullptr; + + friend class TargetTensorPool; + }; + + TargetTensorPool() = default; + + TargetTensorPool(const TargetTensorPool&) = delete; + + TargetTensorPool& operator=(const TargetTensorPool&) = delete; + + TargetTensorPool(TargetTensorPool&&) noexcept = default; + + TargetTensorPool& operator=(TargetTensorPool&&) = delete; + + ~TargetTensorPool() { +#ifndef NDEBUG + for (const auto& entry : entries_) { + assert(!entry->busy && + "`TargetTensorPool` destroyed with an active lease."); + } +#endif + } + + template + Lease Acquire(const TensorLike& tensor) { + static_assert( + noexcept(Adapter::Rebind(std::declval(), + std::declval())), + "`Adapter::Rebind` must be `noexcept`."); + + for (auto& entry : entries_) { + auto* entry_pointer = entry.get(); + + if (entry_pointer->busy) { + continue; + } + + entry_pointer->busy = true; + Adapter::Rebind(*entry_pointer->state, tensor); + + return Lease{entry_pointer}; + } + + auto state = Adapter::Create(tensor); + auto entry = std::make_unique(std::move(state)); + auto* entry_pointer = entry.get(); + entries_.push_back(std::move(entry)); + + return Lease{entry_pointer}; + } + + private: + std::vector> entries_; +}; + +} // namespace infini::ops + +#endif diff --git a/src/torch/pybind11_.cc b/src/torch/pybind11_.cc new file mode 100644 index 000000000..5d9a3aee4 --- /dev/null +++ b/src/torch/pybind11_.cc @@ -0,0 +1,90 @@ +#include "torch/pybind11_.h" + +#include + +#include + +#include "torch/csrc/autograd/python_variable.h" + +namespace infini::ops { + +namespace { + +std::optional TryDataTypeFromAten(at::ScalarType scalar_type) { + switch (scalar_type) { + case at::kChar: + return DataType::kInt8; + case at::kShort: + return DataType::kInt16; + case at::kInt: + return DataType::kInt32; + case at::kLong: + return DataType::kInt64; + case at::kByte: + return DataType::kUInt8; + case at::kBool: + return DataType::kUInt8; + case at::kHalf: + return DataType::kFloat16; + case at::kBFloat16: + return DataType::kBFloat16; + case at::kFloat: + return DataType::kFloat32; + case at::kDouble: + return DataType::kFloat64; +#if TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 4) + case at::ScalarType::UInt16: + return DataType::kUInt16; + case at::ScalarType::UInt32: + return DataType::kUInt32; + case at::ScalarType::UInt64: + return DataType::kUInt64; +#endif + default: + return std::nullopt; + } +} + +std::string DeviceTypeFromAten(const c10::Device& device) { + if (device.type() == c10::kCPU) { + return "cpu"; + } + + if (device.type() == c10::kCUDA) { + return "cuda"; + } + + std::string name{device.str()}; + auto colon{name.find(':')}; + return colon == std::string::npos ? name : name.substr(0, colon); +} + +} // namespace + +std::optional TryAtenTensorMetadataFromPyObject( + void* py_object) { + auto* object{reinterpret_cast(py_object)}; + if (!THPVariable_Check(object)) { + return std::nullopt; + } + + const at::Tensor& tensor{THPVariable_Unpack(object)}; + const c10::Device device{tensor.device()}; + auto dtype{TryDataTypeFromAten(tensor.scalar_type())}; + + if (!dtype.has_value()) { + return std::nullopt; + } + + return AtenTensorMetadata{ + tensor.data_ptr(), + Tensor::Shape(tensor.sizes().begin(), tensor.sizes().end()), + Tensor::Strides(tensor.strides().begin(), tensor.strides().end()), + *dtype, + DeviceTypeFromAten(device), + device.has_index() ? device.index() : 0, + }; +} + +} // namespace infini::ops diff --git a/src/torch/pybind11_.h b/src/torch/pybind11_.h new file mode 100644 index 000000000..c848d04de --- /dev/null +++ b/src/torch/pybind11_.h @@ -0,0 +1,32 @@ +#ifndef INFINI_OPS_TORCH_PYBIND11__H_ +#define INFINI_OPS_TORCH_PYBIND11__H_ + +#include +#include + +#include "data_type.h" +#include "tensor.h" + +namespace infini::ops { + +// Torch-independent metadata returned by the host-compiled source adapter. +struct AtenTensorMetadata { + void* data; + + Tensor::Shape shape; + + Tensor::Strides strides; + + DataType dtype; + + std::string device_type; + + int device_index; +}; + +std::optional TryAtenTensorMetadataFromPyObject( + void* py_object); + +} // namespace infini::ops + +#endif diff --git a/src/torch/tensor_.h b/src/torch/tensor_.h index e199bf704..15e4a9dbf 100644 --- a/src/torch/tensor_.h +++ b/src/torch/tensor_.h @@ -4,7 +4,9 @@ #include #include +#include #include +#include #include "tensor.h" #include "torch/device_.h" @@ -108,6 +110,34 @@ inline at::Tensor ToAtenTensor( return at::from_blob(data, at_shape, at_strides, options); } +template +struct AtenTensorAdapter { + struct State { + at::Tensor tensor; + + c10::Device device; + }; + + using NativeHandle = at::Tensor&; + + static std::unique_ptr Create(const Tensor& tensor) { + auto at_tensor = ToAtenTensor( + const_cast(tensor.data()), tensor.shape(), tensor.strides(), + tensor.dtype(), tensor.device().index()); + auto device = at_tensor.device(); + + return std::make_unique( + State{std::move(at_tensor), std::move(device)}); + } + + static void Rebind(State& state, const Tensor& tensor) noexcept { + state.tensor.storage().set_data_ptr_noswap( + c10::DataPtr(const_cast(tensor.data()), state.device)); + } + + static NativeHandle Native(State& state) { return state.tensor; } +}; + } // namespace infini::ops #endif diff --git a/tests/conftest.py b/tests/conftest.py index 1361c1752..f80396409 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -167,6 +167,7 @@ def _set_random_seed(seed): "tests/test_generate_ninetoothed_ops.py", "tests/test_generate_torch_ops.py", "tests/test_generate_wrappers.py", + "tests/test_target_tensor_pool.py", } _SMOKE_TORCH_OPS = {"abs", "clamp", "exp"} diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index 165511180..fc86ea424 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -49,11 +49,121 @@ def test_schema_self_param_renders_as_input_in_public_cpp_api(): assert "Softmax(const Tensor input, const int64_t dim" in base assert "self_shape_" not in base assert "input_shape_" in base - assert "auto at_self = ToAtenTensor" in source - assert "input_shape_" in source + assert "auto self_lease = input_pool_.Acquire(input);" in source + assert "auto& at_self = self_lease.Native();" in source + assert "self_pool_" not in source assert "at::_softmax_out(at_out, at_self" in source +def test_abs_torch_backend_uses_generic_target_tensor_pools(): + module = _load_generator_module() + op = module._parse_func("abs(Tensor self, *, Tensor(a!) out) -> Tensor(a!)") + + header = module._generate_torch_header("abs", [op]) + source = module._generate_torch_source("abs", [op]) + + assert "using Abs::Abs;" in header + assert "TargetTensorPool> input_pool_;" in header + assert "TargetTensorPool> out_pool_;" in header + assert "auto self_lease = input_pool_.Acquire(input);" in source + assert "auto& at_self = self_lease.Native();" in source + assert "auto out_lease = out_pool_.Acquire(out);" in source + assert "auto& at_out = out_lease.Native();" in source + assert "AtenTensorSlot" not in header + assert "AtenTensorSlot" not in source + assert "ToAtenTensor" not in source + + +def test_pybind_tensor_metadata_bridge_does_not_retain_source_handles(): + root = pathlib.Path(__file__).resolve().parent.parent + pybind_utils = (root / "src" / "pybind11_utils.h").read_text(encoding="utf-8") + bridge_header = (root / "src" / "torch" / "pybind11_.h").read_text(encoding="utf-8") + bridge_source = (root / "src" / "torch" / "pybind11_.cc").read_text( + encoding="utf-8" + ) + + assert '"torch/pybind11_.h"' in pybind_utils + assert "TryAtenTensorMetadataFromPyObject" in pybind_utils + assert "if (metadata.has_value())" in pybind_utils + assert "#ifdef WITH_TORCH" in pybind_utils + assert 'obj.attr("data_ptr")' in pybind_utils + assert "torch/csrc/" not in pybind_utils + assert "torch/csrc/autograd/python_variable.h" not in bridge_header + assert "torch/csrc/autograd/python_variable.h" in bridge_source + assert "THPVariable_Check" in bridge_source + assert "THPVariable_Unpack" in bridge_source + assert "std::optional" in bridge_header + assert "TryAtenTensorMetadataFromPyObject" in bridge_header + assert "case at::kBool:" in bridge_source + assert "throw " not in bridge_source + assert "assert(" not in bridge_source + assert "return std::nullopt;" in bridge_source + assert "if (!THPVariable_Check(object))" in bridge_source + assert "if (!dtype.has_value())" in bridge_source + assert "device.has_index() ? device.index() : 0" in bridge_source + assert "source_handle" not in bridge_header + assert "source_handle" not in bridge_source + assert "make_shared" not in bridge_source + + +def test_pybind_tensor_metadata_bridge_is_scoped_to_python_module(): + root = pathlib.Path(__file__).resolve().parent.parent + cmake = (root / "src" / "CMakeLists.txt").read_text(encoding="utf-8") + + torch_section = cmake.split("if(WITH_TORCH)", maxsplit=1)[1] + torch_section = torch_section.split("\nendif()", maxsplit=1)[0] + + assert "Development" not in torch_section + assert "find_package(Python COMPONENTS Interpreter Development REQUIRED)" in cmake + assert 'list(REMOVE_ITEM TORCH_SOURCES "${TORCH_PYBIND_BRIDGE_SOURCE}")' in cmake + assert "add_library(infini_ops_torch_bridge_obj OBJECT" in cmake + assert "-std=c++17 -fPIC -O2" in cmake + assert '-MMD -MF "${_torch_bridge_dep}"' in cmake + assert 'DEPFILE "${_torch_bridge_dep}"' in cmake + assert '"${TORCH_PYBIND_BRIDGE_HEADER}"' in cmake + assert "target_sources(ops PRIVATE" in cmake + core_bridge = ( + "target_sources(infiniops PRIVATE\n" + " $)" + ) + assert core_bridge not in cmake + assert "find_library(TORCH_PYTHON_LIB torch_python" in cmake + assert "target_link_libraries(ops PRIVATE ${TORCH_PYTHON_LIB})" in cmake + + +def test_torch_header_orders_and_deduplicates_tensor_pools_across_overloads(): + module = _load_generator_module() + tensor_overload = module._parse_func( + "blend.Tensor(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!)" + ) + optional_overload = module._parse_func( + "blend.optional(Tensor self, Tensor? weight=None, Tensor bias, *, " + "Tensor(a!) out) -> Tensor(a!)" + ) + list_overload = module._parse_func( + "blend.list(Tensor self, Tensor[] other, Tensor[] tensors, *, " + "Tensor(a!) out) -> Tensor(a!)" + ) + + header = module._generate_torch_header( + "blend", [tensor_overload, optional_overload, list_overload] + ) + members = [ + "TargetTensorPool> input_pool_;", + "TargetTensorPool> other_pool_;", + "TargetTensorPool> out_pool_;", + "TargetTensorPool> weight_pool_;", + "TargetTensorPool> bias_pool_;", + "std::vector>> other_pools_;", + "std::vector>> tensors_pools_;", + ] + + assert [header.count(member) for member in members] == [1] * len(members) + assert [header.index(member) for member in members] == sorted( + header.index(member) for member in members + ) + + def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): module = _load_generator_module() op = module._parse_func( @@ -73,6 +183,7 @@ def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): ] base = module._generate_base_header("batch_norm_elemt", [op]) + header = module._generate_torch_header("batch_norm_elemt", [op]) source = module._generate_torch_method_source("batch_norm_elemt", op) assert "#include " in base @@ -82,7 +193,31 @@ def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): assert "bool has_bias_" in base assert "c10::optional at_weight" in source assert "c10::optional at_bias" in source - assert "weight->shape()" in source + assert ( + "std::optional>::Lease> weight_lease;" in source + ) + assert ( + "std::optional>::Lease> bias_lease;" in source + ) + assert "weight_lease.emplace(weight_pool_.Acquire(*weight));" in source + assert "at_weight = weight_lease->Native();" in source + assert "bias_lease.emplace(bias_pool_.Acquire(*bias));" in source + assert "at_bias = bias_lease->Native();" in source + assert source.index("weight_lease;") < source.index("at_weight;") + assert source.index("bias_lease;") < source.index("at_bias;") + assert source.index("at_weight;") < source.index("weight_lease.emplace") + assert source.index("at_bias;") < source.index("bias_lease.emplace") + assert source.index("at_weight = weight_lease->Native();") < source.index( + "at::batch_norm_elemt_out" + ) + assert source.index("at_bias = bias_lease->Native();") < source.index( + "at::batch_norm_elemt_out" + ) + assert "weight_pool_;" in header + assert "bias_pool_;" in header + assert "ToAtenTensor" not in source assert "weight_shape_" not in source assert "at::batch_norm_elemt_out" in source assert "at_weight" in source @@ -103,13 +238,40 @@ def test_tensor_list_params_are_exposed_and_forwarded_to_aten(): assert op.is_testable base = module._generate_base_header("stack", [op]) + header = module._generate_torch_header("stack", [op]) source = module._generate_torch_method_source("stack", op) assert "#include " in base assert "std::vector tensors" in base - assert "std::vector at_tensors" in source - assert "at_tensors.reserve(tensors.size())" in source - assert "for (const auto& tensor : tensors)" in source + list_pool = "std::vector>> tensors_pools_;" + lease_vector = ( + "std::vector>::Lease> tensors_leases;" + ) + assert list_pool in header + assert header.count(list_pool) == 1 + assert "tensors_pool_;" not in header + assert "if (tensors_pools_.size() < tensors.size()) {" in source + assert "tensors_pools_.size() != tensors.size()" not in source + assert "tensors_pools_.resize(tensors.size());" in source + assert lease_vector in source + assert "tensors_leases.reserve(tensors.size());" in source + assert "std::vector at_tensors;" in source + assert "at_tensors.reserve(tensors.size());" in source + assert "for (std::size_t i = 0; i < tensors.size(); ++i)" in source + assert "tensors_leases.push_back(tensors_pools_[i].Acquire(tensors[i]));" in source + assert "at_tensors.push_back(tensors_leases.back().Native());" in source + assert source.index("tensors_pools_.resize") < source.index(lease_vector) + assert source.index("tensors_pools_.resize") < source.index( + "tensors_pools_[i].Acquire" + ) + assert source.index(lease_vector) < source.index( + "std::vector at_tensors;" + ) + assert source.index("std::vector at_tensors;") < source.index( + "at::stack_out" + ) + assert "ToAtenTensor" not in source assert "at::stack_out(at_out, at_tensors, dim)" in source @@ -145,6 +307,8 @@ def test_optional_scalar_and_array_params_are_exposed_and_forwarded_to_aten(): upsample_source = module._generate_torch_method_source( "upsample_bicubic2d", upsample ) + quantile_header = module._generate_torch_header("quantile", [quantile]) + upsample_header = module._generate_torch_header("upsample_bicubic2d", [upsample]) assert "c10::optional at_dim" in quantile_source assert "at::quantile_out" in quantile_source @@ -152,6 +316,8 @@ def test_optional_scalar_and_array_params_are_exposed_and_forwarded_to_aten(): assert "c10::optional> at_scale_factors" in upsample_source assert "at::upsample_bicubic2d_out" in upsample_source assert "at_scale_factors" in upsample_source + assert "dim_pool_" not in quantile_header + assert "scale_factors_pool_" not in upsample_header def test_required_scalar_type_params_use_public_data_type(): @@ -204,9 +370,11 @@ def test_existing_base_overload_can_omit_optional_schema_params(): ] source = module._generate_torch_method_source("slow_conv3d", bound) + header = module._generate_torch_header("slow_conv3d", [bound]) assert "std::optional bias" not in source assert "c10::optional{}" in source + assert "bias_pool_" not in header assert "at::slow_conv3d_out" in source @@ -229,7 +397,7 @@ def test_existing_base_overload_can_omit_defaulted_schema_params(): source = module._generate_torch_method_source("add", bound) assert "double alpha" not in source - assert "const auto device_index = out.device().index();" in source + assert "const auto device_index" not in source assert "device_index_)" not in source assert "at::add_out(at_out, at_self, at_other, 1)" in source diff --git a/tests/test_target_tensor_pool.py b/tests/test_target_tensor_pool.py new file mode 100644 index 000000000..fcf2752cd --- /dev/null +++ b/tests/test_target_tensor_pool.py @@ -0,0 +1,225 @@ +import os +import subprocess +import textwrap +from pathlib import Path + +import pytest + + +def test_target_tensor_pool_with_fake_adapter(tmp_path): + repo_root = Path(__file__).resolve().parents[1] + source = tmp_path / "target_tensor_pool_test.cc" + binary = tmp_path / ( + "target_tensor_pool_test.exe" if os.name == "nt" else "target_tensor_pool_test" + ) + source.write_text(_TARGET_TENSOR_POOL_TEST_SOURCE) + + _run( + [ + os.environ.get("CXX") or "c++", + "-std=c++17", + "-Wall", + "-Wextra", + "-Werror", + f"-I{repo_root / 'src'}", + str(source), + "-o", + str(binary), + ], + skip_if_missing=True, + ) + _run([str(binary)]) + + +def _run(command, skip_if_missing=False): + try: + subprocess.run(command, check=True, text=True, capture_output=True) + except FileNotFoundError as error: + if skip_if_missing: + pytest.skip(f"`{command[0]}` is not available: {error}") + + raise AssertionError(f"`{command[0]}` was not produced") from error + except subprocess.CalledProcessError as error: + output = "\n".join((error.stdout, error.stderr)).strip() + raise AssertionError(output) from error + + +_TARGET_TENSOR_POOL_TEST_SOURCE = textwrap.dedent( + r""" + #include "target_tensor_pool.h" + + #include + #include + #include + #include + #include + #include + + struct FakeTensor { + explicit FakeTensor(void* data) : data{data} {} + + void* data; + }; + + struct FakeAdapter { + struct State { + State(int id, void* data) : id{id}, data{data} { + ++states_created; + ++live_states; + } + + State(const State&) = delete; + State& operator=(const State&) = delete; + + ~State() { + ++states_destroyed; + --live_states; + } + + int id; + void* data; + }; + + using NativeHandle = State&; + + static std::unique_ptr Create(const FakeTensor& tensor) { + return std::make_unique(next_id++, tensor.data); + } + + static void Rebind(State& state, const FakeTensor& tensor) noexcept { + state.data = tensor.data; + } + + static NativeHandle Native(State& state) { return state; } + + inline static int next_id = 1; + inline static int states_created = 0; + inline static int states_destroyed = 0; + inline static int live_states = 0; + }; + + using Pool = infini::ops::TargetTensorPool; + using Lease = Pool::Lease; + + static_assert(std::is_default_constructible_v); + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(std::is_nothrow_move_constructible_v); + static_assert(std::is_nothrow_move_assignable_v); + static_assert(std::is_nothrow_destructible_v); + static_assert( + std::is_same_v().Native()), + FakeAdapter::State&>); + static_assert(!std::is_copy_constructible_v); + static_assert(!std::is_copy_assignable_v); + static_assert(std::is_nothrow_move_constructible_v); + static_assert(!std::is_move_assignable_v); + static_assert(noexcept(FakeAdapter::Rebind( + std::declval(), + std::declval()))); + + void TestSequentialReuseAndRebind() { + Pool pool; + int first_data = 1; + int second_data = 2; + int first_id = 0; + const int created_before = FakeAdapter::states_created; + + { + auto first = pool.Acquire(FakeTensor{&first_data}); + first_id = first.Native().id; + assert(first.Native().data == &first_data); + } + + auto reused = pool.Acquire(FakeTensor{&second_data}); + assert(reused.Native().id == first_id); + assert(reused.Native().data == &second_data); + assert(FakeAdapter::states_created == created_before + 1); + } + + void TestNestedAcquireAndStableGrowth() { + Pool pool; + int outer_data = 1; + int nested_data[64] = {}; + int extra_data = 2; + auto outer = pool.Acquire(FakeTensor{&outer_data}); + auto* outer_state = &outer.Native(); + const int outer_id = outer.Native().id; + std::vector nested; + + for (std::size_t i = 0; i < 64; ++i) { + nested.push_back(pool.Acquire(FakeTensor{&nested_data[i]})); + assert(nested.back().Native().id != outer_id); + assert(&outer.Native() == outer_state); + assert(outer.Native().data == &outer_data); + } + + const int created_before_extra = FakeAdapter::states_created; + auto extra = pool.Acquire(FakeTensor{&extra_data}); + assert(FakeAdapter::states_created == created_before_extra + 1); + assert(&outer.Native() == outer_state); + assert(extra.Native().id != outer_id); + } + + void TestLeaseMovesReleaseExactlyOneEntry() { + Pool pool; + int source_data = 1; + int destination_data = 2; + int reused_data = 3; + int nested_data = 4; + int released_data = 5; + auto source = pool.Acquire(FakeTensor{&source_data}); + const int source_id = source.Native().id; + auto destination = pool.Acquire(FakeTensor{&destination_data}); + const int destination_id = destination.Native().id; + + destination = std::move(source); + assert(destination.Native().id == source_id); + + auto reused_destination = pool.Acquire(FakeTensor{&reused_data}); + assert(reused_destination.Native().id == destination_id); + + Lease moved{std::move(destination)}; + auto nested = pool.Acquire(FakeTensor{&nested_data}); + assert(nested.Native().id != source_id); + moved = {}; + + auto reused_source = pool.Acquire(FakeTensor{&released_data}); + assert(reused_source.Native().id == source_id); + assert(reused_source.Native().data == &released_data); + } + + void TestPoolVectorReallocationPreservesLiveEntries() { + int outer_data = 1; + int nested_data = 2; + std::vector pools; + pools.emplace_back(); + auto outer = pools.front().Acquire(FakeTensor{&outer_data}); + auto* outer_state = &outer.Native(); + const int outer_id = outer.Native().id; + + for (int i = 0; i < 64; ++i) { + pools.emplace_back(); + } + + assert(&outer.Native() == outer_state); + assert(outer.Native().data == &outer_data); + auto nested = pools.front().Acquire(FakeTensor{&nested_data}); + assert(nested.Native().id != outer_id); + } + + int main() { + assert(FakeAdapter::live_states == 0); + + TestSequentialReuseAndRebind(); + TestNestedAcquireAndStableGrowth(); + TestLeaseMovesReleaseExactlyOneEntry(); + TestPoolVectorReallocationPreservesLiveEntries(); + + assert(FakeAdapter::live_states == 0); + assert(FakeAdapter::states_created == FakeAdapter::states_destroyed); + + return 0; + } + """ +).lstrip()