[Mlir-commits] [mlir] [mlir][Transform] Reuse bbArgs in FuseIntoContainingOp (PR #135066)

[Oleksandr Alex Zinenko]

================
@@ -718,6 +718,42 @@ static Operation *replaceForAllWithNewSignature(
   return newforallOp;
 }
 
+/// Given two operands coming from a loop iter arg, 'src' and 'dst', return true
+/// if the operand 'src' is equal to 'dst' or equal to a iter arg present in a
+/// outer loop. To determine the second condition, this function iterates
+/// recursively over the enclosing loops, trying to find 'src' in any of the
+/// parent loop's iter args.
+static bool sameOrEquivalentIterArg(Value src, Value dst) {
+  // Base case.
+  if (src == dst)
+    return true;
+
+  // Recursively look for equivalent iter args in enclosing loops.
+  if (auto bbArg = dyn_cast<BlockArgument>(dst)) {
+    Block *parentBlock = bbArg.getOwner();
+    assert(parentBlock && "unlinked block argument");
+
+    // Because we stop doing recursive calls when we find a non loop-like op,
+    // this should never happen.
+    assert(parentBlock->getParentOp() &&
+           "expected block argument with parent operation");
+
+    // Check if parent is loop-like.
+    if (auto parentLoop =
+            dyn_cast<LoopLikeOpInterface>(parentBlock->getParentOp())) {
+      for (auto innerIterArg : parentLoop.getRegionIterArgs()) {
+        OpOperand *operand = parentLoop.getTiedLoopInit(innerIterArg);
+        Value loopBlockArgument =
+            parentLoop->getOperand(operand->getOperandNumber());
+        if (sameOrEquivalentIterArg(src, loopBlockArgument))
----------------
ftynse wrote:

> In such case, I don't see how it can be easily done. The function handles loop nests of arbitrary depth, not only depth 2 (in the MLIR test it is exercising a case of depth 3), so it looks to me like recursive calls is the best here.

By using a `SmallVector<Value> destWorklist` acting as stack, pushing the original `dst` and having a while-non-empty loop that pops the last value, does the logic and potentially pushes new values into the vector?

> Unrelated to this comment, but related to this function: I thought it could make sense to move this function somewhere else? As it could be reused for other things. Not sure where, maybe mlir/lib/Dialect/Linalg/Utils/Utils.cpp?

If that files already exists, it's fine with me. Creating new files will only increase complexity for no good reason (unless there is a second potential user for this immediately).

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