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);
409 linalg::GenericOp op = env.
op();
413 const LoopId numLoops = op.getRank(t);
426 scf::IfOp ifValidLexInsert = builder.
create<scf::IfOp>(
431 Value res = builder.
create<tensor::InsertOp>(loc, rhs, chain, ivs);
432 builder.
create<scf::YieldOp>(loc, res);
435 builder.
create<scf::YieldOp>(loc, chain);
442 builder.
create<tensor::InsertOp>(loc, rhs, chain, ivs));
460 Value isFilled = builder.
create<memref::LoadOp>(loc, filled, index);
461 Value cond = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
467 builder.
create<memref::StoreOp>(loc, tval, filled, index);
468 builder.
create<memref::StoreOp>(loc, index, added, count);
470 Value add = builder.
create<arith::AddIOp>(loc, count, one);
471 builder.
create<scf::YieldOp>(loc, add);
474 builder.
create<scf::YieldOp>(loc, count);
478 builder.
create<memref::StoreOp>(loc, rhs, values, index);
488 linalg::GenericOp op = env.
op();
498 return builder.
create<memref::LoadOp>(op.
getLoc(), ptr, args);
519 linalg::GenericOp op = env.
op();
525 builder.
create<memref::StoreOp>(loc, rhs, ptr, args);
548 builder.
create<scf::YieldOp>(loc, chain);
566 if (
auto arg = dyn_cast<BlockArgument>(e)) {
570 linalg::GenericOp op = env.
op();
571 if (arg.getOwner()->getParentOp() == op) {
577 return rewriter.
create<memref::LoadOp>(op.
getLoc(), ptr, args);
581 if (
auto indexOp = dyn_cast<linalg::IndexOp>(def))
584 if (def->getBlock() == block) {
586 for (
unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
589 i,
relinkBranch(env, rewriter, block, def->getOperand(i)));
602 linalg::GenericOp op = env.
op();
605 const auto kind = exp.
kind;
658 LoopId curr,
bool isStart) {
663 linalg::GenericOp op = env.
op();
665 const auto map = op.getMatchingIndexingMap(&t);
667 const Level lvlRank = stt.getLvlRank();
668 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
669 bool isCurrentLoop = curr == 0;
670 for (
Level l = 0; l < lvlRank; l++) {
682 OpOperand *lhs = op.getDpsInitOperand(0);
729 linalg::GenericOp op = env.
op();
730 OpOperand *lhs = op.getDpsInitOperand(0);
742 auto dynShape = {ShapedType::kDynamic};
743 Type etp = cast<ShapedType>(tensor.
getType()).getElementType();
748 auto r = builder.
create<ExpandOp>(loc,
TypeRange({t1, t2, t3, t4}), tensor);
749 assert(r.getNumResults() == 4);
750 env.
startExpand(r.getResult(0), r.getResult(1), r.getResult(2),
754 for (
LoopId i = 0; i < curr; i++)
761 Value compress = builder.
create<CompressOp>(loc, values, filled, added,
762 count, chain, indices);
783 return isOuter && !isSparse;
791 llvm_unreachable(
"unexpected parallelization strategy");
798 linalg::GenericOp op = env.
op();
799 auto iteratorTypes = op.getIteratorTypesArray();
800 bool isSparse = llvm::any_of(tidLvls, [curr, &env](
TensorLevel tidLvl) {
815 bool tryParallel,
bool needsUniv) {
818 return env.
emitter().enterCoIterationOverTensorsAtLvls(
819 builder, env.
op().getLoc(), tidLvls, reduc, tryParallel, needsUniv);
830 return genCoIteration(env, builder, tidLvls, tryParallel, needsUniv);
839 while (
auto ifOp = dyn_cast_or_null<scf::IfOp>(
864 assert(y == yields.size());
865 builder.
create<scf::YieldOp>(loc, yields);
888 assert(lvl.has_value() &&
isUndefLT(lt));
890 lt = stt.getLvlType(*lvl);
895 assert(lvl.has_value());
898 clause = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
904 cond = cond ? builder.
create<arith::AndIOp>(loc, cond, clause) : clause;
915 scf::IfOp ifOp = builder.
create<scf::IfOp>(loc, types, cond,
true);
930 operands.push_back(
constantI1(builder, env.
op().getLoc(),
true));
942 if (!operands.empty())
943 builder.
create<scf::YieldOp>(env.
op().getLoc(), operands);
954 const BitVector &simple = env.
lat(li).
simple;
956 const std::optional<Level> outLvl = env.
merger().
getLvl(outTid, curr);
958 unsigned numloopCond = 0;
959 bool hasNonUnique =
false;
988 hasNonUnique = !
isUniqueLT(lt) || hasNonUnique;
995 linalg::GenericOp op = env.
op();
996 if (tid >= op.getNumDpsInputs())
1002 if (!stt.hasEncoding())
1006 op.getMatchingIndexingMap(operand).
getResults();
1007 const Level lvlRank = stt.getLvlRank();
1008 assert(affines.size() ==
static_cast<size_t>(lvlRank));
1009 for (
Level l = 0; l < lvlRank; l++) {
1018 if (!isa<AffineConstantExpr>(exp)) {
1019 bool isCurrentLoop =
false;
1041 if (stt.hasEncoding() && stt.isAllDense())
1045 if (numloopCond == 0) {
1054 return numloopCond == 1 && !hasNonUnique;
1074 if (llvm::find(tidLvls, tl) != tidLvls.end())
1076 tidLvls.emplace_back(tl);
1095 linalg::GenericOp op = env.
op();
1096 assert(tid < op.getNumDpsInputs());
1097 OpOperand *input = op.getDpsInputOperands()[tid];
1098 const auto lvlExprs = op.getMatchingIndexingMap(input).
getResults();
1103 const Level lvlRank = enc.getLvlRank();
1104 assert(lvlExprs.size() ==
static_cast<size_t>(lvlRank));
1105 for (
Level l = startLvl; l < lvlRank; l++) {
1107 if (enc.getLvlType(l).hasDenseSemantic() &&
1108 isa<AffineConstantExpr>(lvlExpr))
1123 for (
TensorId tid = 0, e = env.
op().getNumDpsInputs(); tid < e; tid++)
1135 affineTidLvls.emplace_back(tl, exp);
1137 tidLvls.emplace_back(tl);
1157 for (
auto [tidLvl, exp] : affineTidLvls) {
1165 llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvls));
1172 return std::make_pair(loop, isSingleCond);
1177 LatPointId li,
bool needsUniv,
bool isSingleCond) {
1183 }
else if (
auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
1190 env.
emitter().exitCurrentLoop(rewriter, env.
op().getLoc(), reduc);
1191 return std::nullopt;
1226 bool needsUniv =
startLoopSeq(env, rewriter, exp, curr, lts);
1232 const unsigned lsize = env.
set(lts).size();
1233 for (
unsigned i = 0; i < lsize; i++) {
1236 auto [loop, isSingleCond] =
startLoop(env, rewriter, curr, li, needsUniv);
1247 for (
unsigned j = 0;
j < lsize;
j++) {
1252 if (!isSingleCond) {
1253 scf::IfOp ifOp =
genIf(env, rewriter, curr, lj);
1254 genStmt(env, rewriter, ej, curr + 1);
1255 endIf(env, rewriter, ifOp, redInput, cntInput, insInput, validIns);
1257 genStmt(env, rewriter, ej, curr + 1);
1263 needsUniv =
endLoop(env, rewriter, loop, curr, needsUniv, isSingleCond);
1273 linalg::GenericOp op = env.
op();
1274 OpOperand *lhs = op.getDpsInitOperand(0);
1281 bool hasInserts =
false;
1310 if (op.getNumDpsInits() != 1 || !op.hasPureTensorSemantics())
1320 op,
"Loops not yet scheduled, try run --sparse-reinterpret-map "
1321 "before sparsification.");
1329 const unsigned numLoops = op.getNumLoops();
1337 Level maxLvlRank = 0;
1339 if (
auto rtp = dyn_cast<RankedTensorType>(operand.getType())) {
1355 if (op.getNumReductionLoops() > 0) {
1357 assert(isa<linalg::YieldOp>(yield));
1359 if (!isa<arith::AddFOp>(redop) && !isa<complex::AddOp>(redop) &&
1360 !isa<arith::AddIOp>(redop) && !isa<arith::SubFOp>(redop) &&
1361 !isa<complex::SubOp>(redop) && !isa<arith::SubIOp>(redop) &&
1362 !isa<arith::OrIOp>(redop) && !isa<arith::XOrIOp>(redop) &&
1363 !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 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.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a 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.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
@ 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...
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
This class represents an efficient way to signal success or failure.
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Options for the Sparsification pass.
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.