Support SparseXPU backend codegen and basic operators (#1030)

This PR includes the codegen logic for the `SparseXPU` backend as well as some basic operators. --------- Co-authored-by: Yutao Xu <[email protected]>
intel · Nov 14, 2024 · 023194f · 023194f
1 parent e163c4b
commit 023194f
Show file tree

Hide file tree

Showing 17 changed files with 2,388 additions and 84 deletions.
diff --git a/cmake/Codegen.cmake b/cmake/Codegen.cmake
@@ -7,6 +7,7 @@ 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)
+set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp)
 set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
 
 if(WIN32)
@@ -45,6 +46,7 @@ endfunction(GEN_BACKEND)
 
 
 set(RegisterXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp)
+set(RegisterSparseXPU_PATH ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp)
 set(XPUFallback_PATH ${TORCH_XPU_OPS_ROOT}/src/ATen/native/xpu/XPUFallback.template)
 set(XPU_AOTI_INSTALL_DIR ${TORCH_ROOT}/torch/csrc/inductor/aoti_torch/generated/extend)
 function(GEN_XPU file_yaml)
@@ -75,7 +77,7 @@ function(GEN_XPU file_yaml)
     --install-dir ${BUILD_TORCH_XPU_ATEN_GENERATED}
     --per-operator-headers
     --static-dispatch-backend
-    --backend-whitelist=XPU
+    --backend-whitelist XPU SparseXPU
     # --xpu: generate in-tree RegisterXPU.cpp for in-tree OPs
     --xpu
     # --update-aoti-c-shim: generate extend/c_shim_xpu.h
@@ -90,7 +92,8 @@ function(GEN_XPU file_yaml)
     ${REGISTER_FALLBACK_CMD}
     # Codegen post-process
     COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterXPU_PATH}
-    ${SIMPLE_TRACE}
+    COMMAND "${PYTHON_EXECUTABLE}" ${TORCH_XPU_OPS_ROOT}/tools/codegen/remove_headers.py --register_xpu_path ${RegisterSparseXPU_PATH}
+    ${SIMPLE_TRACE} 
     WORKING_DIRECTORY ${TORCH_ROOT}
     DEPENDS
   ${depended_files}
@@ -106,10 +109,11 @@ endfunction(GEN_XPU)
 
 GEN_XPU(
   native_functions.yaml
- ${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions.h
- ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp
- ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h
- ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp
+  ${BUILD_TORCH_XPU_ATEN_GENERATED}/XPUFunctions.h
+  ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterXPU.cpp
+  ${BUILD_TORCH_XPU_ATEN_GENERATED}/RegisterSparseXPU.cpp
+  ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.h
+  ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp
 )
 
 
@@ -119,7 +123,7 @@ GEN_XPU(
 # $TORCH_XPU_OPS_INCLUDE_DIRS, so that "#include <ATen/ops/*.h>" works.
 list(APPEND TORCH_XPU_OPS_INCLUDE_DIRS ${CMAKE_BINARY_DIR}/xpu)
 
-list(APPEND xpu_generated_src ${RegisterXPU_PATH})
+list(APPEND xpu_generated_src ${RegisterXPU_PATH} ${RegisterSparseXPU_PATH})
 list(APPEND xpu_generated_src ${XPU_AOTI_INSTALL_DIR}/c_shim_xpu.cpp)
 add_custom_target(TORCH_XPU_GEN_TARGET DEPENDS ${xpu_generated_src})
 set(ATen_XPU_GEN_SRCS ${xpu_generated_src})
diff --git a/src/ATen/CMakeLists.txt b/src/ATen/CMakeLists.txt
@@ -2,8 +2,8 @@
 
 file(GLOB xpu_h "xpu/*.h")
 file(GLOB xpu_cpp "xpu/*.cpp")
-file(GLOB xpu_native_cpp "native/xpu/*.cpp" "native/sparse/*.cpp" "native/transformers/*.cpp" "native/quantized/*.cpp")
-file(GLOB xpu_sycl "native/xpu/sycl/*.cpp" "native/transformers/sycl/*.cpp" "native/quantized/sycl/*.cpp")
+file(GLOB xpu_native_cpp "native/xpu/*.cpp" "native/sparse/*.cpp" "native/sparse/xpu/*.cpp" "native/transformers/*.cpp" "native/quantized/*.cpp")
+file(GLOB xpu_sycl "native/xpu/sycl/*.cpp" "native/sparse/xpu/sycl/*.cpp" "native/transformers/sycl/*.cpp" "native/quantized/sycl/*.cpp")
 
 list(APPEND ATen_XPU_CPP_SRCS ${xpu_cpp})
 list(APPEND ATen_XPU_NATIVE_CPP_SRCS ${xpu_native_cpp})

diff --git a/src/ATen/native/sparse/SparseTensor.cpp b/src/ATen/native/sparse/SparseTensor.cpp
diff --git a/src/ATen/native/sparse/xpu/SparseBinaryOpIntersection.cpp b/src/ATen/native/sparse/xpu/SparseBinaryOpIntersection.cpp
@@ -0,0 +1,16 @@
+#include <ATen/native/sparse/SparseStubs.h>
+#include <ATen/native/sparse/xpu/sycl/SparseBinaryOpIntersectionKernels.h>
+
+namespace at::native {
+
+REGISTER_XPU_DISPATCH(
+    mul_sparse_sparse_out_stub,
+    &xpu::mul_sparse_sparse_kernel);
+REGISTER_XPU_DISPATCH(
+    sparse_mask_intersection_out_stub,
+    &xpu::sparse_mask_intersection_kernel);
+REGISTER_XPU_DISPATCH(
+    sparse_mask_projection_out_stub,
+    &xpu::sparse_mask_projection_kernel);
+
+} // namespace at::native
diff --git a/src/ATen/native/sparse/xpu/SparseTensor.cpp b/src/ATen/native/sparse/xpu/SparseTensor.cpp
@@ -0,0 +1,14 @@
+#include <ATen/native/sparse/SparseStubs.h>
+#include <ATen/native/sparse/xpu/sycl/SparseTensorKernels.h>
+
+namespace at::native {
+
+using namespace at::sparse;
+
+SparseTensor _coalesce_sparse_xpu(const SparseTensor& self) {
+  return xpu::coalesce_sparse_kernel(self);
+}
+
+REGISTER_XPU_DISPATCH(flatten_indices_stub, &xpu::flatten_indices_kernel);
+
+} // namespace at::native
diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp
@@ -0,0 +1,22 @@
+#include <ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h>
+
+namespace at::native {
+
+using namespace at::sparse;
+
+SparseTensor& add_out_sparse_xpu(
+    const SparseTensor& t,
+    const SparseTensor& src,
+    const Scalar& value,
+    SparseTensor& r_) {
+  return xpu::add_sparse_kernel(t, src, value, r_);
+}
+
+SparseTensor& mul_out_sparse_xpu(
+    const Tensor& t_,
+    const Tensor& src_,
+    SparseTensor& r_) {
+  return xpu::mul_sparse_kernel(t_, src_, r_);
+}
+
+} // namespace at::native