Align directories and codegen with stock PyTorch

.lintrunner.toml

@@ -56,8 +56,8 @@ code = 'CLANGFORMAT'
 include_patterns = [
     'src/aten/*.h',
     'src/aten/*.cpp',
-    'src/aten/sycl/*.h',
-    'src/aten/sycl/*.cpp',
+    'src/ATen/native/xpu/sycl/*.h',
+    'src/ATen/native/xpu/sycl/*.cpp',
     'aten/src/ATen/*.h',
     'aten/src/ATen/mps/**/*.mm',
     'aten/src/ATen/xpu/**/*.h',

CMakeLists.txt

@@ -20,6 +20,8 @@
 
 cmake_minimum_required(VERSION 3.13 FATAL_ERROR)
 
+# message(FATAL_ERROR "In XPU cmakefile ${USE_PER_OPERATOR_HEADERS}")
+
 set(PROJECT_NAME "torch-xpu-ops")
 set(PROJECT_VERSION "2.3.0")
 cmake_policy(SET CMP0048 NEW)

cmake/BuildFlags.cmake

@@ -25,6 +25,11 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" OR CMAKE_CXX_COMPILER_ID STREQUAL "MSVC"
     list(APPEND SYCL_HOST_FLAGS -O0)
   endif(CMAKE_BUILD_TYPE MATCHES Debug)
 
+  if(USE_PER_OPERATOR_HEADERS)
+    list(APPEND SYCL_HOST_FLAGS -DAT_PER_OPERATOR_HEADERS)
+    # target_compile_definitions(torch_xpu_ops PUBLIC -DAT_PER_OPERATOR_HEADERS)
+  endif()
+
   # -- Kernel flags (SYCL_KERNEL_OPTIONS)
   # The fast-math will be enabled by default in SYCL compiler.
   # Refer to [https://clang.llvm.org/docs/UsersManual.html#cmdoption-fno-fast-math]

cmake/Codegen.cmake

@@ -3,7 +3,7 @@ if(Codegen_GPU_cmake_included)
 endif()
 set(Codegen_GPU_cmake_included true)
 
-set(BUILD_TORCH_XPU_ATEN_GENERATED "${CMAKE_BINARY_DIR}/aten/src/ATen/xpu")
+set(BUILD_TORCH_XPU_ATEN_GENERATED "${CMAKE_BINARY_DIR}/xpu/ATen/")
 file(MAKE_DIRECTORY ${BUILD_TORCH_XPU_ATEN_GENERATED})
 
 set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp)
@@ -43,10 +43,53 @@ function(GEN_BACKEND file_yaml)
     )
 endfunction(GEN_BACKEND)
 
-GEN_BACKEND(
-  xpu_functions.yaml
-  XPUNativeFunctions.h
-  RegisterXPU.cpp)
+set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp)
+set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
+function(GEN_XPU file_yaml)
+  set(generated_files "")
+  foreach(f ${ARGN})
+    list(APPEND generated_files "${BUILD_TORCH_XPU_ATEN_GENERATED}/${f}")
+  endforeach()
+  file(GLOB_RECURSE depend_files ${TORCH_XPU_OPS_ROOT}/yaml/${file_yaml})
+  set(CODEGEN_TEMPLATE ${TORCH_XPU_OPS_ROOT}/yaml/)
+
+  # Codegen prepare process
+  execute_process(COMMAND ln -s ${CMAKE_SOURCE_DIR}/aten/src/ATen/templates ${CODEGEN_TEMPLATE}) # soft link to pytorch templates
+
+  add_custom_command(
+    OUTPUT ${generated_files}
+    COMMAND
+    "${PYTHON_EXECUTABLE}" -m torchgen.gen
+    --source-path ${TORCH_XPU_OPS_ROOT}/yaml/
+    --install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
+    --per-operator-headers
+    --static-dispatch-backend
+    --backend-whitelist=XPU
+    COMMAND
+    cat ${XPUFallback_PATH} >> ${RegisterXPU_PATH}
+    # Codegen post-process
+    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_PATH}
+    ${SIMPLE_TRACE} 
+    WORKING_DIRECTORY ${TORCH_ROOT}
+    DEPENDS
+  ${depended_files}
+    ${TORCH_XPU_OPS_ROOT}/yaml/native/${file_yaml}
+    ${XPUFallback_PATH}
+  )
+endfunction(GEN_XPU)
+
+# GEN_BACKEND(
+#   xpu_functions.yaml
+#   XPUNativeFunctions.h
+#   RegisterXPU.cpp)
+
+GEN_XPU(
+  native_functions.yaml
+  XPUFunctions.h
+  RegisterXPU.cpp
+)
+
+
 
 
 list(APPEND xpu_generated_src ${RegisterXPU_PATH})

src/ATen/CMakeLists.txt

@@ -1,5 +1,6 @@
 # ATen XPU sources
 
+
 file(GLOB xpu_cpp "xpu/*.cpp", "native/xpu/*.cpp", "native/sparse/*.cpp")
 file(GLOB xpu_sycl "native/xpu/sycl/*.cpp")
 

src/ATen/native/sparse/SparseTensor.cpp

@@ -5,13 +5,9 @@
 #include <ATen/core/op_registration/adaption.h>
 #include <torch/library.h>
 
-#ifndef AT_PER_OPERATOR_HEADERS
-#include <ATen/Functions.h>
-#include <ATen/NativeFunctions.h>
-#else
 #include <ATen/ops/_nnz_native.h>
 #include <ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_native.h>
-#endif
+#include <ATen/ops/_values_native.h>
 
 namespace at::native::xpu {