llvm/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

//===- Enums.h - Enums for the SparseTensor dialect -------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Typedefs and enums shared between MLIR code for manipulating the
// IR, and the lightweight runtime support library for sparse tensor
// manipulations.  That is, all the enums are used to define the API
// of the runtime library and hence are also needed when generating
// calls into the runtime library.  Moveover, the `LevelType` enum
// is also used as the internal IR encoding of dimension level types,
// to avoid code duplication (e.g., for the predicates).
//
// This file also defines x-macros <https://en.wikipedia.org/wiki/X_Macro>
// so that we can generate variations of the public functions for each
// supported primary- and/or overhead-type.
//
// Because this file defines a library which is a dependency of the
// runtime library itself, this file must not depend on any MLIR internals
// (e.g., operators, attributes, ArrayRefs, etc) lest the runtime library
// inherit those dependencies.
//
//===----------------------------------------------------------------------===//

#ifndef MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H
#define MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H

// NOTE: Client code will need to include "mlir/ExecutionEngine/Float16bits.h"
// if they want to use the `MLIR_SPARSETENSOR_FOREVERY_V` macro.

#include <cassert>
#include <cinttypes>
#include <complex>
#include <optional>
#include <vector>

namespace mlir {
namespace sparse_tensor {

/// This type is used in the public API at all places where MLIR expects
/// values with the built-in type "index".  For now, we simply assume that
/// type is 64-bit, but targets with different "index" bitwidths should
/// link with an alternatively built runtime support library.
index_type;

/// Encoding of overhead types (both position overhead and coordinate
/// overhead), for "overloading" @newSparseTensor.
enum class OverheadType : uint32_t { … };

// This x-macro calls its argument on every overhead type which has
// fixed-width.  It excludes `index_type` because that type is often
// handled specially (e.g., by translating it into the architecture-dependent
// equivalent fixed-width overhead type).
#define MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DO) …

// This x-macro calls its argument on every overhead type, including
// `index_type`.
#define MLIR_SPARSETENSOR_FOREVERY_O(DO) …

// These are not just shorthands but indicate the particular
// implementation used (e.g., as opposed to C99's `complex double`,
// or MLIR's `ComplexType`).
complex64;
complex32;

/// Encoding of the elemental type, for "overloading" @newSparseTensor.
enum class PrimaryType : uint32_t { … };

// This x-macro includes all `V` types.
#define MLIR_SPARSETENSOR_FOREVERY_V(DO) …

// This x-macro includes all `V` types and supports variadic arguments.
#define MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, ...) …

// This x-macro calls its argument on every pair of overhead and `V` types.
#define MLIR_SPARSETENSOR_FOREVERY_V_O(DO) …

constexpr bool isFloatingPrimaryType(PrimaryType valTy) { … }

constexpr bool isIntegralPrimaryType(PrimaryType valTy) { … }

constexpr bool isRealPrimaryType(PrimaryType valTy) { … }

constexpr bool isComplexPrimaryType(PrimaryType valTy) { … }

/// The actions performed by @newSparseTensor.
enum class Action : uint32_t { … };

/// This enum defines all supported storage format without the level properties.
enum class LevelFormat : uint64_t { … };

constexpr bool encPowOfTwo(LevelFormat fmt) { … }

// All LevelFormats must have only one bit set (power of two).
static_assert …;

template <LevelFormat... targets>
constexpr bool isAnyOfFmt(LevelFormat fmt) { … }

/// Returns string representation of the given level format.
constexpr const char *toFormatString(LevelFormat lvlFmt) { … }

/// This enum defines all the nondefault properties for storage formats.
enum class LevelPropNonDefault : uint64_t { … };

/// Returns string representation of the given level properties.
constexpr const char *toPropString(LevelPropNonDefault lvlProp) { … }

/// This enum defines all the sparse representations supportable by
/// the SparseTensor dialect. We use a lightweight encoding to encode
/// the "format" per se (dense, compressed, singleton, loose_compressed,
/// n-out-of-m), the "properties" (ordered, unique) as well as n and m when
/// the format is NOutOfM.
/// The encoding is chosen for performance of the runtime library, and thus may
/// change in future versions; consequently, client code should use the
/// predicate functions defined below, rather than relying on knowledge
/// about the particular binary encoding.
///
/// The `Undef` "format" is a special value used internally for cases
/// where we need to store an undefined or indeterminate `LevelType`.
/// It should not be used externally, since it does not indicate an
/// actual/representable format.

struct LevelType { … };

// For backward-compatibility. TODO: remove below after fully migration.
constexpr uint64_t nToBits(uint64_t n) { … }
constexpr uint64_t mToBits(uint64_t m) { … }

inline std::optional<LevelType>
buildLevelType(LevelFormat lf,
               const std::vector<LevelPropNonDefault> &properties,
               uint64_t n = 0, uint64_t m = 0) { … }
inline std::optional<LevelType> buildLevelType(LevelFormat lf, bool ordered,
                                               bool unique, uint64_t n = 0,
                                               uint64_t m = 0) { … }
inline bool isUndefLT(LevelType lt) { … }
inline bool isDenseLT(LevelType lt) { … }
inline bool isBatchLT(LevelType lt) { … }
inline bool isCompressedLT(LevelType lt) { … }
inline bool isLooseCompressedLT(LevelType lt) { … }
inline bool isSingletonLT(LevelType lt) { … }
inline bool isNOutOfMLT(LevelType lt) { … }
inline bool isOrderedLT(LevelType lt) { … }
inline bool isUniqueLT(LevelType lt) { … }
inline bool isWithCrdLT(LevelType lt) { … }
inline bool isWithPosLT(LevelType lt) { … }
inline bool isValidLT(LevelType lt) { … }
inline std::optional<LevelFormat> getLevelFormat(LevelType lt) { … }
inline uint64_t getN(LevelType lt) { … }
inline uint64_t getM(LevelType lt) { … }
inline bool isValidNOutOfMLT(LevelType lt, uint64_t n, uint64_t m) { … }
inline std::string toMLIRString(LevelType lt) { … }

/// Bit manipulations for affine encoding.
///
/// Note that because the indices in the mappings refer to dimensions
/// and levels (and *not* the sizes of these dimensions and levels), the
/// 64-bit encoding gives ample room for a compact encoding of affine
/// operations in the higher bits. Pure permutations still allow for
/// 60-bit indices. But non-permutations reserve 20-bits for the
/// potential three components (index i, constant, index ii).
///
/// The compact encoding is as follows:
///
///  0xffffffffffffffff
/// |0000      |                        60-bit idx| e.g. i
/// |0001 floor|           20-bit const|20-bit idx| e.g. i floor c
/// |0010 mod  |           20-bit const|20-bit idx| e.g. i mod c
/// |0011 mul  |20-bit idx|20-bit const|20-bit idx| e.g. i + c * ii
///
/// This encoding provides sufficient generality for currently supported
/// sparse tensor types. To generalize this more, we will need to provide
/// a broader encoding scheme for affine functions. Also, the library
/// encoding may be replaced with pure "direct-IR" code in the future.
///
constexpr uint64_t encodeDim(uint64_t i, uint64_t cf, uint64_t cm) { … }
constexpr uint64_t encodeLvl(uint64_t i, uint64_t c, uint64_t ii) { … }
constexpr bool isEncodedFloor(uint64_t v) { … }
constexpr bool isEncodedMod(uint64_t v) { … }
constexpr bool isEncodedMul(uint64_t v) { … }
constexpr uint64_t decodeIndex(uint64_t v) { … }
constexpr uint64_t decodeConst(uint64_t v) { … }
constexpr uint64_t decodeMulc(uint64_t v) { … }
constexpr uint64_t decodeMuli(uint64_t v) { … }

} // namespace sparse_tensor
} // namespace mlir

#endif // MLIR_DIALECT_SPARSETENSOR_IR_ENUMS_H