[mlir][sparse]Make isBlockSparsity more robust

[yinying-lisa-li]

1. A single dimension can either be blocked (with floordiv and mod pair) or non-blocked. Mixing them would be invalid.
2. Block size should be non-zero value.

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Author: Yinying Li (yinying-lisa-li)

Changes

1. A single dimension can either be blocked (with floordiv and mod pair) or non-blocked. Mixing them would be invalid.
2. Block size should be non-zero value.

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+15-2)
• (modified) mlir/test/Dialect/SparseTensor/invalid_encoding.mlir (+45)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 686180c09da724..70f508827e26b4 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -855,16 +855,21 @@ SmallVector<unsigned> mlir::sparse_tensor::getBlockSize(AffineMap dimToLvl) {
 bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
   if (!dimToLvl)
     return false;
+  SmallVector<unsigned> isDimSet;
+  isDimSet.resize(dimToLvl.getNumDims());
   std::map<unsigned, int64_t> coeffientMap;
   for (auto result : dimToLvl.getResults()) {
     if (auto binOp = dyn_cast<AffineBinaryOpExpr>(result)) {
       // Check for "dim op const".
       auto dimOp = dyn_cast<AffineDimExpr>(binOp.getLHS());
       auto conOp = dyn_cast<AffineConstantExpr>(binOp.getRHS());
-      if (!dimOp || !conOp)
+      if (!dimOp || !conOp || conOp.getValue() <= 0)
         return false;
-      // Inspect "dim / const" or "dim % const".
       auto pos = dimOp.getPosition();
+      // Check current dim has not been set before.
+      if (isDimSet[pos] == 1)
+        return false;
+      // Inspect "dim / const" or "dim % const".
       if (binOp.getKind() == AffineExprKind::FloorDiv) {
         // Expect only one floordiv for each dimension.
         if (coeffientMap.find(pos) != coeffientMap.end())
@@ -881,6 +886,14 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
       } else {
         return false;
       }
+    } else if (auto dimOp = dyn_cast<AffineDimExpr>(result)) {
+      auto pos = dimOp.getPosition();
+      isDimSet[pos] = 1;
+      // Expect dim to be non-blocked (without floordiv/mod pair).
+      if (coeffientMap.find(pos) != coeffientMap.end())
+        return false;
+    } else {
+      return false;
     }
   }
   return !coeffientMap.empty();
diff --git a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
index 6514391bae92d9..2d189cc94c15e2 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -254,6 +254,51 @@ func.func private @wrong_order_lvl_decl(%arg0: tensor<?x?xf64, #WrongOrderLvlDec
 
 // -----
 
+// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
+#BSR = #sparse_tensor.encoding<{
+  map = ( i, j ) ->
+  ( i floordiv 2 : dense,
+    j floordiv 3 : compressed,
+    i            : dense,
+    j mod 3      : dense
+  )
+}>
+func.func private @BSR(%arg0: tensor<?x?xf64, #BSR>) {
+  return
+}
+
+// -----
+
+// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
+#BSR = #sparse_tensor.encoding<{
+  map = ( i, j ) ->
+  ( i            : dense,
+    j floordiv 3 : compressed,
+    i floordiv 3 : dense,
+    j mod 3      : dense
+  )
+}>
+func.func private @BSR(%arg0: tensor<?x?xf64, #BSR>) {
+  return
+}
+
+// -----
+
+// expected-error@+1 {{failed to infer lvlToDim from dimToLvl}}
+#BSR = #sparse_tensor.encoding<{
+  map = ( i, j ) ->
+  ( i floordiv -3 : dense,
+    j floordiv -3 : compressed,
+    i mod 3 : dense,
+    j mod 3      : dense
+  )
+}>
+func.func private @BSR(%arg0: tensor<?x?xf64, #BSR>) {
+  return
+}
+
+// -----
+
 // expected-error@+1 {{expected lvlToDim to be an inverse of dimToLvl}}
 #BSR_explicit = #sparse_tensor.encoding<{
   map =