[mlir][linalg][bufferize][NFC] Pass BufferizationState to PostAnalysisStep
Pass BufferizationStep instead of BufferizationAliasInfo. Note: BufferizationState contains BufferizationAliasInfo.
Differential Revision: https://reviews.llvm.org/D114512
GitOrigin-RevId: 8e2214aa60bc836d263fd7cb61be8f595eb05a37
diff --git a/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h b/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
index e03aaea..52befbd 100644
--- a/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
+++ b/include/mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h
@@ -327,10 +327,9 @@
virtual ~PostAnalysisStep() {}
/// Run the post analysis step. This function may modify the IR, but must keep
- /// `aliasInfo` consistent. Newly created operations and operations that
- /// should be re-analyzed must be stored in `newOps`.
- virtual LogicalResult run(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
- DominanceInfo &domInfo,
+ /// `aliasInfo` (inside `state`) consistent. Newly created operations and
+ /// operations that should be re-analyzed must be stored in `newOps`.
+ virtual LogicalResult run(FuncOp funcOp, BufferizationState &state,
SmallVector<Operation *> &newOps) = 0;
};
diff --git a/include/mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h b/include/mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h
index d5ee20b..1858787 100644
--- a/include/mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h
+++ b/include/mlir/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.h
@@ -34,7 +34,7 @@
/// * The result of `rewriteFunc` must usually be analyzed for inplacability.
/// This analysis can be skipped with `skipAnalysis`.
LogicalResult eliminateInitTensors(
- FuncOp funcOp, BufferizationAliasInfo &aliasInfo, DominanceInfo &domInfo,
+ FuncOp funcOp, BufferizationState &state,
std::function<bool(OpOperand &)> anchorMatchFunc,
std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
SmallVector<Operation *> &newOps);
@@ -45,8 +45,7 @@
/// (and some other conditions are met).
struct InsertSliceAnchoredInitTensorEliminationStep
: public InitTensorEliminationStep {
- LogicalResult run(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
- DominanceInfo &domInfo,
+ LogicalResult run(FuncOp funcOp, BufferizationState &state,
SmallVector<Operation *> &newOps) override;
};
diff --git a/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp b/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
index 53eaab5..eb37a54 100644
--- a/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
+++ b/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp
@@ -728,7 +728,7 @@
for (const std::unique_ptr<PostAnalysisStep> &step :
options.postAnalysisSteps) {
SmallVector<Operation *> newOps;
- if (failed(step->run(funcOp, aliasInfo, domInfo, newOps)))
+ if (failed(step->run(funcOp, state, newOps)))
return failure();
// Analyze ops that were created by the PostAnalysisStep.
if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
diff --git a/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp b/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
index 7bea450..b4a439e 100644
--- a/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
+++ b/lib/Dialect/Linalg/ComprehensiveBufferize/LinalgInterfaceImpl.cpp
@@ -407,12 +407,12 @@
/// OpOperand, that eventually ends at a single InitTensorOp.
LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
InitTensorEliminationStep::eliminateInitTensors(
- FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
- DominanceInfo &domInfo,
+ FuncOp funcOp, BufferizationState &state,
std::function<bool(OpOperand &)> anchorMatchFunc,
std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
SmallVector<Operation *> &newOps) {
OpBuilder b(funcOp->getContext());
+ BufferizationAliasInfo &aliasInfo = state.aliasInfo;
WalkResult status = funcOp->walk([&](Operation *op) {
for (OpOperand &operand : op->getOpOperands()) {
@@ -501,17 +501,17 @@
/// out-of-place due to RaW conflicts.
LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
InsertSliceAnchoredInitTensorEliminationStep::run(
- FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
- DominanceInfo &domInfo, SmallVector<Operation *> &newOps) {
+ FuncOp funcOp, BufferizationState &state,
+ SmallVector<Operation *> &newOps) {
return eliminateInitTensors(
- funcOp, aliasInfo, domInfo,
+ funcOp, state,
[&](OpOperand &operand) {
auto insertSliceOp =
dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
if (!insertSliceOp)
return false;
// Only inplace bufferized InsertSliceOps are eligible.
- if (!aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
+ if (!state.aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
return false;
return &operand == &insertSliceOp->getOpOperand(0) /*source*/;
},
