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