Component.RotaryFrequency

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