[mlir][transform] Support results on ForeachOp

[springerm]

Summary

Handles can be yielded from the ForeachOp.

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

let description = [{
     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.
+
+    This op generates as many handles as the terminating YieldOp has operands.
+    For each result, the payload ops of the corresponding YieldOp operand are
+    merged and mapped to the same resulting handle.
   }];
 
   let arguments = (ins PDL_Operation:$target);
-  let results = (outs);
+  let results = (outs Variadic<PDL_Operation>:$results);
   let regions = (region SizedRegion<1>:$body);
-  let assemblyFormat = "$target $body attr-dict";
+  let assemblyFormat = "$target (`->` type($results)^)? $body attr-dict";
+  let hasVerifier = 1;
 
   let extraClassDeclaration = [{
     /// Allow the dialect prefix to be omitted.
     static StringRef getDefaultDialect() { return "transform"; }
 
     BlockArgument getIterationVariable() {
       return getBody().front().getArgument(0);
     }
+
+    transform::YieldOp getYieldOp();
   }];
 }
 
 def GetClosestIsolatedParentOp : TransformDialectOp<"get_closest_isolated_parent",
     [DeclareOpInterfaceMethods<TransformOpInterface>,
      NavigationTransformOpTrait, MemoryEffectsOpInterface]> {
   let summary = "Gets handles to the closest isolated-from-above parents";
   let description = [{

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

 //===----------------------------------------------------------------------===//
 // ForeachOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::ForeachOp::apply(transform::TransformResults &results,
                             transform::TransformState &state) {
   ArrayRef<Operation *> payloadOps = state.getPayloadOps(getTarget());
+  SmallVector<SmallVector<Operation *>> resultOps(getNumResults(), {});
+
   for (Operation *op : payloadOps) {
     auto scope = state.make_region_scope(getBody());
     if (failed(state.mapBlockArguments(getIterationVariable(), {op})))
       return DiagnosedSilenceableFailure::definiteFailure();
 
+    // Execute loop body.
     for (Operation &transform : getBody().front().without_terminator()) {
       DiagnosedSilenceableFailure result = state.applyTransform(
           cast<transform::TransformOpInterface>(transform));
       if (!result.succeeded())
         return result;
     }
+
+    // Append yielded payload ops to result list (if any).
+    for (unsigned i = 0; i < getNumResults(); ++i) {
+      ArrayRef<Operation *> yieldedOps =
+          state.getPayloadOps(getYieldOp().getOperand(i));
+      resultOps[i].append(yieldedOps.begin(), yieldedOps.end());
+    }
   }
+
+  for (unsigned i = 0; i < getNumResults(); ++i)
+    results.set(getResult(i).cast<OpResult>(), resultOps[i]);
+
   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);
   }
+
+  for (Value result : getResults())
+    producesHandle(result, 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());

 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();
 }
 
+transform::YieldOp transform::ForeachOp::getYieldOp() {
+  return cast<transform::YieldOp>(getBody().front().getTerminator());
+}
+
+LogicalResult transform::ForeachOp::verify() {
+  auto yieldOp = getYieldOp();
+  if (getNumResults() != yieldOp.getNumOperands())
+    return emitOpError() << "expects the same number of results as the "
+                            "terminator has operands";
+  for (Value v : yieldOp.getOperands())
+    if (!v.getType().isa<pdl::OperationType>())
+      return yieldOp->emitOpError("expects only PDL_Operation operands");
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // GetClosestIsolatedParentOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::GetClosestIsolatedParentOp::apply(
     transform::TransformResults &results, transform::TransformState &state) {
   SetVector<Operation *> parents;
   for (Operation *target : state.getPayloadOps(getTarget())) {

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

^bb1(%arg1: !pdl.operation):
     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"
     }
   }
 }
+
+// -----
+
+func.func @bar() {
+  scf.execute_region {
+    // expected-remark @below {{transform applied}}
+    %0 = arith.constant 0 : i32
+    scf.yield
+  }
+
+  scf.execute_region {
+    // expected-remark @below {{transform applied}}
+    %1 = arith.constant 1 : i32
+    // expected-remark @below {{transform applied}}
+    %2 = arith.constant 2 : i32
+    scf.yield
+  }
+
+  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"
+  }
+
+  pdl.pattern @execute_region : benefit(1) {
+    %r = pdl.types
+    %0 = pdl.operation "scf.execute_region" -> (%r : !pdl.range<type>)
+    pdl.rewrite %0 with "transform.dialect"
+  }
+
+  transform.sequence %arg0 {
+  ^bb1(%arg1: !pdl.operation):
+    %f = pdl_match @execute_region in %arg1
+    %results = transform.foreach %f -> !pdl.operation {
+    ^bb2(%arg2: !pdl.operation):
+      %g = transform.pdl_match @const in %arg2
+      transform.yield %g : !pdl.operation
+    }
+
+    // expected-remark @below {{3}}
+    transform.test_print_number_of_associated_payload_ir_ops %results
+    transform.test_print_remark_at_operand %results, "transform applied"
+  }
+}