From 2263d85c86c473a024761280f452bc0ee17f34f8 Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik@google.com>
Date: Tue, 26 Sep 2023 12:54:50 -0700
Subject: [PATCH] [mlir][sparse] update BSR specification

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.
---
 .../SparseTensor/IR/SparseTensorAttrDefs.td        | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index e6577aed063ca..9c20c94802a42 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<{