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