sim_lib_interference_runtime/
tensor_bridge.rs

1//! One-copy adapters between reference host fields and canonical Tensors.
2
3use std::sync::Arc;
4
5use sim_kernel::{
6    Cx, DefaultFactory, EagerPolicy, Expr, NumberLiteral, ObjectEncode, Result, Symbol, Value,
7};
8use sim_lib_interference_solve::HostPhasorField;
9use sim_lib_numbers_tensor::{
10    Tensor, TensorLocation, TensorStorage, TypedTensorStorage, domains, tensor_value_class_symbol,
11};
12
13use crate::citizen::{RecordCitizenSpec, encode_field, invalid, next_field, read_construct_parts};
14
15/// Tensor-backed runtime phasor field.
16///
17/// The two component Tensors always have shape `[rows, cols]`, the same real
18/// dtype, and compatible placement. No complex or grid storage is introduced.
19#[derive(Clone)]
20pub struct PhasorFieldDescriptor {
21    /// Row count.
22    pub rows: usize,
23    /// Column count.
24    pub cols: usize,
25    /// Canonical real-component Tensor.
26    pub real: Tensor,
27    /// Canonical imaginary-component Tensor.
28    pub imag: Tensor,
29}
30
31impl PhasorFieldDescriptor {
32    /// Admits two compatible canonical component Tensors.
33    pub fn new(rows: usize, cols: usize, real: Tensor, imag: Tensor) -> Result<Self> {
34        let value = Self {
35            rows,
36            cols,
37            real,
38            imag,
39        };
40        value.validate()?;
41        Ok(value)
42    }
43
44    /// Moves a reference host field into two typed `f64` Tensor storages.
45    ///
46    /// The component buffers are consumed and shared directly by their Tensor;
47    /// there is no intermediate grid or per-cell runtime-value allocation.
48    pub fn from_host(field: HostPhasorField) -> Result<Self> {
49        let (rows, cols, real, imag) = field.into_component_planes();
50        let shape = vec![rows, cols];
51        let real = Tensor::from_storage(
52            shape.clone(),
53            domains::f64(),
54            Arc::new(TypedTensorStorage::<f64>::new(real)),
55        )?;
56        let imag = Tensor::from_storage(
57            shape,
58            domains::f64(),
59            Arc::new(TypedTensorStorage::<f64>::new(imag)),
60        )?;
61        Self::new(rows, cols, real, imag)
62    }
63
64    /// Explicitly materializes both components into one host field.
65    ///
66    /// Each resident Tensor receives exactly one `materialize` call. Host
67    /// storage remains zero-copy until the returned `HostPhasorField` takes its
68    /// owned `f64` component buffers.
69    pub fn materialize_host(&self, cx: &mut Cx) -> Result<HostPhasorField> {
70        self.validate_metadata()?;
71        let real_storage = self.real.materialize()?;
72        let imag_storage = self.imag.materialize()?;
73        let real = materialized_f64(cx, real_storage, self.real.dtype(), "real")?;
74        let imag = materialized_f64(cx, imag_storage, self.imag.dtype(), "imaginary")?;
75        HostPhasorField::from_component_planes(self.rows, self.cols, real, imag)
76            .map_err(|error| invalid("PhasorField", format!("{error:?}")))
77    }
78
79    fn validate_metadata(&self) -> Result<()> {
80        let cells = self
81            .rows
82            .checked_mul(self.cols)
83            .ok_or_else(|| invalid("PhasorField", "rows * cols overflowed"))?;
84        if self.rows == 0 || self.cols == 0 {
85            return Err(invalid("PhasorField", "rows and cols must be non-zero"));
86        }
87        let expected = [self.rows, self.cols];
88        if self.real.shape() != expected || self.imag.shape() != expected {
89            return Err(invalid(
90                "PhasorField",
91                format!(
92                    "component shapes must both be [{}, {}], found {:?} and {:?}",
93                    self.rows,
94                    self.cols,
95                    self.real.shape(),
96                    self.imag.shape()
97                ),
98            ));
99        }
100        if self.real.len() != cells || self.imag.len() != cells {
101            return Err(invalid(
102                "PhasorField",
103                "component cell counts do not match rows * cols",
104            ));
105        }
106        if self.real.dtype() != self.imag.dtype() {
107            return Err(invalid(
108                "PhasorField",
109                "real and imaginary dtypes must match",
110            ));
111        }
112        if !matches_real_dtype(self.real.dtype()) {
113            return Err(invalid(
114                "PhasorField",
115                format!(
116                    "component dtype must be numbers/f32 or numbers/f64, found {}",
117                    self.real.dtype()
118                ),
119            ));
120        }
121        if !compatible_locations(self.real.location(), self.imag.location()) {
122            return Err(invalid(
123                "PhasorField",
124                "component locations must both be host or resident at the same site",
125            ));
126        }
127        Ok(())
128    }
129}
130
131impl core::fmt::Debug for PhasorFieldDescriptor {
132    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
133        formatter
134            .debug_struct("PhasorFieldDescriptor")
135            .field("rows", &self.rows)
136            .field("cols", &self.cols)
137            .field("real_shape", &self.real.shape())
138            .field("imag_shape", &self.imag.shape())
139            .field("dtype", &self.real.dtype())
140            .field("real_location", &self.real.location())
141            .field("imag_location", &self.imag.location())
142            .finish()
143    }
144}
145
146impl PartialEq for PhasorFieldDescriptor {
147    fn eq(&self, other: &Self) -> bool {
148        self.rows == other.rows
149            && self.cols == other.cols
150            && tensor_eq(&self.real, &other.real)
151            && tensor_eq(&self.imag, &other.imag)
152    }
153}
154
155impl RecordCitizenSpec for PhasorFieldDescriptor {
156    const FIELDS: &'static [&'static str] = &["rows", "cols", "real", "imag"];
157
158    fn encode_fields(&self, cx: &mut Cx) -> Result<Vec<Expr>> {
159        Ok(vec![
160            encode_field(&self.rows),
161            encode_field(&self.cols),
162            tensor_expr(cx, &self.real)?,
163            tensor_expr(cx, &self.imag)?,
164        ])
165    }
166
167    fn decode_fields(cx: &mut Cx, fields: Vec<Value>) -> Result<Self> {
168        let mut fields = fields.into_iter();
169        let rows = next_field(cx, &mut fields, "rows")?;
170        let cols = next_field(cx, &mut fields, "cols")?;
171        let real = decode_tensor(
172            cx,
173            fields
174                .next()
175                .ok_or_else(|| invalid("PhasorField", "missing real Tensor"))?,
176            "real",
177        )?;
178        let imag = decode_tensor(
179            cx,
180            fields
181                .next()
182                .ok_or_else(|| invalid("PhasorField", "missing imaginary Tensor"))?,
183            "imag",
184        )?;
185        Self::new(rows, cols, real, imag)
186    }
187
188    fn example() -> Self {
189        Self::from_host(
190            HostPhasorField::from_component_planes(
191                2,
192                2,
193                vec![1.0, 0.0, -1.0, 0.5],
194                vec![0.0, 1.0, 0.0, -0.5],
195            )
196            .expect("example host field"),
197        )
198        .expect("example Tensor field")
199    }
200
201    fn validate(&self) -> Result<()> {
202        self.validate_metadata()?;
203        if self.real.location() == TensorLocation::Host {
204            validate_host_finite(&self.real, "real")?;
205            validate_host_finite(&self.imag, "imaginary")?;
206        }
207        Ok(())
208    }
209}
210
211impl_record_citizen!(PhasorFieldDescriptor, "interference/PhasorField", 4);
212
213pub(crate) fn tensor_expr(cx: &mut Cx, tensor: &Tensor) -> Result<Expr> {
214    match tensor.object_encoding(cx)? {
215        sim_kernel::ObjectEncoding::Constructor { class, args } => Ok(Expr::Extension {
216            tag: Symbol::qualified("citizen", "read-construct"),
217            payload: Box::new(Expr::Vector(
218                std::iter::once(Expr::Symbol(class)).chain(args).collect(),
219            )),
220        }),
221        _ => Err(invalid(
222            "Tensor",
223            "canonical Tensor did not expose a constructor encoding",
224        )),
225    }
226}
227
228pub(crate) fn decode_tensor(cx: &mut Cx, value: Value, field: &'static str) -> Result<Tensor> {
229    let expr = sim_citizen::value_to_expr(cx, value, field)?;
230    decode_tensor_expr(cx, &expr, field)
231}
232
233pub(crate) fn decode_tensor_expr(cx: &mut Cx, expr: &Expr, field: &'static str) -> Result<Tensor> {
234    let (class, args) = read_construct_parts(expr, field)?;
235    if class != tensor_value_class_symbol() {
236        return Err(invalid(
237            field,
238            format!(
239                "expected nested {}, found {class}",
240                tensor_value_class_symbol()
241            ),
242        ));
243    }
244    let [Expr::Symbol(version), shape, data, Expr::Symbol(dtype)] = args else {
245        return Err(invalid(
246            field,
247            "Tensor constructor must contain version, shape, data, and dtype",
248        ));
249    };
250    if *version != Symbol::new("v1") {
251        return Err(invalid(field, "Tensor constructor version must be v1"));
252    }
253    let Expr::List(dimensions) = shape else {
254        return Err(invalid(field, "Tensor shape must be a list"));
255    };
256    let shape = dimensions
257        .iter()
258        .map(|dimension| decode_dimension(dimension, field))
259        .collect::<Result<Vec<_>>>()?;
260    let Expr::List(cells) = data else {
261        return Err(invalid(field, "Tensor data must be a list"));
262    };
263    let cells = cells
264        .iter()
265        .map(|cell| sim_citizen::value_from_expr(cx, cell))
266        .collect::<Result<Vec<_>>>()?;
267    Tensor::new_exact(shape, dtype.clone(), cells)
268}
269
270fn decode_dimension(expr: &Expr, field: &'static str) -> Result<usize> {
271    let Expr::Number(NumberLiteral { domain, canonical }) = expr else {
272        return Err(invalid(field, "Tensor dimensions must be integers"));
273    };
274    if *domain != Symbol::qualified("citizen", "int") {
275        return Err(invalid(field, "Tensor dimensions must use citizen/int"));
276    }
277    canonical
278        .parse::<usize>()
279        .map_err(|_| invalid(field, "Tensor dimension is not a usize"))
280}
281
282fn compatible_locations(real: TensorLocation, imag: TensorLocation) -> bool {
283    match (real, imag) {
284        (TensorLocation::Host, TensorLocation::Host) => true,
285        (
286            TensorLocation::Resident { site: real, .. },
287            TensorLocation::Resident { site: imag, .. },
288        ) => real == imag,
289        _ => false,
290    }
291}
292
293fn matches_real_dtype(dtype: &Symbol) -> bool {
294    *dtype == domains::f32() || *dtype == domains::f64()
295}
296
297fn validate_host_finite(tensor: &Tensor, component: &'static str) -> Result<()> {
298    let mut cx = bare_cx();
299    let storage = tensor.materialize()?;
300    for index in 0..storage.len() {
301        let value = scalar_to_f64(
302            &mut cx,
303            storage.cell(index)?,
304            tensor.dtype(),
305            component,
306            index,
307        )?;
308        if !value.is_finite() {
309            return Err(invalid(
310                "PhasorField",
311                format!("{component} cell {index} must be finite"),
312            ));
313        }
314    }
315    Ok(())
316}
317
318fn materialized_f64(
319    cx: &mut Cx,
320    storage: Arc<dyn TensorStorage>,
321    dtype: &Symbol,
322    component: &'static str,
323) -> Result<Vec<f64>> {
324    (0..storage.len())
325        .map(|index| scalar_to_f64(cx, storage.cell(index)?, dtype, component, index))
326        .collect()
327}
328
329fn scalar_to_f64(
330    cx: &mut Cx,
331    value: Value,
332    dtype: &Symbol,
333    component: &'static str,
334    index: usize,
335) -> Result<f64> {
336    let Expr::Number(number) = value.object().as_expr(cx)? else {
337        return Err(invalid(
338            "PhasorField",
339            format!("{component} cell {index} is not a scalar number"),
340        ));
341    };
342    if number.domain != *dtype {
343        return Err(invalid(
344            "PhasorField",
345            format!(
346                "{component} cell {index} domain {} does not match {dtype}",
347                number.domain
348            ),
349        ));
350    }
351    let value = number.canonical.parse::<f64>().map_err(|_| {
352        invalid(
353            "PhasorField",
354            format!("{component} cell {index} is not a finite real literal"),
355        )
356    })?;
357    if value.is_finite() {
358        Ok(value)
359    } else {
360        Err(invalid(
361            "PhasorField",
362            format!("{component} cell {index} must be finite"),
363        ))
364    }
365}
366
367pub(crate) fn tensor_eq(left: &Tensor, right: &Tensor) -> bool {
368    if left.shape() != right.shape()
369        || left.dtype() != right.dtype()
370        || left.location() != right.location()
371    {
372        return false;
373    }
374    if left.location() != TensorLocation::Host {
375        return Arc::ptr_eq(left.storage(), right.storage());
376    }
377    let mut cx = bare_cx();
378    let (Ok(left), Ok(right)) = (left.cells(), right.cells()) else {
379        return false;
380    };
381    left.len() == right.len()
382        && left.iter().zip(right.iter()).all(|(left, right)| {
383            let left = left.object().as_expr(&mut cx);
384            let right = right.object().as_expr(&mut cx);
385            matches!((left, right), (Ok(left), Ok(right)) if sim_citizen::expr_citizen_eq(&left, &right))
386        })
387}
388
389fn bare_cx() -> Cx {
390    Cx::new(Arc::new(EagerPolicy), Arc::new(DefaultFactory))
391}
sim_lib_interference_runtime/tensor_bridge.rs

sim_lib_interference_runtime/
tensor_bridge.rs
