1use antecedent_core::KernelPolicy;
6
7use crate::view::{BitMaskView, F64VectorView};
8
9#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
11pub enum KernelImpl {
12 Scalar,
14 PortableOptimized,
16 ArchSimd,
18}
19
20#[must_use]
26pub fn arch_simd_available() -> bool {
27 false
28}
29
30#[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#[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#[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#[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
105pub 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
115pub 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#[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#[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#[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
171pub 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
183pub 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
203pub 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#[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#[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#[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}