[Mlir-commits] [mlir] 329c64c - [mlir][sparse] update BSR specification (#67480)

[llvmlistbot at llvm.org]

Author: Aart Bik
Date: 2023-09-26T13:32:57-07:00
New Revision: 329c64cbe3d57104dafd01ed17f723df08034148

URL: https://github.com/llvm/llvm-project/commit/329c64cbe3d57104dafd01ed17f723df08034148
DIFF: https://github.com/llvm/llvm-project/commit/329c64cbe3d57104dafd01ed17f723df08034148.diff

LOG: [mlir][sparse] update BSR specification (#67480)

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.

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

Removed: 
    


################################################################################
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<{