module Data.BFloat16

import Std.Natural

-- Binary32 -> bfloat16, round-to-nearest/ties-to-even, returned in the
-- high half of a binary32 word. Keeping this operation separate from
-- arithmetic makes mixed-precision rounding sites explicit in a program.
-- NaNs stay NaNs even when their entire payload was in the discarded half;
-- infinities and signed zero retain their signs. Inputs are binary32 words.
def bfloat16RoundWord =
  (lambda unrestricted word : Nat .
    (let unrestricted upper = (nat-divide word 65536) in
      (let unrestricted lower = (nat-modulo word 65536) in
        (naturalSelect
          (naturalAnd (naturalEqual (nat-modulo (nat-divide word 8388608) 256) 255)
            (naturalIsZero (naturalIsZero (nat-modulo word 8388608))))
          (nat-multiply (nat-add (nat-multiply (nat-divide upper 128) 128) 64) 65536)
          (nat-multiply
            (nat-add upper
              (naturalOr (nat-less-than 32768 lower)
                (naturalAnd (naturalEqual lower 32768) (nat-modulo upper 2))))
            65536)))))

-- Widen the low sixteen wire bits without rounding. BF16 has the binary32
-- sign and exponent positions with seven explicit fraction bits; appending
-- sixteen zero fraction bits preserves signed zero, subnormals and payloads.
def bfloat16WidenWord =
  (lambda unrestricted half : Nat .
    (naturalMultiply (naturalModuloUnchecked half (naturalPowerOfTwo 16))
      (naturalPowerOfTwo 16)))