Refine the styles and add test cases.

intel · Jan 18, 2025 · 512ef17 · 512ef17
1 parent 0e638ea
commit 512ef17
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 63 deletions.
diff --git a/src/ATen/native/sparse/xpu/sycl/SparseCsrTensorMathKernels.cpp b/src/ATen/native/sparse/xpu/sycl/SparseCsrTensorMathKernels.cpp
@@ -89,57 +89,6 @@ void convert_indices_from_coo_to_csr_xpu(
     sycl_kernel_submit(global_range, local_range, getCurrentSYCLQueue(), functor);
 }
 
-
-// //With reference to coalesce_sparse_kernel
-// void convert_indices_from_coo_to_csr_xpu_kernel(
-//     const Tensor& input,
-//     const int64_t size,
-//     const bool out_int32,
-//     const Tensor& result){
-
-//     int64_t numel = input.numel();
-//     if (numel == 0) {
-//         result.zero_();
-//         return;
-//     }
-
-//     //How to define the global_range and local_range? --numel
-//     int64_t wgroup_size = 64;
-//     int64_t ngroups = (numel + wgroup_size - 1) / wgroup_size;
-//     sycl::range<1> global_range(ngroups * wgroup_size);
-//     sycl::range<1> local_range(wgroup_size);
-
-//     if (out_int32) {
-//         AT_DISPATCH_INTEGRAL_TYPES(
-//             input.scalar_type(), 
-//             "convert_indices_from_coo_to_csr_xpu", [&] {
-//             const scalar_t* data_in = input.data_ptr<scalar_t>();
-//             int* data_out = result.data_ptr<int>();
-//             auto functor = convertIndicesFromCooToCsrXPUFunctor<scalar_t, int>(
-//                 numel,
-//                 data_in,
-//                 data_out, 
-//                 size);
-//             sycl_kernel_submit(
-//             global_range, local_range, getCurrentSYCLQueue(), functor);
-//         });
-//     } else {
-//         AT_DISPATCH_INTEGRAL_TYPES(
-//             input.scalar_type(), 
-//             "convert_indices_from_coo_to_csr_xpu", [&] {
-//             const scalar_t* data_in = input.data_ptr<scalar_t>();
-//             int64_t* data_out = result.data_ptr<int64_t>();
-//             auto functor = convertIndicesFromCooToCsrXPUFunctor<scalar_t, int64_t>(
-//                 numel,
-//                 data_in,
-//                 data_out, 
-//                 size);
-//             sycl_kernel_submit(
-//             global_range, local_range, getCurrentSYCLQueue(), functor);
-//         });
-//     }
-// }
-
 TORCH_IMPL_FUNC(_convert_indices_from_coo_to_csr_structured_xpu) (
     const Tensor& input, 
     const int64_t size, 
@@ -159,17 +108,7 @@ TORCH_IMPL_FUNC(_convert_indices_from_coo_to_csr_structured_xpu) (
               result, input, size);
         });
   }
-
-    // convert_indices_from_coo_to_csr_xpu_kernel(
-    //     input,
-    //     size,
-    //     out_int32,
-    //     result);
-}   
+}
 
 } // namespace at::native
 
-namespace at::native::xpu {
-} // namespace at::native::xpu
-
-
diff --git a/test/xpu/skip_list_common.py b/test/xpu/skip_list_common.py
@@ -35,6 +35,19 @@
         "test_errors_sparse_mul_layout3_xpu",
         "test_out_requires_grad_error_sparse_sampled_addmm_xpu_complex64",
         "test_out_requires_grad_error_sparse_sampled_addmm_xpu_float32",
+        # Similar error as above for to_sparse_csr() operator
+        "test_sparse_csr_from_dense_xpu_bfloat16",
+        "test_sparse_csr_from_dense_xpu_bool",
+        "test_sparse_csr_from_dense_xpu_complex128",
+        "test_sparse_csr_from_dense_xpu_complex64",
+        "test_sparse_csr_from_dense_xpu_float16",
+        "test_sparse_csr_from_dense_xpu_float32",
+        "test_sparse_csr_from_dense_xpu_float64",
+        "test_sparse_csr_from_dense_xpu_int16",
+        "test_sparse_csr_from_dense_xpu_int32",
+        "test_sparse_csr_from_dense_xpu_int64",
+        "test_sparse_csr_from_dense_xpu_int8",
+        "test_sparse_csr_from_dense_xpu_uint8",
 
         # OneDNN issues, https://github.com/intel/torch-xpu-ops/issues/253
         # RuntimeError: Long is not supported in oneDNN!

diff --git a/test/xpu/test_sparse_xpu.py b/test/xpu/test_sparse_xpu.py
@@ -1,7 +1,10 @@
 # Owner(s): ["module: intel"]
-from torch.testing._internal.common_device_type import instantiate_device_type_tests
+from torch.testing._internal.common_device_type import instantiate_device_type_tests, dtypes
+from torch.testing._internal.common_dtype import all_types_and_complex_and
 from torch.testing._internal.common_utils import run_tests
 
+import torch
+
 try:
     from xpu_test_utils import XPUPatchForImport
 except Exception as e:
@@ -10,6 +13,29 @@
 with XPUPatchForImport(False):
     from test_sparse import TestSparse
 
+    #@skipIfTorchDynamo()
+    @dtypes(*all_types_and_complex_and(torch.half, torch.bool, torch.bfloat16))
+    def sparse_csr_from_dense(self, device, dtype):
+        dense = torch.tensor([[4, 5, 0], [0, 0, 0], [1, 0, 0]], dtype=dtype, device=device)
+        sparse = dense.to_sparse_csr()
+        self.assertEqual(torch.tensor([0, 2, 2, 3], dtype=torch.int64), sparse.crow_indices())
+        self.assertEqual(torch.tensor([0, 1, 0], dtype=torch.int64), sparse.col_indices())
+        self.assertEqual(torch.tensor([4, 5, 1], dtype=dtype), sparse.values())
+
+        dense = torch.tensor([[0, 0, 0], [0, 0, 1], [1, 0, 0]], dtype=dtype, device=device)
+        sparse = dense.to_sparse_csr()
+        self.assertEqual(torch.tensor([0, 0, 1, 2], dtype=torch.int64), sparse.crow_indices())
+        self.assertEqual(torch.tensor([2, 0], dtype=torch.int64), sparse.col_indices())
+        self.assertEqual(torch.tensor([1, 1], dtype=dtype), sparse.values())
+
+        dense = torch.tensor([[2, 2, 2], [2, 2, 2], [2, 2, 2]], dtype=dtype, device=device)
+        sparse = dense.to_sparse_csr()
+        self.assertEqual(torch.tensor([0, 3, 6, 9], dtype=torch.int64), sparse.crow_indices())
+        self.assertEqual(torch.tensor([0, 1, 2] * 3, dtype=torch.int64), sparse.col_indices())
+        self.assertEqual(torch.tensor([2] * 9, dtype=dtype), sparse.values())
+
+    TestSparse.test_sparse_csr_from_dense = sparse_csr_from_dense
+
 instantiate_device_type_tests(TestSparse, globals(), only_for="xpu", allow_xpu=True)
 
 if __name__ == "__main__":