[Mlir-commits] [mlir] [mlir][sparse] hoists alloca outside the outermost loop. (PR #70085)

[Peiming Liu]

https://github.com/PeimingLiu created https://github.com/llvm/llvm-project/pull/70085

None

>From 06240a7d17e4b352421ad34b7abfb295117f4326 Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Tue, 24 Oct 2023 18:18:01 +0000
Subject: [PATCH] [mlir][sparse] hoists alloca outside the outermost loop.

---
 .../SparseTensor/Transforms/SparseTensorConversion.cpp    | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index a1f725333530d90..7ce8f71af01a85b 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -592,8 +592,12 @@ class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
     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);
       }