feat: add PyTorch C++ backend for Add and Gemm (#51)

* refactor: extract `TorchDeviceName` into `src/torch/device_.h`

Move the `TorchDeviceName<kDev>` template specializations from
`pybind11_utils.h` into a standalone header so they can be reused
by torch operator implementations without pulling in pybind11.

* feat: add slotted `ActiveImplementationsImpl` for composable registries

Add a third template parameter `N` (slot index, default 0) to
`ActiveImplementationsImpl`. Slot 0 holds the base/device-native
indices, and higher slots let add-on backends register extra indices
without conflicting with existing specializations. `ActiveImplementations`
flattens slots 0-3 via `Flatten`.

* feat: add `WITH_TORCH` CMake option and wrapper generator support

Add `WITH_TORCH` option that finds PyTorch via pip, links against
libtorch, and compiles sources under `src/torch/`. Pass `--with-torch`
to `generate_wrappers.py` so it scans `src/torch/` for operator
specializations.

* feat: add torch tensor conversion utilities

Add `ToAtenDtype()` and `ToAtenTensor<kDev>()` in `src/torch/tensor_.h`
for zero-copy conversion from `infini::ops::Tensor` to `at::Tensor`
via `at::from_blob()`.

* feat: add torch `Add` operator (implementation index 1)

Register torch `Add` via `ActiveImplementationsImpl` slot 1 for all
devices. The implementation uses `at::add_out()` through ATen's
device-generic dispatch.

* feat: add torch `Gemm` operator (implementation index 2)

Register torch `Gemm` via `ActiveImplementationsImpl` slot 1 for all
devices. The implementation uses `at::addmm_out()` / `at::baddbmm_out()`
through ATen's device-generic dispatch.

* feat: add `AUTO_DETECT_BACKENDS` CMake option

Auto-detect PyTorch by attempting `import torch`. When found,
`WITH_TORCH` is enabled automatically.

* fix: replace `find_package(Torch)` with direct path queries

`find_package(Torch)` pulls in Caffe2's cmake config, which calls
`enable_language(CUDA)` and breaks on platforms with non-standard
CUDA toolchains (e.g. Iluvatar). Query include and library paths
directly via `torch.utils.cpp_extension` instead.

* fix(iluvatar): set `CMAKE_CUDA_ARCHITECTURES` before `enable_language(CUDA)`

CMake 4.3+ requires `CMAKE_CUDA_ARCHITECTURES` to be set before
`enable_language(CUDA)` when using non-standard CUDA compilers like
Iluvatar's `clang++`. Without it, CMake fails to detect a default
architecture.

* fix: auto-detect `pybind11_DIR` from Python when not set

When `pybind11` is installed via pip but not in a standard CMake search
path, `find_package(pybind11 CONFIG)` fails. Query `python -m pybind11
--cmakedir` as a fallback to locate the package.

* fix: cross-platform torch compilation and version compatibility

- Iluvatar: set `CMAKE_CUDA_ARCHITECTURES` to `OFF` instead of
  `ivcore20` (CMake 4.3 rejects non-integer architecture names; the
  architecture is already passed via `CMAKE_CUDA_FLAGS`).

- MetaX/Moore: split torch operator headers into declaration-only `.h`
  files and `.cc` implementation files with explicit instantiations.
  Compile the `.cc` files with the system `g++` instead of the vendor
  compiler (`mxcc`/`mcc`), which cannot parse vendor-forked torch
  headers in C++ extension mode.

- Cambricon: guard `UInt16`/`UInt32`/`UInt64` scalar types in
  `ToAtenDtype()` with a `TORCH_VERSION` check (these types require
  PyTorch 2.4+; Cambricon ships torch 2.1).

- Wrapper generator: scan only `.h` files to avoid including `.cc`
  explicit-instantiation files in the generated `ops.cc`.

* fix: resolve build failures on iluvatar, metax, and cambricon

- Iluvatar: move `-x ivcore` from CMAKE_CUDA_FLAGS to compile-only
  options so it doesn't get passed during linking (which caused
  clang++ to re-parse .o files as source code)
- MetaX: add `-DUSE_MACA=1` to g++ flags for torch source
  compilation (MetaX torch fork headers require this define)
- Cambricon: query `torch.compiled_with_cxx11_abi()` and set
  `_GLIBCXX_USE_CXX11_ABI` globally to match torch's ABI setting
  (fixes undefined reference to `c10::Device::Device(std::string)`)

* style: fix comments to comply with `CONTRIBUTING.md`

Use backtick-fenced Markdown syntax for identifiers in comments
and error messages, and ensure comments are complete sentences.

* style: apply `clang-format`

* refactor: auto-detect operator implementations via SFINAE

Replace the manual `ActiveImplementationsImpl` slot system with
`std::is_base_of`-based compile-time detection. A real `Operator`
specialization inherits from `Key` (e.g., `Gemm`), while the primary
template inherits only from `OperatorBase` — SFINAE distinguishes the
two automatically, eliminating the need for `registry.h` files.

* refactor: use `std::index_sequence` for implementation auto-detection

Replace the hand-unrolled `Flatten<..., 0>::type, ..., 3>::type>` with
`std::index_sequence<0..kMaxImplementations>` expansion. Increase
`kMaxImplementations` from 4 to 16.

* refactor: replace preprocessor conditionals with `constexpr` in `tensor_.h`

Use `constexpr int kTorchVersion` and `if constexpr` instead of `#if`
macros for PyTorch version checks. Extract unsigned dtype handling into
`detail::ToAtenUnsignedDataType`.

* docs: add TODO comments for dynamic implementation index parametrization

* fix: use dependent type alias to support `if constexpr` on PyTorch < 2.4

`c10::ScalarType::UInt16` is a non-dependent name resolved at template
definition time.  Introduce `DependentScalarType<kVersion>::type` so
the enum member access becomes dependent and is properly discarded by
`if constexpr` on older PyTorch versions.

* fix(tests): skip torch Gemm on CPU half-precision

ATen `addmm`/`baddbmm` does not support `float16`/`bfloat16` on CPU.

* refactor: rename `ToAtenDtype` to `ToAtenDataType`

Author: voltjia

CMakeLists.txt

@@ -16,7 +16,10 @@ option(WITH_CAMBRICON "Enable Cambricon backend" OFF)
 option(WITH_MOORE "Enable Moore backend" OFF)
 option(WITH_ASCEND "Enable Ascend backend" OFF)
 
+option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
+
 option(AUTO_DETECT_DEVICES "Automatically detect available devices" OFF)
+option(AUTO_DETECT_BACKENDS "Automatically detect available backends" OFF)
 option(GENERATE_PYTHON_BINDINGS "Generate Python bindings" OFF)
 
 if(AUTO_DETECT_DEVICES)
@@ -79,6 +82,72 @@ if(AUTO_DETECT_DEVICES)
     endif()
 endif()
 
+if(AUTO_DETECT_BACKENDS)
+    message(STATUS "Auto-detecting available backends...")
+
+    find_package(Python COMPONENTS Interpreter QUIET)
+
+    if(Python_FOUND)
+        execute_process(
+            COMMAND ${Python_EXECUTABLE} -c "import torch"
+            RESULT_VARIABLE _torch_import_result
+            OUTPUT_QUIET
+            ERROR_QUIET
+        )
+
+        if(_torch_import_result EQUAL 0)
+            set(WITH_TORCH ON)
+            message(STATUS "Auto-detected PyTorch.")
+        endif()
+    endif()
+endif()
+
+if(WITH_TORCH)
+    find_package(Python COMPONENTS Interpreter REQUIRED)
+
+    # Query `torch` paths directly instead of using `find_package(Torch)`,
+    # which pulls in Caffe2's CMake config and may fail on platforms with
+    # non-standard CUDA toolchains.
+    execute_process(
+        COMMAND ${Python_EXECUTABLE} -c "from torch.utils.cpp_extension import include_paths; print(';'.join(include_paths()))"
+        OUTPUT_VARIABLE TORCH_INCLUDE_DIRS
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        RESULT_VARIABLE _torch_result
+    )
+
+    if(NOT _torch_result EQUAL 0)
+        message(FATAL_ERROR "`WITH_TORCH` is `ON` but `torch` is not installed.")
+    endif()
+
+    execute_process(
+        COMMAND ${Python_EXECUTABLE} -c "from torch.utils.cpp_extension import library_paths; print(';'.join(library_paths()))"
+        OUTPUT_VARIABLE _torch_lib_dirs
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+    )
+
+    find_library(TORCH_LIB torch HINTS ${_torch_lib_dirs} REQUIRED)
+    find_library(TORCH_CPU_LIB torch_cpu HINTS ${_torch_lib_dirs} REQUIRED)
+    find_library(C10_LIB c10 HINTS ${_torch_lib_dirs} REQUIRED)
+    set(TORCH_LIBRARIES ${TORCH_LIB} ${TORCH_CPU_LIB} ${C10_LIB})
+
+    # Query the `CXX11` ABI setting that `torch` was compiled with.
+    # A mismatch causes linker errors (e.g. undefined reference to
+    # `c10::Device::Device(std::string const&)`).
+    execute_process(
+        COMMAND ${Python_EXECUTABLE} -c "import torch; print(int(torch.compiled_with_cxx11_abi()))"
+        OUTPUT_VARIABLE TORCH_CXX11_ABI
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        RESULT_VARIABLE _torch_abi_result
+    )
+
+    if(_torch_abi_result EQUAL 0)
+        add_compile_definitions(_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI})
+        message(STATUS "PyTorch CXX11 ABI: ${TORCH_CXX11_ABI}")
+    endif()
+
+    message(STATUS "Found PyTorch: ${TORCH_INCLUDE_DIRS}")
+endif()
+
 include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
 
 # Only one CUDA-like GPU backend can be enabled at a time.
@@ -110,11 +179,23 @@ if(WITH_ILUVATAR)
     else()
         set(CMAKE_CUDA_COMPILER "clang++" CACHE STRING "Iluvatar CUDA compiler (clang++)")
     endif()
-    set(CMAKE_CUDA_FLAGS "-x ivcore -std=c++17 --cuda-gpu-arch=${ILUVATAR_ARCH} -fPIC -Wno-error=unused-variable -Wno-error=unused-private-field -Wno-unused-variable" CACHE STRING "Iluvatar CUDA flags")
+    # `-x ivcore` must not be in `CMAKE_CUDA_FLAGS` — CMake passes those flags
+    # to both compile and link steps. During linking, `-x ivcore` causes
+    # `clang++` to re-parse `.o` files as source code.
+    set(CMAKE_CUDA_FLAGS "--cuda-gpu-arch=${ILUVATAR_ARCH} -fPIC -Wno-error=unused-variable -Wno-error=unused-private-field -Wno-unused-variable" CACHE STRING "Iluvatar CUDA flags")
     set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF CACHE BOOL "Disable RDC for Iluvatar")
+    set(CMAKE_CUDA_ARCHITECTURES OFF CACHE STRING "Iluvatar CUDA architectures (passed via CMAKE_CUDA_FLAGS)")
+    # Iluvatar does not ship `libcudadevrt`, which CMake's compiler test
+    # tries to link. Disable automatic CUDA runtime linking and link
+    # manually via `find_package(CUDAToolkit)` instead.
+    set(CMAKE_CUDA_RUNTIME_LIBRARY NONE)
     message(STATUS "Iluvatar: CUDA compiler ${CMAKE_CUDA_COMPILER}, arch ${ILUVATAR_ARCH}")
     enable_language(CUDA)
     find_package(CUDAToolkit REQUIRED)
+    # Add `-x ivcore` as a compile-only flag so it is not passed during
+    # linking, where it would cause `clang++` to re-parse `.o` files as
+    # source.
+    add_compile_options($<$<COMPILE_LANGUAGE:CUDA>:-x$<SEMICOLON>ivcore>)
 endif()
 
 if(WITH_METAX)

pyproject.toml

@@ -14,6 +14,7 @@ install-dir = "infini"
 
 [tool.scikit-build.cmake.define]
 AUTO_DETECT_DEVICES = "ON"
+AUTO_DETECT_BACKENDS = "ON"
 GENERATE_PYTHON_BINDINGS = "ON"
 
 [tool.pytest.ini_options]

scripts/generate_wrappers.py

@@ -413,7 +413,12 @@ def _snake_to_pascal(snake_str):
     return "".join(word.capitalize() for word in snake_str.split("_"))
 
 
-def _get_all_ops(devices):
+def _get_all_ops(devices, with_torch=False):
+    scan_dirs = set(devices)
+
+    if with_torch:
+        scan_dirs.add("torch")
+
     ops = {}
 
     for file_path in _BASE_DIR.iterdir():
@@ -424,8 +429,8 @@ def _get_all_ops(devices):
 
         ops[op_name] = []
 
-        for file_path in _SRC_DIR.rglob("*"):
-            if not file_path.is_file() or file_path.parent.parent.name not in devices:
+        for file_path in _SRC_DIR.rglob("*.h"):
+            if file_path.parent.parent.name not in scan_dirs:
                 continue
 
             if f"class Operator<{_snake_to_pascal(op_name)}" in file_path.read_text():
@@ -445,6 +450,12 @@ def _get_all_ops(devices):
         help="Devices to use. Please pick from `cpu`, `nvidia`, `cambricon`, `ascend`, `metax`, `moore`, `iluvatar`, `kunlun`, `hygon`, and `qy`. (default: `cpu`)",
     )
 
+    parser.add_argument(
+        "--with-torch",
+        action="store_true",
+        help="Include PyTorch C++ backend implementations.",
+    )
+
     args = parser.parse_args()
 
     _BINDINGS_DIR.mkdir(parents=True, exist_ok=True)
@@ -456,7 +467,7 @@ def _get_all_ops(devices):
     if ops_json.exists():
         ops = json.loads(ops_json.read_text())
     else:
-        ops = _get_all_ops(args.devices)
+        ops = _get_all_ops(args.devices, with_torch=args.with_torch)
 
     header_paths = []
     bind_func_names = []

src/CMakeLists.txt

@@ -16,7 +16,7 @@ if(WITH_CPU)
 
     target_compile_definitions(infiniops PUBLIC WITH_CPU=1)
 
-    find_package(OpenMP REQUIRED)
+    find_package(OpenMP REQUIRED COMPONENTS CXX)
     target_link_libraries(infiniops PRIVATE OpenMP::OpenMP_CXX)
 
     list(APPEND DEVICE_LIST "cpu")
@@ -218,6 +218,69 @@ if(WITH_ASCEND)
     list(APPEND DEVICE_LIST "ascend")
 endif()
 
+if(WITH_TORCH)
+    file(GLOB_RECURSE TORCH_SOURCES CONFIGURE_DEPENDS "torch/*.cc" "torch/*.cpp")
+
+    target_compile_definitions(infiniops PUBLIC WITH_TORCH=1)
+    target_link_libraries(infiniops PUBLIC ${TORCH_LIBRARIES})
+    target_include_directories(infiniops PUBLIC ${TORCH_INCLUDE_DIRS})
+
+    if(WITH_METAX OR WITH_MOORE)
+        # Vendor compilers (`mxcc`/`mcc`) cannot compile vendor-forked `torch`
+        # headers. Compile `torch` sources with the system C++ compiler instead.
+        find_program(SYSTEM_CXX NAMES g++ c++)
+
+        if(NOT SYSTEM_CXX)
+            message(FATAL_ERROR "Could not find system `g++` for `torch` compilation.")
+        endif()
+
+        set(_torch_include_flags "")
+        foreach(_dir ${TORCH_INCLUDE_DIRS})
+            list(APPEND _torch_include_flags "-isystem" "${_dir}")
+        endforeach()
+
+        # Vendor-specific defines required by forked `torch` headers.
+        set(_torch_extra_flags "")
+        if(WITH_METAX)
+            list(APPEND _torch_extra_flags "-DUSE_MACA=1")
+        endif()
+        if(DEFINED TORCH_CXX11_ABI)
+            list(APPEND _torch_extra_flags "-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI}")
+        endif()
+
+        set(TORCH_OBJECT_DIR "${CMAKE_CURRENT_BINARY_DIR}/torch_objs")
+        file(MAKE_DIRECTORY "${TORCH_OBJECT_DIR}")
+
+        set(TORCH_OBJECT_FILES)
+        foreach(_src ${TORCH_SOURCES})
+            file(RELATIVE_PATH _rel ${CMAKE_CURRENT_SOURCE_DIR} ${_src})
+            string(REPLACE "/" "_" _obj_name "${_rel}")
+            string(REPLACE ".cc" ".o" _obj_name "${_obj_name}")
+            string(REPLACE ".cpp" ".o" _obj_name "${_obj_name}")
+            set(_obj "${TORCH_OBJECT_DIR}/${_obj_name}")
+
+            add_custom_command(
+                OUTPUT "${_obj}"
+                COMMAND ${SYSTEM_CXX}
+                    -std=c++17 -fPIC -O2
+                    "-I${CMAKE_CURRENT_SOURCE_DIR}"
+                    ${_torch_include_flags}
+                    ${_torch_extra_flags}
+                    -c "${_src}" -o "${_obj}"
+                DEPENDS "${_src}"
+                COMMENT "Compiling ${_rel} with system C++ compiler"
+            )
+            list(APPEND TORCH_OBJECT_FILES "${_obj}")
+        endforeach()
+
+        set_source_files_properties(${TORCH_OBJECT_FILES}
+            PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE)
+        target_sources(infiniops PRIVATE ${TORCH_OBJECT_FILES})
+    else()
+        target_sources(infiniops PRIVATE ${TORCH_SOURCES})
+    endif()
+endif()
+
 target_include_directories(infiniops PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
 
 if(GENERATE_PYTHON_BINDINGS)
@@ -226,8 +289,14 @@ if(GENERATE_PYTHON_BINDINGS)
     # active device list.  Stale generated files (e.g., committed for one
     # platform) would omit specializations for other enabled backends,
     # causing link-time or runtime failures.
+
+    set(GENERATOR_ARGS --devices ${DEVICE_LIST})
+    if(WITH_TORCH)
+        list(APPEND GENERATOR_ARGS --with-torch)
+    endif()
+
     execute_process(
-        COMMAND ${Python_EXECUTABLE} ${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py --devices ${DEVICE_LIST}
+        COMMAND ${Python_EXECUTABLE} ${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py ${GENERATOR_ARGS}
         WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
         RESULT_VARIABLE script_result
     )
@@ -246,6 +315,20 @@ if(GENERATE_PYTHON_BINDINGS)
     endif()
 
     find_package(Python COMPONENTS Interpreter Development)
+
+    if(NOT pybind11_DIR)
+        execute_process(
+            COMMAND ${Python_EXECUTABLE} -m pybind11 --cmakedir
+            OUTPUT_VARIABLE _pybind11_cmake_dir
+            OUTPUT_STRIP_TRAILING_WHITESPACE
+            RESULT_VARIABLE _pybind11_result
+        )
+
+        if(_pybind11_result EQUAL 0)
+            set(pybind11_DIR "${_pybind11_cmake_dir}" CACHE PATH "pybind11 CMake directory")
+        endif()
+    endif()
+
     find_package(pybind11 CONFIG)
 
     if(PYBIND11_ENABLE_EXTRAS)

src/nvidia/gemm/cublas.h

@@ -3,7 +3,6 @@
 
 #include "cuda/gemm/blas.h"
 #include "nvidia/blas.h"
-#include "nvidia/gemm/registry.h"
 
 namespace infini::ops {
 

src/nvidia/gemm/cublaslt.h

@@ -10,7 +10,6 @@
 
 #include "base/gemm.h"
 #include "nvidia/blas_utils.h"
-#include "nvidia/gemm/registry.h"
 #include "nvidia/runtime_.h"
 
 namespace infini::ops {

src/nvidia/gemm/registry.h

@@ -84,14 +84,9 @@ struct std::equal_to<infini::ops::detail::CacheKey> {
 
 namespace infini::ops {
 
+// Forward declaration — defined after `Operator` using SFINAE auto-detection.
 template <typename Key, Device::Type kDev>
-struct ActiveImplementationsImpl {
-  using type = List<0>;
-};
-
-template <typename Key, Device::Type kDev>
-using ActiveImplementations =
-    typename ActiveImplementationsImpl<Key, kDev>::type;
+struct ActiveImplementations;
 
 class OperatorBase {
  public:
@@ -153,7 +148,7 @@ class Operator : public OperatorBase {
                 }
               },
               "Operator::make(implementation_index)",
-              ActiveImplementations<Key, kDev>{});
+              typename ActiveImplementations<Key, kDev>::type{});
         },
         "Operator::make");
 
@@ -200,7 +195,8 @@ class Operator : public OperatorBase {
         dev_type,
         [&](auto device_tag) {
           constexpr Device::Type kDev = decltype(device_tag)::value;
-          result = detail::ListToVector(ActiveImplementations<Key, kDev>{});
+          result = detail::ListToVector(
+              typename ActiveImplementations<Key, kDev>::type{});
         },
         "Operator::active_implementation_indices");
     return result;
@@ -227,6 +223,45 @@ class Operator : public OperatorBase {
   static constexpr std::size_t implementation_index_{implementation_index};
 };
 
+// Maximum number of implementation slots per (operator, device) pair.
+// Increase this value when adding operators with more implementations.
+constexpr std::size_t kMaxImplementations = 16;
+
+// SFINAE-based implementation detection. A partial specialization
+// `Operator<Key, kDev, N>` inherits from `Key` (the operator base class),
+// while the unspecialized primary template inherits only from `OperatorBase`.
+// `std::is_base_of` distinguishes the two at compile time, eliminating the
+// need for manual `registry.h` files.
+template <typename Key, Device::Type kDev, std::size_t N,
+          bool = std::is_base_of_v<Key, Operator<Key, kDev, N>>>
+struct ActiveImplementationsImpl {
+  using type = List<>;
+};
+
+template <typename Key, Device::Type kDev, std::size_t N>
+struct ActiveImplementationsImpl<Key, kDev, N, true> {
+  using type = List<N>;
+};
+
+namespace detail {
+
+template <typename Key, Device::Type kDev, typename Seq>
+struct ActiveImplementationsHelper;
+
+template <typename Key, Device::Type kDev, std::size_t... ns>
+struct ActiveImplementationsHelper<Key, kDev, std::index_sequence<ns...>> {
+  using type = typename Flatten<
+      typename ActiveImplementationsImpl<Key, kDev, ns>::type...>::type;
+};
+
+}  // namespace detail
+
+template <typename Key, Device::Type kDev>
+struct ActiveImplementations {
+  using type = typename detail::ActiveImplementationsHelper<
+      Key, kDev, std::make_index_sequence<kMaxImplementations>>::type;
+};
+
 }  // namespace infini::ops
 
 #endif