rustytorch/crates/specialized/rtx-cfd/tests/embedded3_gpu_vcycle.rs

//! embedded3 gate 2: the device seven-point V-cycle against the host f32
//! reference. The cheap pin runs on a small channel with a z-cylinder hole;
//! `bench_anchor_size` (ignored) is 378 × 62 × 62 and prints ms per V-cycle.
//!
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test embedded3_gpu_vcycle -- --nocapture [--ignored]`
#![cfg(feature = "cuda")]

use rtx_cfd::solvers::incompressible::embedded3::poisson::device::DeviceVcycle;
use rtx_cfd::solvers::incompressible::embedded3::poisson::{
    Problem, export_hierarchy, vcycle_f32_reference,
};
use rtx_cfd::solvers::incompressible::{MgSmoother, MultigridParameters};
use std::time::Instant;

fn channel(nx: usize, ny: usize, nz: usize, periodic_z: bool, seed: u64) -> Problem {
    let mut p = Problem::new(nx, ny, nz);
    p.periodic_z = periodic_z;
    let h = 0.41 / ny as f64;
    let a = 3.24e-4 * h;
    let hole = |i: usize, j: usize| {
        let (x, y) = ((i as f64 + 0.5) * h, (j as f64 + 0.5) * h);
        (x - 0.2).powi(2) + (y - 0.2).powi(2) < 0.05 * 0.05
    };
    for k in 0..nz {
        for j in 0..ny {
            for i in 0..nx {
                let idx = p.index(k, j, i);
                if hole(i, j) {
                    p.active[idx] = false;
                    continue;
                }
                if i + 1 < nx && !hole(i + 1, j) {
                    p.ae[idx] = a;
                }
                if i > 0 && !hole(i - 1, j) {
                    p.aw[idx] = a;
                }
                if j + 1 < ny && !hole(i, j + 1) {
                    p.an[idx] = a;
                }
                if j > 0 && !hole(i, j - 1) {
                    p.as_[idx] = a;
                }
                if k + 1 < nz || periodic_z {
                    p.at[idx] = a;
                }
                if k > 0 || periodic_z {
                    p.ab[idx] = a;
                }
                if i + 1 == nx {
                    p.extra_diag[idx] = 2.0 * a;
                }
            }
        }
    }
    let mut state = seed | 1;
    for idx in 0..nx * ny * nz {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        p.rhs[idx] = if p.active[idx] {
            1e-3 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
        } else {
            0.0
        };
    }
    p
}

fn compare(p: &Problem, label: &str) -> (f64, f64) {
    let params = MultigridParameters {
        smoother: MgSmoother::RedBlack,
        ..MultigridParameters::default()
    };
    let n = p.nx * p.ny * p.nz;
    let mut z_ref = vec![0.0; n];
    vcycle_f32_reference(p, &params, &p.rhs, &mut z_ref);
    let levels = export_hierarchy(p, &params);
    let mut dev = DeviceVcycle::new(&levels, params.smoother_sweeps.max(1));
    let mut z = vec![0.0; n];
    dev.apply(&p.rhs, &mut z);
    let scale = z_ref.iter().fold(0.0_f64, |m, v| m.max(v.abs()));
    let worst = z
        .iter()
        .zip(&z_ref)
        .fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
    let reps = 20;
    dev.apply(&p.rhs, &mut z);
    let t0 = Instant::now();
    for _ in 0..reps {
        dev.apply(&p.rhs, &mut z);
    }
    let ms = t0.elapsed().as_secs_f64() * 1e3 / reps as f64;
    println!(
        "  {label}: {} levels {:?}; device vs host f32 max |Δz| {worst:.3e} on {scale:.3e}; {ms:.3} ms per V-cycle incl. transfers ({n} cells)",
        levels.len(),
        levels.iter().map(|l| l.cells.len()).collect::<Vec<_>>()
    );
    (worst, scale)
}

#[test]
fn device_vcycle_matches_the_host_reference() {
    for (nz, periodic) in [(1usize, false), (8, true), (8, false)] {
        let (worst, scale) = compare(
            &channel(96, 40, nz, periodic, 3),
            &format!("96×40×{nz} periodic {periodic}"),
        );
        assert!(worst < 1e-4 * scale, "{worst:.3e} of {scale:.3e}");
    }
}

#[test]
#[ignore = "the anchor-size bench (378 × 62 × 62): prints ms per V-cycle"]
fn bench_anchor_size() {
    let (worst, scale) = compare(&channel(378, 62, 62, false, 11), "378×62×62 walls");
    assert!(worst < 1e-4 * scale);
}