1use std::collections::{HashMap, HashSet};
7use std::error::Error as StdError;
8use std::fmt;
9use std::panic::{catch_unwind, AssertUnwindSafe};
10use std::sync::{Arc, Condvar, Mutex, MutexGuard};
11use std::thread::{self, ThreadId};
12
13use smallvec::SmallVec;
14use tenferro_ops::shape_extent::ShapeExtent;
15use tenferro_tensor::{
16 AllocationGroup, DType, DescriptorSlot, GroupError, MemoryKind, Tensor, TensorBackend,
17 TensorRead, TensorValue, TensorView,
18};
19
20use crate::error::ErrorPhase;
21use crate::exec::{
22 DispatchMode, ExecInstruction, ExecProgram, ExecSlot, ExtensionExecutionDispatch,
23};
24use crate::extension_cache::{ExtensionCacheSelector, ExtensionCacheStore};
25use crate::graph::CompiledGraph;
26use crate::runtime::schedule::{
27 EventDependency, ExecutionLocation, ScheduledGraph, ScheduledNode, ScheduledNodeKind,
28 ScheduledTransfer, UnsupportedScheduledNodeError,
29};
30use crate::runtime::{
31 CacheOwnerError, CacheStats, EventDomainError, EventDomainOperation, EventDomainRun,
32 EventToken, InputSignature, PrepareError, PrepareOptions, PreparedOperationPlan, Runtime,
33 RuntimeCacheOwner, SubmissionError, TransferError, TransferProviderContractError,
34 TransferRequest,
35};
36use crate::{Error, Result};
37
38type RuntimeInputReads<'a> = SmallVec<[TensorRead<'a>; 8]>;
39type RuntimeInputShapes<'a> = SmallVec<[&'a [usize]; 8]>;
40type RuntimeShapeScratch = SmallVec<[usize; 8]>;
41
42#[derive(Clone, Copy, Debug)]
43pub(super) enum RuntimeOutputMode {
44 Tensor,
45 Value,
46}
47
48#[derive(Clone)]
54pub struct PreparedCompiledGraph {
55 runtime_id: super::RuntimeId,
56 epoch: super::RuntimeEpoch,
57 program: CompiledGraph,
58 prepared: Arc<super::preparation::PreparedProgram>,
59}
60
61pub struct ExecutionHandle {
63 submission: Arc<InFlightSubmission>,
64}
65
66pub struct ExecutionInputs {
78 group: AllocationGroup,
79 bindings: Box<[DescriptorSlot]>,
80}
81
82impl ExecutionInputs {
83 pub fn new(tensors: Vec<Tensor>) -> Result<Self> {
90 let (group, bindings) = AllocationGroup::from_tensors(tensors).map_err(|error| {
91 Error::runtime_state(
92 "ExecutionInputs::new",
93 ErrorPhase::Execution,
94 error.to_string(),
95 )
96 })?;
97 Ok(Self { group, bindings })
98 }
99
100 pub(crate) fn as_reads(&self) -> Result<Vec<TensorRead<'_>>> {
101 self.group.read_views(&self.bindings).map_err(|error| {
102 Error::runtime_state(
103 "ExecutionInputs::as_reads",
104 ErrorPhase::Execution,
105 error.to_string(),
106 )
107 })
108 }
109}
110
111impl fmt::Debug for ExecutionInputs {
112 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
113 formatter
114 .debug_struct("ExecutionInputs")
115 .field("len", &self.bindings.len())
116 .finish()
117 }
118}
119
120pub struct ScopedReadInputs<'env> {
122 bindings: Box<[ScopedReadBinding<'env>]>,
123}
124
125impl<'env> ScopedReadInputs<'env> {
126 pub fn new(bindings: Vec<TensorView<'env>>) -> Self {
128 Self {
129 bindings: bindings
130 .into_iter()
131 .map(|tensor| ScopedReadBinding { tensor })
132 .collect(),
133 }
134 }
135
136 fn as_reads(&self) -> Vec<TensorRead<'env>> {
137 self.bindings
138 .iter()
139 .map(|binding| TensorRead::from_view(binding.tensor.clone()))
140 .collect()
141 }
142
143 fn has_non_host_provider(&self) -> bool {
144 self.bindings.iter().any(|binding| {
145 !matches!(
146 binding.tensor.placement().memory_kind,
147 MemoryKind::PinnedHost | MemoryKind::UnpinnedHost
148 )
149 })
150 }
151
152 pub fn len(&self) -> usize {
154 self.bindings.len()
155 }
156
157 pub fn is_empty(&self) -> bool {
159 self.bindings.is_empty()
160 }
161}
162
163pub struct ScopedReadBinding<'env> {
165 tensor: TensorView<'env>,
166}
167
168impl<'env> ScopedReadBinding<'env> {
169 pub fn new(tensor: TensorView<'env>) -> Self {
171 Self { tensor }
172 }
173
174 pub fn tensor(&self) -> &TensorView<'env> {
176 &self.tensor
177 }
178}
179
180impl fmt::Debug for ScopedReadInputs<'_> {
181 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
182 formatter
183 .debug_struct("ScopedReadInputs")
184 .field("len", &self.bindings.len())
185 .finish()
186 }
187}
188
189#[allow(clippy::large_enum_variant)]
193#[derive(Debug)]
194pub enum ScopedExecutionOutcome<'env> {
195 Completed(ScopedExecutionBundle<'env>),
197 RetiredFailed {
199 error: Error,
200 inputs: ScopedReadInputs<'env>,
201 },
202}
203
204#[derive(Clone, Debug, PartialEq, Eq)]
206pub struct OutputMetadata {
207 pub dtype: DType,
209 pub shape: Box<[usize]>,
211}
212
213pub struct ScopedExecutionBundle<'env> {
215 owned: AllocationGroup,
216 outputs: Box<[ScopedOutput<'env>]>,
217}
218
219impl fmt::Debug for ScopedExecutionBundle<'_> {
220 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
221 formatter
222 .debug_struct("ScopedExecutionBundle")
223 .field("len", &self.outputs.len())
224 .finish()
225 }
226}
227
228impl<'env> ScopedExecutionBundle<'env> {
229 pub fn output(&self, index: usize) -> std::result::Result<OutputRef<'_>, OutputAccessError> {
237 let output = self
238 .outputs
239 .get(index)
240 .ok_or(OutputAccessError::InvalidOutput { index })?;
241 match output {
242 ScopedOutput::Borrowed(view) => Ok(OutputRef::Tensor(view.clone())),
243 ScopedOutput::Owned(slot) => output_ref_from_group(&self.owned, *slot),
244 ScopedOutput::Metadata(metadata) => Ok(OutputRef::Metadata(metadata)),
245 }
246 }
247
248 #[allow(clippy::result_large_err)]
258 pub fn into_owned_output(
259 self,
260 index: usize,
261 ) -> std::result::Result<Tensor, (Self, ScopedOutputExtractError)> {
262 let slot = match self.outputs.get(index) {
263 Some(ScopedOutput::Owned(slot)) => *slot,
264 Some(ScopedOutput::Borrowed(_)) => {
265 return Err((self, ScopedOutputExtractError::BorrowedOutput))
266 }
267 Some(ScopedOutput::Metadata(_)) => {
268 return Err((self, ScopedOutputExtractError::MetadataOutput))
269 }
270 None => return Err((self, ScopedOutputExtractError::InvalidOutput { index })),
271 };
272 let Self { owned, outputs } = self;
273 match owned.into_tensor(slot) {
274 Ok(tensor) => Ok(tensor),
275 Err((owned, error)) => Err((
276 Self { owned, outputs },
277 ScopedOutputExtractError::Output(OutputExtractError::Group(error)),
278 )),
279 }
280 }
281}
282
283#[allow(clippy::large_enum_variant)]
287#[derive(Debug)]
288pub enum ScopedOutput<'env> {
289 Borrowed(TensorView<'env>),
291 Owned(DescriptorSlot),
293 Metadata(OutputMetadata),
295}
296
297#[derive(Debug, thiserror::Error)]
299pub enum ScopedOutputExtractError {
300 #[error("scoped output index {index} is outside the output set")]
301 InvalidOutput { index: usize },
302 #[error("scoped output is borrowed from the caller")]
303 BorrowedOutput,
304 #[error("scoped output contains metadata only")]
305 MetadataOutput,
306 #[error("scoped owned output cannot be extracted: {0}")]
307 Output(#[from] OutputExtractError),
308}
309
310#[derive(Debug)]
312pub struct ScopedSubmitRejected<'env> {
313 source: Box<Error>,
314 inputs: ScopedReadInputs<'env>,
315}
316
317impl<'env> ScopedSubmitRejected<'env> {
318 pub fn into_parts(self) -> (Error, ScopedReadInputs<'env>) {
320 (*self.source, self.inputs)
321 }
322}
323
324impl fmt::Display for ScopedSubmitRejected<'_> {
325 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
326 write!(
327 formatter,
328 "scoped submission rejected before admission: {}",
329 self.source
330 )
331 }
332}
333
334impl StdError for ScopedSubmitRejected<'_> {
335 fn source(&self) -> Option<&(dyn StdError + 'static)> {
336 Some(self.source.as_ref())
337 }
338}
339
340#[derive(Debug)]
342pub struct ExecutionBundle {
343 group: AllocationGroup,
344 outputs: Box<[DescriptorSlot]>,
345}
346
347#[allow(clippy::large_enum_variant)]
351#[derive(Debug)]
352pub enum OutputRef<'a> {
353 Tensor(TensorView<'a>),
354 Metadata(&'a OutputMetadata),
355}
356
357#[derive(Debug, thiserror::Error)]
358pub enum OutputAccessError {
359 #[error("execution output index {index} is outside the output set")]
360 InvalidOutput { index: usize },
361 #[error("execution output group is invalid: {0}")]
362 Group(#[from] GroupError),
363}
364
365#[derive(Debug, thiserror::Error)]
366pub enum OutputExtractError {
367 #[error("execution output index {index} is outside the output set")]
368 InvalidOutput { index: usize },
369 #[error("execution output cannot be extracted: {0}")]
370 Group(#[from] GroupError),
371}
372
373impl ExecutionBundle {
374 fn from_inputs_and_outputs(inputs: ExecutionInputs, outputs: Vec<Tensor>) -> Result<Self> {
375 let mut group = inputs.group;
376 let mut output_slots = Vec::with_capacity(outputs.len());
377 for output in outputs {
378 let slot = group.append_tensor(output).map_err(|error| {
379 Error::runtime_state(
380 "ExecutionBundle::from_inputs_and_outputs",
381 ErrorPhase::Execution,
382 error.to_string(),
383 )
384 })?;
385 output_slots.push(slot);
386 }
387 Ok(Self {
388 group,
389 outputs: output_slots.into_boxed_slice(),
390 })
391 }
392
393 #[cfg(test)]
394 pub(super) fn from_outputs(outputs: Vec<Tensor>) -> Result<Self> {
395 let (group, bindings) = AllocationGroup::from_tensors(outputs).map_err(|error| {
396 Error::runtime_state(
397 "ExecutionBundle::from_outputs",
398 ErrorPhase::Execution,
399 error.to_string(),
400 )
401 })?;
402 Ok(Self {
403 group,
404 outputs: bindings,
405 })
406 }
407
408 pub fn output(&self, index: usize) -> std::result::Result<OutputRef<'_>, OutputAccessError> {
414 let slot = *self
415 .outputs
416 .get(index)
417 .ok_or(OutputAccessError::InvalidOutput { index })?;
418 let mut views = self.group.read_views(std::slice::from_ref(&slot))?;
419 let view = views
420 .pop()
421 .ok_or(OutputAccessError::InvalidOutput { index })?;
422 match view {
423 TensorRead::View(view) => Ok(OutputRef::Tensor(view)),
424 TensorRead::Tensor(_) => Err(OutputAccessError::InvalidOutput { index }),
425 }
426 }
427
428 #[allow(clippy::result_large_err)]
437 pub fn into_output(
438 self,
439 index: usize,
440 ) -> std::result::Result<Tensor, (Self, OutputExtractError)> {
441 let slot = match self.outputs.get(index).copied() {
442 Some(slot) => slot,
443 None => return Err((self, OutputExtractError::InvalidOutput { index })),
444 };
445 let ExecutionBundle { group, outputs } = self;
446 match group.into_tensor(slot) {
447 Ok(tensor) => Ok(tensor),
448 Err((group, error)) => Err((Self { group, outputs }, OutputExtractError::Group(error))),
449 }
450 }
451}
452
453#[derive(Debug)]
455pub enum ExecutionOutcome {
456 Completed(Box<ExecutionBundle>),
458 RetiredFailed {
461 error: Error,
462 inputs: Box<ExecutionInputs>,
463 },
464 CompletionUnproven {
467 error: Error,
468 diagnostic_keys: Box<[String]>,
469 },
470}
471
472#[derive(Debug)]
488pub enum SubmitError {
489 PreAdmission {
492 source: Box<Error>,
493 inputs: Box<ExecutionInputs>,
494 },
495}
496
497impl SubmitError {
498 pub fn into_pre_admission(self) -> Option<(Error, ExecutionInputs)> {
513 match self {
514 Self::PreAdmission { source, inputs } => Some((*source, *inputs)),
515 }
516 }
517}
518
519impl fmt::Display for SubmitError {
520 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
521 match self {
522 Self::PreAdmission { source, .. } => {
523 write!(formatter, "submission rejected before admission: {source}")
524 }
525 }
526 }
527}
528
529impl StdError for SubmitError {
530 fn source(&self) -> Option<&(dyn StdError + 'static)> {
531 match self {
532 Self::PreAdmission { source, .. } => {
533 source.source().or(Some(source.as_ref()))
537 }
538 }
539 }
540}
541
542impl From<SubmitError> for Error {
543 fn from(error: SubmitError) -> Self {
544 match error {
545 SubmitError::PreAdmission { source, .. } => *source,
546 }
547 }
548}
549
550impl ExecutionHandle {
551 pub fn wait(self) -> Result<ExecutionOutcome> {
559 self.submission.wait()
560 }
561}
562
563impl fmt::Debug for ExecutionHandle {
564 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
565 formatter
566 .debug_struct("ExecutionHandle")
567 .field("pending", &self.submission.is_pending())
568 .finish_non_exhaustive()
569 }
570}
571
572pub(super) struct InFlightSubmission {
573 work: Mutex<Option<InFlightWork>>,
574 completion: Mutex<Option<Result<ExecutionOutcome>>>,
575 completed: Condvar,
576}
577
578struct AdmittedExecution {
579 prepared: PreparedCompiledGraph,
580 inputs: ExecutionInputs,
581}
582
583enum InFlightWork {
584 Admitted(Box<AdmittedExecution>),
585 #[cfg(test)]
586 Test {
587 inputs: Box<ExecutionInputs>,
588 work: Box<dyn FnOnce() -> Result<Vec<Tensor>> + Send>,
589 },
590}
591
592impl InFlightSubmission {
593 fn new(prepared: PreparedCompiledGraph, inputs: ExecutionInputs) -> Self {
594 Self {
595 work: Mutex::new(Some(InFlightWork::Admitted(Box::new(AdmittedExecution {
596 prepared,
597 inputs,
598 })))),
599 completion: Mutex::new(None),
600 completed: Condvar::new(),
601 }
602 }
603
604 #[cfg(test)]
605 pub(super) fn for_test(work: impl FnOnce() -> Result<Vec<Tensor>> + Send + 'static) -> Self {
606 Self {
607 work: Mutex::new(Some(InFlightWork::Test {
608 inputs: Box::new(ExecutionInputs::new(Vec::new()).expect("empty test inputs")),
609 work: Box::new(work),
610 })),
611 completion: Mutex::new(None),
612 completed: Condvar::new(),
613 }
614 }
615
616 fn into_unstarted_inputs(self) -> ExecutionInputs {
617 let work = match self.work.into_inner() {
618 Ok(work) => work,
619 Err(poisoned) => poisoned.into_inner(),
620 };
621 match work {
622 Some(InFlightWork::Admitted(admitted)) => admitted.inputs,
623 #[cfg(test)]
624 Some(InFlightWork::Test { inputs, .. }) => *inputs,
625 None => unreachable!("unstarted submission must still contain its owner"),
626 }
627 }
628
629 pub(super) fn run(&self) {
630 let work = match self.work.lock() {
631 Ok(mut work) => work.take(),
632 Err(poisoned) => poisoned.into_inner().take(),
633 };
634 let result = match work {
635 Some(InFlightWork::Admitted(admitted)) => run_admitted_work(admitted),
636 #[cfg(test)]
637 Some(InFlightWork::Test { work, .. }) => {
638 let result = catch_unwind(AssertUnwindSafe(work));
639 match result {
640 Ok(Ok(outputs)) => ExecutionBundle::from_outputs(outputs)
641 .map(|bundle| ExecutionOutcome::Completed(Box::new(bundle))),
642 Ok(Err(error)) => Err(error),
643 Err(payload) => Ok(ExecutionOutcome::CompletionUnproven {
644 error: Error::runtime_state(
645 "ExecutionHandle::wait",
646 ErrorPhase::Execution,
647 panic_payload_message(payload),
648 ),
649 diagnostic_keys: Box::from(["execution.retirement-unproven".to_owned()]),
650 }),
651 }
652 }
653 None => Err(Error::runtime_state(
654 "Runtime::submit",
655 ErrorPhase::Execution,
656 "in-flight submission work was already consumed",
657 )),
658 };
659 match self.completion.lock() {
660 Ok(mut completion) => {
661 *completion = Some(result);
662 }
663 Err(poisoned) => {
664 *poisoned.into_inner() = Some(Err(Error::runtime_state(
665 "ExecutionHandle::wait",
666 ErrorPhase::Execution,
667 "in-flight completion lock poisoned",
668 )));
669 }
670 }
671 self.completed.notify_all();
672 }
673
674 fn wait(&self) -> Result<ExecutionOutcome> {
675 let mut completion = self.completion.lock().map_err(|_| {
676 Error::runtime_state(
677 "ExecutionHandle::wait",
678 ErrorPhase::Execution,
679 "in-flight completion lock poisoned",
680 )
681 })?;
682 loop {
683 if let Some(result) = completion.take() {
684 return result;
685 }
686 completion = self.completed.wait(completion).map_err(|_| {
687 Error::runtime_state(
688 "ExecutionHandle::wait",
689 ErrorPhase::Execution,
690 "in-flight completion lock poisoned while waiting",
691 )
692 })?;
693 }
694 }
695
696 fn is_pending(&self) -> bool {
697 self.completion
698 .lock()
699 .map_or(true, |completion| completion.is_none())
700 }
701}
702
703fn run_admitted_work(admitted: Box<AdmittedExecution>) -> Result<ExecutionOutcome> {
704 let execution = catch_unwind(AssertUnwindSafe(|| {
705 let input_reads = admitted.inputs.as_reads()?;
706 execute_admitted(&admitted.prepared, &input_reads)
707 }));
708 match execution {
709 Ok(Ok(outputs)) => {
710 let AdmittedExecution { inputs, .. } = *admitted;
711 match ExecutionBundle::from_inputs_and_outputs(inputs, outputs) {
712 Ok(bundle) => Ok(ExecutionOutcome::Completed(Box::new(bundle))),
713 Err(error) => Ok(ExecutionOutcome::CompletionUnproven {
714 error,
715 diagnostic_keys: Box::from(["execution.bundle-build".to_owned()]),
716 }),
717 }
718 }
719 Ok(Err(error)) => {
720 let AdmittedExecution { inputs, .. } = *admitted;
721 Ok(ExecutionOutcome::RetiredFailed {
722 error,
723 inputs: Box::new(inputs),
724 })
725 }
726 Err(payload) => {
727 let error = Error::runtime_state(
728 "ExecutionHandle::wait",
729 ErrorPhase::Execution,
730 panic_payload_message(payload),
731 );
732 Box::leak(admitted);
736 Ok(ExecutionOutcome::CompletionUnproven {
737 error,
738 diagnostic_keys: Box::from(["execution.retirement-unproven".to_owned()]),
739 })
740 }
741 }
742}
743
744pub(super) trait SubmissionSpawner {
745 fn spawn(&self, submission: Arc<InFlightSubmission>) -> std::io::Result<()>;
746}
747
748pub(super) fn spawn_in_flight(
749 submission: Arc<InFlightSubmission>,
750 spawner: &dyn SubmissionSpawner,
751) -> std::result::Result<ExecutionHandle, Box<(ExecutionInputs, Error)>> {
752 if let Err(source) = spawner.spawn(Arc::clone(&submission)) {
753 let submission = match Arc::try_unwrap(submission) {
757 Ok(submission) => submission,
758 Err(_) => unreachable!("failed spawner must not retain submission"),
759 };
760 let inputs = submission.into_unstarted_inputs();
761 let error = Error::runtime_state_source(
762 "Runtime::submit",
763 ErrorPhase::Execution,
764 SubmissionError::WorkerSpawn { source },
765 );
766 return Err(Box::new((inputs, error)));
767 }
768 Ok(ExecutionHandle { submission })
769}
770
771pub(super) struct OsThreadSpawner;
772
773impl SubmissionSpawner for OsThreadSpawner {
774 fn spawn(&self, submission: Arc<InFlightSubmission>) -> std::io::Result<()> {
775 thread::Builder::new()
776 .name("tenferro-runtime-submit".to_string())
777 .spawn(move || submission.run())
778 .map(drop)
779 }
780}
781
782impl fmt::Debug for PreparedCompiledGraph {
783 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
784 formatter
785 .debug_struct("PreparedCompiledGraph")
786 .field("runtime_id", &self.runtime_id)
787 .field("epoch", &self.epoch)
788 .field("program", &self.program)
789 .finish_non_exhaustive()
790 }
791}
792
793fn panic_payload_message(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
794 if let Some(message) = payload.downcast_ref::<&str>() {
795 format!("submitted execution panicked: {message}")
796 } else if let Some(message) = payload.downcast_ref::<String>() {
797 format!("submitted execution panicked: {message}")
798 } else {
799 "submitted execution panicked".to_string()
800 }
801}
802
803#[allow(
804 dead_code,
805 reason = "Phase 5 runtime execution task adds erased dispatch methods"
806)]
807pub(super) trait ErasedTensorBackendExecutor: fmt::Debug + Send + Sync {
808 fn backend_type_name(&self) -> &'static str;
809 fn extension_cache_stats(&self) -> std::result::Result<CacheStats, CacheOwnerError>;
810 fn clear_extension_caches(&self) -> std::result::Result<(), CacheOwnerError>;
811 fn execute(
812 &self,
813 program: &ExecProgram,
814 operations: &[PreparedOperationPlan],
815 inputs: Vec<Tensor>,
816 ) -> Result<Vec<Tensor>>;
817 fn execute_tensor_refs(
818 &self,
819 program: &ExecProgram,
820 operations: &[PreparedOperationPlan],
821 inputs: &[&Tensor],
822 ) -> Result<Vec<Tensor>>;
823 fn execute_values(
824 &self,
825 program: &ExecProgram,
826 operations: &[PreparedOperationPlan],
827 inputs: Vec<Tensor>,
828 ) -> Result<Vec<TensorValue>>;
829 fn execute_value_refs(
830 &self,
831 program: &ExecProgram,
832 operations: &[PreparedOperationPlan],
833 inputs: &[&Tensor],
834 ) -> Result<Vec<TensorValue>>;
835 fn execute_slot_instruction<'input>(
836 &self,
837 instruction_index: usize,
838 instruction: &ExecInstruction,
839 operations: &[PreparedOperationPlan],
840 slots: &mut [Option<ExecSlot<'input>>],
841 output_mode: RuntimeOutputMode,
842 terminal_slots: &[bool],
843 ) -> Result<()>;
844 fn materialize_slot<'input>(&self, slot: ExecSlot<'input>) -> Result<Tensor>;
845 fn materialize_slot_value<'input>(&self, slot: ExecSlot<'input>) -> Result<TensorValue>;
846}
847
848pub(super) fn erased_tensor_backend_executor<B>(backend: B) -> Arc<dyn ErasedTensorBackendExecutor>
849where
850 B: TensorBackend + Send + Sync + 'static,
851{
852 Arc::new(TensorBackendExecutor::<B>::new(backend))
853}
854
855pub(super) fn extension_cache_owner(
856 executor: Arc<dyn ErasedTensorBackendExecutor>,
857) -> Arc<dyn RuntimeCacheOwner> {
858 Arc::new(TensorBackendExtensionCacheOwner { executor })
859}
860
861#[derive(Debug)]
862struct TensorBackendExtensionCacheOwner {
863 executor: Arc<dyn ErasedTensorBackendExecutor>,
864}
865
866impl RuntimeCacheOwner for TensorBackendExtensionCacheOwner {
867 fn cache_stats(&self) -> std::result::Result<CacheStats, CacheOwnerError> {
868 self.executor.extension_cache_stats()
869 }
870
871 fn clear_caches(&self) -> std::result::Result<(), CacheOwnerError> {
872 self.executor.clear_extension_caches()
873 }
874}
875
876#[allow(
877 dead_code,
878 reason = "Phase 5 runtime execution task consumes backend execution state"
879)]
880struct TensorBackendExecutorState<B: TensorBackend + 'static> {
881 backend: B,
882 backend_cache: B::RuntimeCache,
883 extension_caches: ExtensionCacheStore,
884 slot_workspace: Vec<Option<ExecSlot<'static>>>,
885 borrowed_slot_workspace_capacity: usize,
886}
887
888struct TensorBackendExecutor<B: TensorBackend + 'static> {
889 state: Mutex<TensorBackendExecutorSlot<B>>,
890 available: Condvar,
891}
892
893struct TensorBackendExecutorSlot<B: TensorBackend + 'static> {
894 state: Option<TensorBackendExecutorState<B>>,
895 active_thread: Option<ThreadId>,
896 execution_poisoned: bool,
897}
898
899impl<B> TensorBackendExecutor<B>
900where
901 B: TensorBackend + Send + Sync + 'static,
902{
903 fn new(backend: B) -> Self {
904 Self {
905 state: Mutex::new(TensorBackendExecutorSlot {
906 state: Some(TensorBackendExecutorState {
907 backend,
908 backend_cache: B::RuntimeCache::default(),
909 extension_caches: ExtensionCacheStore::new(),
910 slot_workspace: Vec::new(),
911 borrowed_slot_workspace_capacity: 0,
912 }),
913 active_thread: None,
914 execution_poisoned: false,
915 }),
916 available: Condvar::new(),
917 }
918 }
919
920 fn lease_state(&self, caller: &'static str) -> Result<TensorBackendExecutorLease<'_, B>> {
921 let current = thread::current().id();
922 let mut slot = self.lock_slot(caller)?;
923 loop {
924 if slot.execution_poisoned {
925 return Err(Error::runtime_state(
926 caller,
927 ErrorPhase::Execution,
928 "tensor backend executor state poisoned by panic during prior execution",
929 ));
930 }
931 if let Some(state) = slot.state.take() {
932 slot.active_thread = Some(current);
933 return Ok(TensorBackendExecutorLease {
934 executor: self,
935 state: Some(state),
936 });
937 }
938 if slot.active_thread == Some(current) {
939 return Err(Error::runtime_state(
940 caller,
941 ErrorPhase::Execution,
942 "reentrant tensor backend executor call would deadlock",
943 ));
944 }
945 slot = self.available.wait(slot).map_err(|_| {
946 Error::runtime_state(
947 caller,
948 ErrorPhase::Execution,
949 "tensor backend executor state lock poisoned while waiting",
950 )
951 })?;
952 }
953 }
954
955 fn lock_slot(
956 &self,
957 caller: &'static str,
958 ) -> Result<MutexGuard<'_, TensorBackendExecutorSlot<B>>> {
959 self.state.lock().map_err(|_| {
960 Error::runtime_state(
961 caller,
962 ErrorPhase::Execution,
963 "tensor backend executor state lock poisoned",
964 )
965 })
966 }
967}
968
969struct TensorBackendExecutorLease<'a, B: TensorBackend + 'static> {
970 executor: &'a TensorBackendExecutor<B>,
971 state: Option<TensorBackendExecutorState<B>>,
972}
973
974impl<B: TensorBackend + 'static> TensorBackendExecutorLease<'_, B> {
975 fn state_mut(&mut self) -> &mut TensorBackendExecutorState<B> {
976 self.state
977 .as_mut()
978 .expect("executor state lease always owns state before drop")
979 }
980}
981
982impl<B: TensorBackend + 'static> Drop for TensorBackendExecutorLease<'_, B> {
983 fn drop(&mut self) {
984 let Some(state) = self.state.take() else {
985 return;
986 };
987 let panicking = thread::panicking();
988 let mut wake_all_waiters = panicking;
989 match self.executor.state.lock() {
990 Ok(mut slot) => {
991 slot.state = Some(state);
992 slot.active_thread = None;
993 slot.execution_poisoned |= panicking;
994 }
995 Err(poisoned_lock) => {
996 let mut slot = poisoned_lock.into_inner();
997 slot.state = Some(state);
998 slot.active_thread = None;
999 slot.execution_poisoned = true;
1000 wake_all_waiters = true;
1001 }
1002 }
1003 if wake_all_waiters {
1004 self.executor.available.notify_all();
1005 } else {
1006 self.executor.available.notify_one();
1007 }
1008 }
1009}
1010
1011impl<B: TensorBackend + 'static> fmt::Debug for TensorBackendExecutor<B> {
1012 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1013 formatter
1014 .debug_struct("TensorBackendExecutor")
1015 .field("backend_type", &std::any::type_name::<B>())
1016 .field("state_lock_poisoned", &self.state.is_poisoned())
1017 .finish_non_exhaustive()
1018 }
1019}
1020
1021impl<B> ErasedTensorBackendExecutor for TensorBackendExecutor<B>
1022where
1023 B: TensorBackend + Send + Sync + 'static,
1024{
1025 fn backend_type_name(&self) -> &'static str {
1026 std::any::type_name::<B>()
1027 }
1028
1029 fn extension_cache_stats(&self) -> std::result::Result<CacheStats, CacheOwnerError> {
1030 let lease = self
1031 .lease_state("Runtime::extension_cache_stats")
1032 .map_err(cache_owner_error)?;
1033 let state = lease
1034 .state
1035 .as_ref()
1036 .expect("executor state lease always owns state before drop");
1037 Ok(cache_stats_from_tensor_stats(
1038 state.extension_caches.stats(ExtensionCacheSelector::All),
1039 ))
1040 }
1041
1042 fn clear_extension_caches(&self) -> std::result::Result<(), CacheOwnerError> {
1043 let mut lease = self
1044 .lease_state("Runtime::clear_extension_caches")
1045 .map_err(cache_owner_error)?;
1046 lease.state_mut().extension_caches.clear();
1047 Ok(())
1048 }
1049
1050 fn execute(
1051 &self,
1052 program: &ExecProgram,
1053 operations: &[PreparedOperationPlan],
1054 inputs: Vec<Tensor>,
1055 ) -> Result<Vec<Tensor>> {
1056 validate_exec_input_count(program, inputs.len())?;
1057 let mut lease = self.lease_state("Runtime::run_compiled")?;
1058 let TensorBackendExecutorState {
1059 backend,
1060 backend_cache,
1061 extension_caches,
1062 slot_workspace,
1063 borrowed_slot_workspace_capacity: _,
1064 } = lease.state_mut();
1065 let mut extension_dispatch = ExtensionExecutionDispatch {
1066 operations,
1067 caches: extension_caches,
1068 };
1069 crate::segment::eval_exec_segmented_with_cache_and_workspace(
1070 backend,
1071 program,
1072 inputs,
1073 slot_workspace,
1074 backend_cache,
1075 Some(&mut extension_dispatch),
1076 )
1077 }
1078
1079 fn execute_tensor_refs(
1080 &self,
1081 program: &ExecProgram,
1082 operations: &[PreparedOperationPlan],
1083 inputs: &[&Tensor],
1084 ) -> Result<Vec<Tensor>> {
1085 validate_exec_input_count(program, inputs.len())?;
1086 let inputs = inputs
1087 .iter()
1088 .map(|tensor| ExecSlot::Read(TensorRead::from_tensor(tensor)))
1089 .collect();
1090 let mut lease = self.lease_state("Runtime::run_compiled")?;
1091 let TensorBackendExecutorState {
1092 backend,
1093 backend_cache,
1094 extension_caches,
1095 borrowed_slot_workspace_capacity,
1096 ..
1097 } = lease.state_mut();
1098 let mut extension_dispatch = ExtensionExecutionDispatch {
1099 operations,
1100 caches: extension_caches,
1101 };
1102 let mut slot_workspace = Vec::with_capacity(*borrowed_slot_workspace_capacity);
1103 let result = crate::segment::eval_exec_segmented_slots_with_cache_and_workspace(
1104 backend,
1105 program,
1106 inputs,
1107 &mut slot_workspace,
1108 backend_cache,
1109 Some(&mut extension_dispatch),
1110 );
1111 *borrowed_slot_workspace_capacity = slot_workspace.capacity();
1112 result
1113 }
1114
1115 fn execute_values(
1116 &self,
1117 program: &ExecProgram,
1118 operations: &[PreparedOperationPlan],
1119 inputs: Vec<Tensor>,
1120 ) -> Result<Vec<TensorValue>> {
1121 validate_exec_input_count(program, inputs.len())?;
1122 let inputs = inputs.into_iter().map(ExecSlot::Owned).collect();
1123 let mut lease = self.lease_state("Runtime::run_compiled_values")?;
1124 let TensorBackendExecutorState {
1125 backend,
1126 backend_cache,
1127 extension_caches,
1128 slot_workspace,
1129 borrowed_slot_workspace_capacity: _,
1130 } = lease.state_mut();
1131 let mut extension_dispatch = ExtensionExecutionDispatch {
1132 operations,
1133 caches: extension_caches,
1134 };
1135 crate::segment::eval_exec_segmented_slot_values_with_cache_and_workspace(
1136 backend,
1137 program,
1138 inputs,
1139 slot_workspace,
1140 backend_cache,
1141 Some(&mut extension_dispatch),
1142 )
1143 }
1144
1145 fn execute_value_refs(
1146 &self,
1147 program: &ExecProgram,
1148 operations: &[PreparedOperationPlan],
1149 inputs: &[&Tensor],
1150 ) -> Result<Vec<TensorValue>> {
1151 validate_exec_input_count(program, inputs.len())?;
1152 let inputs = inputs
1153 .iter()
1154 .map(|tensor| ExecSlot::Read(TensorRead::from_tensor(tensor)))
1155 .collect();
1156 let mut lease = self.lease_state("Runtime::run_compiled_values")?;
1157 let TensorBackendExecutorState {
1158 backend,
1159 backend_cache,
1160 extension_caches,
1161 borrowed_slot_workspace_capacity,
1162 ..
1163 } = lease.state_mut();
1164 let mut extension_dispatch = ExtensionExecutionDispatch {
1165 operations,
1166 caches: extension_caches,
1167 };
1168 let mut slot_workspace = Vec::with_capacity(*borrowed_slot_workspace_capacity);
1169 let result = crate::segment::eval_exec_segmented_slot_values_with_cache_and_workspace(
1170 backend,
1171 program,
1172 inputs,
1173 &mut slot_workspace,
1174 backend_cache,
1175 Some(&mut extension_dispatch),
1176 );
1177 *borrowed_slot_workspace_capacity = slot_workspace.capacity();
1178 result
1179 }
1180
1181 fn execute_slot_instruction<'input>(
1182 &self,
1183 instruction_index: usize,
1184 instruction: &ExecInstruction,
1185 operations: &[PreparedOperationPlan],
1186 slots: &mut [Option<ExecSlot<'input>>],
1187 output_mode: RuntimeOutputMode,
1188 terminal_slots: &[bool],
1189 ) -> Result<()> {
1190 let mut lease = self.lease_state("Runtime::run_compiled scheduled instruction")?;
1191 let TensorBackendExecutorState {
1192 backend,
1193 backend_cache,
1194 extension_caches,
1195 ..
1196 } = lease.state_mut();
1197 let mut extension_dispatch = ExtensionExecutionDispatch {
1198 operations,
1199 caches: extension_caches,
1200 };
1201
1202 if matches!(output_mode, RuntimeOutputMode::Value)
1203 && backend.with_backend_session(|exec| {
1204 crate::exec::try_execute_terminal_value_instruction(
1205 exec,
1206 slots,
1207 instruction,
1208 terminal_slots,
1209 )
1210 })?
1211 {
1212 } else if crate::exec::is_host_instruction(instruction) {
1214 crate::exec::execute_host_instruction(backend, slots, instruction)?;
1215 } else if crate::exec::is_ffi_instruction(instruction) {
1216 crate::exec::execute_ffi_instruction_cached(
1217 backend,
1218 backend_cache,
1219 slots,
1220 instruction,
1221 DispatchMode::Unsegmented,
1222 Some(instruction_index),
1223 Some(&mut extension_dispatch),
1224 )?;
1225 } else {
1226 let result = backend.with_backend_session(|exec| {
1227 crate::exec::execute_backend_op(exec, slots, instruction)
1228 })?;
1229 slots[instruction.output_slots[0]] = Some(ExecSlot::Owned(result));
1230 }
1231 crate::exec::reclaim_last_use_inputs_backend(slots, instruction, backend);
1232 Ok(())
1233 }
1234
1235 fn materialize_slot<'input>(&self, slot: ExecSlot<'input>) -> Result<Tensor> {
1236 let mut lease = self.lease_state("Runtime::run_compiled collect outputs")?;
1237 let backend = &mut lease.state_mut().backend;
1238 backend.with_backend_session(|exec| slot.into_tensor(exec))
1239 }
1240
1241 fn materialize_slot_value<'input>(&self, slot: ExecSlot<'input>) -> Result<TensorValue> {
1242 let mut lease = self.lease_state("Runtime::run_compiled_values collect outputs")?;
1243 let backend = &mut lease.state_mut().backend;
1244 backend.with_backend_session(|exec| slot.into_value(exec))
1245 }
1246}
1247
1248fn cache_owner_error(error: Error) -> CacheOwnerError {
1249 CacheOwnerError::new(Arc::new(error))
1250}
1251
1252fn cache_stats_from_tensor_stats(stats: tenferro_tensor::CacheStats) -> CacheStats {
1253 CacheStats {
1254 entries: stats.entries,
1255 retained_bytes: stats.retained_bytes,
1256 hits: stats.hits,
1257 misses: stats.misses,
1258 evictions: stats.evictions,
1259 clears: stats.clears,
1260 }
1261}
1262
1263pub(super) fn run_compiled(
1264 runtime: &Runtime,
1265 program: &CompiledGraph,
1266 inputs: &[&Tensor],
1267) -> Result<Vec<Tensor>> {
1268 let inputs = resolve_input_refs(program, inputs)?;
1269 let signature = input_signature_reads(&inputs)?;
1270 let prepared = prepare(runtime, program, &signature)?;
1271 validate_prepared_epoch(runtime, prepared.root().epoch(), "Runtime::run_compiled")?;
1272 execute_scheduled_reads(
1273 prepared.root().staging(),
1274 prepared.root().schedule(),
1275 prepared.operations(),
1276 &inputs,
1277 )
1278}
1279
1280pub(super) fn execute_scoped_read_only<'env>(
1285 runtime: &Runtime,
1286 program: &CompiledGraph,
1287 inputs: ScopedReadInputs<'env>,
1288) -> std::result::Result<ScopedExecutionOutcome<'env>, ScopedSubmitRejected<'env>> {
1289 if inputs.has_non_host_provider() {
1290 return Err(ScopedSubmitRejected {
1291 source: Box::new(Error::unsupported(
1292 "Runtime::execute_scoped_read_only",
1293 ErrorPhase::Execution,
1294 "scoped borrowed execution requires host/CPU storage",
1295 )),
1296 inputs,
1297 });
1298 }
1299
1300 let reads = inputs.as_reads();
1301 let prepared = match prepare_compiled_reads(runtime, program, &reads) {
1302 Ok(prepared) => prepared,
1303 Err(source) => {
1304 return Err(ScopedSubmitRejected {
1305 source: Box::new(source),
1306 inputs,
1307 })
1308 }
1309 };
1310 if !schedule_supports_scoped_execution(prepared.prepared.root().schedule()) {
1311 return Err(ScopedSubmitRejected {
1312 source: Box::new(Error::unsupported(
1313 "Runtime::execute_scoped_read_only",
1314 ErrorPhase::Execution,
1315 "the selected provider does not prove retirement before scoped return",
1316 )),
1317 inputs,
1318 });
1319 }
1320
1321 match execute_scoped_admitted(
1322 prepared.prepared.root().staging(),
1323 prepared.prepared.root().schedule(),
1324 prepared.prepared.operations(),
1325 &reads,
1326 ) {
1327 Ok(bundle) => Ok(ScopedExecutionOutcome::Completed(bundle)),
1328 Err(error) => Ok(ScopedExecutionOutcome::RetiredFailed { error, inputs }),
1329 }
1330}
1331
1332fn execute_scoped_admitted<'env>(
1333 program: &ExecProgram,
1334 schedule: &ScheduledGraph,
1335 operations: &[PreparedOperationPlan],
1336 inputs: &[TensorRead<'env>],
1337) -> Result<ScopedExecutionBundle<'env>> {
1338 validate_exec_input_count(program, inputs.len())?;
1339 crate::exec::validate_exec_program(program, "scoped tensor executor")?;
1340 let input_slots = inputs.iter().cloned().map(ExecSlot::Read).collect();
1341 let mut located = execute_scheduled_slots(
1342 program,
1343 schedule,
1344 operations,
1345 input_slots,
1346 RuntimeOutputMode::Tensor,
1347 )?;
1348 collect_scoped_outputs(program, &mut located)
1349}
1350
1351pub(crate) fn collect_scoped_outputs<'env>(
1352 program: &ExecProgram,
1353 outputs: &mut [Option<LocatedExecSlot<'env>>],
1354) -> Result<ScopedExecutionBundle<'env>> {
1355 let mut owned = AllocationGroup::from_tensors(Vec::new())
1356 .map(|(group, _)| group)
1357 .map_err(|error| {
1358 Error::runtime_state(
1359 "Runtime::collect_scoped_outputs",
1360 ErrorPhase::Execution,
1361 error.to_string(),
1362 )
1363 })?;
1364 let mut scoped = Vec::with_capacity(program.output_slots.len());
1365 for &slot in &program.output_slots {
1366 let located = outputs
1367 .get_mut(slot)
1368 .and_then(Option::take)
1369 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1370 let output = match located.value {
1371 ExecSlot::Read(read) => ScopedOutput::Borrowed(read.tensor_view()),
1372 ExecSlot::Owned(tensor) => {
1373 let slot = owned.append_tensor(tensor).map_err(|error| {
1374 Error::runtime_state(
1375 "Runtime::collect_scoped_outputs",
1376 ErrorPhase::Execution,
1377 error.to_string(),
1378 )
1379 })?;
1380 ScopedOutput::Owned(slot)
1381 }
1382 ExecSlot::Value(value) => {
1383 let (group, slot, dtype, shape) = value.try_into_group_parts().map_err(|_| {
1384 Error::runtime_state(
1385 "Runtime::collect_scoped_outputs",
1386 ErrorPhase::Execution,
1387 "scoped value output was retained by another owner",
1388 )
1389 })?;
1390 let slot = owned.append_group(group, slot).map_err(|error| {
1391 Error::runtime_state(
1392 "Runtime::collect_scoped_outputs",
1393 ErrorPhase::Execution,
1394 error.to_string(),
1395 )
1396 })?;
1397 let _ = (dtype, shape);
1398 ScopedOutput::Owned(slot)
1399 }
1400 };
1401 scoped.push(output);
1402 }
1403 outputs.iter_mut().for_each(|output| *output = None);
1404 Ok(ScopedExecutionBundle {
1405 owned,
1406 outputs: scoped.into_boxed_slice(),
1407 })
1408}
1409
1410fn output_ref_from_group<'a>(
1411 group: &'a AllocationGroup,
1412 slot: DescriptorSlot,
1413) -> std::result::Result<OutputRef<'a>, OutputAccessError> {
1414 let mut views = group.read_views(std::slice::from_ref(&slot))?;
1415 match views.pop() {
1416 Some(TensorRead::View(view)) => Ok(OutputRef::Tensor(view)),
1417 Some(TensorRead::Tensor(_)) | None => Err(OutputAccessError::InvalidOutput {
1418 index: slot.index(),
1419 }),
1420 }
1421}
1422
1423fn schedule_supports_scoped_execution(schedule: &ScheduledGraph) -> bool {
1424 let supports = |location: &ExecutionLocation| {
1425 location
1426 .witness()
1427 .event_domain_driver()
1428 .supports_scoped_read_only()
1429 };
1430 supports(schedule.root_location())
1431 && schedule.input_locations().iter().all(supports)
1432 && schedule.operation_locations().iter().all(supports)
1433 && schedule.nodes().iter().all(|node| {
1434 node.event_domain_witness()
1435 .event_domain_driver()
1436 .supports_scoped_read_only()
1437 })
1438}
1439
1440pub(super) fn prepare_compiled(
1441 runtime: &Runtime,
1442 program: &CompiledGraph,
1443 inputs: &[&Tensor],
1444) -> Result<PreparedCompiledGraph> {
1445 let inputs = resolve_input_refs(program, inputs)?;
1446 let signature = input_signature_reads(&inputs)?;
1447 let prepared = prepare(runtime, program, &signature)?;
1448 validate_prepared_epoch(
1449 runtime,
1450 prepared.root().epoch(),
1451 "Runtime::prepare_compiled",
1452 )?;
1453 Ok(PreparedCompiledGraph {
1454 runtime_id: runtime.id(),
1455 epoch: prepared.root().epoch(),
1456 program: program.clone(),
1457 prepared,
1458 })
1459}
1460
1461pub(super) fn submit(
1462 runtime: &Runtime,
1463 program: &CompiledGraph,
1464 inputs: ExecutionInputs,
1465) -> std::result::Result<ExecutionHandle, SubmitError> {
1466 submit_with_spawner(runtime, program, inputs, &OsThreadSpawner)
1467}
1468
1469pub(super) fn submit_with_spawner(
1470 runtime: &Runtime,
1471 program: &CompiledGraph,
1472 inputs: ExecutionInputs,
1473 spawner: &dyn SubmissionSpawner,
1474) -> std::result::Result<ExecutionHandle, SubmitError> {
1475 let prepared = match prepare_submission(runtime, program, &inputs) {
1476 Ok(prepared) => prepared,
1477 Err(source) => {
1478 return Err(SubmitError::PreAdmission {
1479 source: Box::new(source),
1480 inputs: Box::new(inputs),
1481 })
1482 }
1483 };
1484 let submission = Arc::new(InFlightSubmission::new(prepared, inputs));
1485 match spawn_in_flight(submission, spawner) {
1486 Ok(handle) => Ok(handle),
1487 Err(failure) => {
1488 let (inputs, source) = *failure;
1489 Err(SubmitError::PreAdmission {
1490 source: Box::new(source),
1491 inputs: Box::new(inputs),
1492 })
1493 }
1494 }
1495}
1496
1497pub(super) fn run_prepared(
1498 runtime: &Runtime,
1499 prepared: &PreparedCompiledGraph,
1500 inputs: &[&Tensor],
1501) -> Result<Vec<Tensor>> {
1502 validate_prepared_runtime(runtime, prepared, "Runtime::run_prepared")?;
1503 let inputs = resolve_input_refs(&prepared.program, inputs)?;
1504 execute_scheduled_reads(
1505 prepared.prepared.root().staging(),
1506 prepared.prepared.root().schedule(),
1507 prepared.prepared.operations(),
1508 &inputs,
1509 )
1510}
1511
1512fn prepare_submission(
1513 runtime: &Runtime,
1514 program: &CompiledGraph,
1515 inputs: &ExecutionInputs,
1516) -> Result<PreparedCompiledGraph> {
1517 let input_reads = inputs.as_reads()?;
1518 prepare_compiled_reads(runtime, program, &input_reads)
1519}
1520
1521fn prepare_compiled_reads(
1522 runtime: &Runtime,
1523 program: &CompiledGraph,
1524 inputs: &[TensorRead<'_>],
1525) -> Result<PreparedCompiledGraph> {
1526 validate_ordered_input_metadata_reads(program, inputs)?;
1527 let signature = input_signature_reads(inputs)?;
1528 let prepared = prepare(runtime, program, &signature)?;
1529 validate_prepared_epoch(runtime, prepared.root().epoch(), "Runtime::submit")?;
1530 Ok(PreparedCompiledGraph {
1531 runtime_id: runtime.id(),
1532 epoch: prepared.root().epoch(),
1533 program: program.clone(),
1534 prepared,
1535 })
1536}
1537
1538fn execute_admitted(
1539 prepared: &PreparedCompiledGraph,
1540 inputs: &[TensorRead<'_>],
1541) -> Result<Vec<Tensor>> {
1542 execute_scheduled_reads(
1543 prepared.prepared.root().staging(),
1544 prepared.prepared.root().schedule(),
1545 prepared.prepared.operations(),
1546 inputs,
1547 )
1548}
1549
1550pub(super) fn run_compiled_values(
1551 runtime: &Runtime,
1552 program: &CompiledGraph,
1553 inputs: &[&Tensor],
1554) -> Result<Vec<TensorValue>> {
1555 let inputs = resolve_input_refs(program, inputs)?;
1556 let signature = input_signature_reads(&inputs)?;
1557 let prepared = prepare(runtime, program, &signature)?;
1558 validate_prepared_epoch(
1559 runtime,
1560 prepared.root().epoch(),
1561 "Runtime::run_compiled_values",
1562 )?;
1563 execute_scheduled_value_reads(
1564 prepared.root().staging(),
1565 prepared.root().schedule(),
1566 prepared.operations(),
1567 &inputs,
1568 )
1569}
1570
1571fn validate_prepared_runtime(
1572 runtime: &Runtime,
1573 prepared: &PreparedCompiledGraph,
1574 caller: &'static str,
1575) -> Result<()> {
1576 if runtime.id() != prepared.runtime_id {
1577 return Err(Error::runtime_state(
1578 caller,
1579 ErrorPhase::Execution,
1580 "prepared compiled graph belongs to a different runtime",
1581 ));
1582 }
1583 let epoch = runtime
1584 .epoch()
1585 .map_err(|source| Error::runtime_state_source(caller, ErrorPhase::Execution, source))?;
1586 if epoch != prepared.epoch {
1587 return Err(Error::runtime_state(
1588 caller,
1589 ErrorPhase::Execution,
1590 format!(
1591 "prepared epoch {:?} does not match current epoch {:?}",
1592 prepared.epoch, epoch
1593 ),
1594 ));
1595 }
1596 Ok(())
1597}
1598
1599fn prepare(
1600 runtime: &Runtime,
1601 program: &CompiledGraph,
1602 signature: &InputSignature,
1603) -> Result<Arc<super::preparation::PreparedProgram>> {
1604 runtime
1605 .prepare_compiled_for(program, signature, &PrepareOptions::new())
1606 .map_err(prepare_error)
1607}
1608
1609fn validate_prepared_epoch(
1610 runtime: &Runtime,
1611 prepared_epoch: super::RuntimeEpoch,
1612 caller: &'static str,
1613) -> Result<()> {
1614 let epoch = runtime
1615 .epoch()
1616 .map_err(|source| Error::runtime_state_source(caller, ErrorPhase::Execution, source))?;
1617 if epoch != prepared_epoch {
1618 return Err(Error::runtime_state(
1619 caller,
1620 ErrorPhase::Execution,
1621 format!(
1622 "prepared epoch {:?} does not match current epoch {:?}",
1623 prepared_epoch, epoch
1624 ),
1625 ));
1626 }
1627 Ok(())
1628}
1629
1630fn execute_scheduled_reads(
1631 program: &ExecProgram,
1632 schedule: &ScheduledGraph,
1633 operations: &[PreparedOperationPlan],
1634 inputs: &[TensorRead<'_>],
1635) -> Result<Vec<Tensor>> {
1636 validate_exec_input_count(program, inputs.len())?;
1637 crate::exec::validate_exec_program(program, "scheduled tensor executor")?;
1638 let inputs = inputs.iter().cloned().map(ExecSlot::Read).collect();
1639 execute_scheduled_slots(
1640 program,
1641 schedule,
1642 operations,
1643 inputs,
1644 RuntimeOutputMode::Tensor,
1645 )
1646 .and_then(|mut slots| {
1647 collect_tensor_outputs_with(program, &mut slots, |location, slot| {
1648 location.witness().executor().materialize_slot(slot)
1649 })
1650 })
1651}
1652
1653fn execute_scheduled_value_reads(
1654 program: &ExecProgram,
1655 schedule: &ScheduledGraph,
1656 operations: &[PreparedOperationPlan],
1657 inputs: &[TensorRead<'_>],
1658) -> Result<Vec<TensorValue>> {
1659 validate_exec_input_count(program, inputs.len())?;
1660 crate::exec::validate_exec_program(program, "scheduled value executor")?;
1661 let inputs = inputs.iter().cloned().map(ExecSlot::Read).collect();
1662 execute_scheduled_slots(
1663 program,
1664 schedule,
1665 operations,
1666 inputs,
1667 RuntimeOutputMode::Value,
1668 )
1669 .and_then(|mut slots| {
1670 collect_value_outputs_with(program, &mut slots, |location, slot| {
1671 location.witness().executor().materialize_slot_value(slot)
1672 })
1673 })
1674}
1675
1676fn execute_scheduled_slots<'input>(
1677 program: &ExecProgram,
1678 schedule: &ScheduledGraph,
1679 operations: &[PreparedOperationPlan],
1680 inputs: Vec<ExecSlot<'input>>,
1681 output_mode: RuntimeOutputMode,
1682) -> Result<Vec<Option<LocatedExecSlot<'input>>>> {
1683 let terminal_slots = if matches!(output_mode, RuntimeOutputMode::Value) {
1684 crate::exec::terminal_output_slots(program)
1685 } else {
1686 Vec::new()
1687 };
1688 let mut staged = Vec::new();
1689 let mut located = (0..program.n_slots)
1690 .map(|_| Vec::new())
1691 .collect::<Vec<Vec<LocatedExecSlot<'input>>>>();
1692 let mut event_domains = ScheduledEventDomains::new(schedule)?;
1697 let result = (|| {
1698 crate::exec::initialize_exec_slots_in(program, inputs, &mut staged)?;
1699 if schedule.input_locations().len() != program.input_slots.len() {
1700 return Err(Error::runtime_state(
1701 "Runtime::run_compiled",
1702 ErrorPhase::Execution,
1703 format!(
1704 "prepared schedule has {} input locations for {} inputs",
1705 schedule.input_locations().len(),
1706 program.input_slots.len()
1707 ),
1708 ));
1709 }
1710 for (&slot, location) in program.input_slots.iter().zip(schedule.input_locations()) {
1711 let value = staged
1712 .get_mut(slot)
1713 .and_then(Option::take)
1714 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
1715 validate_runtime_input_ingress(location, &value.as_read(), slot)?;
1716 located[slot].push(LocatedExecSlot {
1717 location: location.clone(),
1718 value,
1719 });
1720 }
1721 for (node_index, node) in schedule.nodes().iter().enumerate() {
1722 match node {
1723 ScheduledNode::Operation(operation_node) => {
1724 let mut launch = || {
1725 let instruction_index = operation_node.instruction_index();
1726 let instruction =
1727 program.instructions.get(instruction_index).ok_or_else(|| {
1728 Error::runtime_state(
1729 "Runtime::run_compiled",
1730 ErrorPhase::Execution,
1731 format!(
1732 "scheduled operation references instruction \
1733 {instruction_index}, but the execution program has {} \
1734 instructions",
1735 program.instructions.len()
1736 ),
1737 )
1738 })?;
1739 let operation = instruction_execution(schedule, instruction)?;
1740 if operation.location() != operation_node.location() {
1741 return Err(Error::runtime_state(
1742 "Runtime::run_compiled",
1743 ErrorPhase::Execution,
1744 format!(
1745 "scheduled instruction {instruction_index} location does not \
1746 match its prepared executor"
1747 ),
1748 ));
1749 }
1750 stage_instruction_inputs(
1751 instruction,
1752 operation.location(),
1753 &mut located,
1754 &mut staged,
1755 )?;
1756 operation.executor().execute_slot_instruction(
1757 instruction_index,
1758 instruction,
1759 operations,
1760 &mut staged,
1761 output_mode,
1762 &terminal_slots,
1763 )?;
1764 validate_instruction_outputs(
1765 instruction_index,
1766 instruction,
1767 operation.location(),
1768 &staged,
1769 )?;
1770 retain_instruction_results(
1771 instruction,
1772 operation.location(),
1773 &mut located,
1774 &mut staged,
1775 )
1776 };
1777 event_domains.enqueue(node_index, node, &mut launch)?;
1778 }
1779 ScheduledNode::Transfer(transfer) => {
1780 let mut launch = || execute_scheduled_transfer(transfer, &mut located);
1781 event_domains.enqueue(node_index, node, &mut launch)?;
1782 }
1783 ScheduledNode::Collective(_) => {
1784 return Err(Error::runtime_state_source(
1785 "Runtime::run_compiled",
1786 ErrorPhase::Execution,
1787 UnsupportedScheduledNodeError {
1788 node_index,
1789 node_kind: ScheduledNodeKind::Collective,
1790 },
1791 ));
1792 }
1793 ScheduledNode::Barrier(_) => {
1794 let mut launch = || Ok(());
1795 event_domains.enqueue(node_index, node, &mut launch)?;
1796 }
1797 }
1798 }
1799 Ok(())
1800 })();
1801 let drain = event_domains.drain();
1802 match (result, drain) {
1803 (Ok(()), Ok(())) => collect_located_outputs(program, &mut located),
1804 (Err(error), Ok(())) | (Ok(()), Err(error)) => {
1805 staged.clear();
1806 located.clear();
1807 Err(error)
1808 }
1809 (Err(primary), Err(cleanup)) => {
1810 staged.clear();
1811 located.clear();
1812 Err(scheduled_execution_cleanup_error(primary, cleanup))
1813 }
1814 }
1815}
1816
1817#[derive(Debug)]
1818pub(crate) struct ScheduledEventDomains {
1819 runs: Vec<RuntimeOwnedEventDomainRun>,
1820 completions: HashMap<EventDependency, Arc<dyn EventToken>>,
1821}
1822
1823#[derive(Debug, thiserror::Error)]
1824#[error(
1825 "scheduled node {node_index} depends on completion {dependency:?}, but no completion token was recorded"
1826)]
1827pub(crate) struct MissingScheduledDependencyCompletionError {
1828 pub(crate) dependency: EventDependency,
1829 pub(crate) node_index: usize,
1830}
1831
1832#[derive(Debug, thiserror::Error)]
1833#[error("scheduled transfer destination {destination:?} already contains value slot {value_slot}")]
1834pub(crate) struct DuplicateTransferDestinationError {
1835 pub(crate) value_slot: usize,
1836 pub(crate) destination: ExecutionLocation,
1837}
1838
1839impl ScheduledEventDomains {
1840 pub(crate) fn new(schedule: &ScheduledGraph) -> Result<Self> {
1841 schedule.preflight().map_err(|source| {
1842 Error::runtime_state_source("Runtime::run_compiled", ErrorPhase::Execution, source)
1843 })?;
1844 let mut drivers = Vec::new();
1845 let mut seen_domains = HashSet::new();
1846 for node in schedule.nodes() {
1847 let domain = node.completion().domain();
1848 if !seen_domains.insert(domain) {
1849 continue;
1850 }
1851 let witness = node.event_domain_witness();
1852 debug_assert_eq!(witness.event_domain_id(), domain);
1853 drivers.push((domain, witness.event_domain_driver().clone()));
1854 }
1855 let mut runs = Vec::with_capacity(drivers.len());
1856 for (domain, driver) in drivers {
1857 let run = RuntimeOwnedEventDomainRun::new(domain, driver.begin_run(domain)?);
1858 let actual = run.domain(EventDomainOperation::BeginRun)?;
1859 if actual != domain {
1860 return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1861 operation: EventDomainOperation::BeginRun,
1862 node_index: None,
1863 expected: domain,
1864 actual,
1865 }));
1866 }
1867 runs.push(run);
1868 }
1869 Ok(Self {
1870 runs,
1871 completions: HashMap::new(),
1872 })
1873 }
1874
1875 #[cfg(test)]
1876 pub(crate) fn for_test(
1877 drivers: Vec<(super::EventDomainId, Arc<dyn super::EventDomainDriver>)>,
1878 ) -> Result<Self> {
1879 let mut runs = Vec::with_capacity(drivers.len());
1880 for (domain, driver) in drivers {
1881 let run = RuntimeOwnedEventDomainRun::new(domain, driver.begin_run(domain)?);
1882 let actual = run.domain(EventDomainOperation::BeginRun)?;
1883 if actual != domain {
1884 return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1885 operation: EventDomainOperation::BeginRun,
1886 node_index: None,
1887 expected: domain,
1888 actual,
1889 }));
1890 }
1891 runs.push(run);
1892 }
1893 Ok(Self {
1894 runs,
1895 completions: HashMap::new(),
1896 })
1897 }
1898
1899 pub(crate) fn enqueue(
1900 &mut self,
1901 node_index: usize,
1902 node: &ScheduledNode,
1903 launch: &mut dyn FnMut() -> Result<()>,
1904 ) -> Result<()> {
1905 let completion = node.completion();
1906 let destination = completion.domain();
1907 let run_index = self.run_index(completion.domain())?;
1908 let actual_preflight_domain = self.runs[run_index].domain(EventDomainOperation::Enqueue)?;
1909 if actual_preflight_domain != destination {
1910 return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1911 operation: EventDomainOperation::Enqueue,
1912 node_index: Some(node_index),
1913 expected: destination,
1914 actual: actual_preflight_domain,
1915 }));
1916 }
1917 let dependencies = self.classify_dependencies(node_index, node, destination)?;
1918 let actual_run_domain = self.runs[run_index].domain(EventDomainOperation::Enqueue)?;
1919 if actual_run_domain != destination {
1920 return Err(event_domain_error(EventDomainError::RunDomainMismatch {
1921 operation: EventDomainOperation::Enqueue,
1922 node_index: Some(node_index),
1923 expected: destination,
1924 actual: actual_run_domain,
1925 }));
1926 }
1927 let completion_event = self.runs[run_index].enqueue(&dependencies, launch)?;
1928 let actual = completion_event.origin();
1929 if actual != completion.domain() {
1930 return Err(event_domain_error(
1931 EventDomainError::CompletionTokenDomainMismatch {
1932 operation: EventDomainOperation::ValidateCompletion,
1933 node_index: Some(node_index),
1934 expected: completion.domain(),
1935 actual,
1936 },
1937 ));
1938 }
1939 self.completions.insert(
1940 EventDependency::from_completion(completion),
1941 completion_event,
1942 );
1943 Ok(())
1944 }
1945
1946 fn classify_dependencies(
1947 &self,
1948 node_index: usize,
1949 node: &ScheduledNode,
1950 destination: super::EventDomainId,
1951 ) -> Result<SmallVec<[Arc<dyn EventToken>; 4]>> {
1952 let mut admitted = SmallVec::with_capacity(node.dependencies().len());
1953 for dependency in node.dependencies() {
1954 let dependency_completion =
1955 self.completions.get(dependency).cloned().ok_or_else(|| {
1956 Error::runtime_state_source(
1957 "Runtime::run_compiled",
1958 ErrorPhase::Execution,
1959 MissingScheduledDependencyCompletionError {
1960 dependency: *dependency,
1961 node_index,
1962 },
1963 )
1964 })?;
1965 let actual = dependency_completion.origin();
1966 if actual != dependency.domain() {
1967 return Err(event_domain_error(
1968 EventDomainError::DependencyDomainMismatch {
1969 operation: match node {
1970 ScheduledNode::Transfer(_) => EventDomainOperation::TransferBridge,
1971 ScheduledNode::Operation(_)
1972 | ScheduledNode::Collective(_)
1973 | ScheduledNode::Barrier(_) => EventDomainOperation::Enqueue,
1974 },
1975 node_index: Some(node_index),
1976 expected: dependency.domain(),
1977 actual,
1978 },
1979 ));
1980 }
1981 match node {
1982 ScheduledNode::Transfer(transfer) => {
1983 let source = transfer.source_event_domain();
1984 if actual == destination {
1985 admitted.push(dependency_completion);
1986 } else if actual == source {
1987 dependency_completion.wait().map_err(|source_error| {
1988 event_domain_error(EventDomainError::DependencyWaitFailed {
1989 operation: EventDomainOperation::TransferBridge,
1990 node_index: Some(node_index),
1991 expected: destination,
1992 actual,
1993 source: Box::new(source_error),
1994 })
1995 })?;
1996 } else {
1997 return Err(event_domain_error(
1998 EventDomainError::DependencyDomainMismatch {
1999 operation: EventDomainOperation::TransferBridge,
2000 node_index: Some(node_index),
2001 expected: source,
2002 actual,
2003 },
2004 ));
2005 }
2006 }
2007 ScheduledNode::Operation(_)
2008 | ScheduledNode::Collective(_)
2009 | ScheduledNode::Barrier(_) => {
2010 if actual != destination {
2011 return Err(event_domain_error(
2012 EventDomainError::DependencyDomainMismatch {
2013 operation: EventDomainOperation::Enqueue,
2014 node_index: Some(node_index),
2015 expected: destination,
2016 actual,
2017 },
2018 ));
2019 }
2020 admitted.push(dependency_completion);
2021 }
2022 }
2023 }
2024 Ok(admitted)
2025 }
2026
2027 fn run_index(&self, domain: super::EventDomainId) -> Result<usize> {
2028 if let Some(index) = self
2029 .runs
2030 .iter()
2031 .position(|run| run.requested_domain() == domain)
2032 {
2033 return Ok(index);
2034 }
2035 Err(missing_event_domain_driver(domain))
2036 }
2037
2038 pub(crate) fn drain(&mut self) -> Result<()> {
2039 let mut failures = Vec::new();
2040 for run in &mut self.runs {
2041 if let Err(error) = run.drain() {
2042 failures.push(error);
2043 }
2044 }
2045 let mut failures = failures.into_iter();
2046 let Some(mut error) = failures.next() else {
2047 return Ok(());
2048 };
2049 for failure in failures {
2050 error = Error::with_suppressed(error, failure);
2051 }
2052 Err(error)
2053 }
2054}
2055
2056#[derive(Debug)]
2057enum RuntimeOwnedEventDomainRunState {
2058 Pending(Box<dyn EventDomainRun>),
2059 Retired,
2060 Failed,
2061}
2062
2063#[derive(Clone, Copy, Debug, Eq, PartialEq)]
2064enum EventDomainRunTerminalState {
2065 Retired,
2066 Failed,
2067}
2068
2069impl fmt::Display for EventDomainRunTerminalState {
2070 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
2071 match self {
2072 Self::Retired => formatter.write_str("retired"),
2073 Self::Failed => formatter.write_str("failed"),
2074 }
2075 }
2076}
2077
2078#[derive(Debug)]
2079struct RuntimeOwnedEventDomainRun {
2080 requested_domain: super::EventDomainId,
2081 state: RuntimeOwnedEventDomainRunState,
2082}
2083
2084impl RuntimeOwnedEventDomainRun {
2085 fn new(requested_domain: super::EventDomainId, inner: Box<dyn EventDomainRun>) -> Self {
2086 Self {
2087 requested_domain,
2088 state: RuntimeOwnedEventDomainRunState::Pending(inner),
2089 }
2090 }
2091
2092 fn requested_domain(&self) -> super::EventDomainId {
2093 self.requested_domain
2094 }
2095
2096 fn domain(&self, operation: EventDomainOperation) -> Result<super::EventDomainId> {
2097 match &self.state {
2098 RuntimeOwnedEventDomainRunState::Pending(run) => Ok(run.domain()),
2099 RuntimeOwnedEventDomainRunState::Retired => Err(event_domain_run_state_error(
2100 operation,
2101 self.requested_domain,
2102 EventDomainRunTerminalState::Retired,
2103 )),
2104 RuntimeOwnedEventDomainRunState::Failed => Err(event_domain_run_state_error(
2105 operation,
2106 self.requested_domain,
2107 EventDomainRunTerminalState::Failed,
2108 )),
2109 }
2110 }
2111
2112 fn enqueue(
2113 &mut self,
2114 dependencies: &[Arc<dyn EventToken>],
2115 launch: &mut dyn FnMut() -> Result<()>,
2116 ) -> Result<Arc<dyn EventToken>> {
2117 match &mut self.state {
2118 RuntimeOwnedEventDomainRunState::Pending(run) => run.enqueue(dependencies, launch),
2119 RuntimeOwnedEventDomainRunState::Retired => Err(event_domain_run_state_error(
2120 EventDomainOperation::Enqueue,
2121 self.requested_domain,
2122 EventDomainRunTerminalState::Retired,
2123 )),
2124 RuntimeOwnedEventDomainRunState::Failed => Err(event_domain_run_state_error(
2125 EventDomainOperation::Enqueue,
2126 self.requested_domain,
2127 EventDomainRunTerminalState::Failed,
2128 )),
2129 }
2130 }
2131
2132 fn drain(&mut self) -> Result<()> {
2133 let run = match std::mem::replace(&mut self.state, RuntimeOwnedEventDomainRunState::Failed)
2134 {
2135 RuntimeOwnedEventDomainRunState::Pending(run) => run,
2136 RuntimeOwnedEventDomainRunState::Retired => {
2137 self.state = RuntimeOwnedEventDomainRunState::Retired;
2138 return Err(event_domain_run_state_error(
2139 EventDomainOperation::Drain,
2140 self.requested_domain,
2141 EventDomainRunTerminalState::Retired,
2142 ));
2143 }
2144 RuntimeOwnedEventDomainRunState::Failed => {
2145 self.state = RuntimeOwnedEventDomainRunState::Failed;
2146 return Err(event_domain_run_state_error(
2147 EventDomainOperation::Drain,
2148 self.requested_domain,
2149 EventDomainRunTerminalState::Failed,
2150 ));
2151 }
2152 };
2153 let domain = self.requested_domain;
2154 let mut run = run;
2155 let drain_result = catch_unwind(AssertUnwindSafe(|| run.drain()));
2156 drop_event_domain_run(run);
2157 match drain_result {
2158 Ok(Ok(())) => {
2159 self.state = RuntimeOwnedEventDomainRunState::Retired;
2160 Ok(())
2161 }
2162 Ok(Err(error)) => {
2163 self.state = RuntimeOwnedEventDomainRunState::Failed;
2164 Err(error)
2165 }
2166 Err(payload) => {
2167 self.state = RuntimeOwnedEventDomainRunState::Failed;
2168 Err(event_domain_error(EventDomainError::DrainPanicked {
2169 operation: EventDomainOperation::Drain,
2170 domain,
2171 message: safe_event_domain_panic_message(payload),
2172 }))
2173 }
2174 }
2175 }
2176}
2177
2178#[derive(Debug, thiserror::Error)]
2181#[error("{operation} used event-domain run {domain:?} after it reached terminal state {state}")]
2182pub(crate) struct EventDomainRunLifecycleError {
2183 operation: EventDomainOperation,
2184 domain: super::EventDomainId,
2185 state: EventDomainRunTerminalState,
2186}
2187
2188fn event_domain_run_state_error(
2189 operation: EventDomainOperation,
2190 domain: super::EventDomainId,
2191 state: EventDomainRunTerminalState,
2192) -> Error {
2193 Error::runtime_state_source(
2194 "Runtime::run_compiled",
2195 ErrorPhase::Execution,
2196 EventDomainRunLifecycleError {
2197 operation,
2198 domain,
2199 state,
2200 },
2201 )
2202}
2203
2204fn drop_event_domain_run(run: Box<dyn EventDomainRun>) {
2205 if catch_unwind(AssertUnwindSafe(|| drop(run))).is_err() {
2206 }
2208}
2209
2210impl Drop for RuntimeOwnedEventDomainRun {
2211 fn drop(&mut self) {
2212 let state = std::mem::replace(&mut self.state, RuntimeOwnedEventDomainRunState::Failed);
2213 if let RuntimeOwnedEventDomainRunState::Pending(run) = state {
2214 drop_event_domain_run(run);
2215 }
2216 }
2217}
2218
2219fn missing_event_domain_driver(domain: super::EventDomainId) -> Error {
2220 Error::from(EventDomainError::MissingDriver { domain })
2221}
2222
2223fn event_domain_error(source: EventDomainError) -> Error {
2224 Error::from(source)
2225}
2226
2227fn safe_event_domain_panic_message(payload: Box<dyn std::any::Any + Send + 'static>) -> String {
2228 match payload.downcast::<&'static str>() {
2229 Ok(message) => (*message).to_owned(),
2230 Err(payload) => match payload.downcast::<String>() {
2231 Ok(message) => *message,
2232 Err(_) => "non-string panic payload".to_owned(),
2233 },
2234 }
2235}
2236
2237fn scheduled_execution_cleanup_error(primary: Error, cleanup: Error) -> Error {
2238 Error::with_suppressed(primary, cleanup)
2239}
2240
2241fn instruction_execution<'a>(
2242 schedule: &'a ScheduledGraph,
2243 instruction: &ExecInstruction,
2244) -> Result<InstructionExecution<'a>> {
2245 let Some(operation_index) = instruction.semantic_operation_index else {
2246 let location = schedule.root_location();
2247 return Ok(InstructionExecution {
2248 witness: location.witness(),
2249 location,
2250 });
2251 };
2252 let location = schedule
2253 .operation_locations()
2254 .get(operation_index)
2255 .ok_or_else(|| {
2256 Error::runtime_state(
2257 "Runtime::run_compiled",
2258 ErrorPhase::Execution,
2259 format!(
2260 "instruction references semantic operation {operation_index}, but prepared schedule has {} operations",
2261 schedule.operation_locations().len()
2262 ),
2263 )
2264 })?;
2265 Ok(InstructionExecution {
2266 witness: location.witness(),
2267 location,
2268 })
2269}
2270
2271struct InstructionExecution<'a> {
2272 witness: &'a super::snapshot::ExecutableEngineSnapshot,
2273 location: &'a ExecutionLocation,
2274}
2275
2276impl InstructionExecution<'_> {
2277 fn executor(&self) -> &Arc<dyn ErasedTensorBackendExecutor> {
2278 self.witness.executor()
2279 }
2280
2281 fn location(&self) -> &ExecutionLocation {
2282 self.location
2283 }
2284}
2285
2286pub(crate) struct LocatedExecSlot<'input> {
2287 pub(crate) location: ExecutionLocation,
2288 pub(crate) value: ExecSlot<'input>,
2289}
2290
2291fn stage_instruction_inputs<'input>(
2292 instruction: &ExecInstruction,
2293 location: &ExecutionLocation,
2294 located: &mut [Vec<LocatedExecSlot<'input>>],
2295 staged: &mut [Option<ExecSlot<'input>>],
2296) -> Result<()> {
2297 for &slot in &instruction.input_slots {
2298 if staged
2299 .get(slot)
2300 .ok_or(tenferro_tensor::Error::MissingValue { slot })?
2301 .is_some()
2302 {
2303 continue;
2304 }
2305 let values = located
2306 .get_mut(slot)
2307 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2308 let value_index = values
2309 .iter()
2310 .position(|value| &value.location == location)
2311 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2312 staged[slot] = Some(values.swap_remove(value_index).value);
2313 }
2314 Ok(())
2315}
2316
2317fn validate_instruction_outputs(
2318 instruction_index: usize,
2319 instruction: &ExecInstruction,
2320 location: &ExecutionLocation,
2321 staged: &[Option<ExecSlot<'_>>],
2322) -> Result<()> {
2323 for &output_slot in &instruction.output_slots {
2324 let output = staged
2325 .get(output_slot)
2326 .and_then(Option::as_ref)
2327 .ok_or(tenferro_tensor::Error::MissingValue { slot: output_slot })?;
2328 let output = output.as_read();
2329 if !location
2330 .witness()
2331 .owns_resident_tensor(&output, location.storage_class())
2332 {
2333 return Err(Error::runtime_state_source(
2334 "Runtime::run_compiled",
2335 ErrorPhase::Execution,
2336 super::EngineExecutionContractError::OutputResidencyMismatch {
2337 instruction_index,
2338 output_slot,
2339 engine_id: location.engine_id().clone(),
2340 storage_class: location.storage_class().clone(),
2341 backend_family: output.backend_family(),
2342 allocation_domain: output.allocation_domain(),
2343 },
2344 ));
2345 }
2346 }
2347 Ok(())
2348}
2349
2350pub(crate) fn retain_instruction_results<'input>(
2351 instruction: &ExecInstruction,
2352 location: &ExecutionLocation,
2353 located: &mut [Vec<LocatedExecSlot<'input>>],
2354 staged: &mut [Option<ExecSlot<'input>>],
2355) -> Result<()> {
2356 for &slot in &instruction.input_slots {
2357 let is_output = instruction.output_slots.contains(&slot);
2358 let is_last_use = instruction
2359 .input_slots
2360 .iter()
2361 .enumerate()
2362 .any(|(index, &candidate)| {
2363 candidate == slot && instruction.last_use.get(index).copied().unwrap_or(false)
2364 });
2365 if is_last_use {
2366 located[slot].clear();
2367 if !is_output {
2368 staged[slot].take();
2369 }
2370 } else if !is_output {
2371 if let Some(value) = staged[slot].take() {
2372 located[slot].push(LocatedExecSlot {
2373 location: location.clone(),
2374 value,
2375 });
2376 }
2377 }
2378 }
2379
2380 for &slot in &instruction.output_slots {
2381 let value = staged
2382 .get_mut(slot)
2383 .and_then(Option::take)
2384 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2385 let values = located
2386 .get_mut(slot)
2387 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2388 values.clear();
2389 values.push(LocatedExecSlot {
2390 location: location.clone(),
2391 value,
2392 });
2393 }
2394 Ok(())
2395}
2396
2397fn validate_runtime_input_ingress(
2398 location: &ExecutionLocation,
2399 input: &TensorRead<'_>,
2400 slot: usize,
2401) -> Result<()> {
2402 let accepted = location
2403 .witness()
2404 .accepts_runtime_input(input, location.storage_class());
2405 if accepted {
2406 return Ok(());
2407 }
2408 Err(Error::runtime_state_source(
2409 "Runtime::run_compiled",
2410 ErrorPhase::Execution,
2411 super::InputIngressContractError::ResidencyMismatch {
2412 input_slot: slot,
2413 ingress_engine_id: location.engine_id().clone(),
2414 ingress_storage_class: location.storage_class().clone(),
2415 placement: input.placement().clone(),
2416 backend_family: input.backend_family(),
2417 allocation_domain: input.allocation_domain(),
2418 },
2419 ))
2420}
2421
2422fn execute_scheduled_transfer<'input>(
2423 transfer: &ScheduledTransfer,
2424 located: &mut [Vec<LocatedExecSlot<'input>>],
2425) -> Result<()> {
2426 let source = transfer.source_location();
2427 let destination = transfer.destination_location();
2428 let provider = transfer.provider();
2429 let values =
2430 located
2431 .get_mut(transfer.value_slot())
2432 .ok_or(tenferro_tensor::Error::MissingValue {
2433 slot: transfer.value_slot(),
2434 })?;
2435 if values.iter().any(|value| &value.location == destination) {
2436 return Err(Error::runtime_state_source(
2437 "Runtime::run_compiled",
2438 ErrorPhase::Execution,
2439 DuplicateTransferDestinationError {
2440 value_slot: transfer.value_slot(),
2441 destination: destination.clone(),
2442 },
2443 ));
2444 }
2445 let transferred = {
2446 let source_value = values
2447 .iter()
2448 .find(|value| &value.location == source)
2449 .ok_or(tenferro_tensor::Error::MissingValue {
2450 slot: transfer.value_slot(),
2451 })?;
2452 let source_read = source_value.value.as_read();
2453 let expected_dtype = source_read.dtype();
2454 let expected_shape = source_read.shape().to_vec();
2455 let transferred =
2456 provider.transfer_blocking(TransferRequest::new(source, destination, source_read))?;
2457 validate_transfer_output(destination, expected_dtype, &expected_shape, &transferred)?;
2458 transferred
2459 };
2460 values.push(LocatedExecSlot {
2461 location: destination.clone(),
2462 value: ExecSlot::Owned(transferred),
2463 });
2464 Ok(())
2465}
2466
2467fn validate_transfer_output(
2468 destination: &ExecutionLocation,
2469 expected_dtype: tenferro_tensor::DType,
2470 expected_shape: &[usize],
2471 output: &Tensor,
2472) -> Result<()> {
2473 let expected_elements = checked_transfer_element_count(expected_shape).map_err(|source| {
2474 Error::runtime_state_source(
2475 "Runtime::run_compiled",
2476 ErrorPhase::Execution,
2477 TransferError::ProviderContract { source },
2478 )
2479 })?;
2480 let contract_error = if output.dtype() != expected_dtype {
2481 Some(TransferProviderContractError::DTypeMismatch {
2482 expected: expected_dtype,
2483 actual: output.dtype(),
2484 })
2485 } else if output.shape() != expected_shape {
2486 Some(TransferProviderContractError::ShapeMismatch {
2487 expected: expected_shape.to_vec(),
2488 actual: output.shape().to_vec(),
2489 })
2490 } else if tensor_buffer_len(output) != expected_elements {
2491 Some(TransferProviderContractError::InvalidBufferLength {
2492 expected: expected_elements,
2493 actual: tensor_buffer_len(output),
2494 })
2495 } else if !destination
2496 .witness()
2497 .accepts_input_placement(output.placement(), destination.storage_class())
2498 {
2499 Some(
2500 TransferProviderContractError::DestinationPlacementMismatch {
2501 destination_engine_id: destination.engine_id().clone(),
2502 destination_storage_class: destination.storage_class().clone(),
2503 actual: output.placement().clone(),
2504 },
2505 )
2506 } else if !destination.witness().owns_resident_tensor(
2507 &TensorRead::from_tensor(output),
2508 destination.storage_class(),
2509 ) {
2510 Some(
2511 TransferProviderContractError::DestinationResidencyMismatch {
2512 destination_engine_id: destination.engine_id().clone(),
2513 destination_storage_class: destination.storage_class().clone(),
2514 actual_backend_family: TensorRead::from_tensor(output).backend_family(),
2515 actual_allocation_domain: TensorRead::from_tensor(output).allocation_domain(),
2516 },
2517 )
2518 } else {
2519 None
2520 };
2521 match contract_error {
2522 None => Ok(()),
2523 Some(source) => Err(Error::runtime_state_source(
2524 "Runtime::run_compiled",
2525 ErrorPhase::Execution,
2526 TransferError::ProviderContract { source },
2527 )),
2528 }
2529}
2530
2531fn checked_transfer_element_count(
2532 shape: &[usize],
2533) -> std::result::Result<usize, TransferProviderContractError> {
2534 tenferro_tensor::validate::checked_shape_product(
2535 "Runtime::run_compiled",
2536 "transfer source shape",
2537 shape,
2538 )
2539 .map_err(|source| TransferProviderContractError::LogicalElementCount { source })
2540}
2541
2542fn tensor_buffer_len(tensor: &Tensor) -> usize {
2543 match tensor {
2544 Tensor::F32(tensor) => tensor.buffer().len(),
2545 Tensor::F64(tensor) => tensor.buffer().len(),
2546 Tensor::I32(tensor) => tensor.buffer().len(),
2547 Tensor::I64(tensor) => tensor.buffer().len(),
2548 Tensor::Bool(tensor) => tensor.buffer().len(),
2549 Tensor::C32(tensor) => tensor.buffer().len(),
2550 Tensor::C64(tensor) => tensor.buffer().len(),
2551 }
2552}
2553
2554#[cfg(test)]
2555mod transfer_validation_tests {
2556 use std::error::Error as _;
2557
2558 use super::checked_transfer_element_count;
2559 use crate::TransferProviderContractError;
2560
2561 #[test]
2562 fn transfer_element_count_overflow_is_typed_and_preserves_source() {
2563 let error = checked_transfer_element_count(&[usize::MAX, 2]).unwrap_err();
2564
2565 assert!(matches!(
2566 error,
2567 TransferProviderContractError::LogicalElementCount { .. }
2568 ));
2569 assert!(error.source().is_some());
2570 }
2571}
2572
2573fn collect_located_outputs<'input>(
2574 program: &ExecProgram,
2575 located: &mut [Vec<LocatedExecSlot<'input>>],
2576) -> Result<Vec<Option<LocatedExecSlot<'input>>>> {
2577 let mut outputs = (0..program.n_slots).map(|_| None).collect::<Vec<_>>();
2578 for &slot in &program.output_slots {
2579 if outputs[slot].is_some() {
2580 continue;
2581 }
2582 let values = located
2583 .get_mut(slot)
2584 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2585 let value = values
2586 .pop()
2587 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2588 values.clear();
2589 outputs[slot] = Some(value);
2590 }
2591 located.iter_mut().for_each(Vec::clear);
2592 Ok(outputs)
2593}
2594
2595pub(super) fn collect_tensor_outputs_with<'input>(
2596 program: &ExecProgram,
2597 outputs: &mut [Option<LocatedExecSlot<'input>>],
2598 mut materialize: impl FnMut(&ExecutionLocation, ExecSlot<'input>) -> Result<Tensor>,
2599) -> Result<Vec<Tensor>> {
2600 program
2601 .output_slots
2602 .iter()
2603 .map(|&slot| {
2604 let located = outputs
2605 .get_mut(slot)
2606 .and_then(Option::take)
2607 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2608 materialize(&located.location, located.value)
2609 })
2610 .collect()
2611}
2612
2613fn collect_value_outputs_with<'input>(
2614 program: &ExecProgram,
2615 outputs: &mut [Option<LocatedExecSlot<'input>>],
2616 mut materialize: impl FnMut(&ExecutionLocation, ExecSlot<'input>) -> Result<TensorValue>,
2617) -> Result<Vec<TensorValue>> {
2618 program
2619 .output_slots
2620 .iter()
2621 .map(|&slot| {
2622 let located = outputs
2623 .get_mut(slot)
2624 .and_then(Option::take)
2625 .ok_or(tenferro_tensor::Error::MissingValue { slot })?;
2626 materialize(&located.location, located.value)
2627 })
2628 .collect()
2629}
2630
2631fn input_signature_reads(inputs: &[TensorRead<'_>]) -> Result<InputSignature> {
2632 InputSignature::from_reads(inputs).map_err(|source| prepare_error(Arc::new(source)))
2633}
2634
2635fn resolve_input_refs<'a>(
2636 program: &'a CompiledGraph,
2637 inputs: &'a [&'a Tensor],
2638) -> Result<RuntimeInputReads<'a>> {
2639 let expected = program.program().inputs().len();
2640 if inputs.len() > expected {
2641 return Err(Error::GraphInputCountMismatch {
2642 expected,
2643 actual: inputs.len(),
2644 });
2645 }
2646 let resolved = if inputs.is_empty() {
2647 semantic_default_inputs(program)?
2648 } else if inputs.len() == expected {
2649 inputs
2650 .iter()
2651 .map(|tensor| TensorRead::from_tensor(tensor))
2652 .collect()
2653 } else {
2654 let mut explicit = inputs.iter();
2655 let mut resolved = RuntimeInputReads::new();
2656 for value in program.program().inputs() {
2657 if let Some(retained) = program.bindings().tensor_ref_for_input(*value) {
2658 resolved.push(retained.tensor_read()?);
2659 } else {
2660 let tensor = explicit.next().ok_or_else(|| {
2661 Error::invalid_argument(
2662 "Runtime::resolve_input_refs",
2663 ErrorPhase::Execution,
2664 "inputs",
2665 format!(
2666 "expected {expected} inputs, or one explicit input for each unbound placeholder"
2667 ),
2668 )
2669 })?;
2670 resolved.push(TensorRead::from_tensor(tensor));
2671 }
2672 }
2673 if explicit.next().is_some() {
2674 return Err(Error::invalid_argument(
2675 "Runtime::resolve_input_refs",
2676 ErrorPhase::Execution,
2677 "inputs",
2678 format!("expected {expected} inputs, received {}", inputs.len()),
2679 ));
2680 }
2681 resolved
2682 };
2683 validate_ordered_input_metadata_reads(program, &resolved)?;
2684 Ok(resolved)
2685}
2686
2687fn semantic_default_inputs(program: &CompiledGraph) -> Result<RuntimeInputReads<'_>> {
2688 program
2689 .program()
2690 .inputs()
2691 .iter()
2692 .enumerate()
2693 .map(|(input_index, value)| {
2694 let retained = program
2695 .bindings()
2696 .tensor_ref_for_input(*value)
2697 .ok_or_else(|| Error::UnboundPlaceholder {
2698 input_key: format!("semantic input {input_index}"),
2699 })?;
2700 retained.tensor_read()
2701 })
2702 .collect()
2703}
2704
2705fn validate_ordered_input_metadata_reads(
2706 program: &CompiledGraph,
2707 inputs: &[TensorRead<'_>],
2708) -> Result<()> {
2709 let actuals = inputs
2710 .iter()
2711 .map(|input| (input.dtype(), input.shape().to_vec()))
2712 .collect::<Vec<_>>();
2713 validate_ordered_input_metadata_values(program, &actuals, "Runtime::submit")
2714}
2715
2716fn validate_ordered_input_metadata_values(
2717 program: &CompiledGraph,
2718 actuals: &[(tenferro_tensor::DType, Vec<usize>)],
2719 caller: &'static str,
2720) -> Result<()> {
2721 let expected = program.input_count();
2722 if actuals.len() != expected {
2723 return Err(Error::GraphInputCountMismatch {
2724 expected,
2725 actual: actuals.len(),
2726 });
2727 }
2728 let input_shapes: RuntimeInputShapes<'_> =
2729 actuals.iter().map(|(_, shape)| shape.as_slice()).collect();
2730 for (input_value, (actual_dtype, actual_shape)) in
2731 program.program().inputs().iter().zip(actuals)
2732 {
2733 let metadata = program
2734 .program()
2735 .value_metadata(*input_value)
2736 .map_err(|source| Error::runtime_state_source(caller, ErrorPhase::Execution, source))?;
2737 if metadata.dtype() != *actual_dtype {
2738 return Err(Error::PlaceholderDtypeMismatch {
2739 expected: metadata.dtype(),
2740 actual: *actual_dtype,
2741 });
2742 }
2743 if metadata.shape().len() != actual_shape.len() {
2744 return Err(Error::PlaceholderRankMismatch {
2745 expected: metadata.shape().len(),
2746 actual: actual_shape.len(),
2747 });
2748 }
2749 let mut expected_shape: RuntimeShapeScratch = actual_shape.iter().copied().collect();
2750 let mut exact_mismatch = false;
2751 for (axis, (extent, actual_size)) in metadata.shape().iter().zip(actual_shape).enumerate() {
2752 match extent {
2753 ShapeExtent::Exact(expression) => {
2754 let expected = expression.eval(&input_shapes).map_err(|source| {
2755 Error::runtime_state_source(caller, ErrorPhase::Execution, source)
2756 })?;
2757 expected_shape[axis] = expected;
2758 exact_mismatch |= expected != *actual_size;
2759 }
2760 ShapeExtent::UpperBound(expression) => {
2761 let bound = expression.eval(&input_shapes).map_err(|source| {
2762 Error::runtime_state_source(caller, ErrorPhase::Execution, source)
2763 })?;
2764 if *actual_size > bound {
2765 return Err(Error::PlaceholderShapeBoundExceeded {
2766 axis,
2767 bound,
2768 actual: *actual_size,
2769 });
2770 }
2771 }
2772 ShapeExtent::Unknown => {}
2773 }
2774 }
2775 if exact_mismatch {
2776 return Err(Error::PlaceholderShapeMismatch {
2777 expected: expected_shape.into_vec(),
2778 actual: actual_shape.clone(),
2779 });
2780 }
2781 }
2782 Ok(())
2783}
2784
2785fn validate_exec_input_count(program: &ExecProgram, actual: usize) -> Result<()> {
2786 let expected = program.input_slots.len();
2787 if actual != expected {
2788 return Err(Error::runtime_state(
2789 "Runtime::run_compiled",
2790 ErrorPhase::Execution,
2791 format!("expected {expected} inputs for execution program, got {actual}"),
2792 ));
2793 }
2794 Ok(())
2795}
2796
2797fn prepare_error(source: Arc<PrepareError>) -> Error {
2798 Error::runtime_state_source(
2799 "Runtime::run_compiled",
2800 ErrorPhase::Execution,
2801 SharedPrepareError(source),
2802 )
2803}
2804
2805#[derive(Clone, Debug)]
2806struct SharedPrepareError(Arc<PrepareError>);
2807
2808impl fmt::Display for SharedPrepareError {
2809 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
2810 write!(formatter, "{}", self.0)
2811 }
2812}
2813
2814impl StdError for SharedPrepareError {
2815 fn source(&self) -> Option<&(dyn StdError + 'static)> {
2816 Some(self.0.as_ref())
2817 }
2818}
2819
2820#[cfg(test)]
2821mod tests {
2822 use std::any::Any;
2823 use std::error::Error as StdError;
2824 use std::hash::Hasher;
2825 use std::num::NonZeroU64;
2826 use std::sync::atomic::{AtomicBool, Ordering};
2827 use std::sync::{Arc, TryLockError, Weak};
2828
2829 use tenferro_cpu::CpuBackend;
2830 use tenferro_ops::dim_expr::DimExpr;
2831 use tenferro_ops::ext_op::ExtensionOp;
2832 use tenferro_ops::SymDim;
2833 use tenferro_tensor::{
2834 BackendSession, BackendSessionHost, DType, Tensor, TensorBackend, TensorRead, TypedTensor,
2835 };
2836
2837 use crate::exec::{ExecInstruction, ExecOp, ExecProgram, ExecSlot};
2838 use crate::runtime::{
2839 CoreCapabilityBundle, EngineId, ErasedExecutionContext, EventDomainDriver, EventDomainId,
2840 ExecutableEngineContract, ExecutionContextIdentity, HardwareClassId, InputIngressContract,
2841 InputSignature, PreparedOperation, PreparedOperationBinding, PreparedOperationExecutor,
2842 PreparedOperationPlan, ProviderDeviceIdentity, ProviderId, RegistrationIdentity,
2843 RuntimeCacheOwner, RuntimeEpoch, RuntimeId, SpecializationProjection,
2844 SpecializationRequirements, StorageClass,
2845 };
2846 use crate::{Error, ErrorPhase, ExtensionCacheStore, Result};
2847
2848 use super::{
2849 DuplicateTransferDestinationError, ErasedTensorBackendExecutor, LocatedExecSlot,
2850 TensorBackendExecutor,
2851 };
2852 use crate::runtime::schedule::{
2853 EventCompletion, EventSlotId, ExecutionLocation, ScheduledTransfer,
2854 };
2855
2856 const LOCK_PROBE_FAMILY: &str = "runtime.lock-probe.v1";
2857 const REENTRANT_PROBE_FAMILY: &str = "runtime.reentrant-probe.v1";
2858
2859 #[derive(Clone, Debug)]
2860 struct LockProbeOp;
2861
2862 impl ExtensionOp for LockProbeOp {
2863 fn family_id(&self) -> &'static str {
2864 LOCK_PROBE_FAMILY
2865 }
2866
2867 fn payload_hash(&self, _hasher: &mut dyn Hasher) {}
2868
2869 fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
2870 other.as_any().downcast_ref::<Self>().is_some()
2871 }
2872
2873 fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
2874 Arc::new(self.clone())
2875 }
2876
2877 fn as_any(&self) -> &dyn Any {
2878 self
2879 }
2880
2881 fn input_count(&self) -> usize {
2882 1
2883 }
2884
2885 fn output_count(&self) -> usize {
2886 1
2887 }
2888
2889 fn infer_output_meta(
2890 &self,
2891 ctx: &mut tenferro_ops::ExtensionShapeContext<'_>,
2892 ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
2893 Ok(vec![(ctx.input_dtype(0)?, ctx.input_shape(0)?.to_vec())])
2894 }
2895 }
2896
2897 #[derive(Clone, Debug)]
2898 struct ReentrantProbeOp;
2899
2900 impl ExtensionOp for ReentrantProbeOp {
2901 fn family_id(&self) -> &'static str {
2902 REENTRANT_PROBE_FAMILY
2903 }
2904
2905 fn payload_hash(&self, _hasher: &mut dyn Hasher) {}
2906
2907 fn payload_eq(&self, other: &dyn ExtensionOp) -> bool {
2908 other.as_any().downcast_ref::<Self>().is_some()
2909 }
2910
2911 fn clone_arc(&self) -> Arc<dyn ExtensionOp> {
2912 Arc::new(self.clone())
2913 }
2914
2915 fn as_any(&self) -> &dyn Any {
2916 self
2917 }
2918
2919 fn input_count(&self) -> usize {
2920 1
2921 }
2922
2923 fn output_count(&self) -> usize {
2924 1
2925 }
2926
2927 fn infer_output_meta(
2928 &self,
2929 ctx: &mut tenferro_ops::ExtensionShapeContext<'_>,
2930 ) -> tenferro_tensor::Result<Vec<(DType, Vec<SymDim>)>> {
2931 Ok(vec![(ctx.input_dtype(0)?, ctx.input_shape(0)?.to_vec())])
2932 }
2933 }
2934
2935 #[derive(Debug)]
2936 struct LockProbePreparedOperation {
2937 binding: PreparedOperationBinding,
2938 specialization: SpecializationProjection,
2939 executor: Weak<TensorBackendExecutor<CpuBackend>>,
2940 observed_unlocked_state: Arc<AtomicBool>,
2941 }
2942
2943 impl PreparedOperation for LockProbePreparedOperation {
2944 fn binding(&self) -> &PreparedOperationBinding {
2945 &self.binding
2946 }
2947
2948 fn specialization(&self) -> &SpecializationProjection {
2949 &self.specialization
2950 }
2951
2952 fn retained_bytes(&self) -> usize {
2953 0
2954 }
2955 }
2956
2957 impl PreparedOperationExecutor for LockProbePreparedOperation {
2958 fn execute(
2959 &self,
2960 context: &mut ErasedExecutionContext<'_>,
2961 _extension_caches: &mut ExtensionCacheStore,
2962 inputs: &[TensorRead<'_>],
2963 ) -> Result<Vec<Tensor>> {
2964 let executor = self.executor.upgrade().expect("executor still alive");
2965 let unlocked = match executor.state.try_lock() {
2966 Ok(_guard) => true,
2967 Err(TryLockError::WouldBlock) => false,
2968 Err(TryLockError::Poisoned(_)) => false,
2969 };
2970 self.observed_unlocked_state
2971 .store(unlocked, Ordering::SeqCst);
2972 let backend = context
2973 .downcast_mut::<CpuBackend>(self.binding.context_identity())
2974 .map_err(|source| {
2975 Error::runtime_state_source("lock_probe", ErrorPhase::Execution, source)
2976 })?;
2977 Ok(vec![backend.with_backend_session(|exec| {
2978 exec.to_contiguous_read(inputs[0].clone())
2979 })?])
2980 }
2981 }
2982
2983 #[derive(Debug)]
2984 struct ReentrantProbePreparedOperation {
2985 binding: PreparedOperationBinding,
2986 specialization: SpecializationProjection,
2987 executor: Weak<TensorBackendExecutor<CpuBackend>>,
2988 observed_reentrant_error: Arc<AtomicBool>,
2989 }
2990
2991 #[derive(Debug)]
2992 struct SessionProbePreparedOperation {
2993 binding: PreparedOperationBinding,
2994 specialization: SpecializationProjection,
2995 observed_session: Arc<AtomicBool>,
2996 }
2997
2998 impl PreparedOperation for SessionProbePreparedOperation {
2999 fn binding(&self) -> &PreparedOperationBinding {
3000 &self.binding
3001 }
3002
3003 fn specialization(&self) -> &SpecializationProjection {
3004 &self.specialization
3005 }
3006
3007 fn retained_bytes(&self) -> usize {
3008 0
3009 }
3010 }
3011
3012 impl PreparedOperationExecutor for SessionProbePreparedOperation {
3013 fn execute(
3014 &self,
3015 _context: &mut ErasedExecutionContext<'_>,
3016 _extension_caches: &mut ExtensionCacheStore,
3017 _inputs: &[TensorRead<'_>],
3018 ) -> Result<Vec<Tensor>> {
3019 Err(Error::unsupported(
3020 "session_probe",
3021 ErrorPhase::Execution,
3022 "session probe must use the scheduler-owned session path",
3023 ))
3024 }
3025
3026 fn supports_session(&self) -> bool {
3027 true
3028 }
3029
3030 fn execute_in_session(
3031 &self,
3032 session: &mut dyn BackendSession,
3033 _extension_caches: &mut ExtensionCacheStore,
3034 inputs: &[TensorRead<'_>],
3035 ) -> Result<Vec<Tensor>> {
3036 self.observed_session.store(true, Ordering::SeqCst);
3037 Ok(vec![session.to_contiguous_read(inputs[0].clone())?])
3038 }
3039 }
3040
3041 impl ReentrantProbePreparedOperation {
3042 fn probe_reentrant_call(&self, input: Tensor) {
3043 let executor = self.executor.upgrade().expect("executor still alive");
3044 let nested = ErasedTensorBackendExecutor::execute(
3045 executor.as_ref(),
3046 &passthrough_program(),
3047 &[],
3048 vec![input],
3049 );
3050 let observed = nested.is_err_and(|error| {
3051 error
3052 .to_string()
3053 .contains("reentrant tensor backend executor call would deadlock")
3054 });
3055 self.observed_reentrant_error
3056 .store(observed, Ordering::SeqCst);
3057 }
3058 }
3059
3060 impl PreparedOperation for ReentrantProbePreparedOperation {
3061 fn binding(&self) -> &PreparedOperationBinding {
3062 &self.binding
3063 }
3064
3065 fn specialization(&self) -> &SpecializationProjection {
3066 &self.specialization
3067 }
3068
3069 fn retained_bytes(&self) -> usize {
3070 0
3071 }
3072 }
3073
3074 impl PreparedOperationExecutor for ReentrantProbePreparedOperation {
3075 fn execute(
3076 &self,
3077 context: &mut ErasedExecutionContext<'_>,
3078 _extension_caches: &mut ExtensionCacheStore,
3079 inputs: &[TensorRead<'_>],
3080 ) -> Result<Vec<Tensor>> {
3081 let backend = context
3082 .downcast_mut::<CpuBackend>(self.binding.context_identity())
3083 .map_err(|source| {
3084 Error::runtime_state_source("reentrant_probe", ErrorPhase::Execution, source)
3085 })?;
3086 let materialized =
3087 backend.with_backend_session(|exec| exec.to_contiguous_read(inputs[0].clone()))?;
3088 self.probe_reentrant_call(materialized.duplicate()?);
3089 Ok(vec![materialized])
3090 }
3091 }
3092
3093 fn lock_probe_program() -> ExecProgram {
3094 ExecProgram {
3095 instructions: vec![ExecInstruction {
3096 op: ExecOp::Extension(Arc::new(LockProbeOp)),
3097 semantic_operation_index: Some(0),
3098 input_slots: vec![0],
3099 output_slots: vec![1],
3100 dtype: DType::F64,
3101 output_shapes: vec![vec![DimExpr::InputDim {
3102 input_idx: 0,
3103 axis: 0,
3104 }]]
3105 .into(),
3106 output_extents: vec![vec![]].into(),
3107 last_use: vec![false],
3108 }],
3109 input_slots: vec![0],
3110 output_slots: vec![1],
3111 n_slots: 2,
3112 shape_guards: vec![],
3113 }
3114 }
3115
3116 fn partial_session_probe_program() -> ExecProgram {
3117 let output_shape = || {
3118 vec![DimExpr::InputDim {
3119 input_idx: 0,
3120 axis: 0,
3121 }]
3122 };
3123 ExecProgram {
3124 instructions: vec![
3125 ExecInstruction {
3126 op: ExecOp::Extension(Arc::new(LockProbeOp)),
3127 semantic_operation_index: Some(0),
3128 input_slots: vec![0],
3129 output_slots: vec![1],
3130 dtype: DType::F64,
3131 output_shapes: vec![output_shape()].into(),
3132 output_extents: vec![vec![]].into(),
3133 last_use: vec![false],
3134 },
3135 ExecInstruction {
3136 op: ExecOp::Extension(Arc::new(LockProbeOp)),
3137 semantic_operation_index: Some(1),
3138 input_slots: vec![1],
3139 output_slots: vec![2],
3140 dtype: DType::F64,
3141 output_shapes: vec![output_shape()].into(),
3142 output_extents: vec![vec![]].into(),
3143 last_use: vec![false],
3144 },
3145 ],
3146 input_slots: vec![0],
3147 output_slots: vec![2],
3148 n_slots: 3,
3149 shape_guards: vec![],
3150 }
3151 }
3152
3153 fn reentrant_probe_program() -> ExecProgram {
3154 ExecProgram {
3155 instructions: vec![ExecInstruction {
3156 op: ExecOp::Extension(Arc::new(ReentrantProbeOp)),
3157 semantic_operation_index: Some(0),
3158 input_slots: vec![0],
3159 output_slots: vec![1],
3160 dtype: DType::F64,
3161 output_shapes: vec![vec![DimExpr::InputDim {
3162 input_idx: 0,
3163 axis: 0,
3164 }]]
3165 .into(),
3166 output_extents: vec![vec![]].into(),
3167 last_use: vec![false],
3168 }],
3169 input_slots: vec![0],
3170 output_slots: vec![1],
3171 n_slots: 2,
3172 shape_guards: vec![],
3173 }
3174 }
3175
3176 fn passthrough_program() -> ExecProgram {
3177 ExecProgram {
3178 instructions: vec![],
3179 input_slots: vec![0],
3180 output_slots: vec![0],
3181 n_slots: 1,
3182 shape_guards: vec![],
3183 }
3184 }
3185
3186 fn f64_zeros(shape: Vec<usize>) -> Tensor {
3187 Tensor::F64(TypedTensor::zeros(shape).unwrap())
3188 }
3189
3190 fn nz(value: u64) -> NonZeroU64 {
3191 NonZeroU64::new(value).unwrap_or(NonZeroU64::MIN)
3192 }
3193
3194 fn probe_binding() -> PreparedOperationBinding {
3195 PreparedOperationBinding::new(
3196 RuntimeId::from_nonzero(nz(1)),
3197 RuntimeEpoch::from_nonzero(nz(2)),
3198 EngineId::new("tenferro.cpu").unwrap(),
3199 RegistrationIdentity::new(nz(3), nz(4)),
3200 ExecutionContextIdentity::of::<CpuBackend>(),
3201 HardwareClassId::new("tenferro.cpu.host").unwrap(),
3202 )
3203 }
3204
3205 fn probe_specialization() -> SpecializationProjection {
3206 SpecializationRequirements::polymorphic(0)
3207 .project(&InputSignature::new(Vec::new()))
3208 .unwrap()
3209 }
3210
3211 #[test]
3212 fn tensor_backend_executor_releases_state_lock_during_extension_execution() {
3213 let executor = Arc::new(TensorBackendExecutor::<CpuBackend>::new(CpuBackend::new()));
3214 let observed_unlocked_state = Arc::new(AtomicBool::new(false));
3215 let prepared = Arc::new(LockProbePreparedOperation {
3216 binding: probe_binding(),
3217 specialization: probe_specialization(),
3218 executor: Arc::downgrade(&executor),
3219 observed_unlocked_state: Arc::clone(&observed_unlocked_state),
3220 });
3221 let operations = vec![PreparedOperationPlan::executable(
3222 prepared.clone(),
3223 prepared,
3224 )];
3225
3226 let output = ErasedTensorBackendExecutor::execute(
3227 executor.as_ref(),
3228 &lock_probe_program(),
3229 &operations,
3230 vec![f64_zeros(vec![2])],
3231 )
3232 .expect("extension executes");
3233
3234 assert_eq!(output[0].shape(), &[2]);
3235 assert!(
3236 observed_unlocked_state.load(Ordering::SeqCst),
3237 "executor state lock must not be held while extension runtime callbacks execute"
3238 );
3239 }
3240
3241 #[test]
3242 fn scheduled_execution_cleanup_error_preserves_primary_error() {
3243 let primary = Error::runtime_state(
3244 "primary-execution",
3245 ErrorPhase::Execution,
3246 "primary execution failure",
3247 );
3248 let cleanup = Error::runtime_state(
3249 "event-domain-cleanup",
3250 ErrorPhase::Execution,
3251 "cleanup failure",
3252 );
3253 let primary_display = primary.to_string();
3254
3255 let combined = super::scheduled_execution_cleanup_error(primary, cleanup);
3256 let primary_error = combined.primary().expect("primary execution error");
3257 let cleanup_error = combined.suppressed().expect("cleanup error");
3258 assert_eq!(primary_error.to_string(), primary_display);
3259 assert_eq!(
3260 cleanup_error.to_string(),
3261 "event-domain-cleanup (Execution): runtime state failure: cleanup failure"
3262 );
3263 assert_eq!(
3264 StdError::source(&combined)
3265 .expect("primary error in the standard source chain")
3266 .to_string(),
3267 primary_display
3268 );
3269 }
3270
3271 #[test]
3272 fn duplicate_scheduled_transfer_destination_reports_typed_fields() {
3273 let domain = EventDomainId::runtime_created_for_test(
3274 RuntimeId::from_nonzero(nz(1)),
3275 RuntimeEpoch::from_nonzero(nz(1)),
3276 RegistrationIdentity::new(nz(1), nz(1)),
3277 );
3278 let source = ExecutionLocation::new(
3279 EngineId::new("tenferro-test.transfer-source").expect("source engine"),
3280 ProviderDeviceIdentity::new(
3281 ProviderId::new("tenferro-test.transfer").expect("provider id"),
3282 "source",
3283 )
3284 .expect("source provider target"),
3285 domain,
3286 StorageClass::new("tenferro-test.transfer-source").expect("source storage"),
3287 );
3288 let destination = ExecutionLocation::new(
3289 EngineId::new("tenferro-test.transfer-destination").expect("destination engine"),
3290 ProviderDeviceIdentity::new(
3291 ProviderId::new("tenferro-test.transfer").expect("provider id"),
3292 "destination",
3293 )
3294 .expect("destination provider target"),
3295 domain,
3296 StorageClass::new("tenferro-test.transfer-destination").expect("destination storage"),
3297 );
3298 let transfer = ScheduledTransfer::new(
3299 3,
3300 source,
3301 destination.clone(),
3302 [],
3303 EventCompletion::new(domain, EventSlotId::new(0), 0),
3304 );
3305 let mut located = vec![
3306 Vec::new(),
3307 Vec::new(),
3308 Vec::new(),
3309 vec![LocatedExecSlot {
3310 location: destination.clone(),
3311 value: ExecSlot::Owned(f64_zeros(vec![1])),
3312 }],
3313 ];
3314
3315 let error = super::execute_scheduled_transfer(&transfer, &mut located)
3316 .expect_err("duplicate transfer destination");
3317 let Error::RuntimeStateSource { source, .. } = error else {
3318 panic!("duplicate transfer destination must retain a typed source");
3319 };
3320 let duplicate = source
3321 .downcast_ref::<DuplicateTransferDestinationError>()
3322 .expect("typed duplicate transfer destination source");
3323 assert_eq!(duplicate.value_slot, 3);
3324 assert_eq!(duplicate.destination, destination);
3325 }
3326
3327 #[test]
3328 fn tensor_backend_executor_reentrant_call_returns_error_instead_of_deadlocking() {
3329 let executor = Arc::new(TensorBackendExecutor::<CpuBackend>::new(CpuBackend::new()));
3330 let observed_reentrant_error = Arc::new(AtomicBool::new(false));
3331 let prepared = Arc::new(ReentrantProbePreparedOperation {
3332 binding: probe_binding(),
3333 specialization: probe_specialization(),
3334 executor: Arc::downgrade(&executor),
3335 observed_reentrant_error: Arc::clone(&observed_reentrant_error),
3336 });
3337 let operations = vec![PreparedOperationPlan::executable(
3338 prepared.clone(),
3339 prepared,
3340 )];
3341
3342 let output = ErasedTensorBackendExecutor::execute(
3343 executor.as_ref(),
3344 &reentrant_probe_program(),
3345 &operations,
3346 vec![f64_zeros(vec![2])],
3347 )
3348 .expect("outer extension executes");
3349
3350 assert_eq!(output[0].shape(), &[2]);
3351 assert!(
3352 observed_reentrant_error.load(Ordering::SeqCst),
3353 "same-thread reentrant executor call must fail immediately instead of deadlocking"
3354 );
3355 }
3356
3357 #[test]
3358 fn tensor_backend_executor_dispatches_session_capable_extension_in_one_session_path() {
3359 let executor = TensorBackendExecutor::<CpuBackend>::new(CpuBackend::new());
3360 let observed_session = Arc::new(AtomicBool::new(false));
3361 let prepared = Arc::new(SessionProbePreparedOperation {
3362 binding: probe_binding(),
3363 specialization: probe_specialization(),
3364 observed_session: Arc::clone(&observed_session),
3365 });
3366 let operations = vec![PreparedOperationPlan::executable(
3367 prepared.clone(),
3368 prepared,
3369 )];
3370
3371 let output = ErasedTensorBackendExecutor::execute(
3372 &executor,
3373 &lock_probe_program(),
3374 &operations,
3375 vec![f64_zeros(vec![2])],
3376 )
3377 .expect("session-capable extension executes");
3378
3379 assert_eq!(output[0].shape(), &[2]);
3380 assert!(observed_session.load(Ordering::SeqCst));
3381 }
3382
3383 #[test]
3384 fn tensor_backend_executor_batches_session_capable_region_after_boundary() {
3385 let executor = Arc::new(TensorBackendExecutor::<CpuBackend>::new(CpuBackend::new()));
3386 let observed_unlocked_state = Arc::new(AtomicBool::new(false));
3387 let observed_session = Arc::new(AtomicBool::new(false));
3388 let ordinary = Arc::new(LockProbePreparedOperation {
3389 binding: probe_binding(),
3390 specialization: probe_specialization(),
3391 executor: Arc::downgrade(&executor),
3392 observed_unlocked_state: Arc::clone(&observed_unlocked_state),
3393 });
3394 let session = Arc::new(SessionProbePreparedOperation {
3395 binding: probe_binding(),
3396 specialization: probe_specialization(),
3397 observed_session: Arc::clone(&observed_session),
3398 });
3399 let operations = vec![
3400 PreparedOperationPlan::executable(ordinary.clone(), ordinary),
3401 PreparedOperationPlan::executable(session.clone(), session),
3402 ];
3403
3404 let output = ErasedTensorBackendExecutor::execute(
3405 executor.as_ref(),
3406 &partial_session_probe_program(),
3407 &operations,
3408 vec![f64_zeros(vec![2])],
3409 )
3410 .expect("mixed session regions execute");
3411
3412 assert_eq!(output[0].shape(), &[2]);
3413 assert!(observed_unlocked_state.load(Ordering::SeqCst));
3414 assert!(observed_session.load(Ordering::SeqCst));
3415
3416 let values = ErasedTensorBackendExecutor::execute_values(
3417 executor.as_ref(),
3418 &partial_session_probe_program(),
3419 &operations,
3420 vec![f64_zeros(vec![2])],
3421 )
3422 .expect("mixed session regions execute in value mode");
3423 assert_eq!(values[0].shape(), &[2]);
3424 }
3425
3426 #[test]
3427 fn tensor_backend_executor_bridge_does_not_require_clone_source_contract() {
3428 type ContractConstructor<B> = fn(
3429 ProviderDeviceIdentity,
3430 CoreCapabilityBundle,
3431 B,
3432 Arc<dyn EventDomainDriver>,
3433 InputIngressContract,
3434 Option<Arc<dyn RuntimeCacheOwner>>,
3435 ) -> ExecutableEngineContract;
3436
3437 fn factory_accepts_backend_without_clone_bound<B>()
3438 where
3439 B: TensorBackend + Send + Sync + 'static,
3440 {
3441 let _factory: fn(B) -> Arc<dyn ErasedTensorBackendExecutor> =
3442 super::erased_tensor_backend_executor::<B>;
3443 }
3444
3445 fn contract_accepts_backend_without_clone_bound<B>()
3446 where
3447 B: TensorBackend + Send + Sync + 'static,
3448 {
3449 let _constructor: ContractConstructor<B> = ExecutableEngineContract::new::<B>;
3450 }
3451
3452 factory_accepts_backend_without_clone_bound::<CpuBackend>();
3453 contract_accepts_backend_without_clone_bound::<CpuBackend>();
3454 }
3455}