oxicuda_levelzero/
spirv_nn.rs

1//! SPIR-V compute kernel generators for neural-network operations.
2//!
3//! This module extends the Level Zero SPIR-V generator with Conv2D and
4//! scaled dot-product attention kernels.  Both use the OpenCL execution
5//! model (Kernel + Physical64 + Addresses) so they can be consumed by
6//! `zeModuleCreate`.
7
8use super::spirv::{
9    FUNCTION_CONTROL_NONE, OP_F_ADD, OP_F_DIV, OP_F_MUL, OP_F_SUB, OP_I_ADD, OP_I_MUL, OP_U_DIV,
10    OP_U_LESS_THAN, OP_U_MOD, OPENCL_EXP, OPENCL_FMAX, STORAGE_CLASS_FUNCTION, SpvModule,
11    emit_preamble, load_gid_x,
12};
13
14/// SPIR-V opcode for `OpISub` (integer subtract).
15const OP_I_SUB: u32 = 130;
16/// SPIR-V opcode for `OpFOrdGreaterThan` (ordered float >).
17const OP_F_ORD_GT: u32 = 188;
18
19// ─── Conv2D compute kernel ──────────────────────────────────
20
21/// Generate an OpenCL SPIR-V compute kernel for 2-D convolution (NCHW layout).
22///
23/// Each work-item computes one output element.
24///
25/// Kernel parameters (all passed via `zeKernelSetArgumentValue`):
26///
27/// ```text
28/// (CrossWorkgroup float* input,
29///  CrossWorkgroup float* filter,
30///  CrossWorkgroup float* output)
31/// ```
32///
33/// All dimension constants are baked in as `OpConstant`.
34#[allow(clippy::too_many_arguments)]
35pub fn conv2d_spirv(
36    n: u32,
37    c_in: u32,
38    h_in: u32,
39    w_in: u32,
40    k_out: u32,
41    fh: u32,
42    fw: u32,
43    oh: u32,
44    ow: u32,
45    stride_h: u32,
46    stride_w: u32,
47    pad_h: u32,
48    pad_w: u32,
49) -> Vec<u32> {
50    let mut m = SpvModule::new();
51    let b = emit_preamble(&mut m, "main");
52
53    let fn_ty = m.alloc_id();
54    let p_input = m.alloc_id();
55    let p_filter = m.alloc_id();
56    let p_output = m.alloc_id();
57
58    // Function type: void(float*, float*, float*)
59    m.emit_type_function(
60        fn_ty,
61        b.ty_void,
62        &[
63            b.ty_ptr_cross_float,
64            b.ty_ptr_cross_float,
65            b.ty_ptr_cross_float,
66        ],
67    );
68
69    // Emit constants for dimensions
70    let c_n = m.alloc_id();
71    m.emit_constant_u32(b.ty_uint, c_n, n);
72    let c_c_in = m.alloc_id();
73    m.emit_constant_u32(b.ty_uint, c_c_in, c_in);
74    let c_h_in = m.alloc_id();
75    m.emit_constant_u32(b.ty_uint, c_h_in, h_in);
76    let c_w_in = m.alloc_id();
77    m.emit_constant_u32(b.ty_uint, c_w_in, w_in);
78    let c_k_out = m.alloc_id();
79    m.emit_constant_u32(b.ty_uint, c_k_out, k_out);
80    let c_fh = m.alloc_id();
81    m.emit_constant_u32(b.ty_uint, c_fh, fh);
82    let c_fw = m.alloc_id();
83    m.emit_constant_u32(b.ty_uint, c_fw, fw);
84    let c_oh = m.alloc_id();
85    m.emit_constant_u32(b.ty_uint, c_oh, oh);
86    let c_ow = m.alloc_id();
87    m.emit_constant_u32(b.ty_uint, c_ow, ow);
88    let c_stride_h = m.alloc_id();
89    m.emit_constant_u32(b.ty_uint, c_stride_h, stride_h);
90    let c_stride_w = m.alloc_id();
91    m.emit_constant_u32(b.ty_uint, c_stride_w, stride_w);
92    let c_pad_h = m.alloc_id();
93    m.emit_constant_u32(b.ty_uint, c_pad_h, pad_h);
94    let c_pad_w = m.alloc_id();
95    m.emit_constant_u32(b.ty_uint, c_pad_w, pad_w);
96
97    // Labels
98    let label_entry = m.alloc_id();
99    let label_body = m.alloc_id();
100    let label_merge = m.alloc_id();
101
102    // Function
103    m.emit_function(b.ty_void, b.main_fn, FUNCTION_CONTROL_NONE, fn_ty);
104    m.emit_function_parameter(b.ty_ptr_cross_float, p_input);
105    m.emit_function_parameter(b.ty_ptr_cross_float, p_filter);
106    m.emit_function_parameter(b.ty_ptr_cross_float, p_output);
107    m.emit_label(label_entry);
108
109    let gid = load_gid_x(&mut m, &b);
110
111    // total = n * k_out * oh * ow
112    let t1 = m.alloc_id();
113    m.emit(OP_I_MUL, &[b.ty_uint, t1, c_n, c_k_out]);
114    let t2 = m.alloc_id();
115    m.emit(OP_I_MUL, &[b.ty_uint, t2, t1, c_oh]);
116    let total = m.alloc_id();
117    m.emit(OP_I_MUL, &[b.ty_uint, total, t2, c_ow]);
118
119    let cond = m.alloc_id();
120    m.emit(OP_U_LESS_THAN, &[b.ty_bool, cond, gid, total]);
121    m.emit_selection_merge(label_merge);
122    m.emit_branch_conditional(cond, label_body, label_merge);
123
124    m.emit_label(label_body);
125
126    // Decompose gid -> (b_idx, kf, oy, ox)
127    let ox = m.alloc_id();
128    m.emit(OP_U_MOD, &[b.ty_uint, ox, gid, c_ow]);
129    let tmp1 = m.alloc_id();
130    m.emit(OP_U_DIV, &[b.ty_uint, tmp1, gid, c_ow]);
131    let oy = m.alloc_id();
132    m.emit(OP_U_MOD, &[b.ty_uint, oy, tmp1, c_oh]);
133    let tmp2 = m.alloc_id();
134    m.emit(OP_U_DIV, &[b.ty_uint, tmp2, tmp1, c_oh]);
135    let kf = m.alloc_id();
136    m.emit(OP_U_MOD, &[b.ty_uint, kf, tmp2, c_k_out]);
137    let b_idx = m.alloc_id();
138    m.emit(OP_U_DIV, &[b.ty_uint, b_idx, tmp2, c_k_out]);
139
140    // Accumulator variable
141    let var_acc = m.alloc_id();
142    m.emit_variable(b.ty_ptr_func_float, var_acc, STORAGE_CLASS_FUNCTION);
143    m.emit_store(var_acc, b.c_float_0);
144
145    // Flatten ci * fh * fw
146    let flat_total_id = m.alloc_id();
147    let flat_t1 = m.alloc_id();
148    m.emit(OP_I_MUL, &[b.ty_uint, flat_t1, c_c_in, c_fh]);
149    m.emit(OP_I_MUL, &[b.ty_uint, flat_total_id, flat_t1, c_fw]);
150
151    let var_flat = m.alloc_id();
152    m.emit_variable(b.ty_ptr_func_uint, var_flat, STORAGE_CLASS_FUNCTION);
153    m.emit_store(var_flat, b.c_uint_0);
154
155    let lbl_loop_hdr = m.alloc_id();
156    let lbl_loop_body = m.alloc_id();
157    let lbl_loop_cont = m.alloc_id();
158    let lbl_loop_merge = m.alloc_id();
159
160    m.emit_branch(lbl_loop_hdr);
161
162    // ── Loop header ──
163    m.emit_label(lbl_loop_hdr);
164    let flat_val = m.alloc_id();
165    m.emit_load(b.ty_uint, flat_val, var_flat);
166    let loop_cond = m.alloc_id();
167    m.emit(
168        OP_U_LESS_THAN,
169        &[b.ty_bool, loop_cond, flat_val, flat_total_id],
170    );
171    m.emit_loop_merge(lbl_loop_merge, lbl_loop_cont);
172    m.emit_branch_conditional(loop_cond, lbl_loop_body, lbl_loop_merge);
173
174    // ── Loop body ──
175    m.emit_label(lbl_loop_body);
176
177    // Decompose flat_val -> (ci, fy, fx)
178    let fx = m.alloc_id();
179    m.emit(OP_U_MOD, &[b.ty_uint, fx, flat_val, c_fw]);
180    let ftmp1 = m.alloc_id();
181    m.emit(OP_U_DIV, &[b.ty_uint, ftmp1, flat_val, c_fw]);
182    let fy = m.alloc_id();
183    m.emit(OP_U_MOD, &[b.ty_uint, fy, ftmp1, c_fh]);
184    let ci = m.alloc_id();
185    m.emit(OP_U_DIV, &[b.ty_uint, ci, ftmp1, c_fh]);
186
187    // iy_raw = oy * stride_h + fy, ix_raw = ox * stride_w + fx
188    let oy_sh = m.alloc_id();
189    m.emit(OP_I_MUL, &[b.ty_uint, oy_sh, oy, c_stride_h]);
190    let iy_raw = m.alloc_id();
191    m.emit(OP_I_ADD, &[b.ty_uint, iy_raw, oy_sh, fy]);
192    let ox_sw = m.alloc_id();
193    m.emit(OP_I_MUL, &[b.ty_uint, ox_sw, ox, c_stride_w]);
194    let ix_raw = m.alloc_id();
195    m.emit(OP_I_ADD, &[b.ty_uint, ix_raw, ox_sw, fx]);
196
197    // Bounds check: iy_raw >= pad_h  &&  (iy_raw - pad_h) < h_in
198    //           &&  ix_raw >= pad_w  &&  (ix_raw - pad_w) < w_in
199    let lbl_skip = m.alloc_id();
200
201    let iy_lt_pad = m.alloc_id();
202    m.emit(OP_U_LESS_THAN, &[b.ty_bool, iy_lt_pad, iy_raw, c_pad_h]);
203    let lbl_iy_ok = m.alloc_id();
204    m.emit_selection_merge(lbl_skip);
205    m.emit_branch_conditional(iy_lt_pad, lbl_skip, lbl_iy_ok);
206
207    m.emit_label(lbl_iy_ok);
208    let iy_real = m.alloc_id();
209    m.emit(OP_I_SUB, &[b.ty_uint, iy_real, iy_raw, c_pad_h]);
210    let iy_in_bounds = m.alloc_id();
211    m.emit(OP_U_LESS_THAN, &[b.ty_bool, iy_in_bounds, iy_real, c_h_in]);
212    let lbl_ix_check = m.alloc_id();
213    m.emit_selection_merge(lbl_skip);
214    m.emit_branch_conditional(iy_in_bounds, lbl_ix_check, lbl_skip);
215
216    m.emit_label(lbl_ix_check);
217    let ix_lt_pad = m.alloc_id();
218    m.emit(OP_U_LESS_THAN, &[b.ty_bool, ix_lt_pad, ix_raw, c_pad_w]);
219    let lbl_ix_ok = m.alloc_id();
220    m.emit_selection_merge(lbl_skip);
221    m.emit_branch_conditional(ix_lt_pad, lbl_skip, lbl_ix_ok);
222
223    m.emit_label(lbl_ix_ok);
224    let ix_real = m.alloc_id();
225    m.emit(OP_I_SUB, &[b.ty_uint, ix_real, ix_raw, c_pad_w]);
226    let ix_in_bounds = m.alloc_id();
227    m.emit(OP_U_LESS_THAN, &[b.ty_bool, ix_in_bounds, ix_real, c_w_in]);
228    let lbl_accum = m.alloc_id();
229    m.emit_selection_merge(lbl_skip);
230    m.emit_branch_conditional(ix_in_bounds, lbl_accum, lbl_skip);
231
232    m.emit_label(lbl_accum);
233
234    // input_idx = ((b_idx * c_in + ci) * h_in + iy_real) * w_in + ix_real
235    let in1 = m.alloc_id();
236    m.emit(OP_I_MUL, &[b.ty_uint, in1, b_idx, c_c_in]);
237    let in2 = m.alloc_id();
238    m.emit(OP_I_ADD, &[b.ty_uint, in2, in1, ci]);
239    let in3 = m.alloc_id();
240    m.emit(OP_I_MUL, &[b.ty_uint, in3, in2, c_h_in]);
241    let in4 = m.alloc_id();
242    m.emit(OP_I_ADD, &[b.ty_uint, in4, in3, iy_real]);
243    let in5 = m.alloc_id();
244    m.emit(OP_I_MUL, &[b.ty_uint, in5, in4, c_w_in]);
245    let in_idx = m.alloc_id();
246    m.emit(OP_I_ADD, &[b.ty_uint, in_idx, in5, ix_real]);
247
248    let inp_ptr = m.alloc_id();
249    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, inp_ptr, p_input, in_idx);
250    let inp_val = m.alloc_id();
251    m.emit_load(b.ty_float, inp_val, inp_ptr);
252
253    // filter_idx = ((kf * c_in + ci) * fh + fy) * fw + fx
254    let f1 = m.alloc_id();
255    m.emit(OP_I_MUL, &[b.ty_uint, f1, kf, c_c_in]);
256    let f2 = m.alloc_id();
257    m.emit(OP_I_ADD, &[b.ty_uint, f2, f1, ci]);
258    let f3 = m.alloc_id();
259    m.emit(OP_I_MUL, &[b.ty_uint, f3, f2, c_fh]);
260    let f4 = m.alloc_id();
261    m.emit(OP_I_ADD, &[b.ty_uint, f4, f3, fy]);
262    let f5 = m.alloc_id();
263    m.emit(OP_I_MUL, &[b.ty_uint, f5, f4, c_fw]);
264    let flt_idx = m.alloc_id();
265    m.emit(OP_I_ADD, &[b.ty_uint, flt_idx, f5, fx]);
266
267    let flt_ptr = m.alloc_id();
268    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, flt_ptr, p_filter, flt_idx);
269    let flt_val = m.alloc_id();
270    m.emit_load(b.ty_float, flt_val, flt_ptr);
271
272    // acc += inp * flt
273    let prod = m.alloc_id();
274    m.emit(OP_F_MUL, &[b.ty_float, prod, inp_val, flt_val]);
275    let old_acc = m.alloc_id();
276    m.emit_load(b.ty_float, old_acc, var_acc);
277    let new_acc = m.alloc_id();
278    m.emit(OP_F_ADD, &[b.ty_float, new_acc, old_acc, prod]);
279    m.emit_store(var_acc, new_acc);
280
281    m.emit_branch(lbl_skip);
282
283    m.emit_label(lbl_skip);
284    m.emit_branch(lbl_loop_cont);
285
286    // ── Loop continue ──
287    m.emit_label(lbl_loop_cont);
288    let flat_inc = m.alloc_id();
289    m.emit(OP_I_ADD, &[b.ty_uint, flat_inc, flat_val, b.c_uint_1]);
290    m.emit_store(var_flat, flat_inc);
291    m.emit_branch(lbl_loop_hdr);
292
293    // ── Loop merge: store result ──
294    m.emit_label(lbl_loop_merge);
295
296    let final_acc = m.alloc_id();
297    m.emit_load(b.ty_float, final_acc, var_acc);
298
299    let out_ptr = m.alloc_id();
300    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, out_ptr, p_output, gid);
301    m.emit_store(out_ptr, final_acc);
302
303    m.emit_branch(label_merge);
304
305    m.emit_label(label_merge);
306    m.emit_return();
307    m.emit_function_end();
308
309    m.finalize()
310}
311
312// ─── Attention compute kernel ───────────────────────────────
313
314/// Generate an OpenCL SPIR-V compute kernel for scaled dot-product attention.
315///
316/// Each work-item handles one (batch_head, query_position) pair.
317///
318/// Kernel parameters:
319///
320/// ```text
321/// (CrossWorkgroup float* Q,
322///  CrossWorkgroup float* K,
323///  CrossWorkgroup float* V,
324///  CrossWorkgroup float* O)
325/// ```
326///
327/// Dimension constants are baked in as `OpConstant`.
328#[allow(clippy::too_many_arguments)]
329pub fn attention_spirv(
330    batch_heads: u32,
331    seq_q: u32,
332    seq_kv: u32,
333    head_dim: u32,
334    scale: f32,
335    causal: bool,
336) -> Vec<u32> {
337    let mut m = SpvModule::new();
338    let b = emit_preamble(&mut m, "main");
339
340    let fn_ty = m.alloc_id();
341    let p_q = m.alloc_id();
342    let p_k = m.alloc_id();
343    let p_v = m.alloc_id();
344    let p_o = m.alloc_id();
345
346    // Function type: void(float*, float*, float*, float*)
347    m.emit_type_function(
348        fn_ty,
349        b.ty_void,
350        &[
351            b.ty_ptr_cross_float,
352            b.ty_ptr_cross_float,
353            b.ty_ptr_cross_float,
354            b.ty_ptr_cross_float,
355        ],
356    );
357
358    // Constants
359    let c_batch_heads = m.alloc_id();
360    m.emit_constant_u32(b.ty_uint, c_batch_heads, batch_heads);
361    let c_seq_q = m.alloc_id();
362    m.emit_constant_u32(b.ty_uint, c_seq_q, seq_q);
363    let c_seq_kv = m.alloc_id();
364    m.emit_constant_u32(b.ty_uint, c_seq_kv, seq_kv);
365    let c_head_dim = m.alloc_id();
366    m.emit_constant_u32(b.ty_uint, c_head_dim, head_dim);
367    let c_scale = m.alloc_id();
368    m.emit_constant_f32(b.ty_float, c_scale, scale);
369    let c_neg_inf = m.alloc_id();
370    m.emit_constant_f32(b.ty_float, c_neg_inf, f32::NEG_INFINITY);
371    // Stride: seq_kv * head_dim
372    let c_skv_hd = m.alloc_id();
373    m.emit_constant_u32(b.ty_uint, c_skv_hd, seq_kv * head_dim);
374
375    // Labels
376    let label_entry = m.alloc_id();
377    let label_body = m.alloc_id();
378    let label_merge = m.alloc_id();
379
380    m.emit_function(b.ty_void, b.main_fn, FUNCTION_CONTROL_NONE, fn_ty);
381    m.emit_function_parameter(b.ty_ptr_cross_float, p_q);
382    m.emit_function_parameter(b.ty_ptr_cross_float, p_k);
383    m.emit_function_parameter(b.ty_ptr_cross_float, p_v);
384    m.emit_function_parameter(b.ty_ptr_cross_float, p_o);
385    m.emit_label(label_entry);
386
387    let gid = load_gid_x(&mut m, &b);
388
389    // total = batch_heads * seq_q
390    let total = m.alloc_id();
391    m.emit(OP_I_MUL, &[b.ty_uint, total, c_batch_heads, c_seq_q]);
392
393    let cond = m.alloc_id();
394    m.emit(OP_U_LESS_THAN, &[b.ty_bool, cond, gid, total]);
395    m.emit_selection_merge(label_merge);
396    m.emit_branch_conditional(cond, label_body, label_merge);
397
398    m.emit_label(label_body);
399
400    // bh = gid / seq_q, sq = gid % seq_q
401    let bh = m.alloc_id();
402    m.emit(OP_U_DIV, &[b.ty_uint, bh, gid, c_seq_q]);
403    let sq = m.alloc_id();
404    m.emit(OP_U_MOD, &[b.ty_uint, sq, gid, c_seq_q]);
405
406    // q_base = gid * head_dim
407    let q_base = m.alloc_id();
408    m.emit(OP_I_MUL, &[b.ty_uint, q_base, gid, c_head_dim]);
409    // kv_base = bh * seq_kv * head_dim
410    let kv_base = m.alloc_id();
411    m.emit(OP_I_MUL, &[b.ty_uint, kv_base, bh, c_skv_hd]);
412
413    // var_max_score
414    let var_max = m.alloc_id();
415    m.emit_variable(b.ty_ptr_func_float, var_max, STORAGE_CLASS_FUNCTION);
416    m.emit_store(var_max, c_neg_inf);
417
418    // ── Pass 1: find max score ──
419    emit_score_pass(
420        &mut m, &b, causal, sq, c_seq_kv, c_head_dim, c_scale, q_base, kv_base, p_q, p_k, var_max,
421        true, None, None, p_v, p_o,
422    );
423
424    let final_max = m.alloc_id();
425    m.emit_load(b.ty_float, final_max, var_max);
426
427    // ── Zero-initialize O[q_base .. q_base+head_dim) ──
428    // Pass 2 accumulates weighted-V into O via read-modify-write, so O must
429    // start at 0. Each work-item owns a disjoint O slice (q_base = gid*head_dim)
430    // so this is race-free; do NOT rely on the host pre-zeroing the buffer.
431    let var_dz = m.alloc_id();
432    m.emit_variable(b.ty_ptr_func_uint, var_dz, STORAGE_CLASS_FUNCTION);
433    m.emit_store(var_dz, b.c_uint_0);
434
435    let lbl_dz_hdr = m.alloc_id();
436    let lbl_dz_body = m.alloc_id();
437    let lbl_dz_cont = m.alloc_id();
438    let lbl_dz_merge = m.alloc_id();
439
440    m.emit_branch(lbl_dz_hdr);
441
442    m.emit_label(lbl_dz_hdr);
443    let dz_val = m.alloc_id();
444    m.emit_load(b.ty_uint, dz_val, var_dz);
445    let dz_cond = m.alloc_id();
446    m.emit(OP_U_LESS_THAN, &[b.ty_bool, dz_cond, dz_val, c_head_dim]);
447    m.emit_loop_merge(lbl_dz_merge, lbl_dz_cont);
448    m.emit_branch_conditional(dz_cond, lbl_dz_body, lbl_dz_merge);
449
450    m.emit_label(lbl_dz_body);
451    let oz_idx = m.alloc_id();
452    m.emit(OP_I_ADD, &[b.ty_uint, oz_idx, q_base, dz_val]);
453    let oz_ptr = m.alloc_id();
454    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, oz_ptr, p_o, oz_idx);
455    m.emit_store(oz_ptr, b.c_float_0);
456    m.emit_branch(lbl_dz_cont);
457
458    m.emit_label(lbl_dz_cont);
459    let dz_inc = m.alloc_id();
460    m.emit(OP_I_ADD, &[b.ty_uint, dz_inc, dz_val, b.c_uint_1]);
461    m.emit_store(var_dz, dz_inc);
462    m.emit_branch(lbl_dz_hdr);
463
464    m.emit_label(lbl_dz_merge);
465
466    // ── Pass 2: accumulate exp-weighted V ──
467    let var_sum_exp = m.alloc_id();
468    m.emit_variable(b.ty_ptr_func_float, var_sum_exp, STORAGE_CLASS_FUNCTION);
469    m.emit_store(var_sum_exp, b.c_float_0);
470
471    emit_score_pass(
472        &mut m,
473        &b,
474        causal,
475        sq,
476        c_seq_kv,
477        c_head_dim,
478        c_scale,
479        q_base,
480        kv_base,
481        p_q,
482        p_k,
483        var_sum_exp,
484        false,
485        Some(final_max),
486        Some(p_o),
487        p_v,
488        p_o,
489    );
490
491    // Normalize: O[o_base+d] /= sum_exp if sum_exp > 0
492    let sum_final = m.alloc_id();
493    m.emit_load(b.ty_float, sum_final, var_sum_exp);
494
495    let sum_gt_zero = m.alloc_id();
496    m.emit(
497        OP_F_ORD_GT,
498        &[b.ty_bool, sum_gt_zero, sum_final, b.c_float_0],
499    );
500
501    let lbl_norm = m.alloc_id();
502    let lbl_norm_merge = m.alloc_id();
503    m.emit_selection_merge(lbl_norm_merge);
504    m.emit_branch_conditional(sum_gt_zero, lbl_norm, lbl_norm_merge);
505
506    m.emit_label(lbl_norm);
507
508    // Normalize loop
509    let var_d4 = m.alloc_id();
510    m.emit_variable(b.ty_ptr_func_uint, var_d4, STORAGE_CLASS_FUNCTION);
511    m.emit_store(var_d4, b.c_uint_0);
512
513    let lbl_d4_hdr = m.alloc_id();
514    let lbl_d4_body = m.alloc_id();
515    let lbl_d4_cont = m.alloc_id();
516    let lbl_d4_merge = m.alloc_id();
517
518    m.emit_branch(lbl_d4_hdr);
519
520    m.emit_label(lbl_d4_hdr);
521    let d4_val = m.alloc_id();
522    m.emit_load(b.ty_uint, d4_val, var_d4);
523    let d4_cond = m.alloc_id();
524    m.emit(OP_U_LESS_THAN, &[b.ty_bool, d4_cond, d4_val, c_head_dim]);
525    m.emit_loop_merge(lbl_d4_merge, lbl_d4_cont);
526    m.emit_branch_conditional(d4_cond, lbl_d4_body, lbl_d4_merge);
527
528    m.emit_label(lbl_d4_body);
529    let o4_idx = m.alloc_id();
530    m.emit(OP_I_ADD, &[b.ty_uint, o4_idx, q_base, d4_val]);
531    let o4_ptr = m.alloc_id();
532    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, o4_ptr, p_o, o4_idx);
533    let o4_val = m.alloc_id();
534    m.emit_load(b.ty_float, o4_val, o4_ptr);
535    let o4_normed = m.alloc_id();
536    m.emit(OP_F_DIV, &[b.ty_float, o4_normed, o4_val, sum_final]);
537    m.emit_store(o4_ptr, o4_normed);
538
539    m.emit_branch(lbl_d4_cont);
540    m.emit_label(lbl_d4_cont);
541    let d4_inc = m.alloc_id();
542    m.emit(OP_I_ADD, &[b.ty_uint, d4_inc, d4_val, b.c_uint_1]);
543    m.emit_store(var_d4, d4_inc);
544    m.emit_branch(lbl_d4_hdr);
545
546    m.emit_label(lbl_d4_merge);
547
548    m.emit_branch(lbl_norm_merge);
549    m.emit_label(lbl_norm_merge);
550
551    m.emit_branch(label_merge);
552
553    m.emit_label(label_merge);
554    m.emit_return();
555    m.emit_function_end();
556
557    m.finalize()
558}
559
560/// Emit a score-computation pass (used for both max-finding and accumulation).
561///
562/// When `is_max_pass` is true, updates `accum_var` with fmax(accum, score).
563/// When false, uses `max_val` to compute `exp(score - max)`, adds to `accum_var`,
564/// and accumulates weighted V into `o_buf`.
565#[allow(clippy::too_many_arguments)]
566fn emit_score_pass(
567    m: &mut SpvModule,
568    b: &super::spirv::BaseIds,
569    causal: bool,
570    sq: u32,
571    c_seq_kv: u32,
572    c_head_dim: u32,
573    c_scale: u32,
574    q_base: u32,
575    kv_base: u32,
576    p_q: u32,
577    p_k: u32,
578    accum_var: u32,
579    is_max_pass: bool,
580    max_val: Option<u32>,
581    o_buf: Option<u32>,
582    p_v: u32,
583    _p_o_unused: u32,
584) {
585    let var_sk = m.alloc_id();
586    m.emit_variable(b.ty_ptr_func_uint, var_sk, STORAGE_CLASS_FUNCTION);
587    m.emit_store(var_sk, b.c_uint_0);
588
589    let lbl_hdr = m.alloc_id();
590    let lbl_body = m.alloc_id();
591    let lbl_cont = m.alloc_id();
592    let lbl_merge = m.alloc_id();
593
594    m.emit_branch(lbl_hdr);
595
596    m.emit_label(lbl_hdr);
597    let sk_val = m.alloc_id();
598    m.emit_load(b.ty_uint, sk_val, var_sk);
599    let cond = m.alloc_id();
600    m.emit(OP_U_LESS_THAN, &[b.ty_bool, cond, sk_val, c_seq_kv]);
601    m.emit_loop_merge(lbl_merge, lbl_cont);
602    m.emit_branch_conditional(cond, lbl_body, lbl_merge);
603
604    m.emit_label(lbl_body);
605
606    let lbl_compute = m.alloc_id();
607    let lbl_skip = m.alloc_id();
608    if causal {
609        let sk_gt_sq = m.alloc_id();
610        m.emit(OP_U_LESS_THAN, &[b.ty_bool, sk_gt_sq, sq, sk_val]);
611        m.emit_selection_merge(lbl_skip);
612        m.emit_branch_conditional(sk_gt_sq, lbl_skip, lbl_compute);
613    } else {
614        m.emit_branch(lbl_compute);
615    }
616
617    m.emit_label(lbl_compute);
618
619    // k_off = kv_base + sk * head_dim
620    let sk_hd = m.alloc_id();
621    m.emit(OP_I_MUL, &[b.ty_uint, sk_hd, sk_val, c_head_dim]);
622    let k_off = m.alloc_id();
623    m.emit(OP_I_ADD, &[b.ty_uint, k_off, kv_base, sk_hd]);
624
625    // Inner dot product loop
626    let var_d = m.alloc_id();
627    m.emit_variable(b.ty_ptr_func_uint, var_d, STORAGE_CLASS_FUNCTION);
628    m.emit_store(var_d, b.c_uint_0);
629    let var_dot = m.alloc_id();
630    m.emit_variable(b.ty_ptr_func_float, var_dot, STORAGE_CLASS_FUNCTION);
631    m.emit_store(var_dot, b.c_float_0);
632
633    let lbl_d_hdr = m.alloc_id();
634    let lbl_d_body = m.alloc_id();
635    let lbl_d_cont = m.alloc_id();
636    let lbl_d_merge = m.alloc_id();
637
638    m.emit_branch(lbl_d_hdr);
639
640    m.emit_label(lbl_d_hdr);
641    let d_val = m.alloc_id();
642    m.emit_load(b.ty_uint, d_val, var_d);
643    let d_cond = m.alloc_id();
644    m.emit(OP_U_LESS_THAN, &[b.ty_bool, d_cond, d_val, c_head_dim]);
645    m.emit_loop_merge(lbl_d_merge, lbl_d_cont);
646    m.emit_branch_conditional(d_cond, lbl_d_body, lbl_d_merge);
647
648    m.emit_label(lbl_d_body);
649    let q_idx = m.alloc_id();
650    m.emit(OP_I_ADD, &[b.ty_uint, q_idx, q_base, d_val]);
651    let q_ptr = m.alloc_id();
652    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, q_ptr, p_q, q_idx);
653    let q_val = m.alloc_id();
654    m.emit_load(b.ty_float, q_val, q_ptr);
655    let k_idx = m.alloc_id();
656    m.emit(OP_I_ADD, &[b.ty_uint, k_idx, k_off, d_val]);
657    let k_ptr = m.alloc_id();
658    m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, k_ptr, p_k, k_idx);
659    let k_val = m.alloc_id();
660    m.emit_load(b.ty_float, k_val, k_ptr);
661
662    let qk_prod = m.alloc_id();
663    m.emit(OP_F_MUL, &[b.ty_float, qk_prod, q_val, k_val]);
664    let old_dot = m.alloc_id();
665    m.emit_load(b.ty_float, old_dot, var_dot);
666    let new_dot = m.alloc_id();
667    m.emit(OP_F_ADD, &[b.ty_float, new_dot, old_dot, qk_prod]);
668    m.emit_store(var_dot, new_dot);
669
670    m.emit_branch(lbl_d_cont);
671    m.emit_label(lbl_d_cont);
672    let d_inc = m.alloc_id();
673    m.emit(OP_I_ADD, &[b.ty_uint, d_inc, d_val, b.c_uint_1]);
674    m.emit_store(var_d, d_inc);
675    m.emit_branch(lbl_d_hdr);
676
677    m.emit_label(lbl_d_merge);
678
679    // score = dot * scale
680    let dot_final = m.alloc_id();
681    m.emit_load(b.ty_float, dot_final, var_dot);
682    let score = m.alloc_id();
683    m.emit(OP_F_MUL, &[b.ty_float, score, dot_final, c_scale]);
684
685    if is_max_pass {
686        // accum = fmax(accum, score)
687        let old_acc = m.alloc_id();
688        m.emit_load(b.ty_float, old_acc, accum_var);
689        let new_acc = m.alloc_id();
690        m.emit_opencl_ext(
691            b.opencl_ext,
692            b.ty_float,
693            new_acc,
694            OPENCL_FMAX,
695            &[old_acc, score],
696        );
697        m.emit_store(accum_var, new_acc);
698    } else {
699        // w = exp(score - max_score)
700        let max_id = max_val.unwrap_or(b.c_float_0);
701        let score_shifted = m.alloc_id();
702        m.emit(OP_F_SUB, &[b.ty_float, score_shifted, score, max_id]);
703        let w = m.alloc_id();
704        m.emit_opencl_ext(b.opencl_ext, b.ty_float, w, OPENCL_EXP, &[score_shifted]);
705
706        // sum_exp += w
707        let old_sum = m.alloc_id();
708        m.emit_load(b.ty_float, old_sum, accum_var);
709        let new_sum = m.alloc_id();
710        m.emit(OP_F_ADD, &[b.ty_float, new_sum, old_sum, w]);
711        m.emit_store(accum_var, new_sum);
712
713        // Accumulate weighted V
714        if let Some(o_buf_id) = o_buf {
715            let v_off = m.alloc_id();
716            m.emit(OP_I_ADD, &[b.ty_uint, v_off, kv_base, sk_hd]);
717
718            let var_d3 = m.alloc_id();
719            m.emit_variable(b.ty_ptr_func_uint, var_d3, STORAGE_CLASS_FUNCTION);
720            m.emit_store(var_d3, b.c_uint_0);
721
722            let lbl_d3_hdr = m.alloc_id();
723            let lbl_d3_body = m.alloc_id();
724            let lbl_d3_cont = m.alloc_id();
725            let lbl_d3_merge = m.alloc_id();
726
727            m.emit_branch(lbl_d3_hdr);
728
729            m.emit_label(lbl_d3_hdr);
730            let d3_val = m.alloc_id();
731            m.emit_load(b.ty_uint, d3_val, var_d3);
732            let d3_cond = m.alloc_id();
733            m.emit(OP_U_LESS_THAN, &[b.ty_bool, d3_cond, d3_val, c_head_dim]);
734            m.emit_loop_merge(lbl_d3_merge, lbl_d3_cont);
735            m.emit_branch_conditional(d3_cond, lbl_d3_body, lbl_d3_merge);
736
737            m.emit_label(lbl_d3_body);
738            let v_idx = m.alloc_id();
739            m.emit(OP_I_ADD, &[b.ty_uint, v_idx, v_off, d3_val]);
740            let v_ptr = m.alloc_id();
741            m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, v_ptr, p_v, v_idx);
742            let v_val = m.alloc_id();
743            m.emit_load(b.ty_float, v_val, v_ptr);
744            let wv = m.alloc_id();
745            m.emit(OP_F_MUL, &[b.ty_float, wv, w, v_val]);
746
747            let o_idx = m.alloc_id();
748            m.emit(OP_I_ADD, &[b.ty_uint, o_idx, q_base, d3_val]);
749            let o_ptr = m.alloc_id();
750            m.emit_in_bounds_ptr_access_chain(b.ty_ptr_cross_float, o_ptr, o_buf_id, o_idx);
751            let o_old = m.alloc_id();
752            m.emit_load(b.ty_float, o_old, o_ptr);
753            let o_new = m.alloc_id();
754            m.emit(OP_F_ADD, &[b.ty_float, o_new, o_old, wv]);
755            m.emit_store(o_ptr, o_new);
756
757            m.emit_branch(lbl_d3_cont);
758            m.emit_label(lbl_d3_cont);
759            let d3_inc = m.alloc_id();
760            m.emit(OP_I_ADD, &[b.ty_uint, d3_inc, d3_val, b.c_uint_1]);
761            m.emit_store(var_d3, d3_inc);
762            m.emit_branch(lbl_d3_hdr);
763
764            m.emit_label(lbl_d3_merge);
765        }
766    }
767
768    m.emit_branch(lbl_skip);
769    m.emit_label(lbl_skip);
770    m.emit_branch(lbl_cont);
771
772    m.emit_label(lbl_cont);
773    let sk_inc = m.alloc_id();
774    m.emit(OP_I_ADD, &[b.ty_uint, sk_inc, sk_val, b.c_uint_1]);
775    m.emit_store(var_sk, sk_inc);
776    m.emit_branch(lbl_hdr);
777
778    m.emit_label(lbl_merge);
779}
780
781// ─── Tests ──────────────────────────────────────────────────
782
783#[cfg(test)]
784mod tests {
785    use super::*;
786    use crate::spirv::SPIRV_MAGIC;
787
788    fn check_valid_spirv(words: &[u32]) {
789        assert!(words.len() >= 5, "too short for SPIR-V header");
790        assert_eq!(words[0], SPIRV_MAGIC, "bad magic");
791        assert!(words[3] > 0, "ID bound must be > 0");
792        assert_eq!(words[4], 0, "schema must be 0");
793    }
794
795    #[test]
796    fn conv2d_spirv_valid() {
797        let words = conv2d_spirv(1, 3, 8, 8, 16, 3, 3, 6, 6, 1, 1, 0, 0);
798        check_valid_spirv(&words);
799    }
800
801    #[test]
802    fn conv2d_spirv_with_padding() {
803        let words = conv2d_spirv(2, 1, 5, 5, 4, 3, 3, 5, 5, 1, 1, 1, 1);
804        check_valid_spirv(&words);
805    }
806
807    #[test]
808    fn conv2d_spirv_1x1() {
809        let words = conv2d_spirv(1, 3, 4, 4, 8, 1, 1, 4, 4, 1, 1, 0, 0);
810        check_valid_spirv(&words);
811    }
812
813    #[test]
814    fn attention_spirv_valid() {
815        let words = attention_spirv(2, 4, 4, 8, 0.125, false);
816        check_valid_spirv(&words);
817    }
818
819    #[test]
820    fn attention_spirv_causal() {
821        let words = attention_spirv(1, 8, 8, 16, 0.25, true);
822        check_valid_spirv(&words);
823    }
824
825    #[test]
826    fn attention_spirv_magic_number() {
827        let words = attention_spirv(1, 4, 4, 8, 0.125, false);
828        assert_eq!(words[0], 0x07230203);
829    }
830
831    #[test]
832    fn conv2d_spirv_magic_number() {
833        let words = conv2d_spirv(1, 1, 4, 4, 1, 1, 1, 4, 4, 1, 1, 0, 0);
834        assert_eq!(words[0], 0x07230203);
835    }
836}
oxicuda_levelzero/spirv_nn.rs

oxicuda_levelzero/
spirv_nn.rs
