diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 7f2df3f..b28d6f8 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -39,6 +39,7 @@ inductive Dtype where | float8_e4m3 | float8_e3m4 | float8_e5m2 +| float8_e2m5 -- spec followed is here: https://gfloat.readthedocs.io/en/latest/formats.html (P3109_8p6) | float16 | bfloat16 | float32 @@ -65,6 +66,7 @@ def gen : Gen Dtype := Gen.elements [ float8_e4m3, float8_e3m4, float8_e5m2, + float8_e2m5, float16, bfloat16, float32, @@ -90,6 +92,7 @@ instance : ToString Dtype where | float8_e4m3 => "float8_e4m3fn" | float8_e3m4 => "float8_e3m4" | float8_e5m2 => "float8_e5m2" -- no fn since e5m2 has infinity + | float8_e2m5 => "float8_e2m5" -- no fn since e2m5 has infinity | float16 => "float16" | bfloat16 => "bfloat16" | float32 => "float32" @@ -97,7 +100,7 @@ instance : ToString Dtype where def isOneByte (x : Dtype) : Bool := match x with -| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => true +| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 | float8_e2m5 => true | _ => false def isMultiByte (x : Dtype) : Bool := ! x.isOneByte @@ -134,7 +137,7 @@ def intMax (x : Dtype) : Int := match x with -- Added float16 and bfloat16 so bitwise op know to reject it def isFloat (x : Dtype) : Bool := match x with -| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => true +| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 | .float8_e2m5 => true | _ => false --! Number of bytes used by each element of the given dtype @@ -142,7 +145,7 @@ def itemsize (x : Dtype) : Nat := match x with | float64 | int64 | uint64 => 8 | float32 | int32 | uint32 => 4 | bfloat16 | float16 | int16 | uint16 => 2 -| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => 1 +| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 | float8_e2m5 => 1 -- Previously this was inline in join with a recursive swap, -- but adding more fp8 types made the match too large. Lean needs to prove @@ -195,6 +198,19 @@ private def joinOrdered (x y : Dtype) : Option Dtype := -- diverges from numpy | .float8_e3m4, .float8_e4m3 => none | .float8_e3m4, _ => none + -- fp8_e2m5 has inf (like e5m2), bias = 1, max = 3.875, and 5 mantissa bits + -- promotes with bool/int8/uint8 to fp8_e2m5, with fp32 to fp32, with fp64 to fp64 + | .float8_e2m5, .bool + | .float8_e2m5, .int8 + | .float8_e2m5, .uint8 => float8_e2m5 + | .float8_e2m5, .float32 => float32 + | .float8_e2m5, .float64 => float64 + -- e2m5 vs other fp8 variants + | .float8_e2m5, .float8_e4m3 + | .float8_e2m5, .float8_e3m4 + | .float8_e2m5, .float8_e5m2 => none + -- e2m5 with fp16/bf16/int16 gives none + | .float8_e2m5, _ => none | .float32, .float64 => float64 | .float32, _ | _, .float32 => none @@ -316,6 +332,14 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with | .float8_e3m4, .float32 | .float8_e3m4, .float64 => true | .float8_e3m4, _ => false +-- A cast is lossless if every representable value in the source format can be represented exactly in the target. 2 conditions need to hold: +-- The targets max >= source max (no overflow) +-- The targets mantissa bits must be >= source's mantissa bits (no rounding) +| .float8_e2m5 , .float16 +| .float8_e2m5 , .bfloat16 +| .float8_e2m5, .float32 +| .float8_e2m5, .float64 => true +| .float8_e2m5, _ => false | .float32, .float32 | .float32, .float64 => true | .float32, _ => false @@ -362,6 +386,7 @@ private def maxSafeNat : Dtype -> Option Nat | .float8_e4m3 => maxSafeNatForFloat8e4m3 | .float8_e3m4 => maxSafeNatForFloat8e3m4 | .float8_e5m2 => maxSafeNatForFloat8e5m2 +| .float8_e2m5 => maxSafeNatForFloat8e2m5 | .float16 => maxSafeNatForFloat16 | .bfloat16 => maxSafeNatForBFloat16 | .float32 => maxSafeNatForFloat32 @@ -384,6 +409,7 @@ private def minSafeInt : Dtype -> Option Int | .float8_e4m3 => some (-maxSafeNatForFloat8e4m3) | .float8_e3m4 => some (-maxSafeNatForFloat8e3m4) | .float8_e5m2 => some (-maxSafeNatForFloat8e5m2) +| .float8_e2m5 => some (-maxSafeNatForFloat8e2m5) | .float16 => some (-maxSafeNatForFloat16) | .bfloat16 => some (-maxSafeNatForBFloat16) | .float32 => some (-maxSafeNatForFloat32) @@ -424,11 +450,22 @@ def decodeFloat8E3M4 (arr : ByteArray) : Err Float32 := private def encodeFloat8E3M4 (f : Float32) : ByteArray := ByteArray.mk #[f.toFloat8E3M4Bits] +-- Decode 1-byte fp8_e2m5 to Float32. +-- Centralizes the size check so callers don't need inline guards. +def decodeFloat8E2M5 (arr : ByteArray) : Err Float32 := + if arr.size != 1 then .error "decoder: expected 1 byte for float8_e2m5" + else .ok (arr.data[0]!.toFloat32FromFloat8E2M5) + +-- Encode Float32 to 1-byte fp8_e2m5. +private def encodeFloat8E2M5 (f : Float32) : ByteArray := + ByteArray.mk #[f.toFloat8E2M5Bits] + -- Dispatch fp8 decode by dtype private def decodeFloat8 (dtype : Dtype) (arr : ByteArray) : Err Float32 := match dtype with | .float8_e4m3 => decodeFloat8E4M3 arr | .float8_e5m2 => decodeFloat8E5M2 arr | .float8_e3m4 => decodeFloat8E3M4 arr + | .float8_e2m5 => decodeFloat8E2M5 arr | _ => .error "decoder: expected float8 type" -- Dispatch fp8 encode by dtype @@ -436,6 +473,7 @@ private def encodeFloat8 (dtype : Dtype) (f : Float32) : Err ByteArray := match | .float8_e4m3 => .ok (encodeFloat8E4M3 f) | .float8_e5m2 => .ok (encodeFloat8E5M2 f) | .float8_e3m4 => .ok (encodeFloat8E3M4 f) + | .float8_e2m5 => .ok (encodeFloat8E2M5 f) | _ => .error "encoder: expected float8 type" def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with @@ -451,6 +489,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float8_e3m4 => encodeFloat8E3M4 n.toFloat32 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 +| .float8_e2m5 => encodeFloat8E2M5 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -562,6 +601,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 | .float8_e3m4 => encodeFloat8E3M4 n.toFloat32 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 +| .float8_e2m5 => encodeFloat8E2M5 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits | .float32 => toLEByteArray n.toFloat32 @@ -719,6 +759,10 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y return encodeFloat8E5M2 (x + y) + | .float8_e2m5 => do + let x <- decodeFloat8E2M5 x + let y <- decodeFloat8E2M5 y + return encodeFloat8E2M5 (x + y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -754,6 +798,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y return encodeFloat8E5M2 (x - y) + | .float8_e2m5 => do + let x <- decodeFloat8E2M5 x + let y <- decodeFloat8E2M5 y + return encodeFloat8E2M5 (x - y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -790,6 +838,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y return encodeFloat8E5M2 (x * y) + | .float8_e2m5 => do + let x <- decodeFloat8E2M5 x + let y <- decodeFloat8E2M5 y + return encodeFloat8E2M5 (x * y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -826,6 +878,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y return encodeFloat8E5M2 (x / y) + | .float8_e2m5 => do + let x <- decodeFloat8E2M5 x + let y <- decodeFloat8E2M5 y + return encodeFloat8E2M5 (x / y) | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -861,6 +917,9 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .float8_e5m2 => do let f <- decodeFloat8E5M2 x return encodeFloat8E5M2 f.abs + | .float8_e2m5 => do + let f <- decodeFloat8E2M5 x + return encodeFloat8E2M5 f.abs | .float16 | .bfloat16 => do let x <- dtype.decodeFloat16OrBFloat16 x @@ -900,6 +959,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | float8_e5m2 => do let f <- decodeFloat8E5M2 x return f == 0 +| float8_e2m5 => do + let f <- decodeFloat8E2M5 x + return f == 0 | float16 | bfloat16 => do let f <- dtype.decodeFloat16OrBFloat16 x @@ -978,55 +1040,59 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err | .float16, .bfloat16 | .bfloat16, .float16 => do let f <- decodeFloat16OrBFloat16 fromDtype data encodeFloat16OrBFloat16 toDtype f - -- fp8 to unsigned integers | .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64 | .float8_e5m2, .uint8 | .float8_e5m2, .uint16 | .float8_e5m2, .uint32 | .float8_e5m2, .uint64 - | .float8_e3m4, .uint8 | .float8_e3m4, .uint16 | .float8_e3m4, .uint32 | .float8_e3m4, .uint64 => do + | .float8_e3m4, .uint8 | .float8_e3m4, .uint16 | .float8_e3m4, .uint32 | .float8_e3m4, .uint64 + | .float8_e2m5, .uint8 | .float8_e2m5, .uint16 | .float8_e2m5, .uint32 | .float8_e2m5, .uint64 => do let f <- decodeFloat8 fromDtype data return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) -- fp8 to signed integers | .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64 | .float8_e5m2, .int8 | .float8_e5m2, .int16 | .float8_e5m2, .int32 | .float8_e5m2, .int64 - | .float8_e3m4, .int8 | .float8_e3m4, .int16 | .float8_e3m4, .int32 | .float8_e3m4, .int64 => do + | .float8_e3m4, .int8 | .float8_e3m4, .int16 | .float8_e3m4, .int32 | .float8_e3m4, .int64 + | .float8_e2m5, .int8 | .float8_e2m5, .int16 | .float8_e2m5, .int32 | .float8_e2m5, .int64 => do let f <- decodeFloat8 fromDtype data return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) -- fp8 to float32 - | .float8_e4m3, .float32 | .float8_e5m2, .float32 | .float8_e3m4, .float32 => do + | .float8_e4m3, .float32 | .float8_e5m2, .float32 | .float8_e3m4, .float32 | .float8_e2m5, .float32 => do let f <- decodeFloat8 fromDtype data return toLEByteArray f -- fp8 to float64 - | .float8_e4m3, .float64 | .float8_e5m2, .float64 | .float8_e3m4, .float64 => do + | .float8_e4m3, .float64 | .float8_e5m2, .float64 | .float8_e3m4, .float64 | .float8_e2m5, .float64 => do let f <- decodeFloat8 fromDtype data return toLEByteArray f.toFloat -- fp8 to fp16/bf16 | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 | .float8_e5m2, .float16 | .float8_e5m2, .bfloat16 - | .float8_e3m4, .float16 | .float8_e3m4, .bfloat16 => do + | .float8_e3m4, .float16 | .float8_e3m4, .bfloat16 + | .float8_e2m5, .float16 | .float8_e2m5, .bfloat16 => do let f <- decodeFloat8 fromDtype data encodeFloat16OrBFloat16 toDtype f -- fp8 to fp8 (cross-format) - | .float8_e4m3, .float8_e5m2 | .float8_e4m3, .float8_e3m4 - | .float8_e5m2, .float8_e4m3 | .float8_e5m2, .float8_e3m4 - | .float8_e3m4, .float8_e4m3 | .float8_e3m4, .float8_e5m2 => do + | .float8_e4m3, .float8_e5m2 | .float8_e4m3, .float8_e3m4 | .float8_e4m3, .float8_e2m5 + | .float8_e5m2, .float8_e4m3 | .float8_e5m2, .float8_e3m4 | .float8_e5m2, .float8_e2m5 + | .float8_e3m4, .float8_e4m3 | .float8_e3m4, .float8_e5m2 | .float8_e3m4, .float8_e2m5 + | .float8_e2m5, .float8_e4m3 | .float8_e2m5, .float8_e3m4 | .float8_e2m5, .float8_e5m2 => do let f <- decodeFloat8 fromDtype data encodeFloat8 toDtype f -- float32 -> fp8 - | .float32, .float8_e4m3 | .float32, .float8_e5m2 | .float32, .float8_e3m4 => do + | .float32, .float8_e4m3 | .float32, .float8_e5m2 | .float32, .float8_e3m4 | .float32, .float8_e2m5 => do let f <- Float32.ofLEByteArray data encodeFloat8 toDtype f -- float64 -> fp8 (rounds twice via fp32, can disagree with ml_dtypes at interior values) - | .float64, .float8_e4m3 | .float64, .float8_e5m2 | .float64, .float8_e3m4 => do + | .float64, .float8_e4m3 | .float64, .float8_e5m2 | .float64, .float8_e3m4 | .float64, .float8_e2m5 => do let f <- Float.ofLEByteArray data encodeFloat8 toDtype f.toFloat32 -- fp16/bf16 -> fp8 | .float16, .float8_e4m3 | .bfloat16, .float8_e4m3 | .float16, .float8_e5m2 | .bfloat16, .float8_e5m2 - | .float16, .float8_e3m4 | .bfloat16, .float8_e3m4 => do + | .float16, .float8_e3m4 | .bfloat16, .float8_e3m4 + | .float16, .float8_e2m5 | .bfloat16, .float8_e2m5 => do let f <- decodeFloat16OrBFloat16 fromDtype data encodeFloat8 toDtype f - | .float8_e3m4, .float8_e3m4 | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 + | .float8_e2m5, .float8_e2m5 | .float8_e3m4, .float8_e3m4 | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible @@ -1096,6 +1162,10 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo (dtype : Dtype) (data : ByteArray) : Err ByteArray := do if data.size != dtype.itemsize then throw "incorrect byte count" else match dtype with + | .float8_e2m5 => do + let f <- decodeFloat8E2M5 data + let x <- f32 f + return encodeFloat8E2M5 x | .float8_e5m2 => do let f <- decodeFloat8E5M2 data let x <- f32 f @@ -1163,7 +1233,7 @@ def tanh : Dtype -> ByteArray -> Err ByteArray := def tanh! (dtype : Dtype) (data : ByteArray) : ByteArray := get! $ tanh dtype data private def shift (f : UInt64 -> UInt64 -> UInt64) (dtype : Dtype) (bits : ByteArray) (shiftAmount : ByteArray) : Err ByteArray := match dtype with -| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => throw "shifts not supported at float type" +| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 | .float8_e2m5 => throw "shifts not supported at float type" | .bool => throw "In NumPy, bool shifts are cast to int64. This seems arbitrary so please cast (e.g. with astype) before you shift." | .uint64 | .int64 | .uint32 | .int32 | .uint16 | .int16 | .uint8 | .int8 => let k := dtype.itemsize diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 1bc96d8..7c819a2 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -36,6 +36,7 @@ private def float16MantissaBits : Nat := 10 private def bfloat16MantissaBits : Nat := 7 private def float8e4m3MantissaBits : Nat := 3 private def float8e5m2MantissaBits : Nat := 2 +private def float8e2m5MantissaBits : Nat := 5 -- Add 1 to the mantissa length because of the implicit leading 1 def maxSafeNatForFloat32 : Nat := Nat.pow 2 (float32MantissaBits + 1) @@ -48,6 +49,9 @@ def maxSafeNatForFloat8e5m2 : Nat := Nat.pow 2 (float8e5m2MantissaBits + 1) -- is too small (max value is 15.5). The format overflows to inf before losing -- integer precision def maxSafeNatForFloat8e3m4 : Nat := 15 +-- The formula gives 64 for e2m5 but the exponent range is too small (max value = 3.875) +-- The format overflows to inf at 4.0 so the largest integer with a lossless round-trip is 3 +def maxSafeNatForFloat8e2m5 : Nat := 3 def _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF @@ -654,6 +658,116 @@ def _root_.Float32.toFloat8E3M4Bits (f : Float32) : UInt8 := -- Negative overflow #guard (Float32.ofBits 0xC1800000).toFloat8E3M4Bits == (240 : UInt8) -- -16.0 → -inf +-- Decode fp8_e2m5 (P3109_8p6) to Float32 +-- Format: sign-magnitude, 8 bits total. Positive codes 0-127, negative codes 128-255. +-- Byte 0 = +0, Byte 127 = +inf, Byte 128 = NaN, Byte 255 = -inf +-- No negative zero. Within each half: 2-bit exp + 5-bit mant, bias=1 +-- Source: https://gfloat.readthedocs.io/en/latest/formats.html (p3109_8p6) +-- Properties: max=3.875, min normal=0.5, min subnormal=0.015625, 2 infs, 1 NaN +def _root_.UInt8.toFloat32FromFloat8E2M5 (bits : UInt8) : Float32 := + let code := bits.toUInt32 + -- Special cases: +0, +inf, NaN, -inf + if code == 0 then Float32.ofBits 0x00000000 -- +0 + else if code == 127 then Float32.ofBits 0x7F800000 -- +inf + else if code == 128 then Float32.ofBits 0x7FC00000 -- NaN (the single NaN encoding) + else if code == 255 then Float32.ofBits 0xFF800000 -- -inf + else + -- Sign-magnitude: bit 7 is sign, bits 6..0 are magnitude (1-126 or 129-254) + let sign := code >>> 7 -- 0 or 1 + let magnitude := code &&& 0x7F -- 1..126 (special codes already handled) + let sign32 := sign <<< 31 + -- Decompose magnitude into exp (bits 6..5 of magnitude) and mant (bits 4..0) + let exp := (magnitude >>> 5) &&& 0x3 -- 2-bit exponent + let mant := magnitude &&& 0x1F -- 5-bit mantissa + if exp == 0 then + -- Subnormal: value = mant × 2^(-6) + -- Normalize: find leading bit position in mant (0-indexed) + let p := if mant >= 16 then 4 + else if mant >= 8 then 3 + else if mant >= 4 then 2 + else if mant >= 2 then 1 + else 0 + -- fp32 exponent: real exponent is (p - 6), biased = p + 121 + let fp32Exp := (p + 121).toUInt32 + -- fp32 mantissa: remove leading 1, shift into 23-bit field + let fp32Mant := (mant - (1 <<< p.toUInt32)) <<< (23 - p).toUInt32 + Float32.ofBits (sign32 ||| (fp32Exp <<< 23) ||| fp32Mant) + else + -- Normal: value = (1 + mant/32) × 2^(exp - 2) + -- fp32 exponent = (exp - 2) + 127 = exp + 125 + let fp32Exp := (exp + 125) + let fp32Mant := mant <<< 18 + Float32.ofBits (sign32 ||| (fp32Exp <<< 23) ||| fp32Mant) + +-- Encode fp32 to fp8_e2m5 (P3109_8p6) +-- Format: sign-magnitude. +inf=127, -inf=255, NaN=128, +0=0, no -0. +-- Uses round-to-nearest-even. Overflow maps to ±inf. +def _root_.Float32.toFloat8E2M5Bits (f : Float32) : UInt8 := + let bits := f.toBits + let sign := (bits >>> 31) &&& 1 + let exp := (bits >>> 23) &&& 0xFF + let mant := bits &&& 0x7FFFFF + let sign8 := sign.toUInt8 <<< 7 + if exp == 0xFF then + if mant == 0 then + -- ±inf → byte 127 (+inf) or 255 (-inf) + sign8 ||| 0x7F + else + -- NaN → byte 128 (the single NaN encoding) + 128 + else if exp == 0 then + -- fp32 zero or subnormal → too small for e2m5, flush to +0 (no -0 in this format) + 0 + else + -- Normal fp32. Rebias exponent from fp32 (127) to e2m5 (bias=1). + let realExp : Int := exp.toNat - 127 + let fullMant := mant ||| 0x800000 + if realExp > 1 then + -- Overflow → ±inf + sign8 ||| 0x7F + else if realExp >= -1 then + -- Normal e2m5 range: realExp in [-1, 0, 1], exp field = realExp + 2 (gives 1, 2, or 3) + let e2m5Exp := (realExp + 2).toNat + -- Truncate fp32 mantissa (23 bits) to 5 bits: shift right by 18 + let truncated := mant >>> 18 + let roundBit := (mant >>> 17) &&& 1 + let stickyBits := mant &&& 0x1FFFF + let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) + then truncated + 1 else truncated + -- If rounding overflows mantissa (> 0x1F = 31), bump exponent + let (finalExp, finalMant) := if rounded > 0x1F then + (e2m5Exp + 1, (0 : UInt32)) + else (e2m5Exp, rounded) + -- Construct magnitude (7 bits: 2-bit exp + 5-bit mant) + let magnitude := (finalExp.toUInt8 <<< 5) ||| finalMant.toUInt8 + -- magnitude=127 would be +inf, magnitude>=127 means overflow + if magnitude >= 127 then + sign8 ||| 0x7F -- ±inf + else + sign8 ||| magnitude + else + -- Subnormal in e2m5: realExp < -1 + -- Subnormal value = mant × 2^(-6) + -- Need: result_mant = fullMant >> (17 - realExp) + let totalShift := (17 - realExp).toNat + if totalShift >= 25 then + -- All bits shifted away -> flush to +0 + 0 + else + let shifted := fullMant >>> totalShift.toUInt32 + let roundBit := (fullMant >>> (totalShift.toUInt32 - 1)) &&& 1 + let stickyMask := (1 <<< (totalShift.toUInt32 - 1)) - 1 + let stickyBits := fullMant &&& stickyMask + let rounded := if roundBit == 1 && (stickyBits != 0 || shifted &&& 1 == 1) + then shifted + 1 else shifted + -- If rounded up to 32, becomes smallest normal (exp=1, mant=0) + if rounded >= 32 then + sign8 ||| (1 : UInt8) <<< 5 + else if rounded == 0 then + 0 -- flush to +0 + else + sign8 ||| rounded.toUInt8 + section Test #guard ( @@ -804,6 +918,41 @@ warning: declaration uses 'sorry' let f := bits.toFloat32FromFloat8E3M4 f.toFloat8E3M4Bits == bits || f != f +-- e2m5 decode tests (verified against P3109_8p6 spec from gfloat) +-- takes byte 0, passes it through the decoder, and checks that the result == fp32(+0 = 0x00000000) +#guard (0 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x00000000 -- +0 (S=0, exp=00, mant=00000) +-- Takes byte 128 (0b10000000 = sign=1, exp=00, mant=00000), decodes to NaN (the single NaN encoding in P3109_8p6) +#guard (128 : UInt8).toFloat32FromFloat8E2M5.toBits == 0x7FC00000 -- NaN (the single NaN at byte 128) +-- Takes byte 1 (0b00000001 = sign=0, exp=00, mant=00001), decodes to smallest subnormal = 0.015625 +#guard (1 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x3C800000 -- 0.015625 (min subnormal per P3109) +-- Takes byte 32 (0b00100000 = sign=0, exp=01, mant=00000), decodes to min normal = 0.5 +#guard (32 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x3F000000 -- 0.5 (min normal) +-- Takes byte 63 (0b00111111 = sign=0, exp=01, mant=11111), decodes to (1 + 31/32) × 2^(-1) = 0.984375 +#guard (63 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x3F7C0000 -- 0.984375 (normal value) +-- Takes byte 126 (0b01111110 = sign=0, exp=11, mant=11110), decodes to (1 + 30/32) × 2^1 = 3.875 +#guard (126 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x40780000 -- 3.875 (max normal) +-- Takes byte 96 (0b01100000 = exp=3, mant=0), decodes to 2.0 (not inf — byte 127 is +inf) +#guard (96 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x40000000 -- 2.0 (exp=3, mant=0 is normal, not inf) +-- Note: NaN has to be skipped here +-- Takes byte 64 (0b01000000 = sign=0, exp=10, mant=00000), decodes to (1+0) × 2^0 = 1.0 +#guard (64 : UInt8).toFloat32FromFloat8E2M5 == Float32.ofBits 0x3F800000 -- 1.0 +-- Encode boundary tests near max (verified against gfloat TiesToEven) +#guard (Float32.ofBits 0x407A0000).toFloat8E2M5Bits == (126 : UInt8) -- 3.90625 -> 3.875 (midpoint, ties to even) +#guard (Float32.ofBits 0x407C0000).toFloat8E2M5Bits == (127 : UInt8) -- 3.9375 -> +inf +-- Negative boundary +#guard (Float32.ofBits 0xC07A0000).toFloat8E2M5Bits == (254 : UInt8) -- -3.90625 -> -3.875 (midpoint, ties to even) +#guard (Float32.ofBits 0xC07C0000).toFloat8E2M5Bits == (255 : UInt8) -- -3.9375 -> -inf + +-- Exhaustive e2m5 round-trip: decode -> encode for all 256 byte values. +-- Byte 128 is the single NaN encoding — excluded because Lean's Float32 BEq +-- doesn't implement IEEE NaN semantics. +#guard (List.range 256).all fun i => + let bits := i.toUInt8 + if bits == 128 then true + else + let f := bits.toFloat32FromFloat8E2M5 + f.toFloat8E2M5Bits == bits + end Test end TensorLib diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 5c8f085..02a5d5a 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -114,7 +114,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with -- float8_e3m4 serializes as "V1" in ml_dtypes, same as e4m3. -- The npy format cannot distinguish between fp8 subtypes that use V1. -- Reading " "V1" +| .float8_e4m3 | .float8_e3m4 | .float8_e2m5 => "V1" | .float8_e5m2 => "f1" | .float16 => "f2" | .bfloat16 => "V2" @@ -433,8 +433,8 @@ end Save -- unreachable normally since toNpy blocks e3m4 before reaching -- save!, and parseFile maps V1 to e4m3. Guards against hand-constructed Ndarrays. def Ndarray.save! (arr : Ndarray) (file : System.FilePath) : IO Unit := - if arr.header.descr.name == .float8_e3m4 then - throw $ IO.userError "float8_e3m4 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" + if arr.header.descr.name == .float8_e3m4 || arr.header.descr.name == .float8_e2m5 then + throw $ IO.userError "float8_e3m4/float8_e2m5 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" else IO.FS.writeBinFile file arr.toByteArray! diff --git a/TensorLib/Tensor.lean b/TensorLib/Tensor.lean index 84e22f8..4af80fd 100644 --- a/TensorLib/Tensor.lean +++ b/TensorLib/Tensor.lean @@ -597,6 +597,7 @@ def toNatTree! (arr : Tensor) : Format.Tree Nat := get! $ toNatTree arr def toFloat32Tree (arr : Tensor) : Err (Format.Tree Float32) := do let t <- arr.toByteArrayTree match arr.dtype with + | .float8_e2m5 => t.mapM (fun b => Dtype.decodeFloat8E2M5 b) | .float8_e5m2 => t.mapM (fun b => Dtype.decodeFloat8E5M2 b) | .float8_e4m3 => t.mapM (fun b => Dtype.decodeFloat8E4M3 b) | .float8_e3m4 => t.mapM (fun b => Dtype.decodeFloat8E3M4 b) @@ -609,6 +610,7 @@ def toFloat32Tree! (arr : Tensor) : Format.Tree Float32 := get! $ toFloat32Tree def toFloat64Tree (arr : Tensor) : Err (Format.Tree Float) := do let t <- arr.toByteArrayTree match arr.dtype with + | .float8_e2m5 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E2M5 b; return f.toFloat) | .float8_e5m2 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E5M2 b; return f.toFloat) | .float8_e4m3 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E4M3 b; return f.toFloat) | .float8_e3m4 => t.mapM (fun b => do let f <- Dtype.decodeFloat8E3M4 b; return f.toFloat) @@ -681,7 +683,7 @@ def toNpy (arr : Tensor) : Err Npy.Ndarray := -- Our guard helps to surface an explicit error during write instead of allowing a -- silent round-trip corruption — without it, a user could save an e3m4 tensor, load it back, -- and get wrong values (interpreted as e4m3) with no indication anything went wrong. - if arr.dtype == .float8_e3m4 then .error "float8_e3m4 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" + if arr.dtype == .float8_e3m4 || arr.dtype == .float8_e2m5 then .error "float8_e3m4/float8_e2m5 cannot be saved to npy: format uses V1 which is indistinguishable from float8_e4m3" else let arr := if arr.isTriviallyReshapable then arr else arr.copy let descr := Npy.Dtype.mk arr.dtype Npy.ByteOrder.littleEndian @@ -771,6 +773,7 @@ open TensorLib.Tensor.Format.Tree #guard match (Tensor.zeros .float8_e3m4 (Shape.mk [2])).toNpy with | .error _ => true | .ok _ => false -- toNpy accepts e4m3 (not blocked) #guard match (Tensor.zeros .float8_e4m3 (Shape.mk [2])).toNpy with | .ok _ => true | .error _ => false +#guard match (Tensor.zeros .float8_e2m5 (Shape.mk [2])).toNpy with | .error _ => true | .ok _ => false end Test diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index bf449d0..2fb4990 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -712,6 +712,70 @@ private def testFloat8E3M4EdgeCases : IO Bool := do return checks.all id +-- E2M5: 1 sign + 2 exponent + 5 mantissa, bias=1, P3109 (has inf and NaN) +-- Expected values verified against gfloat (p3109_8p6) manually using python3.12 +-- Unlike ml_dtypes (used for e4m3/e3m4/e5m2), gfloat doesn't integrate with numpy arrays +-- or produce .npy files, so we can't use saveNumpyArray as an oracle at test time. +-- Instead, values were verified once manually and hardcoded here. +private def testFloat8E2M5EdgeCases : IO Bool := do + let mut checks : List Bool := [] + + -- Arithmetic: 1.0 (byte 64) + 0.5 (byte 32) = 1.5 (byte 80) + let a := toLEByteArray (64 : UInt8) -- e2m5 encoding of 1.0 + let b := toLEByteArray (32 : UInt8) -- e2m5 encoding of 0.5 + let negA := toLEByteArray (208 : UInt8) -- e2m5 encoding of -1.5 + + let pass <- checkBitsU8 "fp8_e2m5 add (1.0 + 0.5 = 1.5)" 80 (Dtype.add .float8_e2m5 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e2m5 sub (1.5 - 0.5 = 1.0)" 64 (Dtype.sub .float8_e2m5 (toLEByteArray (80 : UInt8)) b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e2m5 mul (1.5 * 0.5 = 0.75)" 48 (Dtype.mul .float8_e2m5 (toLEByteArray (80 : UInt8)) b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e2m5 div (1.5 / 0.5 = 3.0)" 112 (Dtype.div .float8_e2m5 (toLEByteArray (80 : UInt8)) b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e2m5 abs (-1.5) = 1.5" 80 (Dtype.abs .float8_e2m5 negA) + checks := pass :: checks + + -- Casting: e2m5(1.0) -> int8 = 1 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e2m5 a .int8) + let pass := castToI8 == toLEByteArray (1 : Int8) + IO.println s!"fp8_e2m5 cast to int8 (1.0 -> 1): {pass}" + checks := pass :: checks + + -- e2m5 NaN (byte 128) -> bool = true (NaN is not zero) + let nanByte := toLEByteArray (128 : UInt8) -- NaN in e2m5 + let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e2m5 nanByte .bool) + let pass := castToBool == ByteArray.mk #[1] + IO.println s!"fp8_e2m5 NaN to bool (true): {pass}" + checks := pass :: checks + + -- fp32(1.5) -> e2m5 = byte 80 + let f32_1_5 := toLEByteArray (1.5 : Float32) + let castToE2m5 <- IO.ofExcept (Dtype.castOverflow .float32 f32_1_5 .float8_e2m5) + let pass := castToE2m5 == toLEByteArray (80 : UInt8) + IO.println s!"fp8_e2m5 fp32 to e2m5 (1.5): {pass}" + checks := pass :: checks + + -- Overflow: 4.0 -> inf (byte 127) + let f32_4 := toLEByteArray (4.0 : Float32) + let castOverflow <- IO.ofExcept (Dtype.castOverflow .float32 f32_4 .float8_e2m5) + let pass := castOverflow == toLEByteArray (127 : UInt8) + IO.println s!"fp8_e2m5 overflow (4.0 -> inf): {pass}" + checks := pass :: checks + + -- +inf preserved + let f32_inf := toLEByteArray (Float32.ofBits 0x7F800000) + let castInf <- IO.ofExcept (Dtype.castOverflow .float32 f32_inf .float8_e2m5) + let pass := castInf == toLEByteArray (127 : UInt8) + IO.println s!"fp8_e2m5 +inf -> +inf: {pass}" + checks := pass :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && @@ -719,7 +783,8 @@ def runAllTests : IO Bool := do (<- testBFloat16EdgeCases) && (<- testFloat8E4M3EdgeCases) && (<- testFloat8E5M2EdgeCases) && - (<- testFloat8E3M4EdgeCases) + (<- testFloat8E3M4EdgeCases) && + (<- testFloat8E2M5EdgeCases) end Test end TensorLib