diff --git a/mlir/docs/DataLayout.md b/mlir/docs/DataLayout.md new file mode 100644 --- /dev/null +++ b/mlir/docs/DataLayout.md @@ -0,0 +1,238 @@ +# Data Layout Modeling + +Data layout information allows the compiler to answer questions related to how a +value of a particular type is stored in memory. For example, the size of a value +or its address alignment requirements. It enables, among others, the generation +of various linear memory addressing schemes for containers of abstract types and +deeper reasoning about vectors. + +The data layout subsystem is designed to scale to MLIR's open type and operation +system. At the top level, it consists of: + +* attribute interfaces that can be implemented by concrete data layout + specifications; +* type interfaces that should be implemented by types subject to data layout; +* operation interfaces that must be implemented by operations that can serve + as data layout scopes (e.g., modules); +* and dialect interfaces for data layout properties unrelated to specific + types. + +Built-in types are handled specially to decrease the overall query cost. + +## Usage + +### Scoping + +Following MLIR's nested structure, data layout properties are _scoped_ to +regions belonging to specific operations that implement the +`DataLayoutOpInterface`. Such scoping operations partially control the data +layout properties and may have attributes that affect them, typically organized +in a data layout specification. + +Types may have a different data layout in different scopes, including scopes +that are nested in other scopes such as modules contained in other modules. At +the same time, within the given scope excluding any nested scope, a given type +has fixed data layout properties. Types are also expected to have a default, +"natural" data layout in case they are used outside of any operation that +provides data layout scope for them. This ensure data layout queries always have +a valid result. + +### Compatibility and Transformations + +The information necessary to compute layout properties can be combined from +nested scopes. For example, an outer scope can define layout properties for a +subset of types while inner scopes define them for a disjoint subset, or scopes +can progressively relax alignment requirements on a type. This mechanism is +supported by the notion of data layout _compatibility_: the layout defined in a +nested scope is expected to be compatible with that of the outer scope. MLIR +does not prescribe what compatibility means for particular ops and types but +provides hooks for them to provide target- and type-specific checks. For +example, one may want to only allow relaxation of alignment constraints (i.e., +smaller alignment) in nested modules or, alternatively, one may require nested +modules to fully redefine all constraints of the outer scope. + +Data layout compatibility is also relevant during IR transformation. Any +transformation that affects the data layout scoping operation is expected to +maintain data layout compatibility. It is under responsibility of the +transformation to ensure it is indeed the case. + +### Queries + +Data layout property queries can be performed on the special object -- +`DataLayout` -- which can be created for the given scoping operation. These +objects allow one to interface with the data layout infrastructure and query +properties of given types in the scope of the object. The signature of +`DataLayout` class is as follows. + +```c++ +class DataLayout { +public: + explicit DataLayout(DataLayoutOpInterface scope); + + unsigned getTypeSize(Type type) const; + unsigned getTypeABIAlignment(Type type) const; + unsigned getTypePreferredAlignment(Type type) const; +}; +``` + +The user can construct the `DataLayout` object for the scope of interest. Since +the data layout properties are fixed in the scope, they will be computed only +once upon first request and cached for further use. Therefore, +`DataLayout(op.getParentOfType()).getTypeSize(type)` is +considered an anti-pattern since it discards the cache after use. Because of +caching, a `DataLayout` object returns valid results as long as the data layout +properties of enclosing scopes remain the same, that is, as long as none of the +ancestor operations are modified in a way that affects data layout. After such a +modification, the user is expected to create a fresh `DataLayout` object. To aid +with this, `DataLayout` asserts that the scope remains identical if MLIR is +compiled with assertions enabled. + +## Custom Implementations + +Extensibility of the data layout modeling is provided through a set of MLIR +[Interfaces](Interfaces.md). + +### Data Layout Specifications + +Data layout specification is an [attribute](LangRef.md#attributes) that is +conceptually a collection of key-value pairs called data layout specification +_entries_. Data layout specification attributes implement the +`DataLayoutSpecInterface`, described below. Each entry is itself an attribute +that implements the `DataLayoutEntryInterface`. Entries have a key, either a +`Type` or an `Identifier`, and a value. Keys are used to associate entries with +specific types or dialects: when handling a data layout properties request, a +type or a dialect can only see the specification entries relevant to them and +must go through the supplied `DataLayout` object for any recursive query. This +supports and enforces better composability because types cannot (and should not) +understand layout details of other types. Entry values are arbitrary attributes, +specific to the type. + +For example, a data layout specification may be an actual list of pairs with +simple custom syntax resembling the following: + +``` +#my_dialect.layout_spec< + #my_dialect.layout_entry, + #my_dialect.layout_entry<"my_dialect.endianness", "little">, + #my_dialect.layout_entry> +``` + +The exact details of the specification and entry attributes, as well as their +syntax, are up to implementations. + +We use the notion of _type class_ throughout the data layout subsystem. It +corresponds to the C++ class of the given type, e.g., `IntegerType` for built-in +integers. MLIR does not have a mechanism to represent type classes in the IR. +Instead, data layout entries contain specific _instances_ of a type class, for +example, `IntegerType{signedness=signless, bitwidth=8}` (or `i8` in the IR) or +`IntegerType{signedness=unsigned, bitwidth=32}` (or `ui32` in the IR). When +handling a data layout property query, a type class will be supplied with _all_ +entries with keys belonging to this type class. For example, `IntegerType` will +see the entries for `i8`, `si16` and `ui32`, but will _not_ see those for `f32` +or `memref` (neither will `MemRefType` see the entry for `i32`). This +allows for type-specific "interpolation" behavior where a type class can compute +data layout properties of _any_ specific type instance given properties of other +instances. Using integers as an example again, their alignment could be computed +by taking that of the closest from above integer type with power-of-two +bitwidth. + +[include "Interfaces/DataLayoutAttrInterface.md"] + +### Data Layout Scoping Operations + +Operations that define a scope for data layout queries, and that can be used to +create a `DataLayout` object, are expected to implement the +`DataLayoutOpInterface`. Such ops must provide at least a way of obtaining the +data layout specification. The specification need not be necessarily attached to +the operation as an attribute and may be constructed on-the-fly; it is only +fetched once per `DataLayout` object and cached. Such ops may also provide +custom handlers for data layout queries that provide results without forwarding +the queries down to specific types or post-processing the results returned by +types in target- or scope-specific ways. These custom handlers make it possible +for scoping operations to (re)define data layout properties for types without +having to modify the types themselves, e.g., when types are defined in another +dialect. + +[include "Interfaces/DataLayoutOpInterface.md"] + +### Types with Data Layout + +Type classes that intend to handle data layout queries themselves are expected +to implement the `DataLayoutTypeInterface`. This interface provides overridable +hooks for each data layout query. Each of these hooks is supplied with the type +instance, a `DataLayout` object suitable for recursive queries, and a list of +data layout queries relevant for the type class. It is expected to provide a +valid result even if the list of entries is empty. These hooks do not have +access to the operation in the scope of which the query is handled and should +use the supplied entries instead. + +[include "Interfaces/DataLayoutTypeInterface.md"] + +### Dialects with Data Layout Identifiers + +For data layout entries that are not related to a particular type class, the key +of the entry is an Identifier that belongs to some dialect. In this case, the +dialect is expected to implement the `DataLayoutDialectInterface`. This dialect +provides hooks for verifying the validity of the entry value attributes and for +and the compatibility of nested entries. + +### Query Dispatch + +The overall flow of a data layout property query is as follows. + +- The user constructs a `DataLayout` at the given scope. The constructor + fetches the data layout specification and combines it with those of + enclosing scopes (layouts are expected to be compatible). +- The user calls `DataLayout::query(Type ty)`. +- If `DataLayout` has a cached response, this response is returned + immediately. +- Otherwise, the query is handed down by `DataLayout` to + `DataLayoutOpInterface::query(ty, *this, relevantEntries)` where the + relevant entries are computed as described above. +- Unless the `query` hook is reimplemented by the op interface, the query is + handled further down to `DataLayoutTypeInterface::query(dataLayout, + relevantEntries)` after casting `ty` to the type interface. If the type does + not implement the interface, an unrecoverable fatal error is produced. +- The type is expected to always provide the response, which is returned up + the call stack and cached by the `DataLayout.` + +## Default Implementation + +The default implementation of the data layout interfaces directly handles +queries for a subset of built-in types. + +### Built-in Types + +The following describes the default properties of built-in types. + +The size of built-in integers and floats in bytes is computed as +`ceildiv(bitwidth, 8)`. The ABI alignment of integer types with bitwidth below +64 and of the float types is the closest from above power-of-two number of +bytes. The ABI alignment of integer types with bitwidth 64 and above is 4 bytes +(32 bits). + +The size of built-in vectors is computed by first rounding their number of +elements in the _innermost_ dimension to the closest power-of-two from above, +then getting the total number of elements, and finally multiplying it with the +element size. For example, `vector<3xi32>` and `vector<4xi32>` have the same +size. So do `vector<2x3xf32>` and `vector<2x4xf32>`, but `vector<3x4xf32>` and +`vector<4x4xf32>` have different sizes. The ABI and preferred alignment of +vector types is computed by taking the innermost dimension of the vector, +rounding it up to the closest power-of-two, taking a product of that with +element size in bytes, and rounding the result up again to the closest +power-of-two. + +Note: these values are selected for consistency with the +[default data layout in LLVM](https://llvm.org/docs/LangRef.html#data-layout), +which MLIR assumed until the introduction of proper data layout modeling, and +with the +[modeling of n-D vectors](https://mlir.llvm.org/docs/Dialects/Vector/#deeperdive). +They **may change** in the future. + +### DLTI Dialect + +The [DLTI](Dialects/DLTI.md) dialect provides the attributes implementing +`DataLayoutSpecInterface` and `DataLayoutEntryInterface`, as well as a dialect +attribute that can be used to attach the specification to a given operation. The +verifier of this attribute triggers those of the specification and checks the +compatiblity of nested specifications. diff --git a/mlir/include/mlir/Dialect/CMakeLists.txt b/mlir/include/mlir/Dialect/CMakeLists.txt --- a/mlir/include/mlir/Dialect/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/CMakeLists.txt @@ -4,6 +4,7 @@ add_subdirectory(ArmSVE) add_subdirectory(AVX512) add_subdirectory(Complex) +add_subdirectory(DLTI) add_subdirectory(GPU) add_subdirectory(Math) add_subdirectory(Linalg) diff --git a/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/CMakeLists.txt @@ -0,0 +1,2 @@ +add_mlir_dialect(DLTI dlti) +add_mlir_doc(DLTI -gen-dialect-doc DLTIDialect Dialects/) diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.h b/mlir/include/mlir/Dialect/DLTI/DLTI.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/DLTI.h @@ -0,0 +1,112 @@ +//===- DLTI.h - Data Layout and Target Info MLIR Dialect --------*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Defines the dialect containing the objects pertaining to target information. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_DLTI_DLTI_H +#define MLIR_DIALECT_DLTI_DLTI_H + +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Dialect.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" + +namespace mlir { +namespace impl { +class DataLayoutEntryStorage; +class DataLayoutSpecStorage; +} // namespace impl + +//===----------------------------------------------------------------------===// +// DataLayoutEntryAttr +//===----------------------------------------------------------------------===// + +/// A data layout entry attribute is a key-value pair where the key is a type or +/// an identifier and the value is another attribute. These entries form a data +/// layout specification. +class DataLayoutEntryAttr + : public Attribute::AttrBase { +public: + using Base::Base; + + /// The keyword used for this attribute in custom syntax. + constexpr const static llvm::StringLiteral kAttrKeyword = "dl_entry"; + + /// Returns the entry with the given key and value. + static DataLayoutEntryAttr get(Identifier key, Attribute value); + static DataLayoutEntryAttr get(Type key, Attribute value); + + /// Returns the key of this entry. + DataLayoutEntryKey getKey() const; + + /// Returns the value of this entry. + Attribute getValue() const; + + /// Parses an instance of this attribute. + static DataLayoutEntryAttr parse(DialectAsmParser &parser); + + /// Prints this attribute. + void print(DialectAsmPrinter &os) const; +}; + +//===----------------------------------------------------------------------===// +// DataLayoutSpecAttr +//===----------------------------------------------------------------------===// + +/// A data layout specification is a list of entries that specify (partial) data +/// layout information. It is expected to be attached to operations that serve +/// as scopes for data layout requests. +class DataLayoutSpecAttr + : public Attribute::AttrBase { +public: + using Base::Base; + + /// The keyword used for this attribute in custom syntax. + constexpr const static StringLiteral kAttrKeyword = "dl_spec"; + + /// Returns the specification containing the given list of keys. + static DataLayoutSpecAttr get(MLIRContext *ctx, + ArrayRef entries); + + /// Returns the specification containing the given list of keys. If the list + /// contains duplicate keys or is otherwise invalid, reports errors using the + /// given callback and returns null. + static DataLayoutSpecAttr + getChecked(function_ref emitError, MLIRContext *context, + ArrayRef entries); + + /// Checks that the given list of entries does not contain duplicate keys. + static LogicalResult verify(function_ref emitError, + ArrayRef entries); + + /// Combines this specification with `specs`, enclosing specifications listed + /// from outermost to innermost. This overwrites the older entries with the + /// same key as the newer entries if the entries are compatible. Returns null + /// if the specifications are not compatible. + DataLayoutSpecAttr combineWith(ArrayRef specs) const; + + /// Returns the list of entries. + DataLayoutEntryListRef getEntries() const; + + /// Parses an instance of this attribute. + static DataLayoutSpecAttr parse(DialectAsmParser &parser); + + /// Prints this attribute. + void print(DialectAsmPrinter &os) const; +}; + +} // namespace mlir + +#include "mlir/Dialect/DLTI/DLTIDialect.h.inc" + +#endif // MLIR_DIALECT_DLTI_DLTI_H diff --git a/mlir/include/mlir/Dialect/DLTI/DLTI.td b/mlir/include/mlir/Dialect/DLTI/DLTI.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/DLTI.td @@ -0,0 +1,14 @@ +//===- DLTI.td - Data Layout and Target Info Dialect --------*- tablegen -*-==// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef DLTI_TD +#define DLTI_TD + +include "mlir/Dialect/DLTI/DLTIBase.td" + +#endif // DLTI_TD diff --git a/mlir/include/mlir/Dialect/DLTI/DLTIBase.td b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/DLTIBase.td @@ -0,0 +1,62 @@ +//===- DLTIBase.td - Target information dialect base defs ---*- tablegen -*-==// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef DLTI_BASE +#define DLTI_BASE + +include "mlir/IR/OpBase.td" + +def DLTI_Dialect : Dialect { + let name = "dlti"; + let cppNamespace = "::mlir"; + let hasOperationAttrVerify = 1; + + let description = [{ + The Data Layout and Target Information (DLTI) dialect is intended to hold + attributes and other components pertaining to descriptions of in-memory + data layout and compilation targets. + }]; + + let extraClassDeclaration = [{ + constexpr const static ::llvm::StringLiteral + kDataLayoutAttrName = "dlti.dl_spec"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutEndiannessKey = "dlti.endianness"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutEndiannessBig = "big"; + + constexpr const static ::llvm::StringLiteral + kDataLayoutEndiannessLittle = "little"; + }]; +} + +def DLTI_DataLayoutEntryAttr : DialectAttr< + DLTI_Dialect, + CPred<"$_self.isa<::mlir::DataLayoutEntryAttr>()">, + "Target data layout entry"> { + let storageType = "::mlir::DataLayoutEntryAttr"; + let returnType = "::mlir::DataLayoutEntryAttr"; + let convertFromStorage = "$_self"; +} + +def DLTI_DataLayoutSpecAttr : DialectAttr< + DLTI_Dialect, + CPred<"$_self.isa<::mlir::DataLayoutSpecAttr>()">, + "Target data layout specification"> { + let storageType = "::mlir::DataLayoutSpecAttr"; + let returnType = "::mlir::DataLayoutSpecAttr"; + let convertFromStorage = "$_self"; +} + +def HasDefaultDLTIDataLayout : NativeOpTrait<"HasDefaultDLTIDataLayout"> { + let cppNamespace = "::mlir"; +} + +#endif // DLTI_BASE diff --git a/mlir/include/mlir/Dialect/DLTI/Traits.h b/mlir/include/mlir/Dialect/DLTI/Traits.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/DLTI/Traits.h @@ -0,0 +1,47 @@ +//===- Traits.h - Trait Declaration for MLIR Target dialect -----*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Defines the dialect containing the objects pertaining to target information. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DIALECT_TARGET_TRAITS_H +#define MLIR_DIALECT_TARGET_TRAITS_H + +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" + +namespace mlir { +class DataLayoutSpecAttr; + +namespace impl { +LogicalResult verifyHasDefaultDLTIDataLayoutTrait(Operation *op); +DataLayoutSpecInterface getDataLayoutSpec(Operation *op); +} // namespace impl + +/// Trait to be used by operations willing to use the implementation of the +/// data layout interfaces provided by the Target dialect. +template +class HasDefaultDLTIDataLayout + : public OpTrait::TraitBase { +public: + /// Verifies that the operation to which this trait is attached is valid for + /// the trait, i.e., that it implements the data layout operation interface. + static LogicalResult verifyTrait(Operation *op) { + return impl::verifyHasDefaultDLTIDataLayoutTrait(op); + } + + /// Returns the data layout specification as provided by the Target dialect + /// specification attribute. + DataLayoutSpecInterface getDataLayoutSpec() { + return impl::getDataLayoutSpec(this->getOperation()); + } +}; +} // namespace mlir + +#endif // MLIR_DIALECT_TARGET_TRAITS_H 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 @@ -240,7 +240,7 @@ string description = ?; // A list of dialects this dialect will load on construction as dependencies. - // These are dialects that this dialect may involved in canonicalization + // These are dialects that this dialect may involve in canonicalization // pattern or interfaces. list dependentDialects = []; diff --git a/mlir/include/mlir/InitAllDialects.h b/mlir/include/mlir/InitAllDialects.h --- a/mlir/include/mlir/InitAllDialects.h +++ b/mlir/include/mlir/InitAllDialects.h @@ -20,6 +20,7 @@ #include "mlir/Dialect/ArmSVE/ArmSVEDialect.h" #include "mlir/Dialect/Async/IR/Async.h" #include "mlir/Dialect/Complex/IR/Complex.h" +#include "mlir/Dialect/DLTI/DLTI.h" #include "mlir/Dialect/GPU/GPUDialect.h" #include "mlir/Dialect/LLVMIR/LLVMArmSVEDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" @@ -53,6 +54,7 @@ async::AsyncDialect, avx512::AVX512Dialect, complex::ComplexDialect, + DLTIDialect, gpu::GPUDialect, LLVM::LLVMDialect, LLVM::LLVMArmSVEDialect, diff --git a/mlir/include/mlir/Interfaces/CMakeLists.txt b/mlir/include/mlir/Interfaces/CMakeLists.txt --- a/mlir/include/mlir/Interfaces/CMakeLists.txt +++ b/mlir/include/mlir/Interfaces/CMakeLists.txt @@ -9,3 +9,16 @@ add_mlir_interface(VectorInterfaces) add_mlir_interface(ViewLikeInterface) +set(LLVM_TARGET_DEFINITIONS DataLayoutInterfaces.td) +mlir_tablegen(DataLayoutAttrInterface.h.inc -gen-attr-interface-decls) +mlir_tablegen(DataLayoutAttrInterface.cpp.inc -gen-attr-interface-defs) +mlir_tablegen(DataLayoutAttrInterface.md -gen-attr-interface-docs) +mlir_tablegen(DataLayoutOpInterface.h.inc -gen-op-interface-decls) +mlir_tablegen(DataLayoutOpInterface.cpp.inc -gen-op-interface-defs) +mlir_tablegen(DataLayoutOpInterface.md -gen-op-interface-docs) +mlir_tablegen(DataLayoutTypeInterface.h.inc -gen-type-interface-decls) +mlir_tablegen(DataLayoutTypeInterface.cpp.inc -gen-type-interface-defs) +mlir_tablegen(DataLayoutTypeInterface.md -gen-type-interface-docs) +add_public_tablegen_target(MLIRDataLayoutInterfacesIncGen) +add_dependencies(mlir-generic-headers MLIRDataLayoutInterfacesIncGen) + diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.h @@ -0,0 +1,172 @@ +//===- DataLayoutInterfaces.h - Data Layout Interface Decls -----*- C++ -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Defines the interfaces for the data layout specification, operations to which +// they can be attached, types subject to data layout and dialects containing +// data layout entries. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_INTERFACES_DATALAYOUTINTERFACES_H +#define MLIR_INTERFACES_DATALAYOUTINTERFACES_H + +#include "mlir/IR/DialectInterface.h" +#include "mlir/IR/OpDefinition.h" +#include "llvm/ADT/DenseMap.h" + +namespace mlir { +class DataLayout; +class DataLayoutEntryInterface; +using DataLayoutEntryKey = llvm::PointerUnion; +// Using explicit SmallVector size because we cannot infer the size from the +// forward declaration, and we need the typedef in the actual declaration. +using DataLayoutEntryList = llvm::SmallVector; +using DataLayoutEntryListRef = llvm::ArrayRef; +class DataLayoutOpInterface; +class DataLayoutSpecInterface; + +namespace detail { +/// Default handler for the type size request. Computes results for built-in +/// types and dispatches to the DataLayoutTypeInterface for other types. +unsigned getDefaultTypeSize(Type type, const DataLayout &dataLayout, + ArrayRef params); + +/// Default handler for the required alignemnt request. Computes results for +/// built-in types and dispatches to the DataLayoutTypeInterface for other +/// types. +unsigned getDefaultABIAlignment(Type type, const DataLayout &dataLayout, + ArrayRef params); + +/// Default handler for the preferred alignemnt request. Computes results for +/// built-in types and dispatches to the DataLayoutTypeInterface for other +/// types. +unsigned +getDefaultPreferredAlignment(Type type, const DataLayout &dataLayout, + ArrayRef params); + +/// Given a list of data layout entries, returns a new list containing the +/// entries with keys having the given type ID, i.e. belonging to the same type +/// class. +DataLayoutEntryList filterEntriesForType(DataLayoutEntryListRef entries, + TypeID typeID); + +/// Given a list of data layout entries, returns the entry that has the given +/// identifier as key, if such an entry exists in the list. +DataLayoutEntryInterface +filterEntryForIdentifier(DataLayoutEntryListRef entries, Identifier id); + +/// Verifies that the operation implementing the data layout interface is valid. +/// This calls the verifier of the spec attribute and checks if the layout is +/// compatible with specs attached to the enclosing operations. +LogicalResult verifyDataLayoutOp(DataLayoutOpInterface op); + +/// Verifies that a data layout spec is valid. This dispatches to individual +/// entry verifiers, and then to the verifiers implemented by the relevant type +/// and dialect interfaces for type and identifier keys respectively. +LogicalResult verifyDataLayoutSpec(DataLayoutSpecInterface spec, Location loc); +} // namespace detail +} // namespace mlir + +#include "mlir/Interfaces/DataLayoutAttrInterface.h.inc" +#include "mlir/Interfaces/DataLayoutOpInterface.h.inc" +#include "mlir/Interfaces/DataLayoutTypeInterface.h.inc" + +namespace mlir { + +//===----------------------------------------------------------------------===// +// DataLayoutDialectInterface +//===----------------------------------------------------------------------===// + +/// An interface to be implemented by dialects that can have identifiers in the +/// data layout specification entries. Provides hooks for verifying the entry +/// validity and combining two entries. +class DataLayoutDialectInterface + : public DialectInterface::Base { +public: + DataLayoutDialectInterface(Dialect *dialect) : Base(dialect) {} + + /// Checks whether the given data layout entry is valid and reports any errors + /// at the provided location. Derived classes should override this. + virtual LogicalResult verifyEntry(DataLayoutEntryInterface entry, + Location loc) const { + return success(); + } + + /// Default implementation of entry combination that combines identical + /// entries and returns null otherwise. + static DataLayoutEntryInterface + defaultCombine(DataLayoutEntryInterface outer, + DataLayoutEntryInterface inner) { + if (!outer || outer == inner) + return inner; + return {}; + } + + /// Combines two entries with identifiers that belong to this dialect. Returns + /// the combined entry or null if the entries are not compatible. Derived + /// classes likely need to reimplement this. + virtual DataLayoutEntryInterface + combine(DataLayoutEntryInterface outer, + DataLayoutEntryInterface inner) const { + return defaultCombine(outer, inner); + } +}; + +//===----------------------------------------------------------------------===// +// DataLayout +//===----------------------------------------------------------------------===// + +/// The main mechanism for performing data layout queries. Instances of this +/// class can be created for an operation implementing DataLayoutOpInterface. +/// Upon construction, a layout spec combining that of the given operation with +/// all its ancestors will be computed and used to handle further requests. For +/// efficiency, results to all requests will be cached in this object. +/// Therefore, if the data layout spec for the scoping operation, or any of the +/// enclosing operations, changes, the cache is no longer valid. The user is +/// responsible creating a new DataLayout object after any spec change. In debug +/// mode, the cache validity is being checked in every request. +class DataLayout { +public: + explicit DataLayout(DataLayoutOpInterface op); + + /// Returns the size of the given type in the current scope. + unsigned getTypeSize(Type t) const; + + /// Returns the required alignment of the given type in the current scope. + unsigned getTypeABIAlignment(Type t) const; + + /// Returns the preferred of the given type in the current scope. + unsigned getTypePreferredAlignment(Type t) const; + +private: + /// Combined layout spec at the given scope. + const DataLayoutSpecInterface originalLayout; + +#ifndef NDEBUG + /// List of enclosing layout specs. + SmallVector layoutStack; +#endif + + /// Asserts that the cache is still valid. Expensive in debug mode. No-op in + /// release mode. + void checkValid() const; + + /// Operation defining the scope of requests. + // TODO: this is mutable because the generated interface method are not const. + // Update the generator to support const methods and change this to const. + mutable DataLayoutOpInterface scope; + + /// Caches for individual requests. + mutable DenseMap sizes; + mutable DenseMap abiAlignments; + mutable DenseMap preferredAlignments; +}; + +} // namespace mlir + +#endif // MLIR_INTERFACES_DATALAYOUTINTERFACES_H diff --git a/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Interfaces/DataLayoutInterfaces.td @@ -0,0 +1,327 @@ +//===- DataLayoutInterfaces.td - Data layout interfaces ----*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// Defines the interfaces for the data layout specification, operations to which +// they can be attached, and types that are subject to data layout. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_DATALAYOUTINTERFACES +#define MLIR_DATALAYOUTINTERFACES + +include "mlir/IR/OpBase.td" + +//===----------------------------------------------------------------------===// +// Attribute interfaces +//===----------------------------------------------------------------------===// + +def DataLayoutEntryInterface : AttrInterface<"DataLayoutEntryInterface"> { + let cppNamespace = "::mlir"; + + let description = [{ + Attribute interface describing an entry in a data layout specification. + + A data layout specification entry is a key-value pair. Its key is either a + type, when the entry is related to a type or a class of types, or an + identifier, when it is not. `DataLayoutEntryKey` is an alias allowing one + to use both key types. Its value is an arbitrary attribute that is + interpreted either by the type for type keys or by the dialect containing + the identifier for identifier keys. The interface provides a hook that + can be used by specific implementations to delegate the verification of + attribute fitness for a particular key to the relevant type or dialect. + }]; + + let methods = [ + InterfaceMethod< + /*description=*/"Returns the key of the this layout entry.", + /*retTy=*/"::mlir::DataLayoutEntryKey", + /*methodName=*/"getKey", + /*args=*/(ins) + >, + InterfaceMethod< + /*description=*/"Returns the value of this layout entry.", + /*retTy=*/"::mlir::Attribute", + /*methodName=*/"getValue", + /*args=*/(ins) + >, + InterfaceMethod< + /*description=*/"Checks that the entry is well-formed, reports errors " + "at the provided location.", + /*retTy=*/"::mlir::LogicalResult", + /*methodName=*/"verifyEntry", + /*args=*/(ins "::mlir::Location":$loc), + /*methodBody=*/"", + /*defaultImplementation=*/[{ return ::mlir::success(); }] + > + ]; + + let extraClassDeclaration = [{ + /// Returns `true` if the key of this entry is a type. + bool isTypeEntry() { + return getKey().is<::mlir::Type>(); + } + }]; +} + +def DataLayoutSpecInterface : AttrInterface<"DataLayoutSpecInterface"> { + let cppNamespace = "::mlir"; + + let description = [{ + Attribute interface describing a data layout specification. + + A data layout specification is seen as a sequence of entries, each of which + is an attribute implementing the data layout entry interface. It assumes + a contiguous underlying storage for entries. The interface provides a hook + for implementations to verify the well-formedness of the specification, + with a default implementation that verifies the absence of entries with + duplicate keys and the well-formedness of each individual entry before + dispatching to the type or dialect the entry is associated with. + + Data layout specifications may need to be combined in case they appear on + nested operations subject to layout, or to ensure the validity of layout + modification. Concerete specification attributes must implement the + corresponding hook. + }]; + // The underlying storage being contiguous may be revised in the future, but + // care must be taken to avoid materializing or copying the entire list of + // entries. + + let methods = [ + InterfaceMethod< + /*description=*/"Combines the current layout with the given list of " + "layouts, provided from the outermost (oldest) to the " + "innermost (newest). Returns null on failure.", + /*retTy=*/"::mlir::DataLayoutSpecInterface", + /*methodName=*/"combineWith", + /*args=*/(ins "::llvm::ArrayRef":$specs) + >, + InterfaceMethod< + /*description=*/"Returns the list of layout entries.", + /*retTy=*/"::mlir::DataLayoutEntryListRef", + /*methodName=*/"getEntries", + /*args=*/(ins) + >, + // Implementations may override this if they have an efficient lookup + // mechanism. + InterfaceMethod< + /*description=*/"Returns a copy of the entries related to a specific " + "type class regardles of type parameters. ", + /*retTy=*/"::mlir::DataLayoutEntryList", + /*methodName=*/"getSpecForType", + /*args=*/(ins "::mlir::TypeID":$type), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::filterEntriesForType($_attr.getEntries(), type); + }] + >, + // Implementations may override this if they have an efficient lookup + // mechanism. + InterfaceMethod< + /*description=*/"Returns the entry related to the given identifier, if " + "present.", + /*retTy=*/"::mlir::DataLayoutEntryInterface", + /*methodName=*/"getSpecForIdentifier", + /*args=*/(ins "::mlir::Identifier":$identifier), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::filterEntryForIdentifier($_attr.getEntries(), + identifier); + }] + >, + InterfaceMethod< + /*description=*/"Verifies the validity of the specification and reports " + "any errors at the given location.", + /*retTy=*/"::mlir::LogicalResult", + /*methodName=*/"verifySpec", + /*args=*/(ins "::mlir::Location":$loc), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::verifyDataLayoutSpec($_attr, loc); + }] + >, + ]; + + let extraClassDeclaration = [{ + /// Returns a copy of the entries related to the type specified as template + /// parameter. + template + DataLayoutEntryList getSpecForType() { + return getSpecForType(TypeID::get()); + } + + /// Populates the given maps with lists of entries grouped by the type or + /// identifier they are associated with. Users are not expected to call this + /// method directly. + void bucketEntriesByType( + ::llvm::DenseMap<::mlir::TypeID, ::mlir::DataLayoutEntryList> &types, + ::llvm::DenseMap<::mlir::Identifier, + ::mlir::DataLayoutEntryInterface> &ids); + }]; +} + +//===----------------------------------------------------------------------===// +// Operation interface +//===----------------------------------------------------------------------===// + +def DataLayoutOpInterface : OpInterface<"DataLayoutOpInterface"> { + let cppNamespace = "::mlir"; + + let description = [{ + Interface for operations that can have a data layout specification attached. + + The `DataLayout` object, which can be used for data layout queries, can be + constructed for such operations. The absence of a data layout specification + must be handled without failing. + + Concrete operations must implement the hook returning the data layout + specification. They may optionally override the methods used in data layout + queries, default implementations of which provide predefined answers for + built-in types and dispatch to the type interface for all other types. These + methods must be idempotent, that is return the same result on repeated + queries with the same parameters. They are declared static and therefore + have no access to the operation or its attributes. Instead, they receive a + list of data layout entries relevant to the request. The entries are known + to have passed the spec and entry verifier. + }]; + + let methods = [ + InterfaceMethod< + /*description=*/"Returns the data layout specification for this op, or " + "null if it does not exist.", + /*retTy=*/"DataLayoutSpecInterface", + /*methodName=*/"getDataLayoutSpec", + /*args=*/(ins) + >, + StaticInterfaceMethod< + /*description=*/"Returns the size of the given type computed using the " + "relevant entries. The data layout object can be used " + "for recursive queries.", + /*retTy=*/"unsigned", + /*methodName=*/"getTypeSize", + /*args=*/(ins "::mlir::Type":$type, + "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultTypeSize(type, dataLayout, params); + }] + >, + StaticInterfaceMethod< + /*description=*/"Returns the alignment required by the ABI for the given " + "type computed using the relevant entries. The data " + "layout object can be used for recursive queries.", + /*retTy=*/"unsigned", + /*methodName=*/"getTypeABIAlignment", + /*args=*/(ins "::mlir::Type":$type, + "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultABIAlignment(type, dataLayout, params); + }] + >, + StaticInterfaceMethod< + /*description=*/"Returns the alignment preferred by the given type " + "computed using the relevant entries. The data layout" + "object can be used for recursive queries.", + /*retTy=*/"unsigned", + /*methodName=*/"getTypePreferredAlignment", + /*args=*/(ins "::mlir::Type":$type, + "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return ::mlir::detail::getDefaultPreferredAlignment(type, dataLayout, + params); + }] + >, + ]; + + let verify = [{ return ::mlir::detail::verifyDataLayoutOp($_op); }]; +} + +//===----------------------------------------------------------------------===// +// Type interface +//===----------------------------------------------------------------------===// + +def DataLayoutTypeInterface : TypeInterface<"DataLayoutTypeInterface"> { + let cppNamespace = "::mlir"; + + let description = [{ + Interface for types subject to data layout. + + Types willing to be supported by the data layout subsystem should implement + this interface by providing implementations of functions querying their + size, required and preferred alignment. Each of these functions accepts as + arguments a data layout object that can be used to perform recursive queries + in the same scope, and a list of data layout entries relevant to this type. + Specifically, the entries are those that have as key _any instance_ of the + same type class as the current type. For example, if IntegerType had + implemented this interface, it would have received the entries with keys i1, + i2, i8, etc. regardless of the bitwidth of this type. This mechanism allows + types to "interpolate" the results in a type-specific way instead of listing + all possible types in the specification. + + The list of entries may be empty, in which case the type must provide a + reasonable default value. The entries in the list are known to have passed + the spec and the entry verifiers, as well as the type-specifid verifier if + provided. + + In case of nested layout specs or spec changes, the type can override a hook + indicating whether the outer (old) and the inner (new) spec are compatible. + }]; + + let methods = [ + InterfaceMethod< + /*description=*/"Returns the size of this type in bytes.", + /*retTy=*/"unsigned", + /*methodName=*/"getTypeSize", + /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params) + >, + InterfaceMethod< + /*description=*/"Returns the ABI-required alignment for this type, " + "in bytes", + /*retTy=*/"unsigned", + /*methodName=*/"getABIAlignment", + /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params) + >, + InterfaceMethod< + /*description=*/"Returns the preferred alignemnt for this type, " + "in bytes.", + /*retTy=*/"unsigned", + /*methodName=*/"getPreferredAlignment", + /*args=*/(ins "const ::mlir::DataLayout &":$dataLayout, + "::mlir::DataLayoutEntryListRef":$params) + >, + InterfaceMethod< + /*desc=*/"Returns true if the two lists of entries are compatible, that " + "is, that `newLayout` spec entries can be nested in an op with " + "`oldLayout` spec entries.", + /*retTy=*/"bool", + /*methodName=*/"areCompatible", + /*args=*/(ins "::mlir::DataLayoutEntryListRef":$oldLayout, + "::mlir::DataLayoutEntryListRef":$newLayout), + /*methodBody=*/"", + /*defaultImplementation=*/[{ return true; }] + >, + InterfaceMethod< + /*description=*/"Verifies that the given list of entries is valid for " + "this type.", + /*retTy=*/"::mlir::LogicalResult", + /*methodName=*/"verifyEntries", + /*args=*/(ins "::mlir::DataLayoutEntryListRef":$entries, + "::mlir::Location":$loc), + /*methodBody=*/"", + /*defaultImplementation=*/[{ return ::mlir::success(); }] + >, + ]; +} + +#endif // MLIR_DATALAYOUTINTERFACES diff --git a/mlir/lib/Dialect/CMakeLists.txt b/mlir/lib/Dialect/CMakeLists.txt --- a/mlir/lib/Dialect/CMakeLists.txt +++ b/mlir/lib/Dialect/CMakeLists.txt @@ -4,6 +4,7 @@ add_subdirectory(Async) add_subdirectory(AVX512) add_subdirectory(Complex) +add_subdirectory(DLTI) add_subdirectory(GPU) add_subdirectory(Linalg) add_subdirectory(LLVMIR) diff --git a/mlir/lib/Dialect/DLTI/CMakeLists.txt b/mlir/lib/Dialect/DLTI/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/DLTI/CMakeLists.txt @@ -0,0 +1,11 @@ +add_mlir_dialect_library(MLIRDLTI + DLTI.cpp + Traits.cpp + + DEPENDS + MLIRDLTIIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRDataLayoutInterfaces + ) diff --git a/mlir/lib/Dialect/DLTI/DLTI.cpp b/mlir/lib/Dialect/DLTI/DLTI.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/DLTI/DLTI.cpp @@ -0,0 +1,377 @@ +//===- DLTI.cpp - Data Layout And Target Info MLIR Dialect Implementation -===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinDialect.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/DialectImplementation.h" +#include "llvm/ADT/TypeSwitch.h" + +using namespace mlir; + +//===----------------------------------------------------------------------===// +// DataLayoutEntryAttr +//===----------------------------------------------------------------------===// +// +constexpr const StringLiteral mlir::DataLayoutEntryAttr::kAttrKeyword; + +namespace mlir { +namespace impl { +class DataLayoutEntryStorage : public AttributeStorage { +public: + using KeyTy = std::pair; + + DataLayoutEntryStorage(DataLayoutEntryKey entryKey, Attribute value) + : entryKey(entryKey), value(value) {} + + static DataLayoutEntryStorage *construct(AttributeStorageAllocator &allocator, + const KeyTy &key) { + return new (allocator.allocate()) + DataLayoutEntryStorage(key.first, key.second); + } + + bool operator==(const KeyTy &other) const { + return other.first == entryKey && other.second == value; + } + + DataLayoutEntryKey entryKey; + Attribute value; +}; +} // namespace impl +} // namespace mlir + +DataLayoutEntryAttr DataLayoutEntryAttr::get(Identifier key, Attribute value) { + return Base::get(key.getContext(), key, value); +} + +DataLayoutEntryAttr DataLayoutEntryAttr::get(Type key, Attribute value) { + return Base::get(key.getContext(), key, value); +} + +DataLayoutEntryKey DataLayoutEntryAttr::getKey() const { + return getImpl()->entryKey; +} + +Attribute DataLayoutEntryAttr::getValue() const { return getImpl()->value; } + +/// Parses an attribute with syntax: +/// attr ::= `#target.` `dl_entry` `<` (type | quoted-string) `,` attr `>` +DataLayoutEntryAttr DataLayoutEntryAttr::parse(DialectAsmParser &parser) { + if (failed(parser.parseLess())) + return {}; + + Type type = nullptr; + StringRef identifier; + llvm::SMLoc idLoc = parser.getCurrentLocation(); + OptionalParseResult parsedType = parser.parseOptionalType(type); + if (parsedType.hasValue() && failed(parsedType.getValue())) + return {}; + if (!parsedType.hasValue()) { + OptionalParseResult parsedString = parser.parseOptionalString(&identifier); + if (!parsedString.hasValue() || failed(parsedString.getValue())) { + parser.emitError(idLoc) << "expected a type or a quoted string"; + return {}; + } + } + + Attribute value; + if (failed(parser.parseComma()) || failed(parser.parseAttribute(value)) || + failed(parser.parseGreater())) + return {}; + + return type ? get(type, value) + : get(parser.getBuilder().getIdentifier(identifier), value); +} + +void DataLayoutEntryAttr::print(DialectAsmPrinter &os) const { + os << DataLayoutEntryAttr::kAttrKeyword << "<"; + if (auto type = getKey().dyn_cast()) + os << type; + else + os << "\"" << getKey().get().strref() << "\""; + os << ", " << getValue() << ">"; +} + +//===----------------------------------------------------------------------===// +// DataLayoutSpecAttr +//===----------------------------------------------------------------------===// +// +constexpr const StringLiteral mlir::DataLayoutSpecAttr::kAttrKeyword; + +namespace mlir { +namespace impl { +class DataLayoutSpecStorage : public AttributeStorage { +public: + using KeyTy = ArrayRef; + + DataLayoutSpecStorage(ArrayRef entries) + : entries(entries) {} + + bool operator==(const KeyTy &key) const { return key == entries; } + + static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator, + const KeyTy &key) { + return new (allocator.allocate()) + DataLayoutSpecStorage(allocator.copyInto(key)); + } + + ArrayRef entries; +}; +} // namespace impl +} // namespace mlir + +DataLayoutSpecAttr +DataLayoutSpecAttr::get(MLIRContext *ctx, + ArrayRef entries) { + return Base::get(ctx, entries); +} + +DataLayoutSpecAttr +DataLayoutSpecAttr::getChecked(function_ref emitError, + MLIRContext *context, + ArrayRef entries) { + return Base::getChecked(emitError, context, entries); +} + +LogicalResult +DataLayoutSpecAttr::verify(function_ref emitError, + ArrayRef entries) { + DenseSet types; + DenseSet ids; + for (DataLayoutEntryInterface entry : entries) { + if (auto type = entry.getKey().dyn_cast()) { + if (!types.insert(type).second) + return emitError() << "repeated layout entry key: " << type; + } else { + auto id = entry.getKey().get(); + if (!ids.insert(id).second) + return emitError() << "repeated layout entry key: " << id; + } + } + return success(); +} + +/// Given a list of old and a list of new entries, overwrites old entries with +/// new ones if they have matching keys, appends new entries to the old entry +/// list otherwise. +static void +overwriteDuplicateEntries(SmallVectorImpl &oldEntries, + ArrayRef newEntries) { + unsigned oldEntriesSize = oldEntries.size(); + for (DataLayoutEntryInterface entry : newEntries) { + // We expect a small (dozens) number of entries, so it is practically + // cheaper to iterate over the list linearly rather than to create an + // auxiliary hashmap to avoid duplication. Also note that we never need to + // check for duplicate keys the values that were added from `newEntries`. + bool replaced = false; + for (unsigned i = 0; i < oldEntriesSize; ++i) { + if (oldEntries[i].getKey() == entry.getKey()) { + oldEntries[i] = entry; + replaced = true; + break; + } + } + if (!replaced) + oldEntries.push_back(entry); + } +} + +/// Combines a data layout spec into the given lists of entries organized by +/// type class and identifier, overwriting them if necessary. Fails to combine +/// if the two entries with identical keys are not compatible. +static LogicalResult +combineOneSpec(DataLayoutSpecInterface spec, + DenseMap &entriesForType, + DenseMap &entriesForID) { + // A missing spec should be fine. + if (!spec) + return success(); + + DenseMap newEntriesForType; + DenseMap newEntriesForID; + spec.bucketEntriesByType(newEntriesForType, newEntriesForID); + + // Try overwriting the old entries with the new ones. + for (const auto &kvp : newEntriesForType) { + if (!entriesForType.count(kvp.first)) { + entriesForType[kvp.first] = std::move(kvp.second); + continue; + } + + Type typeSample = kvp.second.front().getKey().get(); + assert(&typeSample.getDialect() != + typeSample.getContext()->getLoadedDialect() && + "unexpected data layout entry for built-in type"); + + auto interface = typeSample.cast(); + if (!interface.areCompatible(entriesForType.lookup(kvp.first), kvp.second)) + return failure(); + + overwriteDuplicateEntries(entriesForType[kvp.first], kvp.second); + } + + for (const auto &kvp : newEntriesForID) { + Identifier id = kvp.second.getKey().get(); + Dialect *dialect = id.getDialect(); + if (!entriesForID.count(id)) { + entriesForID[id] = kvp.second; + continue; + } + + // Attempt to combine the enties using the dialect interface. If the + // dialect is not loaded for some reason, use the default combinator + // that conservatively accepts identical entries only. + entriesForID[id] = + dialect ? dialect->getRegisteredInterface() + ->combine(entriesForID[id], kvp.second) + : DataLayoutDialectInterface::defaultCombine(entriesForID[id], + kvp.second); + if (!entriesForID[id]) + return failure(); + } + + return success(); +} + +DataLayoutSpecAttr +DataLayoutSpecAttr::combineWith(ArrayRef specs) const { + // Only combine with attributes of the same kind. + // TODO: reconsider this when the need arises. + if (llvm::any_of(specs, [](DataLayoutSpecInterface spec) { + return !spec.isa(); + })) + return {}; + + // Combine all specs in order, with `this` being the last one. + DenseMap entriesForType; + DenseMap entriesForID; + for (DataLayoutSpecInterface spec : specs) + if (failed(combineOneSpec(spec, entriesForType, entriesForID))) + return nullptr; + if (failed(combineOneSpec(*this, entriesForType, entriesForID))) + return nullptr; + + // Rebuild the linear list of entries. + SmallVector entries; + llvm::append_range(entries, llvm::make_second_range(entriesForID)); + for (const auto &kvp : entriesForType) + llvm::append_range(entries, kvp.getSecond()); + + return DataLayoutSpecAttr::get(getContext(), entries); +} + +DataLayoutEntryListRef DataLayoutSpecAttr::getEntries() const { + return getImpl()->entries; +} + +/// Parses an attribute with syntax +/// attr ::= `#target.` `dl_spec` `<` attr-list? `>` +/// attr-list ::= attr +/// | attr `,` attr-list +DataLayoutSpecAttr DataLayoutSpecAttr::parse(DialectAsmParser &parser) { + if (failed(parser.parseLess())) + return {}; + + // Empty spec. + if (succeeded(parser.parseOptionalGreater())) + return get(parser.getBuilder().getContext(), {}); + + SmallVector entries; + do { + entries.emplace_back(); + if (failed(parser.parseAttribute(entries.back()))) + return {}; + } while (succeeded(parser.parseOptionalComma())); + + if (failed(parser.parseGreater())) + return {}; + return getChecked([&] { return parser.emitError(parser.getNameLoc()); }, + parser.getBuilder().getContext(), entries); +} + +void DataLayoutSpecAttr::print(DialectAsmPrinter &os) const { + os << DataLayoutSpecAttr::kAttrKeyword << "<"; + llvm::interleaveComma(getEntries(), os); + os << ">"; +} + +//===----------------------------------------------------------------------===// +// DLTIDialect +//===----------------------------------------------------------------------===// + +constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutAttrName; +constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessKey; +constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessBig; +constexpr const StringLiteral mlir::DLTIDialect::kDataLayoutEndiannessLittle; + +namespace { +class TargetDataLayoutInterface : public DataLayoutDialectInterface { +public: + using DataLayoutDialectInterface::DataLayoutDialectInterface; + + LogicalResult verifyEntry(DataLayoutEntryInterface entry, + Location loc) const final { + StringRef entryName = entry.getKey().get().strref(); + if (entryName == DLTIDialect::kDataLayoutEndiannessKey) { + auto value = entry.getValue().dyn_cast(); + if (value && + (value.getValue() == DLTIDialect::kDataLayoutEndiannessBig || + value.getValue() == DLTIDialect::kDataLayoutEndiannessLittle)) + return success(); + return emitError(loc) << "'" << entryName + << "' data layout entry is expected to be either '" + << DLTIDialect::kDataLayoutEndiannessBig << "' or '" + << DLTIDialect::kDataLayoutEndiannessLittle << "'"; + } + return emitError(loc) << "unknown data layout entry name: " << entryName; + } +}; +} // namespace + +void DLTIDialect::initialize() { + addAttributes(); + addInterfaces(); +} + +Attribute DLTIDialect::parseAttribute(DialectAsmParser &parser, + Type type) const { + StringRef attrKind; + if (parser.parseKeyword(&attrKind)) + return {}; + + if (attrKind == DataLayoutEntryAttr::kAttrKeyword) + return DataLayoutEntryAttr::parse(parser); + if (attrKind == DataLayoutSpecAttr::kAttrKeyword) + return DataLayoutSpecAttr::parse(parser); + + parser.emitError(parser.getNameLoc(), "unknown attrribute type: ") + << attrKind; + return {}; +} + +void DLTIDialect::printAttribute(Attribute attr, DialectAsmPrinter &os) const { + llvm::TypeSwitch(attr) + .Case( + [&](auto a) { a.print(os); }) + .Default([](Attribute) { llvm_unreachable("unknown attribute kind"); }); +} + +LogicalResult DLTIDialect::verifyOperationAttribute(Operation *op, + NamedAttribute attr) { + if (attr.first == DLTIDialect::kDataLayoutAttrName) { + if (!attr.second.isa()) { + return op->emitError() << "'" << DLTIDialect::kDataLayoutAttrName + << "' is expected to be a #dlti.dl_spec attribute"; + } + return success(); + } + + return op->emitError() << "attribute '" << attr.first + << "' not supported by dialect"; +} diff --git a/mlir/lib/Dialect/DLTI/Traits.cpp b/mlir/lib/Dialect/DLTI/Traits.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Dialect/DLTI/Traits.cpp @@ -0,0 +1,29 @@ +//===- Traits.cpp - Traits for MLIR DLTI dialect --------------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Dialect/DLTI/Traits.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" + +using namespace mlir; + +LogicalResult mlir::impl::verifyHasDefaultDLTIDataLayoutTrait(Operation *op) { + // TODO: consider having trait inheritance so that HasDefaultDLTIDataLayout + // trait can inherit DataLayoutOpInterface::Trait and enforce the validity of + // the assertion below. + assert( + isa(op) && + "HasDefaultDLTIDataLayout trait unexpectedly attached to an op that does " + "not implement DataLayoutOpInterface"); + return success(); +} + +DataLayoutSpecInterface mlir::impl::getDataLayoutSpec(Operation *op) { + return op->getAttrOfType( + DLTIDialect::kDataLayoutAttrName); +} diff --git a/mlir/lib/Interfaces/CMakeLists.txt b/mlir/lib/Interfaces/CMakeLists.txt --- a/mlir/lib/Interfaces/CMakeLists.txt +++ b/mlir/lib/Interfaces/CMakeLists.txt @@ -3,6 +3,7 @@ CastInterfaces.cpp ControlFlowInterfaces.cpp CopyOpInterface.cpp + DataLayoutInterfaces.cpp DerivedAttributeOpInterface.cpp InferTypeOpInterface.cpp LoopLikeInterface.cpp @@ -31,6 +32,7 @@ add_mlir_interface_library(CastInterfaces) add_mlir_interface_library(ControlFlowInterfaces) add_mlir_interface_library(CopyOpInterface) +add_mlir_interface_library(DataLayoutInterfaces) add_mlir_interface_library(DerivedAttributeOpInterface) add_mlir_interface_library(InferTypeOpInterface) add_mlir_interface_library(LoopLikeInterface) diff --git a/mlir/lib/Interfaces/DataLayoutInterfaces.cpp b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Interfaces/DataLayoutInterfaces.cpp @@ -0,0 +1,319 @@ +//===- DataLayoutInterfaces.cpp - Data Layout Interface Implementation ----===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Interfaces/DataLayoutInterfaces.h" +#include "mlir/IR/BuiltinDialect.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Operation.h" + +using namespace mlir; + +//===----------------------------------------------------------------------===// +// Default implementations +//===----------------------------------------------------------------------===// + +/// Reports that the given type is missing the data layout information and +/// exits. +static LLVM_ATTRIBUTE_NORETURN void reportMissingDataLayout(Type type) { + std::string message; + llvm::raw_string_ostream os(message); + os << "neither the scoping op nor the type class provide data layout " + "information for " + << type; + llvm::report_fatal_error(os.str()); +} + +unsigned +mlir::detail::getDefaultTypeSize(Type type, const DataLayout &dataLayout, + ArrayRef params) { + if (type.isa()) + return llvm::divideCeil(type.getIntOrFloatBitWidth(), 8); + + // Sizes of vector types are rounded up to those of types with closest + // power-of-two number of elements. + // TODO: make this extensible. + if (auto vecType = type.dyn_cast()) + return llvm::PowerOf2Ceil(vecType.getNumElements()) * + dataLayout.getTypeSize(vecType.getElementType()); + + if (auto typeInterface = type.dyn_cast()) + return typeInterface.getTypeSize(dataLayout, params); + + reportMissingDataLayout(type); +} + +unsigned mlir::detail::getDefaultABIAlignment( + Type type, const DataLayout &dataLayout, + ArrayRef params) { + // Natural alignment is the closest power-of-two number above. + if (type.isa()) + return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type)); + + if (auto intType = type.dyn_cast()) { + return intType.getWidth() < 64 + ? llvm::PowerOf2Ceil(llvm::divideCeil(intType.getWidth(), 8)) + : 4; + } + + if (auto typeInterface = type.dyn_cast()) + return typeInterface.getABIAlignment(dataLayout, params); + + reportMissingDataLayout(type); +} + +unsigned mlir::detail::getDefaultPreferredAlignment( + Type type, const DataLayout &dataLayout, + ArrayRef params) { + // Preferred alignment is same as natural for floats and vectors. + if (type.isa()) + return dataLayout.getTypeABIAlignment(type); + + // Preferred alignment is the cloest power-of-two number above for integers + // (ABI alignment may be smaller). + if (auto intType = type.dyn_cast()) + return llvm::PowerOf2Ceil(dataLayout.getTypeSize(type)); + + if (auto typeInterface = type.dyn_cast()) + return typeInterface.getPreferredAlignment(dataLayout, params); + + reportMissingDataLayout(type); +} + +DataLayoutEntryList +mlir::detail::filterEntriesForType(DataLayoutEntryListRef entries, + TypeID typeID) { + return llvm::to_vector<4>(llvm::make_filter_range( + entries, [typeID](DataLayoutEntryInterface entry) { + auto type = entry.getKey().dyn_cast(); + return type && type.getTypeID() == typeID; + })); +} + +DataLayoutEntryInterface +mlir::detail::filterEntryForIdentifier(DataLayoutEntryListRef entries, + Identifier id) { + const auto *it = llvm::find_if(entries, [id](DataLayoutEntryInterface entry) { + if (!entry.getKey().is()) + return false; + return entry.getKey().get() == id; + }); + return it == entries.end() ? DataLayoutEntryInterface() : *it; +} + +/// Populates `opsWithLayout` with the list of proper ancestors of `leaf` that +/// implement the `DataLayoutOpInterface`. +static void findProperAscendantsWithLayout( + Operation *leaf, SmallVectorImpl &opsWithLayout) { + if (!leaf) + return; + + while (auto opLayout = leaf->getParentOfType()) { + opsWithLayout.push_back(opLayout); + leaf = opLayout; + } +} + +/// Returns a layout spec that is a combination of the layout specs attached +/// to the given operation and all its ancestors. +static DataLayoutSpecInterface +getCombinedDataLayout(DataLayoutOpInterface leaf) { + if (!leaf) + return {}; + + SmallVector opsWithLayout; + findProperAscendantsWithLayout(leaf, opsWithLayout); + + // Fast track if there are no ancestors. + if (opsWithLayout.empty()) + return leaf.getDataLayoutSpec(); + + // Create the list of non-null specs (null/missing specs can be safely + // ignored) from the outermost to the innermost. + SmallVector specs; + specs.reserve(opsWithLayout.size()); + for (DataLayoutOpInterface op : llvm::reverse(opsWithLayout)) + if (DataLayoutSpecInterface current = op.getDataLayoutSpec()) + specs.push_back(current); + + // Combine the specs using the innermost as anchor. + if (DataLayoutSpecInterface current = leaf.getDataLayoutSpec()) + return current.combineWith(specs); + if (specs.empty()) + return {}; + return specs.back().combineWith(llvm::makeArrayRef(specs).drop_back()); +} + +LogicalResult mlir::detail::verifyDataLayoutOp(DataLayoutOpInterface op) { + DataLayoutSpecInterface spec = op.getDataLayoutSpec(); + // The layout specification may be missing and it's fine. + if (!spec) + return success(); + + if (failed(spec.verifySpec(op.getLoc()))) + return failure(); + if (!getCombinedDataLayout(op)) { + InFlightDiagnostic diag = + op.emitError() + << "data layout is not a refinement of the layouts in enclosing ops"; + SmallVector opsWithLayout; + findProperAscendantsWithLayout(op, opsWithLayout); + for (DataLayoutOpInterface parent : opsWithLayout) + diag.attachNote(parent.getLoc()) << "enclosing op with data layout"; + return diag; + } + return success(); +} + +//===----------------------------------------------------------------------===// +// DataLayout +//===----------------------------------------------------------------------===// + +mlir::DataLayout::DataLayout(DataLayoutOpInterface op) + : originalLayout(getCombinedDataLayout(op)), scope(op) { + if (!originalLayout) { + assert((!op || !op.getDataLayoutSpec()) && + "could not compute layout information for an op (failed to " + "combine attributes?)"); + } + +#ifndef NDEBUG + SmallVector opsWithLayout; + findProperAscendantsWithLayout(op, opsWithLayout); + layoutStack = llvm::to_vector<2>( + llvm::map_range(opsWithLayout, [](DataLayoutOpInterface iface) { + return iface.getDataLayoutSpec(); + })); +#endif +} + +void mlir::DataLayout::checkValid() const { +#ifndef NDEBUG + SmallVector opsWithLayout; + findProperAscendantsWithLayout(scope, opsWithLayout); + assert(opsWithLayout.size() == layoutStack.size() && + "data layout object used, but no longer valid due to the change in " + "number of nested layouts"); + for (auto pair : llvm::zip(opsWithLayout, layoutStack)) { + Attribute newLayout = std::get<0>(pair).getDataLayoutSpec(); + Attribute origLayout = std::get<1>(pair); + assert(newLayout == origLayout && + "data layout object used, but no longer valid " + "due to the change in layout attributes"); + } +#endif + assert(((!scope && !this->originalLayout) || + (scope && this->originalLayout == getCombinedDataLayout(scope))) && + "data layout object used, but no longer valid due to the change in " + "layout spec"); +} + +/// Looks up the value for the given type key in the given cache. If there is no +/// such value in the cache, compute it using the given callback and put it in +/// the cache before returning. +static unsigned cachedLookup(Type t, DenseMap &cache, + function_ref compute) { + auto it = cache.find(t); + if (it != cache.end()) + return it->second; + + auto result = cache.try_emplace(t, compute(t)); + return result.first->second; +} + +unsigned mlir::DataLayout::getTypeSize(Type t) const { + checkValid(); + return cachedLookup(t, sizes, [&](Type ty) { + return (scope && originalLayout) + ? scope.getTypeSize( + ty, *this, originalLayout.getSpecForType(ty.getTypeID())) + : detail::getDefaultTypeSize(ty, *this, {}); + }); +} + +unsigned mlir::DataLayout::getTypeABIAlignment(Type t) const { + checkValid(); + return cachedLookup(t, abiAlignments, [&](Type ty) { + return (scope && originalLayout) + ? scope.getTypeABIAlignment( + ty, *this, originalLayout.getSpecForType(ty.getTypeID())) + : detail::getDefaultABIAlignment(ty, *this, {}); + }); +} + +unsigned mlir::DataLayout::getTypePreferredAlignment(Type t) const { + checkValid(); + return cachedLookup(t, preferredAlignments, [&](Type ty) { + return (scope && originalLayout) + ? scope.getTypePreferredAlignment( + ty, *this, originalLayout.getSpecForType(ty.getTypeID())) + : detail::getDefaultPreferredAlignment(ty, *this, {}); + }); +} + +//===----------------------------------------------------------------------===// +// DataLayoutSpecInterface +//===----------------------------------------------------------------------===// + +void DataLayoutSpecInterface::bucketEntriesByType( + DenseMap &types, + DenseMap &ids) { + for (DataLayoutEntryInterface entry : getEntries()) { + if (auto type = entry.getKey().dyn_cast()) + types[type.getTypeID()].push_back(entry); + else + ids[entry.getKey().get()] = entry; + } +} + +LogicalResult mlir::detail::verifyDataLayoutSpec(DataLayoutSpecInterface spec, + Location loc) { + // First, verify individual entries. + for (DataLayoutEntryInterface entry : spec.getEntries()) + if (failed(entry.verifyEntry(loc))) + return failure(); + + // Second, dispatch verifications of entry groups to types or dialects they + // are are associated with. + DenseMap types; + DenseMap ids; + spec.bucketEntriesByType(types, ids); + + for (const auto &kvp : types) { + auto sampleType = kvp.second.front().getKey().get(); + if (isa(&sampleType.getDialect())) + return emitError(loc) << "unexpected data layout for a built-in type"; + + auto dlType = sampleType.dyn_cast(); + if (!dlType) + return emitError(loc) + << "data layout specified for a type that does not support it"; + if (failed(dlType.verifyEntries(kvp.second, loc))) + return failure(); + } + + for (const auto &kvp : ids) { + Identifier identifier = kvp.second.getKey().get(); + Dialect *dialect = identifier.getDialect(); + + // Ignore attributes that belong to an unknown dialect, the dialect may + // actually implement the relevant interface but we don't know about that. + if (!dialect) + continue; + + const auto *iface = + dialect->getRegisteredInterface(); + if (failed(iface->verifyEntry(kvp.second, loc))) + return failure(); + } + + return success(); +} + +#include "mlir/Interfaces/DataLayoutAttrInterface.cpp.inc" +#include "mlir/Interfaces/DataLayoutOpInterface.cpp.inc" +#include "mlir/Interfaces/DataLayoutTypeInterface.cpp.inc" diff --git a/mlir/test/Dialect/DLTI/invalid.mlir b/mlir/test/Dialect/DLTI/invalid.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/DLTI/invalid.mlir @@ -0,0 +1,73 @@ +// RUN: mlir-opt -split-input-file -verify-diagnostics %s + +// expected-error@below {{attribute 'dlti.unknown' not supported by dialect}} +"test.unknown_op"() { dlti.unknown } : () -> () + +// ----- + +// expected-error@below {{'dlti.dl_spec' is expected to be a #dlti.dl_spec attribute}} +"test.unknown_op"() { dlti.dl_spec = 42 } : () -> () + +// ----- + +// expected-error@below {{invalid kind of attribute specified}} +"test.unknown_op"() { dlti.dl_spec = #dlti.dl_spec<[]> } : () -> () + +// ----- + +// expected-error@below {{expected a type or a quoted string}} +"test.unknown_op"() { test.unknown_attr = #dlti.dl_entry<42, 42> } : () -> () + +// ----- + +// expected-error@below {{repeated layout entry key: test.id}} +"test.unknown_op"() { test.unknown_attr = #dlti.dl_spec< + #dlti.dl_entry<"test.id", 42>, + #dlti.dl_entry<"test.id", 43> +>} : () -> () + +// ----- + +// expected-error@below {{repeated layout entry key: 'i32'}} +"test.unknown_op"() { test.unknown_attr = #dlti.dl_spec< + #dlti.dl_entry, + #dlti.dl_entry +>} : () -> () + +// ----- + +// expected-error@below {{unknown attrribute type: unknown}} +"test.unknown_op"() { test.unknown_attr = #dlti.unknown } : () -> () + +// ----- + +// expected-error@below {{unknown data layout entry name: dlti.unknown_id}} +"test.op_with_data_layout"() ({ +}) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.unknown_id", 42>> } : () -> () + +// ----- + +// expected-error@below {{'dlti.endianness' data layout entry is expected to be either 'big' or 'little'}} +"test.op_with_data_layout"() ({ +}) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.endianness", "some">> } : () -> () + +// ----- + +// Mismatching entries don't combine. +"test.op_with_data_layout"() ({ + // expected-error@below {{data layout is not a refinement of the layouts in enclosing ops}} + // expected-note@above {{enclosing op with data layout}} + "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () + "test.maybe_terminator_op"() : () -> () +}) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 33>> } : () -> () + +// ----- + +// Layout not supported for built-in types. +// expected-error@below {{unexpected data layout for a built-in type}} +"test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry> } : () -> () + +// ----- + +// expected-error@below {{data layout specified for a type that does not support it}} +"test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry> } : () -> () diff --git a/mlir/test/Dialect/DLTI/roundtrip.mlir b/mlir/test/Dialect/DLTI/roundtrip.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Dialect/DLTI/roundtrip.mlir @@ -0,0 +1,53 @@ +// RUN: mlir-opt %s | mlir-opt | FileCheck %s + +// Round-tripping the syntax. + +"test.unknown_op"() { + // CHECK: #dlti.dl_entry<"test.identifier", 42 : i64> + test.unknown_attr_1 = #dlti.dl_entry<"test.identifier", 42 : i64>, + // CHECK: #dlti.dl_entry<"test.identifier", affine_map<(d0) -> (d0)>> + test.unknown_attr_2 = #dlti.dl_entry<"test.identifier", affine_map<(d0) -> (d0)>>, + // CHECK: #dlti.dl_entry + test.unknown_attr_3 = #dlti.dl_entry, + // CHECK: #dlti.dl_entry, ["string", 10]> + test.unknown_attr_4 = #dlti.dl_entry, ["string", 10]>, + // CHECK: #dlti.dl_spec<> + test.unknown_attr_5 = #dlti.dl_spec<>, + // CHECK: #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>> + test.unknown_attr_6 = #dlti.dl_spec<#dlti.dl_entry<"test.id", 42 : i32>>, + // CHECK: #dlti.dl_spec< + // CHECK: #dlti.dl_entry<"test.id1", 43 : index> + // CHECK: #dlti.dl_entry<"test.id2", 44 : index> + // CHECK: #dlti.dl_entry<"test.id3", 45 : index>> + test.unknown_attr_7 = #dlti.dl_spec< + #dlti.dl_entry<"test.id1", 43 : index>, + #dlti.dl_entry<"test.id2", 44 : index>, + #dlti.dl_entry<"test.id3", 45 : index>> +} : () -> () + +// +// Supported cases where we shouldn't fail. No need to file-check these, not +// triggering an error or an assertion is enough. +// + +// Should not fail on missing spec. +"test.op_with_data_layout"() : () -> () + +// Should not fail on empty spec. +"test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<> }: () -> () + +// Should not fail on nested compatible layouts. +"test.op_with_data_layout"() ({ + "test.op_with_data_layout"() { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () + "test.maybe_terminator_op"() : () -> () +}) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () + +// Should not fail on deeper nested compatible layouts. +"test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() + { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () + "test.maybe_terminator_op"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () + "test.maybe_terminator_op"() : () -> () +}) { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"unknown.unknown", 32>> } : () -> () diff --git a/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Interfaces/DataLayoutInterfaces/query.mlir @@ -0,0 +1,172 @@ +// RUN: mlir-opt --test-data-layout-query %s | FileCheck %s + +// CHECK-LABEL: @no_layout_builtin +func @no_layout_builtin() { + // CHECK: alignment = 4 + // CHECK: preferred = 4 + // CHECK: size = 4 + "test.data_layout_query"() : () -> i32 + // CHECK: alignment = 8 + // CHECK: preferred = 8 + // CHECK: size = 8 + "test.data_layout_query"() : () -> f64 + return +} + +// CHECK-LABEL: @no_layout_custom +func @no_layout_custom() { + // CHECK: alignment = 1 + // CHECK: preferred = 1 + // CHECK: size = 1 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + return +} + +// CHECK-LABEL: @layout_op_no_layout +func @layout_op_no_layout() { + "test.op_with_data_layout"() ({ + // CHECK: alignment = 1 + // CHECK: preferred = 1 + // CHECK: size = 1 + "test.data_layout_query"() : () -> !test.test_type_with_layout<1000> + "test.maybe_terminator"() : () -> () + }) : () -> () + return +} + +// CHECK-LABEL: @layout_op +func @layout_op() { + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 1 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + return +} + +// Make sure the outer op with layout may be missing the spec. +// CHECK-LABEL: @nested_inner_only +func @nested_inner_only() { + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 1 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + "test.maybe_terminator"() : () -> () + }) : () -> () + return +} + +// Make sure the inner op with layout may be missing the spec. +// CHECK-LABEL: @nested_outer_only +func @nested_outer_only() { + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 1 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) : () -> () + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + return +} + +// CHECK-LABEL: @nested_middle_only +func @nested_middle_only() { + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 1 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) : () -> () + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + "test.maybe_terminator"() : () -> () + }) : () -> () + return +} + +// CHECK-LABEL: @nested_combine_with_missing +func @nested_combine_with_missing() { + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 30 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) : () -> () + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + // CHECK: alignment = 1 + // CHECK: preferred = 30 + // CHECK: size = 42 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 42]>, + #dlti.dl_entry, ["preferred", 30]> + >}: () -> () + return +} + +// CHECK-LABEL: @nested_combine_all +func @nested_combine_all() { + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + "test.op_with_data_layout"() ({ + // CHECK: alignment = 20 + // CHECK: preferred = 30 + // CHECK: size = 3 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 3]>, + #dlti.dl_entry, ["preferred", 30]> + >} : () -> () + // CHECK: alignment = 20 + // CHECK: preferred = 30 + // CHECK: size = 10 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 10]>, + #dlti.dl_entry, ["alignment", 20]> + >} : () -> () + // CHECK: alignment = 1 + // CHECK: preferred = 30 + // CHECK: size = 42 + "test.data_layout_query"() : () -> !test.test_type_with_layout<10> + "test.maybe_terminator"() : () -> () + }) { dlti.dl_spec = #dlti.dl_spec< + #dlti.dl_entry, ["size", 42]>, + #dlti.dl_entry, ["preferred", 30]> + >}: () -> () + return +} diff --git a/mlir/test/lib/Dialect/Test/CMakeLists.txt b/mlir/test/lib/Dialect/Test/CMakeLists.txt --- a/mlir/test/lib/Dialect/Test/CMakeLists.txt +++ b/mlir/test/lib/Dialect/Test/CMakeLists.txt @@ -25,7 +25,7 @@ set(LLVM_TARGET_DEFINITIONS TestOps.td) mlir_tablegen(TestOps.h.inc -gen-op-decls) mlir_tablegen(TestOps.cpp.inc -gen-op-defs) -mlir_tablegen(TestOpsDialect.h.inc -gen-dialect-decls) +mlir_tablegen(TestOpsDialect.h.inc -gen-dialect-decls -dialect=test) mlir_tablegen(TestOpEnums.h.inc -gen-enum-decls) mlir_tablegen(TestOpEnums.cpp.inc -gen-enum-defs) mlir_tablegen(TestOpStructs.h.inc -gen-struct-attr-decls) @@ -52,8 +52,10 @@ LINK_LIBS PUBLIC MLIRControlFlowInterfaces + MLIRDataLayoutInterfaces MLIRDerivedAttributeOpInterface MLIRDialect + MLIRDLTI MLIRIR MLIRInferTypeOpInterface MLIRLinalgTransforms diff --git a/mlir/test/lib/Dialect/Test/TestDialect.h b/mlir/test/lib/Dialect/Test/TestDialect.h --- a/mlir/test/lib/Dialect/Test/TestDialect.h +++ b/mlir/test/lib/Dialect/Test/TestDialect.h @@ -15,6 +15,8 @@ #define MLIR_TESTDIALECT_H #include "TestInterfaces.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Dialect/DLTI/Traits.h" #include "mlir/Dialect/Traits.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypes.h" @@ -30,6 +32,10 @@ #include "mlir/Interfaces/InferTypeOpInterface.h" #include "mlir/Interfaces/SideEffectInterfaces.h" +namespace mlir { +class DLTIDialect; +} // namespace mlir + #include "TestOpEnums.h.inc" #include "TestOpInterfaces.h.inc" #include "TestOpStructs.h.inc" diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp --- a/mlir/test/lib/Dialect/Test/TestDialect.cpp +++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp @@ -9,6 +9,7 @@ #include "TestDialect.h" #include "TestAttributes.h" #include "TestTypes.h" +#include "mlir/Dialect/DLTI/DLTI.h" #include "mlir/Dialect/StandardOps/IR/Ops.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/DialectImplementation.h" @@ -175,7 +176,7 @@ >(); addInterfaces(); - addTypes(); diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td --- a/mlir/test/lib/Dialect/Test/TestOps.td +++ b/mlir/test/lib/Dialect/Test/TestOps.td @@ -9,6 +9,7 @@ #ifndef TEST_OPS #define TEST_OPS +include "mlir/Dialect/DLTI/DLTIBase.td" include "mlir/IR/OpBase.td" include "mlir/IR/OpAsmInterface.td" include "mlir/IR/RegionKindInterface.td" @@ -16,6 +17,7 @@ include "mlir/Interfaces/CallInterfaces.td" include "mlir/Interfaces/ControlFlowInterfaces.td" include "mlir/Interfaces/CopyOpInterface.td" +include "mlir/Interfaces/DataLayoutInterfaces.td" include "mlir/Interfaces/InferTypeOpInterface.td" include "mlir/Interfaces/SideEffectInterfaces.td" include "TestInterfaces.td" @@ -27,6 +29,7 @@ let hasOperationAttrVerify = 1; let hasRegionArgAttrVerify = 1; let hasRegionResultAttrVerify = 1; + let dependentDialects = ["::mlir::DLTIDialect"]; let extraClassDeclaration = [{ Attribute parseAttribute(DialectAsmParser &parser, @@ -2008,4 +2011,26 @@ let results = (outs TupleOf<[AnyType]>); } +//===----------------------------------------------------------------------===// +// Test Target DataLayout +//===----------------------------------------------------------------------===// + +def OpWithDataLayoutOp : TEST_Op<"op_with_data_layout", + [HasDefaultDLTIDataLayout, DataLayoutOpInterface]> { + let summary = + "An op that uses DataLayout implementation from the Target dialect"; + let regions = (region VariadicRegion:$regions); +} + +def DataLayoutQueryOp : TEST_Op<"data_layout_query"> { + let summary = "A token op recognized by data layout query test pass"; + let description = [{ + The data layout query pass pattern-matches this op and attaches to it an + array attribute containing the result of data layout query of the result + type of this op. + }]; + + let results = (outs AnyType:$res); +} + #endif // TEST_OPS diff --git a/mlir/test/lib/Dialect/Test/TestTypes.h b/mlir/test/lib/Dialect/Test/TestTypes.h --- a/mlir/test/lib/Dialect/Test/TestTypes.h +++ b/mlir/test/lib/Dialect/Test/TestTypes.h @@ -21,6 +21,7 @@ #include "mlir/IR/DialectImplementation.h" #include "mlir/IR/Operation.h" #include "mlir/IR/Types.h" +#include "mlir/Interfaces/DataLayoutInterfaces.h" namespace mlir { namespace test { @@ -107,6 +108,62 @@ StringRef getName() { return getImpl()->name; } }; +struct TestTypeWithLayoutStorage : public TypeStorage { + using KeyTy = unsigned; + + explicit TestTypeWithLayoutStorage(unsigned key) : key(key) {} + bool operator==(const KeyTy &other) const { return other == key; } + + static TestTypeWithLayoutStorage *construct(TypeStorageAllocator &allocator, + const KeyTy &key) { + return new (allocator.allocate()) + TestTypeWithLayoutStorage(key); + } + + unsigned key; +}; + +class TestTypeWithLayout + : public Type::TypeBase { +public: + using Base::Base; + + static TestTypeWithLayout get(MLIRContext *ctx, unsigned key) { + return Base::get(ctx, key); + } + + unsigned getKey() { return getImpl()->key; } + + unsigned getTypeSize(const DataLayout &dataLayout, + DataLayoutEntryListRef params) const { + return extractKind(params, "size"); + } + + unsigned getABIAlignment(const DataLayout &dataLayout, + DataLayoutEntryListRef params) const { + return extractKind(params, "alignment"); + } + + unsigned getPreferredAlignment(const DataLayout &dataLayout, + DataLayoutEntryListRef params) const { + return extractKind(params, "preferred"); + } + + bool areCompatible(DataLayoutEntryListRef oldLayout, + DataLayoutEntryListRef newLayout) const { + unsigned old = extractKind(oldLayout, "alignment"); + return old == 1 || extractKind(newLayout, "alignment") <= old; + } + + LogicalResult verifyEntries(DataLayoutEntryListRef params, + Location loc) const; + +private: + unsigned extractKind(DataLayoutEntryListRef params, + StringRef expectedKind) const; +}; + } // namespace test } // namespace mlir diff --git a/mlir/test/lib/Dialect/Test/TestTypes.cpp b/mlir/test/lib/Dialect/Test/TestTypes.cpp --- a/mlir/test/lib/Dialect/Test/TestTypes.cpp +++ b/mlir/test/lib/Dialect/Test/TestTypes.cpp @@ -128,6 +128,38 @@ #define GET_TYPEDEF_CLASSES #include "TestTypeDefs.cpp.inc" +LogicalResult TestTypeWithLayout::verifyEntries(DataLayoutEntryListRef params, + Location loc) const { + for (DataLayoutEntryInterface entry : params) { + // This is for testing purposes only, so assert well-formedness. + assert(entry.isTypeEntry() && "unexpected identifier entry"); + assert(entry.getKey().get().isa() && + "wrong type passed in"); + auto array = entry.getValue().dyn_cast(); + assert(array && array.getValue().size() == 2 && + "expected array of two elements"); + auto kind = array.getValue().front().dyn_cast(); + (void)kind; + assert(kind && + (kind.getValue() == "size" || kind.getValue() == "alignment" || + kind.getValue() == "preferred") && + "unexpected kind"); + assert(array.getValue().back().isa()); + } + return success(); +} + +unsigned TestTypeWithLayout::extractKind(DataLayoutEntryListRef params, + StringRef expectedKind) const { + for (DataLayoutEntryInterface entry : params) { + ArrayRef pair = entry.getValue().cast().getValue(); + StringRef kind = pair.front().cast().getValue(); + if (kind == expectedKind) + return pair.back().cast().getValue().getZExtValue(); + } + return 1; +} + //===----------------------------------------------------------------------===// // TestDialect //===----------------------------------------------------------------------===// @@ -147,8 +179,19 @@ if (typeTag == "test_type") return TestType::get(parser.getBuilder().getContext()); - if (typeTag != "test_rec") + if (typeTag == "test_type_with_layout") { + unsigned val; + if (parser.parseLess() || parser.parseInteger(val) || + parser.parseGreater()) { + return Type(); + } + return TestTypeWithLayout::get(parser.getBuilder().getContext(), val); + } + + if (typeTag != "test_rec") { + parser.emitError(parser.getNameLoc()) << "unknown type!"; return Type(); + } StringRef name; if (parser.parseLess() || parser.parseKeyword(&name)) @@ -189,6 +232,11 @@ return; } + if (auto t = type.dyn_cast()) { + printer << "test_type_with_layout<" << t.getKey() << ">"; + return; + } + auto rec = type.cast(); printer << "test_rec<" << rec.getName(); if (!stack.contains(rec)) { diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt --- a/mlir/test/lib/Transforms/CMakeLists.txt +++ b/mlir/test/lib/Transforms/CMakeLists.txt @@ -9,6 +9,7 @@ TestConvertCallOp.cpp TestConvertGPUKernelToCubin.cpp TestConvertGPUKernelToHsaco.cpp + TestDataLayoutQuery.cpp TestDominance.cpp TestDynamicPipeline.cpp TestLoopFusion.cpp diff --git a/mlir/test/lib/Transforms/TestDataLayoutQuery.cpp b/mlir/test/lib/Transforms/TestDataLayoutQuery.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/Transforms/TestDataLayoutQuery.cpp @@ -0,0 +1,60 @@ +//===- TestDataLayoutQuery.cpp - Test Data Layout Queries -----------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "TestDialect.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/Pass/Pass.h" + +using namespace mlir; + +namespace { + +/// A pass that finds "test.data_layout_query" operations and attaches to them +/// attributes containing the results of data layout queries for operation +/// result types. +struct TestDataLayoutQuery + : public PassWrapper { + void runOnFunction() override { + FuncOp func = getFunction(); + Builder builder(func.getContext()); + DenseMap layouts; + + func.walk([&](test::DataLayoutQueryOp op) { + // Skip the ops with already processed in a deeper call. + if (op->getAttr("size")) + return; + + auto scope = op->getParentOfType(); + if (!layouts.count(scope)) { + layouts.try_emplace( + scope, scope ? cast(scope.getOperation()) + : nullptr); + } + + const DataLayout &layout = layouts.find(scope)->getSecond(); + unsigned size = layout.getTypeSize(op.getType()); + unsigned alignment = layout.getTypeABIAlignment(op.getType()); + unsigned preferred = layout.getTypePreferredAlignment(op.getType()); + op->setAttrs( + {builder.getNamedAttr("size", builder.getIndexAttr(size)), + builder.getNamedAttr("alignment", builder.getIndexAttr(alignment)), + builder.getNamedAttr("preferred", builder.getIndexAttr(preferred))}); + }); + } +}; +} // namespace + +namespace mlir { +namespace test { +void registerTestDataLayoutQuery() { + PassRegistration("test-data-layout-query", + "Test data layout queries"); +} +} // namespace test +} // namespace mlir diff --git a/mlir/test/mlir-opt/commandline.mlir b/mlir/test/mlir-opt/commandline.mlir --- a/mlir/test/mlir-opt/commandline.mlir +++ b/mlir/test/mlir-opt/commandline.mlir @@ -7,6 +7,7 @@ // CHECK-NEXT: async // CHECK-NEXT: avx512 // CHECK-NEXT: complex +// CHECK-NEXT: dlti // CHECK-NEXT: gpu // CHECK-NEXT: linalg // CHECK-NEXT: llvm diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -66,6 +66,7 @@ void registerTestConvVectorization(); void registerTestGpuSerializeToCubinPass(); void registerTestConvertGPUKernelToHsacoPass(); +void registerTestDataLayoutQuery(); void registerTestDecomposeCallGraphTypes(); void registerTestDialect(DialectRegistry &); void registerTestDominancePass(); @@ -143,6 +144,7 @@ #endif test::registerTestConvVectorization(); test::registerTestDecomposeCallGraphTypes(); + test::registerTestDataLayoutQuery(); test::registerTestDominancePass(); test::registerTestDynamicPipelinePass(); test::registerTestExpandTanhPass(); diff --git a/mlir/unittests/CMakeLists.txt b/mlir/unittests/CMakeLists.txt --- a/mlir/unittests/CMakeLists.txt +++ b/mlir/unittests/CMakeLists.txt @@ -7,6 +7,7 @@ add_subdirectory(Analysis) add_subdirectory(Dialect) add_subdirectory(ExecutionEngine) +add_subdirectory(Interfaces) add_subdirectory(IR) add_subdirectory(Pass) add_subdirectory(SDBM) diff --git a/mlir/unittests/Interfaces/CMakeLists.txt b/mlir/unittests/Interfaces/CMakeLists.txt new file mode 100644 --- /dev/null +++ b/mlir/unittests/Interfaces/CMakeLists.txt @@ -0,0 +1,11 @@ +add_mlir_unittest(MLIRInterfacesTests + DataLayoutInterfacesTest.cpp +) + +target_link_libraries(MLIRInterfacesTests + PRIVATE + MLIRDataLayoutInterfaces + MLIRDLTI + MLIRParser + LLVMSupport +) diff --git a/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp new file mode 100644 --- /dev/null +++ b/mlir/unittests/Interfaces/DataLayoutInterfacesTest.cpp @@ -0,0 +1,357 @@ +//===- DataLayoutInterfacesTest.cpp - Unit Tests for Data Layouts ---------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "mlir/Interfaces/DataLayoutInterfaces.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/Parser.h" + +#include + +using namespace mlir; + +namespace { +constexpr static llvm::StringLiteral kAttrName = "dltest.layout"; + +/// Trivial array storage for the custom data layout spec attribute, just a list +/// of entries. +class DataLayoutSpecStorage : public AttributeStorage { +public: + using KeyTy = ArrayRef; + + DataLayoutSpecStorage(ArrayRef entries) + : entries(entries) {} + + bool operator==(const KeyTy &key) const { return key == entries; } + + static DataLayoutSpecStorage *construct(AttributeStorageAllocator &allocator, + const KeyTy &key) { + return new (allocator.allocate()) + DataLayoutSpecStorage(allocator.copyInto(key)); + } + + ArrayRef entries; +}; + +/// Simple data layout spec containing a list of entries that always verifies +/// as valid. +struct CustomDataLayoutSpec + : public Attribute::AttrBase { + using Base::Base; + static CustomDataLayoutSpec get(MLIRContext *ctx, + ArrayRef entries) { + return Base::get(ctx, entries); + } + CustomDataLayoutSpec + combineWith(ArrayRef specs) const { + return *this; + } + DataLayoutEntryListRef getEntries() const { return getImpl()->entries; } + LogicalResult verifySpec(Location loc) { return success(); } +}; + +/// A type subject to data layout that exits the program if it is queried more +/// than once. Handy to check if the cache works. +struct SingleQueryType + : public Type::TypeBase { + using Base::Base; + + static SingleQueryType get(MLIRContext *ctx) { return Base::get(ctx); } + + unsigned getTypeSize(const DataLayout &layout, + DataLayoutEntryListRef params) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return 1; + } + + unsigned getABIAlignment(const DataLayout &layout, + DataLayoutEntryListRef params) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return 2; + } + + unsigned getPreferredAlignment(const DataLayout &layout, + DataLayoutEntryListRef params) { + static bool executed = false; + if (executed) + llvm::report_fatal_error("repeated call"); + + executed = true; + return 4; + } +}; + +/// A types that is not subject to data layout. +struct TypeNoLayout : public Type::TypeBase { + using Base::Base; + + static TypeNoLayout get(MLIRContext *ctx) { return Base::get(ctx); } +}; + +/// An op that serves as scope for data layout queries with the relevant +/// attribute attached. This can handle data layout requests for the built-in +/// types itself. +struct OpWithLayout : public Op { + using Op::Op; + + static StringRef getOperationName() { return "dltest.op_with_layout"; } + + DataLayoutSpecInterface getDataLayoutSpec() { + return getOperation()->getAttrOfType(kAttrName); + } + + static unsigned getTypeSize(Type type, const DataLayout &dataLayout, + DataLayoutEntryListRef params) { + // Make a recursive query. + if (type.isa()) + return dataLayout.getTypeSize( + IntegerType::get(type.getContext(), type.getIntOrFloatBitWidth())); + + // Handle built-in types that are not handled by the default process. + if (auto iType = type.dyn_cast()) { + for (DataLayoutEntryInterface entry : params) + if (entry.getKey().dyn_cast() == type) + return entry.getValue().cast().getValue().getZExtValue(); + return iType.getIntOrFloatBitWidth(); + } + + // Use the default process for everything else. + return detail::getDefaultTypeSize(type, dataLayout, params); + } + + static unsigned getTypeABIAlignment(Type type, const DataLayout &dataLayout, + DataLayoutEntryListRef params) { + return llvm::PowerOf2Ceil(getTypeSize(type, dataLayout, params)); + } + + static unsigned getTypePreferredAlignment(Type type, + const DataLayout &dataLayout, + DataLayoutEntryListRef params) { + return 2 * getTypeABIAlignment(type, dataLayout, params); + } +}; + +/// A dialect putting all the above together. +struct DLTestDialect : Dialect { + explicit DLTestDialect(MLIRContext *ctx) + : Dialect(getDialectNamespace(), ctx, TypeID::get()) { + ctx->getOrLoadDialect(); + addAttributes(); + addOperations(); + addTypes(); + } + static StringRef getDialectNamespace() { return "dltest"; } + + void printAttribute(Attribute attr, + DialectAsmPrinter &printer) const override { + printer << "spec<"; + llvm::interleaveComma(attr.cast().getEntries(), + printer); + printer << ">"; + } + + Attribute parseAttribute(DialectAsmParser &parser, Type type) const override { + bool ok = + succeeded(parser.parseKeyword("spec")) && succeeded(parser.parseLess()); + (void)ok; + assert(ok); + if (succeeded(parser.parseOptionalGreater())) + return CustomDataLayoutSpec::get(parser.getBuilder().getContext(), {}); + + SmallVector entries; + do { + entries.emplace_back(); + ok = succeeded(parser.parseAttribute(entries.back())); + assert(ok); + } while (succeeded(parser.parseOptionalComma())); + ok = succeeded(parser.parseGreater()); + assert(ok); + return CustomDataLayoutSpec::get(parser.getBuilder().getContext(), entries); + } + + void printType(Type type, DialectAsmPrinter &printer) const override { + if (type.isa()) + printer << "single_query"; + else + printer << "no_layout"; + } + + Type parseType(DialectAsmParser &parser) const override { + bool ok = succeeded(parser.parseKeyword("single_query")); + (void)ok; + assert(ok); + return SingleQueryType::get(parser.getBuilder().getContext()); + } +}; + +} // end namespace + +TEST(DataLayout, FallbackDefault) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 6u); + EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 2u); + EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u); + EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 2u); + EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 8u); + EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 2u); +} + +TEST(DataLayout, EmptySpec) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() { dltest.layout = #dltest.spec< > } : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 42u); + EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 16u); + EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 64u); + EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 16u); + EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 128u); + EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 32u); +} + +TEST(DataLayout, SpecWithEntries) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() { dltest.layout = #dltest.spec< + #dlti.dl_entry, + #dlti.dl_entry +> } : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 42)), 5u); + EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u); + EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 42)), 8u); + EXPECT_EQ(layout.getTypeABIAlignment(Float16Type::get(&ctx)), 8u); + EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 42)), 16u); + EXPECT_EQ(layout.getTypePreferredAlignment(Float16Type::get(&ctx)), 16u); + + EXPECT_EQ(layout.getTypeSize(IntegerType::get(&ctx, 32)), 32u); + EXPECT_EQ(layout.getTypeSize(Float32Type::get(&ctx)), 32u); + EXPECT_EQ(layout.getTypeABIAlignment(IntegerType::get(&ctx, 32)), 32u); + EXPECT_EQ(layout.getTypeABIAlignment(Float32Type::get(&ctx)), 32u); + EXPECT_EQ(layout.getTypePreferredAlignment(IntegerType::get(&ctx, 32)), 64u); + EXPECT_EQ(layout.getTypePreferredAlignment(Float32Type::get(&ctx)), 64u); +} + +TEST(DataLayout, Caching) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + + unsigned sum = 0; + sum += layout.getTypeSize(SingleQueryType::get(&ctx)); + // The second call should hit the cache. If it does not, the function in + // SingleQueryType will be called and will abort the process. + sum += layout.getTypeSize(SingleQueryType::get(&ctx)); + // Make sure the complier doesn't optimize away the query code. + EXPECT_EQ(sum, 2u); + + // A fresh data layout has a new cache, so the call to it should be dispatched + // down to the type and abort the proces. + DataLayout second(op); + ASSERT_DEATH(second.getTypeSize(SingleQueryType::get(&ctx)), "repeated call"); +} + +TEST(DataLayout, CacheInvalidation) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() { dltest.layout = #dltest.spec< + #dlti.dl_entry, + #dlti.dl_entry +> } : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + + // Normal query is fine. + EXPECT_EQ(layout.getTypeSize(Float16Type::get(&ctx)), 6u); + + // Replace the data layout spec with a new, empty spec. + op->setAttr(kAttrName, CustomDataLayoutSpec::get(&ctx, {})); + + // Data layout is no longer valid and should trigger assertion when queried. +#ifndef NDEBUG + ASSERT_DEATH(layout.getTypeSize(Float16Type::get(&ctx)), "no longer valid"); +#endif +} + +TEST(DataLayout, UnimplementedTypeInterface) { + const char *ir = R"MLIR( +"dltest.op_with_layout"() { dltest.layout = #dltest.spec<> } : () -> () + )MLIR"; + + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(registry); + + OwningModuleRef module = parseSourceString(ir, &ctx); + auto op = + cast(module->getBody()->getOperations().front()); + DataLayout layout(op); + + ASSERT_DEATH(layout.getTypeSize(TypeNoLayout::get(&ctx)), + "neither the scoping op nor the type class provide data layout " + "information"); +}