[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

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

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

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/BufferizationBase.td

let extraClassDeclaration = [{
     /// Attribute name used to mark the bufferization layout for region
     /// arguments during One-Shot Module Bufferize.
     constexpr const static ::llvm::StringLiteral
         kBufferLayoutAttrName = "bufferization.buffer_layout";
 
     /// Attribute name used to mark escaping behavior of buffer allocations.
     constexpr const static ::llvm::StringLiteral
         kEscapeAttrName = "bufferization.escape";
+
+    /// Attribute name used to indicate the memory space of buffer allocations.
+    constexpr const static ::llvm::StringLiteral
+        kMemorySpaceAttrName = "bufferization.memory_space";
   }];
   let hasOperationAttrVerify = 1;
   let emitAccessorPrefix = kEmitAccessorPrefix_Prefixed;
 }
 
 #endif // BUFFERIZATION_BASE

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/BufferizationDialect.cpp

for (const auto &it : llvm::enumerate(arrayAttr)) {
         return op->emitError()
                << "'" << kEscapeAttrName << "' only valid for tensor results";
       if (!bufferizableOp.bufferizesToAllocation(op->getOpResult(it.index())))
         return op->emitError() << "'" << kEscapeAttrName
                                << "' only valid for allocation results";
     }
     return success();
   }
+  if (attr.getName() == kMemorySpaceAttrName) {
+    auto arrayAttr = attr.getValue().dyn_cast<ArrayAttr>();
+    if (!arrayAttr)
+      return op->emitError() << "'" << kMemorySpaceAttrName
+                             << "' is expected to be an int array attribute";
+    if (arrayAttr.size() != op->getNumResults())
+      return op->emitError()
+             << "'" << kMemorySpaceAttrName
+             << "' has wrong number of elements, expected "
+             << op->getNumResults() << ", got " << arrayAttr.size();
+    auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
+    if (!bufferizableOp)
+      return op->emitError() << "'" << kMemorySpaceAttrName
+                             << "' only valid on bufferizable ops";
+    for (const auto &it : llvm::enumerate(arrayAttr)) {
+      auto attr = it.value();
+      auto intAttr = attr.dyn_cast<IntegerAttr>();
+      if (!intAttr)
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' is expected to be an int array attribute";
+      bool isTensor = op->getResult(it.index()).getType().isa<TensorType>();
+      // Memory space value of -1 is used for non-tensor results.
+      if (isTensor && intAttr.getValue().getSExtValue() == -1)
+        continue;
+      if (intAttr.getValue().getSExtValue() < 0)
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' must not contain negative values";
+      if (!isTensor)
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' only valid for tensor results";
+      if (!bufferizableOp.bufferizesToAllocation(op->getOpResult(it.index())))
+        return op->emitError() << "'" << kMemorySpaceAttrName
+                               << "' only valid for allocation results";
+    }
+    return success();
+  }
 
   return op->emitError() << "attribute '" << attr.getName()
                          << "' not supported by the bufferization dialect";
 }

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 (op->hasAttr(BufferizationDialect::kMemorySpaceAttrName)) {
+    memorySpace = op->getAttrOfType<ArrayAttr>(
+                        BufferizationDialect::kMemorySpaceAttrName)[0]
+                      .cast<IntegerAttr>()
+                      .getValue()
+                      .getZExtValue();
+  } 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))

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() {bufferization.memory_space = [1]} : 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 @escape_attr_non_bufferizable(%m0: memref<?xf32>) {
   // expected-error @+1{{'bufferization.escape' only valid on bufferizable ops}}
   %0 = memref.cast %m0 {bufferization.escape = [true]} : memref<?xf32> to memref<10xf32>
   return
-}
+}
+
+// -----
+
+func.func @alloc_tensor_invalid_memory_space_attr(%sz: index) {
+  // expected-error @+1{{'bufferization.memory_space' is expected to be an int array attribute}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.memory_space = "foo"} : tensor<?xf32>
+  return
+}
+
+// -----
+
+func.func @alloc_tensor_invalid_memory_space_attr_size(%sz: index) {
+  // expected-error @+1{{'bufferization.memory_space' has wrong number of elements, expected 1, got 2}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.memory_space = [0, 0]} : tensor<?xf32>
+  return
+}
+
+// -----
+
+func.func @memory_space_attr_non_allocating(%t0: tensor<?xf32>) {
+  // expected-error @+1{{'bufferization.memory_space' only valid for allocation results}}
+  %0 = tensor.extract_slice %t0[0][5][1] {bufferization.memory_space = [0]} : tensor<?xf32> to tensor<5xf32>
+  return
+}
+
+// -----
+
+func.func @ememory_space_attr_non_bufferizable(%m0: memref<?xf32>) {
+  // expected-error @+1{{'bufferization.memory_space' only valid on bufferizable ops}}
+  %0 = memref.cast %m0 {bufferization.memory_space = [0]} : memref<?xf32> to memref<10xf32>
+  return
+}
+
+// -----
+
+func.func @alloc_tensor_invalid_attr(%sz: index) {
+  // expected-error @+1{{attribute '"bufferization.foo"' not supported by the bufferization dialect}}
+  %0 = bufferization.alloc_tensor(%sz) {bufferization.foo = "foo"} : tensor<?xf32>
+  return
+}