[mlir] PromoteBuffersToStackPass - Copy attributes of original AllocOp

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D143185

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
index 5710d3a..83b2ef6 100644 (file)
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferOptimizations.cpp
@@ -356,7 +356,7 @@ public:
       Operation *allocOp = alloc.getDefiningOp();
       Operation *alloca = builder.create<memref::AllocaOp>(
           alloc.getLoc(), alloc.getType().cast<MemRefType>(),
-          allocOp->getOperands());
+          allocOp->getOperands(), allocOp->getAttrs());
 
       // Replace the original alloc by a newly created alloca.
       allocOp->replaceAllUsesWith(alloca);

diff --git a/mlir/test/Transforms/promote-buffers-to-stack.mlir b/mlir/test/Transforms/promote-buffers-to-stack.mlir
index 243fbb1..f7f2d2e 100644 (file)
--- a/mlir/test/Transforms/promote-buffers-to-stack.mlir
+++ b/mlir/test/Transforms/promote-buffers-to-stack.mlir
@@ -603,11 +603,11 @@ func.func @indexElementType() {
 // CHECK-LABEL: func @bigIndexElementType
 module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<index, 256>>} {
   func.func @bigIndexElementType() {
-    %0 = memref.alloc() : memref<4xindex>
+    %0 = memref.alloc() {alignment = 64 : i64, custom_attr} : memref<4xindex>
     return
   }
 }
-// DEFINDEX-NEXT: memref.alloca()
-// LOWLIMIT-NEXT: memref.alloc()
-// RANK-NEXT: memref.alloca()
+// DEFINDEX-NEXT: memref.alloca() {alignment = 64 : i64, custom_attr}
+// LOWLIMIT-NEXT: memref.alloc() {alignment = 64 : i64, custom_attr}
+// RANK-NEXT: memref.alloca() {alignment = 64 : i64, custom_attr}
 // CHECK-NEXT: return