[mlir][bufferize][NFC] Minor code and comment cleanups

[springerm]

Summary

Depends On D135055

[springerm]

rebase

[springerm]

rebase

[springerm]

update

[springerm]

update

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

public:
   OneShotAnalysisState(const OneShotAnalysisState &) = delete;
 
   ~OneShotAnalysisState() override = default;
 
   static bool classof(const AnalysisState *base) {
     return base->getType() == TypeID::get<OneShotAnalysisState>();
   }
 
+  /// Return a reference to the BufferizationOptions.
+  const OneShotBufferizationOptions &getOptions() const {
+    return static_cast<const OneShotBufferizationOptions &>(
+        AnalysisState::getOptions());
+  }
+
   /// Return a reference to the BufferizationAliasInfo.
   BufferizationAliasInfo &getAliasInfo() { return aliasInfo; }
 
   /// Return `true` if the given OpResult has been decided to bufferize inplace.
   bool isInPlace(OpOperand &opOperand) const override;
 
   /// 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

 //===- OneShotAnalysis.cpp - One-Shot (Single Pass) Analysis --------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
-// One-Shot Analysis analyzes function bodies. Function boundaries (FuncOp
-// bbArgs, CallOps, ReturnOps) are treated as "unknown" ops.
-// ModuleBufferization.cpp is an extension of One-Shot Analysis for simple
-// call graphs.
+// One-Shot Analysis analyzes function bodies. By default, function boundaries
+// (FuncOp bbArgs, CallOps, ReturnOps) are treated as "unknown" ops.
+// OneShotModuleBufferization.cpp is an extension of One-Shot Analysis for
+// simple call graphs without loops.
 //
-// One-Shot Bufferize consists of two phases.
+// One-Shot Bufferize consists of three phases.
 //
-// 1. Analyze ops to decide which OpResults can bufferize inplace, i.e., without
-//    inserting buffer copies. The analysis queries op bufferization semantics
-//    via `BufferizableOpInterface`.
-// 2. Bufferize ops by calling `BufferizableOpInterface::bufferize`. This
-//    function does not generate buffer copies for OpResults that were decided
-//    to bufferize inplace during the analysis phase.
+// 1. Analyze ops to decide which OpOperands can bufferize inplace, i.e.,
+//    without inserting buffer copies. The analysis queries op bufferization
+//    semantics via `BufferizableOpInterface`.
+// 2. Insert copies for OpOperands that were decided to bufferize out-of-place
+//    in tensor land during `TensorCopyInsertion`.
+// 3. Bufferize ops by calling `BufferizableOpInterface::bufferize`.
 //
-// This file contains only the analysis. The actual bufferization is implemented
-// via `bufferizeOp` (Bufferize.h). For convenience, this file also contains a
-// helper function `runOneShotBufferize` that analyzes an op (and its nested
-// ops) and then bufferizes it.
+// This file contains only the analysis. For convenience, this file also
+// contains a helper function `runOneShotBufferize` that analyzes an op (and its
+// nested ops) and then bufferizes it.
 //
 // Inplace bufferization decisions are passed from the analysis to the
-// bufferization phase via `AnalysisState` and `BufferizationAliasInfo`.
-// They can be printed for debugging purposes with `testAnalysisOnly`.
+// `TensorCopyInsertion` phase via `AnalysisState`. They can be printed for
+// debugging purposes with `testAnalysisOnly`.
 //
 // Ops that do not implement `BufferizableOpInterface` can be analyzed but are
 // treated conservatively. E.g., the analysis has to assume that their tensor
 // OpOperands bufferize to memory writes. While such ops can be analyzed, they
 // are not bufferized and remain in the IR. to_tensor and to_memref ops are
 // inserted at the bufferization boundary.
 //
 // This analysis caters to high-performance codegen where buffer reuse is deemed

 
 using namespace mlir;
 using namespace mlir::bufferization;
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific attribute manipulation.
-// These are for testing and debugging only. Bufferization information is
-// stored in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR
-// is annotated with the results of the analysis (copied from
-// BufferizationAliasInfo), so that they can be checked in tests.
+// These are for testing and debugging only. Bufferization information is stored
+// in BufferizationAliasInfo. When run with `testAnalysisOnly`, the IR is
+// annotated with the results of the analysis, so that they can be checked in
+// tests.
 //===----------------------------------------------------------------------===//
 
-/// Attribute marker to specify op results that can be bufferized inPlace.
-constexpr StringLiteral kInPlaceResultsAttrName = "__inplace_operands_attr__";
+/// Attribute marker to specify op operands that bufferize in-place.
+constexpr StringLiteral kInPlaceOperandsAttrName = "__inplace_operands_attr__";
 
 /// Mark whether OpOperand will be bufferized inplace.
 static void setInPlaceOpOperand(OpOperand &opOperand, bool inPlace) {
   Operation *op = opOperand.getOwner();
-  auto attr =
-      op->getAttr(kInPlaceResultsAttrName).dyn_cast_or_null<ArrayAttr>();
   SmallVector<StringRef> inPlaceVector;
-  if (attr) {
-    inPlaceVector = SmallVector<StringRef>(
-        llvm::to_vector<4>(attr.getAsValueRange<StringAttr>()));
+  if (auto attr = op->getAttr(kInPlaceOperandsAttrName)) {
+    inPlaceVector = SmallVector<StringRef>(llvm::to_vector<4>(
+        attr.cast<ArrayAttr>().getAsValueRange<StringAttr>()));
   } else {
     inPlaceVector = SmallVector<StringRef>(op->getNumOperands(), "none");
     for (OpOperand &opOperand : op->getOpOperands())
       if (opOperand.get().getType().isa<TensorType>())
         inPlaceVector[opOperand.getOperandNumber()] = "false";
   }
-
   inPlaceVector[opOperand.getOperandNumber()] = inPlace ? "true" : "false";
-  op->setAttr(kInPlaceResultsAttrName,
+  op->setAttr(kInPlaceOperandsAttrName,
               OpBuilder(op).getStrArrayAttr(inPlaceVector));
 }
 
 //===----------------------------------------------------------------------===//
 // BufferizationAliasInfo
 //===----------------------------------------------------------------------===//
 
 BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
for (OpOperand &opOperand : op->getOpOperands())
         if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
           if (failed(bufferizableInPlaceAnalysisImpl(opOperand, aliasInfo,
                                                      state, domInfo)))
             return failure();
     return success();
   };
 
   OneShotBufferizationOptions::AnalysisHeuristic heuristic =
-      static_cast<const OneShotBufferizationOptions &>(state.getOptions())
-          .analysisHeuristic;
+      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)
checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
 
   return success(!walkResult.wasInterrupted());
 }
 
 /// Annotate the IR with the result of the analysis. For testing/debugging only.
 static void
 annotateOpsWithBufferizationMarkers(Operation *op,
                                     const BufferizationAliasInfo &aliasInfo,
-                                    AnalysisState &state) {
-  op->walk([&](Operation *op) {
-    if (auto bufferizableOp = state.getOptions().dynCastBufferizableOp(op))
-      for (OpOperand &opOperand : op->getOpOperands())
+                                    const BufferizationOptions &options) {
+  // Add __inplace_operands_attr__.
+  op->walk([&](BufferizableOpInterface bufferizableOp) {
+    if (options.isOpAllowed(bufferizableOp.getOperation()))
+      for (OpOperand &opOperand : bufferizableOp->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>())
           setInPlaceOpOperand(opOperand, aliasInfo.isInPlace(opOperand));
   });
 }
 
 /// Assert that IR is in destination-passing style. I.e., every value that is
 /// returned or yielded from a block is:
 /// * aliasing a bbArg of that block or a parent block, or
assertDestinationPassingStyle(Operation *op, AnalysisState &state,
 
   return status;
 }
 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
-  const auto &options =
-      static_cast<const OneShotBufferizationOptions &>(state.getOptions());
+  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();
LogicalResult bufferization::analyzeOp(Operation *op,
   op->walk([&](Operation *op) {
     if (BufferizableOpInterface bufferizableOp =
             options.dynCastBufferizableOp(op))
       failedAnalysis |= failed(bufferizableOp.verifyAnalysis(state));
   });
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
-    annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
+    annotateOpsWithBufferizationMarkers(op, aliasInfo, options);
 
   return success(!failedAnalysis);
 }
 
 LogicalResult
 bufferization::runOneShotBufferize(Operation *op,
                                    const OneShotBufferizationOptions &options) {
   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.
-    OneShotAnalysisState state(op, options);
     if (failed(insertTensorCopies(op, options)))
       return failure();
   }
   if (options.testAnalysisOnly)
     return success();
   return bufferizeOp(op, options, /*copyBeforeWrite=*/options.copyBeforeWrite);
 }

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) {
-  OneShotBufferizationOptions options =
-      static_cast<const OneShotBufferizationOptions &>(state.getOptions());
-  assert(options.bufferizeFunctionBoundaries &&
+  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.