diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index 2da8707..7f2df3f 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -37,6 +37,7 @@ inductive Dtype where | uint32 | uint64 | float8_e4m3 +| float8_e3m4 | float8_e5m2 | float16 | bfloat16 @@ -62,6 +63,7 @@ def gen : Gen Dtype := Gen.elements [ uint32, uint64, float8_e4m3, + float8_e3m4, float8_e5m2, float16, bfloat16, @@ -86,6 +88,7 @@ instance : ToString Dtype where | uint32 => "uint32" | uint64 => "uint64" | float8_e4m3 => "float8_e4m3fn" + | float8_e3m4 => "float8_e3m4" | float8_e5m2 => "float8_e5m2" -- no fn since e5m2 has infinity | float16 => "float16" | bfloat16 => "bfloat16" @@ -94,7 +97,7 @@ instance : ToString Dtype where def isOneByte (x : Dtype) : Bool := match x with -| bool | int8 | uint8 | float8_e4m3 | float8_e5m2 => true +| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => true | _ => false def isMultiByte (x : Dtype) : Bool := ! x.isOneByte @@ -131,7 +134,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_e5m2 => true +| .float16 | .bfloat16 | .float32 | .float64 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => true | _ => false --! Number of bytes used by each element of the given dtype @@ -139,7 +142,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_e5m2 => 1 +| bool | int8 | uint8 | float8_e4m3 | float8_e3m4 | float8_e5m2 => 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 @@ -166,6 +169,8 @@ private def joinOrdered (x y : Dtype) : Option Dtype := | .float8_e4m3, .bool | .float8_e4m3, .int8 | .float8_e4m3, .uint8 => float8_e4m3 + -- e4m3 has more range, e3m4 has more mantissa but neither dominates so no safe common type exists. This diverges from numpy + | .float8_e4m3, .float8_e3m4 => none | .float8_e4m3, _ => none | .float8_e5m2, .float32 => float32 | .float8_e5m2, .float64 => float64 @@ -180,6 +185,16 @@ private def joinOrdered (x y : Dtype) : Option Dtype := | .float8_e5m2, .int8 => float8_e5m2 | .float8_e5m2, .uint8 => float8_e5m2 | .float8_e5m2, _ => none + -- e3m4 follows e4m3's promotion rules + -- numpy: result_type(e3m4, int16) = none, matching e4m3 behavior. + | .float8_e3m4, .float32 => float32 + | .float8_e3m4, .float64 => float64 + | .float8_e3m4, .bool + | .float8_e3m4, .int8 + | .float8_e3m4, .uint8 => float8_e3m4 + -- diverges from numpy + | .float8_e3m4, .float8_e4m3 => none + | .float8_e3m4, _ => none | .float32, .float64 => float64 | .float32, _ | _, .float32 => none @@ -294,6 +309,13 @@ def lossless (fromDtype toDtype : Dtype) : Bool := match fromDtype, toDtype with | .float8_e5m2, .float32 | .float8_e5m2, .float64 => true | .float8_e5m2, _ => false +-- Note: np.can_cast(e3m4, e4m3) returns true but this is incorrect since precision is lost +| .float8_e3m4, .float8_e3m4 +| .float8_e3m4, .float16 +| .float8_e3m4, .bfloat16 +| .float8_e3m4, .float32 +| .float8_e3m4, .float64 => true +| .float8_e3m4, _ => false | .float32, .float32 | .float32, .float64 => true | .float32, _ => false @@ -338,6 +360,7 @@ private def maxSafeNat : Dtype -> Option Nat | .uint64 => some 0xFFFFFFFFFFFFFFFF | .int64 => some 0x7FFFFFFFFFFFFFFF | .float8_e4m3 => maxSafeNatForFloat8e4m3 +| .float8_e3m4 => maxSafeNatForFloat8e3m4 | .float8_e5m2 => maxSafeNatForFloat8e5m2 | .float16 => maxSafeNatForFloat16 | .bfloat16 => maxSafeNatForBFloat16 @@ -359,6 +382,7 @@ private def minSafeInt : Dtype -> Option Int | .int32 => some (-0x80000000) | .int64 => some (-0x8000000000000000) | .float8_e4m3 => some (-maxSafeNatForFloat8e4m3) +| .float8_e3m4 => some (-maxSafeNatForFloat8e3m4) | .float8_e5m2 => some (-maxSafeNatForFloat8e5m2) | .float16 => some (-maxSafeNatForFloat16) | .bfloat16 => some (-maxSafeNatForBFloat16) @@ -390,6 +414,30 @@ def decodeFloat8E5M2 (arr : ByteArray) : Err Float32 := private def encodeFloat8E5M2 (f : Float32) : ByteArray := ByteArray.mk #[f.toFloat8E5M2Bits] +-- Decode 1-byte fp8_e3m4 to Fp32. +-- Centralizes the size check so callers don't need inline guards. +def decodeFloat8E3M4 (arr : ByteArray) : Err Float32 := + if arr.size != 1 then .error "decoder: expected 1 byte for float8_e3m4" + else .ok (arr.data[0]!.toFloat32FromFloat8E3M4) + +-- Encode fp32 to 1 byte fp8_e3m4 +private def encodeFloat8E3M4 (f : Float32) : ByteArray := + ByteArray.mk #[f.toFloat8E3M4Bits] + +-- 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 + | _ => .error "decoder: expected float8 type" + +-- Dispatch fp8 encode by dtype +private def encodeFloat8 (dtype : Dtype) (f : Float32) : Err ByteArray := match dtype with + | .float8_e4m3 => .ok (encodeFloat8E4M3 f) + | .float8_e5m2 => .ok (encodeFloat8E5M2 f) + | .float8_e3m4 => .ok (encodeFloat8E3M4 f) + | _ => .error "encoder: expected float8 type" + def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype with | .bool => toLEByteArray (if n == 0 then 0 else 1).toUInt8 | .uint8 => toLEByteArray n.toUInt8 @@ -401,6 +449,7 @@ def byteArrayOfNatOverflow (dtype : Dtype) (n : Nat) : ByteArray := match dtype | .uint64 => toLEByteArray n.toUInt64 | .int64 => toLEByteArray n.toInt64 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 +| .float8_e3m4 => encodeFloat8E3M4 n.toFloat32 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits @@ -511,6 +560,7 @@ private def byteArrayOfIntOverflow (dtype : Dtype) (n : Int) : ByteArray := matc | .uint32 | .int32 => toLEByteArray n.toInt32 | .uint64 | .int64 => toLEByteArray n.toInt64 | .float8_e4m3 => encodeFloat8E4M3 n.toFloat32 +| .float8_e3m4 => encodeFloat8E3M4 n.toFloat32 | .float8_e5m2 => encodeFloat8E5M2 n.toFloat32 | .float16 => toLEByteArray n.toFloat32.toFloat16Bits | .bfloat16 => toLEByteArray n.toFloat32.toBFloat16Bits @@ -639,8 +689,6 @@ private def byteArrayToBFloat16RoundTrip (dtype : Dtype) (f : Float32) : Bool := #guard bfloat16.byteArrayToBFloat16RoundTrip (-0) #guard bfloat16.byteArrayToBFloat16RoundTrip 256 - - -- Add produces the IEEE754 result because fp32 (p=24) satisfies the innocuous double rounding condition (theoreom 20) -- https://hal.science/hal-01091186v1/document /- @@ -663,6 +711,10 @@ def add (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x + y) + | .float8_e3m4 => do + let x <- decodeFloat8E3M4 x + let y <- decodeFloat8E3M4 y + return encodeFloat8E3M4 (x + y) | .float8_e5m2 => do let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y @@ -694,6 +746,10 @@ def sub (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x - y) + | .float8_e3m4 => do + let x <- decodeFloat8E3M4 x + let y <- decodeFloat8E3M4 y + return encodeFloat8E3M4 (x - y) | .float8_e5m2 => do let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y @@ -726,6 +782,10 @@ def mul (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x * y) + | .float8_e3m4 => do + let x <- decodeFloat8E3M4 x + let y <- decodeFloat8E3M4 y + return encodeFloat8E3M4 (x * y) | .float8_e5m2 => do let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y @@ -758,6 +818,10 @@ def div (dtype : Dtype) (x y : ByteArray) : Err ByteArray := let x <- decodeFloat8E4M3 x let y <- decodeFloat8E4M3 y return encodeFloat8E4M3 (x / y) + | .float8_e3m4 => do + let x <- decodeFloat8E3M4 x + let y <- decodeFloat8E3M4 y + return encodeFloat8E3M4 (x / y) | .float8_e5m2 => do let x <- decodeFloat8E5M2 x let y <- decodeFloat8E5M2 y @@ -791,6 +855,9 @@ def abs (dtype : Dtype) (x : ByteArray) : Err ByteArray := do | .float8_e4m3 => do let f <- decodeFloat8E4M3 x return encodeFloat8E4M3 f.abs + | .float8_e3m4 => do + let f <- decodeFloat8E3M4 x + return encodeFloat8E3M4 f.abs | .float8_e5m2 => do let f <- decodeFloat8E5M2 x return encodeFloat8E5M2 f.abs @@ -827,6 +894,9 @@ def isZero (dtype : Dtype) (x : ByteArray) : Err Bool := match dtype with | float8_e4m3 => do let f <- decodeFloat8E4M3 x return f == 0 +| float8_e3m4 => do + let f <- decodeFloat8E3M4 x + return f == 0 | float8_e5m2 => do let f <- decodeFloat8E5M2 x return f == 0 @@ -909,90 +979,55 @@ def castOverflow (fromDtype : Dtype) (data : ByteArray) (toDtype : Dtype) : Err let f <- decodeFloat16OrBFloat16 fromDtype data encodeFloat16OrBFloat16 toDtype f - -- float8_e4m3 to unsigned integers - | .float8_e4m3, .uint8 | .float8_e4m3, .uint16 | .float8_e4m3, .uint32 | .float8_e4m3, .uint64 => do - let f <- decodeFloat8E4M3 data - return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) - -- float8_e4m3 to signed integers - | .float8_e4m3, .int8 | .float8_e4m3, .int16 | .float8_e4m3, .int32 | .float8_e4m3, .int64 => do - let f <- decodeFloat8E4M3 data - return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) - -- float8_e4m3 to float32 - | .float8_e4m3, .float32 => do - let f <- decodeFloat8E4M3 data - return toLEByteArray f - -- float8_e4m3 to float64 - | .float8_e4m3, .float64 => do - let f <- decodeFloat8E4M3 data - return toLEByteArray f.toFloat - -- float8_e4m3 to fp16/bf16 - | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 => do - let f <- decodeFloat8E4M3 data - encodeFloat16OrBFloat16 toDtype f - -- float32 -> float8_e4m3 - | .float32, .float8_e4m3 => do - let f <- Float32.ofLEByteArray data - return encodeFloat8E4M3 f - -- float64 -> float8_e4m3: rounds twice via fp32. Can disagree with ml_dtypes at the overflow edge (eg: 464.00000000000006) - | .float64, .float8_e4m3 => do - let f <- Float.ofLEByteArray data - return encodeFloat8E4M3 f.toFloat32 - -- fp16/bf16 -> float8_e4m3 - | .float16, .float8_e4m3 | .bfloat16, .float8_e4m3 => do - let f <- decodeFloat16OrBFloat16 fromDtype data - return encodeFloat8E4M3 f - - -- float8_e5m2 to unsigned integers - | .float8_e5m2, .uint8 - | .float8_e5m2, .uint16 - | .float8_e5m2, .uint32 - | .float8_e5m2, .uint64 => do - let f <- decodeFloat8E5M2 data + -- 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 + let f <- decodeFloat8 fromDtype data return toDtype.byteArrayOfNatOverflow (saturatingNatOfFloat32 toDtype f) - -- float8_e5m2 to signed integers - | .float8_e5m2, .int8 - | .float8_e5m2, .int16 - | .float8_e5m2, .int32 - | .float8_e5m2, .int64 => do - let f <- decodeFloat8E5M2 data + -- 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 + let f <- decodeFloat8 fromDtype data return toDtype.byteArrayOfIntOverflow (saturatingIntOfFloat32 toDtype f) - -- float8_e5m2 to float32 - | .float8_e5m2, .float32 => do - let f <- decodeFloat8E5M2 data + -- fp8 to float32 + | .float8_e4m3, .float32 | .float8_e5m2, .float32 | .float8_e3m4, .float32 => do + let f <- decodeFloat8 fromDtype data return toLEByteArray f - -- float8_e5m2 to float64 - | .float8_e5m2, .float64 => do - let f <- decodeFloat8E5M2 data + -- fp8 to float64 + | .float8_e4m3, .float64 | .float8_e5m2, .float64 | .float8_e3m4, .float64 => do + let f <- decodeFloat8 fromDtype data return toLEByteArray f.toFloat - -- float8_e5m2 to fp16/bf16 - | .float8_e5m2, .float16 - | .float8_e5m2, .bfloat16 => do - let f <- decodeFloat8E5M2 data + -- fp8 to fp16/bf16 + | .float8_e4m3, .float16 | .float8_e4m3, .bfloat16 + | .float8_e5m2, .float16 | .float8_e5m2, .bfloat16 + | .float8_e3m4, .float16 | .float8_e3m4, .bfloat16 => do + let f <- decodeFloat8 fromDtype data encodeFloat16OrBFloat16 toDtype f - -- float8_e5m2 to float8_e4m3 - | .float8_e5m2, .float8_e4m3 => do - let f <- decodeFloat8E5M2 data - return encodeFloat8E4M3 f - -- float32 -> float8_e5m2 - | .float32, .float8_e5m2 => do + -- 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 + let f <- decodeFloat8 fromDtype data + encodeFloat8 toDtype f + -- float32 -> fp8 + | .float32, .float8_e4m3 | .float32, .float8_e5m2 | .float32, .float8_e3m4 => do let f <- Float32.ofLEByteArray data - return encodeFloat8E5M2 f - -- float64 -> float8_e5m2 - | .float64, .float8_e5m2 => do + 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 let f <- Float.ofLEByteArray data - return encodeFloat8E5M2 f.toFloat32 - -- fp16/bf16 -> float8_e5m2 - | .float16, .float8_e5m2 - | .bfloat16, .float8_e5m2 => do + 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 let f <- decodeFloat16OrBFloat16 fromDtype data - return encodeFloat8E5M2 f - -- float8_e4m3 -> float8_e5m2 - | .float8_e4m3, .float8_e5m2 => do - let f <- decodeFloat8E4M3 data - return encodeFloat8E5M2 f + encodeFloat8 toDtype f - - | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible + | .float8_e3m4, .float8_e3m4 | .float8_e5m2, .float8_e5m2 | .float8_e4m3, .float8_e4m3 + | .float16, .float16 | .bfloat16, .bfloat16 | .float32, .float32 | .float64, .float64 => impossible def isZero! (dtype : Dtype) (x : ByteArray) : Bool := get! $ dtype.isZero x @@ -1069,6 +1104,10 @@ private def liftFloatUnop (f32 : Float32 -> Err Float32) (f64 : Float -> Err Flo let f <- decodeFloat8E4M3 data let x <- f32 f return encodeFloat8E4M3 x + | .float8_e3m4 => do + let f <- decodeFloat8E3M4 data + let x <- f32 f + return encodeFloat8E3M4 x | .float16 | .bfloat16 => do let f <- decodeFloat16OrBFloat16 dtype data let x <- f32 f @@ -1124,7 +1163,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_e5m2 => throw "shifts not supported at float type" +| .float32 | .float64 | .bfloat16 | .float16 | .float8_e4m3 | .float8_e3m4 | .float8_e5m2 => 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 @@ -1421,6 +1460,18 @@ example (a b : UInt8) : let xb := toLEByteArray b Dtype.add .float8_e4m3 xa xb == Dtype.add .float8_e4m3 xb xa := by plausible +-- Property: e3m4 addition is commutative (a + b == b + a) +/-- +info: Unable to find a counter-example +--- +warning: declaration uses 'sorry' +-/ +#guard_msgs in +example (a b : UInt8) : + let xa := toLEByteArray a + let xb := toLEByteArray b + Dtype.add .float8_e3m4 xa xb == Dtype.add .float8_e3m4 xb xa := by plausible + -- PBT for join commutativity -- Since joinOrdered requires both arguments to be listed for same size types (the swap guard is triggered only when sizes differ) -- This PBT catches any missing direction that would silently return none instead of promoting diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index a1931e3..1bc96d8 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -44,6 +44,10 @@ def maxSafeNatForFloat16 : Nat := Nat.pow 2 (float16MantissaBits + 1) def maxSafeNatForBFloat16 : Nat := Nat.pow 2 (bfloat16MantissaBits + 1) def maxSafeNatForFloat8e4m3 : Nat := Nat.pow 2 (float8e4m3MantissaBits + 1) def maxSafeNatForFloat8e5m2 : Nat := Nat.pow 2 (float8e5m2MantissaBits + 1) +-- The formula doesn't apply here because e3m4's exponent range +-- is too small (max value is 15.5). The format overflows to inf before losing +-- integer precision +def maxSafeNatForFloat8e3m4 : Nat := 15 def _root_.Float32.minValue : Float32 := Float32.ofBits 0xFF7FFFFF def _root_.Float32.maxValue : Float32 := Float32.ofBits 0x7F7FFFFF @@ -519,6 +523,137 @@ def _root_.Float32.toFloat8E5M2Bits (f : Float32) : UInt8 := #guard (Float32.ofBits 0x7F800000).toFloat8E5M2Bits == (124 : UInt8) -- +inf #guard (Float32.ofBits 0xFF800000).toFloat8E5M2Bits == (252 : UInt8) -- -inf +-- Convert float8_e3m4 bits (stored as UInt8) to Float32. +-- E3M4: 1 sign + 3 exponent + 4 mantissa, bias = 3, IEEE-like (has inf and NaN). +-- Max value: +-15.5, smallest subnormal: 2^(-6), smallest normal: 2^(-2). +def _root_.UInt8.toFloat32FromFloat8E3M4 (bits : UInt8) : Float32 := + let sign := (bits >>> 7) &&& 1 -- bit 7: sign + let exp := (bits >>> 4) &&& 0x7 -- bits 6..4: exponent (3 bits) + let mant := bits &&& 0xF -- bits 3..0: mantissa (4 bits) + let sign32 := sign.toUInt32 <<< 31 + if exp == 0 then + if mant == 0 then + -- ±zero + Float32.ofBits sign32 + else + -- Subnormal: value = (-1)^sign × mant × 2^(1 - 3 - 4) = mant × 2^(-6) + let f := Float32.ofNat mant.toNat + let scale := Float32.ofBits 0x3C800000 -- 2^(-6) in fp32 + let result := f * scale + if sign == 1 then Float32.ofBits (result.toBits ||| 0x80000000) + else result + else if exp == 0x7 then + if mant == 0 then + -- ±infinity + Float32.ofBits (sign32 ||| 0x7F800000) + else + -- NaN + Float32.ofBits (sign32 ||| 0x7FC00000) + else + -- Normal: rebias exponent from e3m4 (bias=3) to fp32 (bias=127) + -- Shift mantissa from 4 bits to 23 bits (left-pad with 19 zeros) + let newExp := exp.toUInt32 - 3 + 127 + Float32.ofBits (sign32 ||| newExp <<< 23 ||| mant.toUInt32 <<< 19) + +-- Convert Float32 to float8_e3m4 bits (UInt8). +-- E3M4: 1 sign + 3 exponent + 4 mantissa, bias = 3. +-- Overflow maps to +-inf (IEEE-like). Round-to-nearest-even. +def _root_.Float32.toFloat8E3M4Bits (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 -> +-inf in e3m4 + sign8 ||| 0x70 + else + -- NaN -> quiet NaN in e3m4 (exp=7, mant=0b1000 - quiet bit is MSB of mant) + sign8 ||| 0x78 + else if exp == 0 then + -- fp32 zero or subnormal -> too small for e3m4, flush to zero + sign8 + else + -- Normal fp32. Rebias exponent from fp32 (127) to e3m4 (3). + let realExp : Int := exp.toNat - 127 + let fullMant := mant ||| 0x800000 + if realExp > 3 then + -- Overflow → ±inf + sign8 ||| 0x70 + else if realExp >= -2 then + -- Normal e3m4 range: realExp in [-2, 3], exp field 1..6 + let e3m4Exp := (realExp + 3).toNat + -- Truncate fp32 mantissa (23 bits) to 4 bits: shift right by 19 + let truncated := mant >>> 19 + let roundBit := (mant >>> 18) &&& 1 + let stickyBits := mant &&& 0x3FFFF + let rounded := if roundBit == 1 && (stickyBits != 0 || truncated &&& 1 == 1) + then truncated + 1 else truncated + -- If rounding overflows mantissa (0x10), bump exponent + let (finalExp, finalMant) := if rounded > 0xF then + (e3m4Exp + 1, (0 : UInt32)) + else (e3m4Exp, rounded) + -- If exponent overflows to 7 with mant=0, that's inf + if finalExp >= 7 then + sign8 ||| 0x70 + else + sign8 ||| (finalExp.toUInt8 <<< 4) ||| finalMant.toUInt8 + else + -- Subnormal in e3m4: realExp < -2 + -- Subnormal value = mant × 2^(-6), need: + -- mant = fullMant * 2^(realExp) / 2^(-6) / 2^23 = fullMant >> (17 - realExp) + let totalShift := (17 - realExp).toNat + if totalShift >= 25 then + sign8 + else + let shifted := fullMant >>> totalShift.toUInt32 + -- totalShift >= 20 in this branch (realExp < -2), so always > 1 + 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 16, becomes smallest normal (exp=1, mant=0) + if rounded >= 16 then + sign8 ||| (1 : UInt8) <<< 4 + else + sign8 ||| rounded.toUInt8 + +-- e3m4 decode tests (verified against ml_dtypes) +#guard (0 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x00000000 -- +0 +#guard (128 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x80000000 -- -0 +#guard (48 : UInt8).toFloat32FromFloat8E3M4 == 1.0 -- 1.0 +#guard (176 : UInt8).toFloat32FromFloat8E3M4 == -1.0 -- -1.0 +#guard (64 : UInt8).toFloat32FromFloat8E3M4 == 2.0 -- 2.0 +#guard (111 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x41780000 -- max (15.5) +#guard (239 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0xC1780000 -- -max (-15.5) +#guard (112 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x7F800000 -- +inf +#guard (240 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0xFF800000 -- -inf +#guard (1 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x3C800000 -- smallest subnormal +#guard (16 : UInt8).toFloat32FromFloat8E3M4 == Float32.ofBits 0x3E800000 -- smallest normal +#guard (56 : UInt8).toFloat32FromFloat8E3M4 == 1.5 -- 1.5 +-- NaN: f != f +#guard (113 : UInt8).toFloat32FromFloat8E3M4 != (113 : UInt8).toFloat32FromFloat8E3M4 -- NaN +#guard (127 : UInt8).toFloat32FromFloat8E3M4 != (127 : UInt8).toFloat32FromFloat8E3M4 -- NaN + +-- e3m4 encode tests (verified against ml_dtypes) +#guard (Float32.ofBits 0x00000000).toFloat8E3M4Bits == (0 : UInt8) -- +0 +#guard (Float32.ofBits 0x80000000).toFloat8E3M4Bits == (128 : UInt8) -- -0 +#guard (Float32.ofNat 1).toFloat8E3M4Bits == (48 : UInt8) -- 1.0 +#guard (Float32.ofNat 2).toFloat8E3M4Bits == (64 : UInt8) -- 2.0 +#guard (Float32.ofNat 3).toFloat8E3M4Bits == (72 : UInt8) -- 3.0 +#guard (Float32.ofBits 0x41780000).toFloat8E3M4Bits == (111 : UInt8) -- 15.5 (max) +#guard (Float32.ofBits 0x7F800000).toFloat8E3M4Bits == (112 : UInt8) -- +inf +#guard (Float32.ofBits 0xFF800000).toFloat8E3M4Bits == (240 : UInt8) -- -inf +#guard (Float32.ofBits 0x7FC00000).toFloat8E3M4Bits == (0x78 : UInt8) -- +NaN +-- Subnormal encode (verified against ml_dtypes) +#guard (Float32.ofBits 0x3C800000).toFloat8E3M4Bits == (1 : UInt8) -- 0.015625 (smallest subnormal) +#guard (Float32.ofBits 0x3D000000).toFloat8E3M4Bits == (2 : UInt8) -- 0.03125 +#guard (Float32.ofBits 0x3D800000).toFloat8E3M4Bits == (4 : UInt8) -- 0.0625 +-- Negative overflow +#guard (Float32.ofBits 0xC1800000).toFloat8E3M4Bits == (240 : UInt8) -- -16.0 → -inf + section Test #guard ( @@ -662,6 +797,13 @@ warning: declaration uses 'sorry' let f := bits.toFloat32FromFloat8E5M2 f.toFloat8E5M2Bits == bits ∨ f != f := by plausible +-- Exhaustive e3m4 round-trip: decode -> encode for all 256 byte values. +-- NaN patterns are excluded (f != f) since Lean normalizes NaN bits. +#guard (List.range 256).all fun i => + let bits := i.toUInt8 + let f := bits.toFloat32FromFloat8E3M4 + f.toFloat8E3M4Bits == bits || f != f + end Test end TensorLib diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 842ec34..5c8f085 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -111,7 +111,10 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with | .uint16 => "u2" | .uint32 => "u4" | .uint64 => "u8" -| .float8_e4m3 => "V1" +-- 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_e5m2 => "f1" | .float16 => "f2" | .bfloat16 => "V2" @@ -127,8 +130,10 @@ def fromNpyString (s : String) : Err Dtype := -- We only recognize " 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) | .float16 => t.mapM (fun b => Dtype.byteArrayToFloat16 .float16 b) | .bfloat16 => t.mapM (fun b => Dtype.byteArrayToBFloat16 .bfloat16 b) | _ => t.mapM ( fun b => Float32.ofLEByteArray b) @@ -610,6 +611,7 @@ def toFloat64Tree (arr : Tensor) : Err (Format.Tree Float) := do match arr.dtype with | .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) | .float16 => t.mapM (fun b => do let f <- Dtype.byteArrayToFloat16 .float16 b; return f.toFloat) | .bfloat16 => t.mapM (fun b => do let f <- Dtype.byteArrayToBFloat16 .bfloat16 b; return f.toFloat) | .float32 => t.mapM (fun b => do let f <- Float32.ofLEByteArray b; return f.toFloat) @@ -668,14 +670,26 @@ def ofNpy (arr : Npy.Ndarray) : Err Tensor := do If we have a non-trivial view, we will need a copy, since strides and start positions are not included in the .npy file format -/ -def toNpy (arr : Tensor) : Npy.Ndarray := - let arr := if arr.isTriviallyReshapable then arr else arr.copy - let descr := Npy.Dtype.mk arr.dtype Npy.ByteOrder.littleEndian - let shape := arr.shape - let header : Npy.Header := { descr := descr, shape := shape } - let data := arr.data - let startIndex := 0 - { header, data, startIndex } +def toNpy (arr : Tensor) : Err Npy.Ndarray := + -- fp8_e3m4 and fp8_e4m3 both serialize to " true | .ok _ => false +-- toNpy accepts e4m3 (not blocked) +#guard match (Tensor.zeros .float8_e4m3 (Shape.mk [2])).toNpy with | .ok _ => true | .error _ => false + end Test end Tensor diff --git a/TensorLib/Test.lean b/TensorLib/Test.lean index 5daedcc..bf449d0 100644 --- a/TensorLib/Test.lean +++ b/TensorLib/Test.lean @@ -589,13 +589,137 @@ private def testFloat8E5M2EdgeCases : IO Bool := do return checks.all id +-- float8_e3m4 edge cases: decode from npy, arithmetic, and casting +-- E3M4: 1 sign + 3 exponent + 4 mantissa, bias=3, IEEE-like (has inf and NaN) +-- verified against ml_dtypes outputs +private def testFloat8E3M4EdgeCases : IO Bool := do + let file <- saveNumpyArray "np.array([15.5, 0.1, -0.0, 0.015625, 1.5, 0.5, 8.0, 2.0]).astype(__import__('ml_dtypes').float8_e3m4)" + let npy <- Npy.parseFile file + let arr <- IO.ofExcept (Tensor.ofNpy npy) + let _ <- IO.FS.removeFile file + -- Decode 1-byte e3m4 element at offset using extract + let decode (offset : Nat) : Err Float32 := + Dtype.decodeFloat8E3M4 (arr.data.extract offset (offset + 1)) + let mut checks : List Bool := [] + + -- max representable value (15.5) + let v0 <- IO.ofExcept (decode 0) + let pass := v0 == Float32.ofBits 0x41780000 + IO.println s!"fp8_e3m4 v0 (15.5): {pass}" + checks := pass :: checks + + -- 0.1 rounded in e3m4 + let v1 <- IO.ofExcept (decode 1) + let diff := v1 - 0.09375 + let pass := diff < 0.01 && diff > -0.01 + IO.println s!"fp8_e3m4 v1 (0.1 ~ 0.09375): {pass}" + checks := pass :: checks + + -- -0 + let v2 <- IO.ofExcept (decode 2) + let pass := v2 == 0.0 + IO.println s!"fp8_e3m4 v2 (-0): {pass}" + checks := pass :: checks + + -- -0 sign bit preserved + let pass := v2.toBits == 0x80000000 + IO.println s!"fp8_e3m4 v2 (-0 sign preserved): {pass}" + checks := pass :: checks + + -- smallest subnormal: 2^(-6) + let v3 <- IO.ofExcept (decode 3) + let pass := v3 == Float32.ofBits 0x3C800000 + IO.println s!"fp8_e3m4 v3 (smallest subnormal): {pass}" + checks := pass :: checks + + -- 1.5 + let v4 <- IO.ofExcept (decode 4) + let pass := v4 == 1.5 + IO.println s!"fp8_e3m4 v4 (1.5): {pass}" + checks := pass :: checks + + -- 0.5 + let v5 <- IO.ofExcept (decode 5) + let pass := v5 == 0.5 + IO.println s!"fp8_e3m4 v5 (0.5): {pass}" + checks := pass :: checks + + -- 8 = (within range, maxSafeForNat for e3m4 is 15) + let v6 <- IO.ofExcept (decode 6) + let pass := v6 == Float32.ofNat 8 + IO.println s!"fp8_e3m4 v6 (8): {pass}" + checks := pass :: checks + + -- 2.0 + let v7 <- IO.ofExcept (decode 7) + let pass := v7 == 2.0 + IO.println s!"fp8_e3m4 v7 (2.0): {pass}" + checks := pass :: checks + + -- Arithmetic: 1.5 and 2.0 + let a := toLEByteArray (56 : UInt8) -- e3m4 encoding of 1.5 + let b := toLEByteArray (64 : UInt8) -- e3m4 encoding of 2.0 + let negA := toLEByteArray (184 : UInt8) -- e3m4 encoding of -1.5 + + let pass <- checkBitsU8 "fp8_e3m4 add (1.5 + 2.0 = 3.5)" 76 (Dtype.add .float8_e3m4 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e3m4 sub (1.5 - 2.0 = -0.5)" 160 (Dtype.sub .float8_e3m4 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e3m4 mul (1.5 * 2.0 = 3.0)" 72 (Dtype.mul .float8_e3m4 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e3m4 div (1.5 / 2.0 = 0.75)" 40 (Dtype.div .float8_e3m4 a b) + checks := pass :: checks + + let pass <- checkBitsU8 "fp8_e3m4 abs (-1.5) = 1.5" 56 (Dtype.abs .float8_e3m4 negA) + checks := pass :: checks + + -- Casting: e3m4(1.5) -> int8 = 1 + let castToI8 <- IO.ofExcept (Dtype.castOverflow .float8_e3m4 a .int8) + let pass := castToI8.toNat == 1 + IO.println s!"fp8_e3m4 cast to int8 (1.5 -> 1): {pass}" + checks := pass :: checks + + -- e3m4(-0) -> bool = false + let negZero := toLEByteArray (128 : UInt8) + let castToBool <- IO.ofExcept (Dtype.castOverflow .float8_e3m4 negZero .bool) + let pass := castToBool == ByteArray.mk #[0] + IO.println s!"fp8_e3m4 -0 to bool (false): {pass}" + checks := pass :: checks + + -- fp32(2.0) -> e3m4 = bits 64 + let f32_2 := toLEByteArray (Float32.ofNat 2) + let castToE3m4 <- IO.ofExcept (Dtype.castOverflow .float32 f32_2 .float8_e3m4) + let pass := castToE3m4 == toLEByteArray (64 : UInt8) + IO.println s!"fp8_e3m4 fp32 to e3m4 (2.0): {pass}" + checks := pass :: checks + + -- Overflow to inf: 16 -> inf (bits 112) + let f32_16 := toLEByteArray (Float32.ofNat 16) + let castOverflow <- IO.ofExcept (Dtype.castOverflow .float32 f32_16 .float8_e3m4) + let pass := castOverflow == toLEByteArray (112 : UInt8) + IO.println s!"fp8_e3m4 overflow (16 -> inf): {pass}" + checks := pass :: checks + + -- +inf preserved + let f32_inf := toLEByteArray (Float32.ofBits 0x7F800000) + let castInf <- IO.ofExcept (Dtype.castOverflow .float32 f32_inf .float8_e3m4) + let pass := castInf == toLEByteArray (112 : UInt8) + IO.println s!"fp8_e3m4 +inf -> +inf: {pass}" + checks := pass :: checks + + return checks.all id + def runAllTests : IO Bool := do return (<- testTensorElementBV Dtype.uint16) && (<- testTensorElementBV Dtype.uint32) && (<- testFloat16EdgeCases) && (<- testBFloat16EdgeCases) && (<- testFloat8E4M3EdgeCases) && - (<- testFloat8E5M2EdgeCases) + (<- testFloat8E5M2EdgeCases) && + (<- testFloat8E3M4EdgeCases) end Test end TensorLib