diff --git a/mlir/include/mlir/IR/BuiltinTypes.h b/mlir/include/mlir/IR/BuiltinTypes.h --- a/mlir/include/mlir/IR/BuiltinTypes.h +++ b/mlir/include/mlir/IR/BuiltinTypes.h @@ -22,22 +22,8 @@ class Identifier; class IndexType; class IntegerType; -class Location; -class MLIRContext; class TypeRange; -namespace detail { - -struct BaseMemRefTypeStorage; -struct MemRefTypeStorage; -struct RankedTensorTypeStorage; -struct ShapedTypeStorage; -struct UnrankedMemRefTypeStorage; -struct UnrankedTensorTypeStorage; -struct VectorTypeStorage; - -} // namespace detail - //===----------------------------------------------------------------------===// // FloatType //===----------------------------------------------------------------------===// @@ -78,7 +64,6 @@ /// from ShapedType. class ShapedType : public Type { public: - using ImplType = detail::ShapedTypeStorage; using Type::Type; // TODO: merge these two special values in a single one used everywhere. @@ -157,46 +142,6 @@ } }; -//===----------------------------------------------------------------------===// -// VectorType -//===----------------------------------------------------------------------===// - -/// Vector types represent multi-dimensional SIMD vectors, and have a fixed -/// known constant shape with one or more dimension. -class VectorType - : public Type::TypeBase { -public: - using Base::Base; - using Base::getChecked; - - /// Get or create a new VectorType of the provided shape and element type. - /// Assumes the arguments define a well-formed VectorType. - static VectorType get(ArrayRef shape, Type elementType); - - /// Get or create a new VectorType of the provided shape and element type. If - /// the VectorType defined by the arguments would be ill-formed, an error is - /// emitted to `emitError` and a null type is returned. - static VectorType getChecked(function_ref emitError, - ArrayRef shape, Type elementType); - - /// Verify the construction of a vector type. - static LogicalResult verify(function_ref emitError, - ArrayRef shape, Type elementType); - - /// Returns true of the given type can be used as an element of a vector type. - /// In particular, vectors can consist of integer or float primitives. - static bool isValidElementType(Type t) { - return t.isa(); - } - - ArrayRef getShape() const; - - /// Get or create a new VectorType with the same shape as `this` and an - /// element type of bitwidth scaled by `scale`. - /// Return null if the scaled element type cannot be represented. - VectorType scaleElementBitwidth(unsigned scale); -}; - //===----------------------------------------------------------------------===// // TensorType //===----------------------------------------------------------------------===// @@ -214,68 +159,6 @@ static bool classof(Type type); }; -//===----------------------------------------------------------------------===// -// RankedTensorType - -/// Ranked tensor types represent multi-dimensional arrays that have a shape -/// with a fixed number of dimensions. Each shape element can be a non-negative -/// integer or unknown (represented by -1). -class RankedTensorType - : public Type::TypeBase { -public: - using Base::Base; - using Base::getChecked; - - /// Get or create a new RankedTensorType of the provided shape and element - /// type. Assumes the arguments define a well-formed type. - static RankedTensorType get(ArrayRef shape, Type elementType); - - /// Get or create a new RankedTensorType of the provided shape and element - /// type. If the RankedTensorType defined by the arguments would be - /// ill-formed, an error is emitted to `emitError` and a null type is - /// returned. - static RankedTensorType - getChecked(function_ref emitError, - ArrayRef shape, Type elementType); - - /// Verify the construction of a ranked tensor type. - static LogicalResult verify(function_ref emitError, - ArrayRef shape, Type elementType); - - ArrayRef getShape() const; -}; - -//===----------------------------------------------------------------------===// -// UnrankedTensorType - -/// Unranked tensor types represent multi-dimensional arrays that have an -/// unknown shape. -class UnrankedTensorType - : public Type::TypeBase { -public: - using Base::Base; - using Base::getChecked; - - /// Get or create a new UnrankedTensorType of the provided shape and element - /// type. Assumes the arguments define a well-formed type. - static UnrankedTensorType get(Type elementType); - - /// Get or create a new UnrankedTensorType of the provided shape and element - /// type. If the RankedTensorType defined by the arguments would be - /// ill-formed, an error is emitted to `emitError` and a null type is - /// returned. - static UnrankedTensorType - getChecked(function_ref emitError, Type elementType); - - /// Verify the construction of a unranked tensor type. - static LogicalResult verify(function_ref emitError, - Type elementType); - - ArrayRef getShape() const { return llvm::None; } -}; - //===----------------------------------------------------------------------===// // BaseMemRefType //===----------------------------------------------------------------------===// @@ -283,7 +166,6 @@ /// Base MemRef for Ranked and Unranked variants class BaseMemRefType : public ShapedType { public: - using ImplType = detail::BaseMemRefTypeStorage; using ShapedType::ShapedType; /// Return true if the specified element type is ok in a memref. @@ -296,149 +178,69 @@ unsigned getMemorySpaceAsInt() const; }; +} // end namespace mlir + +//===----------------------------------------------------------------------===// +// Tablegen Type Declarations +//===----------------------------------------------------------------------===// + +#define GET_TYPEDEF_CLASSES +#include "mlir/IR/BuiltinTypes.h.inc" + +namespace mlir { //===----------------------------------------------------------------------===// // MemRefType +//===----------------------------------------------------------------------===// -/// MemRef types represent a region of memory that have a shape with a fixed -/// number of dimensions. Each shape element can be a non-negative integer or -/// unknown (represented by -1). MemRef types also have an affine map -/// composition, represented as an array AffineMap pointers. -class MemRefType : public Type::TypeBase { +/// This is a builder type that keeps local references to arguments. Arguments +/// that are passed into the builder must out-live the builder. +class MemRefType::Builder { public: - /// This is a builder type that keeps local references to arguments. Arguments - /// that are passed into the builder must out-live the builder. - class Builder { - public: - // Build from another MemRefType. - explicit Builder(MemRefType other) - : shape(other.getShape()), elementType(other.getElementType()), - affineMaps(other.getAffineMaps()), - memorySpace(other.getMemorySpaceAsInt()) {} - - // Build from scratch. - Builder(ArrayRef shape, Type elementType) - : shape(shape), elementType(elementType), affineMaps(), memorySpace(0) { - } - - Builder &setShape(ArrayRef newShape) { - shape = newShape; - return *this; - } - - Builder &setElementType(Type newElementType) { - elementType = newElementType; - return *this; - } - - Builder &setAffineMaps(ArrayRef newAffineMaps) { - affineMaps = newAffineMaps; - return *this; - } - - Builder &setMemorySpace(unsigned newMemorySpace) { - memorySpace = newMemorySpace; - return *this; - } - - operator MemRefType() { - return MemRefType::get(shape, elementType, affineMaps, memorySpace); - } - - private: - ArrayRef shape; - Type elementType; - ArrayRef affineMaps; - unsigned memorySpace; - }; - - using Base::Base; - using Base::getChecked; - - /// Get or create a new MemRefType based on shape, element type, affine - /// map composition, and memory space. Assumes the arguments define a - /// well-formed MemRef type. Use getChecked to gracefully handle MemRefType - /// construction failures. - static MemRefType get(ArrayRef shape, Type elementType, - ArrayRef affineMapComposition = {}, - unsigned memorySpace = 0); - - /// Get or create a new MemRefType based on shape, element type, affine - /// map composition, and memory space. If the MemRefType defined by the - /// arguments would be ill-formed, an error is emitted to `emitError` and a - /// null type is returned. - static MemRefType getChecked(function_ref emitError, - ArrayRef shape, Type elementType, - ArrayRef affineMapComposition, - unsigned memorySpace); - - ArrayRef getShape() const; + // Build from another MemRefType. + explicit Builder(MemRefType other) + : shape(other.getShape()), elementType(other.getElementType()), + affineMaps(other.getAffineMaps()), + memorySpace(other.getMemorySpaceAsInt()) {} + + // Build from scratch. + Builder(ArrayRef shape, Type elementType) + : shape(shape), elementType(elementType), affineMaps(), memorySpace(0) {} + + Builder &setShape(ArrayRef newShape) { + shape = newShape; + return *this; + } - /// Returns an array of affine map pointers representing the memref affine - /// map composition. - ArrayRef getAffineMaps() const; + Builder &setElementType(Type newElementType) { + elementType = newElementType; + return *this; + } - // TODO: merge these two special values in a single one used everywhere. - // Unfortunately, uses of `-1` have crept deep into the codebase now and are - // hard to track. - static int64_t getDynamicStrideOrOffset() { - return ShapedType::kDynamicStrideOrOffset; + Builder &setAffineMaps(ArrayRef newAffineMaps) { + affineMaps = newAffineMaps; + return *this; } -private: - /// Get or create a new MemRefType defined by the arguments. If the resulting - /// type would be ill-formed, return nullptr. - static MemRefType getImpl(ArrayRef shape, Type elementType, - ArrayRef affineMapComposition, - unsigned memorySpace, - function_ref emitError); - using Base::getImpl; -}; + Builder &setMemorySpace(unsigned newMemorySpace) { + memorySpace = newMemorySpace; + return *this; + } -//===----------------------------------------------------------------------===// -// UnrankedMemRefType + operator MemRefType() { + return MemRefType::get(shape, elementType, affineMaps, memorySpace); + } -/// Unranked MemRef type represent multi-dimensional MemRefs that -/// have an unknown rank. -class UnrankedMemRefType - : public Type::TypeBase { -public: - using Base::Base; - using Base::getChecked; - - /// Get or create a new UnrankedMemRefType of the provided element - /// type and memory space - static UnrankedMemRefType get(Type elementType, unsigned memorySpace); - - /// Get or create a new UnrankedMemRefType of the provided element - /// type and memory space. If the UnrankedMemRefType defined by the arguments - /// would be ill-formed, an error is emitted to `emitError` and a null type is - /// returned. - static UnrankedMemRefType - getChecked(function_ref emitError, Type elementType, - unsigned memorySpace); - - /// Verify the construction of a unranked memref type. - static LogicalResult verify(function_ref emitError, - Type elementType, unsigned memorySpace); - - ArrayRef getShape() const { return llvm::None; } +private: + ArrayRef shape; + Type elementType; + ArrayRef affineMaps; + unsigned memorySpace; }; -} // end namespace mlir - -//===----------------------------------------------------------------------===// -// Tablegen Type Declarations -//===----------------------------------------------------------------------===// - -#define GET_TYPEDEF_CLASSES -#include "mlir/IR/BuiltinTypes.h.inc" //===----------------------------------------------------------------------===// // Deferred Method Definitions //===----------------------------------------------------------------------===// -namespace mlir { inline bool BaseMemRefType::classof(Type type) { return type.isa(); } diff --git a/mlir/include/mlir/IR/BuiltinTypes.td b/mlir/include/mlir/IR/BuiltinTypes.td --- a/mlir/include/mlir/IR/BuiltinTypes.td +++ b/mlir/include/mlir/IR/BuiltinTypes.td @@ -21,7 +21,8 @@ // remove the definitions in OpBase.td, and repoint users to this file instead. // Base class for Builtin dialect types. -class Builtin_Type : TypeDef { +class Builtin_Type + : TypeDef { let mnemonic = ?; } @@ -250,6 +251,316 @@ }]; } +//===----------------------------------------------------------------------===// +// MemRefType +//===----------------------------------------------------------------------===// + +def Builtin_MemRef : Builtin_Type<"MemRef", "BaseMemRefType"> { + let summary = "Shaped reference to a region of memory"; + let description = [{ + Syntax: + + ``` + memref-type ::= `memref` `<` dimension-list-ranked type + (`,` layout-specification)? (`,` memory-space)? `>` + + stride-list ::= `[` (dimension (`,` dimension)*)? `]` + strided-layout ::= `offset:` dimension `,` `strides: ` stride-list + semi-affine-map-composition ::= (semi-affine-map `,` )* semi-affine-map + layout-specification ::= semi-affine-map-composition | strided-layout + memory-space ::= integer-literal /* | TODO: address-space-id */ + ``` + + A `memref` type is a reference to a region of memory (similar to a buffer + pointer, but more powerful). The buffer pointed to by a memref can be + allocated, aliased and deallocated. A memref can be used to read and write + data from/to the memory region which it references. Memref types use the + same shape specifier as tensor types. Note that `memref`, + `memref<0 x f32>`, `memref<1 x 0 x f32>`, and `memref<0 x 1 x f32>` are all + different types. + + A `memref` is allowed to have an unknown rank (e.g. `memref<*xf32>`). The + purpose of unranked memrefs is to allow external library functions to + receive memref arguments of any rank without versioning the functions based + on the rank. Other uses of this type are disallowed or will have undefined + behavior. + + ##### Codegen of Unranked Memref + + Using unranked memref in codegen besides the case mentioned above is highly + discouraged. Codegen is concerned with generating loop nests and specialized + instructions for high-performance, unranked memref is concerned with hiding + the rank and thus, the number of enclosing loops required to iterate over + the data. However, if there is a need to code-gen unranked memref, one + possible path is to cast into a static ranked type based on the dynamic + rank. Another possible path is to emit a single while loop conditioned on a + linear index and perform delinearization of the linear index to a dynamic + array containing the (unranked) indices. While this is possible, it is + expected to not be a good idea to perform this during codegen as the cost + of the translations is expected to be prohibitive and optimizations at this + level are not expected to be worthwhile. If expressiveness is the main + concern, irrespective of performance, passing unranked memrefs to an + external C++ library and implementing rank-agnostic logic there is expected + to be significantly simpler. + + Unranked memrefs may provide expressiveness gains in the future and help + bridge the gap with unranked tensors. Unranked memrefs will not be expected + to be exposed to codegen but one may query the rank of an unranked memref + (a special op will be needed for this purpose) and perform a switch and cast + to a ranked memref as a prerequisite to codegen. + + Example: + + ```mlir + // With static ranks, we need a function for each possible argument type + %A = alloc() : memref<16x32xf32> + %B = alloc() : memref<16x32x64xf32> + call @helper_2D(%A) : (memref<16x32xf32>)->() + call @helper_3D(%B) : (memref<16x32x64xf32>)->() + + // With unknown rank, the functions can be unified under one unranked type + %A = alloc() : memref<16x32xf32> + %B = alloc() : memref<16x32x64xf32> + // Remove rank info + %A_u = memref_cast %A : memref<16x32xf32> -> memref<*xf32> + %B_u = memref_cast %B : memref<16x32x64xf32> -> memref<*xf32> + // call same function with dynamic ranks + call @helper(%A_u) : (memref<*xf32>)->() + call @helper(%B_u) : (memref<*xf32>)->() + ``` + + The core syntax and representation of a layout specification is a + [semi-affine map](Dialects/Affine.md#semi-affine-maps). Additionally, + syntactic sugar is supported to make certain layout specifications more + intuitive to read. For the moment, a `memref` supports parsing a strided + form which is converted to a semi-affine map automatically. + + The memory space of a memref is specified by a target-specific integer + index. If no memory space is specified, then the default memory space (0) + is used. The default space is target specific but always at index 0. + + TODO: MLIR will eventually have target-dialects which allow symbolic use of + memory hierarchy names (e.g. L3, L2, L1, ...) but we have not spec'd the + details of that mechanism yet. Until then, this document pretends that it + is valid to refer to these memories by `bare-id`. + + The notionally dynamic value of a memref value includes the address of the + buffer allocated, as well as the symbols referred to by the shape, layout + map, and index maps. + + Examples of memref static type + + ```mlir + // Identity index/layout map + #identity = affine_map<(d0, d1) -> (d0, d1)> + + // Column major layout. + #col_major = affine_map<(d0, d1, d2) -> (d2, d1, d0)> + + // A 2-d tiled layout with tiles of size 128 x 256. + #tiled_2d_128x256 = affine_map<(d0, d1) -> (d0 div 128, d1 div 256, d0 mod 128, d1 mod 256)> + + // A tiled data layout with non-constant tile sizes. + #tiled_dynamic = affine_map<(d0, d1)[s0, s1] -> (d0 floordiv s0, d1 floordiv s1, + d0 mod s0, d1 mod s1)> + + // A layout that yields a padding on two at either end of the minor dimension. + #padded = affine_map<(d0, d1) -> (d0, (d1 + 2) floordiv 2, (d1 + 2) mod 2)> + + + // The dimension list "16x32" defines the following 2D index space: + // + // { (i, j) : 0 <= i < 16, 0 <= j < 32 } + // + memref<16x32xf32, #identity> + + // The dimension list "16x4x?" defines the following 3D index space: + // + // { (i, j, k) : 0 <= i < 16, 0 <= j < 4, 0 <= k < N } + // + // where N is a symbol which represents the runtime value of the size of + // the third dimension. + // + // %N here binds to the size of the third dimension. + %A = alloc(%N) : memref<16x4x?xf32, #col_major> + + // A 2-d dynamic shaped memref that also has a dynamically sized tiled + // layout. The memref index space is of size %M x %N, while %B1 and %B2 + // bind to the symbols s0, s1 respectively of the layout map #tiled_dynamic. + // Data tiles of size %B1 x %B2 in the logical space will be stored + // contiguously in memory. The allocation size will be + // (%M ceildiv %B1) * %B1 * (%N ceildiv %B2) * %B2 f32 elements. + %T = alloc(%M, %N) [%B1, %B2] : memref + + // A memref that has a two-element padding at either end. The allocation + // size will fit 16 * 64 float elements of data. + %P = alloc() : memref<16x64xf32, #padded> + + // Affine map with symbol 's0' used as offset for the first dimension. + #imapS = affine_map<(d0, d1) [s0] -> (d0 + s0, d1)> + // Allocate memref and bind the following symbols: + // '%n' is bound to the dynamic second dimension of the memref type. + // '%o' is bound to the symbol 's0' in the affine map of the memref type. + %n = ... + %o = ... + %A = alloc (%n)[%o] : <16x?xf32, #imapS> + ``` + + ##### Index Space + + A memref dimension list defines an index space within which the memref can + be indexed to access data. + + ##### Index + + Data is accessed through a memref type using a multidimensional index into + the multidimensional index space defined by the memref's dimension list. + + Examples + + ```mlir + // Allocates a memref with 2D index space: + // { (i, j) : 0 <= i < 16, 0 <= j < 32 } + %A = alloc() : memref<16x32xf32, #imapA> + + // Loads data from memref '%A' using a 2D index: (%i, %j) + %v = load %A[%i, %j] : memref<16x32xf32, #imapA> + ``` + + ##### Index Map + + An index map is a one-to-one + [semi-affine map](Dialects/Affine.md#semi-affine-maps) that transforms a + multidimensional index from one index space to another. For example, the + following figure shows an index map which maps a 2-dimensional index from a + 2x2 index space to a 3x3 index space, using symbols `S0` and `S1` as + offsets. + + ![Index Map Example](/includes/img/index-map.svg) + + The number of domain dimensions and range dimensions of an index map can be + different, but must match the number of dimensions of the input and output + index spaces on which the map operates. The index space is always + non-negative and integral. In addition, an index map must specify the size + of each of its range dimensions onto which it maps. Index map symbols must + be listed in order with symbols for dynamic dimension sizes first, followed + by other required symbols. + + ##### Layout Map + + A layout map is a [semi-affine map](Dialects/Affine.md#semi-affine-maps) + which encodes logical to physical index space mapping, by mapping input + dimensions to their ordering from most-major (slowest varying) to most-minor + (fastest varying). Therefore, an identity layout map corresponds to a + row-major layout. Identity layout maps do not contribute to the MemRef type + identification and are discarded on construction. That is, a type with an + explicit identity map is `memref(i,j)>` is strictly the + same as the one without layout maps, `memref`. + + Layout map examples: + + ```mlir + // MxN matrix stored in row major layout in memory: + #layout_map_row_major = (i, j) -> (i, j) + + // MxN matrix stored in column major layout in memory: + #layout_map_col_major = (i, j) -> (j, i) + + // MxN matrix stored in a 2-d blocked/tiled layout with 64x64 tiles. + #layout_tiled = (i, j) -> (i floordiv 64, j floordiv 64, i mod 64, j mod 64) + ``` + + ##### Affine Map Composition + + A memref specifies a semi-affine map composition as part of its type. A + semi-affine map composition is a composition of semi-affine maps beginning + with zero or more index maps, and ending with a layout map. The composition + must be conformant: the number of dimensions of the range of one map, must + match the number of dimensions of the domain of the next map in the + composition. + + The semi-affine map composition specified in the memref type, maps from + accesses used to index the memref in load/store operations to other index + spaces (i.e. logical to physical index mapping). Each of the + [semi-affine maps](Affine.md) and thus its composition is required + to be one-to-one. + + The semi-affine map composition can be used in dependence analysis, memory + access pattern analysis, and for performance optimizations like + vectorization, copy elision and in-place updates. If an affine map + composition is not specified for the memref, the identity affine map is + assumed. + + ##### Strided MemRef + + A memref may specify strides as part of its type. A stride specification is + a list of integer values that are either static or `?` (dynamic case). + Strides encode the distance, in number of elements, in (linear) memory + between successive entries along a particular dimension. A stride + specification is syntactic sugar for an equivalent strided memref + representation using semi-affine maps. For example, + `memref<42x16xf32, offset: 33, strides: [1, 64]>` specifies a non-contiguous + memory region of `42` by `16` `f32` elements such that: + + 1. the minimal size of the enclosing memory region must be + `33 + 42 * 1 + 16 * 64 = 1066` elements; + 2. the address calculation for accessing element `(i, j)` computes + `33 + i + 64 * j` + 3. the distance between two consecutive elements along the inner dimension + is `1` element and the distance between two consecutive elements along + the outer dimension is `64` elements. + + This corresponds to a column major view of the memory region and is + internally represented as the type + `memref<42x16xf32, (i, j) -> (33 + i + 64 * j)>`. + + The specification of strides must not alias: given an n-D strided memref, + indices `(i1, ..., in)` and `(j1, ..., jn)` may not refer to the same memory + address unless `i1 == j1, ..., in == jn`. + + Strided memrefs represent a view abstraction over preallocated data. They + are constructed with special ops, yet to be introduced. Strided memrefs are + a special subclass of memrefs with generic semi-affine map and correspond to + a normalized memref descriptor when lowering to LLVM. + }]; + let parameters = (ins + ArrayRefParameter<"int64_t">:$shape, + "Type":$elementType, + ArrayRefParameter<"AffineMap">:$affineMaps, + "unsigned":$memorySpaceAsInt + ); + + let builders = [ + TypeBuilderWithInferredContext<(ins + "ArrayRef":$shape, "Type":$elementType, + CArg<"ArrayRef", "{}">:$affineMaps, + CArg<"unsigned", "0">:$memorySpace + ), [{ + // Drop identity maps from the composition. This may lead to the + // composition becoming empty, which is interpreted as an implicit + // identity. + auto nonIdentityMaps = llvm::make_filter_range(affineMaps, + [](AffineMap map) { return !map.isIdentity(); }); + return $_get(elementType.getContext(), shape, elementType, + llvm::to_vector<4>(nonIdentityMaps), memorySpace); + }]> + ]; + let extraClassDeclaration = [{ + /// This is a builder type that keeps local references to arguments. + /// Arguments that are passed into the builder must out-live the builder. + class Builder; + + // TODO: merge these two special values in a single one used everywhere. + // Unfortunately, uses of `-1` have crept deep into the codebase now and are + // hard to track. + static int64_t getDynamicStrideOrOffset() { + return ShapedType::kDynamicStrideOrOffset; + } + }]; + let skipDefaultBuilders = 1; + let genVerifyDecl = 1; +} + //===----------------------------------------------------------------------===// // NoneType //===----------------------------------------------------------------------===// @@ -305,6 +616,79 @@ let genVerifyDecl = 1; } +//===----------------------------------------------------------------------===// +// RankedTensorType +//===----------------------------------------------------------------------===// + +def Builtin_RankedTensor : Builtin_Type<"RankedTensor", "TensorType"> { + let summary = "Multi-dimensional array with a fixed number of dimensions"; + let description = [{ + Syntax: + + ``` + tensor-type ::= `tensor` `<` dimension-list type `>` + dimension-list ::= (dimension `x`)* + dimension ::= `?` | decimal-literal + ``` + + Values with tensor type represents aggregate N-dimensional data values, and + have a known element type and a fixed rank with a list of dimensions. Each + dimension may be a static non-negative decimal constant or be dynamically + determined (indicated by `?`). + + The runtime representation of the MLIR tensor type is intentionally + abstracted - you cannot control layout or get a pointer to the data. For + low level buffer access, MLIR has a [`memref` type](#memref-type). This + abstracted runtime representation holds both the tensor data values as well + as information about the (potentially dynamic) shape of the tensor. The + [`dim` operation](Dialects/Standard.md#dim-operation) returns the size of a + dimension from a value of tensor type. + + Note: hexadecimal integer literals are not allowed in tensor type + declarations to avoid confusion between `0xf32` and `0 x f32`. Zero sizes + are allowed in tensors and treated as other sizes, e.g., + `tensor<0 x 1 x i32>` and `tensor<1 x 0 x i32>` are different types. Since + zero sizes are not allowed in some other types, such tensors should be + optimized away before lowering tensors to vectors. + + Examples: + + ```mlir + // Known rank but unknown dimensions. + tensor + + // Partially known dimensions. + tensor + + // Full static shape. + tensor<17 x 4 x 13 x 4 x f32> + + // Tensor with rank zero. Represents a scalar. + tensor + + // Zero-element dimensions are allowed. + tensor<0 x 42 x f32> + + // Zero-element tensor of f32 type (hexadecimal literals not allowed here). + tensor<0xf32> + ``` + }]; + let parameters = (ins + ArrayRefParameter<"int64_t">:$shape, + "Type":$elementType + ); + + let builders = [ + TypeBuilderWithInferredContext<(ins + "ArrayRef":$shape, "Type":$elementType + ), [{ + return $_get(elementType.getContext(), shape, elementType); + }]> + ]; + let skipDefaultBuilders = 1; + let genVerifyDecl = 1; +} + //===----------------------------------------------------------------------===// // TupleType //===----------------------------------------------------------------------===// @@ -372,4 +756,149 @@ }]; } +//===----------------------------------------------------------------------===// +// UnrankedMemRefType +//===----------------------------------------------------------------------===// + +def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", "BaseMemRefType"> { + let summary = "Shaped reference, with unknown rank, to a region of memory"; + let description = [{ + Syntax: + + ``` + unranked-memref-type ::= `memref` `<*x` type (`,` memory-space)? `>` + memory-space ::= integer-literal /* | TODO: address-space-id */ + ``` + + A `memref` type with an unknown rank (e.g. `memref<*xf32>`). The purpose of + unranked memrefs is to allow external library functions to receive memref + arguments of any rank without versioning the functions based on the rank. + Other uses of this type are disallowed or will have undefined behavior. + + See [MemRefType](#builtin_memref-memreftype) for more information on + memref types. + + Examples: + + ```mlir + memref<*f32> + + // An unranked memref with a memory space of 10. + memref<*f32, 10> + ``` + }]; + let parameters = (ins "Type":$elementType, "unsigned":$memorySpaceAsInt); + + let builders = [ + TypeBuilderWithInferredContext<(ins "Type":$elementType, + "unsigned":$memorySpace), [{ + return $_get(elementType.getContext(), elementType, memorySpace); + }]> + ]; + let extraClassDeclaration = [{ + ArrayRef getShape() const { return llvm::None; } + }]; + let skipDefaultBuilders = 1; + let genVerifyDecl = 1; +} + +//===----------------------------------------------------------------------===// +// UnrankedTensorType +//===----------------------------------------------------------------------===// + +def Builtin_UnrankedTensor : Builtin_Type<"UnrankedTensor", "TensorType"> { + let summary = "Multi-dimensional array with unknown dimensions"; + let description = [{ + Syntax: + + ``` + tensor-type ::= `tensor` `<` `*` `x` type `>` + ``` + + An unranked tensor is a type of tensor in which the set of dimensions have + unknown rank. See [RankedTensorType](#builtin_rankedtensor-rankedtensortype) + for more information on tensor types. + + Examples: + + ```mlir + tensor<*f32> + ``` + }]; + let parameters = (ins "Type":$elementType); + + let builders = [ + TypeBuilderWithInferredContext<(ins "Type":$elementType), [{ + return $_get(elementType.getContext(), elementType); + }]> + ]; + let extraClassDeclaration = [{ + ArrayRef getShape() const { return llvm::None; } + }]; + let skipDefaultBuilders = 1; + let genVerifyDecl = 1; +} + +//===----------------------------------------------------------------------===// +// VectorType +//===----------------------------------------------------------------------===// + +def Builtin_Vector : Builtin_Type<"Vector", "ShapedType"> { + let summary = "Multi-dimensional SIMD vector type"; + let description = [{ + Syntax: + + ``` + vector-type ::= `vector` `<` static-dimension-list vector-element-type `>` + vector-element-type ::= float-type | integer-type + + static-dimension-list ::= (decimal-literal `x`)+ + ``` + + The vector type represents a SIMD style vector, used by target-specific + operation sets like AVX. While the most common use is for 1D vectors (e.g. + vector<16 x f32>) we also support multidimensional registers on targets that + support them (like TPUs). + + Vector shapes must be positive decimal integers. + + Note: hexadecimal integer literals are not allowed in vector type + declarations, `vector<0x42xi32>` is invalid because it is interpreted as a + 2D vector with shape `(0, 42)` and zero shapes are not allowed. + + + Examples: + + ```mlir + vector<3x42xi32> + ``` + }]; + let parameters = (ins + ArrayRefParameter<"int64_t">:$shape, + "Type":$elementType + ); + + let builders = [ + TypeBuilderWithInferredContext<(ins + "ArrayRef":$shape, "Type":$elementType + ), [{ + return $_get(elementType.getContext(), shape, elementType); + }]> + ]; + let extraClassDeclaration = [{ + /// Returns true of the given type can be used as an element of a vector + /// type. In particular, vectors can consist of integer or float primitives. + static bool isValidElementType(Type t) { + return t.isa(); + } + + /// Get or create a new VectorType with the same shape as `this` and an + /// element type of bitwidth scaled by `scale`. + /// Return null if the scaled element type cannot be represented. + VectorType scaleElementBitwidth(unsigned scale); + }]; + let skipDefaultBuilders = 1; + let genVerifyDecl = 1; +} + #endif // BUILTIN_TYPES diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td --- a/mlir/include/mlir/IR/OpBase.td +++ b/mlir/include/mlir/IR/OpBase.td @@ -2684,13 +2684,13 @@ class TypeParameter : AttrOrTypeParameter; // For StringRefs, which require allocation. -class StringRefParameter : +class StringRefParameter : AttrOrTypeParameter<"::llvm::StringRef", desc> { let allocator = [{$_dst = $_allocator.copyInto($_self);}]; } // For standard ArrayRefs, which require allocation. -class ArrayRefParameter : +class ArrayRefParameter : AttrOrTypeParameter<"::llvm::ArrayRef<" # arrayOf # ">", desc> { let allocator = [{$_dst = $_allocator.copyInto($_self);}]; } diff --git a/mlir/lib/IR/BuiltinTypes.cpp b/mlir/lib/IR/BuiltinTypes.cpp --- a/mlir/lib/IR/BuiltinTypes.cpp +++ b/mlir/lib/IR/BuiltinTypes.cpp @@ -16,6 +16,7 @@ #include "llvm/ADT/BitVector.h" #include "llvm/ADT/Sequence.h" #include "llvm/ADT/Twine.h" +#include "llvm/ADT/TypeSwitch.h" using namespace mlir; using namespace mlir::detail; @@ -266,7 +267,9 @@ } Type ShapedType::getElementType() const { - return static_cast(impl)->elementType; + return TypeSwitch(*this) + .Case([](auto ty) { return ty.getElementType(); }); } unsigned ShapedType::getElementTypeBitWidth() const { @@ -357,16 +360,6 @@ // VectorType //===----------------------------------------------------------------------===// -VectorType VectorType::get(ArrayRef shape, Type elementType) { - return Base::get(elementType.getContext(), shape, elementType); -} - -VectorType VectorType::getChecked(function_ref emitError, - ArrayRef shape, Type elementType) { - return Base::getChecked(emitError, elementType.getContext(), shape, - elementType); -} - LogicalResult VectorType::verify(function_ref emitError, ArrayRef shape, Type elementType) { if (shape.empty()) @@ -381,8 +374,6 @@ return success(); } -ArrayRef VectorType::getShape() const { return getImpl()->getShape(); } - VectorType VectorType::scaleElementBitwidth(unsigned scale) { if (!scale) return VectorType(); @@ -422,46 +413,19 @@ // RankedTensorType //===----------------------------------------------------------------------===// -RankedTensorType RankedTensorType::get(ArrayRef shape, - Type elementType) { - return Base::get(elementType.getContext(), shape, elementType); -} - -RankedTensorType -RankedTensorType::getChecked(function_ref emitError, - ArrayRef shape, Type elementType) { - return Base::getChecked(emitError, elementType.getContext(), shape, - elementType); -} - LogicalResult RankedTensorType::verify(function_ref emitError, ArrayRef shape, Type elementType) { - for (int64_t s : shape) { + for (int64_t s : shape) if (s < -1) return emitError() << "invalid tensor dimension size"; - } return checkTensorElementType(emitError, elementType); } -ArrayRef RankedTensorType::getShape() const { - return getImpl()->getShape(); -} - //===----------------------------------------------------------------------===// // UnrankedTensorType //===----------------------------------------------------------------------===// -UnrankedTensorType UnrankedTensorType::get(Type elementType) { - return Base::get(elementType.getContext(), elementType); -} - -UnrankedTensorType -UnrankedTensorType::getChecked(function_ref emitError, - Type elementType) { - return Base::getChecked(emitError, elementType.getContext(), elementType); -} - LogicalResult UnrankedTensorType::verify(function_ref emitError, Type elementType) { @@ -473,115 +437,50 @@ //===----------------------------------------------------------------------===// unsigned BaseMemRefType::getMemorySpaceAsInt() const { - return static_cast(impl)->memorySpace; + if (auto rankedMemRefTy = dyn_cast()) + return rankedMemRefTy.getMemorySpaceAsInt(); + return cast().getMemorySpaceAsInt(); } //===----------------------------------------------------------------------===// // MemRefType //===----------------------------------------------------------------------===// -/// Get or create a new MemRefType based on shape, element type, affine -/// map composition, and memory space. Assumes the arguments define a -/// well-formed MemRef type. Use getChecked to gracefully handle MemRefType -/// construction failures. -MemRefType MemRefType::get(ArrayRef shape, Type elementType, - ArrayRef affineMapComposition, - unsigned memorySpace) { - auto result = - getImpl(shape, elementType, affineMapComposition, memorySpace, [=] { - return emitError(UnknownLoc::get(elementType.getContext())); - }); - assert(result && "Failed to construct instance of MemRefType."); - return result; -} - -/// Get or create a new MemRefType based on shape, element type, affine -/// map composition, and memory space declared at the given location. -/// If the location is unknown, the last argument should be an instance of -/// UnknownLoc. If the MemRefType defined by the arguments would be -/// ill-formed, emits errors (to the handler registered with the context or to -/// the error stream) and returns nullptr. -MemRefType MemRefType::getChecked(function_ref emitError, - ArrayRef shape, Type elementType, - ArrayRef affineMapComposition, - unsigned memorySpace) { - return getImpl(shape, elementType, affineMapComposition, memorySpace, - emitError); -} - -/// Get or create a new MemRefType defined by the arguments. If the resulting -/// type would be ill-formed, return nullptr. If the location is provided, -/// emit detailed error messages. To emit errors when the location is unknown, -/// pass in an instance of UnknownLoc. -MemRefType MemRefType::getImpl(ArrayRef shape, Type elementType, - ArrayRef affineMapComposition, - unsigned memorySpace, - function_ref emitError) { - auto *context = elementType.getContext(); - +LogicalResult MemRefType::verify(function_ref emitError, + ArrayRef shape, Type elementType, + ArrayRef affineMapComposition, + unsigned memorySpace) { if (!BaseMemRefType::isValidElementType(elementType)) - return (emitError() << "invalid memref element type", MemRefType()); + return emitError() << "invalid memref element type"; - for (int64_t s : shape) { // Negative sizes are not allowed except for `-1` that means dynamic size. + for (int64_t s : shape) if (s < -1) - return (emitError() << "invalid memref size", MemRefType()); - } + return emitError() << "invalid memref size"; // Check that the structure of the composition is valid, i.e. that each // subsequent affine map has as many inputs as the previous map has results. // Take the dimensionality of the MemRef for the first map. - auto dim = shape.size(); - unsigned i = 0; - for (const auto &affineMap : affineMapComposition) { - if (affineMap.getNumDims() != dim) { - emitError() << "memref affine map dimension mismatch between " - << (i == 0 ? Twine("memref rank") : "affine map " + Twine(i)) - << " and affine map" << i + 1 << ": " << dim - << " != " << affineMap.getNumDims(); - return nullptr; - } - - dim = affineMap.getNumResults(); - ++i; - } - - // Drop identity maps from the composition. - // This may lead to the composition becoming empty, which is interpreted as an - // implicit identity. - SmallVector cleanedAffineMapComposition; - for (const auto &map : affineMapComposition) { - if (map.isIdentity()) + size_t dim = shape.size(); + for (auto it : llvm::enumerate(affineMapComposition)) { + AffineMap map = it.value(); + if (map.getNumDims() == dim) { + dim = map.getNumResults(); continue; - cleanedAffineMapComposition.push_back(map); + } + return emitError() << "memref affine map dimension mismatch between " + << (it.index() == 0 ? Twine("memref rank") + : "affine map " + Twine(it.index())) + << " and affine map" << it.index() + 1 << ": " << dim + << " != " << map.getNumDims(); } - - return Base::get(context, shape, elementType, cleanedAffineMapComposition, - memorySpace); -} - -ArrayRef MemRefType::getShape() const { return getImpl()->getShape(); } - -ArrayRef MemRefType::getAffineMaps() const { - return getImpl()->getAffineMaps(); + return success(); } //===----------------------------------------------------------------------===// // UnrankedMemRefType //===----------------------------------------------------------------------===// -UnrankedMemRefType UnrankedMemRefType::get(Type elementType, - unsigned memorySpace) { - return Base::get(elementType.getContext(), elementType, memorySpace); -} - -UnrankedMemRefType -UnrankedMemRefType::getChecked(function_ref emitError, - Type elementType, unsigned memorySpace) { - return Base::getChecked(emitError, elementType.getContext(), elementType, - memorySpace); -} - LogicalResult UnrankedMemRefType::verify(function_ref emitError, Type elementType, unsigned memorySpace) { diff --git a/mlir/lib/IR/TypeDetail.h b/mlir/lib/IR/TypeDetail.h --- a/mlir/lib/IR/TypeDetail.h +++ b/mlir/lib/IR/TypeDetail.h @@ -23,8 +23,6 @@ namespace mlir { -class MLIRContext; - namespace detail { /// Integer Type Storage and Uniquing. @@ -96,175 +94,6 @@ Type const *inputsAndResults; }; -/// Shaped Type Storage. -struct ShapedTypeStorage : public TypeStorage { - ShapedTypeStorage(Type elementTy) : elementType(elementTy) {} - - /// The hash key used for uniquing. - using KeyTy = Type; - bool operator==(const KeyTy &key) const { return key == elementType; } - - Type elementType; -}; - -/// Vector Type Storage and Uniquing. -struct VectorTypeStorage : public ShapedTypeStorage { - VectorTypeStorage(unsigned shapeSize, Type elementTy, - const int64_t *shapeElements) - : ShapedTypeStorage(elementTy), shapeElements(shapeElements), - shapeSize(shapeSize) {} - - /// The hash key used for uniquing. - using KeyTy = std::pair, Type>; - bool operator==(const KeyTy &key) const { - return key == KeyTy(getShape(), elementType); - } - - /// Construction. - static VectorTypeStorage *construct(TypeStorageAllocator &allocator, - const KeyTy &key) { - // Copy the shape into the bump pointer. - ArrayRef shape = allocator.copyInto(key.first); - - // Initialize the memory using placement new. - return new (allocator.allocate()) - VectorTypeStorage(shape.size(), key.second, shape.data()); - } - - ArrayRef getShape() const { - return ArrayRef(shapeElements, shapeSize); - } - - const int64_t *shapeElements; - unsigned shapeSize; -}; - -struct RankedTensorTypeStorage : public ShapedTypeStorage { - RankedTensorTypeStorage(unsigned shapeSize, Type elementTy, - const int64_t *shapeElements) - : ShapedTypeStorage(elementTy), shapeElements(shapeElements), - shapeSize(shapeSize) {} - - /// The hash key used for uniquing. - using KeyTy = std::pair, Type>; - bool operator==(const KeyTy &key) const { - return key == KeyTy(getShape(), elementType); - } - - /// Construction. - static RankedTensorTypeStorage *construct(TypeStorageAllocator &allocator, - const KeyTy &key) { - // Copy the shape into the bump pointer. - ArrayRef shape = allocator.copyInto(key.first); - - // Initialize the memory using placement new. - return new (allocator.allocate()) - RankedTensorTypeStorage(shape.size(), key.second, shape.data()); - } - - ArrayRef getShape() const { - return ArrayRef(shapeElements, shapeSize); - } - - const int64_t *shapeElements; - unsigned shapeSize; -}; - -struct UnrankedTensorTypeStorage : public ShapedTypeStorage { - using ShapedTypeStorage::KeyTy; - using ShapedTypeStorage::ShapedTypeStorage; - - /// Construction. - static UnrankedTensorTypeStorage *construct(TypeStorageAllocator &allocator, - Type elementTy) { - return new (allocator.allocate()) - UnrankedTensorTypeStorage(elementTy); - } -}; - -struct BaseMemRefTypeStorage : public ShapedTypeStorage { - BaseMemRefTypeStorage(Type elementType, unsigned memorySpace) - : ShapedTypeStorage(elementType), memorySpace(memorySpace) {} - - /// Memory space in which data referenced by memref resides. - const unsigned memorySpace; -}; - -struct MemRefTypeStorage : public BaseMemRefTypeStorage { - MemRefTypeStorage(unsigned shapeSize, Type elementType, - const int64_t *shapeElements, const unsigned numAffineMaps, - AffineMap const *affineMapList, const unsigned memorySpace) - : BaseMemRefTypeStorage(elementType, memorySpace), - shapeElements(shapeElements), shapeSize(shapeSize), - numAffineMaps(numAffineMaps), affineMapList(affineMapList) {} - - /// The hash key used for uniquing. - // MemRefs are uniqued based on their shape, element type, affine map - // composition, and memory space. - using KeyTy = - std::tuple, Type, ArrayRef, unsigned>; - bool operator==(const KeyTy &key) const { - return key == KeyTy(getShape(), elementType, getAffineMaps(), memorySpace); - } - - /// Construction. - static MemRefTypeStorage *construct(TypeStorageAllocator &allocator, - const KeyTy &key) { - // Copy the shape into the bump pointer. - ArrayRef shape = allocator.copyInto(std::get<0>(key)); - - // Copy the affine map composition into the bump pointer. - ArrayRef affineMapComposition = - allocator.copyInto(std::get<2>(key)); - - // Initialize the memory using placement new. - return new (allocator.allocate()) - MemRefTypeStorage(shape.size(), std::get<1>(key), shape.data(), - affineMapComposition.size(), - affineMapComposition.data(), std::get<3>(key)); - } - - ArrayRef getShape() const { - return ArrayRef(shapeElements, shapeSize); - } - - ArrayRef getAffineMaps() const { - return ArrayRef(affineMapList, numAffineMaps); - } - - /// An array of integers which stores the shape dimension sizes. - const int64_t *shapeElements; - /// The number of shape elements. - unsigned shapeSize; - /// The number of affine maps in the 'affineMapList' array. - const unsigned numAffineMaps; - /// List of affine maps in the memref's layout/index map composition. - AffineMap const *affineMapList; -}; - -/// Unranked MemRef is a MemRef with unknown rank. -/// Only element type and memory space are known -struct UnrankedMemRefTypeStorage : public BaseMemRefTypeStorage { - - UnrankedMemRefTypeStorage(Type elementTy, const unsigned memorySpace) - : BaseMemRefTypeStorage(elementTy, memorySpace) {} - - /// The hash key used for uniquing. - using KeyTy = std::tuple; - bool operator==(const KeyTy &key) const { - return key == KeyTy(elementType, memorySpace); - } - - /// Construction. - static UnrankedMemRefTypeStorage *construct(TypeStorageAllocator &allocator, - const KeyTy &key) { - - // Initialize the memory using placement new. - return new (allocator.allocate()) - UnrankedMemRefTypeStorage(std::get<0>(key), std::get<1>(key)); - } -}; - /// A type representing a collection of other types. struct TupleTypeStorage final : public TypeStorage,