From 3e2fd3d5559be55485ef5dec4722f5b9cdc96a87 Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 18 Nov 2025 16:40:18 -0800 Subject: [PATCH 01/19] sparse addmm --- .../native/sparse/xpu/SparseTensorMath.cpp | 59 +++++++++++++++++++ test/xpu/test_sparse_xpu.py | 6 +- yaml/native/native_functions.yaml | 19 ++++++ 3 files changed, 83 insertions(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 954ab406c8..311a1bca3a 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -1,4 +1,7 @@ #include +#include + +#include namespace at::native { @@ -26,4 +29,60 @@ Tensor _sparse_sum_backward_xpu( return xpu::_sparse_sum_backward_kernel(grad_, input_, dims_to_sum); } +Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTensor& sparse_, const Tensor& dense, const Scalar& beta, const Scalar& alpha) { + TORCH_CHECK(t.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'self' to be XPU, but got CPU"); + TORCH_CHECK(r_.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'out' to be XPU, but got CPU"); + TORCH_CHECK(sparse_.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'mat1' to be XPU, but got CPU"); + TORCH_CHECK(dense.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'mat2' to be XPU, but got CPU"); + + TORCH_CHECK(xpu::check_device({sparse_, r_, t, dense})); + + TORCH_CHECK(dense.dim() == 2, "addmm: 2D tensor expected, got ", dense.dim(), "D tensor"); + TORCH_CHECK(sparse_.sparse_dim() == 2, "addmm: expected first two dims to be sparse (indices has size 2 at first dim), but got ", sparse_.sparse_dim(), " sparse dims"); + // no need to check dense_dim because dense_dim + sparse_dim = dim + + // change sparse to dense + Tensor mat1_dense = at::native::sparse_to_dense(sparse_); + // calcaulte + at::native::addmm_out_xpu(t, mat1_dense, dense, beta, alpha, r_); + + return r_; +} + +Tensor s_addmm_sparse_dense_xpu( + const Tensor& t, + const SparseTensor& sparse, + const Tensor& dense, + const Scalar& beta, + const Scalar& alpha +) { + Tensor r = at::empty({0}, t.options()); + s_addmm_out_sparse_dense_xpu(r, t, sparse, dense, beta, alpha); + return r; +} + + +Tensor& addmm_out_sparse_dense_xpu( + const Tensor& self, + const SparseTensor& mat1, + const Tensor& mat2, + const Scalar& beta, + const Scalar& alpha, + Tensor& result +) { + std::cout << "Call addmm_out_sparse_dense_xpu" << std::endl; + c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); + return s_addmm_sparse_dense_xpu(*b_self, mat1, mat2, beta, alpha); +} + +Tensor& s_addmm_sparse_dense_xpu_( + Tensor& t, + const SparseTensor& sparse, + const Tensor& dense, + const Scalar& beta, + const Scalar& alpha +) { + return s_addmm_out_sparse_dense_xpu(t, t, sparse, dense, beta, alpha); +} + } // namespace at::native diff --git a/test/xpu/test_sparse_xpu.py b/test/xpu/test_sparse_xpu.py index 98b3c5c643..4e5539c33b 100644 --- a/test/xpu/test_sparse_xpu.py +++ b/test/xpu/test_sparse_xpu.py @@ -2073,8 +2073,12 @@ def test_shape(di, dj, dk, nnz): @precisionOverride({torch.bfloat16: 5e-2, torch.float16: 5e-2}) @dtypes(torch.double, torch.cdouble, torch.bfloat16, torch.float16) @dtypesIfMPS(torch.float32, torch.complex64, torch.bfloat16, torch.float16) + # @skipXPUIf( + # True, + # "addmm sprase xpu not supported yet, see https://github.com/intel/torch-xpu-ops/issues/2211", + # ) @skipXPUIf( - True, + False, "addmm sprase xpu not supported yet, see https://github.com/intel/torch-xpu-ops/issues/2211", ) def test_sparse_addmm(self, device, dtype, coalesced): diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml index a3281791de..7849da5d49 100644 --- a/yaml/native/native_functions.yaml +++ b/yaml/native/native_functions.yaml @@ -9463,3 +9463,22 @@ variants: function, method - func: inverse.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + structured_delegate: addmm.out + variants: function, method + dispatch: + SparseXPU: addmm_sparse_dense_xpu + +- func: addmm.out(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + SparseXPU: addmm_out_sparse_dense_xpu + +- func: addmm_(Tensor(a!) self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + structured_delegate: addmm.out + variants: method + dispatch: + # Warning! For whatever reason, the inplace sparse addmm is NON + # broadcasting + SparseXPU: s_addmm_sparse_dense_xpu_ From 60361f879422e685b23ecf3dab1d6d6dca2c1a4e Mon Sep 17 00:00:00 2001 From: jenniew Date: Wed, 19 Nov 2025 00:06:08 -0800 Subject: [PATCH 02/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 311a1bca3a..138d9eb2fa 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include From 14ffb2031d17bdb4f8d32872b7a53666cdfa6858 Mon Sep 17 00:00:00 2001 From: jenniew Date: Wed, 19 Nov 2025 11:15:30 -0800 Subject: [PATCH 03/19] update --- .../xpu/sycl/SparseTensorMathKernels.cpp | 26 +++++++++---------- .../sparse/xpu/sycl/SparseTensorMathKernels.h | 14 ++++++++++ 2 files changed, 27 insertions(+), 13 deletions(-) diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp index 7cc7aefa6f..218b9571b4 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp @@ -193,19 +193,19 @@ struct SparseElementwiseKernelScalarFunctor { } // namespace apply -// Check if every tensor in a list of tensors matches the current -// device. -inline bool check_device(ArrayRef ts) { - if (ts.empty()) { - return true; - } - Device curDevice = Device(kXPU, c10::xpu::current_device()); - for (const Tensor& t : ts) { - if (t.device() != curDevice) - return false; - } - return true; -} +// // Check if every tensor in a list of tensors matches the current +// // device. +// inline bool check_device(ArrayRef ts) { +// if (ts.empty()) { +// return true; +// } +// Device curDevice = Device(kXPU, c10::xpu::current_device()); +// for (const Tensor& t : ts) { +// if (t.device() != curDevice) +// return false; +// } +// return true; +// } Tensor& add_out_dense_sparse_kernel( Tensor& r_, diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h index 1bcd08d40d..ec9e67c0be 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h @@ -23,4 +23,18 @@ TORCH_XPU_API Tensor _sparse_sum_backward_kernel( const SparseTensor& input_, IntArrayRef dims_to_sum); +// Check if every tensor in a list of tensors matches the current +// device. +inline bool check_device(ArrayRef ts) { + if (ts.empty()) { + return true; + } + Device curDevice = Device(kXPU, c10::xpu::current_device()); + for (const Tensor& t : ts) { + if (t.device() != curDevice) + return false; + } + return true; +} + } // namespace at::native::xpu From 6ddb883a6c87a3d47648ffb13f37d6e99ccda290 Mon Sep 17 00:00:00 2001 From: jenniew Date: Wed, 19 Nov 2025 11:41:07 -0800 Subject: [PATCH 04/19] update --- .../native/sparse/xpu/SparseTensorMath.cpp | 2 +- .../xpu/sycl/SparseTensorMathKernels.cpp | 26 +++++++++---------- .../sparse/xpu/sycl/SparseTensorMathKernels.h | 26 +++++++++---------- 3 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 138d9eb2fa..8a428013fa 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -35,7 +35,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe TORCH_CHECK(sparse_.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'mat1' to be XPU, but got CPU"); TORCH_CHECK(dense.is_xpu(), "Expected all tensors to be on the same device. addmm: expected 'mat2' to be XPU, but got CPU"); - TORCH_CHECK(xpu::check_device({sparse_, r_, t, dense})); + // TORCH_CHECK(xpu::check_device({sparse_, r_, t, dense})); TORCH_CHECK(dense.dim() == 2, "addmm: 2D tensor expected, got ", dense.dim(), "D tensor"); TORCH_CHECK(sparse_.sparse_dim() == 2, "addmm: expected first two dims to be sparse (indices has size 2 at first dim), but got ", sparse_.sparse_dim(), " sparse dims"); diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp index 218b9571b4..7cc7aefa6f 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.cpp @@ -193,19 +193,19 @@ struct SparseElementwiseKernelScalarFunctor { } // namespace apply -// // Check if every tensor in a list of tensors matches the current -// // device. -// inline bool check_device(ArrayRef ts) { -// if (ts.empty()) { -// return true; -// } -// Device curDevice = Device(kXPU, c10::xpu::current_device()); -// for (const Tensor& t : ts) { -// if (t.device() != curDevice) -// return false; -// } -// return true; -// } +// Check if every tensor in a list of tensors matches the current +// device. +inline bool check_device(ArrayRef ts) { + if (ts.empty()) { + return true; + } + Device curDevice = Device(kXPU, c10::xpu::current_device()); + for (const Tensor& t : ts) { + if (t.device() != curDevice) + return false; + } + return true; +} Tensor& add_out_dense_sparse_kernel( Tensor& r_, diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h index ec9e67c0be..fa32fb7019 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h @@ -23,18 +23,18 @@ TORCH_XPU_API Tensor _sparse_sum_backward_kernel( const SparseTensor& input_, IntArrayRef dims_to_sum); -// Check if every tensor in a list of tensors matches the current -// device. -inline bool check_device(ArrayRef ts) { - if (ts.empty()) { - return true; - } - Device curDevice = Device(kXPU, c10::xpu::current_device()); - for (const Tensor& t : ts) { - if (t.device() != curDevice) - return false; - } - return true; -} +// // Check if every tensor in a list of tensors matches the current +// // device. +// inline bool check_device(ArrayRef ts) { +// if (ts.empty()) { +// return true; +// } +// Device curDevice = Device(kXPU, c10::xpu::current_device()); +// for (const Tensor& t : ts) { +// if (t.device() != curDevice) +// return false; +// } +// return true; +// } } // namespace at::native::xpu From d92a628c64b3f55a2da5d10265ac8f7280889045 Mon Sep 17 00:00:00 2001 From: jenniew Date: Wed, 19 Nov 2025 12:11:31 -0800 Subject: [PATCH 05/19] update at::addmm_out --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 8a428013fa..27f8c8a9e5 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -43,8 +43,8 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe // change sparse to dense Tensor mat1_dense = at::native::sparse_to_dense(sparse_); - // calcaulte - at::native::addmm_out_xpu(t, mat1_dense, dense, beta, alpha, r_); + // calculate + at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); return r_; } From f413e27c258212f9c05275a04b5bd9f5a08434ad Mon Sep 17 00:00:00 2001 From: jenniew Date: Wed, 19 Nov 2025 12:31:11 -0800 Subject: [PATCH 06/19] add include files --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 27f8c8a9e5..d2d2379494 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -1,6 +1,15 @@ #include #include +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +#include +#endif + +#include + #include namespace at::native { From b20065f03cfee9b7f56387070825aa9f9841b0e9 Mon Sep 17 00:00:00 2001 From: jenniew Date: Thu, 20 Nov 2025 14:31:14 -0800 Subject: [PATCH 07/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index d2d2379494..ab490dded2 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -51,7 +51,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe // no need to check dense_dim because dense_dim + sparse_dim = dim // change sparse to dense - Tensor mat1_dense = at::native::sparse_to_dense(sparse_); + Tensor mat1_dense = at::native::sparse_to_dense(sparse_，{}, {}); // calculate at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); @@ -80,6 +80,17 @@ Tensor& addmm_out_sparse_dense_xpu( Tensor& result ) { std::cout << "Call addmm_out_sparse_dense_xpu" << std::endl; + c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); + return s_addmm_out_sparse_dense_xpu(result, *b_self, mat1, mat2, beta, alpha); +} + +Tensor addmm_sparse_dense_xpu( + const Tensor& self, + const SparseTensor& mat1, + const Tensor& mat2, + const Scalar& beta, + const Scalar& alpha +) { c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); return s_addmm_sparse_dense_xpu(*b_self, mat1, mat2, beta, alpha); } From 8b17d68655d05e505d249945306988f55ec6224f Mon Sep 17 00:00:00 2001 From: jenniew Date: Thu, 20 Nov 2025 14:46:51 -0800 Subject: [PATCH 08/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index ab490dded2..e424afd83a 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -51,7 +51,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe // no need to check dense_dim because dense_dim + sparse_dim = dim // change sparse to dense - Tensor mat1_dense = at::native::sparse_to_dense(sparse_，{}, {}); + Tensor mat1_dense = at::native::sparse_to_dense(sparse_, {}, {}); // calculate at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); From f8808c97513833c3f62593224e12df6183462967 Mon Sep 17 00:00:00 2001 From: jenniew Date: Thu, 20 Nov 2025 15:11:33 -0800 Subject: [PATCH 09/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index e424afd83a..8483dbc647 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -51,6 +51,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe // no need to check dense_dim because dense_dim + sparse_dim = dim // change sparse to dense + std::cout << "Call at::native::sparse_to_dense" << std::endl; Tensor mat1_dense = at::native::sparse_to_dense(sparse_, {}, {}); // calculate at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); @@ -91,6 +92,7 @@ Tensor addmm_sparse_dense_xpu( const Scalar& beta, const Scalar& alpha ) { + std::cout << "Call addmm_sparse_dense_xpu" << std::endl; c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); return s_addmm_sparse_dense_xpu(*b_self, mat1, mat2, beta, alpha); } @@ -102,6 +104,7 @@ Tensor& s_addmm_sparse_dense_xpu_( const Scalar& beta, const Scalar& alpha ) { + std::cout << "Call s_addmm_sparse_dense_xpu_" << std::endl; return s_addmm_out_sparse_dense_xpu(t, t, sparse, dense, beta, alpha); } From 81c48329208b987f164be52e3e91f91c3336f3c6 Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 25 Nov 2025 12:34:49 -0800 Subject: [PATCH 10/19] change to _to_dense --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 8483dbc647..acaccb31e8 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -1,5 +1,5 @@ #include -#include +// #include #ifndef AT_PER_OPERATOR_HEADERS #include @@ -52,7 +52,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe // change sparse to dense std::cout << "Call at::native::sparse_to_dense" << std::endl; - Tensor mat1_dense = at::native::sparse_to_dense(sparse_, {}, {}); + Tensor mat1_dense = sparse_._to_dense(std::nullopt, std::nullopt); // calculate at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); From 05ec3297bc51e9df166c7bf64f6f2c23296141da Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 25 Nov 2025 16:25:21 -0800 Subject: [PATCH 11/19] update --- .../native/sparse/xpu/SparseTensorMath.cpp | 5 ---- .../sparse/xpu/sycl/SparseTensorMathKernels.h | 26 +++++++++---------- 2 files changed, 13 insertions(+), 18 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index acaccb31e8..1934f33bc8 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -10,8 +10,6 @@ #include -#include - namespace at::native { using namespace at::sparse; @@ -50,10 +48,7 @@ Tensor& s_addmm_out_sparse_dense_xpu(Tensor& r_, const Tensor& t, const SparseTe TORCH_CHECK(sparse_.sparse_dim() == 2, "addmm: expected first two dims to be sparse (indices has size 2 at first dim), but got ", sparse_.sparse_dim(), " sparse dims"); // no need to check dense_dim because dense_dim + sparse_dim = dim - // change sparse to dense - std::cout << "Call at::native::sparse_to_dense" << std::endl; Tensor mat1_dense = sparse_._to_dense(std::nullopt, std::nullopt); - // calculate at::addmm_out(r_, t, mat1_dense, dense, beta, alpha); return r_; diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h index fa32fb7019..ec9e67c0be 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h @@ -23,18 +23,18 @@ TORCH_XPU_API Tensor _sparse_sum_backward_kernel( const SparseTensor& input_, IntArrayRef dims_to_sum); -// // Check if every tensor in a list of tensors matches the current -// // device. -// inline bool check_device(ArrayRef ts) { -// if (ts.empty()) { -// return true; -// } -// Device curDevice = Device(kXPU, c10::xpu::current_device()); -// for (const Tensor& t : ts) { -// if (t.device() != curDevice) -// return false; -// } -// return true; -// } +// Check if every tensor in a list of tensors matches the current +// device. +inline bool check_device(ArrayRef ts) { + if (ts.empty()) { + return true; + } + Device curDevice = Device(kXPU, c10::xpu::current_device()); + for (const Tensor& t : ts) { + if (t.device() != curDevice) + return false; + } + return true; +} } // namespace at::native::xpu From 73733093390926457eb9297289f9153ed7c498ca Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 25 Nov 2025 16:26:25 -0800 Subject: [PATCH 12/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 1934f33bc8..23671451b9 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -75,7 +75,6 @@ Tensor& addmm_out_sparse_dense_xpu( const Scalar& alpha, Tensor& result ) { - std::cout << "Call addmm_out_sparse_dense_xpu" << std::endl; c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); return s_addmm_out_sparse_dense_xpu(result, *b_self, mat1, mat2, beta, alpha); } @@ -87,7 +86,6 @@ Tensor addmm_sparse_dense_xpu( const Scalar& beta, const Scalar& alpha ) { - std::cout << "Call addmm_sparse_dense_xpu" << std::endl; c10::MaybeOwned b_self = expand_size(self, {mat1.size(0), mat2.size(1)}, "addmm_out"); return s_addmm_sparse_dense_xpu(*b_self, mat1, mat2, beta, alpha); } @@ -99,7 +97,6 @@ Tensor& s_addmm_sparse_dense_xpu_( const Scalar& beta, const Scalar& alpha ) { - std::cout << "Call s_addmm_sparse_dense_xpu_" << std::endl; return s_addmm_out_sparse_dense_xpu(t, t, sparse, dense, beta, alpha); } From fbb360f3d66556a532c389fd1d86dea8dec90691 Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 25 Nov 2025 16:28:01 -0800 Subject: [PATCH 13/19] update --- test/xpu/test_sparse_xpu.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/test/xpu/test_sparse_xpu.py b/test/xpu/test_sparse_xpu.py index 4e5539c33b..42aa76552e 100644 --- a/test/xpu/test_sparse_xpu.py +++ b/test/xpu/test_sparse_xpu.py @@ -2073,14 +2073,6 @@ def test_shape(di, dj, dk, nnz): @precisionOverride({torch.bfloat16: 5e-2, torch.float16: 5e-2}) @dtypes(torch.double, torch.cdouble, torch.bfloat16, torch.float16) @dtypesIfMPS(torch.float32, torch.complex64, torch.bfloat16, torch.float16) - # @skipXPUIf( - # True, - # "addmm sprase xpu not supported yet, see https://github.com/intel/torch-xpu-ops/issues/2211", - # ) - @skipXPUIf( - False, - "addmm sprase xpu not supported yet, see https://github.com/intel/torch-xpu-ops/issues/2211", - ) def test_sparse_addmm(self, device, dtype, coalesced): if (dtype is torch.bfloat16 or dtype is torch.float16) and device.startswith( "cuda" From cd31ed10466242c6f56b28c6e28f08dd88833482 Mon Sep 17 00:00:00 2001 From: jenniew Date: Tue, 25 Nov 2025 16:31:54 -0800 Subject: [PATCH 14/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 1 - .../sparse/xpu/sycl/SparseTensorMathKernels.h | 14 -------------- 2 files changed, 15 deletions(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 23671451b9..146169af21 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -1,5 +1,4 @@ #include -// #include #ifndef AT_PER_OPERATOR_HEADERS #include diff --git a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h index ec9e67c0be..1bcd08d40d 100644 --- a/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h +++ b/src/ATen/native/sparse/xpu/sycl/SparseTensorMathKernels.h @@ -23,18 +23,4 @@ TORCH_XPU_API Tensor _sparse_sum_backward_kernel( const SparseTensor& input_, IntArrayRef dims_to_sum); -// Check if every tensor in a list of tensors matches the current -// device. -inline bool check_device(ArrayRef ts) { - if (ts.empty()) { - return true; - } - Device curDevice = Device(kXPU, c10::xpu::current_device()); - for (const Tensor& t : ts) { - if (t.device() != curDevice) - return false; - } - return true; -} - } // namespace at::native::xpu From 7ce0059f56399f49bb3434db153c91970fdafe4a Mon Sep 17 00:00:00 2001 From: jenniew Date: Mon, 1 Dec 2025 10:56:01 -0800 Subject: [PATCH 15/19] sparse mm --- .../native/sparse/xpu/SparseTensorMath.cpp | 29 +++++++++++++++++++ test/xpu/test_sparse_xpu.py | 1 + yaml/native/native_functions.yaml | 17 +++++++++++ 3 files changed, 47 insertions(+) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 146169af21..8d1464d9db 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -99,4 +99,33 @@ Tensor& s_addmm_sparse_dense_xpu_( return s_addmm_out_sparse_dense_xpu(t, t, sparse, dense, beta, alpha); } +Tensor sparse_sparse_matmul_xpu(const Tensor& mat1_, const Tensor& mat2_) { + TORCH_INTERNAL_ASSERT(mat1_.is_sparse()); + TORCH_INTERNAL_ASSERT(mat2_.is_sparse()); + TORCH_CHECK(mat1_.dim() == 2); + TORCH_CHECK(mat2_.dim() == 2); + TORCH_CHECK(mat1_.dense_dim() == 0, "sparse_mm: scalar values expected, mat1 got ", mat1_.dense_dim(), "D values"); + TORCH_CHECK(mat2_.dense_dim() == 0, "sparse_mm: scalar values expected, mat2 got ", mat2_.dense_dim(), "D values"); + + TORCH_CHECK( + mat1_.size(1) == mat2_.size(0), "mat1 and mat2 shapes cannot be multiplied (", + mat1_.size(0), "x", mat1_.size(1), " and ", mat2_.size(0), "x", mat2_.size(1), ")"); + + TORCH_CHECK(mat1_.scalar_type() == mat2_.scalar_type(), + "mat1 dtype ", mat1_.scalar_type(), " does not match mat2 dtype ", mat2_.scalar_type()); + + // convert to dense + Tensor mat1_dense = mat1_._to_dense(std::nullopt, std::nullopt); + Tensor mat2_dense = mat2_._to_dense(std::nullopt, std::nullopt); + + Tensor output_dense = at::matmul(mat1_dense, mat2_dense); + // convert back to sparse + Tensor output_sparse = output_dense._to_sparse(mat1.layout()); + + return output_sparse; + + // auto output = at::native::empty_like(mat1_); + // output.sparse_resize_and_clear_({mat1_.size(0), mat2_.size(1)}, mat1_.sparse_dim(), 0); +} + } // namespace at::native diff --git a/test/xpu/test_sparse_xpu.py b/test/xpu/test_sparse_xpu.py index 42aa76552e..1ca7e1b69c 100644 --- a/test/xpu/test_sparse_xpu.py +++ b/test/xpu/test_sparse_xpu.py @@ -2073,6 +2073,7 @@ def test_shape(di, dj, dk, nnz): @precisionOverride({torch.bfloat16: 5e-2, torch.float16: 5e-2}) @dtypes(torch.double, torch.cdouble, torch.bfloat16, torch.float16) @dtypesIfMPS(torch.float32, torch.complex64, torch.bfloat16, torch.float16) + @skipXPUIf(False, "https://github.com/intel/torch-xpu-ops/issues/2211") def test_sparse_addmm(self, device, dtype, coalesced): if (dtype is torch.bfloat16 or dtype is torch.float16) and device.startswith( "cuda" diff --git a/yaml/native/native_functions.yaml b/yaml/native/native_functions.yaml index 7849da5d49..f71fac04b1 100644 --- a/yaml/native/native_functions.yaml +++ b/yaml/native/native_functions.yaml @@ -9482,3 +9482,20 @@ # Warning! For whatever reason, the inplace sparse addmm is NON # broadcasting SparseXPU: s_addmm_sparse_dense_xpu_ + +- func: mm(Tensor self, Tensor mat2) -> Tensor + structured_delegate: mm.out + variants: function, method + dispatch: + SparseXPU: _sparse_mm + tags: core + +- func: mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + SparseXPU: _sparse_mm_out + +- func: _sparse_sparse_matmul(Tensor self, Tensor other) -> Tensor + dispatch: + SparseXPU: sparse_sparse_matmul_xpu + autogen: _sparse_sparse_matmul.out \ No newline at end of file From b13168f20002a5135ece1f2a190dc7bf37d34dff Mon Sep 17 00:00:00 2001 From: jenniew Date: Mon, 1 Dec 2025 12:09:45 -0800 Subject: [PATCH 16/19] include matmul --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 8d1464d9db..55a2c96e92 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -5,6 +5,7 @@ #include #else #include +#include #endif #include From a2a5280c7423a8d5cb4b8bfa47b92379cc7754f8 Mon Sep 17 00:00:00 2001 From: jenniew Date: Mon, 1 Dec 2025 18:43:59 -0800 Subject: [PATCH 17/19] update --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 55a2c96e92..2591c08fe4 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -121,7 +121,7 @@ Tensor sparse_sparse_matmul_xpu(const Tensor& mat1_, const Tensor& mat2_) { Tensor output_dense = at::matmul(mat1_dense, mat2_dense); // convert back to sparse - Tensor output_sparse = output_dense._to_sparse(mat1.layout()); + Tensor output_sparse = output_dense._to_sparse(mat1_.layout()); return output_sparse; From 2af45810c6a37331eeb0fca2d16e83b383b2e324 Mon Sep 17 00:00:00 2001 From: "Cui, Yifeng" Date: Wed, 3 Dec 2025 21:39:04 +0800 Subject: [PATCH 18/19] Update src/ATen/native/sparse/xpu/SparseTensorMath.cpp to fix lint error --- src/ATen/native/sparse/xpu/SparseTensorMath.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp index 2591c08fe4..d2f610332d 100644 --- a/src/ATen/native/sparse/xpu/SparseTensorMath.cpp +++ b/src/ATen/native/sparse/xpu/SparseTensorMath.cpp @@ -114,7 +114,7 @@ Tensor sparse_sparse_matmul_xpu(const Tensor& mat1_, const Tensor& mat2_) { TORCH_CHECK(mat1_.scalar_type() == mat2_.scalar_type(), "mat1 dtype ", mat1_.scalar_type(), " does not match mat2 dtype ", mat2_.scalar_type()); - + // convert to dense Tensor mat1_dense = mat1_._to_dense(std::nullopt, std::nullopt); Tensor mat2_dense = mat2_._to_dense(std::nullopt, std::nullopt); From f448be1d62b92bf657b8ac6aa9c7cfa4c13d18fb Mon Sep 17 00:00:00 2001 From: jenniew Date: Thu, 4 Dec 2025 21:29:04 -0800 Subject: [PATCH 19/19] enable tests --- test/xpu/test_sparse_xpu.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/test/xpu/test_sparse_xpu.py b/test/xpu/test_sparse_xpu.py index 1ca7e1b69c..d973e5e9ea 100644 --- a/test/xpu/test_sparse_xpu.py +++ b/test/xpu/test_sparse_xpu.py @@ -2121,7 +2121,7 @@ def fn(S, D1, D2, beta=beta, alpha=alpha): @unittest.skipIf( TEST_WITH_CROSSREF, "generator unsupported triggers assertion error" ) - @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/2211") + @skipXPUIf(False, "https://github.com/intel/torch-xpu-ops/issues/2211") def test_sparse_mm(self, device, dtype, coalesced): def test_shape(d1, d2, d3, nnz, transposed): if transposed: @@ -2197,7 +2197,7 @@ def test_shape(sparse_dims, nnz, with_shape): @coalescedonoff @dtypes(torch.double) @dtypesIfMPS(torch.float32) - @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/2211") + @skipXPUIf(False, "https://github.com/intel/torch-xpu-ops/issues/2211") def test_dsmm(self, device, dtype, coalesced): def test_shape(di, dj, dk, nnz): x = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced)[0] @@ -3363,7 +3363,7 @@ def test_asin_arcsin(self, device, dtype, coalesced): @expectedFailureMPS @dtypes(torch.double) @dtypesIfMPS(torch.float32) - @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/2211") + @skipXPUIf(False, "https://github.com/intel/torch-xpu-ops/issues/2211") def test_mv(self, device, dtype, coalesced): def test_shape(di, dj, dk, nnz): x, _, _ = self._gen_sparse(2, nnz, [di, dj], dtype, device, coalesced) @@ -4759,7 +4759,7 @@ def test_log_softmax_float(self, device, dtype): *[torch.half], *[torch.bfloat16], torch.complex64, *[torch.complex128] ) ) - @skipXPUIf(True, "https://github.com/intel/torch-xpu-ops/issues/2211") + @skipXPUIf(False, "https://github.com/intel/torch-xpu-ops/issues/2211") @unittest.skipIf(TEST_WITH_CROSSREF, "not working with fake tensor") @precisionOverride( {