perf: reuse ATen tensors and optimize PyTorch fallback wrappers

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=<path>` | InfiniRT install prefix containing `include/` and `lib/` | `$INFINI_RT_ROOT` |
 

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<void*>({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<kDev>(\n"
-            f"      {data_expr}, {api_name}_shape_, {api_name}_strides_,\n"
-            f"      {api_name}_type_, device_index);"
+            f"  auto at_{param.name} = AtenFromTensor<kDev>({api_name});"
         )
 
     for schema_index, param in enumerate(op.params):

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)
             $<TARGET_PROPERTY:infiniops,INCLUDE_DIRECTORIES>
             ${INFINI_RT_INCLUDE_DIRS}
             ${TORCH_INCLUDE_DIRS}
+            ${Python_INCLUDE_DIRS}
             ${PROJECT_SOURCE_DIR}/generated)
         target_compile_definitions(infini_ops_torch_objs PRIVATE
             $<TARGET_PROPERTY:infiniops,COMPILE_DEFINITIONS>)
         target_compile_options(infini_ops_torch_objs PRIVATE
             $<TARGET_PROPERTY:infiniops,COMPILE_OPTIONS>
-            -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
             $<TARGET_OBJECTS:infini_ops_torch_objs>)
+
+        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
+            $<TARGET_PROPERTY:infiniops,INCLUDE_DIRECTORIES>
+            ${INFINI_RT_INCLUDE_DIRS}
+            ${TORCH_INCLUDE_DIRS}
+            ${Python_INCLUDE_DIRS}
+            ${PROJECT_SOURCE_DIR}/generated)
+        target_compile_definitions(infini_ops_torch_bridge_obj PRIVATE
+            $<TARGET_PROPERTY:infiniops,COMPILE_DEFINITIONS>)
+        target_compile_options(infini_ops_torch_bridge_obj PRIVATE
+            $<TARGET_PROPERTY:infiniops,COMPILE_OPTIONS>)
+        target_sources(infiniops PRIVATE
+            $<TARGET_OBJECTS:infini_ops_torch_bridge_obj>)
     endif()
 endif()
 

src/pybind11_utils.h

@@ -5,9 +5,11 @@
 #include <pybind11/stl.h>
 
 #include <algorithm>
+#include <memory>
 
 #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<void*>(obj.attr("data_ptr")().cast<std::uintptr_t>())};
-
-  auto shape{obj.attr("shape").cast<typename Tensor::Shape>()};
-
-  auto dtype_str{py::str(obj.attr("dtype")).cast<std::string>()};
-  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<typename Tensor::Strides>()};
-
-  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<Tensor> OptionalTensorFromPybind11Handle(

src/tensor.h

@@ -3,10 +3,50 @@
 
 #include <infini/rt.h>
 
+#include <memory>
+
 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<const void>`) so the core Tensor has no torch dependency.
+  void set_source_handle(std::shared_ptr<const void> handle) {
+    source_handle_ = std::move(handle);
+  }
+
+  const std::shared_ptr<const void>& source_handle() const {
+    return source_handle_;
+  }
+
+ private:
+  std::shared_ptr<const void> source_handle_;
+};
 
 }  // namespace infini::ops
 
+template <>
+struct std::hash<infini::ops::Tensor> {
+  std::size_t operator()(const infini::ops::Tensor& tensor) const {
+    return std::hash<infini::rt::TensorView>{}(
+        static_cast<const infini::rt::TensorView&>(tensor));
+  }
+};
+
+template <>
+struct std::equal_to<infini::ops::Tensor> {
+  bool operator()(const infini::ops::Tensor& a,
+                  const infini::ops::Tensor& b) const {
+    return std::equal_to<infini::rt::TensorView>{}(
+        static_cast<const infini::rt::TensorView&>(a),
+        static_cast<const infini::rt::TensorView&>(b));
+  }
+};
+
 #endif

src/torch/pybind11_.cc

@@ -0,0 +1,79 @@
+#include "torch/pybind11_.h"
+
+#include <torch/version.h>
+
+#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<at::Tensor>(
+      THPVariable_Unpack(reinterpret_cast<PyObject*>(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<const void>(holder),
+  };
+}
+
+}  // namespace infini::ops

src/torch/pybind11_.h

@@ -0,0 +1,28 @@
+#ifndef INFINI_OPS_TORCH_PYBIND11__H_
+#define INFINI_OPS_TORCH_PYBIND11__H_
+
+#include <memory>
+#include <string>
+
+#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<const void> source_handle;
+};
+
+AtenTensorMetadata AtenTensorMetadataFromPyObject(void* py_object);
+
+}  // namespace infini::ops
+
+#endif

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 <Device::Type kDev>
+inline at::Tensor AtenFromTensor(const Tensor& t) {
+  if (t.source_handle()) {
+    return *std::static_pointer_cast<at::Tensor>(
+        std::const_pointer_cast<void>(t.source_handle()));
+  }
+
+  return ToAtenTensor<kDev>(const_cast<void*>(t.data()), t.shape(), t.strides(),
+                            t.dtype(), t.device().index());
+}
+
 }  // namespace infini::ops
 
 #endif