[mlir] Support bufferization of arith.constant to memref.global with memory space

[mamrami]

rebase

[mamrami]

Fix bug

[springerm]

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h

129

nit: Can we make this optional: Attribute memorySpace = {}

[mamrami]

CR fix

mlir/include/mlir/Dialect/Bufferization/Transforms/BufferUtils.h

protected:
   /// about alloc and dealloc positions.
   Liveness liveness;
 };
 
 // Create a global op for the given tensor-valued constant in the program.
 // Globals are created lazily at the top of the enclosing ModuleOp with pretty
 // names. Duplicates are avoided.
 FailureOr<memref::GlobalOp> getGlobalFor(arith::ConstantOp constantOp,
-                                         uint64_t alignment);
+                                         uint64_t alignment,
+                                         Attribute memorySpace = {});
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERUTILS_H

mlir/lib/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.cpp

 //===- BufferizableOpInterfaceImpl.cpp - Impl. of BufferizableOpInterface -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/Transforms/BufferizableOpInterfaceImpl.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/Transforms/BufferUtils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Operation.h"
 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 namespace {
 /// Bufferization of arith.constant. Replace with memref.get_global.
 struct ConstantOpInterface
     : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     arith::ConstantOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto constantOp = cast<arith::ConstantOp>(op);
 
-    // TODO: Implement memory space for this op. E.g., by adding a memory_space
-    // attribute to ConstantOp.
-    if (options.defaultMemorySpace != Attribute())
-      return op->emitError("memory space not implemented yet");
+    Attribute memorySpace;
+    if (options.defaultMemorySpace.has_value())
+      memorySpace = *options.defaultMemorySpace;
+    else
+      return constantOp->emitError("could not infer memory space");
 
     // Only ranked tensors are supported.
     if (!constantOp.getType().isa<RankedTensorType>())
       return failure();
 
     // Only constants inside a module are supported.
     auto moduleOp = constantOp->getParentOfType<ModuleOp>();
     if (!moduleOp)
       return failure();
 
     // Create global memory segment and replace tensor with memref pointing to
     // that memory segment.
     FailureOr<memref::GlobalOp> globalOp =
-        getGlobalFor(constantOp, options.bufferAlignment);
+        getGlobalFor(constantOp, options.bufferAlignment, memorySpace);
     if (failed(globalOp))
       return failure();
     memref::GlobalOp globalMemref = *globalOp;
     replaceOpWithNewBufferizedOp<memref::GetGlobalOp>(
         rewriter, op, globalMemref.getType(), globalMemref.getName());
 
     return success();
   }

mlir/lib/Dialect/Bufferization/Transforms/BufferUtils.cpp

bool BufferPlacementTransformationBase::isLoop(Operation *op) {
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementTransformationBase
 //===----------------------------------------------------------------------===//
 
 FailureOr<memref::GlobalOp>
-bufferization::getGlobalFor(arith::ConstantOp constantOp, uint64_t alignment) {
+bufferization::getGlobalFor(arith::ConstantOp constantOp, uint64_t alignment,
+                            Attribute memorySpace) {
   auto type = constantOp.getType().cast<RankedTensorType>();
   auto moduleOp = constantOp->getParentOfType<ModuleOp>();
   if (!moduleOp)
     return failure();
 
   // If we already have a global for this constant value, no need to do
   // anything else.
   for (Operation &op : moduleOp.getRegion().getOps()) {
bufferization::getGlobalFor(arith::ConstantOp constantOp, uint64_t alignment,
   os << "x" << type.getElementType();
 
   // Add an optional alignment to the global memref.
   IntegerAttr memrefAlignment =
       alignment > 0 ? IntegerAttr::get(globalBuilder.getI64Type(), alignment)
                     : IntegerAttr();
 
   BufferizeTypeConverter typeConverter;
+  auto memrefType = typeConverter.convertType(type).cast<MemRefType>();
+  if (memorySpace)
+    memrefType = MemRefType::Builder(memrefType).setMemorySpace(memorySpace);
   auto global = globalBuilder.create<memref::GlobalOp>(
       constantOp.getLoc(), (Twine("__constant_") + os.str()).str(),
       /*sym_visibility=*/globalBuilder.getStringAttr("private"),
-      /*type=*/typeConverter.convertType(type).cast<MemRefType>(),
+      /*type=*/memrefType,
       /*initial_value=*/constantOp.getValue().cast<ElementsAttr>(),
       /*constant=*/true,
       /*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());
   return global;
 }

mlir/test/Dialect/Arith/one-shot-bufferize-memory-space-invalid.mlir

 // RUN: mlir-opt %s -one-shot-bufferize="must-infer-memory-space" -split-input-file -verify-diagnostics
 
 func.func @inconsistent_memory_space_arith_select(%c: i1) -> tensor<10xf32> {
   // Selecting tensors with different memory spaces. Such IR cannot be
   // bufferized.
   %0 = bufferization.alloc_tensor() {memory_space = 0 : ui64} : tensor<10xf32>
   %1 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<10xf32>
   // expected-error @+2 {{inconsistent memory space on true/false operands}}
   // expected-error @+1 {{failed to bufferize op}}
   %r = arith.select %c, %0, %1 : tensor<10xf32>
   func.return %r : tensor<10xf32>
 }
 
 // -----
 
-func.func @constant_memory_space(%idx: index, %v: i32) -> tensor<3xi32> {
-  // expected-error @+2 {{memory space not implemented yet}}
+func.func @unknown_memory_space(%idx: index, %v: i32) -> tensor<3xi32> {
+  // expected-error @+2 {{could not infer memory space}}
   // expected-error @+1 {{failed to bufferize op}}
   %cst = arith.constant dense<[5, 1000, 20]> : tensor<3xi32>
   %0 = tensor.insert %v into %cst[%idx] : tensor<3xi32>
   return %0 : tensor<3xi32>
 }
 No newline at end of file