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

[llvmlistbot at llvm.org]

Author: Peiming Liu
Date: 2023-10-24T11:22:50-07:00
New Revision: 43961264c810deeb063444f7a83ff86cfe52066d

URL: https://github.com/llvm/llvm-project/commit/43961264c810deeb063444f7a83ff86cfe52066d
DIFF: https://github.com/llvm/llvm-project/commit/43961264c810deeb063444f7a83ff86cfe52066d.diff

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

Removed: 
    


################################################################################
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);
       }