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/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

Merger &	merger ()

LoopEmitter &	emitter ()

void	startEmit ()

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	getLoopDepth () 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 loop identified by the given `LoopId`. More...

bool	hasSparseOutput () const

bool	isSparseOutput (OpOperand *o) const

Value	getInsertionChain () const

void	updateInsertionChain (Value chain)

bool	atExpandLevel (OpOperand *o, unsigned rank, LoopOrd 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	setValidLexInsert (Value val)

void	clearValidLexInsert ()

Value	getValidLexInsert () const

void	startCustomReduc (ExprId exp)

bool	isCustomReduc () const

Value	getCustomRedId ()

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 44 of file CodegenEnv.cpp.

Member Function Documentation

◆ atExpandLevel()

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

Definition at line 210 of file CodegenEnv.cpp.

◆ clearValidLexInsert()

void CodegenEnv::clearValidLexInsert ( )

Definition at line 266 of file CodegenEnv.cpp.

References isReduc().

◆ emitter()

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

inline

Definition at line 53 of file CodegenEnv.h.

Referenced by endLoop(), endLoopSeq(), genCoIteration(), genIf(), genSubscript(), startLoop(), and startLoopSeq().

◆ endCustomReduc()

void CodegenEnv::endCustomReduc ( )

Definition at line 281 of file CodegenEnv.cpp.

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

◆ endExpand()

void CodegenEnv::endExpand ( )

Definition at line 229 of file CodegenEnv.cpp.

◆ endReduc()

Value CodegenEnv::endReduc ( )

Definition at line 252 of file CodegenEnv.cpp.

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

◆ exp()

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

inline

Definition at line 77 of file CodegenEnv.h.

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

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

◆ 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 112 of file CodegenEnv.cpp.

References isExpand(), isReduc(), setValidLexInsert(), updateExpandCount(), updateInsertionChain(), and updateReduc().

Referenced by endLoop(), and genCoIteration().

◆ getCustomRedId()

Value CodegenEnv::getCustomRedId ( )

Definition at line 276 of file CodegenEnv.cpp.

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

◆ getExpandAdded()

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

inline

Definition at line 143 of file CodegenEnv.h.

◆ getExpandCount()

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

inline

Definition at line 144 of file CodegenEnv.h.

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

◆ getExpandFilled()

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

inline

Definition at line 142 of file CodegenEnv.h.

◆ getExpandValues()

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

inline

Definition at line 141 of file CodegenEnv.h.

◆ 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 133 of file CodegenEnv.h.

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

◆ getLoopDepth()

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

inline

Definition at line 111 of file CodegenEnv.h.

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

◆ getLoopNum()

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

inline

Definition at line 85 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 loop identified by the given LoopId.

This method handles application of the topological sort in order to convert the LoopId into the corresponding LoopOrd.

Definition at line 196 of file CodegenEnv.cpp.

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

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

◆ getReduc()

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

inline

Definition at line 154 of file CodegenEnv.h.

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

◆ getValidLexInsert()

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

inline

Definition at line 158 of file CodegenEnv.h.

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

◆ hasSparseOutput()

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

inline

Definition at line 130 of file CodegenEnv.h.

Referenced by isParallelFor().

◆ initTensorExp()

LogicalResult CodegenEnv::initTensorExp ( )

Definition at line 52 of file CodegenEnv.cpp.

References mlir::sparse_tensor::Merger::buildTensorExpFromLinalg(), mlir::failure(), isAdmissibleTensorExp(), op(), and mlir::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 145 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 161 of file CodegenEnv.h.

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

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

◆ isExpand()

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

inline

Definition at line 139 of file CodegenEnv.h.

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

◆ isReduc()

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

inline

Definition at line 152 of file CodegenEnv.h.

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

Referenced by clearValidLexInsert(), endIf(), endLoop(), endReduc(), finalizeWhileOp(), genIf(), genLoopBoundary(), genTensorStore(), setValidLexInsert(), startReduc(), and updateReduc().

◆ isSparseOutput()

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

inline

Definition at line 131 of file CodegenEnv.h.

◆ lat()

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

inline

Definition at line 78 of file CodegenEnv.h.

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

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

◆ lt() [1/2]

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

inline

Definition at line 80 of file CodegenEnv.h.

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

◆ lt() [2/2]

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

inline

Definition at line 83 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 71 of file CodegenEnv.h.

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

Referenced by genIndex().

◆ makeTensorId()

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

inlineconstexpr

Definition at line 68 of file CodegenEnv.h.

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

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

◆ makeTensorLevel() [1/2]

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

inline

Definition at line 100 of file CodegenEnv.h.

References makeTensorLevel().

◆ makeTensorLevel() [2/2]

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

inline

Definition at line 91 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 makeTensorLevel(), startEmit(), startLoopSeq(), and translateBitsToTidLvlPairs().

◆ makeTensorLoopId()

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

inlineconstexpr

Definition at line 74 of file CodegenEnv.h.

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

◆ merger()

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

inline

Definition at line 52 of file CodegenEnv.h.

Referenced by findSparseAnnotations(), genIf(), genStmt(), startEmit(), startLoop(), startLoopSeq(), and translateBitsToTidLvlPairs().

◆ op()

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

inline

Definition at line 50 of file CodegenEnv.h.

Referenced by endIf(), endLoop(), endLoopSeq(), finalizeWhileOp(), findSparseAnnotations(), genCoIteration(), genIf(), genIndex(), genInitConstantDenseAddress(), genSubscript(), getCustomRedId(), initTensorExp(), startLoop(), and startLoopSeq().

◆ options()

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

inline

Definition at line 51 of file CodegenEnv.h.

Referenced by isParallelFor().

◆ set()

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

inline

Definition at line 79 of file CodegenEnv.h.

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

Referenced by genStmt(), and startLoopSeq().

◆ setValidLexInsert()

void CodegenEnv::setValidLexInsert ( Value val )

Definition at line 261 of file CodegenEnv.cpp.

References isReduc().

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

◆ startCustomReduc()

void CodegenEnv::startCustomReduc ( ExprId exp )

Definition at line 271 of file CodegenEnv.cpp.

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

◆ startEmit()

void CodegenEnv::startEmit ( )

Definition at line 62 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(), makeTensorLevel(), merger(), mlir::sparse_tensor::Merger::setHasSparseOut(), and sortDependentLoops().

◆ startExpand()

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

Definition at line 215 of file CodegenEnv.cpp.

◆ startReduc()

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

Definition at line 238 of file CodegenEnv.cpp.

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

◆ unpackTensorLevel()

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

inline

Definition at line 103 of file CodegenEnv.h.

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

◆ unpackTensorLevelRange()

template<class ContainerTy >

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

inline

Definition at line 107 of file CodegenEnv.h.

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

Referenced by startLoop().

◆ updateExpandCount()

void CodegenEnv::updateExpandCount ( Value count )

Definition at line 224 of file CodegenEnv.cpp.

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

◆ updateInsertionChain()

void CodegenEnv::updateInsertionChain ( Value chain )

Definition at line 204 of file CodegenEnv.cpp.

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

◆ updateReduc()

void CodegenEnv::updateReduc ( Value val )

Definition at line 245 of file CodegenEnv.cpp.

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

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

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

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

Public Member Functions

Detailed Description

Constructor & Destructor Documentation

◆ CodegenEnv()

Member Function Documentation

◆ atExpandLevel()

◆ clearValidLexInsert()

◆ emitter()

◆ endCustomReduc()

◆ endExpand()

◆ endReduc()

◆ exp()

◆ genLoopBoundary()

◆ getCustomRedId()

◆ getExpandAdded()

◆ getExpandCount()

◆ getExpandFilled()

◆ getExpandValues()

◆ getExprId()

◆ getInsertionChain()

◆ getLoopDepth()

◆ getLoopNum()

◆ getLoopVar()

◆ getReduc()

◆ getValidLexInsert()

◆ hasSparseOutput()

◆ initTensorExp()

◆ isAdmissibleTensorExp()

◆ isCustomReduc()

◆ isExpand()

◆ isReduc()

◆ isSparseOutput()

◆ lat()

◆ lt() [1/2]

◆ lt() [2/2]

◆ makeLoopId()

◆ makeTensorId()

◆ makeTensorLevel() [1/2]

◆ makeTensorLevel() [2/2]

◆ makeTensorLoopId()

◆ merger()

◆ op()

◆ options()

◆ set()

◆ setValidLexInsert()

◆ startCustomReduc()

◆ startEmit()

◆ startExpand()

◆ startReduc()

◆ unpackTensorLevel()

◆ unpackTensorLevelRange()

◆ updateExpandCount()

◆ updateInsertionChain()

◆ updateReduc()