[Mlir-commits] [mlir] [mlir][transform] Add HoistLoopInvariantCodeMotionOp (PR #163687)

[lonely eagle]

https://github.com/linuxlonelyeagle created https://github.com/llvm/llvm-project/pull/163687

HoistLoopInvariantSubsetsOp already exists upstream, and there is equally good reason for HoistLoopInvariantCodeMotionOp to exist. HoistLoopInvariantSubsetsOp cannot solve all problems; when an extractOp has a dynamic shape, we need to use HoistLoopInvariantCodeMotionOp to advance the dynamic size first, and then use HoistLoopInvariantSubsetsOp.

>From 6530bf051f8d122d6381f0af58f7080963595d16 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 16 Oct 2025 02:57:15 +0000
Subject: [PATCH] add HoistLoopInvariantCodeMotionOp

---
 .../LoopExtension/LoopExtensionOps.td         | 56 +++++++++++++++++++
 .../LoopExtension/LoopExtensionOps.cpp        | 18 ++++++
 .../Transform/test-loop-transforms.mlir       | 34 +++++++++++
 3 files changed, 108 insertions(+)

diff --git a/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.td b/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.td
index 885b55811e62c..9854c8db2ebf3 100644
--- a/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.td
+++ b/mlir/include/mlir/Dialect/Transform/LoopExtension/LoopExtensionOps.td
@@ -73,4 +73,60 @@ def HoistLoopInvariantSubsetsOp
   }];
 }
 
+def HoistLoopInvariantCodeMotionOp
+    : TransformDialectOp<"loop.hoist_loop_invariant_code_motion",
+        [TransformOpInterface, TransformEachOpTrait,
+         DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
+         ReportTrackingListenerFailuresOpTrait]> {
+  let summary = "Hoist loop invariant subset ops";
+  let description = [{
+    This transform hoists loop-invariant instructions out of the targeted
+    loop-like op.
+
+    Example:
+    ```
+    %m = memref.alloc() : memref<10xf32>
+    %cf7 = arith.constant 7.0 : f32
+    %cf8 = arith.constant 8.0 : f32
+
+    affine.for %arg0 = 0 to 10 {
+      affine.for %arg1 = 0 to 10 {
+        %v0 = arith.addf %cf7, %cf8 : f32
+      }
+    }
+    ```
+    Is transformed to:
+    ```
+    %alloc = memref.alloc() : memref<10xf32>
+    %cst = arith.constant 7.000000e+00 : f32
+    %cst_0 = arith.constant 8.000000e+00 : f32
+    %0 = arith.addf %cst, %cst_0 : f32
+    affine.for %arg0 = 0 to 10 {
+    }
+    affine.for %arg0 = 0 to 10 {
+    }
+    ```
+
+    loop-invariant instructions are hoisted only if they are pure ops and
+    they are ancestors of the parent regionOp of all operands.
+
+    This transform reads the target handle and modifies the payload. This
+    transform does not invalidate any handles, but loop-like ops are replaced
+    with new loop-like ops when a loop-invariant op is hoisted. The transform
+    rewriter updates all handles accordingly.
+  }];
+
+  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 loopLikeOp,
+      ::mlir::transform::ApplyToEachResultList &results,
+      ::mlir::transform::TransformState &state);
+  }];
+}
+
 #endif // MLIR_DIALECT_TRANSFORM_LOOPEXTENSION_LOOPEXTENSIONOPS
diff --git a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp
index 95870e8ef87be..de50e5aab4510 100644
--- a/mlir/lib/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp
+++ b/mlir/lib/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp
@@ -32,3 +32,21 @@ void transform::HoistLoopInvariantSubsetsOp::getEffects(
   transform::onlyReadsHandle(getTargetMutable(), effects);
   transform::modifiesPayload(effects);
 }
+
+//===----------------------------------------------------------------------===//
+// HoistLoopInvariantCodeMotionOp
+//===----------------------------------------------------------------------===//
+
+DiagnosedSilenceableFailure transform::HoistLoopInvariantCodeMotionOp::applyToOne(
+    transform::TransformRewriter &rewriter, LoopLikeOpInterface loopLikeOp,
+    transform::ApplyToEachResultList &results,
+    transform::TransformState &state) {
+  (void)moveLoopInvariantCode(loopLikeOp);
+  return DiagnosedSilenceableFailure::success();
+}
+
+void transform::HoistLoopInvariantCodeMotionOp::getEffects(
+    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+  transform::onlyReadsHandle(getTargetMutable(), effects);
+  transform::modifiesPayload(effects);
+}
diff --git a/mlir/test/Dialect/Transform/test-loop-transforms.mlir b/mlir/test/Dialect/Transform/test-loop-transforms.mlir
index 833dd738f79ed..c3d21c406172d 100644
--- a/mlir/test/Dialect/Transform/test-loop-transforms.mlir
+++ b/mlir/test/Dialect/Transform/test-loop-transforms.mlir
@@ -79,3 +79,37 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+func.func @nested_loops_both_having_invariant_code_transform() {
+  %m = memref.alloc() : memref<10xf32>
+  %cf7 = arith.constant 7.0 : f32
+  %cf8 = arith.constant 8.0 : f32
+  affine.for %arg0 = 0 to 10 {
+    %v0 = arith.addf %cf7, %cf8 : f32
+    affine.for %arg1 = 0 to 10 {
+      %v1 = arith.addf %v0, %cf8 : f32
+      affine.store %v0, %m[%arg0] : memref<10xf32>
+    }
+  }
+
+  // CHECK: memref.alloc() : memref<10xf32>
+  // CHECK-NEXT: %[[CST0:.*]] = arith.constant 7.000000e+00 : f32
+  // CHECK-NEXT: %[[CST1:.*]] = arith.constant 8.000000e+00 : f32
+  // CHECK-NEXT: %[[ADD0:.*]] = arith.addf %[[CST0]], %[[CST1]] : f32
+  // CHECK-NEXT: arith.addf %[[ADD0]], %[[CST1]] : f32
+  // CHECK-NEXT: affine.for
+  // CHECK-NEXT: affine.for
+  // CHECK-NEXT: affine.store
+  return
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(
+       %arg0: !transform.any_op {transform.readonly}) {
+    %loop = transform.structured.match interface{LoopLikeInterface} in %arg0 : (!transform.any_op) -> !transform.any_op 
+    transform.loop.hoist_loop_invariant_code_motion %loop : !transform.any_op
+    transform.yield
+  }
+}