[Mlir-commits] [mlir] [mlir][bufferization] Add `BufferizableOpInterface::hasTensorSemantics` (PR #75273)
Matthias Springer
llvmlistbot at llvm.org
Tue Jan 16 01:02:26 PST 2024
https://github.com/matthias-springer updated https://github.com/llvm/llvm-project/pull/75273
>From 9cb1e5d95e0a69952755ad6caca481cfd8e4ebae Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Tue, 16 Jan 2024 09:01:22 +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 ++++++++++++++
.../IR/BufferizableOpInterface.cpp | 23 +++++++++++++++++
.../Bufferization/Transforms/Bufferize.cpp | 25 -------------------
.../FuncBufferizableOpInterfaceImpl.cpp | 22 ++++++++++++++++
.../Transforms/OneShotAnalysis.cpp | 7 ------
.../Transforms/OneShotModuleBufferize.cpp | 9 +++++++
7 files changed, 83 insertions(+), 32 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h
index 7c09a43f96397b7..63e2d19e68ef97c 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 fd1ceb68af5dd91..007c05adc30b5f8 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/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 1e8dc4387ed4f0d..a0bb8715f2c5613 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,20 @@ bufferization::detail::unknownGetAliasingValues(OpOperand &opOperand) {
r.addAlias({bbArg, BufferRelation::Unknown, /*isDefinite=*/false});
return r;
}
+
+bool bufferization::detail::defaultHasTensorSemantics(Operation *op) {
+ auto isaTensor = [](Type t) { return isa<TensorType>(t); };
+ 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 (any_of(op->getResultTypes(), isaTensor))
+ return true;
+ return any_of(op->getOperandTypes(), isaTensor);
+}
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp
index f2125feeda54159..3f1626a6af34d49 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 3a8c397c02a8096..07cd1f90b17df4c 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/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 498db1182317a96..fba9cd87306381d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -1030,13 +1030,6 @@ OneShotAnalysisState::analyzeSingleOp(Operation *op,
return success();
}
-/// Return true if the given op has a tensor result or a tensor operand.
-static bool hasTensorSemantics(Operation *op) {
- bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
- bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
- return hasTensorResult || hasTensorOperand;
-}
-
/// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
static void equivalenceAnalysis(SmallVector<Operation *> &ops,
OneShotAnalysisState &state) {
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp
index 1404ed8f43f9643..aeda995fd585aac 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