Skip to content

refactor InfiniOps cpu runtime through InfiniRT #644

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
spike-zhu wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

refactor InfiniOps cpu runtime through InfiniRT #644

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,10 @@ 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)
option(USE_EXISTING_GENERATED_WRAPPERS
Copy link
Copy Markdown
Collaborator

@voltjia voltjia Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个功能的确是需要的，但是不应该跟着这个 PR 提，应该单独拆分出一个 PR。而且这个看上去有点一刀切，应该做成那种通过识别 hash 等方式来决定是否重新生成，然后自动开启。

"Build from existing generated wrapper sources instead of regenerating them" OFF)
option(INFINIOPS_MINIMAL_ADD_BINDINGS
Copy link
Copy Markdown
Collaborator

@voltjia voltjia Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个是用来干啥的？如果只是开发时测试用途的话，这里就先不要了。

"Build a minimal Python module exposing only the add op" OFF)

set(_DEFAULT_HYGON_DTK_ROOT "/opt/dtk")

Expand Down Expand Up @@ -476,6 +480,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()

Expand All @@ -487,6 +492,11 @@ 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 TARGET infinirt)
add_subdirectory("${INFINIRT_SOURCE_DIR}" "${CMAKE_BINARY_DIR}/InfiniRT")
endif()

add_subdirectory(src)

if(NOT GENERATE_PYTHON_BINDINGS)
Expand Down
84 changes: 81 additions & 3 deletions scripts/generate_wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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
Copy link
Copy Markdown
Collaborator

@voltjia voltjia Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

需要注意核对 CONTRIBUTING.md，比如这里 return 上面应该空一行。



def _generate_pybind11(operator):
optional_tensor_params = _find_optional_tensor_params(operator.name)
vector_tensor_params = _find_vector_tensor_params(operator.name)
Expand Down Expand Up @@ -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)

Expand All @@ -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 = []

Expand All @@ -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}")

Expand Down Expand Up @@ -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"
)
Expand Down
67 changes: 64 additions & 3 deletions src/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Copy link
Copy Markdown
Collaborator

@voltjia voltjia Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个 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)
Expand All @@ -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)

Expand Down Expand Up @@ -510,7 +519,9 @@ target_include_directories(infiniops
$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
)

if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
if((GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
AND NOT USE_EXISTING_GENERATED_WRAPPERS
AND NOT INFINIOPS_MINIMAL_ADD_BINDINGS)
find_package(Python COMPONENTS Interpreter REQUIRED)
# Always regenerate wrappers so emitted call instantiations and bindings
# match the active device list. Stale generated files would omit
Expand All @@ -526,7 +537,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
)
Expand All @@ -536,6 +551,15 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS OR GENERATE_PYTHON_BINDINGS)
else()
message(STATUS "Generating wrappers - done")
endif()
elseif(USE_EXISTING_GENERATED_WRAPPERS)
if(NOT EXISTS "${PROJECT_SOURCE_DIR}/generated/include" OR
NOT EXISTS "${PROJECT_SOURCE_DIR}/generated/src" OR
NOT EXISTS "${PROJECT_SOURCE_DIR}/generated/bindings")
message(FATAL_ERROR
"`USE_EXISTING_GENERATED_WRAPPERS` is ON but generated wrapper "
"sources are missing under `${PROJECT_SOURCE_DIR}/generated`.")
endif()
message(STATUS "Using existing generated wrapper sources")
endif()

if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
Expand Down Expand Up @@ -606,6 +630,37 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
endif()

if(GENERATE_PYTHON_BINDINGS)
if(INFINIOPS_MINIMAL_ADD_BINDINGS)
find_package(Python COMPONENTS Interpreter Development REQUIRED)

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 REQUIRED)
pybind11_add_module(ops python/add_module.cc)
target_include_directories(ops PRIVATE
${PROJECT_SOURCE_DIR}
${PROJECT_SOURCE_DIR}/include
${PROJECT_SOURCE_DIR}/generated
${PROJECT_SOURCE_DIR}/generated/include
${CMAKE_CURRENT_SOURCE_DIR}
$<TARGET_PROPERTY:infiniops,INCLUDE_DIRECTORIES>)
target_compile_definitions(ops PRIVATE
$<TARGET_PROPERTY:infiniops,COMPILE_DEFINITIONS>)
target_link_libraries(ops PRIVATE infiniops)
return()
endif()

file(GLOB_RECURSE PYBIND11_SOURCES CONFIGURE_DEPENDS
"${PROJECT_SOURCE_DIR}/generated/bindings/*.cc")

Expand Down Expand Up @@ -792,7 +847,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 .)
Expand All @@ -811,6 +866,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}
)
Expand Down
Loading
Loading