Merge branch 'main' into chunk-cat

intel · Nov 18, 2024 · 131e546 · 131e546
2 parents db8a471 + f9c7682
commit 131e546
Show file tree

Hide file tree

Showing 79 changed files with 4,113 additions and 900 deletions.
diff --git a/.github/workflows/nightly_ondemand.yml b/.github/workflows/nightly_ondemand.yml
@@ -82,7 +82,11 @@ jobs:
     uses: ./.github/workflows/_linux_ut.yml
     with:
       abi: 0
+      keep_torch_xpu_ops: ${{ github.event_name == 'schedule' && 'false' || inputs.keep_torch_xpu_ops }}
       ut: op_regression,op_regression_dev1,op_extended,op_ut,torch_xpu
+      pytorch: ${{ github.event_name == 'schedule' && 'main' || inputs.pytorch }}
+      python: ${{ github.event_name == 'schedule' && '3.10' || inputs.python }}
+      triton: ${{ github.event_name == 'schedule' && '' || inputs.triton }}
       runner: linux.idc.xpu
 
   Linux-Nightly-Ondemand-E2E-Tests:

diff --git a/.github/workflows/nightly_ondemand_rolling.yml b/.github/workflows/nightly_ondemand_rolling.yml
@@ -82,7 +82,10 @@ jobs:
     uses: ./.github/workflows/_linux_ut.yml
     with:
       abi: 0
+      keep_torch_xpu_ops: ${{ github.event_name == 'schedule' && 'false' || inputs.keep_torch_xpu_ops }}
       ut: op_regression,op_regression_dev1,op_extended,op_ut,torch_xpu
+      pytorch: ${{ github.event_name == 'schedule' && 'main' || inputs.pytorch }}
+      python: ${{ github.event_name == 'schedule' && '3.10' || inputs.python }}
       driver: rolling
       runner: pvc_rolling
 

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -24,7 +24,7 @@ set(PROJECT_NAME "torch-xpu-ops")
 set(PROJECT_VERSION "2.3.0")
 # Avoid SYCL compiler error
 if(NOT WIN32)
-  string(APPEND CMAKE_CXX_FLAGS " -Wno-error=comment")
+  string(APPEND CMAKE_CXX_FLAGS " -Wno-error")
 endif()
 
 cmake_policy(SET CMP0048 NEW)

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