gpu_bench/
gpu_bench.rs

1//! GPU 对 CPU 朴素实现（`pagesift_infer::ops`）的测速：SGEMM 与几种 OCR 模型里常见的卷积形状。
2//!
3//!   cargo run --release -p pagesift-gpu --example gpu_bench
4//!
5//! 每项给三个时间：CPU（单线程朴素实现，预热 1 次后取 3 次平均）；GPU 内核（输入、权重、输出已在显存，
6//! CUDA 事件计时，预热 1 次后取 20 次平均；卷积含 im2col 临时显存的分配与释放）；GPU 端到端（含上传输入与权重、下载结果，墙钟计时，
7//! 预热 1 次后取 5 次平均）。倍数 = CPU / GPU。
8
9use std::time::Instant;
10
11use pagesift_gpu::{ConvShape, Gpu};
12use pagesift_infer::{Tensor, ops};
13
14struct Rng(u64);
15
16impl Rng {
17    fn tensor(&mut self, shape: Vec<usize>) -> Tensor {
18        let n = shape.iter().product();
19        let data = (0..n)
20            .map(|_| {
21                self.0 ^= self.0 << 13;
22                self.0 ^= self.0 >> 7;
23                self.0 ^= self.0 << 17;
24                (self.0 >> 40) as f32 / (1u64 << 24) as f32 - 0.5
25            })
26            .collect();
27        Tensor::new(shape, data)
28    }
29}
30
31/// 预热一次后跑 `reps` 次，返回平均毫秒。
32fn wall_ms(reps: usize, mut f: impl FnMut()) -> f64 {
33    f();
34    let t = Instant::now();
35    for _ in 0..reps {
36        f();
37    }
38    t.elapsed().as_secs_f64() * 1e3 / reps as f64
39}
40
41/// 预热一次后跑 `reps` 次，用 CUDA 事件量 GPU 上的平均毫秒。
42fn gpu_ms(g: &Gpu, reps: usize, mut f: impl FnMut()) -> f64 {
43    f();
44    g.synchronize().unwrap();
45    let t = g.timer().unwrap();
46    for _ in 0..reps {
47        f();
48    }
49    t.stop_ms().unwrap() as f64 / reps as f64
50}
51
52fn row(name: &str, gflop: f64, cpu: f64, kern: f64, e2e: f64) {
53    println!(
54        "| {name} | {cpu:.2} | {kern:.3} | {e2e:.2} | {:.0} | {:.0} | {:.1} | {:.0} |",
55        cpu / kern,
56        cpu / e2e,
57        gflop / (cpu / 1e3),
58        gflop / (kern / 1e3),
59    );
60}
61
62fn main() {
63    let g = match Gpu::open() {
64        Ok(g) => g,
65        Err(e) => {
66            println!("没有可用的 GPU：{e}");
67            return;
68        }
69    };
70    println!("GPU：{}（算力 {:?}）\n", g.name(), g.compute_capability());
71    println!(
72        "| 项目 | CPU ms | GPU 内核 ms | GPU 端到端 ms | 倍数（内核） | 倍数（端到端） | CPU GFLOPS | GPU GFLOPS |"
73    );
74    println!("|---|---:|---:|---:|---:|---:|---:|---:|");
75    let mut rng = Rng(0x9e37_79b9_7f4a_7c15);
76
77    for s in [512, 1024, 2048] {
78        let a = rng.tensor(vec![s, s]);
79        let b = rng.tensor(vec![s, s]);
80        let gflop = 2.0 * (s as f64).powi(3) / 1e9;
81        let cpu = wall_ms(3, || drop(ops::matmul(&a, &b)));
82        let (da, db) = (g.upload(&a.data).unwrap(), g.upload(&b.data).unwrap());
83        let mut dc = g.alloc(s * s).unwrap();
84        let kern = gpu_ms(&g, 20, || {
85            g.gemm_device(&da, &db, &mut dc, s, s, s).unwrap()
86        });
87        let e2e = wall_ms(5, || drop(g.matmul(&a, &b).unwrap()));
88        row(&format!("SGEMM {s}³"), gflop, cpu, kern, e2e);
89    }
90
91    // (说明, [C, H, W], [O, kh, kw], stride, pad, groups)
92    type Conv = (
93        &'static str,
94        [usize; 3],
95        [usize; 3],
96        (usize, usize),
97        (usize, usize),
98        usize,
99    );
100    let convs: [Conv; 6] = [
101        (
102            "3×3 s2 首层 3→32 @640²",
103            [3, 640, 640],
104            [32, 3, 3],
105            (2, 2),
106            (1, 1),
107            1,
108        ),
109        (
110            "1×1 64→128 @160²",
111            [64, 160, 160],
112            [128, 1, 1],
113            (1, 1),
114            (0, 0),
115            1,
116        ),
117        (
118            "3×3 64→64 @80²",
119            [64, 80, 80],
120            [64, 3, 3],
121            (1, 1),
122            (1, 1),
123            1,
124        ),
125        (
126            "3×3 s(2,1) 64→128 @48×320",
127            [64, 48, 320],
128            [128, 3, 3],
129            (2, 1),
130            (1, 1),
131            1,
132        ),
133        (
134            "深度 3×3 128 @160²",
135            [128, 160, 160],
136            [128, 3, 3],
137            (1, 1),
138            (1, 1),
139            128,
140        ),
141        (
142            "深度 7×7 96 @80²",
143            [96, 80, 80],
144            [96, 7, 7],
145            (1, 1),
146            (3, 3),
147            96,
148        ),
149    ];
150    for (name, [c, h, w], [o, kh, kw], stride, pad, groups) in convs {
151        let x = rng.tensor(vec![c, h, w]);
152        let wt = rng.tensor(vec![o, c / groups, kh, kw]);
153        let bias: Vec<f32> = rng.tensor(vec![o]).data;
154        let s = ConvShape {
155            c,
156            h,
157            w,
158            o,
159            kh,
160            kw,
161            stride,
162            pad,
163            dilation: (1, 1),
164            groups,
165        };
166        let (oh, ow) = s.out_hw();
167        let gflop = 2.0 * (o * oh * ow * (c / groups) * kh * kw) as f64 / 1e9;
168        let cpu = wall_ms(3, || {
169            drop(ops::conv2d(
170                &x,
171                &wt,
172                Some(&bias),
173                stride,
174                pad,
175                (1, 1),
176                groups,
177            ))
178        });
179        let (dx, dw, db) = (
180            g.upload(&x.data).unwrap(),
181            g.upload(&wt.data).unwrap(),
182            g.upload(&bias).unwrap(),
183        );
184        let mut out = g.alloc(o * oh * ow).unwrap();
185        let kern = gpu_ms(&g, 20, || {
186            g.conv2d_device_into(&dx, &dw, Some(&db), &s, &mut out)
187                .unwrap()
188        });
189        let e2e = wall_ms(5, || {
190            drop(
191                g.conv2d(&x, &wt, Some(&bias), stride, pad, (1, 1), groups)
192                    .unwrap(),
193            )
194        });
195        row(name, gflop, cpu, kern, e2e);
196    }
197}
gpu_bench/gpu_bench.rs

gpu_bench/
gpu_bench.rs
