[mlir][bufferize][NFC] Add statistics to OneShotBufferizePass

[springerm]

Summary

Print statistics about the number of alloc/deallocs and in-place/out-of-place bufferization.

mlir/include/mlir/Dialect/Bufferization/Transforms/Bufferize.h

 
 namespace mlir {
 namespace bufferization {
 
 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 {
 public:
   BufferizeTypeConverter();
 };
 
 /// Marks ops used by bufferization for type conversion materializations as

 /// because read-after-write conflicts may materialize during bufferization.
 /// `copyBeforeWrite = false` is safe only if the input IR is guaranteed to
 /// *not* require any out-of-place bufferization.
 ///
 /// Note: This function bufferizes ops without utilizing analysis results. It
 /// 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();
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h

 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 
 namespace mlir {
 namespace bufferization {
 
 struct OneShotBufferizationOptions;
 class BufferizationAliasInfo;
+struct BufferizationStatistics;
 class OneShotAnalysisState;
 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
   enum class AnalysisHeuristic { BottomUp, TopDown };
 
   OneShotBufferizationOptions() = default;
 
public:
   void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
 
   /// Mark a value as in-place bufferized.
   void markInPlace(OpOperand &o) { inplaceBufferized.insert(&o); }
 
   /// 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
   /// comparison but it's not like UnionFind needs ordering anyway.
   struct ValueComparator {
     bool operator()(const Value &lhs, const Value &rhs) const {
       return lhs.getImpl() < rhs.getImpl();
     }
private:
 
   /// Auxiliary structure to store all the equivalent buffer classes. Equivalent
   /// buffer information is "must be" conservative: Only if two values are
   /// guaranteed to be equivalent at runtime, they said to be equivalent. It is
   /// possible that, in the presence of branches, it cannot be determined
   /// statically if two values are equivalent. In that case, the values are
   /// considered to be not equivalent.
   llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
+
+  // Bufferization statistics.
+  int64_t statNumTensorOutOfPlace = 0;
+  int64_t statNumTensorInPlace = 0;
 };
 
 /// State for analysis-enabled bufferization. This class keeps track of alias
 /// (via BufferizationAliasInfo) to decide if tensor OpOperands should bufferize
 /// in-place.
 class OneShotAnalysisState : public AnalysisState {
 public:
   OneShotAnalysisState(Operation *op,
private:
 
   /// Extensions attached to the state, identified by the TypeID of their type.
   /// Only one extension of any given type is allowed.
   DenseMap<TypeID, std::unique_ptr<Extension>> extensions;
 };
 
 /// 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
 
 MLIR_DECLARE_EXPLICIT_TYPE_ID(mlir::bufferization::OneShotAnalysisState)
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H

mlir/include/mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h

 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 
 namespace mlir {
 
 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`.
 ///
 /// Note: This function does not run One-Shot Analysis. No buffer copies are
 /// 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);
 
 /// Run One-Shot Module Bufferization on the given module. Performs a simple
 /// function call analysis to determine which function arguments are
 /// inplaceable. Then analyzes and bufferizes FuncOps one-by-one with One-Shot
 /// Bufferize.
 LogicalResult runOneShotModuleBufferize(
     ModuleOp moduleOp,
-    const bufferization::OneShotBufferizationOptions &options);
+    const bufferization::OneShotBufferizationOptions &options,
+    BufferizationStatistics *statistics = nullptr);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

Option<"printConflicts", "print-conflicts", "bool",
             /*default=*/"false",
            "Test only: Annotate IR with RaW conflicts. Requires "
            "test-analysis-only.">,
     Option<"unknownTypeConversion", "unknown-type-conversion", "std::string",
            /*default=*/"\"fully-dynamic-layout-map\"",
            "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"> {
   let summary = "Promotes heap-based allocations to automatically managed "
                 "stack-based allocations";
   let description = [{
     This pass implements a simple algorithm to convert heap-based memory
     allocations to stack-based ones. It uses a built-in heuristic to decide

mlir/include/mlir/Dialect/Bufferization/Transforms/Transforms.h

 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_TRANSFORMS_H
 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/IR/Operation.h"
 
 namespace mlir {
 namespace bufferization {
 class AnalysisState;
+struct BufferizationStatistics;
 struct OneShotBufferizationOptions;
 
 /// A function that matches anchor OpOperands for tensor::EmptyOp elimination.
 /// If an OpOperand is matched, the function should populate the SmallVector
 /// with all values that are needed during `RewriteFn` to produce the
 /// replacement value.
 using AnchorMatchFn = std::function<bool(OpOperand &, SmallVector<Value> &)>;
 

 /// (and some other conditions are met).
 LogicalResult insertSliceAnchoredEmptyTensorEliminationStep(
     RewriterBase &rewriter, Operation *op, bufferization::AnalysisState &state);
 
 /// Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
 /// 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
 /// additional buffer allocations.
 LogicalResult insertTensorCopies(Operation *op, const AnalysisState &state);
 
 /// Populate patterns to lower tensor.empty ops to bufferization.alloc_tensor
 /// ops.
 void populateEmptyTensorToAllocTensorPattern(RewritePatternSet &patterns);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Bufferization/Transforms/Bufferize.cpp

if (!options) {
         // No filter specified: All other ops are allowed.
         return true;
       };
       opt.opFilter.allowOperation(filterFn);
     } else {
       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;
 
     OpPassManager cleanupPipeline("builtin.module");
     cleanupPipeline.addPass(createCanonicalizerPass());
     cleanupPipeline.addPass(createCSEPass());
     cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
     (void)runPipeline(cleanupPipeline, moduleOp);

 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
 class BufferizationRewriter : public IRRewriter {
 public:
   BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps,
                         DenseSet<Operation *> &toMemrefOps,
                         SmallVector<Operation *> &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 {
     IRRewriter::notifyOperationRemoved(op);
     erasedOps.insert(op);
     // Erase if present.
     toMemrefOps.erase(op);
   }
 
   void notifyOperationInserted(Operation *op) override {
     IRRewriter::notifyOperationInserted(op);
     erasedOps.erase(op);
 
+    // Gather statistics about allocs and deallocs.
+    if (statistics) {
+      if (auto sideEffectingOp = dyn_cast<MemoryEffectOpInterface>(op)) {
+        statistics->numBufferAlloc += static_cast<int64_t>(
+            sideEffectingOp.hasEffect<MemoryEffects::Allocate>());
+        statistics->numBufferDealloc += static_cast<int64_t>(
+            sideEffectingOp.hasEffect<MemoryEffects::Free>());
+      }
+    }
+
     // Keep track of to_memref ops.
     if (isa<ToMemrefOp>(op)) {
       toMemrefOps.insert(op);
       return;
     }
 
     // Skip to_tensor ops.
     if (isa<ToTensorOp>(op))
private:
   SmallVector<Operation *> &worklist;
 
   /// The analysis state. Used for debug assertions and access to the
   /// bufferization options.
   const AnalysisState analysisState;
 
   /// 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)))
       return failure();
   }
 
   // Keep track of to_memref ops.
   DenseSet<Operation *> toMemrefOps;
if (hasTensorSemantics(op) && !isa<func::FuncOp>(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
 
   // 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.
     if (erasedOps.contains(nextOp))
       continue;
     // Skip ops that are not bufferizable or not allowed.
     auto bufferizableOp = options.dynCastBufferizableOp(nextOp);
     if (!bufferizableOp)

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 /// Return `true` if a value was marked as in-place bufferized.
 bool BufferizationAliasInfo::isInPlace(OpOperand &operand) const {
   return inplaceBufferized.contains(&operand);
 }
 
 /// 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`.
 void BufferizationAliasInfo::applyOnEquivalenceClass(
     Value v, function_ref<void(Value)> fun) const {
   auto leaderIt = equivalentInfo.findLeader(v);
   for (auto mit = leaderIt, meit = equivalentInfo.member_end(); mit != meit;
        ++mit) {
OneShotAnalysisState::OneShotAnalysisState(
     Operation *op, const OneShotBufferizationOptions &options)
     : AnalysisState(options, TypeID::get<OneShotAnalysisState>()),
       aliasInfo(op) {
   // Set up alias sets for OpResults that must bufferize in-place. This should
   // be done before making any other bufferization decisions.
   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<TensorType>())
-        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();
   });
 }
 
 bool OneShotAnalysisState::isInPlace(OpOperand &opOperand) const {
   return aliasInfo.isInPlace(opOperand);
 }
 
op->walk([&](Operation *returnOp) {
 
     return WalkResult::advance();
   });
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
-                                       OneShotAnalysisState &state) {
+                                       OneShotAnalysisState &state,
+                                       BufferizationStatistics *statistics) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const OneShotBufferizationOptions &options = state.getOptions();
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   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;
   if (!options.allowReturnAllocs)
     failedAnalysis |= failed(assertNoAllocsReturned(op, options, aliasInfo));
 
   // Gather some extra analysis data.
   state.gatherYieldedTensors(op);
LogicalResult bufferization::analyzeOp(Operation *op,
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, options);
 
   return success(!failedAnalysis);
 }
 
 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);
 }

mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

static void foldMemRefCasts(func::FuncOp funcOp) {
 
   auto newFuncType = FunctionType::get(
       funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
   funcOp.setType(newFuncType);
 }
 
 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);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Analyze ops.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Now analyzing function.
     funcState.startFunctionAnalysis(funcOp);
 
     // Gather equivalence info for CallOps.
     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.
     if (failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) ||
         failed(funcOpBbArgReadWriteAnalysis(funcOp, state, funcState)))
       return failure();
 
     // Mark op as fully analyzed.
     funcState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
   }
 
   return success();
 }
 
 void mlir::bufferization::removeBufferizationAttributesInModule(
     ModuleOp moduleOp) {
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 }
 
 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());
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Bufferize functions.
   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 ==
         LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }
 
 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();
 }

mlir/lib/Dialect/Bufferization/Transforms/TensorCopyInsertion.cpp

for (OpOperand &opOperand : bufferizableOp->getOpOperands()) {
         use->set(tensorCopy);
     }
 
     return WalkResult::advance();
   });
 }
 
 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);
 
   OneShotAnalysisState state(op, options);
   // Run normal One-Shot Bufferize analysis or One-Shot Module Bufferize
   // analysis depending on whether function boundary bufferization is enabled or
   // not.
   if (options.bufferizeFunctionBoundaries) {
-    if (failed(analyzeModuleOp(cast<ModuleOp>(op), state)))
+    if (failed(analyzeModuleOp(cast<ModuleOp>(op), state, statistics)))
       return failure();
   } else {
-    if (failed(analyzeOp(op, state)))
+    if (failed(analyzeOp(op, state, statistics)))
       return failure();
   }
 
   if (options.testAnalysisOnly)
     return success();
 
   return insertTensorCopies(op, state);
 }

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-pass-statistics.mlir

+// 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
+}