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

Public Types
using	OutputUpdater
	Optional callback function to setup dense output tensors when initializing the loop emitter (e.g., to fill a dense output with zeros).
using	SynTensorBoundSetter
	Optional callback function to set the bound for the synthetic tensor, which essentially is the dense loop bound.
using	DependentLvlGetter

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.
	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.
Value	genAffine (OpBuilder &builder, Location loc, AffineExpr a)
	Generates code to compute an affine expression whose variables are LoopIds (i.e., cast<AffineDimExpr>(a).getPosition() is a valid LoopId).
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.
void	exitCurrentLoopSeq (OpBuilder &builder, Location loc)
	Exits the current loop sequence, this will reset universal index to 0.
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.
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.
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).
auto	getLoopIVsRange () const
	Get the range of values for all induction variables.
SmallVector< Value >	getLoopIVs () const
	Fills the out-parameter with the loop induction variables for all loops in the current loop-stack.
LoopId	getCurrentDepth () const
	Gets the current depth of the loop-stack.
Value	getLoopIV (LoopId n) const
	Gets loop induction variable for the given loop.
unsigned	getNumManifestTensors () const
	Gets the total number of manifest tensors (excluding the synthetic tensor).
unsigned	getNumTensors () const
	Gets the total number of tensors that loopEmitter is operating on.
TensorId	getSynTensorId () const
	Gets the TensorId for synthetic tensor.
TensorId	getOutTensorId () const
	Gets the TensorId for output tensor.
TensorLevel	makeTensorLevel (TensorId t, Level l) const
	Compresses a TensorId and Level into a TensorLevel.
std::pair< TensorId, Level >	unpackTensorLevel (TensorLevel tidLvl) const
	De-compresses a TensorLevel back to a pair of TensorId and Level.
template<class ContainerTy>
auto	unpackTensorLevelRange (ContainerTy &&c) const
	Converts a range of TensorLevel to a range of std::pair<TensorId, Level>
SmallVector< Value >	getValPosits (TensorId tid) const
	Getters.
Value	getCoord (TensorId tid, Level lvl) const
const std::vector< Value > &	getValBuffer () const

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

Detailed Description

Definition at line 55 of file LoopEmitter.h.

Member Typedef Documentation

◆ DependentLvlGetter

using mlir::sparse_tensor::LoopEmitter::DependentLvlGetter

Initial value:

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

mlir::sparse_tensor::TensorId

unsigned TensorId

Tensor identifiers, chosen to be the BlockArgument::getArgNumber of the value passed to Merger::build...

Definition Merger.h:35

mlir::sparse_tensor::Level

uint64_t Level

The type of level identifiers and level-ranks.

Definition SparseTensor.h:42

mlir::function_ref

llvm::function_ref< Fn > function_ref

Definition LLVM.h:152

Definition at line 77 of file LoopEmitter.h.

◆ OutputUpdater

using mlir::sparse_tensor::LoopEmitter::OutputUpdater

Initial value:

function_ref<Value(OpBuilder &builder, Location loc,

Value memref, Value tensor)>

mlir::Location

This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...

Definition Location.h:76

mlir::OpBuilder

This class helps build Operations.

Definition Builders.h:207

mlir::Value

This class represents an instance of an SSA value in the MLIR system, representing a computable value...

Definition Value.h:96

mlir::memref

Definition Passes.h:26

mlir::tensor

Definition BufferizationTransformOps.h:19

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

Initial value:

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

References mlir::kFunctional.

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

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

Referenced by genCoIteration().

◆ enterCurrentCoIterationCase()

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

Definition at line 539 of file LoopEmitter.cpp.

References mlir::Block::addArguments(), mlir::sparse_tensor::I64BitSet::bits(), mlir::Region::emplaceBlock(), 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 411 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 858 of file LoopEmitter.cpp.

References mlir::Block::back(), mlir::Block::empty(), mlir::Operation::getNumResults(), mlir::Operation::getResults(), mlir::kSparseIterator, mlir::OpBuilder::setInsertionPoint(), mlir::OpBuilder::setInsertionPointAfter(), and mlir::OpBuilder::setInsertionPointToEnd().

Referenced by endLoop().

◆ exitCurrentLoopSeq()

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

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

Definition at line 428 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., cast<AffineDimExpr>(a).getPosition() is a valid LoopId).

Definition at line 439 of file LoopEmitter.cpp.

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

Referenced by genAffine(), genSubscript(), and locateLvlAtAffineAddress().

◆ getCoord()

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

inline

Definition at line 238 of file LoopEmitter.h.

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

◆ getLoopEmitterLoopAttrName()

constexpr 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 genExpand().

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

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

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

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

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

References getNumManifestTensors(), mlir::sparse_tensor::getSparseTensorType(), getSynTensorId(), 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 692 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(), nullptr, 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().

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

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