rust_particle_sim_SIMD

use std::time::Instant;
use rand::prelude::*;

const NUM_PARTICLES: usize = 100000;
const SIMD_WIDTH: usize = 8; // AVX equivalent for 256-bit
const GRAVITY: f32 = -9.81;
const DELTA_TIME: f32 = 0.01;
const FIXED_SEED: u64 = 12345;

// Structure to hold particle data
struct ParticleData {
    positions_x: Vec<f32>,
    positions_y: Vec<f32>,
    positions_z: Vec<f32>,
    velocities_x: Vec<f32>,
    velocities_y: Vec<f32>,
    velocities_z: Vec<f32>,
    accelerations_x: Vec<f32>,
    accelerations_y: Vec<f32>,
    accelerations_z: Vec<f32>,
    masses: Vec<f32>,
}

fn allocate_particle_data() -> ParticleData {
    ParticleData {
        positions_x: vec![0.0; NUM_PARTICLES],
        positions_y: vec![0.0; NUM_PARTICLES],
        positions_z: vec![0.0; NUM_PARTICLES],
        velocities_x: vec![0.0; NUM_PARTICLES],
        velocities_y: vec![0.0; NUM_PARTICLES],
        velocities_z: vec![0.0; NUM_PARTICLES],
        accelerations_x: vec![0.0; NUM_PARTICLES],
        accelerations_y: vec![0.0; NUM_PARTICLES],
        accelerations_z: vec![0.0; NUM_PARTICLES],
        masses: vec![0.0; NUM_PARTICLES],
    }
}

fn initialize_particles(data: &mut ParticleData) {
    let mut rng = rand::rngs::StdRng::seed_from_u64(FIXED_SEED);

    for i in 0..NUM_PARTICLES {
        data.positions_x[i] = rng.gen::<f32>() * 10.0;
        data.positions_y[i] = rng.gen::<f32>() * 10.0;
        data.positions_z[i] = rng.gen::<f32>() * 10.0;

        data.velocities_x[i] = rng.gen::<f32>() * 2.0 - 1.0;
        data.velocities_y[i] = rng.gen::<f32>() * 2.0 - 1.0;
        data.velocities_z[i] = rng.gen::<f32>() * 2.0 - 1.0;

        data.accelerations_x[i] = 0.0;
        data.accelerations_y[i] = 0.0;
        data.accelerations_z[i] = 0.0;

        data.masses[i] = rng.gen::<f32>() * 2.0 + 1.0;
    }
}

//#[cfg(target_arch = "x86_64")]
fn update_particles_simd(data: &mut ParticleData) {
    use core::arch::x86_64::{
        __m256, _mm256_add_ps, _mm256_mul_ps, _mm256_set1_ps, _mm256_loadu_ps, _mm256_storeu_ps
    };

    let gravity_vector: __m256 = unsafe { _mm256_set1_ps(GRAVITY) };
    let delta_time_vector: __m256 = unsafe { _mm256_set1_ps(DELTA_TIME) };

    for i in (0..NUM_PARTICLES).step_by(SIMD_WIDTH) {
        unsafe {
            // Load data
            let px: __m256 = _mm256_loadu_ps(data.positions_x[i..].as_ptr());
            let py: __m256 = _mm256_loadu_ps(data.positions_y[i..].as_ptr());
            let pz: __m256 = _mm256_loadu_ps(data.positions_z[i..].as_ptr());

            let vx: __m256 = _mm256_loadu_ps(data.velocities_x[i..].as_ptr());
            let vy: __m256 = _mm256_loadu_ps(data.velocities_y[i..].as_ptr());
            let vz: __m256 = _mm256_loadu_ps(data.velocities_z[i..].as_ptr());

            let ax: __m256 = _mm256_loadu_ps(data.accelerations_x[i..].as_ptr());
            let ay: __m256 = _mm256_loadu_ps(data.accelerations_y[i..].as_ptr());
            let mut az: __m256 = _mm256_loadu_ps(data.accelerations_z[i..].as_ptr());

            // Update accelerations (gravity)
            az = _mm256_add_ps(az, gravity_vector);

            // Update velocities
            let nvx = _mm256_add_ps(vx, _mm256_mul_ps(ax, delta_time_vector));
            let nvy = _mm256_add_ps(vy, _mm256_mul_ps(ay, delta_time_vector));
            let nvz = _mm256_add_ps(vz, _mm256_mul_ps(az, delta_time_vector));

            // Update positions
            let npx = _mm256_add_ps(px, _mm256_mul_ps(nvx, delta_time_vector));
            let npy = _mm256_add_ps(py, _mm256_mul_ps(nvy, delta_time_vector));
            let npz = _mm256_add_ps(pz, _mm256_mul_ps(nvz, delta_time_vector));

            // Store results
            _mm256_storeu_ps(data.positions_x[i..].as_mut_ptr(), npx);
            _mm256_storeu_ps(data.positions_y[i..].as_mut_ptr(), npy);
            _mm256_storeu_ps(data.positions_z[i..].as_mut_ptr(), npz);
            _mm256_storeu_ps(data.velocities_x[i..].as_mut_ptr(), nvx);
            _mm256_storeu_ps(data.velocities_y[i..].as_mut_ptr(), nvy);
            _mm256_storeu_ps(data.velocities_z[i..].as_mut_ptr(), nvz);
            _mm256_storeu_ps(data.accelerations_x[i..].as_mut_ptr(), ax);
            _mm256_storeu_ps(data.accelerations_y[i..].as_mut_ptr(), ay);
            _mm256_storeu_ps(data.accelerations_z[i..].as_mut_ptr(), az);
        }
    }
}

fn main() {
    let mut particle_data = allocate_particle_data();
    initialize_particles(&mut particle_data);

    let num_steps: i32 = 100;

     // Time SIMD version
    println!("Starting SIMD test...");
    let start_simd = Instant::now();
    for _ in 0..num_steps {
        update_particles_simd(&mut particle_data);
    }
    let end_simd = Instant::now();
    let time_simd = (end_simd - start_simd).as_secs_f64() * 1000.0;
    println!("SIMD Time: {:.2} ms", time_simd);

    println!("First position Rust: x={:.4}, y={:.4}, z={:.4}",
        particle_data.positions_x[0],
        particle_data.positions_y[0],
        particle_data.positions_z[0]
    );
}