[mlir][sparse] update BSR specification

[aartbik]

Makes outer level dense, so we get the common
block-column index way of storing blocks. Also
gives an example of a dim2lvl/lvl2dim map.

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Changes

Makes outer level dense, so we get the common
block-column index way of storing blocks. Also
gives an example of a dim2lvl/lvl2dim map.

---

Full diff: https://github.com/llvm/llvm-project/pull/67480.diff

1 Files Affected:

• (modified) mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td (+13-1)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index e6577aed063ca7f..9c20c94802a4267 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -217,7 +217,7 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     // Block sparse row storage (2x3 blocks).
     #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
@@ -225,6 +225,18 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
     }>
     ... tensor<20x30xf32, #BSR> ...
 
+    // 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,
+              j = jb * 3 + jj) ->
+      ( ib = i floordiv 2 : dense,
+        jb = j floordiv 3 : compressed,
+        ii = i mod 2 : dense,
+        jj = j mod 3 : dense)
+    }>
+    ... tensor<20x30xf32, #BSR_explicit> ...
+
     // CSR slice (offset = 0, size = 4, stride = 1 on the first dimension;
     // offset = 0, size = 8, and a dynamic stride on the second dimension).
     #CSR_SLICE = #sparse_tensor.encoding<{