Radix-4 forward butterfly

ntt_roots_correct, the radix2Pass correctness lemmas, and the forward radix-4 butterfly (butterfly4_forward_*) culminating in butterfly4_forward_ZMod_combined. Split out of the former SolutionHelpers.lean.

def ntt_roots_correct (m : ℕ) (roots : Vector UInt32 m) : Prop

A root table is correct for an m-point NTT when every entry at index len/2 + j holds the twiddle factor ω_{len}^j in the Montgomery domain (multiplied by R), where ω_{len} = primRoot ^ ((mod64 - 1) / len) is a primitive len-th root of unity and R = montR1 = 2^32 mod mod64.

Equations

• One or more equations did not get rendered due to their size.

theorem pow2_dvd_is_pow2 (len N : ℕ) (hdvd : len ∣ 2 ^ N) (hlen : 2 ≤ len) : ∃ (k : ℕ), len = 2 ^ (k + 1)

theorem ensure_roots_zmod_cast (k j : ℕ) : have w := primRoot.toNat ^ ((mod64.toNat - 1) / 2 ^ (k + 1)) % mod64.toNat; ↑(w ^ j * montR1.toNat % mod64.toNat) = ↑primRoot.toNat ^ ((mod64.toNat - 1) / 2 ^ (k + 1) * j) * ↑montR1.toNat

theorem ensure_roots_ntt_correct {m : ℕ} (hm_pow2 : m.isPowerOfTwo) (hm_dvd : m ∣ mod64.toNat - 1) : ntt_roots_correct m (ensureRoots m)

ensureRoots m produces a correct NTT root table for an m-point NTT.

theorem to_mont_dft_factor {m : ℕ} (v : Vector UInt32 m) (hv_bound : (v.all fun (x : UInt32) => decide (x < mod32)) = true) (ω : ZMod mod32.toNat) (k : ℕ) : ∑ j : Fin m, ↑(toMont v[↑j]).toNat * ω ^ (↑j * k) = ↑montR1.toNat * ∑ j : Fin m, ↑v[↑j].toNat * ω ^ (↑j * k)

theorem radix2Pass_preserves {n : ℕ} (k i : ℕ) (a : Vector UInt32 n) (idx : ℕ) (hidx : idx < n) (h_below : idx < 2 * i ∨ idx ≥ 2 * (i + k)) : (radix2Pass k i a).get ⟨idx, hidx⟩ = a.get ⟨idx, hidx⟩

theorem radix2Pass_get_lo {n : ℕ} (k i : ℕ) (a : Vector UInt32 n) (p : ℕ) (hp : 2 * p + 1 < n) (hpi : i ≤ p) (hpk : p < i + k) (hk : i + k ≤ n / 2) : (radix2Pass k i a).get ⟨2 * p, ⋯⟩ = addMod32 (a.get ⟨2 * p, ⋯⟩) (a.get ⟨2 * p + 1, hp⟩)

theorem radix2Pass_get_hi {n : ℕ} (k i : ℕ) (a : Vector UInt32 n) (p : ℕ) (hp : 2 * p + 1 < n) (hpi : i ≤ p) (hpk : p < i + k) (hk : i + k ≤ n / 2) : (radix2Pass k i a).get ⟨2 * p + 1, hp⟩ = subMod32 (a.get ⟨2 * p, ⋯⟩) (a.get ⟨2 * p + 1, hp⟩)

theorem radix2Pass_ZMod_pair {n : ℕ} (a : Vector UInt32 n) (ha : (a.all fun (x : UInt32) => decide (x < mod32)) = true) (p : ℕ) (hp : 2 * p + 1 < n) : have result := radix2Pass (n / 2) 0 a; ↑(result.get ⟨2 * p, ⋯⟩).toNat = ↑(a.get ⟨2 * p, ⋯⟩).toNat + ↑(a.get ⟨2 * p + 1, hp⟩).toNat ∧ ↑(result.get ⟨2 * p + 1, hp⟩).toNat = ↑(a.get ⟨2 * p, ⋯⟩).toNat - ↑(a.get ⟨2 * p + 1, hp⟩).toNat

theorem ω_is_primitive_half {m : UInt64} (n : ℕ) (hm_eq : m.toNat = 2 ^ n) (hm_dvd : m.toNat ∣ mod64.toNat - 1) : have ω := ↑primRoot.toNat ^ ((mod64.toNat - 1) / m.toNat); n = 0 ∨ ω ^ 2 ^ (n - 1) = -1

theorem ω_is_primitive_half_nat (m n : ℕ) (hm_eq : m = 2 ^ n) (hm_dvd : m ∣ mod64.toNat - 1) : n = 0 ∨ (↑primRoot.toNat ^ ((mod64.toNat - 1) / m)) ^ 2 ^ (n - 1) = -1

Nat version of ω_is_primitive_half: works when m : ℕ instead of UInt64.

theorem ntt_roots_correct_at {N : ℕ} (roots : Vector UInt32 N) (hroots : ntt_roots_correct N roots) (len j idx : ℕ) (h2 : 2 ≤ len) (hdvd : len ∣ N) (hj : j < len / 2) (hidx : idx < N) (heq : idx = len / 2 + j) : ↑roots[idx].toNat = ↑primRoot.toNat ^ ((mod64.toNat - 1) / len * j) * ↑montR1.toNat

theorem vector_getD_eq_getElem {α : Type u_1} {N : ℕ} (v : Vector α N) (idx : ℕ) (h : idx < N) (default : α) : v.getD idx default = v[idx]

theorem butterfly4_forward_getElem_pos0 {N : ℕ} (roots a : Vector UInt32 N) (s len i2 j2 : ℕ) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (ht1 : s + j2 < N) (ht2 : len + j2 < N) (ht3 : len + j2 + s < N) (hne10 : i2 + j2 + s ≠ i2 + j2) (hne20 : i2 + len + j2 ≠ i2 + j2) (hne30 : i2 + len + j2 + s ≠ i2 + j2) : have t1 := roots[s + j2]; have t2 := roots[len + j2]; have aB := montMul a[i2 + j2 + s] t1; have aD := montMul a[i2 + len + j2 + s] t1; have P := addMod32 a[i2 + j2] aB; have R_v := addMod32 a[i2 + len + j2] aD; have t2R := montMul t2 R_v; (butterfly4 a false roots s len i2 j2)[i2 + j2] = addMod32 P t2R

theorem butterfly4_forward_getElem_pos2 {N : ℕ} (roots a : Vector UInt32 N) (s len i2 j2 : ℕ) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (ht1 : s + j2 < N) (ht2 : len + j2 < N) (ht3 : len + j2 + s < N) (hne12 : i2 + j2 + s ≠ i2 + len + j2) (hne32 : i2 + len + j2 + s ≠ i2 + len + j2) : have t1 := roots[s + j2]; have t2 := roots[len + j2]; have aB := montMul a[i2 + j2 + s] t1; have aD := montMul a[i2 + len + j2 + s] t1; have P := addMod32 a[i2 + j2] aB; have R_v := addMod32 a[i2 + len + j2] aD; have t2R := montMul t2 R_v; (butterfly4 a false roots s len i2 j2)[i2 + len + j2] = subMod32 P t2R

theorem butterfly4_forward_getElem_pos1 {N : ℕ} (roots a : Vector UInt32 N) (s len i2 j2 : ℕ) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (ht1 : s + j2 < N) (ht2 : len + j2 < N) (ht3 : len + j2 + s < N) (hne31 : i2 + len + j2 + s ≠ i2 + j2 + s) : have t1 := roots[s + j2]; have t3 := roots[len + j2 + s]; have aB := montMul a[i2 + j2 + s] t1; have aD := montMul a[i2 + len + j2 + s] t1; have Q := subMod32 a[i2 + j2] aB; have S_v := subMod32 a[i2 + len + j2] aD; have t3S := montMul t3 S_v; (butterfly4 a false roots s len i2 j2)[i2 + j2 + s] = addMod32 Q t3S

theorem butterfly4_forward_getElem_pos3 {N : ℕ} (roots a : Vector UInt32 N) (s len i2 j2 : ℕ) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (ht1 : s + j2 < N) (ht2 : len + j2 < N) (ht3 : len + j2 + s < N) : have t1 := roots[s + j2]; have t3 := roots[len + j2 + s]; have aB := montMul a[i2 + j2 + s] t1; have aD := montMul a[i2 + len + j2 + s] t1; have Q := subMod32 a[i2 + j2] aB; have S_v := subMod32 a[i2 + len + j2] aD; have t3S := montMul t3 S_v; (butterfly4 a false roots s len i2 j2)[i2 + len + j2 + s] = subMod32 Q t3S

theorem butterfly4_getElem_ne {N : ℕ} (roots a : Vector UInt32 N) (s len i2 j2 : ℕ) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (i : ℕ) (hi : i < N) (hne0 : i2 + j2 ≠ i) (hne1 : i2 + j2 + s ≠ i) (hne2 : i2 + len + j2 ≠ i) (hne3 : i2 + len + j2 + s ≠ i) : (butterfly4 a false roots s len i2 j2)[i] = a[i]

butterfly4 only modifies the four butterfly positions; all other positions are unchanged.

theorem radix4Inner_getElem_ne {N : ℕ} (roots a : Vector UInt32 N) (s len i2 k j2_start hi : ℕ) (hlt : hi < N) (hbnd : ∀ (j2 : ℕ), j2_start ≤ j2 → j2 < j2_start + k → i2 + j2 < N ∧ i2 + j2 + s < N ∧ i2 + len + j2 < N ∧ i2 + len + j2 + s < N) (hne : ∀ (j2 : ℕ), j2_start ≤ j2 → j2 < j2_start + k → i2 + j2 ≠ hi ∧ i2 + j2 + s ≠ hi ∧ i2 + len + j2 ≠ hi ∧ i2 + len + j2 + s ≠ hi) : (radix4Inner false roots s len i2 k j2_start a)[hi] = a[hi]

radix4Inner leaves position hi unchanged if it is not any of the four butterfly positions for any j2 in the processed range [j2_start, j2_start + k).

theorem radix4Middle_getElem_ne {N : ℕ} (roots a : Vector UInt32 N) (s len k b_start hi : ℕ) (hlt : hi < N) (hbnd : ∀ (b : ℕ), b_start ≤ b → b < b_start + k → ∀ j2 < s, have i2 := b * 2 * len; i2 + j2 < N ∧ i2 + j2 + s < N ∧ i2 + len + j2 < N ∧ i2 + len + j2 + s < N) (hne : ∀ (b : ℕ), b_start ≤ b → b < b_start + k → ∀ j2 < s, have i2 := b * 2 * len; i2 + j2 ≠ hi ∧ i2 + j2 + s ≠ hi ∧ i2 + len + j2 ≠ hi ∧ i2 + len + j2 + s ≠ hi) : (radix4Middle false roots s len k b_start a)[hi] = a[hi]

radix4Middle leaves position hi unchanged if it is not any butterfly position for any group b in [b_start, b_start + k) and any j2 < s.

theorem radix4Inner_comp {N : ℕ} (roots : Vector UInt32 N) (inverse : Bool) (s len i2 k1 k2 j2_start : ℕ) (a : Vector UInt32 N) : radix4Inner inverse roots s len i2 (k1 + k2) j2_start a = radix4Inner inverse roots s len i2 k2 (j2_start + k1) (radix4Inner inverse roots s len i2 k1 j2_start a)

theorem butterfly4_forward_ZMod_combined {N : ℕ} (roots a : Vector UInt32 N) (ha : (a.all fun (x : UInt32) => decide (x < mod32)) = true) (hroots : ntt_roots_correct N roots) (hroots_bnd : (roots.all fun (x : UInt32) => decide (x < mod32)) = true) (s len i2 j2 : ℕ) (hlen : len = 2 * s) (hlen_dvd : 2 * len ∣ N) (hj2 : j2 < s) (hbnd0 : i2 + j2 < N) (hbnd1 : i2 + j2 + s < N) (hbnd2 : i2 + len + j2 < N) (hbnd3 : i2 + len + j2 + s < N) (hN_le : N ≤ 2 ^ 64) : have τ₁ := ↑primRoot.toNat ^ ((mod64.toNat - 1) / len * j2); have τ₂ := ↑primRoot.toNat ^ ((mod64.toNat - 1) / (2 * len) * j2); have τ₃ := ↑primRoot.toNat ^ ((mod64.toNat - 1) / (2 * len) * (s + j2)); have r := butterfly4 a false roots s len i2 j2; have A₀ := ↑a[i2 + j2].toNat; have A₁ := ↑a[i2 + j2 + s].toNat; have A₂ := ↑a[i2 + len + j2].toNat; have A₃ := ↑a[i2 + len + j2 + s].toNat; ↑r[i2 + j2].toNat = A₀ + τ₁ * A₁ + τ₂ * (A₂ + τ₁ * A₃) ∧ ↑r[i2 + len + j2].toNat = A₀ + τ₁ * A₁ - τ₂ * (A₂ + τ₁ * A₃) ∧ ↑r[i2 + j2 + s].toNat = A₀ - τ₁ * A₁ + τ₃ * (A₂ - τ₁ * A₃) ∧ ↑r[i2 + len + j2 + s].toNat = A₀ - τ₁ * A₁ - τ₃ * (A₂ - τ₁ * A₃)