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

tenferro_runtime/runtime/
schedule.rs
