cherenkov/qwen4_exp/gpu/
dispatch.rs

1//! Buffer binding, dispatch geometry, dense projections, and transfers.
2
3use super::*;
4
5impl Gpu<'_> {
6    // ---- small helpers ----
7
8    pub(super) fn bind(&self, enc: &Enc, index: usize, buf: &Buf, offset: usize) {
9        unsafe { enc.setBuffer_offset_atIndex(Some(buf), offset, index) };
10    }
11
12    pub(super) fn dispatch(
13        &self,
14        enc: &Enc,
15        pso: &Pso,
16        setup: impl FnOnce(&Enc),
17        grid: usize,
18        tg: usize,
19        threadgroups: bool,
20    ) {
21        enc.setComputePipelineState(pso);
22        setup(enc);
23        self.dispatch_count.set(self.dispatch_count.get() + 1);
24
25        let g = MTLSize {
26            width: grid,
27            height: 1,
28            depth: 1,
29        };
30        let t = MTLSize {
31            width: tg,
32            height: 1,
33            depth: 1,
34        };
35
36        if threadgroups {
37            enc.dispatchThreadgroups_threadsPerThreadgroup(g, t);
38        } else {
39            enc.dispatchThreads_threadsPerThreadgroup(g, t);
40        }
41    }
42
43    /// Half even/odd streams + group sums of `nb` rows of `x` (f32,
44    /// `in_dim` wide, starting at byte offset `x_off`) into `set`.
45    pub(super) fn prep_h(
46        &self,
47        enc: &Enc,
48        x: &Buf,
49        x_off: usize,
50        in_dim: u32,
51        nb: usize,
52        set: &HalfSet,
53    ) {
54        let n2 = in_dim / 2;
55        let nbu = nb as u32;
56
57        self.dispatch(
58            enc,
59            &self.pipes.prep_h,
60            |e| {
61                self.bind(e, 0, x, x_off);
62                self.bind(e, 1, &set.xe, 0);
63                self.bind(e, 2, &set.xo, 0);
64                set_bytes(e, 3, &n2);
65                set_bytes(e, 4, &nbu);
66                self.bind(e, 5, &set.xsum, 0);
67            },
68            nb * n2 as usize,
69            256,
70            false,
71        );
72    }
73
74    /// `y[b] = W x[b]` over `nb` prepped rows (chunks of up to 8 rows).
75    pub(super) fn qmv_h(&self, enc: &Enc, q: &Q, y: &Buf, nb: usize, set: &HalfSet) {
76        let p = QmvParams {
77            out_dim: q.out,
78            in_dim: q.inp,
79        };
80        let mut r0 = 0;
81
82        while r0 < nb {
83            let n = (nb - r0).min(8);
84            let pipe = if n <= 3 {
85                &self.pipes.qmv_h
86            } else {
87                &self.pipes.qmv_hn
88            };
89            let nu = n as u32;
90            let x_off = r0 * (q.inp as usize / 2) * 2;
91            let s_off = r0 * (q.inp as usize / 32) * 4;
92            let y_off = r0 * q.out as usize * 4;
93
94            self.dispatch(
95                enc,
96                pipe,
97                |e| {
98                    self.bind(e, 0, &self.dense, q.w);
99                    self.bind(e, 1, &self.dense, q.s);
100                    self.bind(e, 2, &self.dense, q.b);
101                    self.bind(e, 3, &set.xe, x_off);
102                    self.bind(e, 4, &set.xo, x_off);
103                    self.bind(e, 5, y, y_off);
104                    set_bytes(e, 6, &p);
105                    set_bytes(e, 7, &nu);
106                    self.bind(e, 8, &set.xsum, s_off);
107                },
108                (q.out as usize).div_ceil(4),
109                128,
110                true,
111            );
112
113            r0 += n;
114        }
115    }
116
117    pub(super) fn zero(&self, enc: &Enc, x: &Buf, n: u32) {
118        self.dispatch(
119            enc,
120            &self.pipes.zero,
121            |e| {
122                self.bind(e, 0, x, 0);
123                set_bytes(e, 1, &n);
124            },
125            n as usize,
126            256,
127            false,
128        );
129    }
130
131    pub(super) fn add(&self, enc: &Enc, x: &Buf, r: &Buf, n: usize) {
132        self.dispatch(
133            enc,
134            &self.pipes.add,
135            |e| {
136                self.bind(e, 0, x, 0);
137                self.bind(e, 1, r, 0);
138            },
139            n,
140            256,
141            false,
142        );
143    }
144
145    pub(super) fn read_u32(&self, buf: &Buf, n: usize) -> Vec<u32> {
146        let ptr = buf.contents().cast::<u32>();
147
148        unsafe { std::slice::from_raw_parts(ptr.as_ptr(), n) }.to_vec()
149    }
150
151    pub(super) fn read_f32(&self, buf: &Buf, n: usize) -> Vec<f32> {
152        let ptr = buf.contents().cast::<f32>();
153
154        unsafe { std::slice::from_raw_parts(ptr.as_ptr(), n) }.to_vec()
155    }
156}
cherenkov/qwen4_exp/gpu/dispatch.rs

cherenkov/qwen4_exp/gpu/
dispatch.rs
