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](OpDefinitions.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### Parametric Traits 104 105The above demonstrates the definition of a simple self-contained trait. It is 106also often useful to provide some static parameters to the trait to control its 107behavior. Given that the definition of the trait class is rigid, i.e. we must 108have a single template argument for the concrete object, the templates for the 109parameters will need to be split out. An example is shown below: 110 111```c++ 112template <int Parameter> 113class MyParametricTrait { 114public: 115 template <typename ConcreteType> 116 class Impl : public TraitBase<ConcreteType, Impl> { 117 // Inside of 'Impl' we have full access to the template parameters 118 // specified above. 119 }; 120}; 121``` 122 123## Attaching a Trait 124 125Traits may be used when defining a derived object type, by simply appending the 126name of the trait class to the end of the base object class operation type: 127 128```c++ 129/// Here we define 'MyAttr' along with the 'MyTrait' and `MyParametric trait 130/// classes we defined previously. 131class MyAttr : public Attribute::AttrBase<MyAttr, ..., MyTrait, MyParametricTrait<10>::Impl> {}; 132/// Here we define 'MyOp' along with the 'MyTrait' and `MyParametric trait 133/// classes we defined previously. 134class MyOp : public Op<MyOp, MyTrait, MyParametricTrait<10>::Impl> {}; 135/// Here we define 'MyType' along with the 'MyTrait' and `MyParametric trait 136/// classes we defined previously. 137class MyType : public Type::TypeBase<MyType, ..., MyTrait, MyParametricTrait<10>::Impl> {}; 138``` 139 140### Attaching Operation Traits in ODS 141 142To use an operation trait in the [ODS](OpDefinitions.md) framework, we need to 143provide a definition of the trait class. This can be done using the 144`NativeOpTrait` and `ParamNativeOpTrait` classes. `ParamNativeOpTrait` provides 145a mechanism in which to specify arguments to a parametric trait class with an 146internal `Impl`. 147 148```tablegen 149// The argument is the c++ trait class name. 150def MyTrait : NativeOpTrait<"MyTrait">; 151 152// The first argument is the parent c++ class name. The second argument is a 153// string containing the parameter list. 154class MyParametricTrait<int prop> 155 : NativeOpTrait<"MyParametricTrait", !cast<string>(!head(parameters))>; 156``` 157 158These can then be used in the `traits` list of an op definition: 159 160```tablegen 161def OpWithInferTypeInterfaceOp : Op<...[MyTrait, MyParametricTrait<10>]> { ... } 162``` 163 164See the documentation on [operation definitions](OpDefinitions.md) for more 165details. 166 167## Using a Trait 168 169Traits may be used to provide additional methods, static fields, or other 170information directly on the concrete object. `Traits` internally become `Base` 171classes of the concrete operation, so all of these are directly accessible. To 172expose this information opaquely to transformations and analyses, 173[`interfaces`](Interfaces.md) may be used. 174 175To query if a specific object contains a specific trait, the `hasTrait<>` method 176may be used. This takes as a template parameter the trait class, which is the 177same as the one passed when attaching the trait to an operation. 178 179```c++ 180Operation *op = ..; 181if (op->hasTrait<MyTrait>() || op->hasTrait<MyParametricTrait<10>::Impl>()) 182 ...; 183``` 184 185## Operation Traits List 186 187MLIR provides a suite of traits that provide various functionalities that are 188common across many different operations. Below is a list of some key traits that 189may be used directly by any dialect. The format of the header for each trait 190section goes as follows: 191 192* `Header` 193 - (`C++ class` -- `ODS class`(if applicable)) 194 195### AffineScope 196 197* `OpTrait::AffineScope` -- `AffineScope` 198 199This trait is carried by region holding operations that define a new scope for 200the purposes of polyhedral optimization and the affine dialect in particular. 201Any SSA values of 'index' type that either dominate such operations, or are 202defined at the top-level of such operations, or appear as region arguments for 203such operations automatically become valid symbols for the polyhedral scope 204defined by that operation. As a result, such SSA values could be used as the 205operands or index operands of various affine dialect operations like affine.for, 206affine.load, and affine.store. The polyhedral scope defined by an operation with 207this trait includes all operations in its region excluding operations that are 208nested inside of other operations that themselves have this trait. 209 210### AutomaticAllocationScope 211 212* `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope` 213 214This trait is carried by region holding operations that define a new scope for 215automatic allocation. Such allocations are automatically freed when control is 216transferred back from the regions of such operations. As an example, allocations 217performed by 218[`memref.alloca`](Dialects/MemRef.md/#memrefalloca-mlirmemrefallocaop) are 219automatically freed when control leaves the region of its closest surrounding op 220that has the trait AutomaticAllocationScope. 221 222### Broadcastable 223 224* `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape` 225 226This trait adds the property that the operation is known to have 227[broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) 228operands and its result types' shape is the broadcast compatible with the shape 229of the broadcasted operands. Specifically, starting from the most varying 230dimension, each dimension pair of the two operands' shapes should either be the 231same or one of them is one. Also, the result shape should have the corresponding 232dimension equal to the larger one, if known. Shapes are checked partially if 233ranks or dimensions are not known. For example, an op with `tensor<?x2xf32>` and 234`tensor<2xf32>` as operand types and `tensor<3x2xf32>` as the result type is 235broadcast-compatible. 236 237This trait requires that the operands are either vector or tensor types. 238 239### Commutative 240 241* `OpTrait::IsCommutative` -- `Commutative` 242 243This trait adds the property that the operation is commutative, i.e. `X op Y == 244Y op X` 245 246### ElementwiseMappable 247 248* `OpTrait::ElementwiseMappable` -- `ElementwiseMappable` 249 250This trait tags scalar ops that also can be applied to vectors/tensors, with 251their semantics on vectors/tensors being elementwise application. This trait 252establishes a set of properties that allow reasoning about / converting between 253scalar/vector/tensor code. These same properties allow blanket implementations 254of various analyses/transformations for all `ElementwiseMappable` ops. 255 256Note: Not all ops that are "elementwise" in some abstract sense satisfy this 257trait. In particular, broadcasting behavior is not allowed. See the comments on 258`OpTrait::ElementwiseMappable` for the precise requirements. 259 260### HasParent 261 262* `OpTrait::HasParent<typename ParentOpType>` -- `HasParent<string op>` or 263 `ParentOneOf<list<string> opList>` 264 265This trait provides APIs and verifiers for operations that can only be nested 266within regions that are attached to operations of `ParentOpType`. 267 268### IsolatedFromAbove 269 270* `OpTrait::IsIsolatedFromAbove` -- `IsolatedFromAbove` 271 272This trait signals that the regions of an operations are known to be isolated 273from above. This trait asserts that the regions of an operation will not 274capture, or reference, SSA values defined above the region scope. This means 275that the following is invalid if `foo.region_op` is defined as 276`IsolatedFromAbove`: 277 278```mlir 279%result = arith.constant 10 : i32 280foo.region_op { 281 foo.yield %result : i32 282} 283``` 284 285This trait is an important structural property of the IR, and enables operations 286to have [passes](PassManagement.md) scheduled under them. 287 288### MemRefsNormalizable 289 290* `OpTrait::MemRefsNormalizable` -- `MemRefsNormalizable` 291 292This trait is used to flag operations that consume or produce values of `MemRef` 293type where those references can be 'normalized'. In cases where an associated 294`MemRef` has a non-identity memory-layout specification, such normalizable 295operations can be modified so that the `MemRef` has an identity layout 296specification. This can be implemented by associating the operation with its own 297index expression that can express the equivalent of the memory-layout 298specification of the MemRef type. See [the -normalize-memrefs pass]. 299(https://mlir.llvm.org/docs/Passes/#-normalize-memrefs-normalize-memrefs) 300 301### Single Block Region 302 303* `OpTrait::SingleBlock` -- `SingleBlock` 304 305This trait provides APIs and verifiers for operations with regions that have a 306single block. 307 308### Single Block with Implicit Terminator 309 310* `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` -- 311 `SingleBlockImplicitTerminator<string op>` 312 313This trait implies the `SingleBlock` above, but adds the additional requirement 314that the single block must terminate with `TerminatorOpType`. 315 316### SymbolTable 317 318* `OpTrait::SymbolTable` -- `SymbolTable` 319 320This trait is used for operations that define a 321[`SymbolTable`](SymbolsAndSymbolTables.md#symbol-table). 322 323### Terminator 324 325* `OpTrait::IsTerminator` -- `Terminator` 326 327This trait provides verification and functionality for operations that are known 328to be [terminators](LangRef.md#terminator-operations). 329 330* `OpTrait::NoTerminator` -- `NoTerminator` 331 332This trait removes the requirement on regions held by an operation to have 333[terminator operations](LangRef.md#terminator-operations) at the end of a block. 334This requires that these regions have a single block. An example of operation 335using this trait is the top-level `ModuleOp`. 336