perf(torch): reuse ATen wrappers in generated operators

scripts/generate_torch_ops.py

@@ -1168,38 +1168,60 @@ 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<AtenTensorAdapter<kDev>> "
+                    f"{api_name}_pool_;"
+                )
+            else:
+                pool_members.append(
+                    "  mutable std::vector<"
+                    "TargetTensorPool<AtenTensorAdapter<kDev>>> "
+                    f"{api_name}_pools_;"
+                )
 
     return _TORCH_HEADER_TEMPLATE.format(
         name_uc=name.upper(),
         name=name,
         pascal=pascal,
         op_type=op_type,
         op_calls=op_calls,
+        pool_members="\n\n".join(pool_members),
         slot=_PYTORCH_SLOT,
     )
 
 
 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("{}")
 
     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<void*>({api_name}->data())"
-
-            return (
-                f"ToAtenTensor<kDev>({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<typename "
+                "TargetTensorPool<AtenTensorAdapter<kDev>>::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<typename "
+                "TargetTensorPool<AtenTensorAdapter<kDev>>::Lease> "
+                f"{param.name}_leases;"
+            )
+            conversion_lines.append(
+                f"  {param.name}_leases.reserve({api_name}.size());"
+            )
             conversion_lines.append(f"  std::vector<at::Tensor> 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<kDev>("
-                "const_cast<void*>(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<void*>({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<kDev>(\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_<DEV>` 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

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)
         "$<$<COMPILE_LANGUAGE:CXX,CUDA>:-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
+                $<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>
+                -O2)
+            target_sources(ops PRIVATE
+                $<TARGET_OBJECTS:infini_ops_torch_bridge_obj>)
+        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

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<void*>(obj.attr("data_ptr")().cast<std::uintptr_t>())};