sim_lib_interference_compute/
lower.rs

1//! Complete preflight followed by canonical Tensor request composition.
2
3use std::sync::Arc;
4
5use sim_kernel::{Cx, Symbol, Value};
6use sim_lib_interference_core::{InterferenceProblem, SamplingPlane};
7use sim_lib_numbers_tensor::{
8    CpuTensorExecutor, SubmissionEvidence, Tensor, TensorExecution, TensorExecutor,
9    TensorExecutorCard, TensorMeta, TensorOp, TensorRequest, TypedTensorStorage,
10    active_tensor_executor, add_op_symbol, cos_op_symbol, div_op_symbol, domains, exp_op_symbol,
11    mul_op_symbol, sin_op_symbol, sqrt_op_symbol, sub_op_symbol,
12};
13
14use crate::{
15    LoweringError, PhaseBudget, PlaneTile, PlaneTileConstants, PointTileConstants, PreflightCheck,
16    SourcePhaseEstimate, SourceTileConstants, TilePlan, TileProfile,
17    constants::prepare_source_constants,
18    coordinates::{prepare_local_coordinates, scalar},
19    preflight::{finite_f32, finite_nonzero_f32, validate_executor},
20};
21
22struct PreparedTile {
23    tile: PlaneTile,
24    local: [Tensor; 3],
25    sources: Vec<SourceTileConstants>,
26}
27
28/// One completed tile with separate canonical real and imaginary f32 Tensors.
29#[derive(Clone)]
30pub struct LoweredTile {
31    tile: PlaneTile,
32    real: Tensor,
33    imaginary: Tensor,
34    submissions: Vec<SubmissionEvidence>,
35}
36
37impl LoweredTile {
38    /// Returns the global placement and physical center of this tile.
39    pub fn tile(&self) -> PlaneTile {
40        self.tile
41    }
42
43    /// Returns the real component Tensor with shape `[tile.rows, tile.columns]`.
44    pub fn real(&self) -> &Tensor {
45        &self.real
46    }
47
48    /// Returns the imaginary component Tensor with shape `[tile.rows, tile.columns]`.
49    pub fn imaginary(&self) -> &Tensor {
50        &self.imaginary
51    }
52
53    /// Returns executor flush evidence in deterministic submission order.
54    pub fn submissions(&self) -> &[SubmissionEvidence] {
55        &self.submissions
56    }
57}
58
59/// A completely admitted lowering bound to one canonical Tensor executor.
60///
61/// Construction performs all provider, geometry, shape, allocation, and
62/// accuracy checks and creates every local-coordinate Tensor. Calling
63/// [`execute`](Self::execute) never selects a different executor.
64pub struct LoweringPlan {
65    executor: Arc<dyn TensorExecutor>,
66    executor_card: TensorExecutorCard,
67    nil_attributes: Value,
68    tile_plan: TilePlan,
69    prepared: Vec<PreparedTile>,
70    max_phase_estimate: SourcePhaseEstimate,
71    phase_budget: PhaseBudget,
72    wavenumber: f32,
73    attenuation: f32,
74}
75
76impl LoweringPlan {
77    /// Preflights a complete request against the active environment executor.
78    ///
79    /// When the environment has no active executor, the canonical CPU executor
80    /// is captured. No Tensor request is submitted by this constructor.
81    pub fn preflight(
82        cx: &mut Cx,
83        problem: &InterferenceProblem,
84        plane: SamplingPlane,
85        budget: PhaseBudget,
86        profile: TileProfile,
87    ) -> Result<Self, LoweringError> {
88        let executor =
89            active_tensor_executor(cx).unwrap_or_else(|| Arc::new(CpuTensorExecutor::new()));
90        Self::preflight_with_executor(cx, problem, plane, budget, profile, executor)
91    }
92
93    pub(crate) fn preflight_with_executor(
94        cx: &mut Cx,
95        problem: &InterferenceProblem,
96        plane: SamplingPlane,
97        budget: PhaseBudget,
98        profile: TileProfile,
99        executor: Arc<dyn TensorExecutor>,
100    ) -> Result<Self, LoweringError> {
101        budget.validate()?;
102        profile.validate()?;
103        let executor_card = executor.card();
104        validate_executor(&executor_card)?;
105
106        let wave = problem.wavenumber();
107        let wavenumber = finite_nonzero_f32("real wavenumber", wave.real_radians_per_metre())?;
108        let attenuation = finite_f32("attenuation coefficient", wave.imaginary_nepers_per_metre())?;
109        let tile_plan = TilePlan::new(plane, wave.real_radians_per_metre(), budget, profile)?;
110        let nil_attributes = cx.factory().nil().map_err(|error| {
111            LoweringError::new(
112                PreflightCheck::Shape,
113                format!("cannot construct Tensor operation attributes: {error}"),
114            )
115        })?;
116        let mut prepared = Vec::with_capacity(tile_plan.tiles().len());
117        let mut max_phase_estimate = SourcePhaseEstimate::default();
118        for tile in tile_plan.tiles().iter().copied() {
119            let local = prepare_local_coordinates(plane, tile)?;
120            let mut sources = Vec::with_capacity(problem.sources.len());
121            for source in &problem.sources {
122                let constants = prepare_source_constants(
123                    problem,
124                    source,
125                    tile.center(),
126                    &local.exact,
127                    &local.lowered,
128                    budget,
129                )?;
130                merge_estimate(&mut max_phase_estimate, constants.estimate());
131                sources.push(constants);
132            }
133            prepared.push(PreparedTile {
134                tile,
135                local: local.tensors,
136                sources,
137            });
138        }
139        Ok(Self {
140            executor,
141            executor_card,
142            nil_attributes,
143            tile_plan,
144            prepared,
145            max_phase_estimate,
146            phase_budget: budget,
147            wavenumber,
148            attenuation,
149        })
150    }
151
152    /// Returns the captured executor descriptor.
153    pub fn executor_card(&self) -> &TensorExecutorCard {
154        &self.executor_card
155    }
156
157    /// Returns the complete tile plan.
158    pub fn tile_plan(&self) -> &TilePlan {
159        &self.tile_plan
160    }
161
162    /// Returns the maximum estimate across every source and tile.
163    pub fn max_phase_estimate(&self) -> SourcePhaseEstimate {
164        self.max_phase_estimate
165    }
166
167    /// Returns the admitted residual-phase and predicted-error limits.
168    pub fn phase_budget(&self) -> PhaseBudget {
169        self.phase_budget
170    }
171
172    /// Returns the prepared source constants for one tile.
173    pub fn source_constants(&self, tile_index: usize) -> Option<&[SourceTileConstants]> {
174        self.prepared
175            .get(tile_index)
176            .map(|prepared| prepared.sources.as_slice())
177    }
178
179    /// Executes every prepared tile through the captured executor.
180    ///
181    /// Provider errors after accepted work are returned as execution failures;
182    /// this method never restarts a tile on another executor.
183    pub fn execute(&self, cx: &mut Cx) -> Result<Vec<LoweredTile>, LoweringError> {
184        let mut lowered = Vec::with_capacity(self.prepared.len());
185        for prepared in &self.prepared {
186            let mut source_fields = Vec::with_capacity(prepared.sources.len());
187            let mut submissions = Vec::new();
188            for source in &prepared.sources {
189                source_fields.push(self.lower_source(cx, prepared, *source)?);
190                submissions.push(self.flush()?);
191            }
192            let (real, imaginary) =
193                self.pairwise_accumulate(cx, prepared.tile, source_fields, &mut submissions)?;
194            lowered.push(LoweredTile {
195                tile: prepared.tile,
196                real,
197                imaginary,
198                submissions,
199            });
200        }
201        Ok(lowered)
202    }
203
204    pub(crate) fn execute_with_uploaded_inputs(
205        &self,
206        cx: &mut Cx,
207    ) -> Result<(Vec<LoweredTile>, usize), LoweringError> {
208        let mut lowered = Vec::with_capacity(self.prepared.len());
209        let mut uploads = 0_usize;
210        for prepared in &self.prepared {
211            let zero = Tensor::from_storage(
212                prepared.tile.shape().to_vec(),
213                domains::f32(),
214                Arc::new(TypedTensorStorage::<f32>::new(vec![
215                    0.0;
216                    prepared.tile.rows()
217                        * prepared
218                            .tile
219                            .columns()
220                ])),
221            )
222            .map_err(|error| {
223                LoweringError::new(
224                    PreflightCheck::Execution,
225                    format!("cannot construct resident-upload zero Tensor: {error}"),
226                )
227            })?;
228            let local = [
229                self.binary(
230                    cx,
231                    add_op_symbol(),
232                    &prepared.local[0],
233                    &zero,
234                    &prepared.tile.shape(),
235                )?,
236                self.binary(
237                    cx,
238                    add_op_symbol(),
239                    &prepared.local[1],
240                    &zero,
241                    &prepared.tile.shape(),
242                )?,
243                self.binary(
244                    cx,
245                    add_op_symbol(),
246                    &prepared.local[2],
247                    &zero,
248                    &prepared.tile.shape(),
249                )?,
250            ];
251            uploads = uploads.checked_add(local.len()).ok_or_else(|| {
252                LoweringError::new(
253                    PreflightCheck::Execution,
254                    "resident input upload count overflowed usize",
255                )
256            })?;
257            let uploaded = PreparedTile {
258                tile: prepared.tile,
259                local,
260                sources: prepared.sources.clone(),
261            };
262            let mut source_fields = Vec::with_capacity(uploaded.sources.len());
263            let mut submissions = vec![self.flush()?];
264            for source in &uploaded.sources {
265                source_fields.push(self.lower_source(cx, &uploaded, *source)?);
266                submissions.push(self.flush()?);
267            }
268            let (real, imaginary) =
269                self.pairwise_accumulate(cx, uploaded.tile, source_fields, &mut submissions)?;
270            lowered.push(LoweredTile {
271                tile: uploaded.tile,
272                real,
273                imaginary,
274                submissions,
275            });
276        }
277        Ok((lowered, uploads))
278    }
279
280    fn lower_source(
281        &self,
282        cx: &mut Cx,
283        tile: &PreparedTile,
284        source: SourceTileConstants,
285    ) -> Result<(Tensor, Tensor), LoweringError> {
286        let delta = match source {
287            SourceTileConstants::Point { constants, .. } => {
288                self.point_delta(cx, tile, constants)?
289            }
290            SourceTileConstants::ForwardPlane { constants, .. } => {
291                self.plane_delta(cx, tile, constants)?
292            }
293        };
294        let shape = tile.tile.shape().to_vec();
295        let psi = self.binary(
296            cx,
297            mul_op_symbol(),
298            &delta,
299            &scalar(self.wavenumber)?,
300            &shape,
301        )?;
302        let phase_sin = self.unary(cx, sin_op_symbol(), &psi, &shape)?;
303        let phase_cos = self.unary(cx, cos_op_symbol(), &psi, &shape)?;
304        let (anchor_cos, anchor_sin, gain0) = match source {
305            SourceTileConstants::Point { constants, .. } => {
306                (constants.phase_cos, constants.phase_sin, constants.gain0)
307            }
308            SourceTileConstants::ForwardPlane { constants, .. } => {
309                (constants.phase_cos, constants.phase_sin, constants.gain0)
310            }
311        };
312        let cos_cos = self.binary(
313            cx,
314            mul_op_symbol(),
315            &phase_cos,
316            &scalar(anchor_cos)?,
317            &shape,
318        )?;
319        let sin_sin = self.binary(
320            cx,
321            mul_op_symbol(),
322            &phase_sin,
323            &scalar(anchor_sin)?,
324            &shape,
325        )?;
326        let real_phase = self.binary(cx, sub_op_symbol(), &cos_cos, &sin_sin, &shape)?;
327        let sin_cos = self.binary(
328            cx,
329            mul_op_symbol(),
330            &phase_sin,
331            &scalar(anchor_cos)?,
332            &shape,
333        )?;
334        let cos_sin = self.binary(
335            cx,
336            mul_op_symbol(),
337            &phase_cos,
338            &scalar(anchor_sin)?,
339            &shape,
340        )?;
341        let imaginary_phase = self.binary(cx, add_op_symbol(), &sin_cos, &cos_sin, &shape)?;
342        let attenuation_argument = self.binary(
343            cx,
344            mul_op_symbol(),
345            &delta,
346            &scalar(-self.attenuation)?,
347            &shape,
348        )?;
349        let attenuation = self.unary(cx, exp_op_symbol(), &attenuation_argument, &shape)?;
350        let gain = match source {
351            SourceTileConstants::Point { constants, .. } => {
352                let ratio_offset =
353                    self.binary(cx, mul_op_symbol(), &delta, &scalar(constants.rho)?, &shape)?;
354                let ratio_denominator =
355                    self.binary(cx, add_op_symbol(), &scalar(1.0)?, &ratio_offset, &shape)?;
356                let center_scaled =
357                    self.binary(cx, mul_op_symbol(), &attenuation, &scalar(gain0)?, &shape)?;
358                self.binary(
359                    cx,
360                    div_op_symbol(),
361                    &center_scaled,
362                    &ratio_denominator,
363                    &shape,
364                )?
365            }
366            SourceTileConstants::ForwardPlane { .. } => {
367                self.binary(cx, mul_op_symbol(), &attenuation, &scalar(gain0)?, &shape)?
368            }
369        };
370        Ok((
371            self.binary(cx, mul_op_symbol(), &gain, &real_phase, &shape)?,
372            self.binary(cx, mul_op_symbol(), &gain, &imaginary_phase, &shape)?,
373        ))
374    }
375
376    fn point_delta(
377        &self,
378        cx: &mut Cx,
379        tile: &PreparedTile,
380        constants: PointTileConstants,
381    ) -> Result<Tensor, LoweringError> {
382        let shape = tile.tile.shape().to_vec();
383        let a = self.dot_local(cx, tile, constants.n0)?;
384        let b = self.dot_local_tensors(cx, tile)?;
385        let two_a = self.binary(cx, mul_op_symbol(), &a, &scalar(2.0)?, &shape)?;
386        let two_a_rho =
387            self.binary(cx, mul_op_symbol(), &two_a, &scalar(constants.rho)?, &shape)?;
388        let b_rho_squared = self.binary(
389            cx,
390            mul_op_symbol(),
391            &b,
392            &scalar(constants.rho * constants.rho)?,
393            &shape,
394        )?;
395        let z = self.binary(cx, add_op_symbol(), &two_a_rho, &b_rho_squared, &shape)?;
396        let b_rho = self.binary(cx, mul_op_symbol(), &b, &scalar(constants.rho)?, &shape)?;
397        let numerator = self.binary(cx, add_op_symbol(), &two_a, &b_rho, &shape)?;
398        let sqrt_argument = self.binary(cx, add_op_symbol(), &scalar(1.0)?, &z, &shape)?;
399        let root = self.unary(cx, sqrt_op_symbol(), &sqrt_argument, &shape)?;
400        let denominator = self.binary(cx, add_op_symbol(), &root, &scalar(1.0)?, &shape)?;
401        self.binary(cx, div_op_symbol(), &numerator, &denominator, &shape)
402    }
403
404    fn plane_delta(
405        &self,
406        cx: &mut Cx,
407        tile: &PreparedTile,
408        constants: PlaneTileConstants,
409    ) -> Result<Tensor, LoweringError> {
410        self.dot_local(cx, tile, constants.direction)
411    }
412
413    fn dot_local(
414        &self,
415        cx: &mut Cx,
416        tile: &PreparedTile,
417        vector: [f32; 3],
418    ) -> Result<Tensor, LoweringError> {
419        let shape = tile.tile.shape().to_vec();
420        let x = self.binary(
421            cx,
422            mul_op_symbol(),
423            &tile.local[0],
424            &scalar(vector[0])?,
425            &shape,
426        )?;
427        let y = self.binary(
428            cx,
429            mul_op_symbol(),
430            &tile.local[1],
431            &scalar(vector[1])?,
432            &shape,
433        )?;
434        let z = self.binary(
435            cx,
436            mul_op_symbol(),
437            &tile.local[2],
438            &scalar(vector[2])?,
439            &shape,
440        )?;
441        let xy = self.binary(cx, add_op_symbol(), &x, &y, &shape)?;
442        self.binary(cx, add_op_symbol(), &xy, &z, &shape)
443    }
444
445    fn dot_local_tensors(&self, cx: &mut Cx, tile: &PreparedTile) -> Result<Tensor, LoweringError> {
446        let shape = tile.tile.shape().to_vec();
447        let x = self.binary(cx, mul_op_symbol(), &tile.local[0], &tile.local[0], &shape)?;
448        let y = self.binary(cx, mul_op_symbol(), &tile.local[1], &tile.local[1], &shape)?;
449        let z = self.binary(cx, mul_op_symbol(), &tile.local[2], &tile.local[2], &shape)?;
450        let xy = self.binary(cx, add_op_symbol(), &x, &y, &shape)?;
451        self.binary(cx, add_op_symbol(), &xy, &z, &shape)
452    }
453
454    fn pairwise_accumulate(
455        &self,
456        cx: &mut Cx,
457        tile: PlaneTile,
458        mut fields: Vec<(Tensor, Tensor)>,
459        submissions: &mut Vec<SubmissionEvidence>,
460    ) -> Result<(Tensor, Tensor), LoweringError> {
461        let shape = tile.shape().to_vec();
462        while fields.len() > 1 {
463            let mut next = Vec::with_capacity(fields.len().div_ceil(2));
464            let mut pairs = fields.into_iter();
465            while let Some(left) = pairs.next() {
466                if let Some(right) = pairs.next() {
467                    next.push((
468                        self.binary(cx, add_op_symbol(), &left.0, &right.0, &shape)?,
469                        self.binary(cx, add_op_symbol(), &left.1, &right.1, &shape)?,
470                    ));
471                } else {
472                    next.push(left);
473                }
474            }
475            submissions.push(self.flush()?);
476            fields = next;
477        }
478        fields.into_iter().next().ok_or_else(|| {
479            LoweringError::new(
480                PreflightCheck::Execution,
481                "source accumulation received no source fields",
482            )
483        })
484    }
485
486    fn binary(
487        &self,
488        cx: &mut Cx,
489        symbol: Symbol,
490        left: &Tensor,
491        right: &Tensor,
492        shape: &[usize],
493    ) -> Result<Tensor, LoweringError> {
494        self.submit(cx, symbol, vec![left.clone(), right.clone()], shape)
495    }
496
497    fn unary(
498        &self,
499        cx: &mut Cx,
500        symbol: Symbol,
501        input: &Tensor,
502        shape: &[usize],
503    ) -> Result<Tensor, LoweringError> {
504        self.submit(cx, symbol, vec![input.clone()], shape)
505    }
506
507    fn submit(
508        &self,
509        cx: &mut Cx,
510        symbol: Symbol,
511        inputs: Vec<Tensor>,
512        shape: &[usize],
513    ) -> Result<Tensor, LoweringError> {
514        let request = TensorRequest::new(
515            TensorOp::new(symbol.clone(), self.nil_attributes.clone()),
516            inputs,
517            TensorMeta::new(shape.to_vec(), domains::f32()),
518        );
519        let tensor = match self.executor.execute(cx, request).map_err(|error| {
520            LoweringError::new(
521                PreflightCheck::Execution,
522                format!(
523                    "executor {} failed {symbol}: {error}",
524                    self.executor_card.symbol
525                ),
526            )
527        })? {
528            TensorExecution::Complete(tensor) => tensor,
529            TensorExecution::Unsupported { reason } => {
530                return Err(LoweringError::new(
531                    PreflightCheck::Execution,
532                    format!(
533                        "executor {} advertised then declined {symbol}: {reason}",
534                        self.executor_card.symbol
535                    ),
536                ));
537            }
538        };
539        if tensor.shape() != shape || tensor.dtype() != &domains::f32() {
540            return Err(LoweringError::new(
541                PreflightCheck::Execution,
542                format!(
543                    "executor {} returned shape {:?} dtype {} for {symbol}, expected {shape:?} numbers/f32",
544                    self.executor_card.symbol,
545                    tensor.shape(),
546                    tensor.dtype()
547                ),
548            ));
549        }
550        Ok(tensor)
551    }
552
553    fn flush(&self) -> Result<SubmissionEvidence, LoweringError> {
554        self.executor.flush().map_err(|error| {
555            LoweringError::new(
556                PreflightCheck::Execution,
557                format!(
558                    "executor {} flush failed: {error}",
559                    self.executor_card.symbol
560                ),
561            )
562        })
563    }
564}
565
566fn merge_estimate(total: &mut SourcePhaseEstimate, next: SourcePhaseEstimate) {
567    total.max_abs_residual_phase_rad = total
568        .max_abs_residual_phase_rad
569        .max(next.max_abs_residual_phase_rad);
570    total.max_predicted_geometry_error_rad = total
571        .max_predicted_geometry_error_rad
572        .max(next.max_predicted_geometry_error_rad);
573    total.max_predicted_roundoff_error_rad = total
574        .max_predicted_roundoff_error_rad
575        .max(next.max_predicted_roundoff_error_rad);
576}
sim_lib_interference_compute/lower.rs

sim_lib_interference_compute/
lower.rs
