Implement some simple affine expr canonicalization/simplification.
- fold constants when possible.
- for a mul expression, canonicalize to always keep the LHS as the
constant/symbolic term, and similarly, the RHS for an add expression to keep
it closer to the mathematical form. (Eg: f(x) = 3*x + 5)); other similar simplifications;
- verify binary op expressions at creation time.
TODO: we can completely drop AffineSubExpr, and instead use add and mul by -1.
This way something like x - 4 and -4 + x get canonicalized to x + -1 * 4
instead of being x - 4 and x + -4. (The other alternative if wanted to retain
AffineSubExpr would be to simplify x + -1*y to x - y and x + <neg number> to x
- <pos number>).
PiperOrigin-RevId: 204240258
diff --git a/lib/IR/AffineExpr.cpp b/lib/IR/AffineExpr.cpp
index e894f49..54acedf 100644
--- a/lib/IR/AffineExpr.cpp
+++ b/lib/IR/AffineExpr.cpp
@@ -17,9 +17,44 @@
#include "mlir/IR/AffineExpr.h"
#include "mlir/Support/STLExtras.h"
+#include "third_party/llvm/llvm/include/llvm/ADT/STLExtras.h"
using namespace mlir;
+AffineBinaryOpExpr::AffineBinaryOpExpr(Kind kind, AffineExpr *lhs,
+ AffineExpr *rhs)
+ : AffineExpr(kind), lhs(lhs), rhs(rhs) {
+ // We verify affine op expr forms at construction time.
+ switch (kind) {
+ case Kind::Add:
+ assert(!isa<AffineConstantExpr>(lhs));
+ // TODO (more verification)
+ break;
+ case Kind::Sub:
+ // TODO (verification)
+ break;
+ case Kind::Mul:
+ assert(!isa<AffineConstantExpr>(lhs));
+ assert(rhs->isSymbolic());
+ // TODO (more verification)
+ break;
+ case Kind::FloorDiv:
+ assert(rhs->isSymbolic());
+ // TODO (more verification)
+ break;
+ case Kind::CeilDiv:
+ assert(rhs->isSymbolic());
+ // TODO (more verification)
+ break;
+ case Kind::Mod:
+ assert(rhs->isSymbolic());
+ // TODO (more verification)
+ break;
+ default:
+ llvm_unreachable("unexpected binary affine expr");
+ }
+}
+
/// Returns true if this expression is made out of only symbols and
/// constants (no dimensional identifiers).
bool AffineExpr::isSymbolic() const {