[mlir][sparse] Fix errors in doc and tests (#68641)

llvm · Oct 10, 2023 · 14d0cd6 · 14d0cd6
1 parent ac0dda8
commit 14d0cd6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 23 deletions.
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -133,9 +133,11 @@ 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 +178,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 +227,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 +265,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
@@ -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
 }