Skip to main content

tenferro_runtime/runtime/
execution.rs

1//! Runtime-owned compiled-graph execution boundary.
2//!
3//! [`crate::Runtime::prepare`] turns a compiled graph into a [`Prepared`]
4//! program for a set of engine targets, and [`crate::Runtime::run`] runs it on
5//! the engines the caller passes in. The runtime holds no engine and no
6//! execution state between calls.
7
8use num_complex::{Complex32, Complex64};
9use std::collections::{HashMap, HashSet};
10use std::error::Error as StdError;
11use std::fmt;
12use std::panic::AssertUnwindSafe;
13use std::sync::Arc;
14
15use smallvec::SmallVec;
16use tenferro_ops::shape_extent::ShapeExtent;
17use tenferro_tensor::{DType, Engine, Tensor, TensorRead};
18use tinyvec::TinyVec;
19
20use crate::error::ErrorPhase;
21use crate::exec::{ExecInstruction, ExecProgram, ExecSlot, ExtensionExecutionDispatch};
22use crate::extension_cache::ExtensionCacheStore;
23use crate::graph::CompiledGraph;
24use crate::runtime::schedule::{
25    EventDependency, ExecutionLocation, ScheduledGraph, ScheduledNode, ScheduledNodeKind,
26    ScheduledTransfer, UnsupportedScheduledNodeError,
27};
28use crate::runtime::{
29    EngineId, EngineTarget, EventDomainDriver, EventDomainError, EventDomainOperation,
30    EventDomainRun, EventToken, ExecutionEngine, InputSignature, Output, PrepareError,
31    PrepareOptions, PreparedOperationPlan, ProgramPlacementConstraint, Runtime, TransferError,
32    TransferProviderContractError, TransferRequest,
33};
34use crate::{Error, Result};
35
36// `TensorRead` has no `Default`, so `tinyvec::Array` cannot hold it inline.
37// Keep the short input-read list in a `SmallVec`, which does not require
38// `Default` and stays inline for the common case.
39type RuntimeInputReads<'a> = SmallVec<[TensorRead<'a>; 8]>;
40type RuntimeInputShapes<'a> = TinyVec<[&'a [usize]; 8]>;
41type RuntimeShapeScratch = TinyVec<[usize; 8]>;
42
43/// A compiled graph prepared for a set of engine targets.
44///
45/// A prepared program holds the execution plan and the requirements each
46/// target slot puts on its engine; it holds no engine, stream or event state,
47/// so any engines that report the same targets can run it, and several callers
48/// can run it at the same time on their own engines. It is tied to the runtime
49/// epoch that created it and becomes stale after runtime reconfiguration.
50///
51/// # Examples
52///
53/// ```
54/// use tenferro_runtime::Prepared;
55///
56/// fn target_count(prepared: &Prepared) -> usize {
57///     prepared.targets().len()
58/// }
59/// ```
60#[derive(Clone)]
61pub struct Prepared {
62    runtime_id: super::RuntimeId,
63    epoch: super::RuntimeEpoch,
64    program: CompiledGraph,
65    prepared: Arc<super::preparation::PreparedProgram>,
66    /// Whether this handle's preparation ran the metadata check on the graph's
67    /// *retained* inputs, i.e. on the values an empty-input run would reuse.
68    ///
69    /// Preparation that replaced every input validated the replacements only,
70    /// so its prepared handle must re-check the defaults at execution time.
71    defaults_validated: bool,
72    /// The engine targets, in slot order, that a run must bind.
73    targets: Arc<[EngineTarget]>,
74    /// The input signature this handle was prepared for.
75    ///
76    /// A run whose inputs produce the same signature is exactly the run the
77    /// preparation validated: the projection, the arity, and the shape guards
78    /// all passed on this signature then, so re-deriving and re-comparing the
79    /// specialization per run is redundant (#2186).
80    signature: Arc<super::InputSignature>,
81}
82
83impl fmt::Debug for Prepared {
84    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
85        formatter
86            .debug_struct("Prepared")
87            .field("runtime_id", &self.runtime_id)
88            .field("epoch", &self.epoch)
89            .field("targets", &self.targets)
90            .finish_non_exhaustive()
91    }
92}
93
94impl Prepared {
95    /// The engine targets this program was prepared for, in slot order.
96    ///
97    /// [`crate::Runtime::run`] binds slot `i` to the `i`-th engine it receives.
98    ///
99    /// # Examples
100    ///
101    /// ```
102    /// use tenferro_runtime::Prepared;
103    ///
104    /// fn first_engine(prepared: &Prepared) -> Option<&str> {
105    ///     prepared.targets().first().map(|target| target.engine_id().as_str())
106    /// }
107    /// ```
108    #[must_use]
109    pub fn targets(&self) -> &[EngineTarget] {
110        &self.targets
111    }
112
113    /// How executions must treat this handle's retained inputs.
114    ///
115    /// Skipping the check is sound only when preparation itself ran it on those
116    /// values; preparation with a complete replacement input list did not.
117    fn bound_input_check(&self) -> BoundInputCheck {
118        if self.defaults_validated {
119            BoundInputCheck::AlreadyValidatedAtPrepare
120        } else {
121            BoundInputCheck::Validate
122        }
123    }
124
125    /// Number of planned elementwise regions and the instructions they cover,
126    /// for execution-path parity tests.
127    #[doc(hidden)]
128    #[must_use]
129    pub fn elementwise_region_summary(&self) -> (usize, usize) {
130        super::region::region_summary(self.prepared.root().regions())
131    }
132
133    /// Number of planned elementwise regions executed as one fused command, and
134    /// the number that fell back to per-instruction dispatch.
135    #[doc(hidden)]
136    #[must_use]
137    pub fn elementwise_region_execution_counts(&self) -> (usize, usize) {
138        let counters = self.prepared.root().region_counters();
139        (counters.fused(), counters.fallbacks())
140    }
141
142    /// Number of successful scheduled event-domain submissions observed by this
143    /// prepared program's production executor.
144    ///
145    /// This is runtime submission evidence, rather than a plan-derived census.
146    #[doc(hidden)]
147    #[must_use]
148    pub fn execution_submission_count(&self) -> usize {
149        self.prepared.root().region_counters().submissions()
150    }
151
152    /// Plan-derived command count for the production execution path, for
153    /// execution-path tests: scheduled operations that no region covers, plus
154    /// one command per planned region.
155    ///
156    /// This is a plan census, not a runtime submission count.
157    #[doc(hidden)]
158    #[must_use]
159    pub fn execution_command_count(&self) -> usize {
160        let root = self.prepared.root();
161        super::region::command_count(root.schedule(), root.regions())
162    }
163}
164
165/// Run an extension callback, marking the engine broken if it unwinds.
166fn run_extension_callback<R>(
167    engine: &mut dyn ExecutionEngine,
168    callback: impl FnOnce(&mut dyn ExecutionEngine) -> Result<R>,
169) -> Result<R> {
170    match std::panic::catch_unwind(AssertUnwindSafe(|| callback(engine))) {
171        Ok(result) => result,
172        Err(payload) => {
173            engine.mark_broken();
174            std::panic::resume_unwind(payload)
175        }
176    }
177}
178
179/// Run one scheduled instruction on `exec`.
180///
181/// An extension instruction is the only one that borrows the engine's
182/// extension caches; it takes the owner fallback when its executor cannot run
183/// in the engine's session. Every other instruction executes through the
184/// engine's typed operations and never touches the caches, so a program
185/// without extension commands lends none (#2226).
186fn execute_slot_instruction_on_engine<'input>(
187    engine: &mut dyn ExecutionEngine,
188    extension_caches: Option<&mut ExtensionCacheStore>,
189    operations: &[PreparedOperationPlan],
190    instruction_index: usize,
191    instruction: &ExecInstruction,
192    slots: &mut [Option<ExecSlot<'input>>],
193) -> Result<()> {
194    if matches!(&instruction.op, crate::exec::ExecOp::Extension(_)) {
195        let caches = extension_caches.ok_or_else(|| {
196            Error::runtime_state(
197                "Runtime::run",
198                ErrorPhase::Execution,
199                "engine extension caches were not taken before an extension command",
200            )
201        })?;
202        let mut extension_dispatch = ExtensionExecutionDispatch { operations, caches };
203        if crate::exec::needs_owner_extension_fallback(instruction, Some(&extension_dispatch)) {
204            // A prepared operation that executes against the concrete engine
205            // type instead of the engine operations receives the engine itself.
206            let identity = engine.target().context_identity();
207            run_extension_callback(engine, |engine| {
208                crate::exec::execute_owner_extension_fallback(
209                    engine.as_any_mut(),
210                    identity,
211                    slots,
212                    instruction,
213                    Some(&mut extension_dispatch),
214                )
215            })?;
216        } else {
217            run_extension_callback(engine, |engine| {
218                crate::exec::execute_ffi_instruction_exec(
219                    engine,
220                    slots,
221                    instruction,
222                    Some(instruction_index),
223                    Some(&mut extension_dispatch),
224                )
225            })?;
226        }
227        crate::exec::reclaim_last_use_inputs_exec(slots, instruction, engine);
228        return Ok(());
229    }
230    if crate::exec::is_host_instruction(instruction) {
231        crate::exec::execute_host_instruction_exec(engine, slots, instruction)?;
232    } else if crate::exec::is_ffi_instruction(instruction) {
233        crate::exec::execute_ffi_instruction_exec(
234            engine,
235            slots,
236            instruction,
237            Some(instruction_index),
238            None,
239        )?;
240    } else if let Some(view) = crate::exec::execute_backend_view_op(slots, instruction) {
241        slots[instruction.output_slots[0]] = Some(view);
242    } else {
243        let result = crate::exec::execute_backend_op(engine, slots, instruction)?;
244        slots[instruction.output_slots[0]] = Some(ExecSlot::Owned(result));
245    }
246    crate::exec::reclaim_last_use_inputs_exec(slots, instruction, engine);
247    Ok(())
248}
249
250/// Run one planned elementwise region on `exec`.
251fn execute_elementwise_fusion_slots_on_engine<'input>(
252    exec: &mut dyn Engine,
253    input_slots: &[usize],
254    instruction_count: usize,
255    plan: &tenferro_tensor::backend::ElementwiseFusionPlan,
256    slots: &mut [Option<ExecSlot<'input>>],
257    output_slots: &[usize],
258) -> Result<bool> {
259    // The fused entry point takes owned tensors. A borrowed view input is
260    // materialized once when the region is long enough to amortize that
261    // copy; otherwise the per-instruction path, which resolves views itself,
262    // stays in charge.
263    let mut views: Vec<usize> = Vec::new();
264    for &slot in input_slots {
265        let value = slots
266            .get(slot)
267            .and_then(Option::as_ref)
268            .ok_or_else(|| crate::Error::from(tenferro_tensor::Error::MissingValue { slot }))?;
269        if value.as_tensor("elementwise_region").is_err() {
270            views.push(slot);
271        }
272    }
273    if !views.is_empty() {
274        if instruction_count < super::region::VIEW_COPY_MIN_INSTRUCTIONS {
275            return Ok(false);
276        }
277        for &slot in &views {
278            let tensor = {
279                let read = slots[slot].as_ref().map(ExecSlot::as_read).ok_or_else(|| {
280                    crate::Error::from(tenferro_tensor::Error::MissingValue { slot })
281                })?;
282                exec.to_contiguous_read(read)?
283            };
284            slots[slot] = Some(ExecSlot::Owned(tensor));
285        }
286    }
287
288    let mut inputs: Vec<&Tensor> = Vec::with_capacity(input_slots.len());
289    for &slot in input_slots {
290        let value = slots
291            .get(slot)
292            .and_then(Option::as_ref)
293            .ok_or_else(|| crate::Error::from(tenferro_tensor::Error::MissingValue { slot }))?;
294        inputs.push(value.as_tensor("elementwise_region")?);
295    }
296
297    if !crate::segment::plan_matches_input_dtypes(plan, &inputs) {
298        return Ok(false);
299    }
300
301    let outputs = exec.execute_elementwise_fusion(&inputs, plan)?;
302    let Some(outputs) = outputs else {
303        return Ok(false);
304    };
305    if outputs.len() != output_slots.len() {
306        return Err(crate::Error::Internal(format!(
307            "fused elementwise region produced {} outputs for {} slots",
308            outputs.len(),
309            output_slots.len()
310        )));
311    }
312    for (&slot, tensor) in output_slots.iter().zip(outputs) {
313        slots[slot] = Some(ExecSlot::Owned(tensor));
314    }
315    Ok(true)
316}
317
318/// Prepare `program` for the engines described by `targets`.
319pub(super) fn prepare(
320    runtime: &Runtime,
321    program: &CompiledGraph,
322    inputs: &[TensorRead<'_>],
323    targets: &[EngineTarget],
324) -> Result<Prepared> {
325    const CALLER: &str = "Runtime::prepare";
326    if targets.is_empty() {
327        return Err(Error::invalid_argument(
328            CALLER,
329            ErrorPhase::Execution,
330            "targets",
331            "prepare needs at least one engine target",
332        ));
333    }
334    let snapshot = runtime
335        .snapshot()
336        .map_err(|source| Error::runtime_state_source(CALLER, ErrorPhase::Execution, source))?;
337    let mut engine_ids = Vec::with_capacity(targets.len());
338    for (index, target) in targets.iter().enumerate() {
339        if targets[..index]
340            .iter()
341            .any(|earlier| earlier.engine_id() == target.engine_id())
342        {
343            return Err(Error::invalid_argument(
344                CALLER,
345                ErrorPhase::Execution,
346                "targets",
347                format!(
348                    "engine {:?} is named by more than one target",
349                    target.engine_id().as_str()
350                ),
351            ));
352        }
353        let view = snapshot.engine(target.engine_id()).ok_or_else(|| {
354            Error::runtime_state(
355                CALLER,
356                ErrorPhase::Execution,
357                format!(
358                    "no engine {:?} is registered with this runtime",
359                    target.engine_id().as_str()
360                ),
361            )
362        })?;
363        if view.context_identity() != target.context_identity() {
364            return Err(Error::runtime_state(
365                CALLER,
366                ErrorPhase::Execution,
367                format!(
368                    "engine {:?} is registered for context {:?}, but the target reports {:?}",
369                    target.engine_id().as_str(),
370                    view.context_identity(),
371                    target.context_identity()
372                ),
373            ));
374        }
375        engine_ids.push(target.engine_id().clone());
376    }
377    let defaults_validated = retained_inputs_validated(program, inputs.len());
378    let inputs = resolve_input_reads(program, inputs, BoundInputCheck::Validate)?;
379    let signature = input_signature_reads(&inputs)?;
380    // An extension whose output extent was symbolic at compile time is re-inferred
381    // against the concrete signature here, so its own typed validation surfaces at
382    // preparation instead of at execution (#2434).
383    super::preparation::validate_resolved_extension_extents(program.frozen(), &signature)?;
384    let placement = ProgramPlacementConstraint::new(engine_ids, None)
385        .map_err(|source| Error::runtime_state_source(CALLER, ErrorPhase::Execution, source))?;
386    let options = PrepareOptions::new().with_placement(placement);
387    let prepared = runtime
388        .prepare_compiled_for(program, &signature, &options)
389        .map_err(prepare_error)?;
390    validate_prepared_epoch(runtime, prepared.root().epoch(), CALLER)?;
391    Ok(Prepared {
392        runtime_id: runtime.id(),
393        epoch: prepared.root().epoch(),
394        program: program.clone(),
395        prepared,
396        defaults_validated,
397        targets: targets.into(),
398        signature: Arc::new(signature),
399    })
400}
401
402/// Whether resolving `input_count` inputs makes [`resolve_input_reads`] check
403/// the graph's retained inputs.
404///
405/// Empty inputs resolve to the retained values, and a partial list still
406/// resolves every bound input, so both validate them. A list that covers every
407/// input replaces all of them, which leaves the retained values unchecked.
408fn retained_inputs_validated(program: &CompiledGraph, input_count: usize) -> bool {
409    let expected = program.program().inputs().len();
410    input_count != expected || expected == 0
411}
412
413/// Run `prepared` on the engines the caller passes in.
414pub(super) fn run<'a>(
415    runtime: &Runtime,
416    engines: &mut [&mut dyn ExecutionEngine],
417    prepared: &'a Prepared,
418    inputs: &[TensorRead<'a>],
419) -> Result<Vec<Output<'a>>> {
420    const CALLER: &str = "Runtime::run";
421    validate_prepared_runtime(runtime, prepared, CALLER)?;
422    let mut engines = EngineSet::bind(engines, &prepared.targets, CALLER)?;
423    let inputs = resolve_input_reads(&prepared.program, inputs, prepared.bound_input_check())?;
424    let signature = input_signature_reads(&inputs)?;
425    if *prepared.signature != signature {
426        super::preparation::validate_execution_signature(
427            prepared.prepared.as_ref(),
428            prepared.program.frozen(),
429            &signature,
430        )
431        .map_err(prepare_error)?;
432    }
433    let root = prepared.prepared.root();
434    engines.take_caches(root.needs_extension_caches());
435    let result = execute_scheduled(root, prepared.prepared.operations(), inputs, &mut engines);
436    engines.restore_caches();
437    result
438}
439
440fn validate_prepared_runtime(
441    runtime: &Runtime,
442    prepared: &Prepared,
443    caller: &'static str,
444) -> Result<()> {
445    if runtime.id() != prepared.runtime_id {
446        return Err(Error::runtime_state(
447            caller,
448            ErrorPhase::Execution,
449            "prepared program belongs to a different runtime",
450        ));
451    }
452    validate_prepared_epoch(runtime, prepared.epoch, caller)
453}
454
455fn validate_prepared_epoch(
456    runtime: &Runtime,
457    prepared_epoch: super::RuntimeEpoch,
458    caller: &'static str,
459) -> Result<()> {
460    let epoch = runtime
461        .epoch()
462        .map_err(|source| Error::runtime_state_source(caller, ErrorPhase::Execution, source))?;
463    if epoch != prepared_epoch {
464        return Err(Error::runtime_state(
465            caller,
466            ErrorPhase::Execution,
467            format!(
468                "prepared epoch {:?} does not match current epoch {:?}",
469                prepared_epoch, epoch
470            ),
471        ));
472    }
473    Ok(())
474}
475
476/// The engines of one run, bound to the prepared targets.
477///
478/// Slot `i` of the prepared program runs on `engines[i]`; [`Self::bind`]
479/// checks once, before any launch, that every engine reports the target it is
480/// bound to. Extension caches are moved out of the engines for the duration of
481/// the run so an instruction can borrow an engine and its caches together.
482struct EngineSet<'e, 'r> {
483    engines: &'e mut [&'r mut dyn ExecutionEngine],
484    targets: &'e [EngineTarget],
485    caches: Vec<ExtensionCacheStore>,
486}
487
488impl<'e, 'r> EngineSet<'e, 'r> {
489    fn bind(
490        engines: &'e mut [&'r mut dyn ExecutionEngine],
491        targets: &'e [EngineTarget],
492        caller: &'static str,
493    ) -> Result<Self> {
494        if engines.len() != targets.len() {
495            return Err(Error::invalid_argument(
496                caller,
497                ErrorPhase::Execution,
498                "engines",
499                format!(
500                    "the program was prepared for {} engines, but {} were passed",
501                    targets.len(),
502                    engines.len()
503                ),
504            ));
505        }
506        for (index, (engine, target)) in engines.iter().zip(targets).enumerate() {
507            let actual = engine.target();
508            if &actual != target {
509                return Err(Error::invalid_argument(
510                    caller,
511                    ErrorPhase::Execution,
512                    "engines",
513                    format!(
514                        "engine {index} reports target {actual:?}, but the program was prepared for {target:?}"
515                    ),
516                ));
517            }
518        }
519        Ok(Self {
520            engines,
521            targets,
522            caches: Vec::new(),
523        })
524    }
525
526    /// Lend every engine's extension caches to this run.
527    ///
528    /// Only a program with an extension command reads them; a run of any other
529    /// program leaves the caches in place and skips the exchange.
530    fn take_caches(&mut self, needed: bool) {
531        if needed {
532            self.caches = self
533                .engines
534                .iter_mut()
535                .map(|engine| engine.take_extension_caches())
536                .collect();
537        }
538    }
539
540    fn restore_caches(&mut self) {
541        for (engine, caches) in self.engines.iter_mut().zip(self.caches.drain(..)) {
542            engine.restore_extension_caches(caches);
543        }
544    }
545
546    fn index_of(&self, engine_id: &EngineId) -> Result<usize> {
547        self.targets
548            .iter()
549            .position(|target| target.engine_id() == engine_id)
550            .ok_or_else(|| {
551                Error::runtime_state(
552                    "Runtime::run",
553                    ErrorPhase::Execution,
554                    format!(
555                        "the prepared schedule uses engine {:?}, which is not among the run's engines",
556                        engine_id.as_str()
557                    ),
558                )
559            })
560    }
561
562    fn engine(&mut self, engine_id: &EngineId) -> Result<&mut dyn ExecutionEngine> {
563        let index = self.index_of(engine_id)?;
564        Ok(&mut *self.engines[index])
565    }
566
567    fn driver(&self, engine_id: &EngineId) -> Result<Arc<dyn EventDomainDriver>> {
568        let index = self.index_of(engine_id)?;
569        Ok(self.engines[index].event_domain_driver())
570    }
571
572    fn execute_slot_instruction<'input>(
573        &mut self,
574        location: &ExecutionLocation,
575        instruction_index: usize,
576        instruction: &ExecInstruction,
577        operations: &[PreparedOperationPlan],
578        slots: &mut [Option<ExecSlot<'input>>],
579    ) -> Result<()> {
580        let index = self.index_of(location.engine_id())?;
581        let engine = &mut *self.engines[index];
582        execute_slot_instruction_on_engine(
583            engine,
584            self.caches.get_mut(index),
585            operations,
586            instruction_index,
587            instruction,
588            slots,
589        )
590    }
591
592    fn execute_elementwise_fusion_slots<'input>(
593        &mut self,
594        location: &ExecutionLocation,
595        input_slots: &[usize],
596        instruction_count: usize,
597        plan: &tenferro_tensor::backend::ElementwiseFusionPlan,
598        slots: &mut [Option<ExecSlot<'input>>],
599        output_slots: &[usize],
600    ) -> Result<bool> {
601        let index = self.index_of(location.engine_id())?;
602        execute_elementwise_fusion_slots_on_engine(
603            &mut *self.engines[index],
604            input_slots,
605            instruction_count,
606            plan,
607            slots,
608            output_slots,
609        )
610    }
611
612    /// Hand every event-domain run to the engine that owns its domain.
613    fn retire(&mut self, event_domains: ScheduledEventDomains) -> Result<()> {
614        let mut failures = Vec::new();
615        for (engine_id, run) in event_domains.into_pending_runs() {
616            let result = match self.index_of(&engine_id) {
617                Ok(index) => self.engines[index].retire_run(run),
618                Err(error) => {
619                    drop_event_domain_run(run);
620                    Err(error)
621                }
622            };
623            if let Err(error) = result {
624                failures.push(error);
625            }
626        }
627        let mut failures = failures.into_iter();
628        let Some(mut error) = failures.next() else {
629            return Ok(());
630        };
631        for failure in failures {
632            error = Error::with_suppressed(error, failure);
633        }
634        Err(error)
635    }
636}
637
638impl Drop for EngineSet<'_, '_> {
639    fn drop(&mut self) {
640        self.restore_caches();
641    }
642}
643
644fn execute_scheduled<'input>(
645    root: &super::preparation::PreparedProgramRoot,
646    operations: &[PreparedOperationPlan],
647    inputs: RuntimeInputReads<'input>,
648    engines: &mut EngineSet<'_, '_>,
649) -> Result<Vec<Output<'input>>> {
650    let program = root.staging();
651    let schedule = root.schedule();
652    // The immutable program was validated when the root was prepared; only
653    // the run's own input list is checked here.
654    validate_exec_input_count(program, inputs.len())?;
655    let mut event_domains = Some(ScheduledEventDomains::new(schedule, engines)?);
656    let inputs = inputs.into_iter().map(ExecSlot::Read).collect();
657    let flat = schedule.flat_location();
658    let executed = execute_scheduled_slots_into(
659        program,
660        schedule,
661        flat,
662        root.regions(),
663        root.node_regions(),
664        root.region_counters(),
665        operations,
666        inputs,
667        engines,
668        &mut event_domains,
669    );
670    let Some(event_domains) = event_domains.take() else {
671        return executed.and_then(|_| {
672            Err(Error::runtime_state(
673                "Runtime::run",
674                ErrorPhase::Execution,
675                "scheduled execution consumed its event domains before collecting outputs",
676            ))
677        });
678    };
679    let collected = executed
680        .and_then(|slots| collect_outputs(program, slots, flat, root.output_duplicates(), engines));
681    match collected {
682        Ok(outputs) => {
683            engines.retire(event_domains)?;
684            Ok(outputs)
685        }
686        Err(primary) => {
687            // On failure, nothing the run produced is handed out; retire the
688            // run now so its device work cannot outlive the error.
689            let mut event_domains = event_domains;
690            match event_domains.drain() {
691                Ok(()) => Err(primary),
692                Err(cleanup) => Err(scheduled_execution_cleanup_error(primary, cleanup)),
693            }
694        }
695    }
696}
697
698/// Turn the run's output slots into [`Output`]s, in program output order.
699///
700/// A slot that still reads one of the run's inputs (the input itself or a view
701/// of it) is returned borrowed, without a copy. Owned results, which are
702/// always compact, are returned as they are. A program may name one output slot
703/// more than once (`duplicates`, computed at prepare): the first position takes
704/// the value, a repeated borrowed read is cloned shallowly, and a repeated
705/// owned result is an explicit engine copy on the value's location.
706fn collect_outputs<'input>(
707    program: &ExecProgram,
708    slots: RunSlots<'input>,
709    flat: Option<&ExecutionLocation>,
710    duplicates: &[Option<usize>],
711    engines: &mut EngineSet<'_, '_>,
712) -> Result<Vec<Output<'input>>> {
713    let RunSlots {
714        mut staged,
715        mut located,
716    } = slots;
717    let mut collected: Vec<Output<'input>> = Vec::with_capacity(program.output_slots.len());
718    // The location of each collected owned output, for a repeated position's
719    // engine copy; the flat run has one location for every value.
720    let mut locations: Vec<Option<ExecutionLocation>> = Vec::new();
721    for (position, &slot) in program.output_slots.iter().enumerate() {
722        if let Some(first) = duplicates.get(position).copied().flatten() {
723            let output = match collected.get(first) {
724                Some(Output::Borrowed(read)) => Output::Borrowed(read.clone()),
725                Some(Output::Owned(tensor)) => {
726                    let location = match flat {
727                        Some(location) => location,
728                        None => locations
729                            .get(first)
730                            .and_then(Option::as_ref)
731                            .ok_or(tenferro_tensor::Error::MissingValue { slot })?,
732                    };
733                    Output::Owned(duplicate_owned_output(engines, location, tensor)?)
734                }
735                None => return Err(tenferro_tensor::Error::MissingValue { slot }.into()),
736            };
737            collected.push(output);
738            if flat.is_none() {
739                locations.push(None);
740            }
741            continue;
742        }
743        let value = match (flat, &mut located) {
744            (Some(_), _) => staged
745                .get_mut(slot)
746                .and_then(Option::take)
747                .ok_or(tenferro_tensor::Error::MissingValue { slot })?,
748            (None, Some(located)) => {
749                let values = located
750                    .get_mut(slot)
751                    .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
752                let value = values
753                    .pop()
754                    .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
755                values.clear();
756                locations.push(Some(value.location));
757                value.value
758            }
759            (None, None) => {
760                return Err(Error::runtime_state(
761                    "Runtime::run",
762                    ErrorPhase::Execution,
763                    "a multi-location run lost its value lists".to_string(),
764                ))
765            }
766        };
767        collected.push(match value {
768            ExecSlot::Read(read) => Output::Borrowed(read),
769            ExecSlot::Owned(tensor) => Output::Owned(tensor),
770        });
771    }
772    Ok(collected)
773}
774
775fn duplicate_owned_output(
776    engines: &mut EngineSet<'_, '_>,
777    location: &ExecutionLocation,
778    tensor: &Tensor,
779) -> Result<Tensor> {
780    let engine = engines.engine(location.engine_id())?;
781    engine
782        .to_contiguous_read(TensorRead::from_tensor(tensor))
783        .map_err(Error::from)
784}
785
786/// One run's value store.
787///
788/// A run has either one location or several. With one (`flat`) a slot holds at
789/// most one value and always at that location, so `staged` is the slot array
790/// itself: a command neither stages its inputs nor republishes its results.
791/// With several, `located` keeps one value list per slot per location and a
792/// command stages its inputs into `staged` and publishes its results back.
793struct RunSlots<'input> {
794    /// The array commands execute against.
795    staged: Vec<Option<ExecSlot<'input>>>,
796    /// `Some` exactly when the run visits more than one location.
797    located: Option<Vec<Vec<LocatedExecSlot<'input>>>>,
798}
799
800impl<'input> RunSlots<'input> {
801    fn new(n_slots: usize, flat: bool) -> Self {
802        Self {
803            staged: Vec::new(),
804            located: (!flat).then(|| (0..n_slots).map(|_| Vec::new()).collect()),
805        }
806    }
807
808    /// Place the run's inputs and check they are readable at their location.
809    fn initialize(
810        &mut self,
811        program: &ExecProgram,
812        inputs: Vec<ExecSlot<'input>>,
813        schedule: &ScheduledGraph,
814    ) -> Result<()> {
815        crate::exec::initialize_exec_slots_in(program, inputs, &mut self.staged)?;
816        if schedule.input_locations().len() != program.input_slots.len() {
817            return Err(Error::runtime_state(
818                "Runtime::run",
819                ErrorPhase::Execution,
820                format!(
821                    "prepared schedule has {} input locations for {} inputs",
822                    schedule.input_locations().len(),
823                    program.input_slots.len()
824                ),
825            ));
826        }
827        for (&slot, location) in program.input_slots.iter().zip(schedule.input_locations()) {
828            match &mut self.located {
829                Some(located) => {
830                    let value = self
831                        .staged
832                        .get_mut(slot)
833                        .and_then(Option::take)
834                        .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
835                    validate_runtime_input_ingress(location, &value.as_read(), slot)?;
836                    located[slot].push(LocatedExecSlot {
837                        location: location.clone(),
838                        value,
839                    });
840                }
841                None => {
842                    let value = self
843                        .staged
844                        .get(slot)
845                        .and_then(Option::as_ref)
846                        .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
847                    validate_runtime_input_ingress(location, &value.as_read(), slot)?;
848                }
849            }
850        }
851        Ok(())
852    }
853
854    fn execution(&mut self) -> &mut [Option<ExecSlot<'input>>] {
855        &mut self.staged
856    }
857
858    fn execution_ref(&self) -> &[Option<ExecSlot<'input>>] {
859        &self.staged
860    }
861
862    /// Move one command's inputs into the execution array.
863    fn stage_command(
864        &mut self,
865        instruction: &ExecInstruction,
866        location: &ExecutionLocation,
867    ) -> Result<()> {
868        match &mut self.located {
869            Some(located) => {
870                stage_instruction_inputs(instruction, location, located, &mut self.staged)
871            }
872            None => Ok(()),
873        }
874    }
875
876    /// Publish one command's results, releasing the inputs it consumed last.
877    fn publish_command(
878        &mut self,
879        instruction: &ExecInstruction,
880        location: &ExecutionLocation,
881    ) -> Result<()> {
882        match &mut self.located {
883            Some(located) => {
884                retain_instruction_results(instruction, location, located, &mut self.staged)
885            }
886            None => Ok(()),
887        }
888    }
889
890    /// Move a region's inputs into the execution array.
891    fn stage_region(
892        &mut self,
893        region: &super::region::ElementwiseRegion,
894        location: &ExecutionLocation,
895    ) -> Result<()> {
896        let Some(located) = &mut self.located else {
897            return Ok(());
898        };
899        for &slot in &region.input_slots {
900            stage_slot_input(slot, location, located, &mut self.staged)?;
901        }
902        Ok(())
903    }
904
905    /// Publish a fused region's live-outs, releasing the inputs whose last use
906    /// is inside the region.
907    fn publish_region(
908        &mut self,
909        region: &super::region::ElementwiseRegion,
910        location: &ExecutionLocation,
911    ) -> Result<()> {
912        match &mut self.located {
913            Some(located) => retain_region_results(region, location, located, &mut self.staged),
914            None => {
915                for (&slot, &last_use) in
916                    region.input_slots.iter().zip(region.input_last_use.iter())
917                {
918                    if last_use {
919                        self.staged[slot] = None;
920                    }
921                }
922                Ok(())
923            }
924        }
925    }
926
927    /// The per-location value lists a transfer node needs.
928    fn transfers(&mut self) -> Result<&mut [Vec<LocatedExecSlot<'input>>]> {
929        self.located.as_deref_mut().ok_or_else(|| {
930            Error::runtime_state(
931                "Runtime::run",
932                ErrorPhase::Execution,
933                "a single-location run must not contain a transfer node".to_string(),
934            )
935        })
936    }
937}
938
939/// The location `instruction` executes at, proven equal to its node's.
940///
941/// A flat schedule proved at prepare that every operation sits at its one
942/// location (`ScheduledGraph::flat_location`), so the per-command comparison
943/// only runs for a multi-location schedule.
944fn command_location<'a>(
945    schedule: &'a ScheduledGraph,
946    flat: Option<&'a ExecutionLocation>,
947    instruction: &ExecInstruction,
948    node_location: &ExecutionLocation,
949) -> Result<&'a ExecutionLocation> {
950    if let Some(location) = flat {
951        return Ok(location);
952    }
953    let location = instruction_execution(schedule, instruction)?.location();
954    if location != node_location {
955        return Err(Error::runtime_state(
956            "Runtime::run",
957            ErrorPhase::Execution,
958            "a prepared command's location does not match its scheduled node".to_string(),
959        ));
960    }
961    Ok(location)
962}
963
964#[allow(clippy::too_many_arguments)]
965fn execute_scheduled_slots_into<'input>(
966    program: &ExecProgram,
967    schedule: &ScheduledGraph,
968    flat: Option<&ExecutionLocation>,
969    regions: &[super::region::ElementwiseRegion],
970    node_regions: &[super::region::NodeRegion],
971    region_counters: &super::region::RegionExecutionCounters,
972    operations: &[PreparedOperationPlan],
973    inputs: Vec<ExecSlot<'input>>,
974    engines: &mut EngineSet<'_, '_>,
975    event_domains: &mut Option<ScheduledEventDomains>,
976) -> Result<RunSlots<'input>> {
977    let domains = event_domains.as_mut().ok_or_else(|| {
978        Error::runtime_state(
979            "Runtime::run",
980            ErrorPhase::Execution,
981            "scheduled execution has no event-domain run to submit into",
982        )
983    })?;
984    let mut slots = RunSlots::new(program.n_slots, flat.is_some());
985    // INVARIANT: the immutable schedule was validated during preparation;
986    // driver preflight completes before input ingress or any operation launch.
987    // A provider panic is caught after this body and drops the event-domain
988    // run before this frame's staged tensor storage drops.
989    // Planned elementwise regions (`node_regions`, computed at prepare): the
990    // start node executes the whole region as one fused command and the
991    // covered nodes are skipped.
992    let result = std::panic::catch_unwind(AssertUnwindSafe(|| -> Result<RunSlots<'input>> {
993        (|| {
994            slots.initialize(program, inputs, schedule)?;
995            for (node_index, node) in schedule.nodes().iter().enumerate() {
996                match node {
997                    ScheduledNode::Operation(operation_node) => {
998                        let role = node_regions
999                            .get(node_index)
1000                            .copied()
1001                            .unwrap_or(super::region::NodeRegion::Command);
1002                        if let super::region::NodeRegion::Start(region_index) = role {
1003                            let region = &regions[region_index];
1004                            let first_instruction = program
1005                                    .instructions
1006                                    .get(region.instruction_range.start)
1007                                    .ok_or_else(|| {
1008                                        Error::runtime_state(
1009                                            "Runtime::run",
1010                                            ErrorPhase::Execution,
1011                                            format!(
1012                                                "elementwise region references instruction {}, but the execution program has {} instructions",
1013                                                region.instruction_range.start,
1014                                                program.instructions.len()
1015                                            ),
1016                                        )
1017                                    })?;
1018                            let location = command_location(
1019                                schedule,
1020                                flat,
1021                                first_instruction,
1022                                operation_node.location(),
1023                            )?;
1024                            let mut launch = || {
1025                                slots.stage_region(region, location)?;
1026                                let applied = engines.execute_elementwise_fusion_slots(
1027                                    location,
1028                                    &region.input_slots,
1029                                    region.instructions.len(),
1030                                    &region.plan,
1031                                    slots.execution(),
1032                                    &region.output_slots,
1033                                )?;
1034                                if applied {
1035                                    region_counters.record_fused();
1036                                    validate_region_outputs(
1037                                        region,
1038                                        location,
1039                                        slots.execution_ref(),
1040                                    )?;
1041                                    return slots.publish_region(region, location);
1042                                }
1043                                // The backend rejected the fusion: execute the
1044                                // region's instructions one by one.
1045                                region_counters.record_fallback();
1046                                for (offset, instruction) in region.instructions.iter().enumerate()
1047                                {
1048                                    let instruction_index = region.instruction_range.start + offset;
1049                                    let member_location = command_location(
1050                                        schedule,
1051                                        flat,
1052                                        instruction,
1053                                        operation_node.location(),
1054                                    )?;
1055                                    slots.stage_command(instruction, member_location)?;
1056                                    engines.execute_slot_instruction(
1057                                        member_location,
1058                                        instruction_index,
1059                                        instruction,
1060                                        operations,
1061                                        slots.execution(),
1062                                    )?;
1063                                    validate_instruction_outputs(
1064                                        instruction_index,
1065                                        instruction,
1066                                        member_location,
1067                                        slots.execution_ref(),
1068                                    )?;
1069                                    slots.publish_command(instruction, member_location)?;
1070                                }
1071                                Ok(())
1072                            };
1073                            domains.enqueue(node_index, node, &mut launch)?;
1074                            region_counters.record_submission();
1075                            // Every covered node's completion resolves to the region's
1076                            // completion token, which was recorded after the region's
1077                            // last command.
1078                            let completion = EventDependency::from_completion(node.completion());
1079                            if let Some(token) = domains.completions.get(&completion).cloned() {
1080                                for &covered in &region.node_indices {
1081                                    if covered == node_index {
1082                                        continue;
1083                                    }
1084                                    let covered_node = schedule.nodes().get(covered).ok_or_else(|| {
1085                                            Error::runtime_state(
1086                                                "Runtime::run",
1087                                                ErrorPhase::Execution,
1088                                                format!(
1089                                                    "elementwise region references schedule node {covered}, but the schedule has {} nodes",
1090                                                    schedule.nodes().len()
1091                                                ),
1092                                            )
1093                                        })?;
1094                                    domains.completions.insert(
1095                                        EventDependency::from_completion(covered_node.completion()),
1096                                        Arc::clone(&token),
1097                                    );
1098                                }
1099                            }
1100                            continue;
1101                        }
1102                        if role == super::region::NodeRegion::Covered {
1103                            // Executed as part of its region.
1104                            continue;
1105                        }
1106                        let mut launch = || {
1107                            let instruction_index = operation_node.instruction_index();
1108                            let instruction = program
1109                                    .instructions
1110                                    .get(instruction_index)
1111                                    .ok_or_else(|| {
1112                                        Error::runtime_state(
1113                                            "Runtime::run",
1114                                            ErrorPhase::Execution,
1115                                            format!(
1116                                                "scheduled operation references instruction {instruction_index}, but the execution program has {} instructions",
1117                                                program.instructions.len()
1118                                            ),
1119                                        )
1120                                    })?;
1121                            let location = command_location(
1122                                schedule,
1123                                flat,
1124                                instruction,
1125                                operation_node.location(),
1126                            )?;
1127                            slots.stage_command(instruction, location)?;
1128                            engines.execute_slot_instruction(
1129                                location,
1130                                instruction_index,
1131                                instruction,
1132                                operations,
1133                                slots.execution(),
1134                            )?;
1135                            validate_instruction_outputs(
1136                                instruction_index,
1137                                instruction,
1138                                location,
1139                                slots.execution_ref(),
1140                            )?;
1141                            slots.publish_command(instruction, location)
1142                        };
1143                        domains.enqueue(node_index, node, &mut launch)?;
1144                        region_counters.record_submission();
1145                    }
1146                    ScheduledNode::Transfer(transfer) => {
1147                        let mut launch =
1148                            || execute_scheduled_transfer(transfer, slots.transfers()?);
1149                        domains.enqueue(node_index, node, &mut launch)?;
1150                        region_counters.record_submission();
1151                    }
1152                    ScheduledNode::Collective(_) => {
1153                        return Err(Error::runtime_state_source(
1154                            "Runtime::run",
1155                            ErrorPhase::Execution,
1156                            UnsupportedScheduledNodeError {
1157                                node_index,
1158                                node_kind: ScheduledNodeKind::Collective,
1159                            },
1160                        ));
1161                    }
1162                    ScheduledNode::Barrier(_) => {
1163                        let mut launch = || Ok(());
1164                        domains.enqueue(node_index, node, &mut launch)?;
1165                        region_counters.record_submission();
1166                    }
1167                }
1168            }
1169            Ok(())
1170        })()?;
1171        Ok(slots)
1172    }));
1173    match result {
1174        Ok(slots) => slots,
1175        Err(payload) => {
1176            // A provider launch may unwind through the scheduler. Drop the
1177            // event-domain run before this frame's staged tensor storage drops,
1178            // preserving the same cleanup order as the synchronous wrapper.
1179            event_domains.take();
1180            std::panic::resume_unwind(payload)
1181        }
1182    }
1183}
1184
1185#[derive(Debug)]
1186pub(crate) struct ScheduledEventDomains {
1187    runs: Vec<RuntimeOwnedEventDomainRun>,
1188    completions: HashMap<EventDependency, Arc<dyn EventToken>>,
1189    /// Whether the single run reported that submission order alone orders its
1190    /// nodes. The scheduler then submits each launch without resolving
1191    /// dependencies or recording a completion token for it.
1192    coalesced: bool,
1193}
1194
1195#[derive(Debug, thiserror::Error)]
1196#[error(
1197    "scheduled node {node_index} depends on completion {dependency:?}, but no completion token was recorded"
1198)]
1199pub(crate) struct MissingScheduledDependencyCompletionError {
1200    pub(crate) dependency: EventDependency,
1201    pub(crate) node_index: usize,
1202}
1203
1204#[derive(Debug, thiserror::Error)]
1205#[error("scheduled transfer destination {destination:?} already contains value slot {value_slot}")]
1206pub(crate) struct DuplicateTransferDestinationError {
1207    pub(crate) value_slot: usize,
1208    pub(crate) destination: ExecutionLocation,
1209}
1210
1211impl ScheduledEventDomains {
1212    /// Number of completion tokens recorded so far.
1213    #[cfg(test)]
1214    pub(crate) fn completion_count_for_test(&self) -> usize {
1215        self.completions.len()
1216    }
1217
1218    fn new(schedule: &ScheduledGraph, engines: &EngineSet<'_, '_>) -> Result<Self> {
1219        Self::build(schedule, |engine_id| engines.driver(engine_id))
1220    }
1221
1222    /// Build the runs a schedule needs with an explicit driver lookup.
1223    ///
1224    /// Test-only: production always takes each run's driver from the engine the
1225    /// schedule placed the work on.
1226    #[cfg(test)]
1227    pub(crate) fn new_for_test(
1228        schedule: &ScheduledGraph,
1229        drivers: &[(EngineId, Arc<dyn super::EventDomainDriver>)],
1230    ) -> Result<Self> {
1231        Self::build(schedule, |engine_id| {
1232            drivers
1233                .iter()
1234                .find(|(id, _)| id == engine_id)
1235                .map(|(_, driver)| Arc::clone(driver))
1236                .ok_or_else(|| {
1237                    Error::runtime_state(
1238                        "ScheduledEventDomains::new_for_test",
1239                        ErrorPhase::Execution,
1240                        format!("no test driver for engine {:?}", engine_id.as_str()),
1241                    )
1242                })
1243        })
1244    }
1245
1246    fn build(
1247        schedule: &ScheduledGraph,
1248        driver_for: impl Fn(&EngineId) -> Result<Arc<dyn super::EventDomainDriver>>,
1249    ) -> Result<Self> {
1250        schedule.preflight().map_err(|source| {
1251            Error::runtime_state_source("Runtime::run", ErrorPhase::Execution, source)
1252        })?;
1253        // A coalescing-eligible schedule names exactly one event domain: every
1254        // location, completion, and dependency carries the root's, and no
1255        // transfer node exists. Seed that single run directly instead of
1256        // scanning every location into a driver table only to find one entry.
1257        if schedule.coalescing_eligible() {
1258            let domain = schedule.root_location().event_domain_id();
1259            let engine_id = schedule.root_location().engine_id().clone();
1260            let mut run = RuntimeOwnedEventDomainRun::new(
1261                domain,
1262                engine_id.clone(),
1263                driver_for(&engine_id)?.begin_run(domain)?,
1264            );
1265            let actual = run.domain(EventDomainOperation::BeginRun)?;
1266            if actual != domain {
1267                return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1268                    operation: EventDomainOperation::BeginRun,
1269                    node_index: None,
1270                    expected: domain,
1271                    actual,
1272                }));
1273            }
1274            let coalesced = run.enable_single_domain_coalescing();
1275            return Ok(Self {
1276                runs: vec![run],
1277                completions: HashMap::new(),
1278                coalesced,
1279            });
1280        }
1281        fn add_domain(
1282            drivers: &mut Vec<(super::EventDomainId, EngineId)>,
1283            seen_domains: &mut HashSet<super::EventDomainId>,
1284            domain: super::EventDomainId,
1285            engine_id: &EngineId,
1286        ) {
1287            if seen_domains.insert(domain) {
1288                drivers.push((domain, engine_id.clone()));
1289            }
1290        }
1291        let mut drivers = Vec::new();
1292        let mut seen_domains = HashSet::new();
1293        // Seed a run for every scheduled location, not only for scheduled nodes.
1294        // Output materialization launches work on the output location's stream
1295        // after the nodes have run, and a zero-instruction graph (an identity
1296        // graph) has no node to seed a run at all; without the location runs
1297        // there would be nothing to retire that launch.
1298        add_domain(
1299            &mut drivers,
1300            &mut seen_domains,
1301            schedule.root_location().event_domain_id(),
1302            schedule.root_location().engine_id(),
1303        );
1304        for location in schedule.input_locations() {
1305            add_domain(
1306                &mut drivers,
1307                &mut seen_domains,
1308                location.event_domain_id(),
1309                location.engine_id(),
1310            );
1311        }
1312        for location in schedule.operation_locations() {
1313            add_domain(
1314                &mut drivers,
1315                &mut seen_domains,
1316                location.event_domain_id(),
1317                location.engine_id(),
1318            );
1319        }
1320        for node in schedule.nodes() {
1321            add_domain(
1322                &mut drivers,
1323                &mut seen_domains,
1324                node.completion().domain(),
1325                node.event_domain_witness().engine_id(),
1326            );
1327        }
1328        let mut runs = Vec::with_capacity(drivers.len());
1329        for (domain, engine_id) in drivers {
1330            let driver = driver_for(&engine_id)?;
1331            let run = RuntimeOwnedEventDomainRun::new(domain, engine_id, driver.begin_run(domain)?);
1332            let actual = run.domain(EventDomainOperation::BeginRun)?;
1333            if actual != domain {
1334                return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1335                    operation: EventDomainOperation::BeginRun,
1336                    node_index: None,
1337                    expected: domain,
1338                    actual,
1339                }));
1340            }
1341            runs.push(run);
1342        }
1343        Ok(Self {
1344            runs,
1345            completions: HashMap::new(),
1346            coalesced: false,
1347        })
1348    }
1349
1350    #[cfg(test)]
1351    pub(crate) fn for_test(
1352        drivers: Vec<(super::EventDomainId, Arc<dyn super::EventDomainDriver>)>,
1353    ) -> Result<Self> {
1354        let mut runs = Vec::with_capacity(drivers.len());
1355        for (domain, driver) in drivers {
1356            let engine_id =
1357                EngineId::new("tenferro-test.event-domain-engine").expect("test engine id");
1358            let run = RuntimeOwnedEventDomainRun::new(domain, engine_id, driver.begin_run(domain)?);
1359            let actual = run.domain(EventDomainOperation::BeginRun)?;
1360            if actual != domain {
1361                return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1362                    operation: EventDomainOperation::BeginRun,
1363                    node_index: None,
1364                    expected: domain,
1365                    actual,
1366                }));
1367            }
1368            runs.push(run);
1369        }
1370        Ok(Self {
1371            runs,
1372            completions: HashMap::new(),
1373            coalesced: false,
1374        })
1375    }
1376
1377    pub(crate) fn enqueue(
1378        &mut self,
1379        node_index: usize,
1380        node: &ScheduledNode,
1381        launch: &mut dyn FnMut() -> Result<()>,
1382    ) -> Result<()> {
1383        let completion = node.completion();
1384        if self.coalesced {
1385            // The run reported that submission order alone orders its nodes
1386            // (#2186): no dependency token exists that a later node could read,
1387            // so submit the launch through the run's own submission context
1388            // without resolving dependencies or recording a completion.
1389            debug_assert_eq!(self.runs.len(), 1);
1390            self.runs[0].enqueue(&[], launch)?;
1391            return Ok(());
1392        }
1393        let destination = completion.domain();
1394        let run_index = self.run_index(completion.domain())?;
1395        let actual_preflight_domain = self.runs[run_index].domain(EventDomainOperation::Enqueue)?;
1396        if actual_preflight_domain != destination {
1397            return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1398                operation: EventDomainOperation::Enqueue,
1399                node_index: Some(node_index),
1400                expected: destination,
1401                actual: actual_preflight_domain,
1402            }));
1403        }
1404        let dependencies = self.classify_dependencies(node_index, node, destination)?;
1405        let actual_run_domain = self.runs[run_index].domain(EventDomainOperation::Enqueue)?;
1406        if actual_run_domain != destination {
1407            return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1408                operation: EventDomainOperation::Enqueue,
1409                node_index: Some(node_index),
1410                expected: destination,
1411                actual: actual_run_domain,
1412            }));
1413        }
1414        let completion_event = self.runs[run_index].enqueue(&dependencies, launch)?;
1415        let actual = completion_event.origin();
1416        if actual != completion.domain() {
1417            return Err(event_domain_error(
1418                EventDomainError::CompletionTokenDomainMismatch {
1419                    operation: EventDomainOperation::ValidateCompletion,
1420                    node_index: Some(node_index),
1421                    expected: completion.domain(),
1422                    actual,
1423                },
1424            ));
1425        }
1426        self.completions.insert(
1427            EventDependency::from_completion(completion),
1428            completion_event,
1429        );
1430        Ok(())
1431    }
1432
1433    fn classify_dependencies(
1434        &self,
1435        node_index: usize,
1436        node: &ScheduledNode,
1437        destination: super::EventDomainId,
1438    ) -> Result<SmallVec<[Arc<dyn EventToken>; 4]>> {
1439        // `Arc<dyn EventToken>` has no `Default`; `tinyvec::Array` cannot hold it.
1440        let mut admitted = SmallVec::with_capacity(node.dependencies().len());
1441        for dependency in node.dependencies() {
1442            let dependency_completion =
1443                self.completions.get(dependency).cloned().ok_or_else(|| {
1444                    Error::runtime_state_source(
1445                        "Runtime::run",
1446                        ErrorPhase::Execution,
1447                        MissingScheduledDependencyCompletionError {
1448                            dependency: *dependency,
1449                            node_index,
1450                        },
1451                    )
1452                })?;
1453            let actual = dependency_completion.origin();
1454            if actual != dependency.domain() {
1455                return Err(event_domain_error(
1456                    EventDomainError::DependencyDomainMismatch {
1457                        operation: match node {
1458                            ScheduledNode::Transfer(_) => EventDomainOperation::TransferBridge,
1459                            ScheduledNode::Operation(_)
1460                            | ScheduledNode::Collective(_)
1461                            | ScheduledNode::Barrier(_) => EventDomainOperation::Enqueue,
1462                        },
1463                        node_index: Some(node_index),
1464                        expected: dependency.domain(),
1465                        actual,
1466                    },
1467                ));
1468            }
1469            match node {
1470                ScheduledNode::Transfer(transfer) => {
1471                    let source = transfer.source_event_domain();
1472                    if actual == destination {
1473                        admitted.push(dependency_completion);
1474                    } else if actual == source {
1475                        dependency_completion.wait().map_err(|source_error| {
1476                            event_domain_error(EventDomainError::DependencyWaitFailed {
1477                                operation: EventDomainOperation::TransferBridge,
1478                                node_index: Some(node_index),
1479                                expected: destination,
1480                                actual,
1481                                source: Box::new(source_error),
1482                            })
1483                        })?;
1484                    } else {
1485                        return Err(event_domain_error(
1486                            EventDomainError::DependencyDomainMismatch {
1487                                operation: EventDomainOperation::TransferBridge,
1488                                node_index: Some(node_index),
1489                                expected: source,
1490                                actual,
1491                            },
1492                        ));
1493                    }
1494                }
1495                ScheduledNode::Operation(_)
1496                | ScheduledNode::Collective(_)
1497                | ScheduledNode::Barrier(_) => {
1498                    if actual != destination {
1499                        return Err(event_domain_error(
1500                            EventDomainError::DependencyDomainMismatch {
1501                                operation: EventDomainOperation::Enqueue,
1502                                node_index: Some(node_index),
1503                                expected: destination,
1504                                actual,
1505                            },
1506                        ));
1507                    }
1508                    admitted.push(dependency_completion);
1509                }
1510            }
1511        }
1512        Ok(admitted)
1513    }
1514
1515    fn run_index(&self, domain: super::EventDomainId) -> Result<usize> {
1516        if let Some(index) = self
1517            .runs
1518            .iter()
1519            .position(|run| run.requested_domain() == domain)
1520        {
1521            return Ok(index);
1522        }
1523        Err(missing_event_domain_driver(domain))
1524    }
1525
1526    /// Release every pending run, each with the engine that owns it.
1527    fn into_pending_runs(self) -> impl Iterator<Item = (EngineId, Box<dyn EventDomainRun>)> {
1528        self.runs
1529            .into_iter()
1530            .filter_map(RuntimeOwnedEventDomainRun::into_pending)
1531    }
1532
1533    pub(crate) fn drain(&mut self) -> Result<()> {
1534        let mut failures = Vec::new();
1535        for run in &mut self.runs {
1536            if let Err(error) = run.drain() {
1537                failures.push(error);
1538            }
1539        }
1540        let mut failures = failures.into_iter();
1541        let Some(mut error) = failures.next() else {
1542            return Ok(());
1543        };
1544        for failure in failures {
1545            error = Error::with_suppressed(error, failure);
1546        }
1547        Err(error)
1548    }
1549}
1550
1551#[derive(Debug)]
1552enum RuntimeOwnedEventDomainRunState {
1553    Pending(Box<dyn EventDomainRun>),
1554    Retired,
1555    Failed,
1556}
1557
1558#[derive(Clone, Copy, Debug, Eq, PartialEq)]
1559enum EventDomainRunTerminalState {
1560    Retired,
1561    Failed,
1562}
1563
1564impl fmt::Display for EventDomainRunTerminalState {
1565    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1566        match self {
1567            Self::Retired => formatter.write_str("retired"),
1568            Self::Failed => formatter.write_str("failed"),
1569        }
1570    }
1571}
1572
1573#[derive(Debug)]
1574struct RuntimeOwnedEventDomainRun {
1575    requested_domain: super::EventDomainId,
1576    /// The engine whose driver opened this run; it takes the run over once the
1577    /// outputs are collected.
1578    engine_id: EngineId,
1579    state: RuntimeOwnedEventDomainRunState,
1580}
1581
1582impl RuntimeOwnedEventDomainRun {
1583    fn new(
1584        requested_domain: super::EventDomainId,
1585        engine_id: EngineId,
1586        inner: Box<dyn EventDomainRun>,
1587    ) -> Self {
1588        Self {
1589            requested_domain,
1590            engine_id,
1591            state: RuntimeOwnedEventDomainRunState::Pending(inner),
1592        }
1593    }
1594
1595    /// Release the pending run to its engine, or `None` once it is terminal.
1596    fn into_pending(mut self) -> Option<(EngineId, Box<dyn EventDomainRun>)> {
1597        match std::mem::replace(&mut self.state, RuntimeOwnedEventDomainRunState::Retired) {
1598            RuntimeOwnedEventDomainRunState::Pending(run) => Some((self.engine_id.clone(), run)),
1599            RuntimeOwnedEventDomainRunState::Retired | RuntimeOwnedEventDomainRunState::Failed => {
1600                None
1601            }
1602        }
1603    }
1604
1605    fn requested_domain(&self) -> super::EventDomainId {
1606        self.requested_domain
1607    }
1608
1609    /// Ask the provider run to accept coalesced single-domain submission.
1610    ///
1611    /// A run that keeps the default per-node contract returns `false`, and a
1612    /// retired or failed run is never asked.
1613    fn enable_single_domain_coalescing(&mut self) -> bool {
1614        match &mut self.state {
1615            RuntimeOwnedEventDomainRunState::Pending(run) => run.enable_single_domain_coalescing(),
1616            RuntimeOwnedEventDomainRunState::Retired | RuntimeOwnedEventDomainRunState::Failed => {
1617                false
1618            }
1619        }
1620    }
1621
1622    fn domain(&self, operation: EventDomainOperation) -> Result<super::EventDomainId> {
1623        match &self.state {
1624            RuntimeOwnedEventDomainRunState::Pending(run) => Ok(run.domain()),
1625            RuntimeOwnedEventDomainRunState::Retired => Err(event_domain_run_state_error(
1626                operation,
1627                self.requested_domain,
1628                EventDomainRunTerminalState::Retired,
1629            )),
1630            RuntimeOwnedEventDomainRunState::Failed => Err(event_domain_run_state_error(
1631                operation,
1632                self.requested_domain,
1633                EventDomainRunTerminalState::Failed,
1634            )),
1635        }
1636    }
1637
1638    fn enqueue(
1639        &mut self,
1640        dependencies: &[Arc<dyn EventToken>],
1641        launch: &mut dyn FnMut() -> Result<()>,
1642    ) -> Result<Arc<dyn EventToken>> {
1643        match &mut self.state {
1644            RuntimeOwnedEventDomainRunState::Pending(run) => run.enqueue(dependencies, launch),
1645            RuntimeOwnedEventDomainRunState::Retired => Err(event_domain_run_state_error(
1646                EventDomainOperation::Enqueue,
1647                self.requested_domain,
1648                EventDomainRunTerminalState::Retired,
1649            )),
1650            RuntimeOwnedEventDomainRunState::Failed => Err(event_domain_run_state_error(
1651                EventDomainOperation::Enqueue,
1652                self.requested_domain,
1653                EventDomainRunTerminalState::Failed,
1654            )),
1655        }
1656    }
1657
1658    fn drain(&mut self) -> Result<()> {
1659        let run = match std::mem::replace(&mut self.state, RuntimeOwnedEventDomainRunState::Failed)
1660        {
1661            RuntimeOwnedEventDomainRunState::Pending(run) => run,
1662            RuntimeOwnedEventDomainRunState::Retired => {
1663                self.state = RuntimeOwnedEventDomainRunState::Retired;
1664                return Err(event_domain_run_state_error(
1665                    EventDomainOperation::Drain,
1666                    self.requested_domain,
1667                    EventDomainRunTerminalState::Retired,
1668                ));
1669            }
1670            RuntimeOwnedEventDomainRunState::Failed => {
1671                self.state = RuntimeOwnedEventDomainRunState::Failed;
1672                return Err(event_domain_run_state_error(
1673                    EventDomainOperation::Drain,
1674                    self.requested_domain,
1675                    EventDomainRunTerminalState::Failed,
1676                ));
1677            }
1678        };
1679        let domain = self.requested_domain;
1680        let mut run = run;
1681        let drain_result = std::panic::catch_unwind(AssertUnwindSafe(|| run.drain()));
1682        drop_event_domain_run(run);
1683        match drain_result {
1684            Ok(Ok(())) => {
1685                self.state = RuntimeOwnedEventDomainRunState::Retired;
1686                Ok(())
1687            }
1688            Ok(Err(error)) => {
1689                self.state = RuntimeOwnedEventDomainRunState::Failed;
1690                Err(error)
1691            }
1692            Err(payload) => {
1693                self.state = RuntimeOwnedEventDomainRunState::Failed;
1694                Err(event_domain_error(EventDomainError::DrainPanicked {
1695                    operation: EventDomainOperation::Drain,
1696                    domain,
1697                    message: safe_event_domain_panic_message(payload),
1698                }))
1699            }
1700        }
1701    }
1702}
1703
1704/// Internal error returned when a runtime-owned event run is used after
1705/// explicit retirement or failure.
1706#[derive(Debug, thiserror::Error)]
1707#[error("{operation} used event-domain run {domain:?} after it reached terminal state {state}")]
1708pub(crate) struct EventDomainRunLifecycleError {
1709    operation: EventDomainOperation,
1710    domain: super::EventDomainId,
1711    state: EventDomainRunTerminalState,
1712}
1713
1714fn event_domain_run_state_error(
1715    operation: EventDomainOperation,
1716    domain: super::EventDomainId,
1717    state: EventDomainRunTerminalState,
1718) -> Error {
1719    Error::runtime_state_source(
1720        "Runtime::run",
1721        ErrorPhase::Execution,
1722        EventDomainRunLifecycleError {
1723            operation,
1724            domain,
1725            state,
1726        },
1727    )
1728}
1729
1730fn drop_event_domain_run(run: Box<dyn EventDomainRun>) {
1731    if std::panic::catch_unwind(AssertUnwindSafe(|| drop(run))).is_err() {
1732        // The provider run has been consumed; its cleanup panic is contained.
1733    }
1734}
1735
1736impl Drop for RuntimeOwnedEventDomainRun {
1737    fn drop(&mut self) {
1738        let state = std::mem::replace(&mut self.state, RuntimeOwnedEventDomainRunState::Failed);
1739        if let RuntimeOwnedEventDomainRunState::Pending(run) = state {
1740            drop_event_domain_run(run);
1741        }
1742    }
1743}
1744
1745fn missing_event_domain_driver(domain: super::EventDomainId) -> Error {
1746    Error::from(EventDomainError::MissingDriver { domain })
1747}
1748
1749fn event_domain_error(source: EventDomainError) -> Error {
1750    Error::from(source)
1751}
1752
1753fn safe_event_domain_panic_message(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
1754    match payload.downcast::<&'static str>() {
1755        Ok(message) => (*message).to_owned(),
1756        Err(payload) => match payload.downcast::<String>() {
1757            Ok(message) => *message,
1758            Err(_) => "non-string panic payload".to_owned(),
1759        },
1760    }
1761}
1762
1763fn scheduled_execution_cleanup_error(primary: Error, cleanup: Error) -> Error {
1764    Error::with_suppressed(primary, cleanup)
1765}
1766
1767fn instruction_execution<'a>(
1768    schedule: &'a ScheduledGraph,
1769    instruction: &ExecInstruction,
1770) -> Result<InstructionExecution<'a>> {
1771    let Some(operation_index) = instruction.semantic_operation_index else {
1772        let location = schedule.root_location();
1773        return Ok(InstructionExecution { location });
1774    };
1775    let location = schedule
1776        .operation_locations()
1777        .get(operation_index)
1778        .ok_or_else(|| {
1779            Error::runtime_state(
1780                "Runtime::run",
1781                ErrorPhase::Execution,
1782                format!(
1783                    "instruction references semantic operation {operation_index}, but prepared schedule has {} operations",
1784                    schedule.operation_locations().len()
1785                ),
1786            )
1787        })?;
1788    Ok(InstructionExecution { location })
1789}
1790
1791struct InstructionExecution<'a> {
1792    location: &'a ExecutionLocation,
1793}
1794
1795impl<'a> InstructionExecution<'a> {
1796    fn location(&self) -> &'a ExecutionLocation {
1797        self.location
1798    }
1799}
1800
1801pub(crate) struct LocatedExecSlot<'input> {
1802    pub(crate) location: ExecutionLocation,
1803    pub(crate) value: ExecSlot<'input>,
1804}
1805
1806fn stage_instruction_inputs<'input>(
1807    instruction: &ExecInstruction,
1808    location: &ExecutionLocation,
1809    located: &mut [Vec<LocatedExecSlot<'input>>],
1810    staged: &mut [Option<ExecSlot<'input>>],
1811) -> Result<()> {
1812    for &slot in &instruction.input_slots {
1813        stage_slot_input(slot, location, located, staged)?;
1814    }
1815    Ok(())
1816}
1817
1818/// Move one slot's value from `located` into `staged`.
1819fn stage_slot_input<'input>(
1820    slot: usize,
1821    location: &ExecutionLocation,
1822    located: &mut [Vec<LocatedExecSlot<'input>>],
1823    staged: &mut [Option<ExecSlot<'input>>],
1824) -> Result<()> {
1825    if staged
1826        .get(slot)
1827        .ok_or(tenferro_tensor::Error::MissingValue { slot })?
1828        .is_some()
1829    {
1830        return Ok(());
1831    }
1832    let values = located
1833        .get_mut(slot)
1834        .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1835    let value_index = values
1836        .iter()
1837        .position(|value| &value.location == location)
1838        .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1839    staged[slot] = Some(values.swap_remove(value_index).value);
1840    Ok(())
1841}
1842
1843/// Validate every live-out of a fused elementwise region: present and resident
1844/// on the region's location.
1845fn validate_region_outputs(
1846    region: &super::region::ElementwiseRegion,
1847    location: &ExecutionLocation,
1848    staged: &[Option<ExecSlot<'_>>],
1849) -> Result<()> {
1850    for &output_slot in &region.output_slots {
1851        let output = staged
1852            .get(output_slot)
1853            .and_then(Option::as_ref)
1854            .ok_or(tenferro_tensor::Error::MissingValue { slot: output_slot })?;
1855        let output = output.as_read();
1856        if !location
1857            .witness()
1858            .owns_resident_tensor(&output, location.storage_class())
1859        {
1860            return Err(Error::runtime_state_source(
1861                "Runtime::run",
1862                ErrorPhase::Execution,
1863                super::EngineExecutionContractError::OutputResidencyMismatch {
1864                    instruction_index: region.instruction_range.start,
1865                    output_slot,
1866                    engine_id: location.engine_id().clone(),
1867                    storage_class: location.storage_class().clone(),
1868                    backend_family: output.backend_family(),
1869                    allocation_domain: output.allocation_domain(),
1870                },
1871            ));
1872        }
1873    }
1874    Ok(())
1875}
1876
1877/// Publish a fused region's results: release region inputs whose last use is
1878/// inside the region and move every live-out into `located`.
1879fn retain_region_results<'input>(
1880    region: &super::region::ElementwiseRegion,
1881    location: &ExecutionLocation,
1882    located: &mut [Vec<LocatedExecSlot<'input>>],
1883    staged: &mut [Option<ExecSlot<'input>>],
1884) -> Result<()> {
1885    for (&slot, &last_use) in region.input_slots.iter().zip(region.input_last_use.iter()) {
1886        if last_use {
1887            located[slot].clear();
1888            staged[slot].take();
1889        }
1890    }
1891    for &slot in &region.output_slots {
1892        let value = staged
1893            .get_mut(slot)
1894            .and_then(Option::take)
1895            .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1896        let values = located
1897            .get_mut(slot)
1898            .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1899        values.clear();
1900        values.push(LocatedExecSlot {
1901            location: location.clone(),
1902            value,
1903        });
1904    }
1905    Ok(())
1906}
1907
1908fn validate_instruction_outputs(
1909    instruction_index: usize,
1910    instruction: &ExecInstruction,
1911    location: &ExecutionLocation,
1912    staged: &[Option<ExecSlot<'_>>],
1913) -> Result<()> {
1914    for &output_slot in &instruction.output_slots {
1915        let output = staged
1916            .get(output_slot)
1917            .and_then(Option::as_ref)
1918            .ok_or(tenferro_tensor::Error::MissingValue { slot: output_slot })?;
1919        let output = output.as_read();
1920        if !location
1921            .witness()
1922            .owns_resident_tensor(&output, location.storage_class())
1923        {
1924            return Err(Error::runtime_state_source(
1925                "Runtime::run",
1926                ErrorPhase::Execution,
1927                super::EngineExecutionContractError::OutputResidencyMismatch {
1928                    instruction_index,
1929                    output_slot,
1930                    engine_id: location.engine_id().clone(),
1931                    storage_class: location.storage_class().clone(),
1932                    backend_family: output.backend_family(),
1933                    allocation_domain: output.allocation_domain(),
1934                },
1935            ));
1936        }
1937    }
1938    Ok(())
1939}
1940
1941pub(crate) fn retain_instruction_results<'input>(
1942    instruction: &ExecInstruction,
1943    location: &ExecutionLocation,
1944    located: &mut [Vec<LocatedExecSlot<'input>>],
1945    staged: &mut [Option<ExecSlot<'input>>],
1946) -> Result<()> {
1947    for &slot in &instruction.input_slots {
1948        let is_output = instruction.output_slots.contains(&slot);
1949        let is_last_use = instruction
1950            .input_slots
1951            .iter()
1952            .enumerate()
1953            .any(|(index, &candidate)| {
1954                candidate == slot && instruction.last_use.get(index).copied().unwrap_or(false)
1955            });
1956        if is_last_use {
1957            located[slot].clear();
1958            if !is_output {
1959                staged[slot].take();
1960            }
1961        } else if !is_output && let Some(value) = staged[slot].take() {
1962            located[slot].push(LocatedExecSlot {
1963                location: location.clone(),
1964                value,
1965            });
1966        }
1967    }
1968
1969    for &slot in &instruction.output_slots {
1970        let value = staged
1971            .get_mut(slot)
1972            .and_then(Option::take)
1973            .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1974        let values = located
1975            .get_mut(slot)
1976            .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1977        values.clear();
1978        values.push(LocatedExecSlot {
1979            location: location.clone(),
1980            value,
1981        });
1982    }
1983    Ok(())
1984}
1985
1986fn validate_runtime_input_ingress(
1987    location: &ExecutionLocation,
1988    input: &TensorRead<'_>,
1989    slot: usize,
1990) -> Result<()> {
1991    let accepted = location
1992        .witness()
1993        .accepts_runtime_input(input, location.storage_class());
1994    if accepted {
1995        return Ok(());
1996    }
1997    Err(Error::runtime_state_source(
1998        "Runtime::run",
1999        ErrorPhase::Execution,
2000        super::InputIngressContractError::ResidencyMismatch {
2001            input_slot: slot,
2002            ingress_engine_id: location.engine_id().clone(),
2003            ingress_storage_class: location.storage_class().clone(),
2004            placement: input.placement().clone(),
2005            backend_family: input.backend_family(),
2006            allocation_domain: input.allocation_domain(),
2007        },
2008    ))
2009}
2010
2011fn execute_scheduled_transfer<'input>(
2012    transfer: &ScheduledTransfer,
2013    located: &mut [Vec<LocatedExecSlot<'input>>],
2014) -> Result<()> {
2015    let source = transfer.source_location();
2016    let destination = transfer.destination_location();
2017    let provider = transfer.provider();
2018    let values =
2019        located
2020            .get_mut(transfer.value_slot())
2021            .ok_or(tenferro_tensor::Error::MissingValue {
2022                slot: transfer.value_slot(),
2023            })?;
2024    if values.iter().any(|value| &value.location == destination) {
2025        return Err(Error::runtime_state_source(
2026            "Runtime::run",
2027            ErrorPhase::Execution,
2028            DuplicateTransferDestinationError {
2029                value_slot: transfer.value_slot(),
2030                destination: destination.clone(),
2031            },
2032        ));
2033    }
2034    let transferred = {
2035        let source_value = values
2036            .iter()
2037            .find(|value| &value.location == source)
2038            .ok_or(tenferro_tensor::Error::MissingValue {
2039                slot: transfer.value_slot(),
2040            })?;
2041        let source_read = source_value.value.as_read();
2042        let expected_dtype = source_read.dtype();
2043        let expected_shape = source_read.shape().to_vec();
2044        let transferred =
2045            provider.transfer_blocking(TransferRequest::new(source, destination, source_read))?;
2046        validate_transfer_output(destination, expected_dtype, &expected_shape, &transferred)?;
2047        transferred
2048    };
2049    values.push(LocatedExecSlot {
2050        location: destination.clone(),
2051        value: ExecSlot::Owned(transferred),
2052    });
2053    Ok(())
2054}
2055
2056fn validate_transfer_output(
2057    destination: &ExecutionLocation,
2058    expected_dtype: tenferro_tensor::DType,
2059    expected_shape: &[usize],
2060    output: &Tensor,
2061) -> Result<()> {
2062    let expected_elements = checked_transfer_element_count(expected_shape).map_err(|source| {
2063        Error::runtime_state_source(
2064            "Runtime::run",
2065            ErrorPhase::Execution,
2066            TransferError::ProviderContract { source },
2067        )
2068    })?;
2069    let contract_error = if output.dtype() != expected_dtype {
2070        Some(TransferProviderContractError::DTypeMismatch {
2071            expected: expected_dtype,
2072            actual: output.dtype(),
2073        })
2074    } else if output.shape() != expected_shape {
2075        Some(TransferProviderContractError::ShapeMismatch {
2076            expected: expected_shape.to_vec(),
2077            actual: output.shape().to_vec(),
2078        })
2079    } else if tensor_buffer_len(output) != expected_elements {
2080        Some(TransferProviderContractError::InvalidBufferLength {
2081            expected: expected_elements,
2082            actual: tensor_buffer_len(output),
2083        })
2084    } else if !destination
2085        .witness()
2086        .accepts_input_placement(output.placement(), destination.storage_class())
2087    {
2088        Some(
2089            TransferProviderContractError::DestinationPlacementMismatch {
2090                destination_engine_id: destination.engine_id().clone(),
2091                destination_storage_class: destination.storage_class().clone(),
2092                actual: output.placement().clone(),
2093            },
2094        )
2095    } else if !destination.witness().owns_resident_tensor(
2096        &TensorRead::from_tensor(output),
2097        destination.storage_class(),
2098    ) {
2099        Some(
2100            TransferProviderContractError::DestinationResidencyMismatch {
2101                destination_engine_id: destination.engine_id().clone(),
2102                destination_storage_class: destination.storage_class().clone(),
2103                actual_backend_family: TensorRead::from_tensor(output).backend_family(),
2104                actual_allocation_domain: TensorRead::from_tensor(output).allocation_domain(),
2105            },
2106        )
2107    } else {
2108        None
2109    };
2110    match contract_error {
2111        None => Ok(()),
2112        Some(source) => Err(Error::runtime_state_source(
2113            "Runtime::run",
2114            ErrorPhase::Execution,
2115            TransferError::ProviderContract { source },
2116        )),
2117    }
2118}
2119
2120fn checked_transfer_element_count(
2121    shape: &[usize],
2122) -> std::result::Result<usize, TransferProviderContractError> {
2123    tenferro_tensor::validate::checked_shape_product("Runtime::run", "transfer source shape", shape)
2124        .map_err(|source| TransferProviderContractError::LogicalElementCount { source })
2125}
2126
2127fn tensor_buffer_len(tensor: &Tensor) -> usize {
2128    match tensor.dtype() {
2129        DType::F32 => tensor.as_typed::<f32>().map_or(0, |t| t.buffer().len()),
2130        DType::F64 => tensor.as_typed::<f64>().map_or(0, |t| t.buffer().len()),
2131        DType::I32 => tensor.as_typed::<i32>().map_or(0, |t| t.buffer().len()),
2132        DType::I64 => tensor.as_typed::<i64>().map_or(0, |t| t.buffer().len()),
2133        DType::Bool => tensor.as_typed::<bool>().map_or(0, |t| t.buffer().len()),
2134        DType::C32 => tensor
2135            .as_typed::<Complex32>()
2136            .map_or(0, |t| t.buffer().len()),
2137        DType::C64 => tensor
2138            .as_typed::<Complex64>()
2139            .map_or(0, |t| t.buffer().len()),
2140        // A caller-owned payload has no runtime buffer length to report, which is
2141        // also what the accessor falls back to when the tag table and the runtime
2142        // dtype disagree.
2143        DType::External(_) => 0,
2144    }
2145}
2146
2147#[cfg(test)]
2148mod transfer_validation_tests {
2149    use std::error::Error as _;
2150
2151    use super::checked_transfer_element_count;
2152    use crate::TransferProviderContractError;
2153
2154    #[test]
2155    fn transfer_element_count_overflow_is_typed_and_preserves_source() {
2156        let error = checked_transfer_element_count(&[usize::MAX, 2]).unwrap_err();
2157
2158        assert!(matches!(
2159            error,
2160            TransferProviderContractError::LogicalElementCount { .. }
2161        ));
2162        assert!(error.source().is_some());
2163    }
2164}
2165
2166fn input_signature_reads(inputs: &[TensorRead<'_>]) -> Result<InputSignature> {
2167    InputSignature::from_reads(inputs).map_err(|source| prepare_error(Arc::new(source)))
2168}
2169
2170/// Whether [`resolve_input_reads`] must re-validate the graph's retained inputs.
2171///
2172/// A retained input's value is the program's own tensor and cannot change
2173/// between calls, so a check that already ran on it need not run again. That is
2174/// a fact about the resolution, not about preparedness:
2175/// `Prepared::defaults_validated` records whether preparation
2176/// itself checked these values, and only then may an execution skip the check.
2177#[derive(Clone, Copy, Debug, Eq, PartialEq)]
2178enum BoundInputCheck {
2179    Validate,
2180    AlreadyValidatedAtPrepare,
2181}
2182
2183fn resolve_input_reads<'a>(
2184    program: &'a CompiledGraph,
2185    inputs: &[TensorRead<'a>],
2186    bound_inputs: BoundInputCheck,
2187) -> Result<RuntimeInputReads<'a>> {
2188    let expected = program.program().inputs().len();
2189    if inputs.len() > expected {
2190        return Err(Error::GraphInputCountMismatch {
2191            expected,
2192            actual: inputs.len(),
2193        });
2194    }
2195    if inputs.is_empty() {
2196        let resolved = semantic_default_inputs(program)?;
2197        if matches!(bound_inputs, BoundInputCheck::AlreadyValidatedAtPrepare) {
2198            return Ok(resolved);
2199        }
2200        validate_ordered_input_metadata_reads(program, &resolved)?;
2201        return Ok(resolved);
2202    }
2203    let resolved = if inputs.len() == expected {
2204        inputs.iter().cloned().collect()
2205    } else {
2206        let mut explicit = inputs.iter();
2207        let mut resolved = RuntimeInputReads::new();
2208        for value in program.program().inputs() {
2209            if let Some(retained) = program.bindings().tensor_ref_for_input(*value) {
2210                resolved.push(retained.tensor_read()?);
2211            } else {
2212                let read = explicit.next().ok_or_else(|| {
2213                    Error::invalid_argument(
2214                        "Runtime::run",
2215                        ErrorPhase::Execution,
2216                        "inputs",
2217                        format!(
2218                            "expected {expected} inputs, or one explicit input for each unbound placeholder"
2219                        ),
2220                    )
2221                })?;
2222                resolved.push(read.clone());
2223            }
2224        }
2225        if explicit.next().is_some() {
2226            return Err(Error::invalid_argument(
2227                "Runtime::run",
2228                ErrorPhase::Execution,
2229                "inputs",
2230                format!("expected {expected} inputs, received {}", inputs.len()),
2231            ));
2232        }
2233        resolved
2234    };
2235    validate_ordered_input_metadata_reads(program, &resolved)?;
2236    Ok(resolved)
2237}
2238
2239fn semantic_default_inputs(program: &CompiledGraph) -> Result<RuntimeInputReads<'_>> {
2240    program
2241        .program()
2242        .inputs()
2243        .iter()
2244        .enumerate()
2245        .map(|(input_index, value)| {
2246            let retained = program
2247                .bindings()
2248                .tensor_ref_for_input(*value)
2249                .ok_or_else(|| Error::UnboundPlaceholder {
2250                    input_key: format!("semantic input {input_index}"),
2251                })?;
2252            retained.tensor_read()
2253        })
2254        .collect()
2255}
2256
2257fn validate_ordered_input_metadata_reads(
2258    program: &CompiledGraph,
2259    inputs: &[TensorRead<'_>],
2260) -> Result<()> {
2261    let actuals = inputs
2262        .iter()
2263        .map(|input| (input.dtype(), input.shape().to_vec()))
2264        .collect::<Vec<_>>();
2265    validate_ordered_input_metadata_values(program, &actuals, "Runtime::submit")
2266}
2267
2268fn validate_ordered_input_metadata_values(
2269    program: &CompiledGraph,
2270    actuals: &[(tenferro_tensor::DType, Vec<usize>)],
2271    caller: &'static str,
2272) -> Result<()> {
2273    let expected = program.input_count();
2274    if actuals.len() != expected {
2275        return Err(Error::GraphInputCountMismatch {
2276            expected,
2277            actual: actuals.len(),
2278        });
2279    }
2280    let input_shapes: RuntimeInputShapes<'_> =
2281        actuals.iter().map(|(_, shape)| shape.as_slice()).collect();
2282    for (input_value, (actual_dtype, actual_shape)) in
2283        program.program().inputs().iter().zip(actuals)
2284    {
2285        let metadata = program
2286            .program()
2287            .value_metadata(*input_value)
2288            .map_err(|source| Error::runtime_state_source(caller, ErrorPhase::Execution, source))?;
2289        if metadata.dtype() != *actual_dtype {
2290            return Err(Error::PlaceholderDtypeMismatch {
2291                expected: metadata.dtype(),
2292                actual: *actual_dtype,
2293            });
2294        }
2295        if metadata.shape().len() != actual_shape.len() {
2296            return Err(Error::PlaceholderRankMismatch {
2297                expected: metadata.shape().len(),
2298                actual: actual_shape.len(),
2299            });
2300        }
2301        let mut expected_shape: RuntimeShapeScratch = actual_shape.iter().copied().collect();
2302        let mut exact_mismatch = false;
2303        for (axis, (extent, actual_size)) in metadata.shape().iter().zip(actual_shape).enumerate() {
2304            match extent {
2305                ShapeExtent::Exact(expression) => {
2306                    let expected = expression.eval(&input_shapes).map_err(|source| {
2307                        Error::runtime_state_source(caller, ErrorPhase::Execution, source)
2308                    })?;
2309                    expected_shape[axis] = expected;
2310                    exact_mismatch |= expected != *actual_size;
2311                }
2312                ShapeExtent::UpperBound(expression) => {
2313                    let bound = expression.eval(&input_shapes).map_err(|source| {
2314                        Error::runtime_state_source(caller, ErrorPhase::Execution, source)
2315                    })?;
2316                    if *actual_size > bound {
2317                        return Err(Error::PlaceholderShapeBoundExceeded {
2318                            axis,
2319                            bound,
2320                            actual: *actual_size,
2321                        });
2322                    }
2323                }
2324                ShapeExtent::Unknown => {}
2325            }
2326        }
2327        if exact_mismatch {
2328            return Err(Error::PlaceholderShapeMismatch {
2329                expected: expected_shape.into_vec(),
2330                actual: actual_shape.clone(),
2331            });
2332        }
2333    }
2334    Ok(())
2335}
2336
2337fn validate_exec_input_count(program: &ExecProgram, actual: usize) -> Result<()> {
2338    let expected = program.input_slots.len();
2339    if actual != expected {
2340        return Err(Error::runtime_state(
2341            "Runtime::run",
2342            ErrorPhase::Execution,
2343            format!("expected {expected} inputs for execution program, got {actual}"),
2344        ));
2345    }
2346    Ok(())
2347}
2348
2349fn prepare_error(source: Arc<PrepareError>) -> Error {
2350    Error::runtime_state_source(
2351        "Runtime::run",
2352        ErrorPhase::Execution,
2353        SharedPrepareError(source),
2354    )
2355}
2356
2357#[derive(Clone, Debug)]
2358struct SharedPrepareError(Arc<PrepareError>);
2359
2360impl fmt::Display for SharedPrepareError {
2361    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
2362        write!(formatter, "{}", self.0)
2363    }
2364}
2365
2366impl StdError for SharedPrepareError {
2367    fn source(&self) -> Option<&(dyn StdError + 'static)> {
2368        Some(self.0.as_ref())
2369    }
2370}