Add alignment attribute to memref.global

[dpotop]

Summary

I have added an alignment attribute to memref.global.
Naturally, I have also:

• changed the documentation of memref.global to reflect this
• added verification code that checks that the attribute value (if present) is a positive power of 2
• added code to StandardToLLVM.cpp ensuring that the alignment attribute is correctly passed during lowering.
• added one test case for the lowering to LLVM and one test case checking that the verification code functions.

[rriddle]

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

2303–2308

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

1117–1118

1119–1120

1121

Print what the alignment is here.

[dpotop]

Add alignment attribute to memref.global

Applying code simplifications suggested by rriddle.

[dpotop]

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

1119–1120

Since I have to also print it, I preferred first extracting it to avoid calling the * operator twice.

[dpotop]

Executing git pull --rebase to correct a testing error.

[dpotop]

Minor formatting based on git-clang-format.

[ftynse]

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

508–509

Nit: I wouldn't expect the documentation of an op to describe how it is lowered to a specific dialect.

mlir/test/Conversion/StandardToLLVM/standard-to-llvm.mlir

243

Nit: we use "translation" when we create a different IR, e.g. LLVM IR, and "conversion" when we stay within MLIR.

[ftynse]

This has been subsumed by another patch and landed.

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

let description = [{
     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>
     ```
+
+    A global variable may specify an alignment (a power of 2 positive
+    integer). This value is passed during lowering to the `llvm.mlir.global`
+    operation.
+
+    Example:
+
+    ```mlir
+    // Require alignment on 64 byte barrier
+    memref.global "private" @x : memref<2xf32> = dense<0.0,2.0>  {alignment = 64}
   }];
 
   let arguments = (ins
       SymbolNameAttr:$sym_name,
       OptionalAttr<StrAttr>:$sym_visibility,
       TypeAttr:$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/StandardToLLVM/StandardToLLVM.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);
+
     rewriter.replaceOpWithNewOp<LLVM::GlobalOp>(
         global, arrayTy, global.constant(), linkage, global.sym_name(),
-        initialValue, /*alignment=*/0, type.getMemorySpaceAsInt());
+        initialValue, alignment, type.getMemorySpaceAsInt());
     return success();
   }
 };
 
 /// GetGlobalMemrefOp is lowered into a Memref descriptor with the pointer to
 /// the first element stashed into the descriptor. This reuses
 /// `AllocLikeOpLowering` to reuse the Memref descriptor construction.
 struct GetGlobalMemrefOpLowering : public AllocLikeOpLowering {
   GetGlobalMemrefOpLowering(LLVMTypeConverter &converter)

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 value = alignAttr.getValue();
+    if (!llvm::isPowerOf2_64(value))
+      return op->emitError()
+             << "alignment attribute value " << value << " 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),
       /*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/StandardToLLVM/standard-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 global variable alignment translation
+// CHECK: llvm.mlir.global private @gv4(dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]> : tensor<4xf32>) {alignment = 64 : i64} : !llvm.array<4 x f32>
+memref.global "private" @gv4 : memref<4xf32> = dense<[0.0, 1.0, 2.0, 3.0]> { alignment = 64 }
+
+
 // This should not trigger an assertion by creating an LLVM::CallOp with a
 // nullptr result type.
 
 // CHECK-LABEL: @call_zero_result_func
 func @call_zero_result_func() {
   // CHECK: call @zero_result_func
   call @zero_result_func() : () -> ()
   return

mlir/test/Dialect/MemRef/invalid.mlir

 
 // -----
 
 // expected-error @+1 {{inferred shape of elements literal ([2]) does not match type ([2, 2])}}
 memref.global @foo : memref<2x2xf32> = dense<[0.0, 1.0]>
 
 // -----
 
+// expected-error @+1 {{alignment attribute value 63 is not a power of 2}}
+memref.global "private" @foo : memref<4xf32> = dense<[0.0, 1.0, 2.0, 3.0]> { alignment = 63 }
+
+// -----
+
 // expected-error @+1 {{expected valid '@'-identifier for symbol name}}
 memref.global "private" "public" @foo : memref<2x2xf32>  = "foo"
 
 // -----
 
 // expected-error @+1 {{expected valid '@'-identifier for symbol name}}
 memref.global constant external @foo : memref<2x2xf32>  = "foo"