[mlir][transform] Add ForeachOp to transform dialect

[springerm]

Summary

This op "unbatches" an op handle and executes the loop body for each payload op.

Depends On D130244

[nicolasvasilache]

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

290

Seems dangerous to potentially fail silently if the IR may be left in an invalid state

mlir/test/Dialect/Transform/ops.mlir

61

nit: nl

[springerm]

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

290

The transform may fail silently but it still fails. In that case, the entire ForeachOp will fail (silently). It is my understanding that this will "silent failure" will continue to bubble up until it is caught by an op that handles the failure (e.g., the alternatives op, by restoring the pre-transform IR). If no op handles the failure, the entire transform will fail.

So I think it is OK to fail silently here and we do not have to worry about invalid IR.

[springerm]

address comments

[ftynse]

I would consider describing the default "batched" mode in https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td and mentioning this operation as a way of obtaining "unbatched" mode.

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

124

I would consider returning as many handles as transform.yield has operands, with each handle corresponding to a list of individual payload operations yielded by each iteration. This would make this op composable with functional-style composition of transformations. Not necessary for this commit.

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

290

Indeed. One of the big points of DiagnosedSilenceableFailure is to report errors in the place that has the most information and just bubble up the error. Therefore, container ops such as this one are not supposed to add extra error messages, leaf ops are, or even better the actual transformation code.

mlir/test/Dialect/Transform/test-interpreter.mlir

604–608

We need a mechanism to check that transformations are indeed applied to one op at a time. Here, the output of the "batched" mode will be indistinguishable. We have TestPrintNumberOfAssociatedPayloadIROps that could help here.

[springerm]

address comments

mlir/include/mlir/Dialect/Transform/IR/TransformDialect.td

let description = [{
     ```
 
     The values defined by operations in this dialect correspond to (groups of)
     operations in the payload IR. In the example above, `%0` corresponds to the
     set of loops found in the payload IR that satisfy the condition, and `%1`
     correspond to groups of outer and inner loops, respectively, produced by
     the tiling transformation.
 
+    A Transform IR value such as `%0` may be associated with multiple payload
+    operations. This is conceptually a set of operations and no assumptions
+    should be made about the order of ops. Most Transform IR ops support
+    operand values that are mapped to multiple operations. They usually apply
+    the respective transformation for every mapped op ("batched execution").
+    Deviations from this convention are described in the documentation of
+    Transform IR ops.
+
     Overall, Transform IR ops are expected to be contained in a single top-level
     op. Such top-level ops specify how to apply the transformations described
     by the operations they contain, e.g., `transform.sequence` executes
     transformations one by one and fails if any of them fails. Such ops are
     expected to have the `PossibleTopLevelTransformOpTrait` and may be used
     without arguments.
 
     ## Dialect Extension Mechanism

mlir/include/mlir/Dialect/Transform/IR/TransformOps.td

def AlternativesOp : TransformDialectOp<"alternatives",
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region VariadicRegion<SizedRegion<1>>:$alternatives);
 
   let assemblyFormat =
     "($scope^)? (`->` type($results)^)? attr-dict-with-keyword regions";
   let hasVerifier = 1;
 }
 
+def ForeachOp : TransformDialectOp<"foreach",
+    [DeclareOpInterfaceMethods<TransformOpInterface>,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     DeclareOpInterfaceMethods<RegionBranchOpInterface, [
+         "getSuccessorRegions", "getSuccessorEntryOperands"]>,
+     SingleBlockImplicitTerminator<"::mlir::transform::YieldOp">
+    ]> {
+  let summary = "Executes the body for each payload op";
+  let description = [{
+    This op has exactly one region with exactly one block ("body"). The body is
+    executed for each payload op that is associated to the target operand in an
+    unbatched fashion. I.e., the block argument ("iteration variable") is always
+    mapped to exactly one payload op.
+
+    This op always reads the target handle. Furthermore, it consumes the handle
+    if there is a transform op in the body that consumes the iteration variable.
+    This op does not return anything.
+
+    The transformations inside the body are applied in order of their
+    appearance. During application, if any transformation in the sequence fails,
+    the entire sequence fails immediately leaving the payload IR in potentially
+    invalid state, i.e., this operation offers no transformation rollback
+    capabilities.
+  }];
+
+  let arguments = (ins PDL_Operation:$target);
+  let results = (outs);
+  let regions = (region SizedRegion<1>:$body);
+  let assemblyFormat = "$target $body attr-dict";
+
+  let extraClassDeclaration = [{
+    /// Allow the dialect prefix to be omitted.
+    static StringRef getDefaultDialect() { return "transform"; }
+
+    BlockArgument getIterationVariable() {
+      return getBody().front().getArgument(0);
+    }
+  }];
+}
+
 def GetClosestIsolatedParentOp : TransformDialectOp<"get_closest_isolated_parent",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
   let summary = "Gets handles to the closest isolated-from-above parents";
   let description = [{
     The handles defined by this Transform op correspond to the closest isolated
     from above ancestor of the Payload IR operations associated with its
     operand. If any of the given Payload IR ops has no such parent (unlikely as

mlir/lib/Dialect/Transform/IR/TransformOps.cpp

if (terminator->getOperands().getTypes() != getResults().getTypes()) {
       return diag;
     }
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
+// ForeachOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::ForeachOp::apply(transform::TransformResults &results,
+                            transform::TransformState &state) {
+  ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  for (Operation *op : payloadOps) {
+    auto scope = state.make_region_scope(getBody());
+    if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
+      return DiagnosedSilenceableFailure::definiteFailure();
+
+    for (Operation &transform : getBody().front().without_terminator()) {
+      DiagnosedSilenceableFailure result = state.applyTransform(
+          cast<transform::TransformOpInterface>(transform));
+      if (!result.succeeded())
+        return result;
+    }
+  }
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::ForeachOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  BlockArgument iterVar = getIterationVariable();
+  if (any_of(getBody().front().without_terminator(), [&](Operation &op) {
+        return isHandleConsumed(iterVar, cast<TransformOpInterface>(&op));
+      })) {
+    consumesHandle(getTarget(), effects);
+  } else {
+    onlyReadsHandle(getTarget(), effects);
+  }
+}
+
+void transform::ForeachOp::getSuccessorRegions(
+    Optional<unsigned> index, ArrayRef<Attribute> operands,
+    SmallVectorImpl<RegionSuccessor> &regions) {
+  Region *bodyRegion = &getBody();
+  if (!index) {
+    regions.emplace_back(bodyRegion, bodyRegion->getArguments());
+    return;
+  }
+
+  // Branch back to the region or the parent.
+  assert(*index == 0 && "unexpected region index");
+  regions.emplace_back(bodyRegion, bodyRegion->getArguments());
+  regions.emplace_back();
+}
+
+OperandRange
+transform::ForeachOp::getSuccessorEntryOperands(Optional<unsigned> index) {
+  // The iteration variable op handle is mapped to a subset (one op to be
+  // precise) of the payload ops of the ForeachOp operand.
+  assert(index && *index == 0 && "unexpected region index");
+  return getOperation()->getOperands();
+}
+
+//===----------------------------------------------------------------------===//
 // GetClosestIsolatedParentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::GetClosestIsolatedParentOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Operation *parent =

mlir/test/Dialect/Transform/ops-invalid.mlir

 
 // -----
 
 // expected-error @below {{expects the entry block to have one argument of type '!pdl.operation'}}
 transform.alternatives {
 ^bb0:
   transform.yield
 }
+
+// -----
+
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  // expected-error @below {{result #0 has more than one potential consumer}}
+  %0 = test_produce_param_or_forward_operand 42
+  // expected-note @below {{used here as operand #0}}
+  transform.foreach %0 {
+  ^bb1(%arg1: !pdl.operation):
+    transform.test_consume_operand %arg1
+  }
+  // expected-note @below {{used here as operand #0}}
+  transform.test_consume_operand %0
+}

mlir/test/Dialect/Transform/ops.mlir

^bb0(%arg0: !pdl.operation):
   } : !pdl.operation
   transform.sequence %0 {
   ^bb2(%arg2: !pdl.operation):
   }
   transform.sequence %0 {
   ^bb3(%arg3: !pdl.operation):
   }
 }
+
+// CHECK: transform.sequence
+// CHECK: foreach
+transform.sequence {
+^bb0(%arg0: !pdl.operation):
+  transform.foreach %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+  }
+}

mlir/test/Dialect/Transform/test-interpreter.mlir

^bb0(%arg1: !pdl.operation):
       // expected-remark @below {{2}}
       test_print_number_of_associated_payload_ir_ops %1
       %2 = replicate num(%0) %1
       // expected-remark @below {{4}}
       test_print_number_of_associated_payload_ir_ops %2
     }
   }
 }
+
+// -----
+
+func.func @bar() {
+  // expected-remark @below {{transform applied}}
+  %0 = arith.constant 0 : i32
+  // expected-remark @below {{transform applied}}
+  %1 = arith.constant 1 : i32
+  return
+}
+
+transform.with_pdl_patterns {
+^bb0(%arg0: !pdl.operation):
+  pdl.pattern @const : benefit(1) {
+    %r = pdl.types
+    %0 = pdl.operation "arith.constant" -> (%r : !pdl.range<type>)
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %f = pdl_match @const in %arg1
+    transform.foreach %f {
+    ^bb2(%arg2: !pdl.operation):
+      // expected-remark @below {{1}}
+      transform.test_print_number_of_associated_payload_ir_ops %arg2
+      transform.test_print_remark_at_operand %arg2, "transform applied"
+    }
+  }
+}