[Mlir-commits] [mlir] [mlir][sparse] add dim/lvl information to sparse_tensor.print (PR #83913)

Aart Bik llvmlistbot at llvm.org
Mon Mar 4 13:50:17 PST 2024


https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/83913

More information is more testing!
Also adjusts already migrated integration tests

>From 7d18612f1e8b974832cd5e9b7c6337152982d54a Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 4 Mar 2024 13:47:10 -0800
Subject: [PATCH] [mlir][sparse] add dim/lvl information to sparse_tensor.print

More information is more testing!
Also adjusts already migrated integration tests
---
 .../Transforms/SparseTensorRewriting.cpp      | 28 +++++++++++++++++++
 .../Dialect/SparseTensor/CPU/block.mlir       |  6 ++++
 .../SparseTensor/CPU/block_majors.mlir        |  8 ++++++
 .../SparseTensor/CPU/dense_output.mlir        |  2 ++
 .../SparseTensor/CPU/dense_output_bf16.mlir   |  2 ++
 .../SparseTensor/CPU/dense_output_f16.mlir    |  2 ++
 .../SparseTensor/CPU/sparse_insert_3d.mlir    |  8 ++++++
 .../SparseTensor/CPU/sparse_loose.mlir        |  2 ++
 .../SparseTensor/CPU/sparse_matmul.mlir       | 28 +++++++++++++++++++
 .../SparseTensor/CPU/sparse_matmul_slice.mlir | 10 +++++++
 .../SparseTensor/CPU/sparse_matrix_ops.mlir   | 12 ++++++++
 .../SparseTensor/CPU/sparse_print.mlir        | 24 ++++++++++++++++
 .../SparseTensor/CPU/sparse_re_im.mlir        |  4 +++
 .../GPU/CUDA/sparse-gemm-lib.mlir             |  2 ++
 .../GPU/CUDA/sparse-sampled-matmul-lib.mlir   |  4 +++
 .../GPU/CUDA/sparse-sddmm-lib.mlir            |  4 +++
 16 files changed, 146 insertions(+)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 6ff21468e05764..ddac602b9a43e9 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -616,6 +616,11 @@ struct PrintRewriter : public OpRewritePattern<PrintOp> {
     rewriter.create<vector::PrintOp>(
         loc, rewriter.getStringAttr("---- Sparse Tensor ----\nnse = "));
     rewriter.create<vector::PrintOp>(loc, nse);
+    // Print run-time contents for dim/lvl sizes.
+    rewriter.create<vector::PrintOp>(loc, rewriter.getStringAttr("dim = "));
+    printSizes(rewriter, loc, tensor, stt.getDimRank(), /*isDim=*/true);
+    rewriter.create<vector::PrintOp>(loc, rewriter.getStringAttr("lvl = "));
+    printSizes(rewriter, loc, tensor, stt.getLvlRank(), /*isDim=*/false);
     // Use the "codegen" foreach loop construct to iterate over
     // all typical sparse tensor components for printing.
     foreachFieldAndTypeInSparseTensor(stt, [&rewriter, &loc, &tensor,
@@ -693,6 +698,29 @@ struct PrintRewriter : public OpRewritePattern<PrintOp> {
     rewriter.create<vector::PrintOp>(loc, vector::PrintPunctuation::Close);
     rewriter.create<vector::PrintOp>(loc, vector::PrintPunctuation::NewLine);
   }
+
+  // Helper method to print run-time lvl/dim sizes.
+  static void printSizes(PatternRewriter &rewriter, Location loc, Value tensor,
+                         unsigned size, bool isDim) {
+    // Open bracket.
+    rewriter.create<vector::PrintOp>(loc, vector::PrintPunctuation::Open);
+    // Print unrolled contents (dimop requires constant value).
+    for (unsigned i = 0; i < size; i++) {
+      auto idx = constantIndex(rewriter, loc, i);
+      Value val;
+      if (isDim)
+        val = rewriter.create<tensor::DimOp>(loc, tensor, idx);
+      else
+        val = rewriter.create<LvlOp>(loc, tensor, idx);
+      rewriter.create<vector::PrintOp>(
+          loc, val,
+          i != size - 1 ? vector::PrintPunctuation::Comma
+                        : vector::PrintPunctuation::NoPunctuation);
+    }
+    // Close bracket and end of line.
+    rewriter.create<vector::PrintOp>(loc, vector::PrintPunctuation::Close);
+    rewriter.create<vector::PrintOp>(loc, vector::PrintPunctuation::NewLine);
+  }
 };
 
 /// Sparse rewriting rule for sparse-to-sparse reshape operator.
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir
index 1184d407541b6f..f79e7e68f382d3 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/block.mlir
@@ -91,6 +91,8 @@ module {
 
     // CHECK:   ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 6 )
+    // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
     // CHECK-NEXT: crd[1] : ( 0, 2, 1,
     // CHECK-NEXT: values : ( 1, 2, 0, 3, 4, 0, 0, 5, 6, 7, 8, 0,
@@ -99,6 +101,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 2, 3, 2, 2 )
+    // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
     // CHECK-NEXT: crd[1] : ( 0, 2, 1
     // CHECK-NEXT: values : ( 1, 2, 0, 3, 4, 0, 0, 5, 6, 7, 8, 0,
@@ -109,6 +113,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 6 )
+    // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
     // CHECK-NEXT: crd[1] : ( 0, 2, 1,
     // CHECK-NEXT: values : ( 3, 6, 0, 9, 12, 0, 0, 15, 18, 21, 24, 0,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/block_majors.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/block_majors.mlir
index f8e83b5019679f..3534e7d15207c4 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/block_majors.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/block_majors.mlir
@@ -106,6 +106,8 @@ module {
   //
   // CHECK: ---- Sparse Tensor ----
   // CHECK-NEXT: nse = 24
+  // CHECK-NEXT: dim = ( 6, 16 )
+  // CHECK-NEXT: lvl = ( 2, 4, 3, 4 )
   // CHECK-NEXT: pos[1] : ( 0, 1, 2,
   // CHECK-NEXT: crd[1] : ( 0, 2,
   // CHECK-NEXT: values : ( 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7,
@@ -130,6 +132,8 @@ module {
   //
   // CHECK-NEXT: ---- Sparse Tensor ----
   // CHECK-NEXT: nse = 24
+  // CHECK-NEXT: dim = ( 6, 16 )
+  // CHECK-NEXT: lvl = ( 2, 4, 4, 3 )
   // CHECK-NEXT: pos[1] : ( 0, 1, 2,
   // CHECK-NEXT: crd[1] : ( 0, 2,
   // CHECK-NEXT: values : ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7,
@@ -154,6 +158,8 @@ module {
   //
   // CHECK-NEXT: ---- Sparse Tensor ----
   // CHECK-NEXT: nse = 24
+  // CHECK-NEXT: dim = ( 6, 16 )
+  // CHECK-NEXT: lvl = ( 4, 2, 3, 4 )
   // CHECK-NEXT: pos[1] : ( 0, 1, 1, 2, 2,
   // CHECK-NEXT: crd[1] : ( 0, 1,
   // CHECK-NEXT: values : ( 1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 5, 0, 0, 0, 0, 0, 0, 0, 0, 6, 7,
@@ -178,6 +184,8 @@ module {
   //
   // CHECK-NEXT: ---- Sparse Tensor ----
   // CHECK-NEXT: nse = 24
+  // CHECK-NEXT: dim = ( 6, 16 )
+  // CHECK-NEXT: lvl = ( 4, 2, 4, 3 )
   // CHECK-NEXT: pos[1] : ( 0, 1, 1, 2, 2,
   // CHECK-NEXT: crd[1] : ( 0, 1,
   // CHECK-NEXT: values : ( 1, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 5, 0, 0, 0, 0, 6, 0, 0, 7,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir
index c6ee0ce0705021..b2bbc64f1688ed 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output.mlir
@@ -106,6 +106,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 25
+    // CHECK-NEXT: dim = ( 5, 5 )
+    // CHECK-NEXT: lvl = ( 5, 5 )
     // CHECK-NEXT: values : ( 2, 0, 0, 2.8, 0, 0, 4, 0, 0, 5, 0, 0, 6, 0, 0, 8.2, 0, 0, 8, 0, 0, 10.4, 0, 0, 10,
     // CHECK-NEXT: ----
     //
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
index 0b34ff581016da..ca9df03c69eed6 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_bf16.mlir
@@ -93,6 +93,8 @@ module {
     //
     // CHECK: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 32 )
+    // CHECK-NEXT: lvl = ( 32 )
     // CHECK-NEXT: values : ( 1, 11, 0, 2, 13, 0, 0, 0, 0, 0, 14, 3, 0, 0, 0, 0, 15, 4, 16, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9,
     // CHECK-NEXT: ----
     //
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir
index 495682169c2909..4f5e6ddd48d8ed 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/dense_output_f16.mlir
@@ -94,6 +94,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 32 )
+    // CHECK-NEXT: lvl = ( 32 )
     // CHECK-NEXT: values : ( 1, 11, 0, 2, 13, 0, 0, 0, 0, 0, 14, 3, 0, 0, 0, 0, 15, 4, 16, 0, 5, 6, 0, 0, 0, 0, 0, 0, 7, 8, 0, 9,
     // CHECK-NEXT: ----
     //
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
index 3a32ff28527001..1917fd987c5d7d 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_insert_3d.mlir
@@ -62,6 +62,8 @@ module {
 
     // CHECK: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 5, 4, 3 )
+    // CHECK-NEXT: lvl = ( 5, 4, 3 )
     // CHECK-NEXT: pos[0] : ( 0, 2
     // CHECK-NEXT: crd[0] : ( 3, 4
     // CHECK-NEXT: pos[2] : ( 0, 2, 2, 2, 3, 3, 3, 4, 5
@@ -79,6 +81,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 5, 4, 3 )
+    // CHECK-NEXT: lvl = ( 5, 4, 3 )
     // CHECK-NEXT: pos[0] : ( 0, 2
     // CHECK-NEXT: crd[0] : ( 3, 4
     // CHECK-NEXT: pos[1] : ( 0, 2, 4
@@ -96,6 +100,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 5, 4, 3 )
+    // CHECK-NEXT: lvl = ( 5, 4, 3 )
     // CHECK-NEXT: pos[0] : ( 0, 2
     // CHECK-NEXT: crd[0] : ( 3, 4
     // CHECK-NEXT: pos[1] : ( 0, 3, 5
@@ -113,6 +119,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 5, 4, 3 )
+    // CHECK-NEXT: lvl = ( 5, 4, 3 )
     // CHECK-NEXT: pos[1] : ( 0, 0, 0, 0, 3, 5
     // CHECK-NEXT: crd[1] : ( 0, 1, 0, 2, 3, 1, 2, 2, 3, 2
     // CHECK-NEXT: values : ( 1.1, 2.2, 3.3, 4.4, 5.5
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
index e1f062121b12f9..c05a9f5742690c 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_loose.mlir
@@ -41,6 +41,8 @@ module {
     //
     // CHECK:   ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 17
+    // CHECK-NEXT: dim = ( 5, 4 )
+    // CHECK-NEXT: lvl = ( 5, 4 )
     // 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
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
index 863e1c62370e32..e505559037a99f 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul.mlir
@@ -144,6 +144,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // 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
@@ -154,6 +156,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 8, 16, 24, 32
@@ -166,6 +170,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // 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
@@ -176,6 +182,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 4
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 3
     // CHECK-NEXT: crd[0] : ( 0, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 4
@@ -188,6 +196,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 8, 4 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // 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
@@ -198,6 +208,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 8, 4 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // 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
@@ -210,6 +222,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: dim = ( 8, 4 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // 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
@@ -220,6 +234,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: dim = ( 8, 4 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // 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
@@ -271,6 +287,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 16
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -281,6 +299,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 16
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12, 16
@@ -302,6 +322,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -312,6 +334,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 3
     // CHECK-NEXT: crd[0] : ( 0, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 4, 8, 12
@@ -333,6 +357,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -343,6 +369,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 3
     // CHECK-NEXT: crd[0] : ( 0, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 5
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 b95f72e1a4799f..58e96d1fa51faa 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matmul_slice.mlir
@@ -172,6 +172,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 3
     // CHECK-NEXT: crd[0] : ( 0, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 5
@@ -192,6 +194,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -204,6 +208,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -220,6 +226,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
@@ -236,6 +244,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: dim = ( 4, 4 )
+    // CHECK-NEXT: lvl = ( 4, 4 )
     // 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
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
index 2cef46f4cb1546..8ea26fa3efdf51 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matrix_ops.mlir
@@ -161,6 +161,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 9
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 6, 9
@@ -173,6 +175,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 6
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 3
     // CHECK-NEXT: crd[0] : ( 0, 1, 2
     // CHECK-NEXT: pos[1] : ( 0, 2, 4, 6
@@ -185,6 +189,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 9
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 6, 9
@@ -197,6 +203,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 9
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 6, 9
@@ -209,6 +217,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 13
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 4
     // CHECK-NEXT: crd[0] : ( 0, 1, 2, 3
     // CHECK-NEXT: pos[1] : ( 0, 3, 6, 10, 13
@@ -221,6 +231,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 2
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 2
     // CHECK-NEXT: crd[0] : ( 0, 2
     // CHECK-NEXT: pos[1] : ( 0, 1, 2
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_print.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_print.mlir
index 79728fdb0f8cde..98d76ba350cbd9 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_print.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_print.mlir
@@ -137,12 +137,16 @@ module {
 
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 5, 0, 0,
     // CHECK-NEXT: ----
     sparse_tensor.print %XO : tensor<4x8xi32, #AllDense>
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 32
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // CHECK-NEXT: values : ( 1, 0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 3, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0, 0, 0,
     // CHECK-NEXT: ----
     sparse_tensor.print %XT : tensor<4x8xi32, #AllDenseT>
@@ -160,6 +164,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 2, 2, 5,
     // CHECK-NEXT: crd[1] : ( 0, 2, 2, 3, 5,
     // CHECK-NEXT: values : ( 1, 2, 3, 4, 5,
@@ -168,6 +174,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 4, 8 )
     // CHECK-NEXT: pos[0] : ( 0, 2,
     // CHECK-NEXT: crd[0] : ( 0, 3,
     // CHECK-NEXT: pos[1] : ( 0, 2, 5,
@@ -178,6 +186,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // CHECK-NEXT: pos[1] : ( 0, 1, 1, 3, 4, 4, 5, 5, 5,
     // CHECK-NEXT: crd[1] : ( 0, 0, 3, 3, 3,
     // CHECK-NEXT: values : ( 1, 2, 3, 4, 5,
@@ -186,6 +196,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 8, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 4,
     // CHECK-NEXT: crd[0] : ( 0, 2, 3, 5,
     // CHECK-NEXT: pos[1] : ( 0, 1, 3, 4, 5,
@@ -196,6 +208,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 2,
     // CHECK-NEXT: crd[0] : ( 0, 1,
     // CHECK-NEXT: pos[1] : ( 0, 1, 3,
@@ -206,6 +220,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 4, 2 )
     // CHECK-NEXT: pos[0] : ( 0, 2,
     // CHECK-NEXT: crd[0] : ( 0, 1,
     // CHECK-NEXT: pos[1] : ( 0, 1, 3,
@@ -216,6 +232,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )
     // CHECK-NEXT: pos[0] : ( 0, 2,
     // CHECK-NEXT: crd[0] : ( 0, 1,
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
@@ -226,6 +244,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 4, 2 )
     // CHECK-NEXT: pos[0] : ( 0, 2,
     // CHECK-NEXT: crd[0] : ( 0, 1,
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
@@ -236,6 +256,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )
     // CHECK-NEXT: pos[1] : ( 0, 1, 3,
     // CHECK-NEXT: crd[1] : ( 0, 0, 1,
     // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0,
@@ -244,6 +266,8 @@ module {
 
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 24
+    // CHECK-NEXT: dim = ( 4, 8 )
+    // CHECK-NEXT: lvl = ( 2, 2, 2, 4 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
     // CHECK-NEXT: crd[1] : ( 0, 1, 1,
     // CHECK-NEXT: values : ( 1, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 5, 0, 0,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir
index 1860fc1c7027a1..7bacbe3b87e429 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_re_im.mlir
@@ -91,6 +91,8 @@ module {
     //
     // CHECK:    ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: dim = ( 32 )
+    // CHECK-NEXT: lvl = ( 32 )
     // CHECK-NEXT: pos[0] : ( 0, 3,
     // CHECK-NEXT: crd[0] : ( 0, 20, 31,
     // CHECK-NEXT: values : ( 5.13, 3, 5,
@@ -98,6 +100,8 @@ module {
     //
     // CHECK-NEXT: ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 3
+    // CHECK-NEXT: dim = ( 32 )
+    // CHECK-NEXT: lvl = ( 32 )
     // CHECK-NEXT: pos[0] : ( 0, 3,
     // CHECK-NEXT: crd[0] : ( 0, 20, 31,
     // CHECK-NEXT: values : ( 2, 4, 6,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-gemm-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-gemm-lib.mlir
index 984b3614eded56..da78452d94fd6b 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-gemm-lib.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-gemm-lib.mlir
@@ -66,6 +66,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 20
+    // CHECK-NEXT: dim = ( 8, 8 )
+    // CHECK-NEXT: lvl = ( 8, 8 )
     // CHECK-NEXT: pos[1] : ( 0, 5, 5, 6, 7, 8, 12, 16, 20,
     // CHECK-NEXT: crd[1] : ( 0, 1, 2, 6, 7, 2, 3, 4, 1, 2, 6, 7, 1, 2, 6, 7, 1, 2, 6, 7,
     // CHECK-NEXT: values : ( 1, 39, 52, 45, 51, 16, 25, 36, 117, 158, 135, 144, 156, 318, 301, 324, 208, 430, 405, 436,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir
index aad26556b00e1e..3d17b719732ff2 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sampled-matmul-lib.mlir
@@ -115,6 +115,8 @@ module {
     //
     // CHECK:   ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 9
+    // CHECK-NEXT: dim = ( 5, 5 )
+    // CHECK-NEXT: lvl = ( 5, 5 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 4, 5, 7, 9,
     // CHECK-NEXT: crd[1] : ( 0, 3, 1, 4, 2, 0, 3, 1, 4,
     // CHECK-NEXT: values : ( 11, 41.4, 42, 102.5, 93, 44.1, 164, 105.2, 255,
@@ -141,6 +143,8 @@ module {
     //
     // CHECK:     ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 5
+    // CHECK-NEXT: dim = ( 8, 8 )
+    // CHECK-NEXT: lvl = ( 8, 8 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3, 3, 4, 4, 4, 4, 5,
     // CHECK-NEXT: crd[1] : ( 0, 1, 0, 4, 7,
     // CHECK-NEXT: values : ( 17, 18, 19, 20, 21,
diff --git a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir
index bd2c72271c613a..68bb32891f3439 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/GPU/CUDA/sparse-sddmm-lib.mlir
@@ -168,6 +168,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 8
+    // CHECK-NEXT: dim = ( 4, 6 )
+    // CHECK-NEXT: lvl = ( 4, 6 )
     // CHECK-NEXT: pos[1] : ( 0, 3, 5, 7, 8,
     // CHECK-NEXT: crd[1] : ( 0, 1, 4, 1, 5, 2, 3, 2,
     // CHECK-NEXT: values : ( 5, 10, 24, 19, 53, 42, 55, 56,
@@ -175,6 +177,8 @@ module {
     //
     // CHECK:      ---- Sparse Tensor ----
     // CHECK-NEXT: nse = 12
+    // CHECK-NEXT: dim = ( 4, 6 )
+    // CHECK-NEXT: lvl = ( 2, 3, 2, 2 )
     // CHECK-NEXT: pos[1] : ( 0, 2, 3,
     // CHECK-NEXT: crd[1] : ( 0, 2, 1,
     // CHECK-NEXT: values : ( 5, 10, 8, 19, 24, 24, 40, 53, 42, 55, 56, 64,



More information about the Mlir-commits mailing list