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