vyre_driver/
validation.rs

1//! Shared validation caches and launch-geometry checks for concrete drivers.
2
3use std::collections::HashSet;
4use std::hash::BuildHasherDefault;
5
6use rustc_hash::FxHasher;
7use vyre_foundation::ir::{OpId, Program};
8use vyre_foundation::validate::{BackendValidationCapabilities, ValidationOptions};
9
10use crate::{BackendError, DispatchConfig, VyreBackend};
11
12/// Default successful-validation hash entries retained per backend instance.
13pub const DEFAULT_VALIDATION_HASH_ENTRIES: usize = 8192;
14/// Default VSA fingerprints retained per backend instance.
15pub const DEFAULT_VALIDATION_VSA_ENTRIES: usize = 2048;
16/// Default VSA shard count.
17pub const DEFAULT_VALIDATION_VSA_SHARDS: usize = 64;
18
19type ValidationSet = dashmap::DashSet<blake3::Hash, BuildHasherDefault<FxHasher>>;
20
21/// Successful-program validation cache shared by concrete drivers.
22pub struct ValidationCache {
23    hashes: ValidationSet,
24    vsa_hashes: ValidationSet,
25    max_hash_entries: usize,
26    max_vsa_entries: usize,
27    vsa_shards: usize,
28}
29
30impl std::fmt::Debug for ValidationCache {
31    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
32        formatter
33            .debug_struct("ValidationCache")
34            .field("hashes", &self.hashes.len())
35            .field("vsa_hashes", &self.vsa_hashes.len())
36            .field("vsa_shards", &self.vsa_shards)
37            .field("max_hash_entries", &self.max_hash_entries)
38            .field("max_vsa_entries", &self.max_vsa_entries)
39            .finish()
40    }
41}
42
43impl Default for ValidationCache {
44    fn default() -> Self {
45        Self::new(
46            DEFAULT_VALIDATION_HASH_ENTRIES,
47            DEFAULT_VALIDATION_VSA_ENTRIES,
48            DEFAULT_VALIDATION_VSA_SHARDS,
49        )
50    }
51}
52
53impl ValidationCache {
54    /// Create a validation cache with bounded hash and VSA storage.
55    #[must_use]
56    pub fn new(max_hash_entries: usize, max_vsa_entries: usize, vsa_shards: usize) -> Self {
57        let shard_count = vsa_shards.max(1);
58        Self {
59            hashes: dashmap::DashSet::with_hasher(BuildHasherDefault::<FxHasher>::default()),
60            vsa_hashes: dashmap::DashSet::with_capacity_and_hasher(
61                max_vsa_entries.max(1),
62                BuildHasherDefault::<FxHasher>::default(),
63            ),
64            max_hash_entries: max_hash_entries.max(1),
65            max_vsa_entries: max_vsa_entries.max(1),
66            vsa_shards: shard_count,
67        }
68    }
69
70    /// Compute the validation hash for a program.
71    #[must_use]
72    pub fn program_hash(program: &Program) -> blake3::Hash {
73        blake3::Hash::from(program.fingerprint())
74    }
75
76    /// Return whether a validation hash is cached.
77    #[must_use]
78    pub fn contains_hash(&self, hash: &blake3::Hash) -> bool {
79        self.hashes.contains(hash)
80    }
81
82    /// Remember a successful validation hash.
83    pub fn remember_hash(&self, hash: blake3::Hash) {
84        if self.hashes.len() >= self.max_hash_entries {
85            self.hashes.clear();
86        }
87        self.hashes.insert(hash);
88    }
89
90    /// Remember a successful validation hash and its VSA fingerprint.
91    ///
92    /// # Errors
93    ///
94    /// Returns if a VSA shard lock is poisoned.
95    pub fn remember_success(&self, hash: blake3::Hash, vsa: &[u32]) -> Result<(), BackendError> {
96        self.remember_hash(hash);
97        if self.vsa_hashes.len() >= self.max_vsa_entries {
98            self.vsa_hashes.clear();
99        }
100        self.vsa_hashes.insert(vsa_words_hash(vsa));
101        Ok(())
102    }
103
104    /// Clear cached validation state.
105    ///
106    /// # Errors
107    ///
108    /// Returns if a VSA shard lock is poisoned.
109    pub fn clear(&self) -> Result<(), BackendError> {
110        self.hashes.clear();
111        self.vsa_hashes.clear();
112        Ok(())
113    }
114
115    /// Validate `program` once, memoizing the complete backend contract.
116    ///
117    /// This is the shared driver validation path: foundation invariants,
118    /// backend supported-op coverage, program capability requirements, and
119    /// VSA cache insertion all happen in one place. Concrete drivers supply
120    /// only their actual capability values.
121    ///
122    /// # Errors
123    ///
124    /// Returns when validation fails or a VSA shard lock is poisoned.
125    pub fn get_or_validate(
126        &self,
127        program: &Program,
128        validation_options: ValidationOptions<'_>,
129        supported_ops: &HashSet<OpId>,
130        caps: ProgramValidationCaps,
131    ) -> Result<(), BackendError> {
132        let hash = Self::program_hash(program);
133        if self.contains_hash(&hash) || program.is_validated_on(caps.backend_id) {
134            self.remember_hash(hash);
135            return Ok(());
136        }
137
138        validate_program_contract(program, validation_options, supported_ops, caps)?;
139
140        let vsa = crate::launch::program_vsa_fingerprint_words(program);
141        self.remember_success(hash, &vsa)?;
142        program.mark_validated_on(caps.backend_id);
143        Ok(())
144    }
145
146    /// Validate `program` against a concrete backend and cache successful
147    /// results.
148    ///
149    /// This is the canonical driver-owned validation-cache entry point for
150    /// backends that implement both the runtime backend contract and the
151    /// foundation capability-validation contract.
152    ///
153    /// # Errors
154    ///
155    /// Returns when validation fails or cache mutation fails.
156    pub fn get_or_validate_backend<B>(
157        &self,
158        program: &Program,
159        backend: &B,
160    ) -> Result<(), BackendError>
161    where
162        B: VyreBackend + BackendValidationCapabilities,
163    {
164        let validation_options = ValidationOptions::default().with_backend(backend);
165        self.get_or_validate(
166            program,
167            validation_options,
168            backend.supported_ops(),
169            ProgramValidationCaps::from_backend(backend),
170        )
171    }
172}
173
174/// Concrete backend capability values needed for shared program validation.
175#[derive(Debug, Clone, Copy, PartialEq, Eq)]
176pub struct ProgramValidationCaps {
177    /// Stable backend identifier used in diagnostics and validation stamps.
178    pub backend_id: &'static str,
179    /// Native subgroup operations are available and lowered.
180    pub supports_subgroup_ops: bool,
181    /// IEEE binary16 buffers/operations are lowered.
182    pub supports_f16: bool,
183    /// Bfloat16 buffers/operations are lowered.
184    pub supports_bf16: bool,
185    /// Indirect dispatch is lowered.
186    pub supports_indirect_dispatch: bool,
187    /// Distributed collective communication nodes are lowered.
188    pub supports_distributed_collectives: bool,
189    /// `Node::Trap` is lowered with backend-visible trap semantics.
190    pub supports_trap_propagation: bool,
191    /// Maximum supported workgroup dimensions.
192    pub max_workgroup_size: [u32; 3],
193}
194
195impl ProgramValidationCaps {
196    /// Snapshot capability values from a `VyreBackend` trait object.
197    #[must_use]
198    pub fn from_backend(backend: &dyn VyreBackend) -> Self {
199        Self {
200            backend_id: backend.id(),
201            supports_subgroup_ops: backend.supports_subgroup_ops(),
202            supports_f16: backend.supports_f16(),
203            supports_bf16: backend.supports_bf16(),
204            supports_indirect_dispatch: backend.supports_indirect_dispatch(),
205            supports_distributed_collectives: backend.supports_distributed_collectives(),
206            supports_trap_propagation: true,
207            max_workgroup_size: backend.max_workgroup_size(),
208        }
209    }
210}
211
212/// Validate a program against backend-neutral and backend-reported contracts.
213///
214/// # Errors
215///
216/// Returns when foundation validation, supported-op validation, or required
217/// capability checks fail.
218pub fn validate_program_contract(
219    program: &Program,
220    validation_options: ValidationOptions<'_>,
221    supported_ops: &HashSet<OpId>,
222    caps: ProgramValidationCaps,
223) -> Result<(), BackendError> {
224    let lowered_program = if caps.supports_distributed_collectives {
225        None
226    } else {
227        vyre_foundation::transform::collectives::lower_single_rank_collectives(program).map_err(
228            |error| BackendError::InvalidProgram {
229                fix: error.to_string(),
230            },
231        )?
232    };
233    let program = lowered_program.as_ref().unwrap_or(program);
234    let report = vyre_foundation::validate::validate_with_options(program, validation_options);
235    if let Some(first) = report.errors.into_iter().next() {
236        return Err(BackendError::InvalidProgram {
237            fix: first.message.into_owned(),
238        });
239    }
240
241    validate_supported_ops(program, caps.backend_id, supported_ops).map_err(|error| {
242        BackendError::InvalidProgram {
243            fix: error.to_string(),
244        }
245    })?;
246
247    let required = vyre_foundation::program_caps::scan(program);
248    vyre_foundation::program_caps::check_backend_capabilities(
249        caps.backend_id,
250        caps.supports_subgroup_ops,
251        caps.supports_f16,
252        caps.supports_bf16,
253        caps.supports_indirect_dispatch,
254        caps.supports_trap_propagation,
255        caps.supports_distributed_collectives,
256        caps.max_workgroup_size,
257        &required,
258    )
259    .map_err(|error| BackendError::InvalidProgram {
260        fix: error.to_string(),
261    })
262}
263
264fn validate_supported_ops(
265    program: &Program,
266    backend_id: &'static str,
267    supported_ops: &HashSet<OpId>,
268) -> Result<(), vyre_foundation::ir::ValidationError> {
269    struct SupportedOpsBackend<'a> {
270        id: &'static str,
271        ops: &'a HashSet<OpId>,
272    }
273
274    impl crate::backend::Backend for SupportedOpsBackend<'_> {
275        fn id(&self) -> &'static str {
276            self.id
277        }
278
279        fn version(&self) -> &'static str {
280            env!("CARGO_PKG_VERSION")
281        }
282
283        fn supported_ops(&self) -> &HashSet<OpId> {
284            self.ops
285        }
286    }
287
288    crate::backend::validation::validate_program(
289        program,
290        &SupportedOpsBackend {
291            id: backend_id,
292            ops: supported_ops,
293        },
294    )
295}
296
297/// Launch-geometry limits reported by a concrete driver.
298#[derive(Debug, Clone, Copy, PartialEq, Eq)]
299pub struct LaunchGeometryLimits {
300    /// Backend name used in diagnostics.
301    pub backend: &'static str,
302    /// Maximum invocations in one workgroup or block.
303    pub max_threads_per_block: u32,
304    /// Maximum workgroup or block dimensions (x, y, z).
305    pub max_block_dim: [u32; 3],
306    /// Maximum workgroup count per grid dimension.
307    pub max_grid_dim: [u32; 3],
308    /// Maximum threads the device keeps resident on one compute unit (a CUDA
309    /// streaming multiprocessor, a Metal threadgroup-hosting core, and so on),
310    /// or `0` when the backend does not probe this number.
311    ///
312    /// This is the budget that decides how many whole workgroups fit on one
313    /// unit, and the division is integral: a workgroup width that does not
314    /// divide it strands the remainder on every unit for the launch's whole
315    /// lifetime. Backends that leave this `0` opt out of every residency-aware
316    /// decision rather than receiving one derived from a guessed budget.
317    pub max_threads_per_sm: u32,
318}
319
320impl LaunchGeometryLimits {
321    /// Whole workgroups of `workgroup_threads` that stay resident on one
322    /// compute unit, or `None` when this backend reports no per-unit budget.
323    #[must_use]
324    pub fn blocks_per_compute_unit(&self, workgroup_threads: u32) -> Option<u32> {
325        (self.max_threads_per_sm != 0 && workgroup_threads != 0)
326            .then(|| blocks_per_compute_unit(self.max_threads_per_sm, workgroup_threads))
327    }
328
329    /// Threads that stay resident on one compute unit at `workgroup_threads`
330    /// wide, or `None` when this backend reports no per-unit budget.
331    #[must_use]
332    pub fn resident_threads_per_compute_unit(&self, workgroup_threads: u32) -> Option<u32> {
333        (self.max_threads_per_sm != 0 && workgroup_threads != 0)
334            .then(|| resident_threads_per_compute_unit(self.max_threads_per_sm, workgroup_threads))
335    }
336}
337
338/// Whole workgroups of `workgroup_threads` that fit one compute unit's thread
339/// budget.
340///
341/// This is the single definition of the residency division in the workspace.
342/// CUDA's cooperative launch preflight and cold-start launch-width selection
343/// both route through it, because two independent copies of this arithmetic
344/// had already drifted apart once. The division is integral by hardware: a
345/// unit hosts whole workgroups only.
346///
347/// Threads are the only ceiling modelled here. Hardware also caps blocks per
348/// unit independently (CUDA reports it as
349/// `CU_DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR`), so at narrow widths
350/// the real block count is lower than this returns and the shortfall can be
351/// large: where the device caps blocks at 24, a 32-wide group against a
352/// 1536-thread budget measures 24 blocks and 768 resident threads, half what
353/// this function's 48 blocks and 1536 threads predict. A caller that ranks
354/// widths from widest downward never reaches that regime. A caller that
355/// answers "does this declared width fit", such as a cooperative launch
356/// preflight, does, and must clamp by the device-reported block cap before
357/// admitting a grid.
358#[must_use]
359pub fn blocks_per_compute_unit(max_threads_per_unit: u32, workgroup_threads: u32) -> u32 {
360    if workgroup_threads == 0 {
361        return 0;
362    }
363    max_threads_per_unit / workgroup_threads
364}
365
366/// Threads resident on one compute unit at `workgroup_threads` wide, under the
367/// per-unit thread budget alone.
368///
369/// Equal to `blocks_per_compute_unit(..) * workgroup_threads`, so it is at most
370/// `max_threads_per_unit` and falls short of it by exactly the slots the
371/// integral division strands. A width of 1024 against a 1536-thread budget
372/// resolves to one block and 1024 resident threads, leaving 512 slots per unit
373/// unusable for the launch's duration. The block-count caveat on
374/// [`blocks_per_compute_unit`] applies here too.
375#[must_use]
376pub fn resident_threads_per_compute_unit(max_threads_per_unit: u32, workgroup_threads: u32) -> u32 {
377    blocks_per_compute_unit(max_threads_per_unit, workgroup_threads)
378        .saturating_mul(workgroup_threads)
379}
380
381/// Validate workgroup and grid dimensions against backend launch limits.
382///
383/// # Errors
384///
385/// Returns when dimensions are zero, overflow the invocation product, exceed
386/// workgroup limits, or exceed per-axis grid limits.
387pub fn validate_launch_geometry(
388    workgroup: [u32; 3],
389    grid: [u32; 3],
390    limits: LaunchGeometryLimits,
391) -> Result<(), BackendError> {
392    if workgroup.contains(&0) || grid.contains(&0) {
393        return Err(BackendError::InvalidProgram {
394            fix: format!(
395                "Fix: {} workgroup and grid dimensions must all be non-zero.",
396                limits.backend
397            ),
398        });
399    }
400    let threads = workgroup[0]
401        .checked_mul(workgroup[1])
402        .and_then(|xy| xy.checked_mul(workgroup[2]))
403        .ok_or_else(|| BackendError::InvalidProgram {
404            fix: format!(
405                "Fix: {} workgroup dimensions overflowed u32; reduce workgroup_override.",
406                limits.backend
407            ),
408        })?;
409    if threads > limits.max_threads_per_block {
410        return Err(BackendError::InvalidProgram {
411            fix: format!(
412                "Fix: {} workgroup has {threads} threads but device max is {}.",
413                limits.backend, limits.max_threads_per_block
414            ),
415        });
416    }
417    for (axis, &dim) in workgroup.iter().enumerate() {
418        if dim > limits.max_block_dim[axis] {
419            return Err(BackendError::InvalidProgram {
420                fix: format!(
421                    "Fix: {} workgroup axis {axis} requested {} threads but device max is {}.",
422                    limits.backend, dim, limits.max_block_dim[axis]
423                ),
424            });
425        }
426    }
427    for (axis, &dim) in grid.iter().enumerate() {
428        if dim > limits.max_grid_dim[axis] {
429            return Err(BackendError::InvalidProgram {
430                fix: format!(
431                    "Fix: {} grid axis {axis} requested {} workgroups but device max is {}.",
432                    limits.backend, dim, limits.max_grid_dim[axis]
433                ),
434            });
435        }
436    }
437    Ok(())
438}
439
440/// Validate a program's effective workgroup shape against a backend's reported limits.
441///
442/// This is the shared pre-dispatch gate for callers that have a `VyreBackend`
443/// trait object but have not entered a concrete driver yet.
444///
445/// # Errors
446///
447/// Returns when any workgroup axis is zero, exceeds the backend's per-axis
448/// limit, or when total invocations exceed the backend's workgroup limit.
449pub fn validate_program_for_backend(
450    backend: &dyn VyreBackend,
451    program: &Program,
452    config: &DispatchConfig,
453) -> Result<(), BackendError> {
454    let workgroup = config
455        .workgroup_override
456        .unwrap_or(program.workgroup_size());
457    let max_axes = backend.max_workgroup_size();
458    if workgroup.contains(&0) {
459        return Err(BackendError::InvalidProgram {
460            fix: format!(
461                "Fix: backend `{}` cannot dispatch zero-sized workgroup dimensions; set positive workgroup sizes.",
462                backend.id()
463            ),
464        });
465    }
466    for (axis, &dim) in workgroup.iter().enumerate() {
467        if dim > max_axes[axis] {
468            return Err(BackendError::InvalidProgram {
469                fix: format!(
470                    "Fix: backend `{}` workgroup axis {axis} requested {} but max is {}.",
471                    backend.id(),
472                    dim,
473                    max_axes[axis]
474                ),
475            });
476        }
477    }
478    let invocations = workgroup[0]
479        .checked_mul(workgroup[1])
480        .and_then(|xy| xy.checked_mul(workgroup[2]))
481        .ok_or_else(|| BackendError::InvalidProgram {
482            fix: format!(
483                "Fix: backend `{}` workgroup dimensions overflowed u32; reduce workgroup size.",
484                backend.id()
485            ),
486        })?;
487    let max_invocations = backend.max_compute_invocations_per_workgroup();
488    if invocations > max_invocations {
489        return Err(BackendError::InvalidProgram {
490            fix: format!(
491                "Fix: backend `{}` workgroup has {invocations} invocations but max is {max_invocations}.",
492                backend.id()
493            ),
494        });
495    }
496    if let Some(grid) = config.grid_override {
497        let max_workgroups = backend.max_compute_workgroups_per_dimension();
498        if grid.contains(&0) {
499            return Err(BackendError::InvalidProgram {
500                fix: format!(
501                    "Fix: backend `{}` cannot dispatch zero-sized grid dimensions; set positive grid_override values.",
502                    backend.id()
503                ),
504            });
505        }
506        for (axis, &dim) in grid.iter().enumerate() {
507            if dim > max_workgroups {
508                return Err(BackendError::InvalidProgram {
509                    fix: format!(
510                        "Fix: backend `{}` grid_override axis {axis} requested {} workgroups but max is {}.",
511                        backend.id(),
512                        dim,
513                        max_workgroups
514                    ),
515                });
516            }
517        }
518    }
519    Ok(())
520}
521
522fn vsa_words_hash(words: &[u32]) -> blake3::Hash {
523    let mut hasher = blake3::Hasher::new();
524    hasher.update(&(words.len() as u64).to_le_bytes());
525    for word in words {
526        hasher.update(&word.to_le_bytes());
527    }
528    hasher.finalize()
529}
530
531#[cfg(test)]
532mod tests {
533    use super::*;
534
535    #[test]
536    fn validation_cache_records_vsa_without_lock_shards() {
537        let cache = ValidationCache::new(8, 8, 4);
538        let hash = blake3::hash(b"program");
539        cache
540            .remember_success(hash, &[1, 2, 3, 4])
541            .expect("Fix: lock-free VSA cache insertion must not fail");
542
543        assert!(cache.contains_hash(&hash));
544        assert_eq!(cache.vsa_hashes.len(), 1);
545        assert!(format!("{cache:?}").contains("vsa_hashes"));
546    }
547
548    #[test]
549    fn validation_cache_bounds_vsa_hashes_by_clear() {
550        let cache = ValidationCache::new(8, 2, 4);
551        for i in 0..3u32 {
552            cache
553                .remember_success(blake3::hash(&i.to_le_bytes()), &[i])
554                .expect("Fix: VSA cache insertion must stay infallible");
555        }
556        assert!(
557            cache.vsa_hashes.len() <= 2,
558            "Fix: bounded VSA cache must not grow past max entries"
559        );
560    }
561
562    /// The residency division is integral and both of its edges are pinned,
563    /// because this arithmetic now has exactly one definition and CUDA's
564    /// cooperative launch preflight reads it.
565    ///
566    /// A zero width has no meaningful block count and yields zero rather than
567    /// dividing. A width wider than the whole per-unit budget hosts no block at
568    /// all, so it also yields zero: that is a launch the caller must reject,
569    /// not one silently rounded up to a single block. Both match what
570    /// `cooperative_thread_residency_block_limit` did before the arithmetic
571    /// moved here, and a factoring that quietly changed either edge would be
572    /// worse than the duplicate it replaced.
573    #[test]
574    fn residency_division_is_integral_at_both_edges() {
575        assert_eq!(blocks_per_compute_unit(1536, 0), 0);
576        assert_eq!(resident_threads_per_compute_unit(1536, 0), 0);
577        assert_eq!(blocks_per_compute_unit(1536, 2048), 0);
578        assert_eq!(resident_threads_per_compute_unit(1536, 2048), 0);
579        assert_eq!(blocks_per_compute_unit(0, 256), 0);
580        assert_eq!(resident_threads_per_compute_unit(0, 256), 0);
581
582        assert_eq!(blocks_per_compute_unit(1536, 1024), 1);
583        assert_eq!(
584            resident_threads_per_compute_unit(1536, 1024),
585            1024,
586            "Fix: 1024 wide against a 1536-thread unit strands 512 slots. The truncation is the whole point of pinning this."
587        );
588        assert_eq!(blocks_per_compute_unit(1536, 256), 6);
589        assert_eq!(resident_threads_per_compute_unit(1536, 256), 1536);
590    }
591
592    /// A backend that reports no per-unit thread budget answers `unknown`, so
593    /// no residency-aware decision can be derived from a number it never gave.
594    #[test]
595    fn unreported_per_unit_budget_answers_unknown_rather_than_zero() {
596        let reported = LaunchGeometryLimits {
597            backend: "reported",
598            max_threads_per_block: 1024,
599            max_block_dim: [1024, 1024, 64],
600            max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
601            max_threads_per_sm: 1536,
602        };
603        let unreported = LaunchGeometryLimits {
604            max_threads_per_sm: 0,
605            ..reported
606        };
607
608        assert_eq!(reported.blocks_per_compute_unit(256), Some(6));
609        assert_eq!(reported.resident_threads_per_compute_unit(256), Some(1536));
610        assert_eq!(reported.blocks_per_compute_unit(0), None);
611        assert_eq!(unreported.blocks_per_compute_unit(256), None);
612        assert_eq!(unreported.resident_threads_per_compute_unit(256), None);
613    }
614}
vyre_driver/validation.rs

vyre_driver/
validation.rs
