llvm-project/mlir/docs/Traits.md

# Operation Traits

[TOC]

MLIR allows for a truly open operation ecosystem, as any dialect may define
operations that suit a specific level of abstraction. `Traits` are a mechanism
which abstracts implementation details and properties that are common
across many different operations. `Traits` may be used to specify special
properties and constraints of the operation, including whether the operation has
side effects or whether its output has the same type as the input. Some examples
of traits are `Commutative`, `SingleResult`, `Terminator`, etc. See the more
[comprehensive list](#traits) below for more examples of what is possible.

## Defining a Trait

Traits may be defined in C++ by inheriting from the
`OpTrait::TraitBase<ConcreteType, TraitType>` class. This base class takes as
template parameters:

*   ConcreteType
    -   The concrete operation type that this trait was attached to.
*   TraitType
    -   The type of the trait class that is being defined, for use with the
        [`Curiously Recurring Template Pattern`](https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern).

A derived trait class is expected to take a single template that corresponds to
the `ConcreteType`. An example trait definition is shown below:

```c++
template <typename ConcreteType>
class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
};
```

Derived traits may also provide a `verifyTrait` hook, that is called when
verifying the concrete operation. The trait verifiers will currently always be
invoked before the main `Op::verify`.

```c++
template <typename ConcreteType>
class MyTrait : public OpTrait::TraitBase<ConcreteType, MyTrait> {
public:
  /// Override the 'verifyTrait' hook to add additional verification on the
  /// concrete operation.
  static LogicalResult verifyTrait(Operation *op) {
    // ...
  }
};
```

Note: It is generally good practice to define the implementation of the
`verifyTrait` hook out-of-line as a free function when possible to avoid
instantiating the implementation for every concrete operation type.

### Parametric Traits

The above demonstrates the definition of a simple self-contained trait. It is
also often useful to provide some static parameters to the trait to control its
behavior. Given that the definition of the trait class is rigid, i.e. we must
have a single template argument for the concrete operation, the templates for
the parameters will need to be split out. An example is shown below:

```c++
template <int Parameter>
class MyParametricTrait {
public:
  template <typename ConcreteType>
  class Impl : public OpTrait::TraitBase<ConcreteType, Impl> {
    // Inside of 'Impl' we have full access to the template parameters
    // specified above.
  };
};
```

## Attaching a Trait

Traits may be used when defining a derived operation type, by simply adding the
name of the trait class to the `Op` class after the concrete operation type:

```c++
/// Here we define 'MyOp' along with the 'MyTrait' and `MyParametric trait
/// classes we defined previously.
class MyOp : public Op<MyOp, MyTrait, MyParametricTrait<10>::Impl> {};
```

To use a trait in the [ODS](OpDefinitions.md) framework, we need to provide a
definition of the trait class. This can be done using the `NativeOpTrait` and
`ParamNativeOpTrait` classes. `ParamNativeOpTrait` provides a mechanism in which
to specify arguments to a parametric trait class with an internal `Impl`.

```tablegen
// The argument is the c++ trait class name.
def MyTrait : NativeOpTrait<"MyTrait">;

// The first argument is the parent c++ class name. The second argument is a
// string containing the parameter list.
class MyParametricTrait<int prop>
  : NativeOpTrait<"MyParametricTrait", !cast<string>(!head(parameters))>;
```

These can then be used in the `traits` list of an op definition:

```tablegen
def OpWithInferTypeInterfaceOp : Op<...[MyTrait, MyParametricTrait<10>]> { ... }
```

See the documentation on [operation definitions](OpDefinitions.md) for more
details.

## Using a Trait

Traits may be used to provide additional methods, static fields, or other
information directly on the concrete operation. `Traits` internally become
`Base` classes of the concrete operation, so all of these are directly
accessible. To expose this information opaquely to transformations and analyses,
[`interfaces`](Interfaces.md) may be used.

To query if a specific operation contains a specific trait, the `hasTrait<>`
method may be used. This takes as a template parameter the trait class, which is
the same as the one passed when attaching the trait to an operation.

```c++
Operation *op = ..;
if (op->hasTrait<MyTrait>() || op->hasTrait<MyParametricTrait<10>::Impl>())
  ...;
```

## Trait List

MLIR provides a suite of traits that provide various functionalities that are
common across many different operations. Below is a list of some key traits that
may be used directly by any dialect. The format of the header for each trait
section goes as follows:

*   `Header`
    -   (`C++ class` -- `ODS class`(if applicable))

### AffineScope

*   `OpTrait::AffineScope` -- `AffineScope`

This trait is carried by region holding operations that define a new scope for
the purposes of polyhedral optimization and the affine dialect in particular.
Any SSA values of 'index' type that either dominate such operations, or are
defined at the top-level of such operations, or appear as region arguments for
such operations automatically become valid symbols for the polyhedral scope
defined by that operation. As a result, such SSA values could be used as the
operands or index operands of various affine dialect operations like affine.for,
affine.load, and affine.store.  The polyhedral scope defined by an operation
with this trait includes all operations in its region excluding operations that
are nested inside of other operations that themselves have this trait.

### AutomaticAllocationScope

*   `OpTrait::AutomaticAllocationScope` -- `AutomaticAllocationScope`

This trait is carried by region holding operations that define a new scope for
automatic allocation. Such allocations are automatically freed when control is
transferred back from the regions of such operations. As an example, allocations
performed by [`std.alloca`](Dialects/Standard.md#stdalloca-allocaop) are
automatically freed when control leaves the region of its closest surrounding op
that has the trait AutomaticAllocationScope.

### Broadcastable

*   `OpTrait::ResultsBroadcastableShape` -- `ResultsBroadcastableShape`

This trait adds the property that the operation is known to have
[broadcast-compatible](https://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
operands and its result types' shape is the broadcast compatible with the shape
of the broadcasted operands. Specifically, starting from the most varying
dimension, each dimension pair of the two operands' shapes should either be the
same or one of them is one. Also, the result shape should have the corresponding
dimension equal to the larger one, if known. Shapes are checked partially if
ranks or dimensions are not known. For example, an op with `tensor<?x2xf32>` and
`tensor<2xf32>` as operand types and `tensor<3x2xf32>` as the result type is
broadcast-compatible.

This trait requires that the operands are either vector or tensor types.

### Commutative

*   `OpTrait::IsCommutative` -- `Commutative`

This trait adds the property that the operation is commutative, i.e. `X op Y ==
Y op X`

### Function-Like

*   `OpTrait::FunctionLike`

This trait provides APIs for operations that behave like functions. In
particular:

-   Ops must be symbols, i.e. also have the `Symbol` trait;
-   Ops have a single region with multiple blocks that corresponds to the body
    of the function;
-   the absence of a region corresponds to an external function;
-   arguments of the first block of the region are treated as function
    arguments;
-   they can have argument and result attributes that are stored in dictionary
    attributes on the operation itself.

This trait does *NOT* provide type support for the functions, meaning that
concrete Ops must handle the type of the declared or defined function.
`getTypeAttrName()` is a convenience function that returns the name of the
attribute that can be used to store the function type, but the trait makes no
assumption based on it.

### HasParent

*   `OpTrait::HasParent<typename ParentOpType>` -- `HasParent<string op>`

This trait provides APIs and verifiers for operations that can only be nested
within regions that are attached to operations of `ParentOpType`.

### IsolatedFromAbove

*   `OpTrait::IsIsolatedFromAbove` -- `IsolatedFromAbove`

This trait signals that the regions of an operations are known to be isolated
from above. This trait asserts that the regions of an operation will not
capture, or reference, SSA values defined above the region scope. This means
that the following is invalid if `foo.region_op` is defined as
`IsolatedFromAbove`:

```mlir
%result = constant 10 : i32
foo.region_op {
  foo.yield %result : i32
}
```

This trait is an important structural property of the IR, and enables operations
to have [passes](PassManagement.md) scheduled under them.

### Single Block with Implicit Terminator

*   `OpTrait::SingleBlockImplicitTerminator<typename TerminatorOpType>` :
    `SingleBlockImplicitTerminator<string op>`

This trait provides APIs and verifiers for operations with regions that have a
single block that must terminate with `TerminatorOpType`.

### Symbol

*   `OpTrait::Symbol` -- `Symbol`

This trait is used for operations that define a
[`Symbol`](SymbolsAndSymbolTables.md#symbol).

### SymbolTable

*   `OpTrait::SymbolTable` -- `SymbolTable`

This trait is used for operations that define a
[`SymbolTable`](SymbolsAndSymbolTables.md#symbol-table).

### Terminator

*   `OpTrait::IsTerminator` -- `Terminator`

This trait provides verification and functionality for operations that are known
to be [terminators](LangRef.md#terminator-operations).