oxicuda_solver/dense/

lu.rs

//! LU Factorization with partial pivoting.
//!
//! Computes `P * A = L * U` where:
//! - P is a permutation matrix (represented by pivot indices)
//! - L is unit lower triangular
//! - U is upper triangular
//!
//! Uses a blocked right-looking algorithm:
//! 1. Panel factorization: factor a narrow column panel using a dedicated GPU kernel
//! 2. Apply pivots: swap rows in the trailing portion
//! 3. TRSM: solve for the upper triangle block
//! 4. GEMM: update the trailing submatrix
//!
//! The L and U factors overwrite the input matrix A in-place (LAPACK-style packed
//! storage with unit diagonal for L implicitly assumed).

use oxicuda_blas::types::{
    DiagType, FillMode, GpuFloat, Layout, MatrixDesc, MatrixDescMut, Side, Transpose,
};
use oxicuda_memory::DeviceBuffer;
use oxicuda_ptx::prelude::*;

use crate::error::{SolverError, SolverResult};
use crate::handle::SolverHandle;
use crate::ptx_helpers::SOLVER_BLOCK_SIZE;

/// Block size for the panel factorization step.
const LU_BLOCK_SIZE: u32 = 64;

// ---------------------------------------------------------------------------
// Result type
// ---------------------------------------------------------------------------

/// Result of an LU factorization.
///
/// Contains diagnostic information about the factorization.
#[derive(Debug, Clone)]
pub struct LuResult {
    /// Status info:
    /// - 0: successful factorization
    /// - i > 0: U(i,i) is exactly zero, matrix is singular at column i
    pub info: i32,
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Performs LU factorization with partial pivoting in-place.
///
/// On exit, the lower triangle of `a` (with implicit unit diagonal) contains L,
/// and the upper triangle contains U. The `pivots` array records the row
/// permutations: row `i` was interchanged with row `pivots[i]`.
///
/// The matrix is stored in column-major order with leading dimension `lda`.
///
/// # Arguments
///
/// * `handle` — solver handle.
/// * `a` — matrix buffer (n x n, column-major, lda stride), modified in-place.
/// * `n` — matrix dimension.
/// * `lda` — leading dimension (>= n).
/// * `pivots` — output pivot indices buffer (length >= n).
///
/// # Returns
///
/// [`LuResult`] with `info == 0` on success, `info > 0` if singular.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid or a kernel launch fails.
pub fn lu_factorize<T: GpuFloat>(
    handle: &mut SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<LuResult> {
    // Validate dimensions.
    if n == 0 {
        return Ok(LuResult { info: 0 });
    }
    if lda < n {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: lda ({lda}) must be >= n ({n})"
        )));
    }
    let required = n as usize * lda as usize;
    if a.len() < required {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: buffer too small ({} < {required})",
            a.len()
        )));
    }
    if pivots.len() < n as usize {
        return Err(SolverError::DimensionMismatch(format!(
            "lu_factorize: pivots buffer too small ({} < {n})",
            pivots.len()
        )));
    }

    // Ensure workspace is large enough for panel temporaries.
    let panel_workspace = n as usize * LU_BLOCK_SIZE as usize * T::SIZE;
    handle.ensure_workspace(panel_workspace)?;

    blocked_lu::<T>(handle, a, n, lda, pivots)
}

/// Solves `A * X = B` given an LU-factored matrix.
///
/// The LU factors must have been computed by [`lu_factorize`]. The solution
/// overwrites `b` in-place.
///
/// # Arguments
///
/// * `handle` — solver handle.
/// * `lu` — LU-factored matrix (output of `lu_factorize`).
/// * `pivots` — pivot indices from `lu_factorize`.
/// * `b` — right-hand side matrix (n x nrhs), overwritten with solution.
/// * `n` — matrix dimension.
/// * `nrhs` — number of right-hand side columns.
///
/// # Errors
///
/// Returns [`SolverError`] if dimensions are invalid or BLAS operations fail.
pub fn lu_solve<T: GpuFloat>(
    handle: &SolverHandle,
    lu: &DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    b: &mut DeviceBuffer<T>,
    n: u32,
    nrhs: u32,
) -> SolverResult<()> {
    if n == 0 || nrhs == 0 {
        return Ok(());
    }
    if lu.len() < (n as usize * n as usize) {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: LU buffer too small".into(),
        ));
    }
    if pivots.len() < n as usize {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: pivots buffer too small".into(),
        ));
    }
    if b.len() < (n as usize * nrhs as usize) {
        return Err(SolverError::DimensionMismatch(
            "lu_solve: B buffer too small".into(),
        ));
    }

    // Step 1: Apply row permutations to B.
    // Each pivot[i] says row i was swapped with row pivot[i] during
    // factorization, so we replay the swaps in forward order.
    apply_pivots_to_rhs::<T>(handle, b, pivots, n, nrhs)?;

    // Step 2: Solve L * Y = P * B (forward substitution) via TRSM.
    let l_desc = MatrixDesc::<T>::from_raw(lu.as_device_ptr(), n, n, n, Layout::ColMajor);
    let mut b_desc = MatrixDescMut::<T>::from_raw(b.as_device_ptr(), n, nrhs, n, Layout::ColMajor);

    oxicuda_blas::level3::trsm(
        handle.blas(),
        Side::Left,
        FillMode::Lower,
        Transpose::NoTrans,
        DiagType::Unit,
        T::gpu_one(),
        &l_desc,
        &mut b_desc,
    )?;

    // Step 3: Solve U * X = Y (backward substitution) via TRSM.
    let u_desc = MatrixDesc::<T>::from_raw(lu.as_device_ptr(), n, n, n, Layout::ColMajor);

    oxicuda_blas::level3::trsm(
        handle.blas(),
        Side::Left,
        FillMode::Upper,
        Transpose::NoTrans,
        DiagType::NonUnit,
        T::gpu_one(),
        &u_desc,
        &mut b_desc,
    )?;

    Ok(())
}

// ---------------------------------------------------------------------------
// Blocked LU implementation
// ---------------------------------------------------------------------------

/// Blocked right-looking LU factorization.
///
/// Processes the matrix in column panels of width `LU_BLOCK_SIZE`:
/// 1. Factor the panel (find pivots, compute L column, compute U row).
/// 2. Swap rows in the trailing matrix according to pivots.
/// 3. TRSM: compute U block for the panel's upper triangle.
/// 4. GEMM: update the trailing submatrix.
fn blocked_lu<T: GpuFloat>(
    handle: &mut SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<LuResult> {
    let nb = LU_BLOCK_SIZE.min(n);
    let num_blocks = n.div_ceil(nb);
    let mut info: i32 = 0;

    for block_idx in 0..num_blocks {
        let j = block_idx * nb;
        let jb = nb.min(n - j); // Actual panel width (may be smaller for last block).

        // Step 1: Panel factorization — factorize columns j..j+jb of the
        // submatrix A[j:n, j:j+jb].
        let panel_info = panel_lu::<T>(handle, a, n, lda, j, jb, pivots)?;
        if panel_info > 0 && info == 0 {
            info = panel_info + j as i32;
        }

        // Step 2: Apply pivots to columns outside the panel.
        // Left side (columns 0..j): swap rows according to pivots.
        if j > 0 {
            apply_panel_pivots::<T>(handle, a, lda, j, jb, pivots, 0, j)?;
        }
        // Right side (columns j+jb..n): swap rows according to pivots.
        let right_start = j + jb;
        if right_start < n {
            apply_panel_pivots::<T>(handle, a, lda, j, jb, pivots, right_start, n - right_start)?;
        }

        // Step 3: TRSM — solve L[j:j+jb, j:j+jb] * U[j:j+jb, j+jb:n] = A[j:j+jb, j+jb:n].
        if right_start < n {
            let l_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + j as u64 * lda as u64) * T::SIZE as u64,
                jb,
                jb,
                lda,
                Layout::ColMajor,
            );
            let mut u_desc = MatrixDescMut::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + right_start as u64 * lda as u64) * T::SIZE as u64,
                jb,
                n - right_start,
                lda,
                Layout::ColMajor,
            );
            oxicuda_blas::level3::trsm(
                handle.blas(),
                Side::Left,
                FillMode::Lower,
                Transpose::NoTrans,
                DiagType::Unit,
                T::gpu_one(),
                &l_desc,
                &mut u_desc,
            )?;
        }

        // Step 4: GEMM — update trailing matrix:
        // A[j+jb:n, j+jb:n] -= A[j+jb:n, j:j+jb] * A[j:j+jb, j+jb:n]
        let remaining_rows = n.saturating_sub(j + jb);
        let remaining_cols = n.saturating_sub(j + jb);
        if remaining_rows > 0 && remaining_cols > 0 {
            let a21_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + ((j + jb) as u64 + j as u64 * lda as u64) * T::SIZE as u64,
                remaining_rows,
                jb,
                lda,
                Layout::ColMajor,
            );
            let a12_desc = MatrixDesc::<T>::from_raw(
                a.as_device_ptr() + (j as u64 + (j + jb) as u64 * lda as u64) * T::SIZE as u64,
                jb,
                remaining_cols,
                lda,
                Layout::ColMajor,
            );
            let mut a22_desc = MatrixDescMut::<T>::from_raw(
                a.as_device_ptr()
                    + ((j + jb) as u64 + (j + jb) as u64 * lda as u64) * T::SIZE as u64,
                remaining_rows,
                remaining_cols,
                lda,
                Layout::ColMajor,
            );

            // Compute the negative one for alpha.
            let neg_one = T::from_bits_u64({
                let one = T::gpu_one();
                // Negate by XORing the sign bit.
                let bits = one.to_bits_u64();
                if T::SIZE == 4 {
                    bits ^ 0x8000_0000
                } else {
                    bits ^ 0x8000_0000_0000_0000
                }
            });

            oxicuda_blas::level3::gemm_api::gemm(
                handle.blas(),
                Transpose::NoTrans,
                Transpose::NoTrans,
                neg_one,
                &a21_desc,
                &a12_desc,
                T::gpu_one(),
                &mut a22_desc,
            )?;
        }
    }

    Ok(LuResult { info })
}

/// Panel factorization: factorizes columns j..j+jb of A[j:n, j:j+jb].
///
/// This performs unblocked LU within the panel, finding pivots, scaling the
/// column below the pivot, and updating the panel's trailing columns.
///
/// Returns the panel-local info (0 if success, >0 if singular at panel-local column).
fn panel_lu<T: GpuFloat>(
    _handle: &SolverHandle,
    a: &mut DeviceBuffer<T>,
    n: u32,
    lda: u32,
    j: u32,
    jb: u32,
    pivots: &mut DeviceBuffer<i32>,
) -> SolverResult<i32> {
    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_panel_lu::<T>(_handle.sm_version(), jb)?;

    let n_usize = n as usize;
    let lda_usize = lda as usize;
    let j_usize = j as usize;
    let jb_usize = jb as usize;

    let mut a_host = vec![T::gpu_zero(); a.len()];
    a.copy_to_host(&mut a_host)?;

    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    let mut info: i32 = 0;
    let panel_end = (j_usize + jb_usize).min(n_usize);

    for kk in 0..jb_usize {
        let col = j_usize + kk;
        if col >= n_usize {
            break;
        }

        // Pivot search in column `col` over rows col..n-1.
        let mut pivot_row = col;
        let mut max_abs = 0.0_f64;
        for row in col..n_usize {
            let bits = a_host[col * lda_usize + row].to_bits_u64();
            let val = if T::SIZE == 8 {
                f64::from_bits(bits)
            } else {
                f64::from(f32::from_bits(bits as u32))
            };
            let abs = val.abs();
            if abs > max_abs {
                max_abs = abs;
                pivot_row = row;
            }
        }

        piv_host[col] = pivot_row as i32;

        // Swap within panel columns; trailing columns are swapped later.
        if pivot_row != col {
            for c in j_usize..panel_end {
                a_host.swap(c * lda_usize + col, c * lda_usize + pivot_row);
            }
        }

        // Detect singular pivot in the panel (1-based panel-local info).
        let pivot_bits = a_host[col * lda_usize + col].to_bits_u64();
        let pivot_val = if T::SIZE == 8 {
            f64::from_bits(pivot_bits)
        } else {
            f64::from(f32::from_bits(pivot_bits as u32))
        };
        if info == 0 && pivot_val.abs() <= 1e-30 {
            info = (kk + 1) as i32;
            continue;
        }

        // Scale below-diagonal entries in this panel column.
        for row in (col + 1)..n_usize {
            let x_bits = a_host[col * lda_usize + row].to_bits_u64();
            let x = if T::SIZE == 8 {
                f64::from_bits(x_bits)
            } else {
                f64::from(f32::from_bits(x_bits as u32))
            };
            let scaled = x / pivot_val;
            a_host[col * lda_usize + row] = if T::SIZE == 8 {
                T::from_bits_u64(scaled.to_bits())
            } else {
                T::from_bits_u64(u64::from((scaled as f32).to_bits()))
            };
        }

        // Update trailing panel columns.
        for c in (col + 1)..panel_end {
            let uk_bits = a_host[c * lda_usize + col].to_bits_u64();
            let u_kc = if T::SIZE == 8 {
                f64::from_bits(uk_bits)
            } else {
                f64::from(f32::from_bits(uk_bits as u32))
            };
            for row in (col + 1)..n_usize {
                let l_bits = a_host[col * lda_usize + row].to_bits_u64();
                let l_rc = if T::SIZE == 8 {
                    f64::from_bits(l_bits)
                } else {
                    f64::from(f32::from_bits(l_bits as u32))
                };
                let a_bits = a_host[c * lda_usize + row].to_bits_u64();
                let a_rc = if T::SIZE == 8 {
                    f64::from_bits(a_bits)
                } else {
                    f64::from(f32::from_bits(a_bits as u32))
                };
                let updated = a_rc - l_rc * u_kc;
                a_host[c * lda_usize + row] = if T::SIZE == 8 {
                    T::from_bits_u64(updated.to_bits())
                } else {
                    T::from_bits_u64(u64::from((updated as f32).to_bits()))
                };
            }
        }
    }

    a.copy_from_host(&a_host)?;
    pivots.copy_from_host(&piv_host)?;

    Ok(info)
}

/// Applies pivot swaps from panel factorization to columns outside the panel.
///
/// For each pivot in `pivots[j..j+jb]`, swaps rows in the column range
/// `[col_start..col_start+col_count]`.
#[allow(clippy::too_many_arguments)]
fn apply_panel_pivots<T: GpuFloat>(
    _handle: &SolverHandle,
    a: &mut DeviceBuffer<T>,
    lda: u32,
    j: u32,
    jb: u32,
    pivots: &DeviceBuffer<i32>,
    col_start: u32,
    col_count: u32,
) -> SolverResult<()> {
    if col_count == 0 || jb == 0 {
        return Ok(());
    }

    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_pivot_swap::<T>(_handle.sm_version())?;

    let lda_usize = lda as usize;
    let j_usize = j as usize;
    let jb_usize = jb as usize;
    let col_start_usize = col_start as usize;
    let col_end = col_start_usize + col_count as usize;

    let mut a_host = vec![T::gpu_zero(); a.len()];
    a.copy_to_host(&mut a_host)?;
    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    for t in 0..jb_usize {
        let row = j_usize + t;
        if row >= piv_host.len() {
            break;
        }
        let piv = piv_host[row].max(0) as usize;
        if piv >= lda_usize {
            return Err(SolverError::DimensionMismatch(format!(
                "apply_panel_pivots: pivot index out of range ({piv} >= lda {lda_usize})"
            )));
        }
        if piv == row {
            continue;
        }
        for col in col_start_usize..col_end {
            a_host.swap(col * lda_usize + row, col * lda_usize + piv);
        }
    }

    a.copy_from_host(&a_host)?;

    Ok(())
}

/// Applies pivot permutations to the right-hand side B.
fn apply_pivots_to_rhs<T: GpuFloat>(
    _handle: &SolverHandle,
    b: &mut DeviceBuffer<T>,
    pivots: &DeviceBuffer<i32>,
    n: u32,
    nrhs: u32,
) -> SolverResult<()> {
    if n == 0 || nrhs == 0 {
        return Ok(());
    }

    // Keep PTX generation path exercised while host fallback is active.
    let _ = emit_pivot_swap::<T>(_handle.sm_version())?;

    let n_usize = n as usize;
    let nrhs_usize = nrhs as usize;

    let mut b_host = vec![T::gpu_zero(); b.len()];
    b.copy_to_host(&mut b_host)?;
    let mut piv_host = vec![0_i32; pivots.len()];
    pivots.copy_to_host(&mut piv_host)?;

    // Apply all pivots across all RHS columns (column-major, lda = n).
    for row in 0..n_usize {
        if row >= piv_host.len() {
            break;
        }
        let piv = piv_host[row].max(0) as usize;
        if piv >= n_usize {
            return Err(SolverError::DimensionMismatch(format!(
                "apply_pivots_to_rhs: pivot index out of range ({piv} >= n {n_usize})"
            )));
        }
        if piv == row {
            continue;
        }
        for col in 0..nrhs_usize {
            b_host.swap(col * n_usize + row, col * n_usize + piv);
        }
    }

    b.copy_from_host(&b_host)?;

    Ok(())
}

// ---------------------------------------------------------------------------
// PTX kernel generation
// ---------------------------------------------------------------------------

fn panel_lu_name<T: GpuFloat>(block_size: u32) -> String {
    format!("solver_panel_lu_{}_{}", T::NAME, block_size)
}

fn pivot_swap_name<T: GpuFloat>() -> String {
    format!("solver_pivot_swap_{}", T::NAME)
}

/// Emits PTX for a single-CTA panel LU factorization kernel.
///
/// The kernel factorizes a `panel_rows x panel_cols` submatrix in shared memory.
/// Each column is processed sequentially: find pivot (max abs), swap rows,
/// scale below-diagonal elements, and update trailing columns.
fn emit_panel_lu<T: GpuFloat>(sm: SmVersion, panel_cols: u32) -> SolverResult<String> {
    let name = panel_lu_name::<T>(panel_cols);
    let float_ty = T::PTX_TYPE;

    let ptx = KernelBuilder::new(&name)
        .target(sm)
        .max_threads_per_block(SOLVER_BLOCK_SIZE)
        .param("panel_ptr", PtxType::U64)
        .param("pivots_ptr", PtxType::U64)
        .param("panel_rows", PtxType::U32)
        .param("panel_cols", PtxType::U32)
        .param("lda", PtxType::U32)
        .body(move |b| {
            let tid = b.thread_id_x();
            let panel_rows_reg = b.load_param_u32("panel_rows");
            let panel_cols_reg = b.load_param_u32("panel_cols");
            let lda_reg = b.load_param_u32("lda");
            let panel_ptr = b.load_param_u64("panel_ptr");

            // Each thread handles elements in the column below the diagonal.
            // This is a simplified single-CTA panel factorization.
            // For each column k = 0..panel_cols:
            //   1. Find pivot (thread 0 finds max abs in column k, rows k..panel_rows)
            //   2. Swap pivot row with row k
            //   3. Scale elements below diagonal: A[i,k] /= A[k,k] for i > k
            //   4. Update trailing: A[i,j] -= A[i,k] * A[k,j] for i > k, j > k

            // The kernel processes panel_cols columns sequentially.
            // Each column step uses all threads in the CTA cooperatively.
            let _ = (
                tid,
                panel_rows_reg,
                panel_cols_reg,
                lda_reg,
                panel_ptr,
                float_ty,
            );

            b.ret();
        })
        .build()?;

    Ok(ptx)
}

/// Emits PTX for a row-permutation kernel.
///
/// Each thread handles one column: for each pivot in `pivots[j..j+jb]`,
/// swaps rows in columns `col_start..col_start+col_count`.
fn emit_pivot_swap<T: GpuFloat>(sm: SmVersion) -> SolverResult<String> {
    let name = pivot_swap_name::<T>();
    let float_ty = T::PTX_TYPE;

    let ptx = KernelBuilder::new(&name)
        .target(sm)
        .max_threads_per_block(SOLVER_BLOCK_SIZE)
        .param("a_ptr", PtxType::U64)
        .param("pivots_ptr", PtxType::U64)
        .param("j", PtxType::U32)
        .param("jb", PtxType::U32)
        .param("col_start", PtxType::U32)
        .param("col_count", PtxType::U32)
        .param("lda", PtxType::U32)
        .body(move |b| {
            let gid = b.global_thread_id_x();
            let col_count_reg = b.load_param_u32("col_count");

            b.if_lt_u32(gid.clone(), col_count_reg, |b| {
                let a_ptr = b.load_param_u64("a_ptr");
                let col_start = b.load_param_u32("col_start");
                let lda = b.load_param_u32("lda");

                // Compute the actual column index.
                let col_idx = b.add_u32(gid, col_start);

                // Column base address: a_ptr + col_idx * lda * sizeof(T)
                let col_elem_offset = b.mul_lo_u32(col_idx, lda);
                let _col_base = b.byte_offset_addr(a_ptr, col_elem_offset, T::size_u32());

                // In the full implementation, this would loop over pivots[j..j+jb]
                // and swap the corresponding rows.
                let _ = float_ty;
            });

            b.ret();
        })
        .build()?;

    Ok(ptx)
}

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

    // ---------------------------------------------------------------------------
    // CPU reference helpers for LU integration tests
    // ---------------------------------------------------------------------------

    /// Doolittle LU factorization (no pivoting) on a 4×4 f64 matrix.
    ///
    /// Returns (L, U) where L is unit lower triangular and U is upper triangular,
    /// such that A = L * U.
    fn doolittle_lu_4x4(a: &[[f64; 4]; 4]) -> ([[f64; 4]; 4], [[f64; 4]; 4]) {
        let mut l = [[0.0_f64; 4]; 4];
        let mut u = [[0.0_f64; 4]; 4];

        for i in 0..4 {
            l[i][i] = 1.0; // Unit diagonal for L.

            // U row i.
            for j in i..4 {
                let sum: f64 = (0..i).map(|k| l[i][k] * u[k][j]).sum();
                u[i][j] = a[i][j] - sum;
            }

            // L column i (below diagonal).
            for j in (i + 1)..4 {
                let sum: f64 = (0..i).map(|k| l[j][k] * u[k][i]).sum();
                if u[i][i].abs() > 1e-15 {
                    l[j][i] = (a[j][i] - sum) / u[i][i];
                }
            }
        }

        (l, u)
    }

    /// 4×4 matrix multiply (row-major).
    fn matmul_4x4(a: &[[f64; 4]; 4], b: &[[f64; 4]; 4]) -> [[f64; 4]; 4] {
        let mut c = [[0.0_f64; 4]; 4];
        for i in 0..4 {
            for j in 0..4 {
                for k in 0..4 {
                    c[i][j] += a[i][k] * b[k][j];
                }
            }
        }
        c
    }

    // ---------------------------------------------------------------------------
    // LU + GEMM/TRSM integration tests
    // ---------------------------------------------------------------------------

    #[test]
    fn lu_trsm_trailing_update() {
        // Verify Doolittle LU on a 4×4 matrix: A = L * U to tolerance 1e-10.
        let a = [
            [4.0_f64, 3.0, 2.0, 1.0],
            [2.0, 5.0, 3.0, 2.0],
            [1.0, 2.0, 6.0, 3.0],
            [1.0, 1.0, 2.0, 7.0],
        ];
        let (l, u) = doolittle_lu_4x4(&a);

        // L must be unit lower triangular.
        for (i, l_row) in l.iter().enumerate() {
            assert!(
                (l_row[i] - 1.0).abs() < 1e-15,
                "L[{i},{i}] must be 1.0 (unit diagonal)"
            );
            for (j, &val) in l_row.iter().enumerate().filter(|(j, _)| *j > i) {
                assert!(
                    val.abs() < 1e-15,
                    "L[{i},{j}] = {val} must be 0.0 (upper triangle)",
                );
            }
        }

        // U must be upper triangular.
        for (i, u_row) in u.iter().enumerate() {
            for (j, &val) in u_row.iter().enumerate().filter(|(j, _)| *j < i) {
                assert!(
                    val.abs() < 1e-15,
                    "U[{i},{j}] = {val} must be 0.0 (lower triangle)",
                );
            }
        }

        // Reconstruct: L*U must equal A.
        let reconstructed = matmul_4x4(&l, &u);
        for i in 0..4 {
            for j in 0..4 {
                assert!(
                    (reconstructed[i][j] - a[i][j]).abs() < 1e-10,
                    "LU[{i},{j}] = {} ≠ A[{i},{j}] = {} (diff = {})",
                    reconstructed[i][j],
                    a[i][j],
                    (reconstructed[i][j] - a[i][j]).abs()
                );
            }
        }
    }

    #[test]
    fn lu_gemm_rank_update_correctness() {
        // Verify that the GEMM trailing update for k=0 is correct on a 3×3 example.
        //
        // After the first column of LU (k=0):
        //   L[:,0] is computed, U[0,:] is computed.
        //   Trailing update: A[1:3, 1:3] -= L[1:3, 0:1] * U[0:1, 1:3]
        //
        // Use a = [[2, 4, 6], [1, 3, 5], [1, 2, 4]] (simple example).
        let a = [[2.0_f64, 4.0, 6.0], [1.0, 3.0, 5.0], [1.0, 2.0, 4.0]];

        // After first pivot (k=0), L column 0 = [1, a[1,0]/a[0,0], a[2,0]/a[0,0]]
        //                                      = [1, 0.5, 0.5]
        // U row 0 = a[0,:] = [2, 4, 6]
        // Trailing update for A[1:3, 1:3]:
        //   A[1,1] -= L[1,0]*U[0,1] = 3 - 0.5*4 = 1
        //   A[1,2] -= L[1,0]*U[0,2] = 5 - 0.5*6 = 2
        //   A[2,1] -= L[2,0]*U[0,1] = 2 - 0.5*4 = 0
        //   A[2,2] -= L[2,0]*U[0,2] = 4 - 0.5*6 = 1
        let l_col0 = [1.0_f64, a[1][0] / a[0][0], a[2][0] / a[0][0]];
        let u_row0 = [a[0][0], a[0][1], a[0][2]];

        // Trailing submatrix after k=0 update.
        let mut trailing = [[0.0_f64; 2]; 2];
        for i in 0..2 {
            for j in 0..2 {
                trailing[i][j] = a[i + 1][j + 1] - l_col0[i + 1] * u_row0[j + 1];
            }
        }

        assert!(
            (trailing[0][0] - 1.0).abs() < 1e-12,
            "trailing[0,0] should be 1"
        );
        assert!(
            (trailing[0][1] - 2.0).abs() < 1e-12,
            "trailing[0,1] should be 2"
        );
        assert!(trailing[1][0].abs() < 1e-12, "trailing[1,0] should be 0");
        assert!(
            (trailing[1][1] - 1.0).abs() < 1e-12,
            "trailing[1,1] should be 1"
        );
    }

    #[test]
    fn lu_block_size_positive() {
        let block_size = LU_BLOCK_SIZE;
        assert!(block_size > 0);
        assert!(block_size <= 256);
    }

    #[test]
    fn lu_result_info() {
        let result = LuResult { info: 0 };
        assert_eq!(result.info, 0);

        let singular = LuResult { info: 3 };
        assert!(singular.info > 0);
    }

    #[test]
    fn panel_lu_name_format() {
        let name = panel_lu_name::<f32>(64);
        assert!(name.contains("f32"));
        assert!(name.contains("64"));
    }

    #[test]
    fn pivot_swap_name_format() {
        let name = pivot_swap_name::<f64>();
        assert!(name.contains("f64"));
    }

    #[test]
    fn neg_one_f32() {
        let neg = f32::from_bits_u64(f32::gpu_one().to_bits_u64() ^ 0x8000_0000);
        assert!((neg + 1.0).abs() < 1e-10);
    }

    #[test]
    fn neg_one_f64() {
        let neg = f64::from_bits_u64(f64::gpu_one().to_bits_u64() ^ 0x8000_0000_0000_0000);
        assert!((neg + 1.0).abs() < 1e-15);
    }
}