[Reland] Install xpu codegen header to torch/include

cmake/Codegen.cmake

@@ -4,6 +4,7 @@ endif()
 set(Codegen_XPU_cmake_included true)
 
 set(BUILD_TORCH_XPU_ATEN_GENERATED "${CMAKE_BINARY_DIR}/xpu/ATen")
+set(BUILD_TORCH_ATEN_GENERATED "${CMAKE_BINARY_DIR}/aten/src/ATen")
 file(MAKE_DIRECTORY ${BUILD_TORCH_XPU_ATEN_GENERATED})
 
 set(RegisterXPU_GENERATED ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU_0.cpp)
@@ -14,81 +15,74 @@ set(XPUFallback_TEMPLATE ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.t
 set(XPU_AOTI_INSTALL_DIR ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/generated/extend)
 set(XPU_AOTI_SHIM_HEADER ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h)
 set(XPU_AOTI_SHIM_SOURCE ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp)
+set(CODEGEN_XPU_YAML_DIR ${TORCH_XPU_OPS_ROOT}/yaml)
 
+# Codegen prepare process
 if(WIN32)
-  set(FILE_DISPLAY_CMD type)
+  file(TO_NATIVE_PATH "${CODEGEN_XPU_YAML_DIR}/templates" DestPATH)
+  file(TO_NATIVE_PATH "${CMAKE_SOURCE_DIR}/aten/src/ATen/templates" SrcPATH)
+  # Copy pytorch templates
+  execute_process(COMMAND cmd /c xcopy ${SrcPATH} ${DestPATH} /E /H /C /I /Y > nul)
 else()
-  set(FILE_DISPLAY_CMD cat)
+  # soft link to pytorch templates
+  execute_process(COMMAND ln -sf ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_XPU_YAML_DIR})
 endif()
-file(TO_NATIVE_PATH "${RegisterXPU_GENERATED}" RegisterXPU_GENERATED_NATIVE)
-file(TO_NATIVE_PATH "${XPUFallback_TEMPLATE}" XPUFallback_TEMPLATE_NATIVE)
-set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_TEMPLATE_NATIVE} ">>" ${RegisterXPU_GENERATED_NATIVE})
-
-function(GEN_XPU file_yaml)
-  set(generated_files "")
-  foreach(f ${ARGN})
-    list(APPEND generated_files "${f}")
-  endforeach()
-  set(CODEGEN_XPU_YAML_DIR ${TORCH_XPU_OPS_ROOT}/yaml)
 
-  # Codegen prepare process
-  if(WIN32)
-    file(TO_NATIVE_PATH "${CODEGEN_XPU_YAML_DIR}/templates" DestPATH)
-    file(TO_NATIVE_PATH "${CMAKE_SOURCE_DIR}/aten/src/ATen/templates" SrcPATH)
-    execute_process(COMMAND cmd /c xcopy ${SrcPATH} ${DestPATH} /E /H /C /I /Y > nul)
-  else()
-    execute_process(COMMAND ln -sf ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_XPU_YAML_DIR}) # soft link to pytorch templates
-  endif()
+set(XPU_CODEGEN_COMMAND
+  "${Python_EXECUTABLE}" -m torchgen.gen
+  --source-path ${CODEGEN_XPU_YAML_DIR}
+  --install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
+  --per-operator-headers
+  --backend-whitelist XPU SparseXPU SparseCsrXPU NestedTensorXPU
+  --xpu
+)
 
-  set(XPU_CODEGEN_COMMAND
-    "${Python_EXECUTABLE}" -m torchgen.gen
-    --source-path ${CODEGEN_XPU_YAML_DIR}
-    --install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
-    --per-operator-headers
-    --backend-whitelist XPU SparseXPU SparseCsrXPU NestedTensorXPU
-    --xpu
-  )
+set(XPU_INSTALL_HEADER_COMMAND
+  "${Python_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/install_xpu_headers.py
+  --src-header-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
+  --dst-header-dir ${BUILD_TORCH_ATEN_GENERATED}
+)
 
-  add_custom_command(
-    COMMENT "Generating XPU ATen Codegen..."
-    OUTPUT ${generated_files}
-    COMMAND
+# Generate ops_generated_headers.cmake for torch-xpu-ops
+execute_process(
+  COMMAND
     ${XPU_CODEGEN_COMMAND}
-    --static-dispatch-backend
-    # --update-aoti-c-shim: generate extend/c_shim_xpu.h
-    --update-aoti-c-shim
-    # --exten-aoti-c-shim: specifiy the extend/c_shim_xpu
-    # is out-of-tree extention for in-tree c_shim_xpu
-    --extend-aoti-c-shim
-    # --aoti-install-dir: generates c_shim_xpu.h and c_shim_xpu.cpp at
-    # torch/csrc/inductor/aoti_torch/generated/extend/
-    --aoti-install-dir=${XPU_AOTI_INSTALL_DIR}
-    COMMAND
-    ${REGISTER_FALLBACK_CMD}
-    # Codegen post-process
-    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_GENERATED}
-    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_GENERATED}
-    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseCsrXPU_GENERATED}
-    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterNestedTensorXPU_GENERATED}
-    WORKING_DIRECTORY ${TORCH_ROOT}
-    DEPENDS
-    ${CODEGEN_XPU_YAML_DIR}/native/${file_yaml}
-    ${XPUFallback_TEMPLATE}
-  )
+    --generate headers
+    --dry-run
+    --output-dependencies ${BUILD_TORCH_XPU_ATEN_GENERATED}/generated_headers.cmake
+  RESULT_VARIABLE RETURN_VALUE
+  WORKING_DIRECTORY ${TORCH_ROOT}
+)
+
+if(NOT RETURN_VALUE EQUAL 0)
+  message(FATAL_ERROR "Failed to generate ops_generated_headers.cmake for torch-xpu-ops.")
+endif()
+
+# Generate xpu_ops_generated_headers.cmake
+execute_process(
+  COMMAND
+    ${XPU_INSTALL_HEADER_COMMAND}
+    --dry-run
+  RESULT_VARIABLE RETURN_VALUE
+  WORKING_DIRECTORY ${TORCH_ROOT}
+)
+
+if(NOT RETURN_VALUE EQUAL 0)
+  message(FATAL_ERROR "Failed to generate xpu_ops_generated_headers.cmake.")
+endif()
+
+include(${BUILD_TORCH_XPU_ATEN_GENERATED}/xpu_ops_generated_headers.cmake)
 
-  # Post codegen delete the copied templates folder only on Windows.
-  if(WIN32)
-    add_custom_target(DELETE_TEMPLATES ALL DEPENDS ${generated_files})
-    add_custom_command(
-      TARGET DELETE_TEMPLATES
-      POST_BUILD
-      COMMAND ${CMAKE_COMMAND} -E remove_directory "${DestPATH}"
-    )
-  endif()
-endfunction(GEN_XPU)
+if(WIN32)
+  set(FILE_DISPLAY_CMD type)
+else()
+  set(FILE_DISPLAY_CMD cat)
+endif()
+file(TO_NATIVE_PATH "${RegisterXPU_GENERATED}" RegisterXPU_GENERATED_NATIVE)
+file(TO_NATIVE_PATH "${XPUFallback_TEMPLATE}" XPUFallback_TEMPLATE_NATIVE)
+set(REGISTER_FALLBACK_CMD ${FILE_DISPLAY_CMD} ${XPUFallback_TEMPLATE_NATIVE} ">>" ${RegisterXPU_GENERATED_NATIVE})
 
-GEN_XPU(
-  native_functions.yaml
+set(OUTPUT_LIST
   ${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions.h
   ${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions_inl.h
   ${RegisterXPU_GENERATED}
@@ -99,17 +93,49 @@ GEN_XPU(
   ${XPU_AOTI_SHIM_SOURCE}
 )
 
-# The c_shim_xpu.cpp needs include files in ${CMAKE_BINARY_DIR}/xpu/ATen/ops/*.h)
-# The include path is auto generated as "#include <ATen/ops/*.h">
-# To follow the design of aoti codegen, here ${CMAKE_BINARY_DIR}/xpu is added to
-# $TORCH_XPU_OPS_INCLUDE_DIRS, so that "#include <ATen/ops/*.h>" works.
-list(APPEND TORCH_XPU_OPS_INCLUDE_DIRS ${CMAKE_BINARY_DIR}/xpu)
+# Generate torch-xpu-ops codegen
+add_custom_command(
+  COMMENT "Generating XPU ATen Codegen..."
+  OUTPUT ${OUTPUT_LIST}
+  COMMAND
+    ${XPU_CODEGEN_COMMAND}
+    --static-dispatch-backend
+    --update-aoti-c-shim
+    --extend-aoti-c-shim
+    --aoti-install-dir=${XPU_AOTI_INSTALL_DIR}
+  COMMAND
+    ${REGISTER_FALLBACK_CMD}
+  # Codegen post process
+  COMMAND
+    ${XPU_INSTALL_HEADER_COMMAND}
+  DEPENDS
+    ATEN_CPU_FILES_GEN_TARGET
+    ATEN_XPU_FILES_GEN_TARGET
+    ${XPUFallback_TEMPLATE}
+    ${TORCH_XPU_OPS_ROOT}/tools/codegen/install_xpu_headers.py
+    ${BUILD_TORCH_XPU_ATEN_GENERATED}/xpu_ops_generated_headers.cmake
+    ${CODEGEN_XPU_YAML_DIR}/native/native_functions.yaml
+    ${all_python} ${headers_templates}
+    ${TORCH_ROOT}/aten/src/ATen/native/native_functions.yaml
+    ${TORCH_ROOT}/aten/src/ATen/native/tags.yaml
+  WORKING_DIRECTORY ${TORCH_ROOT}
+)
+
+# Codegen post progress
+if(WIN32)
+  add_custom_target(DELETE_TEMPLATES ALL DEPENDS ${OUTPUT_LIST})
+  # Delete the copied templates folder only on Windows.
+  add_custom_command(
+    TARGET DELETE_TEMPLATES
+    POST_BUILD
+    COMMAND ${CMAKE_COMMAND} -E remove_directory "${DestPATH}"
+  )
+endif()
 
-list(APPEND xpu_generated_src
+set(ATen_XPU_GEN_SRCS
   ${RegisterXPU_GENERATED}
   ${RegisterSparseXPU_GENERATED}
   ${RegisterSparseCsrXPU_GENERATED}
   ${RegisterNestedTensorXPU_GENERATED}
   ${XPU_AOTI_SHIM_SOURCE}
 )
-set(ATen_XPU_GEN_SRCS ${xpu_generated_src})

src/ATen/CMakeLists.txt

@@ -21,3 +21,7 @@ set(ATen_XPU_SYCL_SRCS ${ATen_XPU_SYCL_SRCS} PARENT_SCOPE)
 foreach(HEADER  ${xpu_h})
   install(FILES ${HEADER} DESTINATION "${AT_INSTALL_INCLUDE_DIR}/ATen/xpu")
 endforeach()
+
+foreach(HEADER  ${xpu_ops_generated_headers})
+  install(FILES ${HEADER} DESTINATION ${AT_INSTALL_INCLUDE_DIR}/ATen/ops)
+endforeach()

src/ATen/native/sparse/xpu/SparseCsrTensorMath.cpp

@@ -1,7 +1,7 @@
 #include <ATen/native/sparse/SparseStubs.h>
 #include <ATen/native/sparse/xpu/sycl/SparseCsrTensorMathKernels.h>
-#include <xpu/ATen/ops/_convert_indices_from_coo_to_csr_native.h>
-#include <xpu/ATen/ops/_convert_indices_from_csr_to_coo_native.h>
+#include <ATen/ops/_convert_indices_from_coo_to_csr_native.h>
+#include <ATen/ops/_convert_indices_from_csr_to_coo_native.h>
 
 namespace at::native {
 

src/ATen/native/xpu/Activation.cpp

@@ -7,9 +7,9 @@
 #include <ATen/native/TensorIterator.h>
 
 #include <ATen/ops/empty_like.h>
-#include <xpu/ATen/ops/empty.h>
-#include <xpu/ATen/ops/gelu_backward_native.h>
-#include <xpu/ATen/ops/gelu_native.h>
+#include <ATen/ops/empty.h>
+#include <ATen/ops/gelu_backward_native.h>
+#include <ATen/ops/gelu_native.h>
 
 #include <ATen/native/xpu/sycl/ActivationEluKernels.h>
 #include <ATen/native/xpu/sycl/ActivationGeluKernel.h>

src/ATen/native/xpu/AdaptiveAveragePooling2d.cpp

@@ -7,8 +7,8 @@
 
 #include <ATen/ops/mean.h>
 #include <ATen/ops/zeros_like.h>
-#include <xpu/ATen/ops/_adaptive_avg_pool2d_backward_native.h>
-#include <xpu/ATen/ops/_adaptive_avg_pool2d_native.h>
+#include <ATen/ops/_adaptive_avg_pool2d_backward_native.h>
+#include <ATen/ops/_adaptive_avg_pool2d_native.h>
 
 #include <ATen/native/xpu/sycl/AdaptiveAveragePooling2dKernels.h>
 

src/ATen/native/xpu/AdaptiveAveragePooling3d.cpp

@@ -4,8 +4,8 @@
 
 #include <ATen/ops/empty.h>
 #include <ATen/ops/empty_like.h>
-#include <xpu/ATen/ops/adaptive_avg_pool3d_backward_native.h>
-#include <xpu/ATen/ops/adaptive_avg_pool3d_native.h>
+#include <ATen/ops/adaptive_avg_pool3d_backward_native.h>
+#include <ATen/ops/adaptive_avg_pool3d_native.h>
 
 namespace at::native {
 

src/ATen/native/xpu/AdaptiveMaxPooling2d.cpp

@@ -4,8 +4,8 @@
 #include <ATen/native/xpu/sycl/AdaptiveMaxPooling2dKernels.h>
 #include <comm/RegisterUtils.h>
 
-#include <xpu/ATen/ops/adaptive_max_pool2d_backward_native.h>
-#include <xpu/ATen/ops/adaptive_max_pool2d_native.h>
+#include <ATen/ops/adaptive_max_pool2d_backward_native.h>
+#include <ATen/ops/adaptive_max_pool2d_native.h>
 
 namespace at {
 namespace native {

src/ATen/native/xpu/AdaptiveMaxPooling3d.cpp

@@ -4,8 +4,8 @@
 #include <ATen/native/xpu/sycl/AdaptiveMaxPooling3dKernels.h>
 
 #include <ATen/ops/empty.h>
-#include <xpu/ATen/ops/adaptive_max_pool3d_backward_native.h>
-#include <xpu/ATen/ops/adaptive_max_pool3d_native.h>
+#include <ATen/ops/adaptive_max_pool3d_backward_native.h>
+#include <ATen/ops/adaptive_max_pool3d_native.h>
 
 namespace at {
 namespace native {

src/ATen/native/xpu/AveragePool2d.cpp

@@ -5,8 +5,8 @@
 #include <ATen/native/xpu/sycl/AveragePool2dKernels.h>
 #include <comm/RegisterUtils.h>
 
-#include <xpu/ATen/ops/avg_pool2d_backward_native.h>
-#include <xpu/ATen/ops/avg_pool2d_native.h>
+#include <ATen/ops/avg_pool2d_backward_native.h>
+#include <ATen/ops/avg_pool2d_native.h>
 
 namespace at {
 namespace native {

src/ATen/native/xpu/AveragePool3d.cpp

@@ -1,8 +1,8 @@
 #include <ATen/core/Tensor.h>
 #include <ATen/native/xpu/sycl/AveragePool3dKernels.h>
 
-#include <xpu/ATen/ops/avg_pool3d_backward_native.h>
-#include <xpu/ATen/ops/avg_pool3d_native.h>
+#include <ATen/ops/avg_pool3d_backward_native.h>
+#include <ATen/ops/avg_pool3d_native.h>
 
 namespace at {
 namespace native {

src/ATen/native/xpu/BinaryOps.cpp

@@ -4,7 +4,7 @@
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/TensorIterator.h>
 
-#include <xpu/ATen/ops/add_native.h>
+#include <ATen/ops/add_native.h>
 
 #include <ATen/native/xpu/sycl/BinaryBitwiseOpsKernels.h>
 #include <ATen/native/xpu/sycl/BinaryGeometricKernels.h>

src/ATen/native/xpu/Col2Im.cpp

@@ -7,7 +7,7 @@
 #include <ATen/native/xpu/sycl/Col2ImKernel.h>
 
 #include <comm/xpu_aten.h>
-#include <xpu/ATen/ops/col2im_native.h>
+#include <ATen/ops/col2im_native.h>
 
 namespace at::native {
 

src/ATen/native/xpu/DilatedMaxPool2d.cpp

@@ -4,9 +4,9 @@
 #include <ATen/native/xpu/sycl/DilatedMaxPool2d.h>
 #include <comm/RegisterUtils.h>
 
-#include <xpu/ATen/ops/max.h>
-#include <xpu/ATen/ops/max_pool2d_with_indices_backward_native.h>
-#include <xpu/ATen/ops/max_pool2d_with_indices_native.h>
+#include <ATen/ops/max.h>
+#include <ATen/ops/max_pool2d_with_indices_backward_native.h>
+#include <ATen/ops/max_pool2d_with_indices_native.h>
 
 namespace at {
 namespace native {

src/ATen/native/xpu/DilatedMaxPool3d.cpp

@@ -2,8 +2,8 @@
 #include <ATen/native/xpu/sycl/DilatedMaxPool3d.h>
 
 #include <ATen/ops/empty.h>
-#include <xpu/ATen/ops/max_pool3d_with_indices_backward_native.h>
-#include <xpu/ATen/ops/max_pool3d_with_indices_native.h>
+#include <ATen/ops/max_pool3d_with_indices_backward_native.h>
+#include <ATen/ops/max_pool3d_with_indices_native.h>
 namespace at {
 namespace native {
 

src/ATen/native/xpu/Dropout.cpp

@@ -3,8 +3,8 @@
 #include <ATen/native/TensorIterator.h>
 #include <ATen/native/xpu/sycl/DropoutKernels.h>
 
-#include <xpu/ATen/ops/native_dropout_backward_native.h>
-#include <xpu/ATen/ops/native_dropout_native.h>
+#include <ATen/ops/native_dropout_backward_native.h>
+#include <ATen/ops/native_dropout_native.h>
 
 #include <comm/xpu_aten.h>
 

src/ATen/native/xpu/Embedding.cpp

@@ -1,6 +1,6 @@
 #include <ATen/core/op_registration/adaption.h>
 
-#include <xpu/ATen/ops/embedding_dense_backward_native.h>
+#include <ATen/ops/embedding_dense_backward_native.h>
 
 #include <ATen/native/xpu/sycl/EmbeddingKernels.h>
 #include <comm/xpu_aten.h>

src/ATen/native/xpu/EmbeddingBag.cpp

@@ -1,5 +1,5 @@
-#include <xpu/ATen/ops/_embedding_bag_forward_only_native.h>
-#include <xpu/ATen/ops/_embedding_bag_native.h>
+#include <ATen/ops/_embedding_bag_forward_only_native.h>
+#include <ATen/ops/_embedding_bag_native.h>
 
 #include <ATen/native/xpu/sycl/EmbeddingBagKernels.h>
 #include <comm/xpu_aten.h>

src/ATen/native/xpu/Equal.cpp

@@ -1,6 +1,6 @@
 #include <ATen/NamedTensorUtils.h>
 
-#include <xpu/ATen/ops/equal_native.h>
+#include <ATen/ops/equal_native.h>
 
 namespace at {
 namespace xpu {

src/ATen/native/xpu/ForeachOpScalarList.cpp

@@ -16,8 +16,8 @@
 #include <ATen/native/xpu/sycl/ForeachPointwiseOpScalarListKernels.h>
 #include <ATen/native/xpu/sycl/ForeachTernaryOpScalarListKernels.h>
 
-#include <xpu/ATen/ops/_foreach_add_native.h>
-#include <xpu/ATen/ops/_foreach_mul_native.h>
+#include <ATen/ops/_foreach_add_native.h>
+#include <ATen/ops/_foreach_mul_native.h>
 
 namespace at {
 namespace native {