MLIR  19.0.0git
Functions
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 <optional>

Go to the source code of this file.

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< ReassociationIndicesgetReassociationForFlattening (ShapedType srcTp, unsigned batchLvls)
 Creates the reassociation array. More...
 

Function Documentation

◆ allocSchemeForRank()

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

◆ 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

◆ 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

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

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.

References allocSchemeForRank(), mlir::sparse_tensor::constantI1(), mlir::sparse_tensor::constantIndex(), mlir::sparse_tensor::CrdMemRef, mlir::OpBuilder::create(), 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()

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

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