[Mlir-commits] [mlir] [mlir][linalg] Fix exponential walk of tensor.extract indices (PR #222319)

[Jelle Schühmacher]

https://github.com/jschuhmacher created https://github.com/llvm/llvm-project/pull/222319

`isLoopInvariantIdx` showed exponential runtime behavior. Replaced with a worklist and a visited set.

Encountered in iree-org/iree#24886.

Assisted-by: Claude

>From 50d84ab5aa8cd28720112afe533477d8f58fd326 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jelle=20Sch=C3=BChmacher?= <schuehmacher at roofline.ai>
Date: Wed, 9 Sep 2026 11:19:59 +0000
Subject: [PATCH] [mlir][linalg] Fix exponential walk of tensor.extract indices

Encountered in iree-org/iree#24886.

Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
 .../Linalg/Transforms/Vectorization.cpp       | 64 +++++++++++--------
 1 file changed, 39 insertions(+), 25 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index e19f59f0e4478..df0809f5f76ee 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -37,6 +37,7 @@
 #include "mlir/Transforms/RegionUtils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/SmallVectorExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
@@ -977,39 +978,52 @@ static bool isLoopInvariantIdx(LinalgOp &linalgOp, Value &val,
                           [](int64_t dimSize) { return dimSize > 1; }) == 1)) &&
          "n-D vectors are not yet supported");
 
-  // Blocks outside _this_ linalg.generic are effectively loop invariant.
-  // However, analysing block arguments for _this_ linalg.generic Op is a bit
-  // tricky. Just bail out in the latter case.
-  // TODO: We could try analysing the corresponding affine map here.
   auto *block = linalgOp.getBlock();
-  if (isa<BlockArgument>(val))
-    return !llvm::is_contained(block->getArguments(), val);
 
-  Operation *defOp = val.getDefiningOp();
-  assert(defOp && "This is neither a block argument nor an operation result");
+  // A shared DAG has exponentially many paths; a revisit adds nothing.
+  SmallPtrSet<Operation *, 8> visited;
+  SmallVector<Value> worklist{val};
 
-  // IndexOp is loop invariant as long as its result remains constant across
-  // iterations. Note that for dynamic shapes, the corresponding dim will also
-  // be conservatively treated as != 1.
-  if (auto indexOp = dyn_cast<linalg::IndexOp>(defOp)) {
-    return linalgOp.getStaticLoopRanges()[indexOp.getDim()] == 1;
-  }
+  while (!worklist.empty()) {
+    Value v = worklist.pop_back_val();
 
-  auto *ancestor = block->findAncestorOpInBlock(*defOp);
+    // Blocks outside _this_ linalg.generic are effectively loop invariant.
+    // However, analysing block arguments for _this_ linalg.generic Op is a bit
+    // tricky. Just bail out in the latter case.
+    // TODO: We could try analysing the corresponding affine map here.
+    if (isa<BlockArgument>(v)) {
+      if (llvm::is_contained(block->getArguments(), v))
+        return false;
+      continue;
+    }
 
-  // Values define outside `linalgOp` are loop invariant.
-  if (!ancestor)
-    return true;
+    Operation *defOp = v.getDefiningOp();
+    assert(defOp && "This is neither a block argument nor an operation result");
 
-  // Values defined inside `linalgOp`, which are constant, are loop invariant.
-  if (isa<arith::ConstantOp>(ancestor))
-    return true;
+    // IndexOp is loop invariant as long as its result remains constant across
+    // iterations. Note that for dynamic shapes, the corresponding dim will also
+    // be conservatively treated as != 1.
+    if (auto indexOp = dyn_cast<linalg::IndexOp>(defOp)) {
+      if (linalgOp.getStaticLoopRanges()[indexOp.getDim()] != 1)
+        return false;
+      continue;
+    }
 
-  bool result = true;
-  for (auto op : ancestor->getOperands())
-    result &= isLoopInvariantIdx(linalgOp, op, resType);
+    auto *ancestor = block->findAncestorOpInBlock(*defOp);
 
-  return result;
+    // Values define outside `linalgOp` are loop invariant.
+    if (!ancestor)
+      continue;
+
+    // Values defined inside `linalgOp`, which are constant, are loop invariant.
+    if (isa<arith::ConstantOp>(ancestor))
+      continue;
+
+    if (visited.insert(ancestor).second)
+      llvm::append_range(worklist, ancestor->getOperands());
+  }
+
+  return true;
 }
 
 /// Check whether `val` could be used for calculating the trailing index for a