Affine expression analysis and simplification.
Outside of IR/
- simplify a MutableAffineMap by flattening the affine expressions
- add a simplify affine expression pass that uses this analysis
- update the FlatAffineConstraints API (to be used in the next CL)
In IR:
- add isMultipleOf and getKnownGCD for AffineExpr, and make the in-IR
simplication of simplifyMod simpler and more powerful.
- rename the AffineExpr visitor methods to distinguish b/w visiting and
walking, and to simplify API names based on context.
The next CL will use some of these for the loop unrolling/unroll-jam to make
the detection for the need of cleanup loop powerful/non-trivial.
A future CL will finally move this simplification to FlatAffineConstraints to
make it more powerful. For eg., currently, even if a mod expr appearing in a
part of the expression tree can't be simplified, the whole thing won't be
simplified.
PiperOrigin-RevId: 211012256
diff --git a/lib/IR/AffineExpr.cpp b/lib/IR/AffineExpr.cpp
index 6bfbaf5..6393b3e 100644
--- a/lib/IR/AffineExpr.cpp
+++ b/lib/IR/AffineExpr.cpp
@@ -100,3 +100,56 @@
}
}
}
+
+uint64_t AffineExpr::getKnownGcd() const {
+ AffineBinaryOpExpr *binExpr = nullptr;
+ switch (kind) {
+ case Kind::SymbolId:
+ LLVM_FALLTHROUGH;
+ case Kind::DimId:
+ return 1;
+ case Kind::Constant:
+ return std::abs(cast<AffineConstantExpr>(this)->getValue());
+ case Kind::Mul:
+ binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
+ return binExpr->getLHS()->getKnownGcd() * binExpr->getRHS()->getKnownGcd();
+ case Kind::Add:
+ LLVM_FALLTHROUGH;
+ case Kind::FloorDiv:
+ case Kind::CeilDiv:
+ case Kind::Mod:
+ binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
+ return llvm::GreatestCommonDivisor64(binExpr->getLHS()->getKnownGcd(),
+ binExpr->getRHS()->getKnownGcd());
+ }
+}
+
+bool AffineExpr::isMultipleOf(int64_t factor) const {
+ AffineBinaryOpExpr *binExpr = nullptr;
+ uint64_t l, u;
+ switch (kind) {
+ case Kind::SymbolId:
+ LLVM_FALLTHROUGH;
+ case Kind::DimId:
+ return factor * factor == 1;
+ case Kind::Constant:
+ return cast<AffineConstantExpr>(this)->getValue() % factor == 0;
+ case Kind::Mul:
+ binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
+ // It's probably not worth optimizing this further (to not traverse the
+ // whole sub-tree under - it that would require a version of isMultipleOf
+ // that on a 'false' return also returns the known GCD).
+ return (l = binExpr->getLHS()->getKnownGcd()) % factor == 0 ||
+ (u = binExpr->getRHS()->getKnownGcd()) % factor == 0 ||
+ (l * u) % factor == 0;
+ case Kind::Add:
+ case Kind::FloorDiv:
+ case Kind::CeilDiv:
+ case Kind::Mod:
+ binExpr = cast<AffineBinaryOpExpr>(const_cast<AffineExpr *>(this));
+ return llvm::GreatestCommonDivisor64(binExpr->getLHS()->getKnownGcd(),
+ binExpr->getRHS()->getKnownGcd()) %
+ factor ==
+ 0;
+ }
+}