[Mlir-commits] [mlir] [mlir][sparse] Fix errors in doc and tests (PR #68641)

[Yinying Li]

https://github.com/yinying-lisa-li created https://github.com/llvm/llvm-project/pull/68641

None

>From c6b17735dcbe66a45df5083edc68a3e3d3c5af4f Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Mon, 9 Oct 2023 22:07:08 +0000
Subject: [PATCH] [mlir][sparse] Fix errors in doc and tests

---
 .../SparseTensor/IR/SparseTensorAttrDefs.td   | 14 ++++-----
 .../SparseTensor/roundtrip_encoding.mlir      | 30 +++++++++----------
 2 files changed, 21 insertions(+), 23 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index cacc8176c678241..0d7340f663ac45b 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -133,9 +133,9 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     level-expressions collectively define an affine map from dimension-coordinates to
     level-coordinates. The dimension-expressions collectively define the inverse map,
     which only needs to be provided for elaborate cases where it cannot be inferred
-    automatically. Within the sparse storage format, we refer to indices that are
-    stored explicitly as **coordinates** and offsets into the storage format as
-    **positions**.
+    automatically. Each dimension could also have an optional `SparseTensorDimSliceAttr`.
+    Within the sparse storage format, we refer to indices that are stored explicitly
+    as **coordinates** and offsets into the storage format as **positions**.
 
     The supported level-formats are the following:
 
@@ -176,9 +176,6 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
       coordinate over all levels).  The choices are `8`, `16`, `32`,
       `64`, or, the default, `0` to indicate a native bitwidth.
 
-    - An optional array of `SparseTensorDimSliceAttr`, which specifies
-      how the sparse tensor is partitioned on each dimension.
-
     Examples:
 
     ```mlir
@@ -228,7 +225,8 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     // Same block sparse row storage (2x3 blocks) but this time
     // also with a redundant reverse mapping, which can be inferred.
     #BSR_explicit = #sparse_tensor.encoding<{
-      map = ( i = ib * 2 + ii,
+      map = {ib, jb, ii, jj}
+            ( i = ib * 2 + ii,
               j = jb * 3 + jj) ->
       ( ib = i floordiv 2 : dense,
         jb = j floordiv 3 : compressed,
@@ -265,7 +263,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
              j : #sparse_tensor<slice(0, 8, ?)>) ->
             (i : dense, j : compressed)
     }>
-    ... tensor<?x?xf64, #CSC_SLICE> ...
+    ... tensor<?x?xf64, #CSR_SLICE> ...
 
     ```
   }];
diff --git a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
index c4ef50bee01ea2c..ae3805d8b774176 100644
--- a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
@@ -84,18 +84,18 @@ func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>)
 
 // -----
 
-#BCSR = #sparse_tensor.encoding<{
+#BSR = #sparse_tensor.encoding<{
    map = ( i, j ) ->
-      ( i floordiv 2 : compressed,
+      ( i floordiv 2 : dense,
         j floordiv 3 : compressed,
         i mod 2      : dense,
         j mod 3      : dense
       )
 }>
 
-// CHECK-LABEL: func private @sparse_bcsr(
-// CHECK-SAME: tensor<10x60xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
-func.func private @sparse_bcsr(tensor<10x60xf64, #BCSR>)
+// CHECK-LABEL: func private @sparse_bsr(
+// CHECK-SAME: tensor<10x60xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
+func.func private @sparse_bsr(tensor<10x60xf64, #BSR>)
 
 
 // -----
@@ -143,39 +143,39 @@ func.func private @sparse_2_out_of_4(tensor<?x?xf64, #NV_24>)
 
 // -----
 
-#BCSR = #sparse_tensor.encoding<{
+#BSR = #sparse_tensor.encoding<{
   map = ( i, j ) ->
-  ( i floordiv 2 : compressed,
+  ( i floordiv 2 : dense,
     j floordiv 3 : compressed,
     i mod 2      : dense,
     j mod 3      : dense
   )
 }>
 
-// CHECK-LABEL: func private @BCSR(
-// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
-func.func private @BCSR(%arg0: tensor<?x?xf64, #BCSR>) {
+// CHECK-LABEL: func private @BSR(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
+func.func private @BSR(%arg0: tensor<?x?xf64, #BSR>) {
   return
 }
 
 // -----
 
-#BCSR_explicit = #sparse_tensor.encoding<{
+#BSR_explicit = #sparse_tensor.encoding<{
   map =
   {il, jl, ii, jj}
   ( i = il * 2 + ii,
     j = jl * 3 + jj
   ) ->
-  ( il = i floordiv 2 : compressed,
+  ( il = i floordiv 2 : dense,
     jl = j floordiv 3 : compressed,
     ii = i mod 2      : dense,
     jj = j mod 3      : dense
   )
 }>
 
-// CHECK-LABEL: func private @BCSR_explicit(
-// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : compressed, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
-func.func private @BCSR_explicit(%arg0: tensor<?x?xf64, #BCSR_explicit>) {
+// CHECK-LABEL: func private @BSR_explicit(
+// CHECK-SAME: tensor<?x?xf64, #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 3 : compressed, d0 mod 2 : dense, d1 mod 3 : dense) }>>
+func.func private @BSR_explicit(%arg0: tensor<?x?xf64, #BSR_explicit>) {
   return
 }