foxhunt/crates/ml/examples/cuda_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Simple CUDA functionality test to verify compatibility.
//!
//! Uses `MlDevice` and `cudarc` directly — candle has been eliminated.
//! Verifies CUDA context creation, device memory allocation, and
//! host-to-device transfers.
#![allow(unused_crate_dependencies)]

use ml_core::device::MlDevice;

/// Test basic CUDA device creation and memory operations via cudarc.
pub fn test_cuda_basic() -> Result<(), Box<dyn std::error::Error>> {
    println!("Testing CUDA compatibility via MlDevice + cudarc...");

    match MlDevice::cuda(0) {
        Ok(device) => {
            println!("[OK] CUDA device 0 available ({})", device);

            let stream = device
                .cuda_stream()
                .map_err(|e| format!("Failed to get CUDA stream: {e}"))?;

            // Allocate a zeroed GPU buffer (4x4 f32)
            let gpu_zeros: cudarc::driver::CudaSlice<f32> = stream
                .alloc_zeros(16)
                .map_err(|e| format!("alloc_zeros failed: {e}"))?;
            println!("[OK] Allocated 4x4 f32 zero buffer on GPU ({} elements)", 16);

            // Copy host data to GPU: alloc then htod
            let host_data: Vec<f32> = (0..16).map(|i| i as f32).collect();
            let mut gpu_buf: cudarc::driver::CudaSlice<f32> = stream
                .alloc_zeros(16)
                .map_err(|e| format!("alloc for htod failed: {e}"))?;
            stream
                .memcpy_htod(&host_data, &mut gpu_buf)
                .map_err(|e| format!("memcpy_htod failed: {e}"))?;
            println!("[OK] Copied 16 f32 values host -> device");

            // Read back to verify round-trip
            let mut readback = vec![0.0_f32; 16];
            stream
                .memcpy_dtoh(&gpu_buf, &mut readback)
                .map_err(|e| format!("memcpy_dtoh failed: {e}"))?;

            // Verify first and last elements
            let first = readback.first().copied().unwrap_or(f32::NAN);
            let last = readback.last().copied().unwrap_or(f32::NAN);
            if (first - 0.0).abs() > 1e-6 || (last - 15.0).abs() > 1e-6 {
                return Err(format!(
                    "Round-trip mismatch: first={first}, last={last} (expected 0.0, 15.0)"
                )
                .into());
            }
            println!("[OK] Device -> host round-trip verified (first={first}, last={last})");

            // Verify the zeroed buffer reads back as zeros
            let mut zero_readback = vec![1.0_f32; 16];
            stream
                .memcpy_dtoh(&gpu_zeros, &mut zero_readback)
                .map_err(|e| format!("memcpy_dtoh zeros failed: {e}"))?;
            let all_zero = zero_readback.iter().all(|&v| v == 0.0);
            if !all_zero {
                return Err("Zero buffer contained non-zero values".into());
            }
            println!("[OK] Zero-initialized buffer verified");

            Ok(())
        }
        Err(e) => {
            println!("[WARN] CUDA device not available: {}", e);
            println!("This is expected if no GPU is present, but CUDA compilation succeeded");
            Ok(())
        }
    }
}

/// Test that MlDevice reports correct device properties.
pub fn test_device_properties() -> Result<(), Box<dyn std::error::Error>> {
    println!("Testing MlDevice properties...");

    let cpu = MlDevice::Cpu;
    assert!(!cpu.is_cuda());
    assert!(cpu.is_cpu());
    println!("[OK] CPU device: is_cuda=false, is_cpu=true");

    if let Ok(gpu) = MlDevice::cuda(0) {
        assert!(gpu.is_cuda());
        assert!(!gpu.is_cpu());
        println!("[OK] CUDA device: is_cuda=true, is_cpu=false");

        // Verify stream and context are accessible
        let _stream = gpu
            .cuda_stream()
            .map_err(|e| format!("cuda_stream() failed: {e}"))?;
        let _ctx = gpu
            .cuda_context()
            .map_err(|e| format!("cuda_context() failed: {e}"))?;
        println!("[OK] CUDA stream and context accessible");
    } else {
        println!("[WARN] CUDA device properties test skipped (no GPU)");
    }

    Ok(())
}

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

    #[test]
    fn verify_cuda_compilation() {
        // This test just verifies that CUDA code compiles
        // The actual runtime test is optional since CI may not have GPU
        println!("CUDA compilation test passed!");

        // Try to run basic test but don't fail if no GPU
        let _ = test_cuda_basic();
        let _ = test_device_properties();
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    test_cuda_basic()?;
    test_device_properties()?;
    println!("[DONE] CUDA compatibility verification complete!");
    Ok(())
}