Skip to content

Commit 53dbfaf

Browse files
guangyeyguangyey
committed
install codegen header to torch/include
1 parent f94f651 commit 53dbfaf

File tree

1 file changed

+49
-79
lines changed

1 file changed

+49
-79
lines changed

‎cmake/Codegen.cmake

+49-79
Original file line numberDiff line numberDiff line change
@@ -1,89 +1,62 @@
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)
1720
# 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})
21+
string(REPLACE "/" "\\" RegisterXPU_GENERATED_BACKSLASH "${RegisterXPU_GENERATED}")
22+
string(REPLACE "/" "\\" XPUFallback_TEMPLATE_BACKSLASH "${XPUFallback_TEMPLATE}")
23+
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_TEMPLATE_BACKSLASH} ">>" ${RegisterXPU_GENERATED_BACKSLASH})
2124
else()
2225
set(FILE_DISPLAY_CMD cat)
23-
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_PATH} ">>" ${RegisterXPU_PATH})
26+
set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_TEMPLATE} ">>" ${RegisterXPU_GENERATED})
2427
endif()
2528

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)
5629
function(GEN_XPU file_yaml)
5730
set(generated_files "")
5831
foreach(f ${ARGN})
5932
list(APPEND generated_files "${f}")
6033
endforeach()
61-
file(GLOB_RECURSE depend_files ${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml})
62-
set(CODEGEN_TEMPLATE ${TORCH_XPU_OPS_ROOT}/yaml/)
34+
set(CODEGEN_XPU_YAML_DIR ${TORCH_XPU_OPS_ROOT}/yaml)
6335

6436
# Codegen prepare process
6537
if(WIN32)
66-
string(REPLACE "/" "\\" DestPATH "${CODEGEN_TEMPLATE}templates")
38+
string(REPLACE "/" "\\" DestPATH "${CODEGEN_XPU_YAML_DIR}/templates")
6739
string(REPLACE "/" "\\" SrcPATH "${CMAKE_SOURCE_DIR}/aten/src/ATen/templates")
6840
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})
7241
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})
42+
execute_process(COMMAND ln -s ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_XPU_YAML_DIR}) # soft link to pytorch templates
7543
endif()
76-
add_custom_command(
77-
OUTPUT ${generated_files}
78-
COMMAND
44+
45+
set(XPU_CODEGEN_COMMAND
7946
"${PYTHON_EXECUTABLE}" -m torchgen.gen
80-
--source-path ${TORCH_XPU_OPS_ROOT}/yaml/
47+
--source-path ${CODEGEN_XPU_YAML_DIR}
8148
--install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
8249
--per-operator-headers
83-
--static-dispatch-backend
8450
--backend-whitelist XPU SparseXPU SparseCsrXPU NestedTensorXPU
85-
# --xpu: generate in-tree RegisterXPU_0.cpp for in-tree OPs
8651
--xpu
52+
)
53+
54+
add_custom_command(
55+
COMMENT "Generating XPU ATen Codegen..."
56+
OUTPUT ${generated_files}
57+
COMMAND
58+
${XPU_CODEGEN_COMMAND}
59+
--static-dispatch-backend
8760
# --update-aoti-c-shim: generate extend/c_shim_xpu.h
8861
--update-aoti-c-shim
8962
# --exten-aoti-c-shim: specifiy the extend/c_shim_xpu
@@ -95,16 +68,14 @@ function(GEN_XPU file_yaml)
9568
COMMAND
9669
${REGISTER_FALLBACK_CMD}
9770
# 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}
71+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_GENERATED}
72+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_GENERATED}
73+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseCsrXPU_GENERATED}
74+
COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterNestedTensorXPU_GENERATED}
10375
WORKING_DIRECTORY ${TORCH_ROOT}
10476
DEPENDS
105-
${depended_files}
106-
${TORCH_XPU_OPS_ROOT}/yaml/native/${file_yaml}
107-
${XPUFallback_PATH}
77+
${CODEGEN_XPU_YAML_DIR}/native/${file_yaml}
78+
${XPUFallback_TEMPLATE}
10879
)
10980

11081
# Post codegen delete the copied templates folder only on Windows.
@@ -118,30 +89,29 @@ function(GEN_XPU file_yaml)
11889
endif()
11990
endfunction(GEN_XPU)
12091

121-
# GEN_BACKEND(
122-
# xpu_functions.yaml
123-
# XPUNativeFunctions.h
124-
# RegisterXPU_0.cpp)
125-
12692
GEN_XPU(
12793
native_functions.yaml
12894
${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
95+
${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions_inl.h
96+
${RegisterXPU_GENERATED}
97+
${RegisterSparseXPU_GENERATED}
98+
${RegisterSparseCsrXPU_GENERATED}
99+
${RegisterNestedTensorXPU_GENERATED}
100+
${XPU_AOTI_SHIM_HEADER}
101+
${XPU_AOTI_SHIM_SOURCE}
135102
)
136103

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

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})
110+
list(APPEND xpu_generated_src
111+
${RegisterXPU_GENERATED}
112+
${RegisterSparseXPU_GENERATED}
113+
${RegisterSparseCsrXPU_GENERATED}
114+
${RegisterNestedTensorXPU_GENERATED}
115+
${XPU_AOTI_SHIM_SOURCE}
116+
)
147117
set(ATen_XPU_GEN_SRCS ${xpu_generated_src})

0 commit comments

Comments
 (0)