36 #include "llvm/ADT/SmallBitVector.h"
52 const LoopId i = cast<AffineDimExpr>(a).getPosition();
61 auto binOp = cast<AffineBinaryOpExpr>(a);
66 assert(isa<AffineConstantExpr>(a));
90 if (
auto binOp = dyn_cast<AffineBinaryOpExpr>(a)) {
94 return findAffine(merger, tid, lvl, binOp.getLHS(), lt,
false) &&
95 findAffine(merger, tid, lvl, binOp.getRHS(), lt,
false);
121 int64_t coefficient = 1) {
125 if (coefficient <= 0)
136 assert(coefficient == 1);
168 if (isa<AffineConstantExpr>(a))
169 llvm_unreachable(
"Not yet implemented");
171 auto binOp = cast<AffineBinaryOpExpr>(a);
172 auto lhs = binOp.getLHS(), rhs = binOp.getRHS();
173 if (isa<AffineConstantExpr>(rhs))
176 assert(isa<AffineConstantExpr>(lhs) && isa<AffineDimExpr>(rhs));
177 int64_t coefficient = cast<AffineConstantExpr>(lhs).getValue();
178 return findDepIdxSet(merger, tensor, lvl, rhs, lt, isSubExp, coefficient);
181 auto binOp = cast<AffineBinaryOpExpr>(a);
182 return findDepIdxSet(merger, tensor, lvl, binOp.getLHS(), lt,
true) &&
183 findDepIdxSet(merger, tensor, lvl, binOp.getRHS(), lt,
true);
203 const auto rtp = dyn_cast<RankedTensorType>(tensor.
getType());
210 assert(
static_cast<Dimension>(exprs.size()) == lvlRank &&
211 "AffineMap does not have dimension-rank many results");
213 for (
Level l = 0; l < lvlRank; l++) {
232 OpOperand *out = op.getDpsInitOperand(0);
249 bool annotated =
false;
252 const auto map = env.
op().getMatchingIndexingMap(&t);
256 const Level lvlRank = map.getNumResults();
257 assert(!enc || lvlRank == enc.getLvlRank());
258 assert(
static_cast<Level>(env.
op().getRank(&t)) == lvlRank);
267 for (
Level l = 0; l < lvlRank; l++) {
270 if (idxReducBased && needIdxReduc) {
288 linalg::GenericOp op = env.
op();
293 llvm::cast<linalg::LinalgOp>(op.getOperation())
294 .createLoopRanges(builder, loc);
310 OpOperand *lhs = op.getDpsInitOperand(0);
311 assert(lhs->
get() == tensor);
320 bool isInit = op.isInitTensor(lhs);
331 assert(l < loopRange.size());
338 const auto map = env.
op().getMatchingIndexingMap(t);
340 const Level lvlRank = stt.getLvlRank();
341 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
342 const AffineExpr a = map.getResult(lvlRank - 1);
353 const auto map = env.
op().getMatchingIndexingMap(t);
355 if (stt.hasEncoding()) {
358 assert(!pos.empty());
362 const Level lvlRank = stt.getLvlRank();
363 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
364 for (
Level l = 0; l < lvlRank; l++) {
365 const auto lvlExpr = map.getResult(l);
367 args.push_back(lvlCrd);
376 linalg::GenericOp op = env.
op();
391 linalg::GenericOp op = env.
op();
401 Value isFilled = builder.
create<memref::LoadOp>(loc, filled, index);
402 Value valAtIndex = builder.
create<memref::LoadOp>(loc, values, index);
403 return builder.
create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);
408 scf::IfOp condInsert =
409 builder.
create<scf::IfOp>(loc, sparseOut.
getType(), cond,
true);
412 Value res = builder.
create<tensor::InsertOp>(loc, v, sparseOut, ivs);
413 builder.
create<scf::YieldOp>(loc, res);
416 builder.
create<scf::YieldOp>(loc, sparseOut);
419 return condInsert.getResult(0);
425 linalg::GenericOp op = env.
op();
429 const LoopId numLoops = op.getRank(t);
453 sparseOut = builder.
create<tensor::InsertOp>(loc, rhs, chain, ivs);
474 Value isFilled = builder.
create<memref::LoadOp>(loc, filled, index);
475 Value cond = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
481 builder.
create<memref::StoreOp>(loc, tval, filled, index);
482 builder.
create<memref::StoreOp>(loc, index, added, count);
484 Value add = builder.
create<arith::AddIOp>(loc, count, one);
485 builder.
create<scf::YieldOp>(loc, add);
488 builder.
create<scf::YieldOp>(loc, count);
492 builder.
create<memref::StoreOp>(loc, rhs, values, index);
502 linalg::GenericOp op = env.
op();
507 if (
auto explVal = stt.getExplicitVal())
518 return builder.
create<memref::LoadOp>(loc, ptr, args);
539 linalg::GenericOp op = env.
op();
545 builder.
create<memref::StoreOp>(loc, rhs, ptr, args);
568 builder.
create<scf::YieldOp>(loc, chain);
586 if (
auto arg = dyn_cast<BlockArgument>(e)) {
590 linalg::GenericOp op = env.
op();
591 if (arg.getOwner()->getParentOp() == op) {
597 return rewriter.
create<memref::LoadOp>(op.
getLoc(), ptr, args);
601 if (
auto indexOp = dyn_cast<linalg::IndexOp>(def))
604 if (def->getBlock() == block) {
606 for (
unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
609 i,
relinkBranch(env, rewriter, block, def->getOperand(i)));
622 linalg::GenericOp op = env.
op();
625 const auto kind = exp.
kind;
678 LoopId curr,
bool isStart) {
683 linalg::GenericOp op = env.
op();
685 const auto map = op.getMatchingIndexingMap(&t);
687 const Level lvlRank = stt.getLvlRank();
688 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
689 bool isCurrentLoop = curr == 0;
690 for (
Level l = 0; l < lvlRank; l++) {
702 OpOperand *lhs = op.getDpsInitOperand(0);
749 linalg::GenericOp op = env.
op();
750 OpOperand *lhs = op.getDpsInitOperand(0);
762 auto dynShape = {ShapedType::kDynamic};
763 Type etp = cast<ShapedType>(tensor.
getType()).getElementType();
768 auto r = builder.
create<ExpandOp>(loc,
TypeRange({t1, t2, t3, t4}), tensor);
769 assert(r.getNumResults() == 4);
770 env.
startExpand(r.getResult(0), r.getResult(1), r.getResult(2),
774 for (
LoopId i = 0; i < curr; i++)
781 Value compress = builder.
create<CompressOp>(loc, values, filled, added,
782 count, chain, indices);
803 return isOuter && !isSparse;
811 llvm_unreachable(
"unexpected parallelization strategy");
818 linalg::GenericOp op = env.
op();
819 auto iteratorTypes = op.getIteratorTypesArray();
820 bool isSparse = llvm::any_of(tidLvls, [curr, &env](
TensorLevel tidLvl) {
835 bool tryParallel,
bool needsUniv) {
838 return env.
emitter().enterCoIterationOverTensorsAtLvls(
839 builder, env.
op().getLoc(), tidLvls, reduc, tryParallel, needsUniv);
850 return genCoIteration(env, builder, tidLvls, tryParallel, needsUniv);
859 while (
auto ifOp = dyn_cast_or_null<scf::IfOp>(
884 assert(y == yields.size());
885 builder.
create<scf::YieldOp>(loc, yields);
908 assert(lvl.has_value() &&
isUndefLT(lt));
910 lt = stt.getLvlType(*lvl);
915 assert(lvl.has_value());
918 clause = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
924 cond = cond ? builder.
create<arith::AndIOp>(loc, cond, clause) : clause;
935 scf::IfOp ifOp = builder.
create<scf::IfOp>(loc, types, cond,
true);
950 operands.push_back(
constantI1(builder, env.
op().getLoc(),
true));
962 if (!operands.empty())
963 builder.
create<scf::YieldOp>(env.
op().getLoc(), operands);
974 const BitVector &simple = env.
lat(li).
simple;
976 const std::optional<Level> outLvl = env.
merger().
getLvl(outTid, curr);
978 unsigned numloopCond = 0;
979 bool hasNonUnique =
false;
1008 hasNonUnique = !
isUniqueLT(lt) || hasNonUnique;
1015 linalg::GenericOp op = env.
op();
1016 if (tid >= op.getNumDpsInputs())
1022 if (!stt.hasEncoding())
1026 op.getMatchingIndexingMap(operand).
getResults();
1027 const Level lvlRank = stt.getLvlRank();
1028 assert(affines.size() ==
static_cast<size_t>(lvlRank));
1029 for (
Level l = 0; l < lvlRank; l++) {
1038 if (!isa<AffineConstantExpr>(exp)) {
1039 bool isCurrentLoop =
false;
1061 if (stt.hasEncoding() && stt.isAllDense())
1065 if (numloopCond == 0) {
1076 return numloopCond == 1 &&
1098 if (llvm::find(tidLvls, tl) != tidLvls.end())
1100 tidLvls.emplace_back(tl);
1119 linalg::GenericOp op = env.
op();
1120 assert(tid < op.getNumDpsInputs());
1121 OpOperand *input = op.getDpsInputOperands()[tid];
1122 const auto lvlExprs = op.getMatchingIndexingMap(input).
getResults();
1127 const Level lvlRank = enc.getLvlRank();
1128 assert(lvlExprs.size() ==
static_cast<size_t>(lvlRank));
1129 for (
Level l = startLvl; l < lvlRank; l++) {
1131 if (enc.getLvlType(l).hasDenseSemantic() &&
1132 isa<AffineConstantExpr>(lvlExpr))
1147 for (
TensorId tid = 0, e = env.
op().getNumDpsInputs(); tid < e; tid++)
1159 affineTidLvls.emplace_back(tl, exp);
1161 tidLvls.emplace_back(tl);
1181 for (
auto [tidLvl, exp] : affineTidLvls) {
1189 llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvls));
1196 return std::make_pair(loop, isSingleCond);
1201 LatPointId li,
bool needsUniv,
bool isSingleCond) {
1207 }
else if (
auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
1214 env.
emitter().exitCurrentLoop(rewriter, env.
op().getLoc(), reduc);
1215 return std::nullopt;
1250 bool needsUniv =
startLoopSeq(env, rewriter, exp, curr, lts);
1256 const unsigned lsize = env.
set(lts).size();
1257 for (
unsigned i = 0; i < lsize; i++) {
1260 auto [loop, isSingleCond] =
startLoop(env, rewriter, curr, li, needsUniv);
1271 for (
unsigned j = 0;
j < lsize;
j++) {
1276 if (!isSingleCond) {
1277 scf::IfOp ifOp =
genIf(env, rewriter, curr, lj);
1278 genStmt(env, rewriter, ej, curr + 1);
1279 endIf(env, rewriter, ifOp, redInput, cntInput, insInput, validIns);
1281 genStmt(env, rewriter, ej, curr + 1);
1287 needsUniv =
endLoop(env, rewriter, loop, curr, needsUniv, isSingleCond);
1297 linalg::GenericOp op = env.
op();
1298 OpOperand *lhs = op.getDpsInitOperand(0);
1305 bool hasInserts =
false;
1331 LogicalResult matchAndRewrite(linalg::GenericOp op,
1334 if (op.getNumDpsInits() != 1 || !op.hasPureTensorSemantics())
1344 op,
"Loops not yet scheduled, try run --sparse-reinterpret-map "
1345 "before sparsification.");
1353 const unsigned numLoops = op.getNumLoops();
1361 Level maxLvlRank = 0;
1363 if (
auto rtp = dyn_cast<RankedTensorType>(operand.getType())) {
1379 if (op.getNumReductionLoops() > 0) {
1381 assert(isa<linalg::YieldOp>(yield));
1383 if (!isa<arith::AddFOp>(redop) && !isa<complex::AddOp>(redop) &&
1384 !isa<arith::AddIOp>(redop) && !isa<arith::SubFOp>(redop) &&
1385 !isa<complex::SubOp>(redop) && !isa<arith::SubIOp>(redop) &&
1386 !isa<arith::OrIOp>(redop) && !isa<arith::XOrIOp>(redop) &&
1387 !isa<ReduceOp>(redop)) {
static llvm::ManagedStatic< PassManagerOptions > options
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
static std::pair< Operation *, bool > startLoop(CodegenEnv &env, OpBuilder &builder, LoopId curr, LatPointId li, bool needsUniv)
Starts a single loop in current sequence.
static unsigned getNumNonTrivialIdxExpOnSparseLvls(AffineMap map, Value tensor)
Gets the total number of compound affine expressions in the getMatchingIndexingMap for the given tens...
static Operation * genLoop(CodegenEnv &env, OpBuilder &builder, LoopId curr, bool needsUniv, ArrayRef< TensorLevel > tidLvls)
Generates a for-loop or a while-loop, depending on whether it implements singleton iteration or co-it...
static Value genInsertionLoadReduce(CodegenEnv &env, OpBuilder &builder, OpOperand *t)
Generates insertion code to implement dynamic tensor load for reduction.
static bool isInvariantAffine(AffineExpr a, LoopId curr, bool &isCurrentLoop)
Returns true iff affine expression is invariant.
static bool findDepIdxSet(Merger &merger, TensorId tensor, Level lvl, AffineExpr a, LevelType lt, bool isSubExp=false, int64_t coefficient=1)
Helper method to inspect affine expressions for index variable reduction based codegen.
static scf::IfOp genIf(CodegenEnv &env, OpBuilder &builder, LoopId curr, LatPointId p)
Generates a single if-statement within a while-loop.
static Value genSubscript(CodegenEnv &env, OpBuilder &builder, OpOperand *t, SmallVectorImpl< Value > &args)
Generates subscript for load/store on a dense or sparse tensor.
static Value genConditionalInsert(Location loc, OpBuilder &builder, Value cond, Value sparseOut, ValueRange ivs, Value v)
static void genExpand(CodegenEnv &env, OpBuilder &builder, LoopId curr, bool isStart)
Generates an expanded access pattern in innermost dimension.
static void genConstantDenseAddressFromLevel(CodegenEnv &env, OpBuilder &builder, TensorId tid, Level startLvl)
static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopId curr, LatSetId lts)
Starts a loop sequence at given level.
static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopId curr, bool isStart)
Hoists loop invariant tensor loads for which indices have been exhausted.
static void endLoopSeq(CodegenEnv &env, OpBuilder &builder, unsigned exp, unsigned at)
Ends a loop sequence at given level.
static bool isParallelFor(CodegenEnv &env, bool isOuter, bool isSparse)
Returns parallelization strategy.
static bool findAffine(Merger &merger, TensorId tid, Level lvl, AffineExpr a, LevelType lt, bool setLvlFormat=true)
Helper method to inspect affine expressions.
static Operation * genCoIteration(CodegenEnv &env, OpBuilder &builder, ArrayRef< TensorLevel > tidLvls, bool tryParallel, bool needsUniv)
Emit a loop to coiterate over the list of tensor levels.
static bool findSparseAnnotations(CodegenEnv &env, bool idxReducBased)
Helper method to inspect sparse encodings in the tensor types.
static bool getAllTidLvlsInLatPoints(CodegenEnv &env, LatPointId li, LoopId curr, llvm::function_ref< void(TensorLevel, AffineExpr)> callback)
static Value genInsertionLoad(CodegenEnv &env, OpBuilder &builder, OpOperand *t)
Generates insertion code to implement dynamic tensor load.
static bool hasNonTrivialAffineOnSparseOut(linalg::GenericOp op)
static void endIf(CodegenEnv &env, OpBuilder &builder, scf::IfOp ifOp, Value redInput, Value cntInput, Value insInput, Value validIns)
Generates end of true branch of if-statement within a while-loop.
static void genStmt(CodegenEnv &env, RewriterBase &rewriter, ExprId exp, LoopId curr)
Recursively generates code while computing iteration lattices in order to manage the complexity of im...
static void genInsertionStore(CodegenEnv &env, OpBuilder &builder, OpOperand *t, Value rhs)
Generates insertion code to implement dynamic tensor store.
static void genTensorStore(CodegenEnv &env, OpBuilder &builder, ExprId exp, Value rhs)
Generates a store on a dense or sparse tensor.
static Value relinkBranch(CodegenEnv &env, RewriterBase &rewriter, Block *block, Value e)
Semi-ring branches are simply inlined by the sparsifier.
static void genBuffers(CodegenEnv &env, OpBuilder &builder)
Local bufferization of all dense and sparse data structures.
static void genResult(CodegenEnv &env, RewriterBase &rewriter)
Converts the result computed by the sparse kernel into the required form.
static bool shouldTryParallize(CodegenEnv &env, LoopId curr, ArrayRef< TensorLevel > tidLvls)
Whether or not the current loop being generated should be parallized (if possible) according to the c...
static bool translateBitsToTidLvlPairs(CodegenEnv &env, LatPointId li, LoopId curr, SmallVectorImpl< TensorLevel > &tidLvls, SmallVectorImpl< std::pair< TensorLevel, AffineExpr >> &affineTidLvls)
Returns true if the lattice bit can be iterated by a for loop.
static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e)
Recursively generates tensor expression.
static void genInitConstantDenseAddress(CodegenEnv &env, RewriterBase &rewriter)
static Value genTensorLoad(CodegenEnv &env, OpBuilder &builder, ExprId exp)
Generates a load on a dense or sparse tensor.
static Value genInvariantValue(CodegenEnv &env, ExprId exp)
Generates an invariant value.
static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop, LatPointId li, bool needsUniv, bool isSingleCond)
Ends a single loop in current sequence. Returns new values for needsUniv.
static void finalizeWhileOp(CodegenEnv &env, OpBuilder &builder, bool needsUniv)
Generates the induction structure for a while-loop.
static Value genIndex(CodegenEnv &env, OpOperand *t)
Generates index for load/store on sparse tensor.
Base type for affine expression.
AffineExprKind getKind() const
Return the classification for this type.
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
ArrayRef< AffineExpr > getResults() const
Block represents an ordered list of Operations.
Operation * getTerminator()
Get the terminator operation of this block.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
IRValueT get() const
Return the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
void setInsertionPointAfter(Operation *op)
Sets the insertion point to the node after the specified operation, which will cause subsequent inser...
Block * getInsertionBlock() const
Return the block the current insertion point belongs to.
This class represents an operand of an operation.
unsigned getOperandNumber()
Return which operand this is in the OpOperand list of the Operation.
Operation is the basic unit of execution within MLIR.
Value getOperand(unsigned idx)
OpOperand & getOpOperand(unsigned idx)
bool hasAttr(StringAttr name)
Return true if the operation has an attribute with the provided name, false otherwise.
Location getLoc()
The source location the operation was defined or derived from.
unsigned getNumOperands()
Region & getRegion(unsigned index)
Returns the region held by this operation at position 'index'.
MutableArrayRef< OpOperand > getOpOperands()
operand_range getOperands()
Returns an iterator on the underlying Value's.
result_range getResults()
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
MLIRContext * getContext() const
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&...args)
Add an instance of each of the pattern types 'Ts' to the pattern list with the given arguments.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
std::enable_if_t<!std::is_convertible< CallbackT, Twine >::value, LogicalResult > notifyMatchFailure(Location loc, CallbackT &&reasonCallback)
Used to notify the listener that the IR failed to be rewritten because of a match failure,...
void modifyOpInPlace(Operation *root, CallableT &&callable)
This method is a utility wrapper around an in-place modification of an operation.
OpTy replaceOpWithNewOp(Operation *op, Args &&...args)
Replace the results of the given (original) op with a new op that is created without verification (re...
This class provides an abstraction over the various different ranges of value types.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
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...
Type getType() const
Return the type of this value.
Block * getParentBlock()
Return the Block in which this Value is defined.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
The code generation environment class aggregates a number of data structures that are needed during t...
void startReduc(ExprId exp, Value val)
void updateValidLexInsert(Value val)
const SparsificationOptions & options() const
Value getInsertionChain() const
std::optional< Operation * > genLoopBoundary(function_ref< std::optional< Operation * >(MutableArrayRef< Value > parameters)> callback)
Generates loop boundary statements (entering/exiting loops).
ArrayRef< LatPointId > set(LatSetId s) const
bool atExpandLevel(OpOperand *o, unsigned rank, LoopId n) const
bool isCustomReduc() const
unsigned getCurrentDepth() const
std::pair< TensorId, Level > unpackTensorLevel(TensorLevel tl) const
Value getExpandValues() const
TensorLevel makeTensorLevel(TensorId t, Level l) const
const LatPoint & lat(LatPointId l) const
constexpr TensorId makeTensorId(unsigned t) const
void startExpand(Value values, Value filled, Value added, Value count)
bool hasSparseOutput() const
unsigned getLoopNum() const
void updateInsertionChain(Value chain)
Value getExpandCount() const
void startCustomReduc(ExprId exp)
linalg::GenericOp op() const
Value getLoopVar(LoopId i) const
Returns the induction-variable for the given loop.
Value getExpandFilled() const
LogicalResult initTensorExp()
void startEmit(SparseEmitStrategy emitStrategy)
auto unpackTensorLevelRange(ContainerTy &&c) const
Value getExpandAdded() const
const TensorExp & exp(ExprId e) const
void updateExpandCount(Value count)
void updateReduc(Value val)
Value getValidLexInsert() const
bool isSparseOutput(OpOperand *o) const
void startValidLexInsert(Value val)
constexpr LoopId makeLoopId(unsigned i) const
Value getCustomRedId() const
LevelType lt(TensorId t, LoopId i) const
bool isValidLexInsert() const
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....
Value getLoopIV(LoopId n) const
Gets loop induction variable for the given loop.
SmallVector< Value > getValPosits(TensorId tid) const
Getters.
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.
void exitCurrentLoopSeq(OpBuilder &builder, Location loc)
Exits the current loop sequence, this will reset universal index to 0.
Value getCoord(TensorId tid, Level lvl) const
A class to handle all iteration lattice operations.
std::optional< Level > getLvl(TensorId t, LoopId i) const
Gets the level number of the the tth tensor on ith loop.
LatSetId buildLattices(ExprId e, LoopId i)
Builds the iteration lattices in a bottom-up traversal given the remaining tensor (sub)expression and...
constexpr LoopId makeLoopId(unsigned i) const
Safely converts the argument to a loop identifier.
void setLevelAndType(TensorId t, LoopId i, Level lvl, LevelType lt)
Sets the level number and level-type of the tth tensor on ith loop.
void foreachTensorLoopId(LatPointId p, ForeachTensorLoopIdCallback callback) const
Iterates over a set of TensorLoopIds, invoking the callback for each TensorLoopId and passing it the ...
LatSetId optimizeSet(LatSetId s)
Optimizes the iteration lattice points in the given set.
void setLoopDependentTensorLevel(LoopId i, TensorId t, Level lvl, LevelType lt, unsigned coefficient)
Establishes the two-way map that i <-> <t, lvl, lt>.
bool hasAnySparse(const BitVector &bits) const
Returns true if any TensorLoopId in the bitvector corresponds to sparse level-type.
void clearExprValue(ExprId e)
Clears the value associated with the expression.
constexpr TensorId getSynTensorID() const
Gets the synthetic tensor's identifier (used for all invariant tensor expressions).
bool latGT(LatPointId p0, LatPointId p1) const
Returns true if p0 > p1.
constexpr LoopId loop(TensorLoopId b) const
Gets the loop-identifier of the TensorLoopId.
constexpr TensorId getOutTensorID() const
Gets the output tensor's identifier.
LevelType getLvlType(TensorId t, LoopId i) const
Gets the level-type of the tth tensor on ith loop.
Value buildExp(RewriterBase &rewriter, Location loc, ExprId e, Value v0, Value v1) const
Rebuilds SSA format from a tensor expression.
void setExprValue(ExprId e, Value v)
Sets the expression to have the associated value.
bool hasDependentLvl(LoopId i, TensorId t)
Whether the loop has dependent slice.
A wrapper around RankedTensorType, which has three goals:
bool hasEncoding() const
Returns true for tensors which have an encoding, and false for those which do not.
bool isAllDense() const
Returns true for tensors where every level is dense.
Level getLvlRank() const
Returns the level-rank.
LevelType getLvlType(Level l) const
static constexpr unsigned kInvalidId
A constant serving as the canonically invalid identifier, regardless of the identifier type.
bool isUniqueLT(LevelType lt)
Value constantIndex(OpBuilder &builder, Location loc, int64_t i)
Generates a constant of index type.
Value constantZero(OpBuilder &builder, Location loc, Type tp)
Generates a 0-valued constant of the given type.
unsigned LatSetId
LatSet identifiers.
uint64_t Dimension
The type of dimension identifiers and dimension-ranks.
unsigned TensorLoopId
A compressed representation of std::pair<TensorId, LoopId>.
uint64_t Level
The type of level identifiers and level-ranks.
unsigned LoopId
Loop identifiers.
Value constantI1(OpBuilder &builder, Location loc, bool b)
Generates a constant of i1 type.
Value genValFromAttr(OpBuilder &builder, Location loc, Attribute attr)
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
bool hasAnySparseType(TypeRange types)
Returns true iff the type range has any sparse tensor type.
Value genIsNonzero(OpBuilder &builder, Location loc, Value v)
Generates the comparison v != 0 where v is of numeric type.
bool isUndefLT(LevelType lt)
bool isDenseLT(LevelType lt)
bool hasAnyNonIdentityOperandsOrResults(Operation *op)
Returns true iff MLIR operation has any sparse tensor with non-identity dim2lvl maps.
SparseTensorType getSparseTensorType(Value val)
Convenience methods to obtain a SparseTensorType from a Value.
unsigned ExprId
TensorExp identifiers.
unsigned LatPointId
LatPoint identifiers.
unsigned TensorId
Tensor identifiers, chosen to be the BlockArgument::getArgNumber of the value passed to Merger::build...
Include the generated interface declarations.
@ Mul
RHS of mul is always a constant or a symbolic expression.
@ DimId
Dimensional identifier.
@ Constant
Constant integer.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
void populateSparsificationPatterns(RewritePatternSet &patterns, const SparsificationOptions &options=SparsificationOptions())
Sets up sparsification rewriting rules with the given options.
Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc, OpFoldResult ofr)
Converts an OpFoldResult to a Value.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Options for the Sparsification pass.
SparseEmitStrategy sparseEmitStrategy
SparseParallelizationStrategy parallelizationStrategy
ExprId exp
Identifier of the tensor expression.
BitVector simple
Simplified conjunction of TensorLoopId as bitvector.
This enum defines all the sparse representations supportable by the SparseTensor dialect.
constexpr bool hasSparseSemantic() const
Check if the LevelType is considered to be sparse.
constexpr bool hasDenseSemantic() const
Check if the LevelType is considered to be dense-like.
Tensor expression. Represents an MLIR expression in tensor index notation.
LoopId loop
kLoopVar expressions simply have a loop identifier.
Value val
Direct link to IR for an invariant or the destination value (to infer destination type) of a cast ope...
Children children
All other expressions hold the ExprIds of their children.
TensorId tensor
kTensor expressions simply have a tensor identifier.
Kind kind
Tensor expression kind.
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.