tenferro_runtime/runtime/
schedule.rs

1//! Runtime-owned scheduled graph boundary.
2//!
3//! This representation remains crate-private. Later phases attach native
4//! GPU/XLA dispatch and asynchronous completion to the same node families.
5
6use std::collections::HashSet;
7use std::hash::{Hash, Hasher};
8#[cfg(test)]
9use std::num::NonZeroU64;
10use std::sync::Arc;
11
12use super::snapshot::ExecutableEngineSnapshot;
13use super::{
14    FrozenTransferRegistry, RegistrationIdentity, ResolvedTransferEndpoint, ResolvedTransferRoute,
15    RuntimeEpoch, RuntimeId,
16};
17use crate::exec::ExecProgram;
18use crate::{EngineId, StorageClass, TransferEndpoint};
19
20/// Provenance-qualified identity of one frozen direct-engine event domain.
21///
22/// An event domain is identified by the runtime that owns it, the immutable
23/// configuration epoch that published it, and the frozen registration identity
24/// of the direct engine slot. It is not a reusable slot number.
25///
26/// # Examples
27///
28/// ```
29/// use tenferro_runtime::EventDomainId;
30///
31/// fn inspect(domain: EventDomainId) {
32///     let _runtime = domain.runtime_id();
33///     let _epoch = domain.epoch();
34///     let _registration = domain.registration_identity();
35/// }
36/// ```
37#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
38pub struct EventDomainId {
39    runtime_id: RuntimeId,
40    epoch: RuntimeEpoch,
41    registration_identity: RegistrationIdentity,
42}
43
44impl EventDomainId {
45    pub(crate) const fn new(
46        runtime_id: RuntimeId,
47        epoch: RuntimeEpoch,
48        registration_identity: RegistrationIdentity,
49    ) -> Self {
50        Self {
51            runtime_id,
52            epoch,
53            registration_identity,
54        }
55    }
56
57    /// Return the runtime that owns this event domain.
58    ///
59    /// # Examples
60    ///
61    /// ```
62    /// use tenferro_runtime::EventDomainId;
63    ///
64    /// fn inspect(domain: EventDomainId) {
65    ///     let _runtime = domain.runtime_id();
66    /// }
67    /// ```
68    #[must_use]
69    pub const fn runtime_id(self) -> RuntimeId {
70        self.runtime_id
71    }
72
73    /// Return the immutable configuration epoch that published this domain.
74    ///
75    /// # Examples
76    ///
77    /// ```
78    /// use tenferro_runtime::EventDomainId;
79    ///
80    /// fn inspect(domain: EventDomainId) {
81    ///     let _epoch = domain.epoch();
82    /// }
83    /// ```
84    #[must_use]
85    pub const fn epoch(self) -> RuntimeEpoch {
86        self.epoch
87    }
88
89    /// Return the frozen direct-engine registration identity for this domain.
90    ///
91    /// # Examples
92    ///
93    /// ```
94    /// use tenferro_runtime::EventDomainId;
95    ///
96    /// fn inspect(domain: EventDomainId) {
97    ///     let _registration = domain.registration_identity();
98    /// }
99    /// ```
100    #[must_use]
101    pub const fn registration_identity(self) -> RegistrationIdentity {
102        self.registration_identity
103    }
104
105    #[cfg(test)]
106    pub(crate) const fn runtime_created_for_test(
107        runtime_id: RuntimeId,
108        epoch: RuntimeEpoch,
109        registration_identity: RegistrationIdentity,
110    ) -> Self {
111        Self::new(runtime_id, epoch, registration_identity)
112    }
113}
114
115#[derive(Clone, Debug)]
116pub(crate) struct ExecutionLocation {
117    resolved_endpoint: ResolvedTransferEndpoint,
118    witness: Arc<ExecutableEngineSnapshot>,
119}
120
121impl ExecutionLocation {
122    pub(super) fn from_witness(
123        witness: Arc<ExecutableEngineSnapshot>,
124        storage_class: StorageClass,
125    ) -> Self {
126        Self {
127            resolved_endpoint: ResolvedTransferEndpoint::new(
128                TransferEndpoint::new(witness.engine_id().clone(), storage_class),
129                witness.provider_device_identity().clone(),
130                witness.event_domain_id(),
131            ),
132            witness,
133        }
134    }
135
136    #[cfg(test)]
137    pub(crate) fn new(
138        engine_id: EngineId,
139        provider_device_identity: super::ProviderDeviceIdentity,
140        event_domain_id: EventDomainId,
141        storage_class: StorageClass,
142    ) -> Self {
143        Self::from_witness(
144            super::snapshot::ExecutableEngineSnapshot::for_test(
145                engine_id,
146                provider_device_identity,
147                event_domain_id,
148                storage_class.clone(),
149            ),
150            storage_class,
151        )
152    }
153
154    #[cfg(test)]
155    fn for_test_with_driver(
156        domain: EventDomainId,
157        provider_device_identity: super::ProviderDeviceIdentity,
158        event_domain_driver: Arc<dyn super::EventDomainDriver>,
159    ) -> Self {
160        let engine_id = EngineId::new("tenferro-test.schedule-engine").expect("test engine id");
161        let storage_class =
162            StorageClass::new("tenferro-test.schedule-storage").expect("test storage class");
163        Self::from_test_with_driver(
164            engine_id,
165            provider_device_identity,
166            domain,
167            storage_class,
168            event_domain_driver,
169        )
170    }
171
172    #[cfg(test)]
173    fn from_test_with_driver(
174        engine_id: EngineId,
175        provider_device_identity: super::ProviderDeviceIdentity,
176        event_domain_id: EventDomainId,
177        storage_class: StorageClass,
178        event_domain_driver: Arc<dyn super::EventDomainDriver>,
179    ) -> Self {
180        Self::from_witness(
181            super::snapshot::ExecutableEngineSnapshot::for_test_with_driver(
182                engine_id,
183                provider_device_identity,
184                event_domain_id,
185                storage_class.clone(),
186                event_domain_driver,
187            ),
188            storage_class,
189        )
190    }
191
192    pub(crate) fn engine_id(&self) -> &EngineId {
193        self.resolved_endpoint.logical().engine_id()
194    }
195
196    pub(crate) fn endpoint(&self) -> &TransferEndpoint {
197        self.resolved_endpoint.logical()
198    }
199
200    pub(crate) fn event_domain_id(&self) -> EventDomainId {
201        self.resolved_endpoint.event_domain_id()
202    }
203
204    pub(crate) fn provider_device_identity(&self) -> &super::ProviderDeviceIdentity {
205        self.resolved_endpoint.provider_device_identity()
206    }
207
208    pub(crate) fn resolved_endpoint(&self) -> &ResolvedTransferEndpoint {
209        &self.resolved_endpoint
210    }
211
212    pub(crate) fn storage_class(&self) -> &StorageClass {
213        self.resolved_endpoint.logical().storage_class()
214    }
215
216    pub(super) fn witness(&self) -> &Arc<ExecutableEngineSnapshot> {
217        &self.witness
218    }
219
220    #[cfg(test)]
221    fn for_test(
222        domain: EventDomainId,
223        provider_device_identity: super::ProviderDeviceIdentity,
224    ) -> Self {
225        Self::new(
226            EngineId::new("tenferro-test.schedule-engine").expect("test engine id"),
227            provider_device_identity,
228            domain,
229            StorageClass::new("tenferro-test.schedule-storage").expect("test storage class"),
230        )
231    }
232}
233
234impl PartialEq for ExecutionLocation {
235    fn eq(&self, other: &Self) -> bool {
236        self.resolved_endpoint == other.resolved_endpoint
237    }
238}
239
240impl Eq for ExecutionLocation {}
241
242impl Hash for ExecutionLocation {
243    fn hash<H: Hasher>(&self, state: &mut H) {
244        self.resolved_endpoint.hash(state);
245    }
246}
247
248#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
249pub(crate) struct EventSlotId(u32);
250
251impl EventSlotId {
252    pub(crate) const fn new(value: u32) -> Self {
253        Self(value)
254    }
255}
256
257#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
258pub(crate) struct EventDependency {
259    domain: EventDomainId,
260    slot: EventSlotId,
261    generation: u64,
262}
263
264impl EventDependency {
265    pub(crate) fn new(domain: EventDomainId, slot: EventSlotId, generation: u64) -> Self {
266        Self {
267            domain,
268            slot,
269            generation,
270        }
271    }
272
273    pub(crate) fn domain(&self) -> EventDomainId {
274        self.domain
275    }
276
277    pub(crate) fn from_completion(completion: EventCompletion) -> Self {
278        Self::new(completion.domain, completion.slot, completion.generation)
279    }
280}
281
282#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
283pub(crate) struct EventCompletion {
284    domain: EventDomainId,
285    slot: EventSlotId,
286    generation: u64,
287}
288
289impl EventCompletion {
290    pub(crate) fn new(domain: EventDomainId, slot: EventSlotId, generation: u64) -> Self {
291        Self {
292            domain,
293            slot,
294            generation,
295        }
296    }
297
298    pub(crate) fn domain(&self) -> EventDomainId {
299        self.domain
300    }
301}
302
303#[derive(Clone, Debug)]
304pub(crate) struct ScheduledOperation {
305    instruction_index: usize,
306    location: ExecutionLocation,
307    input_values: Box<[usize]>,
308    output_values: Box<[usize]>,
309    dependencies: Box<[EventDependency]>,
310    completion: EventCompletion,
311}
312
313impl ScheduledOperation {
314    pub(crate) fn new(
315        instruction_index: usize,
316        location: ExecutionLocation,
317        input_values: impl Into<Box<[usize]>>,
318        output_values: impl Into<Box<[usize]>>,
319        dependencies: impl Into<Box<[EventDependency]>>,
320        completion: EventCompletion,
321    ) -> Self {
322        Self {
323            instruction_index,
324            location,
325            input_values: input_values.into(),
326            output_values: output_values.into(),
327            dependencies: dependencies.into(),
328            completion,
329        }
330    }
331
332    #[cfg(test)]
333    pub(crate) fn for_test(
334        domain: EventDomainId,
335        provider_device_identity: super::ProviderDeviceIdentity,
336    ) -> Self {
337        Self::new(
338            0,
339            ExecutionLocation::for_test(domain, provider_device_identity),
340            [],
341            [],
342            [],
343            EventCompletion::new(domain, EventSlotId::new(0), 0),
344        )
345    }
346
347    pub(crate) fn instruction_index(&self) -> usize {
348        self.instruction_index
349    }
350
351    pub(crate) fn location(&self) -> &ExecutionLocation {
352        &self.location
353    }
354
355    pub(crate) fn dependencies(&self) -> &[EventDependency] {
356        &self.dependencies
357    }
358
359    pub(crate) fn completion(&self) -> EventCompletion {
360        self.completion
361    }
362
363    fn retained_bytes(&self) -> Option<usize> {
364        checked_sum([
365            self.input_values
366                .len()
367                .checked_mul(std::mem::size_of::<usize>())?,
368            self.output_values
369                .len()
370                .checked_mul(std::mem::size_of::<usize>())?,
371            self.dependencies
372                .len()
373                .checked_mul(std::mem::size_of::<EventDependency>())?,
374        ])
375    }
376}
377
378#[derive(Clone)]
379pub(crate) struct ScheduledTransfer {
380    value_slot: usize,
381    source_location: ExecutionLocation,
382    destination_location: ExecutionLocation,
383    provider: Arc<dyn super::TransferProvider>,
384    dependencies: Box<[EventDependency]>,
385    completion: EventCompletion,
386}
387
388impl ScheduledTransfer {
389    pub(crate) fn with_provider(
390        value_slot: usize,
391        source_location: ExecutionLocation,
392        destination_location: ExecutionLocation,
393        provider: Arc<dyn super::TransferProvider>,
394        dependencies: impl Into<Box<[EventDependency]>>,
395        completion: EventCompletion,
396    ) -> Self {
397        Self {
398            value_slot,
399            source_location,
400            destination_location,
401            provider,
402            dependencies: dependencies.into(),
403            completion,
404        }
405    }
406
407    #[cfg(test)]
408    pub(crate) fn new(
409        value_slot: usize,
410        source_location: ExecutionLocation,
411        destination_location: ExecutionLocation,
412        dependencies: impl Into<Box<[EventDependency]>>,
413        completion: EventCompletion,
414    ) -> Self {
415        Self::with_provider(
416            value_slot,
417            source_location,
418            destination_location,
419            Arc::new(TestScheduledTransferProvider),
420            dependencies,
421            completion,
422        )
423    }
424
425    #[cfg(test)]
426    pub(crate) fn for_test(
427        source_event_domain: EventDomainId,
428        destination_event_domain: EventDomainId,
429        source_provider_device_identity: super::ProviderDeviceIdentity,
430        destination_provider_device_identity: super::ProviderDeviceIdentity,
431    ) -> Self {
432        let source_location =
433            ExecutionLocation::for_test(source_event_domain, source_provider_device_identity);
434        let destination_location = ExecutionLocation::for_test(
435            destination_event_domain,
436            destination_provider_device_identity,
437        );
438        Self::new(
439            0,
440            source_location,
441            destination_location,
442            [EventDependency::new(
443                source_event_domain,
444                EventSlotId::new(0),
445                0,
446            )],
447            EventCompletion::new(destination_event_domain, EventSlotId::new(0), 0),
448        )
449    }
450
451    pub(crate) fn source_event_domain(&self) -> EventDomainId {
452        self.source_location.event_domain_id()
453    }
454
455    #[cfg(test)]
456    pub(crate) fn destination_event_domain(&self) -> EventDomainId {
457        self.destination_location.event_domain_id()
458    }
459
460    pub(crate) fn value_slot(&self) -> usize {
461        self.value_slot
462    }
463
464    pub(crate) fn source_location(&self) -> &ExecutionLocation {
465        &self.source_location
466    }
467
468    pub(crate) fn destination_location(&self) -> &ExecutionLocation {
469        &self.destination_location
470    }
471
472    pub(crate) fn provider(&self) -> &Arc<dyn super::TransferProvider> {
473        &self.provider
474    }
475
476    pub(crate) fn dependencies(&self) -> &[EventDependency] {
477        &self.dependencies
478    }
479
480    pub(crate) fn completion(&self) -> EventCompletion {
481        self.completion
482    }
483
484    fn retained_bytes(&self) -> Option<usize> {
485        self.dependencies
486            .len()
487            .checked_mul(std::mem::size_of::<EventDependency>())
488    }
489}
490
491impl std::fmt::Debug for ScheduledTransfer {
492    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
493        formatter
494            .debug_struct("ScheduledTransfer")
495            .field("value_slot", &self.value_slot)
496            .field("source_location", &self.source_location)
497            .field("destination_location", &self.destination_location)
498            .field("dependencies", &self.dependencies)
499            .field("completion", &self.completion)
500            .finish_non_exhaustive()
501    }
502}
503
504#[derive(Clone, Debug)]
505pub(crate) struct ScheduledCollective {
506    location: ExecutionLocation,
507    dependencies: Box<[EventDependency]>,
508    completion: EventCompletion,
509}
510
511impl ScheduledCollective {
512    #[cfg(test)]
513    pub(crate) fn unsupported_for_test(
514        event_domain_driver: Arc<dyn super::EventDomainDriver>,
515    ) -> Self {
516        let domain = EventDomainId::runtime_created_for_test(
517            RuntimeId::from_nonzero(NonZeroU64::new(1).expect("runtime id")),
518            RuntimeEpoch::from_nonzero(NonZeroU64::new(1).expect("runtime epoch")),
519            RegistrationIdentity::new(
520                NonZeroU64::new(1).expect("registration issuer"),
521                NonZeroU64::new(1).expect("registration ordinal"),
522            ),
523        );
524        let location = ExecutionLocation::for_test_with_driver(
525            domain,
526            super::ProviderDeviceIdentity::new(
527                super::ProviderId::new("tenferro.test.schedule").expect("test provider id"),
528                "collective",
529            )
530            .expect("test provider target"),
531            event_domain_driver,
532        );
533        Self {
534            location,
535            dependencies: Box::new([]),
536            completion: EventCompletion::new(domain, EventSlotId::new(0), 0),
537        }
538    }
539
540    pub(crate) fn location(&self) -> &ExecutionLocation {
541        &self.location
542    }
543
544    pub(crate) fn completion(&self) -> EventCompletion {
545        self.completion
546    }
547
548    fn dependencies(&self) -> &[EventDependency] {
549        &self.dependencies
550    }
551
552    fn retained_bytes(&self) -> Option<usize> {
553        self.dependencies
554            .len()
555            .checked_mul(std::mem::size_of::<EventDependency>())
556    }
557}
558
559#[derive(Clone, Debug)]
560pub(crate) struct ScheduledBarrier {
561    location: ExecutionLocation,
562    dependencies: Box<[EventDependency]>,
563    completion: EventCompletion,
564}
565
566impl ScheduledBarrier {
567    pub(crate) fn location(&self) -> &ExecutionLocation {
568        &self.location
569    }
570
571    fn dependencies(&self) -> &[EventDependency] {
572        &self.dependencies
573    }
574
575    pub(crate) fn completion(&self) -> EventCompletion {
576        self.completion
577    }
578
579    fn retained_bytes(&self) -> Option<usize> {
580        self.dependencies
581            .len()
582            .checked_mul(std::mem::size_of::<EventDependency>())
583    }
584}
585
586#[derive(Clone, Debug)]
587pub(crate) enum ScheduledNode {
588    Operation(ScheduledOperation),
589    Transfer(ScheduledTransfer),
590    // INVARIANT: collective nodes remain representation-only until the
591    // explicitly deferred collective scheduler work lands.
592    #[allow(
593        dead_code,
594        reason = "collective scheduling remains representation-only in this scoped change"
595    )]
596    Collective(ScheduledCollective),
597    // INVARIANT: execution supports explicit barriers, but schedule construction
598    // does not emit them until a later scheduling-policy change requires one.
599    #[allow(
600        dead_code,
601        reason = "barrier construction is deferred until scheduling policy emits explicit barriers"
602    )]
603    Barrier(ScheduledBarrier),
604}
605
606impl ScheduledNode {
607    pub(crate) fn completion(&self) -> EventCompletion {
608        match self {
609            Self::Operation(node) => node.completion(),
610            Self::Transfer(node) => node.completion(),
611            Self::Collective(node) => node.completion(),
612            Self::Barrier(node) => node.completion(),
613        }
614    }
615
616    pub(crate) fn dependencies(&self) -> &[EventDependency] {
617        match self {
618            Self::Operation(node) => node.dependencies(),
619            Self::Transfer(node) => node.dependencies(),
620            Self::Collective(node) => node.dependencies(),
621            Self::Barrier(node) => node.dependencies(),
622        }
623    }
624
625    pub(super) fn event_domain_witness(&self) -> &Arc<ExecutableEngineSnapshot> {
626        match self {
627            Self::Operation(node) => node.location().witness(),
628            Self::Transfer(node) => node.destination_location().witness(),
629            Self::Collective(node) => node.location().witness(),
630            Self::Barrier(node) => node.location().witness(),
631        }
632    }
633
634    fn retained_bytes(&self) -> Option<usize> {
635        match self {
636            Self::Operation(node) => node.retained_bytes(),
637            Self::Transfer(node) => node.retained_bytes(),
638            Self::Collective(node) => node.retained_bytes(),
639            Self::Barrier(node) => node.retained_bytes(),
640        }
641    }
642}
643
644#[derive(Clone, Copy, Debug, Eq, PartialEq)]
645pub(crate) enum ScheduledNodeKind {
646    Collective,
647}
648
649#[derive(Debug, thiserror::Error)]
650#[error("scheduled node {node_index} of kind {node_kind:?} is unsupported during execution")]
651pub(crate) struct UnsupportedScheduledNodeError {
652    pub(crate) node_index: usize,
653    pub(crate) node_kind: ScheduledNodeKind,
654}
655
656#[derive(Clone, Debug)]
657pub(crate) struct ScheduledGraph {
658    nodes: Box<[ScheduledNode]>,
659    input_slots: Box<[usize]>,
660    output_slots: Box<[usize]>,
661    value_count: usize,
662    root_location: ExecutionLocation,
663    input_locations: Box<[ExecutionLocation]>,
664    operation_locations: Box<[ExecutionLocation]>,
665}
666
667impl ScheduledGraph {
668    pub(crate) fn from_exec_program(
669        program: &ExecProgram,
670        root_location: ExecutionLocation,
671        input_locations: &[ExecutionLocation],
672        operation_locations: &[ExecutionLocation],
673        transfer_registry: &FrozenTransferRegistry,
674    ) -> Result<Self, ScheduleBuildError> {
675        let mut nodes = Vec::with_capacity(program.instructions.len());
676        let mut available = vec![Vec::<AvailableValue>::new(); program.n_slots];
677        if input_locations.len() != program.input_slots.len() {
678            return Err(ScheduleBuildError::InputLocationCountMismatch {
679                expected: program.input_slots.len(),
680                actual: input_locations.len(),
681            });
682        }
683        for (&slot, location) in program.input_slots.iter().zip(input_locations) {
684            let values =
685                available
686                    .get_mut(slot)
687                    .ok_or(ScheduleBuildError::ValueSlotOutOfBounds {
688                        slot,
689                        value_count: program.n_slots,
690                    })?;
691            values.push(AvailableValue {
692                location: location.clone(),
693                completion: None,
694            });
695        }
696
697        for (instruction_index, instruction) in program.instructions.iter().enumerate() {
698            let location = match instruction.semantic_operation_index {
699                Some(operation_index) => operation_locations.get(operation_index).cloned().ok_or(
700                    ScheduleBuildError::MissingOperationLocation {
701                        instruction_index,
702                        operation_index,
703                    },
704                )?,
705                None => root_location.clone(),
706            };
707
708            for &slot in &instruction.input_slots {
709                let values =
710                    available
711                        .get(slot)
712                        .ok_or(ScheduleBuildError::ValueSlotOutOfBounds {
713                            slot,
714                            value_count: program.n_slots,
715                        })?;
716                if values.iter().any(|value| value.location == location) {
717                    continue;
718                }
719                let source = values
720                    .iter()
721                    .find_map(|value| {
722                        let route = ResolvedTransferRoute::new(
723                            value.location.resolved_endpoint().clone(),
724                            location.resolved_endpoint().clone(),
725                        );
726                        transfer_registry
727                            .get(&route)
728                            .map(|provider| (value.clone(), Arc::clone(provider)))
729                    })
730                    .ok_or_else(|| {
731                        if values.is_empty() {
732                            ScheduleBuildError::ValueUnavailable {
733                                instruction_index,
734                                slot,
735                            }
736                        } else {
737                            ScheduleBuildError::MissingTransferProvider {
738                                instruction_index,
739                                slot,
740                                destination_endpoint: location.endpoint().clone(),
741                                available_source_endpoints: values
742                                    .iter()
743                                    .map(|value| value.location.endpoint().clone())
744                                    .collect(),
745                            }
746                        }
747                    })?;
748                let completion = event_completion(&nodes, location.event_domain_id())?;
749                let dependencies = source
750                    .0
751                    .completion
752                    .map(EventDependency::from_completion)
753                    .into_iter()
754                    .collect::<Vec<_>>();
755                nodes.push(ScheduledNode::Transfer(ScheduledTransfer::with_provider(
756                    slot,
757                    source.0.location,
758                    location.clone(),
759                    source.1,
760                    dependencies,
761                    completion,
762                )));
763                available[slot].push(AvailableValue {
764                    location: location.clone(),
765                    completion: Some(completion),
766                });
767            }
768
769            let dependencies = operation_dependencies(
770                &available,
771                instruction_index,
772                &instruction.input_slots,
773                &location,
774            )?;
775            let completion = event_completion(&nodes, location.event_domain_id())?;
776            nodes.push(ScheduledNode::Operation(ScheduledOperation::new(
777                instruction_index,
778                location.clone(),
779                instruction.input_slots.clone(),
780                instruction.output_slots.clone(),
781                dependencies,
782                completion,
783            )));
784
785            for (input_index, &slot) in instruction.input_slots.iter().enumerate() {
786                if instruction
787                    .last_use
788                    .get(input_index)
789                    .copied()
790                    .unwrap_or(false)
791                {
792                    available[slot].clear();
793                }
794            }
795            for &slot in &instruction.output_slots {
796                let values =
797                    available
798                        .get_mut(slot)
799                        .ok_or(ScheduleBuildError::ValueSlotOutOfBounds {
800                            slot,
801                            value_count: program.n_slots,
802                        })?;
803                values.clear();
804                values.push(AvailableValue {
805                    location: location.clone(),
806                    completion: Some(completion),
807                });
808            }
809        }
810
811        let graph = Self {
812            nodes: nodes.into_boxed_slice(),
813            input_slots: program.input_slots.clone().into_boxed_slice(),
814            output_slots: program.output_slots.clone().into_boxed_slice(),
815            value_count: program.n_slots,
816            root_location,
817            input_locations: input_locations.to_vec().into_boxed_slice(),
818            operation_locations: operation_locations.to_vec().into_boxed_slice(),
819        };
820        graph
821            .validate()
822            .map_err(|source| ScheduleBuildError::InvalidSchedule { source })?;
823        Ok(graph)
824    }
825
826    #[cfg(test)]
827    pub(crate) fn for_test(nodes: Vec<ScheduledNode>) -> Self {
828        let value_count = nodes.len();
829        let root_location = nodes
830            .iter()
831            .find_map(|node| match node {
832                ScheduledNode::Operation(operation) => Some(operation.location().clone()),
833                ScheduledNode::Transfer(transfer) => Some(transfer.destination_location().clone()),
834                ScheduledNode::Collective(_) | ScheduledNode::Barrier(_) => None,
835            })
836            .unwrap_or_else(|| {
837                ExecutionLocation::new(
838                    EngineId::new("tenferro-test.schedule-fallback-engine")
839                        .expect("test engine id"),
840                    super::ProviderDeviceIdentity::new(
841                        super::ProviderId::new("tenferro.test.schedule").expect("test provider id"),
842                        "fallback",
843                    )
844                    .expect("test provider target"),
845                    EventDomainId::runtime_created_for_test(
846                        RuntimeId::from_nonzero(NonZeroU64::new(1).expect("runtime id")),
847                        RuntimeEpoch::from_nonzero(NonZeroU64::new(1).expect("runtime epoch")),
848                        RegistrationIdentity::new(
849                            NonZeroU64::new(1).expect("registration issuer"),
850                            NonZeroU64::new(1).expect("registration ordinal"),
851                        ),
852                    ),
853                    StorageClass::new("tenferro-test.schedule-fallback-storage")
854                        .expect("test storage class"),
855                )
856            });
857        Self {
858            nodes: nodes.into_boxed_slice(),
859            input_slots: Box::new([]),
860            output_slots: Box::new([]),
861            value_count,
862            root_location,
863            input_locations: Box::new([]),
864            operation_locations: Box::new([]),
865        }
866    }
867
868    pub(crate) fn validate(&self) -> Result<(), ScheduleValidationError> {
869        let mut known_completions = HashSet::with_capacity(self.nodes.len());
870        for (index, node) in self.nodes.iter().enumerate() {
871            match node {
872                ScheduledNode::Operation(operation) => {
873                    if operation.completion().domain() != operation.location().event_domain_id() {
874                        return Err(ScheduleValidationError::CompletionEventDomainMismatch {
875                            index,
876                        });
877                    }
878                }
879                ScheduledNode::Transfer(transfer) => {
880                    if transfer.source_location == transfer.destination_location {
881                        return Err(ScheduleValidationError::TransferSameLocation { index });
882                    }
883                    if transfer.completion().domain()
884                        != transfer.destination_location().event_domain_id()
885                    {
886                        return Err(ScheduleValidationError::CompletionEventDomainMismatch {
887                            index,
888                        });
889                    }
890                    if transfer.dependencies().iter().any(|dependency| {
891                        dependency.domain() != transfer.source_location().event_domain_id()
892                            && dependency.domain()
893                                != transfer.destination_location().event_domain_id()
894                    }) {
895                        return Err(ScheduleValidationError::DependencyEventDomainMismatch {
896                            index,
897                        });
898                    }
899                }
900                ScheduledNode::Collective(collective) => {
901                    if collective.completion().domain() != collective.location().event_domain_id() {
902                        return Err(ScheduleValidationError::CompletionEventDomainMismatch {
903                            index,
904                        });
905                    }
906                }
907                ScheduledNode::Barrier(barrier) => {
908                    if barrier.completion().domain() != barrier.location().event_domain_id() {
909                        return Err(ScheduleValidationError::CompletionEventDomainMismatch {
910                            index,
911                        });
912                    }
913                }
914            }
915            if node
916                .dependencies()
917                .iter()
918                .any(|dependency| !known_completions.contains(dependency))
919            {
920                return Err(ScheduleValidationError::DependencyNotPriorCompletion { index });
921            }
922            if !known_completions.insert(EventDependency::from_completion(node.completion())) {
923                return Err(ScheduleValidationError::DuplicateCompletion { index });
924            }
925        }
926        Ok(())
927    }
928
929    #[cfg(test)]
930    pub(crate) fn contains_collective(&self) -> bool {
931        self.nodes
932            .iter()
933            .any(|node| matches!(node, ScheduledNode::Collective(_)))
934    }
935
936    pub(crate) fn nodes(&self) -> &[ScheduledNode] {
937        &self.nodes
938    }
939
940    pub(crate) fn root_location(&self) -> &ExecutionLocation {
941        &self.root_location
942    }
943
944    pub(crate) fn input_locations(&self) -> &[ExecutionLocation] {
945        &self.input_locations
946    }
947
948    pub(crate) fn operation_locations(&self) -> &[ExecutionLocation] {
949        &self.operation_locations
950    }
951
952    /// Reject non-executable node kinds before any provider event-domain run
953    /// is acquired. Preparation currently emits operations, transfers, and
954    /// no-op barriers; collective execution remains unsupported.
955    pub(crate) fn preflight(&self) -> Result<(), UnsupportedScheduledNodeError> {
956        if let Some(index) = self
957            .nodes
958            .iter()
959            .position(|node| matches!(node, ScheduledNode::Collective(_)))
960        {
961            return Err(UnsupportedScheduledNodeError {
962                node_index: index,
963                node_kind: ScheduledNodeKind::Collective,
964            });
965        }
966        Ok(())
967    }
968
969    pub(crate) fn retained_bytes(&self) -> Option<usize> {
970        let node_payload_bytes = self
971            .nodes
972            .iter()
973            .try_fold(0usize, |sum, node| sum.checked_add(node.retained_bytes()?))?;
974        checked_sum([
975            std::mem::size_of::<ScheduledGraph>(),
976            self.nodes
977                .len()
978                .checked_mul(std::mem::size_of::<ScheduledNode>())?,
979            node_payload_bytes,
980            self.input_slots
981                .len()
982                .checked_mul(std::mem::size_of::<usize>())?,
983            self.output_slots
984                .len()
985                .checked_mul(std::mem::size_of::<usize>())?,
986            self.value_count.checked_mul(std::mem::size_of::<usize>())?,
987            self.input_locations
988                .len()
989                .checked_mul(std::mem::size_of::<ExecutionLocation>())?,
990            self.operation_locations
991                .len()
992                .checked_mul(std::mem::size_of::<ExecutionLocation>())?,
993        ])
994    }
995
996    #[cfg(test)]
997    pub(crate) fn transfers_for_test(&self) -> impl Iterator<Item = &ScheduledTransfer> {
998        self.nodes.iter().filter_map(|node| match node {
999            ScheduledNode::Transfer(transfer) => Some(transfer),
1000            _ => None,
1001        })
1002    }
1003
1004    #[cfg(test)]
1005    pub(crate) fn nodes_for_test(&self) -> &[ScheduledNode] {
1006        &self.nodes
1007    }
1008}
1009
1010#[derive(Debug, thiserror::Error)]
1011pub(crate) enum ScheduleValidationError {
1012    #[error("transfer node {index} uses the same source and destination location")]
1013    TransferSameLocation { index: usize },
1014    #[error("schedule node {index} completion uses the wrong event domain")]
1015    CompletionEventDomainMismatch { index: usize },
1016    #[error("schedule node {index} dependency uses the wrong event domain")]
1017    DependencyEventDomainMismatch { index: usize },
1018    #[error("schedule node {index} dependency does not refer to a prior completion")]
1019    DependencyNotPriorCompletion { index: usize },
1020    #[error("schedule node {index} reuses a prior completion identity")]
1021    DuplicateCompletion { index: usize },
1022}
1023
1024#[derive(Debug, thiserror::Error)]
1025pub(crate) enum ScheduleBuildError {
1026    #[error("schedule construction produced an invalid immutable schedule")]
1027    InvalidSchedule {
1028        #[source]
1029        source: ScheduleValidationError,
1030    },
1031    #[error("schedule has {actual} input locations for {expected} program inputs")]
1032    InputLocationCountMismatch { expected: usize, actual: usize },
1033    #[error(
1034        "instruction {instruction_index} references semantic operation {operation_index}, \
1035         but that operation has no execution location"
1036    )]
1037    MissingOperationLocation {
1038        instruction_index: usize,
1039        operation_index: usize,
1040    },
1041    #[error("instruction {instruction_index} requires unavailable value slot {slot}")]
1042    ValueUnavailable {
1043        instruction_index: usize,
1044        slot: usize,
1045    },
1046    #[error(
1047        "instruction {instruction_index} has no direct transfer provider for value slot {slot} \
1048         from {available_source_endpoints:?} to {destination_endpoint:?}"
1049    )]
1050    MissingTransferProvider {
1051        instruction_index: usize,
1052        slot: usize,
1053        destination_endpoint: TransferEndpoint,
1054        available_source_endpoints: Vec<TransferEndpoint>,
1055    },
1056    #[error("value slot {slot} is outside schedule value count {value_count}")]
1057    ValueSlotOutOfBounds { slot: usize, value_count: usize },
1058    #[error("scheduled node count exceeds the event-slot identity space")]
1059    EventSlotExhausted,
1060}
1061
1062#[cfg(test)]
1063#[derive(Debug)]
1064struct TestScheduledTransferProvider;
1065
1066#[cfg(test)]
1067impl super::TransferProvider for TestScheduledTransferProvider {
1068    fn transfer_blocking(
1069        &self,
1070        _request: super::TransferRequest<'_>,
1071    ) -> crate::Result<tenferro_tensor::Tensor> {
1072        Err(crate::Error::Internal(
1073            "scheduled test transfer provider".into(),
1074        ))
1075    }
1076}
1077
1078fn checked_sum(values: impl IntoIterator<Item = usize>) -> Option<usize> {
1079    values
1080        .into_iter()
1081        .try_fold(0usize, |sum, value| sum.checked_add(value))
1082}
1083
1084#[derive(Clone)]
1085struct AvailableValue {
1086    location: ExecutionLocation,
1087    completion: Option<EventCompletion>,
1088}
1089
1090fn operation_dependencies(
1091    available: &[Vec<AvailableValue>],
1092    instruction_index: usize,
1093    input_slots: &[usize],
1094    location: &ExecutionLocation,
1095) -> Result<Vec<EventDependency>, ScheduleBuildError> {
1096    let mut dependencies = Vec::with_capacity(input_slots.len());
1097    let mut seen = HashSet::with_capacity(input_slots.len());
1098    for &slot in input_slots {
1099        let values = available
1100            .get(slot)
1101            .ok_or(ScheduleBuildError::ValueSlotOutOfBounds {
1102                slot,
1103                value_count: available.len(),
1104            })?;
1105        let value = values
1106            .iter()
1107            .find(|value| &value.location == location)
1108            .ok_or(ScheduleBuildError::ValueUnavailable {
1109                instruction_index,
1110                slot,
1111            })?;
1112        if let Some(completion) = value.completion {
1113            let dependency = EventDependency::from_completion(completion);
1114            if seen.insert(dependency) {
1115                dependencies.push(dependency);
1116            }
1117        }
1118    }
1119    Ok(dependencies)
1120}
1121
1122fn event_completion(
1123    nodes: &[ScheduledNode],
1124    domain: EventDomainId,
1125) -> Result<EventCompletion, ScheduleBuildError> {
1126    let slot = u32::try_from(nodes.len()).map_err(|_| ScheduleBuildError::EventSlotExhausted)?;
1127    Ok(EventCompletion::new(domain, EventSlotId::new(slot), 0))
1128}
tenferro_runtime/runtime/schedule.rs

tenferro_runtime/runtime/
schedule.rs
