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

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

Differential Revision: https://reviews.llvm.org/D139538

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();
 

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<Value, ValueComparator> 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

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

@@ -15,20 +15,23 @@ 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);
@@ -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

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"> {

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

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<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 {
@@ -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<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);
@@ -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.

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<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();
   });
 }
@@ -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);
 }

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();
 }