Skip to main content

eredu_runtime/
backend.rs

1//! Narrow capability contracts implemented by execution backends.
2
3use eredu_checkpoint::{
4    recipe::DerivedWeightRecipe,
5    store::{CheckpointLease, CheckpointSource},
6};
7use eredu_core::Completion;
8use eredu_nn::NeuralBackend;
9
10/// Submits backend-native work and retains values through exact completion.
11pub trait SubmissionBackend: NeuralBackend {
12    /// Backend executor, queue, stream, or equivalent submission context.
13    type Executor: ?Sized;
14    /// Owned executor used for an independently schedulable graph lane.
15    type OwnedExecutor: std::borrow::Borrow<Self::Executor>;
16    /// Exact completion object for one submission.
17    type Completion: Completion;
18
19    /// Creates independently schedulable executors on the same backend device.
20    fn fork_executors(
21        executor: &Self::Executor,
22        count: usize,
23    ) -> Result<Vec<Self::OwnedExecutor>, <Self::Completion as Completion>::Error>;
24
25    /// Submits evaluation of backend-native values on one executor.
26    fn submit<'a, I>(
27        executor: &Self::Executor,
28        values: I,
29    ) -> Result<Self::Completion, <Self::Completion as Completion>::Error>
30    where
31        Self::Tensor: 'a,
32        I: IntoIterator<Item = &'a Self::Tensor>;
33
34    /// Orders future work on `executor` after an exact producer completion.
35    fn order_after(
36        completion: &Self::Completion,
37        executor: &Self::Executor,
38    ) -> Result<(), <Self::Completion as Completion>::Error>;
39
40    /// Retains an owned value until `completion` has completed exactly.
41    fn retain_until_complete<T: Send + 'static>(
42        executor: &Self::Executor,
43        completion: &Self::Completion,
44        value: T,
45    ) -> Result<(), <Self::Completion as Completion>::Error>;
46}
47
48/// Materializes and binds checkpoint data to backend-native parameter slots.
49pub trait ParameterBackend: NeuralBackend {
50    /// One backend-native parameter slot.
51    type Parameter: 'static;
52    /// Materialized backend-native checkpoint weight.
53    type MaterializedWeight;
54    /// Backend context used only while realizing checkpoint parameters.
55    type MaterializationContext: ?Sized;
56    /// In-flight guard retaining encoded sources through exact realization completion.
57    type Materialization;
58    /// Backend-specific loading failure.
59    type ParameterError: std::error::Error + Send + Sync + 'static;
60
61    /// Lowers one format-preserving encoded lease into a native weight.
62    fn materialize(
63        lease: CheckpointLease,
64        context: &Self::MaterializationContext,
65    ) -> Result<Self::Materialization, Self::ParameterError>;
66
67    /// Lowers a validated neutral recipe directly into a native weight.
68    fn materialize_recipe(
69        recipe: &DerivedWeightRecipe,
70        source: &dyn CheckpointSource,
71        context: &Self::MaterializationContext,
72    ) -> Result<Self::Materialization, Self::ParameterError>;
73
74    /// Borrows the native weight retained by an in-flight materialization.
75    fn materialized_weight(materialization: &Self::Materialization) -> &Self::MaterializedWeight;
76
77    /// Waits for this exact realization and releases its encoded source lease.
78    fn finish_materialization(
79        materialization: Self::Materialization,
80    ) -> Result<Self::MaterializedWeight, Self::ParameterError>;
81
82    /// Creates another native handle to identical materialized storage without
83    /// rereading or rematerializing checkpoint data.
84    fn share_materialized_weight(
85        weight: &Self::MaterializedWeight,
86    ) -> Result<Self::MaterializedWeight, Self::ParameterError>;
87
88    /// Validates destination shape/storage compatibility without publication.
89    fn validate_bind(
90        parameter: &Self::Parameter,
91        weight: &Self::MaterializedWeight,
92    ) -> Result<(), Self::ParameterError>;
93
94    /// Binds one materialized weight to its destination parameter.
95    ///
96    /// After successful [`Self::validate_bind`] on unchanged arguments this
97    /// operation must not fail, allowing orchestration to validate an entire
98    /// atomic unit before publishing any destination.
99    fn bind(
100        parameter: &mut Self::Parameter,
101        weight: Self::MaterializedWeight,
102    ) -> Result<(), Self::ParameterError>;
103}
104
105/// Promotes and demotes backend-native storage without changing its semantics.
106pub trait TransferBackend: SubmissionBackend + ParameterBackend {
107    /// Backend-owned host representation.
108    type HostBuffer;
109    /// In-flight transfer guard retaining all source and destination storage.
110    type Transfer: Completion<Error = Self::TransferError>;
111    /// Backend-specific transfer failure.
112    type TransferError: std::error::Error + Send + Sync + 'static;
113
114    /// Promotes host storage into a materialized execution weight.
115    fn promote(
116        executor: &Self::Executor,
117        host: &Self::HostBuffer,
118    ) -> Result<(Self::MaterializedWeight, Self::Transfer), Self::TransferError>;
119
120    /// Demotes a materialized execution weight into backend-owned host storage.
121    fn demote(
122        executor: &Self::Executor,
123        weight: &Self::MaterializedWeight,
124    ) -> Result<(Self::HostBuffer, Self::Transfer), Self::TransferError>;
125}
126
127/// Collective operations available to distributed runtime policies.
128pub trait CollectiveBackend: SubmissionBackend {
129    /// Backend-native collective group.
130    type Group: ?Sized;
131    /// Backend-specific collective failure.
132    type CollectiveError: std::error::Error + Send + Sync + 'static;
133
134    /// Reduces a tensor across the selected group.
135    fn all_reduce(
136        value: Self::Tensor,
137        group: &Self::Group,
138        executor: &Self::Executor,
139    ) -> Result<Self::Tensor, Self::CollectiveError>;
140
141    /// Gathers a tensor across the selected group.
142    fn all_gather(
143        value: Self::Tensor,
144        group: &Self::Group,
145        executor: &Self::Executor,
146    ) -> Result<Self::Tensor, Self::CollectiveError>;
147
148    /// Exchanges tensor partitions across the selected group.
149    fn all_to_all(
150        value: Self::Tensor,
151        group: &Self::Group,
152        executor: &Self::Executor,
153    ) -> Result<Self::Tensor, Self::CollectiveError>;
154}