fix: export operator call instantiations

voltjia · voltjia · commit ca40c6fd5751 · 2026-05-28T17:48:23.000+08:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -34,6 +34,8 @@ option(BUILD_CUSTOM_KERNEL "Build custom AscendC kernel PyTorch extension (requi
 
 option(AUTO_DETECT_DEVICES "Automatically detect available devices" OFF)
 option(AUTO_DETECT_BACKENDS "Automatically detect available backends" OFF)
+option(GENERATE_OPERATOR_CALL_INSTANTIATIONS
+    "Generate explicit operator call instantiations" ON)
 option(GENERATE_PYTHON_BINDINGS "Generate Python bindings" OFF)
 
 set(_DEFAULT_HYGON_DTK_ROOT "/opt/dtk")
diff --git a/include/infini/ops.h b/include/infini/ops.h
@@ -0,0 +1,8 @@
+#ifndef INFINI_OPS_H_
+#define INFINI_OPS_H_
+
+#ifdef __cplusplus
+#include <infini/operator_call_instantiations.h>
+#endif
+
+#endif  // INFINI_OPS_H_
diff --git a/scripts/generate_wrappers.py b/scripts/generate_wrappers.py
@@ -702,10 +702,113 @@ def _generate_generated_dispatch_source(impl_paths, definitions):
 """
 
 
+def _strip_top_level_const(type_spelling):
+    type_spelling = " ".join(type_spelling.split())
+
+    while type_spelling.startswith("const "):
+        type_spelling = type_spelling[len("const ") :]
+
+    return type_spelling
+
+
+def _generate_operator_call_instantiation_entries(operator):
+    def _generate_template_arguments(node):
+        return ", ".join(
+            _strip_top_level_const(arg.type.spelling)
+            for arg in node.get_arguments()
+            if arg.spelling != "stream"
+        )
+
+    def _generate_parameters(node):
+        return ", ".join(
+            f"const {_strip_top_level_const(arg.type.spelling)}& {arg.spelling}"
+            for arg in node.get_arguments()
+            if arg.spelling != "stream"
+        )
+
+    def _append_optional_params(prefix, params):
+        if params:
+            return f"{prefix}, {params}"
+
+        return prefix
+
+    pascal_case_op_name = _snake_to_pascal(operator.name)
+    declarations = []
+    definitions = []
+
+    for call in operator.calls:
+        template_arguments = _generate_template_arguments(call)
+        params = _generate_parameters(call)
+        function_params = _append_optional_params(
+            "const Handle& handle, const Config& config", params
+        )
+        instantiation = (
+            f"Operator<{pascal_case_op_name}>::Call<{template_arguments}>"
+            f"({function_params})"
+        )
+
+        declarations.append(f"extern template auto {instantiation};")
+        definitions.append(f"template auto {instantiation};")
+
+    return declarations, definitions
+
+
+def _generate_operator_call_instantiation_header(op_names, declarations):
+    header_base_includes = "\n".join(
+        f'#include "base/{op_name}.h"' for op_name in op_names
+    )
+
+    return f"""#ifndef INFINI_OPS_OPERATOR_CALL_INSTANTIATIONS_H_
+#define INFINI_OPS_OPERATOR_CALL_INSTANTIATIONS_H_
+
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include "config.h"
+#include "handle.h"
+#include "operator.h"
+
+{header_base_includes}
+
+namespace infini::ops {{
+
+{chr(10).join(declarations)}
+
+}}  // namespace infini::ops
+
+#endif
+"""
+
+
+def _generate_operator_call_instantiation_source(devices, impl_paths, definitions):
+    device_includes = "\n".join(
+        f'#include "{path}"' for path in _device_marker_headers(devices)
+    )
+    impl_includes = "\n".join(
+        f'#include "{_to_include_path(impl_path)}"' for impl_path in impl_paths
+    )
+
+    return f"""#include "infini/operator_call_instantiations.h"
+
+// clang-format off
+{device_includes}
+{impl_includes}
+// clang-format on
+
+namespace infini::ops {{
+
+{chr(10).join(definitions)}
+
+}}  // namespace infini::ops
+"""
+
+
 def _device_marker_headers(devices):
     paths = {
         "cpu": "native/cpu/device_.h",
         "nvidia": "native/cuda/nvidia/device_.h",
+        "hygon": "native/cuda/hygon/device_.h",
         "cambricon": "native/cambricon/device_.h",
         "ascend": "native/ascend/device_.h",
         "metax": "native/cuda/metax/device_.h",
@@ -821,6 +924,10 @@ def _generate_op_artifacts(item):
     dispatch_declarations, dispatch_definitions = _generate_generated_dispatch_entries(
         operator
     )
+    (
+        call_instantiation_declarations,
+        call_instantiation_definitions,
+    ) = _generate_operator_call_instantiation_entries(operator)
 
     return {
         "op_name": op_name,
@@ -832,6 +939,8 @@ def _generate_op_artifacts(item):
         "legacy_c_header": legacy_c_header,
         "dispatch_declarations": dispatch_declarations,
         "dispatch_definitions": dispatch_definitions,
+        "call_instantiation_declarations": call_instantiation_declarations,
+        "call_instantiation_definitions": call_instantiation_definitions,
         "impl_paths": impl_paths,
     }
 
@@ -929,6 +1038,11 @@ def _dispatch_gen_batch_size():
         for artifact in artifacts
         for declaration in artifact["dispatch_declarations"]
     ]
+    call_instantiation_declarations = [
+        declaration
+        for artifact in artifacts
+        for declaration in artifact["call_instantiation_declarations"]
+    ]
     use_monolithic_bindings = _use_monolithic_bindings()
     op_includes = []
 
@@ -958,6 +1072,14 @@ def _dispatch_gen_batch_size():
     )
     (_BINDINGS_DIR / "generated_dispatch.h").write_text(dispatch_header)
 
+    call_instantiation_header = _generate_operator_call_instantiation_header(
+        op_names, call_instantiation_declarations
+    )
+    (_INCLUDE_DIR / "infini").mkdir(exist_ok=True)
+    (_INCLUDE_DIR / "infini" / "operator_call_instantiations.h").write_text(
+        call_instantiation_header
+    )
+
     dispatch_batch_size = _dispatch_gen_batch_size()
 
     for dispatch_batch_index, start in enumerate(
@@ -979,6 +1101,28 @@ def _dispatch_gen_batch_size():
             dispatch_source
         )
 
+    for call_instantiation_batch_index, start in enumerate(
+        range(0, len(artifacts), dispatch_batch_size)
+    ):
+        batch = artifacts[start : start + dispatch_batch_size]
+        impl_paths = list(
+            dict.fromkeys(
+                impl_path for artifact in batch for impl_path in artifact["impl_paths"]
+            )
+        )
+        definitions = [
+            definition
+            for artifact in batch
+            for definition in artifact["call_instantiation_definitions"]
+        ]
+        call_instantiation_source = _generate_operator_call_instantiation_source(
+            args.devices, impl_paths, definitions
+        )
+        (
+            _GENERATED_SRC_DIR
+            / f"operator_call_instantiations_{call_instantiation_batch_index}.cc"
+        ).write_text(call_instantiation_source)
+
     bind_func_calls = "\n".join(
         f"{bind_func_name}(m);" for bind_func_name in bind_func_names
     )
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -1,5 +1,7 @@
 add_library(infiniops SHARED)
 
+include(GNUInstallDirs)
+
 file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc")
 target_sources(infiniops PRIVATE ${BASE_SRCS})
 
@@ -500,14 +502,20 @@ if(WITH_TORCH)
     endif()
 endif()
 
-target_include_directories(infiniops PUBLIC ${CMAKE_CURRENT_SOURCE_DIR})
+target_include_directories(infiniops
+    PUBLIC
+        $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>
+        $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/include>
+        $<BUILD_INTERFACE:${PROJECT_SOURCE_DIR}/generated/include>
+        $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
+)
 
-if(GENERATE_PYTHON_BINDINGS)
+if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
     find_package(Python COMPONENTS Interpreter REQUIRED)
-    # Always regenerate bindings so the included kernel headers match the
-    # 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.
+    # Always regenerate wrappers so emitted call instantiations and bindings
+    # match the active device list.  Stale generated files would omit
+    # specializations for enabled backends, causing link-time or runtime
+    # failures.
 
     set(GENERATOR_ARGS --devices ${DEVICE_LIST})
     if(WITH_TORCH)
@@ -528,7 +536,76 @@ if(GENERATE_PYTHON_BINDINGS)
     else()
         message(STATUS "Generating wrappers - done")
     endif()
+endif()
+
+if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
+    file(GLOB_RECURSE OPERATOR_CALL_INSTANTIATION_SOURCES CONFIGURE_DEPENDS
+        "${PROJECT_SOURCE_DIR}/generated/src/operator_call_instantiations_*.cc")
+
+    if(WITH_NVIDIA OR WITH_HYGON)
+        set_source_files_properties(${OPERATOR_CALL_INSTANTIATION_SOURCES}
+            PROPERTIES LANGUAGE CUDA)
+        target_sources(infiniops PRIVATE ${OPERATOR_CALL_INSTANTIATION_SOURCES})
+    elseif(WITH_ILUVATAR)
+        set(_iluvatar_call_instantiation_include_flags
+            "-I${CMAKE_CURRENT_SOURCE_DIR}"
+            "-I${PROJECT_SOURCE_DIR}"
+            "-I${PROJECT_SOURCE_DIR}/generated"
+            "-I${PROJECT_SOURCE_DIR}/generated/include")
+        foreach(_dir IN LISTS TORCH_INCLUDE_DIRS CUDAToolkit_INCLUDE_DIRS)
+            list(APPEND _iluvatar_call_instantiation_include_flags "-I${_dir}")
+        endforeach()
+
+        set(_iluvatar_call_instantiation_defs -DWITH_ILUVATAR=1)
+        if(WITH_CPU)
+            list(APPEND _iluvatar_call_instantiation_defs -DWITH_CPU=1)
+        endif()
+        if(WITH_TORCH)
+            list(APPEND _iluvatar_call_instantiation_defs -DWITH_TORCH=1)
+        endif()
+        if(DEFINED TORCH_CXX11_ABI)
+            list(APPEND _iluvatar_call_instantiation_defs
+                "-D_GLIBCXX_USE_CXX11_ABI=${TORCH_CXX11_ABI}")
+        endif()
 
+        set(ILUVATAR_CALL_INSTANTIATION_OBJECTS)
+        set(_iluvatar_call_instantiation_object_dir
+            "${CMAKE_CURRENT_BINARY_DIR}/iluvatar_call_instantiation_objs")
+        foreach(_src IN LISTS OPERATOR_CALL_INSTANTIATION_SOURCES)
+            get_filename_component(_name "${_src}" NAME_WE)
+            set(_obj "${_iluvatar_call_instantiation_object_dir}/${_name}.o")
+            set(_dep "${_obj}.d")
+            set(_depfile_arg)
+            if(CMAKE_GENERATOR MATCHES "Ninja")
+                set(_depfile_arg DEPFILE "${_dep}")
+            endif()
+            add_custom_command(
+                OUTPUT "${_obj}"
+                COMMAND ${CMAKE_COMMAND} -E make_directory
+                    "${_iluvatar_call_instantiation_object_dir}"
+                COMMAND ${ILUVATAR_CUDA_COMPILER}
+                    ${_iluvatar_call_instantiation_defs}
+                    ${_iluvatar_call_instantiation_include_flags}
+                    ${ILUVATAR_CUDA_FLAGS}
+                    -MMD -MF "${_dep}"
+                    -c "${_src}" -o "${_obj}"
+                DEPENDS "${_src}"
+                ${_depfile_arg}
+                COMMENT "Compiling ${_name}.cc with CoreX clang++"
+                VERBATIM
+            )
+            list(APPEND ILUVATAR_CALL_INSTANTIATION_OBJECTS "${_obj}")
+        endforeach()
+
+        set_source_files_properties(${ILUVATAR_CALL_INSTANTIATION_OBJECTS}
+            PROPERTIES EXTERNAL_OBJECT TRUE GENERATED TRUE)
+        target_sources(infiniops PRIVATE ${ILUVATAR_CALL_INSTANTIATION_OBJECTS})
+    else()
+        target_sources(infiniops PRIVATE ${OPERATOR_CALL_INSTANTIATION_SOURCES})
+    endif()
+endif()
+
+if(GENERATE_PYTHON_BINDINGS)
     file(GLOB_RECURSE PYBIND11_SOURCES CONFIGURE_DEPENDS
         "${PROJECT_SOURCE_DIR}/generated/bindings/*.cc")
 
@@ -727,3 +804,44 @@ if(GENERATE_PYTHON_BINDINGS)
             DESTINATION .)
     endif()
 endif()
+
+install(TARGETS infiniops
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+    RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
+)
+
+install(DIRECTORY ${PROJECT_SOURCE_DIR}/include/
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+)
+
+if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
+    install(FILES
+        ${PROJECT_SOURCE_DIR}/generated/include/infini/operator_call_instantiations.h
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/infini
+    )
+endif()
+
+file(GLOB INFINIOPS_PUBLIC_CORE_HEADERS CONFIGURE_DEPENDS
+    "${CMAKE_CURRENT_SOURCE_DIR}/*.h")
+
+install(FILES ${INFINIOPS_PUBLIC_CORE_HEADERS}
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
+)
+
+install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/base/
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/base
+    FILES_MATCHING PATTERN "*.h"
+)
+
+install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/common/
+    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/common
+    FILES_MATCHING PATTERN "*.h"
+)
+
+if(EXISTS ${PROJECT_SOURCE_DIR}/generated/base)
+    install(DIRECTORY ${PROJECT_SOURCE_DIR}/generated/base/
+        DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/base
+        FILES_MATCHING PATTERN "*.h"
+    )
+endif()
diff --git a/src/operator.h b/src/operator.h
diff --git a/tests/test_cpp_api.py b/tests/test_cpp_api.py
