#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 SmallVector< Value >	flattenValues (ArrayRef< ValueRange > values)
	Flatten the given value ranges into a single vector of values. 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 99 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 137 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 167 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 150 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 66 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 81 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(), kind, mlir::sparse_tensor::MutSparseTensorDescriptor::setMemRefField(), mlir::sparse_tensor::MutSparseTensorDescriptor::setSpecifierField(), and mlir::sparse_tensor::toSpecifierKind().

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

◆ flattenValues()

static SmallVector<Value> flattenValues ( ArrayRef< ValueRange > values )

static

Flatten the given value ranges into a single vector of values.

Definition at line 43 of file SparseTensorCodegen.cpp.

◆ 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 268 of file SparseTensorCodegen.cpp.

◆ genEndInsert()

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

static

Generates insertion finalization code.

Definition at line 358 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 51 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 405 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 57 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 426 of file SparseTensorCodegen.cpp.

Functions

Function Documentation

◆ allocSchemeForRank()

◆ createAllocation()

◆ createAllocFields()

◆ createDimSizes()

◆ createFor()

◆ createPushback()

◆ flattenValues()

◆ genCompressed()

◆ genEndInsert()

◆ genLoad()

◆ genSliceToSize()

◆ genStore()

◆ getReassociationForFlattening()