Chapter 4: Matching Payload with Transform Operations

Check the continuously-tested version of MLIR files under mlir/test/Examples/transform/Ch4.

Up until now, we were applying transform dialect scripts under the assumption that specific payload operations are identified by the caller when the transform dialect interpreter is invoked. This may be seen as contrary to the idea of driving transformations from a dialect since the transformation targets must be identified through mechanisms external to the transform dialect interpreter, for example, when invoking the interpreter programmatically in C++ or through pass arguments as seen in previous chapters. It also adds practical overhead due to increased interaction with the interpreter in C++, and cognitive overhead of manipulating two interfaces at once. To remedy this, Transform dialect proposes a subset of operations for matching payload operations that need to be transformed.

Match operations are simply transform operations with some additional guarantees. In particular, they are not expected to modify the payload IR and are expected to fail if their operands (typically payload operation handles) are not associated with payload IR objects having desired properties, such as operation names or kinds of arguments. Using simple combinator operations, it becomes possible to set up a higher-level match and rewrite infrastructure directly within the transform dialect.

Simple match ¶

Let us reconsider the “fully connected layer” example from Chapter 1, reproduced below for convenience.

// Original function to optimize.
func.func @fc_relu(%lhs: tensor<512x512xf32>, %rhs: tensor<512x512xf32>,
                   %bias: tensor<512x512xf32>, %output: tensor<512x512xf32>)
                   -> tensor<512x512xf32> {
  // Matrix-matrix multiplication.
  %matmul = linalg.matmul
            ins(%lhs, %rhs: tensor<512x512xf32>, tensor<512x512xf32>)
            outs(%output: tensor<512x512xf32>) -> tensor<512x512xf32>

  // Elementwise addition.
  %biased = linalg.elemwise_binary { fun = #linalg.binary_fn<add> }
    ins(%matmul, %bias : tensor<512x512xf32>, tensor<512x512xf32>)
    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>

  // Elementwise max with 0 (ReLU).
  %c0f = arith.constant 0.0 : f32
  %relued = linalg.elemwise_binary { fun = #linalg.binary_fn<max_signed> }
    ins(%biased, %c0f : tensor<512x512xf32>, f32)
    outs(%output : tensor<512x512xf32>) -> tensor<512x512xf32>
  func.return %relued : tensor<512x512xf32>
}

In Chapter 1, we were calling the test transform interpreter pass with additional arguments, bind-first-extra-to-ops=linalg.matmul bind-second-extra-to-ops=linalg.elemwise_binary, to provide initial associations for operation handles. Instead, we can use match operations to discover relevant operations in the payload IR. Match operations can be combined with “regular” transform operations using, e.g., the transform.collect_matching combinator operation that leverages the concept of named sequences to organize matchers.

// The module containing named sequences must have an attribute allowing them
// to enable verification.
module @transforms attributes { transform.with_named_sequence } {
  // Entry point. This takes as the only argument the root operation (typically
  // pass root) given to the transform interpreter.
  transform.named_sequence @__transform_main(
      %root: !transform.any_op {transform.readonly}) {
    // Collect operations that match the criteria specified in named sequence.
    // If the named sequence fails with a silenceable failure, silences it (the
    // message is forwarded to the debug stream). If the named sequence
    // succeeds, appends its results to the results of this operation.
    %elemwise = transform.collect_matching @match_elemwise in %root
      : (!transform.any_op) -> !transform.any_op
    %matmul = transform.collect_matching @match_matmul in %root
      : (!transform.any_op) -> !transform.any_op
    transform.include @print_elemwise failures(propagate)  (%elemwise)
      : (!transform.any_op) -> ()
    transform.include @print_matmul failures(propagate)  (%matmul)
      : (!transform.any_op) -> ()

    transform.yield
  }

  // This is a matcher sequence. It is given an operation to match and the
  // match is considered successful unless any nested operation produces a
  // failure. The values yielded by this operation will be forwarded to the
  // rewriter sequence on success.
  transform.named_sequence @match_elemwise(
      %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
    transform.match.operation_name %entry ["linalg.elemwise_binary"]
      : !transform.any_op
    transform.yield %entry : !transform.any_op
  }
  transform.named_sequence @match_matmul(
      %entry: !transform.any_op {transform.readonly}) -> !transform.any_op {
    transform.match.operation_name %entry ["linalg.matmul"] : !transform.any_op
    transform.yield %entry : !transform.any_op
  }

  // This is a rewriter sequence.
  transform.named_sequence @print_elemwise(
      %elemwise_binary: !transform.any_op {transform.readonly}) {
    transform.debug.emit_remark_at
      %elemwise_binary, "elementwise binary" : !transform.any_op
    transform.yield
  }
  transform.named_sequence @print_matmul(
      %matmul: !transform.any_op {transform.readonly}) {
    transform.debug.emit_remark_at %matmul, "matmul" : !transform.any_op
    transform.yield
  }
}

This script can be executed using the non-test interpreter pass running on the root operation of the translation unit without additional flags: mlir-opt --transform-interpreter. It will emit corresponding remarks at linalg.elemwise_binary and linalg.matmul operations. In debug builds, the infrastructure provides a convenient method to understand the matching process by passing -debug-only=transform-matcher to mlir-opt or a derived tool. It will print the silenceable failure messages produced by the match operations into the debug stream, for example:

<...>
[transform-matcher] matching %0 = linalg.matmul ins(%arg0, %arg1 : tensor<512x512xf32>, tensor<512x512xf32>) outs(%arg3 : tensor<512x512xf32>) -> tensor<512x512xf32> @0x5622eee08410
[transform-matcher] matcher match_elemwise failed: wrong operation name
<...>

This is now sufficient to run the rest of the transform script from Chapter 1, substituting %arg1 with %matmul and %arg2 with %elemwise.

Matching Chains of Operations ¶

The matcher above remains naive as it matches all operations of the certain kind under the payload root. These operations may or may not be related, and may, for example, belong to different functions. Even if they are in a single function, if there are multiple groups of such operations, we wouldn’t be able to differentiate them with this approach. In reality, we want to match a specific group of operations where a matmul operation produces a result that is used by an elementwise operation, which in turn feeds another elementwise operation in a similar way.

This can be achieved using the following matcher sequence.

// This is also a matcher sequence. It is similarly given an operation to
// match and nested operations must succeed in order for a match to be deemed
// successful. It starts matching from the last operation in the use-def chain
// and goes back because each operand (use) has exactly one definition.
transform.named_sequence @match_matmul_elemwise(
    %last: !transform.any_op {transform.readonly})
    -> (!transform.any_op, !transform.any_op, !transform.any_op) {
  // The last operation must be an elementwise binary.
  transform.match.operation_name %last ["linalg.elemwise_binary"]
    : !transform.any_op
  // Its first operand must be defined by another operation, to which we
  // will get a handle here. We are guaranteed that the first operand exists
  // because we know the operation is binary, but even in absence of such a
  // guarantee, this operation would have produced a silenceable failure when
  // `%last` does not have enough operands.
  %middle = transform.get_producer_of_operand %last[0]
    : (!transform.any_op) -> !transform.any_op
  // The defining operation must itself be an elementwise binary.
  transform.match.operation_name %middle ["linalg.elemwise_binary"]
    : !transform.any_op
  // And the first operand of that operation must be defined by yet another
  // operation.
  %matmul = transform.get_producer_of_operand %middle[0]
    : (!transform.any_op) -> !transform.any_op
  // And that operation is a matmul.
  transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op
  // We will yield the handles to the matmul and the two elementwise
  // operations separately.
  transform.yield %matmul, %middle, %last
    : !transform.any_op, !transform.any_op, !transform.any_op
}

This matcher is applicable in presence of other elemwise and matmul operations and will return the triple of related operations rather than operations in the order in which they are found. It can be exercised similarly to the previous incarnation, as follows.

// Alternative entry point.
transform.named_sequence @__transform_main(
    %root: !transform.any_op {transform.readonly}) {
  // Collect groups of operations that match the criteria specified in the
  // named sequence.
  %matmul, %el1, %el2 = transform.collect_matching @match_matmul_elemwise in %root
    : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
  %elemwise = transform.merge_handles %el1, %el2 : !transform.any_op

  transform.include @print_elemwise failures(propagate)  (%elemwise)
    : (!transform.any_op) -> ()
  transform.include @print_matmul failures(propagate)  (%matmul)
    : (!transform.any_op) -> ()

  transform.yield
}

Defining Match Operations ¶

The matcher of a chain of operations is correct in presence of other operations, but is still insufficiently robust for many cases of interest. In particular, using transform.get_producer_of_operand %last[0] requires that the first operand of elementwise operations is produced by another operation. The same transformation strategy may however apply regardless of the operand position: many binary operations are associative. Let us use this opportunity to introduce a new match operation. Specifically, we would like this operation to succeed if any of the operands satisfies certain conditions that can be expressed as other match operations. We also want it to return some of the state and the position of the matched operand in the operand list.

Match operations are defined similarly to other transform operations, with the only difference of additionally implementing the MatchOpInterface. Note that this interface has no additional methods (though it may add some eventually) and is only used as a verification contract that the operation is intended for matching and will not attempt to transform the payload. The minimal definition of our operation is as follows.

// Define the new operation. By convention, prefix its name with `match`
// followed by the name of the dialect extension.
def HasOperandSatisfyingOp : TransformDialectOp<"match.my.has_operand_satisfying",
    [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
     DeclareOpInterfaceMethods<TransformOpInterface>,
     // Indicate that the operation implements MatchOpInterface in addition to
     // the TransformOpInterface. This interface is only used as a tag at this
     // point and has no methods that are mandatory to implement.
     MatchOpInterface,
     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">]> {
  let summary = "Succeed if any of the operands matches all nested criteria";
  let arguments = (ins TransformHandleTypeInterface:$op);
  let results = (outs TransformParamTypeInterface:$position,
                      Variadic<Transform_AnyHandleOrParamType>:$results);

  // Match operations can be arbitrarily complex, e.g., containing regions.
  let regions = (region SizedRegion<1>:$body);
  let hasVerifier = 1;
  let assemblyFormat = [{
    $op `:` functional-type($op, results) attr-dict-with-keyword $body
  }];
}

It takes as argument the handle associated with the payload operations whose operands it will match, has an associated single-block region containing the match criteria, and returns the position of the matched operand as well as any other transform value yielded from the body on the successful match.

The matching logic is implemented in the apply method of the TransformOpInterface and is easily composable with other transform operations. All facilities for managing the interpreter state and recursively entering the blocks are available in the same way as they are for “regular” transform operations. Match operations are expected to return a silenceable failure to indicate failure to match, and to immediately propagate definite failures. If they have nested operations, they are expected to handle and, in most cases, silence the silenceable failures produced when applying those operations. For our operation, the matching is essentially a loop iterating over all operands of the (single) payload operation and applying nested transform ops until they all succeed for one of the operands.

// Matcher ops implement `apply` similarly to other transform ops. They are not
// expected to modify payload, but use the tri-state result to signal failure or
// success to match, as well as potential irrecoverable errors.
mlir::DiagnosedSilenceableFailure
mlir::transform::HasOperandSatisfyingOp::apply(
    mlir::transform::TransformRewriter &rewriter,
    mlir::transform::TransformResults &results,
    mlir::transform::TransformState &state) {
  // For simplicity, only handle a single payload op. Actual implementations
  // can use `SingleOpMatcher` trait to simplify implementation and document
  // this expectation.
  auto payloadOps = state.getPayloadOps(getOp());
  if (!llvm::hasSingleElement(payloadOps))
    return emitSilenceableError() << "expected single payload";

  // Iterate over all operands of the payload op to see if they can be matched
  // using the body of this op.
  Operation *payload = *payloadOps.begin();
  for (OpOperand &operand : payload->getOpOperands()) {
    // Create a scope for transform values defined in the body. This corresponds
    // to the syntactic scope of the region attached to this op. Any values
    // associated with payloads from now on will be automatically dissociated
    // when this object is destroyed, i.e. at the end of the iteration.
    // Associate the block argument handle with the operand.
    auto matchScope = state.make_region_scope(getBody());
    if (failed(state.mapBlockArgument(getBody().getArgument(0),
                                      {operand.get()}))) {
      return DiagnosedSilenceableFailure::definiteFailure();
    }

    // Iterate over all nested matchers with the current mapping and see if they
    // succeed.
    bool matchSucceeded = true;
    for (Operation &matcher : getBody().front().without_terminator()) {
      // Matcher ops are applied similarly to any other transform op.
      DiagnosedSilenceableFailure diag =
          state.applyTransform(cast<TransformOpInterface>(matcher));

      // Definite failures are immediately propagated as they are irrecoverable.
      if (diag.isDefiniteFailure())
        return diag;

      // On success, keep checking the remaining conditions.
      if (diag.succeeded())
        continue;

      // Report failure-to-match for debugging purposes and stop matching this
      // operand.
      assert(diag.isSilenceableFailure());
      DEBUG_MATCHER(DBGS_MATCHER()
                    << "failed to match operand #" << operand.getOperandNumber()
                    << ": " << diag.getMessage());
      (void)diag.silence();
      matchSucceeded = false;
      break;
    }
    // If failed to match this operand, try other operands.
    if (!matchSucceeded)
      continue;

    // If we reached this point, the matching succeeded for the current operand.
    // Remap the values associated with terminator operands to be associated
    // with op results, and also map the parameter result to the operand's
    // position. Note that it is safe to do here despite the end of the scope
    // as `results` are integrated into `state` by the interpreter after `apply`
    // returns rather than immediately.
    SmallVector<SmallVector<MappedValue>> yieldedMappings;
    transform::detail::prepareValueMappings(
        yieldedMappings, getBody().front().getTerminator()->getOperands(),
        state);
    results.setParams(getPosition().cast<OpResult>(),
                      {rewriter.getI32IntegerAttr(operand.getOperandNumber())});
    for (auto &&[result, mapping] : llvm::zip(getResults(), yieldedMappings))
      results.setMappedValues(result, mapping);
    return DiagnosedSilenceableFailure::success();
  }

  // If we reached this point, none of the operands succeeded the match.
  return emitSilenceableError()
         << "none of the operands satisfied the conditions";
}

By convention, operations implementing MatchOpInterface must not modify payload IR and must therefore specify that they only read operand handles and payload as their effects.

void transform::CollectMatchingOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  onlyReadsHandle(getRoot(), effects);
  producesHandle(getResults(), effects);
  onlyReadsPayload(effects);
}

This operation can now be included in a transform dialect extension, loaded and used in our matcher. Specifically, we will use it to indicate that either of the operands of the “max” elementwise operation in our example can be produced by the previous elementwise operation. The previous operation will still require the matmul to produce the first operand for simplicity. The updated matcher sequence looks as follows.

transform.named_sequence @match_matmul_elemwise(
    %last: !transform.any_op {transform.readonly})
    -> (!transform.any_op, !transform.any_op, !transform.any_op,
        !transform.param<i32>) {
  // The last operation must be an elementwise binary.
  transform.match.operation_name %last ["linalg.elemwise_binary"]
    : !transform.any_op

  // One of its operands must be defined by another operation, to which we
  // will get a handle here. This is achieved thanks to a newly defined
  // operation that tries to match operands one by one using the match
  // operations nested in its region.
  %pos, %middle = transform.match.my.has_operand_satisfying %last
      : (!transform.any_op) -> (!transform.param<i32>, !transform.any_op) {
  ^bb0(%operand: !transform.any_value):
    // The operand must be defined by an operation.
    %def = transform.get_defining_op %operand
      : (!transform.any_value) -> !transform.any_op
    // The defining operation must itself be an elementwise binary.
    transform.match.operation_name %def ["linalg.elemwise_binary"]
      : !transform.any_op
    transform.yield %def : !transform.any_op
  }

  // And the first operand of that operation must be defined by yet another
  // operation.
  %matmul = transform.get_producer_of_operand %middle[0]
    : (!transform.any_op) -> !transform.any_op
  // And that operation is a matmul.
  transform.match.operation_name %matmul ["linalg.matmul"] : !transform.any_op
  // We will yield the handles to the matmul and the two elementwise
  // operations separately.
  transform.yield %matmul, %middle, %last, %pos
    : !transform.any_op, !transform.any_op, !transform.any_op,
      !transform.param<i32>
}

This achieves the desired effect and matches both max(add(matmul(...), bias), 0) and max(0, add(matmul(...), bias)) in the same values. The %pos value is a transform dialect parameter, which is used to store lists of entities known to be constant throughout the transform application. Most often, parameters are numeric values, but they can generally be any MLIR attributes.

In order to demonstrate that groups of operations are matched independently of each other, let us use the transform.foreach_match operation that allows one to implement a simple high-level pattern rewriting approach within the transform dialect (for advanced or lower-level pattern rewriting, consider PDL(L) or C++ rewriting APIs). It maps a matcher named sequence to an action named sequence, and the latter gets invoked whenever the former succeeds.

// Traverses the payload IR associated with the operand handle, invoking
// @match_matmul_elemwise on each of the operations. If the named sequence
// succeeds, i.e., if none of the nested match (transform) operations
// produced a silenceable failure, invokes @print_matmul_elemwise and
// forwards the values yielded as arguments of the new invocation. If the
// named sequence fails with a silenceable failure, silences it (the message
// is forwarded to the debug stream). Definite failures are propagated
// immediately and unconditionally, as usual.
transform.foreach_match in %root
  @match_matmul_elemwise -> @print_matmul_elemwise
  : (!transform.any_op) -> !transform.any_op

The @print_matmul_elemwise named sequence, available in multiple.mlir, will use the parameter with the position of the operand to differentiate the two groups.

Matchers for Inferred Features ¶

The matcher sequences described above, although useful to drive transformations from within the transform dialect interpreter, are rather basic since they mostly rely on operation names and use-def chains. Alternative implementations using APIs or various declarative rewrite rules are barely less expressive and sometimes more concise. The real power of transform dialect matcher ops lies in the possibility to define matchers of inferred properties of payloads, i.e., properties that are not directly accessible as an attribute of an operation or any straightforward relation between IR components.

The utility of such matchers can be easily demonstrated by slightly modifying our original example. If matrix multiplication is expressed as a special case of tensor contraction using linalg.generic instead of linalg.matmul, the operation name-based matcher no longer applies. Yet such a representation is very common and can appear both in the original input and during the course of transformation, e.g., where a higher-dimensional contraction is decomposed into loops around a matrix multiplication.

In order to be a (potentially transposed) matrix multiplication, the linalg.generic operation must have the following features:

Total rank of 3.
Two inputs accessed as projected permutation of iteration dimensions.
One output accessed as projected permutation of iteration dimensions.
Iteration dimensions can be subdivided into LHS parallel, RHS parallel and reduction dimensions.
The body block consists of a multiplication and an addition.

Most of these features can be derived from the properties of the operation, e.g., the total rank corresponds to the number of entries in the iterators attribute, but almost none of them are immediately accessible in the IR or in any declarative form, which is usually limited to checking the presence or the exact match of an attribute or a type. The transform dialect allows these features to be implemented in the apply method of a matcher op and reused across multiple matching cases. For structured linear algebra payload operations, many such match operations are readily available in the structured extension. They are sufficient to implement a matrix multiplication matcher using the features listed above almost verbatim.

transform.named_sequence @match_generic_matmul(
    %candidate: !transform.any_op {transform.readonly}) -> !transform.any_op {
  // Match a structured linear algebra operation.
  transform.match.structured %candidate : !transform.any_op {
  ^bb0(%c: !transform.any_op):
    // With a rank equal to 3.
    %rank = transform.match.structured.rank %c
      : (!transform.any_op) -> !transform.param<i64>
    %c3 = transform.param.constant 3 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %rank, %c3 : !transform.param<i64>

    // With 2 inputs.
    %n_ins = transform.match.structured.num_inputs %c
      : (!transform.any_op) -> !transform.param<i64>
    %c2 = transform.param.constant 2 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_ins, %c2 : !transform.param<i64>

    // With 1 output (note that structured ops in destination passing style
    // has as many inits as outputs).
    %n_inits = transform.match.structured.num_inits %c
      : (!transform.any_op) -> !transform.param<i64>
    %c1 = transform.param.constant 1 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_inits, %c1 : !transform.param<i64>

    // All inputs and inits are accessed with a projected permutation.
    transform.match.structured.input %c[all] {projected_permutation}
      : !transform.any_op
    transform.match.structured.init %c[0] {projected_permutation}
      : !transform.any_op

    // The body is a mulf/addf contraction with appropriate dimensions.
    transform.match.structured.body %c
      { contraction = ["arith.mulf", "arith.addf"] } : !transform.any_op
    %batch, %lhs, %rhs, %reduction =
    transform.match.structured.classify_contraction_dims %c
      : (!transform.any_op)
      -> (!transform.param<i64>, !transform.param<i64>, !transform.param<i64>,
          !transform.param<i64>)


    // There is one of lhs, rhs and reduction dimensions and zero batch
    // dimensions.
    %n_batch = transform.num_associations %batch
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_lhs = transform.num_associations %lhs
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_rhs = transform.num_associations %rhs
      : (!transform.param<i64>) -> !transform.param<i64>
    %n_reduction = transform.num_associations %reduction
      : (!transform.param<i64>) -> !transform.param<i64>
    %c0 = transform.param.constant 0 : i64 -> !transform.param<i64>
    transform.match.param.cmpi eq %n_batch, %c0 : !transform.param<i64>
    transform.match.param.cmpi eq %n_lhs, %c1 : !transform.param<i64>
    transform.match.param.cmpi eq %n_rhs, %c1 : !transform.param<i64>
    transform.match.param.cmpi eq %n_reduction, %c1 : !transform.param<i64>
  }
  transform.yield %candidate : !transform.any_op
}

While this example leverages the contraction-specific matchers that have a rather non-trivial C++ implementation, the transform dialect is sufficiently flexible to implement this reasoning directly if desired. One could, for example, obtain the access map of each input as a parameter and extract the accessed dimensions as other parameters that can be compared with each other to ensure the subscripts are m,k for LHS, k,n for RHS and m,n for the init/result given the m,n,k notation for loops.

Appendix: Autogenerated Documentation ¶

`transform.match.my.has_operand_satisfying` (transform::HasOperandSatisfyingOp) ¶

Succeed if any of the operands matches all nested criteria

Syntax:

operation ::= `transform.match.my.has_operand_satisfying` $op `:` functional-type($op, results) attr-dict-with-keyword $body

Traits: SingleBlockImplicitTerminator<::mlir::transform::YieldOp>, SingleBlock

Interfaces: MatchOpInterface, MemoryEffectOpInterface, TransformOpInterface

Operands: ¶

Operand	Description
`op`	TransformHandleTypeInterface instance

Results: ¶

Result	Description
`position`	TransformParamTypeInterface instance
`results`	variadic of any transform handle or parameter