[mlir] Add alignment attribute to memref.global

[ezhulenev]

Summary

Revived https://reviews.llvm.org/D102435

Add alignment attribute to memref.global and propagate it to llvm global in memref->llvm lowering

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     given global variable will always return the same memref descriptor).
 
     Example:
 
     ```mlir
     // Private variable with an initial value.
     memref.global "private" @x : memref<2xf32> = dense<0.0,2.0>
 
+    // Private variable with an initial value and an alignment (power of 2).
+    memref.global "private" @x : memref<2xf32> = dense<0.0,2.0> {alignment = 64}
+
     // Declaration of an external variable.
     memref.global "private" @y : memref<4xi32>
 
     // Uninitialized externally visible variable.
     memref.global @z : memref<3xf16> = uninitialized
 
     // Externally visible constant variable.
     memref.global constant @c : memref<2xi32> = dense<1, 4>
     ```
   }];
 
   let arguments = (ins
       SymbolNameAttr:$sym_name,
       OptionalAttr<StrAttr>:$sym_visibility,
       MemRefTypeAttr:$type,
       OptionalAttr<AnyAttr>:$initial_value,
-      UnitAttr:$constant
+      UnitAttr:$constant,
+      OptionalAttr<I64Attr>:$alignment
   );
 
   let assemblyFormat = [{
        ($sym_visibility^)?
        (`constant` $constant^)?
        $sym_name `:`
        custom<GlobalMemrefOpTypeAndInitialValue>($type, $initial_value)
        attr-dict

mlir/lib/Conversion/MemRefToLLVM/MemRefToLLVM.cpp

if (!global.isExternal() && !global.isUninitialized()) {
       initialValue = elementsAttr;
 
       // For scalar memrefs, the global variable created is of the element type,
       // so unpack the elements attribute to extract the value.
       if (type.getRank() == 0)
         initialValue = elementsAttr.getValue({});
     }
 
+    uint64_t alignment = global.alignment().getValueOr(0);
+
     auto newGlobal = rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         global, arrayTy, global.constant(), linkage, global.sym_name(),
-        initialValue, /*alignment=*/0, type.getMemorySpaceAsInt());
+        initialValue, alignment, type.getMemorySpaceAsInt());
     if (!global.isExternal() && global.isUninitialized()) {
       Block *blk = new Block();
       newGlobal.getInitializerRegion().push_back(blk);
       rewriter.setInsertionPointToStart(blk);
       Value undef[] = {
           rewriter.create<LLVM::UndefOp>(global.getLoc(), arrayTy)};
       rewriter.create<LLVM::ReturnOp>(global.getLoc(), undef);
     }

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

if (initValue.isa<ElementsAttr>()) {
       Type initType = initValue.getType();
       Type tensorType = getTensorTypeFromMemRefType(memrefType);
       if (initType != tensorType)
         return op.emitOpError("initial value expected to be of type ")
                << tensorType << ", but was of type " << initType;
     }
   }
 
+  if (Optional<uint64_t> alignAttr = op.alignment()) {
+    uint64_t alignment = alignAttr.getValue();
+
+    if (!llvm::isPowerOf2_64(alignment))
+      return op->emitError() << "alignment attribute value " << alignment
+                             << " is not a power of 2";
+  }
+
   // TODO: verify visibility for declarations.
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // GetGlobalOp
 //===----------------------------------------------------------------------===//
 

mlir/lib/Dialect/StandardOps/Transforms/TensorConstantBufferize.cpp

memref::GlobalOp GlobalCreator::getGlobalFor(ConstantOp constantOp) {
   interleave(type.getShape(), os, "x");
   os << "x" << type.getElementType();
 
   auto global = globalBuilder.create<memref::GlobalOp>(
       constantOp.getLoc(), (Twine("__constant_") + os.str()).str(),
       /*sym_visibility=*/globalBuilder.getStringAttr("private"),
       /*type=*/typeConverter.convertType(type).cast<MemRefType>(),
       /*initial_value=*/constantOp.getValue().cast<ElementsAttr>(),
-      /*constant=*/true);
+      /*constant=*/true,
+      /*alignment=*/IntegerAttr());
   symbolTable.insert(global);
   // The symbol table inserts at the end of the module, but globals are a bit
   // nicer if they are at the beginning.
   global->moveBefore(&moduleOp.front());
   globals[constantOp.getValue()] = global;
   return global;
 }
 

mlir/test/Conversion/MemRefToLLVM/memref-to-llvm.mlir

func @get_gv3_memref() {
   // CHECK: llvm.insertvalue %[[DEADBEEFPTR]], {{.*}}[0] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   // CHECK: llvm.insertvalue %[[GEP]], {{.*}}[1] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   // CHECK: %[[OFFSET:.*]] = llvm.mlir.constant(0 : index) : i64
   // CHECK: llvm.insertvalue %[[OFFSET]], {{.*}}[2] : !llvm.struct<(ptr<f32>, ptr<f32>, i64)>
   %0 = memref.get_global @gv3 : memref<f32>
   return
 }
 
+// Test scalar memref with an alignment.
+// CHECK: llvm.mlir.global private @gv4(1.000000e+00 : f32) {alignment = 64 : i64} : f32
+memref.global "private" @gv4 : memref<f32> = dense<1.0> {alignment = 64}
+
 // -----
 
 func @collapse_shape_static(%arg0: memref<1x3x4x1x5xf32>) -> memref<3x4x5xf32> {
   %0 = memref.collapse_shape %arg0 [[0, 1], [2], [3, 4]] :
     memref<1x3x4x1x5xf32> into memref<3x4x5xf32>
   return %0 : memref<3x4x5xf32>
 }
 // CHECK-LABEL: func @collapse_shape_static

mlir/test/Dialect/MemRef/invalid.mlir

 func @mismatched_types() {
   // expected-error @+1 {{result type 'memref<3xf32>' does not match type 'memref<3xi32>' of the global memref @gv}}
   %0 = memref.get_global @gv : memref<3xf32>
   return
 }
 
 // -----
 
+// expected-error @+1 {{alignment attribute value 63 is not a power of 2}}
+memref.global "private" @gv : memref<4xf32> = dense<1.0> { alignment = 63 }
+
+// -----
+
 func @copy_different_shape(%arg0: memref<2xf32>, %arg1: memref<3xf32>) {
   // expected-error @+1 {{op requires the same shape for all operands}}
   memref.copy %arg0, %arg1 : memref<2xf32> to memref<3xf32>
   return
 }
 
 // -----