diff --git a/crates/specialized/rtx-cfd/src/kernels/cuda/piso3.cu b/crates/specialized/rtx-cfd/src/kernels/cuda/piso3.cu new file mode 100644 index 0000000..db1320a --- /dev/null +++ b/crates/specialized/rtx-cfd/src/kernels/cuda/piso3.cu @@ -0,0 +1,449 @@ +/** + * 3D Stage 1, item 5: the PISO step's maps on the device — the three + * predictors (the host `piso_predictor.rs` expression for expression; the + * module is compiled with FMA contraction OFF so the f64 arithmetic is the + * host's), the normal-velocity stamping from per-side tables, the + * continuity source, the corrections, the pressure update and the mass + * imbalance partials. One thread per face / cell. + */ +#define SIDE_VELOCITY 0 +#define SIDE_SLIP 1 +#define SIDE_OUTLET 2 +#define SIDE_PERIODIC 3 +#define SCHEME_UPWIND 0 +#define SCHEME_VAN_ALBADA 1 +#define SCHEME_VAN_LEER 2 + +struct P3Params { + int nx, ny, nz; + int periodic_z; + int bx0, bx1, by0, by1, bz0, bz1; + int scheme; + int pad; + double dx, dy, dz, dt, rho, nu; +}; + +/* Per-side tables: [side][component]; see `BoundaryTables` in piso_device.rs. */ +struct P3Ptrs { + double *u, *v, *w, *p, *uo, *vo, *wo, *us, *vs, *ws, *pp, *sp; + const double *su, *sv, *sw; + const double *bx0u, *bx0v, *bx0w, *bx1u, *bx1v, *bx1w; + const double *by0u, *by0v, *by0w, *by1u, *by1v, *by1w; + const double *bz0u, *bz0v, *bz0w, *bz1u, *bz1v, *bz1w; +}; + +__device__ __forceinline__ int cell3(const P3Params& g, int k, int j, int i) { return (k * g.ny + j) * g.nx + i; } +__device__ __forceinline__ int uf3(const P3Params& g, int k, int j, int i) { return (k * g.ny + j) * (g.nx + 1) + i; } +__device__ __forceinline__ int vf3(const P3Params& g, int k, int j, int i) { return (k * (g.ny + 1) + j) * g.nx + i; } +__device__ __forceinline__ int wf3(const P3Params& g, int k, int j, int i) { return (k * g.ny + j) * g.nx + i; } + +__device__ __forceinline__ double upwind3(double face_velocity, double upstream, double downstream) { + return face_velocity >= 0.0 ? upstream : downstream; +} + +__device__ __forceinline__ double limiter3(int scheme, double r) { + if (scheme == SCHEME_VAN_ALBADA) return r > 0.0 ? (r * r + r) / (r * r + 1.0) : 0.0; + if (scheme == SCHEME_VAN_LEER) return (r + fabs(r)) / (1.0 + fabs(r)); + return 0.0; +} + +/* The limited correction; `has_far` = the far-upwind node exists. */ +__device__ __forceinline__ double face_corr3(int scheme, int has_far, double far, double up, double down) { + if (!has_far) return 0.0; + double denominator = down - up; + if (fabs(denominator) < 1e-300) return 0.0; + double r = (up - far) / denominator; + return 0.5 * limiter3(scheme, r) * denominator; +} + +/* k above / below with the periodic wrap; −1 = wall. */ +__device__ __forceinline__ int k_up3(const P3Params& g, int k) { return k + 1 < g.nz ? k + 1 : (g.periodic_z ? 0 : -1); } +__device__ __forceinline__ int k_dn3(const P3Params& g, int k) { return k > 0 ? k - 1 : (g.periodic_z ? g.nz - 1 : -1); } + +extern "C" __global__ void p3_predict_u(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int nxp = g.nx + 1; + int total = nxp * g.ny * g.nz; + if (t >= total) return; + int i = t % nxp; int j = (t / nxp) % g.ny; int k = t / (nxp * g.ny); + if (i == 0 || i == g.nx) return; + const double *uo = f.uo, *vo = f.vo, *wo = f.wo; + double dx = g.dx, dy = g.dy, dz = g.dz, rho = g.rho, nu = g.nu; + int scheme = g.scheme; + double u_p = uo[uf3(g, k, j, i)]; + double ue_face = 0.5 * (uo[uf3(g, k, j, i)] + uo[uf3(g, k, j, i + 1)]); + double uw_face = 0.5 * (uo[uf3(g, k, j, i - 1)] + uo[uf3(g, k, j, i)]); + int south_is_wall = (j == 0); + int north_is_wall = (j + 1 == g.ny); + double vn_face = 0.5 * (vo[vf3(g, k, j + 1, i - 1)] + vo[vf3(g, k, j + 1, i)]); + double vs_face = 0.5 * (vo[vf3(g, k, j, i - 1)] + vo[vf3(g, k, j, i)]); + double beyond_north = (g.by1 == SIDE_VELOCITY) ? f.by1u[k * nxp + i] : u_p; + double beyond_south = (g.by0 == SIDE_VELOCITY) ? f.by0u[k * nxp + i] : u_p; + double conv_x = (ue_face * upwind3(ue_face, uo[uf3(g, k, j, i)], uo[uf3(g, k, j, i + 1)]) + - uw_face * upwind3(uw_face, uo[uf3(g, k, j, i - 1)], uo[uf3(g, k, j, i)])) / dx; + double conv_y = (vn_face * (north_is_wall ? upwind3(vn_face, u_p, beyond_north) + : upwind3(vn_face, uo[uf3(g, k, j, i)], uo[uf3(g, k, j + 1, i)])) + - vs_face * (south_is_wall ? upwind3(vs_face, beyond_south, u_p) + : upwind3(vs_face, uo[uf3(g, k, j - 1, i)], uo[uf3(g, k, j, i)]))) / dy; + if (scheme != SCHEME_UPWIND) { + double delta_e, delta_w, delta_n, delta_s; + if (ue_face >= 0.0) delta_e = face_corr3(scheme, 1, uo[uf3(g, k, j, i - 1)], uo[uf3(g, k, j, i)], uo[uf3(g, k, j, i + 1)]); + else { int has = (i + 2 <= g.nx); delta_e = face_corr3(scheme, has, has ? uo[uf3(g, k, j, i + 2)] : 0.0, uo[uf3(g, k, j, i + 1)], uo[uf3(g, k, j, i)]); } + if (uw_face >= 0.0) { int has = (i >= 2); delta_w = face_corr3(scheme, has, has ? uo[uf3(g, k, j, i - 2)] : 0.0, uo[uf3(g, k, j, i - 1)], uo[uf3(g, k, j, i)]); } + else delta_w = face_corr3(scheme, 1, uo[uf3(g, k, j, i + 1)], uo[uf3(g, k, j, i)], uo[uf3(g, k, j, i - 1)]); + if (north_is_wall) delta_n = 0.0; + else if (vn_face >= 0.0) { int has = (j >= 1); delta_n = face_corr3(scheme, has, has ? uo[uf3(g, k, j - 1, i)] : 0.0, uo[uf3(g, k, j, i)], uo[uf3(g, k, j + 1, i)]); } + else { int has = (j + 2 < g.ny); delta_n = face_corr3(scheme, has, has ? uo[uf3(g, k, j + 2, i)] : 0.0, uo[uf3(g, k, j + 1, i)], uo[uf3(g, k, j, i)]); } + if (south_is_wall) delta_s = 0.0; + else if (vs_face >= 0.0) { int has = (j >= 2); delta_s = face_corr3(scheme, has, has ? uo[uf3(g, k, j - 2, i)] : 0.0, uo[uf3(g, k, j - 1, i)], uo[uf3(g, k, j, i)]); } + else { int has = (j + 1 < g.ny); delta_s = face_corr3(scheme, has, has ? uo[uf3(g, k, j + 1, i)] : 0.0, uo[uf3(g, k, j, i)], uo[uf3(g, k, j - 1, i)]); } + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + } + double diff_x = nu * (uo[uf3(g, k, j, i + 1)] - 2.0 * u_p + uo[uf3(g, k, j, i - 1)]) / (dx * dx); + double flux_north = north_is_wall ? (g.by1 == SIDE_VELOCITY ? nu * (f.by1u[k * nxp + i] - u_p) / (0.5 * dy) : 0.0) + : nu * (uo[uf3(g, k, j + 1, i)] - u_p) / dy; + double flux_south = south_is_wall ? (g.by0 == SIDE_VELOCITY ? nu * (u_p - f.by0u[k * nxp + i]) / (0.5 * dy) : 0.0) + : nu * (u_p - uo[uf3(g, k, j - 1, i)]) / dy; + double diff_y = (flux_north - flux_south) / dy; + double pressure_gradient = -(f.p[cell3(g, k, j, i)] - f.p[cell3(g, k, j, i - 1)]) / (rho * dx); + double body_force = f.su[uf3(g, k, j, i)] / rho; + double rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; + /* z terms */ + int ku = k_up3(g, k), kd = k_dn3(g, k); + int top_is_wall = (ku < 0), bottom_is_wall = (kd < 0); + double wt_face = 0.5 * (wo[wf3(g, k + 1, j, i - 1)] + wo[wf3(g, k + 1, j, i)]); + double wb_face = 0.5 * (wo[wf3(g, k, j, i - 1)] + wo[wf3(g, k, j, i)]); + double beyond_top = (g.bz1 == SIDE_VELOCITY) ? f.bz1u[j * nxp + i] : u_p; + double beyond_bottom = (g.bz0 == SIDE_VELOCITY) ? f.bz0u[j * nxp + i] : u_p; + double u_up = top_is_wall ? u_p : uo[uf3(g, ku, j, i)]; + double u_dn = bottom_is_wall ? u_p : uo[uf3(g, kd, j, i)]; + double conv_z = (wt_face * (top_is_wall ? upwind3(wt_face, u_p, beyond_top) : upwind3(wt_face, u_p, u_up)) + - wb_face * (bottom_is_wall ? upwind3(wb_face, beyond_bottom, u_p) : upwind3(wb_face, u_dn, u_p))) / dz; + if (scheme != SCHEME_UPWIND) { + int ku2 = top_is_wall ? -1 : k_up3(g, ku); + int kd2 = bottom_is_wall ? -1 : k_dn3(g, kd); + double far_up2 = ku2 >= 0 ? uo[uf3(g, ku2, j, i)] : 0.0; + double far_dn2 = kd2 >= 0 ? uo[uf3(g, kd2, j, i)] : 0.0; + double delta_t, delta_b; + if (top_is_wall) delta_t = 0.0; + else if (wt_face >= 0.0) delta_t = face_corr3(scheme, !bottom_is_wall, u_dn, u_p, u_up); + else delta_t = face_corr3(scheme, ku2 >= 0, far_up2, u_up, u_p); + if (bottom_is_wall) delta_b = 0.0; + else if (wb_face >= 0.0) delta_b = face_corr3(scheme, kd2 >= 0, far_dn2, u_dn, u_p); + else delta_b = face_corr3(scheme, !top_is_wall, u_up, u_p, u_dn); + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + } + double flux_top = top_is_wall ? (g.bz1 == SIDE_VELOCITY ? nu * (beyond_top - u_p) / (0.5 * dz) : 0.0) : nu * (u_up - u_p) / dz; + double flux_bottom = bottom_is_wall ? (g.bz0 == SIDE_VELOCITY ? nu * (u_p - beyond_bottom) / (0.5 * dz) : 0.0) : nu * (u_p - u_dn) / dz; + double diff_z = (flux_top - flux_bottom) / dz; + double rhs = rhs_2d - conv_z + diff_z; + f.u[uf3(g, k, j, i)] = uo[uf3(g, k, j, i)] + g.dt * rhs; +} + +extern "C" __global__ void p3_predict_v(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int nyp = g.ny + 1; + int total = g.nx * nyp * g.nz; + if (t >= total) return; + int i = t % g.nx; int j = (t / g.nx) % nyp; int k = t / (g.nx * nyp); + if (j == 0 || j == g.ny) return; + const double *uo = f.uo, *vo = f.vo, *wo = f.wo; + double dx = g.dx, dy = g.dy, dz = g.dz, rho = g.rho, nu = g.nu; + int scheme = g.scheme; + double v_p = vo[vf3(g, k, j, i)]; + double vn_face = 0.5 * (vo[vf3(g, k, j, i)] + vo[vf3(g, k, j + 1, i)]); + double vs_face = 0.5 * (vo[vf3(g, k, j - 1, i)] + vo[vf3(g, k, j, i)]); + int west_is_wall = (i == 0), east_is_wall = (i + 1 == g.nx); + double ue_face = 0.5 * (uo[uf3(g, k, j - 1, i + 1)] + uo[uf3(g, k, j, i + 1)]); + double uw_face = 0.5 * (uo[uf3(g, k, j - 1, i)] + uo[uf3(g, k, j, i)]); + double beyond_east = (g.bx1 == SIDE_VELOCITY) ? f.bx1v[k * nyp + j] : v_p; + double beyond_west = (g.bx0 == SIDE_VELOCITY) ? f.bx0v[k * nyp + j] : v_p; + double conv_y = (vn_face * upwind3(vn_face, vo[vf3(g, k, j, i)], vo[vf3(g, k, j + 1, i)]) + - vs_face * upwind3(vs_face, vo[vf3(g, k, j - 1, i)], vo[vf3(g, k, j, i)])) / dy; + double conv_x = (ue_face * (east_is_wall ? upwind3(ue_face, v_p, beyond_east) + : upwind3(ue_face, vo[vf3(g, k, j, i)], vo[vf3(g, k, j, i + 1)])) + - uw_face * (west_is_wall ? upwind3(uw_face, beyond_west, v_p) + : upwind3(uw_face, vo[vf3(g, k, j, i - 1)], vo[vf3(g, k, j, i)]))) / dx; + if (scheme != SCHEME_UPWIND) { + double delta_n, delta_s, delta_e, delta_w; + if (vn_face >= 0.0) delta_n = face_corr3(scheme, 1, vo[vf3(g, k, j - 1, i)], vo[vf3(g, k, j, i)], vo[vf3(g, k, j + 1, i)]); + else { int has = (j + 2 <= g.ny); delta_n = face_corr3(scheme, has, has ? vo[vf3(g, k, j + 2, i)] : 0.0, vo[vf3(g, k, j + 1, i)], vo[vf3(g, k, j, i)]); } + if (vs_face >= 0.0) { int has = (j >= 2); delta_s = face_corr3(scheme, has, has ? vo[vf3(g, k, j - 2, i)] : 0.0, vo[vf3(g, k, j - 1, i)], vo[vf3(g, k, j, i)]); } + else delta_s = face_corr3(scheme, 1, vo[vf3(g, k, j + 1, i)], vo[vf3(g, k, j, i)], vo[vf3(g, k, j - 1, i)]); + if (east_is_wall) delta_e = 0.0; + else if (ue_face >= 0.0) { int has = (i >= 1); delta_e = face_corr3(scheme, has, has ? vo[vf3(g, k, j, i - 1)] : 0.0, vo[vf3(g, k, j, i)], vo[vf3(g, k, j, i + 1)]); } + else { int has = (i + 2 < g.nx); delta_e = face_corr3(scheme, has, has ? vo[vf3(g, k, j, i + 2)] : 0.0, vo[vf3(g, k, j, i + 1)], vo[vf3(g, k, j, i)]); } + if (west_is_wall) delta_w = 0.0; + else if (uw_face >= 0.0) { int has = (i >= 2); delta_w = face_corr3(scheme, has, has ? vo[vf3(g, k, j, i - 2)] : 0.0, vo[vf3(g, k, j, i - 1)], vo[vf3(g, k, j, i)]); } + else { int has = (i + 1 < g.nx); delta_w = face_corr3(scheme, has, has ? vo[vf3(g, k, j, i + 1)] : 0.0, vo[vf3(g, k, j, i)], vo[vf3(g, k, j, i - 1)]); } + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + } + double diff_y = nu * (vo[vf3(g, k, j + 1, i)] - 2.0 * v_p + vo[vf3(g, k, j - 1, i)]) / (dy * dy); + double flux_east = east_is_wall ? (g.bx1 == SIDE_VELOCITY ? nu * (f.bx1v[k * nyp + j] - v_p) / (0.5 * dx) : 0.0) + : nu * (vo[vf3(g, k, j, i + 1)] - v_p) / dx; + double flux_west = west_is_wall ? (g.bx0 == SIDE_VELOCITY ? nu * (v_p - f.bx0v[k * nyp + j]) / (0.5 * dx) : 0.0) + : nu * (v_p - vo[vf3(g, k, j, i - 1)]) / dx; + double diff_x = (flux_east - flux_west) / dx; + double pressure_gradient = -(f.p[cell3(g, k, j, i)] - f.p[cell3(g, k, j - 1, i)]) / (rho * dy); + double body_force = f.sv[vf3(g, k, j, i)] / rho; + double rhs_2d = -conv_x - conv_y + diff_x + diff_y + pressure_gradient + body_force; + /* z terms */ + int ku = k_up3(g, k), kd = k_dn3(g, k); + int top_is_wall = (ku < 0), bottom_is_wall = (kd < 0); + double wt_face = 0.5 * (wo[wf3(g, k + 1, j - 1, i)] + wo[wf3(g, k + 1, j, i)]); + double wb_face = 0.5 * (wo[wf3(g, k, j - 1, i)] + wo[wf3(g, k, j, i)]); + double beyond_top = (g.bz1 == SIDE_VELOCITY) ? f.bz1v[j * g.nx + i] : v_p; + double beyond_bottom = (g.bz0 == SIDE_VELOCITY) ? f.bz0v[j * g.nx + i] : v_p; + double v_up = top_is_wall ? v_p : vo[vf3(g, ku, j, i)]; + double v_dn = bottom_is_wall ? v_p : vo[vf3(g, kd, j, i)]; + double conv_z = (wt_face * (top_is_wall ? upwind3(wt_face, v_p, beyond_top) : upwind3(wt_face, v_p, v_up)) + - wb_face * (bottom_is_wall ? upwind3(wb_face, beyond_bottom, v_p) : upwind3(wb_face, v_dn, v_p))) / dz; + if (scheme != SCHEME_UPWIND) { + int ku2 = top_is_wall ? -1 : k_up3(g, ku); + int kd2 = bottom_is_wall ? -1 : k_dn3(g, kd); + double far_up2 = ku2 >= 0 ? vo[vf3(g, ku2, j, i)] : 0.0; + double far_dn2 = kd2 >= 0 ? vo[vf3(g, kd2, j, i)] : 0.0; + double delta_t, delta_b; + if (top_is_wall) delta_t = 0.0; + else if (wt_face >= 0.0) delta_t = face_corr3(scheme, !bottom_is_wall, v_dn, v_p, v_up); + else delta_t = face_corr3(scheme, ku2 >= 0, far_up2, v_up, v_p); + if (bottom_is_wall) delta_b = 0.0; + else if (wb_face >= 0.0) delta_b = face_corr3(scheme, kd2 >= 0, far_dn2, v_dn, v_p); + else delta_b = face_corr3(scheme, !top_is_wall, v_up, v_p, v_dn); + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + } + double flux_top = top_is_wall ? (g.bz1 == SIDE_VELOCITY ? nu * (beyond_top - v_p) / (0.5 * dz) : 0.0) : nu * (v_up - v_p) / dz; + double flux_bottom = bottom_is_wall ? (g.bz0 == SIDE_VELOCITY ? nu * (v_p - beyond_bottom) / (0.5 * dz) : 0.0) : nu * (v_p - v_dn) / dz; + double diff_z = (flux_top - flux_bottom) / dz; + double rhs = rhs_2d - conv_z + diff_z; + f.v[vf3(g, k, j, i)] = vo[vf3(g, k, j, i)] + g.dt * rhs; +} + +extern "C" __global__ void p3_predict_w(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int total = g.nx * g.ny * (g.nz + 1); + if (t >= total) return; + int i = t % g.nx; int j = (t / g.nx) % g.ny; int k = t / (g.nx * g.ny); + int periodic = g.periodic_z; + if (periodic) { if (k == g.nz) return; } else { if (k == 0 || k == g.nz) return; } + const double *uo = f.uo, *vo = f.vo, *wo = f.wo; + double dx = g.dx, dy = g.dy, dz = g.dz, rho = g.rho, nu = g.nu; + int scheme = g.scheme; + int nz = g.nz; + int k_below = k > 0 ? k - 1 : nz - 1; + int k_above = k % nz; + int w_dn_idx = k > 0 ? wf3(g, k - 1, j, i) : wf3(g, nz - 1, j, i); + int w_up_idx = (k + 1 == nz && periodic) ? wf3(g, 0, j, i) : wf3(g, k + 1, j, i); + double w_p = wo[wf3(g, k, j, i)]; + double wt_face = 0.5 * (wo[wf3(g, k, j, i)] + wo[w_up_idx]); + double wb_face = 0.5 * (wo[w_dn_idx] + wo[wf3(g, k, j, i)]); + int west_is_wall = (i == 0), east_is_wall = (i + 1 == g.nx); + int south_is_wall = (j == 0), north_is_wall = (j + 1 == g.ny); + double ue_face = 0.5 * (uo[uf3(g, k_below, j, i + 1)] + uo[uf3(g, k_above, j, i + 1)]); + double uw_face = 0.5 * (uo[uf3(g, k_below, j, i)] + uo[uf3(g, k_above, j, i)]); + double vn_face = 0.5 * (vo[vf3(g, k_below, j + 1, i)] + vo[vf3(g, k_above, j + 1, i)]); + double vs_face = 0.5 * (vo[vf3(g, k_below, j, i)] + vo[vf3(g, k_above, j, i)]); + double beyond_east = (g.bx1 == SIDE_VELOCITY) ? f.bx1w[k * g.ny + j] : w_p; + double beyond_west = (g.bx0 == SIDE_VELOCITY) ? f.bx0w[k * g.ny + j] : w_p; + double beyond_north = (g.by1 == SIDE_VELOCITY) ? f.by1w[k * g.nx + i] : w_p; + double beyond_south = (g.by0 == SIDE_VELOCITY) ? f.by0w[k * g.nx + i] : w_p; + double conv_z = (wt_face * upwind3(wt_face, w_p, wo[w_up_idx]) - wb_face * upwind3(wb_face, wo[w_dn_idx], w_p)) / dz; + double conv_x = (ue_face * (east_is_wall ? upwind3(ue_face, w_p, beyond_east) : upwind3(ue_face, w_p, wo[wf3(g, k, j, i + 1)])) + - uw_face * (west_is_wall ? upwind3(uw_face, beyond_west, w_p) : upwind3(uw_face, wo[wf3(g, k, j, i - 1)], w_p))) / dx; + double conv_y = (vn_face * (north_is_wall ? upwind3(vn_face, w_p, beyond_north) : upwind3(vn_face, w_p, wo[wf3(g, k, j + 1, i)])) + - vs_face * (south_is_wall ? upwind3(vs_face, beyond_south, w_p) : upwind3(vs_face, wo[wf3(g, k, j - 1, i)], w_p))) / dy; + if (scheme != SCHEME_UPWIND) { + int has_up2, has_dn2; double far_up2 = 0.0, far_dn2 = 0.0; + if (periodic) { has_up2 = 1; far_up2 = wo[wf3(g, (k + 2) % nz, j, i)]; has_dn2 = 1; far_dn2 = wo[wf3(g, (k + nz - 2) % nz, j, i)]; } + else { has_up2 = (k + 2 <= nz); if (has_up2) far_up2 = wo[wf3(g, k + 2, j, i)]; has_dn2 = (k >= 2); if (has_dn2) far_dn2 = wo[wf3(g, k - 2, j, i)]; } + double delta_t = wt_face >= 0.0 ? face_corr3(scheme, 1, wo[w_dn_idx], w_p, wo[w_up_idx]) : face_corr3(scheme, has_up2, far_up2, wo[w_up_idx], w_p); + double delta_b = wb_face >= 0.0 ? face_corr3(scheme, has_dn2, far_dn2, wo[w_dn_idx], w_p) : face_corr3(scheme, 1, wo[w_up_idx], w_p, wo[w_dn_idx]); + double delta_e, delta_w, delta_n, delta_s; + if (east_is_wall) delta_e = 0.0; + else if (ue_face >= 0.0) { int has = (i >= 1); delta_e = face_corr3(scheme, has, has ? wo[wf3(g, k, j, i - 1)] : 0.0, w_p, wo[wf3(g, k, j, i + 1)]); } + else { int has = (i + 2 < g.nx); delta_e = face_corr3(scheme, has, has ? wo[wf3(g, k, j, i + 2)] : 0.0, wo[wf3(g, k, j, i + 1)], w_p); } + if (west_is_wall) delta_w = 0.0; + else if (uw_face >= 0.0) { int has = (i >= 2); delta_w = face_corr3(scheme, has, has ? wo[wf3(g, k, j, i - 2)] : 0.0, wo[wf3(g, k, j, i - 1)], w_p); } + else { int has = (i + 1 < g.nx); delta_w = face_corr3(scheme, has, has ? wo[wf3(g, k, j, i + 1)] : 0.0, w_p, wo[wf3(g, k, j, i - 1)]); } + if (north_is_wall) delta_n = 0.0; + else if (vn_face >= 0.0) { int has = (j >= 1); delta_n = face_corr3(scheme, has, has ? wo[wf3(g, k, j - 1, i)] : 0.0, w_p, wo[wf3(g, k, j + 1, i)]); } + else { int has = (j + 2 < g.ny); delta_n = face_corr3(scheme, has, has ? wo[wf3(g, k, j + 2, i)] : 0.0, wo[wf3(g, k, j + 1, i)], w_p); } + if (south_is_wall) delta_s = 0.0; + else if (vs_face >= 0.0) { int has = (j >= 2); delta_s = face_corr3(scheme, has, has ? wo[wf3(g, k, j - 2, i)] : 0.0, wo[wf3(g, k, j - 1, i)], w_p); } + else { int has = (j + 1 < g.ny); delta_s = face_corr3(scheme, has, has ? wo[wf3(g, k, j + 1, i)] : 0.0, w_p, wo[wf3(g, k, j - 1, i)]); } + conv_z += (wt_face * delta_t - wb_face * delta_b) / dz; + conv_x += (ue_face * delta_e - uw_face * delta_w) / dx; + conv_y += (vn_face * delta_n - vs_face * delta_s) / dy; + } + double diff_z = nu * (wo[w_up_idx] - 2.0 * w_p + wo[w_dn_idx]) / (dz * dz); + double flux_east = east_is_wall ? (g.bx1 == SIDE_VELOCITY ? nu * (beyond_east - w_p) / (0.5 * dx) : 0.0) : nu * (wo[wf3(g, k, j, i + 1)] - w_p) / dx; + double flux_west = west_is_wall ? (g.bx0 == SIDE_VELOCITY ? nu * (w_p - beyond_west) / (0.5 * dx) : 0.0) : nu * (w_p - wo[wf3(g, k, j, i - 1)]) / dx; + double diff_x = (flux_east - flux_west) / dx; + double flux_north = north_is_wall ? (g.by1 == SIDE_VELOCITY ? nu * (beyond_north - w_p) / (0.5 * dy) : 0.0) : nu * (wo[wf3(g, k, j + 1, i)] - w_p) / dy; + double flux_south = south_is_wall ? (g.by0 == SIDE_VELOCITY ? nu * (w_p - beyond_south) / (0.5 * dy) : 0.0) : nu * (w_p - wo[wf3(g, k, j - 1, i)]) / dy; + double diff_y = (flux_north - flux_south) / dy; + double pressure_gradient = -(f.p[cell3(g, k_above, j, i)] - f.p[cell3(g, k_below, j, i)]) / (rho * dz); + double body_force = f.sw[wf3(g, k, j, i)] / rho; + double rhs = -conv_x - conv_y - conv_z + diff_x + diff_y + diff_z + pressure_gradient + body_force; + f.w[wf3(g, k, j, i)] = wo[wf3(g, k, j, i)] + g.dt * rhs; +} + +/* After the predictor: the periodic copy w[nz] = w[0], the outlet + * zero-gradient faces, and the normal stamping from the tables. One thread + * per (k, j) for the x sides, (k, i) for the y sides, (j, i) for the z + * sides — three kernels. */ +extern "C" __global__ void p3_sides_x(P3Params g, P3Ptrs f, int stamp) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= g.ny * g.nz) return; + int j = t % g.ny, k = t / g.ny; + if (g.bx0 == SIDE_OUTLET) f.u[uf3(g, k, j, 0)] = f.u[uf3(g, k, j, 1)]; + else if (stamp) f.u[uf3(g, k, j, 0)] = f.bx0u[k * g.ny + j]; + if (g.bx1 == SIDE_OUTLET) f.u[uf3(g, k, j, g.nx)] = f.u[uf3(g, k, j, g.nx - 1)]; + else if (stamp) f.u[uf3(g, k, j, g.nx)] = f.bx1u[k * g.ny + j]; +} + +extern "C" __global__ void p3_sides_y(P3Params g, P3Ptrs f, int stamp) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= g.nx * g.nz) return; + int i = t % g.nx, k = t / g.nx; + if (g.by0 == SIDE_OUTLET) f.v[vf3(g, k, 0, i)] = f.v[vf3(g, k, 1, i)]; + else if (stamp) f.v[vf3(g, k, 0, i)] = f.by0v[k * g.nx + i]; + if (g.by1 == SIDE_OUTLET) f.v[vf3(g, k, g.ny, i)] = f.v[vf3(g, k, g.ny - 1, i)]; + else if (stamp) f.v[vf3(g, k, g.ny, i)] = f.by1v[k * g.nx + i]; +} + +extern "C" __global__ void p3_sides_z(P3Params g, P3Ptrs f, int stamp) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= g.nx * g.ny) return; + int i = t % g.nx, j = t / g.nx; + if (g.periodic_z) { f.w[wf3(g, g.nz, j, i)] = f.w[wf3(g, 0, j, i)]; return; } + if (g.bz0 == SIDE_OUTLET) f.w[wf3(g, 0, j, i)] = f.w[wf3(g, 1, j, i)]; + else if (stamp) f.w[wf3(g, 0, j, i)] = f.bz0w[j * g.nx + i]; + if (g.bz1 == SIDE_OUTLET) f.w[wf3(g, g.nz, j, i)] = f.w[wf3(g, g.nz - 1, j, i)]; + else if (stamp) f.w[wf3(g, g.nz, j, i)] = f.bz1w[j * g.nx + i]; +} + +/* sp = −ρ Σ (flux out) over the cells; partial[block] = Σ |flux| (source scale). */ +extern "C" __global__ void p3_divergence(P3Params g, P3Ptrs f, double* __restrict__ partial) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int total = g.nx * g.ny * g.nz; + double v = 0.0; + if (t < total) { + int i = t % g.nx; int j = (t / g.nx) % g.ny; int k = t / (g.nx * g.ny); + double divergence_flux = g.rho + * ((f.us[uf3(g, k, j, i + 1)] - f.us[uf3(g, k, j, i)]) * (g.dy * g.dz) + + (f.vs[vf3(g, k, j + 1, i)] - f.vs[vf3(g, k, j, i)]) * (g.dx * g.dz) + + (f.ws[wf3(g, k + 1, j, i)] - f.ws[wf3(g, k, j, i)]) * (g.dx * g.dy)); + f.sp[t] = -divergence_flux; + v = fabs(divergence_flux); + } + __shared__ double sh[256]; + sh[threadIdx.x] = v; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s]; __syncthreads(); } + if (threadIdx.x == 0) partial[blockIdx.x] = sh[0]; +} + +/* The corrections from p' (interior faces, outlet faces against 0 outside). */ +extern "C" __global__ void p3_correct_u(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int nxp = g.nx + 1; + if (t >= nxp * g.ny * g.nz) return; + int i = t % nxp; int j = (t / nxp) % g.ny; int k = t / (nxp * g.ny); + double c = g.dt / g.rho; + if (i >= 1 && i < g.nx) { + double dp_dx = (f.pp[cell3(g, k, j, i)] - f.pp[cell3(g, k, j, i - 1)]) / g.dx; + f.u[t] = f.us[t] - c * dp_dx; + } else if (i == 0 && g.bx0 == SIDE_OUTLET) { + double dp_dx = (f.pp[cell3(g, k, j, 0)] - 0.0) / (0.5 * g.dx); + f.u[t] = f.us[t] - c * dp_dx; + } else if (i == g.nx && g.bx1 == SIDE_OUTLET) { + double dp_dx = (0.0 - f.pp[cell3(g, k, j, g.nx - 1)]) / (0.5 * g.dx); + f.u[t] = f.us[t] - c * dp_dx; + } +} + +extern "C" __global__ void p3_correct_v(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int nyp = g.ny + 1; + if (t >= g.nx * nyp * g.nz) return; + int i = t % g.nx; int j = (t / g.nx) % nyp; int k = t / (g.nx * nyp); + double c = g.dt / g.rho; + if (j >= 1 && j < g.ny) { + double dp_dy = (f.pp[cell3(g, k, j, i)] - f.pp[cell3(g, k, j - 1, i)]) / g.dy; + f.v[t] = f.vs[t] - c * dp_dy; + } else if (j == 0 && g.by0 == SIDE_OUTLET) { + double dp_dy = (f.pp[cell3(g, k, 0, i)] - 0.0) / (0.5 * g.dy); + f.v[t] = f.vs[t] - c * dp_dy; + } else if (j == g.ny && g.by1 == SIDE_OUTLET) { + double dp_dy = (0.0 - f.pp[cell3(g, k, g.ny - 1, i)]) / (0.5 * g.dy); + f.v[t] = f.vs[t] - c * dp_dy; + } +} + +extern "C" __global__ void p3_correct_w(P3Params g, P3Ptrs f) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + if (t >= g.nx * g.ny * (g.nz + 1)) return; + int i = t % g.nx; int j = (t / g.nx) % g.ny; int k = t / (g.nx * g.ny); + double c = g.dt / g.rho; + if (g.periodic_z) { + if (k == g.nz) return; + int below = k > 0 ? k - 1 : g.nz - 1; + double dp_dz = (f.pp[cell3(g, k, j, i)] - f.pp[cell3(g, below, j, i)]) / g.dz; + f.w[t] = f.ws[t] - c * dp_dz; + return; + } + if (k >= 1 && k < g.nz) { + double dp_dz = (f.pp[cell3(g, k, j, i)] - f.pp[cell3(g, k - 1, j, i)]) / g.dz; + f.w[t] = f.ws[t] - c * dp_dz; + } else if (k == 0 && g.bz0 == SIDE_OUTLET) { + double dp_dz = (f.pp[cell3(g, 0, j, i)] - 0.0) / (0.5 * g.dz); + f.w[t] = f.ws[t] - c * dp_dz; + } else if (k == g.nz && g.bz1 == SIDE_OUTLET) { + double dp_dz = (0.0 - f.pp[cell3(g, g.nz - 1, j, i)]) / (0.5 * g.dz); + f.w[t] = f.ws[t] - c * dp_dz; + } +} + +/* p += p'; partial[block] = Σ |mass imbalance| of the corrected field. */ +extern "C" __global__ void p3_add_p_and_imbalance(P3Params g, P3Ptrs f, double* __restrict__ partial) +{ + int t = blockIdx.x * blockDim.x + threadIdx.x; + int total = g.nx * g.ny * g.nz; + double v = 0.0; + if (t < total) { + int i = t % g.nx; int j = (t / g.nx) % g.ny; int k = t / (g.nx * g.ny); + f.p[t] += f.pp[t]; + double divergence_flux = g.rho + * ((f.u[uf3(g, k, j, i + 1)] - f.u[uf3(g, k, j, i)]) * (g.dy * g.dz) + + (f.v[vf3(g, k, j + 1, i)] - f.v[vf3(g, k, j, i)]) * (g.dx * g.dz) + + (f.w[wf3(g, k + 1, j, i)] - f.w[wf3(g, k, j, i)]) * (g.dx * g.dy)); + v = fabs(divergence_flux); + } + __shared__ double sh[256]; + sh[threadIdx.x] = v; + __syncthreads(); + for (int s = 128; s > 0; s >>= 1) { if (threadIdx.x < s) sh[threadIdx.x] += sh[threadIdx.x + s]; __syncthreads(); } + if (threadIdx.x == 0) partial[blockIdx.x] = sh[0]; +} + +/* Σ partial in index order (one thread) — the fixed-order final sum. */ +extern "C" __global__ void p3_reduce(int n, const double* __restrict__ partial, double* __restrict__ out) +{ + if (blockIdx.x * blockDim.x + threadIdx.x != 0) return; + double s = 0.0; + for (int i = 0; i < n; ++i) s += partial[i]; + out[0] = s; +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs index e6fff1f..aa257d8 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/mod.rs @@ -8,6 +8,8 @@ //! Layout: cells are `(k, j, i)` row-major, `cell = (k * ny + j) * nx + i`. pub mod flow_field; +#[cfg(feature = "cuda")] +pub mod piso_device; pub mod piso_host; mod piso_predictor; pub mod poisson; diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_device.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_device.rs new file mode 100644 index 0000000..f95b68b --- /dev/null +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_device.rs @@ -0,0 +1,580 @@ +//! Item 5: the PISO step device-resident (`piso3.cu` + the device CG). The +//! fields live on the device; the host holds the solver's parameters and +//! functions, evaluates the boundary tables per step (kB) and the steady +//! momentum source once, and reads back scalars. `download` mirrors the +//! fields into a `FlowField3D` at instants. Compiled with FMA contraction +//! off so the predictors are the host's arithmetic to the bit. + +use super::Grid3; +use super::flow_field::FlowField3D; +use super::piso_host::{Piso3Result, Piso3Solver, SideBoundary3}; +use super::poisson::device::runtime3; +use super::poisson::device_cg::DeviceCg3; +use crate::solvers::incompressible::poisson::MultigridParameters; +use crate::solvers::incompressible::simple::ConvectionScheme; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, DevicePtr, DeviceRepr, LaunchConfig, PushKernelArg, + ValidAsZeroBits, +}; +use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts}; +use std::sync::{Arc, OnceLock}; +use std::time::Instant; + +const KERNELS: &str = include_str!("../../../kernels/cuda/piso3.cu"); + +struct Kernels { + _module: Arc, + predict_u: CudaFunction, + predict_v: CudaFunction, + predict_w: CudaFunction, + sides_x: CudaFunction, + sides_y: CudaFunction, + sides_z: CudaFunction, + divergence: CudaFunction, + correct_u: CudaFunction, + correct_v: CudaFunction, + correct_w: CudaFunction, + add_p: CudaFunction, + reduce: CudaFunction, +} + +static KERNELS_ONCE: OnceLock = OnceLock::new(); + +fn kernels() -> &'static Kernels { + KERNELS_ONCE.get_or_init(|| { + let rt = runtime3(); + let arch = std::env::var("RTX_CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_string()); + let ptx = compile_ptx_with_opts( + KERNELS, + CompileOptions { + arch: Some(Box::leak(arch.into_boxed_str())), + // The host's f64 arithmetic: no fused multiply-add. + options: vec!["--fmad=false".to_string()], + ..Default::default() + }, + ) + .expect("nvrtc: piso3.cu"); + let module = rt.ctx.load_module(ptx).expect("piso3 module"); + let f = |name: &str| module.load_function(name).expect(name); + Kernels { + predict_u: f("p3_predict_u"), + predict_v: f("p3_predict_v"), + predict_w: f("p3_predict_w"), + sides_x: f("p3_sides_x"), + sides_y: f("p3_sides_y"), + sides_z: f("p3_sides_z"), + divergence: f("p3_divergence"), + correct_u: f("p3_correct_u"), + correct_v: f("p3_correct_v"), + correct_w: f("p3_correct_w"), + add_p: f("p3_add_p_and_imbalance"), + reduce: f("p3_reduce"), + _module: module, + } + }) +} + +/// `struct P3Params` in piso3.cu. +#[repr(C)] +#[derive(Clone, Copy)] +struct P3Params { + nx: i32, + ny: i32, + nz: i32, + periodic_z: i32, + bx0: i32, + bx1: i32, + by0: i32, + by1: i32, + bz0: i32, + bz1: i32, + scheme: i32, + pad: i32, + dx: f64, + dy: f64, + dz: f64, + dt: f64, + rho: f64, + nu: f64, +} +unsafe impl DeviceRepr for P3Params {} +unsafe impl ValidAsZeroBits for P3Params {} + +/// `struct P3Ptrs` in piso3.cu: 33 device pointers. +#[repr(C)] +#[derive(Clone, Copy)] +struct P3Ptrs { + ptrs: [u64; 33], +} +unsafe impl DeviceRepr for P3Ptrs {} +unsafe impl ValidAsZeroBits for P3Ptrs {} + +fn side_code(s: SideBoundary3) -> i32 { + match s { + SideBoundary3::Velocity => 0, + SideBoundary3::SlipWall => 1, + SideBoundary3::PressureOutlet => 2, + SideBoundary3::Periodic => 3, + } +} + +fn scheme_code(s: ConvectionScheme) -> i32 { + match s { + ConvectionScheme::Upwind => 0, + ConvectionScheme::TvdVanAlbada => 1, + ConvectionScheme::TvdVanLeer => 2, + } +} + +fn cfg(n_items: usize) -> LaunchConfig { + LaunchConfig { + grid_dim: ((n_items as u32).div_ceil(256).max(1), 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Step timers (`RTX_PROFILE`): nanoseconds per phase and the step count. +#[derive(Debug, Clone, Copy, Default)] +pub struct StepTimers3 { + pub predictor_ns: u64, + pub poisson_ns: u64, + pub apply_ns: u64, + pub transfer_ns: u64, + pub steps: u64, + pub cg_iterations: u64, +} + +pub struct Piso3Device { + pub solver: Piso3Solver, + grid: Grid3, + u: CudaSlice, + v: CudaSlice, + w: CudaSlice, + p: CudaSlice, + u_old: CudaSlice, + v_old: CudaSlice, + w_old: CudaSlice, + u_star: CudaSlice, + v_star: CudaSlice, + w_star: CudaSlice, + p_prime: CudaSlice, + sp: CudaSlice, + su: CudaSlice, + sv: CudaSlice, + sw: CudaSlice, + tables: Vec>, + partial: CudaSlice, + scalar: CudaSlice, + n_blocks: usize, + cg: Option, + cg_dt: f64, + timers: Option, + initialized: bool, +} + +impl Piso3Device { + /// Allocates the device fields for `grid`; the momentum source is + /// tabulated at `t = 0` (steady sources only in Stage 1). + pub fn new(solver: Piso3Solver, grid: Grid3) -> Self { + let rt = runtime3(); + let nu = (grid.nx + 1) * grid.ny * grid.nz; + let nv = grid.nx * (grid.ny + 1) * grid.nz; + let nw = grid.nx * grid.ny * (grid.nz + 1); + let nc = grid.cells(); + let zeros = |n: usize| rt.stream.alloc_zeros::(n).expect("alloc"); + let (su, sv, sw) = solver.source_tables(grid, 0.0); + let up = |v: &Vec| -> CudaSlice { + rt.stream + .memcpy_stod(if v.is_empty() { &[0.0f64][..] } else { v }) + .expect("upload") + }; + let tables: Vec> = solver + .boundary_tables(grid, 0.0) + .iter() + .flat_map(|side| side.iter().map(up).collect::>()) + .collect(); + let n_blocks = nc.div_ceil(256).max(1); + let timers = std::env::var("RTX_PROFILE") + .is_ok() + .then(StepTimers3::default); + Self { + solver, + grid, + u: zeros(nu), + v: zeros(nv), + w: zeros(nw), + p: zeros(nc), + u_old: zeros(nu), + v_old: zeros(nv), + w_old: zeros(nw), + u_star: zeros(nu), + v_star: zeros(nv), + w_star: zeros(nw), + p_prime: zeros(nc), + sp: zeros(nc), + su: up(&su), + sv: up(&sv), + sw: up(&sw), + tables, + partial: zeros(n_blocks), + scalar: zeros(1), + n_blocks, + cg: None, + cg_dt: 0.0, + timers, + initialized: false, + } + } + + pub fn timers(&self) -> Option { + self.timers + } + + pub fn grid(&self) -> Grid3 { + self.grid + } + + /// The host field onto the device (u, v, w, p; p' too for a warm start). + pub fn upload(&mut self, field: &FlowField3D) { + let rt = runtime3(); + assert_eq!(field.grid, self.grid); + rt.stream.memcpy_htod(&field.u, &mut self.u).expect("u"); + rt.stream.memcpy_htod(&field.v, &mut self.v).expect("v"); + rt.stream.memcpy_htod(&field.w, &mut self.w).expect("w"); + rt.stream.memcpy_htod(&field.p, &mut self.p).expect("p"); + rt.stream + .memcpy_htod(&field.p_prime, &mut self.p_prime) + .expect("p'"); + rt.stream.synchronize().expect("sync"); + } + + /// The device field into the host mirror. + pub fn download(&self, field: &mut FlowField3D) { + let rt = runtime3(); + assert_eq!(field.grid, self.grid); + rt.stream.memcpy_dtoh(&self.u, &mut field.u).expect("u"); + rt.stream.memcpy_dtoh(&self.v, &mut field.v).expect("v"); + rt.stream.memcpy_dtoh(&self.w, &mut field.w).expect("w"); + rt.stream.memcpy_dtoh(&self.p, &mut field.p).expect("p"); + rt.stream + .memcpy_dtoh(&self.p_prime, &mut field.p_prime) + .expect("p'"); + rt.stream.memcpy_dtoh(&self.sp, &mut field.sp).expect("sp"); + rt.stream.synchronize().expect("sync"); + } + + fn params(&self, dt: f64) -> P3Params { + let g = self.grid; + let b = self.solver.params.boundaries; + P3Params { + nx: g.nx as i32, + ny: g.ny as i32, + nz: g.nz as i32, + periodic_z: i32::from(b.z0 == SideBoundary3::Periodic), + bx0: side_code(b.x0), + bx1: side_code(b.x1), + by0: side_code(b.y0), + by1: side_code(b.y1), + bz0: side_code(b.z0), + bz1: side_code(b.z1), + scheme: scheme_code(self.solver.params.convection_scheme), + pad: 0, + dx: g.dx, + dy: g.dy, + dz: g.dz, + dt, + rho: self.solver.fluid.density, + nu: self.solver.fluid.viscosity / self.solver.fluid.density, + } + } + + fn ptrs(&self) -> P3Ptrs { + let rt = runtime3(); + let s = &rt.stream; + let p = |x: &CudaSlice| x.device_ptr(s).0; + let mut ptrs = [0u64; 33]; + let base = [ + &self.u, + &self.v, + &self.w, + &self.p, + &self.u_old, + &self.v_old, + &self.w_old, + &self.u_star, + &self.v_star, + &self.w_star, + &self.p_prime, + &self.sp, + &self.su, + &self.sv, + &self.sw, + ]; + for (k, b) in base.iter().enumerate() { + ptrs[k] = p(b); + } + for (k, t) in self.tables.iter().enumerate() { + ptrs[15 + k] = p(t); + } + P3Ptrs { ptrs } + } + + fn upload_tables(&mut self, t: f64) { + let rt = runtime3(); + let host = self.solver.boundary_tables(self.grid, t); + let mut k = 0; + for side in &host { + for comp in side { + if !comp.is_empty() { + rt.stream + .memcpy_htod(comp, &mut self.tables[k]) + .expect("table"); + } + k += 1; + } + } + } + + fn reduce(&mut self) -> f64 { + let rt = runtime3(); + let k = kernels(); + let nb = self.n_blocks as i32; + unsafe { + rt.stream + .launch_builder(&k.reduce) + .arg(&nb) + .arg(&self.partial) + .arg(&mut self.scalar) + .launch(LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + .expect("p3_reduce"); + } + let mut one = vec![0.0f64]; + rt.stream + .memcpy_dtoh(&self.scalar, &mut one) + .expect("scalar"); + rt.stream.synchronize().expect("sync"); + one[0] + } + + fn launch_sides(&self, prm: P3Params, ptrs: P3Ptrs, stamp: i32) { + let rt = runtime3(); + let k = kernels(); + let g = self.grid; + unsafe { + rt.stream + .launch_builder(&k.sides_x) + .arg(&prm) + .arg(&ptrs) + .arg(&stamp) + .launch(cfg(g.ny * g.nz)) + .expect("sides_x"); + rt.stream + .launch_builder(&k.sides_y) + .arg(&prm) + .arg(&ptrs) + .arg(&stamp) + .launch(cfg(g.nx * g.nz)) + .expect("sides_y"); + rt.stream + .launch_builder(&k.sides_z) + .arg(&prm) + .arg(&ptrs) + .arg(&stamp) + .launch(cfg(g.nx * g.ny)) + .expect("sides_z"); + } + } + + /// Stamp the `t = time` boundary data (lazily by the first step). + pub fn initialize(&mut self) { + let t = self.solver.time(); + self.upload_tables(t); + let prm = self.params(1.0); + let ptrs = self.ptrs(); + self.launch_sides(prm, ptrs, 1); + self.initialized = true; + } + + /// One step of `dt` on the device. + pub fn advance(&mut self, dt: f64) -> Piso3Result { + assert!(dt > 0.0 && dt.is_finite()); + if !self.initialized { + self.initialize(); + } + let rt = runtime3(); + let k = kernels(); + let g = self.grid; + let t_old = self.solver.time(); + let t_new = t_old + dt; + let t0 = Instant::now(); + // History shift. + rt.stream + .memcpy_dtod(&self.u, &mut self.u_old) + .expect("u_old"); + rt.stream + .memcpy_dtod(&self.v, &mut self.v_old) + .expect("v_old"); + rt.stream + .memcpy_dtod(&self.w, &mut self.w_old) + .expect("w_old"); + // Predictor with the t_old tables. + self.upload_tables(t_old); + let prm = self.params(dt); + let ptrs = self.ptrs(); + let nu = (g.nx + 1) * g.ny * g.nz; + let nv = g.nx * (g.ny + 1) * g.nz; + let nw = g.nx * g.ny * (g.nz + 1); + unsafe { + rt.stream + .launch_builder(&k.predict_u) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nu)) + .expect("predict_u"); + rt.stream + .launch_builder(&k.predict_v) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nv)) + .expect("predict_v"); + rt.stream + .launch_builder(&k.predict_w) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nw)) + .expect("predict_w"); + } + // Outlet zero-gradient + periodic copy (no stamping), then u* = u. + self.launch_sides(prm, ptrs, 0); + // The new interval's boundary data. + self.upload_tables(t_new); + self.launch_sides(prm, ptrs, 1); + rt.stream + .memcpy_dtod(&self.u, &mut self.u_star) + .expect("u*"); + rt.stream + .memcpy_dtod(&self.v, &mut self.v_star) + .expect("v*"); + rt.stream + .memcpy_dtod(&self.w, &mut self.w_star) + .expect("w*"); + rt.stream.synchronize().expect("sync"); + let t_pred = t0.elapsed(); + // The operator (a cache keyed on dt; no body: constant otherwise). + let t1 = Instant::now(); + if self.cg.is_none() || self.cg_dt != dt { + let problem = self.solver.poisson_operator(g, dt); + let params = MultigridParameters { + precision: self.solver.params.poisson_precision, + smoother: self.solver.params.poisson_smoother, + ..MultigridParameters::default() + }; + self.cg = Some(DeviceCg3::new(&problem, ¶ms)); + self.cg_dt = dt; + } + let anchor = self.solver.anchor_cell(g); + let mut total = 0; + let mut final_residual = f64::INFINITY; + let mut cg_iterations = 0; + let mut t_poisson = std::time::Duration::ZERO; + let mut t_apply = std::time::Duration::ZERO; + for corrector in 0..self.solver.params.corrector_steps.max(1) { + let tp = Instant::now(); + unsafe { + rt.stream + .launch_builder(&k.divergence) + .arg(&prm) + .arg(&ptrs) + .arg(&mut self.partial) + .launch(cfg(g.cells())) + .expect("divergence"); + } + let source_scale = self.reduce(); + let inner_stop = self.solver.inner_stop(g, source_scale); + if corrector > 0 { + rt.stream.memset_zeros(&mut self.p_prime).expect("p' = 0"); + } + let sol = { + let cg = self.cg.as_mut().expect("cg"); + cg.solve_device(&self.sp, &mut self.p_prime, inner_stop, anchor, 0) + }; + cg_iterations += sol.iterations; + t_poisson += tp.elapsed(); + let ta = Instant::now(); + unsafe { + rt.stream + .launch_builder(&k.correct_u) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nu)) + .expect("correct_u"); + rt.stream + .launch_builder(&k.correct_v) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nv)) + .expect("correct_v"); + rt.stream + .launch_builder(&k.correct_w) + .arg(&prm) + .arg(&ptrs) + .launch(cfg(nw)) + .expect("correct_w"); + } + if prm.periodic_z != 0 { + self.launch_sides(prm, ptrs, 0); + } + unsafe { + rt.stream + .launch_builder(&k.add_p) + .arg(&prm) + .arg(&ptrs) + .arg(&mut self.partial) + .launch(cfg(g.cells())) + .expect("add_p"); + } + let imbalance = self.reduce(); + let reference_flux = self.solver.reference_flux(g); + let mass_residual = if reference_flux > 0.0 { + imbalance / reference_flux + } else { + imbalance + }; + final_residual = mass_residual; + total += 1; + t_apply += ta.elapsed(); + if mass_residual < self.solver.params.tolerance { + break; + } + rt.stream + .memcpy_dtod(&self.u, &mut self.u_star) + .expect("u*"); + rt.stream + .memcpy_dtod(&self.v, &mut self.v_star) + .expect("v*"); + rt.stream + .memcpy_dtod(&self.w, &mut self.w_star) + .expect("w*"); + } + let _ = t1; + self.solver.set_time(t_new); + if let Some(tm) = self.timers.as_mut() { + tm.predictor_ns += t_pred.as_nanos() as u64; + tm.poisson_ns += t_poisson.as_nanos() as u64; + tm.apply_ns += t_apply.as_nanos() as u64; + tm.steps += 1; + tm.cg_iterations += cg_iterations as u64; + } + Piso3Result { + converged: final_residual < self.solver.params.tolerance, + corrector_steps_performed: total, + final_residual, + poisson_iterations: cg_iterations, + } + } +} diff --git a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs index 6eb9f3e..54040d1 100644 --- a/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs +++ b/crates/specialized/rtx-cfd/src/solvers/incompressible/three_d/piso_host.rs @@ -315,12 +315,11 @@ impl Piso3Solver { field.copy_to_starred(); } - /// The pressure-correction system (the 2D `poisson_problem` with the z + /// The pressure-correction OPERATOR (the 2D `poisson_problem` with the z /// faces): coefficient `dt · A / δ` across every fluid interior face, /// zero across prescribed ones, the outlet's Dirichlet half a cell out - /// as `extra_diag`, the mass imbalance as the right-hand side. - pub(crate) fn poisson_problem(&self, field: &FlowField3D, dt: f64) -> PoissonProblem3D { - let g = field.grid; + /// as `extra_diag`; the right-hand side left zero. + pub(crate) fn poisson_operator(&self, g: Grid3, dt: f64) -> PoissonProblem3D { let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); let b = self.params.boundaries; let outlet = SideBoundary3::PressureOutlet; @@ -385,14 +384,155 @@ impl Piso3Solver { problem.ab[idx] = at_interior; } problem.extra_diag[idx] = extra; - problem.rhs[idx] = field.sp[idx]; } } } problem } - fn reference_flux(&self, g: Grid3) -> f64 { + /// The operator with `field.sp` as the right-hand side. + pub(crate) fn poisson_problem(&self, field: &FlowField3D, dt: f64) -> PoissonProblem3D { + let mut problem = self.poisson_operator(field.grid, dt); + problem.rhs.copy_from_slice(&field.sp); + problem + } + + /// The anchor cell of a pure-Neumann projection (the first fluid cell, + /// the 2D `(1, 1)` at `k = 0`), or `None` with an outlet. + pub(crate) fn anchor_cell(&self, g: Grid3) -> Option { + (!self.params.boundaries.any_outlet()).then_some(g.cell(0, 1, 1)) + } + + /// The inner stop of a projection from the source scale (the 2D rule). + pub(crate) fn inner_stop(&self, g: Grid3, source_scale: f64) -> f64 { + (self.params.inner_stop_factor * source_scale) + .max(0.1 * self.params.tolerance * self.reference_flux(g)) + + 1e-14 + } + + /// The boundary function on the six sides at every point a predictor or + /// the stamping reads (`[side][component]`, see the device driver): + /// x sides: u at `(k, j)`, v at `(k, j = 0..=ny)`, w at `(k = 0..=nz, j)`; + /// y sides: u at `(k, i = 0..=nx)`, v at `(k, i)`, w at `(k = 0..=nz, i)`; + /// z sides: u at `(j, i = 0..=nx)`, v at `(j = 0..=ny, i)`, w at `(j, i)`. + #[allow(clippy::type_complexity)] + pub fn boundary_tables(&self, g: Grid3, t: f64) -> [[Vec; 3]; 6] { + let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); + let xc = |i: usize| (i as f64 + 0.5) * dx; + let yc = |j: usize| (j as f64 + 0.5) * dy; + let zc = |k: usize| (k as f64 + 0.5) * dz; + let mut out: [[Vec; 3]; 6] = Default::default(); + for (s, x) in [(0usize, 0.0), (1, nx as f64 * dx)] { + let mut u = vec![0.0; ny * nz]; + let mut v = vec![0.0; (ny + 1) * nz]; + let mut w = vec![0.0; ny * (nz + 1)]; + for k in 0..nz { + for j in 0..ny { + u[k * ny + j] = self.boundary(x, yc(j), zc(k), t).0; + } + for j in 0..=ny { + v[k * (ny + 1) + j] = self.boundary(x, j as f64 * dy, zc(k), t).1; + } + } + for k in 0..=nz { + for j in 0..ny { + w[k * ny + j] = self.boundary(x, yc(j), k as f64 * dz, t).2; + } + } + out[s] = [u, v, w]; + } + for (s, y) in [(2usize, 0.0), (3, ny as f64 * dy)] { + let mut u = vec![0.0; (nx + 1) * nz]; + let mut v = vec![0.0; nx * nz]; + let mut w = vec![0.0; nx * (nz + 1)]; + for k in 0..nz { + for i in 0..=nx { + u[k * (nx + 1) + i] = self.boundary(i as f64 * dx, y, zc(k), t).0; + } + for i in 0..nx { + v[k * nx + i] = self.boundary(xc(i), y, zc(k), t).1; + } + } + for k in 0..=nz { + for i in 0..nx { + w[k * nx + i] = self.boundary(xc(i), y, k as f64 * dz, t).2; + } + } + out[s] = [u, v, w]; + } + for (s, z) in [(4usize, 0.0), (5, nz as f64 * dz)] { + let mut u = vec![0.0; (nx + 1) * ny]; + let mut v = vec![0.0; nx * (ny + 1)]; + let mut w = vec![0.0; nx * ny]; + for j in 0..ny { + for i in 0..=nx { + u[j * (nx + 1) + i] = self.boundary(i as f64 * dx, yc(j), z, t).0; + } + for i in 0..nx { + w[j * nx + i] = self.boundary(xc(i), yc(j), z, t).2; + } + } + for j in 0..=ny { + for i in 0..nx { + v[j * nx + i] = self.boundary(xc(i), j as f64 * dy, z, t).1; + } + } + out[s] = [u, v, w]; + } + out + } + + /// The momentum source at every u, v, w face at time `t`. + pub fn source_tables(&self, g: Grid3, t: f64) -> (Vec, Vec, Vec) { + let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); + let mut su = vec![0.0; (nx + 1) * ny * nz]; + let mut sv = vec![0.0; nx * (ny + 1) * nz]; + let mut sw = vec![0.0; nx * ny * (nz + 1)]; + if let Some(f) = self.momentum_source.as_ref() { + for k in 0..nz { + for j in 0..ny { + for i in 0..=nx { + su[g.uface(k, j, i)] = f( + i as f64 * dx, + (j as f64 + 0.5) * dy, + (k as f64 + 0.5) * dz, + t, + ) + .0; + } + } + } + for k in 0..nz { + for j in 0..=ny { + for i in 0..nx { + sv[g.vface(k, j, i)] = f( + (i as f64 + 0.5) * dx, + j as f64 * dy, + (k as f64 + 0.5) * dz, + t, + ) + .1; + } + } + } + for k in 0..=nz { + for j in 0..ny { + for i in 0..nx { + sw[g.wface(k, j, i)] = f( + (i as f64 + 0.5) * dx, + (j as f64 + 0.5) * dy, + k as f64 * dz, + t, + ) + .2; + } + } + } + } + (su, sv, sw) + } + + pub(crate) fn reference_flux(&self, g: Grid3) -> f64 { self.fluid.density * self.fluid.reference_velocity * self.fluid.reference_length @@ -409,8 +549,6 @@ impl Piso3Solver { let g = field.grid; let (nx, ny, nz, dx, dy, dz) = (g.nx, g.ny, g.nz, g.dx, g.dy, g.dz); let rho = self.fluid.density; - let b = self.params.boundaries; - let any_outlet = b.any_outlet(); let mut source_scale = 0.0; for k in 0..nz { for j in 0..ny { @@ -434,10 +572,7 @@ impl Piso3Solver { } } } - let reference_flux = self.reference_flux(g); - let inner_stop = (self.params.inner_stop_factor * source_scale) - .max(0.1 * self.params.tolerance * reference_flux) - + 1e-14; + let inner_stop = self.inner_stop(g, source_scale); let problem = self.poisson_problem(field, dt); let mut p_prime = vec![0.0; g.cells()]; if warm_start { @@ -452,8 +587,7 @@ impl Piso3Solver { } } } - // The anchor: the first fluid cell (the 2D `(1, 1)` at k = 0). - let anchor_cell = (!any_outlet).then_some(g.cell(0, 1, 1)); + let anchor_cell = self.anchor_cell(g); let params = MultigridParameters { precision: self.params.poisson_precision, smoother: self.params.poisson_smoother, diff --git a/crates/specialized/rtx-cfd/tests/three_d_device_step.rs b/crates/specialized/rtx-cfd/tests/three_d_device_step.rs new file mode 100644 index 0000000..ffb6a60 --- /dev/null +++ b/crates/specialized/rtx-cfd/tests/three_d_device_step.rs @@ -0,0 +1,561 @@ +//! 3D Stage 1, gates 4–6 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 three_d_device_step -- --nocapture` +#![cfg(feature = "cuda")] + +use rtx_cfd::solvers::incompressible::three_d::piso_device::Piso3Device; +use rtx_cfd::solvers::incompressible::three_d::{ + Boundaries3, FlowField3D, Fluid3, Grid3, Piso3Parameters, Piso3Solver, SideBoundary3, +}; +use rtx_cfd::solvers::incompressible::{ConvectionScheme, MgSmoother}; +use std::f64::consts::PI; + +const RHO: f64 = 1.0; + +fn fluid(mu: f64) -> Fluid3 { + Fluid3 { + 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: SideBoundary3, tight: bool) -> Piso3Parameters { + Piso3Parameters { + corrector_steps: 2, + tolerance: if tight { 1e-12 } else { 1e-8 }, + inner_stop_factor: if tight { 1e-6 } else { 1e-2 }, + boundaries: Boundaries3 { + z0: z, + z1: z, + ..Boundaries3::default() + }, + // The device V-cycle is red-black; the host uses the same so the + // preconditioners match. + poisson_smoother: MgSmoother::RedBlack, + convection_scheme: scheme, + ..Piso3Parameters::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: &FlowField3D, b: &FlowField3D) -> (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, +) -> (Piso3Solver, Piso3Solver, Grid3, 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 = Piso3Solver::new(fluid(mu), params(scheme, SideBoundary3::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(), + Grid3 { + nx: n, + ny: n, + nz: n, + dx: h, + dy: h, + dz: h, + }, + dt, + ) +} + +fn run_pair( + host: Piso3Solver, + dev: Piso3Solver, + g: Grid3, + dt: f64, + steps: usize, + label: &str, +) -> (f64, f64) { + let mut host = host; + let mut fh = FlowField3D::new(g); + let mut device = Piso3Device::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 = FlowField3D::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 = Grid3 { + nx: n, + ny: n, + nz: n, + dx: h, + dy: h, + dz: h, + }; + for tight in [false, true] { + let mk = || { + let mut s = Piso3Solver::new( + fluid(NU * RHO), + params(ConvectionScheme::Upwind, SideBoundary3::Velocity, tight), + ); + s.set_boundary_velocity(|x, y, z, t| exact(x, y, z, t)); + s + }; + let mut host = mk(); + let mut fh = FlowField3D::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 = Piso3Device::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 = FlowField3D::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 { + 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, SideBoundary3::SlipWall, false), + (1, 1.0, SideBoundary3::SlipWall, true), + (4, h, SideBoundary3::Periodic, true), + ] { + let g = Grid3 { + nx: n, + ny: n, + nz, + dx: h, + dy: h, + dz, + }; + let mk = || { + let mut s = Piso3Solver::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 = FlowField3D::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 = Piso3Device::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 = FlowField3D::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: &FlowField3D, 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 = Grid3 { + nx: n, + ny: n, + nz: 1, + dx: h, + dy: h, + dz, + }; + let mk = || { + let mut s = Piso3Solver::new( + fluid(MU_P), + params(ConvectionScheme::Upwind, SideBoundary3::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 = FlowField3D::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 = Piso3Device::new(mk(), g); + device.upload(&fh); + let mut fd = FlowField3D::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 = Grid3 { + 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 = Piso3Solver::new( + fluid(mu), + params( + ConvectionScheme::TvdVanAlbada, + SideBoundary3::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 = Piso3Device::new(s, g); + let f = FlowField3D::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 + ); + } +}