khtst-rust / src /simd.rs
v0.9.0: análise/prova/teste das críticas Qwen3.8 lote 2 (r-34) — SIMD com despacho em runtime bit-a-bit (X-1), oráculo candle_nn::optim::AdamW (nosso 1,53× mais rápido), Thompson Sampling (regret 0,0019 vs UCB1 0,0156), GNG por entropia com piso (T-G1'), JEPA de fração adaptativa (bit-a-bit com 1 braço), configs tipadas, entrada numpy (2008× na fronteira), avaliação com dados reais — 104/104 gates, 111/111 cargo
d444358 verified 11 days ago
Raw History Blame Contribute Delete
30.3 kB
//! v0.9.0 — Núcleos SIMD com despacho em RUNTIME (crítica Qwen3.8 nº 4 +
//! plano de ação itens 1/3; erratas e provas em r-34 §4/§9).
//!
//! **ERRATA-1 (da proposta original)**: `RUSTFLAGS="-C target-feature=
//! +avx512f,..."` NÃO cabe no `[profile.release]` do Cargo.toml (só em
//! `.cargo/config.toml` ou variável de ambiente) e, num repo PUBLICADO,
//! produz binários que travam com SIGILL em toda CPU sem AVX-512. O padrão
//! SOTA — o mesmo usado pelo crate `gemm` que o candle usa para matmul — é:
//! detecção de feature EM RUNTIME (`is_x86_feature_detected!`) + kernel
//! especializado `#[target_feature]` + fallback escalar. Portabilidade e
//! velocidade simultâneas; a velocidade deixa de ser um compromisso com a
//! corretude em outras máquinas.
//!
//! **ERRATA-2**: `#[repr(C, align(64))]` numa struct que CONTÉM `Vec` alinha
//! a STRUCT (na pilha), não o buffer de HEAP do `Vec` — o buffer continua com
//! o alinhamento do alocador global (glibc: 16 bytes). Para alinhar o heap é
//! preciso alocação explícita com `Layout::from_size_align(.., 64)` —
//! implementado em [`BufAlinhado64`] e verificado por endereço no teste
//! `simd_enderecos`.
//!
//! **EXATIDÃO (Teo X-1, r-34 §4.3)**: os kernels "exatos" NÃO usam FMA nem
//! aproximações (`rsqrt`/`rcp`) e preservam a ORDEM das operações do código
//! escalar que substituem ⇒ saída BIT-A-BIT idêntica (IEEE-754: mesma
//! sequência de operações ⇒ mesmo resultado; verificado por `to_bits()` nos
//! testes `adamw_bit_a_bit`, `escalar_bit_a_bit`, `distancias_bit_a_bit`).
//! O kernel `distancias4_rapidas` usa redução horizontal vetorizada
//! (ordem de soma diferente) — documentado como ε-diferente, nunca usado
//! onde a igualdade bit-a-bit é garantia.
//!
//! **Onde NÃO vetorizamos (Amdahl honesto)**: `norma_global` é uma REDUÇÃO —
//! vetorizar muda a ordem de soma e quebraria a bit-identidade; fica escalar
//! (custo marginal no passo). Reduções em f64 com associação diferente são a
//! linha exata entre "otimização segura" e "mudança silenciosa de resultados".

use std::alloc::{alloc, dealloc, Layout};
use std::sync::atomic::{AtomicUsize, Ordering};

// ---------------------------------------------------------------------------
// Contadores de telemetria (painel_v090)
// ---------------------------------------------------------------------------
static CHAMADAS_AVX: AtomicUsize = AtomicUsize::new(0);
static CHAMADAS_ESCALAR: AtomicUsize = AtomicUsize::new(0);

/// Relatório do despacho SIMD (para telemetria/painel_v090).
pub fn relatorio() -> serde_json::Value {
    serde_json::json!({
        "avx512_disponivel": tem_avx512(),
        "chamadas_avx": CHAMADAS_AVX.load(Ordering::Relaxed),
        "chamadas_escalar": CHAMADAS_ESCALAR.load(Ordering::Relaxed),
    })
}

/// AVX-512 (F+DQ+BW+VL) disponível nesta CPU?
pub fn tem_avx512() -> bool {
    #[cfg(target_arch = "x86_64")]
    {
        std::arch::is_x86_feature_detected!("avx512f")
            && std::arch::is_x86_feature_detected!("avx512dq")
            && std::arch::is_x86_feature_detected!("avx512bw")
            && std::arch::is_x86_feature_detected!("avx512vl")
    }
    #[cfg(not(target_arch = "x86_64"))]
    {
        false
    }
}

// ---------------------------------------------------------------------------
// Kernel 1: passo AdamW FUNDIDO (elemento-a-elemento, sem reduções)
// ---------------------------------------------------------------------------

/// Um passo AdamW sobre UM tensor achatado — réplica exata (ordem de ops) do
/// loop escalar do `AdamWManual::passo` da v0.8.0:
///
/// ```text
/// m[j] = β₁·m[j] + (1−β₁)·g[j]
/// v[j] = β₂·v[j] + ((1−β₂)·g[j])·g[j]        // assoc. esquerda, como o escalar
/// mh   = m[j]/bc₁ ; vh = v[j]/bc₂
/// out[j] = w[j] − lr·(mh/(√vh+ε) + wd·w[j])
/// ```
///
/// Sem FMA, sem rcp/rsqrt — `add/mul/div/sqrt` apenas, todas exatas por
/// IEEE-754 ⇒ **bit-a-bit idêntico** ao escalar (Teo X-1).
pub fn passo_adamw(
    w: &[f32],
    g: &[f32],
    m: &mut [f32],
    v: &mut [f32],
    out: &mut [f32],
    lr: f64,
    beta1: f64,
    beta2: f64,
    eps: f64,
    wd: f64,
    bc1: f64,
    bc2: f64,
) {
    let n = w.len().min(g.len()).min(m.len()).min(v.len()).min(out.len());
    let (lr, b1, b2, eps, wd, bc1, bc2) = (
        lr as f32,
        beta1 as f32,
        beta2 as f32,
        eps as f32,
        wd as f32,
        bc1 as f32,
        bc2 as f32,
    );
    // réplica EXATA do escalar v0.8.0: `(1.0 - self.beta1 as f32)` — a
    // subtração é em f32 (depois do cast), não em f64 — trocar isso quebra
    // a bit-identidade no último ulp.
    let omb1 = 1.0f32 - b1;
    let omb2 = 1.0f32 - b2;
    let usar_avx = tem_avx512() && n >= 64;
    if usar_avx {
        CHAMADAS_AVX.fetch_add(1, Ordering::Relaxed);
        unsafe { passo_adamw_avx(w, g, m, v, out, lr, b1, b2, eps, wd, bc1, bc2, omb1, omb2) };
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
    }
    // cauda + caminho escalar: MESMO código nas duas situações ⇒ bits idênticos
    let ini = if usar_avx { n & !15 } else { 0 };
    for j in ini..n {
        m[j] = b1 * m[j] + omb1 * g[j];
        v[j] = b2 * v[j] + omb2 * g[j] * g[j];
        let mh = m[j] / bc1;
        let vh = v[j] / bc2;
        out[j] = w[j] - lr * (mh / (vh.sqrt() + eps) + wd * w[j]);
    }
}

/// Caminho 100% escalar (referência de bench A/B e de leitura) — mesmo corpo
/// da cauda; SEM despacho. Exposto para medição honesta do ganho AVX-512.
pub fn passo_adamw_escalar(
    w: &[f32], g: &[f32], m: &mut [f32], v: &mut [f32], out: &mut [f32],
    lr: f64, beta1: f64, beta2: f64, eps: f64, wd: f64, bc1: f64, bc2: f64,
) {
    let n = w.len().min(g.len()).min(m.len()).min(v.len()).min(out.len());
    let (lrf, b1, b2, epsf, wdf, bc1f, bc2f) = (
        lr as f32, beta1 as f32, beta2 as f32, eps as f32, wd as f32, bc1 as f32, bc2 as f32,
    );
    let omb1 = 1.0f32 - b1;
    let omb2 = 1.0f32 - b2;
    for j in 0..n {
        m[j] = b1 * m[j] + omb1 * g[j];
        v[j] = b2 * v[j] + omb2 * g[j] * g[j];
        let mh = m[j] / bc1f;
        let vh = v[j] / bc2f;
        out[j] = w[j] - lrf * (mh / (vh.sqrt() + epsf) + wdf * w[j]);
    }
}

/// Kernel AVX-512 f32 (16 lanes). Processa `n & !15` elementos; a cauda fica
/// para o loop escalar do chamador (mesma matemática ⇒ bits idênticos).
///
/// # Segurança
/// Requer AVX-512 (F) — checado por [`tem_avx512`] no chamador. Fatias de 16.
#[target_feature(enable = "avx512f,avx512dq")]
unsafe fn passo_adamw_avx(
    w: &[f32],
    g: &[f32],
    m: &mut [f32],
    v: &mut [f32],
    out: &mut [f32],
    lr: f32,
    b1: f32,
    b2: f32,
    eps: f32,
    wd: f32,
    bc1: f32,
    bc2: f32,
    omb1: f32,
    omb2: f32,
) {
    use std::arch::x86_64::*;
    let n = w.len().min(g.len()).min(m.len()).min(v.len()).min(out.len());
    let v_lr = _mm512_set1_ps(lr);
    let v_b1 = _mm512_set1_ps(b1);
    let v_b2 = _mm512_set1_ps(b2);
    let v_eps = _mm512_set1_ps(eps);
    let v_wd = _mm512_set1_ps(wd);
    let v_bc1 = _mm512_set1_ps(bc1);
    let v_bc2 = _mm512_set1_ps(bc2);
    let v_omb1 = _mm512_set1_ps(omb1);
    let v_omb2 = _mm512_set1_ps(omb2);
    let mut j = 0usize;
    while j + 16 <= n {
        let vw = _mm512_loadu_ps(w.as_ptr().add(j));
        let vg = _mm512_loadu_ps(g.as_ptr().add(j));
        let vm = _mm512_loadu_ps(m.as_ptr().add(j));
        let vv = _mm512_loadu_ps(v.as_ptr().add(j));
        // m = b1*m + omb1*g ; v = b2*v + (omb2*g)*g   [sem FMA — ordem exata]
        let t1 = _mm512_mul_ps(v_omb1, vg);
        let nm = _mm512_add_ps(_mm512_mul_ps(v_b1, vm), t1);
        let t2 = _mm512_mul_ps(v_omb2, vg);
        let nv = _mm512_add_ps(_mm512_mul_ps(v_b2, vv), _mm512_mul_ps(t2, vg));
        // mh = m/bc1 ; vh = v/bc2 ; ajuste = mh/(sqrt(vh)+eps) + wd*w
        let mh = _mm512_div_ps(nm, v_bc1);
        let vh = _mm512_div_ps(nv, v_bc2);
        let raiz = _mm512_sqrt_ps(vh); // IEEE-exata (não é rsqrt aproximado)
        let denom = _mm512_add_ps(raiz, v_eps);
        let quoc = _mm512_div_ps(mh, denom);
        let decaimento = _mm512_mul_ps(v_wd, vw);
        let ajuste = _mm512_add_ps(quoc, decaimento);
        let passo = _mm512_mul_ps(v_lr, ajuste);
        let novo = _mm512_sub_ps(vw, passo);
        _mm512_storeu_ps(m.as_mut_ptr().add(j), nm);
        _mm512_storeu_ps(v.as_mut_ptr().add(j), nv);
        _mm512_storeu_ps(out.as_mut_ptr().add(j), novo);
        j += 16;
    }
}

// ---------------------------------------------------------------------------
// Kernel 2: escala elemento-a-elemento (clip de gradiente)
// ---------------------------------------------------------------------------

/// `x[i] *= f` para todo i — elemento-a-elemento ⇒ bit-a-bit idêntico ao
/// loop escalar do `AdamWManual::clipar` (uma única multiplicação por elemento).
pub fn escalar_slice(xs: &mut [f32], f: f32) {
    let n = xs.len();
    if tem_avx512() && n >= 64 {
        CHAMADAS_AVX.fetch_add(1, Ordering::Relaxed);
        unsafe { escalar_slice_avx(xs, f) };
        let ini = n & !15;
        for x in xs[ini..].iter_mut() {
            *x *= f;
        }
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        for x in xs.iter_mut() {
            *x *= f;
        }
    }
}

#[target_feature(enable = "avx512f,avx512dq")]
unsafe fn escalar_slice_avx(xs: &mut [f32], f: f32) {
    use std::arch::x86_64::*;
    let vf = _mm512_set1_ps(f);
    let n = xs.len();
    let mut j = 0usize;
    while j + 16 <= n {
        let x = _mm512_loadu_ps(xs.as_ptr().add(j));
        _mm512_storeu_ps(xs.as_mut_ptr().add(j), _mm512_mul_ps(x, vf));
        j += 16;
    }
}

// ---------------------------------------------------------------------------
// Kernel 3: distâncias² 4D em layout SoA (crítica Qwen3.8 B.3)
// ---------------------------------------------------------------------------

/// Layout SoA column-major: `soa[j*c + i]` = coordenada j do protótipo i
/// (d = 4). Construído por [`soa_de_prototipos`].
///
/// **Teo X-2**: o termo (x_j − w_ij)² é computado por lane com sub/mul (sem
/// FMA) e a soma por protótipo percorre j = 0,1,2,3 na MESMA ordem do
/// `distancia2` escalar (fold esquerdo) ⇒ **bit-a-bit idêntico** ao caminho
/// `Vec<Vec<f64>>` da v0.8.0. A diferença é só a localidade: 1 buffer
/// contíguo (sem perseguição de ponteiro por protótipo) + 8 protótipos por
/// instrução no kernel vetorial.
pub fn soa_de_prototipos(ps: &[Vec<f64>]) -> (Vec<f64>, usize, usize) {
    let c = ps.len();
    let d = ps.first().map(|p| p.len()).unwrap_or(0);
    let mut soa = vec![0.0f64; d * c];
    for (i, p) in ps.iter().enumerate() {
        for (j, x) in p.iter().enumerate() {
            if j < d {
                soa[j * c + i] = *x;
            }
        }
    }
    (soa, c, d)
}

/// Distâncias² de `x` a todos os `c` protótipos (modo EXATO — bit-a-bit).
/// `out` é preenchido com `distancia2(x, protótipo_i)` para todo i.
pub fn distancias4_exatas(soa: &[f64], c: usize, x: &[f64], out: &mut [f64]) {
    debug_assert_eq!(soa.len(), 4 * c);
    let usar_avx = tem_avx512() && c >= 64 && x.len() == 4;
    if usar_avx {
        CHAMADAS_AVX.fetch_add(1, Ordering::Relaxed);
        unsafe { distancias4_exatas_avx(soa, c, x, out) };
        let ini = c & !7;
        distancias4_escalar(soa, c, x, out, ini);
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        distancias4_escalar(soa, c, x, out, 0);
    }
}

/// Trecho escalar — a MESMA ordem de soma do `distancia2` (fold esquerdo em j).
fn distancias4_escalar(soa: &[f64], c: usize, x: &[f64], out: &mut [f64], ini: usize) {
    let d = if c == 0 { 0 } else { soa.len() / c };
    for i in ini..c {
        let mut s = 0.0f64;
        for j in 0..d {
            let dx = x[j] - soa[j * c + i];
            s += dx * dx;
        }
        out[i] = s;
    }
}

/// Redução horizontal exata: acumuladores por j são ARMAZENADOS e somados por
/// protótipo na ordem ((q0+q1)+q2)+q3 — a ordem do fold escalar.
#[target_feature(enable = "avx512f,avx512dq")]
unsafe fn distancias4_exatas_avx(soa: &[f64], c: usize, x: &[f64], out: &mut [f64]) {
    use std::arch::x86_64::*;
    let mut buf = [[0f64; 8]; 4];
    let mut i = 0usize;
    while i + 8 <= c {
        let mut quads = [_mm512_setzero_pd(); 4];
        for (j, q) in quads.iter_mut().enumerate() {
            let wj = _mm512_loadu_pd(soa.as_ptr().add(j * c + i));
            let xj = _mm512_set1_pd(x[j]);
            let dx = _mm512_sub_pd(xj, wj);
            *q = _mm512_mul_pd(dx, dx); // sem FMA
        }
        for j in 0..4 {
            _mm512_storeu_pd(buf[j].as_mut_ptr(), quads[j]);
        }
        for k in 0..8 {
            out[i + k] = ((buf[0][k] + buf[1][k]) + buf[2][k]) + buf[3][k];
        }
        i += 8;
    }
}

/// Modo RÁPIDO: redução horizontal vetorizada `(q0+q1)+(q2+q3)` — ordem de
/// associação DIFERENTE do escalar ⇒ ε-diferença (≈ ulp), documentada;
/// usar apenas onde igualdade bit-a-bit não é garantia (ex.: sondagem).
pub fn distancias4_rapidas(soa: &[f64], c: usize, x: &[f64], out: &mut [f64]) {
    debug_assert_eq!(soa.len(), 4 * c);
    if tem_avx512() && c >= 64 && x.len() == 4 {
        CHAMADAS_AVX.fetch_add(1, Ordering::Relaxed);
        unsafe {
            use std::arch::x86_64::*;
            let mut i = 0usize;
            while i + 8 <= c {
                let mut quads = [_mm512_setzero_pd(); 4];
                for (j, q) in quads.iter_mut().enumerate() {
                    let wj = _mm512_loadu_pd(soa.as_ptr().add(j * c + i));
                    let xj = _mm512_set1_pd(x[j]);
                    let dx = _mm512_sub_pd(xj, wj);
                    *q = _mm512_mul_pd(dx, dx);
                }
                let r = _mm512_add_pd(
                    _mm512_add_pd(quads[0], quads[1]),
                    _mm512_add_pd(quads[2], quads[3]),
                );
                _mm512_storeu_pd(out.as_mut_ptr().add(i), r);
                i += 8;
            }
        };
        let ini = c & !7;
        distancias4_escalar(soa, c, x, out, ini);
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        distancias4_escalar(soa, c, x, out, 0);
    }
}


/// Layout ROW-MAJOR plano: `rm[i*d + j]` (protótipos contíguos, SEM
/// perseguição de ponteiros) — v0.9.0 (r-34 §6.4): a medição mostrou que o
/// SoA COLUMN-major (`soa[j*c+i]`) é MAIS LENTO que `Vec<Vec>` em d=4
/// (strides de c·8 B entre coordenadas quebram a localidade), enquanto o
/// row-major plano processa 2 protótipos por registrador ZMM (8 f64) com
/// redução horizontal EXATA por par (mesma ordem de soma do escalar).
pub fn rm_de_prototipos(ps: &[Vec<f64>]) -> (Vec<f64>, usize, usize) {
    let c = ps.len();
    let d = ps.first().map(|p| p.len()).unwrap_or(0);
    let mut rm = vec![0.0f64; d * c];
    for (i, p) in ps.iter().enumerate() {
        for (j, x) in p.iter().enumerate() {
            if j < d {
                rm[i * d + j] = *x;
            }
        }
    }
    (rm, c, d)
}

/// Distâncias² d=4 sobre row-major plano — bit-a-bit idêntico ao escalar
/// (sub/mul por lane SEM FMA + redução por protótipo na ordem j=0..3).
pub fn distancias4_rm(rm: &[f64], c: usize, x: &[f64; 4], out: &mut [f64]) {
    if tem_avx512() && c >= 64 {
        CHAMADAS_AVX.fetch_add(1, Ordering::Relaxed);
        unsafe { distancias4_rm_avx(rm, c, x, out) };
        let ini = c & !1;
        distancias4_rm_escalar(rm, c, x, out, ini);
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        distancias4_rm_escalar(rm, c, x, out, 0);
    }
}

fn distancias4_rm_escalar(rm: &[f64], c: usize, x: &[f64; 4], out: &mut [f64], ini: usize) {
    for i in ini..c {
        let b = &rm[i * 4..i * 4 + 4];
        let t0 = x[0] - b[0];
        let t1 = x[1] - b[1];
        let t2 = x[2] - b[2];
        let t3 = x[3] - b[3];
        out[i] = ((t0 * t0 + t1 * t1) + t2 * t2) + t3 * t3;
    }
}

/// 2 protótipos por ZMM: carga de 8 f64, x replicado [x0..x3, x0..x3],
/// sub/mul por lane; redução por PAR na ordem exata do escalar.
#[target_feature(enable = "avx512f,avx512dq")]
unsafe fn distancias4_rm_avx(rm: &[f64], c: usize, x: &[f64; 4], out: &mut [f64]) {
    use std::arch::x86_64::*;
    let xv = _mm512_setr_pd(x[0], x[1], x[2], x[3], x[0], x[1], x[2], x[3]);
    let mut i = 0usize;
    let mut par = [0f64; 8];
    while i + 2 <= c {
        let w = _mm512_loadu_pd(rm.as_ptr().add(i * 4));
        let dif = _mm512_sub_pd(xv, w);
        let q = _mm512_mul_pd(dif, dif); // lanes: [q0p, q1p, q2p, q3p, q0q, ..., q3q]
        _mm512_storeu_pd(par.as_mut_ptr(), q);
        let (p, r) = (i, i + 1);
        // par JÁ contém os quadrados (q = dif*dif) — só SOMA, na ordem exata
        out[p] = ((par[0] + par[1]) + par[2]) + par[3];
        out[r] = ((par[4] + par[5]) + par[6]) + par[7];
        i += 2;
    }
}


/// Distâncias² row-major (despacho: d==4 → kernel AVX-512 2-protótipos/ZMM;
/// senão escalar genérico com a mesma ordem de soma) — preenche `buf`.
pub fn distancias_rm(rm: &[f64], c: usize, d: usize, x: &[f64], buf: &mut Vec<f64>) {
    buf.clear();
    buf.resize(c, 0.0);
    if d == 4 && x.len() == 4 {
        let x4 = [x[0], x[1], x[2], x[3]];
        distancias4_rm(rm, c, &x4, buf);
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        for i in 0..c {
            let mut s = 0.0;
            for j in 0..d {
                let dx = x[j] - rm[i * d + j];
                s += dx * dx;
            }
            buf[i] = s;
        }
    }
}

/// BMU sobre row-major plano — equivalente bit-a-bit a `melhor_indice(ps, x)`
/// (d=4; teste `distancias_bit_a_bit_rm`).
pub fn melhor_indice_rm_buf(rm: &[f64], c: usize, x: &[f64], buf: &mut Vec<f64>) -> usize {
    if c == 0 {
        return 0;
    }
    buf.clear();
    buf.resize(c, 0.0);
    if x.len() == 4 {
        let x4 = [x[0], x[1], x[2], x[3]];
        distancias4_rm(rm, c, &x4, buf);
    } else {
        let d = if c == 0 { 0 } else { rm.len() / c };
        for i in 0..c {
            let mut s = 0.0;
            for j in 0..d {
                let dx = x[j] - rm[i * d + j];
                s += dx * dx;
            }
            buf[i] = s;
        }
    }
    argmin_primeiro(buf)
}

/// argmin com desempate pelo MENOR índice — semântica idêntica à do
/// `melhor_indice` escalar (`dd < dm` estrito mantém o primeiro mínimo).
pub fn argmin_primeiro(xs: &[f64]) -> usize {
    let mut melhor = 0usize;
    let mut dm = f64::INFINITY;
    for (i, x) in xs.iter().enumerate() {
        if *x < dm {
            dm = *x;
            melhor = i;
        }
    }
    melhor
}

/// Distâncias² para TODO d (despacho: d==4 → kernel AVX-512 exato; senão
/// escalar com a MESMA ordem de soma do `distancia2`) — preenche `buf`
/// (reutilizado entre chamadas: zero alocação por amostra).
pub fn distancias_soa(soa: &[f64], c: usize, d: usize, x: &[f64], buf: &mut Vec<f64>) {
    buf.clear();
    buf.resize(c, 0.0);
    if d == 4 {
        distancias4_exatas(soa, c, x, buf);
    } else {
        CHAMADAS_ESCALAR.fetch_add(1, Ordering::Relaxed);
        distancias4_escalar(soa, c, x, buf, 0);
    }
}

/// BMU sobre layout SoA — equivalente bit-a-bit a `melhor_indice(ps, x)`
/// (Teo X-2; teste `distancias_bit_a_bit`). Versão com buffer reutilizado.
pub fn melhor_indice_soa_buf(soa: &[f64], c: usize, d: usize, x: &[f64], buf: &mut Vec<f64>) -> usize {
    if c == 0 {
        return 0;
    }
    distancias_soa(soa, c, d, x, buf);
    argmin_primeiro(buf)
}

/// BMU sobre layout SoA — equivalente bit-a-bit a `melhor_indice(ps, x)`
/// (Teo X-2; teste `distancias_bit_a_bit`).
pub fn melhor_indice_soa(soa: &[f64], c: usize, d: usize, x: &[f64]) -> usize {
    let mut buf = Vec::with_capacity(c.max(1));
    melhor_indice_soa_buf(soa, c, d, x, &mut buf)
}

// ---------------------------------------------------------------------------
// Buffer de heap com alinhamento GARANTIDO de 64 B (errata-2 da proposta)
// ---------------------------------------------------------------------------

/// Buffer f64 com alinhamento de heap GARANTIDO em 64 bytes (linha de cache /
// registrador ZMM). Esta é a forma CORRETA de obter o que a proposta
/// `#[repr(align(64))]` em struct-com-Vec prometia (e não entrega: isso
/// alinha a struct na pilha, não o buffer — ver teste `simd_enderecos`).
pub struct BufAlinhado64 {
    ptr: *mut f64,
    len: usize,
    layout: Layout,
}

impl BufAlinhado64 {
    /// Aloca `len` f64 zerados, alinhados a 64 B.
    pub fn novo_zeros(len: usize) -> Self {
        let len_b = (len * std::mem::size_of::<f64>()).max(1);
        let layout = Layout::from_size_align(len_b, 64).expect("layout 64B inválido");
        let ptr = unsafe { alloc(layout) };
        assert!(!ptr.is_null(), "alocação alinhada falhou");
        let buf = unsafe { std::slice::from_raw_parts_mut(ptr as *mut f64, len.max(1)) };
        buf.fill(0.0);
        Self { ptr: ptr as *mut f64, len: len.max(1), layout }
    }

    /// Cópia alinhada de uma fatia.
    pub fn de_fatia(xs: &[f64]) -> Self {
        let mut b = Self::novo_zeros(xs.len());
        b.as_mut_slice().copy_from_slice(xs);
        b
    }

    pub fn as_slice(&self) -> &[f64] {
        if self.len == 0 {
            &[]
        } else {
            unsafe { std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }

    pub fn as_mut_slice(&mut self) -> &mut [f64] {
        if self.len == 0 {
            &mut []
        } else {
            unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }

    /// Endereço do buffer (para verificação de alinhamento em testes/bench).
    pub fn endereco(&self) -> usize {
        self.ptr as usize
    }
}

impl Drop for BufAlinhado64 {
    fn drop(&mut self) {
        unsafe { dealloc(self.ptr as *mut u8, self.layout) };
    }
}

impl std::fmt::Debug for BufAlinhado64 {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("BufAlinhado64")
            .field("len", &self.len)
            .field("endereco", &format!("{:#x}", self.endereco()))
            .finish()
    }
}

// O ponteiro é ownership exclusivo (sem compartilhamento); Send/Sync seguros.
unsafe impl Send for BufAlinhado64 {}
unsafe impl Sync for BufAlinhado64 {}

// ---------------------------------------------------------------------------
// Testes — falsificam: FMA escondido, ordem de soma trocada, alinhamento falso
// ---------------------------------------------------------------------------

#[cfg(test)]
mod testes {
    use super::*;

    fn rng_vec(n: usize, semente: u64) -> Vec<f32> {
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        (0..n).map(|_| rng.normal(0.0, 1.0) as f32).collect()
    }

    /// Teo X-1 (falsificação): qualquer FMA/rsqrt/ordem diferente quebra os bits.
    #[test]
    fn adamw_bit_a_bit() {
        let n = 1000usize;
        let w = rng_vec(n, 1);
        let g = rng_vec(n, 2);
        let mut m1 = rng_vec(n, 3);
        let mut v1 = rng_vec(n, 4);
        let mut m2 = m1.clone();
        let mut v2 = v1.clone();
        let mut o1 = vec![0f32; n];
        let mut o2 = vec![0f32; n];
        // referência: loop escalar IDÊNTICO ao que estava em treino.rs v0.8.0
        // (atenção: `(1.0 - b1 as f32)` — subtração em f32 pós-cast)
        let (lr, b1, b2, eps, wd, bc1, bc2) = (3e-4f64, 0.9, 0.999, 1e-8, 0.01, 0.91, 0.001);
        let (lrf, b1f, b2f, epsf, wdf, bc1f, bc2f) = (
            lr as f32,
            b1 as f32,
            b2 as f32,
            eps as f32,
            wd as f32,
            bc1 as f32,
            bc2 as f32,
        );
        for j in 0..n {
            m1[j] = b1f * m1[j] + (1.0f32 - b1f) * g[j];
            v1[j] = b2f * v1[j] + (1.0f32 - b2f) * g[j] * g[j];
            let mh = m1[j] / bc1f;
            let vh = v1[j] / bc2f;
            o1[j] = w[j] - lrf * (mh / (vh.sqrt() + epsf) + wdf * w[j]);
        }
        // caminho SIMD (com cauda escalar interna)
        passo_adamw(&w, &g, &mut m2, &mut v2, &mut o2, lr, b1, b2, eps, wd, bc1, bc2);
        for j in 0..n {
            assert_eq!(o1[j].to_bits(), o2[j].to_bits(), "saída divergiu em {j}");
            assert_eq!(m1[j].to_bits(), m2[j].to_bits(), "m divergiu em {j}");
            assert_eq!(v1[j].to_bits(), v2[j].to_bits(), "v divergiu em {j}");
        }
    }

    #[test]
    fn escalar_bit_a_bit() {
        let n = 500usize;
        let mut a = rng_vec(n, 5);
        let mut b = a.clone();
        let f = 0.37f32;
        escalar_slice(&mut a, f);
        for x in b.iter_mut() {
            *x *= f;
        }
        for j in 0..n {
            assert_eq!(a[j].to_bits(), b[j].to_bits());
        }
    }

    /// Teo X-2 (falsificação): ordem de soma ou FMA no kernel SoA quebra bits.
    #[test]
    fn distancias_bit_a_bit() {
        let (c, d) = (300usize, 4usize);
        let mut rng = crate::rng::RngDeterministico::novo(11);
        let ps: Vec<Vec<f64>> = (0..c)
            .map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect())
            .collect();
        let x: Vec<f64> = (0..d).map(|_| rng.normal(0.0, 2.0)).collect();
        let (soa, c2, d2) = soa_de_prototipos(&ps);
        assert_eq!((c2, d2), (c, d));
        // melhor_indice escalar da v0.8.0 (réplica)
        let dist_esc = |a: &[f64], b: &[f64]| -> f64 {
            a.iter().zip(b).map(|(p, q)| (p - q) * (p - q)).sum()
        };
        let mut melhor = 0usize;
        let mut dm = f64::INFINITY;
        for (i, p) in ps.iter().enumerate() {
            let dd = dist_esc(&x, p);
            if dd < dm {
                dm = dd;
                melhor = i;
            }
        }
        // caminho SoA (AVX-512 quando disponível; cauda escalar)
        let obtido = melhor_indice_soa(&soa, c, d, &x);
        assert_eq!(obtido, melhor, "BMU divergiu");
        // distâncias completas bit-a-bit
        let mut out = vec![0.0f64; c];
        distancias4_exatas(&soa, c, &x, &mut out);
        for (i, p) in ps.iter().enumerate() {
            assert_eq!(
                out[i].to_bits(),
                dist_esc(&x, p).to_bits(),
                "distância divergiu no protótipo {i}"
            );
        }
        // modo rápido: ε-próximo e mesmo BMU (ordem de associação diferente)
        let mut rap = vec![0.0f64; c];
        distancias4_rapidas(&soa, c, &x, &mut rap);
        for i in 0..c {
            let esc = dist_esc(&x, &ps[i]);
            assert!(
                (rap[i] - esc).abs() <= 1e-9 * (1.0 + esc.abs()),
                "modo rápido longe demais: {} vs {}",
                rap[i],
                esc
            );
        }
        assert_eq!(argmin_primeiro(&rap), melhor);
    }


    /// Teo X-2' (row-major): o kernel de 2-protótipos-por-ZMM preserva bits
    /// contra o escalar `distancia2`/`melhor_indice`.
    #[test]
    fn distancias_bit_a_bit_rm() {
        let (c, d) = (500usize, 4usize);
        let mut rng = crate::rng::RngDeterministico::novo(19);
        let ps: Vec<Vec<f64>> = (0..c)
            .map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect())
            .collect();
        let x: Vec<f64> = (0..d).map(|_| rng.normal(0.0, 2.0)).collect();
        let (rm, c2, d2) = rm_de_prototipos(&ps);
        assert_eq!((c2, d2), (c, d));
        let dist_esc = |a: &[f64], b: &[f64]| -> f64 {
            a.iter().zip(b).map(|(p, q)| (p - q) * (p - q)).sum()
        };
        let mut melhor = 0usize;
        let mut dm = f64::INFINITY;
        for (i, p) in ps.iter().enumerate() {
            let dd = dist_esc(&x, p);
            if dd < dm {
                dm = dd;
                melhor = i;
            }
        }
        let mut buf = Vec::with_capacity(c);
        let obtido = melhor_indice_rm_buf(&rm, c, &x, &mut buf);
        assert_eq!(obtido, melhor, "BMU rm divergiu");
        for (i, p) in ps.iter().enumerate() {
            assert_eq!(buf[i].to_bits(), dist_esc(&x, p).to_bits(), "rm dist divergiu em {i}");
        }
    }

    /// Errata-2 empirica: o buffer dedicado É alinhado a 64; `Vec` comum NÃO
    /// tem essa garantia (apenas o alinhamento do tipo/allocador).
    #[test]
    fn simd_enderecos() {
        for len in [1usize, 7, 64, 257, 1000, 4096] {
            let b = BufAlinhado64::novo_zeros(len);
            assert_eq!(b.endereco() % 64, 0, "alinhamento 64B violado (len={len})");
            assert_eq!(b.as_slice().len(), len.max(1));
        }
        // cópia preserva conteúdo e alinhamento
        let orig: Vec<f64> = (0..100).map(|i| i as f64 * 0.5).collect();
        let b = BufAlinhado64::de_fatia(&orig);
        assert_eq!(b.endereco() % 64, 0);
        assert_eq!(b.as_slice(), orig.as_slice());
        // Vec<f64> do alocador global: SEM garantia de 64B (documenta a errata;
        // não afirmamos que sempre viola — afirmamos que NÃO é garantido)
        let mut violacoes = 0;
        for _ in 0..32 {
            let v: Vec<f64> = (0..1000).map(|i| i as f64).collect();
            if v.as_ptr() as usize % 64 != 0 {
                violacoes += 1;
            }
        }
        // com glibc, alocações grandes frequentemente caem em páginas alinhadas
        // — o ponto é a GARANTIA, não a frequência; apenas registramos.
        eprintln!("[simd_enderecos] Vec<f64> não-alinhado em {violacoes}/32 alocações");
    }

    #[test]
    fn despacho_conta_chamadas() {
        let antes = relatorio();
        let b = BufAlinhado64::novo_zeros(128);
        let x = [0.1f64, 0.2, 0.3, 0.4];
        let mut out = vec![0.0f64; 128];
        distancias4_exatas(b.as_slice(), 128, &x, &mut out);
        let depois = relatorio();
        let n0 = antes["chamadas_avx"].as_u64().unwrap_or(0)
            + antes["chamadas_escalar"].as_u64().unwrap_or(0);
        let n1 = depois["chamadas_avx"].as_u64().unwrap_or(0)
            + depois["chamadas_escalar"].as_u64().unwrap_or(0);
        assert!(n1 > n0, "contadores de despacho não incrementaram");
    }
}