[SCF][MemRef] Enable SCF.Parallel Lowering to use Scope Op

[wsmoses]

Summary

As discussed in https://reviews.llvm.org/D119743 scf.parallel would continuously stack allocate since the alloca op was placd in the wsloop rather than the omp.parallel. This PR is the second stage of the fix for that problem. Specifically, we now introduce an alloca scope around the inlined body of the scf.parallel and enable a canonicalization to hoist the allocations to the surrounding allocation scope (e.g. omp.parallel).

[mehdi_amini]

mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp

398–418

Please run clang-format?

[wsmoses]

Clang format

[ftynse]

Conversion changes look okay. Canonicalization changes less so. An explicit allocation scope forces deallocation at its end, and it looks like both canonicalizations would change the point where deallocation happens. It's okay to merge two nested allocation scopes if they are perfectly nested, otherwise the allocations from the nested scope will be live longer than before. Similarly, hoisting can change the lifetime of an allocation. I'm fine with such lifetime-extending transformations in a separate pass, but not as canonicalizations (which are not expected to change semantics).

Specifically about the hoisting pattern, I don't completely follow what it does and there are no tests specifically for it. It seems to be moving allocations from a scope operation a parent scope operation, which may not be always desirable.

mlir/include/mlir/Conversion/Passes.td

504–505

80 cols plz

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

2604

Let's not spuriously move the code around. Op implementations were sorted alphabetically AFAICS.

2659

Nit: make static + add documentation.

2660–2662

This is too optimistic. If the operation, e.g., an unregistered one, doesn't implement the interface, it may still perform the allocation. We can only be sure that it does not allocate if it has the interface and does not declare the effect with it.

2664

Why limit to the first result only?

2672

Nit: /// for top-level comments.

2680–2684

Could you extract the lambda or the walk from the if?

2699

This pattern needs dedicated tests.

2704–2712

This makes assumptions about the nesting structure of ops. In particular, it seems to assume there is always a parent op with the scope trait. This is not necessarily the case, e.g., one can have a custom function op that doesn't have the trait for whatever reason. Please check that the parent op exists before querying its properties.

2724–2725

You shouldn't be needing to iterate over all regions of a parent op, you know which one of those contains the current op.

[wsmoses]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

2699

The two region-based scopeMerge2 and scopeMerge3 tests should test this canonicalizer (and ensure it does and/or doesn't run when relevant).

[wsmoses]

Address comments

[wsmoses]

Fix bug

[ftynse]

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

296–298

I can't parse the "last non terminating block in a region", blocks are not terminators, ops are.

353

if (lastParentWithoutScope == nullptr) return failure();

362–363

while (lastParentWIthoutScope->getParentOp() && the-curent-condition) or factor it out to a variable

mlir/test/Dialect/MemRef/canonicalize.mlir

567

Nit: we likely don't care about these operations being on consecutive lines, so no need for -NEXT here and below.

[wsmoses]

Address comments

mlir/include/mlir/Conversion/Passes.td

 //===----------------------------------------------------------------------===//
 // SCFToOpenMP
 //===----------------------------------------------------------------------===//
 
 def ConvertSCFToOpenMP : Pass<"convert-scf-to-openmp", "ModuleOp"> {
   let summary = "Convert SCF parallel loop to OpenMP parallel + workshare "
                 "constructs.";
   let constructor = "mlir::createConvertSCFToOpenMPPass()";
-  let dependentDialects = ["omp::OpenMPDialect", "LLVM::LLVMDialect"];
+  let dependentDialects = ["omp::OpenMPDialect", "LLVM::LLVMDialect",
+                           "memref::MemRefDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // SCFToSPIRV
 //===----------------------------------------------------------------------===//
 
 def SCFToSPIRV : Pass<"convert-scf-to-spirv", "ModuleOp"> {
   let summary = "Convert SCF dialect to SPIR-V dialect.";

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     If `memref.alloca_scope` returns no value, the `memref.alloca_scope.return ` can
     be left out, and will be inserted implicitly.
   }];
 
   let results = (outs Variadic<AnyType>:$results);
   let regions = (region SizedRegion<1>:$bodyRegion);
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // AllocaScopeReturnOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_AllocaScopeReturnOp : MemRef_Op<"alloca_scope.return",
       [HasParent<"AllocaScopeOp">,

mlir/lib/Conversion/SCFToOpenMP/SCFToOpenMP.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "../PassDetail.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
LogicalResult matchAndRewrite(scf::ParallelOp parallelOp,
 
     // Allocate reduction variables. Make sure the we don't overflow the stack
     // with local `alloca`s by saving and restoring the stack pointer.
     Location loc = parallelOp.getLoc();
     Value one = rewriter.create<LLVM::ConstantOp>(
         loc, rewriter.getIntegerType(64), rewriter.getI64IntegerAttr(1));
     SmallVector<Value> reductionVariables;
     reductionVariables.reserve(parallelOp.getNumReductions());
-    Value token = rewriter.create<LLVM::StackSaveOp>(
-        loc, LLVM::LLVMPointerType::get(rewriter.getIntegerType(8)));
     for (Value init : parallelOp.getInitVals()) {
       assert((LLVM::isCompatibleType(init.getType()) ||
               init.getType().isa<LLVM::PointerElementTypeInterface>()) &&
              "cannot create a reduction variable if the type is not an LLVM "
              "pointer element");
       Value storage = rewriter.create<LLVM::AllocaOp>(
           loc, LLVM::LLVMPointerType::get(init.getType()), one, 0);
       rewriter.create<LLVM::StoreOp>(loc, init, storage);
for (auto pair :
       rewriter.setInsertionPoint(reduceOp);
       rewriter.replaceOpWithNewOp<omp::ReductionOp>(
           reduceOp, reduceOp.getOperand(), std::get<1>(pair));
     }
 
     // Create the parallel wrapper.
     auto ompParallel = rewriter.create<omp::ParallelOp>(loc);
     {
+
       OpBuilder::InsertionGuard guard(rewriter);
       rewriter.createBlock(&ompParallel.region());
 
-      // Replace SCF yield with OpenMP yield.
       {
-        OpBuilder::InsertionGuard innerGuard(rewriter);
-        rewriter.setInsertionPointToEnd(parallelOp.getBody());
-        assert(llvm::hasSingleElement(parallelOp.getRegion()) &&
-               "expected scf.parallel to have one block");
-        rewriter.replaceOpWithNewOp<omp::YieldOp>(
-            parallelOp.getBody()->getTerminator(), ValueRange());
-      }
+        auto scope = rewriter.create<memref::AllocaScopeOp>(parallelOp.getLoc(),
+                                                            TypeRange());
+        rewriter.create<omp::TerminatorOp>(loc);
+        OpBuilder::InsertionGuard allocaGuard(rewriter);
+        rewriter.createBlock(&scope.getBodyRegion());
+        rewriter.setInsertionPointToStart(&scope.getBodyRegion().front());
 
         // Replace the loop.
         auto loop = rewriter.create<omp::WsLoopOp>(
             parallelOp.getLoc(), parallelOp.getLowerBound(),
             parallelOp.getUpperBound(), parallelOp.getStep());
-      rewriter.create<omp::TerminatorOp>(loc);
+        rewriter.create<memref::AllocaScopeReturnOp>(loc);
 
         rewriter.inlineRegionBefore(parallelOp.getRegion(), loop.region(),
                                     loop.region().begin());
         if (!reductionVariables.empty()) {
           loop.reductionsAttr(
               ArrayAttr::get(rewriter.getContext(), reductionDeclSymbols));
           loop.reduction_varsMutable().append(reductionVariables);
         }
       }
+    }
 
     // Load loop results.
     SmallVector<Value> results;
     results.reserve(reductionVariables.size());
     for (Value variable : reductionVariables) {
       Value res = rewriter.create<LLVM::LoadOp>(loc, variable);
       results.push_back(res);
     }
     rewriter.replaceOp(parallelOp, results);
 
-    rewriter.create<LLVM::StackRestoreOp>(loc, token);
     return success();
   }
 };
 
 /// Applies the conversion patterns in the given function.
 static LogicalResult applyPatterns(ModuleOp module) {
   ConversionTarget target(*module.getContext());
   target.addIllegalOp<scf::ReduceOp, scf::ReduceReturnOp, scf::ParallelOp>();
-  target.addLegalDialect<omp::OpenMPDialect, LLVM::LLVMDialect>();
+  target.addLegalDialect<omp::OpenMPDialect, LLVM::LLVMDialect,
+                         memref::MemRefDialect>();
 
   RewritePatternSet patterns(module.getContext());
   patterns.add<ParallelOpLowering>(module.getContext());
   FrozenRewritePatternSet frozen(std::move(patterns));
   return applyPartialConversion(module, target, frozen);
 }
 
 /// A pass converting SCF operations to OpenMP operations.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace mlir::memref;
 
 /// Materialize a single constant operation from a given attribute value with
void AllocaScopeOp::getSuccessorRegions(
   if (index.hasValue()) {
     regions.push_back(RegionSuccessor(getResults()));
     return;
   }
 
   regions.push_back(RegionSuccessor(&bodyRegion()));
 }
 
+/// Given an operation, return whether this op is guaranteed to
+/// allocate an AutomaticAllocationScopeResource
+static bool isGuaranteedAutomaticAllocationScope(Operation *op) {
+  MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op);
+  if (!interface)
+    return false;
+  for (auto res : op->getResults()) {
+    if (auto effect =
+            interface.getEffectOnValue<MemoryEffects::Allocate>(res)) {
+      if (isa<SideEffects::AutomaticAllocationScopeResource>(
+              effect->getResource()))
+        return true;
+    }
+  }
+  return false;
+}
+
+/// Given an operation, return whether this op could to
+/// allocate an AutomaticAllocationScopeResource
+static bool isPotentialAutomaticAllocationScope(Operation *op) {
+  MemoryEffectOpInterface interface = dyn_cast<MemoryEffectOpInterface>(op);
+  if (!interface)
+    return true;
+  for (auto res : op->getResults()) {
+    if (auto effect =
+            interface.getEffectOnValue<MemoryEffects::Allocate>(res)) {
+      if (isa<SideEffects::AutomaticAllocationScopeResource>(
+              effect->getResource()))
+        return true;
+    }
+  }
+  return false;
+}
+
+/// Return whether this op is the last non terminating op
+/// in a region. That is to say, it is in a one-block region
+/// and is only followed by a terminator. This prevents
+/// extending the lifetime of allocations.
+static bool lastNonTerminatorInRegion(Operation *op) {
+  return op->getNextNode() == op->getBlock()->getTerminator() &&
+         op->getParentRegion()->getBlocks().size() == 1;
+}
+
+/// Inline an AllocaScopeOp if either the direct parent is an allocation scope
+/// or it contains no allocation.
+struct AllocaScopeInliner : public OpRewritePattern<AllocaScopeOp> {
+  using OpRewritePattern<AllocaScopeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AllocaScopeOp op,
+                                PatternRewriter &rewriter) const override {
+    if (!op->getParentOp()->hasTrait<OpTrait::AutomaticAllocationScope>()) {
+      bool hasPotentialAlloca =
+          op->walk([&](Operation *alloc) {
+              if (isPotentialAutomaticAllocationScope(alloc))
+                return WalkResult::interrupt();
+              return WalkResult::skip();
+            }).wasInterrupted();
+      if (hasPotentialAlloca)
+        return failure();
+    }
+
+    // Only apply to if this is this last non-terminator
+    // op in the block (lest lifetime be extended) of a one
+    // block region
+    if (!lastNonTerminatorInRegion(op))
+      return failure();
+
+    Block *block = &op.getRegion().front();
+    Operation *terminator = block->getTerminator();
+    ValueRange results = terminator->getOperands();
+    rewriter.mergeBlockBefore(block, op);
+    rewriter.replaceOp(op, results);
+    rewriter.eraseOp(terminator);
+    return success();
+  }
+};
+
+/// Move allocations into an allocation scope, if it is legal to
+/// move them (e.g. their operands are available at the location
+/// the op would be moved to).
+struct AllocaScopeHoister : public OpRewritePattern<AllocaScopeOp> {
+  using OpRewritePattern<AllocaScopeOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(AllocaScopeOp op,
+                                PatternRewriter &rewriter) const override {
+
+    if (!op->getParentWithTrait<OpTrait::AutomaticAllocationScope>())
+      return failure();
+
+    Operation *lastParentWithoutScope = op->getParentOp();
+
+    if (!lastParentWithoutScope ||
+        lastParentWithoutScope->hasTrait<OpTrait::AutomaticAllocationScope>())
+      return failure();
+
+    // Only apply to if this is this last non-terminator
+    // op in the block (lest lifetime be extended) of a one
+    // block region
+    if (!lastNonTerminatorInRegion(op) ||
+        !lastNonTerminatorInRegion(lastParentWithoutScope))
+      return failure();
+
+    while (!lastParentWithoutScope->getParentOp()
+                ->hasTrait<OpTrait::AutomaticAllocationScope>()) {
+      lastParentWithoutScope = lastParentWithoutScope->getParentOp();
+      if (!lastParentWithoutScope ||
+          !lastNonTerminatorInRegion(lastParentWithoutScope))
+        return failure();
+    }
+    Operation *scope = lastParentWithoutScope->getParentOp();
+    assert(scope->hasTrait<OpTrait::AutomaticAllocationScope>());
+
+    Region *containingRegion = nullptr;
+    for (auto &r : lastParentWithoutScope->getRegions()) {
+      if (r.isAncestor(op->getParentRegion())) {
+        assert(containingRegion == nullptr &&
+               "only one region can contain the op");
+        containingRegion = &r;
+      }
+    }
+    assert(containingRegion && "op must be contained in a region");
+
+    SmallVector<Operation *> toHoist;
+    op->walk([&](Operation *alloc) {
+      if (!isGuaranteedAutomaticAllocationScope(alloc))
+        return WalkResult::skip();
+
+      // If any operand is not defined before the location of
+      // lastParentWithoutScope (i.e. where we would hoist to), skip.
+      if (llvm::any_of(alloc->getOperands(), [&](Value v) {
+            return containingRegion->isAncestor(v.getParentRegion());
+          }))
+        return WalkResult::skip();
+      toHoist.push_back(alloc);
+      return WalkResult::advance();
+    });
+
+    if (!toHoist.size())
+      return failure();
+    rewriter.setInsertionPoint(lastParentWithoutScope);
+    for (auto op : toHoist) {
+      auto cloned = rewriter.clone(*op);
+      rewriter.replaceOp(op, cloned->getResults());
+    }
+    return success();
+  }
+};
+
+void AllocaScopeOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                                MLIRContext *context) {
+  results.add<AllocaScopeInliner, AllocaScopeHoister>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // AssumeAlignmentOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult AssumeAlignmentOp::verify() {
   if (!llvm::isPowerOf2_32(alignment()))
     return emitOpError("alignment must be power of 2");
   return success();

 OpFoldResult AtomicRMWOp::fold(ArrayRef<Attribute> operands) {
   /// atomicrmw(memrefcast) -> atomicrmw
   if (succeeded(foldMemRefCast(*this, value())))
     return getResult();
   return OpFoldResult();
 }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/MemRef/IR/MemRefOps.cpp.inc"

mlir/test/Conversion/SCFToOpenMP/reductions.mlir

 // CHECK: %[[RHS:.*]] = llvm.load %[[ARG1]]
 // CHECK: llvm.atomicrmw fadd %[[ARG0]], %[[RHS]] monotonic
 
 // CHECK-LABEL: @reduction1
 func @reduction1(%arg0 : index, %arg1 : index, %arg2 : index,
                  %arg3 : index, %arg4 : index) {
   // CHECK: %[[CST:.*]] = arith.constant 0.0
   // CHECK: %[[ONE:.*]] = llvm.mlir.constant(1
-  // CHECK: llvm.intr.stacksave
   // CHECK: %[[BUF:.*]] = llvm.alloca %[[ONE]] x f32
   // CHECK: llvm.store %[[CST]], %[[BUF]]
   %step = arith.constant 1 : index
   %zero = arith.constant 0.0 : f32
   // CHECK: omp.parallel
+  // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop
   // CHECK-SAME: reduction(@[[$REDF]] -> %[[BUF]]
   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
                             step (%arg4, %step) init (%zero) -> (f32) {
     // CHECK: %[[CST_INNER:.*]] = arith.constant 1.0
     %one = arith.constant 1.0 : f32
     // CHECK: omp.reduction %[[CST_INNER]], %[[BUF]]
     scf.reduce(%one) : f32 {
     ^bb0(%lhs : f32, %rhs: f32):
       %res = arith.addf %lhs, %rhs : f32
       scf.reduce.return %res : f32
     }
     // CHECK: omp.yield
   }
   // CHECK: omp.terminator
   // CHECK: llvm.load %[[BUF]]
-  // CHECK: llvm.intr.stackrestore
   return
 }
 
 // -----
 
 // Only check the declaration here, the rest is same as above.
 // CHECK: omp.reduction.declare @{{.*}} : f32
 
func @reduction4(%arg0 : index, %arg1 : index, %arg2 : index,
   // CHECK: %[[IONE:.*]] = arith.constant 1
   %ione = arith.constant 1 : i64
   // CHECK: %[[BUF1:.*]] = llvm.alloca %{{.*}} x f32
   // CHECK: llvm.store %[[ZERO]], %[[BUF1]]
   // CHECK: %[[BUF2:.*]] = llvm.alloca %{{.*}} x i64
   // CHECK: llvm.store %[[IONE]], %[[BUF2]]
 
   // CHECK: omp.parallel
+  // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop
   // CHECK-SAME: reduction(@[[$REDF1]] -> %[[BUF1]]
   // CHECK-SAME:           @[[$REDF2]] -> %[[BUF2]]
   %res:2 = scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
                         step (%arg4, %step) init (%zero, %ione) -> (f32, i64) {
     %one = arith.constant 1.0 : f32
     // CHECK: omp.reduction %{{.*}}, %[[BUF1]]
     scf.reduce(%one) : f32 {

mlir/test/Conversion/SCFToOpenMP/scf-to-openmp.mlir

 // RUN: mlir-opt -convert-scf-to-openmp %s | FileCheck %s
 
 // CHECK-LABEL: @parallel
 func @parallel(%arg0: index, %arg1: index, %arg2: index,
           %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
+  // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop (%[[LVAR1:.*]], %[[LVAR2:.*]]) : index = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
   scf.parallel (%i, %j) = (%arg0, %arg1) to (%arg2, %arg3) step (%arg4, %arg5) {
     // CHECK: "test.payload"(%[[LVAR1]], %[[LVAR2]]) : (index, index) -> ()
     "test.payload"(%i, %j) : (index, index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }
 
 // CHECK-LABEL: @nested_loops
 func @nested_loops(%arg0: index, %arg1: index, %arg2: index,
                    %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
+    // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop (%[[LVAR_OUT1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) {
   scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) {
     // CHECK: omp.parallel
+    // CHECK: memref.alloca_scope
     // CHECK: omp.wsloop (%[[LVAR_IN1:.*]]) : index = (%arg1) to (%arg3) step (%arg5) {
     scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) {
       // CHECK: "test.payload"(%[[LVAR_OUT1]], %[[LVAR_IN1]]) : (index, index) -> ()
       "test.payload"(%i, %j) : (index, index) -> ()
       // CHECK: }
     }
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }
 
 // CHECK-LABEL: @adjacent_loops
 func @adjacent_loops(%arg0: index, %arg1: index, %arg2: index,
                      %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel {
+  // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop (%[[LVAR_AL1:.*]]) : index = (%arg0) to (%arg2) step (%arg4) {
   scf.parallel (%i) = (%arg0) to (%arg2) step (%arg4) {
     // CHECK: "test.payload1"(%[[LVAR_AL1]]) : (index) -> ()
     "test.payload1"(%i) : (index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
 
   // CHECK: omp.parallel {
+  // CHECK: memref.alloca_scope
   // CHECK: omp.wsloop (%[[LVAR_AL2:.*]]) : index = (%arg1) to (%arg3) step (%arg5) {
   scf.parallel (%j) = (%arg1) to (%arg3) step (%arg5) {
     // CHECK: "test.payload2"(%[[LVAR_AL2]]) : (index) -> ()
     "test.payload2"(%j) : (index) -> ()
     // CHECK:   omp.yield
     // CHECK: }
   }
   // CHECK:   omp.terminator
   // CHECK: }
   return
 }

mlir/test/Dialect/MemRef/canonicalize.mlir

 
 func @self_copy(%m1: memref<?xf32>) {
   memref.copy %m1, %m1 : memref<?xf32> to memref<?xf32>
   return
 }
 
 // CHECK-LABEL: func @self_copy
 //  CHECK-NEXT:   return
+
+// -----
+
+func @scopeMerge() {
+  memref.alloca_scope {
+    %cnt = "test.count"() : () -> index
+    %a = memref.alloca(%cnt) : memref<?xi64>
+    "test.use"(%a) : (memref<?xi64>) -> ()
+  }
+  return
+}
+// CHECK:   func @scopeMerge() {
+// CHECK-NOT: alloca_scope
+// CHECK:     %[[cnt:.+]] = "test.count"() : () -> index
+// CHECK:     %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
+// CHECK:     "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
+// CHECK:     return
+
+func @scopeMerge2() {
+  "test.region"() ({
+    memref.alloca_scope {
+      %cnt = "test.count"() : () -> index
+      %a = memref.alloca(%cnt) : memref<?xi64>
+      "test.use"(%a) : (memref<?xi64>) -> ()
+    }
+    "test.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK:   func @scopeMerge2() {
+// CHECK:     "test.region"() ({
+// CHECK:       memref.alloca_scope {
+// CHECK:         %[[cnt:.+]] = "test.count"() : () -> index
+// CHECK:         %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
+// CHECK:         "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
+// CHECK:       }
+// CHECK:       "test.terminator"() : () -> ()
+// CHECK:     }) : () -> ()
+// CHECK:     return
+// CHECK:   }
+
+func @scopeMerge3() {
+  %cnt = "test.count"() : () -> index
+  "test.region"() ({
+    memref.alloca_scope {
+      %a = memref.alloca(%cnt) : memref<?xi64>
+      "test.use"(%a) : (memref<?xi64>) -> ()
+    }
+    "test.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK:   func @scopeMerge3() {
+// CHECK:     %[[cnt:.+]] = "test.count"() : () -> index
+// CHECK:     %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
+// CHECK:     "test.region"() ({
+// CHECK:       memref.alloca_scope {
+// CHECK:         "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
+// CHECK:       }
+// CHECK:       "test.terminator"() : () -> ()
+// CHECK:     }) : () -> ()
+// CHECK:     return
+// CHECK:   }
+
+func @scopeMerge4() {
+  %cnt = "test.count"() : () -> index
+  "test.region"() ({
+    memref.alloca_scope {
+      %a = memref.alloca(%cnt) : memref<?xi64>
+      "test.use"(%a) : (memref<?xi64>) -> ()
+    }
+    "test.op"() : () -> ()
+    "test.terminator"() : () -> ()
+  }) : () -> ()
+  return
+}
+
+// CHECK:   func @scopeMerge4() {
+// CHECK:     %[[cnt:.+]] = "test.count"() : () -> index
+// CHECK:     "test.region"() ({
+// CHECK:       memref.alloca_scope {
+// CHECK:         %[[alloc:.+]] = memref.alloca(%[[cnt]]) : memref<?xi64>
+// CHECK:         "test.use"(%[[alloc]]) : (memref<?xi64>) -> ()
+// CHECK:       }
+// CHECK:       "test.op"() : () -> ()
+// CHECK:       "test.terminator"() : () -> ()
+// CHECK:     }) : () -> ()
+// CHECK:     return
+// CHECK:   }