use anyhow::Result;
#[derive(Debug, Clone)]
pub enum SIMDVector {
F32x4([f32; 4]),
F32x8([f32; 8]),
F32x16([f32; 16]),
I32x4([i32; 4]),
I32x8([i32; 8]),
U8x16([u8; 16]),
U8x32([u8; 32]),
}
impl SIMDVector {
pub fn splat_f32(value: f32, size: usize) -> Self {
match size {
4 => SIMDVector::F32x4([value; 4]),
8 => SIMDVector::F32x8([value; 8]),
16 => SIMDVector::F32x16([value; 16]),
_ => SIMDVector::F32x4([value; 4]),
}
}
pub fn len(&self) -> usize {
match self {
SIMDVector::F32x4(_) => 4,
SIMDVector::F32x8(_) => 8,
SIMDVector::F32x16(_) => 16,
SIMDVector::I32x4(_) => 4,
SIMDVector::I32x8(_) => 8,
SIMDVector::U8x16(_) => 16,
SIMDVector::U8x32(_) => 32,
}
}
}
pub trait SIMDOperations {
fn add(&self, other: &Self) -> Self;
fn sub(&self, other: &Self) -> Self;
fn mul(&self, other: &Self) -> Self;
fn div(&self, other: &Self) -> Self;
fn max(&self, other: &Self) -> Self;
fn min(&self, other: &Self) -> Self;
fn sqrt(&self) -> Self;
fn abs(&self) -> Self;
fn sum(&self) -> f32;
fn dot_product(&self, other: &Self) -> f32;
}
impl SIMDOperations for SIMDVector {
fn add(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4(simd_add_f32x4(*a, *b))
}
(SIMDVector::F32x8(a), SIMDVector::F32x8(b)) => {
let lo = simd_add_f32x4([a[0], a[1], a[2], a[3]], [b[0], b[1], b[2], b[3]]);
let hi = simd_add_f32x4([a[4], a[5], a[6], a[7]], [b[4], b[5], b[6], b[7]]);
SIMDVector::F32x8([lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]])
}
(SIMDVector::I32x4(a), SIMDVector::I32x4(b)) => {
SIMDVector::I32x4([a[0] + b[0], a[1] + b[1], a[2] + b[2], a[3] + b[3]])
}
(SIMDVector::I32x8(a), SIMDVector::I32x8(b)) => SIMDVector::I32x8([
a[0] + b[0],
a[1] + b[1],
a[2] + b[2],
a[3] + b[3],
a[4] + b[4],
a[5] + b[5],
a[6] + b[6],
a[7] + b[7],
]),
(SIMDVector::U8x16(a), SIMDVector::U8x16(b)) => {
let mut r = [0u8; 16];
for i in 0..16 {
r[i] = a[i].wrapping_add(b[i]);
}
SIMDVector::U8x16(r)
}
(SIMDVector::U8x32(a), SIMDVector::U8x32(b)) => {
let mut r = [0u8; 32];
for i in 0..32 {
r[i] = a[i].wrapping_add(b[i]);
}
SIMDVector::U8x32(r)
}
_ => unreachable!("type mismatch"),
}
}
fn sub(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4(simd_sub_f32x4(*a, *b))
}
(SIMDVector::I32x4(a), SIMDVector::I32x4(b)) => {
SIMDVector::I32x4([a[0] - b[0], a[1] - b[1], a[2] - b[2], a[3] - b[3]])
}
_ => unreachable!("type mismatch"),
}
}
fn mul(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4(simd_mul_f32x4(*a, *b))
}
(SIMDVector::F32x8(a), SIMDVector::F32x8(b)) => {
let lo = simd_mul_f32x4([a[0], a[1], a[2], a[3]], [b[0], b[1], b[2], b[3]]);
let hi = simd_mul_f32x4([a[4], a[5], a[6], a[7]], [b[4], b[5], b[6], b[7]]);
SIMDVector::F32x8([lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]])
}
_ => unreachable!("type mismatch"),
}
}
fn div(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4([a[0] / b[0], a[1] / b[1], a[2] / b[2], a[3] / b[3]])
}
_ => unreachable!("type mismatch"),
}
}
fn max(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4(simd_max_f32x4(*a, *b))
}
(SIMDVector::I32x4(a), SIMDVector::I32x4(b)) => SIMDVector::I32x4([
a[0].max(b[0]),
a[1].max(b[1]),
a[2].max(b[2]),
a[3].max(b[3]),
]),
_ => unreachable!("type mismatch"),
}
}
fn min(&self, other: &Self) -> Self {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
SIMDVector::F32x4(simd_min_f32x4(*a, *b))
}
_ => unreachable!("type mismatch"),
}
}
fn sqrt(&self) -> Self {
match self {
SIMDVector::F32x4(a) => SIMDVector::F32x4(simd_sqrt_f32x4(*a)),
_ => unreachable!("type mismatch"),
}
}
fn abs(&self) -> Self {
match self {
SIMDVector::F32x4(a) => {
SIMDVector::F32x4([a[0].abs(), a[1].abs(), a[2].abs(), a[3].abs()])
}
_ => unreachable!("type mismatch"),
}
}
fn sum(&self) -> f32 {
match self {
SIMDVector::F32x4(a) => a[0] + a[1] + a[2] + a[3],
SIMDVector::F32x8(a) => a[0] + a[1] + a[2] + a[3] + a[4] + a[5] + a[6] + a[7],
_ => unreachable!("type mismatch"),
}
}
fn dot_product(&self, other: &Self) -> f32 {
match (self, other) {
(SIMDVector::F32x4(a), SIMDVector::F32x4(b)) => {
let m = simd_mul_f32x4(*a, *b);
m[0] + m[1] + m[2] + m[3]
}
(SIMDVector::F32x8(a), SIMDVector::F32x8(b)) => {
let lo = simd_mul_f32x4([a[0], a[1], a[2], a[3]], [b[0], b[1], b[2], b[3]]);
let hi = simd_mul_f32x4([a[4], a[5], a[6], a[7]], [b[4], b[5], b[6], b[7]]);
lo[0] + lo[1] + lo[2] + lo[3] + hi[0] + hi[1] + hi[2] + hi[3]
}
_ => unreachable!("type mismatch"),
}
}
}
fn simd_add_f32x4(a: [f32; 4], b: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_add_ps;
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_storeu_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vb = _mm_loadu_ps(b.as_ptr());
let vr = _mm_add_ps(va, vb);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
#[cfg(all(feature = "simd", target_arch = "aarch64"))]
{
unsafe {
use std::arch::aarch64::vaddq_f32;
use std::arch::aarch64::vld1q_f32;
use std::arch::aarch64::vst1q_f32;
let va = vld1q_f32(a.as_ptr());
let vb = vld1q_f32(b.as_ptr());
let vr = vaddq_f32(va, vb);
let mut r = [0.0; 4];
vst1q_f32(r.as_mut_ptr(), vr);
return r;
}
}
[a[0] + b[0], a[1] + b[1], a[2] + b[2], a[3] + b[3]]
}
fn simd_sub_f32x4(a: [f32; 4], b: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_storeu_ps;
use std::arch::x86_64::_mm_sub_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vb = _mm_loadu_ps(b.as_ptr());
let vr = _mm_sub_ps(va, vb);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
#[cfg(all(feature = "simd", target_arch = "aarch64"))]
{
unsafe {
use std::arch::aarch64::vld1q_f32;
use std::arch::aarch64::vst1q_f32;
use std::arch::aarch64::vsubq_f32;
let va = vld1q_f32(a.as_ptr());
let vb = vld1q_f32(b.as_ptr());
let vr = vsubq_f32(va, vb);
let mut r = [0.0; 4];
vst1q_f32(r.as_mut_ptr(), vr);
return r;
}
}
[a[0] - b[0], a[1] - b[1], a[2] - b[2], a[3] - b[3]]
}
fn simd_mul_f32x4(a: [f32; 4], b: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_mul_ps;
use std::arch::x86_64::_mm_storeu_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vb = _mm_loadu_ps(b.as_ptr());
let vr = _mm_mul_ps(va, vb);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
#[cfg(all(feature = "simd", target_arch = "aarch64"))]
{
unsafe {
use std::arch::aarch64::vld1q_f32;
use std::arch::aarch64::vmulq_f32;
use std::arch::aarch64::vst1q_f32;
let va = vld1q_f32(a.as_ptr());
let vb = vld1q_f32(b.as_ptr());
let vr = vmulq_f32(va, vb);
let mut r = [0.0; 4];
vst1q_f32(r.as_mut_ptr(), vr);
return r;
}
}
[a[0] * b[0], a[1] * b[1], a[2] * b[2], a[3] * b[3]]
}
fn simd_max_f32x4(a: [f32; 4], b: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_max_ps;
use std::arch::x86_64::_mm_storeu_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vb = _mm_loadu_ps(b.as_ptr());
let vr = _mm_max_ps(va, vb);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
#[cfg(all(feature = "simd", target_arch = "aarch64"))]
{
unsafe {
use std::arch::aarch64::vld1q_f32;
use std::arch::aarch64::vmaxq_f32;
use std::arch::aarch64::vst1q_f32;
let va = vld1q_f32(a.as_ptr());
let vb = vld1q_f32(b.as_ptr());
let vr = vmaxq_f32(va, vb);
let mut r = [0.0; 4];
vst1q_f32(r.as_mut_ptr(), vr);
return r;
}
}
[
a[0].max(b[0]),
a[1].max(b[1]),
a[2].max(b[2]),
a[3].max(b[3]),
]
}
fn simd_min_f32x4(a: [f32; 4], b: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_min_ps;
use std::arch::x86_64::_mm_storeu_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vb = _mm_loadu_ps(b.as_ptr());
let vr = _mm_min_ps(va, vb);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
#[cfg(all(feature = "simd", target_arch = "aarch64"))]
{
unsafe {
use std::arch::aarch64::vld1q_f32;
use std::arch::aarch64::vminq_f32;
use std::arch::aarch64::vst1q_f32;
let va = vld1q_f32(a.as_ptr());
let vb = vld1q_f32(b.as_ptr());
let vr = vminq_f32(va, vb);
let mut r = [0.0; 4];
vst1q_f32(r.as_mut_ptr(), vr);
return r;
}
}
[
a[0].min(b[0]),
a[1].min(b[1]),
a[2].min(b[2]),
a[3].min(b[3]),
]
}
fn simd_sqrt_f32x4(a: [f32; 4]) -> [f32; 4] {
#[cfg(all(feature = "simd", target_arch = "x86_64"))]
{
if std::is_x86_feature_detected!("sse4.1") {
unsafe {
use std::arch::x86_64::_mm_loadu_ps;
use std::arch::x86_64::_mm_sqrt_ps;
use std::arch::x86_64::_mm_storeu_ps;
let va = _mm_loadu_ps(a.as_ptr());
let vr = _mm_sqrt_ps(va);
let mut r = [0.0; 4];
_mm_storeu_ps(r.as_mut_ptr(), vr);
return r;
}
}
}
[a[0].sqrt(), a[1].sqrt(), a[2].sqrt(), a[3].sqrt()]
}
pub struct SIMDMatrix {
data: Vec<f32>,
rows: usize,
cols: usize,
}
impl SIMDMatrix {
pub fn new(rows: usize, cols: usize) -> Self {
Self {
data: vec![0.0; rows * cols],
rows,
cols,
}
}
pub fn from_vec(data: Vec<f32>, rows: usize, cols: usize) -> Self {
assert_eq!(data.len(), rows * cols);
Self { data, rows, cols }
}
pub fn rows(&self) -> usize {
self.rows
}
pub fn cols(&self) -> usize {
self.cols
}
pub fn matmul_simd(&self, other: &SIMDMatrix) -> Result<SIMDMatrix> {
if self.cols != other.rows {
return Err(anyhow::anyhow!(
"Matrix dimensions don't match for multiplication"
));
}
let mut result = SIMDMatrix::new(self.rows, other.cols);
for i in 0..self.rows {
for j in 0..other.cols {
let mut sum = 0.0;
let mut k = 0;
while k + 4 <= self.cols {
let a_simd = self.get_row_simd(i, k, 4);
let b_simd = other.get_col_simd(k, j, 4);
sum += a_simd.dot_product(&b_simd);
k += 4;
}
while k < self.cols {
sum += self.get(i, k) * other.get(k, j);
k += 1;
}
result.set(i, j, sum);
}
}
Ok(result)
}
pub fn matmul_blocked(&self, other: &SIMDMatrix, block_size: usize) -> Result<SIMDMatrix> {
if self.cols != other.rows {
return Err(anyhow::anyhow!(
"Matrix dimensions don't match for multiplication"
));
}
let mut result = SIMDMatrix::new(self.rows, other.cols);
for i in (0..self.rows).step_by(block_size) {
let i_end = (i + block_size).min(self.rows);
for k in (0..self.cols).step_by(block_size) {
let k_end = (k + block_size).min(self.cols);
for j in (0..other.cols).step_by(block_size) {
let j_end = (j + block_size).min(other.cols);
for ii in i..i_end {
for kk in k..k_end {
let aik = self.get(ii, kk);
for jj in j..j_end {
let old = result.get(ii, jj);
result.set(ii, jj, old + aik * other.get(kk, jj));
}
}
}
}
}
}
Ok(result)
}
pub fn conv2d_simd(
&self,
kernel: &SIMDMatrix,
stride: (usize, usize),
padding: (usize, usize),
) -> Result<SIMDMatrix> {
let output_rows = (self.rows + 2 * padding.0 - kernel.rows) / stride.0 + 1;
let output_cols = (self.cols + 2 * padding.1 - kernel.cols) / stride.1 + 1;
let mut result = SIMDMatrix::new(output_rows, output_cols);
for i in 0..output_rows {
for j in 0..output_cols {
let mut sum = 0.0;
for ki in 0..kernel.rows {
for kj in 0..kernel.cols {
let input_i = i * stride.0 + ki;
let input_j = j * stride.1 + kj;
if input_i >= padding.0
&& input_i < self.rows + padding.0
&& input_j >= padding.1
&& input_j < self.cols + padding.1
{
let actual_i = input_i - padding.0;
let actual_j = input_j - padding.1;
sum += self.get(actual_i, actual_j) * kernel.get(ki, kj);
}
}
}
result.set(i, j, sum);
}
}
Ok(result)
}
pub fn batch_norm_simd(
&mut self,
mean: &[f32],
variance: &[f32],
gamma: &[f32],
beta: &[f32],
epsilon: f32,
) -> Result<()> {
for i in 0..self.rows {
for j in 0..self.cols {
let idx = j % mean.len();
let normalized = (self.get(i, j) - mean[idx]) / (variance[idx] + epsilon).sqrt();
let output = gamma[idx] * normalized + beta[idx];
self.set(i, j, output);
}
}
Ok(())
}
pub fn relu_simd(&mut self) {
for i in 0..self.rows {
for j in 0..self.cols {
let val = self.get(i, j);
self.set(i, j, if val > 0.0 { val } else { 0.0 });
}
}
}
pub fn softmax_simd(&mut self) -> Result<()> {
for i in 0..self.rows {
let mut max_val = f32::NEG_INFINITY;
for j in 0..self.cols {
max_val = max_val.max(self.get(i, j));
}
let mut sum = 0.0;
for j in 0..self.cols {
let exp_val = (self.get(i, j) - max_val).exp();
self.set(i, j, exp_val);
sum += exp_val;
}
for j in 0..self.cols {
self.set(i, j, self.get(i, j) / sum);
}
}
Ok(())
}
fn get(&self, row: usize, col: usize) -> f32 {
self.data[row * self.cols + col]
}
fn set(&mut self, row: usize, col: usize, value: f32) {
self.data[row * self.cols + col] = value;
}
fn get_row_simd(&self, row: usize, start_col: usize, count: usize) -> SIMDVector {
let mut data = [0.0; 4];
for i in 0..count.min(4) {
data[i] = self.get(row, start_col + i);
}
SIMDVector::F32x4(data)
}
fn get_col_simd(&self, start_row: usize, col: usize, count: usize) -> SIMDVector {
let mut data = [0.0; 4];
for i in 0..count.min(4) {
data[i] = self.get(start_row + i, col);
}
SIMDVector::F32x4(data)
}
}
pub struct SIMDImageProcessor;
impl SIMDImageProcessor {
pub fn gaussian_blur_simd(
&self,
image: &[f32],
width: usize,
height: usize,
kernel_size: usize,
sigma: f32,
) -> Result<Vec<f32>> {
let kernel = self.generate_gaussian_kernel(kernel_size, sigma)?;
let mut result = vec![0.0; image.len()];
let padding = kernel_size / 2;
for y in 0..height {
for x in 0..width {
let mut sum = 0.0;
let mut weight_sum = 0.0;
for ky in 0..kernel_size {
for kx in 0..kernel_size {
let img_y = y as i32 + ky as i32 - padding as i32;
let img_x = x as i32 + kx as i32 - padding as i32;
if img_y >= 0 && img_y < height as i32 && img_x >= 0 && img_x < width as i32
{
let pixel_idx = (img_y as usize * width + img_x as usize) as usize;
let kernel_val = kernel[ky * kernel_size + kx];
sum += image[pixel_idx] * kernel_val;
weight_sum += kernel_val;
}
}
}
if weight_sum > 0.0 {
result[y * width + x] = sum / weight_sum;
} else {
result[y * width + x] = image[y * width + x];
}
}
}
Ok(result)
}
pub fn sobel_edge_detection_simd(
&self,
image: &[f32],
width: usize,
height: usize,
) -> Result<Vec<f32>> {
let sobel_x = vec![-1.0, 0.0, 1.0, -2.0, 0.0, 2.0, -1.0, 0.0, 1.0];
let sobel_y = vec![-1.0, -2.0, -1.0, 0.0, 0.0, 0.0, 1.0, 2.0, 1.0];
let gx = self.apply_kernel_simd(image, width, height, &sobel_x, 3)?;
let gy = self.apply_kernel_simd(image, width, height, &sobel_y, 3)?;
let mut result = vec![0.0; image.len()];
for i in 0..image.len() {
result[i] = (gx[i] * gx[i] + gy[i] * gy[i]).sqrt();
}
Ok(result)
}
fn apply_kernel_simd(
&self,
image: &[f32],
width: usize,
height: usize,
kernel: &[f32],
kernel_size: usize,
) -> Result<Vec<f32>> {
let mut result = vec![0.0; image.len()];
let padding = kernel_size / 2;
for y in 0..height {
for x in 0..width {
let mut sum = 0.0;
for ky in 0..kernel_size {
for kx in 0..kernel_size {
let img_y = y as i32 + ky as i32 - padding as i32;
let img_x = x as i32 + kx as i32 - padding as i32;
if img_y >= 0 && img_y < height as i32 && img_x >= 0 && img_x < width as i32
{
let pixel_idx = img_y as usize * width + img_x as usize;
let kernel_idx = ky * kernel_size + kx;
sum += image[pixel_idx] * kernel[kernel_idx];
}
}
}
result[y * width + x] = sum;
}
}
Ok(result)
}
fn generate_gaussian_kernel(&self, size: usize, sigma: f32) -> Result<Vec<f32>> {
let mut kernel = vec![0.0; size * size];
let center = size as f32 / 2.0;
let mut sum = 0.0;
for y in 0..size {
for x in 0..size {
let dx = x as f32 - center;
let dy = y as f32 - center;
let distance_squared = dx * dx + dy * dy;
let value = (-distance_squared / (2.0 * sigma * sigma)).exp();
kernel[y * size + x] = value;
sum += value;
}
}
for val in &mut kernel {
*val /= sum;
}
Ok(kernel)
}
}
pub struct SIMDNeuralNetwork;
impl SIMDNeuralNetwork {
pub fn conv2d_forward_simd(
&self,
input: &[f32],
input_shape: (usize, usize, usize), weights: &[f32],
weight_shape: (usize, usize, usize, usize), bias: &[f32],
stride: (usize, usize),
padding: (usize, usize),
) -> Result<Vec<f32>> {
let (input_h, input_w, input_c) = input_shape;
let (out_c, in_c, kernel_h, kernel_w) = weight_shape;
let output_h = (input_h + 2 * padding.0 - kernel_h) / stride.0 + 1;
let output_w = (input_w + 2 * padding.1 - kernel_w) / stride.1 + 1;
let mut output = vec![0.0; out_c * output_h * output_w];
for oc in 0..out_c {
for oh in 0..output_h {
for ow in 0..output_w {
let mut sum = bias[oc];
for ic in 0..in_c {
for kh in 0..kernel_h {
for kw in 0..kernel_w {
let input_h_idx = oh * stride.0 + kh;
let input_w_idx = ow * stride.1 + kw;
if input_h_idx >= padding.0
&& input_h_idx < input_h + padding.0
&& input_w_idx >= padding.1
&& input_w_idx < input_w + padding.1
{
let actual_h = input_h_idx - padding.0;
let actual_w = input_w_idx - padding.1;
let input_idx =
actual_h * input_w * input_c + actual_w * input_c + ic;
let weight_idx = oc * in_c * kernel_h * kernel_w
+ ic * kernel_h * kernel_w
+ kh * kernel_w
+ kw;
sum += input[input_idx] * weights[weight_idx];
}
}
}
}
let output_idx = oc * output_h * output_w + oh * output_w + ow;
output[output_idx] = sum;
}
}
}
Ok(output)
}
pub fn max_pool2d_simd(
&self,
input: &[f32],
input_shape: (usize, usize, usize), kernel_size: (usize, usize),
stride: (usize, usize),
) -> Result<Vec<f32>> {
let (input_h, input_w, channels) = input_shape;
let (kernel_h, kernel_w) = kernel_size;
let output_h = (input_h - kernel_h) / stride.0 + 1;
let output_w = (input_w - kernel_w) / stride.1 + 1;
let mut output = vec![0.0; channels * output_h * output_w];
for c in 0..channels {
for oh in 0..output_h {
for ow in 0..output_w {
let mut max_val = f32::NEG_INFINITY;
for kh in 0..kernel_h {
for kw in 0..kernel_w {
let input_h_idx = oh * stride.0 + kh;
let input_w_idx = ow * stride.1 + kw;
let input_idx =
input_h_idx * input_w * channels + input_w_idx * channels + c;
max_val = max_val.max(input[input_idx]);
}
}
let output_idx = c * output_h * output_w + oh * output_w + ow;
output[output_idx] = max_val;
}
}
}
Ok(output)
}
pub fn batch_norm_simd(
&self,
input: &[f32],
gamma: &[f32],
beta: &[f32],
mean: &[f32],
variance: &[f32],
epsilon: f32,
) -> Result<Vec<f32>> {
let mut output = vec![0.0; input.len()];
let channels = gamma.len();
for (i, &val) in input.iter().enumerate() {
let c = i % channels;
let normalized = (val - mean[c]) / (variance[c] + epsilon).sqrt();
output[i] = gamma[c] * normalized + beta[c];
}
Ok(output)
}
}
pub struct SIMDUtils;
impl SIMDUtils {
pub fn is_simd_available() -> bool {
#[cfg(target_arch = "x86_64")]
{
std::is_x86_feature_detected!("sse4.1")
}
#[cfg(target_arch = "aarch64")]
{
true
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
{
false
}
}
pub fn get_optimal_vector_size() -> usize {
#[cfg(target_arch = "x86_64")]
{
if std::is_x86_feature_detected!("avx2") {
8
} else if std::is_x86_feature_detected!("sse4.1") {
4
} else {
1
}
}
#[cfg(target_arch = "aarch64")]
{
4
}
#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
{
1
}
}
pub fn align_data(data: &[f32], alignment: usize) -> Vec<f32> {
let mut aligned = data.to_vec();
while aligned.len() % alignment != 0 {
aligned.push(0.0);
}
aligned
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_simd_vector_add() {
let a = SIMDVector::F32x4([1.0, 2.0, 3.0, 4.0]);
let b = SIMDVector::F32x4([4.0, 3.0, 2.0, 1.0]);
let r = a.add(&b);
if let SIMDVector::F32x4(data) = r {
assert_eq!(data, [5.0, 5.0, 5.0, 5.0]);
} else {
panic!("wrong variant");
}
}
#[test]
fn test_simd_vector_mul() {
let a = SIMDVector::F32x4([2.0, 3.0, 4.0, 5.0]);
let b = SIMDVector::F32x4([3.0, 2.0, 1.0, 0.0]);
let r = a.mul(&b);
if let SIMDVector::F32x4(data) = r {
assert_eq!(data, [6.0, 6.0, 4.0, 0.0]);
} else {
panic!("wrong variant");
}
}
#[test]
fn test_simd_matrix_matmul() {
let a = SIMDMatrix::from_vec(vec![1.0, 2.0, 3.0, 4.0], 2, 2);
let b = SIMDMatrix::from_vec(vec![5.0, 6.0, 7.0, 8.0], 2, 2);
let r = a.matmul_simd(&b).unwrap();
assert_eq!(r.rows, 2);
assert_eq!(r.cols, 2);
assert!((r.get(0, 0) - 19.0).abs() < 1e-6);
assert!((r.get(0, 1) - 22.0).abs() < 1e-6);
assert!((r.get(1, 0) - 43.0).abs() < 1e-6);
assert!((r.get(1, 1) - 50.0).abs() < 1e-6);
}
#[test]
fn test_simd_utils() {
let avail = SIMDUtils::is_simd_available();
assert!(avail || cfg!(not(any(target_arch = "x86_64", target_arch = "aarch64"))));
}
}