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());
365 const Level lvlRank = stt.getLvlRank();
366 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
367 for (
Level l = 0; l < lvlRank; l++) {
368 const auto lvlExpr = map.getResult(l);
370 args.push_back(lvlCrd);
379 linalg::GenericOp op = env.
op();
394 linalg::GenericOp op = env.
op();
404 Value isFilled = builder.
create<memref::LoadOp>(loc, filled, index);
405 Value valAtIndex = builder.
create<memref::LoadOp>(loc, values, index);
406 return builder.
create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);
411 scf::IfOp condInsert =
412 builder.
create<scf::IfOp>(loc, sparseOut.
getType(), cond,
true);
415 Value res = builder.
create<tensor::InsertOp>(loc, v, sparseOut, ivs);
416 builder.
create<scf::YieldOp>(loc, res);
419 builder.
create<scf::YieldOp>(loc, sparseOut);
422 return condInsert.getResult(0);
428 linalg::GenericOp op = env.
op();
432 const LoopId numLoops = op.getRank(t);
456 sparseOut = builder.
create<tensor::InsertOp>(loc, rhs, chain, ivs);
477 Value isFilled = builder.
create<memref::LoadOp>(loc, filled, index);
478 Value cond = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
484 builder.
create<memref::StoreOp>(loc, tval, filled, index);
485 builder.
create<memref::StoreOp>(loc, index, added, count);
487 Value add = builder.
create<arith::AddIOp>(loc, count, one);
488 builder.
create<scf::YieldOp>(loc, add);
491 builder.
create<scf::YieldOp>(loc, count);
495 builder.
create<memref::StoreOp>(loc, rhs, values, index);
505 linalg::GenericOp op = env.
op();
510 if (
auto explVal = stt.getExplicitVal())
522 if (llvm::isa<TensorType>(ptr.
getType())) {
526 return builder.
create<ExtractValOp>(loc, ptr, args.front());
528 return builder.
create<memref::LoadOp>(loc, ptr, args);
549 linalg::GenericOp op = env.
op();
555 builder.
create<memref::StoreOp>(loc, rhs, ptr, args);
578 builder.
create<scf::YieldOp>(loc, chain);
596 if (
auto arg = dyn_cast<BlockArgument>(e)) {
600 linalg::GenericOp op = env.
op();
601 if (arg.getOwner()->getParentOp() == op) {
607 return rewriter.
create<memref::LoadOp>(op.
getLoc(), ptr, args);
611 if (
auto indexOp = dyn_cast<linalg::IndexOp>(def))
614 if (def->getBlock() == block) {
616 for (
unsigned i = 0, n = def->getNumOperands(); i < n; i++) {
619 i,
relinkBranch(env, rewriter, block, def->getOperand(i)));
632 linalg::GenericOp op = env.
op();
635 const auto kind = exp.
kind;
688 LoopId curr,
bool isStart) {
693 linalg::GenericOp op = env.
op();
695 const auto map = op.getMatchingIndexingMap(&t);
697 const Level lvlRank = stt.getLvlRank();
698 assert(
static_cast<Level>(map.getNumResults()) == lvlRank);
699 bool isCurrentLoop = curr == 0;
700 for (
Level l = 0; l < lvlRank; l++) {
712 OpOperand *lhs = op.getDpsInitOperand(0);
759 linalg::GenericOp op = env.
op();
760 OpOperand *lhs = op.getDpsInitOperand(0);
772 auto dynShape = {ShapedType::kDynamic};
773 Type etp = cast<ShapedType>(tensor.
getType()).getElementType();
778 auto r = builder.
create<ExpandOp>(loc,
TypeRange({t1, t2, t3, t4}), tensor);
779 assert(r.getNumResults() == 4);
780 env.
startExpand(r.getResult(0), r.getResult(1), r.getResult(2),
784 for (
LoopId i = 0; i < curr; i++)
791 Value compress = builder.
create<CompressOp>(loc, values, filled, added,
792 count, chain, indices);
813 return isOuter && !isSparse;
821 llvm_unreachable(
"unexpected parallelization strategy");
828 linalg::GenericOp op = env.
op();
829 auto iteratorTypes = op.getIteratorTypesArray();
830 bool isSparse = llvm::any_of(tidLvls, [curr, &env](
TensorLevel tidLvl) {
845 unsigned numCases,
bool tryParallel,
849 return env.
emitter().enterCoIterationOverTensorsAtLvls(
850 builder, env.
op().getLoc(), tidLvls, numCases, reduc, tryParallel,
860 unsigned numCases,
bool needsUniv,
863 return genCoIteration(env, builder, tidLvls, numCases, tryParallel,
873 while (
auto ifOp = dyn_cast_or_null<scf::IfOp>(
898 assert(y == yields.size());
899 builder.
create<scf::YieldOp>(loc, yields);
912 assert(allCase == curCase || env.
merger().
latGT(allCase, curCase));
920 if (curCaseBits.test(set))
942 assert(lvl.has_value() &&
isUndefLT(lt));
944 lt = stt.getLvlType(*lvl);
949 assert(lvl.has_value());
952 clause = builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
958 cond = cond ? builder.
create<arith::AndIOp>(loc, cond, clause) : clause;
969 scf::IfOp ifOp = builder.
create<scf::IfOp>(loc, types, cond,
true);
984 operands.push_back(
constantI1(builder, env.
op().getLoc(),
true));
996 if (!operands.empty())
997 builder.
create<scf::YieldOp>(env.
op().getLoc(), operands);
1008 const BitVector &simple = env.
lat(li).
simple;
1010 const std::optional<Level> outLvl = env.
merger().
getLvl(outTid, curr);
1012 unsigned numloopCond = 0;
1013 bool hasNonUnique =
false;
1042 hasNonUnique = !
isUniqueLT(lt) || hasNonUnique;
1049 linalg::GenericOp op = env.
op();
1050 if (tid >= op.getNumDpsInputs())
1056 if (!stt.hasEncoding())
1060 op.getMatchingIndexingMap(operand).
getResults();
1061 const Level lvlRank = stt.getLvlRank();
1062 assert(affines.size() ==
static_cast<size_t>(lvlRank));
1063 for (
Level l = 0; l < lvlRank; l++) {
1072 if (!isa<AffineConstantExpr>(exp)) {
1073 bool isCurrentLoop =
false;
1095 if (stt.hasEncoding() && stt.isAllDense())
1099 if (numloopCond == 0) {
1110 return numloopCond == 1 &&
1132 if (llvm::find(tidLvls, tl) != tidLvls.end())
1134 tidLvls.emplace_back(tl);
1153 linalg::GenericOp op = env.
op();
1154 assert(tid < op.getNumDpsInputs());
1155 OpOperand *input = op.getDpsInputOperands()[tid];
1156 const auto lvlExprs = op.getMatchingIndexingMap(input).
getResults();
1161 const Level lvlRank = enc.getLvlRank();
1162 assert(lvlExprs.size() ==
static_cast<size_t>(lvlRank));
1163 for (
Level l = startLvl; l < lvlRank; l++) {
1165 if (enc.getLvlType(l).hasDenseSemantic() &&
1166 isa<AffineConstantExpr>(lvlExpr))
1181 for (
TensorId tid = 0, e = env.
op().getNumDpsInputs(); tid < e; tid++)
1193 affineTidLvls.emplace_back(tl, exp);
1195 tidLvls.emplace_back(tl);
1216 Operation *loop =
genLoop(env, builder, curr, numCases, needsUniv, tidLvls);
1218 for (
auto [tidLvl, exp] : affineTidLvls) {
1226 llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvls));
1233 return std::make_pair(loop, isSingleCond);
1238 LatPointId li,
bool needsUniv,
bool isSingleCond) {
1244 }
else if (
auto whileOp = dyn_cast<scf::WhileOp>(loop)) {
1251 env.
emitter().exitCurrentLoop(rewriter, env.
op().getLoc(), reduc);
1252 return std::nullopt;
1287 bool needsUniv =
startLoopSeq(env, rewriter, exp, curr, lts);
1291 const unsigned lsize = env.
set(lts).size();
1295 auto [loop, isSingleCond] =
1296 startLoop(env, rewriter, curr, li, lsize, needsUniv);
1297 assert(isSingleCond == llvm::isa<IterateOp>(loop));
1301 for (
unsigned j = 0;
j < lsize;
j++) {
1305 if (!isSingleCond) {
1308 genStmt(env, rewriter, ej, curr + 1);
1310 assert(reduc.empty() &&
"Not Implemented");
1311 rewriter.
create<sparse_tensor::YieldOp>(env.
op().getLoc());
1312 return std::nullopt;
1316 genStmt(env, rewriter, ej, curr + 1);
1320 needsUniv =
endLoop(env, rewriter, loop, curr, needsUniv, isSingleCond);
1323 for (
unsigned i = 0; i < lsize; i++) {
1326 auto [loop, isSingleCond] =
1327 startLoop(env, rewriter, curr, li, lsize, needsUniv);
1338 for (
unsigned j = 0;
j < lsize;
j++) {
1343 if (!isSingleCond) {
1344 scf::IfOp ifOp =
genIf(env, rewriter, curr, lj);
1345 genStmt(env, rewriter, ej, curr + 1);
1346 endIf(env, rewriter, ifOp, redInput, cntInput, insInput, validIns);
1348 genStmt(env, rewriter, ej, curr + 1);
1354 needsUniv =
endLoop(env, rewriter, loop, curr, needsUniv, isSingleCond);
1365 linalg::GenericOp op = env.
op();
1366 OpOperand *lhs = op.getDpsInitOperand(0);
1373 bool hasInserts =
false;
1399 LogicalResult matchAndRewrite(linalg::GenericOp op,
1402 if (op.getNumDpsInits() != 1 || !op.hasPureTensorSemantics())
1412 op,
"Loops not yet scheduled, try run --sparse-reinterpret-map "
1413 "before sparsification.");
1421 const unsigned numLoops = op.getNumLoops();
1429 Level maxLvlRank = 0;
1431 if (
auto rtp = dyn_cast<RankedTensorType>(operand.getType())) {
1447 if (op.getNumReductionLoops() > 0) {
1449 assert(isa<linalg::YieldOp>(yield));
1451 if (!isa<arith::AddFOp>(redop) && !isa<complex::AddOp>(redop) &&
1452 !isa<arith::AddIOp>(redop) && !isa<arith::SubFOp>(redop) &&
1453 !isa<complex::SubOp>(redop) && !isa<arith::SubIOp>(redop) &&
1454 !isa<arith::OrIOp>(redop) && !isa<arith::XOrIOp>(redop) &&
1455 !isa<ReduceOp>(redop)) {
static llvm::ManagedStatic< PassManagerOptions > options
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
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, unsigned numCases, 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 std::pair< Operation *, bool > startLoop(CodegenEnv &env, OpBuilder &builder, LoopId curr, LatPointId li, unsigned numCases, bool needsUniv)
Starts a single loop in current sequence.
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 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 void genCoIterationCase(CodegenEnv &env, OpBuilder &builder, unsigned caseIdx, LatPointId allCase, LatPointId curCase, MutableArrayRef< Value > reduc)
Generates a case region in the coiterate operation.
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)
bool generatingSparseIterator() const
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
A simple wrapper to encode a bitset of (at most 64) levels, currently used by sparse_tensor....
I64BitSet & set(unsigned i)
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....
Region * enterCurrentCoIterationCase(OpBuilder &builder, Location loc, I64BitSet caseBit, unsigned caseIdx, MutableArrayRef< Value > reduc)
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.
const LatPoint & lat(LatPointId p) const
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
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
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)
std::pair< Operation *, Value > genCoIteration(OpBuilder &builder, Location loc, ArrayRef< SparseIterator * > iters, MutableArrayRef< Value > reduc, Value uniIdx, bool userReducFirst=false)
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.