#include "Utils/CodegenUtils.h"
#include "Utils/SparseTensorDescriptor.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SparseTensor/IR/Enums.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/Tensor/IR/Tensor.h"
#include "mlir/Transforms/DialectConversion.h"
#include <optional>

Functions
static void	flattenOperands (ValueRange operands, SmallVectorImpl< Value > &flattened)
	Flattens a list of operands that may contain sparse tensors. More...

static Value	genLoad (OpBuilder &builder, Location loc, Value mem, Value idx)
	Generates a load with proper `index` typing. More...

static void	genStore (OpBuilder &builder, Location loc, Value val, Value mem, Value idx)
	Generates a store with proper `index` typing and proper value. More...

static scf::ForOp	createFor (OpBuilder &builder, Location loc, Value upper, MutableArrayRef< Value > fields, Value lower=Value())
	Creates a straightforward counting for-loop. More...

static void	createPushback (OpBuilder &builder, Location loc, MutSparseTensorDescriptor desc, SparseTensorFieldKind kind, std::optional< Level > lvl, Value value, Value repeat=Value())
	Creates a push back operation. More...

static void	allocSchemeForRank (OpBuilder &builder, Location loc, MutSparseTensorDescriptor desc, Level startLvl)
	Generates code that allocates a sparse storage scheme for given rank. More...

static Value	createAllocation (OpBuilder &builder, Location loc, MemRefType memRefType, Value sz, bool enableInit)
	Creates allocation operation. More...

static void	createDimSizes (OpBuilder &builder, Location loc, SparseTensorType stt, ValueRange dynSizes, SmallVectorImpl< Value > &dimSizesValues)
	Creates the dim sizes array, filling in from dynamic sizes. More...

static void	createAllocFields (OpBuilder &builder, Location loc, SparseTensorType stt, bool enableInit, Value sizeHint, SmallVectorImpl< Value > &lvlSizesValues, SmallVectorImpl< Value > &fields)
	Creates allocation for each field in sparse tensor type. More...

static Value	genCompressed (OpBuilder &builder, Location loc, MutSparseTensorDescriptor desc, ValueRange lvlCoords, Value, Value parentPos, Level lvl)
	Helper method that generates block specific to compressed case: More...

static void	genEndInsert (OpBuilder &builder, Location loc, SparseTensorDescriptor desc)
	Generates insertion finalization code. More...

static Value	genSliceToSize (OpBuilder &builder, Location loc, Value mem, Value sz)
	Generates a subview into the sizes. More...

static SmallVector< ReassociationIndices >	getReassociationForFlattening (ShapedType srcTp, unsigned batchLvls)
	Creates the reassociation array. More...

Function Documentation

◆ allocSchemeForRank()

static void allocSchemeForRank	(	OpBuilder &	builder,
		Location	loc,
		MutSparseTensorDescriptor	desc,
		Level	startLvl
	)

static

Generates code that allocates a sparse storage scheme for given rank.

Definition at line 112 of file SparseTensorCodegen.cpp.

Referenced by createAllocFields(), and genCompressed().

◆ createAllocation()

static Value createAllocation	(	OpBuilder &	builder,
		Location	loc,
		MemRefType	memRefType,
		Value	sz,
		bool	enableInit
	)

static

Creates allocation operation.

Definition at line 150 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::constantZero(), and mlir::OpBuilder::create().

Referenced by createAllocFields().

◆ createAllocFields()

static void createAllocFields	(	OpBuilder &	builder,
		Location	loc,
		SparseTensorType	stt,
		bool	enableInit,
		Value	sizeHint,
		SmallVectorImpl< Value > &	lvlSizesValues,
		SmallVectorImpl< Value > &	fields
	)

static

Creates allocation for each field in sparse tensor type.

Note that for all dynamic memrefs in the sparse tensor stroage layout, the memory size is really the capacity of the "vector", while the actual size resides in the sizes array.

Definition at line 180 of file SparseTensorCodegen.cpp.

◆ createDimSizes()

static void createDimSizes	(	OpBuilder &	builder,
		Location	loc,
		SparseTensorType	stt,
		ValueRange	dynSizes,
		SmallVectorImpl< Value > &	dimSizesValues
	)

static

Creates the dim sizes array, filling in from dynamic sizes.

Definition at line 163 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::SparseTensorType::getDimRank(), and mlir::sparse_tensor::SparseTensorType::getDimShape().

◆ createFor()

static scf::ForOp createFor	(	OpBuilder &	builder,
		Location	loc,
		Value	upper,
		MutableArrayRef< Value >	fields,
		Value	lower = `Value()`
	)

static

Creates a straightforward counting for-loop.

Definition at line 79 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::constantOne(), mlir::sparse_tensor::constantZero(), mlir::OpBuilder::create(), mlir::Builder::getIndexType(), and mlir::OpBuilder::setInsertionPointToStart().

Referenced by genEndInsert().

◆ createPushback()

static void createPushback	(	OpBuilder &	builder,
		Location	loc,
		MutSparseTensorDescriptor	desc,
		SparseTensorFieldKind	kind,
		std::optional< Level >	lvl,
		Value	value,
		Value	repeat = `Value()`
	)

static

Creates a push back operation.

Definition at line 94 of file SparseTensorCodegen.cpp.

References mlir::OpBuilder::create(), mlir::sparse_tensor::genCast(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getMemRefElementType(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getMemRefField(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getSpecifierField(), mlir::sparse_tensor::MutSparseTensorDescriptor::setMemRefField(), mlir::sparse_tensor::MutSparseTensorDescriptor::setSpecifierField(), and mlir::sparse_tensor::toSpecifierKind().

Referenced by allocSchemeForRank(), createAllocFields(), and genCompressed().

◆ flattenOperands()

static void flattenOperands	(	ValueRange	operands,
		SmallVectorImpl< Value > &	flattened
	)

static

Flattens a list of operands that may contain sparse tensors.

Definition at line 43 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::getSparseTensorEncoding(), and mlir::sparse_tensor::getTuple().

◆ genCompressed()

static Value genCompressed	(	OpBuilder &	builder,
		Location	loc,
		MutSparseTensorDescriptor	desc,
		ValueRange	lvlCoords,
		Value	,
		Value	parentPos,
		Level	lvl
	)

static

Helper method that generates block specific to compressed case:

// given: parentPos = posCursor[lvl-1] pstart = desc.positions[lvl][parentPos] pstop = desc.positions[lvl][parentPos+1] plast = pstop - 1 msz = desc.coordinates[lvl].size() if (pstart < pstop) { isPresent = (desc.coordinates[lvl][plast] == lvlCoords[lvl]) } else { // first insertion isPresent = false desc.positions[lvl][parentPos] = msz } if (isPresent) { // coordinate is already present pnext = plast } else { desc.coordinates[lvl].push_back(lvlCoords[lvl]) desc.positions[lvl][parentPos+1] = msz+1 pnext = msz <prepare level lvl+1> } posCursor[lvl] = pnext

Definition at line 281 of file SparseTensorCodegen.cpp.

◆ genEndInsert()

static void genEndInsert	(	OpBuilder &	builder,
		Location	loc,
		SparseTensorDescriptor	desc
	)

static

Generates insertion finalization code.

Definition at line 371 of file SparseTensorCodegen.cpp.

◆ genLoad()

static Value genLoad	(	OpBuilder &	builder,
		Location	loc,
		Value	mem,
		Value	idx
	)

static

Generates a load with proper index typing.

Definition at line 64 of file SparseTensorCodegen.cpp.

References mlir::OpBuilder::create(), mlir::sparse_tensor::genCast(), and mlir::Builder::getIndexType().

Referenced by genCompressed(), and genEndInsert().

◆ genSliceToSize()

static Value genSliceToSize	(	OpBuilder &	builder,
		Location	loc,
		Value	mem,
		Value	sz
	)

static

Generates a subview into the sizes.

Definition at line 418 of file SparseTensorCodegen.cpp.

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

◆ genStore()

static void genStore	(	OpBuilder &	builder,
		Location	loc,
		Value	val,
		Value	mem,
		Value	idx
	)

static

Generates a store with proper index typing and proper value.

Definition at line 70 of file SparseTensorCodegen.cpp.

References mlir::OpBuilder::create(), mlir::sparse_tensor::genCast(), mlir::Builder::getIndexType(), and mlir::Value::getType().

Referenced by genCompressed(), and genEndInsert().

◆ getReassociationForFlattening()

static SmallVector<ReassociationIndices> getReassociationForFlattening	(	ShapedType	srcTp,
		unsigned	batchLvls
	)

static

Creates the reassociation array.

Definition at line 439 of file SparseTensorCodegen.cpp.

Functions

Function Documentation

◆ allocSchemeForRank()

◆ createAllocation()

◆ createAllocFields()

◆ createDimSizes()

◆ createFor()

◆ createPushback()

◆ flattenOperands()

◆ genCompressed()

◆ genEndInsert()

◆ genLoad()

◆ genSliceToSize()

◆ genStore()

◆ getReassociationForFlattening()