MLIR  18.0.0git
Functions
Sparsification.cpp File Reference
#include "CodegenEnv.h"
#include "CodegenUtils.h"
#include "LoopEmitter.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
#include "mlir/Dialect/SparseTensor/Utils/Merger.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/AffineExprVisitor.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/TensorEncoding.h"
#include "llvm/ADT/SmallBitVector.h"
#include <optional>

Go to the source code of this file.

Functions

static bool isInvariantAffine (AffineExpr a, ArrayRef< LoopId > loopStack, LoopId ldx, bool &isAtLoop)
 Determines if affine expression is invariant. More...
 
static bool isInvariantAffine (CodegenEnv &env, AffineExpr a, LoopId ldx, bool &isAtLoop)
 Determines if affine expression is invariant. More...
 
static AffineMap permute (CodegenEnv &env, AffineMap m)
 Helper method to construct a permuted dimension ordering that adheres to the given topological sort. More...
 
static bool findAffine (Merger &merger, TensorId tid, Level lvl, AffineExpr a, DimLevelType dlt, LoopId &filterLdx, bool setLvlFormat=true)
 Helper method to inspect affine expressions. More...
 
static bool findDepIdxSet (Merger &merger, TensorId tensor, Level lvl, AffineExpr a, DimLevelType dlt, bool isSubExp=false, int64_t coefficient=1)
 Helper method to inspect affine expressions for index variable reduction based codegen. More...
 
static unsigned getNumNonTrivialIdxExpOnSparseLvls (AffineMap map, Value tensor)
 Get the total number of compound affine expressions in the getMatchingIndexingMap for the given tensor. More...
 
static unsigned getNumNonTrivialIdxExpOnSparseLvls (linalg::GenericOp op)
 Get the total number of sparse levels with compound affine expressions, summed over all operands of the GenericOp. More...
 
static bool hasNonTrivialAffineOnSparseOut (linalg::GenericOp op)
 
static bool findSparseAnnotations (CodegenEnv &env, bool idxReducBased)
 Helper method to inspect sparse encodings in the tensor types. More...
 
static bool topSortOptimal (CodegenEnv &env, ArrayRef< utils::IteratorType > iteratorTypes, std::vector< unsigned > &inDegree, std::vector< std::vector< bool >> &adjM)
 A helper to compute a topological sort. More...
 
static void addIterOrdering (LoopId f, LoopId t, std::vector< std::vector< bool >> &adjM, std::vector< unsigned > &inDegree)
 
static void addAffineOrderings (std::vector< std::vector< bool >> &adjM, std::vector< unsigned > &inDegree, AffineExpr a, AffineExpr b, std::optional< LoopId > fidx, std::optional< LoopId > tidx)
 Helper method to add all constraints from the indices in one affine expression before all indices in the other affine expression. More...
 
static void tryRelaxAffineConstraints (linalg::GenericOp op, std::optional< LoopId > &fldx, AffineExpr &fa, std::optional< LoopId > &tldx, AffineExpr &ta)
 
static void addFilterLoopBasedConstraints (CodegenEnv &env, OpOperand &t, OpOperand *skip, SortMask mask, std::vector< std::vector< bool >> &adjM, std::vector< unsigned > &inDegree)
 
static void addSliceBasedConstraints (CodegenEnv &env, OpOperand &t, OpOperand *skip, SortMask mask, std::vector< std::vector< bool >> &adjM, std::vector< unsigned > &inDegree)
 
static bool computeIterationGraph (CodegenEnv &env, SortMask mask, OpOperand *skip, bool idxReducBased=false)
 Computes a topologically sorted iteration graph for the linalg operation. More...
 
static void genBuffers (CodegenEnv &env, OpBuilder &builder)
 Local bufferization of all dense and sparse data structures. More...
 
static Value genIndex (CodegenEnv &env, OpOperand *t)
 Generates index for load/store on sparse tensor. More...
 
static Value genSubscript (CodegenEnv &env, OpBuilder &builder, OpOperand *t, SmallVectorImpl< Value > &args)
 Generates subscript for load/store on a dense or sparse tensor. More...
 
static Value genInsertionLoad (CodegenEnv &env, OpBuilder &builder, OpOperand *t)
 Generates insertion code to implement dynamic tensor load. More...
 
static Value genInsertionLoadReduce (CodegenEnv &env, OpBuilder &builder, OpOperand *t)
 Generates insertion code to implement dynamic tensor load for reduction. More...
 
static void genInsertionStore (CodegenEnv &env, OpBuilder &builder, OpOperand *t, Value rhs)
 Generates insertion code to implement dynamic tensor store. More...
 
static Value genTensorLoad (CodegenEnv &env, OpBuilder &builder, ExprId exp)
 Generates a load on a dense or sparse tensor. More...
 
static void genTensorStore (CodegenEnv &env, OpBuilder &builder, ExprId exp, Value rhs)
 Generates a store on a dense or sparse tensor. More...
 
static Value genInvariantValue (CodegenEnv &env, ExprId exp)
 Generates an invariant value. More...
 
static Value relinkBranch (CodegenEnv &env, RewriterBase &rewriter, Block *block, Value e, LoopId ldx)
 Semi-ring branches are simply inlined by the sparse compiler. More...
 
static Value genExp (CodegenEnv &env, RewriterBase &rewriter, ExprId e, LoopId ldx)
 Recursively generates tensor expression. More...
 
static void genInvariants (CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopId ldx, bool atStart)
 Hoists loop invariant tensor loads for which indices have been exhausted. More...
 
static void genExpand (CodegenEnv &env, OpBuilder &builder, LoopOrd at, bool atStart)
 Generates an expanded access pattern in innermost dimension. More...
 
static bool isParallelFor (CodegenEnv &env, bool isOuter, bool isSparse)
 Returns parallelization strategy. More...
 
static bool shouldTryParallize (CodegenEnv &env, LoopId ldx, bool isOuter, ArrayRef< TensorLevel > tidLvls)
 Whether or not the current loop being generated should be parallized (if possible) according to the configuration. More...
 
static OperationgenFilterLoop (CodegenEnv &env, OpBuilder &builder, LoopId ldx, TensorLevel tidLvl)
 Generates a "filter loop" on the given tid level to locate a coordinate that is of the same value as evaluated by the affine expression in its matching indexing map. More...
 
static OperationgenCoIteration (CodegenEnv &env, OpBuilder &builder, LoopId idx, ArrayRef< TensorLevel > tidLvls, bool tryParallel, bool needsUniv)
 Emit a loop to coiterate over the list of tensor levels. More...
 
static OperationgenLoop (CodegenEnv &env, OpBuilder &builder, LoopOrd at, bool needsUniv, ArrayRef< TensorLevel > tidLvls)
 Generates a for-loop or a while-loop, depending on whether it implements singleton iteration or co-iteration over the given conjunction. More...
 
static void finalizeWhileOp (CodegenEnv &env, OpBuilder &builder, LoopId idx, bool needsUniv)
 Generates the induction structure for a while-loop. More...
 
static scf::IfOp genIf (CodegenEnv &env, OpBuilder &builder, LoopId ldx, LatPointId p)
 Generates a single if-statement within a while-loop. More...
 
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. More...
 
static bool startLoopSeq (CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopOrd at, LoopId idx, LoopId ldx, LatSetId lts)
 Starts a loop sequence at given level. More...
 
static void genConstantDenseAddressFromLevel (CodegenEnv &env, OpBuilder &builder, TensorId tid, Level startLvl)
 
static void genInitConstantDenseAddress (CodegenEnv &env, RewriterBase &rewriter)
 
static bool translateBitsToTidLvlPairs (CodegenEnv &env, LatPointId li, LoopId ldx, SmallVectorImpl< TensorLevel > &tidLvls, SmallVectorImpl< std::pair< TensorLevel, AffineExpr >> &affineTidLvls)
 Return true if the lattices bit can be iterated by a for loop. More...
 
static std::pair< Operation *, bool > startLoop (CodegenEnv &env, OpBuilder &builder, LoopOrd at, LatPointId li, bool needsUniv)
 Starts a single loop in current sequence. More...
 
static bool endLoop (CodegenEnv &env, RewriterBase &rewriter, Operation *loop, LoopId idx, LatPointId li, bool needsUniv, bool isSingleCond)
 Ends a single loop in current sequence. Returns new values for needsUniv. More...
 
static void endLoopSeq (CodegenEnv &env, OpBuilder &builder, unsigned exp, unsigned at, unsigned idx, unsigned ldx)
 Ends a loop sequence at given level. More...
 
static void genStmt (CodegenEnv &env, RewriterBase &rewriter, ExprId exp, LoopOrd at)
 Recursively generates code while computing iteration lattices in order to manage the complexity of implementing co-iteration over unions and intersections of sparse iterations spaces. More...
 
static void genResult (CodegenEnv &env, RewriterBase &rewriter)
 Converts the result computed by the sparse kernel into the required form. More...
 

Function Documentation

◆ addAffineOrderings()

static void addAffineOrderings ( std::vector< std::vector< bool >> &  adjM,
std::vector< unsigned > &  inDegree,
AffineExpr  a,
AffineExpr  b,
std::optional< LoopId fidx,
std::optional< LoopId tidx 
)
static

Helper method to add all constraints from the indices in one affine expression before all indices in the other affine expression.

For example i0+i1 < i2+i3+1 yields i0<i2, i0<i3, i1<i2, and i1<i3. The affine expression a is empty iff fidx have a value, leading to b = (i0 + i1) < fidx => i0 < fidx, i1 < fidx. The affine expression b is empty iff tidx have a value, leading to tidx < a = (i0 + i1) => tidx < i0, tidx < i1.

The inDegree is indexed by LoopId, and the adjM is indexed by (LoopId,LoopId).

Definition at line 554 of file Sparsification.cpp.

References mlir::Add, addIterOrdering(), mlir::DimId, and mlir::Mul.

Referenced by addFilterLoopBasedConstraints().

◆ addFilterLoopBasedConstraints()

static void addFilterLoopBasedConstraints ( CodegenEnv env,
OpOperand t,
OpOperand skip,
SortMask  mask,
std::vector< std::vector< bool >> &  adjM,
std::vector< unsigned > &  inDegree 
)
static

◆ addIterOrdering()

static void addIterOrdering ( LoopId  f,
LoopId  t,
std::vector< std::vector< bool >> &  adjM,
std::vector< unsigned > &  inDegree 
)
static

◆ addSliceBasedConstraints()

static void addSliceBasedConstraints ( CodegenEnv env,
OpOperand t,
OpOperand skip,
SortMask  mask,
std::vector< std::vector< bool >> &  adjM,
std::vector< unsigned > &  inDegree 
)
static

◆ computeIterationGraph()

static bool computeIterationGraph ( CodegenEnv env,
SortMask  mask,
OpOperand skip,
bool  idxReducBased = false 
)
static

Computes a topologically sorted iteration graph for the linalg operation.

Ensures all tensors are visited in natural index order. This is essential for sparse storage formats since these only support access along fixed dimensions. Even for dense storage formats, however, the natural index order yields innermost unit-stride access with better spatial locality.

Definition at line 791 of file Sparsification.cpp.

References addFilterLoopBasedConstraints(), addIterOrdering(), addSliceBasedConstraints(), mlir::sparse_tensor::CodegenEnv::dlt(), mlir::sparse_tensor::Merger::getNumLoops(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::sparse_tensor::isCompressedDLT(), mlir::sparse_tensor::isCompressedWithHiDLT(), mlir::sparse_tensor::isDenseDLT(), mlir::sparse_tensor::isSingletonDLT(), mlir::sparse_tensor::isUndefDLT(), mlir::sparse_tensor::CodegenEnv::makeTensorId(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), mlir::sparse_tensor::CodegenEnv::topSortClear(), and topSortOptimal().

◆ endIf()

static void endIf ( CodegenEnv env,
OpBuilder builder,
scf::IfOp  ifOp,
Value  redInput,
Value  cntInput,
Value  insInput,
Value  validIns 
)
static

◆ endLoop()

static bool endLoop ( CodegenEnv env,
RewriterBase rewriter,
Operation loop,
LoopId  idx,
LatPointId  li,
bool  needsUniv,
bool  isSingleCond 
)
static

◆ endLoopSeq()

static void endLoopSeq ( CodegenEnv env,
OpBuilder builder,
unsigned  exp,
unsigned  at,
unsigned  idx,
unsigned  ldx 
)
static

◆ finalizeWhileOp()

static void finalizeWhileOp ( CodegenEnv env,
OpBuilder builder,
LoopId  idx,
bool  needsUniv 
)
static

◆ findAffine()

static bool findAffine ( Merger merger,
TensorId  tid,
Level  lvl,
AffineExpr  a,
DimLevelType  dlt,
LoopId filterLdx,
bool  setLvlFormat = true 
)
static

Helper method to inspect affine expressions.

Rejects cases where the same index is used more than once. Also rejects compound affine expressions in sparse dimensions. filterIdx stores the current filter loop idx should be used for the next compound affine sparse level, and it will be incremented by one when used.

Definition at line 229 of file Sparsification.cpp.

References mlir::Add, mlir::AffineExpr::cast(), mlir::Constant, mlir::DimId, mlir::AffineExpr::dyn_cast(), mlir::AffineExpr::getKind(), mlir::sparse_tensor::Merger::getLvlType(), mlir::AffineDimExpr::getPosition(), mlir::sparse_tensor::isDenseDLT(), mlir::sparse_tensor::isUndefDLT(), mlir::sparse_tensor::Merger::makeLoopId(), mlir::Mul, and mlir::sparse_tensor::Merger::setLevelAndType().

Referenced by findSparseAnnotations().

◆ findDepIdxSet()

static bool findDepIdxSet ( Merger merger,
TensorId  tensor,
Level  lvl,
AffineExpr  a,
DimLevelType  dlt,
bool  isSubExp = false,
int64_t  coefficient = 1 
)
static

Helper method to inspect affine expressions for index variable reduction based codegen.

It finds the dependent index set for all tensor levels in the current expression we are generating.

For example, when handling A[i+j][j+k], we build the two way mapping in merger between (tensor, level) pairs and their dependent index variable set: A_0 <=> [i, j] and A_1 <=> [j, k]

It rejects cases (returns false) 1st, when the same index is used more than once, e.g., A[i+j][i] 2nd, when multiplication is used in the non-trivial index expression. 3rd, when a constant operand is used in the non-trivial index expression.

TODO: constant should be easy to handle.

Definition at line 284 of file Sparsification.cpp.

References mlir::Add, mlir::AffineExpr::cast(), mlir::Constant, mlir::DimId, mlir::AffineExpr::getKind(), mlir::AffineBinaryOpExpr::getLHS(), mlir::sparse_tensor::Merger::getLvlType(), mlir::AffineDimExpr::getPosition(), mlir::sparse_tensor::Merger::hasDependentLvl(), mlir::AffineExpr::isa(), mlir::sparse_tensor::isUndefDLT(), mlir::sparse_tensor::Merger::makeLoopId(), mlir::Mul, mlir::sparse_tensor::Merger::setLevelAndType(), and mlir::sparse_tensor::Merger::setLoopDependentTensorLevel().

Referenced by findSparseAnnotations().

◆ findSparseAnnotations()

static bool findSparseAnnotations ( CodegenEnv env,
bool  idxReducBased 
)
static

Helper method to inspect sparse encodings in the tensor types.

Fills the per-dimension sparsity information for all tensors. Returns true if the sparse annotations and affine subscript expressions of all tensors are admissible. Returns false if no annotations are found or inadmissible constructs occur. We currently support two different ways to handle non-trivial index expression on sparse tensors, and they accept different affine expressions. When using filter-loop-based approach, it accept (almost) arbitrary affine index expression on sparse tensor but it is much less efficient, and will be gradually removed from the codebase. When using dependent index reducton-based approach, it currently only supports affine addition index expression.

Definition at line 430 of file Sparsification.cpp.

References findAffine(), findDepIdxSet(), mlir::sparse_tensor::Merger::getNumLoops(), getNumNonTrivialIdxExpOnSparseLvls(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::sparse_tensor::Merger::getStartingFilterLoopId(), mlir::sparse_tensor::CodegenEnv::makeTensorId(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::sparse_tensor::toOrigDim().

◆ genBuffers()

static void genBuffers ( CodegenEnv env,
OpBuilder builder 
)
static

Local bufferization of all dense and sparse data structures.

Generates buffer for the output tensor. Note that all sparse kernels assume that when all elements are written to (viz. x(i) = y(i) * z(i)), the output buffer is already initialized to all zeroes and only nonzeroes values are computed and written out. For updates (viz. x(i) += y(i) * z(i)), only nonzeroes values are used for the updates and no assumption on the original contents of the output buffer is necessary.

Definition at line 847 of file Sparsification.cpp.

◆ genCoIteration()

static Operation* genCoIteration ( CodegenEnv env,
OpBuilder builder,
LoopId  idx,
ArrayRef< TensorLevel tidLvls,
bool  tryParallel,
bool  needsUniv 
)
static

Emit a loop to coiterate over the list of tensor levels.

The generated loop can either be a for loop or while loop depending on whether there is at most one sparse level in the list.

Definition at line 1421 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::sparse_tensor::CodegenEnv::genLoopBoundary(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genLoop().

◆ genConstantDenseAddressFromLevel()

static void genConstantDenseAddressFromLevel ( CodegenEnv env,
OpBuilder builder,
TensorId  tid,
Level  startLvl 
)
static

Definition at line 1616 of file Sparsification.cpp.

Referenced by genInitConstantDenseAddress(), and startLoop().

◆ genExp()

static Value genExp ( CodegenEnv env,
RewriterBase rewriter,
ExprId  e,
LoopId  ldx 
)
static

Recursively generates tensor expression.

Definition at line 1181 of file Sparsification.cpp.

References mlir::sparse_tensor::detail::kInvalidId.

Referenced by genStmt().

◆ genExpand()

static void genExpand ( CodegenEnv env,
OpBuilder builder,
LoopOrd  at,
bool  atStart 
)
static

Generates an expanded access pattern in innermost dimension.

Definition at line 1311 of file Sparsification.cpp.

Referenced by endLoopSeq(), and startLoopSeq().

◆ genFilterLoop()

static Operation* genFilterLoop ( CodegenEnv env,
OpBuilder builder,
LoopId  ldx,
TensorLevel  tidLvl 
)
static

Generates a "filter loop" on the given tid level to locate a coordinate that is of the same value as evaluated by the affine expression in its matching indexing map.

Definition at line 1398 of file Sparsification.cpp.

Referenced by genLoop().

◆ genIf()

static scf::IfOp genIf ( CodegenEnv env,
OpBuilder builder,
LoopId  ldx,
LatPointId  p 
)
static

◆ genIndex()

static Value genIndex ( CodegenEnv env,
OpOperand t 
)
static

◆ genInitConstantDenseAddress()

static void genInitConstantDenseAddress ( CodegenEnv env,
RewriterBase rewriter 
)
static

◆ genInsertionLoad()

static Value genInsertionLoad ( CodegenEnv env,
OpBuilder builder,
OpOperand t 
)
static

Generates insertion code to implement dynamic tensor load.

Definition at line 957 of file Sparsification.cpp.

◆ genInsertionLoadReduce()

static Value genInsertionLoadReduce ( CodegenEnv env,
OpBuilder builder,
OpOperand t 
)
static

Generates insertion code to implement dynamic tensor load for reduction.

Definition at line 972 of file Sparsification.cpp.

◆ genInsertionStore()

static void genInsertionStore ( CodegenEnv env,
OpBuilder builder,
OpOperand t,
Value  rhs 
)
static

Generates insertion code to implement dynamic tensor store.

Definition at line 990 of file Sparsification.cpp.

◆ genInvariants()

static void genInvariants ( CodegenEnv env,
OpBuilder builder,
ExprId  exp,
LoopId  ldx,
bool  atStart 
)
static

Hoists loop invariant tensor loads for which indices have been exhausted.

Definition at line 1241 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::exp(), mlir::sparse_tensor::TensorExp::kind, mlir::sparse_tensor::detail::kInvalidId, and mlir::sparse_tensor::TensorExp::kTensor.

Referenced by endLoopSeq(), and startLoopSeq().

◆ genInvariantValue()

static Value genInvariantValue ( CodegenEnv env,
ExprId  exp 
)
inlinestatic

Generates an invariant value.

Definition at line 1138 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::exp(), and mlir::sparse_tensor::TensorExp::val.

◆ genLoop()

static Operation* genLoop ( CodegenEnv env,
OpBuilder builder,
LoopOrd  at,
bool  needsUniv,
ArrayRef< TensorLevel tidLvls 
)
static

Generates a for-loop or a while-loop, depending on whether it implements singleton iteration or co-iteration over the given conjunction.

Definition at line 1437 of file Sparsification.cpp.

References genCoIteration(), genFilterLoop(), mlir::sparse_tensor::Merger::isFilterLoop(), mlir::sparse_tensor::CodegenEnv::merger(), shouldTryParallize(), and mlir::sparse_tensor::CodegenEnv::topSortAt().

Referenced by startLoop().

◆ genResult()

static void genResult ( CodegenEnv env,
RewriterBase rewriter 
)
static

Converts the result computed by the sparse kernel into the required form.

Definition at line 1905 of file Sparsification.cpp.

◆ genStmt()

static void genStmt ( CodegenEnv env,
RewriterBase rewriter,
ExprId  exp,
LoopOrd  at 
)
static

◆ genSubscript()

static Value genSubscript ( CodegenEnv env,
OpBuilder builder,
OpOperand t,
SmallVectorImpl< Value > &  args 
)
static

◆ genTensorLoad()

static Value genTensorLoad ( CodegenEnv env,
OpBuilder builder,
ExprId  exp 
)
static

Generates a load on a dense or sparse tensor.

Definition at line 1064 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::exp(), and mlir::sparse_tensor::TensorExp::val.

◆ genTensorStore()

static void genTensorStore ( CodegenEnv env,
OpBuilder builder,
ExprId  exp,
Value  rhs 
)
static

◆ getNumNonTrivialIdxExpOnSparseLvls() [1/2]

static unsigned getNumNonTrivialIdxExpOnSparseLvls ( AffineMap  map,
Value  tensor 
)
static

Get the total number of compound affine expressions in the getMatchingIndexingMap for the given tensor.

For the following inputs:

map = (d0, d1, d2) => (d0 + d1 : compressed, d2 : compressed)

Returns 1 (because the first level is compressed and its corresponding indexing-expression is d0 + d1)

Definition at line 358 of file Sparsification.cpp.

References mlir::sparse_tensor::SparseTensorType::getDimRank(), mlir::sparse_tensor::SparseTensorType::getEncoding(), mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::AffineMap::getResults(), mlir::Value::getType(), mlir::sparse_tensor::SparseTensorType::isDenseLvl(), and mlir::sparse_tensor::toOrigDim().

Referenced by addSliceBasedConstraints(), and findSparseAnnotations().

◆ getNumNonTrivialIdxExpOnSparseLvls() [2/2]

static unsigned getNumNonTrivialIdxExpOnSparseLvls ( linalg::GenericOp  op)
static

Get the total number of sparse levels with compound affine expressions, summed over all operands of the GenericOp.

Definition at line 402 of file Sparsification.cpp.

◆ hasNonTrivialAffineOnSparseOut()

static bool hasNonTrivialAffineOnSparseOut ( linalg::GenericOp  op)
static

Definition at line 410 of file Sparsification.cpp.

◆ isInvariantAffine() [1/2]

static bool isInvariantAffine ( AffineExpr  a,
ArrayRef< LoopId loopStack,
LoopId  ldx,
bool &  isAtLoop 
)
static

Determines if affine expression is invariant.

Definition at line 136 of file Sparsification.cpp.

References mlir::Add, mlir::AffineExpr::cast(), mlir::DimId, mlir::AffineExpr::getKind(), mlir::AffineExpr::isa(), and mlir::Mul.

Referenced by isInvariantAffine(), and permute().

◆ isInvariantAffine() [2/2]

static bool isInvariantAffine ( CodegenEnv env,
AffineExpr  a,
LoopId  ldx,
bool &  isAtLoop 
)
static

Determines if affine expression is invariant.

Definition at line 168 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::getCurrentLoopStack(), and isInvariantAffine().

◆ isParallelFor()

static bool isParallelFor ( CodegenEnv env,
bool  isOuter,
bool  isSparse 
)
static

Returns parallelization strategy.

Any implicit loop in the Linalg operation that is marked "parallel" is a candidate. Whether it is actually converted to a parallel operation depends on the requested strategy.

Definition at line 1355 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::hasSparseOutput(), mlir::sparse_tensor::CodegenEnv::isExpand(), mlir::kAnyStorageAnyLoop, mlir::kAnyStorageOuterLoop, mlir::kDenseAnyLoop, mlir::kDenseOuterLoop, mlir::kNone, mlir::sparse_tensor::CodegenEnv::options(), and mlir::SparsificationOptions::parallelizationStrategy.

◆ permute()

static AffineMap permute ( CodegenEnv env,
AffineMap  m 
)
static

◆ relinkBranch()

static Value relinkBranch ( CodegenEnv env,
RewriterBase rewriter,
Block block,
Value  e,
LoopId  ldx 
)
static

Semi-ring branches are simply inlined by the sparse compiler.

Prior analysis has verified that all computations are "local" to the inlined branch or otherwise invariantly defined outside the loop nest, with the exception of index computations, which need to be relinked to actual inlined cloned code.

Definition at line 1147 of file Sparsification.cpp.

◆ shouldTryParallize()

static bool shouldTryParallize ( CodegenEnv env,
LoopId  ldx,
bool  isOuter,
ArrayRef< TensorLevel tidLvls 
)
static

Whether or not the current loop being generated should be parallized (if possible) according to the configuration.

Definition at line 1380 of file Sparsification.cpp.

Referenced by genLoop().

◆ startLoop()

static std::pair<Operation *, bool> startLoop ( CodegenEnv env,
OpBuilder builder,
LoopOrd  at,
LatPointId  li,
bool  needsUniv 
)
static

◆ startLoopSeq()

static bool startLoopSeq ( CodegenEnv env,
OpBuilder builder,
ExprId  exp,
LoopOrd  at,
LoopId  idx,
LoopId  ldx,
LatSetId  lts 
)
static

◆ topSortOptimal()

static bool topSortOptimal ( CodegenEnv env,
ArrayRef< utils::IteratorType >  iteratorTypes,
std::vector< unsigned > &  inDegree,
std::vector< std::vector< bool >> &  adjM 
)
static

A helper to compute a topological sort.

O(n^2) time complexity as we use adj matrix for the graph. The sorted result will put the first Reduction iterator to the latest possible LoopOrd.

The inDegree is indexed by LoopId, and the adjM is indexed by (LoopId,LoopId).

Definition at line 477 of file Sparsification.cpp.

References mlir::sparse_tensor::Merger::getNumLoops(), mlir::sparse_tensor::Merger::isFilterLoop(), mlir::linalg::isReductionIterator(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::topSortPushBack(), and mlir::sparse_tensor::CodegenEnv::topSortSize().

Referenced by computeIterationGraph().

◆ translateBitsToTidLvlPairs()

static bool translateBitsToTidLvlPairs ( CodegenEnv env,
LatPointId  li,
LoopId  ldx,
SmallVectorImpl< TensorLevel > &  tidLvls,
SmallVectorImpl< std::pair< TensorLevel, AffineExpr >> &  affineTidLvls 
)
static

◆ tryRelaxAffineConstraints()

static void tryRelaxAffineConstraints ( linalg::GenericOp  op,
std::optional< LoopId > &  fldx,
AffineExpr fa,
std::optional< LoopId > &  tldx,
AffineExpr ta 
)
static

Definition at line 594 of file Sparsification.cpp.

Referenced by addFilterLoopBasedConstraints().