Skip to content

Commit 88c24a6

Browse files
guangyeyguangyey
committed
install codegen header to torch/include
1 parent f94f651 commit 88c24a6

File tree

1 file changed

+48
-80
lines changed

1 file changed

+48
-80
lines changed

‎cmake/Codegen.cmake

+48-80
Original file line numberDiff line numberDiff line change
@@ -1,89 +1,60 @@
1-
if(Codegen_GPU_cmake_included)
1+
if(Codegen_XPU_cmake_included)
22
return()
33
endif()
4-
set(Codegen_GPU_cmake_included true)
4+
set(Codegen_XPU_cmake_included true)
55

6-
set(BUILD_TORCH_XPU_ATEN_GENERATED "${CMAKE_BINARY_DIR}/xpu/ATen/")
6+
set(BUILD_TORCH_XPU_ATEN_GENERATED "${CMAKE_BINARY_DIR}/xpu/ATen")
77
file(MAKE_DIRECTORY ${BUILD_TORCH_XPU_ATEN_GENERATED})
88

9-
set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU_0.cpp)
10-
set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU_0.cpp)
11-
set(RegisterSparseCsrXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseCsrXPU_0.cpp)
12-
set(RegisterNestedTensorXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU_0.cpp)
13-
set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
9+
set(RegisterXPU_GENERATED ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU_0.cpp)
10+
set(RegisterSparseXPU_GENERATED ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU_0.cpp)
11+
set(RegisterSparseCsrXPU_GENERATED ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseCsrXPU_0.cpp)
12+
set(RegisterNestedTensorXPU_GENERATED ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU_0.cpp)
13+
set(XPUFallback_TEMPLATE ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
14+
set(XPU_AOTI_INSTALL_DIR ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/generated/extend)
15+
set(XPU_AOTI_SHIM_HEADER ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h)
16+
set(XPU_AOTI_SHIM_SOURCE ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp)
1417

1518
if(WIN32)
1619
set(FILE_DISPLAY_CMD type)
17-
# replace forward slash with back slash for compatibility with 'type' command on Windows
18-
string(REPLACE "/" "\\" RegisterXPU_PATH_BACKSLASH "${RegisterXPU_PATH}")
19-
string(REPLACE "/" "\\" XPUFallback_PATH_BACKSLASH "${XPUFallback_PATH}")
20-
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_PATH_BACKSLASH} ">>" ${RegisterXPU_PATH_BACKSLASH})
2120
else()
2221
set(FILE_DISPLAY_CMD cat)
23-
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_PATH} ">>" ${RegisterXPU_PATH})
2422
endif()
23+
file(TO_NATIVE_PATH "${RegisterXPU_GENERATED}" RegisterXPU_GENERATED_NATIVE)
24+
file(TO_NATIVE_PATH "${XPUFallback_TEMPLATE}" XPUFallback_TEMPLATE_NATIVE)
25+
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_TEMPLATE_NATIVE} ">>" ${RegisterXPU_GENERATED_NATIVE})
2526

26-
function(GEN_BACKEND file_yaml)
27-
set(generated_files "")
28-
foreach(f ${ARGN})
29-
list(APPEND generated_files "${BUILD_TORCH_XPU_ATEN_GENERATED}/${f}")
30-
endforeach()
31-
file(GLOB_RECURSE depended_files ${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml})
32-
add_custom_command(
33-
OUTPUT ${generated_files}
34-
COMMAND
35-
"${PYTHON_EXECUTABLE}" -m torchgen.gen_backend_stubs
36-
--output_dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
37-
--source_yaml ${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml}
38-
COMMAND
39-
${REGISTER_FALLBACK_CMD}
40-
${SIMPLE_TRACE}
41-
WORKING_DIRECTORY ${TORCH_ROOT}
42-
DEPENDS
43-
${depended_files}
44-
${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml}
45-
${XPUFallback_PATH}
46-
)
47-
endfunction(GEN_BACKEND)
48-
49-
50-
set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU_0.cpp)
51-
set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU_0.cpp)
52-
set(RegisterSparseCsrXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseCsrXPU_0.cpp)
53-
set(RegisterNestedTensorXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU_0.cpp)
54-
set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
55-
set(XPU_AOTI_INSTALL_DIR ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/generated/extend)
5627
function(GEN_XPU file_yaml)
5728
set(generated_files "")
5829
foreach(f ${ARGN})
5930
list(APPEND generated_files "${f}")
6031
endforeach()
61-
file(GLOB_RECURSE depend_files ${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml})
62-
set(CODEGEN_TEMPLATE ${TORCH_XPU_OPS_ROOT}/yaml/)
32+
set(CODEGEN_XPU_YAML_DIR ${TORCH_XPU_OPS_ROOT}/yaml)
6333

6434
# Codegen prepare process
6535
if(WIN32)
66-
string(REPLACE "/" "\\" DestPATH "${CODEGEN_TEMPLATE}templates")
36+
string(REPLACE "/" "\\" DestPATH "${CODEGEN_XPU_YAML_DIR}/templates")
6737
string(REPLACE "/" "\\" SrcPATH "${CMAKE_SOURCE_DIR}/aten/src/ATen/templates")
6838
execute_process(COMMAND cmd /c xcopy ${SrcPATH} ${DestPATH} /E /H /C /I /Y > nul)
69-
string(REPLACE "/" "\\" RegisterXPU_PATH_BACKSLASH "${RegisterXPU_PATH}")
70-
string(REPLACE "/" "\\" XPUFallback_PATH_BACKSLASH "${XPUFallback_PATH}")
71-
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_PATH_BACKSLASH} ">>" ${RegisterXPU_PATH_BACKSLASH})
7239
else()
73-
execute_process(COMMAND ln -s ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_TEMPLATE}) # soft link to pytorch templates
74-
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_PATH} ">>" ${RegisterXPU_PATH})
40+
execute_process(COMMAND ln -s ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_XPU_YAML_DIR}) # soft link to pytorch templates
7541
endif()
76-
add_custom_command(
77-
OUTPUT ${generated_files}
78-
COMMAND
42+
43+
set(XPU_CODEGEN_COMMAND
7944
"${PYTHON_EXECUTABLE}" -m torchgen.gen
80-
--source-path ${TORCH_XPU_OPS_ROOT}/yaml/
45+
--source-path ${CODEGEN_XPU_YAML_DIR}
8146
--install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
8247
--per-operator-headers
83-
--static-dispatch-backend
8448
--backend-whitelist XPU SparseXPU SparseCsrXPU NestedTensorXPU
85-
# --xpu: generate in-tree RegisterXPU_0.cpp for in-tree OPs
8649
--xpu
50+
)
51+
52+
add_custom_command(
53+
COMMENT "Generating XPU ATen Codegen..."
54+
OUTPUT ${generated_files}
55+
COMMAND
56+
${XPU_CODEGEN_COMMAND}
57+
--static-dispatch-backend
8758
# --update-aoti-c-shim: generate extend/c_shim_xpu.h
8859
--update-aoti-c-shim
8960
# --exten-aoti-c-shim: specifiy the extend/c_shim_xpu
@@ -95,16 +66,14 @@ function(GEN_XPU file_yaml)
9566
COMMAND
9667
${REGISTER_FALLBACK_CMD}
9768
# Codegen post-process
98-
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_PATH}
99-
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_PATH}
100-
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseCsrXPU_PATH}
101-
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterNestedTensorXPU_PATH}
102-
${SIMPLE_TRACE}
69+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_GENERATED}
70+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_GENERATED}
71+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseCsrXPU_GENERATED}
72+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterNestedTensorXPU_GENERATED}
10373
WORKING_DIRECTORY ${TORCH_ROOT}
10474
DEPENDS
105-
${depended_files}
106-
${TORCH_XPU_OPS_ROOT}/yaml/native/${file_yaml}
107-
${XPUFallback_PATH}
75+
${CODEGEN_XPU_YAML_DIR}/native/${file_yaml}
76+
${XPUFallback_TEMPLATE}
10877
)
10978

11079
# Post codegen delete the copied templates folder only on Windows.
@@ -118,30 +87,29 @@ function(GEN_XPU file_yaml)
11887
endif()
11988
endfunction(GEN_XPU)
12089

121-
# GEN_BACKEND(
122-
# xpu_functions.yaml
123-
# XPUNativeFunctions.h
124-
# RegisterXPU_0.cpp)
125-
12690
GEN_XPU(
12791
native_functions.yaml
12892
${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions.h
129-
${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU_0.cpp
130-
${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU_0.cpp
131-
${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseCsrXPU_0.cpp
132-
${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterNestedTensorXPU_0.cpp
133-
${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h
134-
${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp
93+
${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions_inl.h
94+
${RegisterXPU_GENERATED}
95+
${RegisterSparseXPU_GENERATED}
96+
${RegisterSparseCsrXPU_GENERATED}
97+
${RegisterNestedTensorXPU_GENERATED}
98+
${XPU_AOTI_SHIM_HEADER}
99+
${XPU_AOTI_SHIM_SOURCE}
135100
)
136101

137-
138102
# The c_shim_xpu.cpp needs include files in ${CMAKE_BINARY_DIR}/xpu/ATen/ops/*.h)
139103
# The include path is auto generated as "#include <ATen/ops/*.h">
140104
# To follow the design of aoti codegen, here ${CMAKE_BINARY_DIR}/xpu is added to
141105
# $TORCH_XPU_OPS_INCLUDE_DIRS, so that "#include <ATen/ops/*.h>" works.
142106
list(APPEND TORCH_XPU_OPS_INCLUDE_DIRS ${CMAKE_BINARY_DIR}/xpu)
143107

144-
list(APPEND xpu_generated_src ${RegisterXPU_PATH} ${RegisterSparseXPU_PATH} ${RegisterSparseCsrXPU_PATH} ${RegisterNestedTensorXPU_PATH})
145-
list(APPEND xpu_generated_src ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp)
146-
add_custom_target(TORCH_XPU_GEN_TARGET DEPENDS ${xpu_generated_src})
108+
list(APPEND xpu_generated_src
109+
${RegisterXPU_GENERATED}
110+
${RegisterSparseXPU_GENERATED}
111+
${RegisterSparseCsrXPU_GENERATED}
112+
${RegisterNestedTensorXPU_GENERATED}
113+
${XPU_AOTI_SHIM_SOURCE}
114+
)
147115
set(ATen_XPU_GEN_SRCS ${xpu_generated_src})

0 commit comments

Comments
 (0)