[Mlir-commits] [mlir] 6243d7d - [mlir][sparse] fix stack overflow due to memref.alloca in loops (#69786)

[llvmlistbot at llvm.org]

Author: Peiming Liu
Date: 2023-10-20T16:01:36-07:00
New Revision: 6243d7d28b923c9e4f881f2a7ac77c1d13486ab2

URL: https://github.com/llvm/llvm-project/commit/6243d7d28b923c9e4f881f2a7ac77c1d13486ab2
DIFF: https://github.com/llvm/llvm-project/commit/6243d7d28b923c9e4f881f2a7ac77c1d13486ab2.diff

LOG: [mlir][sparse] fix stack overflow due to memref.alloca in loops (#69786)

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 73f5e3eeb7d512e..a1f725333530d90 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -589,8 +589,17 @@ class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
     const auto stt = getSparseTensorType(op.getTensor());
     const auto elemTp = stt.getElementType();
     const Level lvlRank = stt.getLvlRank();
-    auto lvlCoords = genAlloca(rewriter, loc, lvlRank, rewriter.getIndexType());
-    auto vref = genAllocaScalar(rewriter, loc, elemTp);
+    Value lvlCoords, vref;
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      auto loop = op->getParentOfType<LoopLikeOpInterface>();
+      if (loop) {
+        // Hoists alloca outside the loop to avoid stack overflow.
+        rewriter.setInsertionPoint(loop);
+      }
+      lvlCoords = genAlloca(rewriter, loc, lvlRank, rewriter.getIndexType());
+      vref = genAllocaScalar(rewriter, loc, elemTp);
+    }
     storeAll(rewriter, loc, lvlCoords, adaptor.getLvlCoords());
     rewriter.create<memref::StoreOp>(loc, adaptor.getValue(), vref);
     SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};