[Mlir-commits] [mlir] [mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` (PR #75273)
Matthias Springer
llvmlistbot at llvm.org
Tue Dec 12 18:19:01 PST 2023
https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/75273
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.
>From 7ed95b12fd3909de88a84fb6e9f9c5b050546902 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 13 Dec 2023 11:16:29 +0900
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 | 22 +++++++++++++
.../IR/BufferizableOpInterface.cpp | 25 ++++++++++++++
.../Bufferization/Transforms/Bufferize.cpp | 33 +++----------------
.../FuncBufferizableOpInterfaceImpl.cpp | 22 +++++++++++++
.../Transforms/OneShotModuleBufferize.cpp | 9 +++++
6 files changed, 92 insertions(+), 29 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 7c09a43f96397..d3dc5683772e2 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 hasTensorSemanticsForBufferization(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::hasTensorSemanticsForBufferization.
+bool defaultHasTensorSemanticsForBufferization(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 fd1ceb68af5dd..e7445cb4e63da 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td
@@ -575,6 +575,28 @@ 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.
+
+ TODO: This interface method should be called `hasTensorSemantics`, but
+ that name is already in use in `DestinationStyleOpInterface`.
+ }],
+ /*retType=*/"bool",
+ /*methodName=*/"hasTensorSemanticsForBufferization",
+ /*args=*/(ins),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::bufferization::detail
+ ::defaultHasTensorSemanticsForBufferization($_op.getOperation());
+ }]
+ >,
StaticInterfaceMethod<
/*desc=*/[{
Return `true` if the op and this interface implementation supports
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1e8dc4387ed4f..8f577947f8ddb 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::hasTensorSemanticsForBufferization(Operation *op) {
+ if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+ return bufferizableOp.hasTensorSemanticsForBufferization();
+ return detail::defaultHasTensorSemanticsForBufferization(op);
+}
+
void bufferization::replaceOpWithBufferizedValues(RewriterBase &rewriter,
Operation *op,
ValueRange values) {
@@ -989,3 +995,22 @@ 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::defaultHasTensorSemanticsForBufferization(
+ 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 f2125feeda541..79be7bed79db3 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 {
@@ -417,7 +392,7 @@ class BufferizationRewriter : public IRRewriter, public RewriterBase::Listener {
return;
// Skip non-tensor ops.
- if (!hasTensorSemantics(op))
+ if (!hasTensorSemanticsForBufferization(op))
return;
// Skip ops that are not allowed to be bufferized.
@@ -470,7 +445,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
// canonicalize away (or canonicalize to more precise layouts).
SmallVector<Operation *> worklist;
op->walk<WalkOrder::PostOrder>([&](Operation *op) {
- if (hasTensorSemantics(op))
+ if (hasTensorSemanticsForBufferization(op))
worklist.push_back(op);
});
@@ -492,7 +467,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
if (!options.isOpAllowed(nextOp))
continue;
// Skip ops that no longer have tensor semantics.
- if (!hasTensorSemantics(nextOp))
+ if (!hasTensorSemanticsForBufferization(nextOp))
continue;
// Check for unsupported unstructured control flow.
if (!bufferizableOp.supportsUnstructuredControlFlow())
@@ -546,7 +521,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op,
continue;
// Ops that no longer have tensor semantics (because they were updated
// in-place) are allowed.
- if (!hasTensorSemantics(op))
+ if (!hasTensorSemanticsForBufferization(op))
continue;
// Continue ops that are not allowed.
if (!options.isOpAllowed(op))
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.cpp
index 3a8c397c02a80..5aab035f9f30f 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 hasTensorSemanticsForBufferization(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 1404ed8f43f96..aeda995fd585a 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