1# Defining Dialect Attributes and Types
2
3This document is a quickstart to defining dialect specific extensions to the
4[attribute](LangRef.md#attributes) and [type](LangRef.md#type-system) systems in
5MLIR. The main part of this tutorial focuses on defining types, but the
6instructions are nearly identical for defining attributes.
7
8See [MLIR specification](LangRef.md) for more information about MLIR, the
9structure of the IR, operations, etc.
10
11## Types
12
13Types in MLIR (like attributes, locations, and many other things) are
14value-typed. This means that instances of `Type` are passed around by-value, as
15opposed to by-pointer or by-reference. The `Type` class in itself acts as a
16wrapper around an internal storage object that is uniqued within an instance of
17an `MLIRContext`.
18
19### Defining the type class
20
21As described above, `Type` objects in MLIR are value-typed and rely on having an
22implicit internal storage object that holds the actual data for the type. When
23defining a new `Type` it isn't always necessary to define a new storage class.
24So before defining the derived `Type`, it's important to know which of the two
25classes of `Type` we are defining:
26
27Some types are _singleton_ in nature, meaning they have no parameters and only
28ever have one instance, like the [`index` type](LangRef.md#index-type).
29
30Other types are _parametric_, and contain additional information that
31differentiates different instances of the same `Type`. For example the
32[`integer` type](LangRef.md#integer-type) contains a bitwidth, with `i8` and
33`i16` representing different instances of
34[`integer` type](LangRef.md#integer-type). _Parametric_ may also contain a
35mutable component, which can be used, for example, to construct self-referring
36recursive types. The mutable component _cannot_ be used to differentiate
37instances of a type class, so usually such types contain other parametric
38components that serve to identify them.
39
40#### Singleton types
41
42For singleton types, we can jump straight into defining the derived type class.
43Given that only one instance of such types may exist, there is no need to
44provide our own storage class.
45
46```c++
47/// This class defines a simple parameterless singleton type. All derived types
48/// must inherit from the CRTP class 'Type::TypeBase'. It takes as template
49/// parameters the concrete type (SimpleType), the base class to use (Type),
50/// the internal storage class (the default TypeStorage here), and an optional
51/// set of type traits and interfaces(detailed below).
52class SimpleType : public Type::TypeBase<SimpleType, Type, TypeStorage> {
53public:
54  /// Inherit some necessary constructors from 'TypeBase'.
55  using Base::Base;
56
57  /// The `TypeBase` class provides the following utility methods for
58  /// constructing instances of this type:
59  /// static SimpleType get(MLIRContext *ctx);
60};
61```
62
63#### Parametric types
64
65Parametric types are those with additional construction or uniquing constraints,
66that allow for representing multiple different instances of a single class. As
67such, these types require defining a type storage class to contain the
68parametric data.
69
70##### Defining a type storage
71
72Type storage objects contain all of the data necessary to construct and unique a
73parametric type instance. The storage classes must obey the following:
74
75*   Inherit from the base type storage class `TypeStorage`.
76*   Define a type alias, `KeyTy`, that maps to a type that uniquely identifies
77    an instance of the derived type.
78*   Provide a construction method that is used to allocate a new instance of the
79    storage class.
80    -   `static Storage *construct(TypeStorageAllocator &, const KeyTy &key)`
81*   Provide a comparison method between the storage and `KeyTy`.
82    -   `bool operator==(const KeyTy &) const`
83*   Provide a method to generate the `KeyTy` from a list of arguments passed to
84    the uniquer. (Note: This is only necessary if the `KeyTy` cannot be default
85    constructed from these arguments).
86    -   `static KeyTy getKey(Args...&& args)`
87*   Provide a method to hash an instance of the `KeyTy`. (Note: This is not
88    necessary if an `llvm::DenseMapInfo<KeyTy>` specialization exists)
89    -   `static llvm::hash_code hashKey(const KeyTy &)`
90
91Let's look at an example:
92
93```c++
94/// Here we define a storage class for a ComplexType, that holds a non-zero
95/// integer and an integer type.
96struct ComplexTypeStorage : public TypeStorage {
97  ComplexTypeStorage(unsigned nonZeroParam, Type integerType)
98      : nonZeroParam(nonZeroParam), integerType(integerType) {}
99
100  /// The hash key for this storage is a pair of the integer and type params.
101  using KeyTy = std::pair<unsigned, Type>;
102
103  /// Define the comparison function for the key type.
104  bool operator==(const KeyTy &key) const {
105    return key == KeyTy(nonZeroParam, integerType);
106  }
107
108  /// Define a hash function for the key type.
109  /// Note: This isn't necessary because std::pair, unsigned, and Type all have
110  /// hash functions already available.
111  static llvm::hash_code hashKey(const KeyTy &key) {
112    return llvm::hash_combine(key.first, key.second);
113  }
114
115  /// Define a construction function for the key type.
116  /// Note: This isn't necessary because KeyTy can be directly constructed with
117  /// the given parameters.
118  static KeyTy getKey(unsigned nonZeroParam, Type integerType) {
119    return KeyTy(nonZeroParam, integerType);
120  }
121
122  /// Define a construction method for creating a new instance of this storage.
123  static ComplexTypeStorage *construct(TypeStorageAllocator &allocator,
124                                       const KeyTy &key) {
125    return new (allocator.allocate<ComplexTypeStorage>())
126        ComplexTypeStorage(key.first, key.second);
127  }
128
129  /// The parametric data held by the storage class.
130  unsigned nonZeroParam;
131  Type integerType;
132};
133```
134
135##### Type class definition
136
137Now that the storage class has been created, the derived type class can be
138defined. This structure is similar to [singleton types](#singleton-types),
139except that a bit more of the functionality provided by `Type::TypeBase` is put
140to use.
141
142```c++
143/// This class defines a parametric type. All derived types must inherit from
144/// the CRTP class 'Type::TypeBase'. It takes as template parameters the
145/// concrete type (ComplexType), the base class to use (Type), the storage
146/// class (ComplexTypeStorage), and an optional set of traits and
147/// interfaces(detailed below).
148class ComplexType : public Type::TypeBase<ComplexType, Type,
149                                          ComplexTypeStorage> {
150public:
151  /// Inherit some necessary constructors from 'TypeBase'.
152  using Base::Base;
153
154  /// This method is used to get an instance of the 'ComplexType'. This method
155  /// asserts that all of the construction invariants were satisfied. To
156  /// gracefully handle failed construction, getChecked should be used instead.
157  static ComplexType get(unsigned param, Type type) {
158    // Call into a helper 'get' method in 'TypeBase' to get a uniqued instance
159    // of this type. All parameters to the storage class are passed after the
160    // context.
161    return Base::get(type.getContext(), param, type);
162  }
163
164  /// This method is used to get an instance of the 'ComplexType', defined at
165  /// the given location. If any of the construction invariants are invalid,
166  /// errors are emitted with the provided location and a null type is returned.
167  /// Note: This method is completely optional.
168  static ComplexType getChecked(unsigned param, Type type, Location location) {
169    // Call into a helper 'getChecked' method in 'TypeBase' to get a uniqued
170    // instance of this type. All parameters to the storage class are passed
171    // after the location.
172    return Base::getChecked(location, param, type);
173  }
174
175  /// This method is used to verify the construction invariants passed into the
176  /// 'get' and 'getChecked' methods. Note: This method is completely optional.
177  static LogicalResult verifyConstructionInvariants(
178      Location loc, unsigned param, Type type) {
179    // Our type only allows non-zero parameters.
180    if (param == 0)
181      return emitError(loc) << "non-zero parameter passed to 'ComplexType'";
182    // Our type also expects an integer type.
183    if (!type.isa<IntegerType>())
184      return emitError(loc) << "non integer-type passed to 'ComplexType'";
185    return success();
186  }
187
188  /// Return the parameter value.
189  unsigned getParameter() {
190    // 'getImpl' returns a pointer to our internal storage instance.
191    return getImpl()->nonZeroParam;
192  }
193
194  /// Return the integer parameter type.
195  IntegerType getParameterType() {
196    // 'getImpl' returns a pointer to our internal storage instance.
197    return getImpl()->integerType;
198  }
199};
200```
201
202#### Mutable types
203
204Types with a mutable component are special instances of parametric types that
205allow for mutating certain parameters after construction.
206
207##### Defining a type storage
208
209In addition to the requirements for the type storage class for parametric types,
210the storage class for types with a mutable component must additionally obey the
211following.
212
213*   The mutable component must not participate in the storage `KeyTy`.
214*   Provide a mutation method that is used to modify an existing instance of the
215    storage. This method modifies the mutable component based on arguments,
216    using `allocator` for any newly dynamically-allocated storage, and indicates
217    whether the modification was successful.
218    -   `LogicalResult mutate(StorageAllocator &allocator, Args ...&& args)`
219
220Let's define a simple storage for recursive types, where a type is identified by
221its name and may contain another type including itself.
222
223```c++
224/// Here we define a storage class for a RecursiveType that is identified by its
225/// name and contains another type.
226struct RecursiveTypeStorage : public TypeStorage {
227  /// The type is uniquely identified by its name. Note that the contained type
228  /// is _not_ a part of the key.
229  using KeyTy = StringRef;
230
231  /// Construct the storage from the type name. Explicitly initialize the
232  /// containedType to nullptr, which is used as marker for the mutable
233  /// component being not yet initialized.
234  RecursiveTypeStorage(StringRef name) : name(name), containedType(nullptr) {}
235
236  /// Define the comparison function.
237  bool operator==(const KeyTy &key) const { return key == name; }
238
239  /// Define a construction method for creating a new instance of the storage.
240  static RecursiveTypeStorage *construct(StorageAllocator &allocator,
241                                         const KeyTy &key) {
242    // Note that the key string is copied into the allocator to ensure it
243    // remains live as long as the storage itself.
244    return new (allocator.allocate<RecursiveTypeStorage>())
245        RecursiveTypeStorage(allocator.copyInto(key));
246  }
247
248  /// Define a mutation method for changing the type after it is created. In
249  /// many cases, we only want to set the mutable component once and reject
250  /// any further modification, which can be achieved by returning failure from
251  /// this function.
252  LogicalResult mutate(StorageAllocator &, Type body) {
253    // If the contained type has been initialized already, and the call tries
254    // to change it, reject the change.
255    if (containedType && containedType != body)
256      return failure();
257
258    // Change the body successfully.
259    containedType = body;
260    return success();
261  }
262
263  StringRef name;
264  Type containedType;
265};
266```
267
268##### Type class definition
269
270Having defined the storage class, we can define the type class itself.
271`Type::TypeBase` provides a `mutate` method that forwards its arguments to the
272`mutate` method of the storage and ensures the mutation happens safely.
273
274```c++
275class RecursiveType : public Type::TypeBase<RecursiveType, Type,
276                                            RecursiveTypeStorage> {
277public:
278  /// Inherit parent constructors.
279  using Base::Base;
280
281  /// Creates an instance of the Recursive type. This only takes the type name
282  /// and returns the type with uninitialized body.
283  static RecursiveType get(MLIRContext *ctx, StringRef name) {
284    // Call into the base to get a uniqued instance of this type. The parameter
285    // (name) is passed after the context.
286    return Base::get(ctx, name);
287  }
288
289  /// Now we can change the mutable component of the type. This is an instance
290  /// method callable on an already existing RecursiveType.
291  void setBody(Type body) {
292    // Call into the base to mutate the type.
293    LogicalResult result = Base::mutate(body);
294
295    // Most types expect the mutation to always succeed, but types can implement
296    // custom logic for handling mutation failures.
297    assert(succeeded(result) &&
298           "attempting to change the body of an already-initialized type");
299
300    // Avoid unused-variable warning when building without assertions.
301    (void) result;
302  }
303
304  /// Returns the contained type, which may be null if it has not been
305  /// initialized yet.
306  Type getBody() {
307    return getImpl()->containedType;
308  }
309
310  /// Returns the name.
311  StringRef getName() {
312    return getImpl()->name;
313  }
314};
315```
316
317### Registering types with a Dialect
318
319Once the dialect types have been defined, they must then be registered with a
320`Dialect`. This is done via a similar mechanism to
321[operations](LangRef.md#operations), with the `addTypes` method.
322
323```c++
324struct MyDialect : public Dialect {
325  MyDialect(MLIRContext *context) : Dialect(/*name=*/"mydialect", context) {
326    /// Add these defined types to the dialect.
327    addTypes<SimpleType, ComplexType, RecursiveType>();
328  }
329};
330```
331
332### Parsing and Printing
333
334As a final step after registration, a dialect must override the `printType` and
335`parseType` hooks. These enable native support for round-tripping the type in
336the textual `.mlir`.
337
338```c++
339class MyDialect : public Dialect {
340public:
341  /// Parse an instance of a type registered to the dialect.
342  Type parseType(DialectAsmParser &parser) const override;
343
344  /// Print an instance of a type registered to the dialect.
345  void printType(Type type, DialectAsmPrinter &printer) const override;
346};
347```
348
349These methods take an instance of a high-level parser or printer that allows for
350easily implementing the necessary functionality. As described in the
351[MLIR language reference](../../LangRef.md#dialect-types), dialect types are
352generally represented as: `! dialect-namespace < type-data >`, with a pretty
353form available under certain circumstances. The responsibility of our parser and
354printer is to provide the `type-data` bits.
355
356### Traits
357
358Similarly to operations, `Type` classes may attach `Traits` that provide
359additional mixin methods and other data. `Trait` classes may be specified via
360the trailing template argument of the `Type::TypeBase` class. See the main
361[`Trait`](../Traits.md) documentation for more information on defining and using
362traits.
363
364### Interfaces
365
366Similarly to operations, `Type` classes may attach `Interfaces` to provide an
367abstract interface into the type. See the main [`Interface`](../Interfaces.md)
368documentation for more information on defining and using interfaces.
369
370## Attributes
371
372As stated in the introduction, the process for defining dialect attributes is
373nearly identical to that of defining dialect types. That key difference is that
374the things named `*Type` are generally now named `*Attr`.
375
376*   `Type::TypeBase` -> `Attribute::AttrBase`
377*   `TypeStorageAllocator` -> `AttributeStorageAllocator`
378*   `addTypes` -> `addAttributes`
379
380Aside from that, all of the interfaces for uniquing and storage construction are
381all the same.
382