[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());
----------------
ftynse wrote:

`operand` may be null here...

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