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

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 Operation *	genFilterLoop (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 Operation *	genCoIteration (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 Operation *	genLoop (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

Definition at line 628 of file Sparsification.cpp.

Referenced by addSliceBasedConstraints(), and computeIterationGraph().

◆ addIterOrdering()

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

static

Definition at line 535 of file Sparsification.cpp.

Referenced by addAffineOrderings(), addSliceBasedConstraints(), and computeIterationGraph().

◆ addSliceBasedConstraints()

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

static

Definition at line 690 of file Sparsification.cpp.

References addFilterLoopBasedConstraints(), addIterOrdering(), mlir::IROperand< DerivedT, IRValueT >::get(), getNumNonTrivialIdxExpOnSparseLvls(), mlir::AffineDimExpr::getPosition(), mlir::sparse_tensor::getSparseTensorEncoding(), mlir::Value::getType(), mlir::AffineExpr::isa(), mlir::sparse_tensor::CodegenEnv::makeLoopId(), mlir::sparse_tensor::CodegenEnv::op(), and mlir::sparse_tensor::toOrigDim().

Referenced by computeIterationGraph().

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

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

Referenced by genStmt().

◆ endLoop()

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

static

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

Definition at line 1803 of file Sparsification.cpp.

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

Referenced by genStmt().

◆ endLoopSeq()

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

static

Ends a loop sequence at given level.

Definition at line 1828 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,
		LoopId	idx,
		bool	needsUniv
	)

static

Generates the induction structure for a while-loop.

Definition at line 1450 of file Sparsification.cpp.

Referenced by endLoop().

◆ 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

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

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

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

◆ genInitConstantDenseAddress()

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

static

Definition at line 1643 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 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

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

References mlir::sparse_tensor::CodegenEnv::emitter(), mlir::sparse_tensor::LoopEmitter::genAffine(), mlir::IROperand< DerivedT, IRValueT >::get(), mlir::OpOperand::getOperandNumber(), mlir::sparse_tensor::LoopEmitter::getPosits(), mlir::sparse_tensor::getSparseTensorType(), mlir::sparse_tensor::LoopEmitter::getValBuffer(), mlir::sparse_tensor::CodegenEnv::makeTensorId(), and mlir::sparse_tensor::CodegenEnv::op().

◆ 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

Generates a store on a dense or sparse tensor.

Definition at line 1085 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::exp(), mlir::sparse_tensor::CodegenEnv::isReduc(), mlir::sparse_tensor::TensorExp::kBinary, mlir::sparse_tensor::TensorExp::kind, mlir::sparse_tensor::TensorExp::kReduce, mlir::sparse_tensor::TensorExp::kUnary, and mlir::sparse_tensor::CodegenEnv::updateReduc().

Referenced by genStmt().

◆ 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

Helper method to construct a permuted dimension ordering that adheres to the given topological sort.

Definition at line 179 of file Sparsification.cpp.

References mlir::sparse_tensor::CodegenEnv::getLoopStackUpTo(), mlir::AffineMap::getNumDims(), mlir::sparse_tensor::Merger::getNumFilterLoops(), mlir::AffineMap::getNumResults(), mlir::AffineMap::getPermutationMap(), mlir::AffineMap::getResult(), mlir::AffineExpr::isa(), isInvariantAffine(), mlir::sparse_tensor::CodegenEnv::merger(), mlir::sparse_tensor::CodegenEnv::op(), mlir::sparse_tensor::CodegenEnv::topSortAt(), and mlir::sparse_tensor::CodegenEnv::topSortSize().

◆ 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

Starts a single loop in current sequence.

Definition at line 1769 of file Sparsification.cpp.

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

Referenced by genStmt().

◆ startLoopSeq()

static bool startLoopSeq	(	CodegenEnv &	env,
		OpBuilder &	builder,
		ExprId	exp,
		LoopOrd	at,
		LoopId	idx,
		LoopId	ldx,
		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 1570 of file Sparsification.cpp.

Referenced by genStmt().

◆ 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

Return true if the lattices bit can be iterated by a for loop.

Definition at line 1654 of file Sparsification.cpp.

Referenced by startLoop().

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

Functions

Function Documentation

◆ addAffineOrderings()

◆ addFilterLoopBasedConstraints()

◆ addIterOrdering()

◆ addSliceBasedConstraints()

◆ computeIterationGraph()

◆ endIf()

◆ endLoop()

◆ endLoopSeq()

◆ finalizeWhileOp()

◆ findAffine()

◆ findDepIdxSet()

◆ findSparseAnnotations()

◆ genBuffers()

◆ genCoIteration()

◆ genConstantDenseAddressFromLevel()

◆ genExp()

◆ genExpand()

◆ genFilterLoop()

◆ genIf()

◆ genIndex()

◆ genInitConstantDenseAddress()

◆ genInsertionLoad()

◆ genInsertionLoadReduce()

◆ genInsertionStore()

◆ genInvariants()

◆ genInvariantValue()

◆ genLoop()

◆ genResult()

◆ genStmt()

◆ genSubscript()

◆ genTensorLoad()

◆ genTensorStore()

◆ getNumNonTrivialIdxExpOnSparseLvls() [1/2]

◆ getNumNonTrivialIdxExpOnSparseLvls() [2/2]

◆ hasNonTrivialAffineOnSparseOut()

◆ isInvariantAffine() [1/2]

◆ isInvariantAffine() [2/2]

◆ isParallelFor()

◆ permute()

◆ relinkBranch()

◆ shouldTryParallize()

◆ startLoop()

◆ startLoopSeq()

◆ topSortOptimal()

◆ translateBitsToTidLvlPairs()

◆ tryRelaxAffineConstraints()