vyre_driver_cuda/
lib.rs

1//! # vyre-driver-cuda  -  CUDA/PTX backend for vyre
2//!
3//! Implements [`VyreBackend`] via the CUDA driver API through `cudarc`.
4//! Translates vyre `Program` IR into PTX kernels, loads them through
5//! the CUDA driver JIT, and dispatches on NVIDIA GPUs.
6//!
7//! The backend registers itself as `"cuda"` in the vyre backend registry
8//! via `inventory::submit!` so `vyre::registered_backends()` enumerates
9//! it alongside `wgpu`, `spirv`, etc.
10//!
11//! ## Architecture
12//!
13//! ```text
14//!    Program ─► PTX emitter ─► cuModuleLoadData ─► cuLaunchKernel
15//! ```
16//!
17#![deny(missing_docs)]
18// CUDA driver bindings (`cudarc::driver::sys::cu*`) are inherently unsafe FFI;
19// every call site is the boundary between safe vyre code and the CUDA driver
20// API. Allow `unsafe` here so the rest of the workspace can keep
21// `unsafe_code = "deny"` while this backend wraps cudarc properly with
22// per-call Safety: comments enforced by `check_unsafe_justifications.sh`.
23#![allow(unsafe_code)]
24
25mod aot_launcher;
26/// CUDA backend core: device management and dispatch.
27pub mod backend;
28/// PTX code generation from vyre IR.
29pub mod codegen;
30/// CUDA device capability probing.
31pub mod device;
32/// CUDA upload planning for GPU e-graph device images.
33pub mod egraph_device_image;
34/// CUDA launch-wave planning for resident e-graph device images.
35pub mod egraph_kernel_plan;
36mod egraph_readback;
37/// Adapter from frontier-typed IR plans to CUDA frontier wave envelopes.
38pub mod frontier_typed_ir_adapter;
39mod input_identity;
40mod instrumentation;
41/// Cross-process persistent CUDA JIT cache wiring (E4 + E5): configures
42/// the NVIDIA driver's built-in disk cache at backend bring-up so the
43/// JIT-compiled cuBINs persist across runs and are shared across every
44/// vyre process on the host.
45pub mod jit_cache;
46/// Actionable CUDA kernel capability diagnostics.
47pub mod kernel_failure_diagnostics;
48mod materializer;
49/// Bounded CUDA megakernel plan cache keyed by graph, analysis, device, and
50/// runtime pressure buckets.
51pub mod megakernel_plan_cache;
52mod numeric;
53/// Occupancy-aware empirical autotuning (I4): pure estimator that picks
54/// the workgroup size with the highest predicted hardware occupancy from
55/// `(CudaDeviceCaps, KernelResourceUsage)`. The runtime feeds the result
56/// into `AutotuneStore` (I3) so subsequent dispatches reuse the choice.
57pub mod occupancy;
58/// Self-hosted optimizer GPU dispatcher  -  runs the
59/// `vyre-self-substrate::optimizer` passes (DCE, CSE, const-fold,
60/// validator) on CUDA. External parity tests reach in via the
61/// `CudaOptimizerDispatcher` re-export below.
62pub mod optimizer;
63mod pipeline;
64/// CUDA profiler range integration for Nsight/NVTX without mandatory NVTX linkage.
65pub mod profiler;
66/// CUDA regex hardware-comparison evidence.
67pub mod regex_hardware_comparison;
68/// Repeated execution over persistent CUDA-resident graph state.
69pub mod resident_graph_session;
70mod stream;
71// Neutral policies are imported from `vyre-driver`; CUDA exports only concrete behavior.
72/// A fixed synthetic device envelope for context-free estimator tests. Not a
73/// probe, and not this machine's values: never derive a hardware decision from it.
74pub mod synthetic_device_caps;
75mod target_compiler;
76/// CUDA execution planning for unified token/fact graph frontier waves.
77pub mod token_fact_frontier_execution;
78/// Adapter from unified token/fact graph layouts to CUDA resident bytes.
79pub mod token_fact_graph_cuda_adapter;
80/// CUDA warp-word bit-parallel automata layout evidence.
81pub mod warp_word_automata;
82
83pub use backend::{
84    CudaBackend, CudaPtxSourceCacheSnapshot, CudaResidentBuffer, CudaStreamOrderedPool,
85    CudaTelemetrySnapshot,
86};
87pub use stream::CudaLaunchResourceCounts;
88/// CUDA megakernel global-barrier minimization for dependency-typed waves.
89pub mod megakernel_barrier_planner;
90pub mod megakernel_scheduler;
91/// Release gate for steady-state CUDA megakernel speedup claims.
92pub mod megakernel_speedup_gate;
93pub use device::{CudaDeviceCaps, CudaDeviceHandle};
94pub use egraph_device_image::{
95    plan_cuda_egraph_device_upload, plan_cuda_egraph_device_upload_from_image,
96    plan_cuda_egraph_device_upload_from_image_ref, CudaEGraphDeviceBorrowedUploadPlan,
97    CudaEGraphDeviceByteLayout, CudaEGraphDeviceByteSpan, CudaEGraphDeviceKernelView,
98    CudaEGraphDeviceUploadError, CudaEGraphDeviceUploadPlan, CudaResidentEGraphDeviceImage,
99};
100pub use egraph_kernel_plan::{
101    collect_cuda_egraph_structural_equivalences, cuda_egraph_canonical_rewrite_kernel_ptx,
102    cuda_egraph_signature_pair_rows, cuda_egraph_signature_refresh_kernel_ptx,
103    cuda_egraph_structural_equivalence_kernel_ptx, pack_cuda_egraph_canonical_rewrite_device_image,
104    pack_cuda_egraph_signature_bucket_device_image, plan_cuda_egraph_kernel_work,
105    plan_cuda_egraph_signature_buckets, plan_cuda_egraph_signature_buckets_from_resident_snapshot,
106    plan_cuda_egraph_signature_buckets_from_signature_snapshot,
107    plan_cuda_egraph_structural_equivalence_launch_artifact,
108    plan_cuda_egraph_structural_equivalence_output, plan_cuda_egraph_structural_equivalences,
109    plan_cuda_egraph_union_compaction, CudaEGraphCanonicalRewrite,
110    CudaEGraphCanonicalRewriteDeviceImage, CudaEGraphCanonicalRewriteKernelPtx,
111    CudaEGraphCanonicalRewriteKernelResult, CudaEGraphFixedPointReadback,
112    CudaEGraphKernelLaunchConfig, CudaEGraphKernelPass, CudaEGraphKernelPlanError,
113    CudaEGraphKernelWave, CudaEGraphKernelWorkPlan, CudaEGraphResidentColumnSnapshot,
114    CudaEGraphResidentSignatureSnapshot, CudaEGraphSignatureBucket,
115    CudaEGraphSignatureBucketDeviceImage, CudaEGraphSignatureBucketPlan,
116    CudaEGraphSignaturePairWave, CudaEGraphSignatureRefreshKernelPtx,
117    CudaEGraphSignatureRefreshKernelResult, CudaEGraphStructuralCanonicalizationFixedPointReport,
118    CudaEGraphStructuralCanonicalizationFixedPointResult,
119    CudaEGraphStructuralCanonicalizationRoundResult, CudaEGraphStructuralEquivalenceKernelPtx,
120    CudaEGraphStructuralEquivalenceKernelResult, CudaEGraphStructuralEquivalenceLaunchArtifact,
121    CudaEGraphStructuralEquivalenceOutputPlan, CudaEGraphStructuralEquivalencePlan,
122    CudaEGraphUnionCompactionPass, CudaEGraphUnionCompactionPlan, CudaEGraphUnionCompactionWave,
123    CUDA_EGRAPH_CANONICAL_REWRITE_KERNEL_ENTRY, CUDA_EGRAPH_CANONICAL_REWRITE_KERNEL_PARAM_COUNT,
124    CUDA_EGRAPH_CANONICAL_REWRITE_RECORD_WORDS, CUDA_EGRAPH_SIGNATURE_BUCKET_RECORD_WORDS,
125    CUDA_EGRAPH_SIGNATURE_REFRESH_KERNEL_ENTRY, CUDA_EGRAPH_SIGNATURE_REFRESH_KERNEL_PARAM_COUNT,
126    CUDA_EGRAPH_STRUCTURAL_EQUIVALENCE_KERNEL_ENTRY,
127    CUDA_EGRAPH_STRUCTURAL_EQUIVALENCE_KERNEL_PARAM_COUNT,
128};
129pub use frontier_typed_ir_adapter::{
130    adapt_frontier_typed_ir_to_cuda, CudaFrontierTypedIrAdapterError, CudaFrontierTypedIrInput,
131};
132pub use kernel_failure_diagnostics::{
133    diagnose_cuda_kernel_launch, diagnose_cuda_kernel_launch_shape,
134    diagnose_cuda_kernel_launch_with_scratch, CudaKernelCapabilityFailure,
135    CudaKernelDeviceEnvelope, CudaKernelLaunchDiagnostic, CudaKernelLaunchDiagnosticRef,
136    CudaKernelLaunchDiagnosticScratch, CudaKernelLaunchEnvelope, CudaKernelLaunchEnvelopeError,
137    CudaKernelLaunchShape, CudaKernelRequirement,
138};
139pub use megakernel_barrier_planner::{
140    plan_cuda_frontier_megakernel_execution, plan_cuda_frontier_megakernel_execution_with_scratch,
141    CudaMegakernelFrontierExecutionPlan, CudaMegakernelFrontierExecutionPlanError,
142};
143pub use megakernel_plan_cache::{
144    CudaMegakernelAnalysisKind, CudaMegakernelCachedPlan, CudaMegakernelDeviceKey,
145    CudaMegakernelPlanCache, CudaMegakernelPlanCacheKey, CudaMegakernelPlanCacheStats,
146};
147pub use megakernel_scheduler::{
148    schedule_megakernel_from_cuda_samples, schedule_megakernel_from_cuda_samples_into,
149    select_cuda_megakernel_topology, CudaMegakernelScheduleSample,
150};
151pub use megakernel_speedup_gate::{
152    format_validated_cuda_megakernel_speedup_evidence_csv,
153    validate_cuda_megakernel_speedup_evidence_csv, validate_cuda_megakernel_speedup_gate,
154    CudaMegakernelSpeedupGateError, CudaMegakernelSpeedupProof, CudaMegakernelSpeedupSample,
155    MEGAKERNEL_SPEEDUP_EVIDENCE_CSV_HEADER,
156};
157pub use optimizer::CudaOptimizerDispatcher;
158pub use regex_hardware_comparison::{
159    cuda_regex_hardware_comparison_evidence, cuda_regex_software_fallback_comparison_evidence,
160    CudaRegexHardwareComparisonEvidence, CUDA_REGEX_HARDWARE_COMPARISON_SCHEMA_VERSION,
161};
162pub use resident_graph_session::{
163    format_validated_cuda_resident_graph_session_evidence_csv, plan_cuda_resident_graph_session,
164    resident_graph_session_speedup_sample, CudaResidentGraphReadback,
165    CudaResidentGraphSessionError, CudaResidentGraphSessionEvidence,
166    CudaResidentGraphSessionEvidenceError, CudaResidentGraphSessionPlan,
167    CudaResidentGraphSessionProfile,
168};
169pub use token_fact_frontier_execution::{
170    plan_cuda_token_fact_frontier_execution, plan_cuda_token_fact_frontier_execution_with_scratch,
171    CudaTokenFactFrontierExecutionError, CudaTokenFactFrontierExecutionPlan,
172};
173pub use token_fact_graph_cuda_adapter::{
174    adapt_token_fact_graph_to_cuda_layout, CudaTokenFactGraphLayout, CudaTokenFactGraphLayoutError,
175    CUDA_TOKEN_FACT_DEGREE_PROFILE_BUCKETS, CUDA_TOKEN_FACT_DEGREE_PROFILE_RANKS,
176};
177pub use warp_word_automata::{
178    plan_cuda_warp_word_automata_layout, CudaWarpWordAutomataLayoutError,
179    CudaWarpWordAutomataLayoutEvidence, CudaWarpWordAutomataLayoutRequest,
180    CudaWarpWordInstructionClass, CUDA_WARP_WORD_AUTOMATA_LAYOUT_SCHEMA_VERSION,
181};
182
183use crate::backend::staging_reserve::reserve_smallvec;
184use smallvec::SmallVec;
185use vyre_driver::{BackendError, BackendRegistration, DispatchConfig, Resource, VyreBackend};
186use vyre_foundation::ir::Program;
187use vyre_foundation::operation::TargetId;
188
189/// Stable backend identifier for registration and conform certificates.
190pub const CUDA_BACKEND_ID: &str = "cuda";
191/// Validated target identity owned by the CUDA driver.
192pub const CUDA_TARGET_ID: TargetId = TargetId::expect_valid(CUDA_BACKEND_ID);
193
194/// CUDA implementation of [`vyre_driver::DeviceBuffer`]. Wraps a
195/// [`backend::CudaResidentBuffer`] handle so consumers can hold a
196/// `Box<dyn DeviceBuffer>` against the CUDA backend without naming
197/// `CudaResidentBuffer` directly.
198///
199/// Lifecycle is explicit-free  -  call
200/// `VyreBackend::free_device_buffer(boxed_buffer)` when done. This
201/// matches the existing CUDA-resident contract and keeps the substrate
202/// free of reference-counted backend handles. A future RAII variant
203/// (Drop-managed via `Arc<CudaBackend>`) can ship as a drop-in
204/// replacement when the backend ownership model accommodates it.
205#[derive(Debug)]
206pub struct CudaDeviceBuffer {
207    backend_id: &'static str,
208    handle: backend::CudaResidentBuffer,
209}
210
211impl vyre_driver::DeviceBuffer for CudaDeviceBuffer {
212    fn backend_id(&self) -> &'static str {
213        self.backend_id
214    }
215
216    fn byte_len(&self) -> usize {
217        self.handle.byte_len
218    }
219
220    fn as_any(&self) -> &dyn std::any::Any {
221        self
222    }
223
224    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
225        self
226    }
227}
228
229/// Factory wrapper for the inventory registration path.
230///
231/// Unlike the SPIR-V backend, the CUDA backend owns a live device handle
232/// and can dispatch programs directly.
233#[derive(Debug)]
234pub struct CudaBackendRegistration {
235    inner: CudaBackend,
236}
237
238type ResolvedUploads<'a> = SmallVec<[(CudaResidentBuffer, &'a [u8]); 8]>;
239type ResolvedOffsetUploads<'a> = SmallVec<[(CudaResidentBuffer, usize, &'a [u8]); 8]>;
240type ResolvedDownloadRanges = SmallVec<[(CudaResidentBuffer, usize, usize); 8]>;
241type ResolvedReadRanges = (
242    SmallVec<[CudaResidentBuffer; 8]>,
243    SmallVec<[crate::backend::output_range::CudaOutputReadback; 8]>,
244);
245
246impl CudaBackendRegistration {
247    /// Wrap an already-acquired [`CudaBackend`] as a [`VyreBackend`] trait object.
248    ///
249    /// The inventory-driven path uses [`cuda_factory`] which acquires its own
250    /// device handle. Callers that already own a [`CudaBackend`] (e.g. so they
251    /// can keep the live device handle for direct API access while also handing
252    /// it to a megakernel) use this constructor instead.
253    #[must_use]
254    pub fn new(inner: CudaBackend) -> Self {
255        Self { inner }
256    }
257
258    /// Borrow the inner [`CudaBackend`] for direct device-API access.
259    #[must_use]
260    pub fn inner(&self) -> &CudaBackend {
261        &self.inner
262    }
263
264    /// Snapshot the CUDA PTX-source cache used before driver module loading.
265    #[must_use]
266    pub fn ptx_source_cache_snapshot(&self) -> CudaPtxSourceCacheSnapshot {
267        self.inner.ptx_source_cache_snapshot()
268    }
269
270    /// Runtime CUDA telemetry counters for release-path performance gates.
271    #[must_use]
272    pub fn telemetry_snapshot(&self) -> CudaTelemetrySnapshot {
273        self.inner.telemetry_snapshot()
274    }
275
276    /// Reset runtime CUDA telemetry counters without clearing backend caches.
277    pub fn reset_telemetry(&self) {
278        self.inner.reset_telemetry();
279    }
280
281    fn resolve_uploads<'a>(
282        &self,
283        uploads: &[(&Resource, &'a [u8])],
284    ) -> Result<ResolvedUploads<'a>, BackendError> {
285        let mut concrete = SmallVec::<[(CudaResidentBuffer, &'a [u8]); 8]>::new();
286        reserve_smallvec(&mut concrete, uploads.len(), "CUDA resident upload handles")?;
287        for (resource, bytes) in uploads {
288            let handle = self.inner.resident_handle_from_resource(resource)?;
289            concrete.push((handle, *bytes));
290        }
291        Ok(concrete)
292    }
293
294    fn resolve_offset_uploads<'a>(
295        &self,
296        uploads: &[(&Resource, usize, &'a [u8])],
297    ) -> Result<ResolvedOffsetUploads<'a>, BackendError> {
298        let mut concrete = SmallVec::<[(CudaResidentBuffer, usize, &'a [u8]); 8]>::new();
299        reserve_smallvec(
300            &mut concrete,
301            uploads.len(),
302            "CUDA resident offset upload handles",
303        )?;
304        for (resource, dst_offset_bytes, bytes) in uploads {
305            let handle = self.inner.resident_handle_from_resource(resource)?;
306            concrete.push((handle, *dst_offset_bytes, *bytes));
307        }
308        Ok(concrete)
309    }
310
311    fn resolve_download_ranges(
312        &self,
313        ranges: &[(&Resource, usize, usize)],
314    ) -> Result<ResolvedDownloadRanges, BackendError> {
315        let mut concrete = SmallVec::<[(CudaResidentBuffer, usize, usize); 8]>::new();
316        reserve_smallvec(
317            &mut concrete,
318            ranges.len(),
319            "CUDA resident download range handles",
320        )?;
321        for (resource, byte_offset, byte_len) in ranges {
322            let handle = self.inner.resident_handle_from_resource(resource)?;
323            concrete.push((handle, *byte_offset, *byte_len));
324        }
325        Ok(concrete)
326    }
327
328    fn resolve_read_ranges(
329        &self,
330        read_ranges: &[vyre_driver::backend::ResidentReadRange<'_>],
331    ) -> Result<ResolvedReadRanges, BackendError> {
332        let mut handles = SmallVec::<[CudaResidentBuffer; 8]>::new();
333        let mut concrete_readbacks =
334            SmallVec::<[crate::backend::output_range::CudaOutputReadback; 8]>::new();
335        reserve_smallvec(
336            &mut handles,
337            read_ranges.len(),
338            "CUDA resident read handles",
339        )?;
340        reserve_smallvec(
341            &mut concrete_readbacks,
342            read_ranges.len(),
343            "CUDA resident readback ranges",
344        )?;
345        for range in read_ranges {
346            handles.push(self.inner.resident_handle_from_resource(range.resource)?);
347            concrete_readbacks.push(crate::backend::output_range::CudaOutputReadback {
348                device_offset: range.byte_offset,
349                byte_len: range.byte_len,
350            });
351        }
352        Ok((handles, concrete_readbacks))
353    }
354
355    fn resolve_step_handle_sets(
356        &self,
357        steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
358        field: &'static str,
359    ) -> Result<SmallVec<[SmallVec<[crate::backend::CudaResidentBuffer; 8]>; 8]>, BackendError>
360    {
361        let mut handle_sets =
362            SmallVec::<[SmallVec<[crate::backend::CudaResidentBuffer; 8]>; 8]>::new();
363        reserve_smallvec(&mut handle_sets, steps.len(), field)?;
364        for step in steps {
365            handle_sets.push(self.inner.resident_handles_from_resources(step.resources)?);
366        }
367        Ok(handle_sets)
368    }
369
370    fn resolve_repeated_step_handle_sets(
371        &self,
372        steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
373        repeat_count: usize,
374    ) -> Result<SmallVec<[SmallVec<[crate::backend::CudaResidentBuffer; 8]>; 8]>, BackendError>
375    {
376        let mut handle_sets =
377            SmallVec::<[SmallVec<[crate::backend::CudaResidentBuffer; 8]>; 8]>::new();
378        let capacity = if repeat_count == 0 { 0 } else { steps.len() };
379        reserve_smallvec(
380            &mut handle_sets,
381            capacity,
382            "CUDA repeated resident repeated handle sets",
383        )?;
384        if repeat_count != 0 {
385            for step in steps {
386                handle_sets.push(self.inner.resident_handles_from_resources(step.resources)?);
387            }
388        }
389        Ok(handle_sets)
390    }
391
392    fn concrete_resident_steps<'program: 'handles, 'handles>(
393        steps: &[vyre_driver::backend::ResidentDispatchStep<'program>],
394        handle_sets: &'handles [SmallVec<[crate::backend::CudaResidentBuffer; 8]>],
395        field: &'static str,
396    ) -> Result<SmallVec<[crate::backend::CudaResidentDispatchStep<'handles>; 8]>, BackendError>
397    {
398        let mut concrete_steps =
399            SmallVec::<[crate::backend::CudaResidentDispatchStep<'handles>; 8]>::new();
400        reserve_smallvec(&mut concrete_steps, handle_sets.len(), field)?;
401        for (step, handles) in steps.iter().zip(handle_sets.iter()) {
402            let mut config = DispatchConfig::default();
403            config.grid_override = step.grid_override;
404            config.workgroup_override = step.workgroup_override;
405            concrete_steps.push(crate::backend::CudaResidentDispatchStep {
406                program: step.program,
407                handles,
408                config,
409            });
410        }
411        Ok(concrete_steps)
412    }
413
414    /// Bytes of transient CUDA device memory currently owned by the transient pool.
415    ///
416    /// This includes checked-out dispatch allocations, compiled-pipeline static parameter
417    /// allocations, and cached transient blocks retained for reuse.
418    ///
419    /// # Errors
420    ///
421    /// Returns [`BackendError`] if allocation accounting cannot be read.
422    pub fn allocated_transient_allocation_bytes(&self) -> Result<usize, BackendError> {
423        self.inner.allocated_transient_allocation_bytes()
424    }
425
426    fn reject_grid_sync_without_native_lowering(
427        &self,
428        program: &Program,
429    ) -> Result<(), BackendError> {
430        if vyre_driver::grid_sync::contains_grid_sync(program) && !self.supports_grid_sync() {
431            return Err(BackendError::UnsupportedFeature {
432                name: "cuda_native_grid_sync_lowering (MemoryOrdering::GridSync requires explicit split routing or native cooperative-grid barrier lowering)"
433                    .to_string(),
434                backend: CUDA_BACKEND_ID.to_string(),
435            });
436        }
437        Ok(())
438    }
439
440    fn validate_program_for_dispatch(&self, program: &Program) -> Result<(), BackendError> {
441        let required = vyre_foundation::program_caps::scan(program);
442        vyre_foundation::program_caps::check_backend_capabilities(
443            CUDA_BACKEND_ID,
444            self.supports_subgroup_ops(),
445            self.supports_f16(),
446            self.supports_bf16(),
447            self.supports_indirect_dispatch(),
448            true,
449            self.supports_distributed_collectives(),
450            self.max_workgroup_size(),
451            &required,
452        )
453        .map_err(|error| BackendError::InvalidProgram {
454            fix: error.to_string(),
455        })?;
456        self.reject_grid_sync_without_native_lowering(program)
457    }
458
459    fn validate_resident_steps_for_dispatch(
460        &self,
461        steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
462    ) -> Result<(), BackendError> {
463        for step in steps {
464            self.validate_program_for_dispatch(step.program)?;
465        }
466        Ok(())
467    }
468}
469
470impl vyre_driver::backend::private::Sealed for CudaBackendRegistration {}
471
472impl VyreBackend for CudaBackendRegistration {
473    fn id(&self) -> &'static str {
474        CUDA_BACKEND_ID
475    }
476
477    fn version(&self) -> &'static str {
478        env!("CARGO_PKG_VERSION")
479    }
480
481    fn dispatch(
482        &self,
483        program: &Program,
484        inputs: &[Vec<u8>],
485        config: &DispatchConfig,
486    ) -> Result<Vec<Vec<u8>>, BackendError> {
487        self.validate_program_for_dispatch(program)?;
488        self.inner.dispatch(program, inputs, config)
489    }
490
491    fn dispatch_async(
492        &self,
493        program: &Program,
494        inputs: &[Vec<u8>],
495        config: &DispatchConfig,
496    ) -> Result<Box<dyn vyre_driver::PendingDispatch>, BackendError> {
497        self.validate_program_for_dispatch(program)?;
498        self.inner.dispatch_async(program, inputs, config)
499    }
500
501    fn dispatch_borrowed_async(
502        &self,
503        program: &Program,
504        inputs: &[&[u8]],
505        config: &DispatchConfig,
506    ) -> Result<Box<dyn vyre_driver::PendingDispatch>, BackendError> {
507        self.validate_program_for_dispatch(program)?;
508        self.inner.dispatch_borrowed_async(program, inputs, config)
509    }
510
511    fn dispatch_borrowed(
512        &self,
513        program: &Program,
514        inputs: &[&[u8]],
515        config: &DispatchConfig,
516    ) -> Result<Vec<Vec<u8>>, BackendError> {
517        self.validate_program_for_dispatch(program)?;
518        self.inner
519            .dispatch_borrowed_async(program, inputs, config)?
520            .await_result()
521    }
522
523    fn dispatch_borrowed_into(
524        &self,
525        program: &Program,
526        inputs: &[&[u8]],
527        config: &DispatchConfig,
528        outputs: &mut vyre_driver::OutputBuffers,
529    ) -> Result<(), BackendError> {
530        self.validate_program_for_dispatch(program)?;
531        self.inner
532            .dispatch_borrowed_async(program, inputs, config)?
533            .await_result_into(outputs)
534    }
535
536    fn dispatch_borrowed_timed(
537        &self,
538        program: &Program,
539        inputs: &[&[u8]],
540        config: &DispatchConfig,
541    ) -> Result<vyre_driver::TimedDispatchResult, BackendError> {
542        self.validate_program_for_dispatch(program)?;
543        self.inner.dispatch_borrowed_timed(program, inputs, config)
544    }
545
546    fn allocate_resident(&self, byte_len: usize) -> Result<Resource, BackendError> {
547        self.inner
548            .allocate_resident(byte_len)
549            .map(|handle| Resource::Resident(handle.handle))
550    }
551
552    fn allocate_device_buffer(
553        &self,
554        byte_len: usize,
555    ) -> Result<Box<dyn vyre_driver::DeviceBuffer>, BackendError> {
556        let handle = self.inner.allocate_resident(byte_len)?;
557        Ok(Box::new(CudaDeviceBuffer {
558            backend_id: CUDA_BACKEND_ID,
559            handle,
560        }))
561    }
562
563    fn upload_device_buffer(
564        &self,
565        buffer: &mut dyn vyre_driver::DeviceBuffer,
566        bytes: &[u8],
567    ) -> Result<(), BackendError> {
568        let backend_id = buffer.backend_id().to_string();
569        let handle = buffer
570            .as_any_mut()
571            .downcast_mut::<CudaDeviceBuffer>()
572            .map(|cuda_buf| cuda_buf.handle)
573            .ok_or_else(|| BackendError::InvalidProgram {
574                fix: format!(
575                    "Fix: upload_device_buffer expected a CudaDeviceBuffer (allocated by `cuda` backend) but got buffer owned by `{backend_id}`."
576                ),
577            })?;
578        self.inner.upload_resident(handle, bytes)
579    }
580
581    fn download_device_buffer(
582        &self,
583        buffer: &dyn vyre_driver::DeviceBuffer,
584    ) -> Result<Vec<u8>, BackendError> {
585        let cuda_buf = buffer
586            .as_any()
587            .downcast_ref::<CudaDeviceBuffer>()
588            .ok_or_else(|| BackendError::InvalidProgram {
589                fix: format!(
590                    "Fix: download_device_buffer expected a CudaDeviceBuffer (allocated by `cuda` backend) but got buffer owned by `{}`.",
591                    buffer.backend_id()
592                ),
593            })?;
594        self.inner.download_resident(cuda_buf.handle)
595    }
596
597    fn free_device_buffer(
598        &self,
599        buffer: Box<dyn vyre_driver::DeviceBuffer>,
600    ) -> Result<(), BackendError> {
601        let backend_id = buffer.backend_id().to_string();
602        let handle = buffer
603            .as_any()
604            .downcast_ref::<CudaDeviceBuffer>()
605            .map(|cuda_buf| cuda_buf.handle)
606            .ok_or_else(|| BackendError::InvalidProgram {
607                fix: format!(
608                    "Fix: free_device_buffer expected a CudaDeviceBuffer but got buffer owned by `{backend_id}`."
609                ),
610            })?;
611        // Drop the Box (releases the wrapper allocation) before freeing
612        // the underlying CUDA-resident allocation. CudaResidentBuffer is
613        // Copy so we already captured the handle.
614        drop(buffer);
615        self.inner.free_resident(handle)
616    }
617
618    fn dispatch_with_device_buffers(
619        &self,
620        program: &Program,
621        inputs: &[&dyn vyre_driver::DeviceBuffer],
622        outputs: &mut [&mut dyn vyre_driver::DeviceBuffer],
623        config: &DispatchConfig,
624    ) -> Result<(), BackendError> {
625        self.validate_program_for_dispatch(program)?;
626        // Convert &[&dyn DeviceBuffer] into &[Resource::Resident(id)]
627        // so we can re-use the existing dispatch_resident_timed path.
628        // Outputs are bound by Resource::Resident as well  -  the kernel
629        // writes results in-place into the device-resident buffers; the
630        // caller reads them via download_device_buffer afterwards.
631        vyre_driver::validate_buffer_ownership(self.id(), inputs.iter().copied())?;
632        vyre_driver::validate_buffer_ownership(
633            self.id(),
634            outputs
635                .iter()
636                .map(|b| &**b as &dyn vyre_driver::DeviceBuffer),
637        )?;
638        let resource_capacity =
639            inputs
640                .len()
641                .checked_add(outputs.len())
642                .ok_or_else(|| BackendError::InvalidProgram {
643                    fix: format!(
644                        "Fix: CUDA borrowed dispatch resource capacity overflowed usize for {} input buffer(s) plus {} output buffer(s); split the dispatch.",
645                        inputs.len(),
646                        outputs.len()
647                    ),
648                })?;
649        let mut handles = SmallVec::<[CudaResidentBuffer; 8]>::new();
650        reserve_smallvec(
651            &mut handles,
652            resource_capacity,
653            "CUDA borrowed dispatch resource handles",
654        )?;
655        for buffer in inputs {
656            let handle = buffer
657                .as_any()
658                .downcast_ref::<CudaDeviceBuffer>()
659                .ok_or_else(|| BackendError::InvalidProgram {
660                    fix: format!(
661                        "Fix: dispatch_with_device_buffers expected CudaDeviceBuffer inputs but got buffer owned by `{}`.",
662                        buffer.backend_id()
663                    ),
664                })?
665                .handle;
666            handles.push(handle);
667        }
668        for buffer in outputs.iter() {
669            let backend_id = buffer.backend_id().to_string();
670            let handle = buffer
671                .as_any()
672                .downcast_ref::<CudaDeviceBuffer>()
673                .ok_or_else(|| BackendError::InvalidProgram {
674                    fix: format!(
675                        "Fix: dispatch_with_device_buffers expected CudaDeviceBuffer outputs but got buffer owned by `{backend_id}`."
676                    ),
677                })?
678                .handle;
679            handles.push(handle);
680        }
681        let _timed = self
682            .inner
683            .dispatch_resident_timed(program, &handles, config)?;
684        Ok(())
685    }
686
687    fn upload_resident(&self, resource: &Resource, bytes: &[u8]) -> Result<(), BackendError> {
688        let handle = self.inner.resident_handle_from_resource(resource)?;
689        self.inner.upload_resident(handle, bytes)
690    }
691
692    fn upload_resident_many(&self, uploads: &[(&Resource, &[u8])]) -> Result<(), BackendError> {
693        let concrete = self.resolve_uploads(uploads)?;
694        self.inner.upload_resident_many(&concrete)
695    }
696
697    fn upload_resident_at(
698        &self,
699        resource: &Resource,
700        dst_offset_bytes: usize,
701        bytes: &[u8],
702    ) -> Result<(), BackendError> {
703        let handle = self.inner.resident_handle_from_resource(resource)?;
704        self.inner
705            .upload_resident_at(handle, dst_offset_bytes, bytes)
706    }
707
708    fn upload_resident_at_many(
709        &self,
710        uploads: &[(&Resource, usize, &[u8])],
711    ) -> Result<(), BackendError> {
712        let concrete = self.resolve_offset_uploads(uploads)?;
713        self.inner.upload_resident_at_many(&concrete)
714    }
715
716    fn download_resident(&self, resource: &Resource) -> Result<Vec<u8>, BackendError> {
717        let handle = self.inner.resident_handle_from_resource(resource)?;
718        self.inner.download_resident(handle)
719    }
720
721    fn download_resident_into(
722        &self,
723        resource: &Resource,
724        out: &mut Vec<u8>,
725    ) -> Result<(), BackendError> {
726        let handle = self.inner.resident_handle_from_resource(resource)?;
727        self.inner.download_resident_into(handle, out)
728    }
729
730    fn download_resident_range(
731        &self,
732        resource: &Resource,
733        byte_offset: usize,
734        byte_len: usize,
735    ) -> Result<Vec<u8>, BackendError> {
736        let handle = self.inner.resident_handle_from_resource(resource)?;
737        self.inner
738            .download_resident_range(handle, byte_offset, byte_len)
739    }
740
741    fn download_resident_range_into(
742        &self,
743        resource: &Resource,
744        byte_offset: usize,
745        byte_len: usize,
746        out: &mut Vec<u8>,
747    ) -> Result<(), BackendError> {
748        let handle = self.inner.resident_handle_from_resource(resource)?;
749        self.inner
750            .download_resident_range_into(handle, byte_offset, byte_len, out)
751    }
752
753    fn download_resident_ranges_into(
754        &self,
755        ranges: &[(&Resource, usize, usize)],
756        outputs: &mut [&mut Vec<u8>],
757    ) -> Result<(), BackendError> {
758        let concrete = self.resolve_download_ranges(ranges)?;
759        self.inner.download_resident_ranges_into(&concrete, outputs)
760    }
761
762    fn free_resident(&self, resource: Resource) -> Result<(), BackendError> {
763        let handle = self.inner.resident_handle_from_resource(&resource)?;
764        self.inner.free_resident(handle)
765    }
766
767    fn dispatch_resident_timed(
768        &self,
769        program: &Program,
770        resources: &[Resource],
771        config: &DispatchConfig,
772    ) -> Result<vyre_driver::TimedDispatchResult, BackendError> {
773        self.validate_program_for_dispatch(program)?;
774        let handles = self.inner.resident_handles_from_resources(resources)?;
775        self.inner
776            .dispatch_resident_timed(program, &handles, config)
777    }
778
779    fn dispatch_resident_async(
780        &self,
781        program: &Program,
782        resources: &[Resource],
783        config: &DispatchConfig,
784    ) -> Result<Box<dyn vyre_driver::PendingDispatch>, BackendError> {
785        self.validate_program_for_dispatch(program)?;
786        let handles = self.inner.resident_handles_from_resources(resources)?;
787        self.inner
788            .dispatch_resident_async(program, &handles, config)
789    }
790
791    fn dispatch_resident_sequence_read_ranges_into(
792        &self,
793        steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
794        read_ranges: &[vyre_driver::backend::ResidentReadRange<'_>],
795        outputs: &mut [&mut Vec<u8>],
796    ) -> Result<(), BackendError> {
797        self.validate_resident_steps_for_dispatch(steps)?;
798        if read_ranges.len() != outputs.len() {
799            return Err(BackendError::InvalidProgram {
800                fix: format!(
801                    "Fix: CUDA resident sequence ranged readback expected matching range/output counts but got {} range(s) and {} output(s).",
802                    read_ranges.len(),
803                    outputs.len()
804                ),
805            });
806        }
807        let handle_sets =
808            self.resolve_step_handle_sets(steps, "CUDA resident sequence handle sets")?;
809        let concrete_steps =
810            Self::concrete_resident_steps(steps, &handle_sets, "CUDA resident sequence steps")?;
811
812        let (read_handles, concrete_readbacks) = self.resolve_read_ranges(read_ranges)?;
813
814        let uploads: [(crate::backend::CudaResidentBuffer, &[u8]); 0] = [];
815        self.inner
816            .upload_resident_many_sequence_read_ranges_borrowed_into(
817                &uploads,
818                &concrete_steps,
819                &read_handles,
820                &concrete_readbacks,
821                outputs,
822            )
823    }
824
825    fn dispatch_resident_repeated_sequence_read_ranges_into(
826        &self,
827        prefix_steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
828        repeated_steps: &[vyre_driver::backend::ResidentDispatchStep<'_>],
829        repeat_count: u32,
830        read_ranges: &[vyre_driver::backend::ResidentReadRange<'_>],
831        outputs: &mut [&mut Vec<u8>],
832    ) -> Result<(), BackendError> {
833        self.validate_resident_steps_for_dispatch(prefix_steps)?;
834        self.validate_resident_steps_for_dispatch(repeated_steps)?;
835        let repeat_count =
836            usize::try_from(repeat_count).map_err(|error| BackendError::InvalidProgram {
837                fix: format!(
838                    "Fix: CUDA repeated resident sequence count does not fit usize: {error}."
839                ),
840            })?;
841        if read_ranges.len() != outputs.len() {
842            return Err(BackendError::InvalidProgram {
843                fix: format!(
844                    "Fix: CUDA repeated resident sequence ranged readback expected matching range/output counts but got {} range(s) and {} output(s).",
845                    read_ranges.len(),
846                    outputs.len()
847                ),
848            });
849        }
850
851        let prefix_handle_sets = self
852            .resolve_step_handle_sets(prefix_steps, "CUDA repeated resident prefix handle sets")?;
853        let repeated_handle_sets =
854            self.resolve_repeated_step_handle_sets(repeated_steps, repeat_count)?;
855        let concrete_prefix = Self::concrete_resident_steps(
856            prefix_steps,
857            &prefix_handle_sets,
858            "CUDA repeated resident prefix steps",
859        )?;
860        let concrete_repeated = Self::concrete_resident_steps(
861            repeated_steps,
862            &repeated_handle_sets,
863            "CUDA repeated resident repeated steps",
864        )?;
865
866        let (read_handles, concrete_readbacks) = self.resolve_read_ranges(read_ranges)?;
867        let uploads: [(crate::backend::CudaResidentBuffer, &[u8]); 0] = [];
868        self.inner
869            .upload_resident_many_repeated_sequence_read_ranges_borrowed_into(
870                &uploads,
871                &concrete_prefix,
872                &concrete_repeated,
873                repeat_count,
874                &read_handles,
875                &concrete_readbacks,
876                outputs,
877            )
878    }
879
880    fn pipeline_cache_snapshot(&self) -> Option<vyre_driver::pipeline::PipelineCacheSnapshot> {
881        Some(self.inner.pipeline_cache_snapshot())
882    }
883
884    fn backend_metric_snapshot(&self) -> Vec<(&'static str, u64)> {
885        let source_cache = self.inner.ptx_source_cache_snapshot();
886        let mut metrics = Vec::new();
887        match u64::try_from(source_cache.entries) {
888            Ok(entries) => metrics.push(("cuda_ptx_source_cache_entries", entries)),
889            Err(source) => {
890                tracing::error!(
891                    "CUDA PTX source cache entry count cannot fit u64: {source}. Fix: shard backend metrics before source-cache cardinality exceeds u64."
892                );
893                metrics.push(("cuda_ptx_source_cache_entries_unrepresentable", 1));
894            }
895        }
896        metrics.push(("cuda_ptx_source_cache_hits", source_cache.hits));
897        metrics.push(("cuda_ptx_source_cache_misses", source_cache.misses));
898        let telemetry = self.inner.telemetry_snapshot();
899        metrics.push(("cuda_graph_launches", telemetry.cuda_graph_launches));
900        metrics.push((
901            "cuda_graph_materialized_cache_hits",
902            telemetry.cuda_graph_materialized_cache_hits,
903        ));
904        metrics.push((
905            "cuda_graph_batched_replay_chunks",
906            telemetry.cuda_graph_batched_replay_chunks,
907        ));
908        metrics.push((
909            "cuda_graph_batched_replay_lanes",
910            telemetry.cuda_graph_batched_replay_lanes,
911        ));
912        metrics.push(("cuda_host_to_device_bytes", telemetry.host_to_device_bytes));
913        metrics.push(("cuda_device_to_host_bytes", telemetry.device_to_host_bytes));
914        metrics.push(("cuda_readback_bytes", telemetry.readback_bytes));
915        metrics.push(("cuda_param_upload_bytes", telemetry.param_upload_bytes));
916        metrics.push(("cuda_kernel_launches", telemetry.kernel_launches));
917        metrics.push(("cuda_sync_points", telemetry.sync_points));
918        metrics.push((
919            "cuda_host_upload_operations",
920            telemetry.host_upload_operations,
921        ));
922        metrics.push((
923            "cuda_device_readback_operations",
924            telemetry.device_readback_operations,
925        ));
926        metrics.push((
927            "cuda_resident_borrowed_fallback_dispatches",
928            telemetry.resident_borrowed_fallback_dispatches,
929        ));
930        metrics.push(("cuda_launched_elements", telemetry.launched_elements));
931        metrics.push(("cuda_wasted_thread_slots", telemetry.wasted_thread_slots));
932        metrics.push((
933            "cuda_logical_thread_utilization_bps",
934            u64::from(telemetry.logical_thread_utilization_bps),
935        ));
936        metrics.push((
937            "cuda_logical_thread_waste_bps",
938            u64::from(telemetry.logical_thread_waste_bps),
939        ));
940        metrics.push((
941            "cuda_logical_elements_per_thread_slot_bps",
942            telemetry.logical_elements_per_thread_slot_bps,
943        ));
944        metrics.push(("cuda_timed_dispatches", telemetry.timed_dispatches));
945        metrics.push((
946            "cuda_timed_device_measurements",
947            telemetry.timed_device_measurements,
948        ));
949        metrics.push((
950            "cuda_timed_dispatches_missing_device_time",
951            telemetry.timed_dispatches_missing_device_time,
952        ));
953        metrics.push(("cuda_timed_wall_ns_total", telemetry.timed_wall_ns_total));
954        metrics.push((
955            "cuda_timed_device_ns_total",
956            telemetry.timed_device_ns_total,
957        ));
958        metrics.push(("cuda_timed_device_ns_max", telemetry.timed_device_ns_max));
959        metrics.push((
960            "cuda_timed_enqueue_ns_total",
961            telemetry.timed_enqueue_ns_total,
962        ));
963        metrics.push(("cuda_timed_wait_ns_total", telemetry.timed_wait_ns_total));
964        metrics
965    }
966
967    fn supports_subgroup_ops(&self) -> bool {
968        self.inner.hardware_supports_subgroup_ops()
969    }
970
971    fn supports_f16(&self) -> bool {
972        self.inner.hardware_supports_f16()
973    }
974
975    fn supports_bf16(&self) -> bool {
976        self.inner.hardware_supports_bf16()
977    }
978
979    fn supports_tensor_cores(&self) -> bool {
980        self.inner.hardware_supports_tensor_cores() && self.inner.lowers_tensor_core_ops()
981    }
982
983    fn supports_async_compute(&self) -> bool {
984        self.inner.hardware_supports_async_compute()
985    }
986
987    fn supports_grid_sync(&self) -> bool {
988        self.inner.supports_grid_sync()
989    }
990
991    fn cooperative_grid_sync_fits(
992        &self,
993        program: &Program,
994        inputs: &[&[u8]],
995        config: &DispatchConfig,
996    ) -> Result<bool, BackendError> {
997        self.inner
998            .cooperative_grid_sync_launch_fits(program, inputs, config)
999    }
1000
1001    fn allows_host_grid_sync_split(&self) -> bool {
1002        false
1003    }
1004
1005    fn supports_resident_dispatch(&self) -> bool {
1006        true
1007    }
1008
1009    fn supports_speculation(&self) -> bool {
1010        false
1011    }
1012
1013    fn max_workgroup_size(&self) -> [u32; 3] {
1014        self.inner.max_block_dim()
1015    }
1016
1017    fn max_compute_workgroups_per_dimension(&self) -> u32 {
1018        self.inner.max_grid_dim()[0]
1019    }
1020
1021    fn max_compute_invocations_per_workgroup(&self) -> u32 {
1022        self.inner.max_threads_per_block()
1023    }
1024
1025    fn subgroup_size(&self) -> Option<u32> {
1026        self.inner.warp_size()
1027    }
1028
1029    fn max_storage_buffer_bytes(&self) -> u64 {
1030        self.inner.device_memory_bytes()
1031    }
1032
1033    fn device_profile(&self) -> vyre_driver::DeviceProfile {
1034        let mut profile = self.inner.caps.to_device_profile();
1035        profile.supports_tensor_cores = self.supports_tensor_cores();
1036        profile.supports_indirect_dispatch = self.supports_indirect_dispatch();
1037        profile
1038    }
1039
1040    fn prepare(&self) -> Result<(), BackendError> {
1041        self.inner.warmup()
1042    }
1043
1044    fn shutdown(&self) -> Result<(), BackendError> {
1045        self.inner.cleanup()
1046    }
1047}
1048
1049/// Factory function for inventory registration.
1050pub fn cuda_factory() -> Result<Box<dyn VyreBackend>, BackendError> {
1051    let backend = CudaBackend::acquire().map_err(|e| BackendError::DispatchFailed {
1052        code: None,
1053        message: format!("CUDA backend acquisition failed: {e}"),
1054    })?;
1055    Ok(Box::new(CudaBackendRegistration { inner: backend }))
1056}
1057
1058/// Op-support set  -  CUDA supports every op the foundation IR defines
1059/// plus hardware intrinsics. Populated at runtime by the conform runner.
1060pub fn cuda_supported_ops() -> &'static std::collections::HashSet<vyre_foundation::ir::OpId> {
1061    vyre_driver::backend::validation::default_supported_ops_with_trap()
1062}
1063
1064fn cuda_semantic_operations() -> &'static std::collections::HashSet<vyre_foundation::ir::OpId> {
1065    vyre_driver::backend::dialect_only_supported_ops()
1066}
1067
1068inventory::submit! {
1069    BackendRegistration {
1070        id: CUDA_BACKEND_ID,
1071        target_id: CUDA_TARGET_ID,
1072        payload_format: Some(target_compiler::CUDA_TARGET_FORMAT),
1073        reference_oracle: false,
1074        factory: cuda_factory,
1075        supported_ops: cuda_supported_ops,
1076        semantic_operations: cuda_semantic_operations,
1077        target_compiler: Some(target_compiler::target_compiler_factory),
1078        materializer: Some(materializer::materializer_factory),
1079    }
1080}
1081
1082// rank 5 - CUDA is the canonical release dispatch backend when linked.
1083inventory::submit! {
1084    vyre_driver::backend::BackendPrecedence {
1085        id: CUDA_BACKEND_ID,
1086        rank: 5,
1087    }
1088}
1089
1090// CUDA owns a live dispatch stack, so conform can prove against it.
1091inventory::submit! {
1092    vyre_driver::backend::BackendCapability {
1093        id: CUDA_BACKEND_ID,
1094        dispatches: true,
1095    }
1096}
1097
1098inventory::submit! {
1099    vyre_driver::aot::AotLauncherEmitter {
1100        target: CUDA_TARGET_ID,
1101        emit: aot_launcher::emit_launcher,
1102    }
1103}
vyre_driver_cuda/lib.rs

vyre_driver_cuda/
lib.rs
