26 namespace bufferization {
27 #define GEN_PASS_DEF_ONESHOTBUFFERIZEPASS
28 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
32 #define DEBUG_TYPE "bufferize"
40 parseHeuristicOption(
const std::string &s) {
45 if (s ==
"bottom-up-from-terminators")
50 llvm_unreachable(
"invalid analysisheuristic option");
53 struct OneShotBufferizePass
54 :
public bufferization::impl::OneShotBufferizePassBase<
55 OneShotBufferizePass> {
58 void runOnOperation()
override {
71 if (mustInferMemorySpace && useEncodingForMemorySpace) {
73 <<
"only one of 'must-infer-memory-space' and "
74 "'use-encoding-for-memory-space' are allowed in "
76 return signalPassFailure();
79 if (mustInferMemorySpace) {
86 if (useEncodingForMemorySpace) {
89 if (
auto rtt = dyn_cast<RankedTensorType>(t))
90 return rtt.getEncoding();
103 LayoutMapOption unknownTypeConversionOption = unknownTypeConversion;
104 if (unknownTypeConversionOption == LayoutMapOption::InferLayoutMap) {
106 "Invalid option: 'infer-layout-map' is not a valid value for "
107 "'unknown-type-conversion'");
108 return signalPassFailure();
113 if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
115 tensorType, memorySpace);
116 assert(unknownTypeConversionOption ==
117 LayoutMapOption::FullyDynamicLayoutMap &&
118 "invalid layout map option");
126 if (this->dialectFilter.hasValue() && !(*this->dialectFilter).empty())
127 return llvm::is_contained(this->dialectFilter,
128 op->getDialect()->getNamespace());
143 "Invalid option: 'copy-before-write' cannot be used with "
144 "'test-analysis-only'");
145 return signalPassFailure();
151 "Invalid option: 'print-conflicts' requires 'test-analysis-only'");
152 return signalPassFailure();
158 "Invalid option: 'dump-alias-sets' requires 'test-analysis-only'");
159 return signalPassFailure();
164 ModuleOp moduleOp = getOperation();
174 "Invalid option: 'no-analysis-func-filter' requires "
175 "'bufferize-function-boundaries'");
176 return signalPassFailure();
191 std::optional<OneShotBufferizationOptions>
options;
208 :
IRRewriter(ctx), erasedOps(erasedOps), toBufferOps(toBufferOps),
209 worklist(worklist), analysisState(
options), statistics(statistics) {
214 void notifyOperationErased(
Operation *op)
override {
215 erasedOps.insert(op);
217 toBufferOps.erase(op);
220 void notifyOperationInserted(
Operation *op, InsertPoint previous)
override {
222 if (previous.isSet())
229 if (
auto sideEffectingOp = dyn_cast<MemoryEffectOpInterface>(op))
230 statistics->numBufferAlloc +=
static_cast<int64_t
>(
235 if (isa<ToBufferOp>(op)) {
236 toBufferOps.insert(op);
241 if (isa<ToTensorOp>(op))
249 auto const &
options = analysisState.getOptions();
254 worklist.push_back(op);
288 op->
walk([&](ToBufferOp toBufferOp) { toBufferOps.insert(toBufferOp); });
300 worklist.push_back(op);
307 BufferizationRewriter rewriter(op->
getContext(), erasedOps, toBufferOps,
308 worklist,
options, statistics);
309 for (
unsigned i = 0; i < worklist.size(); ++i) {
312 if (erasedOps.contains(nextOp))
315 auto bufferizableOp =
options.dynCastBufferizableOp(nextOp);
322 if (!bufferizableOp.supportsUnstructuredControlFlow())
324 if (r.getBlocks().size() > 1)
326 "op or BufferizableOpInterface implementation does not support "
327 "unstructured control flow, but at least one region has multiple "
331 LLVM_DEBUG(llvm::dbgs()
332 <<
"//===-------------------------------------------===//\n"
333 <<
"IR after bufferizing: " << nextOp->
getName() <<
"\n");
334 rewriter.setInsertionPoint(nextOp);
336 bufferizableOp.bufferize(rewriter,
options, bufferizationState))) {
337 LLVM_DEBUG(llvm::dbgs()
338 <<
"failed to bufferize\n"
339 <<
"//===-------------------------------------------===//\n");
340 return nextOp->
emitError(
"failed to bufferize op");
342 LLVM_DEBUG(llvm::dbgs()
344 <<
"\n//===-------------------------------------------===//\n");
348 if (erasedOps.contains(op))
353 rewriter.setInsertionPoint(op);
355 rewriter, cast<ToBufferOp>(op),
options);
360 if (toTensorOp->getUses().empty()) {
361 rewriter.eraseOp(toTensorOp);
374 if (erasedOps.contains(op))
387 if (isa<ToTensorOp, ToBufferOp>(op))
389 return op->
emitError(
"op was not bufferized");
407 auto tensorType = dyn_cast<TensorType>(bbArg.getType());
409 newTypes.push_back(bbArg.getType());
413 FailureOr<BufferLikeType> bufferType =
415 if (failed(bufferType))
417 newTypes.push_back(*bufferType);
421 for (
auto [bbArg, type] : llvm::zip(block->
getArguments(), newTypes)) {
422 if (bbArg.getType() == type)
428 bbArgUses.push_back(&use);
430 Type tensorType = bbArg.getType();
436 if (!bbArgUses.empty()) {
437 Value toTensorOp = rewriter.
create<bufferization::ToTensorOp>(
438 bbArg.getLoc(), tensorType, bbArg);
440 use->set(toTensorOp);
446 auto branchOp = dyn_cast<BranchOpInterface>(op);
448 return op->
emitOpError(
"cannot bufferize ops with block references that "
449 "do not implement BranchOpInterface");
452 assert(it != op->
getSuccessors().end() &&
"could find successor");
453 int64_t successorIdx = std::distance(op->
getSuccessors().begin(), it);
457 for (
auto [operand, type] :
459 if (operand.getType() == type) {
461 newOperands.push_back(operand);
464 FailureOr<BufferLikeType> operandBufferType =
466 if (failed(operandBufferType))
469 Value bufferizedOperand = rewriter.
create<bufferization::ToBufferOp>(
470 operand.getLoc(), *operandBufferType, operand);
473 if (type != *operandBufferType)
474 bufferizedOperand = rewriter.
create<memref::CastOp>(
475 operand.getLoc(), type, bufferizedOperand);
476 newOperands.push_back(bufferizedOperand);
static MLIRContext * getContext(OpFoldResult val)
static llvm::ManagedStatic< PassManagerOptions > options
Base class for generic analysis states.
Attributes are known-constant values of operations.
This class represents an argument of a Block.
Block represents an ordered list of Operations.
BlockArgListType getArguments()
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
user_range getUsers() const
Returns a range of all users.
This class coordinates rewriting a piece of IR outside of a pattern rewrite, providing a way to keep ...
MLIRContext is the top-level object for a collection of MLIR operations.
void assign(ValueRange values)
Assign this range to the given values.
RAII guard to reset the insertion point of the builder when destroyed.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
void setInsertionPointAfterValue(Value val)
Sets the insertion point to the node after the specified value.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
This class represents an operand of an operation.
Operation is the basic unit of execution within MLIR.
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
MLIRContext * getContext()
Return the context this operation is associated with.
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
MutableArrayRef< Region > getRegions()
Returns the regions held by this operation.
OperationName getName()
The name of an operation is the key identifier for it.
SuccessorRange getSuccessors()
use_range getUses()
Returns a range of all uses, which is useful for iterating over all uses.
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "'dim' op " which is convenient for verifiers.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
This class models how operands are forwarded to block arguments in control flow.
MutableOperandRange getMutableForwardedOperands() const
Get the range of operands that are simply forwarded to the successor.
OperandRange getForwardedOperands() const
Get the range of operands that are simply forwarded to the successor.
Tensor types represent multi-dimensional arrays, and have two variants: RankedTensorType and Unranked...
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...
static WalkResult advance()
AnalysisState provides a variety of helper functions for dealing with tensor values.
BufferizationState provides information about the state of the IR during the bufferization process.
void allowOperation()
Allow the given ops.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options, BufferizationState &bufferizationState, BufferizationStatistics *statistics=nullptr)
Bufferize op and its nested ops that implement BufferizableOpInterface.
LogicalResult bufferizeBlockSignature(Block *block, RewriterBase &rewriter, const BufferizationOptions &options, BufferizationState &state)
Bufferize the signature of block and its callers (i.e., ops that have the given block as a successor)...
LogicalResult insertTensorCopies(Operation *op, const OneShotBufferizationOptions &options, const BufferizationState &bufferizationState, BufferizationStatistics *statistics=nullptr)
Resolve RaW and other conflicts by inserting bufferization.alloc_tensor ops.
LogicalResult runOneShotBufferize(Operation *op, const OneShotBufferizationOptions &options, BufferizationState &state, BufferizationStatistics *statistics=nullptr)
Run One-Shot Bufferize on the given op: Analysis + Bufferization.
LogicalResult foldToBufferToTensorPair(RewriterBase &rewriter, ToBufferOp toBuffer, const BufferizationOptions &options)
Try to fold to_buffer(to_tensor(x)).
FailureOr< BufferLikeType > getBufferType(Value value, const BufferizationOptions &options, const BufferizationState &state)
Return the buffer type for a given Value (tensor) after bufferization without bufferizing any IR.
llvm::LogicalResult runOneShotModuleBufferize(ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions &options, BufferizationState &state, BufferizationStatistics *statistics=nullptr)
Run One-Shot Module Bufferization on the given module.
BaseMemRefType getMemRefTypeWithFullyDynamicLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with fully dynamic layout.
bool hasTensorSemantics(Operation *op)
Return "true" if the given op has tensor semantics and should be bufferized.
Include the generated interface declarations.
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
bool isMemoryEffectFree(Operation *op)
Returns true if the given operation is free of memory effects.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
The following effect indicates that the operation allocates from some resource.
Options for BufferizableOpInterface-based bufferization.
bool copyBeforeWrite
If set to true, the analysis is skipped.
void setFunctionBoundaryTypeConversion(LayoutMapOption layoutMapOption)
This function controls buffer types on function signatures.
bool allowUnknownOps
Specifies whether not bufferizable ops are allowed in the input.
unsigned int bufferAlignment
Buffer alignment for new memory allocations.
bool printConflicts
If set to true, the IR is annotated with details about RaW conflicts.
bool testAnalysisOnly
If set to true, does not modify the IR apart from adding attributes (for checking the results of the ...
OpFilter opFilter
A filter that specifies which ops should be bufferized and which ops should be ignored.
bool checkParallelRegions
UnknownTypeConverterFn unknownTypeConverterFn
Type converter from tensors to memrefs.
bool bufferizeFunctionBoundaries
Specifies whether function boundaries (ops in the func dialect) should be bufferized or not.
DefaultMemorySpaceFn defaultMemorySpaceFn
Bufferization statistics for debugging.
int64_t numTensorOutOfPlace
Options for analysis-enabled bufferization.
unsigned analysisFuzzerSeed
Seed for the analysis fuzzer.
bool dumpAliasSets
Specifies whether the tensor IR should be annotated with alias sets.
bool allowReturnAllocsFromLoops
Specifies whether returning newly allocated memrefs from loops should be allowed.
AnalysisHeuristic analysisHeuristic
The heuristic controls the order in which ops are traversed during the analysis.
@ BottomUpFromTerminators
llvm::ArrayRef< std::string > noAnalysisFuncFilter
Specify the functions that should not be analyzed.
std::function< bool(Operation *)> FilterFn
If the filter function evaluates to true, the filter matches.