The code generation environment class aggregates a number of data structures that are needed during the code generation phase of sparsification. More...

#include "Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h"

Public Member Functions
	CodegenEnv (linalg::GenericOp linop, SparsificationOptions opts, unsigned numTensors, unsigned numLoops, unsigned maxRank)
	Constructs a code generation environment which can be passed around during sparsification for bookkeeping together with some consistency asserts.
LogicalResult	initTensorExp ()
ExprId	getExprId () const
linalg::GenericOp	op () const
const SparsificationOptions &	options () const
bool	generatingSparseIterator () const
Merger &	merger ()
LoopEmitter &	emitter ()
void	startEmit (SparseEmitStrategy emitStrategy)
std::optional< Operation * >	genLoopBoundary (function_ref< std::optional< Operation * >(MutableArrayRef< Value > parameters)> callback)
	Generates loop boundary statements (entering/exiting loops).
constexpr TensorId	makeTensorId (unsigned t) const
constexpr LoopId	makeLoopId (unsigned i) const
constexpr TensorLoopId	makeTensorLoopId (unsigned t, unsigned i) const
const TensorExp &	exp (ExprId e) const
const LatPoint &	lat (LatPointId l) const
ArrayRef< LatPointId >	set (LatSetId s) const
LevelType	lt (TensorId t, LoopId i) const
LevelType	lt (TensorLoopId b) const
unsigned	getLoopNum () const
TensorLevel	makeTensorLevel (TensorId t, Level l) const
TensorLevel	makeTensorLevel (std::pair< TensorId, Level > tlPair) const
std::pair< TensorId, Level >	unpackTensorLevel (TensorLevel tl) const
template<class ContainerTy>
auto	unpackTensorLevelRange (ContainerTy &&c) const
unsigned	getCurrentDepth () const
bool	isAdmissibleTensorExp (ExprId e)
	Whether the tensor expression is admissible for codegen.
Value	getLoopVar (LoopId i) const
	Returns the induction-variable for the given loop.
bool	hasSparseOutput () const
bool	isSparseOutput (OpOperand *o) const
Value	getInsertionChain () const
void	updateInsertionChain (Value chain)
bool	atExpandLevel (OpOperand *o, unsigned rank, LoopId n) const
void	startExpand (Value values, Value filled, Value added, Value count)
bool	isExpand () const
void	updateExpandCount (Value count)
Value	getExpandValues () const
Value	getExpandFilled () const
Value	getExpandAdded () const
Value	getExpandCount () const
void	endExpand ()
void	startReduc (ExprId exp, Value val)
bool	isReduc () const
void	updateReduc (Value val)
Value	getReduc () const
Value	endReduc ()
void	startValidLexInsert (Value val)
bool	isValidLexInsert () const
void	updateValidLexInsert (Value val)
Value	getValidLexInsert () const
void	endValidLexInsert ()
void	startCustomReduc (ExprId exp)
bool	isCustomReduc () const
Value	getCustomRedId () const
void	endCustomReduc ()

Detailed Description

The code generation environment class aggregates a number of data structures that are needed during the code generation phase of sparsification.

This environment simplifies passing around such data during sparsification (rather than passing around all the individual compoments where needed). Furthermore, it provides convience methods that keep implementation details transparent to sparsification while asserting on internal consistency.

Definition at line 35 of file CodegenEnv.h.

Constructor & Destructor Documentation

◆ CodegenEnv()

CodegenEnv::CodegenEnv	(	linalg::GenericOp	linop,
		SparsificationOptions	opts,
		unsigned	numTensors,
		unsigned	numLoops,
		unsigned	maxRank )

Constructs a code generation environment which can be passed around during sparsification for bookkeeping together with some consistency asserts.

Definition at line 43 of file CodegenEnv.cpp.

References nullptr.

Member Function Documentation

◆ atExpandLevel()

bool CodegenEnv::atExpandLevel	(	OpOperand *	o,
		unsigned	rank,
		LoopId	n ) const

Definition at line 199 of file CodegenEnv.cpp.

Referenced by genExpand().

◆ emitter()

LoopEmitter & mlir::sparse_tensor::CodegenEnv::emitter ( )

inline

Definition at line 57 of file CodegenEnv.h.

Referenced by endLoop(), endLoopSeq(), genBuffers(), genCoIteration(), genCoIterationCase(), genConstantDenseAddressFromLevel(), genExpand(), genIf(), genInsertionStore(), genResult(), genSubscript(), startLoop(), and startLoopSeq().

◆ endCustomReduc()

void CodegenEnv::endCustomReduc ( )

Definition at line 275 of file CodegenEnv.cpp.

References isCustomReduc(), and mlir::sparse_tensor::detail::kInvalidId.

Referenced by genExp(), and genInvariants().

◆ endExpand()

void CodegenEnv::endExpand ( )

Definition at line 218 of file CodegenEnv.cpp.

Referenced by genExpand().

◆ endReduc()

Value CodegenEnv::endReduc ( )

Definition at line 241 of file CodegenEnv.cpp.

References isReduc(), and mlir::sparse_tensor::detail::kInvalidId.

Referenced by genInvariants().

◆ endValidLexInsert()

void CodegenEnv::endValidLexInsert ( )

Definition at line 260 of file CodegenEnv.cpp.

References isReduc(), and isValidLexInsert().

Referenced by genInvariants().

◆ exp()

const TensorExp & mlir::sparse_tensor::CodegenEnv::exp ( ExprId e ) const

inline

Definition at line 81 of file CodegenEnv.h.

Referenced by genExp(), genInvariants(), genInvariantValue(), genTensorLoad(), genTensorStore(), getCustomRedId(), isAdmissibleTensorExp(), startCustomReduc(), and startReduc().

◆ generatingSparseIterator()

bool mlir::sparse_tensor::CodegenEnv::generatingSparseIterator ( ) const

inline

Definition at line 52 of file CodegenEnv.h.

References mlir::kSparseIterator.

Referenced by genStmt().

◆ genLoopBoundary()

std::optional< Operation * > CodegenEnv::genLoopBoundary ( function_ref< std::optional< Operation * >(MutableArrayRef< Value > parameters)> callback )

Generates loop boundary statements (entering/exiting loops).

The function passes and updates the passed-in parameters.

Definition at line 102 of file CodegenEnv.cpp.

References isExpand(), isReduc(), isValidLexInsert(), updateExpandCount(), updateInsertionChain(), updateReduc(), and updateValidLexInsert().

Referenced by endLoop(), genCoIteration(), and genStmt().

◆ getCurrentDepth()

unsigned mlir::sparse_tensor::CodegenEnv::getCurrentDepth ( ) const

inline

Definition at line 115 of file CodegenEnv.h.

Referenced by genInsertionStore(), genStmt(), and getAllTidLvlsInLatPoints().

◆ getCustomRedId()

Value CodegenEnv::getCustomRedId ( ) const

Definition at line 270 of file CodegenEnv.cpp.

References exp(), isCustomReduc(), and op().

Referenced by genInsertionLoadReduce(), and genInvariants().

◆ getExpandAdded()

Value mlir::sparse_tensor::CodegenEnv::getExpandAdded ( ) const

inline

Definition at line 144 of file CodegenEnv.h.

Referenced by genExpand(), and genInsertionStore().

◆ getExpandCount()

Value mlir::sparse_tensor::CodegenEnv::getExpandCount ( ) const

inline

Definition at line 145 of file CodegenEnv.h.

Referenced by endIf(), finalizeWhileOp(), genExpand(), genInsertionStore(), and genStmt().

◆ getExpandFilled()

Value mlir::sparse_tensor::CodegenEnv::getExpandFilled ( ) const

inline

Definition at line 143 of file CodegenEnv.h.

Referenced by genExpand(), genInsertionLoadReduce(), and genInsertionStore().

◆ getExpandValues()

Value mlir::sparse_tensor::CodegenEnv::getExpandValues ( ) const

inline

Definition at line 142 of file CodegenEnv.h.

Referenced by genExpand(), genInsertionLoad(), genInsertionLoadReduce(), and genInsertionStore().

◆ getExprId()

ExprId mlir::sparse_tensor::CodegenEnv::getExprId ( ) const

inline

Definition at line 48 of file CodegenEnv.h.

◆ getInsertionChain()

Value mlir::sparse_tensor::CodegenEnv::getInsertionChain ( ) const

inline

Definition at line 135 of file CodegenEnv.h.

Referenced by endIf(), finalizeWhileOp(), genExpand(), genIf(), genInsertionStore(), genResult(), genStmt(), and genTensorStore().

◆ getLoopNum()

unsigned mlir::sparse_tensor::CodegenEnv::getLoopNum ( ) const

inline

Definition at line 89 of file CodegenEnv.h.

Referenced by genStmt(), isAdmissibleTensorExp(), and startEmit().

◆ getLoopVar()

Value CodegenEnv::getLoopVar ( LoopId i ) const

Returns the induction-variable for the given loop.

Definition at line 186 of file CodegenEnv.cpp.

Referenced by endLoopSeq(), genExp(), genIf(), genIndex(), relinkBranch(), and startLoopSeq().

◆ getReduc()

Value mlir::sparse_tensor::CodegenEnv::getReduc ( ) const

inline

Definition at line 155 of file CodegenEnv.h.

Referenced by endIf(), finalizeWhileOp(), genIf(), and genStmt().

◆ getValidLexInsert()

Value mlir::sparse_tensor::CodegenEnv::getValidLexInsert ( ) const

inline

Definition at line 161 of file CodegenEnv.h.

Referenced by finalizeWhileOp(), genIf(), genInsertionStore(), and genStmt().

◆ hasSparseOutput()

bool mlir::sparse_tensor::CodegenEnv::hasSparseOutput ( ) const

inline

Definition at line 132 of file CodegenEnv.h.

Referenced by genInvariants(), and isParallelFor().

◆ initTensorExp()

LogicalResult CodegenEnv::initTensorExp ( )

Definition at line 51 of file CodegenEnv.cpp.

References isAdmissibleTensorExp(), op(), and success().

◆ isAdmissibleTensorExp()

bool CodegenEnv::isAdmissibleTensorExp ( ExprId e )

Whether the tensor expression is admissible for codegen.

It also sets the sparseOut if the output tensor is sparse.

Definition at line 135 of file CodegenEnv.cpp.

References exp(), getLoopNum(), mlir::sparse_tensor::getSparseTensorType(), mlir::sparse_tensor::SparseTensorType::isAllDense(), isMaterializing(), mlir::linalg::isReductionIterator(), lhs, and makeTensorId().

Referenced by initTensorExp().

◆ isCustomReduc()

bool mlir::sparse_tensor::CodegenEnv::isCustomReduc ( ) const

inline

Definition at line 165 of file CodegenEnv.h.

References mlir::sparse_tensor::detail::kInvalidId.

Referenced by endCustomReduc(), genInvariants(), genTensorLoad(), getCustomRedId(), and startCustomReduc().

◆ isExpand()

bool mlir::sparse_tensor::CodegenEnv::isExpand ( ) const

inline

Definition at line 140 of file CodegenEnv.h.

Referenced by endIf(), finalizeWhileOp(), genIf(), genInsertionLoad(), genInsertionLoadReduce(), genInsertionStore(), genLoopBoundary(), and isParallelFor().

◆ isReduc()

bool mlir::sparse_tensor::CodegenEnv::isReduc ( ) const

inline

Definition at line 153 of file CodegenEnv.h.

References mlir::sparse_tensor::detail::kInvalidId.

Referenced by endIf(), endLoop(), endReduc(), endValidLexInsert(), finalizeWhileOp(), genExpand(), genIf(), genInvariants(), genLoopBoundary(), genTensorStore(), startReduc(), startValidLexInsert(), updateReduc(), and updateValidLexInsert().

◆ isSparseOutput()

bool mlir::sparse_tensor::CodegenEnv::isSparseOutput ( OpOperand * o ) const

inline

Definition at line 133 of file CodegenEnv.h.

Referenced by genTensorLoad(), and genTensorStore().

◆ isValidLexInsert()

bool mlir::sparse_tensor::CodegenEnv::isValidLexInsert ( ) const

inline

Definition at line 159 of file CodegenEnv.h.

Referenced by endIf(), endLoop(), endValidLexInsert(), finalizeWhileOp(), genIf(), genInsertionStore(), genLoopBoundary(), and startValidLexInsert().

◆ lat()

const LatPoint & mlir::sparse_tensor::CodegenEnv::lat ( LatPointId l ) const

inline

Definition at line 82 of file CodegenEnv.h.

Referenced by genStmt(), getAllTidLvlsInLatPoints(), and startLoopSeq().

◆ lt() [1/2]

LevelType mlir::sparse_tensor::CodegenEnv::lt	(	TensorId	t,
		LoopId	i ) const

inline

Definition at line 84 of file CodegenEnv.h.

Referenced by getAllTidLvlsInLatPoints(), and shouldTryParallize().

◆ lt() [2/2]

LevelType mlir::sparse_tensor::CodegenEnv::lt ( TensorLoopId b ) const

inline

Definition at line 87 of file CodegenEnv.h.

References b.

◆ makeLoopId()

LoopId mlir::sparse_tensor::CodegenEnv::makeLoopId ( unsigned i ) const

inlineconstexpr

Definition at line 75 of file CodegenEnv.h.

Referenced by genIndex(), and relinkBranch().

◆ makeTensorId()

TensorId mlir::sparse_tensor::CodegenEnv::makeTensorId ( unsigned t ) const

inlineconstexpr

Definition at line 72 of file CodegenEnv.h.

Referenced by findSparseAnnotations(), genConstantDenseAddressFromLevel(), genSubscript(), isAdmissibleTensorExp(), relinkBranch(), and startEmit().

◆ makeTensorLevel() [1/2]

TensorLevel mlir::sparse_tensor::CodegenEnv::makeTensorLevel ( std::pair< TensorId, Level > tlPair ) const

inline

Definition at line 104 of file CodegenEnv.h.

References makeTensorLevel().

◆ makeTensorLevel() [2/2]

TensorLevel mlir::sparse_tensor::CodegenEnv::makeTensorLevel	(	TensorId	t,
		Level	l ) const

inline

Definition at line 95 of file CodegenEnv.h.

Referenced by genConstantDenseAddressFromLevel(), getAllTidLvlsInLatPoints(), and makeTensorLevel().

◆ makeTensorLoopId()

TensorLoopId mlir::sparse_tensor::CodegenEnv::makeTensorLoopId	(	unsigned	t,
		unsigned	i ) const

inlineconstexpr

Definition at line 78 of file CodegenEnv.h.

◆ merger()

Merger & mlir::sparse_tensor::CodegenEnv::merger ( )

inline

Definition at line 56 of file CodegenEnv.h.

Referenced by findSparseAnnotations(), genCoIterationCase(), genExp(), genIf(), genInvariants(), genResult(), genStmt(), genTensorStore(), getAllTidLvlsInLatPoints(), startEmit(), startLoop(), and startLoopSeq().

◆ op()

linalg::GenericOp mlir::sparse_tensor::CodegenEnv::op ( ) const

inline

◆ options()

const SparsificationOptions & mlir::sparse_tensor::CodegenEnv::options ( ) const

inline

Definition at line 51 of file CodegenEnv.h.

Referenced by genSubscript(), genTensorLoad(), getAllTidLvlsInLatPoints(), and isParallelFor().

◆ set()

ArrayRef< LatPointId > mlir::sparse_tensor::CodegenEnv::set ( LatSetId s ) const

inline

Definition at line 83 of file CodegenEnv.h.

Referenced by genStmt(), and startLoopSeq().

◆ startCustomReduc()

void CodegenEnv::startCustomReduc ( ExprId exp )

Definition at line 265 of file CodegenEnv.cpp.

References exp(), isCustomReduc(), and mlir::sparse_tensor::detail::kInvalidId.

Referenced by genExp(), and genInvariants().

◆ startEmit()

void CodegenEnv::startEmit ( SparseEmitStrategy emitStrategy )

Definition at line 61 of file CodegenEnv.cpp.

References mlir::sparse_tensor::Merger::getDependentLoops(), getLoopNum(), mlir::sparse_tensor::getSparseTensorEncoding(), makeTensorId(), merger(), and sortDependentLoops().

◆ startExpand()

void CodegenEnv::startExpand	(	Value	values,
		Value	filled,
		Value	added,
		Value	count )

Definition at line 204 of file CodegenEnv.cpp.

Referenced by genExpand().

◆ startReduc()

void CodegenEnv::startReduc	(	ExprId	exp,
		Value	val )

Definition at line 227 of file CodegenEnv.cpp.

References exp(), isReduc(), and mlir::sparse_tensor::detail::kInvalidId.

Referenced by genInvariants().

◆ startValidLexInsert()

void CodegenEnv::startValidLexInsert ( Value val )

Definition at line 250 of file CodegenEnv.cpp.

References isReduc(), and isValidLexInsert().

Referenced by genInvariants().

◆ unpackTensorLevel()

std::pair< TensorId, Level > mlir::sparse_tensor::CodegenEnv::unpackTensorLevel ( TensorLevel tl ) const

inline

Definition at line 107 of file CodegenEnv.h.

Referenced by shouldTryParallize().

◆ unpackTensorLevelRange()

template<class ContainerTy>

auto mlir::sparse_tensor::CodegenEnv::unpackTensorLevelRange ( ContainerTy && c ) const

inline

Definition at line 111 of file CodegenEnv.h.

Referenced by startLoop().

◆ updateExpandCount()

void CodegenEnv::updateExpandCount ( Value count )

Definition at line 213 of file CodegenEnv.cpp.

Referenced by endIf(), finalizeWhileOp(), genInsertionStore(), and genLoopBoundary().

◆ updateInsertionChain()

void CodegenEnv::updateInsertionChain ( Value chain )

Definition at line 194 of file CodegenEnv.cpp.

Referenced by endIf(), finalizeWhileOp(), genExpand(), genInsertionStore(), genLoopBoundary(), and genTensorStore().

◆ updateReduc()

void CodegenEnv::updateReduc ( Value val )

Definition at line 234 of file CodegenEnv.cpp.

References isReduc().

Referenced by endIf(), finalizeWhileOp(), genLoopBoundary(), and genTensorStore().

◆ updateValidLexInsert()

void CodegenEnv::updateValidLexInsert ( Value val )

Definition at line 255 of file CodegenEnv.cpp.

References isReduc().

Referenced by endIf(), endLoop(), finalizeWhileOp(), and genLoopBoundary().

The documentation for this class was generated from the following files:

lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.h
lib/Dialect/SparseTensor/Transforms/Utils/CodegenEnv.cpp

Public Member Functions

Detailed Description

Constructor & Destructor Documentation

◆ CodegenEnv()

Member Function Documentation

◆ atExpandLevel()

◆ emitter()

◆ endCustomReduc()

◆ endExpand()

◆ endReduc()

◆ endValidLexInsert()

◆ exp()

◆ generatingSparseIterator()

◆ genLoopBoundary()

◆ getCurrentDepth()

◆ getCustomRedId()

◆ getExpandAdded()

◆ getExpandCount()

◆ getExpandFilled()

◆ getExpandValues()

◆ getExprId()

◆ getInsertionChain()

◆ getLoopNum()

◆ getLoopVar()

◆ getReduc()

◆ getValidLexInsert()

◆ hasSparseOutput()

◆ initTensorExp()

◆ isAdmissibleTensorExp()

◆ isCustomReduc()

◆ isExpand()

◆ isReduc()

◆ isSparseOutput()

◆ isValidLexInsert()

◆ lat()

◆ lt() [1/2]

◆ lt() [2/2]

◆ makeLoopId()

◆ makeTensorId()

◆ makeTensorLevel() [1/2]

◆ makeTensorLevel() [2/2]

◆ makeTensorLoopId()

◆ merger()

◆ op()

◆ options()

◆ set()

◆ startCustomReduc()

◆ startEmit()

◆ startExpand()

◆ startReduc()

◆ startValidLexInsert()

◆ unpackTensorLevel()

◆ unpackTensorLevelRange()

◆ updateExpandCount()

◆ updateInsertionChain()

◆ updateReduc()

◆ updateValidLexInsert()