// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#ifndef TENZR_FRONTEND_DTYPE_H_
#define TENZR_FRONTEND_DTYPE_H_
#include <xla/pjrt/c/pjrt_c_api.h>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <string>
#include <string_view>
#include <type_traits>
namespace tenzr {
// Narrow element types are owned by Tenzr as bit-compatible PODs; their
// conversions are defined out-of-line in dtype.cpp. This keeps Eigen and
// ml_dtypes out of the public headers (they stay build-only dependencies).
// Each type is trivially copyable, so host buffers move with a plain memcpy.
struct float16 {
uint16_t bits = 0;
float16() = default;
explicit float16(float value);
explicit float16(double value);
operator float() const;
operator double() const;
friend bool operator==(const float16&, const float16&) = default;
};
struct bfloat16 {
uint16_t bits = 0;
bfloat16() = default;
explicit bfloat16(float value);
explicit bfloat16(double value);
operator float() const;
operator double() const;
friend bool operator==(const bfloat16&, const bfloat16&) = default;
};
#define TENZR_DECLARE_NARROW_FLOAT(Type) \
struct Type { \
uint8_t bits = 0; \
Type() = default; \
explicit Type(float value); \
explicit Type(double value); \
operator float() const; \
operator double() const; \
friend bool operator==(const Type&, const Type&) = default; \
}
TENZR_DECLARE_NARROW_FLOAT(float4_e2m1fn);
TENZR_DECLARE_NARROW_FLOAT(float8_e3m4);
TENZR_DECLARE_NARROW_FLOAT(float8_e4m3);
TENZR_DECLARE_NARROW_FLOAT(float8_e4m3fn);
TENZR_DECLARE_NARROW_FLOAT(float8_e4m3fnuz);
TENZR_DECLARE_NARROW_FLOAT(float8_e4m3b11fnuz);
TENZR_DECLARE_NARROW_FLOAT(float8_e5m2);
TENZR_DECLARE_NARROW_FLOAT(float8_e5m2fnuz);
TENZR_DECLARE_NARROW_FLOAT(float8_e8m0fnu);
TENZR_DECLARE_NARROW_FLOAT(float6_e2m3fn);
TENZR_DECLARE_NARROW_FLOAT(float6_e3m2fn);
#undef TENZR_DECLARE_NARROW_FLOAT
#define TENZR_DECLARE_SUBBYTE_INT(Type, Rep) \
struct Type { \
Rep bits = 0; \
Type() = default; \
explicit Type(int64_t value); \
explicit Type(double value); \
operator int64_t() const; \
friend bool operator==(const Type&, const Type&) = default; \
}
TENZR_DECLARE_SUBBYTE_INT(int1, int8_t);
TENZR_DECLARE_SUBBYTE_INT(uint1, uint8_t);
TENZR_DECLARE_SUBBYTE_INT(int2, int8_t);
TENZR_DECLARE_SUBBYTE_INT(uint2, uint8_t);
TENZR_DECLARE_SUBBYTE_INT(int4, int8_t);
TENZR_DECLARE_SUBBYTE_INT(uint4, uint8_t);
#undef TENZR_DECLARE_SUBBYTE_INT
enum class DType : uint8_t {
kBool,
kFloat32,
kFloat64,
kInt32,
kInt64,
kFloat16,
kBFloat16,
kFloat8E4M3FN,
kFloat8E5M2,
kInt8,
kUInt8,
kInt16,
kUInt16,
kUInt32,
kUInt64,
kFloat8E3M4,
kFloat8E4M3,
kFloat8E4M3FNUZ,
kFloat8E4M3B11FNUZ,
kFloat8E5M2FNUZ,
kFloat8E8M0FNU,
kFloat6E2M3FN,
kFloat6E3M2FN,
kFloat4E2M1FN,
kInt1,
kUInt1,
kInt2,
kUInt2,
kInt4,
kUInt4
};
enum class CastingMode : uint8_t { kNo, kEquiv, kSafe, kSameKind, kUnsafe };
size_t SizeOf(DType dtype);
std::string_view ToString(DType dtype);
std::string_view ToMlirType(DType dtype);
PJRT_Buffer_Type ToPjrtBufferType(DType dtype);
DType FromPjrtBufferType(PJRT_Buffer_Type type);
bool IsFloatingPoint(DType dtype);
bool IsIntegral(DType dtype);
bool IsBoolean(DType dtype);
bool IsSignedIntegral(DType dtype);
bool IsUnsignedIntegral(DType dtype);
DType PromoteTypes(DType a, DType b);
bool CanCast(DType from, DType to, CastingMode mode = CastingMode::kSafe);
template <typename T> struct DTypeTraits;
template <> struct DTypeTraits<bool> {
static constexpr DType type = DType::kBool;
};
template <> struct DTypeTraits<int8_t> {
static constexpr DType type = DType::kInt8;
};
template <> struct DTypeTraits<uint8_t> {
static constexpr DType type = DType::kUInt8;
};
template <> struct DTypeTraits<int16_t> {
static constexpr DType type = DType::kInt16;
};
template <> struct DTypeTraits<uint16_t> {
static constexpr DType type = DType::kUInt16;
};
template <> struct DTypeTraits<int32_t> {
static constexpr DType type = DType::kInt32;
};
template <> struct DTypeTraits<uint32_t> {
static constexpr DType type = DType::kUInt32;
};
template <> struct DTypeTraits<int64_t> {
static constexpr DType type = DType::kInt64;
};
template <> struct DTypeTraits<uint64_t> {
static constexpr DType type = DType::kUInt64;
};
#if defined(__GNUC__) || defined(__clang__)
template <> struct DTypeTraits<long long> {
static constexpr DType type = DType::kInt64;
};
template <> struct DTypeTraits<unsigned long long> {
static constexpr DType type = DType::kUInt64;
};
#endif
template <> struct DTypeTraits<float16> {
static constexpr DType type = DType::kFloat16;
};
template <> struct DTypeTraits<bfloat16> {
static constexpr DType type = DType::kBFloat16;
};
template <> struct DTypeTraits<float> {
static constexpr DType type = DType::kFloat32;
};
template <> struct DTypeTraits<double> {
static constexpr DType type = DType::kFloat64;
};
template <> struct DTypeTraits<float8_e4m3fn> {
static constexpr DType type = DType::kFloat8E4M3FN;
};
template <> struct DTypeTraits<float8_e5m2> {
static constexpr DType type = DType::kFloat8E5M2;
};
template <> struct DTypeTraits<float8_e3m4> {
static constexpr DType type = DType::kFloat8E3M4;
};
template <> struct DTypeTraits<float8_e4m3> {
static constexpr DType type = DType::kFloat8E4M3;
};
template <> struct DTypeTraits<float8_e4m3fnuz> {
static constexpr DType type = DType::kFloat8E4M3FNUZ;
};
template <> struct DTypeTraits<float8_e4m3b11fnuz> {
static constexpr DType type = DType::kFloat8E4M3B11FNUZ;
};
template <> struct DTypeTraits<float8_e5m2fnuz> {
static constexpr DType type = DType::kFloat8E5M2FNUZ;
};
template <> struct DTypeTraits<float8_e8m0fnu> {
static constexpr DType type = DType::kFloat8E8M0FNU;
};
template <> struct DTypeTraits<float6_e2m3fn> {
static constexpr DType type = DType::kFloat6E2M3FN;
};
template <> struct DTypeTraits<float6_e3m2fn> {
static constexpr DType type = DType::kFloat6E3M2FN;
};
template <> struct DTypeTraits<float4_e2m1fn> {
static constexpr DType type = DType::kFloat4E2M1FN;
};
template <> struct DTypeTraits<int1> {
static constexpr DType type = DType::kInt1;
};
template <> struct DTypeTraits<uint1> {
static constexpr DType type = DType::kUInt1;
};
template <> struct DTypeTraits<int2> {
static constexpr DType type = DType::kInt2;
};
template <> struct DTypeTraits<uint2> {
static constexpr DType type = DType::kUInt2;
};
template <> struct DTypeTraits<int4> {
static constexpr DType type = DType::kInt4;
};
template <> struct DTypeTraits<uint4> {
static constexpr DType type = DType::kUInt4;
};
template <typename T> inline constexpr DType DtypeOf() { return DTypeTraits<T>::type; }
namespace concepts {
template <typename T>
concept ValidElement = requires {
{ DTypeTraits<std::remove_cvref_t<T>>::type } -> std::same_as<const DType&>;
};
namespace detail {
template <typename T> inline constexpr bool kIsTenzrFloat =
std::is_same_v<T, float> || std::is_same_v<T, double> || std::is_same_v<T, float16> || std::is_same_v<T, bfloat16> ||
std::is_same_v<T, float8_e3m4> || std::is_same_v<T, float8_e4m3> || std::is_same_v<T, float8_e4m3fn> || std::is_same_v<T, float8_e4m3fnuz> ||
std::is_same_v<T, float8_e4m3b11fnuz> || std::is_same_v<T, float8_e5m2> || std::is_same_v<T, float8_e5m2fnuz> || std::is_same_v<T, float8_e8m0fnu> ||
std::is_same_v<T, float6_e2m3fn> || std::is_same_v<T, float6_e3m2fn> || std::is_same_v<T, float4_e2m1fn>;
} // namespace detail
template <typename T>
concept FloatingPoint = ValidElement<T> && detail::kIsTenzrFloat<std::remove_cvref_t<T>>;
template <typename T>
concept Integral = ValidElement<T> && !FloatingPoint<T> && !std::is_same_v<std::remove_cvref_t<T>, bool>;
} // namespace concepts
} // namespace tenzr
#endif // TENZR_FRONTEND_DTYPE_H_