[mlir][sparse] hoists alloca outside the outermost loop. (#70085)

2023-10-24 11:22:50 -07:00
parent f418319730
commit 43961264c8
1 changed files with 6 additions and 2 deletions
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -592,8 +592,12 @@ public:
    Value lvlCoords, vref;
    {
      OpBuilder::InsertionGuard guard(rewriter);
-      auto loop = op->getParentOfType<LoopLikeOpInterface>();
-      if (loop) {
+      Operation *loop = op;
+      // Finds the outermost loop.
+      while (auto l = loop->getParentOfType<LoopLikeOpInterface>())
+        loop = l;
+
+      if (llvm::isa<LoopLikeOpInterface>(loop)) {
        // Hoists alloca outside the loop to avoid stack overflow.
        rewriter.setInsertionPoint(loop);
      }