#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, MutableArrayRef< Value > reduc={}, bool isParallel=false, bool needsUniv=false)
	Emits a co-iteration loop over a set of tensors. More...

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 82 of file LoopEmitter.cpp.

References initialize().

Member Function Documentation

◆ enterCoIterationOverTensorsAtLvls()

Operation * LoopEmitter::enterCoIterationOverTensorsAtLvls	(	OpBuilder &	builder,
		Location	loc,
		ArrayRef< TensorLevel >	tidLvls,
		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 555 of file LoopEmitter.cpp.

References mlir::OpBuilder::getInsertionBlock(), mlir::sparse_tensor::SparseIterator::locate(), mlir::sparse_tensor::SparseIterator::lvl, makeTensorLevel(), and mlir::sparse_tensor::SparseIterator::tid.

◆ 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 353 of file LoopEmitter.cpp.

References C_IDX, 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 767 of file LoopEmitter.cpp.

References mlir::OpBuilder::setInsertionPoint(), 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 368 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 379 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 229 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 168 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 234 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 171 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 163 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 157 of file LoopEmitter.h.

Referenced by 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 181 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 184 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 193 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 190 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 232 of file LoopEmitter.h.

Referenced by genSubscript().

◆ getValPosits()

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

inline

Getters.

Definition at line 223 of file LoopEmitter.h.

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 89 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 181 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::sparse_tensor::makeSparseTensorLevel(), and mlir::sparse_tensor::makeSynLevelAndIterator().

◆ 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 610 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 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 199 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 204 of file LoopEmitter.h.

References getNumTensors().

Referenced by 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 211 of file LoopEmitter.h.

References unpackTensorLevel().

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

◆ enterNewLoopSeq()

◆ exitCurrentLoop()

◆ exitCurrentLoopSeq()

◆ genAffine()

◆ getCoord()

◆ getCurrentDepth()

◆ getLoopEmitterLoopAttrName()

◆ getLoopIV()

◆ getLoopIVs()

◆ getLoopIVsRange()

◆ getNumManifestTensors()

◆ getNumTensors()

◆ getOutTensorId()

◆ getSynTensorId()

◆ getValBuffer()

◆ getValPosits()

◆ initialize()

◆ initializeLoopEmit()

◆ locateLvlAtAffineAddress()

◆ makeTensorLevel()

◆ unpackTensorLevel()

◆ unpackTensorLevelRange()