1use scirs2_core::ndarray::{s, Array1, Array2, ArrayView1, ArrayView2, Axis};
7use scirs2_core::numeric::{Float, FromPrimitive};
8use scirs2_core::parallel_ops::*;
9use scirs2_core::simd_ops::{AutoOptimizer, PlatformCapabilities, SimdUnifiedOps};
10use std::fmt::Debug;
11
12use crate::error::{ClusteringError, Result};
13
14#[derive(Debug, Clone)]
23pub struct SimdClusterConfig {
24 pub simd_threshold: usize,
26 pub enable_parallel: bool,
28 pub parallel_chunk_size: usize,
30 pub block_size: usize,
32 pub force_simd: bool,
34}
35
36impl Default for SimdClusterConfig {
37 fn default() -> Self {
38 Self {
39 simd_threshold: 32,
40 enable_parallel: true,
41 parallel_chunk_size: 512,
42 block_size: 128,
43 force_simd: false,
44 }
45 }
46}
47
48pub(super) fn should_use_simd(n_elements: usize, config: &SimdClusterConfig) -> bool {
50 let caps = PlatformCapabilities::detect();
51 let optimizer = AutoOptimizer::new();
52 config.force_simd
53 || (caps.simd_available
54 && (optimizer.should_use_simd(n_elements) || n_elements >= config.simd_threshold))
55}
56
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub enum SimdDistanceMetric {
64 Euclidean,
66 SquaredEuclidean,
68 Manhattan,
70}
71
72pub fn simd_euclidean_distance<F>(
80 a: ArrayView1<F>,
81 b: ArrayView1<F>,
82 config: Option<&SimdClusterConfig>,
83) -> Result<F>
84where
85 F: Float + FromPrimitive + Debug + SimdUnifiedOps,
86{
87 if a.len() != b.len() {
88 return Err(ClusteringError::InvalidInput(format!(
89 "Vectors must have the same length: got {} and {}",
90 a.len(),
91 b.len()
92 )));
93 }
94 if a.is_empty() {
95 return Ok(F::zero());
96 }
97 let default_config = SimdClusterConfig::default();
98 let cfg = config.unwrap_or(&default_config);
99
100 if should_use_simd(a.len(), cfg) {
101 let diff = F::simd_sub(&a, &b);
102 Ok(F::simd_norm(&diff.view()))
103 } else {
104 let mut sum = F::zero();
105 for i in 0..a.len() {
106 let d = a[i] - b[i];
107 sum = sum + d * d;
108 }
109 Ok(sum.sqrt())
110 }
111}
112
113pub fn simd_squared_euclidean_distance<F>(
122 a: ArrayView1<F>,
123 b: ArrayView1<F>,
124 config: Option<&SimdClusterConfig>,
125) -> Result<F>
126where
127 F: Float + FromPrimitive + Debug + SimdUnifiedOps,
128{
129 if a.len() != b.len() {
130 return Err(ClusteringError::InvalidInput(format!(
131 "Vectors must have the same length: got {} and {}",
132 a.len(),
133 b.len()
134 )));
135 }
136 if a.is_empty() {
137 return Ok(F::zero());
138 }
139 let default_config = SimdClusterConfig::default();
140 let cfg = config.unwrap_or(&default_config);
141
142 if should_use_simd(a.len(), cfg) {
143 let diff = F::simd_sub(&a, &b);
144 let sq = F::simd_mul(&diff.view(), &diff.view());
145 Ok(F::simd_sum(&sq.view()))
146 } else {
147 let mut sum = F::zero();
148 for i in 0..a.len() {
149 let d = a[i] - b[i];
150 sum = sum + d * d;
151 }
152 Ok(sum)
153 }
154}
155
156pub fn simd_manhattan_distance<F>(
162 a: ArrayView1<F>,
163 b: ArrayView1<F>,
164 config: Option<&SimdClusterConfig>,
165) -> Result<F>
166where
167 F: Float + FromPrimitive + Debug + SimdUnifiedOps,
168{
169 if a.len() != b.len() {
170 return Err(ClusteringError::InvalidInput(format!(
171 "Vectors must have the same length: got {} and {}",
172 a.len(),
173 b.len()
174 )));
175 }
176 if a.is_empty() {
177 return Ok(F::zero());
178 }
179 let default_config = SimdClusterConfig::default();
180 let cfg = config.unwrap_or(&default_config);
181
182 if should_use_simd(a.len(), cfg) {
183 let diff = F::simd_sub(&a, &b);
184 let abs_diff = F::simd_abs(&diff.view());
185 Ok(F::simd_sum(&abs_diff.view()))
186 } else {
187 let mut sum = F::zero();
188 for i in 0..a.len() {
189 sum = sum + (a[i] - b[i]).abs();
190 }
191 Ok(sum)
192 }
193}
194
195pub fn simd_distance<F>(
203 a: ArrayView1<F>,
204 b: ArrayView1<F>,
205 metric: SimdDistanceMetric,
206 config: Option<&SimdClusterConfig>,
207) -> Result<F>
208where
209 F: Float + FromPrimitive + Debug + SimdUnifiedOps,
210{
211 match metric {
212 SimdDistanceMetric::Euclidean => simd_euclidean_distance(a, b, config),
213 SimdDistanceMetric::SquaredEuclidean => simd_squared_euclidean_distance(a, b, config),
214 SimdDistanceMetric::Manhattan => simd_manhattan_distance(a, b, config),
215 }
216}
217
218pub fn simd_pairwise_distance_matrix<F>(
241 data: ArrayView2<F>,
242 metric: SimdDistanceMetric,
243 config: Option<&SimdClusterConfig>,
244) -> Result<Array2<F>>
245where
246 F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
247{
248 let n_samples = data.shape()[0];
249 if n_samples == 0 {
250 return Err(ClusteringError::InvalidInput(
251 "Data must have at least one sample".to_string(),
252 ));
253 }
254
255 let default_config = SimdClusterConfig::default();
256 let cfg = config.unwrap_or(&default_config);
257 let mut distances = Array2::zeros((n_samples, n_samples));
258
259 let use_parallel =
260 cfg.enable_parallel && is_parallel_enabled() && n_samples > cfg.parallel_chunk_size;
261
262 if use_parallel {
263 simd_pairwise_matrix_parallel(data, metric, cfg, &mut distances);
264 } else {
265 simd_pairwise_matrix_sequential(data, metric, cfg, &mut distances)?;
266 }
267
268 Ok(distances)
269}
270
271fn simd_pairwise_matrix_sequential<F>(
273 data: ArrayView2<F>,
274 metric: SimdDistanceMetric,
275 config: &SimdClusterConfig,
276 distances: &mut Array2<F>,
277) -> Result<()>
278where
279 F: Float + FromPrimitive + Debug + SimdUnifiedOps,
280{
281 let n_samples = data.shape()[0];
282 let block_size = config.block_size;
283
284 for block_i in (0..n_samples).step_by(block_size) {
285 let end_i = (block_i + block_size).min(n_samples);
286 for block_j in (block_i..n_samples).step_by(block_size) {
287 let end_j = (block_j + block_size).min(n_samples);
288
289 for i in block_i..end_i {
290 let start_j = if block_i == block_j { i + 1 } else { block_j };
291 for j in start_j..end_j {
292 let dist = simd_distance(data.row(i), data.row(j), metric, Some(config))?;
293 distances[[i, j]] = dist;
294 distances[[j, i]] = dist;
295 }
296 }
297 }
298 }
299 Ok(())
300}
301
302fn simd_pairwise_matrix_parallel<F>(
304 data: ArrayView2<F>,
305 metric: SimdDistanceMetric,
306 config: &SimdClusterConfig,
307 distances: &mut Array2<F>,
308) where
309 F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
310{
311 let n_samples = data.shape()[0];
312
313 let row_results: Vec<Vec<(usize, F)>> = (0..n_samples)
315 .into_par_iter()
316 .map(|i| {
317 let mut row_dists = Vec::with_capacity(n_samples - i - 1);
318 for j in (i + 1)..n_samples {
319 let dist = simd_distance(data.row(i), data.row(j), metric, Some(config))
320 .unwrap_or_else(|_| F::zero());
321 row_dists.push((j, dist));
322 }
323 row_dists
324 })
325 .collect();
326
327 for (i, row_dists) in row_results.into_iter().enumerate() {
329 for (j, dist) in row_dists {
330 distances[[i, j]] = dist;
331 distances[[j, i]] = dist;
332 }
333 }
334}
335
336pub fn simd_pairwise_condensed_distances<F>(
356 data: ArrayView2<F>,
357 metric: SimdDistanceMetric,
358 config: Option<&SimdClusterConfig>,
359) -> Result<Array1<F>>
360where
361 F: Float + FromPrimitive + Debug + Send + Sync + SimdUnifiedOps,
362{
363 let n_samples = data.shape()[0];
364 if n_samples == 0 {
365 return Err(ClusteringError::InvalidInput(
366 "Data must have at least one sample".to_string(),
367 ));
368 }
369
370 let n_distances = n_samples * (n_samples - 1) / 2;
371 let mut distances = Array1::zeros(n_distances);
372 let default_config = SimdClusterConfig::default();
373 let cfg = config.unwrap_or(&default_config);
374
375 let mut idx = 0;
376 for i in 0..n_samples {
377 for j in (i + 1)..n_samples {
378 let dist = simd_distance(data.row(i), data.row(j), metric, Some(cfg))?;
379 distances[idx] = dist;
380 idx += 1;
381 }
382 }
383
384 Ok(distances)
385}