From 39ee9fd8c1c2424f76260f62ac223e3aa1760612 Mon Sep 17 00:00:00 2001
From: Haruki Imai <imaihal@jp.ibm.com>
Date: Wed, 7 Apr 2021 19:17:19 +0000
Subject: [PATCH] [mlir] Fixed alignment attribute of alloc constant folding.

When allocLikeOp is updated in alloc constant folding,
alighnment attribute was ignored. This patch fixes it.

Signed-off-by: Haruki Imai <imaihal@jp.ibm.com>

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D99882
---
 mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp |  4 ++--
 mlir/test/Transforms/canonicalize.mlir   | 11 +++++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index ffb5bcf..7e55f4c 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -184,8 +184,8 @@ struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
            newMemRefType.getNumDynamicDims());
 
     // Create and insert the alloc op for the new memref.
-    auto newAlloc = rewriter.create<AllocLikeOp>(alloc.getLoc(), newMemRefType,
-                                                 newOperands, IntegerAttr());
+    auto newAlloc = rewriter.create<AllocLikeOp>(
+        alloc.getLoc(), newMemRefType, newOperands, alloc.alignmentAttr());
     // Insert a cast so we have the same type as the old alloc.
     auto resultCast =
         rewriter.create<CastOp>(alloc.getLoc(), newAlloc, alloc.getType());
diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir
index ba04232..988f7ba 100644
--- a/mlir/test/Transforms/canonicalize.mlir
+++ b/mlir/test/Transforms/canonicalize.mlir
@@ -368,6 +368,17 @@ func @alloc_const_fold() -> memref<?xf32> {
   return %a : memref<?xf32>
 }
 
+// CHECK-LABEL: func @alloc_alignment_const_fold
+func @alloc_alignment_const_fold() -> memref<?xf32> {
+  // CHECK-NEXT: %0 = memref.alloc() {alignment = 4096 : i64} : memref<4xf32>
+  %c4 = constant 4 : index
+  %a = memref.alloc(%c4) {alignment = 4096 : i64} : memref<?xf32>
+
+  // CHECK-NEXT: %1 = memref.cast %0 : memref<4xf32> to memref<?xf32>
+  // CHECK-NEXT: return %1 : memref<?xf32>
+  return %a : memref<?xf32>
+}
+
 // CHECK-LABEL: func @dead_alloc_fold
 func @dead_alloc_fold() {
   // CHECK-NEXT: return
-- 
2.7.4