xref: /llvm-project/mlir/docs/Traits/_index.md (revision 777a6e6f10b2b90496d248b7fa904fce834484be)
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` or `verifyRegionTrait` hook
40that is called when verifying the concrete operation. The difference between
41these two is that whether the verifier needs to access the regions, if so, the
42operations in the regions will be verified before the verification of this
43trait. The [verification order](../DefiningDialects/Operations.md/#verification-ordering)
44determines when a verifier will be invoked.
45
46```c++
47template <typename ConcreteType>
48class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
49public:
50  /// Override the 'verifyTrait' hook to add additional verification on the
51  /// concrete operation.
52  static LogicalResult verifyTrait(Operation *op) {
53    // ...
54  }
55};
56```
57
58Note: It is generally good practice to define the implementation of the
59`verifyTrait` or `verifyRegionTrait` hook out-of-line as a free function when
60possible to avoid instantiating the implementation for every concrete operation
61type.
62
63Operation traits may also provide a `foldTrait` hook that is called when folding
64the concrete operation. The trait folders will only be invoked if the concrete
65operation fold is either not implemented, fails, or performs an in-place fold.
66
67The following signature of fold will be called if it is implemented and the op
68has a single result.
69
70```c++
71template <typename ConcreteType>
72class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
73public:
74  /// Override the 'foldTrait' hook to support trait based folding on the
75  /// concrete operation.
76  static OpFoldResult foldTrait(Operation *op, ArrayRef<Attribute> operands) {
77    // ...
78  }
79};
80```
81
82Otherwise, if the operation has a single result and the above signature is not
83implemented, or the operation has multiple results, then the following signature
84will be used (if implemented):
85
86```c++
87template <typename ConcreteType>
88class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
89public:
90  /// Override the 'foldTrait' hook to support trait based folding on the
91  /// concrete operation.
92  static LogicalResult foldTrait(Operation *op, ArrayRef<Attribute> operands,
93                                 SmallVectorImpl<OpFoldResult> &results) {
94    // ...
95  }
96};
97```
98
99Note: It is generally good practice to define the implementation of the
100`foldTrait` hook out-of-line as a free function when possible to avoid
101instantiating the implementation for every concrete operation type.
102
103### Extra Declarations and Definitions
104A trait may require additional declarations and definitions directly on
105the Operation, Attribute or Type instances which specify that trait.
106The `extraConcreteClassDeclaration` and `extraConcreteClassDefinition`
107fields under the `NativeTrait` class are mechanisms designed for injecting
108code directly into generated C++ Operation, Attribute or Type classes.
109
110Code within the `extraConcreteClassDeclaration` field will be formatted and copied
111into the generated C++ Operation, Attribute or Type class. Code within
112`extraConcreteClassDefinition` will be added to the generated source file inside
113the class’s C++ namespace. The substitution `$cppClass` is replaced by the C++ class
114name.
115
116The intention is to group trait specific logic together and reduce
117redundant extra declarations and definitions on the instances themselves.
118
119### Parametric Traits
120
121The above demonstrates the definition of a simple self-contained trait. It is
122also often useful to provide some static parameters to the trait to control its
123behavior. Given that the definition of the trait class is rigid, i.e. we must
124have a single template argument for the concrete object, the templates for the
125parameters will need to be split out. An example is shown below:
126
127```c++
128template <int Parameter>
129class MyParametricTrait {
130public:
131  template <typename ConcreteType>
132  class Impl : public TraitBase<ConcreteType, Impl> {
133    // Inside of 'Impl' we have full access to the template parameters
134    // specified above.
135  };
136};
137```
138
139## Attaching a Trait
140
141Traits may be used when defining a derived object type, by simply appending the
142name of the trait class to the end of the base object class operation type:
143
144```c++
145/// Here we define 'MyAttr' along with the 'MyTrait' and `MyParametric trait
146/// classes we defined previously.
147class MyAttr : public Attribute::AttrBase<MyAttr, ..., MyTrait, MyParametricTrait<10>::Impl> {};
148/// Here we define 'MyOp' along with the 'MyTrait' and `MyParametric trait
149/// classes we defined previously.
150class MyOp : public Op<MyOp, MyTrait, MyParametricTrait<10>::Impl> {};
151/// Here we define 'MyType' along with the 'MyTrait' and `MyParametric trait
152/// classes we defined previously.
153class MyType : public Type::TypeBase<MyType, ..., MyTrait, MyParametricTrait<10>::Impl> {};
154```
155
156### Attaching Operation Traits in ODS
157
158To use an operation trait in the [ODS](../DefiningDialects/Operations.md) framework, we need to
159provide a definition of the trait class. This can be done using the
160`NativeOpTrait` and `ParamNativeOpTrait` classes. `ParamNativeOpTrait` provides
161a mechanism in which to specify arguments to a parametric trait class with an
162internal `Impl`.
163
164```tablegen
165// The argument is the c++ trait class name.
166def MyTrait : NativeOpTrait<"MyTrait">;
167
168// The first argument is the parent c++ class name. The second argument is a
169// string containing the parameter list.
170class MyParametricTrait<int prop>
171  : NativeOpTrait<"MyParametricTrait", !cast<string>(!head(parameters))>;
172```
173
174These can then be used in the `traits` list of an op definition:
175
176```tablegen
177def OpWithInferTypeInterfaceOp : Op<...[MyTrait, MyParametricTrait<10>]> { ... }
178```
179
180See the documentation on [operation definitions](../DefiningDialects/Operations.md) for more
181details.
182
183## Using a Trait
184
185Traits may be used to provide additional methods, static fields, or other
186information directly on the concrete object. `Traits` internally become `Base`
187classes of the concrete operation, so all of these are directly accessible. To
188expose this information opaquely to transformations and analyses,
189[`interfaces`](../Interfaces.md) may be used.
190
191To query if a specific object contains a specific trait, the `hasTrait<>` method
192may be used. This takes as a template parameter the trait class, which is the
193same as the one passed when attaching the trait to an operation.
194
195```c++
196Operation *op = ..;
197if (op->hasTrait<MyTrait>() || op->hasTrait<MyParametricTrait<10>::Impl>())
198  ...;
199```
200
201## Operation Traits List
202
203MLIR provides a suite of traits that provide various functionalities that are
204common across many different operations. Below is a list of some key traits that
205may be used directly by any dialect. The format of the header for each trait
206section goes as follows:
207
208*   `Header`
209    -   (`C++ class` -- `ODS class`(if applicable))
210
211### AffineScope
212
213*   `OpTrait::AffineScope` -- `AffineScope`
214
215This trait is carried by region holding operations that define a new scope for
216the purposes of polyhedral optimization and the affine dialect in particular.
217Any SSA values of 'index' type that either dominate such operations, or are
218defined at the top-level of such operations, or appear as region arguments for
219such operations automatically become valid symbols for the polyhedral scope
220defined by that operation. As a result, such SSA values could be used as the
221operands or index operands of various affine dialect operations like affine.for,
222affine.load, and affine.store. The polyhedral scope defined by an operation with
223this trait includes all operations in its region excluding operations that are
224nested inside of other operations that themselves have this trait.
225
226### AutomaticAllocationScope
227
228*   `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope`
229
230This trait is carried by region holding operations that define a new scope for
231automatic allocation. Such allocations are automatically freed when control is
232transferred back from the regions of such operations. As an example, allocations
233performed by
234[`memref.alloca`](../Dialects/MemRef.md/#memrefalloca-memrefallocaop) are
235automatically freed when control leaves the region of its closest surrounding op
236that has the trait AutomaticAllocationScope.
237
238### Broadcastable
239
240*   `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape`
241
242This trait adds the property that the operation is known to have
243[broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
244operands and that its result type is compatible with the inferred broadcast shape.
245See [The `Broadcastable` Trait](Broadcastable.md) for details.
246
247### Commutative
248
249*   `OpTrait::IsCommutative` -- `Commutative`
250
251This trait adds the property that the operation is commutative, i.e. `X op Y ==
252Y op X`
253
254### ElementwiseMappable
255
256*   `OpTrait::ElementwiseMappable` -- `ElementwiseMappable`
257
258This trait tags scalar ops that also can be applied to vectors/tensors, with
259their semantics on vectors/tensors being elementwise application. This trait
260establishes a set of properties that allow reasoning about / converting between
261scalar/vector/tensor code. These same properties allow blanket implementations
262of various analyses/transformations for all `ElementwiseMappable` ops.
263
264Note: Not all ops that are "elementwise" in some abstract sense satisfy this
265trait. In particular, broadcasting behavior is not allowed. See the comments on
266`OpTrait::ElementwiseMappable` for the precise requirements.
267
268### HasParent
269
270*   `OpTrait::HasParent<typename ParentOpType>` -- `HasParent<string op>` or
271    `ParentOneOf<list<string> opList>`
272
273This trait provides APIs and verifiers for operations that can only be nested
274within regions that are attached to operations of `ParentOpType`.
275
276### IsolatedFromAbove
277
278*   `OpTrait::IsIsolatedFromAbove` -- `IsolatedFromAbove`
279
280This trait signals that the regions of an operations are known to be isolated
281from above. This trait asserts that the regions of an operation will not
282capture, or reference, SSA values defined above the region scope. This means
283that the following is invalid if `foo.region_op` is defined as
284`IsolatedFromAbove`:
285
286```mlir
287%result = arith.constant 10 : i32
288foo.region_op {
289  foo.yield %result : i32
290}
291```
292
293This trait is an important structural property of the IR, and enables operations
294to have [passes](../PassManagement) scheduled under them.
295
296### MemRefsNormalizable
297
298*   `OpTrait::MemRefsNormalizable` -- `MemRefsNormalizable`
299
300This trait is used to flag operations that consume or produce values of `MemRef`
301type where those references can be 'normalized'. In cases where an associated
302`MemRef` has a non-identity memory-layout specification, such normalizable
303operations can be modified so that the `MemRef` has an identity layout
304specification. This can be implemented by associating the operation with its own
305index expression that can express the equivalent of the memory-layout
306specification of the MemRef type. See [the -normalize-memrefs pass](../Passes.md/#-normalize-memrefs).
307
308### Single Block Region
309
310*   `OpTrait::SingleBlock` -- `SingleBlock`
311
312This trait provides APIs and verifiers for operations with regions that have a
313single block.
314
315### Single Block with Implicit Terminator
316
317*   `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` --
318    `SingleBlockImplicitTerminator<string op>`
319
320This trait implies the `SingleBlock` above, but adds the additional requirement
321that the single block must terminate with `TerminatorOpType`.
322
323### SymbolTable
324
325*   `OpTrait::SymbolTable` -- `SymbolTable`
326
327This trait is used for operations that define a
328[`SymbolTable`](../SymbolsAndSymbolTables.md/#symbol-table).
329
330### Terminator
331
332*   `OpTrait::IsTerminator` -- `Terminator`
333
334This trait provides verification and functionality for operations that are known
335to be [terminators](../LangRef.md/#control-flow-and-ssacfg-regions).
336
337*   `OpTrait::NoTerminator` -- `NoTerminator`
338
339This trait removes the requirement on regions held by an operation to have
340[terminator operations](../LangRef.md/#control-flow-and-ssacfg-regions) at the end of a block.
341This requires that these regions have a single block. An example of operation
342using this trait is the top-level `ModuleOp`.
343