vyre_runtime/megakernel/
execution.rs

1//! Compiled persistent-megakernel handle and dispatch path.
2
3mod persistent_handles;
4mod readback_dispatch;
5mod types;
6
7use super::builder::{build_program_jit_slots, build_program_sharded_slots_shared};
8use super::handlers::OpcodeHandler;
9use super::io;
10use super::planner::MegakernelLaunchGeometry;
11use super::protocol;
12use super::protocol_api::{validate_control_bytes, validate_debug_log_bytes};
13use super::recovery::{
14    backend_error_indicates_device_loss, recover_compiled_pipeline, MegakernelRecoveryDecision,
15    MegakernelRecoveryPolicy,
16};
17use super::staging_reserve::reserve_vec_capacity;
18use crate::PipelineError;
19use arc_swap::ArcSwap;
20use std::sync::Arc;
21use std::time::Instant;
22use vyre_driver::backend::{
23    CompiledPipeline, DispatchConfig, OutputBuffers, Resource, VyreBackend,
24};
25use vyre_foundation::ir::Program;
26
27pub use types::{
28    MegakernelBatchDispatchOutput, MegakernelDispatchOutput, MegakernelDispatchStats,
29    MegakernelResidentBatchScratch, MegakernelResidentHandles,
30};
31
32/// Orchestrated persistent-megakernel handle.
33///
34/// Construct with [`Megakernel::bootstrap`] (default 256 lanes x 1
35/// workgroup) or [`Megakernel::bootstrap_sharded`] for multi-tenant fan-in.
36/// Feed bytecode with [`Megakernel::dispatch`].
37pub struct Megakernel {
38    backend: Arc<dyn VyreBackend>,
39    pipeline: ArcSwap<PipelineSlot>,
40    /// Pipeline identity string, updated atomically on every `recover_after_device_loss`.
41    pipeline_id: ArcSwap<String>,
42    program: Arc<Program>,
43    has_grid_sync: bool,
44    empty_io_queue_bytes: Arc<[u8]>,
45    slot_count: u32,
46    workgroup_size_x: u32,
47    recovery_policy: MegakernelRecoveryPolicy,
48}
49
50struct PipelineSlot {
51    inner: Arc<dyn CompiledPipeline>,
52}
53
54impl Megakernel {
55    /// Default bootstrap: 256 lanes x 1 workgroup, no custom opcodes.
56    ///
57    /// # Errors
58    ///
59    /// Returns [`PipelineError::Backend`] if the backend rejects the program.
60    pub fn bootstrap(backend: Arc<dyn VyreBackend>) -> Result<Self, PipelineError> {
61        Self::bootstrap_sharded(backend, 256, 256, Vec::new())
62    }
63
64    /// Bootstrap with custom opcodes but default sharding.
65    ///
66    /// # Errors
67    ///
68    /// See [`Megakernel::bootstrap`].
69    pub fn bootstrap_with_opcodes(
70        backend: Arc<dyn VyreBackend>,
71        opcodes: Vec<OpcodeHandler>,
72    ) -> Result<Self, PipelineError> {
73        Self::bootstrap_sharded(backend, 256, 256, opcodes)
74    }
75
76    /// Compute worker groups for a megakernel slot geometry without compiling.
77    ///
78    /// # Errors
79    ///
80    /// Returns [`PipelineError::QueueFull`] when the geometry cannot map slots
81    /// to whole workgroups.
82    pub fn worker_groups_for_geometry(
83        slot_count: u32,
84        workgroup_size_x: u32,
85    ) -> Result<u32, PipelineError> {
86        validate_bootstrap_geometry(slot_count, workgroup_size_x)?;
87        Ok(slot_count / workgroup_size_x)
88    }
89
90    /// Full bootstrap with sharding and custom opcodes.
91    ///
92    /// # Errors
93    ///
94    /// Returns [`PipelineError::QueueFull`] when geometry is invalid or
95    /// [`PipelineError::Backend`] from the underlying compile.
96    pub fn bootstrap_sharded(
97        backend: Arc<dyn VyreBackend>,
98        slot_count: u32,
99        workgroup_size_x: u32,
100        opcodes: Vec<OpcodeHandler>,
101    ) -> Result<Self, PipelineError> {
102        validate_bootstrap_geometry(slot_count, workgroup_size_x)?;
103        let program = build_program_sharded_slots_shared(workgroup_size_x, slot_count, &opcodes);
104        Self::compile_bootstrap_shared(backend, slot_count, workgroup_size_x, program)
105    }
106
107    /// JIT compiler bootstrap for fused payload processors.
108    ///
109    /// # Errors
110    ///
111    /// See [`Megakernel::bootstrap_sharded`].
112    pub fn bootstrap_jit(
113        backend: Arc<dyn VyreBackend>,
114        slot_count: u32,
115        workgroup_size_x: u32,
116        payload_processor: &[vyre_foundation::ir::Node],
117    ) -> Result<Self, PipelineError> {
118        validate_bootstrap_geometry(slot_count, workgroup_size_x)?;
119        let program = build_program_jit_slots(workgroup_size_x, slot_count, payload_processor);
120        Self::compile_bootstrap(backend, slot_count, workgroup_size_x, program)
121    }
122
123    fn compile_bootstrap(
124        backend: Arc<dyn VyreBackend>,
125        slot_count: u32,
126        workgroup_size_x: u32,
127        program: Program,
128    ) -> Result<Self, PipelineError> {
129        Self::compile_bootstrap_shared(backend, slot_count, workgroup_size_x, Arc::new(program))
130    }
131
132    fn compile_bootstrap_shared(
133        backend: Arc<dyn VyreBackend>,
134        slot_count: u32,
135        workgroup_size_x: u32,
136        program: Arc<Program>,
137    ) -> Result<Self, PipelineError> {
138        validate_bootstrap_geometry(slot_count, workgroup_size_x)?;
139        let config = DispatchConfig::default();
140        let pipeline = vyre_driver::pipeline::compile_shared(
141            Arc::clone(&backend),
142            Arc::clone(&program),
143            &config,
144        )?;
145        let pipeline_id = pipeline.id().to_string();
146        let has_grid_sync = vyre_driver::grid_sync::contains_grid_sync(&program);
147        let empty_io_queue_bytes =
148            Arc::<[u8]>::from(io::try_encode_empty_io_queue(io::IO_SLOT_COUNT)?.into_boxed_slice());
149        Ok(Self {
150            backend,
151            pipeline: ArcSwap::from(Arc::new(PipelineSlot { inner: pipeline })),
152            pipeline_id: ArcSwap::from(Arc::new(pipeline_id)),
153            program,
154            has_grid_sync,
155            empty_io_queue_bytes,
156            slot_count,
157            workgroup_size_x,
158            recovery_policy: MegakernelRecoveryPolicy::default(),
159        })
160    }
161
162    /// Dispatch a full storage buffer set with an empty IO queue.
163    ///
164    /// # Errors
165    ///
166    /// Returns [`PipelineError`] when protocol buffers are malformed, dispatch
167    /// fails, or device-loss recovery cannot rebuild the compiled pipeline.
168    pub fn dispatch(
169        &self,
170        control_bytes: Vec<u8>,
171        ring_bytes: Vec<u8>,
172        debug_log_bytes: Vec<u8>,
173    ) -> Result<Vec<Vec<u8>>, PipelineError> {
174        self.dispatch_borrowed(&control_bytes, &ring_bytes, &debug_log_bytes)
175    }
176
177    /// Dispatch a borrowed storage buffer set with an empty IO queue.
178    ///
179    /// # Errors
180    ///
181    /// Returns [`PipelineError`] when protocol buffers are malformed, dispatch
182    /// fails, or device-loss recovery cannot rebuild the compiled pipeline.
183    pub fn dispatch_borrowed(
184        &self,
185        control_bytes: &[u8],
186        ring_bytes: &[u8],
187        debug_log_bytes: &[u8],
188    ) -> Result<Vec<Vec<u8>>, PipelineError> {
189        Ok(self
190            .dispatch_borrowed_observed(control_bytes, ring_bytes, debug_log_bytes)?
191            .buffers)
192    }
193
194    /// Dispatch a full storage buffer set and return runtime instrumentation.
195    ///
196    /// # Errors
197    ///
198    /// See [`Megakernel::dispatch`].
199    pub fn dispatch_observed(
200        &self,
201        control_bytes: Vec<u8>,
202        ring_bytes: Vec<u8>,
203        debug_log_bytes: Vec<u8>,
204    ) -> Result<MegakernelDispatchOutput, PipelineError> {
205        self.dispatch_with_io_queue_borrowed_observed(
206            &control_bytes,
207            &ring_bytes,
208            &debug_log_bytes,
209            &self.empty_io_queue_bytes,
210        )
211    }
212
213    /// Dispatch borrowed buffers with an empty IO queue and return runtime
214    /// instrumentation.
215    ///
216    /// # Errors
217    ///
218    /// See [`Megakernel::dispatch_borrowed`].
219    pub fn dispatch_borrowed_observed(
220        &self,
221        control_bytes: &[u8],
222        ring_bytes: &[u8],
223        debug_log_bytes: &[u8],
224    ) -> Result<MegakernelDispatchOutput, PipelineError> {
225        self.dispatch_with_io_queue_borrowed_observed(
226            control_bytes,
227            ring_bytes,
228            debug_log_bytes,
229            &self.empty_io_queue_bytes,
230        )
231    }
232
233    /// Dispatch a full storage buffer set with a caller-supplied `io_queue`.
234    ///
235    /// # Errors
236    ///
237    /// Returns [`PipelineError`] when any protocol buffer is malformed, backend
238    /// dispatch fails, or device-loss recovery cannot rebuild the pipeline.
239    pub fn dispatch_with_io_queue(
240        &self,
241        control_bytes: Vec<u8>,
242        ring_bytes: Vec<u8>,
243        debug_log_bytes: Vec<u8>,
244        io_queue_bytes: Vec<u8>,
245    ) -> Result<Vec<Vec<u8>>, PipelineError> {
246        self.dispatch_with_io_queue_borrowed(
247            &control_bytes,
248            &ring_bytes,
249            &debug_log_bytes,
250            &io_queue_bytes,
251        )
252    }
253
254    /// Dispatch borrowed buffers with a caller-supplied `io_queue`.
255    ///
256    /// # Errors
257    ///
258    /// See [`Megakernel::dispatch_with_io_queue`].
259    pub fn dispatch_with_io_queue_borrowed(
260        &self,
261        control_bytes: &[u8],
262        ring_bytes: &[u8],
263        debug_log_bytes: &[u8],
264        io_queue_bytes: &[u8],
265    ) -> Result<Vec<Vec<u8>>, PipelineError> {
266        Ok(self
267            .dispatch_with_io_queue_borrowed_observed(
268                control_bytes,
269                ring_bytes,
270                debug_log_bytes,
271                io_queue_bytes,
272            )?
273            .buffers)
274    }
275
276    /// Dispatch with a caller-supplied `io_queue` and return instrumentation.
277    ///
278    /// # Errors
279    ///
280    /// See [`Megakernel::dispatch_with_io_queue`].
281    pub fn dispatch_with_io_queue_observed(
282        &self,
283        control_bytes: Vec<u8>,
284        ring_bytes: Vec<u8>,
285        debug_log_bytes: Vec<u8>,
286        io_queue_bytes: Vec<u8>,
287    ) -> Result<MegakernelDispatchOutput, PipelineError> {
288        self.dispatch_with_io_queue_borrowed_observed(
289            &control_bytes,
290            &ring_bytes,
291            &debug_log_bytes,
292            &io_queue_bytes,
293        )
294    }
295
296    /// Dispatch borrowed buffers with a caller-supplied `io_queue` and return
297    /// instrumentation.
298    ///
299    /// # Errors
300    ///
301    /// See [`Megakernel::dispatch_with_io_queue`].
302    pub fn dispatch_with_io_queue_borrowed_observed(
303        &self,
304        control_bytes: &[u8],
305        ring_bytes: &[u8],
306        debug_log_bytes: &[u8],
307        io_queue_bytes: &[u8],
308    ) -> Result<MegakernelDispatchOutput, PipelineError> {
309        let mut buffers = Vec::new();
310        reserve_output_shell(
311            &mut buffers,
312            MegakernelResidentHandles::ABI_RESOURCE_COUNT,
313            "borrowed megakernel output shell",
314        )?;
315        let stats = self.dispatch_with_io_queue_borrowed_into(
316            control_bytes,
317            ring_bytes,
318            debug_log_bytes,
319            io_queue_bytes,
320            &mut buffers,
321        )?;
322        Ok(MegakernelDispatchOutput { buffers, stats })
323    }
324
325    /// Dispatch borrowed buffers with a caller-supplied IO queue, writing
326    /// backend outputs into caller-owned storage.
327    ///
328    /// # Errors
329    ///
330    /// See [`Megakernel::dispatch_with_io_queue_borrowed`].
331    pub fn dispatch_with_io_queue_borrowed_into(
332        &self,
333        control_bytes: &[u8],
334        ring_bytes: &[u8],
335        debug_log_bytes: &[u8],
336        io_queue_bytes: &[u8],
337        outputs: &mut OutputBuffers,
338    ) -> Result<MegakernelDispatchStats, PipelineError> {
339        validate_control_bytes(control_bytes)?;
340        validate_debug_log_bytes(debug_log_bytes)?;
341        io::validate_io_queue_bytes(io_queue_bytes)?;
342        self.validate_ring_bytes(ring_bytes)?;
343
344        let input_bytes = total_len([control_bytes, ring_bytes, debug_log_bytes, io_queue_bytes])?;
345        let inputs = [control_bytes, ring_bytes, debug_log_bytes, io_queue_bytes];
346        let config = self.launch_geometry().dispatch_config(None);
347        let started = Instant::now();
348        let mut recovered = false;
349        match self.dispatch_once_into(&inputs, &config, outputs) {
350            Ok(()) => {}
351            Err(error) if self.recovery_policy.allows_retry(&error) => {
352                self.recover_after_device_loss()?;
353                recovered = true;
354                self.dispatch_once_into(&inputs, &config, outputs)?
355            }
356            Err(error) => return Err(error.into()),
357        }
358        let latency_ns = nanos_u64(started.elapsed().as_nanos())?;
359        let output_bytes = output_bytes(outputs)?;
360        let readback_bytes = output_bytes;
361        let bytes_moved = checked_add_u64(input_bytes, readback_bytes, "megakernel bytes moved")?;
362        let device_allocation_bytes = checked_add_u64(
363            input_bytes,
364            output_bytes,
365            "megakernel host-visible device allocation bytes",
366        )?;
367        let output_buffers = count_u32(outputs.len(), "megakernel output buffer count")?;
368        let device_allocation_events =
369            checked_add_u32(4, output_buffers, "megakernel allocation event count")?;
370        Ok(MegakernelDispatchStats {
371            input_bytes,
372            output_bytes,
373            readback_bytes,
374            bytes_moved,
375            device_allocation_bytes,
376            device_allocation_events,
377            latency_ns,
378            output_buffers,
379            resident_resource_rows: 0,
380            resident_resource_handles: 0,
381            kernel_launches: if recovered { 2 } else { 1 },
382            sync_points: 1,
383            recovered_after_device_loss: recovered,
384        })
385    }
386
387    /// Rebuild the compiled pipeline after device-loss symptoms.
388    ///
389    /// This does not mask the failure: if recompilation fails, the structured
390    /// backend error is returned with the original remediation text intact.
391    ///
392    /// # Errors
393    ///
394    /// Returns [`PipelineError::Backend`] when the backend cannot recompile.
395    pub fn recover_after_device_loss(&self) -> Result<MegakernelRecoveryDecision, PipelineError> {
396        let config = self.launch_geometry().dispatch_config(None);
397        let rebuilt = recover_compiled_pipeline(&self.backend, Arc::clone(&self.program), &config)?;
398        let new_id = rebuilt.id().to_string();
399        self.pipeline
400            .store(Arc::new(PipelineSlot { inner: rebuilt }));
401        self.pipeline_id.store(Arc::new(new_id));
402        Ok(MegakernelRecoveryDecision::RecompiledPipeline)
403    }
404
405    /// Pipeline id from the backend.
406    ///
407    /// Returns the id of the currently loaded compiled pipeline. After a call
408    /// to [`Megakernel::recover_after_device_loss`] this reflects the rebuilt
409    /// pipeline's id rather than the original one.
410    #[must_use]
411    pub fn pipeline_id(&self) -> String {
412        (**self.pipeline_id.load()).clone()
413    }
414
415    /// Slot count this kernel was sharded for.
416    #[must_use]
417    pub const fn slot_count(&self) -> u32 {
418        self.slot_count
419    }
420
421    /// Workgroup size this kernel was compiled for.
422    #[must_use]
423    pub const fn workgroup_size_x(&self) -> u32 {
424        self.workgroup_size_x
425    }
426
427    /// Workgroup count needed to cover every ring slot.
428    #[must_use]
429    pub fn worker_groups(&self) -> u32 {
430        self.slot_count / self.workgroup_size_x
431    }
432
433    pub(super) fn validate_ring_bytes(&self, ring_bytes: &[u8]) -> Result<(), PipelineError> {
434        let expected_ring_bytes = protocol::ring_byte_len(self.slot_count).ok_or_else(|| {
435            PipelineError::Backend(
436                "megakernel ring byte length overflowed usize. Fix: split the ring into smaller dispatch shards."
437                    .to_string(),
438            )
439        })?;
440        if ring_bytes.len() != expected_ring_bytes {
441            return Err(PipelineError::Backend(format!(
442                "megakernel ring buffer has {} bytes, expected {expected_ring_bytes} for {} slots. Fix: build ring bytes with Megakernel::encode_empty_ring(slot_count) for this handle.",
443                ring_bytes.len(),
444                self.slot_count
445            )));
446        }
447        Ok(())
448    }
449
450    pub(super) fn launch_geometry(&self) -> MegakernelLaunchGeometry {
451        MegakernelLaunchGeometry {
452            workgroup_size_x: self.workgroup_size_x,
453            slot_count: self.slot_count,
454            dispatch_grid: [self.slot_count / self.workgroup_size_x, 1, 1],
455        }
456    }
457
458    fn dispatch_once_into(
459        &self,
460        inputs: &[&[u8]; 4],
461        config: &DispatchConfig,
462        outputs: &mut OutputBuffers,
463    ) -> Result<(), vyre_driver::BackendError> {
464        if self.has_grid_sync && !self.backend.supports_grid_sync() {
465            return vyre_driver::grid_sync::dispatch_with_grid_sync_split_into(
466                self.backend.as_ref(),
467                &self.program,
468                inputs,
469                config,
470                outputs,
471            );
472        }
473        let pipeline = self.pipeline.load();
474        pipeline
475            .inner
476            .dispatch_borrowed_into(inputs, config, outputs)
477    }
478
479    fn dispatch_persistent_handles_once_into(
480        &self,
481        inputs: &[Resource; 4],
482        config: &DispatchConfig,
483        outputs: &mut OutputBuffers,
484    ) -> Result<(), vyre_driver::BackendError> {
485        let pipeline = self.pipeline.load();
486        pipeline
487            .inner
488            .dispatch_persistent_handles_into(inputs, config, outputs)
489    }
490
491    fn dispatch_persistent_handle_rows_once_into(
492        &self,
493        rows: &[[Resource; 4]],
494        config: &DispatchConfig,
495        outputs: &mut Vec<OutputBuffers>,
496    ) -> Result<(), vyre_driver::BackendError> {
497        let pipeline = self.pipeline.load();
498        pipeline
499            .inner
500            .dispatch_persistent_handle_rows_into(rows, config, outputs)
501    }
502}
503
504impl MegakernelRecoveryPolicy {
505    fn allows_retry(self, error: &vyre_driver::BackendError) -> bool {
506        self.retry_device_loss_once && backend_error_indicates_device_loss(error)
507    }
508}
509
510fn validate_bootstrap_geometry(
511    slot_count: u32,
512    workgroup_size_x: u32,
513) -> Result<(), PipelineError> {
514    if slot_count == 0 || workgroup_size_x == 0 || slot_count % workgroup_size_x != 0 {
515        return Err(PipelineError::QueueFull {
516            queue: "submission",
517            fix: "slot_count must be a non-zero multiple of workgroup_size_x",
518        });
519    }
520    Ok(())
521}
522
523pub(super) fn total_len<const N: usize>(buffers: [&[u8]; N]) -> Result<u64, PipelineError> {
524    let mut total = 0u64;
525    for buffer in buffers {
526        total = checked_add_u64(
527            total,
528            usize_to_u64(buffer.len(), "megakernel input buffer length")?,
529            "megakernel input byte total",
530        )?;
531    }
532    Ok(total)
533}
534
535pub(super) fn output_bytes(outputs: &[Vec<u8>]) -> Result<u64, PipelineError> {
536    let mut total = 0u64;
537    for output in outputs {
538        total = checked_add_u64(
539            total,
540            usize_to_u64(output.len(), "megakernel output buffer length")?,
541            "megakernel output byte total",
542        )?;
543    }
544    Ok(total)
545}
546
547pub(super) fn nested_output_bytes(outputs: &[Vec<Vec<u8>>]) -> Result<u64, PipelineError> {
548    let mut total = 0u64;
549    for row in outputs {
550        total = checked_add_u64(
551            total,
552            output_bytes(row)?,
553            "megakernel nested output byte total",
554        )?;
555    }
556    Ok(total)
557}
558
559pub(super) fn output_count_u32(outputs: &[Vec<u8>]) -> Result<u32, PipelineError> {
560    count_u32(outputs.len(), "megakernel output buffer count")
561}
562
563pub(super) fn nested_output_count_u32(outputs: &[Vec<Vec<u8>>]) -> Result<u32, PipelineError> {
564    let mut total = 0usize;
565    for row in outputs {
566        total = total.checked_add(row.len()).ok_or_else(|| {
567            PipelineError::Backend(
568                "megakernel nested output buffer count overflowed usize. Fix: split resident rows before dispatch.".to_string(),
569            )
570        })?;
571    }
572    count_u32(total, "megakernel nested output buffer count")
573}
574
575pub(super) fn resident_row_count_u32(rows: usize) -> Result<u32, PipelineError> {
576    count_u32(rows, "megakernel resident resource row count")
577}
578
579pub(super) fn resident_handle_count_u32(rows: usize) -> Result<u32, PipelineError> {
580    let handles = rows
581        .checked_mul(MegakernelResidentHandles::ABI_RESOURCE_COUNT)
582        .ok_or_else(|| {
583            PipelineError::Backend(
584                "megakernel resident resource handle count overflowed usize. Fix: split resident rows before dispatch.".to_string(),
585            )
586        })?;
587    count_u32(handles, "megakernel resident resource handle count")
588}
589
590pub(super) fn reserve_output_shell<T>(
591    out: &mut Vec<T>,
592    capacity: usize,
593    label: &'static str,
594) -> Result<(), PipelineError> {
595    reserve_vec_capacity(out, capacity, label)
596}
597
598pub(super) fn nanos_u64(nanos: u128) -> Result<u64, PipelineError> {
599    u64::try_from(nanos).map_err(|source| {
600        PipelineError::Backend(format!(
601            "megakernel latency cannot fit u64 nanoseconds: {source}. Fix: timeout or split the dispatch before telemetry overflows."
602        ))
603    })
604}
605
606fn usize_to_u64(value: usize, label: &str) -> Result<u64, PipelineError> {
607    u64::try_from(value).map_err(|source| {
608        PipelineError::Backend(format!(
609            "{label} cannot fit u64: {source}. Fix: split the megakernel dispatch before telemetry/accounting."
610        ))
611    })
612}
613
614fn count_u32(value: usize, label: &str) -> Result<u32, PipelineError> {
615    u32::try_from(value).map_err(|source| {
616        PipelineError::Backend(format!(
617            "{label} cannot fit u32: {source}. Fix: split the megakernel dispatch before telemetry/accounting."
618        ))
619    })
620}
621
622fn checked_add_u64(left: u64, right: u64, label: &str) -> Result<u64, PipelineError> {
623    left.checked_add(right).ok_or_else(|| {
624        PipelineError::Backend(format!(
625            "{label} overflowed u64. Fix: split the megakernel dispatch before telemetry/accounting."
626        ))
627    })
628}
629
630fn checked_add_u32(left: u32, right: u32, label: &str) -> Result<u32, PipelineError> {
631    left.checked_add(right).ok_or_else(|| {
632        PipelineError::Backend(format!(
633            "{label} overflowed u32. Fix: split the megakernel dispatch before telemetry/accounting."
634        ))
635    })
636}
637
638#[cfg(test)]
639mod tests;
vyre_runtime/megakernel/execution.rs

vyre_runtime/megakernel/
execution.rs
