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

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

Member Function Documentation

◆ atExpandLevel()

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

Definition at line 200 of file CodegenEnv.cpp.

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

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

◆ endExpand()

void CodegenEnv::endExpand ( )

Definition at line 219 of file CodegenEnv.cpp.

◆ endReduc()

Value CodegenEnv::endReduc ( )

Definition at line 242 of file CodegenEnv.cpp.

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

◆ endValidLexInsert()

void CodegenEnv::endValidLexInsert ( )

Definition at line 261 of file CodegenEnv.cpp.

References isReduc(), and isValidLexInsert().

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

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

Referenced by endLoop(), and genCoIteration().

◆ getCurrentDepth()

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

inline

Definition at line 111 of file CodegenEnv.h.

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

Referenced by genStmt(), and getAllTidLvlsInLatPoints().

◆ getCustomRedId()

Value CodegenEnv::getCustomRedId ( ) const

Definition at line 271 of file CodegenEnv.cpp.

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

◆ getExpandAdded()

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

inline

Definition at line 140 of file CodegenEnv.h.

◆ getExpandCount()

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

inline

Definition at line 141 of file CodegenEnv.h.

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

◆ getExpandFilled()

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

inline

Definition at line 139 of file CodegenEnv.h.

◆ getExpandValues()

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

inline

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

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

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

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

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

◆ getValidLexInsert()

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

inline

Definition at line 157 of file CodegenEnv.h.

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

◆ hasSparseOutput()

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

inline

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

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

◆ isReduc()

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

inline

Definition at line 149 of file CodegenEnv.h.

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

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

◆ isSparseOutput()

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

inline

Definition at line 129 of file CodegenEnv.h.

◆ isValidLexInsert()

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

inline

Definition at line 155 of file CodegenEnv.h.

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

◆ 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(), getAllTidLvlsInLatPoints(), and startLoopSeq().

◆ 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 getAllTidLvlsInLatPoints(), and makeTensorLevel().

◆ 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(), getAllTidLvlsInLatPoints(), startEmit(), startLoop(), and startLoopSeq().

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

◆ startCustomReduc()

void CodegenEnv::startCustomReduc ( ExprId exp )

Definition at line 266 of file CodegenEnv.cpp.

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

◆ startEmit()

void CodegenEnv::startEmit ( SparseEmitStrategy emitStrategy )

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

◆ startExpand()

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

Definition at line 205 of file CodegenEnv.cpp.

◆ startReduc()

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

Definition at line 228 of file CodegenEnv.cpp.

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

◆ startValidLexInsert()

void CodegenEnv::startValidLexInsert ( Value val )

Definition at line 251 of file CodegenEnv.cpp.

References isReduc(), and isValidLexInsert().

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

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

◆ updateInsertionChain()

void CodegenEnv::updateInsertionChain ( Value chain )

Definition at line 195 of file CodegenEnv.cpp.

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

◆ updateReduc()

void CodegenEnv::updateReduc ( Value val )

Definition at line 235 of file CodegenEnv.cpp.

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

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

◆ updateValidLexInsert()

void CodegenEnv::updateValidLexInsert ( Value val )

Definition at line 256 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()

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