[mlir][bufferization] Add `memory_space` op attribute

[springerm]

Summary

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.

Depends On D128137

[nicolasvasilache]

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

115 ↗

(On Diff #438685)

this looks fishy, this should be an attribute directly on the op you are interested in and with proper Tablegen definition most of this code and tests become unnecessary

What's the rationale behind the current approach?

[springerm]

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

115 ↗

(On Diff #438685)

memory_space is an attribute that can be added to any op that may bufferize to an allocation. E.g.:

• bufferization.alloc_tensor
• sparse_tensor.new
• sparse_tensor.convert
• tensor.collapse_shape

By making it a dialect attribute, we can avoid code duplication around those ops, treat them uniformly in the code base, and have some basic verification.

Also, we probably don't want to have this attribute on ops such as tensor.collapse_shape directly (separation of concerns). Ideally, there would be a trait that adds this attribute via an external model registration.... This is not possible in MLIR and this is the next best thing.

[springerm]

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

115 ↗

(On Diff #438685)

For tensor.collapse_shape and sparse_tensor.convert, we usually would not add the attribute because the memory space can be inferred.

But there is another dialect attribute bufferization.escape that is implemented in the same way. And that attr applies to all 4 ops that I mentioned above.

[springerm]

rebase

[springerm]

address comments

[nicolasvasilache]

Thanks! Please update the title before landing.

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

struct BufferizationOptions {
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
   /// Specifies whether function boundaries (ops in the func dialect) should be
   /// 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.
   ///
   /// Note: Deactivating this flag can lead to incorrect bufferization results
   /// when used incorrectly. This flag is useful with
   /// `AlwaysCopyAnalysisState` which bufferizes all writing tensor
   /// OpOperands out-of-place.

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

let description = [{
     `alloc_tensor` is a helper op for bufferization. The operation is provided
     as an anchor that marks the beginning of a new tensor SSA use-def chain. It
     can be used to control in-place bufferization decisions during One-Shot
     Bufferize: The bufferized result of a `bufferization.alloc_tensor` does not
     alias with any other buffer, so it can be used to resolve read-after-write
     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
     another op. Sparse tensor allocations should always be used in a local
     construction operation and never escape the function boundary directly.
 
     Example:
 
     ```mlir
     %c = bufferization.alloc_tensor [%d1, %d2] : tensor<?x?xf32, #SparseMatrix>
     %0 = linalg.matmul
       ins(%a, %b: tensor<?x?xf32, #SparseMatrix>, tensor<?x?xf32, #SparseMatrix>)
       outs(%c: tensor<?x?xf32, #SparseMatrix>) -> tensor<?x?xf32, #SparseMatrix>
     return %0 : tensor<?x?xf32, #SparseMatrix>
     ```
   }];
 
   let arguments = (ins Variadic<Index>:$dynamic_sizes,
-                       Optional<AnyTensor>:$copy);
+                       Optional<AnyTensor>:$copy,
+                       OptionalAttr<UI64Attr>:$memory_space);
 
   let results = (outs AnyTensor:$result);
 
   let extraClassDeclaration = [{
     LogicalResult bufferize(RewriterBase &rewriter,
                             const BufferizationOptions &options);
 
     bool isMemoryWrite(OpResult opResult, const AnalysisState &state);
let extraClassDeclaration = [{
     // and return the shape.
     int64_t getStaticSize(unsigned idx) {
       assert(!isDynamicDim(idx) && "expected static size");
       return getType().getShape()[idx];
     }
   }];
 
   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;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

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

Option<"createDeallocs", "create-deallocs", "bool", /*default=*/"true",
            "Specify if buffers should be deallocated. For compatibility with "
            "core bufferization passes.">,
     ListOption<"dialectFilter", "dialect-filter", "std::string",
                "Restrict bufferization to ops from these dialects.">,
     Option<"functionBoundaryTypeConversion",
            "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">,
     Option<"printConflicts", "print-conflicts", "bool",
             /*default=*/"false",
            "Test only: Annotate IR with RaW conflicts. Requires "
            "test-analysis-only.">,
     Option<"unknownTypeConversion", "unknown-type-conversion", "std::string",

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

LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
   Location loc = getLoc();
 
   // Nothing to do for dead AllocTensorOps.
   if (getOperation()->getUses().empty()) {
     rewriter.eraseOp(getOperation());
     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`");
     populateDynamicDimSizes(rewriter, loc, copyBuffer, dynamicDims);
   }
   FailureOr<Value> alloc =
       options.createAlloc(rewriter, loc, allocType, dynamicDims);
   if (failed(alloc))
for (auto &use : getOperation()->getUses())
         return emitError("sparse tensor allocation should not escape function");
   }
 
   return success();
 }
 
 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 {
 /// Change the type of the result of a `bufferization.alloc_tensor` by making
 /// the result type statically sized along dimension that in the original
 /// operation where defined as dynamic, but the size was defined using a
 /// `constant` op. For example:
 ///

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

if (!options) {
       // Make new bufferization options if none were provided when creating the
       // pass.
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
+      if (mustInferMemorySpace)
+        opt.defaultMemorySpace = None;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
       opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion);
 
       OpFilter::Entry::FilterFn filterFn =
           [&](Operation *op) {
             // Filter may be specified via options.

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

+// 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

func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {
   // CHECK: memref.copy %[[m]], %[[alloc]]
   %0 = bufferization.alloc_tensor() copy(%t) : tensor<5xf32>
   // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
   // CHECK: memref.dealloc %[[alloc]]
   // CHECK: return %[[r]]
   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

 
 func.func @sparse_alloc_call() {
   // expected-error @+1{{sparse tensor allocation should not escape function}}
   %0 = bufferization.alloc_tensor() : tensor<20x40xf32, #DCSR>
   call @foo(%0) : (tensor<20x40xf32, #DCSR>) -> ()
   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
+}
+