Visitar URL original
Move float_to_bits()/float_from_bits() to C++ · NVIDIA/cutile-python@285b3a7 · GitHub
Skip to content

Commit 285b3a7

Browse files
gbonikgbonik
committed
Move float_to_bits()/float_from_bits() to C++
Signed-off-by: Greg Bonik <gbonik@nvidia.com>
1 parent 12660a0 commit 285b3a7

18 files changed

Lines changed: 483 additions & 280 deletions

File tree

‎cext/dtype.cpp‎

Lines changed: 317 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,10 @@
66
#include "vec.h"
77
#include "py.h"
88

9+
#if defined(_MSC_VER)
10+
#include <intrin.h>
11+
#endif
12+
913
static PyObject* g_module;
1014
static PyObject* g_pointee_dtype_pyunicode;
1115
static PyObject* g_memory_space_pyunicode;
@@ -77,6 +81,50 @@ static std::optional<MemorySpace> memory_space_from_pyobject(PyObject* space_obj
7781
}
7882

7983

84+
// ----- FloatFormat constants -----
85+
86+
#define IEEE_LIKE(emin, emax, bitwidth, exp_bits, precision) \
87+
{ emin, emax, bitwidth, exp_bits, precision, FloatFormat::IEEE, true, true }
88+
89+
90+
// The bitwidth of tfloat32 is that of its encoding (19 bits), not its 32-bit container.
91+
static constexpr FloatFormat
92+
k_format_float16 = IEEE_LIKE(-14, 15, 16, 5, 10),
93+
k_format_float32 = IEEE_LIKE(-126, 127, 32, 8, 23),
94+
k_format_float64 = IEEE_LIKE(-1022, 1023, 64, 11, 52),
95+
k_format_bfloat16 = IEEE_LIKE(-126, 127, 16, 8, 7),
96+
k_format_tfloat32 = IEEE_LIKE(-126, 127, 19, 8, 10),
97+
k_format_float8_e5m2 = IEEE_LIKE(-14, 15, 8, 5, 2),
98+
k_format_float8_e4m3fn = { -6, 8, 8, 4, 3, FloatFormat::NanOnlyAllOnes, true, true },
99+
k_format_float8_e5m3fnu = { -14, 16, 8, 5, 3, FloatFormat::NanOnlyAllOnes, false, true },
100+
k_format_float8_e8m0fnu = { -127, 127, 8, 8, 0, FloatFormat::NanOnlyAllOnes, false, false },
101+
k_format_float4_e2m1fn = { 0, 2, 4, 2, 1, FloatFormat::FiniteOnly, true, true },
102+
k_format_float6_e2m3fn = { 0, 2, 6, 2, 3, FloatFormat::FiniteOnly, true, true },
103+
k_format_float6_e3m2fn = { -2, 4, 6, 3, 2, FloatFormat::FiniteOnly, true, true };
104+
105+
#undef IEEE_LIKE
106+
107+
static constexpr FloatFormat float_formats[] = {
108+
#define FLOAT_FORMAT_ENTRY(name, _bitwidth, _signed, _doc) k_format_##name,
109+
FOREACH_UNRESTRICTED_FLOAT_DTYPE(FLOAT_FORMAT_ENTRY)
110+
FOREACH_RESTRICTED_FLOAT_DTYPE(FLOAT_FORMAT_ENTRY)
111+
#undef FLOAT_FORMAT_ENTRY
112+
};
113+
114+
static_assert(std::extent_v<decltype(float_formats)>
115+
== _dtype_detail::restricted_float_last - _dtype_detail::unrestricted_float_first + 1);
116+
117+
static constexpr bool float_format_bitwidths_are_consistent() {
118+
for (const FloatFormat& f : float_formats) {
119+
if (f.have_sign + f.exp_bits + f.precision != f.bitwidth)
120+
return false;
121+
}
122+
return true;
123+
}
124+
125+
static_assert(float_format_bitwidths_are_consistent());
126+
127+
80128
// ----- DType implementation -----
81129

82130
namespace _dtype_detail {
@@ -154,6 +202,11 @@ Integer integer_dtype_max(DType dtype) {
154202
}
155203
}
156204

205+
const FloatFormat& float_dtype_format(DType dtype) {
206+
CHECK(is_float_dtype(dtype));
207+
return float_formats[dtype.dtype_id - _dtype_detail::unrestricted_float_first];
208+
}
209+
157210
DType pointer_dtype(DType pointee_dtype, MemorySpace memory_space, GlobalLock& lock) {
158211
uint32_t space_idx = static_cast<uint32_t>(memory_space);
159212
CHECK(space_idx <= kMemorySpaceMax);
@@ -386,6 +439,218 @@ std::optional<DType> promote_dtypes(DType a, DType b) {
386439
}
387440

388441

442+
// ----- Float encoding -----
443+
444+
static int bit_length(uint64_t x) {
445+
#if defined(_MSC_VER)
446+
unsigned long index;
447+
return _BitScanReverse64(&index, x) ? int(index) + 1 : 0;
448+
#else
449+
return x ? 64 - __builtin_clzll(x) : 0;
450+
#endif
451+
}
452+
453+
static Result<uint64_t> double_nonfinite_to_bits(uint64_t f64_bits, DType dtype) {
454+
const FloatFormat& fmt = float_dtype_format(dtype);
455+
uint64_t f64_sign = f64_bits & 0x8000'0000'0000'0000;
456+
uint64_t sign_bit = f64_sign >> (64 - fmt.bitwidth);
457+
bool is_nan = f64_bits & 0x000f'ffff'ffff'ffff;
458+
uint64_t all_ones_but_sign = (1ull << (fmt.bitwidth - 1)) - 1;
459+
460+
switch (fmt.non_finite_behavior) {
461+
case FloatFormat::NanOnlyAllOnes:
462+
// Infinity also becomes NaN
463+
if (!fmt.have_sign) return (1ull << fmt.bitwidth) - 1;
464+
return sign_bit | all_ones_but_sign;
465+
case FloatFormat::FiniteOnly:
466+
// Infinity saturates to the largest magnitude
467+
if (is_nan)
468+
return raise(PyExc_ValueError, "Cannot represent NaN as ", dtype,
469+
", which has no NaN encoding");
470+
return sign_bit | all_ones_but_sign;
471+
case FloatFormat::IEEE: {
472+
// Keep the sign and the all-ones exponent, and truncate the NaN payload.
473+
// A signaling NaN could lose all of its payload bits and become an infinity,
474+
// so it gets a nonzero payload instead.
475+
uint64_t sign_and_exp = (f64_bits >> (63 - fmt.exp_bits)) << fmt.precision;
476+
uint64_t payload;
477+
if (is_nan && !(f64_bits & (1ull << 51)))
478+
payload = 1;
479+
else
480+
payload = (f64_bits >> (52 - fmt.precision)) & ((1ull << fmt.precision) - 1);
481+
return sign_and_exp | payload;
482+
}
483+
}
484+
CHECK_UNREACHABLE;
485+
}
486+
487+
static Result<uint64_t> raise_negative_unsigned(DType dtype) {
488+
return raise(PyExc_ValueError, "Cannot represent a negative value as ", dtype,
489+
", which is unsigned");
490+
}
491+
492+
493+
static uint64_t first_invalid_float_value(const FloatFormat& fmt) {
494+
switch (fmt.non_finite_behavior) {
495+
case FloatFormat::IEEE:
496+
return ((1ull << fmt.exp_bits) - 1) << fmt.precision;
497+
case FloatFormat::NanOnlyAllOnes:
498+
return (1ull << (fmt.exp_bits + fmt.precision)) - 1;
499+
case FloatFormat::FiniteOnly:
500+
return 1ull << (fmt.exp_bits + fmt.precision);
501+
}
502+
CHECK_UNREACHABLE;
503+
}
504+
505+
// The 64-bit significand is interpreted as a binary fraction in [1, 2):
506+
// 1.xxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx xxxxxxxx
507+
// `--------------------------- 63 bits --------------------------------'
508+
// The top bit of the significand must be explicitly set to 1.
509+
//
510+
// Rounding mode is "to nearest, ties to even".
511+
static Result<uint64_t> encode_float_parts(bool sign, int e, uint64_t significand, DType dtype) {
512+
CHECK(significand >> 63);
513+
const FloatFormat& fmt = float_dtype_format(dtype);
514+
if (sign && !fmt.have_sign)
515+
return raise_negative_unsigned(dtype);
516+
517+
uint64_t f64_inf = (uint64_t(sign) << 63) | 0x7ff0'0000'0000'0000;
518+
if (e > fmt.emax)
519+
return double_nonfinite_to_bits(f64_inf, dtype);
520+
521+
// Round the significand to `fmt.precision` (plus one leading 1) bits
522+
int shift = 63 - fmt.precision;
523+
if (e < fmt.emin) {
524+
if (!fmt.have_zero_and_subnormals)
525+
return 0; // round up to the smallest value
526+
// Subnormal
527+
shift += fmt.emin - e;
528+
e = fmt.emin;
529+
}
530+
531+
if (shift > 64) {
532+
significand = 0; // less than half of the smallest subnormal
533+
} else if (shift == 64) {
534+
// At least half of the smallest subnormal: round up, unless it's a tie (to even zero)
535+
significand = significand > (1ull << 63);
536+
} else {
537+
uint64_t rem = significand & ((1ull << shift) - 1);
538+
uint64_t half = 1ull << (shift - 1);
539+
significand >>= shift;
540+
// Round to nearest, ties to even
541+
if (rem > half || (rem == half && (significand & 1)))
542+
++significand;
543+
}
544+
545+
// These two lines are doing a lot of work. We intentionally encode the biased exponent
546+
// as one less the actual value, so that the explicit leading 1 in the `significand` will
547+
// increment it by one. This automatically takes care of the subnormals, as well as the
548+
// rounding to the next power of two.
549+
uint64_t biased_exp_minus_one = e - fmt.emin - !fmt.have_zero_and_subnormals;
550+
uint64_t enc = (biased_exp_minus_one << fmt.precision) + significand;
551+
552+
// Check that rounding didn't overflow into the non-finite encodings
553+
if (enc >= first_invalid_float_value(fmt))
554+
return double_nonfinite_to_bits(f64_inf, dtype);
555+
556+
return (uint64_t(sign) << (fmt.bitwidth - 1)) | enc;
557+
}
558+
559+
static double bitcast_u64_to_double(uint64_t bits) {
560+
union { uint64_t u64; double dbl; } u = {.u64 = bits};
561+
return u.dbl;
562+
}
563+
564+
static uint64_t bitcast_double_to_u64(double val) {
565+
union { double dbl; uint64_t u64; } u = {.dbl = val};
566+
return u.u64;
567+
}
568+
569+
Result<uint64_t> encode_float(double value, DType dtype) {
570+
uint64_t f64_bits = bitcast_double_to_u64(value);
571+
if (dtype == k_float64)
572+
return f64_bits;
573+
574+
const FloatFormat& fmt = float_dtype_format(dtype);
575+
576+
bool sign = f64_bits >> 63;
577+
if (sign && !fmt.have_sign)
578+
return raise_negative_unsigned(dtype);
579+
580+
int e = (f64_bits >> 52) & 0x7ff;
581+
if (e == 0x7ff)
582+
return double_nonfinite_to_bits(f64_bits, dtype);
583+
584+
// Calculate a 53-bit significand with the bit #52 set
585+
uint64_t significand = f64_bits & 0x000f'ffff'ffff'ffff;
586+
if (e) {
587+
// Normal float64: add the implicit bit
588+
significand |= 0x0010'0000'0000'0000;
589+
e -= 1023;
590+
} else {
591+
if (!significand) {
592+
// Zero or negative zero. Without a zero, round up to the smallest value.
593+
return fmt.have_zero_and_subnormals ? uint64_t(sign) << (fmt.bitwidth - 1) : 0;
594+
}
595+
// Subnormal float64: there is no implicit one, so we need to "normalize" the significand.
596+
int shift = 53 - bit_length(significand);
597+
significand <<= shift;
598+
e = -1022 - shift;
599+
}
600+
// Turn the significand into a 64-bit value with the top bit set,
601+
// so that it's compatible with encode_float_parts().
602+
significand <<= 11;
603+
return encode_float_parts(sign, e, significand, dtype);
604+
}
605+
606+
double decode_float(uint64_t bits, DType dtype) {
607+
uint64_t enc = bits;
608+
if (dtype == k_float64)
609+
return bitcast_u64_to_double(enc);
610+
611+
const FloatFormat& fmt = float_dtype_format(dtype);
612+
uint64_t exp_mask = (1ull << fmt.exp_bits) - 1;
613+
uint64_t biased_exp = (enc >> fmt.precision) & exp_mask;
614+
uint64_t mantissa_mask = (1ull << fmt.precision) - 1;
615+
uint64_t mantissa = enc & mantissa_mask;
616+
617+
uint64_t f64_bits = fmt.have_sign ? (enc << (64 - fmt.bitwidth)) & 0x8000'0000'0000'0000 : 0;
618+
619+
bool nonfinite = biased_exp == exp_mask;
620+
switch (fmt.non_finite_behavior) {
621+
case FloatFormat::IEEE:
622+
break;
623+
case FloatFormat::NanOnlyAllOnes:
624+
if (mantissa != mantissa_mask) nonfinite = false;
625+
break;
626+
case FloatFormat::FiniteOnly:
627+
nonfinite = false;
628+
break;
629+
}
630+
631+
uint64_t f64_exp;
632+
if (nonfinite) {
633+
f64_exp = 0x7ff;
634+
if (fmt.precision == 0)
635+
f64_bits |= 0x0008'0000'0000'0000;
636+
} else if (biased_exp == 0 && fmt.have_zero_and_subnormals) {
637+
if (mantissa != 0) {
638+
// Subnormal number becomes normal
639+
int shift = fmt.precision - bit_length(mantissa) + 1;
640+
mantissa = (mantissa << shift) & mantissa_mask;
641+
f64_exp = fmt.emin + 1023 - shift;
642+
} else {
643+
f64_exp = 0;
644+
}
645+
} else {
646+
f64_exp = (biased_exp - fmt.have_zero_and_subnormals) + fmt.emin + 1023;
647+
}
648+
649+
f64_bits |= (f64_exp << 52) | (mantissa << (52 - fmt.precision));
650+
return bitcast_u64_to_double(f64_bits);
651+
}
652+
653+
389654
// ----- DType Python wrapper object -----
390655

391656
struct DTypeObject {
@@ -654,6 +919,56 @@ static PyObject* py_promote_dtypes(PyObject* module_self,
654919
return Py_NewRef(dtype_to_pyobject(*res, lock));
655920
}
656921

922+
static Result<DType> parse_float_dtype(PyObject* dtype_obj) {
923+
if (!PyObject_TypeCheck(dtype_obj, &DTypeObject::pytype))
924+
return raise(PyExc_TypeError, "Expected a DType object, got ", use_repr(dtype_obj));
925+
DType dtype = py_unwrap<DTypeObject>(dtype_obj).dtype;
926+
if (!is_float_dtype(dtype))
927+
return raise(PyExc_TypeError, "Expected a floating-point dtype, got ", dtype);
928+
return dtype;
929+
}
930+
931+
static PyObject* py_encode_float(PyObject* module_self,
932+
PyObject* const* args, Py_ssize_t nargs) {
933+
if (nargs != 2) {
934+
raise(PyExc_TypeError, "Expected 2 positional arguments");
935+
return nullptr;
936+
}
937+
PyObject *py_value = args[0], *py_dtype = args[1];
938+
double value = PyFloat_AsDouble(py_value);
939+
if (PyErr_Occurred()) return nullptr;
940+
Result<DType> dtype = parse_float_dtype(py_dtype);
941+
if (!dtype.is_ok()) return nullptr;
942+
Result<uint64_t> bits = encode_float(value, *dtype);
943+
if (!bits.is_ok()) return nullptr;
944+
return PyLong_FromUnsignedLongLong(*bits);
945+
}
946+
947+
static PyObject* py_decode_float(PyObject* module_self,
948+
PyObject* const* args, Py_ssize_t nargs) {
949+
if (nargs != 2) {
950+
raise(PyExc_TypeError, "Expected 2 positional arguments");
951+
return nullptr;
952+
}
953+
PyObject *py_bits = args[0], *py_dtype = args[1];
954+
if (!PyLong_Check(py_bits)) {
955+
raise(PyExc_TypeError, "Expected an int, got ", use_repr(py_bits));
956+
return nullptr;
957+
}
958+
Result<DType> dtype = parse_float_dtype(py_dtype);
959+
if (!dtype.is_ok()) return nullptr;
960+
int overflow;
961+
uint64_t bits = pylong_as_overflow_and<uint64_t>(py_bits, &overflow);
962+
if (PyErr_Occurred()) return nullptr;
963+
uint32_t bitwidth = float_dtype_format(*dtype).bitwidth;
964+
if (overflow || (bitwidth < 64 && (bits >> bitwidth))) {
965+
raise(PyExc_ValueError, "Bits ", use_repr(py_bits), " do not fit in ", bitwidth,
966+
" bits of ", *dtype);
967+
return nullptr;
968+
}
969+
return PyFloat_FromDouble(decode_float(bits, *dtype));
970+
}
971+
657972
static Result<DType> parse_pointer_dtype(PyObject* py_pointer, GlobalLock& lock) {
658973
if (!PyObject_TypeCheck(py_pointer, &DTypeObject::pytype))
659974
return raise(PyExc_TypeError, "Expected a DType object for the first argument, got ",
@@ -726,6 +1041,8 @@ static PyMethodDef functions[] = {
7261041
{"integer_dtype_max", py_integer_dtype_max, METH_O, nullptr},
7271042
{"_promote_dtypes", reinterpret_cast<PyCFunction>(py_promote_dtypes),
7281043
METH_FASTCALL, nullptr},
1044+
{"encode_float", reinterpret_cast<PyCFunction>(py_encode_float), METH_FASTCALL, nullptr},
1045+
{"decode_float", reinterpret_cast<PyCFunction>(py_decode_float), METH_FASTCALL, nullptr},
7291046
{"_pointer_pointee_dtype", py_pointer_pointee_dtype, METH_O, nullptr},
7301047
{"_pointer_memory_space", py_pointer_memory_space, METH_O, nullptr},
7311048
{"_foreign_pointer_pointee_dtype", py_foreign_pointer_pointee_dtype, METH_O, nullptr},

0 commit comments

Comments
 (0)