9 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_LOOPEMITTER_H_
10 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_UTILS_LOOPEMITTER_H_
23 namespace sparse_tensor {
90 bool hasOutput =
false,
bool isSparseOut =
false,
95 ValueRange tensors, StringAttr loopTag =
nullptr,
bool hasOutput =
false,
96 bool isSparseOut =
false,
unsigned numLoops = 0,
149 bool needsUniv =
false);
158 return llvm::map_range(loopStack, [](
const LoopInfo &li) {
return li.iv; });
206 return std::make_pair(tidLvl % nt, tidLvl / nt);
210 template <
class ContainerTy>
212 using EltTy = decltype(*c.begin());
213 static_assert(std::is_same_v<llvm::remove_cvref_t<EltTy>,
TensorLevel>,
214 "Must be unpacking a TensorLevel range");
215 return llvm::map_range(std::forward<ContainerTy>(c), [
this](EltTy tl) {
225 Value lastLvlPos = iters[tid].back().back()->getCurPosition().first;
226 batchCrds.push_back(lastLvlPos);
230 return getCurIterator(tid, lvl).
getCrd();
235 return llvm::StringLiteral(
"Emitted from");
245 struct LoopInfo final {
247 Value iv, StringAttr loopTag)
248 : tidLvls(tidLvls), loop(loop), userCodeBlock(userBlock), iv(iv) {
259 Block *
const userCodeBlock;
263 void categorizeIterators(ArrayRef<TensorLevel> tidLvls,
264 SmallVectorImpl<SparseIterator *> &raIters,
265 SmallVectorImpl<SparseIterator *> &spIters);
271 MutableArrayRef<Value>)>;
274 bool shouldIteratedByForLoop(ArrayRef<SparseIterator *> spIters);
280 Value genSparseCrd(OpBuilder &builder, Location loc,
TensorId tid,
285 bool isOutputTensor(
TensorId tid)
const {
289 bool isSparseOutput(
TensorId tid)
const {
290 return isOutputTensor(tid) && isSparseOut;
294 return tid < lvls.size() && lvl < lvls[tid].size();
299 void prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
308 std::pair<Operation *, Value>
309 emitForLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
310 SparseIterator &iter, MutableArrayRef<Value> reduc,
319 std::pair<Operation *, Value>
320 emitWhileLoopOverTensorsAtLvls(OpBuilder &builder, Location loc,
321 ArrayRef<SparseIterator *> iters,
322 MutableArrayRef<Value> reduc,
bool needsUniv);
348 void exitForLoop(RewriterBase &rewriter, Location loc,
349 MutableArrayRef<Value> reduc);
352 void exitWhileLoop(OpBuilder &builder, Location loc,
353 MutableArrayRef<Value> reduc);
359 void initSubSectIterator(OpBuilder &builder, Location loc);
363 return levelReducedDep[tid][lvl];
366 SparseIterator &getCurIterator(
TensorId tid,
Level lvl)
const {
367 if (dependentLvlMap[tid][lvl].empty())
368 return *iters[tid][lvl].back();
370 assert(redDepOnLevel(tid, lvl) >= 1);
371 return *iters[tid][lvl][redDepOnLevel(tid, lvl) - 1];
374 std::unique_ptr<SparseIterator>
375 makeLevelIterator(OpBuilder &builder, Location loc,
TensorId tid,
Level l);
392 std::vector<Value> tensors;
393 std::vector<Value> loopHighs;
394 std::vector<std::vector<std::unique_ptr<SparseTensorLevel>>> lvls;
395 std::vector<std::vector<std::vector<std::unique_ptr<SparseIterator>>>> iters;
396 std::vector<Value> valBuffer;
400 std::vector<std::vector<std::vector<std::pair<LoopId, unsigned>>>>
405 std::vector<std::vector<std::vector<std::pair<Value, unsigned>>>> sliceMeta;
408 std::vector<std::vector<unsigned>> levelReducedDep;
416 std::vector<LoopInfo> loopStack;
420 std::vector<std::pair<Value, std::vector<TensorLevel>>> loopSeqStack;
Base type for affine expression.
Block represents an ordered list of Operations.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
This class helps build Operations.
Operation is the basic unit of execution within MLIR.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
This class provides an abstraction over the different types of ranges over Values.
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
void exitCurrentLoop(RewriterBase &rewriter, Location loc, MutableArrayRef< Value > reduc={})
Generates code to exit the current loop (e.g., generates yields, forwards loop induction variables,...
constexpr static llvm::StringLiteral getLoopEmitterLoopAttrName()
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.
const std::vector< Value > & getValBuffer() const
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 previo...
Value genAffine(OpBuilder &builder, Location loc, AffineExpr a)
Generates code to compute an affine expression whose variables are LoopIds (i.e., a....
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.
TensorId getOutTensorId() const
Gets the TensorId for output tensor.
TensorLevel makeTensorLevel(TensorId t, Level l) const
Compresses a TensorId and Level into a TensorLevel.
unsigned getNumManifestTensors() const
Gets the total number of manifest tensors (excluding the synthetic tensor).
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.
Value getLoopIV(LoopId n) const
Gets loop induction variable for the given loop.
std::pair< TensorId, Level > unpackTensorLevel(TensorLevel tidLvl) const
De-compresses a TensorLevel back to a pair of TensorId and Level.
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.
unsigned getNumTensors() const
Gets the total number of tensors that loopEmitter is operating on.
SmallVector< Value > getLoopIVs() const
Fills the out-parameter with the loop induction variables for all loops in the current loop-stack.
auto getLoopIVsRange() const
Get the range of values for all induction variables.
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.
LoopId getCurrentDepth() const
Gets the current depth of the loop-stack.
void exitCurrentLoopSeq(OpBuilder &builder, Location loc)
Exits the current loop sequence, this will reset universal index to 0.
TensorId getSynTensorId() const
Gets the TensorId for synthetic tensor.
Value getCoord(TensorId tid, Level lvl) const
uint64_t Level
The type of level identifiers and level-ranks.
unsigned LoopId
Loop identifiers.
unsigned TensorId
Tensor identifiers, chosen to be the BlockArgument::getArgNumber of the value passed to Merger::build...
Include the generated interface declarations.
SparseEmitStrategy
Defines a scope for reinterpret map pass.