[mlir][bufferize][NFC] OneShotAnalysis: Expose analysis hooks from AnalysisState

[springerm]

Summary

This is in preparation of reusing the same AnalysisState for tensor.empty elimination and One-Shot Bufferize (to address performance bottlenecks).

Depends On: D143313

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

 //===- OneShotAnalysis.h - One-Shot (Single Pass) Analysis ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include <string>
 
 namespace mlir {
+class DominanceInfo;
+
 namespace bufferization {
 
 struct OneShotBufferizationOptions;
 struct BufferizationStatistics;
 class OneShotAnalysisState;
 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
public:
   }
 
   /// Return a reference to the BufferizationOptions.
   const OneShotBufferizationOptions &getOptions() const {
     return static_cast<const OneShotBufferizationOptions &>(
         AnalysisState::getOptions());
   }
 
+  /// Analyze the given op and its nested ops.
+  LogicalResult analyzeOp(Operation *op, const DominanceInfo &domInfo);
+
+  /// Analyze a single op (without nested ops).
+  LogicalResult analyzeSingleOp(Operation *op, const DominanceInfo &domInfo);
+
   /// Apply `fun` to all the members of the equivalence class of `v`.
   void applyOnEquivalenceClass(Value v, function_ref<void(Value)> fun) const;
 
   /// Apply `fun` to all aliases of `v`.
   void applyOnAliases(Value v, function_ref<void(Value)> fun) const;
 
   /// Return true if `v1` and `v2` bufferize to equivalent buffers.
   bool areEquivalentBufferizedValues(Value v1, Value v2) const override;

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

bufferizableInPlaceAnalysisImpl(OpOperand &operand, OneShotAnalysisState &state,
   else
     state.bufferizeInPlace(operand);
 
   LLVM_DEBUG(llvm::dbgs()
              << "//===-------------------------------------------===//\n");
   return success();
 }
 
-/// Analyze the `ops` to determine which OpOperands are inplaceable. Walk ops in
-/// reverse and bufferize ops greedily. This is a good starter heuristic.
-///
-/// Even if an op does not read or write, it may still create an alias when
-/// bufferized in-place. An example of such ops is tensor.extract_slice.
-///
-/// Rationale for bufferizing `%1 = tensor.extract_slice %0[...]` inplace:
-///
-/// When bufferized out of place, an ExtractSliceOp lowers to alloc + copy. This
-/// cannot change the flow of information for either the source or the
-/// result buffers.
-///
-/// When bufferized inplace, an ExtractSliceOp does not by itself create any
-/// read or write from memory. Instead, it has the effect of merging the alias
-/// sets of the source and the result buffers.
-///
-/// An analysis is required to ensure inplace bufferization would not result in
-/// RaW dependence violations.
-static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
-                                     OneShotAnalysisState &state,
-                                     const DominanceInfo &domInfo,
-                                     unsigned analysisFuzzerSeed = 0) {
-  if (analysisFuzzerSeed) {
-    // This is a fuzzer. For testing purposes only. Randomize the order in which
-    // operations are analyzed. The bufferization quality is likely worse, but
-    // we want to make sure that no assertions are triggered anywhere.
-    std::mt19937 g(analysisFuzzerSeed);
-    llvm::shuffle(ops.begin(), ops.end(), g);
-  }
-
-  // Analyze a single op.
-  auto analyzeOp = [&](Operation *op) {
-    for (OpOperand &opOperand : op->getOpOperands())
-      if (opOperand.get().getType().isa<TensorType>())
-        if (failed(bufferizableInPlaceAnalysisImpl(opOperand, state, domInfo)))
-          return failure();
-    return success();
-  };
-
-  OneShotBufferizationOptions::AnalysisHeuristic heuristic =
-      state.getOptions().analysisHeuristic;
-  if (heuristic == OneShotBufferizationOptions::AnalysisHeuristic::BottomUp) {
-    // Default: Walk ops in reverse for better interference analysis.
-    for (Operation *op : reverse(ops))
-      if (failed(analyzeOp(op)))
-        return failure();
-  } else if (heuristic ==
-             OneShotBufferizationOptions::AnalysisHeuristic::TopDown) {
-    for (Operation *op : ops)
-      if (failed(analyzeOp(op)))
+LogicalResult
+OneShotAnalysisState::analyzeSingleOp(Operation *op,
+                                      const DominanceInfo &domInfo) {
+      if (failed(bufferizableInPlaceAnalysisImpl(opOperand, *this, domInfo)))
         return failure();
-  } else {
-    llvm_unreachable("unsupported heuristic");
-  }
-
   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 all ops that are contained in `op`.
-static LogicalResult inPlaceAnalysis(Operation *op,
-                                     OneShotAnalysisState &state,
-                                     const DominanceInfo &domInfo,
-                                     unsigned analysisFuzzerSeed = 0) {
-  // Collect ops so we can build our own reverse traversal.
-  SmallVector<Operation *> ops;
-  op->walk([&](Operation *op) {
-    // No tensors => no buffers.
-    if (!hasTensorSemantics(op))
-      return;
-    ops.push_back(op);
-  });
-
-  return inPlaceAnalysis(ops, state, domInfo, analysisFuzzerSeed);
-}
-
 /// Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.
 static void equivalenceAnalysis(SmallVector<Operation *> &ops,
                                 OneShotAnalysisState &state) {
   for (Operation *op : ops)
     if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
       for (OpResult opResult : op->getOpResults())
         if (opResult.getType().isa<TensorType>())
           for (OpOperand *opOperand :
op->walk<WalkOrder::PostOrder>([&](Operation *op) {
     if (none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
   equivalenceAnalysis(ops, state);
 }
 
+LogicalResult OneShotAnalysisState::analyzeOp(Operation *op,
+                                              const DominanceInfo &domInfo) {
+  // Collect ops so we can build our own reverse traversal.
+  SmallVector<Operation *> ops;
+  op->walk([&](Operation *op) {
+    // No tensors => no buffers.
+    if (!hasTensorSemantics(op))
+      return;
+    ops.push_back(op);
+  });
+
+  if (getOptions().analysisFuzzerSeed) {
+    // This is a fuzzer. For testing purposes only. Randomize the order in which
+    // operations are analyzed. The bufferization quality is likely worse, but
+    // we want to make sure that no assertions are triggered anywhere.
+    std::mt19937 g(getOptions().analysisFuzzerSeed);
+    llvm::shuffle(ops.begin(), ops.end(), g);
+  }
+
+  OneShotBufferizationOptions::AnalysisHeuristic heuristic =
+      getOptions().analysisHeuristic;
+  if (heuristic == OneShotBufferizationOptions::AnalysisHeuristic::BottomUp) {
+    // Default: Walk ops in reverse for better interference analysis.
+    for (Operation *op : reverse(ops))
+      if (failed(analyzeSingleOp(op, domInfo)))
+        return failure();
+  } else if (heuristic ==
+             OneShotBufferizationOptions::AnalysisHeuristic::TopDown) {
+    for (Operation *op : ops)
+      if (failed(analyzeSingleOp(op, domInfo)))
+        return failure();
+  } else {
+    llvm_unreachable("unsupported heuristic");
+  }
+
+  equivalenceAnalysis(op, *this);
+  return success();
+}
+
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult checkAliasInfoConsistency(Operation *op,
                                                const DominanceInfo &domInfo,
                                                OneShotAnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
 
   WalkResult walkResult = op->walk([&](BufferizableOpInterface op) {
     // Skip ops that are not in the filter.
LogicalResult bufferization::analyzeOp(Operation *op,
                                        BufferizationStatistics *statistics) {
   DominanceInfo domInfo(op);
   const OneShotBufferizationOptions &options = state.getOptions();
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state)))
     return failure();
 
   // If the analysis fails, just return.
-  if (failed(inPlaceAnalysis(op, state, domInfo, options.analysisFuzzerSeed)))
+  if (failed(state.analyzeOp(op, domInfo)))
     return failure();
 
   if (statistics) {
     statistics->numTensorInPlace = state.getStatNumTensorInPlace();
     statistics->numTensorOutOfPlace = state.getStatNumTensorOutOfPlace();
   }
 
-  equivalenceAnalysis(op, state);
-
   bool failedAnalysis = false;
   if (!options.allowReturnAllocs)
     failedAnalysis |= failed(assertNoAllocsReturned(op, state));
 
   // Gather some extra analysis data.
   state.gatherYieldedTensors(op);
   state.gatherUndefinedTensorUses(op);