OneShotModuleBufferize.cpp

Go to the documentation of this file.

//===- OneShotModuleBufferize.cpp - Bufferization across Func. Boundaries
//----===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Module Bufferization is an extension of One-Shot Bufferize that
// bufferizes function boundaries. It provides `BufferizableOpInterface`
// implementations for FuncOp, CallOp and ReturnOp. Although it is named
// Module Bufferization, it may operate on any SymbolTable.
//
// Module Bufferization is run via `runOneShotModuleBufferize(SymbolTableOp,
// ...)`. This function analyzes the given op and determines the order of
// analysis and bufferization: Functions that are called are processed before
// their respective callers.
//
// After analyzing a FuncOp, additional information about its bbArgs is
// gathered and stored in `FuncAnalysisState`.
//
// * `aliasingFuncOpBBArgsAnalysis` determines the equivalent/aliasing bbArgs
// for
//   each tensor return value (if any).
// * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
//   read/written.
//
// Module Bufferization implements the following calling convention.
//
// * In the absence of conflicts within a FuncOp, the FuncOp's bbArgs may always
//   be written to in-place.
// * If a tensor operand of a CallOp is read after the CallOp, the operand of
//   the CallOp must bufferize out-of-place.
//
// Example: The tensor.insert op bufferizes in-place because it is allowed to
// modify the buffer of `%t1` directly. The CallOp in `caller` must bufferize
// out-of-place because `%t0` is modified by the callee but read by the
// tensor.extract op. The analysis of CallOps decides whether an OpOperand must
// bufferize out-of-place based on results of `funcOpBbArgReadWriteAnalysis`.
// ```
// func @callee(%t1 : tensor<?xf32>) -> tensor<?xf32> {
//   %f = ... : f32
//   %0 = tensor.insert %f into %t1[...] : tensor<?xf32>
//   return %0 : tensor<?xf32>
// }
//
// func @caller() -> () {
//   %t0 = ... : tensor<?xf32>
//   %1 = call @callee(%t0) : (tensor<?xf32>) -> (tensor<?xf32>)
//   %2 = tensor.extract %1[...]  : tensor<?xf32>
// }
// ```
//
// Note: If a function is external, `funcOpBbArgReadWriteAnalysis` cannot
// analyze the function body. In such a case, the CallOp analysis conservatively
// assumes that each tensor OpOperand is both read and written.
//
// TODO: Add FuncOp attributes so that bbArgs of external FuncOps can be marked
// as "not reading" and/or "not writing".
 
#include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.h"
#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
#include "mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/SmallVectorExtras.h"
 
using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::bufferization::func_ext;
 
/// A mapping of FuncOps to their callers.
using FuncCallerMap = DenseMap<func::FuncOp, DenseSet<Operation *>>;
 
/// Get or create FuncAnalysisState.
static FuncAnalysisState &
getOrCreateFuncAnalysisState(OneShotAnalysisState &state) {
  auto *result = state.getExtension<FuncAnalysisState>();
  if (result)
    return *result;
  return state.addExtension<FuncAnalysisState>();
}
 
namespace {
 
/// Annotate IR with the results of the analysis. For testing purposes only.
static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
                                          BlockArgument bbArg) {
  const char *kEquivalentArgsAttr = "__equivalent_func_args__";
  Operation *op = returnVal.getOwner();
 
  SmallVector<int64_t> equivBbArgs;
  if (op->hasAttr(kEquivalentArgsAttr)) {
    auto attr = cast<ArrayAttr>(op->getAttr(kEquivalentArgsAttr));
    equivBbArgs = llvm::map_to_vector<4>(attr, [](Attribute a) {
      return cast<IntegerAttr>(a).getValue().getSExtValue();
    });
  } else {
    equivBbArgs.append(op->getNumOperands(), -1);
  }
  equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
  OpBuilder b(op->getContext());
  op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
}
 
/// Store function BlockArguments that are equivalent to/aliasing a returned
/// value in FuncAnalysisState.
static LogicalResult
aliasingFuncOpBBArgsAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                             FuncAnalysisState &funcState) {
  if (funcOp.getBody().empty()) {
    // No function body available. Conservatively assume that every tensor
    // return value may alias with any tensor bbArg.
    FunctionType type = funcOp.getFunctionType();
    for (const auto &inputIt : llvm::enumerate(type.getInputs())) {
      if (!isa<TensorLikeType>(inputIt.value()))
        continue;
      for (const auto &resultIt : llvm::enumerate(type.getResults())) {
        if (!isa<TensorLikeType>(resultIt.value()))
          continue;
        int64_t returnIdx = resultIt.index();
        int64_t bbArgIdx = inputIt.index();
        funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(returnIdx);
      }
    }
    return success();
  }
 
  // Find all func.return ops.
  SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
  // TODO: throw error when there is any non-func.return op that has the
  // ReturnLike trait
  if (returnOps.empty()) {
    return funcOp.emitError("cannot bufferize func.func without func.return");
  }
 
  // Build alias sets. Merge all aliases from all func.return ops.
  for (BlockArgument bbArg : funcOp.getArguments()) {
    if (isa<TensorLikeType>(bbArg.getType())) {
      int64_t bbArgIdx = bbArg.getArgNumber();
      // Store aliases in a set, so that we don't add the same alias twice.
      SetVector<int64_t> aliases;
      for (func::ReturnOp returnOp : returnOps) {
        for (OpOperand &returnVal : returnOp->getOpOperands()) {
          if (isa<TensorLikeType>(returnVal.get().getType())) {
            int64_t returnIdx = returnVal.getOperandNumber();
            if (state.areAliasingBufferizedValues(returnVal.get(), bbArg))
              aliases.insert(returnIdx);
          }
        }
      }
      for (int64_t alias : aliases)
        funcState.aliasingReturnVals[funcOp][bbArgIdx].push_back(alias);
    }
  }
 
  // Build equivalence sets.
  // Helper function that finds an equivalent block argument index for the
  // given OpOperand. Return std::nullopt if no equivalent block argument could
  // be found.
  auto findEquivalentBlockArgIdx =
      [&](OpOperand &opOperand) -> std::optional<int64_t> {
    Value v = opOperand.get();
    if (!isa<TensorLikeType>(v.getType()))
      return std::nullopt;
    for (BlockArgument bbArg : funcOp.getArguments()) {
      if (isa<TensorLikeType>(bbArg.getType())) {
        if (state.areEquivalentBufferizedValues(v, bbArg)) {
          if (state.getOptions().testAnalysisOnly)
            annotateEquivalentReturnBbArg(opOperand, bbArg);
          return bbArg.getArgNumber();
        }
      }
    }
    return std::nullopt;
  };
 
  int64_t numResults = returnOps.front()->getNumOperands();
  for (int64_t i = 0; i < numResults; ++i) {
    // Find the equivalent block argument index for the i-th operand of the
    // first func.return op.
    std::optional<int64_t> maybeEquiv =
        findEquivalentBlockArgIdx(returnOps.front()->getOpOperand(i));
    if (!maybeEquiv.has_value())
      continue;
    int64_t bbArgIdx = *maybeEquiv;
    bool allEquiv = true;
 
    // Check if all other func.return ops have the same equivalent block
    // argument for the i-th operand. In contrast to aliasing information,
    // which is just "merged", equivalence information must match across all
    // func.return ops.
    for (func::ReturnOp returnOp : ArrayRef(returnOps).drop_front()) {
      std::optional<int64_t> maybeEquiv =
          findEquivalentBlockArgIdx(returnOp->getOpOperand(i));
      if (maybeEquiv != bbArgIdx) {
        allEquiv = false;
        break;
      }
    }
 
    // All func.return ops have the same equivalent block argument for the i-th
    // operand.
    if (allEquiv)
      funcState.equivalentFuncArgs[funcOp][i] = bbArgIdx;
  }
 
  return success();
}
 
static void annotateFuncArgAccess(func::FuncOp funcOp, int64_t idx, bool isRead,
                                  bool isWritten) {
  OpBuilder b(funcOp.getContext());
  Attribute accessType;
  if (isRead && isWritten) {
    accessType = b.getStringAttr("read-write");
  } else if (isRead) {
    accessType = b.getStringAttr("read");
  } else if (isWritten) {
    accessType = b.getStringAttr("write");
  } else {
    accessType = b.getStringAttr("none");
  }
  funcOp.setArgAttr(idx, BufferizationDialect::kBufferAccessAttrName,
                    accessType);
}
 
/// Determine which FuncOp bbArgs are read and which are written. When run on a
/// function with unknown ops, we conservatively assume that such ops bufferize
/// to a read + write.
static LogicalResult
funcOpBbArgReadWriteAnalysis(FuncOp funcOp, OneShotAnalysisState &state,
                             FuncAnalysisState &funcState) {
  for (int64_t idx = 0, e = funcOp.getFunctionType().getNumInputs(); idx < e;
       ++idx) {
    // Skip non-tensor arguments.
    if (!isa<TensorLikeType>(funcOp.getFunctionType().getInput(idx)))
      continue;
    bool isRead;
    bool isWritten;
    if (auto accessAttr = funcOp.getArgAttrOfType<StringAttr>(
            idx, BufferizationDialect::kBufferAccessAttrName)) {
      // Buffer access behavior is specified on the function. Skip the analysis.
      StringRef str = accessAttr.getValue();
      isRead = str == "read" || str == "read-write";
      isWritten = str == "write" || str == "read-write";
    } else if (funcOp.getBody().empty()) {
      // If the function has no body, conservatively assume that all args are
      // read + written.
      isRead = true;
      isWritten = true;
    } else {
      // Analyze the body of the function.
      BlockArgument bbArg = funcOp.getArgument(idx);
      isRead = state.isValueRead(bbArg);
      isWritten = state.isValueWritten(bbArg);
    }
 
    if (state.getOptions().testAnalysisOnly)
      annotateFuncArgAccess(funcOp, idx, isRead, isWritten);
    if (isRead)
      funcState.readBbArgs[funcOp].insert(idx);
    if (isWritten)
      funcState.writtenBbArgs[funcOp].insert(idx);
  }
 
  return success();
}
} // namespace
 
/// Remove bufferization attributes on FuncOp arguments.
static void removeBufferizationAttributes(BlockArgument bbArg) {
  auto funcOp = cast<func::FuncOp>(bbArg.getOwner()->getParentOp());
  funcOp.removeArgAttr(bbArg.getArgNumber(),
                       BufferizationDialect::kBufferLayoutAttrName);
  funcOp.removeArgAttr(bbArg.getArgNumber(),
                       BufferizationDialect::kWritableAttrName);
}
 
/// Return the func::FuncOp called by `callOp`.
static func::FuncOp
getCalledFunction(func::CallOp callOp,
                  mlir::SymbolTableCollection &symbolTable) {
  return dyn_cast_or_null<func::FuncOp>(
      callOp.resolveCallableInTable(&symbolTable));
}
 
/// Return "true" if the given function signature has tensor semantics.
static bool hasTensorSignature(func::FuncOp funcOp) {
  return llvm::any_of(funcOp.getFunctionType().getInputs(),
                      llvm::IsaPred<TensorLikeType>) ||
         llvm::any_of(funcOp.getFunctionType().getResults(),
                      llvm::IsaPred<TensorLikeType>);
}
 
/// Store all functions of the `moduleOp` in `orderedFuncOps`, sorted by
/// callee-caller order (i.e., callees without callers first). Store all
/// remaining functions (i.e., the ones that call each other recursively) in
/// `remainingFuncOps`. Does not traverse nested symbol tables.
///
/// Store the map of FuncOp to all its callers in `callerMap`.
///
/// Return `failure()` if we are unable to retrieve the called FuncOp from
/// any func::CallOp.
static LogicalResult getFuncOpsOrderedByCalls(
    Operation *moduleOp, SmallVectorImpl<func::FuncOp> &orderedFuncOps,
    SmallVectorImpl<func::FuncOp> &remainingFuncOps, FuncCallerMap &callerMap,
    SymbolTableCollection &symbolTables) {
  // For each FuncOp, the set of functions called by it (i.e. the union of
  // symbols of all nested func::CallOp).
  DenseMap<func::FuncOp, SetVector<func::FuncOp>> calledBy;
  // For each FuncOp, the number of func::CallOp it contains.
  llvm::MapVector<func::FuncOp, unsigned> numberCallOpsContainedInFuncOp;
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (func::FuncOp funcOp : block.getOps<func::FuncOp>()) {
        // Collect function calls and populate the caller map.
        numberCallOpsContainedInFuncOp[funcOp] = 0;
        WalkResult res = funcOp.walk([&](func::CallOp callOp) -> WalkResult {
          func::FuncOp calledFunction = getCalledFunction(callOp, symbolTables);
          assert(calledFunction && "could not retrieved called func::FuncOp");
          // If the called function does not have any tensors in its signature,
          // then it is not necessary to bufferize the callee before the caller.
          if (!hasTensorSignature(calledFunction))
            return WalkResult::skip();
 
          callerMap[calledFunction].insert(callOp);
          if (calledBy[calledFunction].insert(funcOp)) {
            numberCallOpsContainedInFuncOp[funcOp]++;
          }
          return WalkResult::advance();
        });
        if (res.wasInterrupted())
          return failure();
      }
    }
  }
 
  // Iteratively remove function operations that do not call any of the
  // functions remaining in the callCounter map and add them to ordered list.
  SmallVector<func::FuncOp> worklist;
 
  for (const auto &entry : numberCallOpsContainedInFuncOp) {
    if (entry.second == 0)
      worklist.push_back(entry.first);
  }
 
  while (!worklist.empty()) {
    func::FuncOp func = worklist.pop_back_val();
    orderedFuncOps.push_back(func);
 
    for (func::FuncOp caller : calledBy[func]) {
      auto &count = numberCallOpsContainedInFuncOp[caller];
 
      if (--count == 0)
        worklist.push_back(caller);
    }
 
    numberCallOpsContainedInFuncOp.erase(func);
  }
 
  // Put all other functions in the list of remaining functions. These are
  // functions that call each other circularly.
  for (auto it : numberCallOpsContainedInFuncOp)
    remainingFuncOps.push_back(it.first);
 
  return success();
}
 
/// Helper function that extracts the source from a memref.cast. If the given
/// value is not a memref.cast result, simply returns the given value.
/// Only unpacks casts where the source is at least as specific as the result
/// (i.e., does not unpack casts from unranked to ranked memref, which would
/// downgrade the type).
static Value unpackCast(Value v) {
  auto castOp = v.getDefiningOp<memref::CastOp>();
  if (!castOp)
    return v;
  // Do not unpack a cast from unranked to ranked memref: folding would
  // downgrade the function return type from ranked to unranked.
  if (isa<UnrankedMemRefType>(castOp.getSource().getType()) &&
      isa<MemRefType>(v.getType()))
    return v;
  return castOp.getSource();
}
 
/// Helper function that returns the return types (skipping casts) of the given
/// func.return ops. This function returns as many types as the return ops have
/// operands. If the i-th operand is not the same for all func.return ops, then
/// the i-th returned type is an "empty" type.
static SmallVector<Type> getReturnTypes(SmallVector<func::ReturnOp> returnOps) {
  assert(!returnOps.empty() && "expected at least one ReturnOp");
  int numOperands = returnOps.front()->getNumOperands();
 
  // Helper function that unpacks memref.cast ops and returns the type.
  auto getSourceType = [&](Value v) { return unpackCast(v).getType(); };
 
  SmallVector<Type> result;
  for (int i = 0; i < numOperands; ++i) {
    // Get the type of the i-th operand of the first func.return ops.
    Type t = getSourceType(returnOps.front()->getOperand(i));
 
    // Check if all other func.return ops have a matching operand type.
    for (int j = 1; j < static_cast<int>(returnOps.size()); ++j)
      if (getSourceType(returnOps[j]->getOperand(i)) != t)
        t = Type();
 
    result.push_back(t);
  }
 
  return result;
}
 
/// Fold return values that are memref casts and update function return types.
///
/// During FuncOp bufferization, the exact type of the returned memrefs (if any)
/// is not known yet. Therefore, the bufferization uses memref types with the
/// most generic layout map as function return types. After bufferizing the
/// entire function body, a more concise memref type can potentially be used for
/// the return type of the function.
static void foldMemRefCasts(func::FuncOp funcOp) {
  // There is nothing to do for bodiless ops.
  if (funcOp.getBody().empty())
    return;
 
  // Compute the common result types of all return ops.
  SmallVector<func::ReturnOp> returnOps = getReturnOps(funcOp);
  SmallVector<Type> resultTypes = getReturnTypes(returnOps);
 
  // Remove direct casts.
  for (func::ReturnOp returnOp : returnOps) {
    for (OpOperand &operand : returnOp->getOpOperands()) {
      // Bail if no common result type was found.
      if (resultTypes[operand.getOperandNumber()]) {
        operand.set(unpackCast(operand.get()));
      }
    }
  }
 
  // Fill in the missing result types that were not the same among all
  // func.return ops.
  for (int i = 0; i < static_cast<int>(resultTypes.size()); ++i) {
    if (resultTypes[i])
      continue;
    resultTypes[i] = funcOp.getFunctionType().getResult(i);
  }
 
  // Update the function type.
  auto newFuncType = FunctionType::get(
      funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
  funcOp.setType(newFuncType);
}
 
LogicalResult
mlir::bufferization::analyzeModuleOp(Operation *moduleOp,
                                     OneShotAnalysisState &state,
                                     BufferizationStatistics *statistics) {
  assert(state.getOptions().bufferizeFunctionBoundaries &&
         "expected that function boundary bufferization is activated");
  FuncAnalysisState &funcState = getOrCreateFuncAnalysisState(state);
 
  // A list of non-circular functions in the order in which they are analyzed
  // and bufferized.
  SmallVector<func::FuncOp> orderedFuncOps;
  // A list of all other functions. I.e., functions that call each other
  // recursively. For these, we analyze the function body but not the function
  // boundary.
  SmallVector<func::FuncOp> remainingFuncOps;
 
  // A mapping of FuncOps to their callers.
  FuncCallerMap callerMap;
 
  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps,
                                      remainingFuncOps, callerMap,
                                      funcState.symbolTables)))
    return failure();
 
  // Analyze functions in order. Starting with functions that are not calling
  // any other functions.
  for (func::FuncOp funcOp : orderedFuncOps) {
    if (!state.getOptions().isOpAllowed(funcOp))
      continue;
 
    // Now analyzing function.
    funcState.startFunctionAnalysis(funcOp);
 
    // Analyze funcOp.
    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;
  }
 
  // Analyze all other functions. All function boundary analyses are skipped.
  for (func::FuncOp funcOp : remainingFuncOps) {
    if (!state.getOptions().isOpAllowed(funcOp))
      continue;
 
    // Analyze funcOp.
    if (failed(analyzeOp(funcOp, state, statistics)))
      return failure();
 
    // TODO: We currently skip all function argument analyses for functions
    // that call each other circularly. These analyses do not support recursive
    // calls yet. The `BufferizableOpInterface` implementations of `func`
    // dialect ops return conservative results in the absence of analysis
    // information.
  }
 
  return success();
}
 
void mlir::bufferization::removeBufferizationAttributesInModule(
    Operation *moduleOp) {
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (func::FuncOp funcOp : block.getOps<func::FuncOp>()) {
        for (BlockArgument bbArg : funcOp.getArguments())
          removeBufferizationAttributes(bbArg);
      }
    }
  }
}
 
LogicalResult mlir::bufferization::bufferizeModuleOp(
    Operation *moduleOp, const OneShotBufferizationOptions &options,
    BufferizationState &state, BufferizationStatistics *statistics) {
  assert(options.bufferizeFunctionBoundaries &&
         "expected that function boundary bufferization is activated");
  IRRewriter rewriter(moduleOp->getContext());
 
  // A list of non-circular functions in the order in which they are analyzed
  // and bufferized.
  SmallVector<func::FuncOp> orderedFuncOps;
  // A list of all other functions. I.e., functions that call each other
  // recursively. For these, we analyze the function body but not the function
  // boundary.
  SmallVector<func::FuncOp> remainingFuncOps;
 
  // A mapping of FuncOps to their callers.
  FuncCallerMap callerMap;
 
  // Try to bufferize functions in calling order. I.e., first bufferize
  // functions that do not call other functions. This allows us to infer
  // accurate buffer types for function return values. Functions that call
  // each other recursively are bufferized in an unspecified order at the end.
  // We may use unnecessarily "complex" (in terms of layout map) buffer types.
  if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps,
                                      remainingFuncOps, callerMap,
                                      state.getSymbolTables())))
    return failure();
  llvm::append_range(orderedFuncOps, remainingFuncOps);
 
  // 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 (llvm::is_contained(options.noAnalysisFuncFilter, funcOp.getSymName())) {
      // This function was not analyzed and RaW conflicts were not resolved.
      // Buffer copies must be inserted before every write.
      OneShotBufferizationOptions updatedOptions = options;
      updatedOptions.copyBeforeWrite = true;
      if (failed(bufferizeOp(funcOp, updatedOptions, state, statistics)))
        return failure();
    } else {
      if (failed(bufferizeOp(funcOp, options, state, statistics)))
        return failure();
    }
 
    // Change buffer return types to more precise layout maps.
    if (options.inferFunctionResultLayout)
      foldMemRefCasts(funcOp);
  }
 
  // Bufferize all other ops.
  for (mlir::Region &region : moduleOp->getRegions()) {
    for (mlir::Block &block : region.getBlocks()) {
      for (mlir::Operation &op :
           llvm::make_early_inc_range(block.getOperations())) {
        // Functions were already bufferized.
        if (isa<func::FuncOp>(&op) || op.hasTrait<OpTrait::SymbolTable>())
          continue;
        if (failed(bufferizeOp(&op, options, state, statistics)))
          return failure();
      }
    }
  }
 
  // Post-pass cleanup of function argument attributes.
  removeBufferizationAttributesInModule(moduleOp);
 
  return success();
}
 
LogicalResult mlir::bufferization::runOneShotModuleBufferize(
    Operation *moduleOp, const OneShotBufferizationOptions &options,
    BufferizationState &state, 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) {
    if (options.noAnalysisFuncFilter.empty()) {
      if (failed(insertTensorCopies(moduleOp, options, state, statistics)))
        return failure();
    } else {
      // FuncOps whose names are specified in options.noAnalysisFuncFilter will
      // not be analyzed. Ops in these FuncOps will not be analyzed as well.
      OpFilter::Entry::FilterFn analysisFilterFn = [=](Operation *op) {
        auto func = dyn_cast<func::FuncOp>(op);
        if (!func)
          func = op->getParentOfType<func::FuncOp>();
        if (func)
          return llvm::is_contained(options.noAnalysisFuncFilter,
                                    func.getSymName());
        return false;
      };
      OneShotBufferizationOptions updatedOptions(options);
      updatedOptions.opFilter.denyOperation(analysisFilterFn);
      if (failed(
              insertTensorCopies(moduleOp, updatedOptions, state, statistics)))
        return failure();
    }
  }
  if (options.testAnalysisOnly)
    return success();
  if (failed(bufferizeModuleOp(moduleOp, options, state, statistics)))
    return failure();
  return success();
}