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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188761

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

>From 9a4c91b27029c40f501798b5fe98583469d3b044 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 07:33:18 -0700
Subject: [PATCH] [MLIR][Transforms] Fix two bugs in
 loop-invariant-subset-hoisting

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
---
 mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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())) {