[Mlir-commits] [mlir] c06f01f - [mlir][bufferization] Add `memory_space` op attribute
Matthias Springer
llvmlistbot at llvm.org
Mon Jun 27 03:37:16 PDT 2022
Author: Matthias Springer
Date: 2022-06-27T12:33:26+02:00
New Revision: c06f01ffee8677754a40543ab809c9e3e7c02a77
URL: https://github.com/llvm/llvm-project/commit/c06f01ffee8677754a40543ab809c9e3e7c02a77
DIFF: https://github.com/llvm/llvm-project/commit/c06f01ffee8677754a40543ab809c9e3e7c02a77.diff
LOG: [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
Added:
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-memory-space-invalid.mlir
Modified:
mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
mlir/test/Dialect/Bufferization/invalid.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index f28db1e26c09e..f852e7bcfe961 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/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.
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 6b92904427797..81dadee852e3e 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/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;
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index b8c502d12e1f6..6cf0dfa7ff24d 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/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">,
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 37ff8e9be79e2..e10d96b451d42 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/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 {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index 6c7f9782b316d..f9c809081723e 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/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;
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
index 0000000000000..f9436543698a2
--- /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
index 26faa04696eae..3bc380f59eef4 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/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>
+}
diff --git a/mlir/test/Dialect/Bufferization/invalid.mlir b/mlir/test/Dialect/Bufferization/invalid.mlir
index 28e7be0c06dbb..f78de27951d4a 100644
--- a/mlir/test/Dialect/Bufferization/invalid.mlir
+++ b/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
+}
+
More information about the Mlir-commits
mailing list