Skip to content

[mlir][sparse] Migration to sparse_tensor.print #83377

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 2 commits into from
Feb 29, 2024
Merged

[mlir][sparse] Migration to sparse_tensor.print #83377

merged 2 commits into from
Feb 29, 2024

Conversation

yinying-lisa-li
Copy link
Contributor

Continuous efforts following #83357.

@yinying-lisa-li yinying-lisa-li marked this pull request as ready for review February 29, 2024 18:38
@llvmbot llvmbot added mlir:sparse Sparse compiler in MLIR mlir labels Feb 29, 2024
@llvmbot
Copy link
Member

llvmbot commented Feb 29, 2024 •
edited
Loading

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Yinying Li (yinying-lisa-li)

Changes

Continuous efforts following #83357.

Patch is 28.24 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/83377.diff

4 Files Affected:

  • (modified) mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir (+9-14)
  • (modified) mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir (+138-99)
  • (modified) mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir (+41-43)
  • (modified) mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir (+71-43)
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
index 228d4e5f6f8a1a..e1f062121b12f9 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
@@ -10,7 +10,7 @@
 // DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
 // DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
-// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
@@ -28,7 +28,7 @@
 }>
 
 module {
-  func.func @entry() {
+  func.func @main() {
     %c0 = arith.constant 0 : index
     %f0 = arith.constant 0.0 : f64
     %d = arith.constant dense<[[ 1.0,  2.0,  3.0,  4.0 ],
@@ -39,19 +39,14 @@ module {
     %s = sparse_tensor.convert %d : tensor<5x4xf64> to tensor<5x4xf64, #CSR_hi>
 
     //
-    // CHECK:      ( 0, 4, 4, 8, 8, 9, 9, 13 )
-    // CHECK-NEXT: ( 0, 1, 2, 3, 0, 1, 2, 3, 2, 0, 1, 2, 3, 0, 1, 2, 3 )
-    // CHECK-NEXT: ( 1, 2, 3, 4, 5, 6, 7, 8, 5.5, 9, 10, 11, 12, 13, 14, 15, 16 )
+    // CHECK:   ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 17
+    // CHECK-NEXT: pos[1] : ( 0, 4, 4, 8, 8, 9, 9, 13
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 2, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8, 5.5, 9, 10, 11, 12, 13, 14, 15, 16
+    // CHECK-NEXT: ----
     //
-    %pos = sparse_tensor.positions %s {level = 1 : index } : tensor<5x4xf64, #CSR_hi> to memref<?xindex>
-    %vecp = vector.transfer_read %pos[%c0], %c0 : memref<?xindex>, vector<8xindex>
-    vector.print %vecp : vector<8xindex>
-    %crd = sparse_tensor.coordinates %s {level = 1 : index } : tensor<5x4xf64, #CSR_hi> to memref<?xindex>
-    %vecc = vector.transfer_read %crd[%c0], %c0 : memref<?xindex>, vector<17xindex>
-    vector.print %vecc : vector<17xindex>
-    %val = sparse_tensor.values %s : tensor<5x4xf64, #CSR_hi> to memref<?xf64>
-    %vecv = vector.transfer_read %val[%c0], %f0 : memref<?xf64>, vector<17xf64>
-    vector.print %vecv : vector<17xf64>
+    sparse_tensor.print %s : tensor<5x4xf64, #CSR_hi>
 
     // Release the resources.
     bufferization.dealloc_tensor %s: tensor<5x4xf64, #CSR_hi>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
index fa0dbac269b926..863e1c62370e32 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
@@ -10,7 +10,7 @@
 // DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
 // DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
-// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
@@ -90,7 +90,7 @@ module {
   //
   // Main driver.
   //
-  func.func @entry() {
+  func.func @main() {
     %c0 = arith.constant 0 : index
 
     // Initialize various matrices, dense for stress testing,
@@ -140,33 +140,94 @@ module {
     %b4 = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #DCSR>
 
     //
-    // Sanity check on stored entries before going into the computations.
-    //
-    // CHECK:      32
-    // CHECK-NEXT: 32
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 4
-    // CHECK-NEXT: 32
-    // CHECK-NEXT: 32
-    // CHECK-NEXT: 8
-    // CHECK-NEXT: 8
-    //
-    %noea1 = sparse_tensor.number_of_entries %a1 : tensor<4x8xf64, #CSR>
-    %noea2 = sparse_tensor.number_of_entries %a2 : tensor<4x8xf64, #DCSR>
-    %noea3 = sparse_tensor.number_of_entries %a3 : tensor<4x8xf64, #CSR>
-    %noea4 = sparse_tensor.number_of_entries %a4 : tensor<4x8xf64, #DCSR>
-    %noeb1 = sparse_tensor.number_of_entries %b1 : tensor<8x4xf64, #CSR>
-    %noeb2 = sparse_tensor.number_of_entries %b2 : tensor<8x4xf64, #DCSR>
-    %noeb3 = sparse_tensor.number_of_entries %b3 : tensor<8x4xf64, #CSR>
-    %noeb4 = sparse_tensor.number_of_entries %b4 : tensor<8x4xf64, #DCSR>
-    vector.print %noea1 : index
-    vector.print %noea2 : index
-    vector.print %noea3 : index
-    vector.print %noea4 : index
-    vector.print %noeb1 : index
-    vector.print %noeb2 : index
-    vector.print %noeb3 : index
-    vector.print %noeb4 : index
+    // Sanity check before going into the computations.
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: pos[1] : ( 0, 8, 16, 24, 32
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7
+    // CHECK-NEXT: values : ( 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3, 1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %a1 : tensor<4x8xf64, #CSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: pos[0] : ( 0, 4
+    // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 8, 16, 24, 32
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7, 0, 1, 2, 3, 4, 5, 6, 7
+    // CHECK-NEXT: values : ( 1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1, 1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2, 1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3, 1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %a2 : tensor<4x8xf64, #DCSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: pos[1] : ( 0, 2, 2, 3, 4
+    // CHECK-NEXT: crd[1] : ( 1, 5, 1, 7
+    // CHECK-NEXT: values : ( 2.1, 6.1, 2.3, 1
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %a3 : tensor<4x8xf64, #CSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: pos[0] : ( 0, 3
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 2, 3, 4
+    // CHECK-NEXT: crd[1] : ( 1, 5, 1, 7
+    // CHECK-NEXT: values : ( 2.1, 6.1, 2.3, 1
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %a4 : tensor<4x8xf64, #DCSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12, 16, 20, 24, 28, 32
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 10.1, 11.1, 12.1, 13.1, 10.2, 11.2, 12.2, 13.2, 10.3, 11.3, 12.3, 13.3, 10.4, 11.4, 12.4, 13.4, 10.5, 11.5, 12.5, 13.5, 10.6, 11.6, 12.6, 13.6, 10.7, 11.7, 12.7, 13.7, 10.8, 11.8, 12.8, 13.8
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %b1 : tensor<8x4xf64, #CSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: pos[0] : ( 0, 8
+    // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 4, 5, 6, 7
+    // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12, 16, 20, 24, 28, 32
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 10.1, 11.1, 12.1, 13.1, 10.2, 11.2, 12.2, 13.2, 10.3, 11.3, 12.3, 13.3, 10.4, 11.4, 12.4, 13.4, 10.5, 11.5, 12.5, 13.5, 10.6, 11.6, 12.6, 13.6, 10.7, 11.7, 12.7, 13.7, 10.8, 11.8, 12.8, 13.8
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %b2 : tensor<8x4xf64, #DCSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: pos[1] : ( 0, 1, 2, 3, 4, 4, 5, 6, 8
+    // CHECK-NEXT: crd[1] : ( 3, 2, 1, 0, 1, 2, 2, 3
+    // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %b3 : tensor<8x4xf64, #CSR>
+
+    //
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: pos[0] : ( 0, 7
+    // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3, 5, 6, 7
+    // CHECK-NEXT: pos[1] : ( 0, 1, 2, 3, 4, 5, 6, 8
+    // CHECK-NEXT: crd[1] : ( 3, 2, 1, 0, 1, 2, 2, 3
+    // CHECK-NEXT: values : ( 1, 2, 3, 4, 5, 6, 7, 8
+    // CHECK-NEXT: ----
+    //
+    sparse_tensor.print %b4 : tensor<8x4xf64, #DCSR>
 
     // Call kernels with dense.
     %0 = call @matmul1(%da, %db, %zero)
@@ -208,24 +269,26 @@ module {
     call @printMemrefF64(%u0) : (tensor<*xf64>) -> ()
 
     //
-    // CHECK:      {{\[}}[388.76,   425.56,   462.36,   499.16],
-    // CHECK-NEXT: [397.12,   434.72,   472.32,   509.92],
-    // CHECK-NEXT: [405.48,   443.88,   482.28,   520.68],
-    // CHECK-NEXT: [413.84,   453.04,   492.24,   531.44]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 16
+    // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12, 16
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 388.76, 425.56, 462.36, 499.16, 397.12, 434.72, 472.32, 509.92, 405.48, 443.88, 482.28, 520.68, 413.84, 453.04, 492.24, 531.44
+    // CHECK-NEXT: ----
     //
-    %c1 = sparse_tensor.convert %1 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
-    %c1u = tensor.cast %c1 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c1u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %1 : tensor<4x4xf64, #CSR>
 
     //
-    // CHECK:      {{\[}}[388.76,   425.56,   462.36,   499.16],
-    // CHECK-NEXT: [397.12,   434.72,   472.32,   509.92],
-    // CHECK-NEXT: [405.48,   443.88,   482.28,   520.68],
-    // CHECK-NEXT: [413.84,   453.04,   492.24,   531.44]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 16
+    // CHECK-NEXT: pos[0] : ( 0, 4
+    // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12, 16
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 388.76, 425.56, 462.36, 499.16, 397.12, 434.72, 472.32, 509.92, 405.48, 443.88, 482.28, 520.68, 413.84, 453.04, 492.24, 531.44
+    // CHECK-NEXT: ----
     //
-    %c2 = sparse_tensor.convert %2 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
-    %c2u = tensor.cast %c2 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c2u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %2 : tensor<4x4xf64, #DCSR>
 
     //
     // CHECK:      {{\[}}[86.08,   94.28,   102.48,   110.68],
@@ -237,24 +300,26 @@ module {
     call @printMemrefF64(%u3) : (tensor<*xf64>) -> ()
 
     //
-    // CHECK:      {{\[}}[86.08,   94.28,   102.48,   110.68],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [23.46,   25.76,   28.06,   30.36],
-    // CHECK-NEXT: [10.8,   11.8,   12.8,   13.8]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: pos[1] : ( 0, 4, 4, 8, 12
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 86.08, 94.28, 102.48, 110.68, 23.46, 25.76, 28.06, 30.36, 10.8, 11.8, 12.8, 13.8
+    // CHECK-NEXT: ----
     //
-    %c4 = sparse_tensor.convert %4 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
-    %c4u = tensor.cast %c4 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c4u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %4 : tensor<4x4xf64, #CSR>
 
     //
-    // CHECK:      {{\[}}[86.08,   94.28,   102.48,   110.68],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [23.46,   25.76,   28.06,   30.36],
-    // CHECK-NEXT: [10.8,   11.8,   12.8,   13.8]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: pos[0] : ( 0, 3
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12
+    // CHECK-NEXT: crd[1] : ( 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3
+    // CHECK-NEXT: values : ( 86.08, 94.28, 102.48, 110.68, 23.46, 25.76, 28.06, 30.36, 10.8, 11.8, 12.8, 13.8
+    // CHECK-NEXT: ----
     //
-    %c5 = sparse_tensor.convert %5 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
-    %c5u = tensor.cast %c5 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c5u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %5 : tensor<4x4xf64, #DCSR>
 
     //
     // CHECK:      {{\[}}[0,   30.5,   4.2,   0],
@@ -266,46 +331,26 @@ module {
     call @printMemrefF64(%u6) : (tensor<*xf64>) -> ()
 
     //
-    // CHECK:      {{\[}}[0,   30.5,   4.2,   0],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [0,   0,   4.6,   0],
-    // CHECK-NEXT: [0,   0,   7,   8]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: pos[1] : ( 0, 2, 2, 3, 5
+    // CHECK-NEXT: crd[1] : ( 1, 2, 2, 2, 3
+    // CHECK-NEXT: values : ( 30.5, 4.2, 4.6, 7, 8
+    // CHECK-NEXT: ----
     //
-    %c7 = sparse_tensor.convert %7 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
-    %c7u = tensor.cast %c7 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c7u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %7 : tensor<4x4xf64, #CSR>
 
     //
-    // CHECK:      {{\[}}[0,   30.5,   4.2,   0],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [0,   0,   4.6,   0],
-    // CHECK-NEXT: [0,   0,   7,   8]]
-    //
-    %c8 = sparse_tensor.convert %8 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
-    %c8u = tensor.cast %c8 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c8u) : (tensor<*xf64>) -> ()
-
-    //
-    // Sanity check on nonzeros.
-    //
-    // CHECK: [30.5,  4.2,  4.6,  7,  8{{.*}}]
-    // CHECK: [30.5,  4.2,  4.6,  7,  8{{.*}}]
-    //
-    %val7 = sparse_tensor.values %7 : tensor<4x4xf64, #CSR> to memref<?xf64>
-    %val8 = sparse_tensor.values %8 : tensor<4x4xf64, #DCSR> to memref<?xf64>
-    call @printMemref1dF64(%val7) : (memref<?xf64>) -> ()
-    call @printMemref1dF64(%val8) : (memref<?xf64>) -> ()
-
-    //
-    // Sanity check on stored entries after the computations.
-    //
-    // CHECK-NEXT: 5
-    // CHECK-NEXT: 5
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: pos[0] : ( 0, 3
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 2, 3, 5
+    // CHECK-NEXT: crd[1] : ( 1, 2, 2, 2, 3
+    // CHECK-NEXT: values : ( 30.5, 4.2, 4.6, 7, 8
+    // CHECK-NEXT: ----
     //
-    %noe7 = sparse_tensor.number_of_entries %7 : tensor<4x4xf64, #CSR>
-    %noe8 = sparse_tensor.number_of_entries %8 : tensor<4x4xf64, #DCSR>
-    vector.print %noe7 : index
-    vector.print %noe8 : index
+    sparse_tensor.print %8 : tensor<4x4xf64, #DCSR>
 
     // Release the resources.
     bufferization.dealloc_tensor %a1 : tensor<4x8xf64, #CSR>
@@ -316,12 +361,6 @@ module {
     bufferization.dealloc_tensor %b2 : tensor<8x4xf64, #DCSR>
     bufferization.dealloc_tensor %b3 : tensor<8x4xf64, #CSR>
     bufferization.dealloc_tensor %b4 : tensor<8x4xf64, #DCSR>
-    bufferization.dealloc_tensor %c1 : tensor<4x4xf64>
-    bufferization.dealloc_tensor %c2 : tensor<4x4xf64>
-    bufferization.dealloc_tensor %c4 : tensor<4x4xf64>
-    bufferization.dealloc_tensor %c5 : tensor<4x4xf64>
-    bufferization.dealloc_tensor %c7 : tensor<4x4xf64>
-    bufferization.dealloc_tensor %c8 : tensor<4x4xf64>
     bufferization.dealloc_tensor %0 : tensor<4x4xf64>
     bufferization.dealloc_tensor %1 : tensor<4x4xf64, #CSR>
     bufferization.dealloc_tensor %2 : tensor<4x4xf64, #DCSR>
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
index 96c8a30ade8e42..a7184c044569ca 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
@@ -10,7 +10,7 @@
 // DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
 // DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
 // DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
-// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
 // DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
 // DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs}
 //
@@ -132,7 +132,7 @@ module {
   //
   // Main driver.
   //
-  func.func @entry() {
+  func.func @main() {
     %c_0 = arith.constant 0 : index
     %c_1 = arith.constant 1 : index
     %c_2 = arith.constant 2 : index
@@ -170,14 +170,16 @@ module {
 
     // DCSR test
     //
-    // CHECK:       [0,   30.5,   4.2,   0],
-    // CHECK-NEXT:  [0,   0,   0,   0],
-    // CHECK-NEXT:  [0,   0,   4.6,   0],
-    // CHECK-NEXT:  [0,   0,   7,   8]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: pos[0] : ( 0, 3
+    // CHECK-NEXT: crd[0] : ( 0, 2, 3
+    // CHECK-NEXT: pos[1] : ( 0, 2, 3, 5
+    // CHECK-NEXT: crd[1] : ( 1, 2, 2, 2, 3
+    // CHECK-NEXT: values : ( 30.5, 4.2, 4.6, 7, 8
+    // CHECK-NEXT: ----
     //
-    %c2 = sparse_tensor.convert %2 : tensor<4x4xf64, #DCSR> to tensor<4x4xf64>
-    %c2u = tensor.cast %c2 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c2u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %2 : tensor<4x4xf64, #DCSR>
 
     %t1 = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #CSR>
     %a1 = tensor.extract_slice %t1[0, 0][4, 8][1, 1] : tensor<8x8xf64, #CSR> to tensor<4x8xf64, #CSR_SLICE>
@@ -188,63 +190,64 @@ module {
 
     // CSR test
     //
-    // CHECK:       [0,   30.5,   4.2,   0],
-    // CHECK-NEXT:  [0,   0,   0,   0],
-    // CHECK-NEXT:  [0,   0,   4.6,   0],
-    // CHECK-NEXT:  [0,   0,   7,   8]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: pos[1] : ( 0, 2, 2, 3, 5
+    // CHECK-NEXT: crd[1] : ( 1, 2, 2, 2, 3
+    // CHECK-NEXT: values : ( 30.5, 4.2, 4.6, 7, 8
+    // CHECK-NEXT: ----
     //
-    %c3 = sparse_tensor.convert %3 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
-    %c3u = tensor.cast %c3 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c3u) : (tensor<*xf64>) -> ()
+    sparse_tensor.print %3 : tensor<4x4xf64, #CSR>
+
 
     // slice x slice
     //
-    // CHECK:      [2.3,   0,   0,   0],
-    // CHECK-NEXT: [6.9,   0,   0,   0],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [12.6,   0,   0,   0]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: pos[1] : ( 0, 1, 2, 2, 3
+    // CHECK-NEXT: crd[1] : ( 0, 0, 0
+    // CHECK-NEXT: values : ( 2.3, 6.9, 12.6
+    // CHECK-NEXT: ----
     //
+    sparse_tensor.print %4 : tensor<4x4xf64, #CSR>
+
     %s1 = tensor.extract_slice %tmp[0, 1][4, 4][2, 1] : tensor<8x8xf64, #DCSR> to tensor<4x4xf64, #DCSR_SLICE_1>
     %s2 = tensor.extract_slice %b1[0, 0][4, 4][2, 1] : tensor<8x4xf64, #CSR> to tensor<4x4xf64, #CSR_SLICE_1>
     %4 = call @matmul1(%s2, %s1)
        : (tensor<4x4xf64, #CSR_SLICE_1>,
           tensor<4x4xf64, #DCSR_SLICE_1>) -> tensor<4x4xf64, #CSR>
-    %c4 = sparse_tensor.convert %4 : tensor<4x4xf64, #CSR> to tensor<4x4xf64>
-    %c4u = tensor.cast %c4 : tensor<4x4xf64> to tensor<*xf64>
-    call @printMemrefF64(%c4u) : (tensor<*xf64>) -> ()
 
     // slice coo x slice coo
     //
-    // CHECK:      [2.3,   0,   0,   0],
-    // CHECK-NEXT: [6.9,   0,   0,   0],
-    // CHECK-NEXT: [0,   0,   0,   0],
-    // CHECK-NEXT: [12.6,   0,   0,   0]]
+    // CHECK:      ---- Sparse Tensor ----
+    // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: pos[0] : ( 0, 3
+    // CHECK-NEXT: crd[0] : ( 0, 0, 1, 0, 3, 0
+    // CHECK-NEXT: values : ( 2.3, 6.9, 12.6
+    // CHECK-NEXT: ----
     //
+    sparse_tensor.print %o_coo : tensor<4x4xf64, #COO>
     %t1_coo = sparse_tensor.convert %sa : tensor<8x8xf64> to tensor<8x8xf64, #COO>
     %b1_coo = sparse_tensor.convert %sb : tensor<8x4xf64> to tensor<8x4xf64, #COO>
     %s2_coo = tensor.extract_slice %b1_coo[0, 0][4, 4][2, 1] : tensor<8x4xf64, #COO> to tensor<4x4xf64, #COO_SLICE_1>
     %s1_coo = te...
[truncated]

aartbik
aartbik approved these changes Feb 29, 2024
View reviewed changes
@yinying-lisa-li yinying-lisa-li merged commit 1ca65dd into llvm:main Feb 29, 2024
@yinying-lisa-li yinying-lisa-li deleted the migration branch February 29, 2024 19:14
joker-eph added a commit that referenced this pull request Feb 29, 2024
@joker-eph
Revert "[mlir][sparse] Migration to sparse_tensor.print (#83377)"
43a5e5d 
This reverts commit 1ca65dd.
@joker-eph joker-eph mentioned this pull request Feb 29, 2024
Merged
joker-eph added a commit that referenced this pull request Feb 29, 2024
@joker-eph
Revert "[mlir][sparse] Migration to sparse_tensor.print" (#83499)
71eead5 
Reverts #83377

The test does not pass on the bot.
@yinying-lisa-li yinying-lisa-li restored the migration branch February 29, 2024 23:47
@yinying-lisa-li yinying-lisa-li mentioned this pull request Mar 1, 2024
Merged
yinying-lisa-li added a commit that referenced this pull request Mar 1, 2024
@yinying-lisa-li
[mlir][sparse] Migration to sparse_tensor.print (#83506)
5899599 
Continuous efforts #83357. Previously reverted #83377.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@aartbik aartbik aartbik approved these changes

@PeimingLiu PeimingLiu Awaiting requested review from PeimingLiu PeimingLiu is a code owner

@matthias-springer matthias-springer Awaiting requested review from matthias-springer matthias-springer is a code owner

Assignees
No one assigned
Labels
mlir:sparse Sparse compiler in MLIR mlir
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@yinying-lisa-li @llvmbot @aartbik