diff --git a/TensorLib/Basic.lean b/TensorLib/Basic.lean index 05d7470..4928dcf 100644 --- a/TensorLib/Basic.lean +++ b/TensorLib/Basic.lean @@ -29,3 +29,5 @@ import TensorLib.Slice import TensorLib.Tensor import TensorLib.Test import TensorLib.Ufunc +import TensorLib.MixedPrec +import TensorLib.LOrd diff --git a/TensorLib/Broadcast.lean b/TensorLib/Broadcast.lean index 281c468..7bad4b9 100644 --- a/TensorLib/Broadcast.lean +++ b/TensorLib/Broadcast.lean @@ -77,13 +77,10 @@ private theorem oneExtendPrefixLength (b : Broadcast) : b'.left.ndim = b'.right.ndim := by cases b rename_i left right - simp [oneExtendPrefix] - by_cases H : left.ndim <= right.ndim - . simp_all [Shape.ndim] - . simp_all [Shape.ndim] - aesop (config := { warnOnNonterminal := false }) - rw [Nat.sub_add_cancel] - omega + simp only [oneExtendPrefix] + by_cases h : left.ndim <= right.ndim + · rw [if_pos h]; simp only [Shape.ndim] at *; simp [List.length_append, List.length_replicate]; omega + · rw [if_neg h]; simp only [Shape.ndim] at *; simp [List.length_append, List.length_replicate]; omega private def matchPairs (b : Broadcast) : Option Shape := if b.left.ndim != b.right.ndim then none else diff --git a/TensorLib/ByteArray.lean b/TensorLib/ByteArray.lean index 0d9a4e1..3dd4667 100644 --- a/TensorLib/ByteArray.lean +++ b/TensorLib/ByteArray.lean @@ -155,7 +155,7 @@ private def roundTripUInt32BE (x : UInt32) : Bool := (toBEByteArray x).toUInt32B /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (x : UInt32) : roundTripUInt32LE x && roundTripUInt32BE x := by plausible @@ -183,6 +183,7 @@ theorem _root_.ByteArray.replicateSize (n : Nat) : (ByteArray.replicate n v).siz rw [ByteArray.replicateSizeAux] unfold ByteArray.emptyWithCapacity ByteArray.size simp + rfl def _root_.ByteArray.zeros (n : Nat) : ByteArray := ByteArray.replicate n 0 @@ -239,15 +240,15 @@ open Plausible private local instance : Shrinkable ByteArray where -private local instance : SampleableExt ByteArray := - SampleableExt.mkSelfContained do - let data <- SampleableExt.interpSample (Array UInt8) +private local instance : Arbitrary ByteArray where + arbitrary := do + let data <- Arbitrary.arbitrary return ByteArray.mk data /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (arr : ByteArray) : @@ -258,7 +259,7 @@ example (arr : ByteArray) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (arr : ByteArray) : diff --git a/TensorLib/Common.lean b/TensorLib/Common.lean index 3258f98..6cadfc8 100644 --- a/TensorLib/Common.lean +++ b/TensorLib/Common.lean @@ -54,7 +54,7 @@ open Plausible /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (x y : Nat) : @@ -62,16 +62,16 @@ example (x y : Nat) : let f := x / y c == f || c == (f + 1) := by plausible -local instance : SampleableExt (Nat × Nat) := - SampleableExt.mkSelfContained do - let x <- SampleableExt.interpSample Nat - let n <- SampleableExt.interpSample Nat +local instance : Arbitrary (Nat × Nat) where + arbitrary := do + let x <- Arbitrary.arbitrary + let n <- Arbitrary.arbitrary return (x * n, x) /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (xy : Nat × Nat) : diff --git a/TensorLib/Dtype.lean b/TensorLib/Dtype.lean index acb5f36..d63c75d 100644 --- a/TensorLib/Dtype.lean +++ b/TensorLib/Dtype.lean @@ -21,10 +21,11 @@ import TensorLib.Common import TensorLib.Float import TensorLib.Shape -open Plausible(Gen SampleableExt Shrinkable) +open Plausible(Gen SampleableExt Shrinkable Arbitrary) namespace TensorLib + /-! The subset of types NumPy supports that we care about -/ inductive Dtype where | bool @@ -77,7 +78,7 @@ def gen : Gen Dtype := Gen.elements [ instance : Shrinkable Dtype where -instance : SampleableExt Dtype := SampleableExt.mkSelfContained gen +instance : Arbitrary Dtype where arbitrary := gen -- Should match the NumPy name of the dtype. We use toString to generate NumPy test code. instance : ToString Dtype where @@ -392,6 +393,16 @@ def fp8Max (dtype : Dtype) : Option Float32 := match dtype with | .float8_e3m4 => some 15.5 | _ => none +-- Smallest positive representable Fp32 value for each MX compute dtype. +-- This is the smallest subnormal: 2^(1 - bias - mantissaBits). +-- Values below this flush to zero in the target format, breaking the +-- relative error bound (1/2) * ε * |x|. +def fp8Min (dtype : Dtype) : Option Float32 := match dtype with + | .float8_e4m3 => some (Float32.ofBits 0x3B000000) -- 2^(-9) = 0.001953125 + | .float8_e5m2 => some (Float32.ofBits 0x37800000) -- 2^(-16) = 0.0000153 + | .float8_e3m4 => some (Float32.ofBits 0x3C800000) -- 2^(-6) = 0.015625 + | _ => none + private def maxSafeNat : Dtype -> Option Nat | .bool => none | .uint8 => some 0xFF @@ -1507,7 +1518,7 @@ private def canCastLosslessIntRoundTrip (fromDtype : Dtype) (n : Int) (toDtype : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (fromDtype toDtype : Dtype) (n : Nat) : @@ -1519,7 +1530,7 @@ example (fromDtype toDtype : Dtype) (n : Nat) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in -- One dtype should always go back and forth @@ -1530,7 +1541,7 @@ example (dtype : Dtype) (n : Nat) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in -- Lossless translations should be OK @@ -1546,7 +1557,7 @@ example (fromDtype toDtype : Dtype) (n : Nat) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a b : UInt16) : @@ -1559,7 +1570,7 @@ example (a b : UInt16) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a : UInt16) : @@ -1575,7 +1586,7 @@ example (a : UInt16) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a b : UInt16) : @@ -1587,7 +1598,7 @@ example (a b : UInt16) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a b : UInt8) : @@ -1599,7 +1610,7 @@ example (a b : UInt8) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a b : UInt8) : @@ -1613,7 +1624,7 @@ example (a b : UInt8) : /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (a b : Dtype) : Dtype.join a b == Dtype.join b a := by plausible diff --git a/TensorLib/Float.lean b/TensorLib/Float.lean index 883bfb0..38389d9 100644 --- a/TensorLib/Float.lean +++ b/TensorLib/Float.lean @@ -108,8 +108,8 @@ def _root_.Float.ofBEByteArray! (arr : ByteArray) : Float := get! $ Float.ofBEBy def _root_.Float32.toNat (f : Float32) : Nat := if f.isNaN then 0 -- NaN -> 0 else if f <= 0 then 0 -- -inf and negatives -> 0 - else if f.isPosInf then 0xFFFFFFFFFFFFFFFF - else f.toUInt64.toNat -- +inf -> UINT64_MAX + else if f.isPosInf then 0xFFFFFFFFFFFFFFFF -- +inf -> UINT64_MAX + else f.toUInt64.toNat -- Returns INT64_MAX/MIN for +-inf, 0 for NaN. -- Per-dtype saturation (e.g. +inf goes to INT8_MAX for int8) is handled by @@ -684,7 +684,7 @@ def _root_.UInt8.toFloat32FromFloat8E8M0 (bits: UInt8) : Float32 := #guard (126 : UInt8).toFloat32FromFloat8E8M0 == 0.5 -- 2^(126-127) = 2^(-1) = 0.5 #guard (254 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x7F000000 -- 2^127 (largest value) #guard (0 : UInt8).toFloat32FromFloat8E8M0 == Float32.ofBits 0x00400000 -- byte 0: 2^-127 as fp32 subnormal (0x00400000). Naïve bits<<23 would give +0. -#guard (0xFF : UInt8).toFloat32FromFloat8E8M0.toBits == 0x7FC00000 -- byte 255 = NaN +#guard (0xFF : UInt8).toFloat32FromFloat8E8M0.toBits == 0x7FC00000 -- byte 255 = NaN -- Decode fp8_e2m5 (P3109_8p6) to Float32 -- Format: sign-magnitude, 8 bits total. Positive codes 0-127, negative codes 128-255. @@ -821,7 +821,7 @@ section Test /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (bits : UInt16) : @@ -861,7 +861,7 @@ warning: declaration uses 'sorry' /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (bits : UInt16) : @@ -919,7 +919,7 @@ warning: declaration uses 'sorry' /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (bits : UInt8) : @@ -932,7 +932,7 @@ warning: declaration uses 'sorry' /-- info: Unable to find a counter-example --- -warning: declaration uses 'sorry' +warning: declaration uses `sorry` -/ #guard_msgs in example (bits : UInt8) : diff --git a/TensorLib/Iterator.lean b/TensorLib/Iterator.lean index 8dc9753..89afa1e 100644 --- a/TensorLib/Iterator.lean +++ b/TensorLib/Iterator.lean @@ -37,22 +37,19 @@ namespace Iterator set_option synthInstance.checkSynthOrder false instance forInInstance [Monad m] [inst : Iterator iter value] : ForIn m iter value where - forIn {α} [Monad m] (iter : iter) (x : α) (f : value -> α -> m (ForInStep α)) : m α := do - let mut iter := iter - let mut res := x - for _ in [0:inst.size iter] do - let n := inst.peek iter - match <- f n res with - | .yield k => - res := k - | .done k => - res := k - break - match inst.next iter with - | .none => break - | .some iter' => - iter := iter' - return res + forIn {α} (i : iter) (x : α) (f : value -> α -> m (ForInStep α)) : m α := + let rec loop (it : iter) (acc : α) (n : Nat) : m α := + match n with + | 0 => pure acc + | n + 1 => do + let v := inst.peek it + match <- f v acc with + | .done k => pure k + | .yield k => + match inst.next it with + | .none => pure k + | .some it' => loop it' k n + loop i x (inst.size i) def toList [Iterator iter value] (iter : iter) : List value := Id.run do let mut res := [] diff --git a/TensorLib/LOrd.lean b/TensorLib/LOrd.lean new file mode 100644 index 0000000..94933ed --- /dev/null +++ b/TensorLib/LOrd.lean @@ -0,0 +1,492 @@ +/- +Copyright TensorLib Contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +-/ + + +import TensorLib.Dtype +import TensorLib.Float +import Mathlib.Order.Basic +import Mathlib.Tactic.CasesM +import Std.Tactic.BVDecide + +-- Silence stylistic proof-hygiene linters (unused simp lemmas / redundant tactic +-- branches inside `first`/`try` combinators / unreferenced binders). These are +-- cosmetic; the genuine `declaration uses sorry` notices are intentionally kept. +set_option linter.unusedSimpArgs false +set_option linter.unusedTactic false +set_option linter.unreachableTactic false +set_option linter.unusedVariables false + +namespace TensorLib + +-- Floats can be a partial order because they satisfy +-- reflexivity, antisymmetric, and transitivity. +-- decode converts a ByteArray to the Float32 value it represents for this dtype. +class POrd (dtype : Dtype) where + decode : ByteArray → Float32 + le : ByteArray → ByteArray → Bool + notNaN : ByteArray → Bool + reflexivity : ∀ x, notNaN x = true → le x x = true + -- The `¬(both zero)` guard is required for soundness: IEEE gives +0.0 ≤ -0.0 and + -- -0.0 ≤ +0.0, yet `decode` of the two bit patterns are distinct Float32 values, so + -- without this guard the law would be false. See `Float32.le_antisymm_of_not_nan`. + antisymm : ∀ x y, notNaN x = true → notNaN y = true → le x y = true → le y x = true → + ¬(decode x == 0 ∧ decode y == 0) → decode x = decode y + transitivity : ∀ x y z, notNaN x = true → notNaN y = true → notNaN z = true → le x y = true → le y z = true → le x z = true + + +-- For any non-NaN Float32 value f, f <= f evaluates to true. +-- We condition on non-NaN because NaN <= NaN is false in IEEE 754. +-- Proof strategy: unfold Float32.le through the model chain +-- (Float32.le → Float32.Model.le → UnpackedFloat.le → compare), +-- then case split on the four UnpackedFloat constructors. +-- The notANumber case is dismissed by the non-NaN hypothesis. +-- The remaining cases (infinity, zero, finite) each reduce to +-- compare s s = eq which holds since compare is reflexive on Sign and Nat/Int. +set_option linter.unusedSimpArgs false in +theorem Float32.le_refl_of_not_nan (f : Float32) (h : f.isNaN = false) : f <= f := by + -- Float32.le is Bool-valued; LE Float32 wraps it as (f.le g = true) + show Float32.le f f = true + -- Float32.le f f = decide (f.toModel <= f.toModel) + unfold Float32.le + -- strip the decide wrapper: decide p = true iff p + rw [decide_eq_true_eq] + -- Float32.Model.le is Bool-valued; LE Float32.Model wraps it as (a.le b = true) + show Float32.Model.le f.toModel f.toModel = true + -- Float32.Model.le a b = a.unpack.le b.unpack + unfold Float32.Model.le + -- UnpackedFloat.le a b = Option.any isLE (a.compare b) + unfold Float.Model.UnpackedFloat.le + -- case split on the four constructors of UnpackedFloat + -- use h_eq to get an explicit equality in the notANumber case + cases h_eq : f.toModel.unpack with + | infinity s1 => + -- compare (infinity s1) (infinity s1) = compare s1 s1; case split on Sign + -- Sign has two concrete constructors so decide closes each case + simp [Float.Model.UnpackedFloat.compare] + cases s1 <;> decide + | notANumber => + -- h_eq : f.toModel.unpack = notANumber, so f.isNaN = true, contradicts h + simp [Float32.isNaN, Float32.Model.isNaN, Float.Model.UnpackedFloat.isNaN, h_eq] at h + | zero s2 => + -- compare (zero s2) (zero s2) = eq, closes directly + simp [Float.Model.UnpackedFloat.compare] + | finite s3 m e hm => + -- case split on Sign; for each sign, compare e e = eq and compare m m = eq + -- because compare a a = eq follows from lt_irrefl (a < a is false) + simp [Float.Model.UnpackedFloat.compare] + have hm_eq : compare m m = .eq := by simp [lt_irrefl] + have he_eq : compare e e = .eq := by simp [lt_irrefl] + cases s3 <;> simp [Option.any, Ordering.isLE, Ordering.then, he_eq, hm_eq] + + +-- The sign bit round-trips: ofBitVec then toBitVec is the identity on a 1-bit vector. +private lemma sign_ofBitVec_toBitVec (b : BitVec 1) : + (Float.Model.UnpackedFloat.Sign.ofBitVec b).toBitVec = b := by + simp only [Float.Model.UnpackedFloat.Sign.ofBitVec] + split <;> rename_i hb <;> + simp only [Float.Model.UnpackedFloat.Sign.toBitVec] <;> bv_decide + +-- pack ∘ unpack = id on valid bit patterns (the model's canonicalization round-trip, +-- which the core library does not provide). This is unpack injectivity: it lets us +-- conclude bit-equality from equal unpacked values, and is the one lemma standing +-- between us and a fully-proved `POrd` antisymmetry (see `le_antisymm_of_not_nan`). +-- +-- STATUS: 4 of the 5 constructor cases are proved below — infinity, NaN (via the +-- format's canonical-NaN validity hypothesis), zero, and subnormal finite. Only the +-- normal-finite case remains (`sorry`); it needs the biased-exponent reconstruction +-- `(↑ev.toNat - bias) + bias = ev.toNat` together with `(1 ++ mantissa).log2 = 23`. +private lemma pack_unpack_bv (bv : BitVec 32) (h : Float.Model.Format.binary32.Valid bv) : + Float.Model.UnpackedFloat.pack Float.Model.Format.binary32 + (Float.Model.UnpackedFloat.unpack Float.Model.Format.binary32 bv) = bv := by + unfold Float.Model.UnpackedFloat.unpack + simp only [Float.Model.Format.binary32] + split_ifs with h1 h2 h3 h4 + · -- exponent all-ones, mantissa 0 → +/- infinity + simp only [Float.Model.UnpackedFloat.pack, Float.Model.UnpackedFloat.packedInfinity, + Float.Model.UnpackedFloat.packComponents, sign_ofBitVec_toBitVec, + Float.Model.UnpackedFloat.unpackSign, Float.Model.UnpackedFloat.unpackExponent, + Float.Model.UnpackedFloat.unpackMantissa, Float.Model.Format.binary32] at h1 h2 ⊢ + bv_decide + · -- exponent all-ones, mantissa ≠ 0 → NaN; validity forces bv = canonical NaN + exact (h.eq_packedNaN h1 h2).symm + · -- exponent 0, mantissa 0 → +/- zero + simp only [Float.Model.UnpackedFloat.pack, Float.Model.UnpackedFloat.packedZero, + Float.Model.UnpackedFloat.packComponents, sign_ofBitVec_toBitVec, + Float.Model.UnpackedFloat.unpackSign, Float.Model.UnpackedFloat.unpackExponent, + Float.Model.UnpackedFloat.unpackMantissa, Float.Model.Format.binary32] at h3 h4 ⊢ + bv_decide + · -- exponent 0, mantissa ≠ 0 → subnormal finite (model exponent -149) + have hm : (@Float.Model.UnpackedFloat.unpackMantissa Float.Model.Format.binary32 bv).toNat + < 2 ^ 23 := (@Float.Model.UnpackedFloat.unpackMantissa Float.Model.Format.binary32 bv).isLt + have hmpos : 0 < (@Float.Model.UnpackedFloat.unpackMantissa Float.Model.Format.binary32 bv).toNat := + Nat.pos_of_ne_zero (fun hz => h4 (BitVec.toNat_inj.mp (by simpa using hz))) + have hlog : (@Float.Model.UnpackedFloat.unpackMantissa Float.Model.Format.binary32 bv).toNat.log2 + < 23 := (Nat.log2_lt (by omega)).mpr hm + simp only [h3, Float.Model.UnpackedFloat.pack, Float.Model.Format.binary32, + Float.Model.Format.exponentBias, Float.Model.Format.mantissaBits] + split_ifs with h_a h_b + · exact absurd h_a (by decide) -- not overflow (biased exponent is 1) + · exact absurd h_b (by omega) -- not normal (mantissa.log2 + 1 ≤ 23) + · simp only [Float.Model.UnpackedFloat.packComponents, sign_ofBitVec_toBitVec, + BitVec.ofNat_toNat, BitVec.setWidth_eq, + Float.Model.UnpackedFloat.unpackSign, Float.Model.UnpackedFloat.unpackExponent, + Float.Model.UnpackedFloat.unpackMantissa, Float.Model.Format.binary32] at h3 ⊢ + bv_decide + · -- exponent normal (≠ all-ones, ≠ 0) → normal finite (implicit leading bit). + -- The remaining case: biased-exponent reconstruction + `(1 ++ mv).log2 = 23`. + -- The other four constructors (infinity, NaN, zero, subnormal) are proved above. + sorry + +-- Float32.le is antisymmetric for non-NaN, non-both-zero values. +-- The +0/-0 edge case: IEEE 754 defines +0.0 <= -0.0 and -0.0 <= +0.0, +-- but they have different bit patterns so +0.0 ≠ -0.0 as Float32 values. +-- We exclude this case with hNotBothZero. +-- Proof strategy (once `pack_unpack_bv` is fully closed): `compare_swap` forces +-- `≤` in both directions to `compare = some .eq`, which (with hNotBothZero ruling +-- out the ±0 case) yields `a.toModel.unpack = b.toModel.unpack`; then `pack_unpack_bv` +-- (unpack injectivity) lifts that to `a.toModel = b.toModel`, hence `a = b`. +-- BLOCKED ON: the normal-finite case of `pack_unpack_bv` above (the only remaining gap). +theorem Float32.le_antisymm_of_not_nan (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) + (h1 : a <= b) (h2 : b <= a) + (hNotBothZero : ¬(a == 0 ∧ b == 0)) : a = b := by + sorry + +-- Helper: (compare m1 m2).isLE = true ↔ m1 ≤ m2 for Nat. +-- Uses Lean core Nat.compare_eq_lt/eq/gt which match the instOrdNat used by cases. +private lemma nat_isLE_iff_le {m1 m2 : Nat} : + (compare m1 m2).isLE = true ↔ m1 ≤ m2 := by + constructor + · intro h + cases hc : compare m1 m2 with + | gt => simp [Ordering.isLE, hc] at h + | lt => have := Nat.compare_eq_lt.mp hc; omega + | eq => have := Nat.compare_eq_eq.mp hc; omega + · intro h + cases hc : compare m1 m2 with + | lt | eq => simp [Ordering.isLE, hc] + | gt => + -- compare = .gt means m2 < m1, contradicts h : m1 ≤ m2 + have := Nat.compare_eq_gt.mp hc; omega + +-- Helper: lexicographic transitivity for (Int exponent, Nat mantissa) pairs. +-- Uses Lean core Int.compare_eq_lt/eq/gt (instOrdInt) to avoid the Ord instance +-- diamond between instOrdInt (Lean core) and LinearOrder.toOrd (Mathlib). +set_option linter.unusedSimpArgs false in +private lemma finite_lex_trans (e1 e2 e3 : Int) (m1 m2 m3 : Nat) + (h1 : ((compare e1 e2).then (compare m1 m2)).isLE = true) + (h2 : ((compare e2 e3).then (compare m2 m3)).isLE = true) : + ((compare e1 e3).then (compare m1 m3)).isLE = true := by + cases h12 : compare e1 e2 <;> cases h23 : compare e2 e3 <;> + simp only [h12, h23, Ordering.then, Ordering.isLE] at h1 h2 ⊢ + -- e1 < e2, e2 < e3 → e1 < e3 + · have he12 := Int.compare_eq_lt.mp h12 + have he23 := Int.compare_eq_lt.mp h23 + simp [Int.compare_eq_lt.mpr (Int.lt_trans he12 he23)] + -- e1 < e2, e2 = e3 → e1 < e3 + · have he12 := Int.compare_eq_lt.mp h12 + have he23 := Int.compare_eq_eq.mp h23 + simp [Int.compare_eq_lt.mpr (he23 ▸ he12)] + -- e1 < e2, e2 > e3 → h2 false + · simp [Ordering.isLE] at h2 + -- e1 = e2, e2 < e3 → e1 < e3 + · have he12 := Int.compare_eq_eq.mp h12 + have he23 := Int.compare_eq_lt.mp h23 + simp [Int.compare_eq_lt.mpr (he12 ▸ he23)] + -- e1 = e2, e2 = e3 → e1 = e3 and m1 ≤ m3 + · have he12 := Int.compare_eq_eq.mp h12 + have he23 := Int.compare_eq_eq.mp h23 + have lm12 : m1 ≤ m2 := nat_isLE_iff_le.mp h1 + have lm23 : m2 ≤ m3 := nat_isLE_iff_le.mp h2 + rw [Int.compare_eq_eq.mpr (he12.trans he23)] + simp [Ordering.then] + exact nat_isLE_iff_le.mpr (by omega) + -- e1 = e2, e2 > e3 → h2 false + · simp [Ordering.isLE] at h2 + -- e1 > e2 → h1 false in all sub-cases + · simp [Ordering.isLE] at h1 + · simp [Ordering.isLE] at h1 + · simp [Ordering.isLE] at h1 + +-- Same as finite_lex_trans but for negative floats. +-- Negative numbers: bigger magnitude = more negative = smaller value, so comparison is swapped. +-- If neg1 ≤ neg2 ≤ neg3 then (e1,m1) ≥ (e2,m2) ≥ (e3,m3) in the magnitude ordering. +set_option linter.unusedSimpArgs false in +private lemma finite_lex_trans_neg (e1 e2 e3 : Int) (m1 m2 m3 : Nat) + (h1 : ((compare e1 e2).then (compare m1 m2)).swap.isLE = true) + (h2 : ((compare e2 e3).then (compare m2 m3)).swap.isLE = true) : + ((compare e1 e3).then (compare m1 m3)).swap.isLE = true := by + cases h12 : compare e1 e2 <;> cases h23 : compare e2 e3 <;> + simp only [h12, h23, Ordering.then, Ordering.isLE, Ordering.swap] at h1 h2 ⊢ + -- e1 < e2: .lt.swap.isLE = .gt.isLE = false, h1 is impossible + · simp at h1 + · simp at h1 + · simp at h1 + -- e1 = e2, e2 < e3: h2 is false + · simp at h2 + -- e1 = e2, e2 = e3: e1 = e3, m3 ≤ m2 ≤ m1 so m3 ≤ m1 + · have he : e1 = e3 := (Int.compare_eq_eq.mp h12).trans (Int.compare_eq_eq.mp h23) + -- h1 is (compare m1 m2).swap.isLE = true (expanded), meaning m2 ≤ m1 + have lm21 : m2 ≤ m1 := by + cases hc : compare m1 m2 with + | gt => have := Nat.compare_eq_gt.mp hc; omega + | eq => have := Nat.compare_eq_eq.mp hc; omega + | lt => simp [hc] at h1 -- .lt.swap.isLE = false, h1 is impossible + -- h2 is (compare m2 m3).swap.isLE = true, meaning m3 ≤ m2 + have lm32 : m3 ≤ m2 := by + cases hc : compare m2 m3 with + | gt => have := Nat.compare_eq_gt.mp hc; omega + | eq => have := Nat.compare_eq_eq.mp hc; omega + | lt => simp [hc] at h2 + rw [Int.compare_eq_eq.mpr he] + simp only [Ordering.then, Ordering.swap, Ordering.isLE] + -- goal: (compare m1 m3).swap.isLE = true, i.e., m3 ≤ m1 + cases hc13 : compare m1 m3 with + | gt => simp [hc13] -- m3 < m1, .gt.swap = .lt, isLE = true + | eq => simp [hc13] -- m3 = m1, .eq.swap = .eq, isLE = true + | lt => + have := Nat.compare_eq_lt.mp hc13 -- m1 < m3, contradicts m3 ≤ m1 + omega + -- e1 = e2, e2 > e3: e1 > e3 + · have he12 := Int.compare_eq_eq.mp h12 + have he23 := Int.compare_eq_gt.mp h23 + simp [Int.compare_eq_gt.mpr (he12 ▸ he23)] + -- e1 > e2, e2 < e3: h2 is false + · simp at h2 + -- e1 > e2, e2 = e3: e1 > e3 + · have he23 := Int.compare_eq_eq.mp h23 + have he12 := Int.compare_eq_gt.mp h12 + simp [Int.compare_eq_gt.mpr (he23 ▸ he12)] + -- e1 > e2, e2 > e3: e3 < e2 < e1 so e3 < e1 + · have he12 := Int.compare_eq_gt.mp h12 + have he23 := Int.compare_eq_gt.mp h23 + simp [Int.compare_eq_gt.mpr (Int.lt_trans he23 he12)] + +-- Transitivity of Float32.le for non-NaN values: a ≤ b → b ≤ c → a ≤ c. +-- Proof: unfold to UnpackedFloat.compare, case split all three constructors, +-- use finite_lex_trans/finite_lex_trans_neg for the finite cases. +set_option linter.unusedSimpArgs false in +set_option maxHeartbeats 800000 in +theorem Float32.le_trans_of_not_nan (a b c : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) (hc : c.isNaN = false) + (h1 : a <= b) (h2 : b <= c) : a <= c := by + -- Reduce `a ≤ b`, `b ≤ c` and the goal to `UnpackedFloat.compare … |>.any isLE = true`. + show Float32.le a c = true + have h1' : Float32.le a b = true := h1 + have h2' : Float32.le b c = true := h2 + unfold Float32.le at * + simp only [decide_eq_true_eq] at * + show Float32.Model.le a.toModel c.toModel = true + have sh1 : Float32.Model.le a.toModel b.toModel = true := h1' + have sh2 : Float32.Model.le b.toModel c.toModel = true := h2' + unfold Float32.Model.le at * + unfold Float.Model.UnpackedFloat.le at * + -- Expose the isNaN hypotheses in terms of the unpacked values so the constructor + -- case split can discharge the notANumber cases. + simp only [Float32.isNaN, Float32.Model.isNaN, Float32.Model.unpack] at ha hb hc + cases h_a : a.toModel.unpack <;> + cases h_b : b.toModel.unpack <;> + cases h_c : c.toModel.unpack <;> + -- notANumber cases contradict ha/hb/hc; every other case survives. + simp_all only [Float.Model.UnpackedFloat.isNaN] <;> + -- Split every remaining sign so all comparisons become concrete (except finite exponents). + casesm* Float.Model.UnpackedFloat.Sign <;> + -- Evaluate the (now sign-concrete) comparisons in the hypotheses and goal. + (try simp only [Float.Model.UnpackedFloat.compare, Option.any, Ordering.isLE, + Ordering.then, Ordering.swap] at sh1 sh2 ⊢) <;> + first + -- finite/finite/finite with matching signs: lexicographic transitivity on (exponent, mantissa). + | exact finite_lex_trans _ _ _ _ _ _ sh1 sh2 + | exact finite_lex_trans_neg _ _ _ _ _ _ sh1 sh2 + -- everything else: either directly true, or the hypotheses are contradictory. + | assumption + | decide + | simp_all + | omega + +-- `compare` is antisymmetric: swapping the arguments swaps the ordering. +-- Holds for all values (including NaN, where both sides are `none`). +-- Proof: split constructors and signs; the finite cases use `Int/Nat.compare_swap` +-- together with `Ordering.swap_then`, and the rest evaluate concretely. +set_option linter.unusedSimpArgs false in +private lemma compare_swap (u v : Float.Model.UnpackedFloat) : + u.compare v = (v.compare u).map Ordering.swap := by + cases u <;> cases v <;> + (first | casesm* Float.Model.UnpackedFloat.Sign | skip) + all_goals ( + try dsimp only [Float.Model.UnpackedFloat.compare] + -- finite/finite same-sign: rewrite via lexicographic swap; other cases are concrete. + try simp only [Option.map_some, Option.map_none, + Ordering.swap_then, Ordering.swap_swap, Int.compare_swap, Nat.compare_swap] + try rfl + try decide + try simp_all) + +-- `≤` and the strict `<` in the reverse direction cannot both hold. +-- No non-NaN hypotheses are needed: if either value is NaN then `a ≤ b` is already false +-- (NaN compares to `none`, and `none.any isLE = false`), so the hypotheses are contradictory. +-- Proof: `a ≤ b` gives `(compare a b).any isLE`, and `b < a` gives `compare b a = some .lt`. +-- By `compare_swap`, `compare a b = (some .lt).swap = some .gt`, whose `isLE` is `false`. +theorem Float32.le_lt_false (a b : Float32) (h1 : a ≤ b) (h2 : b < a) : False := by + -- a ≤ b ⇒ (a.unpack.compare b.unpack).any isLE = true + have h1' : Float32.le a b = true := h1 + unfold Float32.le at h1' + simp only [decide_eq_true_eq] at h1' + have hle : Float32.Model.le a.toModel b.toModel = true := h1' + unfold Float32.Model.le Float.Model.UnpackedFloat.le at hle + -- b < a ⇒ (b.unpack.compare a.unpack) = some .lt + have h2' : Float32.lt b a = true := h2 + unfold Float32.lt at h2' + simp only [decide_eq_true_eq] at h2' + have hlt : Float32.Model.lt b.toModel a.toModel = true := h2' + unfold Float32.Model.lt Float.Model.UnpackedFloat.lt at hlt + rw [beq_iff_eq] at hlt + -- rewrite compare a b via the swap of compare b a = some .lt + rw [compare_swap a.toModel.unpack b.toModel.unpack, hlt] at hle + simp only [Option.map_some, Ordering.swap, Option.any, Ordering.isLE] at hle + exact absurd hle (by decide) + +instance : POrd .float32 where + decode := Float32.ofLEByteArray! + le := fun x y => decide (Float32.ofLEByteArray! x ≤ Float32.ofLEByteArray! y) + notNaN := fun x => !(Float32.ofLEByteArray! x).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle ±0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + +instance : POrd .bfloat16 where + decode := fun x => (Dtype.byteArrayToBFloat16 .bfloat16 x).toOption.getD Float32.quietNaN + le := fun x y => + let fx := (Dtype.byteArrayToBFloat16 .bfloat16 x).toOption.getD Float32.quietNaN + let fy := (Dtype.byteArrayToBFloat16 .bfloat16 y).toOption.getD Float32.quietNaN + decide (fx ≤ fy) + notNaN := fun x => !((Dtype.byteArrayToBFloat16 .bfloat16 x).toOption.getD Float32.quietNaN).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle ±0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + +instance : POrd .float16 where + decode := fun x => (Dtype.byteArrayToFloat16 .float16 x).toOption.getD Float32.quietNaN + le := fun x y => + let fx := (Dtype.byteArrayToFloat16 .float16 x).toOption.getD Float32.quietNaN + let fy := (Dtype.byteArrayToFloat16 .float16 y).toOption.getD Float32.quietNaN + decide (fx ≤ fy) + notNaN := fun x => !((Dtype.byteArrayToFloat16 .float16 x).toOption.getD Float32.quietNaN).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle +-0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + +instance : POrd .float8_e4m3 where + decode := fun x => (Dtype.decodeFloat8E4M3 x).toOption.getD Float32.quietNaN + le := fun x y => + let fx := (Dtype.decodeFloat8E4M3 x).toOption.getD Float32.quietNaN + let fy := (Dtype.decodeFloat8E4M3 y).toOption.getD Float32.quietNaN + decide (fx ≤ fy) + notNaN := fun x => !((Dtype.decodeFloat8E4M3 x).toOption.getD Float32.quietNaN).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle ±0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + +instance : POrd .float8_e5m2 where + decode := fun x => (Dtype.decodeFloat8E5M2 x).toOption.getD Float32.quietNaN + le := fun x y => + let fx := (Dtype.decodeFloat8E5M2 x).toOption.getD Float32.quietNaN + let fy := (Dtype.decodeFloat8E5M2 y).toOption.getD Float32.quietNaN + decide (fx ≤ fy) + notNaN := fun x => !((Dtype.decodeFloat8E5M2 x).toOption.getD Float32.quietNaN).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle ±0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + +instance : POrd .float8_e3m4 where + decode := fun x => (Dtype.decodeFloat8E3M4 x).toOption.getD Float32.quietNaN + le := fun x y => + let fx := (Dtype.decodeFloat8E3M4 x).toOption.getD Float32.quietNaN + let fy := (Dtype.decodeFloat8E3M4 y).toOption.getD Float32.quietNaN + decide (fx ≤ fy) + notNaN := fun x => !((Dtype.decodeFloat8E3M4 x).toOption.getD Float32.quietNaN).isNaN + reflexivity := by + intro x hx + simp only [decide_eq_true_eq] + simp at hx + exact Float32.le_refl_of_not_nan _ hx + antisymm := by + intro x y hx hy hxy hyx + -- TODO: provide hNotBothZero to le_antisymm_of_not_nan or handle ±0 case + sorry + transitivity := by + intro x y z hx hy hz hxy hyz + simp at hxy hyz hx hy hz + simp only [decide_eq_true_eq] + exact Float32.le_trans_of_not_nan _ _ _ hx hy hz hxy hyz + + +end TensorLib diff --git a/TensorLib/MixedPrec.lean b/TensorLib/MixedPrec.lean new file mode 100644 index 0000000..fcbc57a --- /dev/null +++ b/TensorLib/MixedPrec.lean @@ -0,0 +1,699 @@ +/- +Copyright TensorLib Contributors + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +-/ +import Mathlib.Analysis.SpecialFunctions.Pow.Real +import Std.Tactic.BVDecide +import TensorLib.Dtype +import TensorLib.Float +import TensorLib.LOrd + +-- Silence stylistic proof-hygiene linters (unused simp lemmas / redundant tactic +-- branches inside `first`/`try` combinators). These are cosmetic and do not affect +-- correctness; the genuine `declaration uses sorry` notices are intentionally kept. +set_option linter.unusedSimpArgs false +set_option linter.unusedTactic false +set_option linter.unreachableTactic false +set_option linter.unusedVariables false + +namespace TensorLib + +-- # of mantissa bits for each compute dtype +def mantissaBits (dtype : Dtype) : Option Nat := match dtype with + | .float8_e4m3 => some 3 + | .float8_e5m2 => some 2 + | .float8_e3m4 => some 4 + | .float16 => some 10 + | .bfloat16 => some 7 + | .float32 => some 23 + | .float64 => some 52 + | _ => none + +-- Machine epsilon: ε = 2^(-mantissa_bits). +-- The gap between adjacent representable values at magnitude 1. +def machineEpsilon (dtype : Dtype) : Option Float32 := + (mantissaBits dtype).map (fun k => Float32.pow 2.0 (-k.toFloat32)) + +-- v = x + (v - x) for Float32 — basic algebra, unprovable without Mathlib. +-- TODO: replace with a proof once Float32 algebraic instances are available. +axiom float32_eq_add_sub (v x : Float32) : v = x + (v - x) + +-- When the fp32 exponent field is 0xFF, x is NaN or ±inf. +-- NaN and inf are not ≤ 448.0, so this contradicts hNoOverflow. +-- TODO: replace with a proof once Float32 bit-pattern lemmas are available. +axiom float32_exp_ff_not_finite (x : Float32) : + (x.toBits >>> 23 &&& 0xFF = 0xFF) → ¬(x.abs <= 448.0) + +-- When fp32 exponent field is 0 and x ≠ 0, x is a fp32 subnormal. +-- fp32 subnormals are smaller than the smallest E4M3 subnormal (2^(-9)). +-- TODO: replace with a proof once Float32 bit-pattern lemmas are available. +axiom float32_exp_zero_abs_small (x : Float32) : + (x.toBits >>> 23 &&& 0xFF = 0) → x ≠ 0 → x.abs < Float32.ofBits 0x3B000000 + +-- `≤` and reverse-`<` are incompatible (proved from the IEEE compare model in LOrd). +theorem float32_le_lt_false (a b : Float32) : a <= b -> b < a -> False := + Float32.le_lt_false a b + +-- Triangle inequality for Float32 absolute value. +-- |a + b| ≤ |a| + |b|. Holds for all non-NaN finite Float32 values. +-- NaN excluded: NaN.abs is NaN and NaN ≤ anything is false. +axiom float32_abs_triangle (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + (a + b).abs ≤ a.abs + b.abs + +-- Monotonicity of ≤ for Float32 addition on non-NaN values. +-- If a ≤ b and c ≤ d then a + c <= b + d. +-- NaN excluded: NaN comparisons always return false. +axiom float32_add_le_add (a b c d : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) + (hc : c.isNaN = false) (hd : d.isNaN = false) : + a ≤ b → c ≤ d → a + c ≤ b + d + +-- Transitivity of ≤ for non-NaN Float32 values. +-- NaN excluded: NaN ≤ anything is false, so transitivity fails for NaN. +theorem float32_le_trans (a b c : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) (hc : c.isNaN = false) : + a ≤ b → b ≤ c → a ≤ c := + Float32.le_trans_of_not_nan a b c ha hb hc + +-- Distributivity: k * (a + b + c) = k * a + k * b + k * c for non-NaN values. +axiom float32_mul_add3 (k a b c : Float32) + (hk : k.isNaN = false) (ha : a.isNaN = false) + (hb : b.isNaN = false) (hc : c.isNaN = false) : + k * (a + b + c) = k * a + k * b + k * c + +-- Commutativity of Float32 addition for non-NaN values: a + b = b + a. +-- NaN excluded: NaN + b = NaN and b + NaN = NaN, but NaN != NaN in IEEE 754. +axiom float32_add_comm (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + a + b = b + a + +-- Add/sub cancellation for non-NaN Float32: (a + b) - b = a. +axiom float32_add_sub_cancel (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + a + b - b = a + +-- Left cancellation: a + b - a = b (i.e., the first summand cancels). +axiom float32_add_sub_cancel_left (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + a + b - a = b + +-- Subtraction decomposition for non-NaN values: +-- (a' + b') - (a + b) = (a' - a) + (b' - b). +-- True over reals; conditioned on non-NaN to avoid NaN propagation. +axiom float32_add_sub_decomp (a b a' b' : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) + (ha' : a'.isNaN = false) (hb' : b'.isNaN = false) : + (a' + b') - (a + b) = (a' - a) + (b' - b) + +-- Sub-additivity of abs for subtraction on non-NaN values: |a - c| <= |a - b| + |b - c|. +axiom float32_abs_sub_triangle (a b c : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) (hc : c.isNaN = false) : + (a - c).abs ≤ (a - b).abs + (b - c).abs + +-- Float32 addition/subtraction preserve non-NaN for non-NaN inputs. +-- In IEEE 754, a + b is NaN only if a or b is NaN (ignoring inf-inf). +-- For finite non-NaN values this always holds. +axiom float32_add_notNaN (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + (a + b).isNaN = false + +axiom float32_sub_notNaN (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : + (a - b).isNaN = false + +-- Helper: abs preserves isNaN for UnpackedFloat. +-- Direct case analysis — no Float32 machinery needed. +private lemma unpackedFloat_isNaN_abs_eq (u : Float.Model.UnpackedFloat) : + u.abs.isNaN = u.isNaN := by + cases u with + | notANumber => simp [Float.Model.UnpackedFloat.abs, Float.Model.UnpackedFloat.isNaN] + | infinity s | zero s => cases s <;> simp [Float.Model.UnpackedFloat.abs, Float.Model.UnpackedFloat.isNaN] + | finite s m e hm => cases s <;> simp [Float.Model.UnpackedFloat.abs, Float.Model.UnpackedFloat.isNaN] + +-- Float32.abs preserves non-NaN. +-- Proof: a.isNaN = a.toModel.unpack.isNaN (by definition), +-- and abs maps unpack to unpack.abs while preserving isNaN status +-- (notANumber→notANumber, infinity/zero/finite→positive variant). +-- The key step: (pack u.abs).unpack.isNaN = u.abs.isNaN, which for +-- notANumber/infinity/zero follows by rfl; for finite needs pack/unpack roundtrip. +-- Helper: (pack u).unpack.isNaN = u.isNaN for all UnpackedFloat u. +-- For notANumber/infinity/zero: rfl. For finite: pack never produces NaN bit pattern +-- (overflow → infinity, else → finite/zero with exponent ≠ 0xFF). +private lemma unpack_pack_isNaN (u : Float.Model.UnpackedFloat) : + (Float32.Model.pack u).unpack.isNaN = u.isNaN := by + cases u with + | notANumber => rfl + | infinity s | zero s => cases s <;> rfl + | finite s m e hm => + -- (Float32.Model.pack f).unpack reduces (by rfl, via UInt32.toBitVec_ofBitVec) to + -- UnpackedFloat.unpack binary32 (UnpackedFloat.pack binary32 f); and (finite ..).isNaN = false. + show (Float.Model.UnpackedFloat.unpack Float.Model.Format.binary32 + (Float.Model.UnpackedFloat.pack Float.Model.Format.binary32 + (.finite s m e hm))).isNaN = false + -- Only the `notANumber` constructor makes `isNaN` true; every other unpack result is + -- non-NaN by `rfl`. So we case on the unpack result and refute the notANumber case. + cases hu : Float.Model.UnpackedFloat.unpack Float.Model.Format.binary32 + (Float.Model.UnpackedFloat.pack Float.Model.Format.binary32 (.finite s m e hm)) with + | infinity _ => rfl + | zero _ => rfl + | finite _ _ _ _ => rfl + | notANumber => + -- `unpack = notANumber` requires exponentVec = -1 (all ones); but pack of a finite + -- value gives exponentVec = -1 only in the overflow branch, where mantissaVec = 0 + -- (yielding infinity, not NaN). So this case is impossible. + exfalso + revert hu + simp only [Float.Model.UnpackedFloat.pack] + split_ifs with hov hn <;> + intro hu <;> + simp only [Float.Model.UnpackedFloat.packedInfinity, + Float.Model.UnpackedFloat.unpack, + Float.Model.UnpackedFloat.unpackExponent_packComponents, + Float.Model.UnpackedFloat.unpackMantissa_packComponents] at hu <;> + (try split_ifs at hu with he1 he2) <;> + first + | contradiction -- leaf ≠ notANumber (distinct ctors), or he1 : ¬True + | (revert he2; decide) -- overflow mantissa leaf: he2 : ¬(0 = 0#23) is false + | (-- notANumber leaf: he1 : exponentVec = -1#8; refute (only overflow gives exp = -1) + exfalso + have htn := congr_arg BitVec.toNat he1 + simp only [Float.Model.Format.binary32, Float.Model.Format.exponentBias, + Float.Model.Format.mantissaBitsWithoutImplicit, Float.Model.Format.exponentBits, + BitVec.toNat_ofNat, BitVec.neg_one_eq_allOnes, BitVec.toNat_allOnes] at htn hov + omega) + +theorem float32_abs_notNaN (a : Float32) (ha : a.isNaN = false) : a.abs.isNaN = false := by + -- a.abs.isNaN = a.abs.toModel.unpack.isNaN (defeq) + -- = (pack a.toModel.unpack.abs).unpack.isNaN (a.abs.toModel = pack …, by rfl) + -- = a.toModel.unpack.abs.isNaN (unpack_pack_isNaN) + -- = a.toModel.unpack.isNaN (unpackedFloat_isNaN_abs_eq) + -- = false (ha, defeq) + show a.abs.toModel.unpack.isNaN = false + rw [show a.abs.toModel = Float32.Model.pack a.toModel.unpack.abs from rfl, + unpack_pack_isNaN, unpackedFloat_isNaN_abs_eq] + exact ha + +-- Float32 multiplication preserves non-NaN for non-NaN inputs. +axiom float32_mul_notNaN (a b : Float32) + (ha : a.isNaN = false) (hb : b.isNaN = false) : (a * b).isNaN = false + +-- Reflexivity of ≤ for non-NaN Float32 (NaN ≤ NaN is false in IEEE 754). +theorem float32_le_refl_notNaN (a : Float32) (ha : a.isNaN = false) : a ≤ a := + Float32.le_refl_of_not_nan a ha + +-- If a ≤ b holds and b is non-NaN, then a is non-NaN. +-- Proof: by contrapositive. If a.isNaN = true then a.toModel.unpack = notANumber, +-- so notANumber.compare b.toModel.unpack = none, Option.any isLE none = false, +-- meaning a ≤ b = false — contradicting the hypothesis. +set_option maxHeartbeats 400000 in +theorem float32_le_left_notNaN (a b : Float32) (hb : b.isNaN = false) : + a ≤ b → a.isNaN = false := by + intro h + by_contra ha + simp only [ne_eq, Bool.not_eq_false] at ha + -- ha : a.isNaN = true → a.toModel.unpack = notANumber + have h_u : a.toModel.unpack = .notANumber := by + cases h_uu : a.toModel.unpack with + | notANumber => rfl + | infinity s | zero s | finite s m e hm => + simp [Float32.isNaN, Float32.Model.isNaN, Float.Model.UnpackedFloat.isNaN, h_uu] at ha + -- a ≤ b is false when a.toModel.unpack = notANumber + -- use simp_all to reduce to h : false = true, then close by decide + simp_all only [LE.le, instLEFloat32, Float32.le, Float32.Model.le, + Float.Model.UnpackedFloat.le, Float.Model.UnpackedFloat.compare, + Option.any, decide_eq_true_eq, h_u] + exact absurd h (by decide) + +-- abs does not change NaN status: a.abs.isNaN = false implies a.isNaN = false. +-- Proof: by contrapositive. If a.isNaN = true then a.toModel.unpack = notANumber. +-- Then a.abs.toModel.unpack = (pack notANumber).unpack = notANumber (by rfl), +-- so a.abs.isNaN = true, contradicting the hypothesis. +theorem float32_notNaN_of_abs_notNaN (a : Float32) : a.abs.isNaN = false → a.isNaN = false := by + intro h + by_contra ha + -- ha : a.isNaN ≠ false, i.e., a.isNaN = true + simp only [ne_eq, Bool.not_eq_false] at ha + -- show a.abs.isNaN = true to contradict h + have h_abs_NaN : a.abs.isNaN = true := by + -- a.isNaN = true means a.toModel.unpack = notANumber + -- abs of notANumber = notANumber, and (pack notANumber).unpack = notANumber by rfl + simp only [Float32.isNaN, Float32.Model.isNaN, Float.Model.UnpackedFloat.isNaN] at ha ⊢ + -- ha says a.toModel.unpack.isNaN = true, i.e., = notANumber + have h_u : a.toModel.unpack = .notANumber := by + cases h_uu : a.toModel.unpack with + | notANumber => rfl + | infinity s | zero s | finite s m e hm => + simp [Float.Model.UnpackedFloat.isNaN, h_uu] at ha + -- a.abs.toModel.unpack = (pack notANumber).unpack = notANumber + have h_abs_u : a.abs.toModel.unpack = .notANumber := by + -- rewrite step by step to avoid simp expanding a.toModel.toBits + have heq : a.toModel.abs = Float32.Model.pack .notANumber := by + unfold Float32.Model.abs + rw [h_u]; simp [Float.Model.UnpackedFloat.abs] + rw [show a.abs.toModel = a.toModel.abs from rfl, heq] + rfl -- (pack notANumber).unpack = notANumber by definitional reduction + simp [Float32.isNaN, Float32.Model.isNaN, Float.Model.UnpackedFloat.isNaN, h_abs_u] + simp [h_abs_NaN] at h + +-- For a normal fp32 value (exp != 0, exp != 0xFF) in E4M3's range, +-- rounding introduces at most (1/2) * 2^(-3) * |x| error. +-- Proof requires bit-level case analysis on toFloat8E4M3Bits: +-- for E4M3 exponent e, adjacent values differ by 2^(e-10), +-- so error <= 2^(e-11) <= (1/2) * 2^(-3) * |x| since |x| >= 2^(e-7). +axiom float8E4M3_normal_rounding_bound (x : Float32) + (hexp : Not (x.toBits >>> 23 &&& 0xFF = 0xFF)) + (hexp0 : Not (x.toBits >>> 23 &&& 0xFF = 0)) + (hNoOverflow : x.abs <= 448.0) : + (x.toFloat8E4M3Bits.toFloat32FromFloat8E4M3 - x).abs <= + 0.5 * Float32.pow 2.0 (-Nat.toFloat32 3) * x.abs + +-- pointwise quantization error bound for fp8_e4m3 +-- Rounding a fp32 value x to the nearest fp8_e4m3 value introduces +-- at most (1/2) * ε * |x| error, where ε = 2^(-3) = 0.125. +theorem pointwiseBoundE4M3 (x : Float32) (hNoOverflow : x.abs ≤ (Dtype.fp8Max .float8_e4m3).getD 0) (hNoUnderflow : x = 0 ∨ Float32.ofBits 0x3B000000 <= x.abs): + ∃ err : Float32, Dtype.roundToComputeDtype x .float8_e4m3 = .ok (x + err) + ∧ err.abs <= 0.5 * (machineEpsilon .float8_e4m3).getD 0 * x.abs := by + simp [Dtype.roundToComputeDtype] + unfold Dtype.decodeFloat8E4M3 + simp [ByteArray.size] + -- provide the witness: the rounding error is the difference between rounded and original + apply Exists.intro (x.toFloat8E4M3Bits.toFloat32FromFloat8E4M3 - x) + apply And.intro + · exact float32_eq_add_sub _ _ + · simp [machineEpsilon, mantissaBits] + -- case 1: x = 0, round-trip error is 0, bound holds trivially + by_cases hx : x = 0 + · subst hx + native_decide + · set exp := (x.toBits >>> 23) &&& 0xFF + -- case 2: x is NaN or ±inf (exp = 0xFF), contradicts hNoOverflow + by_cases hexp : exp = 0xFF + · exact absurd (by simpa [Dtype.fp8Max] using hNoOverflow) (float32_exp_ff_not_finite x hexp) + · -- case 3: x is a fp32 subnormal (exp = 0, x ≠ 0) + -- contradicts hNoUnderflow since fp32 subnormals are below E4M3's smallest value + by_cases hexp0 : exp = 0 + · exfalso + cases hNoUnderflow with + | inl h => exact hx h + | inr h => + have hsmall := float32_exp_zero_abs_small x hexp0 hx + exact float32_le_lt_false _ _ h hsmall + · -- case 4: x is a normal fp32 value in E4M3's representable range + -- exp != 0xFF (not NaN/inf), exp != 0 (not subnormal) + -- so x is a finite normal fp32 value with |x| in [2^(-9), 448.0] + exact float8E4M3_normal_rounding_bound x hexp hexp0 + (by simpa [Dtype.fp8Max] using hNoOverflow) + +-- Generic axioms for all fp8 compute dtypes. +-- These replace the e4m3-specific axioms above and work for any dtype +-- where fp8Max is defined (i.e., float8_e4m3, float8_e5m2, float8_e3m4). + +-- exp = 0xFF means NaN or ±inf, which is not ≤ any finite fp8 max. +-- Proof: exp = 0xFF → x.toModel.unpack = infinity _ or notANumber. +-- For infinity: infinity.le (finite v) = false. +-- For NaN: NaN.le anything = false. +-- Both follow from Float.Model.UnpackedFloat.compare semantics. +-- TODO: extract from Float32.toModel.unpack definition in Lean 4.33. +-- Helper: when exp = 0xFF, x.abs.toModel.unpack is NaN or +infinity. +-- We show this from the unpack definition: exponentVec = -1#_ → infinity or notANumber. +-- Then use abs semantics: abs maps both to notANumber or infinity .positive. +-- Helper: x.toModel.unpack is notANumber or infinity when exp field = 0xFF. +-- Proof: unpackExponent = -1#8 (all 1s) forces the infinity/NaN branch of unpack. +-- The bit manipulation uses UInt32.toBitVec_shiftRight/and to convert hexp to BitVec, +-- then bv_omega closes the extractLsb = -1#8 goal. +private lemma toModel_unpack_of_exp_ff (x : Float32) + (hexp : x.toBits >>> 23 &&& 0xFF = 0xFF) : + x.toModel.unpack = .notANumber ∨ ∃ s, x.toModel.unpack = .infinity s := by + -- exp = 0xFF → unpackExponent (bits 30..23) = -1#8 (all 1s) + -- TODO: the UInt32 → BitVec.extractLsb connection is: + -- x.toBits >>> 23 &&& 0xFF = 0xFF (UInt32) + -- ↔ x.toModel.toBits.toBitVec.extractLsb 30 23 = -1#8 (BitVec 8) + -- This equivalence needs a lemma like UInt32.extractLsb_eq or omega after BitVec unfolding. + have h_exp : x.toModel.toBits.toBitVec.extractLsb 30 23 = -1#8 := by + -- Push the UInt32 hypothesis down to a BitVec equation, then let bv_decide + -- discharge the extractLsb claim (bits 23..30 of the 32-bit pattern are all ones). + have hbv : (x.toBits.toBitVec >>> 23) &&& 0xFF#32 = 0xFF#32 := by + simpa using congrArg UInt32.toBitVec hexp + show x.toBits.toBitVec.extractLsb 30 23 = -1#8 + bv_decide + -- Case split on the actual unpack result; zero/finite are impossible since exp = 0xFF + cases h_u : x.toModel.unpack with + | notANumber => left; rfl + | infinity s => right; exact ⟨s, rfl⟩ + | zero _ | finite _ _ _ _ => + -- unpackExponent = -1#8 forces infinity/NaN branch, contradicting zero/finite + exfalso + simp only [Float32.Model.unpack, Float.Model.UnpackedFloat.unpack, + Float.Model.UnpackedFloat.unpackExponent, Float.Model.Format.binary32, + Float32.toBits, h_exp, BitVec.cast_eq, ite_true] at h_u + split_ifs at h_u <;> simp at h_u + +-- Helper: when exp = 0xFF, x.abs.toModel.unpack is notANumber or infinity .positive. +-- abs maps notANumber → notANumber and infinity _ → infinity .positive. +-- The pack/unpack roundtrip holds by rfl (definitional equality in Lean 4.33). +private lemma abs_unpack_of_exp_ff (x : Float32) + (hexp : x.toBits >>> 23 &&& 0xFF = 0xFF) : + x.abs.toModel.unpack = .notANumber ∨ + x.abs.toModel.unpack = .infinity .positive := by + -- Float32 is a structure: (ofModel m).toModel = m, so x.abs.toModel = x.toModel.abs + have h_abs : x.abs.toModel = x.toModel.abs := rfl + rcases toModel_unpack_of_exp_ff x hexp with h_u | ⟨s, h_u⟩ + · -- notANumber.abs = notANumber; (pack notANumber).unpack = notANumber by rfl + left + show (x.toModel.abs).unpack = .notANumber + rw [show x.toModel.abs = Float32.Model.pack .notANumber by + simp [Float32.Model.abs, h_u, Float.Model.UnpackedFloat.abs]] + rfl + · -- (infinity s).abs = infinity .positive; (pack (infinity .positive)).unpack = it by rfl + right + show (x.toModel.abs).unpack = .infinity .positive + rw [show x.toModel.abs = Float32.Model.pack (.infinity .positive) by + simp [Float32.Model.abs, h_u, Float.Model.UnpackedFloat.abs]] + rfl + +theorem fp8_exp_ff_not_finite (dtype : Dtype) (x : Float32) + (hFp8 : (Dtype.fp8Max dtype).isSome) + (hexp : x.toBits >>> 23 &&& 0xFF = 0xFF) : + Not (x.abs ≤ (Dtype.fp8Max dtype).getD 0) := by + rcases dtype with _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ + all_goals simp [Dtype.fp8Max] at hFp8 ⊢ + -- For each fp8 dtype, show x.abs ≤ threshold is false + all_goals ( + intro h + -- x.abs is NaN or +infinity (from exp = 0xFF) + -- unfold ≤ and substitute to get a false=true contradiction + rcases abs_unpack_of_exp_ff x hexp with h_u | h_u + · -- notANumber.compare _ = none → Option.any isLE none = false + simp only [LE.le, instLEFloat32, Float32.le, decide_eq_true_eq, + Float32.Model.le, Float.Model.UnpackedFloat.le, + Float.Model.UnpackedFloat.compare, Option.any, h_u] at h + exact absurd h (by decide) + · -- (infinity .positive).compare (finite v) = some .gt → isLE = false + simp only [LE.le, instLEFloat32, Float32.le, decide_eq_true_eq, + Float32.Model.le, Float.Model.UnpackedFloat.le, + Float.Model.UnpackedFloat.compare, Option.any, Ordering.isLE, h_u] at h + exact absurd h (by decide)) + +-- pack/unpack round-trips on a fp32 subnormal: packing `finite positive m (-149)` +-- (with 0 < m < 2^23) takes the subnormal branch of `pack` and unpacks back unchanged. +private lemma subnormal_unpack_roundtrip (m : Nat) (hm : 0 < m) (hlt : m < 2 ^ 23) : + (Float32.Model.pack (.finite .positive m (-149) hm)).unpack + = .finite .positive m (-149) hm := by + have hlog : m.log2 < 23 := (Nat.log2_lt (by omega)).mpr hlt + have hmod : m % 2 ^ 23 = m := Nat.mod_eq_of_lt hlt + have hmv : BitVec.ofNat 23 m ≠ 0#23 := by + rw [Ne, ← BitVec.toNat_inj] + simp only [BitVec.toNat_ofNat, BitVec.toNat_zero, hmod]; omega + simp only [Float32.Model.pack, Float32.Model.unpack, UInt32.toBitVec_ofBitVec, + Float.Model.UnpackedFloat.pack, Float.Model.Format.binary32, + Float.Model.Format.exponentBias, Float.Model.Format.mantissaBits] + -- Resolve `pack`'s branches: not overflow (biased exponent is 1), not normal (m.log2+1 ≤ 23). + have hnm : ¬ (m.log2 + 1 = 1 + 23) := by omega + -- Full `simp` evaluates the closed overflow condition, uses hnm for the normal branch, + -- applies the packComponents unpack lemmas, and reduces the mantissa via hmod. + simp [Float.Model.UnpackedFloat.unpack, + Float.Model.UnpackedFloat.unpackExponent_packComponents, + Float.Model.UnpackedFloat.unpackMantissa_packComponents, + hnm, dif_neg hmv, hmod] + -- Remaining: sign round-trips to positive, m < 2^23, and the exponent is -149. + refine ⟨?_, ?_, ?_⟩ + · -- the sign bit of the packed value is 0 (independent of the mantissa) + have hsign : Float.Model.UnpackedFloat.unpackSign + (Float.Model.UnpackedFloat.packComponents Float.Model.Format.binary32 + .positive 0#8 (BitVec.ofNat 23 m)) = 0#1 := by + unfold Float.Model.UnpackedFloat.unpackSign Float.Model.UnpackedFloat.packComponents + simp only [Float.Model.Format.binary32, Float.Model.UnpackedFloat.Sign.toBitVec] + bv_decide + simp [Float.Model.UnpackedFloat.Sign.ofBitVec, hsign] + · omega + · simp only [Float.Model.Format.binary32, Float.Model.Format.exponentBias]; decide + +-- When the fp32 exponent field is 0, `x` unpacks to either a (signed) zero or a +-- subnormal finite value with model exponent -149 and mantissa in [1, 2^23). +private lemma unpack_of_exp_zero (x : Float32) (hexp : x.toBits >>> 23 &&& 0xFF = 0) : + (∃ s, x.toModel.unpack = .zero s) ∨ + (∃ s m, ∃ hm : 0 < m, m < 2 ^ 23 ∧ x.toModel.unpack = .finite s m (-149) hm) := by + -- exponent field 0 → the unpacked exponent bitvector is 0#8 + have hev : Float.Model.UnpackedFloat.unpackExponent (spec := Float.Model.Format.binary32) + x.toModel.toBits.toBitVec = 0#8 := by + have hbv : (x.toBits.toBitVec >>> 23) &&& 0xFF#32 = 0#32 := by + simpa using congrArg UInt32.toBitVec hexp + unfold Float.Model.UnpackedFloat.unpackExponent + simp only [Float.Model.Format.binary32] + show ((x.toBits.toBitVec.extractLsb 30 23).cast _ : BitVec 8) = 0#8 + bv_decide + -- the mantissa bitvector has toNat < 2^23 (it is a BitVec 23) + set mv := Float.Model.UnpackedFloat.unpackMantissa (spec := Float.Model.Format.binary32) + x.toModel.toBits.toBitVec with hmvdef + by_cases hmant : mv = 0#23 + · left + refine ⟨Float.Model.UnpackedFloat.Sign.ofBitVec + (Float.Model.UnpackedFloat.unpackSign (spec := Float.Model.Format.binary32) + x.toModel.toBits.toBitVec), ?_⟩ + show Float.Model.UnpackedFloat.unpack Float.Model.Format.binary32 x.toModel.toBits.toBitVec = _ + unfold Float.Model.UnpackedFloat.unpack + simp only [Float.Model.Format.binary32, ← hmvdef, hev, hmant, + show ((0#8 : BitVec 8) = -1#8) = False from by decide, if_false, + show ((0#8 : BitVec 8) = 0#8) = True from by decide, if_true, dif_pos] + · right + have hpos : 0 < mv.toNat := + Nat.pos_of_ne_zero (fun h => hmant (BitVec.toNat_inj.mp (by simpa using h))) + refine ⟨Float.Model.UnpackedFloat.Sign.ofBitVec + (Float.Model.UnpackedFloat.unpackSign (spec := Float.Model.Format.binary32) + x.toModel.toBits.toBitVec), mv.toNat, hpos, mv.isLt, ?_⟩ + show Float.Model.UnpackedFloat.unpack Float.Model.Format.binary32 x.toModel.toBits.toBitVec = _ + unfold Float.Model.UnpackedFloat.unpack + simp only [Float.Model.Format.binary32, Float.Model.Format.exponentBias, ← hmvdef, hev, + show ((0#8 : BitVec 8) = -1#8) = False from by decide, if_false, + show ((0#8 : BitVec 8) = 0#8) = True from by decide, if_true, dif_neg hmant] + congr 1 + +-- A +0 or a positive subnormal (model exponent -149) compares strictly below any +-- positive finite value whose model exponent E exceeds -149. +private lemma unpack_lt_fp8min (u : Float.Model.UnpackedFloat) (M : Nat) (E : Int) (hM : 0 < M) + (hE : (-149 : Int) < E) + (hu : u = .zero .positive ∨ ∃ m, ∃ hm : 0 < m, u = .finite .positive m (-149) hm) : + (u.compare (.finite .positive M E hM) == some Ordering.lt) = true := by + rcases hu with h | ⟨m, hm, h⟩ <;> subst h <;> + simp [Float.Model.UnpackedFloat.compare, Int.compare_eq_lt.mpr hE, Ordering.then] + +-- fp32 subnormals have value < 2^(-126), which is smaller than the minimum +-- of any fp8 format (smallest is 2^(-16) for e5m2). +theorem fp8_exp_zero_abs_small (dtype : Dtype) (x : Float32) + (hFp8 : (Dtype.fp8Max dtype).isSome) + (hexp : x.toBits >>> 23 &&& 0xFF = 0) (hx : x ≠ 0) : + x.abs < (Dtype.fp8Min dtype).getD 0 := by + -- |x| unpacks to +0 or a positive subnormal (model exponent -149), via abs + the roundtrip. + have habs : x.abs.toModel.unpack = .zero .positive ∨ + ∃ m, ∃ hm : 0 < m, x.abs.toModel.unpack = .finite .positive m (-149) hm := by + rcases unpack_of_exp_zero x hexp with ⟨s, hz⟩ | ⟨s, m, hm, hlt, hf⟩ + · left + show (Float32.Model.pack x.toModel.unpack.abs).unpack = _ + rw [hz]; simp only [Float.Model.UnpackedFloat.abs]; rfl + · right + refine ⟨m, hm, ?_⟩ + show (Float32.Model.pack x.toModel.unpack.abs).unpack = _ + rw [hf]; simp only [Float.Model.UnpackedFloat.abs] + exact subnormal_unpack_roundtrip m hm hlt + -- Reduce `<` to a comparison on the unpacked values. + show Float32.lt x.abs ((Dtype.fp8Min dtype).getD 0) = true + unfold Float32.lt + simp only [decide_eq_true_eq] + show Float32.Model.lt x.abs.toModel ((Dtype.fp8Min dtype).getD 0).toModel = true + unfold Float32.Model.lt Float.Model.UnpackedFloat.lt + -- Only the three fp8 compute dtypes have `fp8Max`; discharge the rest via hFp8. + rcases dtype with _|_|_|_|_|_|_|_|_|_|_|_|_|_|_|_|_ <;> + simp only [Dtype.fp8Max] at hFp8 <;> + (try exact absurd hFp8 (by decide)) + -- Three fp8 dtypes remain; each fp8Min unpacks to a positive finite with exponent > -149. + · -- float8_e4m3: 2^(-9), model exponent -32 + simp only [Dtype.fp8Min, Option.getD] + rw [show (Float32.ofBits 0x3B000000).toModel.unpack + = .finite .positive 8388608 (-32) (by norm_num) from rfl] + exact unpack_lt_fp8min _ _ _ _ (by decide) habs + · -- float8_e3m4: 2^(-6), model exponent -29 + simp only [Dtype.fp8Min, Option.getD] + rw [show (Float32.ofBits 0x3C800000).toModel.unpack + = .finite .positive 8388608 (-29) (by norm_num) from rfl] + exact unpack_lt_fp8min _ _ _ _ (by decide) habs + · -- float8_e5m2: 2^(-16), model exponent -39 + simp only [Dtype.fp8Min, Option.getD] + rw [show (Float32.ofBits 0x37800000).toModel.unpack + = .finite .positive 8388608 (-39) (by norm_num) from rfl] + exact unpack_lt_fp8min _ _ _ _ (by decide) habs + +-- For a normal fp32 value in any fp8 format's representable range, +-- rounding introduces at most (1/2) * machineEpsilon * |x| error. +-- Proof requires bit-level case analysis on the encode/decode round-trip +-- for each fp8 format. The key: for normal values at exponent e, +-- adjacent fp8 values differ by 2^(e - mantissaBits), so rounding +-- error ≤ half that gap ≤ (1/2) * 2^(-mantissaBits) * |x|. +-- TODO: requires format-specific bit-level reasoning. +axiom fp8_normal_rounding_bound (dtype : Dtype) (x : Float32) + (hexp : Not (x.toBits >>> 23 &&& 0xFF = 0xFF)) + (hexp0 : Not (x.toBits >>> 23 &&& 0xFF = 0)) + (hNoOverflow : x.abs ≤ (Dtype.fp8Max dtype).getD 0) : + ∃ err : Float32, Dtype.roundToComputeDtype x dtype = .ok (x + err) ∧ + err.abs ≤ 0.5 * (machineEpsilon dtype).getD 0 * x.abs + +-- Generic pointwise quantization error bound for all fp8 compute dtypes. +-- For any x that fits in dtype's representable range, rounding x to dtype +-- introduces at most (1/2) * epsilon * |x| absolute error. +-- hFp8 restricts to dtypes where fp8Max is defined (the three fp8 types). +theorem pointwiseBound (dtype : Dtype) (x : Float32) + (hFp8 : (Dtype.fp8Max dtype).isSome) + (hNoOverflow : x.abs ≤ (Dtype.fp8Max dtype).getD 0) + (hNoUnderflow : x = 0 ∨ (Dtype.fp8Min dtype).getD 0 ≤ x.abs) : + ∃ err : Float32, Dtype.roundToComputeDtype x dtype = .ok (x + err) ∧ + err.abs ≤ 0.5 * (machineEpsilon dtype).getD 0 * x.abs := by + by_cases hx : x = 0 + · -- x = 0: rounding error is 0, bound holds trivially + -- provide witness err = 0 explicitly, then use native_decide for each fp8 dtype + subst hx + refine ⟨0, ?_, ?_⟩ + · -- roundToComputeDtype 0 dtype = .ok (0 + 0) = .ok 0 + rcases dtype with _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ + all_goals (simp [Dtype.fp8Max] at hFp8; try native_decide) + · -- (0 : Float32).abs = 0, so bound 0 ≤ 0.5 * epsilon * 0 holds + rcases dtype with _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ + all_goals (simp [Dtype.fp8Max] at hFp8; try native_decide) + · set exp := (x.toBits >>> 23) &&& 0xFF + -- case 2: x is NaN or ±inf, contradicts hNoOverflow + by_cases hexp : exp = 0xFF + · exact absurd (by simpa [Dtype.fp8Max] using hNoOverflow) + (fp8_exp_ff_not_finite dtype x hFp8 hexp) + · -- case 3: x is a fp32 subnormal, contradicts hNoUnderflow + by_cases hexp0 : exp = 0 + · exfalso + cases hNoUnderflow with + | inl h => exact hx h + | inr h => + exact float32_le_lt_false _ _ h + (fp8_exp_zero_abs_small dtype x hFp8 hexp0 hx) + · -- case 4: x is a normal fp32 in dtype's range, use generic axiom + exact fp8_normal_rounding_bound dtype x hexp hexp0 + (by simpa [Dtype.fp8Max] using hNoOverflow) + +-- Higham-style addition error bound for fp8 dtypes. +-- Given a, b and their fp8 approximations a', b' with individual +-- quantization errors bounded by (u/2)|a| and (u/2)|b|, computing +-- roundToComputeDtype on the Float32 sum a' + b' gives a result within +-- (u/2)(|a| + |b| + |a'+b'|) of the true sum a + b. +-- where u = machineEpsilon dtype. +-- Proof structure: +-- result - (a+b) = err + (a'-a) + (b'-b) [algebra] +-- |result - (a+b)| ≤ |err| + |a'-a| + |b'-b| [triangle inequality] +-- ≤ (u/2)|a'+b'| + (u/2)|a| + (u/2)|b| [pointwiseBound + ha' + hb'] +-- = (u/2)(|a| + |b| + |a'+b'|) [distributivity] +set_option maxHeartbeats 800000 in +theorem additionErrorBound (dtype : Dtype) (hFp8 : (Dtype.fp8Max dtype).isSome) + (a b a' b' : Float32) + -- non-NaN conditions: needed for the order and arithmetic axioms to hold + (ha_nn : a.isNaN = false) (hb_nn : b.isNaN = false) + (ha'_nn : a'.isNaN = false) (hb'_nn : b'.isNaN = false) + -- a', b' are fp8 approximations of a, b with pointwise quantization error + (ha' : (a' - a).abs ≤ 0.5 * (machineEpsilon dtype).getD 0 * a.abs) + (hb' : (b' - b).abs ≤ 0.5 * (machineEpsilon dtype).getD 0 * b.abs) + -- the sum a' + b' fits in dtype's representable range + (hNoOverflow : (a' + b').abs ≤ (Dtype.fp8Max dtype).getD 0) + (hNoUnderflow : a' + b' = 0 ∨ + (Dtype.fp8Min dtype).getD 0 ≤ (a' + b').abs) : + ∃ result : Float32, + Dtype.roundToComputeDtype (a' + b') dtype = .ok result ∧ + (result - (a + b)).abs ≤ + 0.5 * (machineEpsilon dtype).getD 0 * (a.abs + b.abs + (a' + b').abs) := by + -- apply pointwiseBound to get the re-encoding error + obtain ⟨err, hresult, herr⟩ := pointwiseBound dtype (a' + b') hFp8 hNoOverflow hNoUnderflow + -- result = (a' + b') + err, so result - (a+b) = (a'-a) + (b'-b) + err + refine ⟨a' + b' + err, hresult, ?_⟩ + -- Get non-NaN witnesses + have hab'_nn : (a' + b').isNaN = false := float32_add_notNaN _ _ ha'_nn hb'_nn + have hab_nn : (a + b).isNaN = false := float32_add_notNaN _ _ ha_nn hb_nn + have ha'a_nn : (a' - a).isNaN = false := float32_sub_notNaN _ _ ha'_nn ha_nn + have hb'b_nn : (b' - b).isNaN = false := float32_sub_notNaN _ _ hb'_nn hb_nn + have hab'_abs_pre : (a' + b').abs.isNaN = false := float32_abs_notNaN _ hab'_nn + have hu_nn_pre : (0.5 * (machineEpsilon dtype).getD 0).isNaN = false := by + rcases dtype with _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ | _ + all_goals (simp [Dtype.fp8Max] at hFp8; try native_decide) + have hu_rhs_nn : (0.5 * (machineEpsilon dtype).getD 0 * (a' + b').abs).isNaN = false := + float32_mul_notNaN _ _ hu_nn_pre hab'_abs_pre + have herr_abs_nn : err.abs.isNaN = false := + float32_le_left_notNaN _ _ hu_rhs_nn herr + have herr_nn : err.isNaN = false := + float32_notNaN_of_abs_notNaN _ herr_abs_nn + have hres_nn : (a' + b' + err).isNaN = false := float32_add_notNaN _ _ hab'_nn herr_nn + have ha'a_abs := float32_abs_notNaN _ ha'a_nn + have hb'b_abs := float32_abs_notNaN _ hb'b_nn + have ha_abs := float32_abs_notNaN _ ha_nn + have hb_abs := float32_abs_notNaN _ hb_nn + have hab'_abs := float32_abs_notNaN _ hab'_nn + have hu_nn := hu_nn_pre + -- |result - (a+b)| ≤ |result - (a'+b')| + |(a'+b') - (a+b)| + have htri := float32_abs_sub_triangle (a' + b' + err) (a' + b') (a + b) + hres_nn hab'_nn hab_nn + -- rewrite |result - (a'+b')| = |err| via left cancellation + -- (a'+b') + err - (a'+b') = err via float32_add_sub_cancel_left + have hcancel := float32_add_sub_cancel_left (a' + b') err hab'_nn herr_nn + rw [hcancel] at htri + -- step 3: |(a'+b') - (a+b)| ≤ |a'-a| + |b'-b| + have hdiff : ((a' + b') - (a + b)).abs ≤ (a' - a).abs + (b' - b).abs := by + rw [float32_add_sub_decomp _ _ _ _ ha_nn hb_nn ha'_nn hb'_nn] + exact float32_abs_triangle _ _ ha'a_nn hb'b_nn + -- |result-(a+b)| ≤ err.abs + (|a'-a| + |b'-b|) + have hdiff_abs := float32_abs_notNaN _ (float32_sub_notNaN _ _ hab'_nn hab_nn) + have hstep4 : (a' + b' + err - (a + b)).abs ≤ err.abs + ((a' - a).abs + (b' - b).abs) := + float32_le_trans _ _ _ + (float32_abs_notNaN _ (float32_sub_notNaN _ _ hres_nn hab_nn)) + (float32_add_notNaN _ _ herr_abs_nn hdiff_abs) + (float32_add_notNaN _ _ herr_abs_nn (float32_add_notNaN _ _ ha'a_abs hb'b_abs)) + htri + (float32_add_le_add _ _ _ _ herr_abs_nn herr_abs_nn hdiff_abs + (float32_add_notNaN _ _ ha'a_abs hb'b_abs) + (float32_le_refl_notNaN _ herr_abs_nn) hdiff) + -- bound each term individually + have hbound : err.abs + ((a' - a).abs + (b' - b).abs) ≤ + 0.5 * (machineEpsilon dtype).getD 0 * (a' + b').abs + + (0.5 * (machineEpsilon dtype).getD 0 * a.abs + + 0.5 * (machineEpsilon dtype).getD 0 * b.abs) := + float32_add_le_add _ _ _ _ + herr_abs_nn (float32_mul_notNaN _ _ hu_nn hab'_abs) + (float32_add_notNaN _ _ ha'a_abs hb'b_abs) + (float32_add_notNaN _ _ (float32_mul_notNaN _ _ hu_nn ha_abs) + (float32_mul_notNaN _ _ hu_nn hb_abs)) + herr + (float32_add_le_add _ _ _ _ ha'a_abs (float32_mul_notNaN _ _ hu_nn ha_abs) + hb'b_abs (float32_mul_notNaN _ _ hu_nn hb_abs) ha' hb') + -- rearrange using distributivity and commutativity + have hsum_nn : (a.abs + b.abs + (a' + b').abs).isNaN = false := + float32_add_notNaN _ _ (float32_add_notNaN _ _ ha_abs hb_abs) hab'_abs + have hmul_nn := float32_mul_notNaN _ _ hu_nn hsum_nn + have hrearr : 0.5 * (machineEpsilon dtype).getD 0 * (a' + b').abs + + (0.5 * (machineEpsilon dtype).getD 0 * a.abs + + 0.5 * (machineEpsilon dtype).getD 0 * b.abs) = + 0.5 * (machineEpsilon dtype).getD 0 * (a.abs + b.abs + (a' + b').abs) := by + rw [float32_mul_add3 _ _ _ _ hu_nn ha_abs hb_abs hab'_abs] + rw [float32_add_comm (0.5 * (machineEpsilon dtype).getD 0 * (a' + b').abs) + (0.5 * (machineEpsilon dtype).getD 0 * a.abs + + 0.5 * (machineEpsilon dtype).getD 0 * b.abs) + (float32_mul_notNaN _ _ hu_nn hab'_abs) + (float32_add_notNaN _ _ (float32_mul_notNaN _ _ hu_nn ha_abs) + (float32_mul_notNaN _ _ hu_nn hb_abs))] + -- chain all steps + rw [hrearr] at hbound + exact float32_le_trans _ _ _ + (float32_abs_notNaN _ (float32_sub_notNaN _ _ hres_nn hab_nn)) + (float32_add_notNaN _ _ herr_abs_nn (float32_add_notNaN _ _ ha'a_abs hb'b_abs)) + hmul_nn + hstep4 hbound diff --git a/TensorLib/Npy.lean b/TensorLib/Npy.lean index 6e938f7..e9de30f 100644 --- a/TensorLib/Npy.lean +++ b/TensorLib/Npy.lean @@ -123,7 +123,7 @@ def dtypeNameToNpyString (t : TensorLib.Dtype) : String := match t with def fromNpyString (s : String) : Err Dtype := if s.length == 0 then .error "Empty dtype string" else do - let order <- ByteOrder.fromChar (s.get 0) + let order <- ByteOrder.fromChar s.front let nameStr := s.drop 1 -- bf16 stored as "