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. More...

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). More...

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. More...

Value	getLoopVar (LoopId i) const
	Returns the induction-variable for the given loop. More...

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.

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(), 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.

◆ endReduc()

Value CodegenEnv::endReduc ( )

Definition at line 241 of file CodegenEnv.cpp.

References mlir::sparse_tensor::Merger::clearExprValue(), 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.

References mlir::sparse_tensor::Merger::exp().

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, and mlir::SparsificationOptions::sparseEmitStrategy.

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(), and isValidLexInsert().

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

◆ getCurrentDepth()

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

inline

Definition at line 115 of file CodegenEnv.h.

References mlir::sparse_tensor::LoopEmitter::getCurrentDepth().

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 genInsertionStore().

◆ getExpandCount()

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

inline

Definition at line 145 of file CodegenEnv.h.

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

◆ getExpandFilled()

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

inline

Definition at line 143 of file CodegenEnv.h.

Referenced by genInsertionLoadReduce(), and genInsertionStore().

◆ getExpandValues()

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

inline

Definition at line 142 of file CodegenEnv.h.

Referenced by 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(), genIf(), genInsertionStore(), genResult(), genStmt(), and genTensorStore().

◆ getLoopNum()

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

inline

Definition at line 89 of file CodegenEnv.h.

References mlir::sparse_tensor::Merger::getNumLoops().

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.

References mlir::sparse_tensor::LoopEmitter::getLoopIV().

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 mlir::sparse_tensor::Merger::buildTensorExpFromLinalg(), isAdmissibleTensorExp(), and op().

◆ 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(), mlir::IROperand< DerivedT, IRValueT >::get(), getLoopNum(), mlir::OpOperand::getOperandNumber(), mlir::sparse_tensor::getSparseTensorType(), mlir::sparse_tensor::Merger::hasNegateOnOut(), mlir::sparse_tensor::SparseTensorType::isAllDense(), isMaterializing(), mlir::linalg::isReductionIterator(), mlir::sparse_tensor::Merger::isSingleCondition(), 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.

References mlir::sparse_tensor::Merger::lat().

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.

References mlir::sparse_tensor::Merger::getLvlType().

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 mlir::sparse_tensor::Merger::getLvlType().

◆ makeLoopId()

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

inlineconstexpr

Definition at line 75 of file CodegenEnv.h.

References mlir::sparse_tensor::Merger::makeLoopId().

Referenced by genIndex(), and relinkBranch().

◆ makeTensorId()

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

inlineconstexpr

Definition at line 72 of file CodegenEnv.h.

References mlir::sparse_tensor::Merger::makeTensorId().

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.

References mlir::sparse_tensor::LoopEmitter::getNumManifestTensors(), mlir::sparse_tensor::Merger::getNumTensors(), mlir::sparse_tensor::LoopEmitter::getNumTensors(), mlir::sparse_tensor::Merger::getOutTensorID(), mlir::sparse_tensor::LoopEmitter::getOutTensorId(), mlir::sparse_tensor::Merger::getSynTensorID(), mlir::sparse_tensor::LoopEmitter::getSynTensorId(), and mlir::sparse_tensor::LoopEmitter::makeTensorLevel().

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

◆ makeTensorLoopId()

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

inlineconstexpr

Definition at line 78 of file CodegenEnv.h.

References mlir::sparse_tensor::Merger::makeTensorLoopId().

◆ 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.

References mlir::sparse_tensor::Merger::set().

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::IROperand< DerivedT, IRValueT >::get(), mlir::get(), mlir::sparse_tensor::Merger::getDependentLoops(), getLoopNum(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::sparse_tensor::LoopEmitter::initialize(), makeTensorId(), merger(), mlir::sparse_tensor::Merger::setHasSparseOut(), and sortDependentLoops().

◆ startExpand()

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

Definition at line 204 of file CodegenEnv.cpp.

◆ startReduc()

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

Definition at line 227 of file CodegenEnv.cpp.

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

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.

References mlir::sparse_tensor::LoopEmitter::unpackTensorLevel().

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.

References mlir::sparse_tensor::LoopEmitter::unpackTensorLevelRange().

Referenced by startLoop().

◆ updateExpandCount()

void CodegenEnv::updateExpandCount ( Value count )

Definition at line 213 of file CodegenEnv.cpp.

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

◆ updateInsertionChain()

void CodegenEnv::updateInsertionChain ( Value chain )

Definition at line 194 of file CodegenEnv.cpp.

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

◆ updateReduc()

void CodegenEnv::updateReduc ( Value val )

Definition at line 234 of file CodegenEnv.cpp.

References mlir::sparse_tensor::Merger::clearExprValue(), isReduc(), and mlir::sparse_tensor::Merger::setExprValue().

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

◆ updateValidLexInsert()

void CodegenEnv::updateValidLexInsert ( Value val )

Definition at line 255 of file CodegenEnv.cpp.

References isReduc().

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

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()