Struct Gpu

Source

pub struct Gpu { /* private fields */ }

Expand description

一块 N 卡上的上下文、载好的内核模块和函数表。所有方法取 &self，可以跨线程共享。

Implementations§

Source §

impl Gpu

Source

pub fn alloc(&self, len: usize) -> Result<DeviceBuffer<'_>, String>

分配 len 个 f32 的显存（内容未初始化）。len 为 0 时不调驱动。

Examples found in repository ?

examples/gpu_bench.rs (line 83)

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}

Source

pub fn upload(&self, data: &[f32]) -> Result<DeviceBuffer<'_>, String>

分配显存并拷入 data。

Examples found in repository ?

examples/gpu_bench.rs (line 82)

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}

Source §

impl Gpu

Source

pub fn timer(&self) -> Result<Timer<'_>, String>

建计时器并记下起点。

Examples found in repository ?

examples/gpu_bench.rs (line 45)

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}

Source §

impl Gpu

Source

pub fn gemm_device( &self, a: &DeviceBuffer<'_>, b: &DeviceBuffer<'_>, c: &mut DeviceBuffer<'_>, m: usize, n: usize, k: usize, ) -> Result<(), String>

c[m×n] = a[m×k] · b[k×n]，三者都在显存、行优先。

Examples found in repository ?

examples/gpu_bench.rs (line 85)

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}

Source

pub fn conv2d_device( &self, x: &DeviceBuffer<'_>, w: &DeviceBuffer<'_>, bias: Option<&DeviceBuffer<'_>>, s: &ConvShape, ) -> Result<DeviceBuffer<'_>, String>

二维卷积，输入、权重、偏置、输出都在显存。返回 [o, OH, OW] 的新显存。

Source

pub fn conv2d_device_into( &self, x: &DeviceBuffer<'_>, w: &DeviceBuffer<'_>, bias: Option<&DeviceBuffer<'_>>, s: &ConvShape, out: &mut DeviceBuffer<'_>, ) -> Result<(), String>

同 Gpu::conv2d_device，结果写进调用方给的显存 out（长度 o·OH·OW），省一次分配；连续推理时可以复用输出缓冲。im2col 的临时显存仍在每次调用里分配。

Examples found in repository ?

examples/gpu_bench.rs (line 186)

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}

Source

pub fn activate_device( &self, x: &mut DeviceBuffer<'_>, act: Activation, ) -> Result<(), String>

原地激活，语义同 Activation::apply。

Source

pub fn add_bias_device( &self, x: &mut DeviceBuffer<'_>, bias: &DeviceBuffer<'_>, hw: usize, ) -> Result<(), String>

原地加偏置：x 是 [C, hw]，x[c, i] += bias[c]。

Source

pub fn matmul(&self, a: &Tensor, b: &Tensor) -> Result<Vec<f32>, String>

矩阵乘（同 pagesift_infer::ops::matmul）：a 是 [M, K]，b 是 [K, N]，返回 [M, N] 行优先。

Examples found in repository ?

examples/gpu_bench.rs (line 87)

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}

Source

pub fn conv2d( &self, x: &Tensor, w: &Tensor, bias: Option<&[f32]>, stride: (usize, usize), pad: (usize, usize), dilation: (usize, usize), groups: usize, ) -> Result<Vec<f32>, String>

二维卷积（同 pagesift_infer::ops::conv2d）：x 是 [C, H, W]，w 是 [O, C/groups, kh, kw]；返回 [O, OH, OW] 行优先的数据。

Examples found in repository ?

examples/gpu_bench.rs (line 191)

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}

Source

pub fn activate(&self, x: &mut [f32], act: Activation) -> Result<(), String>

原地激活一段主机内存（同 Activation::apply_all）。

Source §

impl Gpu

Source

pub fn open() -> Result<Gpu, String>

打开 0 号 N 卡，同 Gpu::open_index(0)。

Examples found in repository ?

examples/gpu_bench.rs (line 63)

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}

Source

pub fn open_index(index: usize) -> Result<Gpu, String>

打开第 index 块 N 卡：载入驱动、取主上下文、按算力挑 PTX 并 JIT。没有驱动、没有这块卡、算力低于 6.1、JIT 失败或 PAGESIFT_GPU=off 时返回错误（不 panic）。

Source

pub fn name(&self) -> &str

设备名，例如 “NVIDIA GeForce GTX 1070”。

Examples found in repository ?

examples/gpu_bench.rs (line 70)

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}

Source

pub fn compute_capability(&self) -> (i32, i32)

算力 (主, 次)，例如 (6, 1)。

Examples found in repository ?

examples/gpu_bench.rs (line 70)

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}

Source

pub fn memory_free(&self) -> Result<u64, String>

现在的空闲显存（字节）。

Source

pub fn info(&self) -> Result<GpuInfo, String>

型号、算力、显存总量和空闲量、驱动版本、选中的 PTX 版本。

Source

pub fn synchronize(&self) -> Result<(), String>

等设备上已提交的工作全部完成。

Examples found in repository ?

examples/gpu_bench.rs (line 44)

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}