pagesift_gpu/
ops.rs

1//! 与 `pagesift_infer::ops` 同语义的 GPU 算子。两层入口：
2//! - `*_device`：输入输出都在显存，连续多个算子时不来回拷贝；
3//! - `matmul`、`conv2d`、`activate`：收发主机内存，参数和返回布局同 CPU 版，方便直接替换与对照。
4//!
5//! 卷积的走法：每组一个输入通道（深度卷积）用直接内核；1×1、步长 1、不补边时输入本身就是 GEMM 的
6//! B 矩阵；其余先 im2col 再批量 SGEMM（批次 = groups），最后加偏置。
7
8use pagesift_infer::{Activation, Tensor};
9
10use crate::Gpu;
11use crate::buffer::{Args, DeviceBuffer};
12use crate::driver::DevPtr;
13
14/// 一维内核每块的线程数，与 `kernels/src/lib.rs` 的约定一致。
15pub(crate) const BLOCK_1D: u32 = 256; // [策略] 256 线程/块：Pascal 每 SM 可驻留 2048 线程，8 块就能占满
16/// SGEMM 每块算的 C 子块边长与线程块边长，与 `kernels/src/gemm.rs` 的 BM/BN 和 16×16 线程一致。
17const GEMM_TILE: usize = 64; // [推导] kernels/src/gemm.rs：BM = BN = 64
18const GEMM_SIDE: u32 = 16; // [推导] kernels/src/gemm.rs：每块 16×16 线程
19/// 网格 y、z 维的上限。
20const GRID_YZ_MAX: usize = 65535; // [协议] CUDA 编程指南“算力表”：gridDim.y、gridDim.z ≤ 65535
21
22/// 一次卷积的全部尺寸，与 `pagesift_infer::ops::conv2d` 的参数一一对应：
23/// 输入 `[c, h, w]`，权重 `[o, c/groups, kh, kw]`。
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct ConvShape {
26    pub c: usize,
27    pub h: usize,
28    pub w: usize,
29    pub o: usize,
30    pub kh: usize,
31    pub kw: usize,
32    pub stride: (usize, usize),
33    pub pad: (usize, usize),
34    pub dilation: (usize, usize),
35    pub groups: usize,
36}
37
38impl ConvShape {
39    /// 输出的 (OH, OW)，公式同 PyTorch。
40    pub fn out_hw(&self) -> (usize, usize) {
41        let oh =
42            (self.h + 2 * self.pad.0 - self.dilation.0 * (self.kh - 1) - 1) / self.stride.0 + 1;
43        let ow =
44            (self.w + 2 * self.pad.1 - self.dilation.1 * (self.kw - 1) - 1) / self.stride.1 + 1;
45        (oh, ow)
46    }
47
48    /// 要不要先 im2col：深度卷积走直接内核，1×1、步长 1、不补边时输入本身就是 GEMM 的 B。
49    pub(crate) fn needs_col(&self) -> bool {
50        let depthwise = self.c == self.groups;
51        let direct = self.kh == 1 && self.kw == 1 && self.stride == (1, 1) && self.pad == (0, 0);
52        !depthwise && !direct
53    }
54
55    /// 能不能走隐式卷积：步长 1、无空洞、不分组的普通卷积。输入补边后本身就是 GEMM 的 B，按核位置偏移指针，
56    /// 拆成 kh·kw 次累加式 SGEMM。比 im2col 省下展开的显存（9×9 时差约 80 倍）和搬数据的时间。
57    pub(crate) fn implicit(&self) -> bool {
58        self.groups == 1
59            && self.stride == (1, 1)
60            && self.dilation == (1, 1)
61            && self.c > 1
62            && self.kh * self.kw > 1
63    }
64
65    /// 单个样本 im2col 展开后的元素数。
66    pub(crate) fn col_len(&self) -> usize {
67        let (oh, ow) = self.out_hw();
68        self.c * self.kh * self.kw * oh * ow
69    }
70
71    /// 与 CPU 版相同的形状检查（不合法就 panic，属于调用方的编程错误）。
72    fn check(&self) {
73        assert!(
74            self.groups > 0 && self.c.is_multiple_of(self.groups),
75            "卷积输入通道 {} 不能被 groups {} 整除",
76            self.c,
77            self.groups
78        );
79        assert_eq!(
80            self.o % self.groups,
81            0,
82            "输出通道 {} 不能被 groups {} 整除",
83            self.o,
84            self.groups
85        );
86        assert!(
87            self.stride.0 > 0 && self.stride.1 > 0 && self.dilation.0 > 0 && self.dilation.1 > 0
88        );
89        assert!(self.kh > 0 && self.kw > 0, "卷积核为空");
90        assert!(
91            self.h + 2 * self.pad.0 > self.dilation.0 * (self.kh - 1)
92                && self.w + 2 * self.pad.1 > self.dilation.1 * (self.kw - 1),
93            "卷积核比补边后的输入还大"
94        );
95    }
96}
97
98/// `n` 个 f32 的字节数，用于设备地址的偏移。
99fn bytes_of(n: usize) -> u64 {
100    crate::pool::bytes(n)
101}
102
103pub(crate) fn blocks_1d(n: usize) -> u32 {
104    u32::try_from(n.div_ceil(BLOCK_1D as usize)).expect("一维网格超出 u32")
105}
106
107impl Gpu {
108    /// 批量 SGEMM：对 z ∈ 0..batch，`C_z[m×n] = A_z[m×k] · B_z[k×n]`，`A_z = a + z·sa`（单位：元素）。
109    #[allow(clippy::too_many_arguments)]
110    pub(crate) fn sgemm_raw(
111        &self,
112        a: DevPtr,
113        b: DevPtr,
114        c: DevPtr,
115        (m, n, k): (usize, usize, usize),
116        batch: usize,
117        strides: (usize, usize, usize),
118    ) -> Result<(), String> {
119        self.sgemm_ext(a, b, c, (m, n, k), n, false, batch, strides)
120    }
121
122    /// 同 [`Gpu::sgemm_raw`]，B 的行距是 `ldb`，`accumulate` 时结果累加进 C。
123    #[allow(clippy::too_many_arguments)]
124    pub(crate) fn sgemm_ext(
125        &self,
126        a: DevPtr,
127        b: DevPtr,
128        c: DevPtr,
129        (m, n, k): (usize, usize, usize),
130        ldb: usize,
131        accumulate: bool,
132        batch: usize,
133        (sa, sb, sc): (usize, usize, usize),
134    ) -> Result<(), String> {
135        let gy = m.div_ceil(GEMM_TILE);
136        if gy > GRID_YZ_MAX || batch > GRID_YZ_MAX {
137            return Err(format!("SGEMM grid too large: M = {m}, batch {batch}"));
138        }
139        let grid = (
140            u32::try_from(n.div_ceil(GEMM_TILE)).map_err(|e| e.to_string())?,
141            gy as u32,
142            batch as u32,
143        );
144        let args = Args::default()
145            .raw_ptr(a)
146            .raw_ptr(b)
147            .raw_ptr(c)
148            .u32(m)
149            .u32(n)
150            .u32(k)
151            .u32(ldb)
152            .u32(accumulate as usize)
153            .u64(sa)
154            .u64(sb)
155            .u64(sc);
156        self.launch(self.k.sgemm, grid, (GEMM_SIDE, GEMM_SIDE, 1), args)
157    }
158
159    /// `c[m×n] = a[m×k] · b[k×n]`，三者都在显存、行优先。
160    pub fn gemm_device(
161        &self,
162        a: &DeviceBuffer<'_>,
163        b: &DeviceBuffer<'_>,
164        c: &mut DeviceBuffer<'_>,
165        m: usize,
166        n: usize,
167        k: usize,
168    ) -> Result<(), String> {
169        assert!(
170            a.len() == m * k && b.len() == k * n && c.len() == m * n,
171            "矩阵乘的显存长度与 m、n、k 不符"
172        );
173        self.sgemm_raw(a.ptr(), b.ptr(), c.ptr(), (m, n, k), 1, (0, 0, 0))
174    }
175
176    /// 二维卷积，输入、权重、偏置、输出都在显存。返回 `[o, OH, OW]` 的新显存。
177    pub fn conv2d_device(
178        &self,
179        x: &DeviceBuffer<'_>,
180        w: &DeviceBuffer<'_>,
181        bias: Option<&DeviceBuffer<'_>>,
182        s: &ConvShape,
183    ) -> Result<DeviceBuffer<'_>, String> {
184        let (oh, ow) = s.out_hw();
185        let mut out = self.alloc(s.o * oh * ow)?;
186        self.conv2d_device_into(x, w, bias, s, &mut out)?;
187        Ok(out)
188    }
189
190    /// 同 [`Gpu::conv2d_device`]，结果写进调用方给的显存 `out`（长度 `o·OH·OW`），省一次分配；
191    /// 连续推理时可以复用输出缓冲。im2col 的临时显存仍在每次调用里分配。
192    pub fn conv2d_device_into(
193        &self,
194        x: &DeviceBuffer<'_>,
195        w: &DeviceBuffer<'_>,
196        bias: Option<&DeviceBuffer<'_>>,
197        s: &ConvShape,
198        out: &mut DeviceBuffer<'_>,
199    ) -> Result<(), String> {
200        s.check();
201        let cg = s.c / s.groups;
202        assert_eq!(x.len(), s.c * s.h * s.w, "输入长度与形状不符");
203        assert_eq!(w.len(), s.o * cg * s.kh * s.kw, "权重长度与形状不符");
204        assert!(
205            bias.is_none_or(|b| b.len() == s.o),
206            "偏置长度与输出通道数不符"
207        );
208        let (oh, ow) = s.out_hw();
209        let ohw = oh * ow;
210        assert_eq!(out.len(), s.o * ohw, "输出显存长度与形状不符");
211        // 卷积内核用 u32 拆下标（kernels/src/conv.rs），输出与 im2col 展开的元素数都要小于 2³²
212        let col_len = s.c * s.kh * s.kw * ohw;
213        if out.len() > u32::MAX as usize || col_len > u32::MAX as usize {
214            return Err(format!(
215                "convolution too large: {} output and {col_len} unfolded elements exceed 2^32",
216                out.len()
217            ));
218        }
219        let col = if s.needs_col() {
220            Some(self.alloc(col_len)?)
221        } else {
222            None
223        };
224        self.conv_raw(
225            x.ptr(),
226            w.ptr(),
227            bias.map(|b| b.ptr()),
228            s,
229            1,
230            out.ptr(),
231            col.as_ref().map(|c| c.ptr()),
232        )?;
233        // col 在这里释放：cuMemFree 会等用到它的内核跑完
234        drop(col);
235        Ok(())
236    }
237
238    /// 一批 `n` 个样本的卷积，全部用设备地址：`x` 是 `[n, c, h, w]`，`out` 是 `[n, o, OH, OW]`，
239    /// 权重和偏置各样本共用；`col` 是 im2col 的临时缓冲（`n·col_len` 个元素），不需要时为 None。
240    #[allow(clippy::too_many_arguments)]
241    pub(crate) fn conv_raw(
242        &self,
243        x: DevPtr,
244        w: DevPtr,
245        bias: Option<DevPtr>,
246        s: &ConvShape,
247        n: usize,
248        out: DevPtr,
249        col: Option<DevPtr>,
250    ) -> Result<(), String> {
251        s.check();
252        let cg = s.c / s.groups;
253        let og = s.o / s.groups;
254        let (oh, ow) = s.out_hw();
255        let ohw = oh * ow;
256        let total = n * s.o * ohw;
257        let geom = |a: Args| {
258            a.u32(s.h)
259                .u32(s.w)
260                .u32(s.kh)
261                .u32(s.kw)
262                .u32(s.stride.0)
263                .u32(s.stride.1)
264                .u32(s.pad.0)
265                .u32(s.pad.1)
266                .u32(s.dilation.0)
267                .u32(s.dilation.1)
268                .u32(oh)
269                .u32(ow)
270        };
271        if cg == 1 {
272            // 深度卷积：直接内核，偏置在内核里加；权重、偏置按 oc % o 取，各样本共用
273            let args = Args::default()
274                .raw_ptr(x)
275                .raw_ptr(w)
276                .raw_ptr(bias.unwrap_or(0))
277                .raw_ptr(out)
278                .u32(n * s.o)
279                .u32(og)
280                .u32(s.o)
281                .u32(bias.is_some() as usize);
282            return self.launch(
283                self.k.depthwise,
284                (blocks_1d(total), 1, 1),
285                (BLOCK_1D, 1, 1),
286                geom(args),
287            );
288        }
289        let kk = cg * s.kh * s.kw; // GEMM 的 K：每组一个输出元素用到的输入个数
290        let b = match col {
291            None => x, // 1×1：x 的 [c, h·w] 本身就是 B
292            Some(col) => {
293                let len = n * s.col_len();
294                let args = geom(Args::default().raw_ptr(x).raw_ptr(col).u32(n * s.c));
295                self.launch(
296                    self.k.im2col,
297                    (blocks_1d(len), 1, 1),
298                    (BLOCK_1D, 1, 1),
299                    args,
300                )?;
301                col
302            }
303        };
304        // 组 g：A = 权重第 g·og 行起的 [og, kk]；B = 第 g·kk 行起的 [kk, ohw]；C = 输出第 g·og 个通道起
305        let (sb, sc) = (kk * ohw, og * ohw);
306        if s.groups == 1 {
307            // 批次沿 z 维：权重不动（sa = 0），B、C 按样本走
308            self.sgemm_raw(w, b, out, (og, ohw, kk), n, (0, sb, sc))?;
309        } else {
310            // 分组又有批次：每个样本一次，A 的步长已被组占用
311            for z in 0..n {
312                let off = |per: usize| bytes_of(z * s.groups * per);
313                self.sgemm_raw(
314                    w,
315                    b + off(sb),
316                    out + off(sc),
317                    (og, ohw, kk),
318                    s.groups,
319                    (og * kk, sb, sc),
320                )?;
321            }
322        }
323        if let Some(bias) = bias {
324            let args = Args::default()
325                .raw_ptr(out)
326                .raw_ptr(bias)
327                .u32(ohw)
328                .u32(s.o)
329                .u32(total);
330            self.launch(
331                self.k.add_bias,
332                (blocks_1d(total), 1, 1),
333                (BLOCK_1D, 1, 1),
334                args,
335            )?;
336        }
337        Ok(())
338    }
339
340    /// 原地激活，语义同 `Activation::apply`。
341    pub fn activate_device(&self, x: &mut DeviceBuffer<'_>, act: Activation) -> Result<(), String> {
342        let f = match act {
343            Activation::Identity => return Ok(()),
344            Activation::Relu => self.k.relu,
345            Activation::Silu => self.k.silu,
346            Activation::Gelu => self.k.gelu,
347            Activation::HardSigmoid => self.k.hardsigmoid,
348            Activation::Sigmoid => self.k.sigmoid,
349        };
350        let args = Args::default().ptr(x).u32(x.len());
351        self.launch(f, (blocks_1d(x.len()), 1, 1), (BLOCK_1D, 1, 1), args)
352    }
353
354    /// 原地加偏置：`x` 是 `[C, hw]`，`x[c, i] += bias[c]`。
355    pub fn add_bias_device(
356        &self,
357        x: &mut DeviceBuffer<'_>,
358        bias: &DeviceBuffer<'_>,
359        hw: usize,
360    ) -> Result<(), String> {
361        assert_eq!(x.len(), bias.len() * hw, "偏置长度与张量形状不符");
362        let args = Args::default()
363            .ptr(x)
364            .ptr(bias)
365            .u32(hw)
366            .u32(bias.len())
367            .u32(x.len());
368        self.launch(
369            self.k.add_bias,
370            (blocks_1d(x.len()), 1, 1),
371            (BLOCK_1D, 1, 1),
372            args,
373        )
374    }
375
376    /// 矩阵乘（同 `pagesift_infer::ops::matmul`）：`a` 是 `[M, K]`，`b` 是 `[K, N]`，返回 `[M, N]` 行优先。
377    pub fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Vec<f32>, String> {
378        let (m, k) = a.tc();
379        let (kb, n) = b.tc();
380        assert_eq!(k, kb, "矩阵乘的内维不同");
381        let (da, db) = (self.upload(&a.data)?, self.upload(&b.data)?);
382        let mut dc = self.alloc(m * n)?;
383        self.gemm_device(&da, &db, &mut dc, m, n, k)?;
384        dc.to_vec()
385    }
386
387    /// 二维卷积（同 `pagesift_infer::ops::conv2d`）：`x` 是 `[C, H, W]`，`w` 是 `[O, C/groups, kh, kw]`；
388    /// 返回 `[O, OH, OW]` 行优先的数据。
389    #[allow(clippy::too_many_arguments)]
390    pub fn conv2d(
391        &self,
392        x: &Tensor,
393        w: &Tensor,
394        bias: Option<&[f32]>,
395        stride: (usize, usize),
396        pad: (usize, usize),
397        dilation: (usize, usize),
398        groups: usize,
399    ) -> Result<Vec<f32>, String> {
400        let (c, h, wd) = x.chw();
401        // [协议] 卷积权重是 4 维 [O, C/groups, kh, kw]（PyTorch Conv2d.weight 的布局）
402        assert_eq!(w.shape.len(), 4, "卷积权重要 [O, C/groups, kh, kw]");
403        let s = ConvShape {
404            c,
405            h,
406            w: wd,
407            o: w.shape[0],
408            kh: w.shape[2],
409            kw: w.shape[3],
410            stride,
411            pad,
412            dilation,
413            groups,
414        };
415        assert_eq!(
416            c,
417            w.shape[1] * groups,
418            "卷积输入通道 {c} 与权重 {}×{groups} 不符",
419            w.shape[1]
420        );
421        let dx = self.upload(&x.data)?;
422        let dw = self.upload(&w.data)?;
423        let db = bias.map(|b| self.upload(b)).transpose()?;
424        self.conv2d_device(&dx, &dw, db.as_ref(), &s)?.to_vec()
425    }
426
427    /// 原地激活一段主机内存（同 `Activation::apply_all`）。
428    pub fn activate(&self, x: &mut [f32], act: Activation) -> Result<(), String> {
429        if act == Activation::Identity {
430            return Ok(());
431        }
432        let mut d = self.upload(x)?;
433        self.activate_device(&mut d, act)?;
434        d.copy_to_host(x)
435    }
436}
pagesift_gpu/ops.rs

pagesift_gpu/
ops.rs
