1use eredu_checkpoint::{
4 recipe::DerivedWeightRecipe,
5 store::{CheckpointLease, CheckpointSource},
6};
7use eredu_core::Completion;
8use eredu_nn::NeuralBackend;
9
10pub trait SubmissionBackend: NeuralBackend {
12 type Executor: ?Sized;
14 type OwnedExecutor: std::borrow::Borrow<Self::Executor>;
16 type Completion: Completion;
18
19 fn fork_executors(
21 executor: &Self::Executor,
22 count: usize,
23 ) -> Result<Vec<Self::OwnedExecutor>, <Self::Completion as Completion>::Error>;
24
25 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 fn order_after(
36 completion: &Self::Completion,
37 executor: &Self::Executor,
38 ) -> Result<(), <Self::Completion as Completion>::Error>;
39
40 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
48pub trait ParameterBackend: NeuralBackend {
50 type Parameter: 'static;
52 type MaterializedWeight;
54 type MaterializationContext: ?Sized;
56 type Materialization;
58 type ParameterError: std::error::Error + Send + Sync + 'static;
60
61 fn materialize(
63 lease: CheckpointLease,
64 context: &Self::MaterializationContext,
65 ) -> Result<Self::Materialization, Self::ParameterError>;
66
67 fn materialize_recipe(
69 recipe: &DerivedWeightRecipe,
70 source: &dyn CheckpointSource,
71 context: &Self::MaterializationContext,
72 ) -> Result<Self::Materialization, Self::ParameterError>;
73
74 fn materialized_weight(materialization: &Self::Materialization) -> &Self::MaterializedWeight;
76
77 fn finish_materialization(
79 materialization: Self::Materialization,
80 ) -> Result<Self::MaterializedWeight, Self::ParameterError>;
81
82 fn share_materialized_weight(
85 weight: &Self::MaterializedWeight,
86 ) -> Result<Self::MaterializedWeight, Self::ParameterError>;
87
88 fn validate_bind(
90 parameter: &Self::Parameter,
91 weight: &Self::MaterializedWeight,
92 ) -> Result<(), Self::ParameterError>;
93
94 fn bind(
100 parameter: &mut Self::Parameter,
101 weight: Self::MaterializedWeight,
102 ) -> Result<(), Self::ParameterError>;
103}
104
105pub trait TransferBackend: SubmissionBackend + ParameterBackend {
107 type HostBuffer;
109 type Transfer: Completion<Error = Self::TransferError>;
111 type TransferError: std::error::Error + Send + Sync + 'static;
113
114 fn promote(
116 executor: &Self::Executor,
117 host: &Self::HostBuffer,
118 ) -> Result<(Self::MaterializedWeight, Self::Transfer), Self::TransferError>;
119
120 fn demote(
122 executor: &Self::Executor,
123 weight: &Self::MaterializedWeight,
124 ) -> Result<(Self::HostBuffer, Self::Transfer), Self::TransferError>;
125}
126
127pub trait CollectiveBackend: SubmissionBackend {
129 type Group: ?Sized;
131 type CollectiveError: std::error::Error + Send + Sync + 'static;
133
134 fn all_reduce(
136 value: Self::Tensor,
137 group: &Self::Group,
138 executor: &Self::Executor,
139 ) -> Result<Self::Tensor, Self::CollectiveError>;
140
141 fn all_gather(
143 value: Self::Tensor,
144 group: &Self::Group,
145 executor: &Self::Executor,
146 ) -> Result<Self::Tensor, Self::CollectiveError>;
147
148 fn all_to_all(
150 value: Self::Tensor,
151 group: &Self::Group,
152 executor: &Self::Executor,
153 ) -> Result<Self::Tensor, Self::CollectiveError>;
154}