Skip to main content

antecedent_kernels/
dispatch.rs

1//! Once-per-batch kernel dispatch.
2//!
3//! SPDX-License-Identifier: MIT OR Apache-2.0
4
5use antecedent_core::KernelPolicy;
6
7use crate::view::{BitMaskView, F64VectorView};
8
9/// Selected implementation class for a batch.
10#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
11pub enum KernelImpl {
12    /// Scalar reference.
13    Scalar,
14    /// Portable optimized.
15    PortableOptimized,
16    /// Architecture-specific SIMD (requires `simd-runtime` feature + CPU support).
17    ArchSimd,
18}
19
20/// Whether an arch-SIMD path is compiled in and available on this CPU.
21///
22/// Always `false` until a justified `simd-runtime` kernel lands.
23/// `KernelPolicy::allow_arch_simd` is still consulted by [`select_impl`]; when this
24/// returns false, selection falls through to portable/scalar.
25#[must_use]
26pub fn arch_simd_available() -> bool {
27    false
28}
29
30/// Resolve the implementation for a batch from policy.
31///
32/// Selection order:
33/// 1. `force_scalar`, or neither portable nor arch allowed → [`KernelImpl::Scalar`]
34/// 2. `allow_arch_simd` and compiled `simd-runtime` and CPU support → [`KernelImpl::ArchSimd`]
35/// 3. `allow_portable_optimized` → [`KernelImpl::PortableOptimized`]
36/// 4. else → [`KernelImpl::Scalar`]
37#[must_use]
38pub fn select_impl(policy: &KernelPolicy) -> KernelImpl {
39    if policy.force_scalar {
40        return KernelImpl::Scalar;
41    }
42    let arch_ok = policy.allow_arch_simd && arch_simd_available();
43    let portable_ok = policy.allow_portable_optimized;
44    if !arch_ok && !portable_ok {
45        return KernelImpl::Scalar;
46    }
47    if arch_ok {
48        return KernelImpl::ArchSimd;
49    }
50    if portable_ok { KernelImpl::PortableOptimized } else { KernelImpl::Scalar }
51}
52
53fn portable_or_scalar_reductions(policy: KernelPolicy) -> KernelImpl {
54    match select_impl(&policy) {
55        KernelImpl::ArchSimd => KernelImpl::PortableOptimized,
56        other => other,
57    }
58}
59
60/// Public semantic entry: masked sum.
61#[must_use]
62pub fn masked_sum(
63    policy: &KernelPolicy,
64    x: F64VectorView<'_>,
65    mask: Option<BitMaskView<'_>>,
66) -> f64 {
67    match portable_or_scalar_reductions(*policy) {
68        KernelImpl::Scalar => crate::scalar::masked_sum(x, mask),
69        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
70            crate::portable::masked_sum(x, mask)
71        }
72    }
73}
74
75/// Public semantic entry: masked mean.
76#[must_use]
77pub fn masked_mean(
78    policy: &KernelPolicy,
79    x: F64VectorView<'_>,
80    mask: Option<BitMaskView<'_>>,
81) -> Option<f64> {
82    match portable_or_scalar_reductions(*policy) {
83        KernelImpl::Scalar => crate::scalar::masked_mean(x, mask),
84        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
85            crate::portable::masked_mean(x, mask)
86        }
87    }
88}
89
90/// Public semantic entry: masked population variance.
91#[must_use]
92pub fn masked_variance(
93    policy: &KernelPolicy,
94    x: F64VectorView<'_>,
95    mask: Option<BitMaskView<'_>>,
96) -> Option<f64> {
97    match portable_or_scalar_reductions(*policy) {
98        KernelImpl::Scalar => crate::scalar::masked_variance(x, mask),
99        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
100            crate::portable::masked_variance(x, mask)
101        }
102    }
103}
104
105/// Public semantic entry: gather.
106pub fn gather(policy: &KernelPolicy, src: F64VectorView<'_>, indices: &[usize], out: &mut [f64]) {
107    match portable_or_scalar_reductions(*policy) {
108        KernelImpl::Scalar => crate::scalar::gather(src, indices, out),
109        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
110            crate::portable::gather(src, indices, out);
111        }
112    }
113}
114
115/// Public semantic entry: copy.
116pub fn copy_vec(policy: &KernelPolicy, src: F64VectorView<'_>, dst: &mut [f64]) {
117    match portable_or_scalar_reductions(*policy) {
118        KernelImpl::Scalar => crate::scalar::copy_vec(src, dst),
119        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => crate::portable::copy_vec(src, dst),
120    }
121}
122
123/// Public semantic entry: partial correlation of `x` and `y` given `z_cols`.
124#[must_use]
125pub fn partial_correlation(
126    policy: &KernelPolicy,
127    x: &[f64],
128    y: &[f64],
129    z_cols: &[&[f64]],
130    workspace: &mut crate::parcorr::ParCorrWorkspace,
131) -> Option<f64> {
132    match portable_or_scalar_reductions(*policy) {
133        KernelImpl::Scalar => crate::parcorr::partial_correlation_scalar(x, y, z_cols, workspace),
134        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
135            crate::parcorr::partial_correlation_portable(x, y, z_cols, workspace)
136        }
137    }
138}
139
140/// Public semantic entry: masked population covariance.
141///
142/// Contract: deterministic reduction; `StableFloat` tolerance; beneficial for `n ≳ 64`.
143#[must_use]
144pub fn masked_covariance(
145    policy: &KernelPolicy,
146    x: F64VectorView<'_>,
147    y: F64VectorView<'_>,
148    mask: Option<BitMaskView<'_>>,
149) -> Option<f64> {
150    match portable_or_scalar_reductions(*policy) {
151        KernelImpl::Scalar => crate::scalar::masked_covariance(x, y, mask),
152        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
153            crate::portable::masked_covariance(x, y, mask)
154        }
155    }
156}
157
158/// Public semantic entry: in-place standardization (sample SD).
159///
160/// Contract: `StableFloat`; beneficial for `n ≳ 32`.
161#[must_use]
162pub fn standardize_inplace(policy: &KernelPolicy, x: &mut [f64], eps: f64) -> (f64, f64) {
163    match portable_or_scalar_reductions(*policy) {
164        KernelImpl::Scalar => crate::scalar::standardize_inplace(x, eps),
165        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
166            crate::portable::standardize_inplace(x, eps)
167        }
168    }
169}
170
171/// Public semantic entry: pairwise L1 distance matrix fill.
172///
173/// Contract: exact for finite inputs; beneficial for `n ≳ 64`.
174pub fn pairwise_l1_fill(policy: &KernelPolicy, x: &[f64], out: &mut [f64]) {
175    match portable_or_scalar_reductions(*policy) {
176        KernelImpl::Scalar => crate::scalar::pairwise_l1_fill(x, out),
177        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
178            crate::portable::pairwise_l1_fill(x, out);
179        }
180    }
181}
182
183/// Public semantic entry: contingency table accumulation.
184///
185/// Contract: exact integer counts as `f64`; beneficial for `n ≳ 256`.
186pub fn accumulate_contingency(
187    policy: &KernelPolicy,
188    x_codes: &[u32],
189    y_codes: &[u32],
190    out: &mut [f64],
191    n_y_levels: usize,
192) {
193    match portable_or_scalar_reductions(*policy) {
194        KernelImpl::Scalar => {
195            crate::scalar::accumulate_contingency(x_codes, y_codes, out, n_y_levels);
196        }
197        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
198            crate::portable::accumulate_contingency(x_codes, y_codes, out, n_y_levels);
199        }
200    }
201}
202
203/// Public semantic entry: contingency accumulation over a row subset.
204pub fn accumulate_contingency_rows(
205    policy: &KernelPolicy,
206    x_codes: &[u32],
207    y_codes: &[u32],
208    rows: &[usize],
209    out: &mut [f64],
210    n_y_levels: usize,
211) {
212    match portable_or_scalar_reductions(*policy) {
213        KernelImpl::Scalar => {
214            crate::scalar::accumulate_contingency_rows(x_codes, y_codes, rows, out, n_y_levels);
215        }
216        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
217            crate::portable::accumulate_contingency_rows(x_codes, y_codes, rows, out, n_y_levels);
218        }
219    }
220}
221
222/// Public semantic entry: weighted sum (bootstrap weight accumulation).
223///
224/// Contract: non-finite/negative weights → 0; `StableFloat`; beneficial for `n ≳ 64`.
225#[must_use]
226pub fn weighted_sum(policy: &KernelPolicy, x: &[f64], weights: &[f64]) -> f64 {
227    match portable_or_scalar_reductions(*policy) {
228        KernelImpl::Scalar => crate::scalar::weighted_sum(x, weights),
229        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
230            crate::portable::weighted_sum(x, weights)
231        }
232    }
233}
234
235/// Public semantic entry: weighted mean.
236#[must_use]
237pub fn weighted_mean(policy: &KernelPolicy, x: &[f64], weights: &[f64]) -> Option<f64> {
238    match portable_or_scalar_reductions(*policy) {
239        KernelImpl::Scalar => crate::scalar::weighted_mean(x, weights),
240        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
241            crate::portable::weighted_mean(x, weights)
242        }
243    }
244}
245
246/// Public semantic entry: weighted dot product.
247#[must_use]
248pub fn weighted_dot(policy: &KernelPolicy, x: &[f64], y: &[f64], weights: &[f64]) -> f64 {
249    match portable_or_scalar_reductions(*policy) {
250        KernelImpl::Scalar => crate::scalar::weighted_dot(x, y, weights),
251        KernelImpl::PortableOptimized | KernelImpl::ArchSimd => {
252            crate::portable::weighted_dot(x, y, weights)
253        }
254    }
255}