28 static std::optional<std::pair<Value, Value>>
30 if (
auto constOp = tensor.
getDefiningOp<arith::ConstantOp>()) {
31 if (
auto a = dyn_cast<SparseElementsAttr>(constOp.getValue())) {
32 auto coordinates = builder.
create<arith::ConstantOp>(loc, a.getIndices());
33 auto values = builder.
create<arith::ConstantOp>(loc, a.getValues());
34 return std::make_pair(coordinates, values);
48 for (
unsigned d = 0; d < rank; d++) {
56 return builder.
create<tensor::ExtractOp>(loc, values, k);
70 cvs.append(ivs.begin(), ivs.end());
91 llvm_unreachable(
"Unsupported overhead bitwidth");
97 if (
auto intTp = dyn_cast<IntegerType>(tp))
99 llvm_unreachable(
"Unknown overhead type");
117 llvm_unreachable(
"Unknown OverheadType");
142 #define CASE(ONAME, O) \
143 case OverheadType::kU##ONAME: \
148 llvm_unreachable(
"Unknown OverheadType");
172 if (
auto complexTp = dyn_cast<ComplexType>(elemTp)) {
173 auto complexEltTp = complexTp.getElementType();
174 if (complexEltTp.isF64())
176 if (complexEltTp.isF32())
179 llvm_unreachable(
"Unknown primary type");
184 #define CASE(VNAME, V) \
185 case PrimaryType::k##VNAME: \
190 llvm_unreachable(
"Unknown PrimaryType");
208 if (isa<IndexType>(srcTp) || isa<IndexType>(dstTp))
209 return builder.
create<arith::IndexCastOp>(loc, dstTp, value);
211 const auto srcIntTp = dyn_cast_or_null<IntegerType>(srcTp);
212 const bool isUnsignedCast = srcIntTp ? srcIntTp.isUnsigned() :
false;
218 Value load = builder.
create<memref::LoadOp>(loc, mem, s);
219 if (!isa<IndexType>(load.
getType())) {
229 if (isa<FloatType>(tp))
231 if (isa<IndexType>(tp))
233 if (
auto intTp = dyn_cast<IntegerType>(tp))
235 if (isa<RankedTensorType, VectorType>(tp)) {
236 auto shapedTp = cast<ShapedType>(tp);
237 if (
auto one =
getOneAttr(builder, shapedTp.getElementType()))
240 llvm_unreachable(
"Unsupported attribute type");
247 if (isa<FloatType>(tp))
248 return builder.
create<arith::CmpFOp>(loc, arith::CmpFPredicate::UNE, v,
251 return builder.
create<arith::CmpIOp>(loc, arith::CmpIPredicate::ne, v,
253 if (dyn_cast<ComplexType>(tp))
254 return builder.
create<complex::NotEqualOp>(loc, v, zero);
255 llvm_unreachable(
"Non-numeric type");
263 if (reassociation.size() < srcShape.size()) {
267 for (
unsigned i = start; i < start + map.value().size(); i++) {
268 dstDim = builder.
create<arith::MulIOp>(loc, dstDim, srcShape[i]);
270 dstShape.push_back(dstDim);
271 start = start + map.value().size();
273 assert(start == srcShape.size());
278 assert(reassociation.size() == srcShape.size());
281 for (
unsigned i = 0, size = srcShape.size(); i < size; i++) {
282 const auto &map = reassociation[i];
283 auto srcDim = srcShape[i];
285 for (
unsigned j = start;
j < start + map.size();
j++) {
290 if (staticDstShape[
j] == ShapedType::kDynamic) {
294 for (
unsigned k = start; k < start + map.size(); k++) {
295 if (staticDstShape[k] != ShapedType::kDynamic) {
302 builder.
create<arith::DivUIOp>(loc, srcDim, productVal);
303 dstShape.push_back(dynamicSize);
306 dstShape.push_back(
constantIndex(builder, loc, staticDstShape[
j]));
309 start = start + map.size();
311 assert(start == staticDstShape.size());
319 const unsigned srcRank = srcSizes.size();
320 const unsigned dstRank = dstSizes.size();
321 assert(srcRank == srcCvs.size() &&
"Source rank mismatch");
322 const bool isCollapse = srcRank > dstRank;
323 const ValueRange sizes = isCollapse ? srcSizes : dstSizes;
330 for (
unsigned j = start, end = start + map.value().size();
j < end;
j++) {
331 linear = builder.
create<arith::MulIOp>(loc, linear, sizes[
j]);
338 for (
unsigned j = start, end = start + map.value().size();
j < end;
j++) {
339 linear = builder.
create<arith::DivUIOp>(loc, linear, sizes[
j]);
341 const Value mul = builder.
create<arith::MulIOp>(loc, srcCvs[
j], linear);
342 val = val ? builder.
create<arith::AddIOp>(loc, val, mul) : mul;
344 const Value old = val;
345 val = builder.
create<arith::DivUIOp>(loc, val, linear);
346 assert(dstCvs.size() ==
j);
347 dstCvs.push_back(val);
348 val = builder.
create<arith::RemUIOp>(loc, old, linear);
353 assert(dstCvs.size() == i);
354 dstCvs.push_back(val);
356 start += map.value().size();
359 assert(dstCvs.size() == dstRank);
368 auto func = module.lookupSymbol<func::FuncOp>(result.getAttr());
370 OpBuilder moduleBuilder(module.getBodyRegion());
371 func = moduleBuilder.
create<func::FuncOp>(
372 module.getLoc(), name,
375 if (
static_cast<bool>(emitCInterface))
376 func->
setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),
387 getFunc(module, name, resultType, operands, emitCInterface);
388 return builder.
create<func::CallOp>(loc, resultType, fn, operands);
400 unsigned sz,
Type tp,
bool staticShape) {
403 return builder.
create<memref::AllocaOp>(loc, memTp);
421 const unsigned sz = values.size();
424 for (
unsigned i = 0; i < sz; i++) {
426 builder.
create<memref::StoreOp>(loc, values[i], buffer, idx);
432 RankedTensorType tensorTp,
434 Type elemTp = tensorTp.getElementType();
435 auto shape = tensorTp.getShape();
438 for (
unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {
439 if (shape[i] == ShapedType::kDynamic)
440 dynamicSizes.push_back(sizes[i]);
442 Value mem = builder.
create<memref::AllocOp>(loc, memTp, dynamicSizes);
450 builder.
create<memref::DeallocOp>(loc, buffer);
455 Value val = builder.
create<tensor::ExtractOp>(loc, tensor, ivs);
457 scf::IfOp ifOp = builder.
create<scf::IfOp>(loc, cond,
false);
480 if (splitSrc.has_value()) {
481 const Value srcCoordinates = splitSrc->first;
482 const Value srcValues = splitSrc->second;
490 builder, loc, srcCoordinates, srcValues, cvs,
492 bodyBuilder(builder, loc, val, cvs);
496 for (
unsigned i = 0; i < rank; i++) {
506 bodyBuilder(builder, loc, val, cvs);
534 const auto coordinates = attr.getIndices().getValues<IntegerAttr>();
535 const auto values = attr.getValues().getValues<
Attribute>();
542 using ElementAttr = std::pair<SmallVector<IntegerAttr>,
Attribute>;
546 for (
size_t i = 0, nse = values.size(); i < nse; i++) {
547 elems.emplace_back();
548 elems.back().second = values[i];
549 auto &coords = elems.back().first;
550 coords.reserve(dimRank);
552 coords.push_back(coordinates[i * dimRank + d]);
556 std::sort(elems.begin(), elems.end(),
557 [order, dimRank](
const ElementAttr &lhs,
const ElementAttr &rhs) {
558 const auto &lhsCoords = lhs.first;
559 const auto &rhsCoords = rhs.first;
560 for (Dimension d = 0; d < dimRank; d++) {
564 const Level l = order ? order.getDimPosition(d) : d;
565 if (lhsCoords[l].getInt() == rhsCoords[l].getInt())
567 return lhsCoords[l].getInt() < rhsCoords[l].getInt();
569 if (std::addressof(lhs) == std::addressof(rhs))
571 llvm_unreachable(
"no equal coordinate in sparse element attr");
575 cvs.reserve(dimRank);
576 for (
size_t i = 0, nse = values.size(); i < nse; i++) {
579 for (
Dimension d = 0; d < dimRank; d++) {
580 auto crd = elems[i].first[d].getInt();
585 if (isa<ComplexType>(attr.getElementType())) {
586 auto valAttr = cast<ArrayAttr>(elems[i].second);
587 val = builder.
create<complex::ConstantOp>(loc, attr.getElementType(),
590 auto valAttr = cast<TypedAttr>(elems[i].second);
591 val = builder.
create<arith::ConstantOp>(loc, valAttr);
599 size_t size,
Value mem,
600 size_t offsetIdx,
Value offsetVal) {
602 const auto memTp = cast<MemRefType>(mem.
getType());
603 assert(memTp.getRank() == 1);
604 const DynSize memSh = memTp.getDimSize(0);
605 assert(ShapedType::isDynamic(memSh) || memSh >=
static_cast<DynSize>(size));
606 assert(offsetIdx == 0 || offsetIdx < size);
610 for (
unsigned i = 0; i < size; i++) {
613 if (i == offsetIdx && offsetVal)
614 v = builder.
create<arith::AddIOp>(loc, v, offsetVal);
623 const size_t vsize = vs.size();
624 const auto memTp = cast<MemRefType>(mem.
getType());
625 assert(memTp.getRank() == 1);
626 const DynSize memSh = memTp.getDimSize(0);
627 assert(ShapedType::isDynamic(memSh) || memSh >=
static_cast<DynSize>(vsize));
628 assert(offsetIdx == 0 || offsetIdx < vsize);
632 (offsetIdx == v.index() && offsetVal)
633 ? builder.
create<arith::AddIOp>(loc, v.value(), offsetVal)
635 builder.
create<memref::StoreOp>(loc, w, mem,
641 SparseTensorEncodingAttr enc,
646 const Level lvlRank = enc.getLvlRank();
648 lvlSizes.reserve(lvlRank);
649 for (
Level l = 0; l < lvlRank; l++)
651 lvlSizes.push_back(dimSizes[
toOrigDim(enc, l)]);
652 storeAll(builder, loc, lvlCoords, lvlSizes);
658 lvlCoords = builder.
create<memref::CastOp>(loc, lvlSizesTp, lvlCoords);
663 return builder.
create<memref::ReshapeOp>(loc, resTp, valuesBuffer, lvlCoords);
668 auto tTp = llvm::cast<TensorType>(tensor.
getType());
670 return builder.
create<bufferization::ToMemrefOp>(loc, mTp, tensor)
677 const Type posTp = srcTp.getPosType();
679 return builder.
create<ToPositionsOp>(loc, memTp, tensor,
686 const Type crdTp = srcTp.getCrdType();
688 return builder.
create<ToCoordinatesOp>(loc, memTp, tensor,
695 const Type crdTp = srcTp.getCrdType();
697 return builder.
create<ToCoordinatesBufferOp>(loc, memTp, tensor);
705 return builder.
create<ToValuesOp>(loc, valTp, tensor);
716 assert(enc && enc.isSlice());
717 std::optional<unsigned> offset = enc.getStaticDimSliceOffset(dim);
718 if (offset.has_value())
720 return builder.
create<ToSliceOffsetOp>(loc, tensor, APInt(64, dim));
726 assert(enc && enc.isSlice());
727 std::optional<unsigned> stride = enc.getStaticDimSliceStride(dim);
728 if (stride.has_value())
730 return builder.
create<ToSliceStrideOp>(loc, tensor, APInt(64, dim));
static Value genCoordsAndValueForDense(OpBuilder &builder, Location loc, Value tensor, SmallVectorImpl< Value > &cvs, ValueRange ivs)
Generates code to read the value from tensor[ivs], and open a conditional for whether the value is no...
static Value genCoordsAndValueForSparse(OpBuilder &builder, Location loc, Value coordinates, Value values, SmallVectorImpl< Value > &cvs, Value k, unsigned rank)
Reads coordinates[k][0..rank-1] and value[k], appending the former onto cvs and returning the latter.
static std::optional< std::pair< Value, Value > > genSplitSparseConstant(OpBuilder &builder, Location loc, Value tensor)
If the tensor is a sparse constant, generates and returns the pair of the constants for the coordinat...
#define MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DO)
#define MLIR_SPARSETENSOR_FOREVERY_V(DO)
static int64_t product(ArrayRef< int64_t > vals)
A multi-dimensional affine map Affine map's are immutable like Type's, and they are uniqued.
Attributes are known-constant values of operations.
Type getElementType() const
Returns the element type of this memref type.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
This class is a general helper class for creating context-global objects like types,...
IntegerAttr getIndexAttr(int64_t value)
IntegerAttr getIntegerAttr(Type type, int64_t value)
FloatAttr getFloatAttr(Type type, double value)
IntegerType getIntegerType(unsigned width)
MLIRContext * getContext() const
static DenseElementsAttr get(ShapedType type, ArrayRef< Attribute > values)
Constructs a dense elements attribute from an array of element values.
A symbol reference with a reference path containing a single element.
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.
This class helps build Operations.
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Block * getBlock() const
Returns the current block of the builder.
Operation * create(const OperationState &state)
Creates an operation given the fields represented as an OperationState.
Operation is the basic unit of execution within MLIR.
Operation * getParentOp()
Returns the closest surrounding operation that contains this operation or nullptr if this is a top-le...
OpTy getParentOfType()
Return the closest surrounding parent operation that is of type 'OpTy'.
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
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...
bool isInteger(unsigned width) const
Return true if this is an integer type with the specified width.
bool isIntOrIndex() const
Return true if this is an integer (of any signedness) or an index type.
unsigned getIntOrFloatBitWidth() const
Return the bit width of an integer or a float type, assert failure on other types.
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...
Type getType() const
Return the type of this value.
Operation * getDefiningOp() const
If this value is the result of an operation, return the operation that defines it.
Specialization of arith.constant op that returns an integer of index type.
Value getValMemSize(OpBuilder &builder, Location loc) const
Dimension getDimRank() const
Returns the dimension-rank.
BaseMemRefType getMemRefType(Value value, const BufferizationOptions &options, MemRefLayoutAttrInterface layout={}, Attribute memorySpace=nullptr)
Return a MemRefType to which the type of the given value can be bufferized.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Value createOrFoldDimOp(OpBuilder &b, Location loc, Value val, int64_t dim)
Create one memref::DimOp or tensor::DimOp depending on the type of val.
LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs, ValueRange ubs, ValueRange steps, ValueRange iterArgs, function_ref< ValueVector(OpBuilder &, Location, ValueRange, ValueRange)> bodyBuilder=nullptr)
Creates a perfect nest of "for" loops, i.e.
SmallVector< Value > ValueVector
An owning vector of values, handy to return from functions.
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 allocaBuffer(OpBuilder &builder, Location loc, ValueRange values)
Generates a temporary buffer, initializes it with the given contents, and returns it as type memref<?...
Value reshapeValuesToLevels(OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc, ValueRange dimSizes, Value valuesBuffer, Value lvlCoords)
Reshapes the linear values buffer for an annotated all dense sparse tensor to match the shape of the ...
OverheadType posTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for position overhead storage.
OverheadType
Encoding of overhead types (both position overhead and coordinate overhead), for "overloading" @newSp...
uint64_t Dimension
The type of dimension identifiers, and dimension-ranks.
OverheadType crdTypeEncoding(SparseTensorEncodingAttr enc)
Returns the OverheadType for coordinate overhead storage.
SparseTensorType getSparseTensorType(T t)
Convenience method to abbreviate wrapping getRankedTensorType.
Value genToValues(OpBuilder &builder, Location loc, Value tensor)
Infers the result type and generates ToValuesOp.
Value genValueForDense(OpBuilder &builder, Location loc, Value tensor, ValueRange ivs)
Generates code to read the value from tensor[ivs].
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.
MemRefType get1DMemRefType(Type etp, bool withLayout)
Generates a 1D MemRefType with a dynamic size.
StringRef overheadTypeFunctionSuffix(OverheadType ot)
Convert OverheadType to its function-name suffix.
PrimaryType
Encoding of the elemental type, for "overloading" @newSparseTensor.
RankedTensorType getRankedTensorType(T &&t)
Convenience method to abbreviate casting getType().
int64_t DynSize
The type for individual components of a compile-time shape.
PrimaryType primaryTypeEncoding(Type elemTp)
Converts a primary storage type to its internal type-encoding.
DEPRECATED Dimension toOrigDim(SparseTensorEncodingAttr enc, Level l)
[deprecated] Convenience method to translate the given level to the corresponding dimension.
Operation * getTop(Operation *op)
Scans to top of generated loop.
void genDenseTensorOrSparseConstantIterLoop(OpBuilder &builder, Location loc, Value src, unsigned rank, function_ref< void(OpBuilder &, Location, Value, ValueRange)> bodyBuilder)
Generates the loop structure to iterate over a dense tensor or a sparse tensor constant to support th...
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.
SparseTensorEncodingAttr getSparseTensorEncoding(Type type)
Convenience method to get a sparse encoding attribute from a type.
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.
SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor)
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...
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.
int64_t StaticSize
The type for individual components of a compile-time shape which are known not to be ShapedType::kDyn...
void reshapeCvs(OpBuilder &builder, Location loc, ArrayRef< ReassociationIndices > reassociation, ValueRange srcSizes, ValueRange srcCvs, ValueRange dstSizes, SmallVectorImpl< Value > &dstCvs)
Reshape coordinates during a reshaping operation.
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...
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 ...
void genReshapeDstShape(OpBuilder &builder, Location loc, SmallVectorImpl< Value > &dstShape, ArrayRef< Value > srcShape, ArrayRef< StaticSize > staticDstShape, ArrayRef< ReassociationIndices > reassociation)
Computes the shape of destination tensor of a reshape operator.
This header declares functions that assist transformations in the MemRef dialect.
Value convertScalarToDtype(OpBuilder &b, Location loc, Value operand, Type toType, bool isUnsignedCast)
Converts a scalar value operand to type toType.
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...
Eliminates variable at the specified position using Fourier-Motzkin variable elimination.