llvm-project/mlir/g3doc/OpDefinitions.md

Operation definitions

Motivation

MLIR allows pluggable dialects, and dialects contain, among others, a list of operations. This open and extensible ecosystem leads to the "stringly" type IR problem, e.g., repetitive string comparisons during optimization and analysis passes, unintuitive accessor methods (e.g., generic/error prone GetOperand(3) vs GetStride()) with more generic return types, constructors are verbose, generic, and don't have default arguments, the MLIR assembly format is verbose and unclear, and op verification is:

1. best case: a central string-to-verification-function map,
2. middle case: duplication of verification across the code base, or
3. worst case: no verification functions.

The fix is to support op descriptions, which (in one central place per-dialect) contain everything you need to know about the op, its invariants, traits, textual formatting, etc. This description is also used to generate helper functions and classes to allow analysis/builder/verification/parsing/printing.

Requirements

The op description should as declarative as possible to allow a wide range of tools to work with them and query methods generated from them. In particular this means specifying traits, constraints and shape inference information in a way that is easily analyzable (e.g., avoid opaque calls to C++ functions where possible).

We considered the approaches of several contemporary systems and focused on requirements that were desirable:

• Ops registered using a registry separate from C++ code.

  • Unknown ops are allowed in MLIR, so ops need not be registered. The ability of the compiler to optimize those ops or graphs containing those ops is constrained but correct.
  • The current proposal does not include a runtime op description, but it does not preclude such description, it can be added later.
  • The op registry is essential for generating C++ classes that make manipulating ops, verifying correct construction etc. in C++ easier by providing a typed representation and accessors.

• The op registry will be defined in TableGen and be used to generate C++ classes and utility functions (builder/verifier/parser/printer).

  • TableGen is a modelling specification language used by LLVM's backends and fits in well with trait based modelling. This is an implementation decision and there are alternative ways of doing this. But the specification language is good for the requirements of modelling the traits (as seen from usage in LLVM processor backend modelling) and easy to extend, so a practical choice. If another good option comes up, we will consider it.

• MLIR allows both defined and undefined ops.

  • Defined ops should have fixed semantics and could have a corresponding reference implementation defined using, for example, EDSC.
  • Dialects are under full control of the dialect owner and normally live with the framework of the dialect.

• The op's traits (e.g., commutative) are modelled along with the op in the registry.

• The op's operand/return type constraints are modelled along with the op in the registry (see Type constraints discussion below), this allows (e.g.) optimized concise syntax in textual dumps.

• Behavior of the op is documented along with the op with a summary and a description. The description is written in markdown and extracted for inclusion in the generated LangRef section of the dialect.

• The generic assembly form of printing and parsing is available as normal, but a custom parser and printer can either be specified or automatically generated from an optional string representation showing the mapping of the "assembly" string to operands/type.

  • Parser-level remappings (e.g., eq to enum) will be supported as part of the parser generation.

• Matching patterns are specified separately from the op description.

  • Contrasted with LLVM there is no "base" set of ops that every backend needs to be aware of. Instead there are many different dialects and the transformations/legalizations between these dialects form a graph of transformations.

• Reference implementation may be provided along with the op definition.

  • The reference implementation may be in terms of either standard ops or other reference implementations.

    TODO: document expectation if the dependent op's definition changes.

Operation definition

As an example of the proposal to declare an operation (say tf.Add). This is intended to be a fully contained example, in practice one would create a helper classes that abstract out common functionality (e.g., TF_BinaryOp).

def TF_AddOp : Op<"tf.Add", [Broadcastable, NoSideEffect]>,
               Arguments<(ins TF_Tensor:$x, TF_Tensor:$y)>,
               Results<(outs TF_Tensor:$z)> {
  let summary = "Addition operator";

  let description = [{
    Returns lhs + rhs element-wise.

    The inputs and result must be of the same elemental type.
  }];

  let reference = [{
    auto ivs = makeBindables(lhsShape.size());
    block = edsc::Block({
      For(ivs, 0, lhsShape, 1, {
        result[ivs] = lhs[ivs] + rhs[ivs]
      })});
    }
  }];
}

Operation definitions consists of:

1. Operation name (opName).

   This is a unique identifier used to distinguish this operation vs all others defined in MLIR. This is the equivalent of the mnemonic in assembly language. Operations are within dialects which effectively namespace operations. The C++ class generated for the operation is based on the definition in the TableGen's file's name. E.g., TF_AddOp above would result in a C++ class called AddOp generated in the namespace TF.

2. Summmary and description.

   These are human readable documentation for the operation. Documentation of the operations can be generated from the same source of truth as the operation.

3. Arguments (arguments).

   This is a list of operands (optionally named) and named attributes used to generate builder, accessor functions and verification.

   1. Operands.

      These are the results of other operations and mostly only known at runtime. They can have a fixed type or correspond to set of possible types. See Type constraints specification below.

   2. Attributes.

      These are compile time constant values of the operation.

   3. Natural attributes.

      These attributes affect the behavior of the operations (e.g., padding for convolution);

   4. Derived attributes.

      These attributes are not needed to define the operation but are instead derived from attributes of the operation. E.g., the output shape of type. This is mostly used for convenience interface generation or interaction with other frameworks/translation.

4. Return types.

   The type of the value(s) returned by the operation.

5. Traits.

   Traits of the operations. They are operation properties that affect syntax or semantics. MLIR C++ models various traits in the mlir::OpTrait namespace. In TableGen, we have the corresponding OpTrait class to wrap around any C++ trait symbol and use it in operation definition. For example, NoSideEffect is just a definition that expands to OpTrait<"HasNoSideEffect">; having such a definition makes the trait inside TableGen more integrated and easier to parse as a declarative language.

6. Reference description.

   The description of the operation is encoded as C++ builder using EDSC. This is still under active discussion and will be fleshed out post-prototyping.

7. Custom builder method (builder).

   This is used to generate additional convenience builder methods. For example when defining a C++ builder method that has default values. There are two builder automatically generated based on the arguments and returns types (see op_base.td).

8. Custom printer method.

   The custom printer to invoke when producing the custom assembly form output.

9. Custom verifier code.

   Additional verification to perform in addition to those generated due to operands, attributes, and traits.

10. hasCanonicalizer and hasConstantFolder.

    These boolean fields indicate whether canonicalization patterns or constant folding have been defined for this operation.

For custom parsing and printing

In the operation definition the user can specify custom functions to print or parse the operation.

FIXME: Autogenerating printing/parsing has not been prototyped, and potentially just being able to specify custom printer/parser methods are sufficient. This should presumably be influenced by the design of the assembler/disassembler logic that LLVM backends get for free for machine instructions.

The custom assembly form emitter form of the operation is specified using a string with matching operation name, operands and attributes. With the ability to express additional information that needs to be parsed to build the operation:

tfl.Add $lhs, $rhs {fused_activation_function:
                   $fused_activation_function }: ${type(self)}

1. The output is never shown in the "mnemonics" string as that is fixed form and cannot be altered.

2. Custom parsing of ops may include some punctuation (e.g., parenthesis).

3. The operands/results are added to the created operation in the order that they are shown in the input and output dags.

4. The ${type(self)} operator is used to represent the type of the operator. The type of operands can also be queried.

5. Attributes names are matched to the placeholders in the mnemonic strings. E.g., attribute axis is matched with $axis. Custom parsing for attribute type can be defined along with the attribute definition.

6. The information in the custom assembly form should be sufficient to invoke the builder generated. That may require being able to propagate information (e.g., the $lhs has the same type as the result).

Printing is effectively the inverse of the parsing function generated with the mnemonic string serving as a template.

Type constraints

Constraints are along (at least) three axis: 1) elemental type, 2) rank (including static or dynamic), 3) dimensions. While some ops have no compile time fixed shape (e.g., output shape is dictated by data) we could still have some knowledge of constraints/bounds in the system for that op (e.g., the output of a tf.where is at most the size of the input data). And so there are additional valuable constraints that could be captured even without full knowledge.

Initially the shape inference will be declaratively specified using:

• Constraint on the operands of an operation directly. For example constraining the input type to be tensor/vector elements or that the elemental type be of a specific type (e.g., output of sign is of elemental type i1) or class (e.g., float like).
• Constraints on an operands of an operation. For example, enabling specifying equality constraints on type/constituents of a type (shape and elemental type) between operands and results (e.g., the output type of an add is the same as those of the input operands).

In general there is an input/output transfer function which maps the inputs to the outputs (e.g., given input X and Y [or slices thereof] with these sizes, the output is Z [or this slice thereof]). Such a function could be used to determine the output type (shape) for given input type (shape).

But shape functions are determined by attributes and could be arbitrarily complicated with a wide-range of specification possibilities. Equality relationship are common (e.g., the elemental type of the output matches the primitive type of the inputs, both inputs have exactly the same type [primitive type and shape]) and so these should be easy to specify. Algebraic relationships would also be common (e.g., a concat of [n,m] and [n,m] matrix along axis 0 is [n+n, m] matrix), while some ops only have defined shapes under certain cases (e.g., matrix multiplication of [a,b] and [c,d] is only defined if b == c). As ops are also verified, the shape inference need only specify rules for the allowed cases (e.g., shape inference for matmul can ignore the case where b != c), which would simplify type constraint specification.

Instead of specifying an additional mechanism to specify a shape transfer function, the reference implementation of the operation will be used to derive the shape function. The reference implementation is general and can support the arbitrary computations needed to specify output shapes.

Attribute definition

An attribute is a compile time known constant of an operation. Attributes are required to be known to construct an operation (e.g., the padding behavior is required to fully define the conv2d op). Attributes are defined as having a storage type (corresponding to a derived class of mlir::Attribute), a return type (that corresponds to the C++ type to use in the generation of the helper accessors) as well as method to convert between the internal storage and the helper method. Derived attributes are a special class of attributes that do not have storage but are instead calculated based on the operation and its attributes.

As with types, attributes can have a set of condition that need to be satisfied (e.g., attribute has to be floating point, has to be nonnegative, has to be in a range). This is true both in the specification of operations as well as matching rules (see DAG rewrites).

Rewrite pattern description

MLIR aims to support many graph transformations across multiple levels of representation using declarative patterns. These patterns can be expressed using TableGen as well as dynamically (TBD).

Op DAG pattern rewrites

The most direct pattern supported in MLIR is rewrites of the form (dag of operations) -> (dag of operations) along with constraints (on operands and operations). The matchers require both dialects being matched between to be included in the same TableGen file. Hence pattern matching is normally defined in either a separate file that imports both. Matchers are defined in terms of the TableGen instances rather than mnemonics to allow for better error checking and verification generation.

def : Pat<(TF_LeakyReluOp $arg, F32Attr:$a), (TFL_LeakyReluOp $arg, $a)>;
def : Pat<(TF_ReluOp (TF_AddOp $lhs, $rhs)), (TFL_AddOp $lhs, $rhs, TFL_AF_Relu)>;
def : Pat<(TF_BiasAddOp F32Tensor:$l, F32Tensor:$r),
          (TFL_AddOp $l, $r, TFL_AF_None)>;

In the above examples it was shown how to construct matching rules between two dialects (TensorFlow and TensorFlowLite), showing matching arguments (attributes and operands) as well as matching a DAG pattern of multiple input operations to single output.

1. Matchers can be partially specified on the input (e.g., not all arguments contrained) and so multiple matchers can match the same set of nodes. The most discriminative matcher (as determined by the number of constrained/matching terms) will be selected, if two patterns are equally discriminative then an error will be reported.

2. Matchers between dialects have to be completely specified on the output (i.e., there can be no unspecified attributes of the op generated).

3. Operands and attributes can be further constrained from the op definition (e.g., bias add rule only matches the case where both Tensors have F32 elements).

   1. Attributes can be transformed by transform rules to produce an attribute of a type different than the type matched.

TODO: Add constraints on the matching rules.

TODO: Describe the generation of benefit metric given pattern.