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(source) = report.errors.into_iter().next() {
236        return Err(BackendError::Validation { source });
237    }
238
239    validate_supported_ops(program, caps.backend_id, supported_ops)
240        .map_err(|source| BackendError::Validation { source })?;
241
242    let required = vyre_foundation::program_caps::scan(program);
243    vyre_foundation::program_caps::check_backend_capabilities(
244        caps.backend_id,
245        caps.supports_subgroup_ops,
246        caps.supports_f16,
247        caps.supports_bf16,
248        caps.supports_indirect_dispatch,
249        caps.supports_trap_propagation,
250        caps.supports_distributed_collectives,
251        caps.max_workgroup_size,
252        &required,
253    )
254    .map_err(|error| BackendError::InvalidProgram {
255        fix: error.to_string(),
256    })
257}
258
259fn validate_supported_ops(
260    program: &Program,
261    backend_id: &'static str,
262    supported_ops: &HashSet<OpId>,
263) -> Result<(), vyre_foundation::ir::ValidationError> {
264    struct SupportedOpsBackend<'a> {
265        id: &'static str,
266        ops: &'a HashSet<OpId>,
267    }
268
269    impl crate::backend::Backend for SupportedOpsBackend<'_> {
270        fn id(&self) -> &'static str {
271            self.id
272        }
273
274        fn version(&self) -> &'static str {
275            env!("CARGO_PKG_VERSION")
276        }
277
278        fn supported_ops(&self) -> &HashSet<OpId> {
279            self.ops
280        }
281    }
282
283    crate::backend::validation::validate_program(
284        program,
285        &SupportedOpsBackend {
286            id: backend_id,
287            ops: supported_ops,
288        },
289    )
290}
291
292/// Launch-geometry limits reported by a concrete driver.
293#[derive(Debug, Clone, Copy, PartialEq, Eq)]
294pub struct LaunchGeometryLimits {
295    /// Backend name used in diagnostics.
296    pub backend: &'static str,
297    /// Maximum invocations in one workgroup or block.
298    pub max_threads_per_block: u32,
299    /// Maximum workgroup or block dimensions (x, y, z).
300    pub max_block_dim: [u32; 3],
301    /// Maximum workgroup count per grid dimension.
302    pub max_grid_dim: [u32; 3],
303    /// Maximum threads the device keeps resident on one compute unit (a CUDA
304    /// streaming multiprocessor, a Metal threadgroup-hosting core, and so on),
305    /// or `0` when the backend does not probe this number.
306    ///
307    /// This is the budget that decides how many whole workgroups fit on one
308    /// unit, and the division is integral: a workgroup width that does not
309    /// divide it strands the remainder on every unit for the launch's whole
310    /// lifetime. Backends that leave this `0` opt out of every residency-aware
311    /// decision rather than receiving one derived from a guessed budget.
312    pub max_threads_per_sm: u32,
313}
314
315impl LaunchGeometryLimits {
316    /// Whole workgroups of `workgroup_threads` that stay resident on one
317    /// compute unit, or `None` when this backend reports no per-unit budget.
318    #[must_use]
319    pub fn blocks_per_compute_unit(&self, workgroup_threads: u32) -> Option<u32> {
320        (self.max_threads_per_sm != 0 && workgroup_threads != 0)
321            .then(|| blocks_per_compute_unit(self.max_threads_per_sm, workgroup_threads))
322    }
323
324    /// Threads that stay resident on one compute unit at `workgroup_threads`
325    /// wide, or `None` when this backend reports no per-unit budget.
326    #[must_use]
327    pub fn resident_threads_per_compute_unit(&self, workgroup_threads: u32) -> Option<u32> {
328        (self.max_threads_per_sm != 0 && workgroup_threads != 0)
329            .then(|| resident_threads_per_compute_unit(self.max_threads_per_sm, workgroup_threads))
330    }
331}
332
333/// Whole workgroups of `workgroup_threads` that fit one compute unit's thread
334/// budget.
335///
336/// This is the single definition of the residency division in the workspace.
337/// CUDA's cooperative launch preflight and cold-start launch-width selection
338/// both route through it, because two independent copies of this arithmetic
339/// had already drifted apart once. The division is integral by hardware: a
340/// unit hosts whole workgroups only.
341///
342/// Threads are the only ceiling modelled here. Hardware also caps blocks per
343/// unit independently (CUDA reports it as
344/// `CU_DEVICE_ATTRIBUTE_MAX_BLOCKS_PER_MULTIPROCESSOR`), so at narrow widths
345/// the real block count is lower than this returns and the shortfall can be
346/// large: where the device caps blocks at 24, a 32-wide group against a
347/// 1536-thread budget measures 24 blocks and 768 resident threads, half what
348/// this function's 48 blocks and 1536 threads predict. A caller that ranks
349/// widths from widest downward never reaches that regime. A caller that
350/// answers "does this declared width fit", such as a cooperative launch
351/// preflight, does, and must clamp by the device-reported block cap before
352/// admitting a grid.
353#[must_use]
354pub fn blocks_per_compute_unit(max_threads_per_unit: u32, workgroup_threads: u32) -> u32 {
355    if workgroup_threads == 0 {
356        return 0;
357    }
358    max_threads_per_unit / workgroup_threads
359}
360
361/// Threads resident on one compute unit at `workgroup_threads` wide, under the
362/// per-unit thread budget alone.
363///
364/// Equal to `blocks_per_compute_unit(..) * workgroup_threads`, so it is at most
365/// `max_threads_per_unit` and falls short of it by exactly the slots the
366/// integral division strands. A width of 1024 against a 1536-thread budget
367/// resolves to one block and 1024 resident threads, leaving 512 slots per unit
368/// unusable for the launch's duration. The block-count caveat on
369/// [`blocks_per_compute_unit`] applies here too.
370#[must_use]
371pub fn resident_threads_per_compute_unit(max_threads_per_unit: u32, workgroup_threads: u32) -> u32 {
372    blocks_per_compute_unit(max_threads_per_unit, workgroup_threads)
373        .saturating_mul(workgroup_threads)
374}
375
376/// Validate workgroup and grid dimensions against backend launch limits.
377///
378/// # Errors
379///
380/// Returns when dimensions are zero, overflow the invocation product, exceed
381/// workgroup limits, or exceed per-axis grid limits.
382pub fn validate_launch_geometry(
383    workgroup: [u32; 3],
384    grid: [u32; 3],
385    limits: LaunchGeometryLimits,
386) -> Result<(), BackendError> {
387    if workgroup.contains(&0) || grid.contains(&0) {
388        return Err(BackendError::InvalidProgram {
389            fix: format!(
390                "Fix: {} workgroup and grid dimensions must all be non-zero.",
391                limits.backend
392            ),
393        });
394    }
395    let threads = workgroup[0]
396        .checked_mul(workgroup[1])
397        .and_then(|xy| xy.checked_mul(workgroup[2]))
398        .ok_or_else(|| BackendError::InvalidProgram {
399            fix: format!(
400                "Fix: {} workgroup dimensions overflowed u32; reduce workgroup_override.",
401                limits.backend
402            ),
403        })?;
404    if threads > limits.max_threads_per_block {
405        return Err(BackendError::InvalidProgram {
406            fix: format!(
407                "Fix: {} workgroup has {threads} threads but device max is {}.",
408                limits.backend, limits.max_threads_per_block
409            ),
410        });
411    }
412    for (axis, &dim) in workgroup.iter().enumerate() {
413        if dim > limits.max_block_dim[axis] {
414            return Err(BackendError::InvalidProgram {
415                fix: format!(
416                    "Fix: {} workgroup axis {axis} requested {} threads but device max is {}.",
417                    limits.backend, dim, limits.max_block_dim[axis]
418                ),
419            });
420        }
421    }
422    for (axis, &dim) in grid.iter().enumerate() {
423        if dim > limits.max_grid_dim[axis] {
424            return Err(BackendError::InvalidProgram {
425                fix: format!(
426                    "Fix: {} grid axis {axis} requested {} workgroups but device max is {}.",
427                    limits.backend, dim, limits.max_grid_dim[axis]
428                ),
429            });
430        }
431    }
432    Ok(())
433}
434
435/// Validate a program's effective workgroup shape against a backend's reported limits.
436///
437/// This is the shared pre-dispatch gate for callers that have a `VyreBackend`
438/// trait object but have not entered a concrete driver yet.
439///
440/// # Errors
441///
442/// Returns when any workgroup axis is zero, exceeds the backend's per-axis
443/// limit, or when total invocations exceed the backend's workgroup limit.
444pub fn validate_program_for_backend(
445    backend: &dyn VyreBackend,
446    program: &Program,
447    config: &DispatchConfig,
448) -> Result<(), BackendError> {
449    let workgroup = config
450        .workgroup_override
451        .unwrap_or(program.workgroup_size());
452    let max_axes = backend.max_workgroup_size();
453    if workgroup.contains(&0) {
454        return Err(BackendError::InvalidProgram {
455            fix: format!(
456                "Fix: backend `{}` cannot dispatch zero-sized workgroup dimensions; set positive workgroup sizes.",
457                backend.id()
458            ),
459        });
460    }
461    for (axis, &dim) in workgroup.iter().enumerate() {
462        if dim > max_axes[axis] {
463            return Err(BackendError::InvalidProgram {
464                fix: format!(
465                    "Fix: backend `{}` workgroup axis {axis} requested {} but max is {}.",
466                    backend.id(),
467                    dim,
468                    max_axes[axis]
469                ),
470            });
471        }
472    }
473    let invocations = workgroup[0]
474        .checked_mul(workgroup[1])
475        .and_then(|xy| xy.checked_mul(workgroup[2]))
476        .ok_or_else(|| BackendError::InvalidProgram {
477            fix: format!(
478                "Fix: backend `{}` workgroup dimensions overflowed u32; reduce workgroup size.",
479                backend.id()
480            ),
481        })?;
482    let max_invocations = backend.max_compute_invocations_per_workgroup();
483    if invocations > max_invocations {
484        return Err(BackendError::InvalidProgram {
485            fix: format!(
486                "Fix: backend `{}` workgroup has {invocations} invocations but max is {max_invocations}.",
487                backend.id()
488            ),
489        });
490    }
491    if let Some(grid) = config.grid_override {
492        let max_workgroups = backend.max_compute_workgroups_per_dimension();
493        if grid.contains(&0) {
494            return Err(BackendError::InvalidProgram {
495                fix: format!(
496                    "Fix: backend `{}` cannot dispatch zero-sized grid dimensions; set positive grid_override values.",
497                    backend.id()
498                ),
499            });
500        }
501        for (axis, &dim) in grid.iter().enumerate() {
502            if dim > max_workgroups {
503                return Err(BackendError::InvalidProgram {
504                    fix: format!(
505                        "Fix: backend `{}` grid_override axis {axis} requested {} workgroups but max is {}.",
506                        backend.id(),
507                        dim,
508                        max_workgroups
509                    ),
510                });
511            }
512        }
513    }
514    Ok(())
515}
516
517fn vsa_words_hash(words: &[u32]) -> blake3::Hash {
518    let mut hasher = blake3::Hasher::new();
519    hasher.update(&(words.len() as u64).to_le_bytes());
520    for word in words {
521        hasher.update(&word.to_le_bytes());
522    }
523    hasher.finalize()
524}
525
526#[cfg(test)]
527mod tests {
528    use super::*;
529
530    #[test]
531    fn validation_cache_records_vsa_without_lock_shards() {
532        let cache = ValidationCache::new(8, 8, 4);
533        let hash = blake3::hash(b"program");
534        cache
535            .remember_success(hash, &[1, 2, 3, 4])
536            .expect("Fix: lock-free VSA cache insertion must not fail");
537
538        assert!(cache.contains_hash(&hash));
539        assert_eq!(cache.vsa_hashes.len(), 1);
540        assert!(format!("{cache:?}").contains("vsa_hashes"));
541    }
542
543    #[test]
544    fn validation_cache_bounds_vsa_hashes_by_clear() {
545        let cache = ValidationCache::new(8, 2, 4);
546        for i in 0..3u32 {
547            cache
548                .remember_success(blake3::hash(&i.to_le_bytes()), &[i])
549                .expect("Fix: VSA cache insertion must stay infallible");
550        }
551        assert!(
552            cache.vsa_hashes.len() <= 2,
553            "Fix: bounded VSA cache must not grow past max entries"
554        );
555    }
556
557    /// The residency division is integral and both of its edges are pinned,
558    /// because this arithmetic now has exactly one definition and CUDA's
559    /// cooperative launch preflight reads it.
560    ///
561    /// A zero width has no meaningful block count and yields zero rather than
562    /// dividing. A width wider than the whole per-unit budget hosts no block at
563    /// all, so it also yields zero: that is a launch the caller must reject,
564    /// not one silently rounded up to a single block. Both match what
565    /// `cooperative_thread_residency_block_limit` did before the arithmetic
566    /// moved here, and a factoring that quietly changed either edge would be
567    /// worse than the duplicate it replaced.
568    #[test]
569    fn residency_division_is_integral_at_both_edges() {
570        assert_eq!(blocks_per_compute_unit(1536, 0), 0);
571        assert_eq!(resident_threads_per_compute_unit(1536, 0), 0);
572        assert_eq!(blocks_per_compute_unit(1536, 2048), 0);
573        assert_eq!(resident_threads_per_compute_unit(1536, 2048), 0);
574        assert_eq!(blocks_per_compute_unit(0, 256), 0);
575        assert_eq!(resident_threads_per_compute_unit(0, 256), 0);
576
577        assert_eq!(blocks_per_compute_unit(1536, 1024), 1);
578        assert_eq!(
579            resident_threads_per_compute_unit(1536, 1024),
580            1024,
581            "Fix: 1024 wide against a 1536-thread unit strands 512 slots. The truncation is the whole point of pinning this."
582        );
583        assert_eq!(blocks_per_compute_unit(1536, 256), 6);
584        assert_eq!(resident_threads_per_compute_unit(1536, 256), 1536);
585    }
586
587    /// A backend that reports no per-unit thread budget answers `unknown`, so
588    /// no residency-aware decision can be derived from a number it never gave.
589    #[test]
590    fn unreported_per_unit_budget_answers_unknown_rather_than_zero() {
591        let reported = LaunchGeometryLimits {
592            backend: "reported",
593            max_threads_per_block: 1024,
594            max_block_dim: [1024, 1024, 64],
595            max_grid_dim: [u32::MAX, u32::MAX, u32::MAX],
596            max_threads_per_sm: 1536,
597        };
598        let unreported = LaunchGeometryLimits {
599            max_threads_per_sm: 0,
600            ..reported
601        };
602
603        assert_eq!(reported.blocks_per_compute_unit(256), Some(6));
604        assert_eq!(reported.resident_threads_per_compute_unit(256), Some(1536));
605        assert_eq!(reported.blocks_per_compute_unit(0), None);
606        assert_eq!(unreported.blocks_per_compute_unit(256), None);
607        assert_eq!(unreported.resident_threads_per_compute_unit(256), None);
608    }
609}
vyre_driver/validation.rs

vyre_driver/
validation.rs
