Adds MemRef type and adds support for parsing memref affine map composition.
PiperOrigin-RevId: 204756982
diff --git a/include/mlir/IR/Types.h b/include/mlir/IR/Types.h
index e499e56..cf346e0 100644
--- a/include/mlir/IR/Types.h
+++ b/include/mlir/IR/Types.h
@@ -22,9 +22,10 @@
#include "llvm/ADT/ArrayRef.h"
namespace mlir {
- class MLIRContext;
- class PrimitiveType;
- class IntegerType;
+class AffineMap;
+class MLIRContext;
+class PrimitiveType;
+class IntegerType;
/// Instances of the Type class are immutable, uniqued, immortal, and owned by
/// MLIRContext. As such, they are passed around by raw non-const pointer.
@@ -52,8 +53,7 @@
Vector,
RankedTensor,
UnrankedTensor,
-
- // TODO: MemRef types.
+ MemRef,
};
/// Return the classification for this type.
@@ -279,6 +279,55 @@
UnrankedTensorType(Type *elementType, MLIRContext *context);
};
+/// MemRef types represent a region of memory that have a shape with a fixed
+/// number of dimensions. Each shape element can be a positive integer or
+/// unknown (represented by any negative integer). MemRef types also have an
+/// affine map composition, represented as an array AffineMap pointers.
+// TODO: Use -1 for unknown dimensions (rather than arbitrary negative numbers).
+class MemRefType : public Type {
+public:
+ /// Get or create a new MemRefType based on shape, element type, affine
+ /// map composition, and memory space.
+ static MemRefType *get(ArrayRef<int> shape, Type *elementType,
+ ArrayRef<AffineMap*> affineMapComposition,
+ unsigned memorySpace);
+
+ /// Returns an array of memref shape dimension sizes.
+ ArrayRef<int> getShape() const {
+ return ArrayRef<int>(shapeElements, getSubclassData());
+ }
+
+ /// Returns the elemental type for this memref shape.
+ Type *getElementType() const { return elementType; }
+
+ /// Returns an array of affine map pointers representing the memref affine
+ /// map composition.
+ ArrayRef<AffineMap*> getAffineMaps() const;
+
+ /// Returns the memory space in which data referred to by this memref resides.
+ unsigned getMemorySpace() const { return memorySpace; }
+
+ static bool classof(const Type *type) {
+ return type->getKind() == Kind::MemRef;
+ }
+
+private:
+ /// The type of each scalar element of the memref.
+ Type *elementType;
+ /// An array of integers which stores the shape dimension sizes.
+ const int *shapeElements;
+ /// The number of affine maps in the 'affineMapList' array.
+ unsigned numAffineMaps;
+ /// List of affine maps in affine map composition.
+ AffineMap *const *const affineMapList;
+ /// Memory space in which data referenced by memref resides.
+ unsigned memorySpace;
+
+ MemRefType(ArrayRef<int> shape, Type *elementType,
+ ArrayRef<AffineMap*> affineMapList, unsigned memorySpace,
+ MLIRContext *context);
+};
+
} // end namespace mlir
#endif // MLIR_IR_TYPES_H
diff --git a/lib/IR/AsmPrinter.cpp b/lib/IR/AsmPrinter.cpp
index 2dd2a3b..7347a81 100644
--- a/lib/IR/AsmPrinter.cpp
+++ b/lib/IR/AsmPrinter.cpp
@@ -303,7 +303,10 @@
llvm::errs() << "\n";
}
-void AffineMap::dump() const { print(llvm::errs()); }
+void AffineMap::dump() const {
+ print(llvm::errs());
+ llvm::errs() << "\n";
+}
void AffineExpr::dump() const {
print(llvm::errs());
@@ -393,7 +396,6 @@
os << ")";
if (!isBounded()) {
- os << "\n";
return;
}
@@ -401,7 +403,7 @@
os << " size (";
interleave(getRangeSizes(), [&](AffineExpr *expr) { os << *expr; },
[&]() { os << ", "; });
- os << ")\n";
+ os << ")";
}
void BasicBlock::print(raw_ostream &os) const {
@@ -449,6 +451,7 @@
for (auto *map : affineMapList) {
os << "#" << id++ << " = ";
map->print(os);
+ os << '\n';
}
for (auto *fn : functionList)
fn->print(os);
diff --git a/lib/IR/MLIRContext.cpp b/lib/IR/MLIRContext.cpp
index 4448b16..eb06e94 100644
--- a/lib/IR/MLIRContext.cpp
+++ b/lib/IR/MLIRContext.cpp
@@ -113,6 +113,30 @@
}
};
+struct MemRefTypeKeyInfo : DenseMapInfo<MemRefType*> {
+ // MemRefs are uniqued based on their element type, shape, affine map
+ // composition, and memory space.
+ using KeyTy = std::tuple<Type*, ArrayRef<int>, ArrayRef<AffineMap*>,
+ unsigned>;
+ using DenseMapInfo<MemRefType*>::getHashValue;
+ using DenseMapInfo<MemRefType*>::isEqual;
+
+ static unsigned getHashValue(KeyTy key) {
+ return hash_combine(
+ DenseMapInfo<Type*>::getHashValue(std::get<0>(key)),
+ hash_combine_range(std::get<1>(key).begin(), std::get<1>(key).end()),
+ hash_combine_range(std::get<2>(key).begin(), std::get<2>(key).end()),
+ std::get<3>(key));
+ }
+
+ static bool isEqual(const KeyTy &lhs, const MemRefType *rhs) {
+ if (rhs == getEmptyKey() || rhs == getTombstoneKey())
+ return false;
+ return lhs == std::make_tuple(rhs->getElementType(), rhs->getShape(),
+ rhs->getAffineMaps(), rhs->getMemorySpace());
+ }
+};
+
struct ArrayAttrKeyInfo : DenseMapInfo<ArrayAttr*> {
// Array attributes are uniqued based on their elements.
using KeyTy = ArrayRef<Attribute*>;
@@ -195,6 +219,10 @@
/// Unranked tensor type uniquing.
DenseMap<Type*, UnrankedTensorType*> unrankedTensors;
+ /// MemRef type uniquing.
+ using MemRefTypeSet = DenseSet<MemRefType*, MemRefTypeKeyInfo>;
+ MemRefTypeSet memrefs;
+
// Attribute uniquing.
BoolAttr *boolAttrs[2] = { nullptr };
DenseMap<int64_t, IntegerAttr*> integerAttrs;
@@ -403,6 +431,39 @@
return result;
}
+MemRefType *MemRefType::get(ArrayRef<int> shape, Type *elementType,
+ ArrayRef<AffineMap*> affineMapComposition,
+ unsigned memorySpace) {
+ auto *context = elementType->getContext();
+ auto &impl = context->getImpl();
+
+ // Look to see if we already have this memref type.
+ auto key = std::make_tuple(elementType, shape, affineMapComposition,
+ memorySpace);
+ auto existing = impl.memrefs.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 *result = impl.allocator.Allocate<MemRefType>();
+
+ // Copy the shape into the bump pointer.
+ shape = impl.copyInto(shape);
+
+ // Copy the affine map composition into the bump pointer.
+ // TODO(andydavis) Assert that the structure of the composition is valid.
+ affineMapComposition = impl.copyInto(ArrayRef<AffineMap*>(
+ affineMapComposition));
+
+ // Initialize the memory using placement new.
+ new (result) MemRefType(shape, elementType, affineMapComposition, memorySpace,
+ context);
+ // Cache and return it.
+ return *existing.first = result;
+}
+
//===----------------------------------------------------------------------===//
// Attribute uniquing
//===----------------------------------------------------------------------===//
diff --git a/lib/IR/Types.cpp b/lib/IR/Types.cpp
index 0860772..1b7d1a6 100644
--- a/lib/IR/Types.cpp
+++ b/lib/IR/Types.cpp
@@ -16,6 +16,7 @@
// =============================================================================
#include "mlir/IR/Types.h"
+#include "mlir/IR/AffineMap.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Support/STLExtras.h"
using namespace mlir;
@@ -51,6 +52,19 @@
: TensorType(Kind::UnrankedTensor, elementType, context) {
}
+MemRefType::MemRefType(ArrayRef<int> shape, Type *elementType,
+ ArrayRef<AffineMap*> affineMapList,
+ unsigned memorySpace, MLIRContext *context)
+ : Type(Kind::MemRef, context, shape.size()),
+ elementType(elementType), shapeElements(shape.data()),
+ numAffineMaps(affineMapList.size()), affineMapList(affineMapList.data()),
+ memorySpace(memorySpace) {
+}
+
+ArrayRef<AffineMap*> MemRefType::getAffineMaps() const {
+ return ArrayRef<AffineMap*>(affineMapList, numAffineMaps);
+}
+
void Type::print(raw_ostream &os) const {
switch (getKind()) {
case Kind::AffineInt: os << "affineint"; return;
@@ -109,6 +123,25 @@
os << "tensor<??" << *v->getElementType() << '>';
return;
}
+ case Kind::MemRef: {
+ auto *v = cast<MemRefType>(this);
+ os << "memref<";
+ for (auto dim : v->getShape()) {
+ if (dim < 0)
+ os << '?';
+ else
+ os << dim;
+ os << 'x';
+ }
+ os << *v->getElementType();
+ for (auto map : v->getAffineMaps()) {
+ os << ", ";
+ map->print(os);
+ }
+ os << ", " << v->getMemorySpace();
+ os << '>';
+ return;
+ }
}
}
diff --git a/lib/Parser/Parser.cpp b/lib/Parser/Parser.cpp
index e1d184d..626b92f 100644
--- a/lib/Parser/Parser.cpp
+++ b/lib/Parser/Parser.cpp
@@ -132,9 +132,10 @@
return true;
}
- ParseResult parseCommaSeparatedList(Token::Kind rightToken,
- const std::function<ParseResult()> &parseElement,
- bool allowEmptyList = true);
+ ParseResult parseCommaSeparatedList(
+ Token::Kind rightToken,
+ const std::function<ParseResult()> &parseElement,
+ bool allowEmptyList = true);
// We have two forms of parsing methods - those that return a non-null
// pointer on success, and those that return a ParseResult to indicate whether
@@ -158,6 +159,7 @@
// Polyhedral structures.
AffineMap *parseAffineMapInline();
+ AffineMap *parseAffineMapReference();
// SSA
ParseResult parseSSAUse();
@@ -414,14 +416,50 @@
if (!elementType)
return nullptr;
- // TODO: Parse semi-affine-map-composition.
- // TODO: Parse memory-space.
+ if (!consumeIf(Token::comma))
+ return (emitError("expected ',' in memref type"), nullptr);
- if (!consumeIf(Token::greater))
- return (emitError("expected '>' in memref type"), nullptr);
+ // Parse semi-affine-map-composition.
+ SmallVector<AffineMap*, 2> affineMapComposition;
+ unsigned memorySpace;
+ bool parsedMemorySpace = false;
- // FIXME: Add an IR representation for memref types.
- return builder.getIntegerType(1);
+ auto parseElt = [&]() -> ParseResult {
+ if (getToken().is(Token::integer)) {
+ // Parse memory space.
+ if (parsedMemorySpace)
+ return emitError("multiple memory spaces specified in memref type");
+ auto v = getToken().getUnsignedIntegerValue();
+ if (!v.hasValue())
+ return emitError("invalid memory space in memref type");
+ memorySpace = v.getValue();
+ consumeToken(Token::integer);
+ parsedMemorySpace = true;
+ } else {
+ // Parse affine map.
+ if (parsedMemorySpace)
+ return emitError("affine map after memory space in memref type");
+ auto* affineMap = parseAffineMapReference();
+ if (affineMap == nullptr)
+ return ParseFailure;
+ affineMapComposition.push_back(affineMap);
+ }
+ return ParseSuccess;
+ };
+
+ // Parse comma separated list of affine maps, followed by memory space.
+ if (parseCommaSeparatedList(Token::greater, parseElt,
+ /*allowEmptyList=*/false)) {
+ return nullptr;
+ }
+ // Check that MemRef type specifies at least one affine map in composition.
+ if (affineMapComposition.empty())
+ return (emitError("expected semi-affine-map in memref type"), nullptr);
+ if (!parsedMemorySpace)
+ return (emitError("expected memory space in memref type"), nullptr);
+
+ return MemRefType::get(dimensions, elementType, affineMapComposition,
+ memorySpace);
}
/// Parse a function type.
@@ -1106,6 +1144,20 @@
return AffineMapParser(state).parseAffineMapInline();
}
+AffineMap *Parser::parseAffineMapReference() {
+ if (getToken().is(Token::hash_identifier)) {
+ // Parse affine map identifier and verify that it exists.
+ StringRef affineMapId = getTokenSpelling().drop_front();
+ if (getState().affineMapDefinitions.count(affineMapId) == 0)
+ return (emitError("undefined affine map id '" + affineMapId + "'"),
+ nullptr);
+ consumeToken(Token::hash_identifier);
+ return getState().affineMapDefinitions[affineMapId];
+ }
+ // Try to parse inline affine map.
+ return parseAffineMapInline();
+}
+
//===----------------------------------------------------------------------===//
// SSA
//===----------------------------------------------------------------------===//
diff --git a/test/IR/parser-errors.mlir b/test/IR/parser-errors.mlir
index 27a53bd..ff03b51 100644
--- a/test/IR/parser-errors.mlir
+++ b/test/IR/parser-errors.mlir
@@ -13,6 +13,47 @@
extfunc @nestedtensor(tensor<tensor<i8>>) -> () // expected-error {{expected type}}
// -----
+// Test no comma in memref type.
+// TODO(andydavis) Fix this test if we decide to allow empty affine map to
+// imply identity affine map.
+extfunc @memrefs(memref<2x4xi8>) ; expected-error {{expected ',' in memref type}}
+
+// -----
+// Test no map in memref type.
+extfunc @memrefs(memref<2x4xi8, >) ; expected-error {{expected list element}}
+
+// -----
+// Test non-existent map in memref type.
+extfunc @memrefs(memref<2x4xi8, #map7>) ; expected-error {{undefined affine map id 'map7'}}
+
+// -----
+// Test non hash identifier in memref type.
+extfunc @memrefs(memref<2x4xi8, %map7>) ; expected-error {{expected '(' at start of dimensional identifiers list}}
+
+// -----
+// Test non-existent map in map composition of memref type.
+#map0 = (d0, d1) -> (d0, d1)
+
+extfunc @memrefs(memref<2x4xi8, #map0, #map8>) ; expected-error {{undefined affine map id 'map8'}}
+
+// -----
+// Test multiple memory space error.
+#map0 = (d0, d1) -> (d0, d1)
+extfunc @memrefs(memref<2x4xi8, #map0, 1, 2>) ; expected-error {{multiple memory spaces specified in memref type}}
+
+// -----
+// Test affine map after memory space.
+#map0 = (d0, d1) -> (d0, d1)
+#map1 = (d0, d1) -> (d0, d1)
+
+extfunc @memrefs(memref<2x4xi8, #map0, 1, #map1>) ; expected-error {{affine map after memory space in memref type}}
+
+// -----
+// Test no memory space error.
+#map0 = (d0, d1) -> (d0, d1)
+extfunc @memrefs(memref<2x4xi8, #map0>) ; expected-error {{expected memory space in memref type}}
+
+// -----
cfgfunc @foo()
cfgfunc @bar() // expected-error {{expected '{' in CFG function}}
diff --git a/test/IR/parser.mlir b/test/IR/parser.mlir
index 4ba9aa4..069c267 100644
--- a/test/IR/parser.mlir
+++ b/test/IR/parser.mlir
@@ -3,6 +3,11 @@
//
// RUN: %S/../../mlir-opt %s -o - | FileCheck %s
+#map0 = (d0, d1, d2, d3, d4) [s0] -> (d0, d1, d2, d3, d4)
+#map1 = (d0) -> (d0)
+#map2 = (d0, d1, d2) -> (d0, d1, d2)
+#map3 = (d0, d1, d2) -> (d1, d0, d2)
+#map4 = (d0, d1, d2) -> (d2, d1, d0)
// CHECK: extfunc @foo(i32, i64) -> f32
extfunc @foo(i32, i64) -> f32
@@ -27,12 +32,31 @@
extfunc @tensors(tensor<?? f32>, tensor<?? vector<2x4xf32>>,
tensor<1x?x4x?x?xaffineint>, tensor<i8>)
-// CHECK: extfunc @memrefs(i1, i1)
-extfunc @memrefs(memref<1x?x4x?x?xaffineint>, memref<i8>)
+// TODO(andydavis) Add support to outline affine maps for these cases.
+// CHECK: extfunc @memrefs(memref<1x?x4x?x?xaffineint, (d0, d1, d2, d3, d4) [s0] -> (d0, d1, d2, d3, d4), 0>, memref<i8, (d0) -> (d0), 0>)
+extfunc @memrefs(memref<1x?x4x?x?xaffineint, #map0, 0>, memref<i8, #map1, 0>)
-// CHECK: extfunc @functions((i1, i1) -> (), () -> ())
-extfunc @functions((memref<1x?x4x?x?xaffineint>, memref<i8>) -> (), ()->())
+// Test memref affine map compositions.
+// CHECK: extfunc @memrefs2(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), 1>)
+extfunc @memrefs2(memref<2x4x8xi8, #map2, 1>)
+
+// CHECK: extfunc @memrefs23(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), (d0, d1, d2) -> (d1, d0, d2), 0>)
+extfunc @memrefs23(memref<2x4x8xi8, #map2, #map3, 0>)
+
+// CHECK: extfunc @memrefs234(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), (d0, d1, d2) -> (d1, d0, d2), (d0, d1, d2) -> (d2, d1, d0), 3>)
+extfunc @memrefs234(memref<2x4x8xi8, #map2, #map3, #map4, 3>)
+
+// Test memref inline affine map compositions.
+
+// CHECK: extfunc @memrefs2(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), 0>)
+extfunc @memrefs2(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), 0>)
+
+// CHECK: extfunc @memrefs23(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), (d0, d1, d2) -> (d1, d0, d2), 1>)
+extfunc @memrefs23(memref<2x4x8xi8, (d0, d1, d2) -> (d0, d1, d2), (d0, d1, d2) -> (d1, d0, d2), 1>)
+
+// CHECK: extfunc @functions((memref<1x?x4x?x?xaffineint, (d0, d1, d2, d3, d4) [s0] -> (d0, d1, d2, d3, d4), 0>, memref<i8, (d0) -> (d0), 0>) -> (), () -> ())
+extfunc @functions((memref<1x?x4x?x?xaffineint, #map0, 0>, memref<i8, #map1, 0>) -> (), ()->())
// CHECK-LABEL: cfgfunc @simpleCFG(i32, f32) {
cfgfunc @simpleCFG(i32, f32) {