Files
rustytorch/crates/specialized/rtx-cfd/tests/embedded3_device_step.rs
T
Omar SobhandClaude Fable 5.1 54911b4db3
CI / Build (macos-latest) (push) Waiting to run
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / Format Check (push) Failing after 4s
CI / CI Success (push) Blocked by required conditions
CI / Build (ubuntu-latest) (push) Failing after 3s
CI / Clippy Check (push) Failing after 4s
CI / Build CPU-Only (Explicit) (push) Failing after 3s
Documentation / Build API Documentation (push) Failing after 2s
Documentation / Build User Guide (push) Successful in 4s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m4s
rtx-cfd embedded3 item 7: e3_step.cu + step::device::{DeviceStep, StepTimers} on the shared runtime (FMA off); gate 7 HELD: device = host ≤ 7e-12 under tight tolerances on MMS/Beltrami/Poiseuille, equal CG counts on every step, periodic planes within 2e-16; the default-tolerance differences are the projection's inner stop
Co-Authored-By: Claude Fable 5.1 <[email protected]>
2026-09-17 14:56:46 -05:00

546 lines
18 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
//! embedded3 gate 7 — the step on the DEVICE: the device-resident step against
//! the host step from the same start — the manufactured problem (upwind
//! and TVD), Beltrami with time-dependent boundary data, and Poiseuille
//! at nz = 1 and on the periodic extrusion (plane agreement). The
//! predictors are compiled without FMA contraction, so the only difference
//! from the host is the CG's reduction order: agreement to 1e-12 relative.
//!
//! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --test embedded3_device_step -- --nocapture`
#![cfg(feature = "cuda")]
use rtx_cfd::solvers::incompressible::embedded3::step::device::DeviceStep;
use rtx_cfd::solvers::incompressible::embedded3::{
Boundaries, Field, Fluid, Grid, Parameters, Side, Solver,
};
use rtx_cfd::solvers::incompressible::{ConvectionScheme, MgSmoother};
use std::f64::consts::PI;
const RHO: f64 = 1.0;
fn fluid(mu: f64) -> Fluid {
Fluid {
density: RHO,
viscosity: mu,
reference_velocity: 1.0,
reference_length: 1.0,
}
}
/// `tight` = the identity configuration: inner stop 1e-6 of the source
/// scale AND mass tolerance 1e-12 (its 0.1× floor on the inner stop is
/// what dominates near a steady state); otherwise the defaults.
fn params(scheme: ConvectionScheme, z: Side, tight: bool) -> Parameters {
Parameters {
corrector_steps: 2,
tolerance: if tight { 1e-12 } else { 1e-8 },
inner_stop_factor: if tight { 1e-6 } else { 1e-2 },
boundaries: Boundaries {
z0: z,
z1: z,
..Boundaries::default()
},
// The device V-cycle is red-black; the host uses the same so the
// preconditioners match.
poisson_smoother: MgSmoother::RedBlack,
convection_scheme: scheme,
..Parameters::default()
}
}
/// Max |Δ| over u, v, w between two fields on the velocity scale, and max
/// |Δp| on the pressure scale floored at the dynamic pressure `ρ U²` (a
/// flat pressure field must not inflate a rounding-level difference).
fn compare(a: &Field, b: &Field) -> (f64, f64) {
let mut worst = 0.0_f64;
let mut scale = 0.0_f64;
for (x, y) in
a.u.iter()
.zip(&b.u)
.chain(a.v.iter().zip(&b.v))
.chain(a.w.iter().zip(&b.w))
{
worst = worst.max((x - y).abs());
scale = scale.max(x.abs());
}
let mut worst_p = 0.0_f64;
let mut scale_p = 0.0_f64;
for (x, y) in a.p.iter().zip(&b.p) {
worst_p = worst_p.max((x - y).abs());
scale_p = scale_p.max(x.abs());
}
let scale_p = scale_p.max(RHO * scale * scale);
// One number: the larger of the two relative differences, on the velocity scale.
(worst.max(worst_p / scale_p * scale), scale)
}
// ---- the manufactured solution of three_d_mms.rs ----
fn u3(x: f64, y: f64, z: f64) -> f64 {
(PI * x).sin() * (PI * y).cos() * (PI * z).cos()
}
fn v3(x: f64, y: f64, z: f64) -> f64 {
(PI * x).cos() * (PI * y).sin() * (PI * z).cos()
}
fn w3(x: f64, y: f64, z: f64) -> f64 {
-2.0 * (PI * x).cos() * (PI * y).cos() * (PI * z).sin()
}
fn source3(mu: f64, x: f64, y: f64, z: f64) -> (f64, f64, f64) {
let (sx, cx) = (PI * x).sin_cos();
let (sy, cy) = (PI * y).sin_cos();
let (sz, cz) = (PI * z).sin_cos();
let (u, v, w) = (u3(x, y, z), v3(x, y, z), w3(x, y, z));
let (ux, uy, uz) = (PI * cx * cy * cz, -PI * sx * sy * cz, -PI * sx * cy * sz);
let (vx, vy, vz) = (-PI * sx * sy * cz, PI * cx * cy * cz, -PI * cx * sy * sz);
let (wx, wy, wz) = (
2.0 * PI * sx * cy * sz,
2.0 * PI * cx * sy * sz,
-2.0 * PI * cx * cy * cz,
);
let (px, py, pz) = (PI * cx * sy * sz, PI * sx * cy * sz, PI * sx * sy * cz);
let lap = -3.0 * PI * PI;
(
RHO * (u * ux + v * uy + w * uz) + px - mu * lap * u,
RHO * (u * vx + v * vy + w * vz) + py - mu * lap * v,
RHO * (u * wx + v * wy + w * wz) + pz - mu * lap * w,
)
}
fn boundary3(x: f64, y: f64, z: f64) -> (f64, f64, f64) {
let u = if x <= 0.0 || x >= 1.0 {
0.0
} else {
u3(x, y, z)
};
let v = if y <= 0.0 || y >= 1.0 {
0.0
} else {
v3(x, y, z)
};
let w = if z <= 0.0 || z >= 1.0 {
0.0
} else {
w3(x, y, z)
};
(u, v, w)
}
fn mms_pair(n: usize, scheme: ConvectionScheme, isf: bool) -> (Solver, Solver, Grid, f64) {
let mu = 0.05;
let h = 1.0 / n as f64;
let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h);
let mk = || {
let mut s = Solver::new(fluid(mu), params(scheme, Side::Velocity, isf));
s.set_momentum_source(move |x, y, z, _t| source3(mu, x, y, z));
s.set_boundary_velocity(|x, y, z, _t| boundary3(x, y, z));
s
};
(
mk(),
mk(),
Grid {
nx: n,
ny: n,
nz: n,
dx: h,
dy: h,
dz: h,
},
dt,
)
}
fn run_pair(host: Solver, dev: Solver, g: Grid, dt: f64, steps: usize, label: &str) -> (f64, f64) {
let mut host = host;
let mut fh = Field::new(g);
let mut device = DeviceStep::new(dev, g);
device.upload(&fh);
let mut mismatched = 0usize;
for _ in 0..steps {
let rh = host.advance(&mut fh, dt);
let rd = device.advance(dt);
if rh.poisson_iterations != rd.poisson_iterations {
mismatched += 1;
}
}
let mut fd = Field::new(g);
device.download(&mut fd);
let (worst, scale) = compare(&fh, &fd);
println!(
" {label}: {steps} steps, host vs device max |Δ| {worst:.3e} on a scale of {scale:.3e}; CG iteration counts differ on {mismatched} steps"
);
(worst, scale)
}
/// The CG's stop is a threshold on a reduction; host and device reduce in
/// different orders, so on a marginal step one side takes one more
/// iteration and the answers differ by the inner tolerance. The gate is
/// therefore taken at a TIGHT inner stop (1e-6 of the source scale), where
/// that flip cannot show above 1e-12; the default stop's difference is
/// reported alongside.
#[test]
fn device_step_matches_the_host_step_on_the_manufactured_problem() {
for scheme in [ConvectionScheme::Upwind, ConvectionScheme::TvdVanAlbada] {
let (h, d, g, dt) = mms_pair(12, scheme, false);
run_pair(
h,
d,
g,
dt,
100,
&format!("MMS n 12 {scheme:?} (default tolerances)"),
);
let (h, d, g, dt) = mms_pair(12, scheme, true);
let (worst, scale) = run_pair(h, d, g, dt, 100, &format!("MMS n 12 {scheme:?} (tight)"));
assert!(
worst <= 1e-11 * scale,
"{scheme:?}: {worst:.3e} of {scale:.3e}"
);
}
}
// ---- Beltrami (three_d_beltrami.rs) ----
const NU: f64 = 0.02;
const A: f64 = PI / 4.0;
const D: f64 = PI / 2.0;
fn exact(x: f64, y: f64, z: f64, t: f64) -> (f64, f64, f64) {
let decay = (-D * D * NU * t).exp();
(
-A * ((A * x).exp() * (A * y + D * z).sin() + (A * z).exp() * (A * x + D * y).cos())
* decay,
-A * ((A * y).exp() * (A * z + D * x).sin() + (A * x).exp() * (A * y + D * z).cos())
* decay,
-A * ((A * z).exp() * (A * x + D * y).sin() + (A * y).exp() * (A * z + D * x).cos())
* decay,
)
}
#[test]
fn device_step_matches_the_host_step_on_beltrami() {
let n = 16;
let h = 1.0 / n as f64;
let dt = 0.25 * h * h / (4.0 * NU);
let g = Grid {
nx: n,
ny: n,
nz: n,
dx: h,
dy: h,
dz: h,
};
for tight in [false, true] {
let mk = || {
let mut s = Solver::new(
fluid(NU * RHO),
params(ConvectionScheme::Upwind, Side::Velocity, tight),
);
s.set_boundary_velocity(|x, y, z, t| exact(x, y, z, t));
s
};
let mut host = mk();
let mut fh = Field::new(g);
for k in 0..n {
for j in 0..n {
for i in 0..=n {
fh.u[g.uface(k, j, i)] = exact(
i as f64 * h,
(j as f64 + 0.5) * h,
(k as f64 + 0.5) * h,
0.0,
)
.0;
}
}
}
for k in 0..n {
for j in 0..=n {
for i in 0..n {
fh.v[g.vface(k, j, i)] = exact(
(i as f64 + 0.5) * h,
j as f64 * h,
(k as f64 + 0.5) * h,
0.0,
)
.1;
}
}
}
for k in 0..=n {
for j in 0..n {
for i in 0..n {
fh.w[g.wface(k, j, i)] = exact(
(i as f64 + 0.5) * h,
(j as f64 + 0.5) * h,
k as f64 * h,
0.0,
)
.2;
}
}
}
let mut device = DeviceStep::new(mk(), g);
device.upload(&fh);
let mut mismatched = 0;
for _ in 0..40 {
let rh = host.advance(&mut fh, dt);
let rd = device.advance(dt);
if rh.poisson_iterations != rd.poisson_iterations {
mismatched += 1;
}
}
let mut fd = Field::new(g);
device.download(&mut fd);
let (worst, scale) = compare(&fh, &fd);
println!(
" Beltrami n 16 (tight {tight}): 40 steps (time-dependent tables), host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
);
if tight {
assert!(worst <= 1e-11 * scale, "{worst:.3e} of {scale:.3e}");
}
}
}
// ---- Poiseuille (three_d_poiseuille_identity.rs) ----
const MU_P: f64 = 0.1;
const G: f64 = 0.8;
fn discrete_profile(n: usize) -> Vec<f64> {
let h = 1.0 / n as f64;
let rhs_value = -G * h * h / MU_P;
let mut diag = vec![-2.0; n];
diag[0] = -3.0;
diag[n - 1] = -3.0;
let mut rhs = vec![rhs_value; n];
let upper = vec![1.0; n];
for j in 1..n {
let factor = 1.0 / diag[j - 1];
diag[j] -= factor * upper[j - 1];
rhs[j] -= factor * rhs[j - 1];
}
let mut u = vec![0.0; n];
u[n - 1] = rhs[n - 1] / diag[n - 1];
for j in (0..n - 1).rev() {
u[j] = (rhs[j] - upper[j] * u[j + 1]) / diag[j];
}
u
}
#[test]
fn device_step_matches_the_host_step_on_poiseuille_and_is_z_invariant() {
let n = 16;
let h = 1.0 / n as f64;
let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
for (nz, dz, z, tight) in [
(1usize, 1.0, Side::SlipWall, false),
(1, 1.0, Side::SlipWall, true),
(4, h, Side::Periodic, true),
] {
let g = Grid {
nx: n,
ny: n,
nz,
dx: h,
dy: h,
dz,
};
let mk = || {
let mut s = Solver::new(fluid(MU_P), params(ConvectionScheme::Upwind, z, tight));
s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
let u_hat = discrete_profile(n);
s.set_boundary_velocity(move |x, y, _z, _t| {
if x <= 0.0 || x >= 1.0 {
let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
(u_hat[j], 0.0, 0.0)
} else {
(0.0, 0.0, 0.0)
}
});
s
};
let mut host = mk();
let mut fh = Field::new(g);
let u_hat = discrete_profile(n);
for k in 0..nz {
for (j, &uj) in u_hat.iter().enumerate() {
fh.u[g.uface(k, j, 0)] = uj;
fh.u[g.uface(k, j, n)] = uj;
}
}
let mut device = DeviceStep::new(mk(), g);
device.upload(&fh);
let mut mismatched = 0;
for _ in 0..300 {
let rh = host.advance(&mut fh, dt);
let rd = device.advance(dt);
if rh.poisson_iterations != rd.poisson_iterations {
mismatched += 1;
}
}
let mut fd = Field::new(g);
device.download(&mut fd);
let (worst, scale) = compare(&fh, &fd);
println!(
" Poiseuille n 16 nz {nz} (tight {tight}): 300 steps, host vs device max |Δ| {worst:.3e} on {scale:.3e}; CG counts differ on {mismatched} steps"
);
if tight {
assert!(
worst <= 1e-11 * scale,
"nz {nz}: {worst:.3e} of {scale:.3e}"
);
}
if nz > 1 {
let plane = |f: &Field, k: usize| f.u[k * n * (n + 1)..(k + 1) * n * (n + 1)].to_vec();
let p0 = plane(&fd, 0);
let mut worst_plane = 0.0_f64;
for k in 1..nz {
for (a, b) in plane(&fd, k).iter().zip(&p0) {
worst_plane = worst_plane.max((a - b).abs());
}
}
println!(" Poiseuille nz {nz}: device planes within {worst_plane:.3e} of {scale:.3e}");
assert!(worst_plane <= 1e-11 * scale);
}
}
}
/// Diagnostic: where and when the Poiseuille host/device difference enters.
#[test]
#[ignore = "diagnostic: per-component host/device differences on Poiseuille variants"]
fn poiseuille_difference_diagnostic() {
let n = 16;
let h = 1.0 / n as f64;
let dt = 0.4 * (h * h / (4.0 * MU_P)).min(h);
for (label, dz, inlet, source) in [
("as is (dz 1, inlet profile, source G)", 1.0, true, true),
("dz = h", h, true, true),
("closed box (no inlet), source G", 1.0, false, true),
("inlet profile, no source", 1.0, true, false),
] {
let g = Grid {
nx: n,
ny: n,
nz: 1,
dx: h,
dy: h,
dz,
};
let mk = || {
let mut s = Solver::new(
fluid(MU_P),
params(ConvectionScheme::Upwind, Side::SlipWall, true),
);
if source {
s.set_momentum_source(|_x, _y, _z, _t| (G, 0.0, 0.0));
}
let u_hat = discrete_profile(n);
s.set_boundary_velocity(move |x, y, _z, _t| {
if inlet && (x <= 0.0 || x >= 1.0) {
let j = ((y / h - 0.5).round().max(0.0) as usize).min(n - 1);
(u_hat[j], 0.0, 0.0)
} else {
(0.0, 0.0, 0.0)
}
});
s
};
let mut host = mk();
let mut fh = Field::new(g);
if inlet {
let u_hat = discrete_profile(n);
for (j, &uj) in u_hat.iter().enumerate() {
fh.u[g.uface(0, j, 0)] = uj;
fh.u[g.uface(0, j, n)] = uj;
}
}
let mut device = DeviceStep::new(mk(), g);
device.upload(&fh);
let mut fd = Field::new(g);
for step in 1..=300 {
let rh = host.advance(&mut fh, dt);
let rd = device.advance(dt);
if [1, 2, 10, 100, 300].contains(&step) {
device.download(&mut fd);
let du =
fh.u.iter()
.zip(&fd.u)
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
let dv =
fh.v.iter()
.zip(&fd.v)
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
let dp =
fh.p.iter()
.zip(&fd.p)
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
let dpp = fh
.p_prime
.iter()
.zip(&fd.p_prime)
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
let dsp = fh
.sp
.iter()
.zip(&fd.sp)
.fold(0.0_f64, |m, (a, b)| m.max((a - b).abs()));
println!(
" {label} step {step}: Δu {du:.2e} Δv {dv:.2e} Δp {dp:.2e} Δp' {dpp:.2e} Δsp {dsp:.2e}; CG it host {} / dev {}; correctors {} / {}; residual {:.2e} / {:.2e}",
rh.poisson_iterations,
rd.poisson_iterations,
rh.corrector_steps_performed,
rd.corrector_steps_performed,
rh.final_residual,
rd.final_residual
);
}
}
}
}
/// The device step's cost at the anchor size (378 × 62 × 62, cubic cells,
/// the manufactured source, red-black + device CG), `RTX_PROFILE=1` for the
/// phase split. Recorded, not asserted.
#[test]
#[ignore = "bench: ms per device step at 378×62×62"]
fn bench_device_step_anchor_size() {
let (nx, ny, nz) = (378usize, 62usize, 62usize);
let h = 0.41 / ny as f64;
let mu = 1.0e-3;
let g = Grid {
nx,
ny,
nz,
dx: h,
dy: h,
dz: h,
};
let dt = 0.4 * (h * h / (4.0 * mu / RHO)).min(h / 2.0);
let mut s = Solver::new(
fluid(mu),
params(ConvectionScheme::TvdVanAlbada, Side::Velocity, false),
);
s.set_momentum_source(move |x, y, z, _t| source3(mu, x / 2.5, y / 0.41, z / 0.41));
s.set_boundary_velocity(|x, y, z, _t| boundary3(x / 2.5, y / 0.41, z / 0.41));
let mut device = DeviceStep::new(s, g);
let f = Field::new(g);
device.upload(&f);
device.advance(dt);
let t0 = std::time::Instant::now();
let steps = 20;
let mut it = 0;
for _ in 0..steps {
it += device.advance(dt).poisson_iterations;
}
let ms = t0.elapsed().as_secs_f64() * 1e3 / steps as f64;
println!(
" device step at {nx}×{ny}×{nz} ({} cells): {ms:.1} ms per step, {:.1} CG iterations per step",
g.cells(),
it as f64 / steps as f64
);
if let Some(t) = device.timers() {
let n = t.steps.max(1) as f64;
println!(
" split per step: predictor {:.1} ms, poisson {:.1} ms, apply {:.1} ms ({} steps timed)",
t.predictor_ns as f64 / n / 1e6,
t.poisson_ns as f64 / n / 1e6,
t.apply_ns as f64 / n / 1e6,
t.steps
);
}
}