scirs2-datasets 0.6.3

#![cfg(feature = "cuda")]
//! Optional, off-by-default, NVIDIA-only CUDA acceleration for the dense
//! regression-target matmul at the heart of synthetic dataset generation,
//! calling the pure-Rust `oxicuda-*` library crates directly.
//!
//! This module is ADDITIVE and entirely feature-gated behind the `cuda` feature
//! (off by default). It does not replace or modify the existing GPU subsystem
//! (`src/gpu.rs`, `src/gpu_optimization.rs`) or the wgpu-backed `wgpu` path.
//! Every public function is `f64`-native — there is no silent `f32` downcast at
//! the boundary, which is the advantage of the oxicuda CUDA path over the f32
//! wgpu path.
//!
//! Each entry point degrades safely when no NVIDIA device is present:
//! [`cuda_is_available`](crate::gpu_cuda::cuda_is_available) never panics, and
//! the compute functions return a
//! [`crate::error::DatasetsError::ComputationError`] rather than aborting.
//!
//! # Library-call pattern
//!
//! The regression-target generator at the heart of `make_regression` forms the
//! dense product `y = X · coef` (`[n_samples × n_features] × [n_features] →
//! [n_samples]`), an `O(n_samples · n_features)` matrix×vector multiply. This
//! module hands that product off to the pre-built, battle-tested `oxicuda-blas`
//! GEMM rather than authoring any bespoke kernels, mirroring the proven
//! matrix×vector layout of the Phase-1 interpolate reference
//! (`scirs2-interpolate/src/gpu_cuda.rs::cuda_eval_gemm`) and the optimize
//! Phase-2 reference (`scirs2-optimize/src/gpu_cuda.rs`).
//!
//! # Scoping note
//!
//! This is a standalone additive function, deliberately NOT wired into
//! `make_regression` / `generators/basic.rs` — exactly like the Phase-1 crates
//! kept their `cuda_*` functions standalone.

use crate::error::{DatasetsError, Result};
use oxicuda_blas::types::{Layout, MatrixDesc, MatrixDescMut, Transpose};
use scirs2_core::ndarray::{Array1, ArrayView1, ArrayView2};

/// Returns `true` iff the CUDA driver initializes and at least one NVIDIA device
/// is visible. Never panics.
///
/// On a non-NVIDIA platform such as macOS the CUDA driver fails to initialise
/// (or reports zero devices) and this returns `false`, allowing callers to fall
/// back to the CPU path.
pub fn cuda_is_available() -> bool {
    oxicuda_driver::init().is_ok()
        && oxicuda_driver::device::Device::count()
            .map(|c| c > 0)
            .unwrap_or(false)
}

/// Maps an `oxicuda-blas` error onto the crate's honest `ComputationError`.
fn blas_err(e: oxicuda_blas::error::BlasError) -> DatasetsError {
    DatasetsError::ComputationError(format!("oxicuda-blas: {e}"))
}

/// Maps an `oxicuda` CUDA driver/memory error onto `ComputationError`.
fn cuda_err(e: oxicuda_driver::CudaError) -> DatasetsError {
    DatasetsError::ComputationError(format!("oxicuda CUDA driver: {e}"))
}

/// Initialises the CUDA driver, selects device 0, and builds a shared context
/// wrapped in an `Arc` as required by the oxicuda BLAS handle constructor.
///
/// Returns a [`ComputationError`](DatasetsError::ComputationError) — never a
/// fabricated success — when CUDA is unavailable or no device is present.
fn build_context() -> Result<std::sync::Arc<oxicuda_driver::Context>> {
    oxicuda_driver::init()
        .map_err(|e| DatasetsError::ComputationError(format!("CUDA unavailable: {e}")))?;
    let count = oxicuda_driver::device::Device::count()
        .map_err(|e| DatasetsError::ComputationError(format!("device count: {e}")))?;
    if count <= 0 {
        return Err(DatasetsError::ComputationError(
            "no NVIDIA CUDA device available".into(),
        ));
    }
    let dev = oxicuda_driver::device::Device::get(0).map_err(cuda_err)?;
    Ok(std::sync::Arc::new(
        oxicuda_driver::Context::new(&dev).map_err(cuda_err)?,
    ))
}

/// Compute the regression target `y = X · coef` on the GPU in `f64`.
///
/// This is the dense matmul at the heart of `make_regression`: `X` is the
/// `[n_samples × n_features]` design matrix and `coef` is the length-`n_features`
/// coefficient vector; the returned target vector has length `n_samples`.
///
/// Layout correctness (cannot be runtime-checked on a non-NVIDIA host): ndarray
/// is row-major. We materialize a contiguous row-major copy of `X` via
/// [`as_standard_layout`](scirs2_core::ndarray::ArrayRef::as_standard_layout)
/// and describe it as a [`Layout::RowMajor`] `MatrixDesc`
/// (`n_samples`×`n_features`). Mirroring the proven optimize/interpolate
/// matrix×vector layout, `coef` is described as a column-vector
/// [`Layout::RowMajor`] `MatrixDesc` (`n_features`×1) and the output as a
/// [`Layout::RowMajor`] `MatrixDescMut` (`n_samples`×1) — a single-column
/// matrix is byte-identical under either layout tag, and `oxicuda-blas`'s
/// GEMM dispatcher requires RowMajor on every operand. Under
/// [`Transpose::NoTrans`] (`alpha = 1.0`, `beta = 0.0`) this yields
/// `M = n_samples`, `K = n_features`, `N = 1`.
pub fn cuda_regression_target(x: &ArrayView2<f64>, coef: &ArrayView1<f64>) -> Result<Array1<f64>> {
    let n_samples = x.nrows();
    let n_features = x.ncols();
    if coef.len() != n_features {
        return Err(DatasetsError::ValidationError(format!(
            "cuda_regression_target: coef length {} does not match X column count {n_features}",
            coef.len()
        )));
    }
    if n_samples == 0 {
        return Ok(Array1::zeros(0));
    }
    if n_features == 0 {
        // No features => the linear combination is empty, so every target is 0.
        return Ok(Array1::zeros(n_samples));
    }

    // Contiguous row-major copy so `.as_slice()` is `Some` and the device upload
    // matches the `Layout::RowMajor` descriptor below.
    let x_std = x.as_standard_layout();
    let x_slice = x_std.as_slice().ok_or_else(|| {
        DatasetsError::ComputationError("cuda_regression_target: X not contiguous".into())
    })?;
    let coef_vec: Vec<f64> = coef.iter().copied().collect();

    let ctx = build_context()?;
    let handle = oxicuda_blas::BlasHandle::new(&ctx).map_err(blas_err)?;

    let d_x = oxicuda_memory::DeviceBuffer::from_host(x_slice).map_err(cuda_err)?;
    let d_coef = oxicuda_memory::DeviceBuffer::from_host(&coef_vec).map_err(cuda_err)?;
    let mut d_y = oxicuda_memory::DeviceBuffer::<f64>::alloc(n_samples).map_err(cuda_err)?;

    // X as RowMajor (n_samples×n_features), coef as a RowMajor column
    // (n_features×1), output RowMajor (n_samples×1). A single-column matrix
    // has one possible packed layout regardless of the tag, but oxicuda-blas's
    // GEMM dispatcher hard-rejects any non-RowMajor descriptor, so all three
    // must be tagged RowMajor.
    let a_desc =
        MatrixDesc::from_buffer(&d_x, n_samples as u32, n_features as u32, Layout::RowMajor)
            .map_err(blas_err)?;
    let b_desc = MatrixDesc::from_buffer(&d_coef, n_features as u32, 1, Layout::RowMajor)
        .map_err(blas_err)?;
    let mut c_desc = MatrixDescMut::from_buffer(&mut d_y, n_samples as u32, 1, Layout::RowMajor)
        .map_err(blas_err)?;

    oxicuda_blas::level3::gemm_api::gemm::<f64>(
        &handle,
        Transpose::NoTrans,
        Transpose::NoTrans,
        1.0,
        &a_desc,
        &b_desc,
        0.0,
        &mut c_desc,
    )
    .map_err(blas_err)?;

    // The GEMM is enqueued on the BLAS handle's `CU_STREAM_NON_BLOCKING` stream,
    // whereas `copy_to_host` issues `cuMemcpyDtoH_v2` on the legacy default
    // stream. A non-blocking stream does NOT implicitly synchronise with the
    // default stream, so without an explicit wait the device-to-host copy races
    // the kernel and reads the (uninitialised / zero) output buffer before the
    // GEMM has finished writing it — worse the longer the kernel runs (large K).
    // Synchronise the compute stream so the product is fully materialised first.
    handle.stream().synchronize().map_err(cuda_err)?;

    let mut y = vec![0.0f64; n_samples];
    d_y.copy_to_host(&mut y).map_err(cuda_err)?;
    Ok(Array1::from_vec(y))
}

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

    fn mat(rows: usize, cols: usize, data: Vec<f64>) -> Array2<f64> {
        Array2::from_shape_vec((rows, cols), data).expect("valid test matrix shape")
    }

    #[test]
    fn cuda_regression_target_or_skip() {
        if !cuda_is_available() {
            eprintln!("skipping: no NVIDIA CUDA device");
            assert!(!cuda_is_available());
            return;
        }
        // X (3x2) row-major; coef [1, 1] => row sums [3, 7, 11].
        let x = mat(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
        let coef = Array1::from_vec(vec![1.0, 1.0]);
        let y =
            cuda_regression_target(&x.view(), &coef.view()).expect("cuda_regression_target failed");
        let expected = x.dot(&coef);
        let max_diff = y
            .iter()
            .zip(expected.iter())
            .map(|(g, e)| (g - e).abs())
            .fold(0.0f64, f64::max);
        assert!(max_diff < 1e-9, "max abs diff {max_diff} exceeds 1e-9");
    }

    /// Regression guard for the stream-synchronisation race: a GEMM whose
    /// runtime is long enough (large K) that a `copy_to_host` on the default
    /// stream, issued without waiting on the non-blocking compute stream, would
    /// read the output buffer before the kernel finished writing it. Prior to
    /// the explicit `handle.stream().synchronize()` in `cuda_regression_target`
    /// this returned nondeterministic all-zero output for K >~ 100.
    #[test]
    fn cuda_regression_target_large_k_no_stream_race() {
        if !cuda_is_available() {
            eprintln!("skipping: no NVIDIA CUDA device");
            return;
        }
        // X[i,j] = (i+1), coef[j] = 1 => expected[i] = (i+1)*K, so any all-zero
        // row is unambiguously a missed write rather than a legitimate value.
        let m = 2000usize;
        let k = 600usize;
        let mut data = Vec::with_capacity(m * k);
        for i in 0..m {
            for _j in 0..k {
                data.push((i + 1) as f64);
            }
        }
        let x = mat(m, k, data);
        let coef = Array1::from_elem(k, 1.0);
        // Repeat: the race was intermittent, so a single pass could pass by luck.
        for rep in 0..5 {
            let y = cuda_regression_target(&x.view(), &coef.view()).expect("gpu call");
            let expected = x.dot(&coef);
            let max_diff = y
                .iter()
                .zip(expected.iter())
                .map(|(g, e)| (g - e).abs())
                .fold(0.0f64, f64::max);
            assert!(
                max_diff < 1e-6,
                "rep {rep}: GPU output diverged from CPU matmul (max_diff={max_diff:.3e}); \
                 stream-synchronisation race likely regressed"
            );
        }
    }

    #[test]
    fn cuda_regression_target_shape_mismatch_errors() {
        // Runs without a GPU: validation rejects before any device call.
        // X has 2 columns but coef has length 3.
        let x = mat(2, 2, vec![1.0, 0.0, 0.0, 1.0]);
        let coef_bad = Array1::from_vec(vec![1.0, 2.0, 3.0]);
        assert!(cuda_regression_target(&x.view(), &coef_bad.view()).is_err());
    }
}