[Mlir-commits] [mlir] 31b72b0 - [mlir][sparse]Make isBlockSparsity more robust (#75113)

[llvmlistbot at llvm.org]

Author: Yinying Li
Date: 2023-12-12T13:43:03-05:00
New Revision: 31b72b0742b09c970bea47b70eb535fe964e3c8f

URL: https://github.com/llvm/llvm-project/commit/31b72b0742b09c970bea47b70eb535fe964e3c8f
DIFF: https://github.com/llvm/llvm-project/commit/31b72b0742b09c970bea47b70eb535fe964e3c8f.diff

LOG: [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.

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
    mlir/test/Dialect/SparseTensor/invalid_encoding.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 686180c09da724..6033ebf6897ce0 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/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) {

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 at +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 at +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 at +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 at +1 {{expected lvlToDim to be an inverse of dimToLvl}}
 #BSR_explicit = #sparse_tensor.encoding<{
   map =