#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include <optional>
#include <random>
#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/Iterators.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/SubsetOpInterface.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...

static bool	isReachable (Block from, Block to, ArrayRef< Block * > except)

static bool	canUseOpDominanceDueToRegions (OpOperand uRead, OpOperand uWrite, const SetVector< Value > &definitions, 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 bool	canUseOpDominanceDueToBlocks (OpOperand uRead, OpOperand uWrite, const SetVector< Value > &definitions, 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 bool	canUseOpDominance (OpOperand uRead, OpOperand uWrite, const SetVector< Value > &definitions, AnalysisState &state)

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

static bool	hasEquivalentValueInReverseUseDefChain (AnalysisState &state, Value start, Value other)
	Return 'true' if a tensor that is equivalent to `other` can be found in the reverse use-def chain of `start`. More...

static bool	matchesInsertDestination (const AnalysisState &state, Value value, SubsetInsertionOpInterface subsetOp)
	Return "true" if `value` is originating from a subset that is equivalent to the subset that `subsetOp` inserts into. More...

static bool	areNonConflictingSubsets (OpOperand uRead, OpOperand uConflictingWrite, const AnalysisState &state)
	Return "true" if the given "read" and potentially conflicting "write" are not conflicting due to their subset relationship. 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 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 SmallVector< Operation * >	bottomUpFromTerminatorsHeuristic (Operation *op, const OneShotAnalysisState &state)
	"Bottom-up from terminators" heuristic. More...

static LogicalResult	checkPreBufferizationAssumptions (Operation *op, const DominanceInfo &domInfo, OneShotAnalysisState &state)
	Perform various checks on the input IR to see if it contains IR constructs that are unsupported by One-Shot Bufferize. 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)

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

constexpr StringLiteral	kOpResultAliasSetAttrName

constexpr StringLiteral	kBbArgAliasSetAttrName = "__bbarg_alias_set_attr__"

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "one-shot-analysis"

Definition at line 66 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 447 of file OneShotAnalysis.cpp.

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

◆ annotateNonWritableTensor()

static void annotateNonWritableTensor ( Value value )

static

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

Definition at line 927 of file OneShotAnalysis.cpp.

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

◆ annotateOpsWithAliasSets()

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

static

Definition at line 1291 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 1281 of file OneShotAnalysis.cpp.

◆ areNonConflictingSubsets()

static bool areNonConflictingSubsets	(	OpOperand *	uRead,
		OpOperand *	uConflictingWrite,
		const AnalysisState &	state
	)

static

Return "true" if the given "read" and potentially conflicting "write" are not conflicting due to their subset relationship.

The comments in this function are expressed in terms of tensor.extract_slice/tensor.insert_slice pairs, but apply to any subset ops that implement the SubsetInsertionOpInterface.

Definition at line 522 of file OneShotAnalysis.cpp.

◆ bottomUpFromTerminatorsHeuristic()

static SmallVector<Operation *> bottomUpFromTerminatorsHeuristic	(	Operation *	op,
		const OneShotAnalysisState &	state
	)

static

"Bottom-up from terminators" heuristic.

Definition at line 1100 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 1003 of file OneShotAnalysis.cpp.

◆ canUseOpDominance()

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

static

Definition at line 439 of file OneShotAnalysis.cpp.

◆ canUseOpDominanceDueToBlocks()

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

static

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

Returns false if op dominance cannot be used to due block-based loops within a region.

Refer to the canUseOpDominanceDueToRegions documentation for details on how op domiance is used during RaW conflict detection.

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.

Op dominance cannot be used if there is a path from block(READ) to block(WRITE) and a path from block(WRITE) to block(READ). block(DEF) should not appear on that path.

Definition at line 417 of file OneShotAnalysis.cpp.

◆ canUseOpDominanceDueToRegions()

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

static

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

Returns false if op dominance cannot be used to due region-based loops.

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 371 of file OneShotAnalysis.cpp.

◆ checkPreBufferizationAssumptions()

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

static

Perform various checks on the input IR to see if it contains IR constructs that are unsupported by One-Shot Bufferize.

Definition at line 1203 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 1085 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 1035 of file OneShotAnalysis.cpp.

◆ getAliasingInplaceWrites()

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

static

Definition at line 833 of file OneShotAnalysis.cpp.

◆ getAliasingReads()

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

static

Definition at line 844 of file OneShotAnalysis.cpp.

◆ 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 263 of file OneShotAnalysis.cpp.

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

◆ hasEquivalentValueInReverseUseDefChain()

static bool hasEquivalentValueInReverseUseDefChain	(	AnalysisState &	state,
		Value	start,
		Value	other
	)

static

Return 'true' if a tensor that is equivalent to other can be found in the reverse use-def chain of start.

Note: If an OpOperand bufferizes out of place along that use-def chain, the two tensors may not materialize as equivalent buffers (but separate allocations).

Note: This function also requires that the two tensors have equivalent indexing. I.e., the tensor types do not change along the use-def chain, apart from static <-> dynamic dim casts.

Definition at line 486 of file OneShotAnalysis.cpp.

◆ 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 605 of file OneShotAnalysis.cpp.

◆ isaTensor()

static bool isaTensor ( Type t )

static

Definition at line 71 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 252 of file OneShotAnalysis.cpp.

◆ isReachable()

static bool isReachable	(	Block *	from,
		Block *	to,
		ArrayRef< Block * >	except
	)

static

Definition at line 276 of file OneShotAnalysis.cpp.

References mlir::Block::getSuccessors().

◆ matchesInsertDestination()

static bool matchesInsertDestination	(	const AnalysisState &	state,
		Value	value,
		SubsetInsertionOpInterface	subsetOp
	)

static

Return "true" if value is originating from a subset that is equivalent to the subset that subsetOp inserts into.

Definition at line 500 of file OneShotAnalysis.cpp.

◆ setInPlaceOpOperand()

static void setInPlaceOpOperand	(	OpOperand &	opOperand,
		bool	inPlace
	)

static

Mark whether OpOperand will be bufferized inplace.

Definition at line 90 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 909 of file OneShotAnalysis.cpp.

◆ 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 946 of file OneShotAnalysis.cpp.

Variable Documentation

◆ kBbArgAliasSetAttrName

constexpr StringLiteral kBbArgAliasSetAttrName = "__bbarg_alias_set_attr__"

constexpr

Definition at line 87 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 82 of file OneShotAnalysis.cpp.

◆ kOpResultAliasSetAttrName

constexpr StringLiteral kOpResultAliasSetAttrName

constexpr

Initial value:

=

"__opresult_alias_set_attr__"

Definition at line 84 of file OneShotAnalysis.cpp.

Macros

Functions

Variables

Macro Definition Documentation

◆ DEBUG_TYPE

Function Documentation

◆ annotateConflict()

◆ annotateNonWritableTensor()

◆ annotateOpsWithAliasSets()

◆ annotateOpsWithBufferizationMarkers()

◆ areNonConflictingSubsets()

◆ bottomUpFromTerminatorsHeuristic()

◆ bufferizableInPlaceAnalysisImpl()

◆ canUseOpDominance()

◆ canUseOpDominanceDueToBlocks()

◆ canUseOpDominanceDueToRegions()

◆ checkPreBufferizationAssumptions()

◆ equivalenceAnalysis() [1/2]

◆ equivalenceAnalysis() [2/2]

◆ getAliasingInplaceWrites()

◆ getAliasingReads()

◆ happensBefore()

◆ hasEquivalentValueInReverseUseDefChain()

◆ hasReadAfterWriteInterference()

◆ isaTensor()

◆ isInplaceMemoryWrite()

◆ isReachable()

◆ matchesInsertDestination()

◆ setInPlaceOpOperand()

◆ wouldCreateReadAfterWriteInterference()

◆ wouldCreateWriteToNonWritableBuffer()

Variable Documentation

◆ kBbArgAliasSetAttrName

◆ kInPlaceOperandsAttrName

◆ kOpResultAliasSetAttrName