[Mlir-commits] [mlir] [MLIR][Affine] Check values captured by regions in affine LICM (PR #216605)

Sun Aug 16 14:18:21 PDT 2026

llvmorg-github-actions[bot] wrote:



@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir

Author: Alessandro Potenza (alepot55)

<details>
<summary>Changes</summary>

`affine-loop-invariant-code-motion` hoists an operation out of an `affine.for` while its regions still read a value that stays behind, producing IR that fails the verifier.

`isOpLoopInvariant` walks the bodies of `affine.if`, `affine.for` and `affine.parallel`, and for every other region-carrying operation it inspects `op.getOperands()` only. The values an operation's regions capture are not operands of that operation, so nothing examines them. An `scf.for` whose bounds, step and `iter_args` are all loop invariant therefore looks hoistable even when its body reads a value defined inside the loop.

```mlir
func.func @f(%m: memref<4xi64>, %init: i64, %outside: i64) -> i64 {
  %c0 = arith.constant 0 : i32
  %c4 = arith.constant 4 : i32
  %c2 = arith.constant 2 : i32
  %one = arith.constant 1 : i64
  %r = affine.for %i = 0 to 4 iter_args(%acc = %init) -> (i64) {
    %v = affine.load %m[%i] : memref<4xi64>
    %s = scf.for %j = %c0 to %c4 step %c2 iter_args(%a = %outside) -> (i64) : i32 {
      %o = arith.ori %v, %one : i64
      %n = arith.addi %a, %o : i64
      scf.yield %n : i64
    }
    %acc2 = arith.addi %acc, %s : i64
    affine.yield %acc2 : i64
  }
  return %r : i64
}
```

The input verifies. After the pass:

```
error: operand #0 does not dominate this use
      %o = arith.ori %v, %one : i64
```

The `scf.for` has been moved above the `affine.for` and its body still refers to `%v`, which is defined inside the loop.

This extends the existing operand check to the values the operation's regions capture, using the same three conditions the operand check already applies (the induction variable, the loop's `iter_args`, and a value produced inside the loop by an operation that is not itself being hoisted). Only the check is widened; nothing else about which operations are considered changes.

Two tests: one for the case above, and one where the region reads only values defined outside the loop, which must still be hoisted, so that a check that is too broad does not pass unnoticed.

Fixes #216545


---
Full diff: https://github.com/llvm/llvm-project/pull/216605.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp (+28-15) 
- (modified) mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir (+53) 


``````````diff

diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
index 1887c321e206a..dae8600efe46b 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopInvariantCodeMotion.cpp
@@ -16,6 +16,7 @@
 #include "mlir/Dialect/Affine/Analysis/Utils.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Transforms/RegionUtils.h"
 
 namespace mlir {
 namespace affine {
@@ -109,25 +110,37 @@ static bool isOpLoopInvariant(Operation &op, AffineForOp loop,
 
   // Check operands.
   ValueRange iterArgs = loop.getRegionIterArgs();
-  for (unsigned int i = 0; i < op.getNumOperands(); ++i) {
-    auto *operandSrc = op.getOperand(i).getDefiningOp();
-
+  auto isLoopVariantValue = [&](Value value) {
     // If the loop IV is the operand, this op isn't loop invariant.
-    if (iv == op.getOperand(i))
-      return false;
+    if (iv == value)
+      return true;
 
     // If the one of the iter_args is the operand, this op isn't loop invariant.
-    if (llvm::is_contained(iterArgs, op.getOperand(i)))
-      return false;
+    if (llvm::is_contained(iterArgs, value))
+      return true;
+
+    // If the value was defined in the loop (outside of the if/else region),
+    // and that operation itself wasn't meant to be hoisted, then mark this
+    // operation loop dependent.
+    Operation *operandSrc = value.getDefiningOp();
+    return operandSrc && opsWithUsers.count(operandSrc) &&
+           opsToHoist.count(operandSrc) == 0;
+  };
+
+  if (llvm::any_of(op.getOperands(), isLoopVariantValue))
+    return false;
 
-    if (operandSrc) {
-      // If the value was defined in the loop (outside of the if/else region),
-      // and that operation itself wasn't meant to be hoisted, then mark this
-      // operation loop dependent.
-      if (opsWithUsers.count(operandSrc) && opsToHoist.count(operandSrc) == 0)
-        return false;
-    }
-  }
+  // Check the values the op's regions read from around them. The regions travel
+  // with the op, so a value they capture from inside the loop pins the op in
+  // place just as an operand does. Only the region-carrying ops handled above
+  // have their bodies walked, and even there the walk inspects the nested ops
+  // rather than what they capture.
+  bool capturesLoopVariantValue = false;
+  visitUsedValuesDefinedAbove(op.getRegions(), [&](OpOperand *operand) {
+    capturesLoopVariantValue |= isLoopVariantValue(operand->get());
+  });
+  if (capturesLoopVariantValue)
+    return false;
 
   // If no operand was loop variant, mark this op for motion.
   opsToHoist.insert(&op);
diff --git a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
index 83313f303dbd1..7bd8d1535d4a1 100644
--- a/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
+++ b/mlir/test/Dialect/Affine/affine-loop-invariant-code-motion.mlir
@@ -986,3 +986,56 @@ func.func @unknown_trip_count_store_not_hoisted(%x: i32, %n: index) -> i32 {
   %r = affine.load %alloc[0] : memref<1xi32>
   return %r : i32
 }
+
+// -----
+
+// An op whose regions read a value defined in the loop must stay in the loop.
+// The regions move with the op, so hoisting the `scf.for` below would leave its
+// body referring to `%v`, which stays behind.
+
+// CHECK-LABEL: func @region_capturing_loop_variant_value_not_hoisted
+func.func @region_capturing_loop_variant_value_not_hoisted(%m: memref<4xi64>, %init: i64, %outside: i64) -> i64 {
+  %c0 = arith.constant 0 : i32
+  %c4 = arith.constant 4 : i32
+  %c2 = arith.constant 2 : i32
+  %one = arith.constant 1 : i64
+  // CHECK: affine.for
+  %r = affine.for %i = 0 to 4 iter_args(%acc = %init) -> (i64) {
+    // CHECK-NEXT: affine.load
+    %v = affine.load %m[%i] : memref<4xi64>
+    // CHECK-NEXT: scf.for
+    %s = scf.for %j = %c0 to %c4 step %c2 iter_args(%a = %outside) -> (i64) : i32 {
+      %o = arith.ori %v, %one : i64
+      %n = arith.addi %a, %o : i64
+      scf.yield %n : i64
+    }
+    %acc2 = arith.addi %acc, %s : i64
+    affine.yield %acc2 : i64
+  }
+  return %r : i64
+}
+
+// -----
+
+// The same shape, but the region reads only values defined outside the loop, so
+// the op is still hoisted.
+
+// CHECK-LABEL: func @region_capturing_invariant_value_is_hoisted
+func.func @region_capturing_invariant_value_is_hoisted(%init: i64, %outside: i64) -> i64 {
+  %c0 = arith.constant 0 : i32
+  %c4 = arith.constant 4 : i32
+  %c2 = arith.constant 2 : i32
+  %one = arith.constant 1 : i64
+  // CHECK: scf.for
+  // CHECK: affine.for
+  %r = affine.for %i = 0 to 4 iter_args(%acc = %init) -> (i64) {
+    %s = scf.for %j = %c0 to %c4 step %c2 iter_args(%a = %outside) -> (i64) : i32 {
+      %o = arith.ori %outside, %one : i64
+      %n = arith.addi %a, %o : i64
+      scf.yield %n : i64
+    }
+    %acc2 = arith.addi %acc, %s : i64
+    affine.yield %acc2 : i64
+  }
+  return %r : i64
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/216605
