#[derive(Debug, Clone, Copy, PartialEq)]
pub struct HardwareFeatures {
pub has_avx512: bool,
pub has_avx2: bool,
pub has_sse4_2: bool,
pub threads_support: bool,
pub optimal_vector_width: VectorWidth,
}
impl Default for HardwareFeatures {
fn default() -> Self {
Self {
has_avx512: false,
has_avx2: false,
has_sse4_2: false,
threads_support: true, optimal_vector_width: VectorWidth::Scalar(64),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum VectorWidth {
Avx512(u32),
Avx2(u32),
Sse4(u32),
Scalar(u32),
}
impl VectorWidth {
pub fn bits(&self) -> u32 {
match self {
VectorWidth::Avx512(b) => *b,
VectorWidth::Avx2(b) => *b,
VectorWidth::Sse4(b) => *b,
VectorWidth::Scalar(b) => *b,
}
}
pub fn f32_lanes(&self) -> usize {
(self.bits() / 32) as usize
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SimdCapability {
Avx512,
Avx2,
Sse4,
None,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VectorOperation {
Float32Add,
Float32Mul,
Float32Div,
Float32Sqrt,
Int32Add,
Float32x4Add,
Float32x8Add,
Float32x16Add,
}
static CPU_FEATURES: OnceLock<HardwareFeatures> = OnceLock::new();
pub fn detect_cpu_features() -> HardwareFeatures {
*CPU_FEATURES.get_or_init(|| {
#[cfg(target_arch = "x86_64")]
{
detect_x86_features()
}
#[cfg(target_arch = "aarch64")]
{
detect_arm_features()
}
#[cfg(all(not(target_arch = "x86_64"), not(target_arch = "aarch64")))]
{
HardwareFeatures::default()
}
})
}
#[cfg(target_arch = "x86_64")]
fn detect_x86_features() -> HardwareFeatures {
let has_avx512: _ = is_x86_feature_detected!("avx512f");
let has_avx2: _ = is_x86_feature_detected!("avx2");
let has_sse4_2: _ = is_x86_feature_detected!("sse4.2");
let optimal_vector_width: _ = if has_avx512 {
VectorWidth::Avx512(512)
} else if has_avx2 {
VectorWidth::Avx2(256)
} else if has_sse4_2 {
VectorWidth::Sse4(128)
} else {
VectorWidth::Scalar(64)
};
HardwareFeatures {
has_avx512,
has_avx2,
has_sse4_2,
threads_support: true,
optimal_vector_width,
}
}
#[cfg(target_arch = "aarch64")]
fn detect_arm_features() -> HardwareFeatures {
HardwareFeatures {
has_avx512: false,
has_avx2: false,
has_sse4_2: false,
threads_support: true,
optimal_vector_width: VectorWidth::Sse4(128), }
}
#[derive(Debug, Clone)]
pub struct SimdStats {
pub operations_count: u64,
pub vector_ops_count: u64,
pub speedup_estimate: f64,
pub simd_utilization: f64,
}
impl Default for SimdStats {
fn default() -> Self {
Self {
operations_count: 0,
vector_ops_count: 0,
speedup_estimate: 1.0,
simd_utilization: 0.0,
}
}
}
pub struct SimdEngine {
features: HardwareFeatures,
capability: SimdCapability,
initialized: bool,
operations_count: AtomicU64,
vector_ops_count: AtomicU64,
}
impl SimdEngine {
pub fn new() -> Self {
let features: _ = detect_cpu_features();
let capability: _ = Self::determine_capability(&features);
Self {
features,
capability,
initialized: true,
operations_count: AtomicU64::new(0),
vector_ops_count: AtomicU64::new(0),
}
}
fn determine_capability(features: &HardwareFeatures) -> SimdCapability {
if features.has_avx512 {
SimdCapability::Avx512
} else if features.has_avx2 {
SimdCapability::Avx2
} else if features.has_sse4_2 {
SimdCapability::Sse4
} else {
SimdCapability::None
}
}
pub fn is_initialized(&self) -> bool {
self.initialized
}
pub fn get_features(&self) -> HardwareFeatures {
self.features
}
pub fn get_capability(&self) -> SimdCapability {
self.capability
}
pub fn get_optimal_vector_width(&self) -> VectorWidth {
self.features.optimal_vector_width
}
pub fn supports_operation(&self, op: VectorOperation) -> bool {
match op {
VectorOperation::Float32Add
| VectorOperation::Float32Mul
| VectorOperation::Float32Div
| VectorOperation::Float32Sqrt
| VectorOperation::Int32Add => true, VectorOperation::Float32x4Add => self.capability != SimdCapability::None,
VectorOperation::Float32x8Add => {
matches!(self.capability, SimdCapability::Avx2 | SimdCapability::Avx512)
}
VectorOperation::Float32x16Add => self.capability == SimdCapability::Avx512,
}
}
pub fn estimate_speedup_for_operation(&self, _op: VectorOperation) -> f64 {
match self.capability {
SimdCapability::Avx512 => 16.0,
SimdCapability::Avx2 => 8.0,
SimdCapability::Sse4 => 4.0,
SimdCapability::None => 1.0,
}
}
pub fn get_stats(&self) -> SimdStats {
let ops: _ = self.operations_count.load(Ordering::Relaxed);
let vec_ops: _ = self.vector_ops_count.load(Ordering::Relaxed);
let simd_utilization: _ = if ops > 0 {
vec_ops as f64 / ops as f64
} else {
0.0
};
SimdStats {
operations_count: ops,
vector_ops_count: vec_ops,
speedup_estimate: self.estimate_speedup_for_operation(VectorOperation::Float32Add),
simd_utilization,
}
}
pub fn reset_stats(&self) {
self.operations_count.store(0, Ordering::Relaxed);
self.vector_ops_count.store(0, Ordering::Relaxed);
}
fn increment_ops(&self, is_vector: bool) {
self.operations_count.fetch_add(1, Ordering::Relaxed);
if is_vector {
self.vector_ops_count.fetch_add(1, Ordering::Relaxed);
}
}
pub fn vector_add_f32(&self, a: &[f32], b: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if a.is_empty() || b.is_empty() {
return vec![];
}
let len: _ = a.len().min(b.len());
let mut result = Vec::with_capacity(len);
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
self.vector_add_f32_avx2(a, b, &mut result);
return result;
}
}
result.extend(a.iter().zip(b.iter()).map(|(&x, &y)| x + y));
result
}
#[cfg(target_arch = "x86_64")]
fn vector_add_f32_avx2(&self, a: &[f32], b: &[f32], result: &mut Vec<f32>) {
let len: _ = a.len().min(b.len());
result.reserve(len);
let chunks: _ = len / 8;
let remainder: _ = len % 8;
unsafe {
for i in 0..chunks {
let offset: _ = i * 8;
let va: _ = _mm256_loadu_ps(a.as_ptr().add(offset));
let vb: _ = _mm256_loadu_ps(b.as_ptr().add(offset));
let vr: _ = _mm256_add_ps(va, vb);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vr);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 8;
for i in 0..remainder {
result.push(a[base + i] + b[base + i]);
}
}
pub fn vector_mul_f32(&self, a: &[f32], b: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if a.is_empty() || b.is_empty() {
return vec![];
}
let len: _ = a.len().min(b.len());
let mut result = Vec::with_capacity(len);
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
self.vector_mul_f32_avx2(a, b, &mut result);
return result;
}
}
result.extend(a.iter().zip(b.iter()).map(|(&x, &y)| x * y));
result
}
#[cfg(target_arch = "x86_64")]
fn vector_mul_f32_avx2(&self, a: &[f32], b: &[f32], result: &mut Vec<f32>) {
let len: _ = a.len().min(b.len());
result.reserve(len);
let chunks: _ = len / 8;
let remainder: _ = len % 8;
unsafe {
for i in 0..chunks {
let offset: _ = i * 8;
let va: _ = _mm256_loadu_ps(a.as_ptr().add(offset));
let vb: _ = _mm256_loadu_ps(b.as_ptr().add(offset));
let vr: _ = _mm256_mul_ps(va, vb);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vr);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 8;
for i in 0..remainder {
result.push(a[base + i] * b[base + i]);
}
}
pub fn dot_product_f32(&self, a: &[f32], b: &[f32]) -> f32 {
self.increment_ops(true);
if a.is_empty() || b.is_empty() {
return 0.0;
}
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
return self.dot_product_f32_avx2(a, b);
}
}
a.iter().zip(b.iter()).map(|(&x, &y)| x * y).sum()
}
#[cfg(target_arch = "x86_64")]
fn dot_product_f32_avx2(&self, a: &[f32], b: &[f32]) -> f32 {
let len: _ = a.len().min(b.len());
let chunks: _ = len / 8;
let remainder: _ = len % 8;
let mut sum = 0.0f32;
unsafe {
let mut vsum = _mm256_setzero_ps();
for i in 0..chunks {
let offset: _ = i * 8;
let va: _ = _mm256_loadu_ps(a.as_ptr().add(offset));
let vb: _ = _mm256_loadu_ps(b.as_ptr().add(offset));
let vr: _ = _mm256_mul_ps(va, vb);
vsum = _mm256_add_ps(vsum, vr);
}
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vsum);
sum = tmp.iter().sum();
}
let base: _ = chunks * 8;
for i in 0..remainder {
sum += a[base + i] * b[base + i];
}
sum
}
pub fn vector_sum_f32(&self, data: &[f32]) -> f32 {
self.increment_ops(true);
if data.is_empty() {
return 0.0;
}
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
return self.vector_sum_f32_avx2(data);
}
}
data.iter().sum()
}
#[cfg(target_arch = "x86_64")]
fn vector_sum_f32_avx2(&self, data: &[f32]) -> f32 {
let len: _ = data.len();
let chunks: _ = len / 8;
let remainder: _ = len % 8;
let mut sum = 0.0f32;
unsafe {
let mut vsum = _mm256_setzero_ps();
for i in 0..chunks {
let offset: _ = i * 8;
let v: _ = _mm256_loadu_ps(data.as_ptr().add(offset));
vsum = _mm256_add_ps(vsum, v);
}
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vsum);
sum = tmp.iter().sum();
}
let base: _ = chunks * 8;
for i in 0..remainder {
sum += data[base + i];
}
sum
}
pub fn vector_sqrt_f32(&self, data: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if data.is_empty() {
return vec![];
}
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
return self.vector_sqrt_f32_avx2(data);
}
}
data.iter().map(|&x| x.sqrt()).collect()
}
#[cfg(target_arch = "x86_64")]
fn vector_sqrt_f32_avx2(&self, data: &[f32]) -> Vec<f32> {
let len: _ = data.len();
let mut result = Vec::with_capacity(len);
let chunks: _ = len / 8;
let remainder: _ = len % 8;
unsafe {
for i in 0..chunks {
let offset: _ = i * 8;
let v: _ = _mm256_loadu_ps(data.as_ptr().add(offset));
let vr: _ = _mm256_sqrt_ps(v);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vr);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 8;
for i in 0..remainder {
result.push(data[base + i].sqrt());
}
result
}
pub fn vector_max_f32(&self, data: &[f32]) -> f32 {
self.increment_ops(true);
if data.is_empty() {
return f32::NEG_INFINITY;
}
data.iter().cloned().fold(f32::NEG_INFINITY, f32::max)
}
pub fn vector_min_f32(&self, data: &[f32]) -> f32 {
self.increment_ops(true);
if data.is_empty() {
return f32::INFINITY;
}
data.iter().cloned().fold(f32::INFINITY, f32::min)
}
pub fn fused_multiply_add_f32(&self, a: &[f32], b: &[f32], c: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if a.is_empty() || b.is_empty() || c.is_empty() {
return vec![];
}
let len: _ = a.len().min(b.len()).min(c.len());
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 && is_x86_feature_detected!("fma") {
return self.fma_f32_avx2(a, b, c);
}
}
(0..len).map(|i| a[i] * b[i] + c[i]).collect()
}
#[cfg(target_arch = "x86_64")]
fn fma_f32_avx2(&self, a: &[f32], b: &[f32], c: &[f32]) -> Vec<f32> {
let len: _ = a.len().min(b.len()).min(c.len());
let mut result = Vec::with_capacity(len);
let chunks: _ = len / 8;
let remainder: _ = len % 8;
unsafe {
for i in 0..chunks {
let offset: _ = i * 8;
let va: _ = _mm256_loadu_ps(a.as_ptr().add(offset));
let vb: _ = _mm256_loadu_ps(b.as_ptr().add(offset));
let vc: _ = _mm256_loadu_ps(c.as_ptr().add(offset));
let vr: _ = _mm256_fmadd_ps(va, vb, vc);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), vr);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 8;
for i in 0..remainder {
result.push(a[base + i] * b[base + i] + c[base + i]);
}
result
}
pub fn auto_vectorize(&self, code: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if code.is_empty() {
return vec![];
}
let lane_width: _ = self.features.optimal_vector_width.f32_lanes();
if code.len() < lane_width {
return code.to_vec();
}
if code.len() < 1024 {
return match self.capability {
SimdCapability::Avx512 => self.vector_add_f32_scalar(code, &vec![0.0; code.len()]),
SimdCapability::Avx2 => self.vector_add_f32_scalar(code, &vec![0.0; code.len()]),
SimdCapability::Sse4 => self.vector_add_f32_scalar(code, &vec![0.0; code.len()]),
SimdCapability::None => code.to_vec(),
};
}
self.batch_process_f32(code)
}
pub fn auto_vectorize_loop(&self, iterations: usize, init_val: f32, step: f32) -> Vec<f32> {
self.increment_ops(true);
let mut result = Vec::with_capacity(iterations);
let lane_width: _ = self.features.optimal_vector_width.f32_lanes();
#[cfg(target_arch = "x86_64")]
{
if self.capability != SimdCapability::None && iterations >= lane_width * 4 {
let vector_iters: _ = iterations / lane_width;
let remainder: _ = iterations % lane_width;
unsafe {
let step_v: _ = _mm256_set1_ps(step);
for vector_idx in 0..vector_iters {
let base_val: _ = init_val + vector_idx as f32 * lane_width as f32 * step;
let v: _ = _mm256_set1_ps(base_val);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), v);
for &val in &tmp {
result.push(val + step); }
}
}
for i in 0..remainder {
result.push(init_val + (vector_iters * lane_width + i) as f32 * step);
}
return result;
}
}
for i in 0..iterations {
result.push(init_val + i as f32 * step);
}
result
}
pub fn optimize_data_layout(&self, data: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if data.len() < 64 {
return data.to_vec(); }
let lane_width: _ = self.features.optimal_vector_width.f32_lanes();
let num_vectors: _ = data.len() / lane_width;
let remainder: _ = data.len() % lane_width;
let mut optimized = Vec::with_capacity(data.len());
for chunk_idx in 0..num_vectors {
let chunk_start: _ = chunk_idx * lane_width;
let chunk_end: _ = chunk_start + lane_width;
#[cfg(target_arch = "x86_64")]
{
if self.features.has_avx2 {
unsafe {
let v: _ = _mm256_loadu_ps(data.as_ptr().add(chunk_start));
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), v);
optimized.extend_from_slice(&tmp);
}
} else {
optimized.extend_from_slice(&data[chunk_start..chunk_end]);
}
}
#[cfg(not(target_arch = "x86_64"))]
{
optimized.extend_from_slice(&data[chunk_start..chunk_end]);
}
}
if remainder > 0 {
let start: _ = num_vectors * lane_width;
optimized.extend_from_slice(&data[start..]);
}
optimized
}
pub fn batch_vector_add(&self, batch_a: &[Vec<f32>], batch_b: &[Vec<f32>]) -> Vec<Vec<f32>> {
self.increment_ops(true);
assert_eq!(batch_a.len(), batch_b.len(), "批次大小必须相同");
let batch_size: _ = batch_a.len();
let mut results = Vec::with_capacity(batch_size);
#[cfg(feature = "parallel")]
{
use rayon::prelude::*;
use std::sync::atomic::AtomicU64;
results.par_iter()
.zip(batch_a.par_iter())
.zip(batch_b.par_iter())
.map(|((_, a), b)| self.vector_add_f32(a, b))
.collect()
}
#[cfg(not(feature = "parallel"))]
{
for (a, b) in batch_a.iter().zip(batch_b.iter()) {
results.push(self.vector_add_f32(a, b));
}
results
}
}
pub fn batch_matrix_multiply(&self, matrices: &[(&[f32], &[f32])]) -> Vec<f32> {
self.increment_ops(true);
let mut results = Vec::new();
for (a, b) in matrices {
let len: _ = a.len().min(b.len());
let mut product = Vec::with_capacity(len);
for i in 0..len {
product.push(a[i] * b[i]);
}
results.extend(product);
}
results
}
pub fn batch_reduce(&self, data_batch: &[Vec<f32>]) -> Vec<f32> {
self.increment_ops(true);
let mut results = Vec::with_capacity(data_batch.len());
for data in data_batch {
results.push(self.vector_sum_f32(data));
}
results
}
pub fn batch_process_f32(&self, data: &[f32]) -> Vec<f32> {
self.increment_ops(true);
if data.is_empty() {
return vec![];
}
let cache_line_size: _ = 64;
let vector_width_bytes: _ = self.features.optimal_vector_width.bits() / 8;
let optimal_chunk_size: _ = (cache_line_size / 4).max(vector_width_bytes as usize / 4);
let chunk_size: _ = optimal_chunk_size.max(1024); let num_chunks: _ = (data.len() + chunk_size - 1) / chunk_size;
let mut results = Vec::with_capacity(data.len());
for chunk_idx in 0..num_chunks {
let start: _ = chunk_idx * chunk_size;
let end: _ = (start + chunk_size).min(data.len());
let chunk: _ = &data[start..end];
let processed: _ = self.process_chunk_f32(chunk);
results.extend(processed);
}
results
}
fn process_chunk_f32(&self, chunk: &[f32]) -> Vec<f32> {
match self.capability {
SimdCapability::Avx512 => self.simd_process_chunk_avx512(chunk),
SimdCapability::Avx2 => self.simd_process_chunk_avx2(chunk),
SimdCapability::Sse4 => self.simd_process_chunk_sse4(chunk),
SimdCapability::None => chunk.to_vec(),
}
}
#[cfg(target_arch = "x86_64")]
fn simd_process_chunk_avx512(&self, chunk: &[f32]) -> Vec<f32> {
let len: _ = chunk.len();
let mut result = Vec::with_capacity(len);
let chunks: _ = len / 16;
let remainder: _ = len % 16;
unsafe {
for i in 0..chunks {
let offset: _ = i * 16;
let v: _ = _mm512_loadu_ps(chunk.as_ptr().add(offset));
let processed: _ = _mm512_add_ps(v, v); let mut tmp = [0.0f32; 16];
_mm512_storeu_ps(tmp.as_mut_ptr(), processed);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 16;
for i in 0..remainder {
result.push(chunk[base + i] * 2.0);
}
result
}
#[cfg(target_arch = "x86_64")]
fn simd_process_chunk_avx2(&self, chunk: &[f32]) -> Vec<f32> {
let len: _ = chunk.len();
let mut result = Vec::with_capacity(len);
let chunks: _ = len / 8;
let remainder: _ = len % 8;
unsafe {
for i in 0..chunks {
let offset: _ = i * 8;
let v: _ = _mm256_loadu_ps(chunk.as_ptr().add(offset));
let processed: _ = _mm256_add_ps(v, v);
let mut tmp = [0.0f32; 8];
_mm256_storeu_ps(tmp.as_mut_ptr(), processed);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 8;
for i in 0..remainder {
result.push(chunk[base + i] * 2.0);
}
result
}
#[cfg(target_arch = "x86_64")]
fn simd_process_chunk_sse4(&self, chunk: &[f32]) -> Vec<f32> {
let len: _ = chunk.len();
let mut result = Vec::with_capacity(len);
let chunks: _ = len / 4;
let remainder: _ = len % 4;
unsafe {
for i in 0..chunks {
let offset: _ = i * 4;
let v: _ = _mm_loadu_ps(chunk.as_ptr().add(offset));
let processed: _ = _mm_add_ps(v, v);
let mut tmp = [0.0f32; 4];
_mm_storeu_ps(tmp.as_mut_ptr(), processed);
result.extend_from_slice(&tmp);
}
}
let base: _ = chunks * 4;
for i in 0..remainder {
result.push(chunk[base + i] * 2.0);
}
result
}
#[cfg(not(target_arch = "x86_64"))]
fn simd_process_chunk_avx512(&self, chunk: &[f32]) -> Vec<f32> {
chunk.iter().map(|&x| x * 2.0).collect()
}
#[cfg(not(target_arch = "x86_64"))]
fn simd_process_chunk_avx2(&self, chunk: &[f32]) -> Vec<f32> {
chunk.iter().map(|&x| x * 2.0).collect()
}
#[cfg(not(target_arch = "x86_64"))]
fn simd_process_chunk_sse4(&self, chunk: &[f32]) -> Vec<f32> {
chunk.iter().map(|&x| x * 2.0).collect()
}
fn vector_add_f32_scalar(&self, a: &[f32], b: &[f32]) -> Vec<f32> {
a.iter().zip(b.iter()).map(|(&x, &y)| x + y).collect()
}
}
impl Default for SimdEngine {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::{HashMap, BTreeMap};
use std::sync::atomic::Ordering;
#[test]
fn test_feature_detection() {
let features: _ = detect_cpu_features();
println!("CPU Features: {:?}", features);
}
#[test]
fn test_engine_creation() {
let engine: _ = SimdEngine::new();
assert!(engine.is_initialized());
}
#[test]
fn test_basic_vector_add() {
let engine: _ = SimdEngine::new();
let a: _ = vec![1.0, 2.0, 3.0, 4.0];
let b: _ = vec![4.0, 3.0, 2.0, 1.0];
let result: _ = engine.vector_add_f32(&a, &b);
assert_eq!(result, vec![5.0, 5.0, 5.0, 5.0]);
}
}