henad_compute/entry/
mod.rs

1//! Model entries, the type-erased form in which a host offers a model, and the sets that hold them.
2//!
3//! A [`ModelEntry`] holds a model's declarations and the factory that builds it on any device. The `register_*`
4//! functions make one from an authoring trait. They are generic, so the engine that runs a model is monomorphised in
5//! the crate that calls them. A [`ModelSet`] holds the entries a host offers, each id once.
6#![cfg_attr(
7    all(target_arch = "wasm32", target_feature = "atomics"),
8    expect(
9        clippy::arc_with_non_send_sync,
10        reason = "a factory is not `Send` on wasm with atomics, see `send_sync`"
11    )
12)]
13
14mod set;
15
16pub use set::{ModelLookupError, ModelSet, ModelSetError, ModelSetIter};
17
18use std::any::type_name;
19use std::fmt;
20use std::sync::Arc;
21
22use henad_core::action::ActionDescriptor;
23use henad_core::authoring::model::agent_model::{AgentLanes, AgentModel, NeighborIndex};
24use henad_core::authoring::model::field::FieldLayer;
25use henad_core::authoring::model::gpu_agent_model::GpuAgentModel;
26use henad_core::authoring::model::gpu_grid_model::GpuGridModel;
27use henad_core::authoring::model::grid_model::GridModel;
28use henad_core::authoring::model::network_model::NetworkModel;
29use henad_core::explore::plan::ModelSchema;
30use henad_core::metadata::{Backend, ModelMetadata, Structure};
31use henad_core::model::SimState;
32use henad_core::params::{ParamDescriptor, ParamValue};
33use henad_core::provenance::ModelSource;
34use henad_core::send_sync::{WasmNotSend, WasmNotSync};
35use henad_core::topology::TopologyHint;
36use henad_core::view::StatDescriptor;
37
38use crate::cpu::agent_engine::{AgentModelState, agent_model_param_descriptors};
39use crate::cpu::grid_engine::{GridModelState, grid_model_param_descriptors};
40use crate::cpu::network_engine::{NetworkModelState, network_model_param_descriptors};
41use crate::fault::{BUILDING, Fault, catching};
42use crate::gpu::agent_engine::GpuAgentState;
43use crate::gpu::capacity::Demand;
44use crate::gpu::fault::catching_on;
45use crate::gpu::grid_engine::GpuGridState;
46use crate::gpu::sim_thread::GpuSimState;
47use crate::gpu::{GpuContext, GpuNeeds};
48use crate::simulation::RunSetup;
49
50/// A freshly built simulation state, tagged with the runner that can drive it.
51///
52/// The two arms are not interchangeable. A CPU state steps one tick per call on a
53/// [`SimThread`](crate::cpu::sim_thread::SimThread), and a GPU state encodes many steps into one submission on a
54/// [`GpuSimThread`](crate::gpu::GpuSimThread). A caller picks the runner from the arm without downcasting.
55pub enum ModelState {
56    /// A CPU state, for a [`SimThread`](crate::cpu::sim_thread::SimThread).
57    Cpu(Box<dyn SimState>),
58    /// A GPU state, for a [`GpuSimThread`](crate::gpu::GpuSimThread).
59    Gpu(Box<dyn GpuSimState>),
60}
61
62/// Prints the backend alone, since a state holds the whole simulation.
63impl fmt::Debug for ModelState {
64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
65        match self {
66            Self::Cpu(_) => f.debug_tuple("Cpu").finish_non_exhaustive(),
67            Self::Gpu(_) => f.debug_tuple("Gpu").finish_non_exhaustive(),
68        }
69    }
70}
71
72/// Closure behind a model's factory.
73///
74/// A CPU model ignores the device. Public for [`ModelEntry::wrap_factory`] alone.
75/// The `Option<u64>` is the seed. `None` falls back to the model's default seed.
76#[doc(hidden)]
77pub trait Factory:
78    Fn(&[ParamValue], Option<u64>, Option<&GpuContext>) -> Result<ModelState, Fault> + WasmNotSend + WasmNotSync
79{
80}
81
82impl<F> Factory for F where
83    F: Fn(&[ParamValue], Option<u64>, Option<&GpuContext>) -> Result<ModelState, Fault> + WasmNotSend + WasmNotSync
84{
85}
86
87/// Closure that computes a GPU model's device demand at some params, for a device with some limits.
88pub(crate) trait Capacity: Fn(&[ParamValue], &wgpu::Limits) -> Demand + WasmNotSend + WasmNotSync {}
89
90impl<F: Fn(&[ParamValue], &wgpu::Limits) -> Demand + WasmNotSend + WasmNotSync> Capacity for F {}
91
92/// One model a host can offer: its declarations, and the factory that builds it on any device.
93///
94/// Cheap to clone. Clones share the declarations and the factory.
95#[derive(Clone)]
96pub struct ModelEntry {
97    parts: Arc<EntryParts>,
98}
99
100#[derive(Clone)]
101struct EntryParts {
102    id: String,
103    name: String,
104    description: String,
105    param_descriptors: Vec<ParamDescriptor>,
106    stat_descriptors: Vec<StatDescriptor>,
107    action_descriptors: Vec<ActionDescriptor>,
108    topology_hint: TopologyHint,
109    metadata: ModelMetadata,
110    gpu_needs: Option<GpuNeeds>,
111    source: ModelSource,
112    factory: Arc<dyn Factory>,
113    /// Closure that computes the device demand, `None` for a CPU model. A CPU model allocates on the host and has no
114    /// device limit to miss.
115    capacity: Option<Arc<dyn Capacity>>,
116}
117
118/// Prints the declarations, and leaves out the factory and the capacity closures.
119impl fmt::Debug for ModelEntry {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        let parts = &self.parts;
122        f.debug_struct("ModelEntry")
123            .field("id", &parts.id)
124            .field("name", &parts.name)
125            .field("description", &parts.description)
126            .field("param_descriptors", &parts.param_descriptors)
127            .field("stat_descriptors", &parts.stat_descriptors)
128            .field("action_descriptors", &parts.action_descriptors)
129            .field("topology_hint", &parts.topology_hint)
130            .field("metadata", &parts.metadata)
131            .field("gpu_needs", &parts.gpu_needs)
132            .field("source", &parts.source)
133            .finish_non_exhaustive()
134    }
135}
136
137impl ModelEntry {
138    /// Stable id that identifies the model in a model set, on the command line and in a spec file.
139    pub fn id(&self) -> &str {
140        &self.parts.id
141    }
142
143    /// Name shown in the UI.
144    pub fn name(&self) -> &str {
145        &self.parts.name
146    }
147
148    /// One-line description shown in the UI.
149    pub fn description(&self) -> &str {
150        &self.parts.description
151    }
152
153    /// Every parameter, in the order a build reads its values.
154    pub fn param_descriptors(&self) -> &[ParamDescriptor] {
155        &self.parts.param_descriptors
156    }
157
158    /// Stats the model reports, in the order its state returns them.
159    pub fn stat_descriptors(&self) -> &[StatDescriptor] {
160        &self.parts.stat_descriptors
161    }
162
163    /// One-off steps the model offers, in the order the Parameters panel draws their buttons.
164    pub fn action_descriptors(&self) -> &[ActionDescriptor] {
165        &self.parts.action_descriptors
166    }
167
168    /// Display layers the model presents.
169    pub fn topology_hint(&self) -> TopologyHint {
170        self.parts.topology_hint
171    }
172
173    /// Declared facts about the model, derived from its trait consts.
174    pub fn metadata(&self) -> &ModelMetadata {
175        &self.parts.metadata
176    }
177
178    /// Declarations a plan checks a spec against, readable with no device.
179    pub fn schema(&self) -> ModelSchema<'_> {
180        ModelSchema {
181            id: &self.parts.id,
182            params: &self.parts.param_descriptors,
183            stats: &self.parts.stat_descriptors,
184            actions: &self.parts.action_descriptors,
185        }
186    }
187
188    /// Device needs above the WebGPU baseline. `None` for a CPU model.
189    pub fn gpu_needs(&self) -> Option<GpuNeeds> {
190        self.parts.gpu_needs
191    }
192
193    /// Origin of the model's code: its type path always, and the registering crate's build once a set records it.
194    pub fn source(&self) -> &ModelSource {
195        &self.parts.source
196    }
197
198    /// Returns the position of parameter `id` in [`Self::param_descriptors`].
199    pub fn param_index(&self, id: &str) -> Option<usize> {
200        self.parts.param_descriptors.iter().position(|desc| desc.id == id)
201    }
202
203    /// Returns the position of action `id` in [`Self::action_descriptors`].
204    pub fn action_index(&self, id: &str) -> Option<usize> {
205        self.parts.action_descriptors.iter().position(|action| action.id == id)
206    }
207
208    /// Returns a setup at the declared defaults, the default seed and no actions.
209    pub fn setup(&self) -> RunSetup {
210        RunSetup::new(self.clone())
211    }
212
213    /// Builds the model from positional values, unchecked. [`RunSetup`] is the checked path.
214    ///
215    /// A CPU model ignores `gpu`, and a GPU model builds on it.
216    ///
217    /// # Errors
218    ///
219    /// Returns a [`Fault`] when the build panics, the device rejects it, or a GPU model receives no device.
220    pub fn build(
221        &self,
222        params: &[ParamValue],
223        seed: Option<u64>,
224        gpu: Option<&GpuContext>,
225    ) -> Result<ModelState, Fault> {
226        (self.parts.factory)(params, seed, gpu)
227    }
228
229    /// Returns the device resources the model would allocate at `params` on a device with `limits`.
230    ///
231    /// Returns `None` for a CPU model. It allocates on the host and knows its footprint only once built.
232    pub fn demand(&self, params: &[ParamValue], limits: &wgpu::Limits) -> Option<Demand> {
233        self.parts.capacity.as_ref().map(|capacity| capacity(params, limits))
234    }
235
236    /// Returns the reasons a device with `limits` cannot build the model at `params`, empty when it can.
237    pub fn shortfalls(&self, params: &[ParamValue], limits: &wgpu::Limits) -> Vec<String> {
238        self.demand(params, limits)
239            .map_or_else(Vec::new, |demand| demand.shortfalls(limits))
240    }
241
242    /// Returns the entry with its factory wrapped, as henad-explore's test harnesses do to inject a fault.
243    #[doc(hidden)]
244    pub fn wrap_factory(self, wrap: impl FnOnce(Arc<dyn Factory>) -> Arc<dyn Factory>) -> Self {
245        let mut parts = Arc::unwrap_or_clone(self.parts);
246        parts.factory = wrap(parts.factory);
247        Self { parts: Arc::new(parts) }
248    }
249
250    /// Returns the entry with its source replaced by `source`.
251    fn with_source(mut self, source: ModelSource) -> Self {
252        Arc::make_mut(&mut self.parts).source = source;
253        self
254    }
255}
256
257/// Returns the fault a GPU model raises when it receives no device.
258fn no_device(id: &str) -> Fault {
259    Fault::refused(
260        BUILDING,
261        format!("model '{id}' runs on the GPU and needs a device to build on"),
262    )
263}
264
265/// Returns an entry for the CPU grid model `M`.
266pub fn register_grid_model<M: GridModel>() -> ModelEntry {
267    ModelEntry {
268        parts: Arc::new(EntryParts {
269            id: M::ID.to_owned(),
270            name: M::NAME.to_owned(),
271            description: M::DESCRIPTION.to_owned(),
272            param_descriptors: grid_model_param_descriptors::<M>(),
273            stat_descriptors: M::STATS.to_vec(),
274            action_descriptors: M::ACTIONS.to_vec(),
275            topology_hint: TopologyHint::GRID,
276            metadata: ModelMetadata {
277                backend: Backend::Cpu,
278                palette: Some(M::PALETTE),
279                structure: Structure::Grid {
280                    neighborhood: M::NEIGHBORHOOD,
281                },
282                replays_exactly: true,
283            },
284            gpu_needs: None,
285            source: ModelSource::__from_type_path(type_name::<M>()),
286            factory: Arc::new(|params: &[ParamValue], seed: Option<u64>, _gpu: Option<&GpuContext>| {
287                catching(BUILDING, || {
288                    ModelState::Cpu(Box::new(GridModelState::<M>::from_params_seeded(params, seed)))
289                })
290            }),
291            capacity: None,
292        }),
293    }
294}
295
296/// Returns an entry for the CPU agent model `A`.
297pub fn register_agent_model<A: AgentModel>() -> ModelEntry {
298    ModelEntry {
299        parts: Arc::new(EntryParts {
300            id: A::ID.to_owned(),
301            name: A::NAME.to_owned(),
302            description: A::DESCRIPTION.to_owned(),
303            param_descriptors: agent_model_param_descriptors::<A>(),
304            stat_descriptors: A::STATS.to_vec(),
305            action_descriptors: A::ACTIONS.to_vec(),
306            topology_hint: TopologyHint {
307                grid: <A::Field as FieldLayer>::HAS_GRID,
308                agents: true,
309                edges: false,
310            },
311            metadata: ModelMetadata {
312                backend: Backend::Cpu,
313                palette: Some(A::PALETTE),
314                structure: Structure::Agents {
315                    chunk: A::CHUNK,
316                    lanes: <A::Lanes as AgentLanes>::LANES,
317                    index: <A::Index as NeighborIndex>::KIND,
318                    field: <A::Field as FieldLayer>::KIND,
319                },
320                replays_exactly: true,
321            },
322            gpu_needs: None,
323            source: ModelSource::__from_type_path(type_name::<A>()),
324            factory: Arc::new(|params: &[ParamValue], seed: Option<u64>, _gpu: Option<&GpuContext>| {
325                catching(BUILDING, || {
326                    ModelState::Cpu(Box::new(AgentModelState::<A>::from_params_seeded(params, seed)))
327                })
328            }),
329            capacity: None,
330        }),
331    }
332}
333
334/// Returns an entry for the CPU network model `N`.
335pub fn register_network_model<N: NetworkModel>() -> ModelEntry {
336    ModelEntry {
337        parts: Arc::new(EntryParts {
338            id: N::ID.to_owned(),
339            name: N::NAME.to_owned(),
340            description: N::DESCRIPTION.to_owned(),
341            param_descriptors: network_model_param_descriptors::<N>(),
342            stat_descriptors: N::STATS.to_vec(),
343            action_descriptors: N::ACTIONS.to_vec(),
344            topology_hint: TopologyHint::NETWORK,
345            metadata: ModelMetadata {
346                backend: Backend::Cpu,
347                palette: Some(N::PALETTE),
348                structure: Structure::Network {
349                    chunk: N::CHUNK,
350                    lanes: <N::Lanes as AgentLanes>::LANES,
351                    edge_palette: N::EDGE_PALETTE,
352                },
353                replays_exactly: true,
354            },
355            gpu_needs: None,
356            source: ModelSource::__from_type_path(type_name::<N>()),
357            factory: Arc::new(|params: &[ParamValue], seed: Option<u64>, _gpu: Option<&GpuContext>| {
358                catching(BUILDING, || {
359                    ModelState::Cpu(Box::new(NetworkModelState::<N>::from_params_seeded(params, seed)))
360                })
361            }),
362            capacity: None,
363        }),
364    }
365}
366
367/// Returns an entry for the GPU grid model `M`.
368///
369/// Every parameter applies on reload, since the GPU state accepts no live edit.
370pub fn register_gpu_grid_model<M: GpuGridModel>() -> ModelEntry {
371    ModelEntry {
372        parts: Arc::new(EntryParts {
373            id: M::ID.to_owned(),
374            name: M::NAME.to_owned(),
375            description: M::DESCRIPTION.to_owned(),
376            param_descriptors: M::param_descriptors()
377                .into_iter()
378                .map(ParamDescriptor::on_reload)
379                .collect(),
380            stat_descriptors: M::STATS.to_vec(),
381            action_descriptors: M::ACTIONS.iter().map(|action| action.desc).collect(),
382            // The grid reaches the view as a texture instead of a cell buffer.
383            topology_hint: TopologyHint::GRID,
384            metadata: ModelMetadata {
385                backend: Backend::Gpu,
386                palette: Some(M::PALETTE),
387                structure: Structure::GpuGrid {
388                    buffers: M::BUFFERS,
389                    workgroup: M::WORKGROUP_SIZE,
390                },
391                replays_exactly: M::REPLAYS_EXACTLY,
392            },
393            gpu_needs: Some(GpuNeeds::with_storage_buffers(GpuGridState::<M>::max_storage_bindings())),
394            source: ModelSource::__from_type_path(type_name::<M>()),
395            factory: Arc::new(|params: &[ParamValue], seed: Option<u64>, gpu: Option<&GpuContext>| {
396                let ctx = gpu.ok_or_else(|| no_device(M::ID))?;
397                catching_on(ctx, BUILDING, || {
398                    ModelState::Gpu(Box::new(GpuGridState::<M>::new_seeded(ctx, params, seed)))
399                })
400            }),
401            capacity: Some(Arc::new(GpuGridState::<M>::demand)),
402        }),
403    }
404}
405
406/// Returns an entry for the GPU agent model `M`.
407///
408/// Every parameter applies on reload, since the GPU state accepts no live edit.
409pub fn register_gpu_agent_model<M: GpuAgentModel>() -> ModelEntry {
410    ModelEntry {
411        parts: Arc::new(EntryParts {
412            id: M::ID.to_owned(),
413            name: M::NAME.to_owned(),
414            description: M::DESCRIPTION.to_owned(),
415            param_descriptors: M::param_descriptors()
416                .into_iter()
417                .map(ParamDescriptor::on_reload)
418                .collect(),
419            stat_descriptors: M::STATS.to_vec(),
420            action_descriptors: M::ACTIONS.iter().map(|action| action.desc).collect(),
421            // A model that declares a display pass draws a grid layer under its agents.
422            topology_hint: TopologyHint {
423                grid: M::DISPLAY.is_some(),
424                agents: true,
425                edges: false,
426            },
427            metadata: ModelMetadata {
428                backend: Backend::Gpu,
429                // A GPU agent model's shaders write RGBA themselves.
430                palette: None,
431                structure: Structure::GpuAgents {
432                    buffers: M::BUFFERS,
433                    passes: M::STEP_PASSES,
434                    index: M::INDEX,
435                    display: M::DISPLAY.is_some(),
436                    counters: M::COUNTERS,
437                },
438                replays_exactly: M::REPLAYS_EXACTLY,
439            },
440            gpu_needs: Some(GpuNeeds::with_storage_buffers(
441                GpuAgentState::<M>::max_storage_bindings(),
442            )),
443            source: ModelSource::__from_type_path(type_name::<M>()),
444            factory: Arc::new(|params: &[ParamValue], seed: Option<u64>, gpu: Option<&GpuContext>| {
445                let ctx = gpu.ok_or_else(|| no_device(M::ID))?;
446                catching_on(ctx, BUILDING, || {
447                    ModelState::Gpu(Box::new(GpuAgentState::<M>::new_seeded(ctx, params, seed)))
448                })
449            }),
450            capacity: Some(Arc::new(GpuAgentState::<M>::demand)),
451        }),
452    }
453}
henad_compute/entry/mod.rs

henad_compute/entry/
mod.rs
