Skip to content

Commit f995e5c

Browse files
authored
Use Float16Array for to/fromHalf utilities if supported (#161)
1 parent 2d8b903 commit f995e5c

1 file changed

Lines changed: 22 additions & 2 deletions

File tree

src/utils.ts

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,11 @@ import { unindent } from "./dyno/base.js";
1515

1616
const f32buffer = new Float32Array(1);
1717
const u32buffer = new Uint32Array(f32buffer.buffer);
18+
const supportsFloat16Array = "Float16Array" in globalThis;
19+
const f16buffer = supportsFloat16Array
20+
? new globalThis["Float16Array" as keyof typeof globalThis](1)
21+
: null;
22+
const u16buffer = new Uint16Array(f16buffer?.buffer);
1823

1924
// Reinterpret the bits of a float32 as a uint32
2025
export function floatBitsToUint(f: number): number {
@@ -28,8 +33,17 @@ export function uintBitsToFloat(u: number): number {
2833
return f32buffer[0];
2934
}
3035

36+
export const toHalf = supportsFloat16Array ? toHalfNative : toHalfJS;
37+
export const fromHalf = supportsFloat16Array ? fromHalfNative : fromHalfJS;
38+
39+
// Encode a number as a float16, stored as a uint16 number.
40+
function toHalfNative(f: number): number {
41+
f16buffer[0] = f;
42+
return u16buffer[0];
43+
}
44+
3145
// Encode a number as a float16, stored as a uint16 number.
32-
export function toHalf(f: number): number {
46+
function toHalfJS(f: number): number {
3347
// Store the value into the shared Float32 array.
3448
f32buffer[0] = f;
3549
const bits = u32buffer[0];
@@ -75,7 +89,13 @@ export function toHalf(f: number): number {
7589
}
7690

7791
// Convert a float16 stored as a uint16 number back to a float32.
78-
export function fromHalf(h: number): number {
92+
function fromHalfNative(u: number): number {
93+
u16buffer[0] = u;
94+
return f16buffer[0];
95+
}
96+
97+
// Convert a float16 stored as a uint16 number back to a float32.
98+
function fromHalfJS(h: number): number {
7999
// Extract the sign (1 bit), exponent (5 bits), and fraction (10 bits)
80100
const sign = (h >> 15) & 0x1;
81101
const exp = (h >> 10) & 0x1f;

0 commit comments

Comments
 (0)