#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>

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 Operation *	genCoIteration (CodegenEnv &env, OpBuilder &builder, ArrayRef< TensorLevel > tidLvls, unsigned numCases, bool tryParallel, bool needsUniv)
	Emit a loop to coiterate over the list of tensor levels. More...

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-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 void	genCoIterationCase (CodegenEnv &env, OpBuilder &builder, unsigned caseIdx, LatPointId allCase, LatPointId curCase, MutableArrayRef< Value > reduc)
	Generates a case region in the coiterate operation. 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, unsigned numCases, 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

Generates end of true branch of if-statement within a while-loop.

Definition at line 974 of file Sparsification.cpp.

Referenced by genStmt().

◆ endLoop()

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

static

Ends a single loop in current sequence. Returns new values for needsUniv.

Definition at line 1236 of file Sparsification.cpp.

References mlir::sparse_tensor::constantI1(), mlir::sparse_tensor::CodegenEnv::emitter(), finalizeWhileOp(), mlir::sparse_tensor::CodegenEnv::genLoopBoundary(), mlir::sparse_tensor::CodegenEnv::isReduc(), mlir::sparse_tensor::CodegenEnv::isValidLexInsert(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::sparse_tensor::CodegenEnv::updateValidLexInsert().

Referenced by genStmt().

◆ endLoopSeq()

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

static

Ends a loop sequence at given level.

Definition at line 1257 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::sparse_tensor::LoopEmitter::exitCurrentLoopSeq(), genExpand(), genInvariants(), mlir::sparse_tensor::CodegenEnv::getLoopVar(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genStmt().

◆ finalizeWhileOp()

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

static

Generates the induction structure for a while-loop.

Definition at line 867 of file Sparsification.cpp.

Referenced by endLoop().

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

References mlir::sparse_tensor::constantZero(), mlir::OpBuilder::create(), mlir::sparse_tensor::CodegenEnv::emitter(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::getElementTypeOrSelf(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::Value::getType(), mlir::getValueOrCreateConstantIndexOp(), mlir::sparse_tensor::LoopEmitter::initializeLoopEmit(), and mlir::sparse_tensor::CodegenEnv::op().

◆ genCoIteration()

static Operation* genCoIteration	(	CodegenEnv &	env,
		OpBuilder &	builder,
		ArrayRef< TensorLevel >	tidLvls,
		unsigned	numCases,
		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 842 of file Sparsification.cpp.

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

◆ genCoIterationCase()

static void genCoIterationCase	(	CodegenEnv &	env,
		OpBuilder &	builder,
		unsigned	caseIdx,
		LatPointId	allCase,
		LatPointId	curCase,
		MutableArrayRef< Value >	reduc
	)

static

Generates a case region in the coiterate operation.

Computes the subset of iterators that are valid in the current case being generated.

Definition at line 907 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::sparse_tensor::LoopEmitter::enterCurrentCoIterationCase(), mlir::detail::enumerate(), mlir::sparse_tensor::Merger::lat(), mlir::sparse_tensor::Merger::latGT(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), mlir::sparse_tensor::I64BitSet::set(), and mlir::sparse_tensor::LatPoint::simple.

Referenced by genStmt().

◆ genConditionalInsert()

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

static

Definition at line 409 of file Sparsification.cpp.

References mlir::OpBuilder::create(), mlir::Value::getType(), mlir::OpBuilder::setInsertionPointAfter(), and mlir::OpBuilder::setInsertionPointToStart().

Referenced by genInsertionStore().

◆ genConstantDenseAddressFromLevel()

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

static

Definition at line 1148 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::OpOperand::getOperandNumber(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::Value::getType(), mlir::sparse_tensor::LoopEmitter::locateLvlAtAffineAddress(), mlir::sparse_tensor::CodegenEnv::makeTensorId(), mlir::sparse_tensor::CodegenEnv::makeTensorLevel(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genInitConstantDenseAddress(), and startLoop().

◆ genExp()

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

static

Recursively generates tensor expression.

Definition at line 627 of file Sparsification.cpp.

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

References mlir::sparse_tensor::CodegenEnv::atExpandLevel(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::get(), mlir::Builder::getI1Type(), mlir::Builder::getIndexType(), mlir::Value::getType(), mlir::sparse_tensor::CodegenEnv::isReduc(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by endLoopSeq(), and startLoopSeq().

◆ genIf()

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

static

Generates a single if-statement within a while-loop.

Definition at line 927 of file Sparsification.cpp.

Referenced by genStmt().

◆ genIndex()

static Value genIndex	(	CodegenEnv &	env,
		OpOperand *	t
	)

static

Generates index for load/store on sparse tensor.

Definition at line 337 of file Sparsification.cpp.

References mlir::DimId, mlir::IROperand< DerivedT, IRValueT >::get(), mlir::AffineExpr::getKind(), mlir::sparse_tensor::CodegenEnv::getLoopVar(), mlir::sparse_tensor::getSparseTensorType(), mlir::sparse_tensor::CodegenEnv::makeLoopId(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genInsertionLoad(), genInsertionLoadReduce(), and genInsertionStore().

◆ genInitConstantDenseAddress()

static void genInitConstantDenseAddress	(	CodegenEnv &	env,
		RewriterBase &	rewriter
	)

static

Definition at line 1178 of file Sparsification.cpp.

References genConstantDenseAddressFromLevel(), and mlir::sparse_tensor::CodegenEnv::op().

◆ genInsertionLoad()

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

static

Generates insertion code to implement dynamic tensor load.

Definition at line 377 of file Sparsification.cpp.

References mlir::sparse_tensor::constantZero(), mlir::OpBuilder::create(), genIndex(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::getElementTypeOrSelf(), mlir::sparse_tensor::CodegenEnv::getExpandValues(), mlir::Value::getType(), mlir::sparse_tensor::CodegenEnv::isExpand(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genTensorLoad().

◆ genInsertionLoadReduce()

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

static

Generates insertion code to implement dynamic tensor load for reduction.

Definition at line 392 of file Sparsification.cpp.

References mlir::OpBuilder::create(), genIndex(), mlir::sparse_tensor::CodegenEnv::getCustomRedId(), mlir::sparse_tensor::CodegenEnv::getExpandFilled(), mlir::sparse_tensor::CodegenEnv::getExpandValues(), mlir::sparse_tensor::CodegenEnv::isExpand(), and mlir::sparse_tensor::CodegenEnv::op().

Referenced by genTensorLoad().

◆ genInsertionStore()

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

static

Generates insertion code to implement dynamic tensor store.

Definition at line 426 of file Sparsification.cpp.

Referenced by genTensorStore().

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

Referenced by endLoopSeq(), and startLoopSeq().

◆ genInvariantValue()

static Value genInvariantValue	(	CodegenEnv &	env,
		ExprId	exp
	)

inlinestatic

Generates an invariant value.

Definition at line 584 of file Sparsification.cpp.

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

Referenced by genExp().

◆ genLoop()

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

References mlir::sparse_tensor::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 1363 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::sparse_tensor::CodegenEnv::getInsertionChain(), mlir::sparse_tensor::Merger::getOutTensorID(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::Value::getType(), mlir::sparse_tensor::LoopEmitter::getValBuffer(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::RewriterBase::replaceOpWithNewOp().

◆ genStmt()

static void genStmt	(	CodegenEnv &	env,
		RewriterBase &	rewriter,
		ExprId	exp,
		LoopId	curr
	)

static

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.

Definition at line 1270 of file Sparsification.cpp.

◆ genSubscript()

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

static

Generates subscript for load/store on a dense or sparse tensor.

Definition at line 349 of file Sparsification.cpp.

Referenced by genTensorLoad(), genTensorStore(), and relinkBranch().

◆ genTensorLoad()

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

static

◆ genTensorStore()

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

static

Generates a store on a dense or sparse tensor.

Definition at line 531 of file Sparsification.cpp.

Referenced by genInvariants(), and genStmt().

◆ getAllTidLvlsInLatPoints()

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

static

Definition at line 1004 of file Sparsification.cpp.

Referenced by startLoopSeq(), and translateBitsToTidLvlPairs().

◆ 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(), and hasNonTrivialAffineOnSparseOut().

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

References getNumNonTrivialIdxExpOnSparseLvls().

◆ hasNonTrivialAffineOnSparseOut()

static bool hasNonTrivialAffineOnSparseOut ( linalg::GenericOp op )

static

Definition at line 231 of file Sparsification.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), getNumNonTrivialIdxExpOnSparseLvls(), mlir::sparse_tensor::getSparseTensorType(), and mlir::sparse_tensor::SparseTensorType::isAllDense().

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

Referenced by genInvariants(), and getAllTidLvlsInLatPoints().

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

Referenced by shouldTryParallize().

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

References mlir::OpBuilder::create(), genSubscript(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Value::getDefiningOp(), mlir::sparse_tensor::CodegenEnv::getLoopVar(), mlir::sparse_tensor::getSparseTensorType(), mlir::sparse_tensor::SparseTensorType::hasEncoding(), mlir::sparse_tensor::CodegenEnv::makeLoopId(), mlir::sparse_tensor::CodegenEnv::makeTensorId(), mlir::RewriterBase::modifyOpInPlace(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::OpBuilder::setInsertionPoint().

Referenced by genExp().

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

References mlir::sparse_tensor::LevelType::hasSparseSemantic(), isParallelFor(), mlir::sparse_tensor::CodegenEnv::lt(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::sparse_tensor::CodegenEnv::unpackTensorLevel().

Referenced by genLoop().

◆ startLoop()

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

static

Starts a single loop in current sequence.

Definition at line 1199 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), genConstantDenseAddressFromLevel(), genLoop(), mlir::sparse_tensor::Merger::getOutTensorID(), mlir::sparse_tensor::Merger::getSynTensorID(), mlir::sparse_tensor::LoopEmitter::locateLvlAtAffineAddress(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), translateBitsToTidLvlPairs(), and mlir::sparse_tensor::CodegenEnv::unpackTensorLevelRange().

Referenced by genStmt().

◆ startLoopSeq()

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

static

Starts a loop sequence at given level.

Returns true if the universal loop index must be maintained at this level.

Definition at line 1116 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::sparse_tensor::LoopEmitter::enterNewLoopSeq(), genExpand(), genInvariants(), getAllTidLvlsInLatPoints(), mlir::sparse_tensor::CodegenEnv::getLoopVar(), mlir::sparse_tensor::Merger::hasAnySparse(), mlir::sparse_tensor::CodegenEnv::lat(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), mlir::sparse_tensor::CodegenEnv::set(), and mlir::sparse_tensor::LatPoint::simple.

Referenced by genStmt().

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

References getAllTidLvlsInLatPoints().

Referenced by startLoop().

Functions

Function Documentation

◆ endIf()

◆ endLoop()

◆ endLoopSeq()

◆ finalizeWhileOp()

◆ findAffine()

◆ findDepIdxSet()

◆ findSparseAnnotations()

◆ genBuffers()

◆ genCoIteration()

◆ genCoIterationCase()

◆ genConditionalInsert()

◆ genConstantDenseAddressFromLevel()

◆ genExp()

◆ genExpand()

◆ genIf()

◆ genIndex()

◆ genInitConstantDenseAddress()

◆ genInsertionLoad()

◆ genInsertionLoadReduce()

◆ genInsertionStore()

◆ genInvariants()

◆ genInvariantValue()

◆ genLoop()

◆ genResult()

◆ genStmt()

◆ genSubscript()

◆ genTensorLoad()

◆ genTensorStore()

◆ getAllTidLvlsInLatPoints()

◆ getNumNonTrivialIdxExpOnSparseLvls() [1/2]

◆ getNumNonTrivialIdxExpOnSparseLvls() [2/2]

◆ hasNonTrivialAffineOnSparseOut()

◆ isInvariantAffine()

◆ isParallelFor()

◆ relinkBranch()

◆ shouldTryParallize()

◆ startLoop()

◆ startLoopSeq()

◆ translateBitsToTidLvlPairs()