add HoistLoopInvariantCodeMotionOp

Author: linuxlonelyeagle

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 instructions outside of the loop";
+  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

mlir/lib/Dialect/Transform/LoopExtension/LoopExtensionOps.cpp

@@ -32,3 +32,22 @@ 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);
+}

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
+  }
+}