• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Traits
2
3[TOC]
4
5MLIR allows for a truly open ecosystem, as any dialect may define attributes,
6operations, and types that suit a specific level of abstraction. `Traits` are a
7mechanism which abstracts implementation details and properties that are common
8across many different attributes/operations/types/etc.. `Traits` may be used to
9specify special properties and constraints of the object, including whether an
10operation has side effects or that its output has the same type as the input.
11Some examples of operation traits are `Commutative`, `SingleResult`,
12`Terminator`, etc. See the more comprehensive list of
13[operation traits](#operation-traits-list) below for more examples of what is
14possible.
15
16## Defining a Trait
17
18Traits may be defined in C++ by inheriting from the `TraitBase<ConcreteType,
19TraitType>` class for the specific IR type. For attributes, this is
20`AttributeTrait::TraitBase`. For operations, this is `OpTrait::TraitBase`. For
21types, this is `TypeTrait::TraitBase`. This base class takes as template
22parameters:
23
24*   ConcreteType
25    -   The concrete class type that this trait was attached to.
26*   TraitType
27    -   The type of the trait class that is being defined, for use with the
28        [`Curiously Recurring Template Pattern`](https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern).
29
30A derived trait class is expected to take a single template that corresponds to
31the `ConcreteType`. An example trait definition is shown below:
32
33```c++
34template <typename ConcreteType>
35class MyTrait : public TraitBase<ConcreteType, MyTrait> {
36};
37```
38
39Operation traits may also provide a `verifyTrait` hook, that is called when
40verifying the concrete operation. The trait verifiers will currently always be
41invoked before the main `Op::verify`.
42
43```c++
44template <typename ConcreteType>
45class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
46public:
47  /// Override the 'verifyTrait' hook to add additional verification on the
48  /// concrete operation.
49  static LogicalResult verifyTrait(Operation *op) {
50    // ...
51  }
52};
53```
54
55Note: It is generally good practice to define the implementation of the
56`verifyTrait` hook out-of-line as a free function when possible to avoid
57instantiating the implementation for every concrete operation type.
58
59Operation traits may also provide a `foldTrait` hook that is called when
60folding the concrete operation. The trait folders will only be invoked if
61the concrete operation fold is either not implemented, fails, or performs
62an in-place fold.
63
64The following signature of fold will be called if it is implemented
65and the op has a single result.
66
67```c++
68template <typename ConcreteType>
69class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
70public:
71  /// Override the 'foldTrait' hook to support trait based folding on the
72  /// concrete operation.
73  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) { {
74    // ...
75  }
76};
77```
78
79Otherwise, if the operation has a single result and the above signature is
80not implemented, or the operation has multiple results, then the following signature
81will be used (if implemented):
82
83```c++
84template <typename ConcreteType>
85class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
86public:
87  /// Override the 'foldTrait' hook to support trait based folding on the
88  /// concrete operation.
89  static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
90                                 SmallVectorImpl<OpFoldResult> &results) { {
91    // ...
92  }
93};
94```
95
96Note: It is generally good practice to define the implementation of the
97`foldTrait` hook out-of-line as a free function when possible to avoid
98instantiating the implementation for every concrete operation type.
99
100### Parametric Traits
101
102The above demonstrates the definition of a simple self-contained trait. It is
103also often useful to provide some static parameters to the trait to control its
104behavior. Given that the definition of the trait class is rigid, i.e. we must
105have a single template argument for the concrete object, the templates for the
106parameters will need to be split out. An example is shown below:
107
108```c++
109template <int Parameter>
110class MyParametricTrait {
111public:
112  template <typename ConcreteType>
113  class Impl : public TraitBase<ConcreteType, Impl> {
114    // Inside of 'Impl' we have full access to the template parameters
115    // specified above.
116  };
117};
118```
119
120## Attaching a Trait
121
122Traits may be used when defining a derived object type, by simply appending the
123name of the trait class to the end of the base object class operation type:
124
125```c++
126/// Here we define 'MyAttr' along with the 'MyTrait' and `MyParametric trait
127/// classes we defined previously.
128class MyAttr : public Attribute::AttrBase<MyAttr, ..., MyTrait, MyParametricTrait<10>::Impl> {};
129/// Here we define 'MyOp' along with the 'MyTrait' and `MyParametric trait
130/// classes we defined previously.
131class MyOp : public Op<MyOp, MyTrait, MyParametricTrait<10>::Impl> {};
132/// Here we define 'MyType' along with the 'MyTrait' and `MyParametric trait
133/// classes we defined previously.
134class MyType : public Type::TypeBase<MyType, ..., MyTrait, MyParametricTrait<10>::Impl> {};
135```
136
137### Attaching Operation Traits in ODS
138
139To use an operation trait in the [ODS](OpDefinitions.md) framework, we need to
140provide a definition of the trait class. This can be done using the
141`NativeOpTrait` and `ParamNativeOpTrait` classes. `ParamNativeOpTrait` provides
142a mechanism in which to specify arguments to a parametric trait class with an
143internal `Impl`.
144
145```tablegen
146// The argument is the c++ trait class name.
147def MyTrait : NativeOpTrait<"MyTrait">;
148
149// The first argument is the parent c++ class name. The second argument is a
150// string containing the parameter list.
151class MyParametricTrait<int prop>
152  : NativeOpTrait<"MyParametricTrait", !cast<string>(!head(parameters))>;
153```
154
155These can then be used in the `traits` list of an op definition:
156
157```tablegen
158def OpWithInferTypeInterfaceOp : Op<...[MyTrait, MyParametricTrait<10>]> { ... }
159```
160
161See the documentation on [operation definitions](OpDefinitions.md) for more
162details.
163
164## Using a Trait
165
166Traits may be used to provide additional methods, static fields, or other
167information directly on the concrete object. `Traits` internally become `Base`
168classes of the concrete operation, so all of these are directly accessible. To
169expose this information opaquely to transformations and analyses,
170[`interfaces`](Interfaces.md) may be used.
171
172To query if a specific object contains a specific trait, the `hasTrait<>` method
173may be used. This takes as a template parameter the trait class, which is the
174same as the one passed when attaching the trait to an operation.
175
176```c++
177Operation *op = ..;
178if (op->hasTrait<MyTrait>() || op->hasTrait<MyParametricTrait<10>::Impl>())
179  ...;
180```
181
182## Operation Traits List
183
184MLIR provides a suite of traits that provide various functionalities that are
185common across many different operations. Below is a list of some key traits that
186may be used directly by any dialect. The format of the header for each trait
187section goes as follows:
188
189*   `Header`
190    -   (`C++ class` -- `ODS class`(if applicable))
191
192### AffineScope
193
194*   `OpTrait::AffineScope` -- `AffineScope`
195
196This trait is carried by region holding operations that define a new scope for
197the purposes of polyhedral optimization and the affine dialect in particular.
198Any SSA values of 'index' type that either dominate such operations, or are
199defined at the top-level of such operations, or appear as region arguments for
200such operations automatically become valid symbols for the polyhedral scope
201defined by that operation. As a result, such SSA values could be used as the
202operands or index operands of various affine dialect operations like affine.for,
203affine.load, and affine.store.  The polyhedral scope defined by an operation
204with this trait includes all operations in its region excluding operations that
205are nested inside of other operations that themselves have this trait.
206
207### AutomaticAllocationScope
208
209*   `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope`
210
211This trait is carried by region holding operations that define a new scope for
212automatic allocation. Such allocations are automatically freed when control is
213transferred back from the regions of such operations. As an example, allocations
214performed by [`std.alloca`](Dialects/Standard.md#stdalloca-allocaop) are
215automatically freed when control leaves the region of its closest surrounding op
216that has the trait AutomaticAllocationScope.
217
218### Broadcastable
219
220*   `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape`
221
222This trait adds the property that the operation is known to have
223[broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
224operands and its result types' shape is the broadcast compatible with the shape
225of the broadcasted operands. Specifically, starting from the most varying
226dimension, each dimension pair of the two operands' shapes should either be the
227same or one of them is one. Also, the result shape should have the corresponding
228dimension equal to the larger one, if known. Shapes are checked partially if
229ranks or dimensions are not known. For example, an op with `tensor<?x2xf32>` and
230`tensor<2xf32>` as operand types and `tensor<3x2xf32>` as the result type is
231broadcast-compatible.
232
233This trait requires that the operands are either vector or tensor types.
234
235### Commutative
236
237*   `OpTrait::IsCommutative` -- `Commutative`
238
239This trait adds the property that the operation is commutative, i.e. `X op Y ==
240Y op X`
241
242### ElementwiseMappable
243
244* `OpTrait::ElementwiseMappable` -- `ElementwiseMappable`
245
246This trait tags scalar ops that also can be applied to vectors/tensors, with
247their semantics on vectors/tensors being elementwise application. This trait
248establishes a set of properties that allow reasoning about / converting between
249scalar/vector/tensor code. These same properties allow blanket implementations
250of various analyses/transformations for all `ElementwiseMappable` ops.
251
252Note: Not all ops that are "elementwise" in some abstract sense satisfy this
253trait. In particular, broadcasting behavior is not allowed. See the comments on
254`OpTrait::ElementwiseMappable` for the precise requirements.
255
256### Function-Like
257
258*   `OpTrait::FunctionLike`
259
260This trait provides APIs for operations that behave like functions. In
261particular:
262
263-   Ops must be symbols, i.e. also have the `Symbol` trait;
264-   Ops have a single region with multiple blocks that corresponds to the body
265    of the function;
266-   An op with a single empty region corresponds to an external function;
267-   arguments of the first block of the region are treated as function
268    arguments;
269-   they can have argument and result attributes that are stored in dictionary
270    attributes on the operation itself.
271
272This trait provides limited type support for the declared or defined functions.
273The convenience function `getTypeAttrName()` returns the name of an attribute
274that can be used to store the function type. In addition, this trait provides
275`getType` and `setType` helpers to store a `FunctionType` in the attribute named
276by `getTypeAttrName()`.
277
278In general, this trait assumes concrete ops use `FunctionType` under the hood.
279If this is not the case, in order to use the function type support, concrete ops
280must define the following methods, using the same name, to hide the ones defined
281for `FunctionType`: `addBodyBlock`, `getType`, `getTypeWithoutArgsAndResults`
282and `setType`.
283
284### HasParent
285
286*   `OpTrait::HasParent<typename ParentOpType>` -- `HasParent<string op>` or
287    `ParentOneOf<list<string> opList>`
288
289This trait provides APIs and verifiers for operations that can only be nested
290within regions that are attached to operations of `ParentOpType`.
291
292### IsolatedFromAbove
293
294*   `OpTrait::IsIsolatedFromAbove` -- `IsolatedFromAbove`
295
296This trait signals that the regions of an operations are known to be isolated
297from above. This trait asserts that the regions of an operation will not
298capture, or reference, SSA values defined above the region scope. This means
299that the following is invalid if `foo.region_op` is defined as
300`IsolatedFromAbove`:
301
302```mlir
303%result = constant 10 : i32
304foo.region_op {
305  foo.yield %result : i32
306}
307```
308
309This trait is an important structural property of the IR, and enables operations
310to have [passes](PassManagement.md) scheduled under them.
311
312### MemRefsNormalizable
313
314* `OpTrait::MemRefsNormalizable` -- `MemRefsNormalizable`
315
316This trait is used to flag operations that consume or produce
317values of `MemRef` type where those references can be 'normalized'.
318In cases where an associated `MemRef` has a
319non-identity memory-layout specification, such normalizable operations can be
320modified so that the `MemRef` has an identity layout specification.
321This can be implemented by associating the operation with its own
322index expression that can express the equivalent of the memory-layout
323specification of the MemRef type. See [the -normalize-memrefs pass].
324(https://mlir.llvm.org/docs/Passes/#-normalize-memrefs-normalize-memrefs)
325
326### Single Block with Implicit Terminator
327
328*   `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` :
329    `SingleBlockImplicitTerminator<string op>`
330
331This trait provides APIs and verifiers for operations with regions that have a
332single block that must terminate with `TerminatorOpType`.
333
334### SymbolTable
335
336*   `OpTrait::SymbolTable` -- `SymbolTable`
337
338This trait is used for operations that define a
339[`SymbolTable`](SymbolsAndSymbolTables.md#symbol-table).
340
341### Terminator
342
343*   `OpTrait::IsTerminator` -- `Terminator`
344
345This trait provides verification and functionality for operations that are known
346to be [terminators](LangRef.md#terminator-operations).
347