src/frontend/dtype.cc
Namespaces
Defines
Macros Documentation
define TENZR_DEFINE_NARROW_FLOAT
#define TENZR_DEFINE_NARROW_FLOAT(
Type,
MlType
)
Type::Type(float value) : bits(MlType(value).rep()) {} \
Type::Type(double value) : Type(static_cast<float>(value)) {} \
Type::operator float() const { return static_cast<float>(MlType::FromRep(bits)); } \
Type::operator double() const { return static_cast<double>(MlType::FromRep(bits)); }
define TENZR_DEFINE_SUBBYTE_INT
#define TENZR_DEFINE_SUBBYTE_INT(
Type,
MlType,
Rep,
Bits
)
Type::Type(int64_t value) : bits(SubbyteToRep<MlType, Rep>(value, Bits)) {} \
Type::Type(double value) : Type(static_cast<int64_t>(value)) {} \
Type::operator int64_t() const { return SubbyteFromRep<MlType, Rep>(bits); }
Source code
// 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.
#include "tenzr/frontend/dtype.h"
#include <Eigen/Core>
#include <float8.h>
#include <intn.h>
#include <mxfloat.h>
#include "tenzr/common/error.h"
namespace tenzr {
namespace {
// Eigen's 16-bit types are layout-compatible with the uint16_t storage of
// float16/bfloat16; round-trip through their own conversions so rounding
// matches exactly.
template <typename EigenType> uint16_t NarrowFloatToBits(float value) {
const EigenType converted(value);
return Eigen::numext::bit_cast<uint16_t>(converted);
}
template <typename EigenType> float NarrowFloatFromBits(uint16_t bits) {
const EigenType converted = Eigen::numext::bit_cast<EigenType>(bits);
return static_cast<float>(converted);
}
// Store the low `bits` of an ml_dtypes sub-byte integer (upper bits are
// unspecified by ml_dtypes, so mask for a canonical representation).
template <typename MlType, typename Rep> Rep SubbyteToRep(int64_t value, int bits) {
const int64_t masked = static_cast<int64_t>(MlType(static_cast<Rep>(value))) & ((int64_t{1} << bits) - 1);
return static_cast<Rep>(masked);
}
template <typename MlType, typename Rep> int64_t SubbyteFromRep(Rep bits) { return static_cast<int64_t>(MlType(bits)); }
} // namespace
float16::float16(float value) : bits(NarrowFloatToBits<Eigen::half>(value)) {}
float16::float16(double value) : float16(static_cast<float>(value)) {}
float16::operator float() const { return NarrowFloatFromBits<Eigen::half>(bits); }
float16::operator double() const { return static_cast<double>(NarrowFloatFromBits<Eigen::half>(bits)); }
bfloat16::bfloat16(float value) : bits(NarrowFloatToBits<Eigen::bfloat16>(value)) {}
bfloat16::bfloat16(double value) : bfloat16(static_cast<float>(value)) {}
bfloat16::operator float() const { return NarrowFloatFromBits<Eigen::bfloat16>(bits); }
bfloat16::operator double() const { return static_cast<double>(NarrowFloatFromBits<Eigen::bfloat16>(bits)); }
#define TENZR_DEFINE_NARROW_FLOAT(Type, MlType) \
Type::Type(float value) : bits(MlType(value).rep()) {} \
Type::Type(double value) : Type(static_cast<float>(value)) {} \
Type::operator float() const { return static_cast<float>(MlType::FromRep(bits)); } \
Type::operator double() const { return static_cast<double>(MlType::FromRep(bits)); }
TENZR_DEFINE_NARROW_FLOAT(float4_e2m1fn, ml_dtypes::float4_e2m1fn)
TENZR_DEFINE_NARROW_FLOAT(float8_e3m4, ml_dtypes::float8_e3m4)
TENZR_DEFINE_NARROW_FLOAT(float8_e4m3, ml_dtypes::float8_e4m3)
TENZR_DEFINE_NARROW_FLOAT(float8_e4m3fn, ml_dtypes::float8_e4m3fn)
TENZR_DEFINE_NARROW_FLOAT(float8_e4m3fnuz, ml_dtypes::float8_e4m3fnuz)
TENZR_DEFINE_NARROW_FLOAT(float8_e4m3b11fnuz, ml_dtypes::float8_e4m3b11fnuz)
TENZR_DEFINE_NARROW_FLOAT(float8_e5m2, ml_dtypes::float8_e5m2)
TENZR_DEFINE_NARROW_FLOAT(float8_e5m2fnuz, ml_dtypes::float8_e5m2fnuz)
TENZR_DEFINE_NARROW_FLOAT(float8_e8m0fnu, ml_dtypes::float8_e8m0fnu)
TENZR_DEFINE_NARROW_FLOAT(float6_e2m3fn, ml_dtypes::float6_e2m3fn)
TENZR_DEFINE_NARROW_FLOAT(float6_e3m2fn, ml_dtypes::float6_e3m2fn)
#undef TENZR_DEFINE_NARROW_FLOAT
#define TENZR_DEFINE_SUBBYTE_INT(Type, MlType, Rep, Bits) \
Type::Type(int64_t value) : bits(SubbyteToRep<MlType, Rep>(value, Bits)) {} \
Type::Type(double value) : Type(static_cast<int64_t>(value)) {} \
Type::operator int64_t() const { return SubbyteFromRep<MlType, Rep>(bits); }
TENZR_DEFINE_SUBBYTE_INT(int1, ml_dtypes::int1, int8_t, 1)
TENZR_DEFINE_SUBBYTE_INT(uint1, ml_dtypes::uint1, uint8_t, 1)
TENZR_DEFINE_SUBBYTE_INT(int2, ml_dtypes::int2, int8_t, 2)
TENZR_DEFINE_SUBBYTE_INT(uint2, ml_dtypes::uint2, uint8_t, 2)
TENZR_DEFINE_SUBBYTE_INT(int4, ml_dtypes::int4, int8_t, 4)
TENZR_DEFINE_SUBBYTE_INT(uint4, ml_dtypes::uint4, uint8_t, 4)
#undef TENZR_DEFINE_SUBBYTE_INT
size_t SizeOf(DType dtype) {
switch (dtype) {
case DType::kBool:
case DType::kInt8:
case DType::kUInt8:
case DType::kFloat8E4M3FN:
case DType::kFloat8E5M2:
case DType::kFloat8E3M4:
case DType::kFloat8E4M3:
case DType::kFloat8E4M3FNUZ:
case DType::kFloat8E4M3B11FNUZ:
case DType::kFloat8E5M2FNUZ:
case DType::kFloat8E8M0FNU:
case DType::kFloat6E2M3FN:
case DType::kFloat6E3M2FN:
case DType::kFloat4E2M1FN:
case DType::kInt1:
case DType::kUInt1:
case DType::kInt2:
case DType::kUInt2:
case DType::kInt4:
case DType::kUInt4:
return 1;
case DType::kInt16:
case DType::kUInt16:
case DType::kFloat16:
case DType::kBFloat16:
return 2;
case DType::kInt32:
case DType::kUInt32:
case DType::kFloat32:
return 4;
case DType::kInt64:
case DType::kUInt64:
case DType::kFloat64:
return 8;
default:
TENZR_THROW("Unknown DType in size_of");
}
}
std::string_view ToString(DType dtype) {
switch (dtype) {
case DType::kBool:
return "bool";
case DType::kFloat32:
return "float32";
case DType::kFloat64:
return "float64";
case DType::kInt32:
return "int32";
case DType::kInt64:
return "int64";
case DType::kFloat16:
return "float16";
case DType::kBFloat16:
return "bfloat16";
case DType::kFloat8E4M3FN:
return "float8_e4m3fn";
case DType::kFloat8E5M2:
return "float8_e5m2";
case DType::kInt8:
return "int8";
case DType::kUInt8:
return "uint8";
case DType::kInt16:
return "int16";
case DType::kUInt16:
return "uint16";
case DType::kUInt32:
return "uint32";
case DType::kUInt64:
return "uint64";
case DType::kFloat8E3M4:
return "float8_e3m4";
case DType::kFloat8E4M3:
return "float8_e4m3";
case DType::kFloat8E4M3FNUZ:
return "float8_e4m3fnuz";
case DType::kFloat8E4M3B11FNUZ:
return "float8_e4m3b11fnuz";
case DType::kFloat8E5M2FNUZ:
return "float8_e5m2fnuz";
case DType::kFloat8E8M0FNU:
return "float8_e8m0fnu";
case DType::kFloat6E2M3FN:
return "float6_e2m3fn";
case DType::kFloat6E3M2FN:
return "float6_e3m2fn";
case DType::kFloat4E2M1FN:
return "float4_e2m1fn";
case DType::kInt1:
return "int1";
case DType::kUInt1:
return "uint1";
case DType::kInt2:
return "int2";
case DType::kUInt2:
return "uint2";
case DType::kInt4:
return "int4";
case DType::kUInt4:
return "uint4";
default:
return "unknown";
}
}
std::string_view ToMlirType(DType dtype) {
switch (dtype) {
case DType::kBool:
return "i1";
case DType::kFloat32:
return "f32";
case DType::kFloat64:
return "f64";
case DType::kInt32:
return "i32";
case DType::kInt64:
return "i64";
case DType::kFloat16:
return "f16";
case DType::kBFloat16:
return "bf16";
case DType::kFloat8E4M3FN:
return "f8E4M3FN";
case DType::kFloat8E5M2:
return "f8E5M2";
case DType::kInt8:
return "i8";
case DType::kUInt8:
return "ui8";
case DType::kInt16:
return "i16";
case DType::kUInt16:
return "ui16";
case DType::kUInt32:
return "ui32";
case DType::kUInt64:
return "ui64";
case DType::kFloat8E3M4:
return "f8E3M4";
case DType::kFloat8E4M3:
return "f8E4M3";
case DType::kFloat8E4M3FNUZ:
return "f8E4M3FNUZ";
case DType::kFloat8E4M3B11FNUZ:
return "f8E4M3B11FNUZ";
case DType::kFloat8E5M2FNUZ:
return "f8E5M2FNUZ";
case DType::kFloat8E8M0FNU:
return "f8E8M0FNU";
case DType::kFloat6E2M3FN:
return "f6E2M3FN";
case DType::kFloat6E3M2FN:
return "f6E3M2FN";
case DType::kFloat4E2M1FN:
return "f4E2M1FN";
case DType::kInt1:
return "i1";
case DType::kUInt1:
return "ui1";
case DType::kInt2:
return "i2";
case DType::kUInt2:
return "ui2";
case DType::kInt4:
return "i4";
case DType::kUInt4:
return "ui4";
default:
TENZR_THROW("Unsupported DType for MLIR emission: " + std::string(ToString(dtype)));
}
}
PJRT_Buffer_Type ToPjrtBufferType(DType dtype) {
switch (dtype) {
case DType::kBool:
return PJRT_Buffer_Type_PRED;
case DType::kFloat32:
return PJRT_Buffer_Type_F32;
case DType::kFloat64:
return PJRT_Buffer_Type_F64;
case DType::kInt32:
return PJRT_Buffer_Type_S32;
case DType::kInt64:
return PJRT_Buffer_Type_S64;
case DType::kFloat16:
return PJRT_Buffer_Type_F16;
case DType::kBFloat16:
return PJRT_Buffer_Type_BF16;
case DType::kFloat8E4M3FN:
return PJRT_Buffer_Type_F8E4M3FN;
case DType::kFloat8E5M2:
return PJRT_Buffer_Type_F8E5M2;
case DType::kInt8:
return PJRT_Buffer_Type_S8;
case DType::kUInt8:
return PJRT_Buffer_Type_U8;
case DType::kInt16:
return PJRT_Buffer_Type_S16;
case DType::kUInt16:
return PJRT_Buffer_Type_U16;
case DType::kUInt32:
return PJRT_Buffer_Type_U32;
case DType::kUInt64:
return PJRT_Buffer_Type_U64;
case DType::kFloat8E3M4:
return PJRT_Buffer_Type_F8E3M4;
case DType::kFloat8E4M3:
return PJRT_Buffer_Type_F8E4M3;
case DType::kFloat8E4M3FNUZ:
return PJRT_Buffer_Type_F8E4M3FNUZ;
case DType::kFloat8E4M3B11FNUZ:
return PJRT_Buffer_Type_F8E4M3B11FNUZ;
case DType::kFloat8E5M2FNUZ:
return PJRT_Buffer_Type_F8E5M2FNUZ;
case DType::kFloat8E8M0FNU:
return PJRT_Buffer_Type_F8E8M0FNU;
case DType::kFloat6E2M3FN:
return PJRT_Buffer_Type_F6E2M3FN;
case DType::kFloat6E3M2FN:
return PJRT_Buffer_Type_F6E3M2FN;
case DType::kFloat4E2M1FN:
return PJRT_Buffer_Type_F4E2M1FN;
case DType::kInt1:
return PJRT_Buffer_Type_S1;
case DType::kUInt1:
return PJRT_Buffer_Type_U1;
case DType::kInt2:
return PJRT_Buffer_Type_S2;
case DType::kUInt2:
return PJRT_Buffer_Type_U2;
case DType::kInt4:
return PJRT_Buffer_Type_S4;
case DType::kUInt4:
return PJRT_Buffer_Type_U4;
default:
TENZR_THROW("Unsupported DType for PJRT buffer mapping: " + std::string(ToString(dtype)));
}
}
DType FromPjrtBufferType(PJRT_Buffer_Type type) {
switch (type) {
case PJRT_Buffer_Type_PRED:
return DType::kBool;
case PJRT_Buffer_Type_F32:
return DType::kFloat32;
case PJRT_Buffer_Type_F64:
return DType::kFloat64;
case PJRT_Buffer_Type_S32:
return DType::kInt32;
case PJRT_Buffer_Type_S64:
return DType::kInt64;
case PJRT_Buffer_Type_F16:
return DType::kFloat16;
case PJRT_Buffer_Type_BF16:
return DType::kBFloat16;
case PJRT_Buffer_Type_F8E4M3FN:
return DType::kFloat8E4M3FN;
case PJRT_Buffer_Type_F8E5M2:
return DType::kFloat8E5M2;
case PJRT_Buffer_Type_S8:
return DType::kInt8;
case PJRT_Buffer_Type_U8:
return DType::kUInt8;
case PJRT_Buffer_Type_S16:
return DType::kInt16;
case PJRT_Buffer_Type_U16:
return DType::kUInt16;
case PJRT_Buffer_Type_U32:
return DType::kUInt32;
case PJRT_Buffer_Type_U64:
return DType::kUInt64;
case PJRT_Buffer_Type_F8E3M4:
return DType::kFloat8E3M4;
case PJRT_Buffer_Type_F8E4M3:
return DType::kFloat8E4M3;
case PJRT_Buffer_Type_F8E4M3FNUZ:
return DType::kFloat8E4M3FNUZ;
case PJRT_Buffer_Type_F8E4M3B11FNUZ:
return DType::kFloat8E4M3B11FNUZ;
case PJRT_Buffer_Type_F8E5M2FNUZ:
return DType::kFloat8E5M2FNUZ;
case PJRT_Buffer_Type_F8E8M0FNU:
return DType::kFloat8E8M0FNU;
case PJRT_Buffer_Type_F6E2M3FN:
return DType::kFloat6E2M3FN;
case PJRT_Buffer_Type_F6E3M2FN:
return DType::kFloat6E3M2FN;
case PJRT_Buffer_Type_F4E2M1FN:
return DType::kFloat4E2M1FN;
case PJRT_Buffer_Type_S1:
return DType::kInt1;
case PJRT_Buffer_Type_U1:
return DType::kUInt1;
case PJRT_Buffer_Type_S2:
return DType::kInt2;
case PJRT_Buffer_Type_U2:
return DType::kUInt2;
case PJRT_Buffer_Type_S4:
return DType::kInt4;
case PJRT_Buffer_Type_U4:
return DType::kUInt4;
default:
TENZR_THROW("Unsupported PJRT_Buffer_Type for DType mapping: " + std::to_string(type));
}
}
bool IsFloatingPoint(DType dtype) {
return dtype == DType::kFloat32 || dtype == DType::kFloat64 || dtype == DType::kFloat16 || dtype == DType::kBFloat16 || dtype == DType::kFloat8E4M3FN ||
dtype == DType::kFloat8E5M2 || dtype == DType::kFloat8E3M4 || dtype == DType::kFloat8E4M3 || dtype == DType::kFloat8E4M3FNUZ ||
dtype == DType::kFloat8E4M3B11FNUZ || dtype == DType::kFloat8E5M2FNUZ || dtype == DType::kFloat8E8M0FNU || dtype == DType::kFloat6E2M3FN ||
dtype == DType::kFloat6E3M2FN || dtype == DType::kFloat4E2M1FN;
}
bool IsIntegral(DType dtype) {
return dtype == DType::kInt8 || dtype == DType::kUInt8 || dtype == DType::kInt16 || dtype == DType::kUInt16 || dtype == DType::kInt32 ||
dtype == DType::kUInt32 || dtype == DType::kInt64 || dtype == DType::kUInt64 || dtype == DType::kInt1 || dtype == DType::kUInt1 ||
dtype == DType::kInt2 || dtype == DType::kUInt2 || dtype == DType::kInt4 || dtype == DType::kUInt4;
}
bool IsBoolean(DType dtype) { return dtype == DType::kBool; }
bool IsSignedIntegral(DType dtype) {
return dtype == DType::kInt8 || dtype == DType::kInt16 || dtype == DType::kInt32 || dtype == DType::kInt64 || dtype == DType::kInt1 ||
dtype == DType::kInt2 || dtype == DType::kInt4;
}
bool IsUnsignedIntegral(DType dtype) {
return dtype == DType::kUInt8 || dtype == DType::kUInt16 || dtype == DType::kUInt32 || dtype == DType::kUInt64 || dtype == DType::kUInt1 ||
dtype == DType::kUInt2 || dtype == DType::kUInt4;
}
namespace {
bool IsNarrowFloat(DType dtype) { return IsFloatingPoint(dtype) && SizeOf(dtype) == 1; }
DType PromoteIntegers(DType a, DType b) {
if (a == b) {
return a;
}
bool a_signed = IsSignedIntegral(a);
bool b_signed = IsSignedIntegral(b);
size_t size_a = SizeOf(a);
size_t size_b = SizeOf(b);
// Same signedness: larger size wins
if (a_signed == b_signed) {
return (size_a >= size_b) ? a : b;
}
// Mixed signedness: identify signed vs unsigned
DType s = a_signed ? a : b;
DType u = a_signed ? b : a;
size_t size_s = SizeOf(s);
size_t size_u = SizeOf(u);
if (size_s > size_u) {
return s;
}
// size_u >= size_s: promote to next wider signed integer if available
size_t max_size = std::max(size_s, size_u);
if (max_size == 1) {
return DType::kInt16;
}
if (max_size == 2) {
return DType::kInt32;
}
if (max_size == 4) {
return DType::kInt64;
}
// 64-bit int and 64-bit uint: promote to Float64
return DType::kFloat64;
}
DType PromoteFloatAndInt(DType flt, DType in) {
if (flt == DType::kFloat64) {
return DType::kFloat64;
}
size_t size_int = SizeOf(in);
if (flt == DType::kFloat32) {
if (size_int >= 8) {
return DType::kFloat64;
}
return DType::kFloat32;
}
// Float16 or BFloat16
if (size_int >= 4) {
return DType::kFloat64;
}
if (size_int >= 2) {
return DType::kFloat32;
}
return flt;
}
} // namespace
DType PromoteTypes(DType a, DType b) {
if (a == b) {
return a;
}
// Bool yields to any other type
if (a == DType::kBool) {
return b;
}
if (b == DType::kBool) {
return a;
}
// 8-bit floats do not support implicit cross-type promotion (matches JAX / OpenXLA specification)
if (IsNarrowFloat(a) || IsNarrowFloat(b)) {
TENZR_THROW("Types " + std::string(ToString(a)) + " and " + std::string(ToString(b)) +
" have no available implicit promotion path. 8-bit floats require explicit conversion via astype().");
}
bool a_flt = IsFloatingPoint(a);
bool b_flt = IsFloatingPoint(b);
// Both floating point
if (a_flt && b_flt) {
if (a == DType::kFloat64 || b == DType::kFloat64) {
return DType::kFloat64;
}
if (a == DType::kFloat32 || b == DType::kFloat32) {
return DType::kFloat32;
}
// BFloat16 + Float16 -> kFloat32 (fixes F4)
if ((a == DType::kBFloat16 && b == DType::kFloat16) || (a == DType::kFloat16 && b == DType::kBFloat16)) {
return DType::kFloat32;
}
return a;
}
// Both integral
if (!a_flt && !b_flt) {
return PromoteIntegers(a, b);
}
// Mixed floating point and integer
DType flt = a_flt ? a : b;
DType in = a_flt ? b : a;
return PromoteFloatAndInt(flt, in);
}
bool CanCast(DType from, DType to, CastingMode mode) {
if (from == to) {
return true;
}
switch (mode) {
case CastingMode::kNo:
case CastingMode::kEquiv:
return from == to;
case CastingMode::kUnsafe:
return true;
case CastingMode::kSameKind:
if (IsBoolean(from) && IsBoolean(to)) {
return true;
}
if (IsIntegral(from) && IsIntegral(to)) {
return true;
}
if (IsFloatingPoint(from) && IsFloatingPoint(to)) {
return true;
}
return false;
case CastingMode::kSafe: {
// Bool safely casts to everything
if (IsBoolean(from)) {
return true;
}
if (IsBoolean(to)) {
return false;
}
// Both floating point
if (IsFloatingPoint(from) && IsFloatingPoint(to)) {
if (IsNarrowFloat(from)) {
return !IsNarrowFloat(to);
}
if (from == DType::kFloat16) {
return to == DType::kFloat32 || to == DType::kFloat64;
}
if (from == DType::kBFloat16) {
return to == DType::kFloat32 || to == DType::kFloat64;
}
if (from == DType::kFloat32) {
return to == DType::kFloat64;
}
return false;
}
// Both integral
if (IsIntegral(from) && IsIntegral(to)) {
bool from_signed = IsSignedIntegral(from);
bool to_signed = IsSignedIntegral(to);
size_t s_from = SizeOf(from);
size_t s_to = SizeOf(to);
if (from_signed && !to_signed) {
return false;
}
if (!from_signed && to_signed) {
return s_to > s_from;
}
// Same signedness
return s_to >= s_from;
}
// Integer to float
if (IsIntegral(from) && IsFloatingPoint(to)) {
size_t s_from = SizeOf(from);
if (to == DType::kFloat64) {
return s_from <= 4;
}
if (to == DType::kFloat32) {
return s_from <= 2;
}
if (to == DType::kFloat16 || to == DType::kBFloat16) {
return s_from <= 1;
}
return false;
}
// Float to integer is never safe
return false;
}
default:
TENZR_THROW("Unknown CastingMode");
}
}
} // namespace tenzr