eredu_runtime/
mechanism_synthesis.rs

1//! Requirement-driven synthesis of exact backend mechanism capabilities.
2
3use eredu_checkpoint::{LinearFormat, SourceTensorEncoding, StoredDtype};
4use eredu_core::{
5    cache::{StateComponentPolicy, StateResidencyClass},
6    checkpoint::TensorDtype,
7    SessionCapabilities,
8};
9use eredu_nn::NeuralOperatorCapabilities;
10
11use crate::{
12    AddressableStorageCapabilities, BackendMechanismCapabilities, CacheResidencyPolicy,
13    GroupedOperationRequirement, ReplicatedTextParameterRole, ReplicatedTextRequirements,
14    ReplicatedTextSelectionRequest, StateComponentMechanism, StateComponentPlacement, StateLayout,
15    StateMechanismCapabilities, StateStorageDtype, WeightLoweringCapability,
16    WeightLoweringDescriptor, WeightLoweringKind, WeightResidencyMechanism,
17};
18
19/// Collection-independent facilities of a backend's mutable-state implementation.
20#[derive(Debug, Clone, Copy, Default, Eq, PartialEq)]
21pub struct StateLifecycleCapabilities {
22    checkpoint: bool,
23    rollback: bool,
24    reset: bool,
25    prompt_cache: bool,
26    observation_retention: bool,
27}
28
29impl StateLifecycleCapabilities {
30    /// Creates a fail-closed declaration of state lifecycle facilities.
31    pub const fn new() -> Self {
32        Self {
33            checkpoint: false,
34            rollback: false,
35            reset: false,
36            prompt_cache: false,
37            observation_retention: false,
38        }
39    }
40
41    /// Declares exact checkpoint and rollback support.
42    pub const fn with_transactions(mut self, checkpoint: bool, rollback: bool) -> Self {
43        self.checkpoint = checkpoint;
44        self.rollback = rollback;
45        self
46    }
47
48    /// Declares complete reset support.
49    pub const fn with_reset(mut self, supported: bool) -> Self {
50        self.reset = supported;
51        self
52    }
53
54    /// Declares persistence and restoration support for the requested state policy.
55    pub const fn with_prompt_cache(mut self, supported: bool) -> Self {
56        self.prompt_cache = supported;
57        self
58    }
59
60    /// Declares retention of state components through observed submissions.
61    pub const fn with_observation_retention(mut self, supported: bool) -> Self {
62        self.observation_retention = supported;
63        self
64    }
65}
66
67/// Exact backend facts independent of architecture parameter and state collections.
68#[derive(Debug, Clone, Eq, PartialEq)]
69pub struct BackendMechanismFacts {
70    operators: NeuralOperatorCapabilities,
71    weight_residencies: Vec<WeightResidencyMechanism>,
72    state: StateLifecycleCapabilities,
73    session: SessionCapabilities,
74    prompt_cache: bool,
75    exact_completion: bool,
76    grouped_operations: Vec<GroupedOperationRequirement>,
77    indexed_movement: bool,
78    addressable_storage: Option<AddressableStorageCapabilities>,
79}
80
81impl BackendMechanismFacts {
82    /// Declares neural, ordinary residency, and state lifecycle mechanisms.
83    /// Optional facilities remain absent until explicitly declared.
84    pub fn new(
85        operators: NeuralOperatorCapabilities,
86        weight_residencies: impl IntoIterator<Item = WeightResidencyMechanism>,
87        state: StateLifecycleCapabilities,
88    ) -> Self {
89        Self {
90            operators,
91            weight_residencies: weight_residencies.into_iter().collect(),
92            state,
93            session: SessionCapabilities::default(),
94            prompt_cache: false,
95            exact_completion: false,
96            grouped_operations: Vec::new(),
97            indexed_movement: false,
98            addressable_storage: None,
99        }
100    }
101
102    /// Declares the exact capabilities of the resulting session implementation.
103    pub const fn with_session(mut self, session: SessionCapabilities) -> Self {
104        self.session = session;
105        self
106    }
107
108    /// Declares prompt-cache persistence mechanisms.
109    pub const fn with_prompt_cache(mut self, supported: bool) -> Self {
110        self.prompt_cache = supported;
111        self
112    }
113
114    /// Declares exact completion ownership.
115    pub const fn with_exact_completion(mut self, supported: bool) -> Self {
116        self.exact_completion = supported;
117        self
118    }
119
120    /// Declares grouped operation mechanisms.
121    pub fn with_grouped_operations(
122        mut self,
123        operations: impl IntoIterator<Item = GroupedOperationRequirement>,
124    ) -> Self {
125        self.grouped_operations = operations.into_iter().collect();
126        self
127    }
128
129    /// Declares indexed discovery, slicing, remapping, and concatenation.
130    pub const fn with_indexed_movement(mut self, supported: bool) -> Self {
131        self.indexed_movement = supported;
132        self
133    }
134
135    /// Declares independently addressable storage and exact capacity limits.
136    pub const fn with_addressable_storage(
137        mut self,
138        capabilities: AddressableStorageCapabilities,
139    ) -> Self {
140        self.addressable_storage = Some(capabilities);
141        self
142    }
143}
144
145/// Side-effect-free native support queries for neutral requirements.
146///
147/// Implementations report facts about one descriptor or state component. They
148/// must not open payloads, allocate tensors, create native resources, or select
149/// an architecture. Requirement enumeration and semantic validation belong to
150/// [`synthesize_replicated_text_capabilities`].
151pub trait ReplicatedTextMechanismSupport {
152    /// Reports collection-independent facts for the requested state implementation.
153    fn facts(&self, state_policy: &CacheResidencyPolicy) -> BackendMechanismFacts;
154
155    /// Reports native realization of one semantically valid direct descriptor.
156    fn supports_direct(&self, descriptor: &WeightLoweringDescriptor) -> bool;
157
158    /// Reports native realization of one semantically valid transform descriptor.
159    fn supports_transform(&self, descriptor: &WeightLoweringDescriptor) -> bool;
160
161    /// Resolves native floating-state storage from one architecture-selected source dtype.
162    /// Packed embeddings may produce a different native scalar representation.
163    /// This metadata-only query must not inspect names or acquire payloads.
164    fn floating_state_dtype(&self, source: &TensorDtype) -> Option<StateStorageDtype>;
165
166    /// Reports support for the exact component geometry, native storage dtype, and placement.
167    fn supports_state_component(
168        &self,
169        component: &StateComponentPolicy,
170        storage_dtype: StateStorageDtype,
171        placement: StateComponentPlacement,
172    ) -> bool;
173}
174
175/// Enumerates exact lowering and state mechanisms in deterministic requirement order.
176///
177/// Invalid architecture transforms are left for the selector's structured
178/// diagnostic; they never become backend support queries. Absent and tied
179/// parameters introduce no descriptor. Equal descriptors retain their first
180/// supported occurrence, including across primary and auxiliary requirements.
181pub fn synthesize_replicated_text_capabilities(
182    requirements: &ReplicatedTextRequirements,
183    request: &ReplicatedTextSelectionRequest,
184    support: &impl ReplicatedTextMechanismSupport,
185) -> BackendMechanismCapabilities {
186    let facts = support.facts(request.state());
187    let mut weight_lowerings = Vec::new();
188    for parameter in requirements
189        .parameters()
190        .iter()
191        .chain(requirements.auxiliary_parameters())
192    {
193        if !parameter.has_lowering_source() {
194            continue;
195        }
196        let requested = request
197            .quantization()
198            .and_then(|requested| parameter.transform_target(requested).ok().flatten())
199            .map(|target| target.executable());
200        for executable in std::iter::once(parameter.native_executable()).chain(requested) {
201            let Ok(descriptor) = parameter.lowering_descriptor(executable) else {
202                continue;
203            };
204            if parameter.role() == ReplicatedTextParameterRole::LinearWeight
205                && matches!(
206                    descriptor.source(),
207                    SourceTensorEncoding::Safetensors(StoredDtype::U8)
208                )
209                && executable == LinearFormat::Dense
210            {
211                continue;
212            }
213            let kind = if executable == parameter.native_executable()
214                && descriptor.has_valid_direct_geometry()
215                && support.supports_direct(&descriptor)
216            {
217                Some(WeightLoweringKind::Direct)
218            } else if descriptor.has_valid_transform_geometry()
219                && support.supports_transform(&descriptor)
220            {
221                Some(WeightLoweringKind::Transform)
222            } else {
223                None
224            };
225            if let Some(kind) = kind {
226                let capability = WeightLoweringCapability::new(descriptor, kind);
227                if !weight_lowerings.contains(&capability) {
228                    weight_lowerings.push(capability);
229                }
230            }
231        }
232    }
233
234    let floating_dtype = requirements
235        .floating_state_source()
236        .and_then(|source| support.floating_state_dtype(source))
237        .filter(|dtype| dtype.is_floating());
238    let mut state =
239        synthesize_state_components(requirements.state_layout(), facts.state, |component| {
240            let Some(dtype) = StateStorageDtype::resolve(component.dtype(), floating_dtype) else {
241                return (None, None);
242            };
243            let device = support
244                .supports_state_component(component, dtype, StateComponentPlacement::Device)
245                .then_some(StateComponentPlacement::Device);
246            let paged = match component.residency() {
247                StateResidencyClass::SealablePaged => support
248                    .supports_state_component(component, dtype, StateComponentPlacement::Paged)
249                    .then_some(StateComponentPlacement::Paged),
250                StateResidencyClass::AlwaysDeviceMutable
251                | StateResidencyClass::LayerScopedOffloadable => device,
252            };
253            (device, paged)
254        });
255
256    if let (Some(source), Some(dtype)) = (requirements.floating_state_source(), floating_dtype) {
257        state = state.with_floating_state_dtype(source.clone(), dtype);
258    }
259    let capabilities = BackendMechanismCapabilities::new(
260        facts.operators,
261        weight_lowerings,
262        facts.weight_residencies,
263        state,
264    )
265    .with_session(facts.session)
266    .with_prompt_cache(facts.prompt_cache)
267    .with_exact_completion(facts.exact_completion)
268    .with_grouped_operations(facts.grouped_operations)
269    .with_indexed_movement(facts.indexed_movement);
270    match facts.addressable_storage {
271        Some(storage) => capabilities.with_addressable_storage(storage),
272        None => capabilities,
273    }
274}
275
276/// Shared collection traversal; native providers answer only one component at a time.
277pub(crate) fn synthesize_state_components(
278    layout: &StateLayout,
279    lifecycle: StateLifecycleCapabilities,
280    placements: impl Fn(
281        &StateComponentPolicy,
282    ) -> (
283        Option<StateComponentPlacement>,
284        Option<StateComponentPlacement>,
285    ),
286) -> StateMechanismCapabilities {
287    StateMechanismCapabilities::new((0..layout.len()).flat_map(|layer| {
288        let placements = &placements;
289        layout
290            .components(layer)
291            .expect("validated state layout exposes every layer")
292            .iter()
293            .filter_map(move |component| {
294                let (device, paged) = placements(component);
295                (device.is_some() || paged.is_some())
296                    .then(|| StateComponentMechanism::new(layer, component.clone(), device, paged))
297            })
298    }))
299    .with_transactions(lifecycle.checkpoint, lifecycle.rollback)
300    .with_reset(lifecycle.reset)
301    .with_prompt_cache(lifecycle.prompt_cache)
302    .with_observation_retention(lifecycle.observation_retention)
303}
304
305impl WeightLoweringDescriptor {
306    /// Validates source encoding geometry without choosing a native implementation.
307    pub fn has_valid_direct_geometry(&self) -> bool {
308        if self.executable().validate().is_err() {
309            return false;
310        }
311        if self.executable() != LinearFormat::Dense
312            && (self.packed_axis().is_none()
313                || self.packed_axis() != self.logical_shape().len().checked_sub(1))
314        {
315            return false;
316        }
317        let same_unpacked_dimensions = |packed_axis: usize| {
318            self.physical_shape()
319                .iter()
320                .zip(self.logical_shape())
321                .enumerate()
322                .all(|(axis, (physical, logical))| axis == packed_axis || physical == logical)
323        };
324        match self.source() {
325            SourceTensorEncoding::Gguf { ggml_type, .. } => ggml_type
326                .block_and_bytes()
327                .ok()
328                .and_then(|(block, _)| usize::try_from(block).ok())
329                .is_some_and(|block| match self.packed_axis() {
330                    Some(axis) if same_unpacked_dimensions(axis) => {
331                        let physical = self.physical_shape()[axis];
332                        let logical = self.logical_shape()[axis];
333                        physical >= logical
334                            && physical.is_multiple_of(block)
335                            && physical - logical < block
336                    }
337                    Some(_) => false,
338                    None => self.physical_shape() == self.logical_shape(),
339                }),
340            SourceTensorEncoding::Safetensors(StoredDtype::U32)
341            | SourceTensorEncoding::RecipeOutput(StoredDtype::U32) => {
342                let Some(axis) = self.packed_axis() else {
343                    return false;
344                };
345                if !same_unpacked_dimensions(axis) {
346                    return false;
347                }
348                let bits = match self.executable() {
349                    LinearFormat::Affine(format) => usize::try_from(format.bits).ok(),
350                    LinearFormat::MxFp4 => Some(4),
351                    _ => None,
352                };
353                bits.is_some_and(|bits| {
354                    self.physical_shape()[axis]
355                        .checked_mul(32)
356                        .zip(self.logical_shape()[axis].checked_mul(bits))
357                        .is_some_and(|(physical, logical)| physical == logical)
358                }) && self.has_valid_packed_geometry()
359            }
360            SourceTensorEncoding::Safetensors(StoredDtype::U8)
361            | SourceTensorEncoding::RecipeOutput(StoredDtype::U8)
362                if self.executable() == LinearFormat::MxFp4 =>
363            {
364                self.physical_shape() == self.logical_shape()
365            }
366            _ => {
367                self.physical_shape() == self.logical_shape()
368                    && (self.executable() == LinearFormat::Dense
369                        || self.has_valid_packed_geometry())
370            }
371        }
372    }
373
374    /// Validates an unpacked source and the selected transform's packing geometry.
375    pub fn has_valid_transform_geometry(&self) -> bool {
376        self.physical_shape() == self.logical_shape()
377            && self.packed_axis() == self.logical_shape().len().checked_sub(1)
378            && self.has_valid_packed_geometry()
379    }
380
381    fn has_valid_packed_geometry(&self) -> bool {
382        if self.executable() != LinearFormat::Dense
383            && (self.packed_axis().is_none()
384                || self.packed_axis() != self.logical_shape().len().checked_sub(1))
385        {
386            return false;
387        }
388        let Some(extent) = self.packed_extent() else {
389            return false;
390        };
391        match self.executable() {
392            LinearFormat::Affine(format) => {
393                format.validate().is_ok()
394                    && usize::try_from(format.group_size)
395                        .ok()
396                        .is_some_and(|group| extent.is_multiple_of(group))
397            }
398            LinearFormat::MxFp4 => extent.is_multiple_of(32),
399            LinearFormat::GgufIQuant { ggml_type, .. } => ggml_type
400                .block_and_bytes()
401                .ok()
402                .and_then(|(block, _)| usize::try_from(block).ok())
403                .is_some_and(|block| extent.is_multiple_of(block)),
404            LinearFormat::E4M3BlockFp8(format) => format.validate().is_ok(),
405            LinearFormat::Dense => true,
406        }
407    }
408}
409
410#[cfg(test)]
411mod tests;
eredu_runtime/mechanism_synthesis.rs

eredu_runtime/
mechanism_synthesis.rs
