[Mlir-commits] [mlir] [MLIR][Transforms] Fix two bugs in loop-invariant-subset-hoisting (PR #188761)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

Fix two issues in `MatchingSubsets::populateSubsetOpsAtIterArg`:

1. The `collectHoistableOps` parameter was declared but never used when inserting subset ops via `insert(subsetOp)`. As a result, when recursing into nested loops with `collectHoistableOps=false`, the nested loop's subset ops were incorrectly added to the hoistable extraction/insertion pairs of the parent loop. This caused spurious failures in the `allDisjoint` check, preventing valid hoisting when nested loop ops overlapped with outer loop ops. Fix by passing the parameter: `insert(subsetOp, collectHoistableOps)`.

2. In the nested loop handling branch, there was no guard to detect when a value has multiple nested loop uses (i.e., is used as an init arg in more than one nested loop). Without the guard, `nextValue` would be silently overwritten, leading to an incorrect use-def chain traversal. Add `if (nextValue) return failure()` before setting `nextValue` for the nested loop case, mirroring the existing guard for insertion ops.

Fixes #<!-- -->147096

Assisted-by: Claude Code

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


1 Files Affected:

- (modified) mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp (+5-1) 


``````````diff
diff --git a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
index 5f3b04a6d8bce..aef81bf2e15a3 100644
--- a/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp
@@ -271,6 +271,10 @@ MatchingSubsets::populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
         if (failed(populateSubsetOpsAtIterArg(nestedLoop, nestedIterArg,
                                               /*collectHoistableOps=*/false)))
           return failure();
+        // There must be a single use-def chain. Bail if there are multiple
+        // nested loops using this value.
+        if (nextValue)
+          return failure();
         nextValue = nestedLoop.getTiedLoopResult(&use);
         continue;
       }
@@ -278,7 +282,7 @@ MatchingSubsets::populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
       auto subsetOp = dyn_cast<SubsetOpInterface>(use.getOwner());
       if (!subsetOp)
         return failure();
-      insert(subsetOp);
+      insert(subsetOp, collectHoistableOps);
 
       if (auto insertionOp =
               dyn_cast<SubsetInsertionOpInterface>(use.getOwner())) {

``````````

</details>


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