| //===- 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 ®ion : 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 ®ion : 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 ®ion : 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(); | |
| } |