diff --git a/README.md b/README.md index 892503ab6..40677d204 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ InfiniOps extension so `import infini.ops` can load its runtime dependency. | `-DWITH_CAMBRICON=[ON\|OFF]` | Compile the Cambricon implementation | OFF | | `-DWITH_ASCEND=[ON\|OFF]` | Compile the Ascend implementation | OFF | | `-DWITH_TORCH=[ON\|OFF]` | Compile generated PyTorch ATen-backed operators | OFF | +| `-DINFINI_OPS_TORCH_OPT_LEVEL=[0\|1\|2\|3\|s]` | Set optimization level for generated PyTorch wrappers | 3 | | `-DAUTO_DETECT_DEVICES=[ON\|OFF]` | Auto-detect available platforms | ON | | `-DINFINI_RT_ROOT=` | InfiniRT install prefix containing `include/` and `lib/` | `$INFINI_RT_ROOT` | diff --git a/scripts/generate_torch_ops.py b/scripts/generate_torch_ops.py index 9e45a5f4e..f43c6e732 100644 --- a/scripts/generate_torch_ops.py +++ b/scripts/generate_torch_ops.py @@ -1257,15 +1257,10 @@ def _append_optional_conversion(schema_param: Param, api_param: Param) -> None: conversion_lines.append(" }") continue - data_expr = ( - f"{api_name}.data()" - if param.is_mutable_tensor - else f"const_cast({api_name}.data())" - ) + # Reuse the original `at::Tensor` carried by the InfiniOps Tensor when + # present (zero-copy); fall back to `from_blob` only otherwise. 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} = AtenFromTensor({api_name});" ) for schema_index, param in enumerate(op.params): diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index fb022422c..894714065 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -442,7 +442,7 @@ if(WITH_TORCH) # Auto-generate ATen-backed operator wrappers from `scripts/torch_ops.yaml`. # The script writes into `${PROJECT_SOURCE_DIR}/generated/` (gitignored), # which we then glob below alongside any hand-written torch sources. - find_package(Python COMPONENTS Interpreter REQUIRED) + find_package(Python COMPONENTS Interpreter Development REQUIRED) # Pin codegen to the locally installed torch version so vendor # forks (Cambricon's `torch_mlu` 2.1.0, etc.) get a schema whose @@ -489,6 +489,25 @@ if(WITH_TORCH) "${PROJECT_SOURCE_DIR}/generated/torch/*.cpp" ) + # Tensor extraction runs once per input on every Python binding call. Keep + # this small bridge out of the -O0 generated-wrapper target while still + # compiling it with a host compiler on vendor platforms. + set(TORCH_PYBIND_BRIDGE_SOURCE + "${CMAKE_CURRENT_SOURCE_DIR}/torch/pybind11_.cc") + list(REMOVE_ITEM TORCH_SOURCES "${TORCH_PYBIND_BRIDGE_SOURCE}") + + set(INFINI_OPS_TORCH_OPT_LEVEL "3" CACHE STRING + "Optimization level for generated PyTorch wrapper sources") + set_property(CACHE INFINI_OPS_TORCH_OPT_LEVEL + PROPERTY STRINGS 0 1 2 3 s) + if(NOT INFINI_OPS_TORCH_OPT_LEVEL MATCHES "^(0|1|2|3|s)$") + message(FATAL_ERROR + "INFINI_OPS_TORCH_OPT_LEVEL must be one of 0, 1, 2, 3, or s") + endif() + set(_torch_optimization_flag "-O${INFINI_OPS_TORCH_OPT_LEVEL}") + message(STATUS + "Torch wrapper optimization: ${_torch_optimization_flag}") + 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}) @@ -567,7 +586,7 @@ if(WITH_TORCH) endif() set(_torch_include_flags "") - foreach(_dir ${TORCH_INCLUDE_DIRS}) + foreach(_dir ${TORCH_INCLUDE_DIRS} ${Python_INCLUDE_DIRS}) list(APPEND _torch_include_flags "-isystem" "${_dir}") endforeach() @@ -608,7 +627,7 @@ if(WITH_TORCH) add_custom_command( OUTPUT "${_obj}" COMMAND ${SYSTEM_CXX} - -std=c++17 -fPIC -O0 + -std=c++17 -fPIC ${_torch_optimization_flag} "-I${CMAKE_CURRENT_SOURCE_DIR}" "-I${PROJECT_SOURCE_DIR}/generated" ${INFINI_RT_INCLUDE_FLAGS} @@ -623,6 +642,24 @@ if(WITH_TORCH) list(APPEND TORCH_OBJECT_FILES "${_obj}") endforeach() + set(_torch_bridge_obj "${TORCH_OBJECT_DIR}/torch_pybind11_.o") + 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_include_flags} + ${_torch_extra_flags} + -c "${TORCH_PYBIND_BRIDGE_SOURCE}" -o "${_torch_bridge_obj}" + DEPENDS "${TORCH_PYBIND_BRIDGE_SOURCE}" + COMMENT "Compiling torch/pybind11_.cc with system C++ compiler" + JOB_POOL torch_compile + ) + list(APPEND TORCH_OBJECT_FILES "${_torch_bridge_obj}") + set_source_files_properties(${TORCH_OBJECT_FILES} PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE) target_sources(infiniops PRIVATE ${TORCH_OBJECT_FILES}) @@ -640,18 +677,36 @@ if(WITH_TORCH) $ ${INFINI_RT_INCLUDE_DIRS} ${TORCH_INCLUDE_DIRS} + ${Python_INCLUDE_DIRS} ${PROJECT_SOURCE_DIR}/generated) target_compile_definitions(infini_ops_torch_objs PRIVATE $) target_compile_options(infini_ops_torch_objs PRIVATE $ - -O0) + ${_torch_optimization_flag}) if(CMAKE_GENERATOR MATCHES "Ninja") set_target_properties(infini_ops_torch_objs PROPERTIES JOB_POOL_COMPILE torch_compile) endif() target_sources(infiniops PRIVATE $) + + 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 + $) + target_sources(infiniops PRIVATE + $) endif() endif() diff --git a/src/pybind11_utils.h b/src/pybind11_utils.h index 1571e2981..4e227e922 100644 --- a/src/pybind11_utils.h +++ b/src/pybind11_utils.h @@ -5,9 +5,11 @@ #include #include +#include #include "tensor.h" #include "torch/device_.h" +#include "torch/pybind11_.h" namespace py = pybind11; @@ -119,21 +121,13 @@ inline Device DeviceFromPybind11Handle(py::handle obj) { } inline Tensor TensorFromPybind11Handle(py::handle obj) { - auto data{ - reinterpret_cast(obj.attr("data_ptr")().cast())}; - - auto shape{obj.attr("shape").cast()}; - - auto dtype_str{py::str(obj.attr("dtype")).cast()}; - auto pos{dtype_str.find_last_of('.')}; - auto dtype{DataTypeFromString( - pos == std::string::npos ? dtype_str : dtype_str.substr(pos + 1))}; - - auto device{DeviceFromPybind11Handle(obj)}; - - auto strides{obj.attr("stride")().cast()}; - - return Tensor{data, std::move(shape), dtype, device, std::move(strides)}; + auto metadata{AtenTensorMetadataFromPyObject(obj.ptr())}; + Device device{DeviceTypeFromString(metadata.device_type), + metadata.device_index}; + Tensor tensor{metadata.data, std::move(metadata.shape), metadata.dtype, + device, std::move(metadata.strides)}; + tensor.set_source_handle(std::move(metadata.source_handle)); + return tensor; } inline std::optional OptionalTensorFromPybind11Handle( diff --git a/src/tensor.h b/src/tensor.h index 48dc21d54..ee363a7f0 100644 --- a/src/tensor.h +++ b/src/tensor.h @@ -3,10 +3,50 @@ #include +#include + namespace infini::ops { -using Tensor = infini::rt::TensorView; +class Tensor : public infini::rt::TensorView { + public: + using infini::rt::TensorView::TensorView; + + Tensor(const infini::rt::TensorView& tensor) + : infini::rt::TensorView(tensor) {} + + // Opaque handle to the original backend tensor, attached by the pybind layer + // so the PyTorch fallback can reuse its existing tensor wrapper. Type-erased + // (`shared_ptr`) so the core Tensor has no torch dependency. + void set_source_handle(std::shared_ptr handle) { + source_handle_ = std::move(handle); + } + + const std::shared_ptr& source_handle() const { + return source_handle_; + } + + private: + std::shared_ptr source_handle_; +}; } // namespace infini::ops +template <> +struct std::hash { + std::size_t operator()(const infini::ops::Tensor& tensor) const { + return std::hash{}( + static_cast(tensor)); + } +}; + +template <> +struct std::equal_to { + bool operator()(const infini::ops::Tensor& a, + const infini::ops::Tensor& b) const { + return std::equal_to{}( + static_cast(a), + static_cast(b)); + } +}; + #endif diff --git a/src/torch/pybind11_.cc b/src/torch/pybind11_.cc new file mode 100644 index 000000000..a4e10fc86 --- /dev/null +++ b/src/torch/pybind11_.cc @@ -0,0 +1,79 @@ +#include "torch/pybind11_.h" + +#include + +#include "torch/csrc/autograd/python_variable.h" + +namespace infini::ops { + +namespace { + +DataType DataTypeFromAten(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::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: + assert(false && "unsupported at::ScalarType for InfiniOps conversion"); + return DataType::kFloat32; + } +} + +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 + +AtenTensorMetadata AtenTensorMetadataFromPyObject(void* py_object) { + auto holder{std::make_shared( + THPVariable_Unpack(reinterpret_cast(py_object)))}; + const at::Tensor& tensor{*holder}; + const c10::Device device{tensor.device()}; + + return AtenTensorMetadata{ + tensor.data_ptr(), + Tensor::Shape(tensor.sizes().begin(), tensor.sizes().end()), + Tensor::Strides(tensor.strides().begin(), tensor.strides().end()), + DataTypeFromAten(tensor.scalar_type()), + DeviceTypeFromAten(device), + device.index(), + std::static_pointer_cast(holder), + }; +} + +} // namespace infini::ops diff --git a/src/torch/pybind11_.h b/src/torch/pybind11_.h new file mode 100644 index 000000000..acfdf9364 --- /dev/null +++ b/src/torch/pybind11_.h @@ -0,0 +1,28 @@ +#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 bridge. Keeping +// ATen headers out of generated bindings lets vendor compilers build them. +struct AtenTensorMetadata { + void* data; + Tensor::Shape shape; + Tensor::Strides strides; + DataType dtype; + std::string device_type; + int device_index; + std::shared_ptr source_handle; +}; + +AtenTensorMetadata AtenTensorMetadataFromPyObject(void* py_object); + +} // namespace infini::ops + +#endif diff --git a/src/torch/tensor_.h b/src/torch/tensor_.h index 5780babd9..324e58aa6 100644 --- a/src/torch/tensor_.h +++ b/src/torch/tensor_.h @@ -100,6 +100,20 @@ inline at::Tensor ToAtenTensor(void* data, const Tensor::Shape& shape, return at::from_blob(data, at_shape, at_strides, options); } +// Convert an InfiniOps `Tensor` to an `at::Tensor`. Reuse the original +// `at::Tensor` when available to avoid a `from_blob` allocation; otherwise +// reconstruct it from the stored tensor metadata. +template +inline at::Tensor AtenFromTensor(const Tensor& t) { + if (t.source_handle()) { + return *std::static_pointer_cast( + std::const_pointer_cast(t.source_handle())); + } + + return ToAtenTensor(const_cast(t.data()), t.shape(), t.strides(), + t.dtype(), t.device().index()); +} + } // namespace infini::ops #endif diff --git a/tests/test_generate_torch_ops.py b/tests/test_generate_torch_ops.py index 165511180..66b0f9619 100644 --- a/tests/test_generate_torch_ops.py +++ b/tests/test_generate_torch_ops.py @@ -49,11 +49,35 @@ 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 at_self = AtenFromTensor(input);" in source + assert "input_shape_" not in source assert "at::_softmax_out(at_out, at_self" in source +def test_pybind_tensor_bridge_keeps_aten_headers_out_of_generated_bindings(): + root = pathlib.Path(__file__).resolve().parent.parent + pybind_utils = (root / "src" / "pybind11_utils.h").read_text() + bridge_header = (root / "src" / "torch" / "pybind11_.h").read_text() + bridge_source = (root / "src" / "torch" / "pybind11_.cc").read_text() + + assert '"torch/pybind11_.h"' in pybind_utils + assert "torch/csrc/" not in pybind_utils + assert '"torch/tensor_.h"' not in pybind_utils + assert "AtenTensorMetadataFromPyObject" 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_Unpack" in bridge_source + + +def test_pybind_tensor_bridge_uses_an_optimized_host_compiler_target(): + root = pathlib.Path(__file__).resolve().parent.parent + cmake = (root / "src" / "CMakeLists.txt").read_text() + + 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 + + def test_optional_tensor_params_are_exposed_and_forwarded_to_aten(): module = _load_generator_module() op = module._parse_func(