refactor InfiniOps cpu runtime through InfiniRT

CMakeLists.txt

@@ -476,6 +476,7 @@ endif()
 
 # If all other platforms are not enabled, CPU is enabled by default.
 if(NOT WITH_NVIDIA AND NOT WITH_ILUVATAR AND NOT WITH_HYGON AND NOT WITH_METAX AND NOT WITH_MOORE AND NOT WITH_CAMBRICON AND NOT WITH_ASCEND)
+    set(WITH_CPU ON CACHE BOOL "Enable CPU backend" FORCE)
     add_compile_definitions(WITH_CPU=1)
 endif()
 
@@ -487,6 +488,16 @@ if(WITH_HYGON AND NOT EXISTS "${DTK_ROOT}/llvm/lib/LLVMgold.so")
     set(PYBIND11_ENABLE_EXTRAS OFF)
 endif()
 
+set(INFINIRT_SOURCE_DIR "${PROJECT_SOURCE_DIR}/../InfiniRT" CACHE PATH "InfiniRT source directory")
+if(NOT EXISTS "${INFINIRT_SOURCE_DIR}/CMakeLists.txt")
+    message(FATAL_ERROR
+        "InfiniRT not found at `${INFINIRT_SOURCE_DIR}`. "
+        "Set `INFINIRT_SOURCE_DIR` to the InfiniRT source directory.")
+endif()
+if(NOT TARGET infinirt)
+    add_subdirectory("${INFINIRT_SOURCE_DIR}" "${CMAKE_BINARY_DIR}/InfiniRT")
+endif()
+
 add_subdirectory(src)
 
 if(NOT GENERATE_PYTHON_BINDINGS)

scripts/generate_wrappers.py

@@ -32,6 +32,18 @@
 _INDENTATION = "  "
 
 
+def _get_infinirt_include_flags():
+    infinirt_source_dir = pathlib.Path(
+        os.environ.get("INFINIRT_SOURCE_DIR", "../InfiniRT")
+    )
+    infinirt_include_dir = infinirt_source_dir / "src"
+
+    if not infinirt_include_dir.exists():
+        return ()
+
+    return ("-I", str(infinirt_include_dir))
+
+
 @functools.lru_cache(maxsize=1)
 def _get_system_include_flags():
     """Probe the system C++ compiler for default include paths so libclang
@@ -85,7 +97,7 @@ def __call__(self, op_name):
             "src",
             "-I",
             str(_GENERATION_DIR),
-        ) + _get_system_include_flags()
+        ) + _get_infinirt_include_flags() + _get_system_include_flags()
         translation_unit = index.parse(str(_find_base_header(op_name)), args=args)
 
         nodes = tuple(type(self)._find(translation_unit.cursor, op_name))
@@ -160,6 +172,18 @@ def _find_vector_int64_params(op_name):
     return set(re.findall(r"std::vector<int64_t>\s+(\w+)", source))
 
 
+def _find_tensor_params(op_name):
+    source = _find_base_header(op_name).read_text()
+
+    params = set()
+    params.update(
+        re.findall(r"(?:^|[,(]\s*)(?:const\s+)?Tensor\s+(\w+)", source)
+    )
+    params.update(_find_optional_tensor_params(op_name))
+    params.update(_find_vector_tensor_params(op_name))
+    return params
+
+
 def _generate_pybind11(operator):
     optional_tensor_params = _find_optional_tensor_params(operator.name)
     vector_tensor_params = _find_vector_tensor_params(operator.name)
@@ -537,6 +561,7 @@ def _generate_tensor_caster(name, is_data=False):
 
 def _generate_generated_dispatch_entries(operator):
     optional_tensor_params = _find_optional_tensor_params(operator.name)
+    tensor_params = _find_tensor_params(operator.name)
     vector_tensor_params = _find_vector_tensor_params(operator.name)
     vector_int64_params = _find_vector_int64_params(operator.name)
 
@@ -555,6 +580,12 @@ def _is_vector_tensor(arg):
     def _is_vector_int64(arg):
         return arg.spelling in vector_int64_params
 
+    def _is_tensor(arg):
+        if arg.spelling in tensor_params:
+            return True
+
+        return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
+
     def _generate_params(node):
         parts = []
 
@@ -568,6 +599,8 @@ def _generate_params(node):
                 parts.append(f"std::vector<Tensor> {arg.spelling}")
             elif _is_vector_int64(arg):
                 parts.append(f"std::vector<int64_t> {arg.spelling}")
+            elif _is_tensor(arg):
+                parts.append(f"Tensor {arg.spelling}")
             else:
                 parts.append(f"{arg.type.spelling} {arg.spelling}")
 
@@ -712,16 +745,61 @@ def _strip_top_level_const(type_spelling):
 
 
 def _generate_operator_call_instantiation_entries(operator):
+    optional_tensor_params = _find_optional_tensor_params(operator.name)
+    tensor_params = _find_tensor_params(operator.name)
+    vector_tensor_params = _find_vector_tensor_params(operator.name)
+    vector_int64_params = _find_vector_int64_params(operator.name)
+
+    def _is_optional_tensor(arg):
+        if arg.spelling in optional_tensor_params:
+            return True
+
+        return "std::optional" in arg.type.spelling and (
+            "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
+        )
+
+    def _is_vector_tensor(arg):
+        if arg.spelling in vector_tensor_params:
+            return True
+
+        return "std::vector" in arg.type.spelling and (
+            "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
+        )
+
+    def _is_vector_int64(arg):
+        return arg.spelling in vector_int64_params
+
+    def _is_tensor(arg):
+        if arg.spelling in tensor_params:
+            return True
+
+        return "Tensor" in arg.type.spelling or "TensorView" in arg.type.spelling
+
+    def _normalized_type(arg):
+        if _is_optional_tensor(arg):
+            return "std::optional<Tensor>"
+
+        if _is_vector_tensor(arg):
+            return "std::vector<Tensor>"
+
+        if _is_vector_int64(arg):
+            return "std::vector<int64_t>"
+
+        if _is_tensor(arg):
+            return "Tensor"
+
+        return _strip_top_level_const(arg.type.spelling)
+
     def _generate_template_arguments(node):
         return ", ".join(
-            _strip_top_level_const(arg.type.spelling)
+            _normalized_type(arg)
             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}"
+            f"const {_normalized_type(arg)}& {arg.spelling}"
             for arg in node.get_arguments()
             if arg.spelling != "stream"
         )

src/CMakeLists.txt

@@ -3,8 +3,16 @@ add_library(infiniops SHARED)
 include(GNUInstallDirs)
 
 file(GLOB BASE_SRCS CONFIGURE_DEPENDS "*.cc")
+list(FILTER BASE_SRCS EXCLUDE REGEX ".*tensor\\.cc$")
 target_sources(infiniops PRIVATE ${BASE_SRCS})
 
+set(INFINIOPS_EMPTY_SOURCE "${CMAKE_CURRENT_BINARY_DIR}/infiniops_empty.cc")
+file(WRITE "${INFINIOPS_EMPTY_SOURCE}"
+    "namespace infini::ops { void infiniops_link_anchor() {} }\n")
+target_sources(infiniops PRIVATE "${INFINIOPS_EMPTY_SOURCE}")
+
+target_link_libraries(infiniops PUBLIC infinirt)
+
 set(DEVICE_LIST "")
 
 if(WITH_CPU)
@@ -15,6 +23,7 @@ if(WITH_CPU)
 
     file(GLOB_RECURSE CPU_SOURCES CONFIGURE_DEPENDS ${CPU_PATTERNS})
     list(APPEND CORE_SOURCES ${CPU_SOURCES})
+    target_sources(infiniops PRIVATE ${CPU_SOURCES})
 
     target_compile_definitions(infiniops PUBLIC WITH_CPU=1)
 
@@ -526,7 +535,11 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
     endif()
 
     execute_process(
-        COMMAND ${Python_EXECUTABLE} ${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py ${GENERATOR_ARGS}
+        COMMAND ${CMAKE_COMMAND} -E env
+                INFINIRT_SOURCE_DIR=${INFINIRT_SOURCE_DIR}
+                ${Python_EXECUTABLE}
+                ${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py
+                ${GENERATOR_ARGS}
         WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}
         RESULT_VARIABLE script_result
     )
@@ -792,7 +805,7 @@ if(GENERATE_PYTHON_BINDINGS)
     set_target_properties(infiniops PROPERTIES INSTALL_RPATH "${_INFINIOPS_INSTALL_RPATH}")
     set_target_properties(ops PROPERTIES INSTALL_RPATH "${_INFINIOPS_INSTALL_RPATH}")
 
-    install(TARGETS infiniops ops DESTINATION .)
+    install(TARGETS infinirt infiniops ops DESTINATION .)
 
     file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/__init__.py" "")
     install(FILES "${CMAKE_CURRENT_BINARY_DIR}/__init__.py" DESTINATION .)
@@ -811,6 +824,12 @@ install(TARGETS infiniops
     RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
 )
 
+install(TARGETS infinirt
+    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}
 )