MLIR  19.0.0git
Classes | Public Types | Public Member Functions | Static Public Member Functions | List of all members
mlir::sparse_tensor::LoopEmitter Class Reference

#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 LoopIds (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...
 
OperationenterCoIterationOverTensorsAtLvls (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< ValuegetLoopIVs () 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, LevelunpackTensorLevel (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< ValuegetValPosits (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

Definition at line 77 of file LoopEmitter.h.

◆ OutputUpdater

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

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,
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 597 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 395 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 809 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 410 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 421 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 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 
)

◆ locateLvlAtAffineAddress()

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

◆ 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: