Component.RotaryFrequency

1module Component.RotaryFrequency
2
3import Float32Exact
4import FloatLiteralSpec
5import BinaryFloatModel
6import Std.List
7import Std.Natural
8
9-- The Llama half-split rotation uses theta_i = base^(-i / columns).
10-- Keep the exact root and its f32 residual together: the SM86 table writer
11-- uses both words to keep 4096-position angles accurate after reduction.
12def rotaryFrequencyResolution : Nat = (specPow2 96)
13def rotaryFrequencyRoot =
14  (lambda unrestricted base : Nat .
15    (lambda unrestricted columns : Nat .
16      (lambda unrestricted index : Nat .
17        (float32ExactRootKLow columns 1 (specPower base index)
18          rotaryFrequencyResolution))))
19def rotaryFrequencyHigh =
20  (lambda unrestricted base : Nat .
21    (lambda unrestricted columns : Nat .
22      (lambda unrestricted index : Nat .
23        (let unrestricted root = (rotaryFrequencyRoot base columns index) in
24          (float32ExactBetween root rotaryFrequencyResolution
25            (succ root) rotaryFrequencyResolution)))))
26def rotaryFrequencyLow =
27  (lambda unrestricted base : Nat .
28    (lambda unrestricted columns : Nat .
29      (lambda unrestricted index : Nat .
30        (let unrestricted root = (rotaryFrequencyRoot base columns index) in
31          (float32ExactRemainder root (succ root)
32            rotaryFrequencyResolution
33            (rotaryFrequencyHigh base columns index))))))
34def rotaryFrequencyRoundingMagic : Nat =
35  (compile-time (float32ExactRational 12582912 1))
36
37-- Adjacent high/low words let a launch schedule index a column without
38-- reevaluating the large exact root for every scalar patch.
39def rotaryFrequencyWords =
40  (lambda unrestricted base : Nat .
41    (lambda unrestricted columns : Nat .
42      (nat-eliminate
43        (lambda unrestricted current : Nat . (family StdList Nat))
44        (constructor StdList StdListEmpty Nat)
45        (lambda unrestricted predecessor : Nat .
46          (lambda unrestricted rest : (family StdList Nat) .
47            (let unrestricted index =
48              (naturalSaturatingSubtract
49                (naturalSaturatingSubtract columns 1) predecessor) in
50              (constructor StdList StdListCons Nat
51                (rotaryFrequencyHigh base columns index)
52                (constructor StdList StdListCons Nat
53                  (rotaryFrequencyLow base columns index) rest)))))
54        columns)))