SIMD y optimización low-level en Rust: target_feature, std::simd y técnicas de bajo nivel

Cuando los benchmarks muestran que el cuello de botella es el procesamiento puro de datos numéricos, Rust permite descender hasta el nivel de las instrucciones SIMD del procesador. Con std::simd en nightly se obtiene una API portable entre arquitecturas; con std::arch se accede a intrinsics AVX2/SSE4.2 específicos de x86. Ambas rutas pueden coexistir en el mismo codebase.

target_feature y detección en tiempo de compilación

El atributo #[target_feature(enable = "avx2")] compila una función con las instrucciones SIMD habilitadas aunque el target global no las use. La función debe marcarse como unsafe porque invocarla en hardware que no soporte AVX2 produce una instrucción ilegal.

// Comprobar features disponibles en el host de compilación:
// RUSTFLAGS="-C target-cpu=native" cargo build --release

// Antipatrón: habilitar target_feature en todo el binario
// [profile.release]
// rustflags = ["-C", "target-feature=+avx2"]
// Problema: el binario no arrancará en CPUs sin AVX2

// CORRECTO: habilitar solo para funciones concretas
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn suma_avx2(datos: &[f32]) -> f32 {
    // Esta función solo puede llamarse si is_x86_feature_detected!("avx2") == true
    datos.iter().sum()
}

fn suma_segura(datos: &[f32]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    if is_x86_feature_detected!("avx2") {
        return unsafe { suma_avx2(datos) };
    }
    datos.iter().sum() // fallback escalar
}

fn main() {
    let datos: Vec<f32> = (1..=1000).map(|x| x as f32).collect();
    println!("Suma: {}", suma_segura(&datos)); // 500500
}

std::arch: intrinsics AVX2 explícitos

std::arch::x86_64 expone las funciones intrínsecas de Intel directamente. Son los bloques de construcción más bajo nivel disponibles en Rust: cada función corresponde a una instrucción de ensamblador.

#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

/// Suma vectorizada de elementos f32 usando AVX2 (8 elementos por ciclo)
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn suma_f32_avx2(datos: &[f32]) -> f32 {
    let mut acum = _mm256_setzero_ps(); // registro de 8 f32 a cero
    let chunks = datos.chunks_exact(8);
    let resto = chunks.remainder();

    for chunk in chunks {
        let v = _mm256_loadu_ps(chunk.as_ptr()); // cargar 8 f32 del slice
        acum = _mm256_add_ps(acum, v);           // sumar 8 en paralelo
    }

    // Reducir el registro de 8 canales a un escalar
    let mut result = [0f32; 8];
    _mm256_storeu_ps(result.as_mut_ptr(), acum);
    let suma_vec: f32 = result.iter().sum();

    // Sumar el resto escalarmente
    suma_vec + resto.iter().sum::<f32>()
}

fn main() {
    let datos: Vec<f32> = (1..=16).map(|x| x as f32).collect();

    #[cfg(target_arch = "x86_64")]
    let suma = if is_x86_feature_detected!("avx2") {
        unsafe { suma_f32_avx2(&datos) }
    } else {
        datos.iter().sum()
    };

    #[cfg(not(target_arch = "x86_64"))]
    let suma: f32 = datos.iter().sum();

    println!("Suma 1..16 = {suma}"); // 136
}

std::simd en nightly: API portable

std::simd (stabilización en progreso) ofrece una API de alto nivel que el compilador traduce a las instrucciones óptimas para cada arquitectura: AVX2 en x86-64, NEON en ARM, etc.

#![feature(portable_simd)]
use std::simd::{f32x8, SimdFloat};

fn suma_simd_portable(datos: &[f32]) -> f32 {
    let mut acum = f32x8::splat(0.0); // vector de 8 ceros f32
    let chunks = datos.chunks_exact(8);
    let resto = chunks.remainder();

    for chunk in chunks {
        let v = f32x8::from_slice(chunk);
        acum += v;
    }

    acum.reduce_sum() + resto.iter().sum::<f32>()
}

fn multiplicar_simd(a: &[f32], b: &[f32]) -> Vec<f32> {
    assert_eq!(a.len(), b.len());
    let mut resultado = vec![0.0f32; a.len()];
    let n = a.len() / 8 * 8;

    for i in (0..n).step_by(8) {
        let va = f32x8::from_slice(&a[i..]);
        let vb = f32x8::from_slice(&b[i..]);
        (va * vb).copy_to_slice(&mut resultado[i..]);
    }
    for i in n..a.len() {
        resultado[i] = a[i] * b[i];
    }
    resultado
}

repr(align) y prefetch manual

Alinear los datos a los límites de línea de caché (64 bytes) elimina fallos de alineamiento que reducen el rendimiento SIMD. El prefetch manual avisa al procesador de qué datos necesitará a continuación.

// Datos alineados a 64 bytes (tamaño de línea de caché)
#[repr(align(64))]
struct DatosAlineados {
    valores: [f32; 256],
}

// LTO en Cargo.toml para maximizar inlining entre crates:
// [profile.release]
// lto = "fat"
// codegen-units = 1
// opt-level = 3

// Prefetch manual en nightly (raramente necesario, el hardware lo hace bien):
// use std::arch::x86_64::_mm_prefetch;
// unsafe { _mm_prefetch(ptr as *const i8, _MM_HINT_T0); }

fn main() {
    let mut d = DatosAlineados { valores: [1.0; 256] };
    d.valores[0] = 42.0;
    println!("Primer valor: {}", d.valores[0]);
    println!("Alineación: {} bytes", std::mem::align_of::<DatosAlineados>()); // 64
}

La regla práctica: antes de escribir intrinsics AVX2, comprueba si el compilador ya auto-vectoriza el código. Con --release y target-cpu=native, Rust puede vectorizar bucles simples automáticamente. Usa godbolt.org con rustc 1.x -C opt-level=3 -C target-cpu=native para ver el ensamblado generado. Solo desciende a intrinsics manuales cuando la auto-vectorización no sea suficiente y el flamegraph confirme que ese bucle es el cuello de botella.