[mlir][bufferization] Add memory_space op attribute

This attribute is currently supported on AllocTensorOp only. Future changes will add support to other ops. Furthermore, the memory space is not propagated properly in all bufferization patterns and some of the core bufferization infrastructure. This will be addressed in a subsequent change.

Differential Revision: https://reviews.llvm.org/D128274

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

@@ -230,6 +230,11 @@ struct BufferizationOptions {
   /// bufferized or not.
   bool bufferizeFunctionBoundaries = false;
 
+  /// The default memory space that should be used when it cannot be inferred
+  /// from the context. If no default memory space is specified, bufferization
+  /// fails when the memory space cannot be inferred at any point.
+  Optional<unsigned> defaultMemorySpace = 0;
+
   /// Certain ops have aliasing OpOperand/OpResult invariants (e.g., scf.for).
   /// If this flag is set to `false`, those invariants are no longer enforced
   /// with buffer copies.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td

@@ -46,6 +46,10 @@ def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
     conflicts that would have been introduced by the in-place bufferization of
     another op.
 
+    The optional `memory_space` attribute specifies the memory space when
+    bufferizing this op. If `memory_space` is not specified, the default memory
+    space is used during bufferization.
+
     Both dense and sparse tensor types are supported. The result of a
     `bufferization.alloc_tensor` is a tensor value that can be used like any
     other tensor value. In practice, it is often used as the "out" operand of
@@ -64,7 +68,8 @@ def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
-                       Optional<AnyTensor>:$copy);
+                       Optional<AnyTensor>:$copy,
+                       OptionalAttr<UI64Attr>:$memory_space);
 
   let results = (outs AnyTensor:$result);
 
@@ -119,8 +124,12 @@ def Bufferization_AllocTensorOp : Bufferization_Op<"alloc_tensor",
   }];
 
   let builders = [
-    // Build an op without `copy` operand.
+    // Build an op without `copy` or `memory_space`.
     OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes)>,
+
+    // Build an op without `memory_space`.
+    OpBuilder<(ins "RankedTensorType":$type, "ValueRange":$dynamicSizes,
+                   "Value":$copy)>,
   ];
 
   let hasCanonicalizer = 1;

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

@@ -285,6 +285,10 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
            "function-boundary-type-conversion", "std::string",
            /*default=*/"\"infer-layout-map\"",
            "Controls layout maps when bufferizing function signatures.">,
+    Option<"mustInferMemorySpace", "must-infer-memory-space", "bool",
+           /*default=*/"false",
+           "The memory space of an memref types must always be inferred. If "
+           "unset, a default memory space of 0 is used otherwise.">,
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Test only: Only run inplaceability analysis and annotate IR">,

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

@@ -163,12 +163,25 @@ LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
     return success();
   }
 
-  // Create buffer allocation.
+  // Get "copy" buffer.
   Value copyBuffer;
   if (getCopy())
     copyBuffer = getBuffer(rewriter, getCopy(), options);
+
+  // Compute memory space of this allocation.
+  unsigned memorySpace;
+  if (getMemorySpace().hasValue()) {
+    memorySpace = *getMemorySpace();
+  } else if (options.defaultMemorySpace.hasValue()) {
+    memorySpace = *options.defaultMemorySpace;
+  } else {
+    return op->emitError("could not infer memory space");
+  }
+
+  // Create memory allocation.
   auto allocType =
-      MemRefType::get(getType().getShape(), getType().getElementType());
+      MemRefType::get(getType().getShape(), getType().getElementType(),
+                      AffineMap(), memorySpace);
   SmallVector<Value> dynamicDims = getDynamicSizes();
   if (getCopy()) {
     assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`");
@@ -267,7 +280,14 @@ LogicalResult AllocTensorOp::verify() {
 
 void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
                           RankedTensorType type, ValueRange dynamicSizes) {
-  build(builder, result, type, dynamicSizes, /*copy=*/Value());
+  build(builder, result, type, dynamicSizes, /*copy=*/Value(),
+        /*memory_space=*/BoolAttr());
+}
+
+void AllocTensorOp::build(OpBuilder &builder, OperationState &result,
+                          RankedTensorType type, ValueRange dynamicSizes,
+                          Value copy) {
+  build(builder, result, type, dynamicSizes, copy, /*memory_space=*/BoolAttr());
 }
 
 namespace {

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

@@ -187,6 +187,8 @@ struct OneShotBufferizePass
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
+      if (mustInferMemorySpace)
+        opt.defaultMemorySpace = None;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;

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

@@ -0,0 +1,8 @@
+// RUN: mlir-opt %s -allow-unregistered-dialect -one-shot-bufferize="must-infer-memory-space" -split-input-file -verify-diagnostics
+
+func.func @alloc_tensor_without_memory_space() -> tensor<10xf32> {
+  // expected-error @+2 {{could not infer memory space}}
+  // expected-error @+1 {{failed to bufferize op}}
+  %0 = bufferization.alloc_tensor() : tensor<10xf32>
+  return %0 : tensor<10xf32>
+}

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

@@ -138,3 +138,14 @@ func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {
   return %0 : tensor<5xf32>
 }
 
+// -----
+
+// CHECK-LABEL: func @alloc_tensor_with_memory_space()
+func.func @alloc_tensor_with_memory_space() -> tensor<5xf32> {
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>
+  %0 = bufferization.alloc_tensor() {memory_space = 1 : ui64} : tensor<5xf32>
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK: memref.dealloc %[[alloc]]
+  // CHECK: return %[[r]]
+  return %0 : tensor<5xf32>
+}

mlir/test/Dialect/Bufferization/invalid.mlir

@@ -79,3 +79,11 @@ func.func @sparse_alloc_call() {
   return
 }
 
+// -----
+
+func.func @alloc_tensor_invalid_memory_space_attr(%sz: index) {
+  // expected-error @+1{{'bufferization.alloc_tensor' op attribute 'memory_space' failed to satisfy constraint: 64-bit unsigned integer attribute}}
+  %0 = bufferization.alloc_tensor(%sz) {memory_space = "foo"} : tensor<?xf32>
+  return
+}
+