[Mlir-commits] [mlir] 39ee9fd - [mlir] Fixed alignment attribute of alloc constant folding.

[Mehdi Amini]

Author: Haruki Imai
Date: 2021-04-07T19:28:49Z
New Revision: 39ee9fd8c1c2424f76260f62ac223e3aa1760612

URL: https://github.com/llvm/llvm-project/commit/39ee9fd8c1c2424f76260f62ac223e3aa1760612
DIFF: https://github.com/llvm/llvm-project/commit/39ee9fd8c1c2424f76260f62ac223e3aa1760612.diff

LOG: [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 at jp.ibm.com>

Reviewed By: mehdi_amini

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

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Transforms/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index ffb5bcf59b7dc..7e55f4c2877b4 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 ba042327dac41..988f7ba3bfc13 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