[Mlir-commits] [mlir] 3d3e46c - [mlir][sparse] make test for block sparsity more robust (#74798)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Dec 8 11:50:15 PST 2023
Author: Aart Bik
Date: 2023-12-08T11:50:10-08:00
New Revision: 3d3e46cc4db9dd32edc82b7029fb694d5d0316de
URL: https://github.com/llvm/llvm-project/commit/3d3e46cc4db9dd32edc82b7029fb694d5d0316de
DIFF: https://github.com/llvm/llvm-project/commit/3d3e46cc4db9dd32edc82b7029fb694d5d0316de.diff
LOG: [mlir][sparse] make test for block sparsity more robust (#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)
Added:
Modified:
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
Removed:
################################################################################
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;
}
More information about the Mlir-commits
mailing list