SparseTensorCodegen.cpp File Reference

#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 "llvm/ADT/SmallVectorExtras.h"
#include <optional>

Go to the source code of this file.

Functions
static SmallVector< Value >	flattenValues (ArrayRef< ValueRange > values)
	Flatten the given value ranges into a single vector of values.
static Value	genLoad (OpBuilder &builder, Location loc, Value mem, Value idx)
	Generates a load with proper index typing.
static void	genStore (OpBuilder &builder, Location loc, Value val, Value mem, Value idx)
	Generates a store with proper index typing and proper value.
static scf::ForOp	createFor (OpBuilder &builder, Location loc, Value upper, MutableArrayRef< Value > fields, Value lower=Value())
	Creates a straightforward counting for-loop.
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.
static void	allocSchemeForRank (OpBuilder &builder, Location loc, MutSparseTensorDescriptor desc, Level startLvl)
	Generates code that allocates a sparse storage scheme for given rank.
static Value	createAllocation (OpBuilder &builder, Location loc, MemRefType memRefType, Value sz, bool enableInit)
	Creates allocation operation.
static void	createDimSizes (OpBuilder &builder, Location loc, SparseTensorType stt, ValueRange dynSizes, SmallVectorImpl< Value > &dimSizesValues)
	Creates the dim sizes array, filling in from dynamic sizes.
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.
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:
static void	genEndInsert (OpBuilder &builder, Location loc, SparseTensorDescriptor desc)
	Generates insertion finalization code.
static Value	genSliceToSize (OpBuilder &builder, Location loc, Value mem, Value sz)
	Generates a subview into the sizes.
static SmallVector< ReassociationIndices >	getReassociationForFlattening (ShapedType srcTp, unsigned batchLvls)
	Creates the reassociation array.

Function Documentation

allocSchemeForRank()

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

References mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::constantZero(), createPushback(), mlir::sparse_tensor::SparseTensorType::getElementType(), mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getLvlSize(), mlir::sparse_tensor::SparseTensorType::getLvlType(), mlir::sparse_tensor::SparseTensorType::getPosType(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getRankedTensorType(), mlir::sparse_tensor::isCompressedLT(), mlir::sparse_tensor::isDenseLT(), mlir::sparse_tensor::isLooseCompressedLT(), mlir::sparse_tensor::isNOutOfMLT(), mlir::sparse_tensor::isSingletonLT(), mlir::sparse_tensor::PosMemRef, and mlir::sparse_tensor::ValMemRef.

Referenced by createAllocFields(), and genCompressed().

createAllocation()

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

static

Creates allocation operation.

Definition at line 139 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::constantZero().

Referenced by createAllocFields().

createAllocFields()

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

References allocSchemeForRank(), mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::constantZero(), mlir::sparse_tensor::CrdMemRef, createAllocation(), createPushback(), mlir::sparse_tensor::foreachFieldAndTypeInSparseTensor(), mlir::sparse_tensor::SparseTensorType::getAoSCOOStart(), mlir::sparse_tensor::SparseTensorSpecifier::getInitValue(), mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::SparseTensorType::getLvlType(), mlir::sparse_tensor::SparseTensorType::getPosType(), mlir::sparse_tensor::SparseTensorType::isAllDense(), mlir::sparse_tensor::isCompressedLT(), mlir::sparse_tensor::SparseTensorType::isCompressedLvl(), mlir::sparse_tensor::SparseTensorType::isDenseLvl(), mlir::sparse_tensor::isLooseCompressedLT(), mlir::sparse_tensor::PosMemRef, mlir::sparse_tensor::StorageSpec, and mlir::sparse_tensor::ValMemRef.

createDimSizes()

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

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

createFor()

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

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

Referenced by genEndInsert().

createPushback()

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

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

flattenValues()

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

static

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

Definition at line 44 of file SparseTensorCodegen.cpp.

References result.

genCompressed()

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

References allocSchemeForRank(), mlir::sparse_tensor::constantI1(), mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::CrdMemRef, createPushback(), mlir::sparse_tensor::genCast(), genLoad(), genStore(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getCrdMemRefIndexAndStride(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getCrdMemSize(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getField(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getFields(), mlir::Builder::getIndexType(), mlir::Builder::getIntegerType(), mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getMemRefField(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getNumFields(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getPosMemRef(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getRankedTensorType(), mlir::Value::getType(), mlir::sparse_tensor::SparseTensorType::isUniqueLvl(), mlir::sparse_tensor::MutSparseTensorDescriptor::setField(), mlir::OpBuilder::setInsertionPointAfter(), and mlir::OpBuilder::setInsertionPointToStart().

genEndInsert()

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

static

Generates insertion finalization code.

Definition at line 360 of file SparseTensorCodegen.cpp.

References mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::constantZero(), createFor(), genLoad(), genStore(), mlir::sparse_tensor::SparseTensorType::getLvlRank(), mlir::sparse_tensor::SparseTensorType::getLvlType(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getPosMemRef(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getPosMemSize(), mlir::sparse_tensor::SparseTensorType::getPosType(), mlir::sparse_tensor::SparseTensorDescriptorImpl< ValueArrayRef >::getRankedTensorType(), mlir::sparse_tensor::isCompressedLT(), mlir::sparse_tensor::isDenseLT(), mlir::sparse_tensor::isLooseCompressedLT(), mlir::sparse_tensor::isNOutOfMLT(), mlir::sparse_tensor::isSingletonLT(), mlir::OpBuilder::setInsertionPointAfter(), mlir::OpBuilder::setInsertionPointToStart(), and TypeRange.

genLoad()

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

static

Generates a load with proper index typing.

Definition at line 52 of file SparseTensorCodegen.cpp.

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

Referenced by genCompressed(), and genEndInsert().

genSliceToSize()

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

static

Generates a subview into the sizes.

Definition at line 407 of file SparseTensorCodegen.cpp.

References mlir::Value::getType().

genStore()

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

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

Referenced by genCompressed(), and genEndInsert().

getReassociationForFlattening()

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

static

Creates the reassociation array.

Definition at line 428 of file SparseTensorCodegen.cpp.

References push_back().