vyre_driver/
device_signature.rs

1//! Tier-B device signature loader.
2//!
3//! Device signatures are community-extensible TOML records. They describe
4//! architecture facts used by optimizer cost models, tiling, vector packing,
5//! and bank-conflict avoidance without baking a device table into Rust source.
6
7use serde::Deserialize;
8use std::fs;
9use std::path::Path;
10use std::sync::LazyLock;
11
12use crate::DeviceProfile;
13
14const MAX_DEVICE_SIGNATURE_TOML_BYTES: u64 = 256 * 1024;
15
16/// Process-wide memo of the compiled-in signature table.
17///
18/// Backend projections derive a [`DeviceProfile`] on every dispatch, and each
19/// derivation used to reparse this TOML. The input is a `&'static str` fixed
20/// at compile time, so one parse per process is all the information there is.
21///
22/// A parse FAILURE is memoized too, and that is deliberate: the input cannot
23/// change while the process runs, so a retry would reparse the same constant
24/// and fail the same way. Do not "fix" this into a retry.
25static BUILTIN_SIGNATURE_TABLE: LazyLock<Result<DeviceSignatureTable, String>> =
26    LazyLock::new(|| {
27        #[cfg(test)]
28        BUILTIN_SIGNATURE_TABLE_PARSES.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
29        let mut signatures = vec![DeviceSignature::from_toml_str(
30            DeviceSignature::BUILTIN_BLACKWELL_120,
31        )?];
32        signatures.sort_unstable_by(|left, right| left.id.cmp(&right.id));
33        Ok(DeviceSignatureTable { signatures })
34    });
35
36/// Counts how many times the memoized parse actually ran.
37///
38/// This exists so the memo can be pinned by an INVARIANT (the table is built
39/// once no matter how many callers ask for it) rather than by a timing
40/// assertion, which would be flaky on a contended box.
41#[cfg(test)]
42static BUILTIN_SIGNATURE_TABLE_PARSES: std::sync::atomic::AtomicUsize =
43    std::sync::atomic::AtomicUsize::new(0);
44
45/// One parsed device signature record.
46#[derive(Clone, Debug, Deserialize, PartialEq)]
47pub struct DeviceSignature {
48    /// Stable architecture id from the TOML file.
49    pub id: String,
50    /// Human-readable architecture family.
51    pub family: String,
52    /// Backend-reported architecture generation number.
53    #[serde(default)]
54    pub architecture_generation: Option<u32>,
55    /// Case-insensitive device-name fragments that identify this signature.
56    #[serde(default)]
57    pub device_name_contains: Vec<String>,
58    /// Maximum streaming compute units for this architecture family.
59    pub max_sm: u32,
60    /// Native subgroup/warp/wave size.
61    pub warp_size: u32,
62    /// Maximum registers per thread.
63    pub regs_per_thread_max: u32,
64    /// Shared memory per compute unit in KiB.
65    pub shared_mem_per_sm_kb: u32,
66    /// L1 cache size in KiB.
67    pub l1_kb: u32,
68    /// L2 cache size in KiB.
69    pub l2_kb: u32,
70    /// Peak memory bandwidth in GB/s.
71    pub mem_bw_gbps: u32,
72    /// Whether matrix-engine acceleration is available.
73    pub tensor_core_supported: bool,
74    /// Matrix-engine dtype names.
75    #[serde(default)]
76    pub tensor_core_dtypes: Vec<String>,
77    /// Default unroll depth preferred by cost models.
78    pub ideal_unroll_depth: u32,
79    /// Preferred vector pack width in bits.
80    pub ideal_vector_pack_bits: u32,
81    /// Preferred tile shape for workgroup-local kernels.
82    pub ideal_workgroup_tile: [u32; 3],
83    /// Shared-memory bank count.
84    pub bank_count: u32,
85    /// Shared-memory bank width in bytes.
86    pub bank_width_bytes: u32,
87}
88
89/// Collection of signatures loaded from a directory.
90#[derive(Clone, Debug, Default, PartialEq)]
91pub struct DeviceSignatureTable {
92    signatures: Vec<DeviceSignature>,
93}
94
95impl DeviceSignature {
96    /// Built-in Blackwell signature shipped with the crate.
97    pub const BUILTIN_BLACKWELL_120: &'static str = include_str!("../devices/blackwell_120.toml");
98
99    /// Parse a device signature from TOML.
100    ///
101    /// # Errors
102    ///
103    /// Returns an actionable string if the TOML cannot be parsed or violates
104    /// basic invariants required by optimizer consumers.
105    pub fn from_toml_str(source: &str) -> Result<Self, String> {
106        let signature: Self = toml::from_str(source)
107            .map_err(|error| format!("device signature TOML parse failed. Fix: {error}"))?;
108        signature.validate()?;
109        Ok(signature)
110    }
111
112    /// Validate schema invariants that would make cost models unsafe.
113    ///
114    /// # Errors
115    ///
116    /// Returns an actionable error for the first invalid field.
117    pub fn validate(&self) -> Result<(), String> {
118        if self.id.trim().is_empty() {
119            return Err("device signature id is empty. Fix: set a stable id.".to_string());
120        }
121        if self.family.trim().is_empty() {
122            return Err(
123                "device signature family is empty. Fix: set the architecture family.".to_string(),
124            );
125        }
126        if self.warp_size == 0 || !self.warp_size.is_power_of_two() {
127            return Err(format!(
128                "device signature `{}` has invalid warp_size {}. Fix: use a non-zero power of two.",
129                self.id, self.warp_size
130            ));
131        }
132        if self.ideal_vector_pack_bits == 0 || self.ideal_vector_pack_bits % 32 != 0 {
133            return Err(format!(
134                "device signature `{}` has invalid ideal_vector_pack_bits {}. Fix: use a positive multiple of 32.",
135                self.id, self.ideal_vector_pack_bits
136            ));
137        }
138        if self.ideal_workgroup_tile.contains(&0) {
139            return Err(format!(
140                "device signature `{}` has a zero ideal_workgroup_tile axis. Fix: every axis must be positive.",
141                self.id
142            ));
143        }
144        if self.bank_count == 0 || self.bank_width_bytes == 0 {
145            return Err(format!(
146                "device signature `{}` has invalid shared-memory bank metadata. Fix: bank_count and bank_width_bytes must be non-zero.",
147                self.id
148            ));
149        }
150        validate_kib_projection(self.shared_mem_per_sm_kb, "shared_mem_per_sm_kb", &self.id)?;
151        validate_kib_projection(self.l1_kb, "l1_kb", &self.id)?;
152        validate_kib_projection(self.l2_kb, "l2_kb", &self.id)?;
153        Ok(())
154    }
155
156    /// Apply architecture facts to a neutral device profile.
157    #[must_use]
158    pub fn apply_to_profile(&self, mut profile: DeviceProfile) -> DeviceProfile {
159        profile.subgroup_size = self.warp_size;
160        profile.supports_tensor_cores = self.tensor_core_supported;
161        profile.has_subgroup_shuffle = self.warp_size > 0;
162        profile.has_shared_memory |= self.shared_mem_per_sm_kb > 0;
163        if profile.max_shared_memory_bytes == 0 {
164            profile.max_shared_memory_bytes =
165                kib_to_bytes_checked(self.shared_mem_per_sm_kb, "shared_mem_per_sm_kb", &self.id);
166        }
167        profile.compute_units = self.max_sm;
168        profile.regs_per_thread_max = self.regs_per_thread_max;
169        profile.l1_cache_bytes = kib_to_bytes_checked(self.l1_kb, "l1_kb", &self.id);
170        profile.l2_cache_bytes = kib_to_bytes_checked(self.l2_kb, "l2_kb", &self.id);
171        profile.mem_bw_gbps = self.mem_bw_gbps;
172        profile.ideal_unroll_depth = self.ideal_unroll_depth;
173        profile.ideal_vector_pack_bits = self.ideal_vector_pack_bits;
174        profile.ideal_workgroup_tile = self.ideal_workgroup_tile;
175        profile.shared_memory_bank_count = self.bank_count;
176        profile.shared_memory_bank_width_bytes = self.bank_width_bytes;
177        profile
178    }
179
180    /// Return true when this signature should be used for a backend-reported
181    /// architecture generation number.
182    #[must_use]
183    pub fn matches_architecture_generation(&self, generation: u32) -> bool {
184        self.architecture_generation == Some(generation)
185            || self.id.rsplit('_').next().and_then(parse_u32) == Some(generation)
186    }
187
188    /// Return true when the device name carries one of this signature's
189    /// Tier-B aliases.
190    #[must_use]
191    pub fn matches_device_name(&self, device_name: &str) -> bool {
192        let device_name = device_name.to_ascii_lowercase();
193        self.device_name_contains
194            .iter()
195            .any(|needle| device_name.contains(&needle.to_ascii_lowercase()))
196    }
197}
198
199impl DeviceSignatureTable {
200    /// Load the signatures compiled into this crate.
201    ///
202    /// This reads NO files. It parses the [`DeviceSignature::BUILTIN_BLACKWELL_120`]
203    /// string that `include_str!` baked into the binary, sorts by id, and
204    /// wraps. [`Self::load_dir`] is the filesystem path; this is the
205    /// no-filesystem fallback used by backend projections before external
206    /// Tier-B directories are available.
207    ///
208    /// The parse is memoized process-wide, because it is a pure function of
209    /// compile-time constants and backend projections call it on every
210    /// dispatch. Prefer [`Self::builtins_ref`] on a hot path: this clones the
211    /// memoized table so callers that need an owned value keep working.
212    ///
213    /// # Errors
214    ///
215    /// Returns an actionable error when the compiled-in signature is invalid.
216    pub fn builtins() -> Result<Self, String> {
217        Self::builtins_ref().cloned()
218    }
219
220    /// Borrow the process-wide memoized builtin signature table.
221    ///
222    /// # Errors
223    ///
224    /// Returns an actionable error when the compiled-in signature is invalid.
225    pub fn builtins_ref() -> Result<&'static Self, String> {
226        BUILTIN_SIGNATURE_TABLE.as_ref().map_err(Clone::clone)
227    }
228
229    /// Load every `*.toml` signature file in `dir`.
230    ///
231    /// # Errors
232    ///
233    /// Returns an actionable error when the directory cannot be read, a file
234    /// cannot be read, or any signature is invalid.
235    pub fn load_dir(dir: impl AsRef<Path>) -> Result<Self, String> {
236        let dir = dir.as_ref();
237        let entries = fs::read_dir(dir).map_err(|error| {
238            format!(
239                "device signature directory `{}` cannot be read. Fix: create it or pass the correct path: {error}",
240                dir.display()
241            )
242        })?;
243        let mut signatures = Vec::new();
244        for entry in entries {
245            let entry = entry.map_err(|error| {
246                format!(
247                    "device signature directory `{}` contains an unreadable entry. Fix: {error}",
248                    dir.display()
249                )
250            })?;
251            let path = entry.path();
252            if path.extension().and_then(|ext| ext.to_str()) != Some("toml") {
253                continue;
254            }
255            let source = read_device_signature_toml(&path).map_err(|error| {
256                format!(
257                    "device signature file `{}` cannot be read. Fix: {error}",
258                    path.display()
259                )
260            })?;
261            let signature = DeviceSignature::from_toml_str(&source)
262                .map_err(|error| format!("{} in `{}`", error, path.display()))?;
263            signatures.push(signature);
264        }
265        signatures.sort_unstable_by(|left, right| left.id.cmp(&right.id));
266        dedupe_signature_ids(&signatures)?;
267        Ok(Self { signatures })
268    }
269
270    /// Borrow all loaded signatures in stable id order.
271    #[must_use]
272    pub fn signatures(&self) -> &[DeviceSignature] {
273        &self.signatures
274    }
275
276    /// Find a signature by id.
277    #[must_use]
278    pub fn get(&self, id: &str) -> Option<&DeviceSignature> {
279        self.signatures
280            .binary_search_by(|signature| signature.id.as_str().cmp(id))
281            .ok()
282            .and_then(|index| self.signatures.get(index))
283    }
284
285    /// Find the best signature for a backend-reported architecture generation.
286    #[must_use]
287    pub fn find_architecture_generation(&self, generation: u32) -> Option<&DeviceSignature> {
288        self.signatures
289            .iter()
290            .find(|signature| signature.matches_architecture_generation(generation))
291    }
292
293    /// Find the best signature for a backend-reported device name.
294    #[must_use]
295    pub fn find_device_name(&self, device_name: &str) -> Option<&DeviceSignature> {
296        self.signatures
297            .iter()
298            .find(|signature| signature.matches_device_name(device_name))
299    }
300
301    /// Apply a generation signature to `profile` when one is known.
302    #[must_use]
303    pub fn apply_generation_to_profile(
304        &self,
305        generation: u32,
306        profile: DeviceProfile,
307    ) -> DeviceProfile {
308        self.find_architecture_generation(generation)
309            .map_or(profile, |signature| signature.apply_to_profile(profile))
310    }
311
312    /// Apply a device-name signature to `profile` when one is known.
313    #[must_use]
314    pub fn apply_device_name_to_profile(
315        &self,
316        device_name: &str,
317        profile: DeviceProfile,
318    ) -> DeviceProfile {
319        self.find_device_name(device_name)
320            .map_or(profile, |signature| signature.apply_to_profile(profile))
321    }
322}
323
324fn read_device_signature_toml(path: &Path) -> std::io::Result<String> {
325    use std::io::Read as _;
326
327    let mut file = fs::File::open(path)?;
328    let metadata = file.metadata()?;
329    if metadata.len() > MAX_DEVICE_SIGNATURE_TOML_BYTES {
330        return Err(std::io::Error::new(
331            std::io::ErrorKind::InvalidData,
332            format!("device signature TOML exceeds {MAX_DEVICE_SIGNATURE_TOML_BYTES} byte limit"),
333        ));
334    }
335    let mut text = String::with_capacity(metadata.len() as usize);
336    file.by_ref()
337        .take(MAX_DEVICE_SIGNATURE_TOML_BYTES + 1)
338        .read_to_string(&mut text)?;
339    if text.len() as u64 > MAX_DEVICE_SIGNATURE_TOML_BYTES {
340        return Err(std::io::Error::new(
341            std::io::ErrorKind::InvalidData,
342            "device signature TOML exceeded bounded read limit",
343        ));
344    }
345    Ok(text)
346}
347
348fn dedupe_signature_ids(signatures: &[DeviceSignature]) -> Result<(), String> {
349    for pair in signatures.windows(2) {
350        if pair[0].id == pair[1].id {
351            return Err(format!(
352                "duplicate device signature id `{}`. Fix: keep exactly one TOML file per id.",
353                pair[0].id
354            ));
355        }
356    }
357    Ok(())
358}
359
360fn validate_kib_projection(value: u32, field: &str, id: &str) -> Result<(), String> {
361    value.checked_mul(1024).map(|_| ()).ok_or_else(|| {
362        format!(
363            "device signature `{id}` field {field}={value} KiB overflows u32 bytes. Fix: split the architecture record or lower the Tier-B value; silent saturation corrupts GPU resource planning."
364        )
365    })
366}
367
368fn kib_to_bytes_checked(value: u32, field: &str, id: &str) -> u32 {
369    let _ = (field, id);
370    value.saturating_mul(1024)
371}
372
373fn parse_u32(value: &str) -> Option<u32> {
374    let mut out = 0u32;
375    for byte in value.bytes() {
376        if !byte.is_ascii_digit() {
377            return None;
378        }
379        out = out.checked_mul(10)?.checked_add(u32::from(byte - b'0'))?;
380    }
381    Some(out)
382}
383
384#[cfg(test)]
385mod tests {
386    use super::{DeviceSignature, DeviceSignatureTable};
387    use crate::DeviceProfile;
388
389    const SAMPLE: &str = r#"
390id = "sample_arch"
391family = "sample"
392max_sm = 128
393warp_size = 32
394regs_per_thread_max = 255
395shared_mem_per_sm_kb = 128
396l1_kb = 128
397l2_kb = 98304
398mem_bw_gbps = 1700
399tensor_core_supported = true
400tensor_core_dtypes = ["f16", "bf16", "tf32"]
401ideal_unroll_depth = 8
402ideal_vector_pack_bits = 128
403ideal_workgroup_tile = [16, 16, 1]
404bank_count = 32
405bank_width_bytes = 4
406"#;
407
408    #[test]
409    fn parses_and_validates_signature() {
410        let signature = DeviceSignature::from_toml_str(SAMPLE).unwrap();
411
412        assert_eq!(signature.id, "sample_arch");
413        assert_eq!(signature.architecture_generation, None);
414        assert_eq!(signature.warp_size, 32);
415        assert!(signature.tensor_core_supported);
416    }
417
418    #[test]
419    fn rejects_invalid_warp_size() {
420        let err =
421            DeviceSignature::from_toml_str(&SAMPLE.replace("warp_size = 32", "warp_size = 48"))
422                .unwrap_err();
423
424        assert!(err.contains("warp_size"));
425    }
426
427    #[test]
428    fn applies_architecture_facts_to_profile() {
429        let signature = DeviceSignature::from_toml_str(SAMPLE).unwrap();
430        let profile = signature.apply_to_profile(DeviceProfile::conservative("test"));
431
432        assert_eq!(profile.subgroup_size, 32);
433        assert_eq!(profile.max_shared_memory_bytes, 128 * 1024);
434        assert_eq!(profile.compute_units, 128);
435        assert_eq!(profile.ideal_vector_pack_bits, 128);
436        assert_eq!(profile.shared_memory_bank_width_bytes, 4);
437        assert!(profile.supports_tensor_cores);
438    }
439
440    #[test]
441    fn preserves_live_shared_memory_per_workgroup_limit() {
442        let signature = DeviceSignature::from_toml_str(SAMPLE).unwrap();
443        let mut live = DeviceProfile::conservative("native");
444        live.max_shared_memory_bytes = 48 * 1024;
445        let profile = signature.apply_to_profile(live);
446
447        assert_eq!(profile.max_shared_memory_bytes, 48 * 1024);
448        assert!(profile.has_shared_memory);
449        assert_eq!(profile.shared_memory_bank_count, 32);
450    }
451
452    #[test]
453    fn loads_directory_in_id_order() {
454        let dir = tempfile::tempdir().unwrap();
455        std::fs::write(
456            dir.path().join("b.toml"),
457            SAMPLE.replace("sample_arch", "b"),
458        )
459        .unwrap();
460        std::fs::write(
461            dir.path().join("a.toml"),
462            SAMPLE.replace("sample_arch", "a"),
463        )
464        .unwrap();
465
466        let table = DeviceSignatureTable::load_dir(dir.path()).unwrap();
467
468        assert_eq!(table.signatures()[0].id, "a");
469        assert_eq!(table.signatures()[1].id, "b");
470        assert!(table.get("b").is_some());
471    }
472
473    #[test]
474    fn repository_device_signatures_load() {
475        let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../devices");
476        let table = DeviceSignatureTable::load_dir(dir).unwrap();
477
478        assert!(table.get("blackwell_120").is_some());
479    }
480
481    #[test]
482    fn builtins_match_generation_and_device_name() {
483        let table = DeviceSignatureTable::builtins().unwrap();
484        let signature = table.find_architecture_generation(120).unwrap();
485
486        assert_eq!(signature.id, "blackwell_120");
487        assert!(table.find_device_name("RTX 5090").is_some());
488    }
489
490    #[test]
491    fn builtin_signature_materially_projects_planner_fields() {
492        let table = DeviceSignatureTable::builtins().unwrap();
493        let signature = table.find_architecture_generation(120).unwrap();
494        let profile =
495            table.apply_generation_to_profile(120, DeviceProfile::conservative("backend"));
496
497        assert_eq!(profile.ideal_unroll_depth, signature.ideal_unroll_depth);
498        assert_eq!(
499            profile.ideal_vector_pack_bits,
500            signature.ideal_vector_pack_bits
501        );
502        assert_eq!(profile.ideal_workgroup_tile, signature.ideal_workgroup_tile);
503        assert_eq!(profile.shared_memory_bank_count, signature.bank_count);
504    }
505
506    /// The builtin table is parsed ONCE per process, no matter how many
507    /// callers ask for it.
508    ///
509    /// Backend projections call `builtins()` on every dispatch (twice, in the
510    /// CUDA case: once deriving validation capabilities and once deriving
511    /// adapter caps), so a reparse here is a per-dispatch TOML parse on the
512    /// hot path. This pins the invariant rather than a duration, because a
513    /// timing assertion would be flaky on a contended box and would not
514    /// actually say what we mean.
515    ///
516    /// The count is asserted as exactly one rather than "unchanged" so the
517    /// test is independent of whichever test in this binary ran first.
518    #[test]
519    fn builtin_signature_table_is_parsed_once_per_process() {
520        use std::sync::atomic::Ordering;
521
522        let first = DeviceSignatureTable::builtins().expect("Fix: builtin signatures must load");
523        for _ in 0..1_000 {
524            let repeated =
525                DeviceSignatureTable::builtins().expect("Fix: builtin signatures must load");
526            assert_eq!(
527                repeated, first,
528                "Fix: the memoized builtin table must return the same signatures on every call."
529            );
530        }
531        for _ in 0..1_000 {
532            let borrowed = DeviceSignatureTable::builtins_ref()
533                .expect("Fix: builtin signatures must load by reference");
534            assert_eq!(
535                *borrowed, first,
536                "Fix: builtins_ref must borrow the same table that builtins clones."
537            );
538        }
539
540        assert_eq!(
541            super::BUILTIN_SIGNATURE_TABLE_PARSES.load(Ordering::Relaxed),
542            1,
543            "Fix: the compiled-in device signature TOML must be parsed exactly once per process. \
544             2001 calls produced a different parse count, so the memo was bypassed or reset."
545        );
546    }
547}
vyre_driver/device_signature.rs

vyre_driver/
device_signature.rs
