Complete affine expr parsing support
- check for non-affine expressions
- handle negative numbers and negation of id's, expressions
- functions to check if a map is pure affine or semi-affine
- simplify/clean up affine map parsing code
- report more errors messages, more accurate error messages
PiperOrigin-RevId: 203773633
diff --git a/lib/IR/AffineExpr.cpp b/lib/IR/AffineExpr.cpp
index 87901f0..26447a0 100644
--- a/lib/IR/AffineExpr.cpp
+++ b/lib/IR/AffineExpr.cpp
@@ -16,5 +16,66 @@
// =============================================================================
#include "mlir/IR/AffineExpr.h"
+#include "mlir/Support/STLExtras.h"
using namespace mlir;
+
+bool AffineExpr::isSymbolic() const {
+ switch (getKind()) {
+ case Kind::Constant:
+ return true;
+ case Kind::DimId:
+ return false;
+ case Kind::SymbolId:
+ return true;
+
+ case Kind::Add:
+ case Kind::Sub:
+ case Kind::Mul:
+ case Kind::FloorDiv:
+ case Kind::CeilDiv:
+ case Kind::Mod:
+ return cast<AffineBinaryOpExpr>(this)->isSymbolic();
+ }
+}
+
+bool AffineExpr::isPureAffine() const {
+ switch (getKind()) {
+ case Kind::SymbolId:
+ return cast<AffineSymbolExpr>(this)->isPureAffine();
+ case Kind::DimId:
+ return cast<AffineDimExpr>(this)->isPureAffine();
+ case Kind::Constant:
+ return cast<AffineConstantExpr>(this)->isPureAffine();
+ case Kind::Add:
+ return cast<AffineAddExpr>(this)->isPureAffine();
+ case Kind::Sub:
+ return cast<AffineSubExpr>(this)->isPureAffine();
+ case Kind::Mul:
+ return cast<AffineMulExpr>(this)->isPureAffine();
+ case Kind::FloorDiv:
+ return cast<AffineFloorDivExpr>(this)->isPureAffine();
+ case Kind::CeilDiv:
+ return cast<AffineCeilDivExpr>(this)->isPureAffine();
+ case Kind::Mod:
+ return cast<AffineModExpr>(this)->isPureAffine();
+ }
+}
+
+bool AffineMulExpr::isPureAffine() const {
+ return lhsOperand->isPureAffine() && rhsOperand->isPureAffine() &&
+ (isa<AffineConstantExpr>(lhsOperand) ||
+ isa<AffineConstantExpr>(rhsOperand));
+}
+
+bool AffineFloorDivExpr::isPureAffine() const {
+ return lhsOperand->isPureAffine() && isa<AffineConstantExpr>(rhsOperand);
+}
+
+bool AffineCeilDivExpr::isPureAffine() const {
+ return lhsOperand->isPureAffine() && isa<AffineConstantExpr>(rhsOperand);
+}
+
+bool AffineModExpr::isPureAffine() const {
+ return lhsOperand->isPureAffine() && isa<AffineConstantExpr>(rhsOperand);
+}
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index cba3094..ccfd108 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -16,6 +16,7 @@
// =============================================================================
#include "mlir/IR/AffineMap.h"
+#include "mlir/IR/AffineExpr.h"
#include "llvm/ADT/StringRef.h"
using namespace mlir;
@@ -24,3 +25,15 @@
AffineExpr *const *results)
: 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);
+ return nullptr;
+ // TODO(someone): implement more simplification.
+}
+
+// TODO(bondhugula): implement simplify for remaining affine binary op expr's
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index fcfd55f..f2d1e4f 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -282,6 +282,8 @@
print(llvm::errs());
}
+void AffineMap::dump() const { print(llvm::errs()); }
+
void AffineExpr::dump() const {
print(llvm::errs());
llvm::errs() << "\n";
@@ -339,9 +341,6 @@
return cast<AffineCeilDivExpr>(this)->print(os);
case Kind::Mod:
return cast<AffineModExpr>(this)->print(os);
- default:
- os << "<unimplemented expr>";
- return;
}
}
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index 761f804..f919af3 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -171,8 +171,7 @@
// Affine binary op expression uniquing. Figure out uniquing of dimensional
// or symbolic identifiers.
- DenseMap<std::tuple<unsigned, AffineExpr *, AffineExpr *>,
- AffineBinaryOpExpr *>
+ DenseMap<std::tuple<unsigned, AffineExpr *, AffineExpr *>, AffineExpr *>
affineExprs;
/// Integer type uniquing.
@@ -583,68 +582,132 @@
return *existing.first = res;
}
-AffineBinaryOpExpr *AffineBinaryOpExpr::get(AffineExpr::Kind kind,
- AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
+AffineExpr *AffineAddExpr::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, lhsOperand, rhsOperand);
+ auto keyValue = std::make_tuple((unsigned)Kind::Add, lhsOperand, rhsOperand);
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<AffineBinaryOpExpr>();
+ result = impl.allocator.Allocate<AffineAddExpr>();
// Initialize the memory using placement new.
- new (result) AffineBinaryOpExpr(kind, lhsOperand, rhsOperand);
+ new (result) AffineAddExpr(lhsOperand, rhsOperand);
}
return result;
}
-// TODO(bondhugula): complete uniquing of remaining AffineExpr sub-classes.
-AffineAddExpr *AffineAddExpr::get(AffineExpr *lhsOperand,
- AffineExpr *rhsOperand,
- MLIRContext *context) {
- return cast<AffineAddExpr>(
- AffineBinaryOpExpr::get(Kind::Add, lhsOperand, rhsOperand, context));
-}
-
AffineSubExpr *AffineSubExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- return cast<AffineSubExpr>(
- AffineBinaryOpExpr::get(Kind::Sub, lhsOperand, rhsOperand, 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);
}
AffineMulExpr *AffineMulExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- return cast<AffineMulExpr>(
- AffineBinaryOpExpr::get(Kind::Mul, lhsOperand, rhsOperand, 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);
}
AffineFloorDivExpr *AffineFloorDivExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- return cast<AffineFloorDivExpr>(
- AffineBinaryOpExpr::get(Kind::FloorDiv, lhsOperand, rhsOperand, 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);
}
AffineCeilDivExpr *AffineCeilDivExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- return cast<AffineCeilDivExpr>(
- AffineBinaryOpExpr::get(Kind::CeilDiv, lhsOperand, rhsOperand, 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);
}
AffineModExpr *AffineModExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- return cast<AffineModExpr>(
- AffineBinaryOpExpr::get(Kind::Mod, lhsOperand, rhsOperand, context));
+ auto &impl = context->getImpl();
+
+ // 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);
}
AffineDimExpr *AffineDimExpr::get(unsigned position, MLIRContext *context) {