diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h @@ -230,6 +230,11 @@ /// 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 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. diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td --- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td +++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationBase.td @@ -49,6 +49,10 @@ /// allocation (as per BufferizableOpInterface) may have this attribute. 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; diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -285,6 +285,10 @@ "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">, diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp --- a/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationDialect.cpp @@ -115,6 +115,42 @@ } return success(); } + if (attr.getName() == kMemorySpaceAttrName) { + auto arrayAttr = attr.getValue().dyn_cast(); + 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(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(); + if (!intAttr) + return op->emitError() << "'" << kMemorySpaceAttrName + << "' is expected to be an int array attribute"; + bool isTensor = op->getResult(it.index()).getType().isa(); + // 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"; diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp --- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp +++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp @@ -161,12 +161,29 @@ 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( + BufferizationDialect::kMemorySpaceAttrName)[0] + .cast() + .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 dynamicDims = getDynamicSizes(); if (getCopy()) { assert(dynamicDims.empty() && "expected either `copy` or `dynamicDims`"); diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -187,6 +187,8 @@ opt.createDeallocs = createDeallocs; opt.functionBoundaryTypeConversion = parseLayoutMapOption(functionBoundaryTypeConversion); + if (mustInferMemorySpace) + opt.defaultMemorySpace = None; opt.printConflicts = printConflicts; opt.testAnalysisOnly = testAnalysisOnly; opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries; diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-memory-space-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-memory-space-invalid.mlir new file mode 100644 --- /dev/null +++ b/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> +} diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir --- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir @@ -138,3 +138,14 @@ 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> +} diff --git a/mlir/test/Dialect/Bufferization/invalid.mlir b/mlir/test/Dialect/Bufferization/invalid.mlir --- a/mlir/test/Dialect/Bufferization/invalid.mlir +++ b/mlir/test/Dialect/Bufferization/invalid.mlir @@ -54,4 +54,44 @@ // expected-error @+1{{'bufferization.escape' only valid on bufferizable ops}} %0 = memref.cast %m0 {bufferization.escape = [true]} : memref to memref<10xf32> return -} \ No newline at end of file +} + +// ----- + +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 + 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 + return +} + +// ----- + +func.func @memory_space_attr_non_allocating(%t0: tensor) { + // expected-error @+1{{'bufferization.memory_space' only valid for allocation results}} + %0 = tensor.extract_slice %t0[0][5][1] {bufferization.memory_space = [0]} : tensor to tensor<5xf32> + return +} + +// ----- + +func.func @ememory_space_attr_non_bufferizable(%m0: memref) { + // expected-error @+1{{'bufferization.memory_space' only valid on bufferizable ops}} + %0 = memref.cast %m0 {bufferization.memory_space = [0]} : memref 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 + return +}