[mlir] Set alignment in AllocOp of normalizeMemref()

[imaihal]

Summary

AllocOp is updated in normalizeMemref(AllocOp allocOp), but, when the
AllocOp has alignment attribute, it was ignored and updated AllocOp
does not have alignment attribute. This patch fixes it.

[mehdi_amini]

mlir/lib/Transforms/Utils/Utils.cpp

454

Can you remove braces here? The block is trivial

455

Can't there be a value() argument to the alloc in this case as well?

Actually do we need a if or can we always write the first form and the alignmentAttr() will just be null?

[bondhugula]

mlir/lib/Transforms/Utils/Utils.cpp

449

Nit: Please terminate comment with a period. They also have to be complete sentences.

450–455

The following will instead work:

auto newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType, llvm::None,
                                    allocOp.alignmentAttr());

You don't need value() - it only works for static memrefs. If one extends this for dynamic memrefs, one simply can't pass allocOp.value() here.

mlir/test/Transforms/memref-normalize.mlir

150

No need of %{{.*}} =

[imaihal]

Remove redundant code

[imaihal]

Thanks for your comments! I updated.

mlir/lib/Transforms/Utils/Utils.cpp

449

Code becomes simple. so I removed the comment.

450–455

Thanks for your comments! I didn't know llvm:None works. I used value() because this function currently only support static memrefs, but I changed it.

[bondhugula]

Thanks for fixing this.

[bondhugula]

mlir/test/Transforms/memref-normalize.mlir

150

Can you use CHECK-NEXT here?

[imaihal]

mlir/test/Transforms/memref-normalize.mlir

150

Do you mean I need to just replace CHECK with CHECK-NEXT ? Yes, I can do, but why CHECK-NEXT is better?

[imaihal]

Change CHECK to CHECK-NEXT

[imaihal]

Rebase

[imaihal]

Rebase again

[imaihal]

Rebase

[imaihal]

@bondhugula Thanks for your review! I reflected your comments.
I think I don't have right to land this patch. Could you land this if you don't have additional comments?

[bondhugula]

In D83656#2161177, @imaihal wrote:

@bondhugula Thanks for your review! I reflected your comments.
I think I don't have right to land this patch. Could you land this if you don't have additional comments?

Can you check "Done" on the pending comment?

[bondhugula]

mlir/test/Transforms/memref-normalize.mlir

150

CHECK-NEXT would be more stringent - not allowing anything else in between, which is the case here since you don't expect anything in between.

[imaihal]

mlir/test/Transforms/memref-normalize.mlir

150

I see. Thanks!

[bondhugula]

Su

In D83656#2161269, @bondhugula wrote:

In D83656#2161177, @imaihal wrote:

@bondhugula Thanks for your review! I reflected your comments.
I think I don't have right to land this patch. Could you land this if you don't have additional comments?

Sure, I'll commit it.

mlir/lib/Transforms/Utils/Utils.cpp

LogicalResult mlir::normalizeMemRef(AllocOp allocOp) {
 
   auto oldMemRef = allocOp.getResult();
   SmallVector<Value, 4> symbolOperands(allocOp.getSymbolicOperands());
 
   MemRefType newMemRefType =
       MemRefType::Builder(memrefType)
           .setShape(newShape)
           .setAffineMaps(b.getMultiDimIdentityMap(newRank));
-  auto newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType);
+
+  auto newAlloc = b.create<AllocOp>(allocOp.getLoc(), newMemRefType, llvm::None,
+                                    allocOp.alignmentAttr());
 
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
                                       /*extraIndices=*/{},
                                       /*indexRemap=*/layoutMap,
                                       /*extraOperands=*/{},
                                       /*symbolOperands=*/symbolOperands))) {
     // If it failed (due to escapes for example), bail out.
     newAlloc.erase();
     return failure();
   }
   // Replace any uses of the original alloc op and erase it. All remaining uses
   // have to be dealloc's; RAMUW above would've failed otherwise.
   assert(llvm::all_of(oldMemRef.getUsers(),
                       [](Operation *op) { return isa<DeallocOp>(op); }));
   oldMemRef.replaceAllUsesWith(newAlloc);
   allocOp.erase();
   return success();
 }

mlir/test/Transforms/memref-normalize.mlir

func @semi_affine_layout_map(%s0: index, %s1: index) {
   affine.for %i = 0 to 256 {
     affine.for %j = 0 to 1024 {
       // CHECK: memref<256x1024xf32, #map{{[0-9]+}}>
       affine.load %A[%i, %j] : memref<256x1024xf32, affine_map<(d0, d1)[s0, s1] -> (d0*s0 + d1*s1)>>
     }
   }
   return
 }
+
+// CHECK-LABEL: func @alignment
+func @alignment() {
+  %A = alloc() {alignment = 32 : i64}: memref<64x128x256xf32, affine_map<(d0, d1, d2) -> (d2, d0, d1)>>
+  // CHECK-NEXT: alloc() {alignment = 32 : i64} : memref<256x64x128xf32>
+  return
+}