Sketch out affine analysis structures: AffineValueMap, IntegerValueSet,
FlatAffineConstraints, and MutableAffineMap.
All four classes introduced reside in lib/Analysis and are not meant to be
used in the IR (from lib/IR or lib/Parser/). They are all mutable, alloc'ed,
dealloc'ed - although with their fields pointing to immutable affine
expressions (AffineExpr *).
While on this, update simplifyMod to fold mod to a zero when possible.
PiperOrigin-RevId: 209618437
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index 3f640ba..95cd8ff 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -205,20 +205,33 @@
}
return nullptr;
- // TODO(someone): implement more simplification along the lines described in
- // simplifyMod TODO. For eg: 128*N ceildiv 128 is N.
}
AffineExpr *AffineBinaryOpExpr::simplifyMod(AffineExpr *lhs, AffineExpr *rhs,
MLIRContext *context) {
- if (auto *l = dyn_cast<AffineConstantExpr>(lhs))
- if (auto *r = dyn_cast<AffineConstantExpr>(rhs))
- return AffineConstantExpr::get(l->getValue() % r->getValue(), context);
+ auto *lhsConst = dyn_cast<AffineConstantExpr>(lhs);
+ auto *rhsConst = dyn_cast<AffineConstantExpr>(rhs);
+
+ if (lhsConst && rhsConst)
+ return AffineConstantExpr::get(lhsConst->getValue() % rhsConst->getValue(),
+ context);
+
+ // Fold modulo of a multiply with a constant that is a multiple of the
+ // modulo factor to zero. Eg: (i * 128) mod 64 = 0.
+ if (rhsConst) {
+ auto *lBin = dyn_cast<AffineBinaryOpExpr>(lhs);
+ if (lBin && lBin->getKind() == Kind::Mul) {
+ if (auto *lrhs = dyn_cast<AffineConstantExpr>(lBin->getRHS())) {
+ // rhsConst is known to be positive if a constant.
+ if (lrhs->getValue() % rhsConst->getValue() == 0)
+ return AffineConstantExpr::get(0, context);
+ }
+ }
+ }
return nullptr;
- // TODO(someone): implement more simplification; for eg: 2*x mod 2 is 0; (2*x
- // + 1) mod 2 is 1. In general, this can be simplified by using the GCD test
- // iteratively if the RHS of the mod is a small number, or in general using
- // quantifier elimination (add two new variables q and r, and eliminate all
- // variables from the linear system other than r.
+ // TODO(bondhugula): In general, this can be simplified more by using the GCD
+ // test, or in general using quantifier elimination (add two new variables q
+ // and r, and eliminate all variables from the linear system other than r. All
+ // of this can be done through mlir/Analysis/'s FlatAffineConstraints.
}