[Mlir-commits] [mlir] [mlir][sparse] make test for block sparsity more robust (PR #74798)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/74798

For BSR and convolutions, we encounter

(d0, d1, d2, d3) -> ((d0 + d2) floordiv 2, (d1 + d3) floordiv 2, (d0 + d2) mod 2, (d1 + d3) mod 2)

which crashed the current test. Note that an actual test and working code is still to follow (since we need to fix a few other things first)

>From 2a10209924733f309a83e547b82a7e3ee931ee6c Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Thu, 7 Dec 2023 18:19:31 -0800
Subject: [PATCH] [mlir][sparse] make test for block sparsity more robust

For BSR and convolutions, we encounter

(d0, d1, d2, d3) -> ((d0 + d2) floordiv 2, (d1 + d3) floordiv 2, (d0 + d2) mod 2, (d1 + d3) mod 2)

which crashed the current test. Note that an actual test and
working code is still to follow (since we need to fix
a few other things first)
---
 .../SparseTensor/IR/SparseTensorDialect.cpp   | 20 +++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 577dfe5ab2f30..686180c09da72 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -858,22 +858,26 @@ bool mlir::sparse_tensor::isBlockSparsity(AffineMap dimToLvl) {
   std::map<unsigned, int64_t> coeffientMap;
   for (auto result : dimToLvl.getResults()) {
     if (auto binOp = dyn_cast<AffineBinaryOpExpr>(result)) {
-      auto pos = dyn_cast<AffineDimExpr>(binOp.getLHS()).getPosition();
-      if (result.getKind() == AffineExprKind::FloorDiv) {
+      // Check for "dim op const".
+      auto dimOp = dyn_cast<AffineDimExpr>(binOp.getLHS());
+      auto conOp = dyn_cast<AffineConstantExpr>(binOp.getRHS());
+      if (!dimOp || !conOp)
+        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())
           return false;
-        coeffientMap[pos] =
-            dyn_cast<AffineConstantExpr>(binOp.getRHS()).getValue();
-      } else if (result.getKind() == AffineExprKind::Mod) {
+        // Record coefficient of the floordiv.
+        coeffientMap[pos] = conOp.getValue();
+      } else if (binOp.getKind() == AffineExprKind::Mod) {
         // Expect floordiv before mod.
         if (coeffientMap.find(pos) == coeffientMap.end())
           return false;
         // Expect mod to have the same coefficient as floordiv.
-        if (dyn_cast<AffineConstantExpr>(binOp.getRHS()).getValue() !=
-            coeffientMap[pos]) {
+        if (conOp.getValue() != coeffientMap[pos])
           return false;
-        }
       } else {
         return false;
       }