Parsing support for affine maps and affine expressions
A recursive descent parser for affine maps/expressions with operator precedence and
associativity. (While on this, sketch out uniqui'ing functionality for affine maps
and affine binary op expressions (partly).)
PiperOrigin-RevId: 203222063
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index 8631751..cba3094 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -20,10 +20,7 @@
using namespace mlir;
-// TODO(clattner): make this ctor take an LLVMContext. This will eventually
-// copy the elements into the context.
-AffineMap::AffineMap(unsigned dimCount, unsigned symbolCount,
- ArrayRef<AffineExpr *> exprs)
- : numDims(dimCount), numSymbols(symbolCount), exprs(exprs) {
- // TODO(bondhugula)
-}
+AffineMap::AffineMap(unsigned numDims, unsigned numSymbols, unsigned numResults,
+ AffineExpr *const *results)
+ : numDims(numDims), numSymbols(numSymbols), numResults(numResults),
+ results(results) {}
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index b871066..b24eb2f 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -251,13 +251,95 @@
print(llvm::errs());
}
+void AffineExpr::dump() const {
+ print(llvm::errs());
+ llvm::errs() << "\n";
+}
+
+void AffineAddExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " + " << *getRightOperand() << ")";
+}
+
+void AffineSubExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " - " << *getRightOperand() << ")";
+}
+
+void AffineMulExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " * " << *getRightOperand() << ")";
+}
+
+void AffineModExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " mod " << *getRightOperand() << ")";
+}
+
+void AffineFloorDivExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " floordiv " << *getRightOperand() << ")";
+}
+
+void AffineCeilDivExpr::print(raw_ostream &os) const {
+ os << "(" << *getLeftOperand() << " ceildiv " << *getRightOperand() << ")";
+}
+
+void AffineSymbolExpr::print(raw_ostream &os) const {
+ os << "s" << getPosition();
+}
+
+void AffineDimExpr::print(raw_ostream &os) const { os << "d" << getPosition(); }
+
+void AffineConstantExpr::print(raw_ostream &os) const { os << getValue(); }
+
void AffineExpr::print(raw_ostream &os) const {
- // TODO(bondhugula): print out affine expression
+ switch (getKind()) {
+ case Kind::SymbolId:
+ return cast<AffineSymbolExpr>(this)->print(os);
+ case Kind::DimId:
+ return cast<AffineDimExpr>(this)->print(os);
+ case Kind::Constant:
+ return cast<AffineConstantExpr>(this)->print(os);
+ case Kind::Add:
+ return cast<AffineAddExpr>(this)->print(os);
+ case Kind::Sub:
+ return cast<AffineSubExpr>(this)->print(os);
+ case Kind::Mul:
+ return cast<AffineMulExpr>(this)->print(os);
+ case Kind::FloorDiv:
+ return cast<AffineFloorDivExpr>(this)->print(os);
+ case Kind::CeilDiv:
+ return cast<AffineCeilDivExpr>(this)->print(os);
+ case Kind::Mod:
+ return cast<AffineModExpr>(this)->print(os);
+ default:
+ os << "<unimplemented expr>";
+ return;
+ }
}
void AffineMap::print(raw_ostream &os) const {
- // TODO(andydavis) Print out affine map based on dimensionCount and
- // symbolCount: (d0, d1) [S0, S1] -> (d0 + S0, d1 + S1)
+ // Dimension identifiers.
+ os << "(";
+ for (int i = 0; i < (int)getNumDims() - 1; i++)
+ os << "d" << i << ", ";
+ if (getNumDims() >= 1)
+ os << "d" << getNumDims() - 1;
+ os << ")";
+
+ // Symbolic identifiers.
+ if (getNumSymbols() >= 1) {
+ os << " [";
+ for (int i = 0; i < (int)getNumSymbols() - 1; i++)
+ os << "s" << i << ", ";
+ if (getNumSymbols() >= 1)
+ os << "s" << getNumSymbols() - 1;
+ os << "]";
+ }
+
+ // AffineMap should have at least one result.
+ assert(!getResults().empty());
+ // Result affine expressions.
+ os << " -> (";
+ interleave(getResults(), [&](AffineExpr *expr) { os << *expr; },
+ [&]() { os << ", "; });
+ os << ")\n";
}
void BasicBlock::print(raw_ostream &os) const {
@@ -300,8 +382,11 @@
}
void Module::print(raw_ostream &os) const {
- for (auto *map : affineMapList)
+ unsigned id = 0;
+ for (auto *map : affineMapList) {
+ os << "#" << id++ << " = ";
map->print(os);
+ }
for (auto *fn : functionList)
fn->print(os);
}
@@ -309,4 +394,3 @@
void Module::dump() const {
print(llvm::errs());
}
-
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index 85ff432..b6e1faa 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -46,20 +46,27 @@
return lhs == KeyTy(rhs->getInputs(), rhs->getResults());
}
};
+
struct AffineMapKeyInfo : DenseMapInfo<AffineMap *> {
- // Affine maps are uniqued based on their arguments and affine expressions
- using KeyTy = std::pair<unsigned, unsigned>;
+ // Affine maps are uniqued based on their dim/symbol counts and affine
+ // expressions.
+ using KeyTy =
+ std::pair<std::pair<unsigned, unsigned>, ArrayRef<AffineExpr *>>;
using DenseMapInfo<AffineMap *>::getHashValue;
using DenseMapInfo<AffineMap *>::isEqual;
static unsigned getHashValue(KeyTy key) {
- // FIXME(bondhugula): placeholder for now
- return hash_combine(key.first, key.second);
+ return hash_combine(
+ key.first.first, key.first.second,
+ hash_combine_range(key.second.begin(), key.second.end()));
}
- static bool isEqual(const KeyTy &lhs, const FunctionType *rhs) {
- // TODO(bondhugula)
- return false;
+ static bool isEqual(const KeyTy &lhs, const AffineMap *rhs) {
+ if (rhs == getEmptyKey() || rhs == getTombstoneKey())
+ return false;
+ return lhs == KeyTy(std::pair<unsigned, unsigned>(rhs->getNumDims(),
+ rhs->getNumSymbols()),
+ rhs->getResults());
}
};
@@ -120,6 +127,15 @@
using AffineMapSet = DenseSet<AffineMap *, AffineMapKeyInfo>;
AffineMapSet affineMaps;
+ // Affine binary op expression uniquing. We don't need to unique dimensional
+ // or symbolic identifiers.
+ // std::tuple doesn't work with DesnseMap!
+ // DenseSet<AffineBinaryOpExpr *, AffineBinaryOpExprKeyInfo>;
+ // AffineExprSet affineExprs;
+ DenseMap<std::pair<unsigned, std::pair<AffineExpr *, AffineExpr *>>,
+ AffineBinaryOpExpr *>
+ affineExprs;
+
/// Integer type uniquing.
DenseMap<unsigned, IntegerType*> integers;
@@ -325,34 +341,101 @@
return existing.first->second = result;
}
-// TODO(bondhugula,andydavis): unique affine maps based on dim list,
-// symbol list and all affine expressions contained
-AffineMap *AffineMap::get(unsigned dimCount,
- unsigned symbolCount,
- ArrayRef<AffineExpr *> exprs,
+AffineMap *AffineMap::get(unsigned dimCount, unsigned symbolCount,
+ ArrayRef<AffineExpr *> results,
MLIRContext *context) {
- // TODO(bondhugula)
- return new AffineMap(dimCount, symbolCount, exprs);
+ // The number of results can't be zero.
+ assert(!results.empty());
+
+ auto &impl = context->getImpl();
+
+ // Check if we already have this affine map.
+ AffineMapKeyInfo::KeyTy key(
+ std::pair<unsigned, unsigned>(dimCount, symbolCount), results);
+ auto existing = impl.affineMaps.insert_as(nullptr, key);
+
+ // If we already have it, return that value.
+ if (!existing.second)
+ return *existing.first;
+
+ // On the first use, we allocate them into the bump pointer.
+ auto *res = impl.allocator.Allocate<AffineMap>();
+
+ // Copy the results into the bump pointer.
+ results = impl.copyInto(ArrayRef<AffineExpr *>(results));
+
+ // Initialize the memory using placement new.
+ new (res) AffineMap(dimCount, symbolCount, results.size(), results.data());
+
+ // Cache and return it.
+ return *existing.first = res;
}
+// TODO(bondhugula): complete uniqu'ing of remaining AffinExpr sub-classes
AffineBinaryOpExpr *AffineBinaryOpExpr::get(AffineExpr::Kind kind,
AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- // TODO(bondhugula): allocate this through context
- // FIXME
- return new AffineBinaryOpExpr(kind, lhsOperand, rhsOperand);
+ auto &impl = context->getImpl();
+
+ // Check if we already have this affine expression.
+ auto key = std::pair<unsigned, std::pair<AffineExpr *, AffineExpr *>>(
+ (unsigned)kind,
+ std::pair<AffineExpr *, AffineExpr *>(lhsOperand, rhsOperand));
+ auto *&result = impl.affineExprs[key];
+
+ // 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<AffineBinaryOpExpr>();
+
+ // Initialize the memory using placement new.
+ new (result) AffineBinaryOpExpr(kind, lhsOperand, rhsOperand);
+ }
+ return result;
}
AffineAddExpr *AffineAddExpr::get(AffineExpr *lhsOperand,
AffineExpr *rhsOperand,
MLIRContext *context) {
- // TODO(bondhugula): allocate this through context
- // FIXME
- return new AffineAddExpr(lhsOperand, rhsOperand);
+ return cast<AffineAddExpr>(
+ AffineBinaryOpExpr::get(Kind::Add, lhsOperand, rhsOperand, context));
}
-// TODO(bondhugula): add functions for AffineMulExpr, mod, floordiv, ceildiv
+AffineSubExpr *AffineSubExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ return cast<AffineSubExpr>(
+ AffineBinaryOpExpr::get(Kind::Sub, lhsOperand, rhsOperand, context));
+}
+
+AffineMulExpr *AffineMulExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ return cast<AffineMulExpr>(
+ AffineBinaryOpExpr::get(Kind::Mul, lhsOperand, rhsOperand, context));
+}
+
+AffineFloorDivExpr *AffineFloorDivExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ return cast<AffineFloorDivExpr>(
+ AffineBinaryOpExpr::get(Kind::FloorDiv, lhsOperand, rhsOperand, context));
+}
+
+AffineCeilDivExpr *AffineCeilDivExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ return cast<AffineCeilDivExpr>(
+ AffineBinaryOpExpr::get(Kind::CeilDiv, lhsOperand, rhsOperand, context));
+}
+
+AffineModExpr *AffineModExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ return cast<AffineModExpr>(
+ AffineBinaryOpExpr::get(Kind::Mod, lhsOperand, rhsOperand, context));
+}
AffineDimExpr *AffineDimExpr::get(unsigned position, MLIRContext *context) {
// TODO(bondhugula): complete this