#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include <random>
#include <optional>
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/Bufferize.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/AsmState.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"

Macros
#define	DEBUG_TYPE "one-shot-analysis"

Functions
static bool	isaTensor (Type t)

static void	setInPlaceOpOperand (OpOperand &opOperand, bool inPlace)
	Mark whether OpOperand will be bufferized inplace. More...

static bool	isInplaceMemoryWrite (OpOperand &opOperand, const OneShotAnalysisState &state)
	Return true if opOperand has been decided to bufferize in-place. More...

static bool	happensBefore (Operation a, Operation b, const DominanceInfo &domInfo)
	Return true if `a` happens before `b`, i.e., `a` or one of its ancestors properly dominates `b` and `b` is not inside `a`. More...

bool	canUseOpDominance (OpOperand uRead, OpOperand uWrite, const SetVector< Value > &definitions, const AnalysisState &state)
	Return `true` if op dominance can be used to rule out a read-after-write conflicts based on the ordering of ops. More...

static void	annotateConflict (OpOperand uRead, OpOperand uConflictingWrite, Value definition)
	Annotate IR with details about the detected RaW conflict. More...

static bool	hasReadAfterWriteInterference (const DenseSet< OpOperand * > &usesRead, const DenseSet< OpOperand * > &usesWrite, const DominanceInfo &domInfo, OneShotAnalysisState &state)
	Given sets of uses and writes, return true if there is a RaW conflict under the assumption that all given reads/writes alias the same buffer and that all given writes bufferize inplace. More...

static void	getAliasingInplaceWrites (DenseSet< OpOperand * > &res, Value root, const OneShotAnalysisState &state)

static void	getAliasingReads (DenseSet< OpOperand * > &res, Value root, const OneShotAnalysisState &state)

static bool	wouldCreateReadAfterWriteInterference (OpOperand &operand, const DominanceInfo &domInfo, OneShotAnalysisState &state, bool checkConsistencyOnly=false)
	Return true if bufferizing `operand` inplace would create a conflict. More...

static void	annotateNonWritableTensor (Value value)
	Annotate IR with details about the detected non-writability conflict. More...

static bool	wouldCreateWriteToNonWritableBuffer (OpOperand &operand, OneShotAnalysisState &state, bool checkConsistencyOnly=false)
	Return true if bufferizing `operand` inplace would create a write to a non-writable buffer. More...

static LogicalResult	bufferizableInPlaceAnalysisImpl (OpOperand &operand, OneShotAnalysisState &state, const DominanceInfo &domInfo)
	Determine if `operand` can be bufferized in-place. More...

static bool	hasTensorSemantics (Operation *op)
	Return true if the given op has a tensor result or a tensor operand. More...

static void	equivalenceAnalysis (SmallVector< Operation * > &ops, OneShotAnalysisState &state)
	Analyze equivalence of tied OpResult/OpOperand pairs of the given ops. More...

static void	equivalenceAnalysis (Operation *op, OneShotAnalysisState &state)
	Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained in `op`. More...

static LogicalResult	checkAliasInfoConsistency (Operation *op, const DominanceInfo &domInfo, OneShotAnalysisState &state)
	Assert that the current bufferization decisions are consistent. More...

static void	annotateOpsWithBufferizationMarkers (Operation *op, const OneShotAnalysisState &state)
	Annotate the IR with the result of the analysis. For testing/debugging only. More...

static void	annotateOpsWithAliasSets (Operation *op, const OneShotAnalysisState &state)

static LogicalResult	assertNoAllocsReturned (Operation *op, const OneShotAnalysisState &state)
	Assert that every allocation can be deallocated in the same block. More...

Variables
constexpr StringLiteral	kInPlaceOperandsAttrName = "__inplace_operands_attr__"
	Attribute marker to specify op operands that bufferize in-place. More...

constexpr StringLiteral	kAliasSetAttrName = "__alias_set_attr__"

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "one-shot-analysis"

Definition at line 64 of file OneShotAnalysis.cpp.

Function Documentation

◆ annotateConflict()

static void annotateConflict	(	OpOperand *	uRead,
		OpOperand *	uConflictingWrite,
		Value	definition
	)

static

Annotate IR with details about the detected RaW conflict.

Definition at line 417 of file OneShotAnalysis.cpp.

References mlir::Operation::getContext(), mlir::OpOperand::getOperandNumber(), mlir::detail::IROperandBase::getOwner(), mlir::Builder::getUnitAttr(), and mlir::Operation::setAttr().

Referenced by hasReadAfterWriteInterference().

◆ annotateNonWritableTensor()

static void annotateNonWritableTensor ( Value value )

static

Annotate IR with details about the detected non-writability conflict.

Definition at line 716 of file OneShotAnalysis.cpp.

References mlir::Value::getContext(), and mlir::Builder::getUnitAttr().

Referenced by wouldCreateWriteToNonWritableBuffer().

◆ annotateOpsWithAliasSets()

static void annotateOpsWithAliasSets	(	Operation *	op,
		const OneShotAnalysisState &	state
	)

static

Definition at line 991 of file OneShotAnalysis.cpp.

◆ annotateOpsWithBufferizationMarkers()

static void annotateOpsWithBufferizationMarkers	(	Operation *	op,
		const OneShotAnalysisState &	state
	)

static

Annotate the IR with the result of the analysis. For testing/debugging only.

Definition at line 981 of file OneShotAnalysis.cpp.

◆ assertNoAllocsReturned()

static LogicalResult assertNoAllocsReturned	(	Operation *	op,
		const OneShotAnalysisState &	state
	)

static

Assert that every allocation can be deallocated in the same block.

I.e., every value that is returned or yielded from a block is:

guaranteed to be aliasing a bbArg of that block or a parent block, or
guaranteed to be aliasing an OpResult of a op in a parent block.

In that case, buffer deallocation is simple: Every allocated buffer can be deallocated in the same block. Otherwise, the buffer deallocation pass must be run.

Note: The current implementation checks for equivalent values instead of aliasing values, which is stricter than needed. We can currently not check for aliasing values because the analysis is a maybe-alias analysis and we need a must-alias analysis here.

Example:

%0 = "some_op" : tensor<?xf32>
%1 = scf.if %c -> (tensor<?xf32>) {
  scf.yield %0 : tensor<?xf32>
} else {
  %t = linalg.alloc_tensor : tensor<?xf32>
  scf.yield %t : tensor<?xf32>
}

In the above example, the second scf.yield op is problematic because the yielded value t is defined in the same block as the scf.yield op and and bufferizes to a new allocation.

Definition at line 1044 of file OneShotAnalysis.cpp.

◆ bufferizableInPlaceAnalysisImpl()

static LogicalResult bufferizableInPlaceAnalysisImpl	(	OpOperand &	operand,
		OneShotAnalysisState &	state,
		const DominanceInfo &	domInfo
	)

static

Determine if operand can be bufferized in-place.

Definition at line 789 of file OneShotAnalysis.cpp.

References mlir::bufferization::OneShotAnalysisState::bufferizeInPlace(), mlir::bufferization::OneShotAnalysisState::bufferizeOutOfPlace(), mlir::OpOperand::getOperandNumber(), mlir::detail::IROperandBase::getOwner(), mlir::success(), wouldCreateReadAfterWriteInterference(), and wouldCreateWriteToNonWritableBuffer().

◆ canUseOpDominance()

bool canUseOpDominance	(	OpOperand *	uRead,
		OpOperand *	uWrite,
		const SetVector< Value > &	definitions,
		const AnalysisState &	state
	)

Return true if op dominance can be used to rule out a read-after-write conflicts based on the ordering of ops.

Generalized op dominance can often be used to rule out potential conflicts due to "read happens before write". E.g., the following IR is not a RaW conflict because the read happens before the write.

Example 1: %0 = ... : tensor<?xf32> // DEF "reading_op"(%0) : tensor<?xf32> // READ %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE

This is no longer true inside loops (or repetitive regions). In such cases, there may not be a meaningful happensBefore relationship because ops could be executed multiple times. E.g.:

Example 2: %0 = ... : tensor<?xf32> // DEF scf.for ... { "reading_op"(%0) : tensor<?xf32> // READ %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> // WRITE ... }

In the above example, reading_op happens before writing_op according to op dominance. However, both ops may happen multiple times; in particular, the second execution of reading_op happens after the first execution of writing_op. This is problematic because the tensor %0 they operate on (i.e., the "definition") is defined outside of the loop.

On a high-level, there is a potential RaW in a program if there exists a possible program execution such that there is a sequence of DEF, followed by WRITE, followed by READ. Each additional DEF resets the sequence.

E.g.: No conflict: DEF, WRITE, DEF, READ Potential conflict: DEF, READ, WRITE, READ, WRITE

Example 1 has no conflict: DEF, READ, WRITE Example 2 has a potential conflict: DEF, (READ, WRITE)* Example 3: scf.for ... { %0 = ... : tensor<?xf32> "reading_op"(%0) : tensor<?xf32> %1 = "writing_op"(%0) : tensor<?xf32> -> tensor<?xf32> ... } This has no conflict: (DEF, READ, WRITE)*

Example 4: %0 = ... : tensor<?xf32> scf.for ... { scf.for ... { "reading_op"(%0) } %1 = "writing_op"(%0) } This has a potential conflict: DEF, ((READ)*, WRITE)*

Example 5: %0 = ... : tensor<?xf32> scf.for ... { %1 = "writing_op"(%0) } scf.for ... { "reading_op"(%0) } This has a potential conflict: DEF, WRITE*, READ*

The following rules are used to rule out RaW conflicts via ordering of ops:

If the closest enclosing repetitive region of DEF is a proper ancestor of a repetitive region that enclosing both READ and WRITE, we cannot rule out RaW conflict due to the ordering of ops.
Otherwise: There are no loops that interfere with our analysis; for analysis purposes, we can assume that there are no loops/repetitive regions. I.e., we can rule out a RaW conflict if READ happensBefore WRITE or WRITE happensBefore DEF. (Checked in hasReadAfterWriteInterference.)

Definition at line 386 of file OneShotAnalysis.cpp.

References mlir::getEnclosingRepetitiveRegion(), mlir::bufferization::getNextEnclosingRepetitiveRegion(), mlir::detail::IROperandBase::getOwner(), mlir::Region::getParentOp(), mlir::Operation::isAncestor(), and options.

Referenced by hasReadAfterWriteInterference().

◆ checkAliasInfoConsistency()

static LogicalResult checkAliasInfoConsistency	(	Operation *	op,
		const DominanceInfo &	domInfo,
		OneShotAnalysisState &	state
	)

static

Assert that the current bufferization decisions are consistent.

Definition at line 931 of file OneShotAnalysis.cpp.

◆ equivalenceAnalysis() [1/2]

static void equivalenceAnalysis	(	Operation *	op,
		OneShotAnalysisState &	state
	)

static

Analyze equivalence of tied OpResult/OpOperand pairs of all ops contained in op.

Definition at line 878 of file OneShotAnalysis.cpp.

◆ equivalenceAnalysis() [2/2]

static void equivalenceAnalysis	(	SmallVector< Operation * > &	ops,
		OneShotAnalysisState &	state
	)

static

Analyze equivalence of tied OpResult/OpOperand pairs of the given ops.

Definition at line 828 of file OneShotAnalysis.cpp.

◆ getAliasingInplaceWrites()

static void getAliasingInplaceWrites	(	DenseSet< OpOperand * > &	res,
		Value	root,
		const OneShotAnalysisState &	state
	)

static

Definition at line 622 of file OneShotAnalysis.cpp.

References mlir::bufferization::OneShotAnalysisState::applyOnAliases(), mlir::Value::getUses(), and isInplaceMemoryWrite().

Referenced by wouldCreateReadAfterWriteInterference(), and wouldCreateWriteToNonWritableBuffer().

◆ getAliasingReads()

static void getAliasingReads	(	DenseSet< OpOperand * > &	res,
		Value	root,
		const OneShotAnalysisState &	state
	)

static

Definition at line 633 of file OneShotAnalysis.cpp.

References mlir::bufferization::OneShotAnalysisState::applyOnAliases(), mlir::bufferization::AnalysisState::bufferizesToMemoryWrite(), and mlir::Value::getUses().

Referenced by wouldCreateReadAfterWriteInterference().

◆ happensBefore()

static bool happensBefore	(	Operation *	a,
		Operation *	b,
		const DominanceInfo &	domInfo
	)

static

Return true if a happens before b, i.e., a or one of its ancestors properly dominates b and b is not inside a.

Definition at line 299 of file OneShotAnalysis.cpp.

References mlir::Operation::getParentOp(), mlir::Operation::isProperAncestor(), and mlir::DominanceInfo::properlyDominates().

Referenced by hasReadAfterWriteInterference().

◆ hasReadAfterWriteInterference()

static bool hasReadAfterWriteInterference	(	const DenseSet< OpOperand * > &	usesRead,
		const DenseSet< OpOperand * > &	usesWrite,
		const DominanceInfo &	domInfo,
		OneShotAnalysisState &	state
	)

static

Given sets of uses and writes, return true if there is a RaW conflict under the assumption that all given reads/writes alias the same buffer and that all given writes bufferize inplace.

A conflict is: According to SSA use-def chains, a read R is supposed to read the result of a definition W1. But because of bufferization decisions, R actually reads another definition W2.

Definition at line 456 of file OneShotAnalysis.cpp.

References annotateConflict(), canUseOpDominance(), mlir::Block::findAncestorOpInBlock(), mlir::bufferization::OneShotAnalysisState::findDefinitionsCached(), mlir::bufferization::AliasList< T >::getAliases(), mlir::bufferization::AnalysisState::getAliasingOpResults(), mlir::bufferization::AliasList< T >::getNumAliases(), mlir::bufferization::OneShotAnalysisState::getOptions(), happensBefore(), mlir::insideMutuallyExclusiveRegions(), and options.

Referenced by wouldCreateReadAfterWriteInterference().

◆ hasTensorSemantics()

static bool hasTensorSemantics ( Operation * op )

static

Return true if the given op has a tensor result or a tensor operand.

Definition at line 821 of file OneShotAnalysis.cpp.

◆ isaTensor()

static bool isaTensor ( Type t )

static

Definition at line 69 of file OneShotAnalysis.cpp.

◆ isInplaceMemoryWrite()

static bool isInplaceMemoryWrite	(	OpOperand &	opOperand,
		const OneShotAnalysisState &	state
	)

static

Return true if opOperand has been decided to bufferize in-place.

Definition at line 288 of file OneShotAnalysis.cpp.

References mlir::bufferization::AnalysisState::bufferizesToMemoryWrite(), and mlir::bufferization::OneShotAnalysisState::isInPlace().

Referenced by getAliasingInplaceWrites().

◆ setInPlaceOpOperand()

static void setInPlaceOpOperand	(	OpOperand &	opOperand,
		bool	inPlace
	)

static

Mark whether OpOperand will be bufferized inplace.

Definition at line 85 of file OneShotAnalysis.cpp.

◆ wouldCreateReadAfterWriteInterference()

static bool wouldCreateReadAfterWriteInterference	(	OpOperand &	operand,
		const DominanceInfo &	domInfo,
		OneShotAnalysisState &	state,
		bool	checkConsistencyOnly = `false`
	)

static

Return true if bufferizing operand inplace would create a conflict.

A read R and a write W of the same alias set is a conflict if inplace bufferization of W changes the value read by R to a value different from the one that would be expected by tracing back R's origin through SSA use-def chains. A conflict can only be introduced by a new alias and/or an inplace bufferization decision.

Example: %0 = tensor.extract_slice t[...][...][1, 1] {inplace?} %1 = vector.transfer_write v1, t {inplace} : vector<5xf32>, tensor<?xf32> e = tensor.extract_slice %1 %2 = vector.transfer_write v2, %0 {inplace} : vector<6xf32>, tensor<?xf32> %3 = vector.transfer_read e, cst : tensor<?xf32>, vector<7xf32>

In the above example, the two TransferWriteOps have already been decided to bufferize inplace. Bufferizing the ExtractSliceOp inplace would create a conflict because:

According to SSA use-def chains, we expect to read the result of %1.
However, adding an alias {%0, t} would mean that the second TransferWriteOp overwrites the result of the first one. Therefore, the TransferReadOp would no longer be reading the result of %1.

If checkConsistencyOnly is true, this function checks if there is a read-after-write conflict without bufferizing operand inplace. This would indicate a problem with the current inplace bufferization decisions.

Note: If checkConsistencyOnly, this function may be called with a null OpResult. In that case, only the consistency of bufferization decisions involving aliases of the given OpOperand are checked.

Definition at line 698 of file OneShotAnalysis.cpp.

References mlir::bufferization::AnalysisState::bufferizesToMemoryWrite(), mlir::IROperand< DerivedT, IRValueT >::get(), getAliasingInplaceWrites(), mlir::bufferization::AnalysisState::getAliasingOpResults(), getAliasingReads(), and hasReadAfterWriteInterference().

Referenced by bufferizableInPlaceAnalysisImpl().

◆ wouldCreateWriteToNonWritableBuffer()

static bool wouldCreateWriteToNonWritableBuffer	(	OpOperand &	operand,
		OneShotAnalysisState &	state,
		bool	checkConsistencyOnly = `false`
	)

static

Return true if bufferizing operand inplace would create a write to a non-writable buffer.

Definition at line 735 of file OneShotAnalysis.cpp.

References annotateNonWritableTensor(), mlir::bufferization::OneShotAnalysisState::applyOnAliases(), mlir::bufferization::AnalysisState::bufferizesToMemoryWrite(), mlir::IROperand< DerivedT, IRValueT >::get(), getAliasingInplaceWrites(), mlir::bufferization::AnalysisState::getAliasingOpResults(), mlir::bufferization::OneShotAnalysisState::getOptions(), mlir::bufferization::OneShotAnalysisState::isWritable(), and mlir::bufferization::BufferizationOptions::printConflicts.

Referenced by bufferizableInPlaceAnalysisImpl().

Variable Documentation

◆ kAliasSetAttrName

constexpr StringLiteral kAliasSetAttrName = "__alias_set_attr__"

constexpr

Definition at line 82 of file OneShotAnalysis.cpp.

◆ kInPlaceOperandsAttrName

constexpr StringLiteral kInPlaceOperandsAttrName = "__inplace_operands_attr__"

constexpr

Attribute marker to specify op operands that bufferize in-place.

Definition at line 80 of file OneShotAnalysis.cpp.

Macros

Functions

Variables

Macro Definition Documentation

◆ DEBUG_TYPE

Function Documentation

◆ annotateConflict()

◆ annotateNonWritableTensor()

◆ annotateOpsWithAliasSets()

◆ annotateOpsWithBufferizationMarkers()

◆ assertNoAllocsReturned()

◆ bufferizableInPlaceAnalysisImpl()

◆ canUseOpDominance()

◆ checkAliasInfoConsistency()

◆ equivalenceAnalysis() [1/2]

◆ equivalenceAnalysis() [2/2]

◆ getAliasingInplaceWrites()

◆ getAliasingReads()

◆ happensBefore()

◆ hasReadAfterWriteInterference()

◆ hasTensorSemantics()

◆ isaTensor()

◆ isInplaceMemoryWrite()

◆ setInPlaceOpOperand()

◆ wouldCreateReadAfterWriteInterference()

◆ wouldCreateWriteToNonWritableBuffer()

Variable Documentation

◆ kAliasSetAttrName

◆ kInPlaceOperandsAttrName