[Mlir-commits] [mlir] [mlir][sparse] fix stack overflow due to alloca instruction inside loops (PR #69765)

[Peiming Liu]

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

None

>From e93c52708b543b688fa19d9fb18c3003ba78f7fc Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Fri, 20 Oct 2023 19:59:20 +0000
Subject: [PATCH] [mlir][sparse] fix stack overflow due to alloca instruction
 inside loops

---
 mlir/include/mlir/Dialect/SCF/IR/SCFOps.td            |  2 +-
 .../Transforms/SparseTensorConversion.cpp             | 11 ++++++++++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
index 6ac0912f6f706c5..d079b1e5d92eece 100644
--- a/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
+++ b/mlir/include/mlir/Dialect/SCF/IR/SCFOps.td
@@ -120,7 +120,7 @@ def ExecuteRegionOp : SCF_Op<"execute_region", [
 //===----------------------------------------------------------------------===//
 
 def ForOp : SCF_Op<"for",
-      [AutomaticAllocationScope, DeclareOpInterfaceMethods<LoopLikeOpInterface,
+      [DeclareOpInterfaceMethods<LoopLikeOpInterface,
        ["getInits", "getSingleInductionVar", "getSingleLowerBound",
         "getSingleStep", "getSingleUpperBound", "getYieldedValues",
         "promoteIfSingleIteration", "replaceWithAdditionalYields"]>,
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 73f5e3eeb7d512e..64874c49f6641fb 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -589,6 +589,13 @@ class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
     const auto stt = getSparseTensorType(op.getTensor());
     const auto elemTp = stt.getElementType();
     const Level lvlRank = stt.getLvlRank();
+    // We need a alloc scope here as the InsertOp is always generated inside a
+    // loop, which lead to stack overflow.
+    auto allocaScope = rewriter.create<memref::AllocaScopeOp>(
+        loc, adaptor.getTensor().getType());
+    Block *scopeBlock = rewriter.createBlock(&allocaScope.getBodyRegion());
+    rewriter.setInsertionPointToStart(scopeBlock);
+
     auto lvlCoords = genAlloca(rewriter, loc, lvlRank, rewriter.getIndexType());
     auto vref = genAllocaScalar(rewriter, loc, elemTp);
     storeAll(rewriter, loc, lvlCoords, adaptor.getLvlCoords());
@@ -596,7 +603,9 @@ class SparseTensorInsertConverter : public OpConversionPattern<InsertOp> {
     SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {},
                    {adaptor.getTensor(), lvlCoords, vref}, EmitCInterface::On);
-    rewriter.replaceOp(op, adaptor.getTensor());
+    rewriter.create<memref::AllocaScopeReturnOp>(loc, adaptor.getTensor());
+    rewriter.setInsertionPointAfter(allocaScope);
+    rewriter.replaceOp(op, allocaScope.getResult(0));
     return success();
   }
 };