Data.BFloat16

1module Data.BFloat16
2
3import Std.Natural
4
5-- Binary32 -> bfloat16, round-to-nearest/ties-to-even, returned in the
6-- high half of a binary32 word. Keeping this operation separate from
7-- arithmetic makes mixed-precision rounding sites explicit in a program.
8-- NaNs stay NaNs even when their entire payload was in the discarded half;
9-- infinities and signed zero retain their signs. Inputs are binary32 words.
10def bfloat16RoundWord =
11  (lambda unrestricted word : Nat .
12    (let unrestricted upper = (nat-divide word 65536) in
13      (let unrestricted lower = (nat-modulo word 65536) in
14        (naturalSelect
15          (naturalAnd (naturalEqual (nat-modulo (nat-divide word 8388608) 256) 255)
16            (naturalIsZero (naturalIsZero (nat-modulo word 8388608))))
17          (nat-multiply (nat-add (nat-multiply (nat-divide upper 128) 128) 64) 65536)
18          (nat-multiply
19            (nat-add upper
20              (naturalOr (nat-less-than 32768 lower)
21                (naturalAnd (naturalEqual lower 32768) (nat-modulo upper 2))))
22            65536)))))
23
24-- Widen the low sixteen wire bits without rounding. BF16 has the binary32
25-- sign and exponent positions with seven explicit fraction bits; appending
26-- sixteen zero fraction bits preserves signed zero, subnormals and payloads.
27def bfloat16WidenWord =
28  (lambda unrestricted half : Nat .
29    (naturalMultiply (naturalModuloUnchecked half (naturalPowerOfTwo 16))
30      (naturalPowerOfTwo 16)))