[mlir][transform] Add `transform.apply_licm` op

[springerm]

Summary

This op applies loop-invariant code motion to the targeted loop-like op.

[ftynse]

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

1731

Not for this commit, but we should factor out the tests for the _interpreter_, which belong into test-interpreter, from the tests that are exercising various random operations.

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

 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/CastInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 namespace transform {
 
 enum class FailurePropagationMode : uint32_t;
 class FailurePropagationModeAttr;
 
 /// A builder function that populates the body of a SequenceOp.

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

def ApplyCanonicalizationPatternsOp
   let summary = "Populates canonicalization patterns";
   let description = [{
     This op populates all canonicalization patterns of all loaded dialects in
     an `apply_patterns` transform.
   }];
   let assemblyFormat = "attr-dict";
 }
 
+def ApplyLoopInvariantCodeMotionOp : TransformDialectOp<"apply_licm",
+    [TransformOpInterface, TransformEachOpTrait,
+     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+     ReportTrackingListenerFailuresOpTrait]> {
+  let summary = "Move loop-invariant code out of a loop-like op";
+  let description = [{
+    This transform moves side-effect free, loop invariant code out of the
+    targeted loop-like op. The targeted op must implement the
+    `LoopLikeOpInterface`.
+
+    Note: To move invariant ops from a loop nest, this transform must be applied
+    to each loop of the loop nest, starting with the inner-most loop.
+
+    This transform reads the target handle and modifies the payload.
+  }];
+
+  let arguments = (ins TransformHandleTypeInterface:$target);
+  let results = (outs);
+  let assemblyFormat = "`to` $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 ApplyRegisteredPassOp : TransformDialectOp<"apply_registered_pass",
     [TransformOpInterface, TransformEachOpTrait,
      FunctionalStyleTransformOpTrait, MemoryEffectsOpInterface]> {
   let summary = "Applies the specified registered pass";
   let description = [{
     This transform applies the specified pass to the targeted ops. The name of
     the pass is specified as a string attribute, as set during pass
     registration. Optionally, pass options may be specified as a string

mlir/lib/Dialect/Transform/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRTransformDialect
   MatchInterfaces.cpp
   TransformAttrs.cpp
   TransformDialect.cpp
   TransformInterfaces.cpp
   TransformOps.cpp
   TransformTypes.cpp
 
   DEPENDS
   MLIRMatchInterfacesIncGen
   MLIRTransformDialectIncGen
   MLIRTransformInterfacesIncGen
 
   LINK_LIBS PUBLIC
   MLIRCastInterfaces
   MLIRIR
+  MLIRLoopLikeInterface
   MLIRParser
   MLIRPass
   MLIRRewrite
   MLIRSideEffectInterfaces
   MLIRTransforms
   MLIRTransformDialectUtils
   )

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

 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Transforms/CSE.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Debug.h"
 #include <optional>
 
 #define DEBUG_TYPE "transform-dialect"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "] ")
void transform::ApplyCanonicalizationPatternsOp::populatePatterns(
   MLIRContext *ctx = patterns.getContext();
   for (Dialect *dialect : ctx->getLoadedDialects())
     dialect->getCanonicalizationPatterns(patterns);
   for (RegisteredOperationName op : ctx->getRegisteredOperations())
     op.getCanonicalizationPatterns(patterns, ctx);
 }
 
 //===----------------------------------------------------------------------===//
+// ApplyLoopInvariantCodeMotionOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure
+transform::ApplyLoopInvariantCodeMotionOp::applyToOne(
+    transform::TransformRewriter &rewriter, LoopLikeOpInterface target,
+    transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  // Currently, LICM does not remove operations, so we don't need tracking.
+  // If this ever changes, add a LICM entry point that takes a rewriter.
+  moveLoopInvariantCode(target);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::ApplyLoopInvariantCodeMotionOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTarget(), effects);
+  transform::modifiesPayload(effects);
+}
+
+//===----------------------------------------------------------------------===//
 // ApplyRegisteredPassOp
 //===----------------------------------------------------------------------===//
 
 DiagnosedSilenceableFailure transform::ApplyRegisteredPassOp::applyToOne(
     transform::TransformRewriter &rewriter, Operation *target,
     ApplyToEachResultList &results, transform::TransformState &state) {
   // Make sure that this transform is not applied to itself. Modifying the
   // transform IR while it is being interpreted is generally dangerous. Even

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

 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match attributes{original} in %arg1 : (!transform.any_op) -> !transform.any_op
   %1 = transform.structured.match attributes{replacement} in %arg1 : (!transform.any_op) -> !transform.any_op
   test_notify_payload_op_replaced %0, %1 : (!transform.any_op, !transform.any_op) -> ()
   test_print_remark_at_operand %0, "updated handle" : !transform.any_op
 }
 
 // -----
 
 // CHECK-LABEL: func @test_apply_cse()
 //       CHECK:   %[[const:.*]] = arith.constant 0 : index
 //       CHECK:   %[[ex1:.*]] = scf.execute_region -> index {
 //       CHECK:     scf.yield %[[const]]
 //       CHECK:   }
 //       CHECK:   %[[ex2:.*]] = scf.execute_region -> index {
 //       CHECK:     scf.yield %[[const]]
^bb1(%arg1: !transform.any_op):
   // The other handles were also updated.
   test_print_remark_at_operand %elim_first, "eliminated 1" : !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %elim_first : !transform.any_op
   test_print_remark_at_operand %elim_second, "eliminated 2" : !transform.any_op
   // expected-remark @below{{1}}
   test_print_number_of_associated_payload_ir_ops %elim_second : !transform.any_op
 }
+
+// -----
+
+// CHECK-LABEL: func @test_licm(
+//       CHECK:   arith.muli
+//       CHECK:   scf.for {{.*}} {
+//       CHECK:     vector.print
+//       CHECK:   }
+func.func @test_licm(%arg0: index, %arg1: index, %arg2: index) {
+  scf.for %iv = %arg0 to %arg1 step %arg2 {
+    %0 = arith.muli %arg0, %arg1 : index
+    vector.print %0 : index
+  }
+  return
+}
+
+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.apply_licm to %0 : !transform.any_op
+}
+
+// -----
+
+// expected-note @below{{when applied to this op}}
+module {
+  func.func @test_licm_invalid() {
+    return
+  }
+
+  transform.sequence failures(propagate) {
+  ^bb1(%arg1: !transform.any_op):
+    // expected-error @below{{transform applied to the wrong op kind}}
+    transform.apply_licm to %arg1 : !transform.any_op
+  }
+}

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

cc_library(
     srcs = glob(["lib/Dialect/Transform/IR/*.cpp"]),
     hdrs = glob(["include/mlir/Dialect/Transform/IR/*.h"]),
     deps = [
         ":Analysis",
         ":CallOpInterfaces",
         ":CastInterfaces",
         ":ControlFlowInterfaces",
         ":IR",
+        ":LoopLikeInterface",
         ":Pass",
         ":Rewrite",
         ":SideEffectInterfaces",
         ":Support",
         ":TransformDialectEnumsIncGen",
         ":TransformDialectIncGen",
         ":TransformDialectInterfacesIncGen",
         ":TransformDialectMatchInterfacesIncGen",