nx_f32_rope_test.nx source
↩ module page · 80 lines · 2758 B
1// nx_f32_rope_test.nx -- smoke for nx_f32_rope.nx.
2//
3// Two test classes:
4// A) position=0 identity: cos(0)=1, sin(0)=0, so RoPE preserves
5// input exactly bit-for-bit.
6// B) head_dim=2, position=1, base=10000: theta_0 = 1.0,
7// so x' = [x0*cos(1) - x1*sin(1), x0*sin(1) + x1*cos(1)].
8// With x=[1,0]: x' = [cos(1), sin(1)] ~= [0.5403, 0.8415].
9
10import "nx_syscalls.nx"
11import "nx_tier.nx"
12import "nx_f32.nx"
13import "nx_f32_div.nx"
14import "nx_f32_cvt.nx"
15import "nx_f32_exp.nx"
16import "nx_f32_sincos.nx"
17import "nx_f32_rope.nx"
18
19func _ulp_diff_pos(a: i64, b: i64) -> i64 {
20 if a >= b { return a - b }
21 return b - a
22}
23
24func main() -> i64 {
25 var vi: nx_int = 0
26 while vi < NX_F32_ROPE_N_VERDICTS {
27 if nx_f32_rope_verdict_is_valid(vi) != 1 { return 5 + vi }
28 vi = vi + 1
29 }
30
31 // log(10000) f32 ~= 9.2103404 ~= 0x4113411D
32 let log_base: i64 = 0x4113411D
33
34 // ===== Test A: position=0 identity (exact bit preservation) =====
35 let xA: *i64 = sys_mmap(4 * 8) as *i64
36 xA[0] = 0x3F800000 // 1.0
37 xA[1] = 0x40000000 // 2.0
38 xA[2] = 0x40400000 // 3.0
39 xA[3] = 0x40800000 // 4.0
40
41 let vA: nx_int = nx_f32_rope_apply_vector(xA, 4, 0, log_base)
42 if vA != NX_F32_ROPE_OK { return 10 + vA }
43
44 // At position=0: cos(0)=1.0 exact, sin(0)=0 exact.
45 // x_lo' = x_lo * 1 - x_hi * 0 = x_lo
46 // x_hi' = x_lo * 0 + x_hi * 1 = x_hi
47 // Output bit-exact preserved.
48 if xA[0] != 0x3F800000 { return 20 }
49 if xA[1] != 0x40000000 { return 21 }
50 if xA[2] != 0x40400000 { return 22 }
51 if xA[3] != 0x40800000 { return 23 }
52
53 // ===== Test B: head_dim=2, position=1, x=[1,0] =====
54 // theta_0 = 10000^(0/2) = 1.0
55 // m*theta_0 = 1
56 // x'[0] = 1*cos(1) - 0*sin(1) = cos(1) ~= 0.5403 = 0x3F0A5DE7
57 // x'[1] = 1*sin(1) + 0*cos(1) = sin(1) ~= 0.8415 = 0x3F576AA4
58 let xB: *i64 = sys_mmap(2 * 8) as *i64
59 xB[0] = 0x3F800000 // 1.0
60 xB[1] = 0x00000000 // 0.0
61
62 let vB: nx_int = nx_f32_rope_apply_vector(xB, 2, 1, log_base)
63 if vB != NX_F32_ROPE_OK { return 30 }
64
65 // cos(1) f32 nearest = 0x3F0A5DE7 (= 0.5403023)
66 // sin(1) f32 nearest = 0x3F576AA4 (= 0.8414710)
67 // Our v1 sin/cos: ~16384 ULPs. Multiplied by mul+sub for x'[0]:
68 // x'[0] = 1*cos(1) - 0*sin(1) = cos(1) directly (no extra error)
69 // x'[1] = 1*sin(1) + 0*cos(1) = sin(1) directly
70 if _ulp_diff_pos(xB[0], 0x3F0A5DE7) > 32768 { return 31 }
71 if _ulp_diff_pos(xB[1], 0x3F576AA4) > 32768 { return 32 }
72
73 // ===== Test C: bad-dim verdict =====
74 let xC: *i64 = sys_mmap(3 * 8) as *i64
75 xC[0] = 0; xC[1] = 0; xC[2] = 0
76 let vC: nx_int = nx_f32_rope_apply_vector(xC, 3, 0, log_base)
77 if vC != NX_F32_ROPE_ERR_ODD_DIM { return 40 }
78
79 return 0
80}