28 namespace bufferization {
29 #define GEN_PASS_DEF_FINALIZINGBUFFERIZE
30 #define GEN_PASS_DEF_BUFFERIZATIONBUFFERIZE
31 #define GEN_PASS_DEF_ONESHOTBUFFERIZE
32 #include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
36 #define DEBUG_TYPE "bufferize"
47 assert(inputs.size() == 1);
48 assert(isa<BaseMemRefType>(inputs[0].getType()));
49 return builder.
create<bufferization::ToTensorOp>(loc, type, inputs[0]);
68 assert(inputs.size() == 1 &&
"expected exactly one input");
70 if (
auto inputType = dyn_cast<MemRefType>(inputs[0].getType())) {
72 assert(inputType != type &&
"expected different types");
74 auto rankedDestType = dyn_cast<MemRefType>(type);
86 if (isa<TensorType>(inputs[0].getType())) {
88 return builder.
create<bufferization::ToMemrefOp>(loc, type, inputs[0]);
91 llvm_unreachable(
"only tensor/memref input types supported");
97 target.
addLegalOp<bufferization::ToTensorOp, bufferization::ToMemrefOp>();
103 class BufferizeToTensorOp
108 matchAndRewrite(bufferization::ToTensorOp op, OpAdaptor adaptor,
110 rewriter.
replaceOp(op, adaptor.getMemref());
119 class BufferizeToMemrefOp
124 matchAndRewrite(bufferization::ToMemrefOp op, OpAdaptor adaptor,
126 rewriter.
replaceOp(op, adaptor.getTensor());
134 patterns.
add<BufferizeToTensorOp, BufferizeToMemrefOp>(typeConverter,
139 struct FinalizingBufferizePass
140 :
public bufferization::impl::FinalizingBufferizeBase<
141 FinalizingBufferizePass> {
142 using FinalizingBufferizeBase<
143 FinalizingBufferizePass>::FinalizingBufferizeBase;
145 void runOnOperation()
override {
146 auto func = getOperation();
163 target.markUnknownOpDynamicallyLegal(
171 static LayoutMapOption parseLayoutMapOption(
const std::string &s) {
172 if (s ==
"fully-dynamic-layout-map")
173 return LayoutMapOption::FullyDynamicLayoutMap;
174 if (s ==
"identity-layout-map")
175 return LayoutMapOption::IdentityLayoutMap;
176 if (s ==
"infer-layout-map")
177 return LayoutMapOption::InferLayoutMap;
178 llvm_unreachable(
"invalid layout map option");
182 parseHeuristicOption(
const std::string &s) {
183 if (s ==
"bottom-up")
187 if (s ==
"bottom-up-from-terminators")
192 llvm_unreachable(
"invalid analysisheuristic option");
195 struct OneShotBufferizePass
196 :
public bufferization::impl::OneShotBufferizeBase<OneShotBufferizePass> {
197 OneShotBufferizePass() =
default;
204 .
insert<bufferization::BufferizationDialect, memref::MemRefDialect>();
207 void runOnOperation()
override {
219 parseLayoutMapOption(functionBoundaryTypeConversion));
220 if (mustInferMemorySpace) {
222 [](
TensorType t) -> std::optional<Attribute> {
232 LayoutMapOption unknownTypeConversionOption =
233 parseLayoutMapOption(unknownTypeConversion);
234 if (unknownTypeConversionOption == LayoutMapOption::InferLayoutMap) {
236 "Invalid option: 'infer-layout-map' is not a valid value for "
237 "'unknown-type-conversion'");
238 return signalPassFailure();
242 auto tensorType = cast<TensorType>(value.
getType());
243 if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
245 tensorType, memorySpace);
246 assert(unknownTypeConversionOption ==
247 LayoutMapOption::FullyDynamicLayoutMap &&
248 "invalid layout map option");
256 if (this->dialectFilter.hasValue())
257 return llvm::is_contained(this->dialectFilter,
273 "Invalid option: 'copy-before-write' cannot be used with "
274 "'test-analysis-only'");
275 return signalPassFailure();
281 "Invalid option: 'print-conflicts' requires 'test-analysis-only'");
282 return signalPassFailure();
288 "Invalid option: 'dump-alias-sets' requires 'test-analysis-only'");
289 return signalPassFailure();
293 ModuleOp moduleOp = getOperation();
302 "Invalid option: 'no-analysis-func-filter' requires "
303 "'bufferize-function-boundaries'");
304 return signalPassFailure();
319 std::optional<OneShotBufferizationOptions>
options;
324 struct BufferizationBufferizePass
325 :
public bufferization::impl::BufferizationBufferizeBase<
326 BufferizationBufferizePass> {
327 void runOnOperation()
override {
329 options.opFilter.allowDialect<BufferizationDialect>();
337 .
insert<bufferization::BufferizationDialect, memref::MemRefDialect>();
343 return std::make_unique<BufferizationBufferizePass>();
347 return std::make_unique<OneShotBufferizePass>();
352 return std::make_unique<OneShotBufferizePass>(
options);
355 std::unique_ptr<OperationPass<func::FuncOp>>
357 return std::make_unique<FinalizingBufferizePass>();
373 :
IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps),
374 worklist(worklist), analysisState(
options), statistics(statistics) {
379 void notifyOperationErased(
Operation *op)
override {
380 erasedOps.insert(op);
382 toMemrefOps.erase(op);
385 void notifyOperationInserted(
Operation *op, InsertPoint previous)
override {
387 if (previous.isSet())
394 if (
auto sideEffectingOp = dyn_cast<MemoryEffectOpInterface>(op))
400 if (isa<ToMemrefOp>(op)) {
401 toMemrefOps.insert(op);
406 if (isa<ToTensorOp>(op))
414 auto const &
options = analysisState.getOptions();
419 worklist.push_back(op);
452 op->
walk([&](ToMemrefOp toMemrefOp) { toMemrefOps.insert(toMemrefOp); });
464 worklist.push_back(op);
471 BufferizationRewriter rewriter(op->
getContext(), erasedOps, toMemrefOps,
472 worklist,
options, statistics);
473 for (
unsigned i = 0; i < worklist.size(); ++i) {
476 if (erasedOps.contains(nextOp))
479 auto bufferizableOp =
options.dynCastBufferizableOp(nextOp);
486 if (!bufferizableOp.supportsUnstructuredControlFlow())
488 if (r.getBlocks().size() > 1)
490 "op or BufferizableOpInterface implementation does not support "
491 "unstructured control flow, but at least one region has multiple "
495 LLVM_DEBUG(llvm::dbgs()
496 <<
"//===-------------------------------------------===//\n"
497 <<
"IR after bufferizing: " << nextOp->
getName() <<
"\n");
498 rewriter.setInsertionPoint(nextOp);
500 LLVM_DEBUG(llvm::dbgs()
501 <<
"failed to bufferize\n"
502 <<
"//===-------------------------------------------===//\n");
503 return nextOp->
emitError(
"failed to bufferize op");
505 LLVM_DEBUG(llvm::dbgs()
507 <<
"\n//===-------------------------------------------===//\n");
511 if (erasedOps.contains(op))
516 rewriter.setInsertionPoint(op);
518 rewriter, cast<ToMemrefOp>(op),
options);
523 if (toTensorOp->getUses().empty()) {
524 rewriter.eraseOp(toTensorOp);
537 if (erasedOps.contains(op))
550 if (isa<ToTensorOp, ToMemrefOp>(op))
552 return op->
emitError(
"op was not bufferized");
569 auto tensorType = dyn_cast<TensorType>(bbArg.getType());
571 newTypes.push_back(bbArg.getType());
579 newTypes.push_back(*memrefType);
583 for (
auto [bbArg, type] : llvm::zip(block->
getArguments(), newTypes)) {
584 if (bbArg.getType() == type)
590 bbArgUses.push_back(&use);
597 if (!bbArgUses.empty()) {
599 rewriter.
create<bufferization::ToTensorOp>(bbArg.getLoc(), bbArg);
601 use->set(toTensorOp);
607 auto branchOp = dyn_cast<BranchOpInterface>(op);
609 return op->
emitOpError(
"cannot bufferize ops with block references that "
610 "do not implement BranchOpInterface");
613 assert(it != op->
getSuccessors().end() &&
"could find successor");
614 int64_t successorIdx = std::distance(op->
getSuccessors().begin(), it);
618 for (
auto [operand, type] :
620 if (operand.getType() == type) {
622 newOperands.push_back(operand);
627 if (
failed(operandBufferType))
630 Value bufferizedOperand = rewriter.
create<bufferization::ToMemrefOp>(
631 operand.getLoc(), *operandBufferType, operand);
634 if (type != *operandBufferType)
635 bufferizedOperand = rewriter.
create<memref::CastOp>(
636 operand.getLoc(), type, bufferizedOperand);
637 newOperands.push_back(bufferizedOperand);
647 options.allowUnknownOps =
true;
648 options.copyBeforeWrite =
true;
649 options.enforceAliasingInvariants =
false;
653 cast<TensorType>(value.
getType()), memorySpace);
655 options.opFilter.allowDialect<BufferizationDialect>();
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 provides a shared interface for ranked and unranked memref types.
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.
This class implements a pattern rewriter for use with ConversionPatterns.
void replaceOp(Operation *op, ValueRange newValues) override
PatternRewriter hook for replacing an operation.
This class describes a specific conversion target.
void addLegalOp(OperationName op)
Register the given operations as legal.
The DialectRegistry maps a dialect namespace to a constructor for the matching dialect.
StringRef getNamespace() const
This class provides support for representing a failure result, or a valid value of type T.
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 ...
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
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.
This class helps build Operations.
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.
OpConversionPattern is a wrapper around ConversionPattern that allows for matching and rewriting agai...
OpConversionPattern(MLIRContext *context, PatternBenefit benefit=1)
This class represents an operand of an operation.
Operation is the basic unit of execution within MLIR.
Dialect * getDialect()
Return the dialect this operation is associated with, or nullptr if the associated dialect is not loa...
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.
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
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...
void addConversion(FnT &&callback)
Register a conversion function.
bool isLegal(Type type) const
Return true if the given type is legal for this type converter, i.e.
void addArgumentMaterialization(FnT &&callback)
Register a materialization function, which must be convertible to the following form: std::optional<V...
void addSourceMaterialization(FnT &&callback)
This method registers a materialization that will be called when converting a legal type to an illega...
void addTargetMaterialization(FnT &&callback)
This method registers a materialization that will be called when converting type from an illegal,...
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
This class provides an abstraction over the different types of ranges over Values.
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.
static WalkResult advance()
AnalysisState provides a variety of helper functions for dealing with tensor values.
A helper type converter class that automatically populates the relevant materializations and type con...
BufferizeTypeConverter()
Registers conversions into BufferizeTypeConverter.
void allowOperation()
Allow the given ops.
LogicalResult runOneShotBufferize(Operation *op, const OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Run One-Shot Bufferize on the given op: Analysis + Bufferization.
BaseMemRefType getMemRefTypeWithStaticIdentityLayout(TensorType tensorType, Attribute memorySpace=nullptr)
Return a MemRef type with a static identity layout (i.e., no layout map).
FailureOr< Value > castOrReallocMemRefValue(OpBuilder &b, Value value, MemRefType type, const BufferizationOptions &options)
Try to cast the given ranked MemRef-typed value to the given ranked MemRef type.
void populateEliminateBufferizeMaterializationsPatterns(BufferizeTypeConverter &typeConverter, RewritePatternSet &patterns)
Populate patterns to eliminate bufferize materializations.
LogicalResult runOneShotModuleBufferize(ModuleOp moduleOp, const bufferization::OneShotBufferizationOptions &options, BufferizationStatistics *statistics=nullptr)
Run One-Shot Module Bufferization on the given module.
void populateBufferizeMaterializationLegality(ConversionTarget &target)
Marks ops used by bufferization for type conversion materializations as "legal" in the given Conversi...
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.
std::unique_ptr< OperationPass< func::FuncOp > > createFinalizingBufferizePass()
Creates a pass that finalizes a partial bufferization by removing remaining bufferization....
LogicalResult bufferizeBlockSignature(Block *block, RewriterBase &rewriter, const BufferizationOptions &options)
Bufferize the signature of block and its callers (i.e., ops that have the given block as a successor)...
std::unique_ptr< Pass > createOneShotBufferizePass()
Create a pass that bufferizes all ops that implement BufferizableOpInterface with One-Shot Bufferize.
LogicalResult foldToMemrefToTensorPair(RewriterBase &rewriter, ToMemrefOp toMemref, const BufferizationOptions &options)
Try to fold to_memref(to_tensor(x)).
FailureOr< BaseMemRefType > getBufferType(Value value, const BufferizationOptions &options)
Return the buffer type for a given Value (tensor) after bufferization without bufferizing any IR.
BufferizationOptions getPartialBufferizationOptions()
Return BufferizationOptions such that the bufferizeOp behaves like the old (deprecated) partial,...
std::unique_ptr< Pass > createBufferizationBufferizePass()
Create a pass that bufferizes ops from the bufferization dialect.
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.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
LogicalResult applyFullConversion(ArrayRef< Operation * > ops, const ConversionTarget &target, const FrozenRewritePatternSet &patterns, ConversionConfig config=ConversionConfig())
Apply a complete conversion on the given operations, and all nested operations.
bool isMemoryEffectFree(Operation *op)
Returns true if the given operation is free of memory effects.
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.
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.
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.
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.