ferrotensor_backend/
adapters.rs

1//! What GPUs wgpu can see on this machine, and a quick correctness check of
2//! the kernels on the one Ferrotensor uses.
3//!
4//! `ferrotensor info` prints [`list_adapters`] and, with `--check`, runs
5//! [`self_check`]. Both are meant for bug reports and for trying Ferrotensor
6//! on hardware it has not been tested on yet.
7
8use crate::{ComputeEngine, UnifiedTensor, WgpuContext};
9
10/// One adapter wgpu found, with the limits that matter for the kernels.
11#[derive(Clone, Debug, PartialEq)]
12pub struct AdapterReport {
13    /// Adapter name reported by the driver, e.g. `"Apple M2"`.
14    pub name: String,
15    /// GPU vendor from the PCI vendor id, e.g. `"AMD"`, or the id in hex.
16    pub vendor: String,
17    /// `"DiscreteGpu"`, `"IntegratedGpu"`, `"Cpu"` (a software rasteriser), ...
18    pub device_type: String,
19    /// Graphics API: `"Metal"`, `"Vulkan"`, `"Dx12"` or `"Gl"`.
20    pub backend: String,
21    /// Driver name and version, when the backend reports them.
22    pub driver: String,
23    /// Largest single buffer, in bytes.
24    pub max_buffer_size: u64,
25    /// Largest storage buffer a kernel can bind, in bytes.
26    pub max_storage_buffer_binding_size: u32,
27    /// Most threads in one workgroup.
28    pub max_compute_invocations_per_workgroup: u32,
29    /// Workgroup shared memory, in bytes.
30    pub max_compute_workgroup_storage_size: u32,
31    /// Whether the adapter supports `f16` in shaders.
32    pub shader_f16: bool,
33    /// Whether [`WgpuContext::new`] picks this adapter.
34    pub selected: bool,
35}
36
37/// Vendor name for a PCI vendor id.
38pub fn vendor_name(id: u32) -> String {
39    match id {
40        0x1002 | 0x1022 => "AMD".into(),
41        0x10de => "NVIDIA".into(),
42        0x8086 => "Intel".into(),
43        0x106b => "Apple".into(),
44        0x5143 => "Qualcomm".into(),
45        0x13b5 => "Arm".into(),
46        0x1010 => "Imagination".into(),
47        0x14e4 => "Broadcom".into(),
48        0x10005 => "Mesa".into(),
49        0 => "unknown".into(),
50        other => format!("{other:#06x}"),
51    }
52}
53
54fn report(adapter: &wgpu::Adapter, selected: Option<&wgpu::AdapterInfo>) -> AdapterReport {
55    let info = adapter.get_info();
56    let limits = adapter.limits();
57    let driver = [info.driver.trim(), info.driver_info.trim()]
58        .iter()
59        .filter(|s| !s.is_empty())
60        .copied()
61        .collect::<Vec<_>>()
62        .join(" ");
63    AdapterReport {
64        name: info.name.clone(),
65        // Metal reports no PCI vendor id; every Metal GPU is Apple's.
66        vendor: if info.vendor == 0 && info.backend == wgpu::Backend::Metal {
67            "Apple".into()
68        } else {
69            vendor_name(info.vendor)
70        },
71        device_type: format!("{:?}", info.device_type),
72        backend: format!("{:?}", info.backend),
73        driver,
74        max_buffer_size: limits.max_buffer_size,
75        max_storage_buffer_binding_size: limits.max_storage_buffer_binding_size,
76        max_compute_invocations_per_workgroup: limits.max_compute_invocations_per_workgroup,
77        max_compute_workgroup_storage_size: limits.max_compute_workgroup_storage_size,
78        shader_f16: adapter.features().contains(wgpu::Features::SHADER_F16),
79        selected: selected == Some(&info),
80    }
81}
82
83/// Every adapter wgpu can open on this machine, across all backends, with the
84/// one [`WgpuContext::new`] would pick marked `selected`.
85pub fn list_adapters() -> Vec<AdapterReport> {
86    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor {
87        backends: wgpu::Backends::all(),
88        ..Default::default()
89    });
90    let chosen = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
91        power_preference: wgpu::PowerPreference::HighPerformance,
92        compatible_surface: None,
93        force_fallback_adapter: false,
94    }))
95    .map(|a| a.get_info());
96    instance
97        .enumerate_adapters(wgpu::Backends::all())
98        .iter()
99        .map(|a| report(a, chosen.as_ref()))
100        .collect()
101}
102
103/// Result of one kernel compared against a CPU reference.
104#[derive(Clone, Debug, PartialEq)]
105pub struct CheckResult {
106    /// What was run, e.g. `"matmul 37x53 @ 53x29"`.
107    pub name: String,
108    /// Largest absolute difference from the CPU result.
109    pub max_abs_err: f32,
110    /// Largest difference that still counts as a pass.
111    pub tolerance: f32,
112}
113
114impl CheckResult {
115    /// Whether the GPU result is within tolerance.
116    pub fn passed(&self) -> bool {
117        self.max_abs_err.is_finite() && self.max_abs_err <= self.tolerance
118    }
119}
120
121/// Deterministic values in roughly [-1, 1], so the check needs no RNG.
122fn values(n: usize, phase: f32) -> Vec<f32> {
123    (0..n)
124        .map(|i| ((i as f32) * 0.37 + phase).sin() * ((i as f32) * 0.11).cos())
125        .collect()
126}
127
128fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
129    if a.len() != b.len() {
130        return f32::INFINITY;
131    }
132    a.iter()
133        .zip(b)
134        .map(|(x, y)| (x - y).abs())
135        .fold(0.0, f32::max)
136}
137
138/// Run matmul (odd, non-square sizes), softmax and RMSNorm on `ctx`'s GPU
139/// and compare each with a CPU reference computed in f64.
140pub async fn self_check(ctx: &WgpuContext) -> Vec<CheckResult> {
141    let engine = ComputeEngine::new(ctx);
142    let mut results = Vec::new();
143
144    // matmul with sizes that are not multiples of any tile size.
145    let (m, k, n) = (37, 53, 29);
146    let (a, b) = (values(m * k, 0.1), values(k * n, 0.7));
147    let mut want = vec![0f32; m * n];
148    for i in 0..m {
149        for j in 0..n {
150            let s: f64 = (0..k)
151                .map(|p| a[i * k + p] as f64 * b[p * n + j] as f64)
152                .sum();
153            want[i * n + j] = s as f32;
154        }
155    }
156    let (ga, gb) = (
157        UnifiedTensor::new(ctx, &a, &[m, k]),
158        UnifiedTensor::new(ctx, &b, &[k, n]),
159    );
160    let gc = UnifiedTensor::empty(ctx, &[m, n]);
161    engine.matmul(ctx, &ga, &gb, &gc);
162    results.push(CheckResult {
163        name: format!("matmul {m}x{k} @ {k}x{n}"),
164        max_abs_err: max_abs_diff(&gc.to_vec(ctx).await, &want),
165        tolerance: 1e-4,
166    });
167
168    // softmax over rows longer than one workgroup.
169    let (rows, cols) = (7, 1031);
170    let x = values(rows * cols, 1.3);
171    let mut want = vec![0f32; rows * cols];
172    for r in 0..rows {
173        let row = &x[r * cols..(r + 1) * cols];
174        let max = row.iter().copied().fold(f32::NEG_INFINITY, f32::max) as f64;
175        let sum: f64 = row.iter().map(|&v| (v as f64 - max).exp()).sum();
176        for c in 0..cols {
177            want[r * cols + c] = ((row[c] as f64 - max).exp() / sum) as f32;
178        }
179    }
180    let gx = UnifiedTensor::new(ctx, &x, &[rows, cols]);
181    let gy = UnifiedTensor::empty(ctx, &[rows, cols]);
182    engine.softmax(ctx, &gx, &gy);
183    results.push(CheckResult {
184        name: format!("softmax {rows}x{cols}"),
185        max_abs_err: max_abs_diff(&gy.to_vec(ctx).await, &want),
186        tolerance: 1e-6,
187    });
188
189    // RMSNorm with a non-trivial weight.
190    let (rows, cols) = (5, 777);
191    let (x, w) = (values(rows * cols, 2.1), values(cols, 0.4));
192    let mut want = vec![0f32; rows * cols];
193    for r in 0..rows {
194        let row = &x[r * cols..(r + 1) * cols];
195        let ms: f64 = row.iter().map(|&v| (v as f64).powi(2)).sum::<f64>() / cols as f64;
196        let inv = 1.0 / (ms + 1e-5).sqrt();
197        for c in 0..cols {
198            want[r * cols + c] = (row[c] as f64 * inv * w[c] as f64) as f32;
199        }
200    }
201    let (gx, gw) = (
202        UnifiedTensor::new(ctx, &x, &[rows, cols]),
203        UnifiedTensor::new(ctx, &w, &[cols]),
204    );
205    let gy = UnifiedTensor::empty(ctx, &[rows, cols]);
206    engine.rms_norm(ctx, &gx, &gw, &gy);
207    results.push(CheckResult {
208        name: format!("rms_norm {rows}x{cols}"),
209        max_abs_err: max_abs_diff(&gy.to_vec(ctx).await, &want),
210        tolerance: 1e-4,
211    });
212
213    results
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    #[test]
221    fn vendor_ids_map_to_names() {
222        assert_eq!(vendor_name(0x1002), "AMD");
223        assert_eq!(vendor_name(0x10de), "NVIDIA");
224        assert_eq!(vendor_name(0x8086), "Intel");
225        assert_eq!(vendor_name(0x106b), "Apple");
226        assert_eq!(vendor_name(0), "unknown");
227        assert_eq!(vendor_name(0xabcd), "0xabcd");
228    }
229
230    #[test]
231    fn check_result_pass_rule() {
232        let r = |e| CheckResult {
233            name: String::new(),
234            max_abs_err: e,
235            tolerance: 1e-4,
236        };
237        assert!(r(0.0).passed());
238        assert!(r(1e-4).passed());
239        assert!(!r(2e-4).passed());
240        assert!(!r(f32::NAN).passed());
241        assert!(!r(f32::INFINITY).passed());
242    }
243
244    #[test]
245    fn diff_of_mismatched_lengths_fails() {
246        assert_eq!(max_abs_diff(&[1.0], &[1.0, 2.0]), f32::INFINITY);
247        assert_eq!(max_abs_diff(&[1.0, 2.0], &[1.5, 2.0]), 0.5);
248    }
249}
ferrotensor_backend/adapters.rs

ferrotensor_backend/
adapters.rs
