359 lines · cpp
1//===- SparseStorageSpecifierToLLVM.cpp - convert specifier to llvm -------===//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 "Utils/CodegenUtils.h"10 11#include "mlir/Conversion/LLVMCommon/StructBuilder.h"12#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"13#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"14 15#include <optional>16 17using namespace mlir;18using namespace sparse_tensor;19 20namespace {21 22//===----------------------------------------------------------------------===//23// Helper methods.24//===----------------------------------------------------------------------===//25 26static SmallVector<Type, 4> getSpecifierFields(StorageSpecifierType tp) {27 MLIRContext *ctx = tp.getContext();28 auto enc = tp.getEncoding();29 const Level lvlRank = enc.getLvlRank();30 31 SmallVector<Type, 4> result;32 // TODO: how can we get the lowering type for index type in the later pipeline33 // to be consistent? LLVM::StructureType does not allow index fields.34 auto sizeType = IntegerType::get(tp.getContext(), 64);35 auto lvlSizes = LLVM::LLVMArrayType::get(ctx, sizeType, lvlRank);36 auto memSizes = LLVM::LLVMArrayType::get(ctx, sizeType,37 getNumDataFieldsFromEncoding(enc));38 result.push_back(lvlSizes);39 result.push_back(memSizes);40 41 if (enc.isSlice()) {42 // Extra fields are required for the slice information.43 auto dimOffset = LLVM::LLVMArrayType::get(ctx, sizeType, lvlRank);44 auto dimStride = LLVM::LLVMArrayType::get(ctx, sizeType, lvlRank);45 46 result.push_back(dimOffset);47 result.push_back(dimStride);48 }49 50 return result;51}52 53static Type convertSpecifier(StorageSpecifierType tp) {54 return LLVM::LLVMStructType::getLiteral(tp.getContext(),55 getSpecifierFields(tp));56}57 58//===----------------------------------------------------------------------===//59// Specifier struct builder.60//===----------------------------------------------------------------------===//61 62constexpr uint64_t kLvlSizePosInSpecifier = 0;63constexpr uint64_t kMemSizePosInSpecifier = 1;64constexpr uint64_t kDimOffsetPosInSpecifier = 2;65constexpr uint64_t kDimStridePosInSpecifier = 3;66 67class SpecifierStructBuilder : public StructBuilder {68private:69 Value extractField(OpBuilder &builder, Location loc,70 ArrayRef<int64_t> indices) const {71 return genCast(builder, loc,72 LLVM::ExtractValueOp::create(builder, loc, value, indices),73 builder.getIndexType());74 }75 76 void insertField(OpBuilder &builder, Location loc, ArrayRef<int64_t> indices,77 Value v) {78 value = LLVM::InsertValueOp::create(79 builder, loc, value,80 genCast(builder, loc, v, builder.getIntegerType(64)), indices);81 }82 83public:84 explicit SpecifierStructBuilder(Value specifier) : StructBuilder(specifier) {85 assert(value);86 }87 88 // Undef value for dimension sizes, all zero value for memory sizes.89 static Value getInitValue(OpBuilder &builder, Location loc, Type structType,90 Value source);91 92 Value lvlSize(OpBuilder &builder, Location loc, Level lvl) const;93 void setLvlSize(OpBuilder &builder, Location loc, Level lvl, Value size);94 95 Value dimOffset(OpBuilder &builder, Location loc, Dimension dim) const;96 void setDimOffset(OpBuilder &builder, Location loc, Dimension dim,97 Value size);98 99 Value dimStride(OpBuilder &builder, Location loc, Dimension dim) const;100 void setDimStride(OpBuilder &builder, Location loc, Dimension dim,101 Value size);102 103 Value memSize(OpBuilder &builder, Location loc, FieldIndex fidx) const;104 void setMemSize(OpBuilder &builder, Location loc, FieldIndex fidx,105 Value size);106 107 Value memSizeArray(OpBuilder &builder, Location loc) const;108 void setMemSizeArray(OpBuilder &builder, Location loc, Value array);109};110 111Value SpecifierStructBuilder::getInitValue(OpBuilder &builder, Location loc,112 Type structType, Value source) {113 Value metaData = LLVM::PoisonOp::create(builder, loc, structType);114 SpecifierStructBuilder md(metaData);115 if (!source) {116 auto memSizeArrayType =117 cast<LLVM::LLVMArrayType>(cast<LLVM::LLVMStructType>(structType)118 .getBody()[kMemSizePosInSpecifier]);119 120 Value zero = constantZero(builder, loc, memSizeArrayType.getElementType());121 // Fill memSizes array with zero.122 for (int i = 0, e = memSizeArrayType.getNumElements(); i < e; i++)123 md.setMemSize(builder, loc, i, zero);124 } else {125 // We copy non-slice information (memory sizes array) from source126 SpecifierStructBuilder sourceMd(source);127 md.setMemSizeArray(builder, loc, sourceMd.memSizeArray(builder, loc));128 }129 return md;130}131 132/// Builds IR extracting the pos-th offset from the descriptor.133Value SpecifierStructBuilder::dimOffset(OpBuilder &builder, Location loc,134 Dimension dim) const {135 return extractField(136 builder, loc,137 ArrayRef<int64_t>{kDimOffsetPosInSpecifier, static_cast<int64_t>(dim)});138}139 140/// Builds IR inserting the pos-th offset into the descriptor.141void SpecifierStructBuilder::setDimOffset(OpBuilder &builder, Location loc,142 Dimension dim, Value size) {143 insertField(144 builder, loc,145 ArrayRef<int64_t>{kDimOffsetPosInSpecifier, static_cast<int64_t>(dim)},146 size);147}148 149/// Builds IR extracting the `lvl`-th level-size from the descriptor.150Value SpecifierStructBuilder::lvlSize(OpBuilder &builder, Location loc,151 Level lvl) const {152 // This static_cast makes the narrowing of `lvl` explicit, as required153 // by the braces notation for the ctor.154 return extractField(155 builder, loc,156 ArrayRef<int64_t>{kLvlSizePosInSpecifier, static_cast<int64_t>(lvl)});157}158 159/// Builds IR inserting the `lvl`-th level-size into the descriptor.160void SpecifierStructBuilder::setLvlSize(OpBuilder &builder, Location loc,161 Level lvl, Value size) {162 // This static_cast makes the narrowing of `lvl` explicit, as required163 // by the braces notation for the ctor.164 insertField(165 builder, loc,166 ArrayRef<int64_t>{kLvlSizePosInSpecifier, static_cast<int64_t>(lvl)},167 size);168}169 170/// Builds IR extracting the pos-th stride from the descriptor.171Value SpecifierStructBuilder::dimStride(OpBuilder &builder, Location loc,172 Dimension dim) const {173 return extractField(174 builder, loc,175 ArrayRef<int64_t>{kDimStridePosInSpecifier, static_cast<int64_t>(dim)});176}177 178/// Builds IR inserting the pos-th stride into the descriptor.179void SpecifierStructBuilder::setDimStride(OpBuilder &builder, Location loc,180 Dimension dim, Value size) {181 insertField(182 builder, loc,183 ArrayRef<int64_t>{kDimStridePosInSpecifier, static_cast<int64_t>(dim)},184 size);185}186 187/// Builds IR extracting the pos-th memory size into the descriptor.188Value SpecifierStructBuilder::memSize(OpBuilder &builder, Location loc,189 FieldIndex fidx) const {190 return extractField(191 builder, loc,192 ArrayRef<int64_t>{kMemSizePosInSpecifier, static_cast<int64_t>(fidx)});193}194 195/// Builds IR inserting the `fidx`-th memory-size into the descriptor.196void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,197 FieldIndex fidx, Value size) {198 insertField(199 builder, loc,200 ArrayRef<int64_t>{kMemSizePosInSpecifier, static_cast<int64_t>(fidx)},201 size);202}203 204/// Builds IR extracting the memory size array from the descriptor.205Value SpecifierStructBuilder::memSizeArray(OpBuilder &builder,206 Location loc) const {207 return LLVM::ExtractValueOp::create(builder, loc, value,208 kMemSizePosInSpecifier);209}210 211/// Builds IR inserting the memory size array into the descriptor.212void SpecifierStructBuilder::setMemSizeArray(OpBuilder &builder, Location loc,213 Value array) {214 value = LLVM::InsertValueOp::create(builder, loc, value, array,215 kMemSizePosInSpecifier);216}217 218} // namespace219 220//===----------------------------------------------------------------------===//221// The sparse storage specifier type converter (defined in Passes.h).222//===----------------------------------------------------------------------===//223 224StorageSpecifierToLLVMTypeConverter::StorageSpecifierToLLVMTypeConverter() {225 addConversion([](Type type) { return type; });226 addConversion(convertSpecifier);227}228 229//===----------------------------------------------------------------------===//230// Storage specifier conversion rules.231//===----------------------------------------------------------------------===//232 233template <typename Base, typename SourceOp>234class SpecifierGetterSetterOpConverter : public OpConversionPattern<SourceOp> {235public:236 using OpAdaptor = typename SourceOp::Adaptor;237 using OpConversionPattern<SourceOp>::OpConversionPattern;238 239 LogicalResult240 matchAndRewrite(SourceOp op, OpAdaptor adaptor,241 ConversionPatternRewriter &rewriter) const override {242 SpecifierStructBuilder spec(adaptor.getSpecifier());243 switch (op.getSpecifierKind()) {244 case StorageSpecifierKind::LvlSize: {245 Value v = Base::onLvlSize(rewriter, op, spec, (*op.getLevel()));246 rewriter.replaceOp(op, v);247 return success();248 }249 case StorageSpecifierKind::DimOffset: {250 Value v = Base::onDimOffset(rewriter, op, spec, (*op.getLevel()));251 rewriter.replaceOp(op, v);252 return success();253 }254 case StorageSpecifierKind::DimStride: {255 Value v = Base::onDimStride(rewriter, op, spec, (*op.getLevel()));256 rewriter.replaceOp(op, v);257 return success();258 }259 case StorageSpecifierKind::CrdMemSize:260 case StorageSpecifierKind::PosMemSize:261 case StorageSpecifierKind::ValMemSize: {262 auto enc = op.getSpecifier().getType().getEncoding();263 StorageLayout layout(enc);264 std::optional<unsigned> lvl;265 if (op.getLevel())266 lvl = (*op.getLevel());267 unsigned idx =268 layout.getMemRefFieldIndex(toFieldKind(op.getSpecifierKind()), lvl);269 Value v = Base::onMemSize(rewriter, op, spec, idx);270 rewriter.replaceOp(op, v);271 return success();272 }273 }274 llvm_unreachable("unrecognized specifer kind");275 }276};277 278struct StorageSpecifierSetOpConverter279 : public SpecifierGetterSetterOpConverter<StorageSpecifierSetOpConverter,280 SetStorageSpecifierOp> {281 using SpecifierGetterSetterOpConverter::SpecifierGetterSetterOpConverter;282 283 static Value onLvlSize(OpBuilder &builder, SetStorageSpecifierOp op,284 SpecifierStructBuilder &spec, Level lvl) {285 spec.setLvlSize(builder, op.getLoc(), lvl, op.getValue());286 return spec;287 }288 289 static Value onDimOffset(OpBuilder &builder, SetStorageSpecifierOp op,290 SpecifierStructBuilder &spec, Dimension d) {291 spec.setDimOffset(builder, op.getLoc(), d, op.getValue());292 return spec;293 }294 295 static Value onDimStride(OpBuilder &builder, SetStorageSpecifierOp op,296 SpecifierStructBuilder &spec, Dimension d) {297 spec.setDimStride(builder, op.getLoc(), d, op.getValue());298 return spec;299 }300 301 static Value onMemSize(OpBuilder &builder, SetStorageSpecifierOp op,302 SpecifierStructBuilder &spec, FieldIndex fidx) {303 spec.setMemSize(builder, op.getLoc(), fidx, op.getValue());304 return spec;305 }306};307 308struct StorageSpecifierGetOpConverter309 : public SpecifierGetterSetterOpConverter<StorageSpecifierGetOpConverter,310 GetStorageSpecifierOp> {311 using SpecifierGetterSetterOpConverter::SpecifierGetterSetterOpConverter;312 313 static Value onLvlSize(OpBuilder &builder, GetStorageSpecifierOp op,314 SpecifierStructBuilder &spec, Level lvl) {315 return spec.lvlSize(builder, op.getLoc(), lvl);316 }317 318 static Value onDimOffset(OpBuilder &builder, GetStorageSpecifierOp op,319 const SpecifierStructBuilder &spec, Dimension d) {320 return spec.dimOffset(builder, op.getLoc(), d);321 }322 323 static Value onDimStride(OpBuilder &builder, GetStorageSpecifierOp op,324 const SpecifierStructBuilder &spec, Dimension d) {325 return spec.dimStride(builder, op.getLoc(), d);326 }327 328 static Value onMemSize(OpBuilder &builder, GetStorageSpecifierOp op,329 SpecifierStructBuilder &spec, FieldIndex fidx) {330 return spec.memSize(builder, op.getLoc(), fidx);331 }332};333 334struct StorageSpecifierInitOpConverter335 : public OpConversionPattern<StorageSpecifierInitOp> {336public:337 using OpConversionPattern::OpConversionPattern;338 LogicalResult339 matchAndRewrite(StorageSpecifierInitOp op, OpAdaptor adaptor,340 ConversionPatternRewriter &rewriter) const override {341 Type llvmType = getTypeConverter()->convertType(op.getResult().getType());342 rewriter.replaceOp(343 op, SpecifierStructBuilder::getInitValue(344 rewriter, op.getLoc(), llvmType, adaptor.getSource()));345 return success();346 }347};348 349//===----------------------------------------------------------------------===//350// Public method for populating conversion rules.351//===----------------------------------------------------------------------===//352 353void mlir::populateStorageSpecifierToLLVMPatterns(354 const TypeConverter &converter, RewritePatternSet &patterns) {355 patterns.add<StorageSpecifierGetOpConverter, StorageSpecifierSetOpConverter,356 StorageSpecifierInitOpConverter>(converter,357 patterns.getContext());358}359