llvm/mlir/lib/Dialect/SparseTensor/Transforms/SparseStorageSpecifierToLLVM.cpp

//===- SparseStorageSpecifierToLLVM.cpp - convert specifier to llvm -------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "Utils/CodegenUtils.h"

#include "mlir/Conversion/LLVMCommon/StructBuilder.h"
#include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"

#include <optional>

usingnamespacemlir;
usingnamespacesparse_tensor;

namespace {

//===----------------------------------------------------------------------===//
// Helper methods.
//===----------------------------------------------------------------------===//

static SmallVector<Type, 4> getSpecifierFields(StorageSpecifierType tp) { … }

static Type convertSpecifier(StorageSpecifierType tp) { … }

//===----------------------------------------------------------------------===//
// Specifier struct builder.
//===----------------------------------------------------------------------===//

constexpr uint64_t kLvlSizePosInSpecifier = …;
constexpr uint64_t kMemSizePosInSpecifier = …;
constexpr uint64_t kDimOffsetPosInSpecifier = …;
constexpr uint64_t kDimStridePosInSpecifier = …;

class SpecifierStructBuilder : public StructBuilder { … };

Value SpecifierStructBuilder::getInitValue(OpBuilder &builder, Location loc,
                                           Type structType, Value source) { … }

/// Builds IR extracting the pos-th offset from the descriptor.
Value SpecifierStructBuilder::dimOffset(OpBuilder &builder, Location loc,
                                        Dimension dim) const { … }

/// Builds IR inserting the pos-th offset into the descriptor.
void SpecifierStructBuilder::setDimOffset(OpBuilder &builder, Location loc,
                                          Dimension dim, Value size) { … }

/// Builds IR extracting the `lvl`-th level-size from the descriptor.
Value SpecifierStructBuilder::lvlSize(OpBuilder &builder, Location loc,
                                      Level lvl) const { … }

/// Builds IR inserting the `lvl`-th level-size into the descriptor.
void SpecifierStructBuilder::setLvlSize(OpBuilder &builder, Location loc,
                                        Level lvl, Value size) { … }

/// Builds IR extracting the pos-th stride from the descriptor.
Value SpecifierStructBuilder::dimStride(OpBuilder &builder, Location loc,
                                        Dimension dim) const { … }

/// Builds IR inserting the pos-th stride into the descriptor.
void SpecifierStructBuilder::setDimStride(OpBuilder &builder, Location loc,
                                          Dimension dim, Value size) { … }

/// Builds IR extracting the pos-th memory size into the descriptor.
Value SpecifierStructBuilder::memSize(OpBuilder &builder, Location loc,
                                      FieldIndex fidx) const { … }

/// Builds IR inserting the `fidx`-th memory-size into the descriptor.
void SpecifierStructBuilder::setMemSize(OpBuilder &builder, Location loc,
                                        FieldIndex fidx, Value size) { … }

/// Builds IR extracting the memory size array from the descriptor.
Value SpecifierStructBuilder::memSizeArray(OpBuilder &builder,
                                           Location loc) const { … }

/// Builds IR inserting the memory size array into the descriptor.
void SpecifierStructBuilder::setMemSizeArray(OpBuilder &builder, Location loc,
                                             Value array) { … }

} // namespace

//===----------------------------------------------------------------------===//
// The sparse storage specifier type converter (defined in Passes.h).
//===----------------------------------------------------------------------===//

StorageSpecifierToLLVMTypeConverter::StorageSpecifierToLLVMTypeConverter() { … }

//===----------------------------------------------------------------------===//
// Storage specifier conversion rules.
//===----------------------------------------------------------------------===//

template <typename Base, typename SourceOp>
class SpecifierGetterSetterOpConverter : public OpConversionPattern<SourceOp> { … };

struct StorageSpecifierSetOpConverter
    : public SpecifierGetterSetterOpConverter<StorageSpecifierSetOpConverter,
                                              SetStorageSpecifierOp> { … };

struct StorageSpecifierGetOpConverter
    : public SpecifierGetterSetterOpConverter<StorageSpecifierGetOpConverter,
                                              GetStorageSpecifierOp> { … };

struct StorageSpecifierInitOpConverter
    : public OpConversionPattern<StorageSpecifierInitOp> { … };

//===----------------------------------------------------------------------===//
// Public method for populating conversion rules.
//===----------------------------------------------------------------------===//

void mlir::populateStorageSpecifierToLLVMPatterns(
    const TypeConverter &converter, RewritePatternSet &patterns) { … }