[mlir] make transform.loop.outline also return the call handle

[ftynse]

Summary

Outlining is particularly interesting when the outlined function is
replaced with something else, e.g., a microkernel. It is good to have a
handle to the call in this case.

mlir/include/mlir/Dialect/SCF/TransformOps/SCFTransformOps.td

let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 }
 
 def LoopOutlineOp : Op<Transform_Dialect, "loop.outline",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      DeclareOpInterfaceMethods<TransformOpInterface>]> {
   let summary = "Outlines a loop into a named function";
   let description = [{
-     Moves the loop into a separate function with the specified name and
-     replaces the loop in the Payload IR with a call to that function. Takes
-     care of forwarding values that are used in the loop as function arguments.
-     If the operand is associated with more than one loop, each loop will be
-     outlined into a separate function. The provided name is used as a _base_
-     for forming actual function names following SymbolTable auto-renaming
-     scheme to avoid duplicate symbols. Expects that all ops in the Payload IR
-     have a SymbolTable ancestor (typically true because of the top-level
-     module). Returns the handle to the list of outlined functions in the same
-     order as the operand handle.
+    Moves the loop into a separate function with the specified name and replaces
+    the loop in the Payload IR with a call to that function. Takes care of
+    forwarding values that are used in the loop as function arguments. If the
+    operand is associated with more than one loop, each loop will be outlined
+    into a separate function. The provided name is used as a _base_ for forming
+    actual function names following `SymbolTable` auto-renaming scheme to avoid
+    duplicate symbols. Expects that all ops in the Payload IR have a
+    `SymbolTable` ancestor (typically true because of the top-level module).
+
+    #### Return Modes
+
+    Returns a handle to the list of outlined functions and a handle to the
+    corresponding function call operations in the same order as the operand
+    handle.
+
+    Produces a definite failure if outlining failed for any of the targets.
   }];
 
   // Note that despite the name of the transform operation and related utility
   // functions, the actual implementation does not require the operation to be
   // a loop.
   let arguments = (ins TransformHandleTypeInterface:$target,
                    StrAttr:$func_name);
-  let results = (outs TransformHandleTypeInterface:$transformed);
+  let results = (outs TransformHandleTypeInterface:$function,
+                      TransformHandleTypeInterface:$call);
 
   let assemblyFormat =
     "$target attr-dict `:` functional-type(operands, results)";
 }
 
 def LoopPeelOp : Op<Transform_Dialect, "loop.peel",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {

mlir/lib/Dialect/SCF/TransformOps/SCFTransformOps.cpp

static scf::ExecuteRegionOp wrapInExecuteRegion(RewriterBase &b,
   }
   b.replaceOp(op, executeRegionOp.getResults());
   return executeRegionOp;
 }
 
 DiagnosedSilenceableFailure
 transform::LoopOutlineOp::apply(transform::TransformResults &results,
                                 transform::TransformState &state) {
-  SmallVector<Operation *> transformed;
+  SmallVector<Operation *> functions;
+  SmallVector<Operation *> calls;
   DenseMap<Operation *, SymbolTable> symbolTables;
   for (Operation *target : state.getPayloadOps(getTarget())) {
     Location location = target->getLoc();
     Operation *symbolTableOp = SymbolTable::getNearestSymbolTable(target);
     IRRewriter rewriter(getContext());
     scf::ExecuteRegionOp exec = wrapInExecuteRegion(rewriter, target);
     if (!exec) {
       DiagnosedSilenceableFailure diag = emitSilenceableError()
for (Operation *target : state.getPayloadOps(getTarget())) {
 
     if (symbolTableOp) {
       SymbolTable &symbolTable =
           symbolTables.try_emplace(symbolTableOp, symbolTableOp)
               .first->getSecond();
       symbolTable.insert(*outlined);
       call.setCalleeAttr(FlatSymbolRefAttr::get(*outlined));
     }
-    transformed.push_back(*outlined);
+    functions.push_back(*outlined);
+    calls.push_back(call);
   }
-  results.set(getTransformed().cast<OpResult>(), transformed);
+  results.set(getFunction().cast<OpResult>(), functions);
+  results.set(getCall().cast<OpResult>(), calls);
   return DiagnosedSilenceableFailure::success();
 }
 
 //===----------------------------------------------------------------------===//
 // LoopPeelOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure

mlir/python/mlir/dialects/_loop_transform_ops_ext.py

):
     )
 
 
 class LoopOutlineOp:
   """Extension for LoopOutlineOp."""
 
   def __init__(
       self,
-      result_type: Type,
+      function_type: Type,
+      call_type: Type,
       target: Union[Operation, Value],
       *,
       func_name: Union[str, StringAttr],
       ip=None,
       loc=None,
   ):
     super().__init__(
-        result_type,
+        function_type,
+        call_type,
         _get_op_result_or_value(target),
         func_name=(func_name if isinstance(func_name, StringAttr) else
                    StringAttr.get(func_name)),
         ip=ip,
         loc=loc,
     )
 
 

mlir/test/Dialect/SCF/transform-ops-invalid.mlir

func.func @loop_outline_op_multi_region() {
   ^bb0:
     func.call @body() : () -> ()
     scf.yield
   }
   return
 }
 
 transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["scf.while"]} in %arg1 : (!pdl.operation) -> !pdl.operation
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["scf.while"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   // expected-error @below {{failed to outline}}
-  transform.loop.outline %0 {func_name = "foo"} : (!pdl.operation) -> !pdl.operation
+  transform.loop.outline %0 {func_name = "foo"} : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
 }

mlir/test/Dialect/SCF/transform-ops.mlir

func.func @loop_outline_op(%arg0: index, %arg1: index, %arg2: index) {
   // CHECK:   func.call @foo[[SUFFIX]]
   scf.for %j = %arg0 to %arg1 step %arg2 {
     arith.addi %j, %j : index
   }
   return
 }
 
 transform.sequence failures(propagate) {
-^bb1(%arg1: !pdl.operation):
-  %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!pdl.operation) -> !pdl.operation
-  %1 = transform.loop.get_parent_for %0  : (!pdl.operation) -> !transform.op<"scf.for">
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["arith.addi"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %1 = transform.loop.get_parent_for %0  : (!transform.any_op) -> !transform.op<"scf.for">
   // CHECK: = transform.loop.outline %{{.*}}
-  transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> !pdl.operation
+  transform.loop.outline %1 {func_name = "foo"} : (!transform.op<"scf.for">) -> (!transform.any_op, !transform.any_op)
 }
 
 // -----
 
 // CHECK-LABEL: @loop_peel_op
 func.func @loop_peel_op() {
   // CHECK: %[[C0:.+]] = arith.constant 0
   // CHECK: %[[C42:.+]] = arith.constant 42

mlir/test/python/dialects/transform_loop_ext.py

def getParentLoop():
   # CHECK: num_loops = 2
 
 
 @run
 def loopOutline():
   sequence = transform.SequenceOp(transform.FailurePropagationMode.PROPAGATE,
                                   [], transform.OperationType.get("scf.for"))
   with InsertionPoint(sequence.body):
-    loop.LoopOutlineOp(pdl.OperationType.get(), sequence.bodyTarget, func_name="foo")
+    loop.LoopOutlineOp(transform.AnyOpType.get(), transform.AnyOpType.get(), sequence.bodyTarget, func_name="foo")
     transform.YieldOp()
   # CHECK-LABEL: TEST: loopOutline
   # CHECK: = transform.loop.outline %
   # CHECK: func_name = "foo"
 
 
 @run
 def loopPeel():