MLIR  19.0.0git
Functions
Sparsification.cpp File Reference
#include "Utils/CodegenEnv.h"
#include "Utils/CodegenUtils.h"
#include "Utils/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, LoopId curr, bool &isCurrentLoop)
 Returns true iff affine expression is invariant. More...
 
static bool findAffine (Merger &merger, TensorId tid, Level lvl, AffineExpr a, LevelType lt, bool setLvlFormat=true)
 Helper method to inspect affine expressions. More...
 
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. More...
 
static unsigned getNumNonTrivialIdxExpOnSparseLvls (AffineMap map, Value tensor)
 Gets the total number of compound affine expressions in the getMatchingIndexingMap for the given tensor. More...
 
static unsigned getNumNonTrivialIdxExpOnSparseLvls (linalg::GenericOp op)
 Gets 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 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 Value genConditionalInsert (Location loc, OpBuilder &builder, Value cond, Value sparseOut, ValueRange ivs, Value v)
 
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)
 Semi-ring branches are simply inlined by the sparsifier. More...
 
static Value genExp (CodegenEnv &env, RewriterBase &rewriter, ExprId e)
 Recursively generates tensor expression. More...
 
static void genInvariants (CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopId curr, bool isStart)
 Hoists loop invariant tensor loads for which indices have been exhausted. More...
 
static void genExpand (CodegenEnv &env, OpBuilder &builder, LoopId curr, bool isStart)
 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 curr, ArrayRef< TensorLevel > tidLvls)
 Whether or not the current loop being generated should be parallized (if possible) according to the configuration. More...
 
static OperationgenCoIteration (CodegenEnv &env, OpBuilder &builder, 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, LoopId curr, 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, bool needsUniv)
 Generates the induction structure for a while-loop. More...
 
static scf::IfOp genIf (CodegenEnv &env, OpBuilder &builder, LoopId curr, 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 getAllTidLvlsInLatPoints (CodegenEnv &env, LatPointId li, LoopId curr, llvm::function_ref< void(TensorLevel, AffineExpr)> callback)
 
static bool startLoopSeq (CodegenEnv &env, OpBuilder &builder, ExprId exp, LoopId curr, 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 curr, SmallVectorImpl< TensorLevel > &tidLvls, SmallVectorImpl< std::pair< TensorLevel, AffineExpr >> &affineTidLvls)
 Returns true if the lattice bit can be iterated by a for loop. More...
 
static std::pair< Operation *, bool > startLoop (CodegenEnv &env, OpBuilder &builder, LoopId curr, LatPointId li, bool needsUniv)
 Starts a single loop in current sequence. More...
 
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. More...
 
static void endLoopSeq (CodegenEnv &env, OpBuilder &builder, unsigned exp, unsigned at)
 Ends a loop sequence at given level. More...
 
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 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

◆ 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,
LatPointId  li,
bool  needsUniv,
bool  isSingleCond 
)
static

◆ endLoopSeq()

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

◆ finalizeWhileOp()

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

◆ findAffine()

static bool findAffine ( Merger merger,
TensorId  tid,
Level  lvl,
AffineExpr  a,
LevelType  lt,
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.

Definition at line 75 of file Sparsification.cpp.

References mlir::Add, mlir::Constant, mlir::DimId, mlir::AffineExpr::getKind(), mlir::sparse_tensor::Merger::getLvlType(), mlir::sparse_tensor::LevelType::hasDenseSemantic(), mlir::sparse_tensor::isUndefLT(), 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,
LevelType  lt,
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 119 of file Sparsification.cpp.

References mlir::Add, mlir::Constant, mlir::DimId, mlir::AffineExpr::getKind(), mlir::sparse_tensor::Merger::getLvlType(), mlir::sparse_tensor::Merger::hasDependentLvl(), mlir::sparse_tensor::isUndefLT(), 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 dependent index reducton-based approach, it currently only supports affine addition index expression.

Definition at line 248 of file Sparsification.cpp.

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

◆ 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 287 of file Sparsification.cpp.

◆ genCoIteration()

static Operation* genCoIteration ( CodegenEnv env,
OpBuilder builder,
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 833 of file Sparsification.cpp.

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

Referenced by genLoop().

◆ genConditionalInsert()

static Value genConditionalInsert ( Location  loc,
OpBuilder builder,
Value  cond,
Value  sparseOut,
ValueRange  ivs,
Value  v 
)
static

◆ genConstantDenseAddressFromLevel()

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

Definition at line 1115 of file Sparsification.cpp.

Referenced by genInitConstantDenseAddress(), and startLoop().

◆ genExp()

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

Recursively generates tensor expression.

Definition at line 618 of file Sparsification.cpp.

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

Referenced by genStmt().

◆ genExpand()

static void genExpand ( CodegenEnv env,
OpBuilder builder,
LoopId  curr,
bool  isStart 
)
static

Generates an expanded access pattern in innermost dimension.

Definition at line 747 of file Sparsification.cpp.

Referenced by endLoopSeq(), and startLoopSeq().

◆ genIf()

static scf::IfOp genIf ( CodegenEnv env,
OpBuilder builder,
LoopId  curr,
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 374 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 389 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 423 of file Sparsification.cpp.

◆ genInvariants()

static void genInvariants ( CodegenEnv env,
OpBuilder builder,
ExprId  exp,
LoopId  curr,
bool  isStart 
)
static

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

Definition at line 677 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 575 of file Sparsification.cpp.

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

◆ genLoop()

static Operation* genLoop ( CodegenEnv env,
OpBuilder builder,
LoopId  curr,
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 847 of file Sparsification.cpp.

References genCoIteration(), and shouldTryParallize().

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 1296 of file Sparsification.cpp.

◆ genStmt()

static void genStmt ( CodegenEnv env,
RewriterBase rewriter,
ExprId  exp,
LoopId  curr 
)
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 496 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

◆ getAllTidLvlsInLatPoints()

static bool getAllTidLvlsInLatPoints ( CodegenEnv env,
LatPointId  li,
LoopId  curr,
llvm::function_ref< void(TensorLevel, AffineExpr)>  callback 
)
static

◆ getNumNonTrivialIdxExpOnSparseLvls() [1/2]

static unsigned getNumNonTrivialIdxExpOnSparseLvls ( AffineMap  map,
Value  tensor 
)
static

Gets 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 197 of file Sparsification.cpp.

References mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::SparseTensorType::getLvlType(), mlir::AffineMap::getResults(), mlir::Value::getType(), and mlir::sparse_tensor::LevelType::hasDenseSemantic().

Referenced by findSparseAnnotations().

◆ getNumNonTrivialIdxExpOnSparseLvls() [2/2]

static unsigned getNumNonTrivialIdxExpOnSparseLvls ( linalg::GenericOp  op)
static

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

Definition at line 222 of file Sparsification.cpp.

◆ hasNonTrivialAffineOnSparseOut()

static bool hasNonTrivialAffineOnSparseOut ( linalg::GenericOp  op)
static

Definition at line 231 of file Sparsification.cpp.

◆ isInvariantAffine()

static bool isInvariantAffine ( AffineExpr  a,
LoopId  curr,
bool &  isCurrentLoop 
)
static

Returns true iff affine expression is invariant.

Sets the parameter isCurrentLoop when expression just became invariant.

Definition at line 49 of file Sparsification.cpp.

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

◆ 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 791 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.

◆ relinkBranch()

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

Semi-ring branches are simply inlined by the sparsifier.

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 584 of file Sparsification.cpp.

◆ shouldTryParallize()

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

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

Definition at line 816 of file Sparsification.cpp.

Referenced by genLoop().

◆ startLoop()

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

◆ startLoopSeq()

static bool startLoopSeq ( CodegenEnv env,
OpBuilder builder,
ExprId  exp,
LoopId  curr,
LatSetId  lts 
)
static

◆ translateBitsToTidLvlPairs()

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

Returns true if the lattice bit can be iterated by a for loop.

Definition at line 1152 of file Sparsification.cpp.

References getAllTidLvlsInLatPoints().

Referenced by startLoop().