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/include/mlir/IR/AffineExpr.h b/include/mlir/IR/AffineExpr.h
index b561257..a19b0e1 100644
--- a/include/mlir/IR/AffineExpr.h
+++ b/include/mlir/IR/AffineExpr.h
@@ -32,45 +32,45 @@
/// A one-dimensional affine expression.
/// AffineExpression's are immutable (like Type's)
class AffineExpr {
- public:
- enum class Kind {
- Add,
- Sub,
- Mul,
- Mod,
- FloorDiv,
- CeilDiv,
+public:
+ enum class Kind {
+ Add,
+ Sub,
+ Mul,
+ Mod,
+ FloorDiv,
+ CeilDiv,
- /// This is a marker for the last affine binary op. The range of binary
- /// op's is expected to be this element and earlier.
- LAST_AFFINE_BINARY_OP = CeilDiv,
+ /// This is a marker for the last affine binary op. The range of binary
+ /// op's is expected to be this element and earlier.
+ LAST_AFFINE_BINARY_OP = CeilDiv,
- // Constant integer.
- Constant,
- // Dimensional identifier.
- DimId,
- // Symbolic identifier.
- SymbolId,
- };
+ // Constant integer.
+ Constant,
+ // Dimensional identifier.
+ DimId,
+ // Symbolic identifier.
+ SymbolId,
+ };
- /// Return the classification for this type.
- Kind getKind() const { return kind; }
+ /// Return the classification for this type.
+ Kind getKind() const { return kind; }
- void print(raw_ostream &os) const;
- void dump() const;
+ void print(raw_ostream &os) const;
+ void dump() const;
- /// Returns true if this expression is made out of only symbols and
- /// constants (no dimensional identifiers).
- bool isSymbolic() const;
+ /// Returns true if this expression is made out of only symbols and
+ /// constants, i.e., it does not involve dimensional identifiers.
+ bool isSymbolic() const;
- /// Returns true if this is a pure affine expression, i.e., multiplication,
- /// floordiv, ceildiv, and mod is only allowed w.r.t constants.
- bool isPureAffine() const;
+ /// Returns true if this is a pure affine expression, i.e., multiplication,
+ /// floordiv, ceildiv, and mod is only allowed w.r.t constants.
+ bool isPureAffine() const;
- protected:
+protected:
explicit AffineExpr(Kind kind) : kind(kind) {}
- private:
+private:
AffineExpr(const AffineExpr&) = delete;
void operator=(const AffineExpr&) = delete;
@@ -85,52 +85,62 @@
/// Binary affine expression.
class AffineBinaryOpExpr : public AffineExpr {
- public:
- AffineExpr *getLHS() const { return lhs; }
- AffineExpr *getRHS() const { return rhs; }
+public:
+ AffineExpr *getLHS() const { return lhs; }
+ AffineExpr *getRHS() const { return rhs; }
- /// Methods for support type inquiry through isa, cast, and dyn_cast.
- static bool classof(const AffineExpr *expr) {
- return expr->getKind() <= Kind::LAST_AFFINE_BINARY_OP;
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() <= Kind::LAST_AFFINE_BINARY_OP;
}
- protected:
- static AffineBinaryOpExpr *get(Kind kind, AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+protected:
+ static AffineExpr *get(Kind kind, AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
- explicit AffineBinaryOpExpr(Kind kind, AffineExpr *lhs, AffineExpr *rhs)
- : AffineExpr(kind), lhs(lhs), rhs(rhs) {}
+ explicit AffineBinaryOpExpr(Kind kind, AffineExpr *lhs, AffineExpr *rhs);
- AffineExpr *const lhs;
- AffineExpr *const rhs;
+ AffineExpr *const lhs;
+ AffineExpr *const rhs;
+
+private:
+ // Simplification prior to construction of binary affine op expressions.
+ static AffineExpr *simplifyAdd(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
+ static AffineExpr *simplifySub(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
+ static AffineExpr *simplifyMul(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
+ static AffineExpr *simplifyFloorDiv(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
+ static AffineExpr *simplifyCeilDiv(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
+ static AffineExpr *simplifyMod(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
};
/// Binary affine add expression.
class AffineAddExpr : public AffineBinaryOpExpr {
- public:
- static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+public:
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
- /// Methods for support type inquiry through isa, cast, and dyn_cast.
- static bool classof(const AffineExpr *expr) {
- return expr->getKind() == Kind::Add;
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::Add;
}
void print(raw_ostream &os) const;
private:
- /// Simplify the addition of two affine expressions.
- static AffineExpr *simplify(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
-
- explicit AffineAddExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::Add, lhs, rhs) {}
+ // No constructor; use AffineBinaryOpExpr::get
+ AffineAddExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// Binary affine subtract expression.
class AffineSubExpr : public AffineBinaryOpExpr {
public:
- static AffineSubExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
/// Methods for support type inquiry through isa, cast, and dyn_cast.
static bool classof(const AffineExpr *expr) {
@@ -139,15 +149,14 @@
void print(raw_ostream &os) const;
private:
- explicit AffineSubExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::Sub, lhs, rhs) {}
+ AffineSubExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// Binary affine multiplication expression.
class AffineMulExpr : public AffineBinaryOpExpr {
public:
- static AffineMulExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
/// Methods for support type inquiry through isa, cast, and dyn_cast.
static bool classof(const AffineExpr *expr) {
@@ -156,15 +165,14 @@
void print(raw_ostream &os) const;
private:
- explicit AffineMulExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::Mul, lhs, rhs) {}
+ AffineMulExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// Binary affine modulo operation expression.
class AffineModExpr : public AffineBinaryOpExpr {
public:
- static AffineModExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
/// Methods for support type inquiry through isa, cast, and dyn_cast.
static bool classof(const AffineExpr *expr) {
@@ -173,32 +181,30 @@
void print(raw_ostream &os) const;
private:
- explicit AffineModExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::Mod, lhs, rhs) {}
+ AffineModExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// Binary affine floordiv expression.
class AffineFloorDivExpr : public AffineBinaryOpExpr {
- public:
- static AffineFloorDivExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+public:
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
- /// Methods for support type inquiry through isa, cast, and dyn_cast.
- static bool classof(const AffineExpr *expr) {
- return expr->getKind() == Kind::FloorDiv;
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const AffineExpr *expr) {
+ return expr->getKind() == Kind::FloorDiv;
}
void print(raw_ostream &os) const;
private:
- explicit AffineFloorDivExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::FloorDiv, lhs, rhs) {}
+ AffineFloorDivExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// Binary affine ceildiv expression.
class AffineCeilDivExpr : public AffineBinaryOpExpr {
public:
- static AffineCeilDivExpr *get(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context);
+ static AffineExpr *get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context);
/// Methods for support type inquiry through isa, cast, and dyn_cast.
static bool classof(const AffineExpr *expr) {
@@ -207,8 +213,7 @@
void print(raw_ostream &os) const;
private:
- explicit AffineCeilDivExpr(AffineExpr *lhs, AffineExpr *rhs)
- : AffineBinaryOpExpr(Kind::CeilDiv, lhs, rhs) {}
+ AffineCeilDivExpr(AffineExpr *lhs, AffineExpr *rhs) = delete;
};
/// A dimensional identifier appearing in an affine expression.
@@ -242,7 +247,7 @@
/// value. The underlying data is owned by MLIRContext and is thus immortal for
/// almost all clients.
class AffineSymbolExpr : public AffineExpr {
- public:
+public:
static AffineSymbolExpr *get(unsigned position, MLIRContext *context);
unsigned getPosition() const { return position; }
@@ -263,7 +268,7 @@
/// An integer constant appearing in affine expression.
class AffineConstantExpr : public AffineExpr {
- public:
+public:
static AffineConstantExpr *get(int64_t constant, MLIRContext *context);
int64_t getValue() const { return constant; }
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 {
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index ccfd108..822b230 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -18,6 +18,7 @@
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/AffineExpr.h"
#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/MathExtras.h"
using namespace mlir;
@@ -26,14 +27,95 @@
: numDims(numDims), numSymbols(numSymbols), numResults(numResults),
results(results) {}
-AffineExpr *AffineAddExpr::simplify(AffineExpr *lhs, AffineExpr *rhs,
- MLIRContext *context) {
- AffineConstantExpr *l, *r;
- if ((l = dyn_cast<AffineConstantExpr>(lhs)) &&
- (r = dyn_cast<AffineConstantExpr>(rhs)))
- return AffineConstantExpr::get(l->getValue() + r->getValue(), context);
+/// Fold to a constant when possible. Canonicalize so that only the RHS is a
+/// constant. (4 + d0 becomes d0 + 4). If only one of them is a symbolic
+/// expressions, make it the RHS. Return nullptr if it can't be simplified.
+AffineExpr *AffineBinaryOpExpr::simplifyAdd(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);
+
+ if (isa<AffineConstantExpr>(lhs) || (lhs->isSymbolic() && !rhs->isSymbolic()))
+ return AffineAddExpr::get(rhs, lhs, context);
+
return nullptr;
- // TODO(someone): implement more simplification.
+ // TODO(someone): implement more simplification like x + 0 -> x; (x + 2) + 4
+ // -> (x + 6). Do this in a systematic way in conjunction with other
+ // simplifications as opposed to incremental hacks.
}
-// TODO(bondhugula): implement simplify for remaining affine binary op expr's
+AffineExpr *AffineBinaryOpExpr::simplifySub(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);
+
+ return nullptr;
+ // TODO(someone): implement more simplification like mentioned for add.
+}
+
+/// Simplify a multiply expression. Fold it to a constant when possible, and
+/// make the symbolic/constant operand the RHS.
+AffineExpr *AffineBinaryOpExpr::simplifyMul(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);
+
+ assert(lhs->isSymbolic() || rhs->isSymbolic());
+
+ // Canonicalize the mul expression so that the constant/symbolic term is the
+ // RHS. If both the lhs and rhs are symbolic, swap them if the lhs is a
+ // constant. (Note that a constant is trivially symbolic).
+ if (!rhs->isSymbolic() || isa<AffineConstantExpr>(lhs)) {
+ // At least one of them has to be symbolic.
+ return AffineMulExpr::get(rhs, lhs, context);
+ }
+
+ return nullptr;
+ // TODO(someone): implement some more simplification/canonicalization such as
+ // 1*x is same as x, and in general, move it in the form d_i*expr where d_i is
+ // a dimensional identifier. So, 2*(d0 + 4) + s0*d0 becomes (2 + s0)*d0 + 8.
+}
+
+AffineExpr *AffineBinaryOpExpr::simplifyFloorDiv(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);
+
+ return nullptr;
+ // TODO(someone): implement more simplification along the lines described in
+ // simplifyMod TODO. For eg: 128*N floordiv 128 is N.
+}
+
+AffineExpr *AffineBinaryOpExpr::simplifyCeilDiv(AffineExpr *lhs,
+ AffineExpr *rhs,
+ MLIRContext *context) {
+ if (auto *l = dyn_cast<AffineConstantExpr>(lhs))
+ if (auto *r = dyn_cast<AffineConstantExpr>(rhs))
+ return AffineConstantExpr::get(
+ (int64_t)llvm::divideCeil((uint64_t)l->getValue(),
+ (uint64_t)r->getValue()),
+ context);
+
+ 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);
+
+ 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.
+}
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index f919af3..f495bda 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -582,132 +582,93 @@
return *existing.first = res;
}
-AffineExpr *AffineAddExpr::get(AffineExpr *lhsOperand, AffineExpr *rhsOperand,
- MLIRContext *context) {
+/// Return a binary affine op expression with the specified op type and
+/// operands: if it doesn't exist, create it and store it; if it is already
+/// present, return from the list. The stored expressions are unique: they are
+/// constructed and stored in a simplified/canonicalized form. The result after
+/// simplification could be any form of affine expression.
+AffineExpr *AffineBinaryOpExpr::get(AffineExpr::Kind kind, AffineExpr *lhs,
+ AffineExpr *rhs, MLIRContext *context) {
auto &impl = context->getImpl();
// Check if we already have this affine expression.
- auto keyValue = std::make_tuple((unsigned)Kind::Add, lhsOperand, rhsOperand);
+ auto keyValue = std::make_tuple((unsigned)kind, lhs, rhs);
auto *&result = impl.affineExprs[keyValue];
// If we already have it, return that value.
if (result)
return result;
- // Use the simplified expression if it can be simplified.
- result = AffineAddExpr::simplify(lhsOperand, rhsOperand, context);
-
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineAddExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineAddExpr(lhsOperand, rhsOperand);
+ // Simplify the expression if possible.
+ AffineExpr *simplified;
+ switch (kind) {
+ case Kind::Add:
+ simplified = AffineBinaryOpExpr::simplifyAdd(lhs, rhs, context);
+ break;
+ case Kind::Sub:
+ simplified = AffineBinaryOpExpr::simplifySub(lhs, rhs, context);
+ break;
+ case Kind::Mul:
+ simplified = AffineBinaryOpExpr::simplifyMul(lhs, rhs, context);
+ break;
+ case Kind::FloorDiv:
+ simplified = AffineBinaryOpExpr::simplifyFloorDiv(lhs, rhs, context);
+ break;
+ case Kind::CeilDiv:
+ simplified = AffineBinaryOpExpr::simplifyCeilDiv(lhs, rhs, context);
+ break;
+ case Kind::Mod:
+ simplified = AffineBinaryOpExpr::simplifyMod(lhs, rhs, context);
+ break;
+ default:
+ llvm_unreachable("unexpected binary affine expr");
}
+
+ // If simplified to a non-binary affine op expr, don't store it.
+ if (simplified && !isa<AffineBinaryOpExpr>(simplified)) {
+ // 'affineExprs' only contains uniqued AffineBinaryOpExpr's.
+ return simplified;
+ }
+
+ if (simplified)
+ // We know that it's a binary op expression.
+ return result = simplified;
+
+ // On the first use, we allocate them into the bump pointer.
+ result = impl.allocator.Allocate<AffineBinaryOpExpr>();
+ // Initialize the memory using placement new.
+ new (result) AffineBinaryOpExpr(kind, lhs, rhs);
return result;
}
-AffineSubExpr *AffineSubExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- auto &impl = context->getImpl();
-
- // Check if we already have this affine expression.
- auto keyValue = std::make_tuple((unsigned)Kind::Sub, lhsOperand, rhsOperand);
- auto *&result = impl.affineExprs[keyValue];
-
- // If we already have it, return that value.
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineSubExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineSubExpr(lhsOperand, rhsOperand);
- }
- return cast<AffineSubExpr>(result);
+AffineExpr *AffineAddExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::Add, lhs, rhs, context);
}
-AffineMulExpr *AffineMulExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- auto &impl = context->getImpl();
-
- // Check if we already have this affine expression.
- const auto keyValue =
- std::make_tuple((unsigned)Kind::Mul, lhsOperand, rhsOperand);
- auto *&result = impl.affineExprs[keyValue];
-
- // If we already have it, return that value.
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineMulExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineMulExpr(lhsOperand, rhsOperand);
- }
- return cast<AffineMulExpr>(result);
+AffineExpr *AffineSubExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::Sub, lhs, rhs, context);
}
-AffineFloorDivExpr *AffineFloorDivExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- auto &impl = context->getImpl();
-
- // Check if we already have this affine expression.
- auto keyValue =
- std::make_tuple((unsigned)Kind::FloorDiv, lhsOperand, rhsOperand);
- auto *&result = impl.affineExprs[keyValue];
-
- // If we already have it, return that value.
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineFloorDivExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineFloorDivExpr(lhsOperand, rhsOperand);
- }
- return cast<AffineFloorDivExpr>(result);
+AffineExpr *AffineMulExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::Mul, lhs, rhs, context);
}
-AffineCeilDivExpr *AffineCeilDivExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- auto &impl = context->getImpl();
-
- // Check if we already have this affine expression.
- auto keyValue =
- std::make_tuple((unsigned)Kind::CeilDiv, lhsOperand, rhsOperand);
- auto *&result = impl.affineExprs[keyValue];
-
- // If we already have it, return that value.
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineCeilDivExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineCeilDivExpr(lhsOperand, rhsOperand);
- }
- return cast<AffineCeilDivExpr>(result);
+AffineExpr *AffineFloorDivExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::FloorDiv, lhs, rhs, context);
}
-AffineModExpr *AffineModExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- auto &impl = context->getImpl();
+AffineExpr *AffineCeilDivExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::CeilDiv, lhs, rhs, context);
+}
- // Check if we already have this affine expression.
- auto keyValue = std::make_tuple((unsigned)Kind::Mod, lhsOperand, rhsOperand);
- auto *&result = impl.affineExprs[keyValue];
-
- // If we already have it, return that value.
- if (!result) {
- // On the first use, we allocate them into the bump pointer.
- result = impl.allocator.Allocate<AffineModExpr>();
-
- // Initialize the memory using placement new.
- new (result) AffineModExpr(lhsOperand, rhsOperand);
- }
- return cast<AffineModExpr>(result);
+AffineExpr *AffineModExpr::get(AffineExpr *lhs, AffineExpr *rhs,
+ MLIRContext *context) {
+ return AffineBinaryOpExpr::get(Kind::Mod, lhs, rhs, context);
}
AffineDimExpr *AffineDimExpr::get(unsigned position, MLIRContext *context) {
diff --git a/test/IR/parser-affine-map.mlir b/test/IR/parser-affine-map.mlir
index 05db858..7f55f69 100644
--- a/test/IR/parser-affine-map.mlir
+++ b/test/IR/parser-affine-map.mlir
@@ -15,7 +15,7 @@
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + s0), d1)
#hello_world4 = (i, j) [s0] -> (i + s0, j)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> ((1 + d0), d1)
+; CHECK: #{{[0-9]+}} = (d0, d1) -> ((d0 + 1), d1)
#hello_world5 = (i, j) -> (1+i, j)
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + s0), (d1 + 5))
@@ -24,7 +24,7 @@
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((d0 + d1) + s0), d1)
#hello_world7 = (i, j) [s0] -> (i + j + s0, j)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((((5 + d0) + d1) + s0), d1)
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((((d0 + 5) + d1) + s0), d1)
#hello_world8 = (i, j) [s0] -> (5 + i + j + s0, j)
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((d0 + d1) + 5), d1)
@@ -33,13 +33,13 @@
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + (d1 + 5)), d1)
#hello_world10 = (i, j) [s0] -> (i + (j + 5), j)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((2 * d0), (3 * d1))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((d0 * 2), (d1 * 3))
#hello_world11 = (i, j) [s0] -> (2*i, 3*j)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((d0 + (2 * 6)) + (5 * (d1 + (s0 * 3)))), d1)
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((d0 + 12) + ((d1 + (s0 * 3)) * 5)), d1)
#hello_world12 = (i, j) [s0] -> (i + 2*6 + 5*(j+s0*3), j)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((5 * d0) + d1), d1)
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (((d0 * 5) + d1), d1)
#hello_world13 = (i, j) [s0] -> (5*i + j, j)
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((d0 + d1), d1)
@@ -51,22 +51,22 @@
; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (d0, 0)
#hello_world16 = (i, j) [s1] -> (i, 0)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (d0, (s0 * d1))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (d0, (d1 * s0))
#hello_world17 = (i, j) [s0] -> (i, s0*j)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> (d0, ((3 * d0) + d1))
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (d0, ((d0 * 3) + d1))
#hello_world19 = (i, j) -> (i, 3*i + j)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> (d0, (d0 + (3 * d1)))
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (d0, (d0 + (d1 * 3)))
#hello_world20 = (i, j) -> (i, i + 3*j)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (d0, ((2 + (((s0 * s0) * 9) * d0)) + 1))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> (d0, (((d0 * ((s0 * s0) * 9)) + 2) + 1))
#hello_world18 = (i, j) [N] -> (i, 2 + N*N*9*i + 1)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> (1, ((d0 + (3 * d1)) + 5))
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (1, ((d0 + (d1 * 3)) + 5))
#hello_world21 = (i, j) -> (1, i + 3*j + 5)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((5 * s0), ((d0 + (3 * d1)) + (5 * d0)))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0] -> ((s0 * 5), ((d0 + (d1 * 3)) + (d0 * 5)))
#hello_world22 = (i, j) [s0] -> (5*s0, i + 3*j + 5*i)
; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> ((d0 * (s0 * s1)), d1)
@@ -84,13 +84,13 @@
; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> (d0, ((d0 - d1) - 5))
#hello_world29 = (i, j) [s0, s1] -> (i, i - j - 5)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> (d0, ((d0 - (s1 * d1)) + 2))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> (d0, ((d0 - (d1 * s1)) + 2))
#hello_world30 = (i, j) [M, N] -> (i, i - N*j + 2)
-; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> (((-1 * 5) * d0), ((-1 * 3) * d1), (-1 * 2), ((-1 * 1) * (d0 + d1)), ((-1 * 1) * s0))
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> ((d0 * -5), (d1 * -3), -2, ((d0 + d1) * -1), (s0 * -1))
#hello_world32 = (i, j) [s0, s1] -> (-5*i, -3*j, -2, -1*(i+j), -1*s0)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> ((((-1 * 2) + (-1 * 5)) - (-1 * 3)), ((-1 * 1) * d0))
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (-4, (d0 * -1))
#hello_world33 = (i, j) -> (-2+-5-(-3), -1*i)
; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> (d0, (d1 floordiv s0), (d1 mod s0))
@@ -99,5 +99,14 @@
; CHECK: #{{[0-9]+}} = (d0, d1, d2) [s0, s1, s2] -> (((((d0 * s1) * s2) + (d1 * s1)) + d2))
#hello_world35 = (i, j, k) [s0, s1, s2] -> (i*s1*s2 + j*s1 + k)
-; CHECK: #{{[0-9]+}} = (d0, d1) -> (2, 8)
-#hello_world36 = (i, j) -> (1+1, 5+3)
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (8, 4, 1, 3, 2, 4)
+#hello_world36 = (i, j) -> (5+3, 2*2, 8-7, 100 floordiv 32, 5 mod 3, 10 ceildiv 3)
+
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (4, 11, 512, 15)
+#hello_world37 = (i, j) -> (5 mod 3 + 2, 5*3 - 4, 128 * (500 ceildiv 128), 40 floordiv 7 * 3)
+
+; CHECK: #{{[0-9]+}} = (d0, d1) -> (((d0 * 2) + 1), (d1 + 2))
+#hello_world38 = (i, j) -> (1 + i*2, 2 + j)
+
+; CHECK: #{{[0-9]+}} = (d0, d1) [s0, s1] -> ((d0 * s0), (d0 + s0), (d0 + 2), (d1 * 2), (s1 * 2), (s0 + 2))
+#hello_world39 = (i, j) [M, N] -> (i*M, M + i, 2+i, j*2, N*2, 2 + M)