nx_dispatcher_test.nx source
↩ module page · 123 lines · 6241 B
1// nx_dispatcher_test.nx -- smoke for nx_dispatcher.
2
3import "nx_syscalls.nx"
4import "nx_dispatcher.nx"
5
6func main() -> i64 {
7 // 1: enum validity
8 if nx_di_op_is_valid(NX_DI_OP_MATMUL) != 1 { return 1 }
9 if nx_di_op_is_valid(NX_DI_OP_DEQUANTIZE) != 1 { return 2 }
10 if nx_di_op_is_valid(-1) != 0 { return 3 }
11 if nx_di_op_is_valid(11) != 0 { return 4 }
12 if NX_DI_OP_N != 11 { return 5 }
13
14 if nx_di_be_is_valid(NX_DI_BE_CPU_SCALAR) != 1 { return 6 }
15 if nx_di_be_is_valid(NX_DI_BE_NPU) != 1 { return 7 }
16 if nx_di_be_is_valid(-1) != 0 { return 8 }
17 if nx_di_be_is_valid(5) != 0 { return 9 }
18
19 if nx_di_v_is_valid(NX_DI_V_DISPATCHED) != 1 { return 10 }
20 if nx_di_v_is_valid(NX_DI_V_NO_VIABLE_BACKEND) != 1 { return 11 }
21
22 // 2: op classification
23 if nx_di_op_is_compute_heavy(NX_DI_OP_MATMUL) != 1 { return 12 }
24 if nx_di_op_is_compute_heavy(NX_DI_OP_ATTENTION) != 1 { return 13 }
25 if nx_di_op_is_compute_heavy(NX_DI_OP_COPY) != 0 { return 14 }
26
27 if nx_di_op_is_lightweight(NX_DI_OP_ELEMENTWISE_UNARY) != 1 { return 15 }
28 if nx_di_op_is_lightweight(NX_DI_OP_COPY) != 1 { return 16 }
29 if nx_di_op_is_lightweight(NX_DI_OP_MATMUL) != 0 { return 17 }
30
31 // 3: backend classification
32 if nx_di_be_is_gpu(NX_DI_BE_GPU_COMPUTE) != 1 { return 18 }
33 if nx_di_be_is_gpu(NX_DI_BE_GPU_TENSOR_CORE) != 1 { return 19 }
34 if nx_di_be_is_gpu(NX_DI_BE_CPU_SIMD) != 0 { return 20 }
35 if nx_di_be_is_cpu(NX_DI_BE_CPU_SCALAR) != 1 { return 21 }
36 if nx_di_be_is_cpu(NX_DI_BE_CPU_SIMD) != 1 { return 22 }
37 if nx_di_be_is_cpu(NX_DI_BE_GPU_COMPUTE) != 0 { return 23 }
38
39 // 4: Build matmul request, all 4MB inputs already in VRAM (gpu_bytes=4194304)
40 let req_mm: *NxOpRequest = nx_di_req_new(NX_DI_OP_MATMUL, 4194304, 4194304, 0, 0, 1)
41 if req_mm.op_kind != NX_DI_OP_MATMUL { return 24 }
42 if req_mm.input_bytes_total != 4194304 { return 25 }
43
44 // 5: dispatch should pick GPU_TENSOR_CORE (cheapest for matmul, all bytes on GPU)
45 let dec_mm: *NxDispatchDecision = nx_di_dec_new()
46 if nx_di_dispatch(req_mm, dec_mm) != NX_DI_V_DISPATCHED { return 26 }
47 if dec_mm.chosen_backend != NX_DI_BE_GPU_TENSOR_CORE { return 27 }
48 if dec_mm.requires_promotion != 0 { return 28 }
49 if dec_mm.refused_reason != NX_DI_V_DISPATCHED { return 29 }
50
51 // 6: same matmul but all inputs on CPU. GPU still wins (matmul base 1
52 // so heavy compute saving dominates penalty), but check the
53 // dispatcher actually evaluated transfer cost.
54 let req_mm_cpu: *NxOpRequest = nx_di_req_new(NX_DI_OP_MATMUL, 4194304, 0, 4194304, 0, 1)
55 let dec_mm_cpu: *NxDispatchDecision = nx_di_dec_new()
56 nx_di_dispatch(req_mm_cpu, dec_mm_cpu)
57 if nx_di_be_is_gpu(dec_mm_cpu.chosen_backend) != 1 { return 30 }
58 // Cost should be higher than the all-on-GPU case for the same backend
59 let gpu_cost_on_gpu: nx_int = nx_di_estimate_cost_q10(req_mm, NX_DI_BE_GPU_TENSOR_CORE)
60 let gpu_cost_on_cpu: nx_int = nx_di_estimate_cost_q10(req_mm_cpu, NX_DI_BE_GPU_TENSOR_CORE)
61 if gpu_cost_on_cpu <= gpu_cost_on_gpu { return 31 }
62
63 // 7: small elementwise (4KB) already in CPU RAM -- CPU_SIMD should win over GPU
64 // because PCIe transfer dominates the tiny compute savings
65 let req_ew: *NxOpRequest = nx_di_req_new(NX_DI_OP_ELEMENTWISE_UNARY, 4096, 0, 4096, 0, 1)
66 let dec_ew: *NxDispatchDecision = nx_di_dec_new()
67 nx_di_dispatch(req_ew, dec_ew)
68 if dec_ew.chosen_backend != NX_DI_BE_CPU_SIMD { return 32 }
69
70 // 8: copy op -- should pick a CPU backend (all 1's; tie-break picks first
71 // encountered = CPU_SCALAR since we iterate from 0)
72 let req_cp: *NxOpRequest = nx_di_req_new(NX_DI_OP_COPY, 4096, 0, 4096, 0, 1)
73 let dec_cp: *NxDispatchDecision = nx_di_dec_new()
74 nx_di_dispatch(req_cp, dec_cp)
75 if nx_di_be_is_cpu(dec_cp.chosen_backend) != 1 { return 33 }
76
77 // 9: embedding lookup -- CPU wins for small tables
78 let req_em: *NxOpRequest = nx_di_req_new(NX_DI_OP_EMBEDDING_LOOKUP, 1024, 0, 1024, 0, 1)
79 let dec_em: *NxDispatchDecision = nx_di_dec_new()
80 nx_di_dispatch(req_em, dec_em)
81 if nx_di_be_is_cpu(dec_em.chosen_backend) != 1 { return 34 }
82
83 // 10: cold-tier inputs require promotion flag
84 let req_cold: *NxOpRequest = nx_di_req_new(NX_DI_OP_MATMUL, 1048576, 0, 0, 1048576, 1)
85 let dec_cold: *NxDispatchDecision = nx_di_dec_new()
86 nx_di_dispatch(req_cold, dec_cold)
87 if dec_cold.requires_promotion != 1 { return 35 }
88
89 // 11: invalid op rejected in req_new
90 if nx_di_req_new(99, 100, 0, 0, 0, 1) != (0 as *NxOpRequest) { return 36 }
91
92 // 12: invalid op via direct cost estimation -> sentinel
93 let req_bad_op: *NxOpRequest = nx_di_req_new(NX_DI_OP_MATMUL, 100, 100, 0, 0, 1)
94 req_bad_op.op_kind = 99 // poke invalid post-construction
95 let dec_bad: *NxDispatchDecision = nx_di_dec_new()
96 if nx_di_dispatch(req_bad_op, dec_bad) != NX_DI_V_INVALID { return 37 }
97
98 // 13: null guards
99 let null_req: *NxOpRequest = (0 as i64) as *NxOpRequest
100 let null_dec: *NxDispatchDecision = (0 as i64) as *NxDispatchDecision
101 if nx_di_dispatch(null_req, dec_mm) != NX_DI_V_NULL { return 38 }
102 if nx_di_dispatch(req_mm, null_dec) != NX_DI_V_NULL { return 39 }
103 if nx_di_estimate_cost_q10(null_req, NX_DI_BE_GPU_TENSOR_CORE) != 1000000 { return 40 }
104
105 // 14: invalid backend in estimate
106 if nx_di_estimate_cost_q10(req_mm, 99) != 1000000 { return 41 }
107 if nx_di_estimate_cost_q10(req_mm, -1) != 1000000 { return 42 }
108
109 // 15: attention op -- should prefer GPU tensor cores
110 let req_at: *NxOpRequest = nx_di_req_new(NX_DI_OP_ATTENTION, 524288, 524288, 0, 0, 1)
111 let dec_at: *NxDispatchDecision = nx_di_dec_new()
112 nx_di_dispatch(req_at, dec_at)
113 if dec_at.chosen_backend != NX_DI_BE_GPU_TENSOR_CORE { return 43 }
114
115 // 16: rms_norm prefers CPU SIMD (base cost 3 vs GPU 4 + PCIe penalty
116 // if cpu-resident). But check the decision is at least produced.
117 let req_rn: *NxOpRequest = nx_di_req_new(NX_DI_OP_RMS_NORM, 4096, 0, 4096, 0, 1)
118 let dec_rn: *NxDispatchDecision = nx_di_dec_new()
119 if nx_di_dispatch(req_rn, dec_rn) != NX_DI_V_DISPATCHED { return 44 }
120 if dec_rn.chosen_backend != NX_DI_BE_CPU_SIMD { return 45 }
121
122 return 0
123}