[mlir][SCF] Add `loop.promote_if_one_iteration` transform op

[springerm]

Summary

This transform op promotes loops with one iteration. I.e., the loop op is replaced by just the loop body.

Depends On: D154343

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

 //===- SCFTransformOps.h - SCF transformation ops ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SCF_TRANSFORMOPS_SCFTRANSFORMOPS_H
 #define MLIR_DIALECT_SCF_TRANSFORMOPS_SCFTRANSFORMOPS_H
 
 #include "mlir/Bytecode/BytecodeOpInterface.h"
 #include "mlir/Dialect/Transform/IR/TransformInterfaces.h"
 #include "mlir/Dialect/Transform/IR/TransformTypes.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 namespace func {
 class FuncOp;
 } // namespace func
 namespace scf {
 class ForOp;
 class IfOp;

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

let extraClassDeclaration = [{
     ::mlir::DiagnosedSilenceableFailure applyToOne(
         ::mlir::transform::TransformRewriter &rewriter,
         ::mlir::scf::ForOp target,
         ::mlir::transform::ApplyToEachResultList &results,
         ::mlir::transform::TransformState &state);
   }];
 }
 
+def LoopPromoteIfOneIterationOp : Op<Transform_Dialect,
+    "loop.promote_if_one_iteration", [
+        DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+        TransformOpInterface, TransformEachOpTrait]> {
+  let summary = "Promote loop if it has one iteration";
+  let description = [{
+    Promotes the given target loop op if it has a single iteration. I.e., the
+    loop op is removed and only the body remains.
+
+    #### Return modes
+
+    This transform fails if the target is mapped to ops that are loops. Ops are
+    considered loops if they implement the `LoopLikeOpInterface`. Otherwise,
+    this transform always succeeds. The transform consumes the target handle and
+    modifies the payload.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+  let assemblyFormat = "$target attr-dict `:` type($target)";
+
+  let extraClassDeclaration = [{
+    ::mlir::DiagnosedSilenceableFailure applyToOne(
+        ::mlir::transform::TransformRewriter &rewriter,
+        ::mlir::LoopLikeOpInterface target,
+        ::mlir::transform::ApplyToEachResultList &results,
+        ::mlir::transform::TransformState &state);
+  }];
+}
+
 def LoopUnrollOp : Op<Transform_Dialect, "loop.unroll",
     [FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface,
      TransformOpInterface, TransformEachOpTrait]> {
   let summary = "Unrolls the given loop with the given unroll factor";
   let description = [{
     Unrolls each loop associated with the given handle to have up to the given
     number of loop body copies per iteration. If the unroll factor is larger
     than the loop trip count, the latter is used as the unroll factor instead.

mlir/lib/Dialect/SCF/TransformOps/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFTransformOps
   SCFTransformOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SCF/TransformOps
 
   DEPENDS
   MLIRSCFTransformOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
   MLIRFuncDialect
   MLIRIR
+  MLIRLoopLikeInterface
   MLIRSCFDialect
   MLIRSCFTransforms
   MLIRSCFUtils
   MLIRTransformDialect
   MLIRVectorDialect
 )

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

transform::LoopPipelineOp::applyToOne(transform::TransformRewriter &rewriter,
   if (succeeded(patternResult)) {
     results.push_back(*patternResult);
     return DiagnosedSilenceableFailure::success();
   }
   return emitDefaultSilenceableFailure(target);
 }
 
 //===----------------------------------------------------------------------===//
+// LoopPromoteIfOneIterationOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::LoopPromoteIfOneIterationOp::applyToOne(
+    transform::TransformRewriter &rewriter, LoopLikeOpInterface target,
+    transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  (void)target.promoteIfSingleIteration(rewriter);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::LoopPromoteIfOneIterationOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  consumesHandle(getTarget(), effects);
+  modifiesPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
 // LoopUnrollOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure
 transform::LoopUnrollOp::applyToOne(transform::TransformRewriter &rewriter,
                                     Operation *op,
                                     transform::ApplyToEachResultList &results,
                                     transform::TransformState &state) {

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

 
 transform.sequence failures(propagate) {
 ^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 { num_loops = 1, affine = true } : (!transform.any_op) -> !transform.op<"affine.for">
   transform.test_print_remark_at_operand %1, "affine for loop" : !transform.op<"affine.for">
   transform.loop.unroll %1 { factor = 4 } : !transform.op<"affine.for">
 }
+
+// -----
+
+// CHECK-LABEL: func @test_promote_if_one_iteration(
+//   CHECK-NOT:   scf.for
+//       CHECK:   %[[r:.*]] = "test.foo"
+//       CHECK:   return %[[r]]
+func.func @test_promote_if_one_iteration(%a: index) -> index {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = scf.for %j = %c0 to %c1 step %c1 iter_args(%arg0 = %a) -> index {
+    %1 = "test.foo"(%a) : (index) -> (index)
+    scf.yield %1 : index
+  }
+  return %0 : index
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  transform.loop.promote_if_one_iteration %0 : !transform.any_op
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     srcs = glob(["lib/Dialect/SCF/TransformOps/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/SCF/TransformOps/*.h"]),
     includes = ["include"],
     deps = [
         ":AffineDialect",
         ":AffineUtils",
         ":FuncDialect",
         ":IR",
+        ":LoopLikeInterface",
         ":SCFDialect",
         ":SCFTransformOpsIncGen",
         ":SCFTransforms",
         ":SCFUtils",
         ":SideEffectInterfaces",
         ":TransformDialect",
         ":VectorDialect",
         "//llvm:Support",