[Mlir-commits] [mlir] [mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` (PR #75273)
Matthias Springer
llvmlistbot at llvm.org
Wed Jan 10 03:07:22 PST 2024
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/75273
>From 81e929c4a53bb80fc14a13318b0fed6d8e95da15 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 10 Jan 2024 11:06:05 +0000
Subject: [PATCH] [mlir][bufferization] Add
`BufferizableOpInterface::hasTensorSemantics`
Add a new interface method to `BufferizableOpInterface`: `hasTensorSemanticsForBufferization`. This method returns "true" if the op has tensor semantics and should be bufferized.
Until now, we assumed that an op has tensor semantics if it has tensor operands and/or tensor op results. However, there are ops like `ml_program.global` that do not have any results/operands but must still be bufferized (#75103). The new interface method can return "true" for such ops.
This change also decouples `bufferization::bufferizeOp` a bit from the func dialect.
---
.../IR/BufferizableOpInterface.h | 10 ++++++++
.../IR/BufferizableOpInterface.td | 19 ++++++++++++++
.../Bufferization/IR/BufferizationOps.td | 6 +++++
.../IR/BufferizableOpInterface.cpp | 24 ++++++++++++++++++
.../Bufferization/Transforms/Bufferize.cpp | 25 -------------------
.../FuncBufferizableOpInterfaceImpl.cpp | 22 ++++++++++++++++
.../Transforms/OneShotModuleBufferize.cpp | 9 +++++++
7 files changed, 90 insertions(+), 25 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 7c09a43f96397b..63e2d19e68ef97 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
@@ -601,6 +601,12 @@ FailureOr<BaseMemRefType> getBufferType(Value value,
const BufferizationOptions &options,
SmallVector<Value> &invocationStack);
+/// Return "true" if the given op has tensor semantics and should be bufferized.
+/// If the op is bufferizable, the BufferizableOpInterface is queried.
+/// Otherwise, an op has tensor semantics if it has tensor operands, tensor
+/// op results and/or tensor block arguments.
+bool hasTensorSemantics(Operation *op);
+
/// Replace an op with replacement values. The op is deleted. Tensor OpResults
/// must be replaced with memref values.
void replaceOpWithBufferizedValues(RewriterBase &rewriter, Operation *op,
@@ -694,6 +700,10 @@ AliasingOpOperandList unknownGetAliasingOpOperands(Value value);
/// This is the default implementation of getAliasingValues in case the owner
/// op does not implement the BufferizableOpInterface.
AliasingValueList unknownGetAliasingValues(OpOperand &opOperand);
+
+/// This is the default implementation of
+/// BufferizableOpInterface::hasTensorSemantics
+bool defaultHasTensorSemantics(Operation *op);
} // namespace detail
} // namespace bufferization
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
index fd1ceb68af5dd9..007c05adc30b5f 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -575,6 +575,25 @@ def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
return false;
}]
>,
+ InterfaceMethod<
+ /*desc=*/[{
+ Return "true" if the this op has tensor semantics and should be
+ bufferized. By default, ops with tensor operands, tensor op results
+ and/or tensor block arguments have tensor semantics.
+
+ This interface methods can be implemented by ops that should be
+ bufferized but do not have tensor semantics according to the above
+ definition. E.g., this function can return "true" for symbols.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"hasTensorSemantics",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::bufferization::detail
+ ::defaultHasTensorSemantics($_op.getOperation());
+ }]
+ >,
StaticInterfaceMethod<
/*desc=*/[{
Return `true` if the op and this interface implementation supports
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 9dc6afcaab31c8..795a7d79a73a58 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -281,6 +281,12 @@ def Bufferization_MaterializeInDestinationOp
let results = (outs Optional<AnyTensor>:$result);
let extraClassDeclaration = [{
+ // Both `DestinationStyleOpInterface` and `BufferizableOpInterface` define
+ // `hasTensorSemantics`. Both return the same result, but we have to choose
+ // one to disambiguate the method lookup.
+ using DestinationStyleOpInterface::Trait<MaterializeInDestinationOp>
+ ::hasTensorSemantics;
+
LogicalResult bufferize(RewriterBase &rewriter,
const BufferizationOptions &options);
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1e8dc4387ed4f0..aa6b751a8f0f21 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -689,6 +689,12 @@ bufferization::getBufferType(Value value, const BufferizationOptions &options,
*options.defaultMemorySpace);
}
+bool bufferization::hasTensorSemantics(Operation *op) {
+ if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+ return bufferizableOp.hasTensorSemantics();
+ return detail::defaultHasTensorSemantics(op);
+}
+
void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
Operation *op,
ValueRange values) {
@@ -989,3 +995,21 @@ bufferization::detail::unknownGetAliasingValues(OpOperand &opOperand) {
r.addAlias({bbArg, BufferRelation::Unknown, /*isDefinite=*/false});
return r;
}
+
+static bool isaTensor(Type t) { return isa<TensorType>(t); }
+
+bool bufferization::detail::defaultHasTensorSemantics(Operation *op) {
+ bool hasTensorBlockArgument = any_of(op->getRegions(), [](Region &r) {
+ return any_of(r.getBlocks(), [](Block &b) {
+ return any_of(b.getArguments(), [](BlockArgument bbArg) {
+ return isaTensor(bbArg.getType());
+ });
+ });
+ });
+ if (hasTensorBlockArgument)
+ return true;
+
+ bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
+ bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+ return hasTensorResult || hasTensorOperand;
+}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index f2125feeda5415..3f1626a6af34d4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
@@ -350,31 +350,6 @@ mlir::bufferization::createFinalizingBufferizePass() {
// BufferizableOpInterface-based Bufferization
//===----------------------------------------------------------------------===//
-static bool isaTensor(Type t) { return isa<TensorType>(t); }
-
-/// Return true if the given op has a tensor result or a tensor operand.
-static bool hasTensorSemantics(Operation *op) {
- bool hasTensorBlockArgument = any_of(op->getRegions(), [](Region &r) {
- return any_of(r.getBlocks(), [](Block &b) {
- return any_of(b.getArguments(), [](BlockArgument bbArg) {
- return isaTensor(bbArg.getType());
- });
- });
- });
- if (hasTensorBlockArgument)
- return true;
-
- if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
- bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor);
- bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor);
- return hasTensorArg || hasTensorResult;
- }
-
- bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
- bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
- return hasTensorResult || hasTensorOperand;
-}
-
namespace {
/// A rewriter that keeps track of extra information during bufferization.
class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 3a8c397c02a809..07cd1f90b17df4 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
@@ -325,6 +325,28 @@ struct FuncOpInterface
static bool supportsUnstructuredControlFlow() { return true; }
+ bool hasTensorSemantics(Operation *op) const {
+ auto isaTensor = [](Type type) { return isa<TensorType>(type); };
+
+ // A function has tensor semantics if it has tensor arguments/results.
+ auto funcOp = cast<FuncOp>(op);
+ bool hasTensorArg = any_of(funcOp.getArgumentTypes(), isaTensor);
+ bool hasTensorResult = any_of(funcOp.getResultTypes(), isaTensor);
+ if (hasTensorArg || hasTensorResult)
+ return true;
+
+ // It also has tensor semantics if it has tensor block arguments.
+ // TODO: Decouple bufferization of unstructured control flow from
+ // BufferizableOpInterface implementations. We should only care about
+ // region entry block arguments here (which are already covered by the
+ // argument types of the function).
+ for (Block &block : funcOp.getBody())
+ if (any_of(block.getArgumentTypes(), isaTensor))
+ return true;
+
+ return false;
+ }
+
AliasingOpOperandList
getAliasingOpOperands(Operation *op, Value value,
const AnalysisState &state) const {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 1404ed8f43f964..aeda995fd585aa 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
@@ -458,6 +458,15 @@ LogicalResult mlir::bufferization::bufferizeModuleOp(
foldMemRefCasts(funcOp);
}
+ // Bufferize all other ops.
+ for (Operation &op : moduleOp.getOps()) {
+ // Functions were already bufferized.
+ if (isa<func::FuncOp>(&op))
+ continue;
+ if (failed(bufferizeOp(&op, options, statistics)))
+ return failure();
+ }
+
// Post-pass cleanup of function argument attributes.
removeBufferizationAttributesInModule(moduleOp);
More information about the Mlir-commits
mailing list