731 lines · cpp
1//===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "CodegenUtils.h"10 11#include "mlir/Dialect/Bufferization/IR/Bufferization.h"12#include "mlir/Dialect/Linalg/IR/Linalg.h"13#include "mlir/Dialect/Linalg/Utils/Utils.h"14#include "mlir/Dialect/MemRef/IR/MemRef.h"15#include "mlir/Dialect/Tensor/IR/Tensor.h"16#include "mlir/IR/Types.h"17#include "mlir/IR/Value.h"18#include <optional>19 20using namespace mlir;21using namespace mlir::sparse_tensor;22 23//===----------------------------------------------------------------------===//24// ExecutionEngine/SparseTensorUtils helper functions.25//===----------------------------------------------------------------------===//26 27OverheadType mlir::sparse_tensor::overheadTypeEncoding(unsigned width) {28 switch (width) {29 case 64:30 return OverheadType::kU64;31 case 32:32 return OverheadType::kU32;33 case 16:34 return OverheadType::kU16;35 case 8:36 return OverheadType::kU8;37 case 0:38 return OverheadType::kIndex;39 }40 llvm_unreachable("Unsupported overhead bitwidth");41}42 43OverheadType mlir::sparse_tensor::overheadTypeEncoding(Type tp) {44 if (tp.isIndex())45 return OverheadType::kIndex;46 if (auto intTp = dyn_cast<IntegerType>(tp))47 return overheadTypeEncoding(intTp.getWidth());48 llvm_unreachable("Unknown overhead type");49}50 51Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) {52 switch (ot) {53 case OverheadType::kIndex:54 return builder.getIndexType();55 case OverheadType::kU64:56 return builder.getIntegerType(64);57 case OverheadType::kU32:58 return builder.getIntegerType(32);59 case OverheadType::kU16:60 return builder.getIntegerType(16);61 case OverheadType::kU8:62 return builder.getIntegerType(8);63 }64 llvm_unreachable("Unknown OverheadType");65}66 67OverheadType68mlir::sparse_tensor::posTypeEncoding(SparseTensorEncodingAttr enc) {69 return overheadTypeEncoding(enc.getPosWidth());70}71 72OverheadType73mlir::sparse_tensor::crdTypeEncoding(SparseTensorEncodingAttr enc) {74 return overheadTypeEncoding(enc.getCrdWidth());75}76 77// TODO: we ought to add some `static_assert` tests to ensure that the78// `STEA::get{Pos,Crd}Type` methods agree with `getOverheadType(builder,79// {pos,crd}OverheadTypeEncoding(enc))`80 81// TODO: Adjust the naming convention for the constructors of82// `OverheadType` so we can use the `MLIR_SPARSETENSOR_FOREVERY_O` x-macro83// here instead of `MLIR_SPARSETENSOR_FOREVERY_FIXED_O`; to further reduce84// the possibility of typo bugs or things getting out of sync.85StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(OverheadType ot) {86 switch (ot) {87 case OverheadType::kIndex:88 return "0";89#define CASE(ONAME, O) \90 case OverheadType::kU##ONAME: \91 return #ONAME;92 MLIR_SPARSETENSOR_FOREVERY_FIXED_O(CASE)93#undef CASE94 }95 llvm_unreachable("Unknown OverheadType");96}97 98StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(Type tp) {99 return overheadTypeFunctionSuffix(overheadTypeEncoding(tp));100}101 102PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {103 if (elemTp.isF64())104 return PrimaryType::kF64;105 if (elemTp.isF32())106 return PrimaryType::kF32;107 if (elemTp.isF16())108 return PrimaryType::kF16;109 if (elemTp.isBF16())110 return PrimaryType::kBF16;111 if (elemTp.isInteger(64))112 return PrimaryType::kI64;113 if (elemTp.isInteger(32))114 return PrimaryType::kI32;115 if (elemTp.isInteger(16))116 return PrimaryType::kI16;117 if (elemTp.isInteger(8))118 return PrimaryType::kI8;119 if (auto complexTp = dyn_cast<ComplexType>(elemTp)) {120 auto complexEltTp = complexTp.getElementType();121 if (complexEltTp.isF64())122 return PrimaryType::kC64;123 if (complexEltTp.isF32())124 return PrimaryType::kC32;125 }126 llvm_unreachable("Unknown primary type");127}128 129StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(PrimaryType pt) {130 switch (pt) {131#define CASE(VNAME, V) \132 case PrimaryType::k##VNAME: \133 return #VNAME;134 MLIR_SPARSETENSOR_FOREVERY_V(CASE)135#undef CASE136 }137 llvm_unreachable("Unknown PrimaryType");138}139 140StringRef mlir::sparse_tensor::primaryTypeFunctionSuffix(Type elemTp) {141 return primaryTypeFunctionSuffix(primaryTypeEncoding(elemTp));142}143 144//===----------------------------------------------------------------------===//145// Misc code generators.146//===----------------------------------------------------------------------===//147 148Value sparse_tensor::genCast(OpBuilder &builder, Location loc, Value value,149 Type dstTp) {150 const Type srcTp = value.getType();151 if (srcTp == dstTp)152 return value;153 154 // int <=> index155 if (isa<IndexType>(srcTp) || isa<IndexType>(dstTp))156 return arith::IndexCastOp::create(builder, loc, dstTp, value);157 158 const auto srcIntTp = dyn_cast_or_null<IntegerType>(srcTp);159 const bool isUnsignedCast = srcIntTp ? srcIntTp.isUnsigned() : false;160 return mlir::convertScalarToDtype(builder, loc, value, dstTp, isUnsignedCast);161}162 163Value sparse_tensor::genScalarToTensor(OpBuilder &builder, Location loc,164 Value elem, Type dstTp) {165 if (auto rtp = dyn_cast<RankedTensorType>(dstTp)) {166 // Scalars can only be converted to 0-ranked tensors.167 assert(rtp.getRank() == 0);168 elem = sparse_tensor::genCast(builder, loc, elem, rtp.getElementType());169 return tensor::FromElementsOp::create(builder, loc, rtp, elem);170 }171 return sparse_tensor::genCast(builder, loc, elem, dstTp);172}173 174Value sparse_tensor::genIndexLoad(OpBuilder &builder, Location loc, Value mem,175 ValueRange s) {176 Value load = memref::LoadOp::create(builder, loc, mem, s);177 if (!isa<IndexType>(load.getType())) {178 if (load.getType().getIntOrFloatBitWidth() < 64)179 load = arith::ExtUIOp::create(builder, loc, builder.getI64Type(), load);180 load =181 arith::IndexCastOp::create(builder, loc, builder.getIndexType(), load);182 }183 return load;184}185 186mlir::TypedAttr mlir::sparse_tensor::getOneAttr(Builder &builder, Type tp) {187 if (isa<FloatType>(tp))188 return builder.getFloatAttr(tp, 1.0);189 if (isa<IndexType>(tp))190 return builder.getIndexAttr(1);191 if (auto intTp = dyn_cast<IntegerType>(tp))192 return builder.getIntegerAttr(tp, APInt(intTp.getWidth(), 1));193 if (isa<RankedTensorType, VectorType>(tp)) {194 auto shapedTp = cast<ShapedType>(tp);195 if (auto one = getOneAttr(builder, shapedTp.getElementType()))196 return DenseElementsAttr::get(shapedTp, one);197 }198 llvm_unreachable("Unsupported attribute type");199}200 201Value mlir::sparse_tensor::genIsNonzero(OpBuilder &builder, mlir::Location loc,202 Value v) {203 Type tp = v.getType();204 Value zero = constantZero(builder, loc, tp);205 if (isa<FloatType>(tp))206 return arith::CmpFOp::create(builder, loc, arith::CmpFPredicate::UNE, v,207 zero);208 if (tp.isIntOrIndex())209 return arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::ne, v,210 zero);211 if (isa<ComplexType>(tp))212 return complex::NotEqualOp::create(builder, loc, v, zero);213 llvm_unreachable("Non-numeric type");214}215 216void mlir::sparse_tensor::genReshapeDstShape(217 OpBuilder &builder, Location loc, SmallVectorImpl<Value> &dstShape,218 ArrayRef<Value> srcShape, ArrayRef<Size> staticDstShape,219 ArrayRef<ReassociationIndices> reassociation) {220 // Collapse shape.221 if (reassociation.size() < srcShape.size()) {222 unsigned start = 0;223 for (const auto &map : llvm::enumerate(reassociation)) {224 auto dstDim = constantIndex(builder, loc, 1);225 for (unsigned i = start; i < start + map.value().size(); i++) {226 dstDim = arith::MulIOp::create(builder, loc, dstDim, srcShape[i]);227 }228 dstShape.push_back(dstDim);229 start = start + map.value().size();230 }231 assert(start == srcShape.size());232 return;233 }234 235 // Expand shape.236 assert(reassociation.size() == srcShape.size());237 unsigned start = 0;238 // Expand the i-th dimension in srcShape.239 for (unsigned i = 0, size = srcShape.size(); i < size; i++) {240 const auto &map = reassociation[i];241 auto srcDim = srcShape[i];242 // Iterate through dimensions expanded from the i-th dimension.243 for (unsigned j = start; j < start + map.size(); j++) {244 // There can be only one dynamic sized dimension among dimensions245 // expanded from the i-th dimension in srcShape.246 // For example, if srcDim = 8, then the expanded shape could be <2x?x2>,247 // but not <2x?x?>.248 if (staticDstShape[j] == ShapedType::kDynamic) {249 // The expanded dimension has dynamic size. We compute the dimension250 // by dividing srcDim by the product of the static dimensions.251 Size product = 1;252 for (unsigned k = start; k < start + map.size(); k++) {253 if (staticDstShape[k] != ShapedType::kDynamic) {254 product *= staticDstShape[k];255 }256 }257 // Compute the dynamic dimension size.258 Value productVal = constantIndex(builder, loc, product);259 Value dynamicSize =260 arith::DivUIOp::create(builder, loc, srcDim, productVal);261 dstShape.push_back(dynamicSize);262 } else {263 // The expanded dimension is statically known.264 dstShape.push_back(constantIndex(builder, loc, staticDstShape[j]));265 }266 }267 start = start + map.size();268 }269 assert(start == staticDstShape.size());270}271 272void mlir::sparse_tensor::reshapeCvs(273 OpBuilder &builder, Location loc,274 ArrayRef<ReassociationIndices> reassociation, // NOLINT275 ValueRange srcSizes, ValueRange srcCvs, // NOLINT276 ValueRange dstSizes, SmallVectorImpl<Value> &dstCvs) {277 const unsigned srcRank = srcSizes.size();278 const unsigned dstRank = dstSizes.size();279 assert(srcRank == srcCvs.size() && "Source rank mismatch");280 const bool isCollapse = srcRank > dstRank;281 const ValueRange sizes = isCollapse ? srcSizes : dstSizes;282 // Iterate over reassociation map.283 unsigned i = 0;284 unsigned start = 0;285 for (const auto &map : llvm::enumerate(reassociation)) {286 // Prepare strides information in dimension slice.287 Value linear = constantIndex(builder, loc, 1);288 for (unsigned j = start, end = start + map.value().size(); j < end; j++) {289 linear = arith::MulIOp::create(builder, loc, linear, sizes[j]);290 }291 // Start expansion.292 Value val;293 if (!isCollapse)294 val = srcCvs[i];295 // Iterate over dimension slice.296 for (unsigned j = start, end = start + map.value().size(); j < end; j++) {297 linear = arith::DivUIOp::create(builder, loc, linear, sizes[j]);298 if (isCollapse) {299 const Value mul =300 arith::MulIOp::create(builder, loc, srcCvs[j], linear);301 val = val ? arith::AddIOp::create(builder, loc, val, mul) : mul;302 } else {303 const Value old = val;304 val = arith::DivUIOp::create(builder, loc, val, linear);305 assert(dstCvs.size() == j);306 dstCvs.push_back(val);307 val = arith::RemUIOp::create(builder, loc, old, linear);308 }309 }310 // Finalize collapse.311 if (isCollapse) {312 assert(dstCvs.size() == i);313 dstCvs.push_back(val);314 }315 start += map.value().size();316 i++;317 }318 assert(dstCvs.size() == dstRank);319}320 321FlatSymbolRefAttr mlir::sparse_tensor::getFunc(ModuleOp module, StringRef name,322 TypeRange resultType,323 ValueRange operands,324 EmitCInterface emitCInterface) {325 MLIRContext *context = module.getContext();326 auto result = SymbolRefAttr::get(context, name);327 auto func = module.lookupSymbol<func::FuncOp>(result.getAttr());328 if (!func) {329 OpBuilder moduleBuilder(module.getBodyRegion());330 func = func::FuncOp::create(331 moduleBuilder, module.getLoc(), name,332 FunctionType::get(context, operands.getTypes(), resultType));333 func.setPrivate();334 if (static_cast<bool>(emitCInterface))335 func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(),336 UnitAttr::get(context));337 }338 return result;339}340 341func::CallOp mlir::sparse_tensor::createFuncCall(342 OpBuilder &builder, Location loc, StringRef name, TypeRange resultType,343 ValueRange operands, EmitCInterface emitCInterface) {344 auto module = builder.getBlock()->getParentOp()->getParentOfType<ModuleOp>();345 FlatSymbolRefAttr fn =346 getFunc(module, name, resultType, operands, emitCInterface);347 return func::CallOp::create(builder, loc, resultType, fn, operands);348}349 350Type mlir::sparse_tensor::getOpaquePointerType(MLIRContext *ctx) {351 return LLVM::LLVMPointerType::get(ctx);352}353 354Type mlir::sparse_tensor::getOpaquePointerType(Builder &builder) {355 return getOpaquePointerType(builder.getContext());356}357 358Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc,359 unsigned sz, Type tp, bool staticShape) {360 if (staticShape) {361 auto memTp = MemRefType::get({sz}, tp);362 return memref::AllocaOp::create(builder, loc, memTp);363 }364 return genAlloca(builder, loc, constantIndex(builder, loc, sz), tp);365}366 367Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc, Value sz,368 Type tp) {369 auto memTp = MemRefType::get({ShapedType::kDynamic}, tp);370 return memref::AllocaOp::create(builder, loc, memTp, ValueRange{sz});371}372 373Value mlir::sparse_tensor::genAllocaScalar(OpBuilder &builder, Location loc,374 Type tp) {375 return memref::AllocaOp::create(builder, loc, MemRefType::get({}, tp));376}377 378Value mlir::sparse_tensor::allocaBuffer(OpBuilder &builder, Location loc,379 ValueRange values) {380 const unsigned sz = values.size();381 assert(sz >= 1);382 Value buffer = genAlloca(builder, loc, sz, values[0].getType());383 for (unsigned i = 0; i < sz; i++) {384 Value idx = constantIndex(builder, loc, i);385 memref::StoreOp::create(builder, loc, values[i], buffer, idx);386 }387 return buffer;388}389 390Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc,391 RankedTensorType tensorTp,392 ValueRange sizes) {393 Type elemTp = tensorTp.getElementType();394 auto shape = tensorTp.getShape();395 auto memTp = MemRefType::get(shape, elemTp);396 SmallVector<Value> dynamicSizes;397 for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {398 if (shape[i] == ShapedType::kDynamic)399 dynamicSizes.push_back(sizes[i]);400 }401 Value mem = memref::AllocOp::create(builder, loc, memTp, dynamicSizes);402 Value zero = constantZero(builder, loc, elemTp);403 linalg::FillOp::create(builder, loc, ValueRange{zero}, ValueRange{mem});404 return mem;405}406 407void mlir::sparse_tensor::deallocDenseTensor(OpBuilder &builder, Location loc,408 Value buffer) {409 memref::DeallocOp::create(builder, loc, buffer);410}411 412void mlir::sparse_tensor::sizesFromSrc(OpBuilder &builder,413 SmallVectorImpl<Value> &sizes,414 Location loc, Value src) {415 const Dimension dimRank = getSparseTensorType(src).getDimRank();416 for (Dimension d = 0; d < dimRank; d++)417 sizes.push_back(linalg::createOrFoldDimOp(builder, loc, src, d));418}419 420Operation *mlir::sparse_tensor::getTop(Operation *op) {421 for (; isa<scf::ForOp>(op->getParentOp()) ||422 isa<scf::WhileOp>(op->getParentOp()) ||423 isa<scf::ParallelOp>(op->getParentOp()) ||424 isa<scf::IfOp>(op->getParentOp());425 op = op->getParentOp())426 ;427 return op;428}429 430void sparse_tensor::foreachInSparseConstant(431 OpBuilder &builder, Location loc, SparseElementsAttr attr, AffineMap order,432 function_ref<void(ArrayRef<Value>, Value)> callback) {433 if (!order)434 order = builder.getMultiDimIdentityMap(attr.getType().getRank());435 436 auto stt = SparseTensorType(getRankedTensorType(attr));437 const Dimension dimRank = stt.getDimRank();438 const auto coordinates = attr.getIndices().getValues<IntegerAttr>();439 const auto values = attr.getValues().getValues<Attribute>();440 441 // This is like the `Element<V>` class in the runtime library, but for442 // MLIR attributes. In the future we may want to move this out into443 // a proper class definition to help improve code legibility (e.g.,444 // `first` -> `coords`, `second` -> `value`) as well as being able445 // to factor out analogues of `ElementLT<V>` for the sort below, etc.446 using ElementAttr = std::pair<SmallVector<IntegerAttr>, Attribute>;447 448 // Construct the COO from the SparseElementsAttr.449 SmallVector<ElementAttr> elems;450 for (size_t i = 0, nse = values.size(); i < nse; i++) {451 elems.emplace_back();452 elems.back().second = values[i];453 auto &coords = elems.back().first;454 coords.reserve(dimRank);455 for (Dimension d = 0; d < dimRank; d++)456 coords.push_back(coordinates[i * dimRank + d]);457 }458 459 // Sorts the sparse element attribute based on coordinates.460 llvm::sort(elems, [order](const ElementAttr &lhs, const ElementAttr &rhs) {461 if (std::addressof(lhs) == std::addressof(rhs))462 return false;463 464 auto lhsCoords = llvm::map_to_vector(465 lhs.first, [](IntegerAttr i) { return i.getInt(); });466 auto rhsCoords = llvm::map_to_vector(467 rhs.first, [](IntegerAttr i) { return i.getInt(); });468 469 SmallVector<int64_t, 4> lhsLvlCrds = order.compose(lhsCoords);470 SmallVector<int64_t, 4> rhsLvlCrds = order.compose(rhsCoords);471 // Sort the element based on the lvl coordinates.472 for (Level l = 0; l < order.getNumResults(); l++) {473 if (lhsLvlCrds[l] == rhsLvlCrds[l])474 continue;475 return lhsLvlCrds[l] < rhsLvlCrds[l];476 }477 llvm_unreachable("no equal coordinate in sparse element attr");478 });479 480 SmallVector<Value> cvs;481 cvs.reserve(dimRank);482 for (size_t i = 0, nse = values.size(); i < nse; i++) {483 // Remap coordinates.484 cvs.clear();485 for (Dimension d = 0; d < dimRank; d++) {486 auto crd = elems[i].first[d].getInt();487 cvs.push_back(arith::ConstantIndexOp::create(builder, loc, crd));488 }489 // Remap value.490 Value val;491 if (isa<ComplexType>(attr.getElementType())) {492 auto valAttr = cast<ArrayAttr>(elems[i].second);493 val = complex::ConstantOp::create(builder, loc, attr.getElementType(),494 valAttr);495 } else {496 auto valAttr = cast<TypedAttr>(elems[i].second);497 val = arith::ConstantOp::create(builder, loc, valAttr);498 }499 assert(val);500 callback(cvs, val);501 }502}503 504SmallVector<Value> sparse_tensor::loadAll(OpBuilder &builder, Location loc,505 size_t size, Value mem,506 size_t offsetIdx, Value offsetVal) {507#ifndef NDEBUG508 const auto memTp = cast<MemRefType>(mem.getType());509 assert(memTp.getRank() == 1);510 const Size memSh = memTp.getDimSize(0);511 assert(ShapedType::isDynamic(memSh) || memSh >= static_cast<Size>(size));512 assert(offsetIdx == 0 || offsetIdx < size);513#endif // NDEBUG514 SmallVector<Value> vs;515 vs.reserve(size);516 for (unsigned i = 0; i < size; i++) {517 Value v = memref::LoadOp::create(builder, loc, mem,518 constantIndex(builder, loc, i));519 if (i == offsetIdx && offsetVal)520 v = arith::AddIOp::create(builder, loc, v, offsetVal);521 vs.push_back(v);522 }523 return vs;524}525 526void sparse_tensor::storeAll(OpBuilder &builder, Location loc, Value mem,527 ValueRange vs, size_t offsetIdx, Value offsetVal) {528#ifndef NDEBUG529 const size_t vsize = vs.size();530 const auto memTp = cast<MemRefType>(mem.getType());531 assert(memTp.getRank() == 1);532 const Size memSh = memTp.getDimSize(0);533 assert(ShapedType::isDynamic(memSh) || memSh >= static_cast<Size>(vsize));534 assert(offsetIdx == 0 || offsetIdx < vsize);535#endif // NDEBUG536 for (const auto &v : llvm::enumerate(vs)) {537 const Value w =538 (offsetIdx == v.index() && offsetVal)539 ? arith::AddIOp::create(builder, loc, v.value(), offsetVal)540 : v.value();541 memref::StoreOp::create(builder, loc, w, mem,542 constantIndex(builder, loc, v.index()));543 }544}545 546TypedValue<BaseMemRefType>547sparse_tensor::genToMemref(OpBuilder &builder, Location loc, Value tensor) {548 auto tTp = llvm::cast<TensorType>(tensor.getType());549 auto mTp = MemRefType::get(tTp.getShape(), tTp.getElementType());550 return cast<TypedValue<BaseMemRefType>>(551 bufferization::ToBufferOp::create(builder, loc, mTp, tensor).getResult());552}553 554Value sparse_tensor::createOrFoldSliceOffsetOp(OpBuilder &builder, Location loc,555 Value tensor, Dimension dim) {556 auto enc = getSparseTensorEncoding(tensor.getType());557 assert(enc && enc.isSlice());558 std::optional<unsigned> offset = enc.getStaticDimSliceOffset(dim);559 if (offset.has_value())560 return constantIndex(builder, loc, *offset);561 return ToSliceOffsetOp::create(builder, loc, tensor, APInt(64, dim));562}563 564Value sparse_tensor::createOrFoldSliceStrideOp(OpBuilder &builder, Location loc,565 Value tensor, Dimension dim) {566 auto enc = getSparseTensorEncoding(tensor.getType());567 assert(enc && enc.isSlice());568 std::optional<unsigned> stride = enc.getStaticDimSliceStride(dim);569 if (stride.has_value())570 return constantIndex(builder, loc, *stride);571 return ToSliceStrideOp::create(builder, loc, tensor, APInt(64, dim));572}573 574Value sparse_tensor::genReader(OpBuilder &builder, Location loc,575 SparseTensorType stt, Value tensor,576 /*out*/ SmallVectorImpl<Value> &dimSizesValues,577 /*out*/ Value &dimSizesBuffer) {578 // Construct the dimension **shapes** buffer. The buffer contains the static579 // size per dimension, or otherwise a zero for a dynamic size.580 Dimension dimRank = stt.getDimRank();581 dimSizesValues.clear();582 dimSizesValues.reserve(dimRank);583 for (const Size sz : stt.getDimShape()) {584 const auto s = ShapedType::isDynamic(sz) ? 0 : sz;585 dimSizesValues.push_back(constantIndex(builder, loc, s));586 }587 Value dimShapesBuffer = allocaBuffer(builder, loc, dimSizesValues);588 // Create the `CheckedSparseTensorReader`. This reader performs a589 // consistency check on the static sizes, but accepts any size590 // of each dimension with a dynamic size.591 Type opaqueTp = getOpaquePointerType(builder);592 Type eltTp = stt.getElementType();593 Value valTp = constantPrimaryTypeEncoding(builder, loc, eltTp);594 Value reader =595 createFuncCall(builder, loc, "createCheckedSparseTensorReader", opaqueTp,596 {tensor, dimShapesBuffer, valTp}, EmitCInterface::On)597 .getResult(0);598 // For static shapes, the shape buffer can be used right away. For dynamic599 // shapes, use the information from the reader to construct a buffer that600 // supplies the actual size for each dynamic dimension.601 dimSizesBuffer = dimShapesBuffer;602 if (stt.hasDynamicDimShape()) {603 Type indexTp = builder.getIndexType();604 auto memTp = MemRefType::get({ShapedType::kDynamic}, indexTp);605 dimSizesBuffer =606 createFuncCall(builder, loc, "getSparseTensorReaderDimSizes", memTp,607 reader, EmitCInterface::On)608 .getResult(0);609 // Also convert the dim shapes values into dim sizes values, just in case610 // subsequent clients need the values (DCE will remove unused).611 for (Dimension d = 0; d < dimRank; d++) {612 if (stt.isDynamicDim(d))613 dimSizesValues[d] = memref::LoadOp::create(614 builder, loc, dimSizesBuffer, constantIndex(builder, loc, d));615 }616 }617 return reader;618}619 620Value sparse_tensor::genMapBuffers(621 OpBuilder &builder, Location loc, SparseTensorType stt,622 ArrayRef<Value> dimSizesValues, Value dimSizesBuffer,623 /*out*/ SmallVectorImpl<Value> &lvlSizesValues,624 /*out*/ Value &dim2lvlBuffer,625 /*out*/ Value &lvl2dimBuffer) {626 const Dimension dimRank = stt.getDimRank();627 const Level lvlRank = stt.getLvlRank();628 lvlSizesValues.clear();629 lvlSizesValues.reserve(lvlRank);630 // For an identity mapping, the dim2lvl and lvl2dim mappings are631 // identical as are dimSizes and lvlSizes, so buffers are reused632 // as much as possible.633 if (stt.isIdentity()) {634 assert(dimRank == lvlRank);635 SmallVector<Value> iotaValues;636 iotaValues.reserve(lvlRank);637 for (Level l = 0; l < lvlRank; l++) {638 iotaValues.push_back(constantIndex(builder, loc, l));639 lvlSizesValues.push_back(dimSizesValues[l]);640 }641 dim2lvlBuffer = lvl2dimBuffer = allocaBuffer(builder, loc, iotaValues);642 return dimSizesBuffer; // now lvlSizesBuffer643 }644 // Otherwise, some code needs to be generated to set up the buffers.645 // This code deals with permutations as well as non-permutations that646 // arise from rank changing blocking.647 const auto dimToLvl = stt.getDimToLvl();648 const auto lvlToDim = stt.getLvlToDim();649 SmallVector<Value> dim2lvlValues(lvlRank); // for each lvl, expr in dim vars650 SmallVector<Value> lvl2dimValues(dimRank); // for each dim, expr in lvl vars651 // Generate dim2lvl.652 assert(lvlRank == dimToLvl.getNumResults());653 for (Level l = 0; l < lvlRank; l++) {654 AffineExpr exp = dimToLvl.getResult(l);655 // We expect:656 // (1) l = d657 // (2) l = d / c658 // (3) l = d % c659 Dimension d = 0;660 uint64_t cf = 0, cm = 0;661 switch (exp.getKind()) {662 case AffineExprKind::DimId: {663 d = cast<AffineDimExpr>(exp).getPosition();664 break;665 }666 case AffineExprKind::FloorDiv: {667 auto floor = cast<AffineBinaryOpExpr>(exp);668 d = cast<AffineDimExpr>(floor.getLHS()).getPosition();669 cf = cast<AffineConstantExpr>(floor.getRHS()).getValue();670 break;671 }672 case AffineExprKind::Mod: {673 auto mod = cast<AffineBinaryOpExpr>(exp);674 d = cast<AffineDimExpr>(mod.getLHS()).getPosition();675 cm = cast<AffineConstantExpr>(mod.getRHS()).getValue();676 break;677 }678 default:679 llvm::report_fatal_error("unsupported dim2lvl in sparse tensor type");680 }681 dim2lvlValues[l] = constantIndex(builder, loc, encodeDim(d, cf, cm));682 // Compute the level sizes.683 // (1) l = d : size(d)684 // (2) l = d / c : size(d) / c685 // (3) l = d % c : c686 Value lvlSz;687 if (cm == 0) {688 lvlSz = dimSizesValues[d];689 if (cf != 0)690 lvlSz = arith::DivUIOp::create(builder, loc, lvlSz,691 constantIndex(builder, loc, cf));692 } else {693 lvlSz = constantIndex(builder, loc, cm);694 }695 lvlSizesValues.push_back(lvlSz);696 }697 // Generate lvl2dim.698 assert(dimRank == lvlToDim.getNumResults());699 for (Dimension d = 0; d < dimRank; d++) {700 AffineExpr exp = lvlToDim.getResult(d);701 // We expect:702 // (1) d = l703 // (2) d = l' * c + l704 Level l = 0, ll = 0;705 uint64_t c = 0;706 switch (exp.getKind()) {707 case AffineExprKind::DimId: {708 l = cast<AffineDimExpr>(exp).getPosition();709 break;710 }711 case AffineExprKind::Add: {712 // Always mul on lhs, symbol/constant on rhs.713 auto add = cast<AffineBinaryOpExpr>(exp);714 assert(add.getLHS().getKind() == AffineExprKind::Mul);715 auto mul = cast<AffineBinaryOpExpr>(add.getLHS());716 ll = cast<AffineDimExpr>(mul.getLHS()).getPosition();717 c = cast<AffineConstantExpr>(mul.getRHS()).getValue();718 l = cast<AffineDimExpr>(add.getRHS()).getPosition();719 break;720 }721 default:722 llvm::report_fatal_error("unsupported lvl2dim in sparse tensor type");723 }724 lvl2dimValues[d] = constantIndex(builder, loc, encodeLvl(l, c, ll));725 }726 // Return buffers.727 dim2lvlBuffer = allocaBuffer(builder, loc, dim2lvlValues);728 lvl2dimBuffer = allocaBuffer(builder, loc, lvl2dimValues);729 return allocaBuffer(builder, loc, lvlSizesValues); // lvlSizesBuffer730}731