diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h index a2b7f7f5017d7..49e74140626fb 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h @@ -28,6 +28,16 @@ class AnalysisState; struct BufferizationOptions; class OpFilter; +/// Bufferization statistics for debugging. These can be printed after running +/// the OneShotBufferizePass with `-mlir-pass-statistics`. See the pass +/// definition for more details. +struct BufferizationStatistics { + int64_t numBufferAlloc = 0; + int64_t numBufferDealloc = 0; + int64_t numTensorInPlace = 0; + int64_t numTensorOutOfPlace = 0; +}; + /// A helper type converter class that automatically populates the relevant /// materializations and type conversions for bufferization. class BufferizeTypeConverter : public TypeConverter { @@ -65,7 +75,8 @@ void populateEliminateBufferizeMaterializationsPatterns( /// can be used to implement partial bufferization passes. LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options, bool copyBeforeWrite = true, - const OpFilter *opFilter = nullptr); + const OpFilter *opFilter = nullptr, + BufferizationStatistics *statistics = nullptr); BufferizationOptions getPartialBufferizationOptions(); diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h index fefe432c77f99..66e1e8e0355ff 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h @@ -17,6 +17,7 @@ namespace bufferization { struct OneShotBufferizationOptions; class BufferizationAliasInfo; +struct BufferizationStatistics; class OneShotAnalysisState; /// Options for analysis-enabled bufferization. @@ -92,6 +93,9 @@ class BufferizationAliasInfo { /// Return `true` if a value was marked as in-place bufferized. bool isInPlace(OpOperand &opOperand) const; + int64_t getStatNumTensorOutOfPlace() const { return statNumTensorOutOfPlace; } + int64_t getStatNumTensorInPlace() const { return statNumTensorInPlace; } + private: /// llvm::EquivalenceClasses wants comparable elements. This comparator uses /// uses pointer comparison on the defining op. This is a poor man's @@ -124,6 +128,10 @@ class BufferizationAliasInfo { /// statically if two values are equivalent. In that case, the values are /// considered to be not equivalent. llvm::EquivalenceClasses equivalentInfo; + + // Bufferization statistics. + int64_t statNumTensorOutOfPlace = 0; + int64_t statNumTensorInPlace = 0; }; /// State for analysis-enabled bufferization. This class keeps track of alias @@ -284,11 +292,13 @@ class OneShotAnalysisState : public AnalysisState { /// Analyze `op` and its nested ops. Bufferization decisions are stored in /// `state`. -LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state); +LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state, + BufferizationStatistics *statistics = nullptr); /// Run One-Shot Bufferize on the given op: Analysis + Bufferization -LogicalResult runOneShotBufferize(Operation *op, - const OneShotBufferizationOptions &options); +LogicalResult +runOneShotBufferize(Operation *op, const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics = nullptr); } // namespace bufferization } // namespace mlir diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h index f6402b0726c2f..a69e1c2da04ca 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h @@ -15,12 +15,14 @@ struct LogicalResult; class ModuleOp; namespace bufferization { +struct BufferizationStatistics; class OneShotAnalysisState; struct OneShotBufferizationOptions; /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in /// `state`. -LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state); +LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state, + BufferizationStatistics *statistics = nullptr); /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`. /// @@ -28,7 +30,8 @@ LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state); /// inserted unless `options.copyBeforeWrite` is set, in which case buffers are /// copied before every write. LogicalResult bufferizeModuleOp(ModuleOp moduleOp, - const OneShotBufferizationOptions &options); + const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics = nullptr); /// Remove bufferization attributes on every FuncOp arguments in the ModuleOp. void removeBufferizationAttributesInModule(ModuleOp moduleOp); @@ -39,7 +42,8 @@ void removeBufferizationAttributesInModule(ModuleOp moduleOp); /// Bufferize. LogicalResult runOneShotModuleBufferize( ModuleOp moduleOp, - const bufferization::OneShotBufferizationOptions &options); + const bufferization::OneShotBufferizationOptions &options, + BufferizationStatistics *statistics = nullptr); } // namespace bufferization } // namespace mlir diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td index cbf01e4f788ea..cffe3bcb5cfbf 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -317,6 +317,17 @@ def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> { "Controls layout maps for non-inferrable memref types.">, ]; let constructor = "mlir::bufferization::createOneShotBufferizePass()"; + + let statistics = [ + Statistic<"numBufferAlloc", "num-buffer-alloc", + "Number of buffer allocations">, + Statistic<"numBufferDealloc", "num-buffer-dealloc", + "Number of buffer deallocations">, + Statistic<"numTensorInPlace", "num-tensor-in-place", + "Number of in-place tensor OpOperands">, + Statistic<"numTensorOutOfPlace", "num-tensor-out-of-place", + "Number of out-of-place tensor OpOperands">, + ]; } def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "func::FuncOp"> { diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h index 5658f7d86116d..44986a60f0693 100644 --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h @@ -15,6 +15,7 @@ namespace mlir { namespace bufferization { class AnalysisState; +struct BufferizationStatistics; struct OneShotBufferizationOptions; /// A function that matches anchor OpOperands for tensor::EmptyOp elimination. @@ -49,7 +50,8 @@ LogicalResult insertSliceAnchoredEmptyTensorEliminationStep( /// After applying this transform, the IR can be bufferized without inserting /// additional buffer allocations. LogicalResult insertTensorCopies(Operation *op, - const OneShotBufferizationOptions &options); + const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics = nullptr); /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops. /// After applying this transform, the IR can be bufferized without inserting diff --git a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp index 0c8c6faf93c89..7ff9b0c19eab8 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp @@ -245,19 +245,26 @@ struct OneShotBufferizePass opt = *options; } + BufferizationStatistics statistics; ModuleOp moduleOp = getOperation(); if (opt.bufferizeFunctionBoundaries) { - if (failed(runOneShotModuleBufferize(moduleOp, opt))) { + if (failed(runOneShotModuleBufferize(moduleOp, opt, &statistics))) { signalPassFailure(); return; } } else { - if (failed(runOneShotBufferize(moduleOp, opt))) { + if (failed(runOneShotBufferize(moduleOp, opt, &statistics))) { signalPassFailure(); return; } } + // Set pass statistics. + this->numBufferAlloc = statistics.numBufferAlloc; + this->numBufferDealloc = statistics.numBufferDealloc; + this->numTensorInPlace = statistics.numTensorInPlace; + this->numTensorOutOfPlace = statistics.numTensorOutOfPlace; + if (opt.testAnalysisOnly) return; @@ -337,9 +344,11 @@ class BufferizationRewriter : public IRRewriter { DenseSet &toMemrefOps, SmallVector &worklist, const BufferizationOptions &options, - const OpFilter *opFilter) + const OpFilter *opFilter, + BufferizationStatistics *statistics) : IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps), - worklist(worklist), analysisState(options), opFilter(opFilter) {} + worklist(worklist), analysisState(options), opFilter(opFilter), + statistics(statistics) {} protected: void notifyOperationRemoved(Operation *op) override { @@ -353,6 +362,16 @@ class BufferizationRewriter : public IRRewriter { IRRewriter::notifyOperationInserted(op); erasedOps.erase(op); + // Gather statistics about allocs and deallocs. + if (statistics) { + if (auto sideEffectingOp = dyn_cast(op)) { + statistics->numBufferAlloc += static_cast( + sideEffectingOp.hasEffect()); + statistics->numBufferDealloc += static_cast( + sideEffectingOp.hasEffect()); + } + } + // Keep track of to_memref ops. if (isa(op)) { toMemrefOps.insert(op); @@ -392,13 +411,17 @@ class BufferizationRewriter : public IRRewriter { /// An extra op filter for bufferization. const OpFilter *opFilter; + + /// Bufferization statistics for debugging. + BufferizationStatistics *statistics; }; } // namespace LogicalResult bufferization::bufferizeOp(Operation *op, const BufferizationOptions &options, bool copyBeforeWrite, - const OpFilter *opFilter) { + const OpFilter *opFilter, + BufferizationStatistics *statistics) { if (copyBeforeWrite) { AnalysisState state(options); if (failed(insertTensorCopies(op, state))) @@ -434,7 +457,7 @@ LogicalResult bufferization::bufferizeOp(Operation *op, // Bufferize all ops. BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps, - worklist, options, opFilter); + worklist, options, opFilter, statistics); for (unsigned i = 0; i < worklist.size(); ++i) { Operation *nextOp = worklist[i]; // Skip ops that were erased. diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp index 034ff5b3db2d8..cd06899595f4c 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp @@ -143,15 +143,19 @@ bool BufferizationAliasInfo::isInPlace(OpOperand &operand) const { /// Set the inPlace bufferization spec to true. void BufferizationAliasInfo::bufferizeInPlace(OpOperand &operand, AnalysisState &state) { + if (inplaceBufferized.contains(&operand)) + return; markInPlace(operand); for (OpResult result : state.getAliasingOpResult(operand)) aliasInfo.unionSets(result, operand.get()); + ++statNumTensorInPlace; } /// Set the inPlace bufferization spec to false. void BufferizationAliasInfo::bufferizeOutOfPlace(OpOperand &operand) { assert(!inplaceBufferized.contains(&operand) && "OpOperand was already decided to bufferize inplace"); + ++statNumTensorOutOfPlace; } /// Apply `fun` to all the members of the equivalence class of `v`. @@ -198,15 +202,10 @@ OneShotAnalysisState::OneShotAnalysisState( op->walk([&](BufferizableOpInterface bufferizableOp) { if (!options.isOpAllowed(bufferizableOp)) return WalkResult::skip(); - for (OpOperand &opOperand : bufferizableOp->getOpOperands()) { + for (OpOperand &opOperand : bufferizableOp->getOpOperands()) if (opOperand.get().getType().isa()) - if (bufferizableOp.mustBufferizeInPlace(opOperand, *this)) { - for (OpResult opResult : - bufferizableOp.getAliasingOpResult(opOperand, *this)) - aliasInfo.unionAliasSets(opOperand.get(), opResult); - aliasInfo.markInPlace(opOperand); - } - } + if (bufferizableOp.mustBufferizeInPlace(opOperand, *this)) + aliasInfo.bufferizeInPlace(opOperand, *this); return WalkResult::advance(); }); } @@ -1159,7 +1158,8 @@ static LogicalResult assertNoAllocsReturned(Operation *op, } LogicalResult bufferization::analyzeOp(Operation *op, - OneShotAnalysisState &state) { + OneShotAnalysisState &state, + BufferizationStatistics *statistics) { DominanceInfo domInfo(op); BufferizationAliasInfo &aliasInfo = state.getAliasInfo(); const OneShotBufferizationOptions &options = state.getOptions(); @@ -1171,6 +1171,12 @@ LogicalResult bufferization::analyzeOp(Operation *op, if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo, options.analysisFuzzerSeed))) return failure(); + + if (statistics) { + statistics->numTensorInPlace = aliasInfo.getStatNumTensorInPlace(); + statistics->numTensorOutOfPlace = aliasInfo.getStatNumTensorOutOfPlace(); + } + equivalenceAnalysis(op, aliasInfo, state); bool failedAnalysis = false; @@ -1199,15 +1205,17 @@ LogicalResult bufferization::analyzeOp(Operation *op, LogicalResult bufferization::runOneShotBufferize(Operation *op, - const OneShotBufferizationOptions &options) { + const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics) { assert(!(options.copyBeforeWrite && options.testAnalysisOnly) && "invalid combination of bufferization flags"); if (!options.copyBeforeWrite) { // If a buffer is copied before every write, no analysis is needed. - if (failed(insertTensorCopies(op, options))) + if (failed(insertTensorCopies(op, options, statistics))) return failure(); } if (options.testAnalysisOnly) return success(); - return bufferizeOp(op, options, /*copyBeforeWrite=*/options.copyBeforeWrite); + return bufferizeOp(op, options, /*copyBeforeWrite=*/options.copyBeforeWrite, + /*opFilter=*/nullptr, statistics); } diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp index 87cd11657f92a..6584dfda4eb7c 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp @@ -363,7 +363,8 @@ static void foldMemRefCasts(func::FuncOp funcOp) { LogicalResult mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp, - OneShotAnalysisState &state) { + OneShotAnalysisState &state, + BufferizationStatistics *statistics) { assert(state.getOptions().bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state); @@ -387,7 +388,7 @@ mlir::bufferization::analyzeModuleOp(ModuleOp moduleOp, equivalenceAnalysis(funcOp, aliasInfo, state, funcState); // Analyze funcOp. - if (failed(analyzeOp(funcOp, state))) + if (failed(analyzeOp(funcOp, state, statistics))) return failure(); // Run some extra function analyses. @@ -411,7 +412,8 @@ void mlir::bufferization::removeBufferizationAttributesInModule( } LogicalResult mlir::bufferization::bufferizeModuleOp( - ModuleOp moduleOp, const OneShotBufferizationOptions &options) { + ModuleOp moduleOp, const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics) { assert(options.bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); IRRewriter rewriter(moduleOp.getContext()); @@ -429,7 +431,8 @@ LogicalResult mlir::bufferization::bufferizeModuleOp( for (func::FuncOp funcOp : orderedFuncOps) { // Note: It would be good to apply cleanups here but we cannot as aliasInfo // would be invalidated. - if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite))) + if (failed(bufferizeOp(funcOp, options, options.copyBeforeWrite, + /*opFilter=*/nullptr, statistics))) return failure(); // Change buffer return types to more precise layout maps. if (options.functionBoundaryTypeConversion == @@ -444,19 +447,19 @@ LogicalResult mlir::bufferization::bufferizeModuleOp( } LogicalResult mlir::bufferization::runOneShotModuleBufferize( - ModuleOp moduleOp, const OneShotBufferizationOptions &options) { + ModuleOp moduleOp, const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics) { assert(options.bufferizeFunctionBoundaries && "expected that function boundary bufferization is activated"); assert(!(options.copyBeforeWrite && options.testAnalysisOnly) && "invalid combination of bufferization flags"); if (!options.copyBeforeWrite) { - OneShotAnalysisState analysisState(moduleOp, options); - if (failed(insertTensorCopies(moduleOp, options))) + if (failed(insertTensorCopies(moduleOp, options, statistics))) return failure(); } if (options.testAnalysisOnly) return success(); - if (failed(bufferizeModuleOp(moduleOp, options))) + if (failed(bufferizeModuleOp(moduleOp, options, statistics))) return failure(); return success(); } diff --git a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp index 677f0d5a7f4be..4cd19b4efc636 100644 --- a/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp @@ -97,7 +97,8 @@ resolveUsesInRepetitiveRegions(Operation *op, } LogicalResult mlir::bufferization::insertTensorCopies( - Operation *op, const OneShotBufferizationOptions &options) { + Operation *op, const OneShotBufferizationOptions &options, + BufferizationStatistics *statistics) { // Preprocessing: Resolve currently unsupported bufferization cases. resolveUsesInRepetitiveRegions(op, options); @@ -106,10 +107,10 @@ LogicalResult mlir::bufferization::insertTensorCopies( // analysis depending on whether function boundary bufferization is enabled or // not. if (options.bufferizeFunctionBoundaries) { - if (failed(analyzeModuleOp(cast(op), state))) + if (failed(analyzeModuleOp(cast(op), state, statistics))) return failure(); } else { - if (failed(analyzeOp(op, state))) + if (failed(analyzeOp(op, state, statistics))) return failure(); } diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir new file mode 100644 index 0000000000000..68bf33b149bbb --- /dev/null +++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir @@ -0,0 +1,15 @@ +// RUN: mlir-opt %s -one-shot-bufferize="allow-unknown-ops" -mlir-pass-statistics 2>&1 | FileCheck %s + +// CHECK: OneShotBufferize +// CHECK: (S) 1 num-buffer-alloc +// CHECK: (S) 1 num-buffer-dealloc +// CHECK: (S) 1 num-tensor-in-place +// CHECK: (S) 1 num-tensor-out-of-place +func.func @read_after_write_conflict(%cst : f32, %idx : index, %idx2 : index) + -> (f32, f32) { + %t = "test.dummy_op"() : () -> (tensor<10xf32>) + %write = tensor.insert %cst into %t[%idx2] : tensor<10xf32> + %read = "test.some_use"(%t) : (tensor<10xf32>) -> (f32) + %read2 = tensor.extract %write[%idx] : tensor<10xf32> + return %read, %read2 : f32, f32 +}