91 func::ReturnOp returnOp;
92 for (
Block &b : funcOp.getBody()) {
93 if (
auto candidateOp = dyn_cast<func::ReturnOp>(b.getTerminator())) {
96 returnOp = candidateOp;
105 static void annotateEquivalentReturnBbArg(
OpOperand &returnVal,
107 const char *kEquivalentArgsAttr =
"__equivalent_func_args__";
111 if (op->
hasAttr(kEquivalentArgsAttr)) {
112 auto attr = cast<ArrayAttr>(op->
getAttr(kEquivalentArgsAttr));
113 equivBbArgs = llvm::to_vector<4>(llvm::map_range(attr, [](
Attribute a) {
114 return cast<IntegerAttr>(a).getValue().getSExtValue();
122 op->
setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
130 if (funcOp.getBody().empty()) {
133 FunctionType type = funcOp.getFunctionType();
135 if (!isa<TensorType>(inputIt.value()))
138 if (!isa<TensorType>(resultIt.value()))
140 int64_t returnIdx = resultIt.index();
141 int64_t bbArgIdx = inputIt.index();
150 assert(returnOp &&
"expected func with single return op");
152 for (
OpOperand &returnVal : returnOp->getOpOperands())
153 if (isa<RankedTensorType>(returnVal.
get().
getType()))
155 if (isa<RankedTensorType>(bbArg.
getType())) {
158 if (state.areEquivalentBufferizedValues(returnVal.
get(), bbArg)) {
160 if (state.getOptions().testAnalysisOnly)
161 annotateEquivalentReturnBbArg(returnVal, bbArg);
163 if (state.areAliasingBufferizedValues(returnVal.
get(), bbArg))
170 static void annotateFuncArgAccess(func::FuncOp funcOp, int64_t idx,
bool isRead,
174 if (isRead && isWritten) {
175 accessType = b.getStringAttr(
"read-write");
177 accessType = b.getStringAttr(
"read");
178 }
else if (isWritten) {
179 accessType = b.getStringAttr(
"write");
181 accessType = b.getStringAttr(
"none");
183 funcOp.setArgAttr(idx, BufferizationDialect::kBufferAccessAttrName,
193 for (int64_t idx = 0, e = funcOp.getFunctionType().getNumInputs(); idx < e;
196 if (!isa<TensorType>(funcOp.getFunctionType().getInput(idx)))
200 if (
auto accessAttr = funcOp.getArgAttrOfType<StringAttr>(
201 idx, BufferizationDialect::kBufferAccessAttrName)) {
203 StringRef str = accessAttr.getValue();
204 isRead = str ==
"read" || str ==
"read-write";
205 isWritten = str ==
"write" || str ==
"read-write";
206 }
else if (funcOp.getBody().empty()) {
214 isRead = state.isValueRead(bbArg);
215 isWritten = state.isValueWritten(bbArg);
218 if (state.getOptions().testAnalysisOnly)
219 annotateFuncArgAccess(funcOp, idx, isRead, isWritten);
234 BufferizationDialect::kBufferLayoutAttrName);
236 BufferizationDialect::kWritableAttrName);
242 llvm::dyn_cast_if_present<SymbolRefAttr>(callOp.getCallableForCallee());
245 return dyn_cast_or_null<func::FuncOp>(
256 funcOp->walk([&](func::CallOp callOp) {
258 assert(calledFunction &&
"could not retrieved called func::FuncOp");
265 int64_t returnIdx = it.first;
266 int64_t bbargIdx = it.second;
267 if (!state.isInPlace(callOp->getOpOperand(bbargIdx)))
269 Value returnVal = callOp.getResult(returnIdx);
270 Value argVal = callOp->getOperand(bbargIdx);
271 state.unionEquivalenceClasses(returnVal, argVal);
281 return llvm::any_of(funcOp.getFunctionType().getInputs(),
isaTensor) ||
282 llvm::any_of(funcOp.getFunctionType().getResults(),
isaTensor);
300 if (!funcOp.getBody().empty()) {
301 func::ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
303 return funcOp->emitError()
304 <<
"cannot bufferize a FuncOp with tensors and "
305 "without a unique ReturnOp";
309 numberCallOpsContainedInFuncOp[funcOp] = 0;
310 return funcOp.walk([&](func::CallOp callOp) ->
WalkResult {
311 func::FuncOp calledFunction = getCalledFunction(callOp);
312 assert(calledFunction &&
"could not retrieved called func::FuncOp");
315 if (!hasTensorSignature(calledFunction))
316 return WalkResult::skip();
318 callerMap[calledFunction].insert(callOp);
319 if (calledBy[calledFunction].insert(funcOp).second) {
320 numberCallOpsContainedInFuncOp[funcOp]++;
325 if (res.wasInterrupted())
329 while (!numberCallOpsContainedInFuncOp.empty()) {
330 auto it = llvm::find_if(numberCallOpsContainedInFuncOp,
331 [](
auto entry) {
return entry.getSecond() == 0; });
332 if (it == numberCallOpsContainedInFuncOp.end())
333 return moduleOp.emitOpError(
334 "expected callgraph to be free of circular dependencies.");
335 orderedFuncOps.push_back(it->getFirst());
336 for (
auto callee : calledBy[it->getFirst()])
337 numberCallOpsContainedInFuncOp[callee]--;
338 numberCallOpsContainedInFuncOp.erase(it);
351 if (funcOp.getBody().empty())
357 for (
OpOperand &operand : returnOp->getOpOperands()) {
358 if (
auto castOp = operand.get().getDefiningOp<memref::CastOp>()) {
359 operand.set(castOp.getSource());
360 resultTypes.push_back(castOp.getSource().getType());
362 resultTypes.push_back(operand.get().getType());
367 funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
368 funcOp.setType(newFuncType);
375 assert(state.getOptions().bufferizeFunctionBoundaries &&
376 "expected that function boundary bufferization is activated");
389 for (func::FuncOp funcOp : orderedFuncOps) {
390 if (!state.getOptions().isOpAllowed(funcOp))
404 if (
failed(aliasingFuncOpBBArgsAnalysis(funcOp, state, funcState)) ||
405 failed(funcOpBbArgReadWriteAnalysis(funcOp, state, funcState)))
417 moduleOp.walk([&](func::FuncOp op) {
426 assert(
options.bufferizeFunctionBoundaries &&
427 "expected that function boundary bufferization is activated");
440 for (func::FuncOp funcOp : orderedFuncOps) {
444 if (llvm::is_contained(
options.noAnalysisFuncFilter, funcOp.getSymName())) {
457 if (
options.inferFunctionResultLayout)
470 assert(
options.bufferizeFunctionBoundaries &&
471 "expected that function boundary bufferization is activated");
473 "invalid combination of bufferization flags");
474 if (!
options.copyBeforeWrite) {
475 if (
options.noAnalysisFuncFilter.empty()) {
482 auto func = dyn_cast<func::FuncOp>(op);
486 return llvm::is_contained(
options.noAnalysisFuncFilter,
static bool hasTensorSignature(func::FuncOp funcOp)
Return "true" if the given function signature has tensor semantics.
static LogicalResult getFuncOpsOrderedByCalls(ModuleOp moduleOp, SmallVectorImpl< func::FuncOp > &orderedFuncOps, FuncCallerMap &callerMap)
Store all functions of the moduleOp in orderedFuncOps, sorted by callee-caller order (i....
static FuncAnalysisState & getOrCreateFuncAnalysisState(OneShotAnalysisState &state)
Get or create FuncAnalysisState.
static void removeBufferizationAttributes(BlockArgument bbArg)
Remove bufferization attributes on FuncOp arguments.
static void foldMemRefCasts(func::FuncOp funcOp)
Fold return values that are memref casts and update function return types.
static void equivalenceAnalysis(func::FuncOp funcOp, OneShotAnalysisState &state, FuncAnalysisState &funcState)
Gather equivalence info of CallOps.
static func::ReturnOp getAssumedUniqueReturnOp(func::FuncOp funcOp)
Return the unique ReturnOp that terminates funcOp.
static llvm::ManagedStatic< PassManagerOptions > options
Attributes are known-constant values of operations.
This class represents an argument of a Block.
Block * getOwner() const
Returns the block that owns this argument.
unsigned getArgNumber() const
Returns the number of this argument.
Block represents an ordered list of Operations.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
IRValueT get() const
Return the current value being used by this operand.
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
This class helps build Operations.
This class represents an operand of an operation.
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Operation is the basic unit of execution within MLIR.
Attribute getAttr(StringAttr name)
Return the specified attribute if present, null otherwise.
bool hasAttr(StringAttr name)
Return true if the operation has an attribute with the provided name, false otherwise.
MLIRContext * getContext()
Return the context this operation is associated with.
unsigned getNumOperands()
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
static Operation * lookupNearestSymbolFrom(Operation *from, StringAttr symbol)
Returns the operation registered with the given symbol name within the closest parent operation of,...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Type getType() const
Return the type of this value.
A utility result that is used to signal how to proceed with an ongoing walk:
static WalkResult advance()
State for analysis-enabled bufferization.
void denyOperation()
Deny the given ops.
Operation * getOwner() const
Return the owner of this operand.
static func::ReturnOp getAssumedUniqueReturnOp(FuncOp funcOp)
Return the unique ReturnOp that terminates funcOp.
static FuncOp getCalledFunction(CallOpInterface callOp)
Return the FuncOp called by callOp.
LogicalResult analyzeOp(Operation *op, OneShotAnalysisState &state, BufferizationStatistics *statistics=nullptr)
Analyze op and its nested ops.
LogicalResult runOneShotModuleBufferize(ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Run One-Shot Module Bufferization on the given module.
LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Bufferize op and its nested ops that implement BufferizableOpInterface.
LogicalResult insertTensorCopies(Operation *op, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state, BufferizationStatistics *statistics=nullptr)
Analyze moduleOp and its nested ops.
void removeBufferizationAttributesInModule(ModuleOp moduleOp)
Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
LogicalResult bufferizeModuleOp(ModuleOp moduleOp, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Bufferize op and its nested ops that implement BufferizableOpInterface.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Include the generated interface declarations.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
This class represents an efficient way to signal success or failure.
bool copyBeforeWrite
If set to true, the analysis is skipped.
OpFilter opFilter
A filter that specifies which ops should be bufferized and which ops should be ignored.
Bufferization statistics for debugging.
Options for analysis-enabled bufferization.
std::function< bool(Operation *)> FilterFn
If the filter function evaluates to true, the filter matches.
Extra analysis state that is required for bufferization of function boundaries.
DenseMap< FuncOp, IndexMapping > equivalentFuncArgs
A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg indices.
DenseMap< FuncOp, IndexToIndexListMapping > aliasingReturnVals
A mapping of FuncOp BBArg indices to aliasing ReturnOp OpOperand indices.
DenseMap< FuncOp, BbArgIndexSet > readBbArgs
A set of all read BlockArguments of FuncOps.
DenseMap< FuncOp, BbArgIndexSet > writtenBbArgs
A set of all written-to BlockArguments of FuncOps.
DenseMap< FuncOp, FuncOpAnalysisState > analyzedFuncOps
Keep track of which FuncOps are fully analyzed or currently being analyzed.
void startFunctionAnalysis(FuncOp funcOp)
This function is called right before analyzing the given FuncOp.