[mlir] Fixed alighnment attribute of alloc constant folding.

[imaihal]

Summary

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

Signed-off-by: Haruki Imai <[e-mail]>

[imaihal]

@mehdi_amini Thanks for review. I don't have commit access. Could you commit this patch?

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

LogicalResult matchAndRewrite(AllocLikeOp alloc,
 
     // Create new memref type (which will have fewer dynamic dimensions).
     MemRefType newMemRefType =
         MemRefType::Builder(memrefType).setShape(newShapeConstants);
     assert(static_cast<int64_t>(newOperands.size()) ==
            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());
 
     rewriter.replaceOp(alloc, {resultCast});
     return success();
   }
 };

mlir/test/Transforms/canonicalize.mlir

func @alloc_const_fold() -> memref<?xf32> {
   %c4 = constant 4 : index
   %a = memref.alloc(%c4) : 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 @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
   %c4 = constant 4 : index
   %a = memref.alloc(%c4) : memref<?xf32>
   return
 }