[mlir][sparse]Make isBlockSparsity more robust (#75113)

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.

Author: yinying-lisa-li

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

@@ -856,12 +856,13 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
   if (!dimToLvl)
     return false;
   std::map<unsigned, int64_t> coeffientMap;
+  bool hasBlock = false;
   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();
@@ -878,12 +879,21 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
         // Expect mod to have the same coefficient as floordiv.
         if (conOp.getValue() != coeffientMap[pos])
           return false;
+        hasBlock = true;
       } else {
         return false;
       }
+    } else if (auto dimOp = dyn_cast<AffineDimExpr>(result)) {
+      auto pos = dimOp.getPosition();
+      // Expect dim to be unset.
+      if (coeffientMap.find(pos) != coeffientMap.end())
+        return false;
+      coeffientMap[pos] = 0;
+    } else {
+      return false;
     }
   }
-  return !coeffientMap.empty();
+  return hasBlock;
 }
 
 bool mlir::sparse_tensor::hasAnyNonIdentityOperandsOrResults(Operation *op) {

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 =