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

Public Types
using	OutputUpdater = function_ref< Value(OpBuilder &builder, Location loc, Value memref, Value tensor)>
	Optional callback function to setup dense output tensors when initializing the loop emitter (e.g., to fill a dense output with zeros). More...

using	SynTensorBoundSetter = function_ref< Value(OpBuilder &builder, Location loc, Level lvl)>
	Optional callback function to set the bound for the synthetic tensor, which essentially is the dense loop bound. More...

using	DependentLvlGetter = function_ref< std::vector< std::pair< LoopId, unsigned > >(TensorId, Level)>

Public Member Functions
	LoopEmitter ()=default

void	initialize (ValueRange tensors, StringAttr loopTag=nullptr, bool hasOutput=false, bool isSparseOut=false, unsigned numLoops=0, DependentLvlGetter getter=nullptr, SparseEmitStrategy emitStrategy=SparseEmitStrategy::kFunctional)
	Takes an array of input tensors, which the generated loops will iterate over. More...

	LoopEmitter (ValueRange tensors, StringAttr loopTag=nullptr, bool hasOutput=false, bool isSparseOut=false, unsigned numLoops=0, DependentLvlGetter getter=nullptr, SparseEmitStrategy emitStrategy=SparseEmitStrategy::kFunctional)

void	initializeLoopEmit (OpBuilder &builder, Location loc, OutputUpdater updater=nullptr, SynTensorBoundSetter synSetter=nullptr)
	Starts a loop emitting session by generating all the buffers needed for iterating over the tensors. More...

Value	genAffine (OpBuilder &builder, Location loc, AffineExpr a)
	Generates code to compute an affine expression whose variables are `LoopId`s (i.e., `a.cast<AffineDimExpr>().getPosition()` is a valid `LoopId`). More...

void	enterNewLoopSeq (OpBuilder &builder, Location loc, ArrayRef< TensorLevel > tidLvls)
	Enters a new loop sequence, the loops within the same sequence starts from the break points of previous loop instead of starting over from 0. More...

void	exitCurrentLoopSeq (OpBuilder &builder, Location loc)
	Exits the current loop sequence, this will reset universal index to 0. More...

void	locateLvlAtAffineAddress (OpBuilder &builder, Location loc, TensorLevel tidLvl, AffineExpr lvlExpr)
	Emits the address for a dense level based on the value evaluated by the provided affine expression. More...

Operation *	enterCoIterationOverTensorsAtLvls (OpBuilder &builder, Location loc, ArrayRef< TensorLevel > tidLvls, unsigned numCases, MutableArrayRef< Value > reduc={}, bool isParallel=false, bool needsUniv=false)
	Emits a co-iteration loop over a set of tensors. More...

Region *	enterCurrentCoIterationCase (OpBuilder &builder, Location loc, I64BitSet caseBit, unsigned caseIdx, MutableArrayRef< Value > reduc)

void	exitCurrentLoop (RewriterBase &rewriter, Location loc, MutableArrayRef< Value > reduc={})
	Generates code to exit the current loop (e.g., generates yields, forwards loop induction variables, etc). More...

auto	getLoopIVsRange () const
	Get the range of values for all induction variables. More...

SmallVector< Value >	getLoopIVs () const
	Fills the out-parameter with the loop induction variables for all loops in the current loop-stack. More...

LoopId	getCurrentDepth () const
	Gets the current depth of the loop-stack. More...

Value	getLoopIV (LoopId n) const
	Gets loop induction variable for the given loop. More...

unsigned	getNumManifestTensors () const
	Gets the total number of manifest tensors (excluding the synthetic tensor). More...

unsigned	getNumTensors () const
	Gets the total number of tensors that loopEmitter is operating on. More...

TensorId	getSynTensorId () const
	Gets the TensorId for synthetic tensor. More...

TensorId	getOutTensorId () const
	Gets the TensorId for output tensor. More...

TensorLevel	makeTensorLevel (TensorId t, Level l) const
	Compresses a TensorId and Level into a TensorLevel. More...

std::pair< TensorId, Level >	unpackTensorLevel (TensorLevel tidLvl) const
	De-compresses a TensorLevel back to a pair of TensorId and Level. More...

template<class ContainerTy >
auto	unpackTensorLevelRange (ContainerTy &&c) const
	Converts a range of TensorLevel to a range of std::pair<TensorId, Level> More...

SmallVector< Value >	getValPosits (TensorId tid) const
	Getters. More...

Value	getCoord (TensorId tid, Level lvl) const

const std::vector< Value > &	getValBuffer () const

Static Public Member Functions
constexpr static llvm::StringLiteral	getLoopEmitterLoopAttrName ()

Detailed Description

Definition at line 55 of file LoopEmitter.h.

Member Typedef Documentation

◆ DependentLvlGetter

using mlir::sparse_tensor::LoopEmitter::DependentLvlGetter = function_ref<std::vector<std::pair<LoopId, unsigned> >(TensorId, Level)>

Definition at line 77 of file LoopEmitter.h.

◆ OutputUpdater

using mlir::sparse_tensor::LoopEmitter::OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc, Value memref, Value tensor)>

Optional callback function to setup dense output tensors when initializing the loop emitter (e.g., to fill a dense output with zeros).

Definition at line 59 of file LoopEmitter.h.

◆ SynTensorBoundSetter

using mlir::sparse_tensor::LoopEmitter::SynTensorBoundSetter = function_ref<Value(OpBuilder &builder, Location loc, Level lvl)>

Optional callback function to set the bound for the synthetic tensor, which essentially is the dense loop bound.

Definition at line 64 of file LoopEmitter.h.

Constructor & Destructor Documentation

◆ LoopEmitter() [1/2]

mlir::sparse_tensor::LoopEmitter::LoopEmitter ( )

default

◆ LoopEmitter() [2/2]

LoopEmitter::LoopEmitter	(	ValueRange	tensors,
		StringAttr	loopTag = `nullptr`,
		bool	hasOutput = `false`,
		bool	isSparseOut = `false`,
		unsigned	numLoops = `0`,
		DependentLvlGetter	getter = `nullptr`,
		SparseEmitStrategy	emitStrategy = `SparseEmitStrategy::kFunctional`
	)

explicit

Definition at line 116 of file LoopEmitter.cpp.

References initialize().

Member Function Documentation

◆ enterCoIterationOverTensorsAtLvls()

Operation * LoopEmitter::enterCoIterationOverTensorsAtLvls	(	OpBuilder &	builder,
		Location	loc,
		ArrayRef< TensorLevel >	tidLvls,
		unsigned	numCases,
		MutableArrayRef< Value >	reduc = `{}`,
		bool	isParallel = `false`,
		bool	needsUniv = `false`
	)

Emits a co-iteration loop over a set of tensors.

Emits loop over tensor_tid_lvl, it assumes that loops between tensor_tid_[0, lvl - 1] have already been generated. The function will also perform in-place update on the reduc vector to return the reduction variable used inside the generated loop.

Definition at line 595 of file LoopEmitter.cpp.

References copy(), mlir::OpBuilder::create(), mlir::Block::front(), mlir::OpBuilder::getInsertionBlock(), mlir::kSparseIterator, mlir::sparse_tensor::SparseIterator::locate(), mlir::sparse_tensor::SparseIterator::lvl, makeTensorLevel(), mlir::OpBuilder::setInsertionPointToStart(), mlir::sparse_tensor::SparseIterator::tid, unpackTensorLevel(), and unpackTensorLevelRange().

◆ enterCurrentCoIterationCase()

Region * LoopEmitter::enterCurrentCoIterationCase	(	OpBuilder &	builder,
		Location	loc,
		I64BitSet	caseBit,
		unsigned	caseIdx,
		MutableArrayRef< Value >	reduc
	)

Definition at line 542 of file LoopEmitter.cpp.

References mlir::Block::addArguments(), mlir::sparse_tensor::I64BitSet::bits(), copy(), mlir::Region::emplaceBlock(), mlir::detail::enumerate(), mlir::Region::front(), mlir::Builder::getArrayAttr(), mlir::Region::getBlocks(), mlir::Builder::getI64IntegerAttr(), mlir::Builder::getIndexType(), mlir::OpBuilder::setInsertionPointToStart(), and unpackTensorLevelRange().

Referenced by genCoIterationCase().

◆ enterNewLoopSeq()

void LoopEmitter::enterNewLoopSeq	(	OpBuilder &	builder,
		Location	loc,
		ArrayRef< TensorLevel >	tidLvls
	)

Enters a new loop sequence, the loops within the same sequence starts from the break points of previous loop instead of starting over from 0.

e.g., { // loop sequence start. p0 = while(xxx) ... break p0

// Starts loop from p0 for (i = p0; i < end; i++) ... // loop sequence end. }

Definition at line 414 of file LoopEmitter.cpp.

References C_IDX, mlir::kSparseIterator, and unpackTensorLevelRange().

Referenced by startLoopSeq().

◆ exitCurrentLoop()

void LoopEmitter::exitCurrentLoop	(	RewriterBase &	rewriter,
		Location	loc,
		MutableArrayRef< Value >	reduc = `{}`
	)

Generates code to exit the current loop (e.g., generates yields, forwards loop induction variables, etc).

Definition at line 861 of file LoopEmitter.cpp.

References copy(), mlir::OpBuilder::create(), mlir::Operation::getResults(), mlir::kSparseIterator, mlir::OpBuilder::setInsertionPoint(), mlir::OpBuilder::setInsertionPointAfter(), and mlir::OpBuilder::setInsertionPointToEnd().

◆ exitCurrentLoopSeq()

void LoopEmitter::exitCurrentLoopSeq	(	OpBuilder &	builder,
		Location	loc
	)

Exits the current loop sequence, this will reset universal index to 0.

Definition at line 431 of file LoopEmitter.cpp.

References unpackTensorLevelRange().

Referenced by endLoopSeq().

◆ genAffine()

Value LoopEmitter::genAffine	(	OpBuilder &	builder,
		Location	loc,
		AffineExpr	a
	)

Generates code to compute an affine expression whose variables are LoopIds (i.e., a.cast<AffineDimExpr>().getPosition() is a valid LoopId).

Definition at line 442 of file LoopEmitter.cpp.

References mlir::Add, ADDI, C_IDX, mlir::Constant, mlir::DimId, mlir::AffineExpr::getKind(), mlir::Mul, and MULI.

Referenced by genSubscript(), and locateLvlAtAffineAddress().

◆ getCoord()

Value mlir::sparse_tensor::LoopEmitter::getCoord	(	TensorId	tid,
		Level	lvl
	)		const

inline

Definition at line 238 of file LoopEmitter.h.

References mlir::sparse_tensor::SparseIterator::getCrd().

Referenced by genIf().

◆ getCurrentDepth()

LoopId mlir::sparse_tensor::LoopEmitter::getCurrentDepth ( ) const

inline

Gets the current depth of the loop-stack.

Definition at line 172 of file LoopEmitter.h.

References getLoopIVsRange().

Referenced by mlir::sparse_tensor::CodegenEnv::getCurrentDepth(), and getLoopIV().

◆ getLoopEmitterLoopAttrName()

constexpr static llvm::StringLiteral mlir::sparse_tensor::LoopEmitter::getLoopEmitterLoopAttrName ( )

inlinestaticconstexpr

Definition at line 243 of file LoopEmitter.h.

Referenced by finalizeWhileOp().

◆ getLoopIV()

Value mlir::sparse_tensor::LoopEmitter::getLoopIV ( LoopId n ) const

inline

Gets loop induction variable for the given loop.

Definition at line 175 of file LoopEmitter.h.

References getCurrentDepth(), and getLoopIVsRange().

Referenced by mlir::sparse_tensor::CodegenEnv::getLoopVar().

◆ getLoopIVs()

SmallVector<Value> mlir::sparse_tensor::LoopEmitter::getLoopIVs ( ) const

inline

Fills the out-parameter with the loop induction variables for all loops in the current loop-stack.

Definition at line 167 of file LoopEmitter.h.

References getLoopIVsRange().

◆ getLoopIVsRange()

auto mlir::sparse_tensor::LoopEmitter::getLoopIVsRange ( ) const

inline

Get the range of values for all induction variables.

Definition at line 161 of file LoopEmitter.h.

Referenced by genInsertionStore(), getCurrentDepth(), getLoopIV(), and getLoopIVs().

◆ getNumManifestTensors()

unsigned mlir::sparse_tensor::LoopEmitter::getNumManifestTensors ( ) const

inline

Gets the total number of manifest tensors (excluding the synthetic tensor).

Definition at line 185 of file LoopEmitter.h.

Referenced by getNumTensors(), getOutTensorId(), initializeLoopEmit(), and mlir::sparse_tensor::CodegenEnv::makeTensorLevel().

◆ getNumTensors()

unsigned mlir::sparse_tensor::LoopEmitter::getNumTensors ( ) const

inline

Gets the total number of tensors that loopEmitter is operating on.

Definition at line 188 of file LoopEmitter.h.

References getNumManifestTensors().

Referenced by mlir::sparse_tensor::CodegenEnv::makeTensorLevel(), makeTensorLevel(), and unpackTensorLevel().

◆ getOutTensorId()

TensorId mlir::sparse_tensor::LoopEmitter::getOutTensorId ( ) const

inline

Gets the TensorId for output tensor.

Definition at line 197 of file LoopEmitter.h.

References getNumManifestTensors().

Referenced by mlir::sparse_tensor::CodegenEnv::makeTensorLevel().

◆ getSynTensorId()

TensorId mlir::sparse_tensor::LoopEmitter::getSynTensorId ( ) const

inline

Gets the TensorId for synthetic tensor.

Definition at line 194 of file LoopEmitter.h.

Referenced by initializeLoopEmit(), and mlir::sparse_tensor::CodegenEnv::makeTensorLevel().

◆ getValBuffer()

const std::vector<Value>& mlir::sparse_tensor::LoopEmitter::getValBuffer ( ) const

inline

Definition at line 241 of file LoopEmitter.h.

Referenced by genResult(), and genSubscript().

◆ getValPosits()

SmallVector<Value> mlir::sparse_tensor::LoopEmitter::getValPosits ( TensorId tid ) const

inline

Getters.

Definition at line 227 of file LoopEmitter.h.

References mlir::kSparseIterator.

Referenced by genSubscript().

◆ initialize()

void LoopEmitter::initialize	(	ValueRange	tensors,
		StringAttr	loopTag = `nullptr`,
		bool	hasOutput = `false`,
		bool	isSparseOut = `false`,
		unsigned	numLoops = `0`,
		DependentLvlGetter	getter = `nullptr`,
		SparseEmitStrategy	emitStrategy = `SparseEmitStrategy::kFunctional`
	)

Takes an array of input tensors, which the generated loops will iterate over.

Each tensor is given a TensorId (numerically equal to the position of that tensor Value in the array). Setting isSparseOut indicates that the sparse output tensor is empty, so the loop emitter will generate loops over it according to the level-sizes.

Definition at line 123 of file LoopEmitter.cpp.

References mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::getRankedTensorType(), mlir::Value::getType(), and mlir::sparse_tensor::isZeroRankedTensorOrScalar().

Referenced by LoopEmitter(), and mlir::sparse_tensor::CodegenEnv::startEmit().

◆ initializeLoopEmit()

void LoopEmitter::initializeLoopEmit	(	OpBuilder &	builder,
		Location	loc,
		LoopEmitter::OutputUpdater	updater = `nullptr`,
		LoopEmitter::SynTensorBoundSetter	synSetter = `nullptr`
	)

Starts a loop emitting session by generating all the buffers needed for iterating over the tensors.

Definition at line 231 of file LoopEmitter.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::Value::getDefiningOp(), mlir::bufferization::getMemRefTypeWithFullyDynamicLayout(), getNumManifestTensors(), mlir::sparse_tensor::getSparseTensorType(), getSynTensorId(), mlir::Value::getType(), mlir::kSparseIterator, mlir::sparse_tensor::makeSparseTensorLevel(), mlir::sparse_tensor::makeSynLevelAndIterator(), and tryFoldTensors().

Referenced by genBuffers().

◆ locateLvlAtAffineAddress()

void LoopEmitter::locateLvlAtAffineAddress	(	OpBuilder &	builder,
		Location	loc,
		TensorLevel	tidLvl,
		AffineExpr	lvlExpr
	)

Emits the address for a dense level based on the value evaluated by the provided affine expression.

Definition at line 695 of file LoopEmitter.cpp.

References genAffine(), mlir::sparse_tensor::SparseIterator::genInit(), mlir::sparse_tensor::SparseIterator::kind, mlir::sparse_tensor::kTrivial, mlir::sparse_tensor::SparseIterator::locate(), mlir::sparse_tensor::SparseIterator::randomAccessible(), and unpackTensorLevel().

Referenced by genConstantDenseAddressFromLevel(), and startLoop().

◆ makeTensorLevel()

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

inline

Compresses a TensorId and Level into a TensorLevel.

Definition at line 203 of file LoopEmitter.h.

References getNumTensors().

Referenced by enterCoIterationOverTensorsAtLvls(), and mlir::sparse_tensor::CodegenEnv::makeTensorLevel().

◆ unpackTensorLevel()

std::pair<TensorId, Level> mlir::sparse_tensor::LoopEmitter::unpackTensorLevel ( TensorLevel tidLvl ) const

inline

De-compresses a TensorLevel back to a pair of TensorId and Level.

Definition at line 208 of file LoopEmitter.h.

References getNumTensors().

Referenced by enterCoIterationOverTensorsAtLvls(), locateLvlAtAffineAddress(), mlir::sparse_tensor::CodegenEnv::unpackTensorLevel(), and unpackTensorLevelRange().

◆ unpackTensorLevelRange()

template<class ContainerTy >

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

inline

Converts a range of TensorLevel to a range of std::pair<TensorId, Level>

Definition at line 215 of file LoopEmitter.h.

References unpackTensorLevel().

Referenced by enterCoIterationOverTensorsAtLvls(), enterCurrentCoIterationCase(), enterNewLoopSeq(), exitCurrentLoopSeq(), and mlir::sparse_tensor::CodegenEnv::unpackTensorLevelRange().

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

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

Public Types

Public Member Functions

Static Public Member Functions

Detailed Description

Member Typedef Documentation

◆ DependentLvlGetter

◆ OutputUpdater

◆ SynTensorBoundSetter

Constructor & Destructor Documentation

◆ LoopEmitter() [1/2]

◆ LoopEmitter() [2/2]

Member Function Documentation

◆ enterCoIterationOverTensorsAtLvls()

◆ enterCurrentCoIterationCase()

◆ enterNewLoopSeq()

◆ exitCurrentLoop()

◆ exitCurrentLoopSeq()

◆ genAffine()

◆ getCoord()

◆ getCurrentDepth()

◆ getLoopEmitterLoopAttrName()

◆ getLoopIV()

◆ getLoopIVs()

◆ getLoopIVsRange()

◆ getNumManifestTensors()

◆ getNumTensors()

◆ getOutTensorId()

◆ getSynTensorId()

◆ getValBuffer()

◆ getValPosits()

◆ initialize()

◆ initializeLoopEmit()

◆ locateLvlAtAffineAddress()

◆ makeTensorLevel()

◆ unpackTensorLevel()

◆ unpackTensorLevelRange()