[WIP] Sketching IR and parsing support for affine maps, affine expressions
Run test case:
$ mlir-opt test/IR/parser-affine-map.mlir
test/IR/parser-affine-map.mlir:3:30: error: expect '(' at start of map range
#hello_world2 (i, j) [s0] -> i+s0, j)
^
PiperOrigin-RevId: 202736856
diff --git a/lib/IR/AffineExpr.cpp b/lib/IR/AffineExpr.cpp
index 63d8d7d..87901f0 100644
--- a/lib/IR/AffineExpr.cpp
+++ b/lib/IR/AffineExpr.cpp
@@ -18,6 +18,3 @@
#include "mlir/IR/AffineExpr.h"
using namespace mlir;
-
-AffineExpr::AffineExpr() {
-}
diff --git a/lib/IR/AffineMap.cpp b/lib/IR/AffineMap.cpp
index 83d1d23..8631751 100644
--- a/lib/IR/AffineMap.cpp
+++ b/lib/IR/AffineMap.cpp
@@ -17,15 +17,13 @@
#include "mlir/IR/AffineMap.h"
#include "llvm/ADT/StringRef.h"
-#include "llvm/Support/raw_ostream.h"
using namespace mlir;
-AffineMap::AffineMap(unsigned dimCount, unsigned symbolCount)
- : dimCount(dimCount), symbolCount(symbolCount) {
-}
-
-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)
+// 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)
}
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 639000b..067697c 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -20,13 +20,15 @@
//
//===----------------------------------------------------------------------===//
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/CFGFunction.h"
#include "mlir/IR/MLFunction.h"
#include "mlir/IR/Module.h"
#include "mlir/IR/Types.h"
#include "mlir/Support/STLExtras.h"
-#include "llvm/Support/raw_ostream.h"
#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/raw_ostream.h"
using namespace mlir;
@@ -162,6 +164,15 @@
print(llvm::errs());
}
+void AffineExpr::print(raw_ostream &os) const {
+ // TODO(bondhugula): print out affine expression
+}
+
+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)
+}
+
void BasicBlock::print(raw_ostream &os) const {
CFGFunctionState state(getFunction(), os);
state.print();
@@ -208,6 +219,8 @@
}
void Module::print(raw_ostream &os) const {
+ for (auto *map : affineMapList)
+ map->print(os);
for (auto *fn : functionList)
fn->print(os);
}
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index 5f2bd8e..7c1112b 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -17,6 +17,8 @@
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Identifier.h"
+#include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Types.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/DenseSet.h"
@@ -44,6 +46,23 @@
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>;
+ 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);
+ }
+
+ static bool isEqual(const KeyTy &lhs, const FunctionType *rhs) {
+ // TODO(bondhugula)
+ return false;
+ }
+};
+
struct VectorTypeKeyInfo : DenseMapInfo<VectorType*> {
// Vectors are uniqued based on their element type and shape.
using KeyTy = std::pair<Type*, ArrayRef<unsigned>>;
@@ -97,6 +116,10 @@
// Primitive type uniquing.
PrimitiveType *primitives[int(TypeKind::LAST_PRIMITIVE_TYPE)+1] = { nullptr };
+ // Affine map uniquing.
+ using AffineMapSet = DenseSet<AffineMap *, AffineMapKeyInfo>;
+ AffineMapSet affineMaps;
+
/// Function type uniquing.
using FunctionTypeSet = DenseSet<FunctionType*, FunctionTypeKeyInfo>;
FunctionTypeSet functions;
@@ -316,3 +339,52 @@
// Cache and return it.
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,
+ MLIRContext *context) {
+ // TODO(bondhugula)
+ return new AffineMap(dimCount, symbolCount, exprs);
+}
+
+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);
+}
+
+AffineAddExpr *AffineAddExpr::get(AffineExpr *lhsOperand,
+ AffineExpr *rhsOperand,
+ MLIRContext *context) {
+ // TODO(bondhugula): allocate this through context
+ // FIXME
+ return new AffineAddExpr(lhsOperand, rhsOperand);
+}
+
+// TODO(bondhugula): add functions for AffineMulExpr, mod, floordiv, ceildiv
+
+AffineDimExpr *AffineDimExpr::get(unsigned position, MLIRContext *context) {
+ // TODO(bondhugula): complete this
+ // FIXME: this should be POD
+ return new AffineDimExpr(position);
+}
+
+AffineSymbolExpr *AffineSymbolExpr::get(unsigned position,
+ MLIRContext *context) {
+ // TODO(bondhugula): complete this
+ // FIXME: this should be POD
+ return new AffineSymbolExpr(position);
+}
+
+AffineConstantExpr *AffineConstantExpr::get(int64_t constant,
+ MLIRContext *context) {
+ // TODO(bondhugula): complete this
+ // FIXME: this should be POD
+ return new AffineConstantExpr(constant);
+}