[mlir][linalg][bufferize][NFC] Clean up headers and function visibility

[springerm]

Summary

Depends On D113962

[springerm]

update

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

 //===- ComprehensiveBufferize.h - Linalg bufferization pass -----*- 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_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/Value.h"
-#include "llvm/ADT/SetOperations.h"
 
 namespace mlir {
 
-class DominanceInfo;
-class FuncOp;
-class GlobalCreator;
 class ModuleOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 // TODO: from some HW description.
 static constexpr int64_t kBufferAlignments = 128;
 
-struct BufferizationState;
-
-/// Analyze the `ops` to determine which OpResults are inplaceable.
-LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
-                              BufferizationAliasInfo &aliasInfo,
-                              const DominanceInfo &domInfo,
-                              unsigned analysisFuzzerSeed = 0);
-
-// TODO: Do not expose those functions in the header file.
-/// Default allocation function that is used by the comprehensive bufferization
-/// pass. The default currently creates a ranked memref using `memref.alloc`.
-Optional<Value> defaultAllocationFn(OpBuilder &b, Location loc, MemRefType type,
-                                    const SmallVector<Value> &dynShape);
-
-/// Default deallocation function that is used by the comprehensive
-/// bufferization pass. It expects to recieve back the value called from the
-/// `defaultAllocationFn`.
-void defaultDeallocationFn(OpBuilder &b, Location loc, Value allocatedBuffer);
-
-/// Default memory copy function that is used by the comprehensive bufferization
-/// pass. Creates a `linalg.copy` op.
-void defaultMemCpyFn(OpBuilder &b, Location loc, Value from, Value to);
-
 /// Return default allocation callbacks.
 std::unique_ptr<AllocationCallbacks> defaultAllocationCallbacks();
 
 /// Bufferize one particular op.
 /// `bufferizedFunctionTypes` (resp. `globalCreator`) are expected to be
 /// non-null if `op` is a CallOpInterface (resp. GlobalCreator).
 LogicalResult
 bufferizeOp(Operation *op, BufferizationState &state,
             DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes = nullptr);
 
 /// Register external models implemented for the `BufferizableOpInterface`.
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
+/// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   BufferizationOptions();
 
+  // BufferizationOptions cannot be copied.
+  BufferizationOptions(const BufferizationOptions &other) = delete;
+
   /// Register a "post analysis" step. Such steps are executed after the
   /// analysis, but before bufferization.
   template <typename Step, typename... Args>
   void addPostAnalysisStep(Args... args) {
     postAnalysisSteps.emplace_back(
         std::make_unique<Step>(std::forward<Args>(args)...));
   }
 
+  /// Helper functions for allocation, deallocation, memory copying.
   std::unique_ptr<AllocationCallbacks> allocationFns;
+
+  /// Specifies whether returning newly allocated memrefs should be allowed.
+  /// Otherwise, a pass failure is triggered.
   bool allowReturnMemref = false;
+
+  /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
+  /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
+
+  /// If set to `true`, does not modify the IR apart from adding attributes (for
+  /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
+
+  /// Registered post analysis steps.
   std::vector<std::unique_ptr<PostAnalysisStep>> postAnalysisSteps;
 };
 
 LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
                                         const BufferizationOptions &options);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

#endif // NDEBUG
   else
     aliasInfo.bufferizeInPlace(result, operand);
 
   LDBG("Done inplace analysis for result #" << resultNumber << '\n');
 
   return success();
 }
 
-/// Determine if `operand` can be bufferized in-place with one of the op's
-/// results.
+/// Analyze the `ops` to determine which OpResults 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
-bufferizableInPlaceAnalysis(OpOperand &operand,
+static LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                                      BufferizationAliasInfo &aliasInfo,
-                            const DominanceInfo &domInfo) {
-  auto bufferizableOp = dyn_cast<BufferizableOpInterface>(operand.getOwner());
-  if (!bufferizableOp)
-    return success();
-  if (OpResult result = bufferizableOp.getAliasingOpResult(operand))
-    return bufferizableInPlaceAnalysisImpl(operand, result, aliasInfo, domInfo);
-  return success();
-}
-
-/// Analyze the `ops` to determine which OpResults are inplaceable. Walk ops in
-/// reverse and bufferize ops greedily. This is a good starter heuristic.
-/// ExtractSliceOps are interleaved with other ops in traversal order.
-LogicalResult mlir::linalg::comprehensive_bufferize::inPlaceAnalysis(
-    SmallVector<Operation *> &ops, BufferizationAliasInfo &aliasInfo,
-    const DominanceInfo &domInfo, unsigned analysisFuzzerSeed) {
+                                     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);
   }
 
   // Walk ops in reverse for better interference analysis.
   for (Operation *op : reverse(ops))
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
-        if (failed(bufferizableInPlaceAnalysis(opOperand, aliasInfo, domInfo)))
+        if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
+          if (OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand))
+            if (failed(bufferizableInPlaceAnalysisImpl(opOperand, opResult,
+                                                       aliasInfo, domInfo)))
               return failure();
 
   return success();
 }
 
 /// Analyze the `funcOp` body to determine which OpResults are inplaceable.
 static LogicalResult
 inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
                           const DominanceInfo &domInfo,
inPlaceAnalysisFuncOpBody(FuncOp funcOp, BufferizationAliasInfo &aliasInfo,
 
   return res;
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization entry-point for functions.
 //===----------------------------------------------------------------------===//
 
-Optional<Value> mlir::linalg::comprehensive_bufferize::defaultAllocationFn(
-    OpBuilder &b, Location loc, MemRefType type,
-    const SmallVector<Value> &dynShape) {
-  Value allocated = b.create<memref::AllocOp>(
-      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
-  return allocated;
-}
-
-void mlir::linalg::comprehensive_bufferize::defaultDeallocationFn(
-    OpBuilder &b, Location loc, Value allocatedBuffer) {
-  b.create<memref::DeallocOp>(loc, allocatedBuffer);
-}
-
-void mlir::linalg::comprehensive_bufferize::defaultMemCpyFn(OpBuilder &b,
-                                                            Location loc,
-                                                            Value from,
-                                                            Value to) {
-  b.create<memref::CopyOp>(loc, from, to);
-}
-
 LogicalResult mlir::linalg::comprehensive_bufferize::bufferizeOp(
     Operation *op, BufferizationState &state,
     DenseMap<FuncOp, FunctionType> *bufferizedFunctionTypes) {
   OpBuilder b(op->getContext());
 
   // CallOps are handled separately.
   if (auto callOp = dyn_cast<CallOpInterface>(op)) {
     LDBG("Begin bufferize:\n" << callOp << '\n');
#endif // NDEBUG
   moduleOp.walk([&](FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationFuncArguments(bbArg);
   });
 
   return success();
 }
 
+/// Default allocation function that is used by the comprehensive bufferization
+/// pass. The default currently creates a ranked memref using `memref.alloc`.
+static Optional<Value> defaultAllocationFn(OpBuilder &b, Location loc,
+                                           MemRefType type,
+                                           const SmallVector<Value> &dynShape) {
+  Value allocated = b.create<memref::AllocOp>(
+      loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
+  return allocated;
+}
+
+/// Default deallocation function that is used by the comprehensive
+/// bufferization pass. It expects to recieve back the value called from the
+/// `defaultAllocationFn`.
+static void defaultDeallocationFn(OpBuilder &b, Location loc,
+                                  Value allocatedBuffer) {
+  b.create<memref::DeallocOp>(loc, allocatedBuffer);
+}
+
+/// Default memory copy function that is used by the comprehensive bufferization
+/// pass. Creates a `memref.copy` op.
+static void defaultMemCpyFn(OpBuilder &b, Location loc, Value from, Value to) {
+  b.create<memref::CopyOp>(loc, from, to);
+}
+
 std::unique_ptr<AllocationCallbacks>
 mlir::linalg::comprehensive_bufferize::defaultAllocationCallbacks() {
   return std::make_unique<AllocationCallbacks>(
       defaultAllocationFn, defaultDeallocationFn, defaultMemCpyFn,
       createNewAllocDeallocPairForShapedValue);
 }
 
 // Default constructor for BufferizationOptions that sets all allocation