13 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
14 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_
32 namespace sparse_tensor {
77 template <
class SubClass>
81 : retTypes(retTypes), params(params), genCall(genCall) {}
87 return genImplementation(retTypes, params, builder, loc);
90 std::string funcName = getMangledFuncName();
91 ModuleOp module = getParentOpOf<ModuleOp>(builder);
94 auto func = module.lookupSymbol<func::FuncOp>(result.getAttr());
100 func = builder.
create<func::FuncOp>(
105 Block *entryBB = func.addEntryBlock();
110 genImplementation(retTypes, args, builder, loc);
111 builder.
create<func::ReturnOp>(loc, result);
114 func::CallOp call = builder.
create<func::CallOp>(loc, func, params);
119 template <
class OpTp>
125 std::string getMangledFuncName() {
126 return static_cast<SubClass *
>(
this)->getMangledFuncName();
130 SmallVector<Value> genImplementation(TypeRange retTypes, ValueRange params,
131 OpBuilder &builder, Location loc) {
132 return static_cast<SubClass *
>(
this)->genImplementation(retTypes, params,
143 Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy);
158 TypedAttr
getOneAttr(Builder &builder, Type tp);
169 SmallVectorImpl<Value> &dstShape,
170 ArrayRef<Value> srcShape, ArrayRef<Size> staticDstShape,
171 ArrayRef<ReassociationIndices> reassociation);
174 void reshapeCvs(OpBuilder &builder, Location loc,
175 ArrayRef<ReassociationIndices> reassociation,
176 ValueRange srcSizes, ValueRange srcCvs,
177 ValueRange dstSizes, SmallVectorImpl<Value> &dstCvs);
183 FlatSymbolRefAttr
getFunc(ModuleOp module, StringRef name, TypeRange resultType,
188 func::CallOp
createFuncCall(OpBuilder &builder, Location loc, StringRef name,
189 TypeRange resultType, ValueRange operands,
200 Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp);
205 Value genAlloca(OpBuilder &builder, Location loc,
unsigned sz, Type tp,
206 bool staticShape =
false);
222 RankedTensorType tensorTp, ValueRange sizes);
228 void sizesFromSrc(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
229 Location loc, Value src);
237 {ShapedType::kDynamic})
238 : StridedLayoutAttr();
270 Value mem,
size_t offsetIdx = 0,
278 size_t offsetIdx = 0,
Value offsetVal =
Value());
317 Value &dimSizesBuffer);
323 Value &dim2lvlBuffer,
324 Value &lvl2dimBuffer);
340 if (
auto ctp = dyn_cast<ComplexType>(tp)) {
341 auto zeroe = builder.
getZeroAttr(ctp.getElementType());
343 return builder.
create<complex::ConstantOp>(loc, tp, zeroa);
351 if (
auto ctp = dyn_cast<ComplexType>(tp)) {
352 auto zeroe = builder.
getZeroAttr(ctp.getElementType());
353 auto onee =
getOneAttr(builder, ctp.getElementType());
355 return builder.
create<complex::ConstantOp>(loc, tp, zeroa);
392 return constantI32(builder, loc,
static_cast<uint32_t
>(action));
405 SparseTensorEncodingAttr enc) {
412 SparseTensorEncodingAttr enc) {
426 return constantI8(builder, loc,
static_cast<uint8_t
>(lt));
430 auto rtp = dyn_cast<RankedTensorType>(type);
431 return !rtp || rtp.getRank() == 0;
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Block represents an ordered list of Operations.
Region * getParent() const
Provide a 'getParent' method for ilist_node_with_parent methods.
BlockArgListType getArguments()
This class is a general helper class for creating context-global objects like types,...
TypedAttr getZeroAttr(Type type)
ArrayAttr getArrayAttr(ArrayRef< Attribute > value)
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
MLIRContext is the top-level object for a collection of MLIR operations.
RAII guard to reset the insertion point of the builder when destroyed.
This class helps build Operations.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
void setInsertionPoint(Block *block, Block::iterator insertPoint)
Set the insertion point to the specified location.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Block * getInsertionBlock() const
Return the block the current insertion point belongs to.
Operation is the basic unit of execution within MLIR.
result_range getResults()
ParentT getParentOfType()
Find the first parent operation of the given type, or nullptr if there is no ancestor operation.
This class provides an abstraction over the various different ranges of value types.
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
MLIRContext * getContext() const
Return the MLIRContext in which this type was uniqued.
This class provides an abstraction over the different types of ranges over Values.
type_range getTypes() const
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Specialization of arith.constant op that returns an integer of index type.
Specialization of arith.constant op that returns an integer value.
A helper class to simplify lowering operations with/without function calls.
FuncCallOrInlineGenerator(TypeRange retTypes, ValueRange params, bool genCall)
SmallVector< Value > genCallOrInline(OpBuilder &builder, Location loc)
@ Type
An inlay hint that for a type annotation.
TypedAttr getOneAttr(Builder &builder, Type tp)
Generates a 1-valued attribute of the given type.
Value genToCoordinatesBuffer(OpBuilder &builder, Location loc, Value tensor)
Infers the result type and generates ToCoordinatesBufferOp.
FlatSymbolRefAttr getFunc(ModuleOp module, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)
Returns a function reference (first hit also inserts into module).
Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp)
Generates an uninitialized temporary buffer with room for one value of the given type,...
Value constantIndex(OpBuilder &builder, Location loc, int64_t i)
Generates a constant of index type.
void foreachInSparseConstant(OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order, function_ref< void(ArrayRef< Value >, Value)> callback)
Iterate over a sparse constant, generates constantOp for value and coordinates.
Value constantZero(OpBuilder &builder, Location loc, Type tp)
Generates a 0-valued constant of the given type.
Value constantLevelTypeEncoding(OpBuilder &builder, Location loc, LevelType lt)
Generates a constant of the internal dimension level type encoding.
Value allocaBuffer(OpBuilder &builder, Location loc, ValueRange values)
Generates a temporary buffer, initializes it with the given contents, and returns it as type memref<?...
Value constantPosTypeEncoding(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc)
Generates a constant of the internal type-encoding for position overhead storage.
OverheadType posTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for position overhead storage.
Value constantOne(OpBuilder &builder, Location loc, Type tp)
Generates a 1-valued constant of the given type.
OverheadType
Encoding of overhead types (both position overhead and coordinate overhead), for "overloading" @newSp...
Value constantI8(OpBuilder &builder, Location loc, int8_t i)
Generates a constant of i8 type.
Action
The actions performed by @newSparseTensor.
uint64_t Dimension
The type of dimension identifiers and dimension-ranks.
OverheadType crdTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for coordinate overhead storage.
Value genToValues(OpBuilder &builder, Location loc, Value tensor)
Infers the result type and generates ToValuesOp.
uint64_t Level
The type of level identifiers and level-ranks.
TypedValue< BaseMemRefType > genToMemref(OpBuilder &builder, Location loc, Value tensor)
Value genToPositions(OpBuilder &builder, Location loc, Value tensor, Level lvl)
Infers the result type and generates ToPositionsOp.
OverheadType overheadTypeEncoding(unsigned width)
Converts an overhead storage bitwidth to its internal type-encoding.
Value constantI1(OpBuilder &builder, Location loc, bool b)
Generates a constant of i1 type.
Value constantAction(OpBuilder &builder, Location loc, Action action)
Generates a constant of the given Action.
MemRefType get1DMemRefType(Type etp, bool withLayout)
Generates a 1D MemRefType with a dynamic size.
Value constantCrdTypeEncoding(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc)
Generates a constant of the internal type-encoding for coordinate overhead storage.
StringRef overheadTypeFunctionSuffix(OverheadType ot)
Convert OverheadType to its function-name suffix.
PrimaryType
Encoding of the elemental type, for "overloading" @newSparseTensor.
PrimaryType primaryTypeEncoding(Type elemTp)
Converts a primary storage type to its internal type-encoding.
Operation * getTop(Operation *op)
Scans to top of generated loop.
Value createOrFoldSliceStrideOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)
Generates code to retrieve the slice slice for the sparse tensor slice, return a constant if the offs...
Type getOpaquePointerType(MLIRContext *ctx)
Returns the equivalent of void* for opaque arguments to the execution engine.
LevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Value genMapBuffers(OpBuilder &builder, Location loc, SparseTensorType stt, ArrayRef< Value > dimSizesValues, Value dimSizesBuffer, SmallVectorImpl< Value > &lvlSizesValues, Value &dim2lvlBuffer, Value &lvl2dimBuffer)
Generates code to set up the buffer parameters for a map.
Value genIsNonzero(OpBuilder &builder, Location loc, Value v)
Generates the comparison v != 0 where v is of numeric type.
Value genIndexLoad(OpBuilder &builder, Location loc, Value mem, Value s)
Generates a pointer/index load from the sparse storage scheme.
Value genReader(OpBuilder &builder, Location loc, SparseTensorType stt, Value tensor, SmallVectorImpl< Value > &dimSizesValues, Value &dimSizesBuffer)
Generates code that opens a reader and sets the dimension sizes.
Value genScalarToTensor(OpBuilder &builder, Location loc, Value elem, Type dstTp)
Add conversion from scalar to given type (possibly a 0-rank tensor).
bool isZeroRankedTensorOrScalar(Type type)
void deallocDenseTensor(OpBuilder &builder, Location loc, Value buffer)
Generates code to deallocate a dense buffer.
Value genAlloca(OpBuilder &builder, Location loc, Value sz, Type tp)
Generates an uninitialized temporary buffer of the given size and type, but returns it as type memref...
Value constantOverheadTypeEncoding(OpBuilder &builder, Location loc, unsigned width)
Generates a constant of the internal type-encoding for overhead storage.
SmallVector< Value > loadAll(OpBuilder &builder, Location loc, size_t size, Value mem, size_t offsetIdx=0, Value offsetVal=Value())
Loads size-many values from the memref, which must have rank-1 and size greater-or-equal to size.
Value genToCoordinates(OpBuilder &builder, Location loc, Value tensor, Level lvl, Level cooStart)
Infers the result type and generates ToCoordinatesOp.
void genReshapeDstShape(OpBuilder &builder, Location loc, SmallVectorImpl< Value > &dstShape, ArrayRef< Value > srcShape, ArrayRef< Size > staticDstShape, ArrayRef< ReassociationIndices > reassociation)
Computes the shape of destination tensor of a reshape operator.
Value constantI64(OpBuilder &builder, Location loc, int64_t i)
Generates a constant of i64 type.
void reshapeCvs(OpBuilder &builder, Location loc, ArrayRef< ReassociationIndices > reassociation, ValueRange srcSizes, ValueRange srcCvs, ValueRange dstSizes, SmallVectorImpl< Value > &dstCvs)
Reshape coordinates during a reshaping operation.
Value constantI16(OpBuilder &builder, Location loc, int16_t i)
Generates a constant of i16 type.
func::CallOp createFuncCall(OpBuilder &builder, Location loc, StringRef name, TypeRange resultType, ValueRange operands, EmitCInterface emitCInterface)
Creates a CallOp to the function reference returned by getFunc() in the builder's module.
Value genCast(OpBuilder &builder, Location loc, Value value, Type dstTy)
Add type casting between arith and index types when needed.
StringRef primaryTypeFunctionSuffix(PrimaryType pt)
Convert PrimaryType to its function-name suffix.
Value createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc, Value tensor, Dimension dim)
Generates code to retrieve the slice offset for the sparse tensor slice, return a constant if the off...
Value constantI32(OpBuilder &builder, Location loc, int32_t i)
Generates a constant of i32 type.
Value constantPrimaryTypeEncoding(OpBuilder &builder, Location loc, Type elemTp)
Generates a constant of the internal type-encoding for primary storage.
void sizesFromSrc(OpBuilder &builder, SmallVectorImpl< Value > &sizes, Location loc, Value src)
Populates given sizes array from dense tensor or sparse tensor constant.
Type getOverheadType(Builder &builder, OverheadType ot)
Converts the internal type-encoding for overhead storage to an mlir::Type.
Value genValMemSize(OpBuilder &builder, Location loc, Value tensor)
Generates code to retrieve the values size for the sparse tensor.
EmitCInterface
Shorthand aliases for the emitCInterface argument to getFunc(), createFuncCall(), and replaceOpWithFu...
Value allocDenseTensor(OpBuilder &builder, Location loc, RankedTensorType tensorTp, ValueRange sizes)
Generates code to allocate a buffer of the given type, and zero initialize it.
void storeAll(OpBuilder &builder, Location loc, Value mem, ValueRange vs, size_t offsetIdx=0, Value offsetVal=Value())
Stores all the values of vs into the memref mem, which must have rank-1 and size greater-or-equal to ...
Include the generated interface declarations.
std::conditional_t< std::is_same_v< Ty, mlir::Type >, mlir::Value, detail::TypedValue< Ty > > TypedValue
If Ty is mlir::Type this will select Value instead of having a wrapper around it.
auto get(MLIRContext *context, Ts &&...params)
Helper method that injects context only if needed, this helps unify some of the attribute constructio...