[mlir] Add alignment option to constant tensor bufferization pass

[bkramer]

looks good

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.h

 
 /// Add patterns to bufferize tensor constants into global memrefs to the given
 /// pattern list.
 void populateTensorConstantBufferizePatterns(
     GlobalCreator &globalCreator, BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns);
 
 /// Creates an instance of tensor constant bufferization pass.
-std::unique_ptr<Pass> createTensorConstantBufferizePass();
+std::unique_ptr<Pass> createTensorConstantBufferizePass(unsigned alignment = 0);
 
 /// Creates an instance of the StdExpand pass that legalizes Std
 /// dialect ops to be convertible to LLVM. For example,
 /// `std.ceildivi_signed` gets transformed to a number of std operations,
 /// which can be lowered to LLVM; `memref.reshape` gets converted to
 /// `memref_reinterpret_cast`.
 std::unique_ptr<Pass> createStdExpandOpsPass();
 

mlir/include/mlir/Dialect/StandardOps/Transforms/Passes.td

let description = [{
 
     This pass needs to be a module pass because it inserts memref.global
     ops into the module, which cannot be done safely from a function pass due to
     multi-threading. Most other bufferization passes can run in parallel at
     function granularity.
   }];
   let constructor = "mlir::createTensorConstantBufferizePass()";
   let dependentDialects = ["memref::MemRefDialect"];
+  let options = [
+    Option<"alignment", "alignment", "unsigned", /*default=*/"0",
+           "Create global memrefs with a specified alignment">,
+  ];
 }
 
 #endif // MLIR_DIALECT_STANDARD_TRANSFORMS_PASSES

mlir/include/mlir/Transforms/BufferUtils.h

 class GlobalOp;
 } // namespace memref
 
 // Support class to create global ops for tensor-valued constants in the
 // program. Globals are created lazily at the top of the `moduleOp` with pretty
 // names. Duplicates are avoided.
 class GlobalCreator {
 public:
-  explicit GlobalCreator(ModuleOp module) : moduleOp(module) {}
+  GlobalCreator(ModuleOp module, unsigned alignment = 0)
+      : moduleOp(module), alignment(alignment) {}
   memref::GlobalOp getGlobalFor(ConstantOp constantOp);
 
 private:
   ModuleOp moduleOp;
+  unsigned alignment;
   // This could use memref::GlobalOp key but we avoid introducing a new
   // dependence to the memref dialect for this.
   DenseMap<Attribute, Operation *> globals;
 };
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERUTILS_H

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

memref::GlobalOp GlobalCreator::getGlobalFor(ConstantOp constantOp) {
   SymbolTable symbolTable(moduleOp);
 
   // Create a pretty name.
   SmallString<64> buf;
   llvm::raw_svector_ostream os(buf);
   interleave(type.getShape(), os, "x");
   os << "x" << type.getElementType();
 
+  // Add an optional alignment to the global memref.
+  IntegerAttr memrefAlignment =
+      alignment > 0 ? IntegerAttr::get(globalBuilder.getI64Type(), alignment)
+                    : IntegerAttr();
+
   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,
-      /*alignment=*/IntegerAttr());
+      /*alignment=*/memrefAlignment);
   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;
 }
 

 void mlir::populateTensorConstantBufferizePatterns(
     GlobalCreator &globalCreator, BufferizeTypeConverter &typeConverter,
     RewritePatternSet &patterns) {
   patterns.add<BufferizeTensorConstantOp>(globalCreator, typeConverter,
                                           patterns.getContext());
 }
 
 namespace {
-struct TensorConstantBufferizePass
+class TensorConstantBufferizePass
     : public TensorConstantBufferizeBase<TensorConstantBufferizePass> {
+public:
+  explicit TensorConstantBufferizePass(unsigned alignment) {
+    if (alignment)
+      this->alignment = alignment;
+  }
+
   void runOnOperation() override {
     auto module = getOperation();
-    GlobalCreator globals(module);
+    GlobalCreator globals(module, alignment);
 
     auto *context = &getContext();
     BufferizeTypeConverter typeConverter;
     RewritePatternSet patterns(context);
     ConversionTarget target(*context);
 
     target.addLegalDialect<memref::MemRefDialect>();
     populateTensorConstantBufferizePatterns(globals, typeConverter, patterns);
     target.addDynamicallyLegalOp<ConstantOp>(
         [&](ConstantOp op) { return typeConverter.isLegal(op.getType()); });
     if (failed(applyPartialConversion(module, target, std::move(patterns))))
       signalPassFailure();
   }
 };
 } // namespace
 
-std::unique_ptr<Pass> mlir::createTensorConstantBufferizePass() {
-  return std::make_unique<TensorConstantBufferizePass>();
+std::unique_ptr<Pass>
+mlir::createTensorConstantBufferizePass(unsigned alignment) {
+  return std::make_unique<TensorConstantBufferizePass>(alignment);
 }

mlir/test/Dialect/Standard/tensor-constant-bufferize.mlir

 // RUN: mlir-opt %s -tensor-constant-bufferize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -tensor-constant-bufferize=alignment=64 -split-input-file | FileCheck --check-prefix=ALIGNED %s
 
 // CHECK-LABEL: module {
+
 // We check the debug name too since we put some effort into making that readable.
 // The name isn't load-bearing though.
+
 // CHECK: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
+// CHECK-NOT: alignment
+
+// ALIGNED: memref.global "private" constant @__constant_3x4xf32 : memref<3x4xf32> = dense<7.000000e+00>
+// ALIGNED-SAME: {alignment = 64 : i64}
+
 // CHECK: @basic
 func @basic() -> tensor<3x4xf32> {
   // CHECK: %[[MEMREF:.*]] = memref.get_global @__constant_3x4xf32 : memref<3x4xf32>
   // CHECK: %[[TENSOR:.*]] = memref.tensor_load %[[MEMREF]]
   %0 = constant dense<7.0> : tensor<3x4xf32>
   // CHECK: return %[[TENSOR]]
   return %0 : tensor<3x4xf32>
 }