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

[Yinying Li]

https://github.com/yinying-lisa-li updated https://github.com/llvm/llvm-project/pull/75113

>From d85c5f99e39ccca6e80f43dc49f773c80643c363 Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Mon, 11 Dec 2023 23:01:44 +0000
Subject: [PATCH 1/2] [mlir][sparse]Make isBlockSparsity more robust

---
 .../SparseTensor/IR/SparseTensorDialect.cpp   | 17 ++++++-
 .../SparseTensor/invalid_encoding.mlir        | 45 +++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)

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 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 =

>From 5933b1008a441fbce04767f827302968790e3744 Mon Sep 17 00:00:00 2001
From: Yinying Li <yinyingli at google.com>
Date: Mon, 11 Dec 2023 23:44:25 +0000
Subject: [PATCH 2/2] address comments

---
 .../SparseTensor/IR/SparseTensorDialect.cpp       | 15 ++++++---------
 1 file changed, 6 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 70f508827e26b4..622ed13422b84a 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -855,9 +855,8 @@ 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;
+  bool hasBlock = false;
   for (auto result : dimToLvl.getResults()) {
     if (auto binOp = dyn_cast<AffineBinaryOpExpr>(result)) {
       // Check for "dim op const".
@@ -865,11 +864,8 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
       auto conOp = dyn_cast<AffineConstantExpr>(binOp.getRHS());
       if (!dimOp || !conOp || conOp.getValue() <= 0)
         return false;
-      auto pos = dimOp.getPosition();
-      // Check current dim has not been set before.
-      if (isDimSet[pos] == 1)
-        return false;
       // Inspect "dim / const" or "dim % const".
+      auto pos = dimOp.getPosition();
       if (binOp.getKind() == AffineExprKind::FloorDiv) {
         // Expect only one floordiv for each dimension.
         if (coeffientMap.find(pos) != coeffientMap.end())
@@ -883,20 +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();
-      isDimSet[pos] = 1;
-      // Expect dim to be non-blocked (without floordiv/mod pair).
+      // Expect dim to be unset.
       if (coeffientMap.find(pos) != coeffientMap.end())
         return false;
+      coeffientMap[pos] = 0;
     } else {
       return false;
     }
   }
-  return !coeffientMap.empty();
+  return !coeffientMap.empty() && hasBlock;
 }
 
 bool mlir::sparse_tensor::hasAnyNonIdentityOperandsOrResults(Operation *op) {