Skip to content

fix(visco): per-shot gbar for batched models #124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
GeophyAI merged 1 commit into from
Oct 21, 2026
Merged

fix(visco): per-shot gbar for batched models #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 6 additions & 6 deletions src/sweep/csrc/cuda/equations/visco_acoustic2d/backward.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,18 +224,18 @@ void adjoint_damping_extra(AcousticWavefieldTensor& adj, const ViscoSpectral& d,
if (d.disp) {
const Buf X = visco_acoustic2d_real_alias(ws.C1);
visco_ops::binary<visco_ops::MulOp>(X.data_ptr<float>(), d.Gd1.data_ptr<float>(), 1, u_now, 1, n); // Gd1 * u_now
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.Dk2.data_ptr<float>(), ws); // C1.re = L_Dk2
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.Dk2, ws); // C1.re = L_Dk2
visco_ops::real_copy(ws.R1.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n); // R1 = L_Dk2
visco_ops::binary<visco_ops::MulOp>(X.data_ptr<float>(), d.Gd2.data_ptr<float>(), 1, u_now, 1, n); // Gd2 * u_now (L_Dk2 is in R1)
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.Dfrac.data_ptr<float>(), ws); // C1.re = L_Dfrac
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.Dfrac, ws); // C1.re = L_Dfrac
visco_ops::inplace<visco_ops::SubOp>(ws.R1.data_ptr<float>(), visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); // R1 = L_Dk2 - L_Dfrac = e
m = ws.R1.data_ptr<float>(); m_stride = 1;
}
if (d.active) {
const Buf X = visco_acoustic2d_real_alias(ws.C1);
visco_ops::binary<visco_ops::SubOp>(X.data_ptr<float>(), u_prev, 1, u_now, 1, n); // u_prev - u_now
visco_ops::inplace<visco_ops::MulOp>(X.data_ptr<float>(), d.Gp.data_ptr<float>(), 1, n); // == Gp * (u_prev - u_now)
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.kmul.data_ptr<float>(), ws); // C1.re = L
visco_acoustic2d_lop_into(X.data_ptr<float>(), d.kmul, ws); // C1.re = L
if (d.disp)
visco_ops::inplace<visco_ops::AddOp>(ws.R1.data_ptr<float>(), visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n); // e + m == m + e
else {
Expand Down Expand Up @@ -272,7 +272,7 @@ void accumulate_grad_A(float* grad_A,
{
if (!d.active) return;
const int64_t n = ws.C1.numel() / 2;
visco_acoustic2d_lop_into(du, d.kmul.data_ptr<float>(), ws); // C1.re = L(du)
visco_acoustic2d_lop_into(du, d.kmul, ws); // C1.re = L(du)
const Buf P = visco_acoustic2d_real_alias(ws.C0);
visco_ops::binary<visco_ops::MulOp>(P.data_ptr<float>(), lam, 1,
visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n);
Expand All @@ -294,10 +294,10 @@ void accumulate_grad_disp(float* grad_B1, float* grad_B2,
// add_(P, Scalar(dt2)): the double dt*dt, loaded as the float alpha
const float dt2 = static_cast<float>(static_cast<double>(dt) * static_cast<double>(dt));
const Buf P = visco_acoustic2d_real_alias(ws.C0);
visco_acoustic2d_lop_into(u_it, d.Dk2.data_ptr<float>(), ws);
visco_acoustic2d_lop_into(u_it, d.Dk2, ws);
visco_ops::binary<visco_ops::MulOp>(P.data_ptr<float>(), lam, 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n);
visco_ops::axpy(grad_B1, P.data_ptr<float>(), 1, dt2, n);
visco_acoustic2d_lop_into(u_it, d.Dfrac.data_ptr<float>(), ws);
visco_acoustic2d_lop_into(u_it, d.Dfrac, ws);
visco_ops::binary<visco_ops::MulOp>(P.data_ptr<float>(), lam, 1, visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n);
visco_ops::axpy(grad_B2, P.data_ptr<float>(), 1, -dt2, n);
}
Expand Down
45 changes: 27 additions & 18 deletions src/sweep/csrc/cuda/equations/visco_acoustic2d/kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,20 +112,25 @@ struct ViscoSpectral {
Buf Gd1, Gd2; // dt^2 * B1, dt^2 * B2
};

// A grid is shared by the batch, (nz, nx), or one per batch element,
// (B, nz, nx): per-shot models whose frozen exponents differ (the filter
// strides it per batch index, see visco_acoustic2d_lop_into).
inline void visco_acoustic2d_check_grid(
const Buf& g, int nz, int nx, const char* name)
const Buf& g, int64_t batch, int nz, int nx, const char* name)
{
SWEEP_CHECK(g.dim() == 2 && g.size(0) == nz && g.size(1) == nx,
const bool shared = g.dim() == 2 && g.size(0) == nz && g.size(1) == nx;
const bool per_batch = g.dim() == 3 && g.size(0) == batch && g.size(1) == nz && g.size(2) == nx;
SWEEP_CHECK(shared || per_batch,
"visco eq_aux grid ", name,
" must be (nz_runtime, nx_runtime) = (", nz, ", ", nx,
"), got ", g.sizes());
") or (B, nz_runtime, nx_runtime) with B = ", batch, ", got ", g.sizes());
SWEEP_CHECK(g.is_cuda() && g.scalar_type() == BoundaryDtype::FP32,
"visco eq_aux grid ", name, " must be a float32 CUDA tensor");
}

// The term switches and filter grids from the eq_aux composition (no tables).
inline ViscoSpectral visco_acoustic2d_spectral_grids(
const BufList& eq_aux, int nz, int nx)
const BufList& eq_aux, int64_t batch, int nz, int nx)
{
ViscoSpectral s;
const size_t n = eq_aux.size();
Expand All @@ -134,13 +139,13 @@ inline ViscoSpectral visco_acoustic2d_spectral_grids(
s.disp = (n >= 2);
if (s.active) {
s.kmul = eq_aux[0];
visco_acoustic2d_check_grid(s.kmul, nz, nx, "D_loss");
visco_acoustic2d_check_grid(s.kmul, batch, nz, nx, "D_loss");
}
if (s.disp) {
s.Dk2 = eq_aux[n - 2];
s.Dfrac = eq_aux[n - 1];
visco_acoustic2d_check_grid(s.Dk2, nz, nx, "D_k2");
visco_acoustic2d_check_grid(s.Dfrac, nz, nx, "D_frac");
visco_acoustic2d_check_grid(s.Dk2, batch, nz, nx, "D_k2");
visco_acoustic2d_check_grid(s.Dfrac, batch, nz, nx, "D_frac");
}
return s;
}
Expand All @@ -161,7 +166,7 @@ inline ViscoSpectral visco_acoustic2d_make_spectral_from(
{
SWEEP_CHECK(models.size() == 4, what, ": visco_acoustic2d expects the prepared models "
"(vp_step, B1, B2, A); got ", models.size());
ViscoSpectral s = visco_acoustic2d_spectral_grids(eq_aux, nz, nx);
ViscoSpectral s = visco_acoustic2d_spectral_grids(eq_aux, models[0].size(0), nz, nx);
const float dt2 = dt * dt; // float on the host: the scalar of the replaced ``model * (dt * dt)`` products
const auto c = derived::visco_coefficients(
derived_models, models[1], models[2], models[3], s.active, s.disp, mode, dt, dt2, what);
Expand Down Expand Up @@ -195,7 +200,7 @@ inline ViscoSpectral visco_acoustic2d_make_spectral_from(
// builds the same CuFFTParams and issues the same three calls, with the
// work area taken from the pool. The same plan on the same data gives the
// same bits.
// * kmul * F, a float (nz, nx) grid times the complex (B, 1, nz, nx)
// * kmul * F, a float (nz, nx) [or (B, nz, nx)] grid times the complex (B, 1, nz, nx)
// spectrum: mul_kernel_cuda on the common dtype complex64, kmul loaded as
// (k, 0) and c10::complex operator* evaluated as (a*c - b*d, a*d + b*c)
// with (a, b) = (k, 0). Here: at::mul_out(C0, kmul, F) -- functional mul
Expand Down Expand Up @@ -461,16 +466,21 @@ inline ViscoScratch visco_acoustic2d_bind_scratch(
// is consumed by the promote before C1 is overwritten) but NOT in C0's.
// ---------------------------------------------------------------------------
// ``x``: a contiguous float grid of the slot's shape (a wavefield, a real
// alias, a history row); ``kmul``: the (nz, nx) filter. The result is the
// alias, a history row); ``kmul``: the (nz, nx) filter, or (B, nz, nx) with
// one filter per batch element (cmul_grid indexes it modulo its own size, so
// the shared grid is read exactly as before). The result is the
// real part of C1 -- visco_acoustic2d_real_ptr(ws.C1) at VISCO_REAL_STRIDE --
// valid until C1 is next written.
inline void visco_acoustic2d_lop_into(const float* x, const float* kmul, ViscoScratch& ws)
inline void visco_acoustic2d_lop_into(const float* x, const Buf& kmul, ViscoScratch& ws)
{
const int64_t n = ws.C1.numel() / 2;
const int64_t ng = ws.C1.size(-3) * ws.C1.size(-2); // (nz, nx) of the (..., nz, nx, 2) slot
const int64_t ng = kmul.numel();
SWEEP_CHECK(ng == ws.C1.size(-3) * ws.C1.size(-2) || ng == n,
"visco filter grid of ", ng, " cells matches neither one (nz, nx) plane nor the "
"whole ", n, "-cell batch of the spectral slot");
visco_ops::promote(visco_acoustic2d_complex_ptr(ws.C0), x, n);
ws.fft->forward(ws.C0, ws.C1, ws.fft_ws);
visco_ops::cmul_grid(visco_acoustic2d_complex_ptr(ws.C0), kmul, visco_acoustic2d_complex_ptr(ws.C1), n, ng);
visco_ops::cmul_grid(visco_acoustic2d_complex_ptr(ws.C0), kmul.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n, ng);
ws.fft->inverse(ws.C0, ws.C1, ws.fft_ws);
}

Expand All @@ -485,7 +495,7 @@ inline void visco_acoustic2d_lop_into(const float* x, const float* kmul, ViscoSc
// subtracted in place, as before.
inline void visco_acoustic2d_apply_damping_into(
AcousticWavefieldTensor& wf,
const Buf& kmul, // (nz, nx) D_loss grid
const Buf& kmul, // (nz, nx) or (B, nz, nx) D_loss grid
const Buf& dt2A, // (B, 1, nz, nx) dt^2 * A (derived table)
float dt, ViscoScratch& ws)
{
Expand All @@ -495,7 +505,7 @@ inline void visco_acoustic2d_apply_damping_into(
wf.u_now_t.template data_ptr<float>(), 1,
wf.u_prev_t.template data_ptr<float>(), 1, n);
visco_ops::scale(dudt.data_ptr<float>(), 1.0f / dt, n); // div_(dt): times the reciprocal
visco_acoustic2d_lop_into(dudt.data_ptr<float>(), kmul.data_ptr<float>(), ws); // C1.re = L; dudt is gone
visco_acoustic2d_lop_into(dudt.data_ptr<float>(), kmul, ws); // C1.re = L; dudt is gone
const Buf R = visco_acoustic2d_real_alias(ws.C0);
visco_ops::binary<visco_ops::MulOp>(R.data_ptr<float>(), dt2A.data_ptr<float>(), 1,
visco_acoustic2d_real_ptr(ws.C1), VISCO_REAL_STRIDE, n);
Expand Down Expand Up @@ -527,17 +537,16 @@ inline void visco_acoustic2d_apply_spectral_into(
SWEEP_CHECK(ws.C2.defined(), "visco_acoustic2d: the dispersion pipeline needs the C2 slot "
"(bind the scratch with dispersion on, in forward / ckpt / recursive mode)");
const int64_t n = ws.C1.numel() / 2;
const int64_t ng = ws.C1.size(-3) * ws.C1.size(-2);
float2* C0 = visco_acoustic2d_complex_ptr(ws.C0);
visco_ops::promote(C0, wf.u_now_t.template data_ptr<float>(), n);
ws.fft->forward(ws.C0, ws.C1, ws.fft_ws); // C1 = F, alive across both products
visco_ops::cmul_grid(C0, s.Dk2.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n, ng); // C0 = Dk2 * F
visco_ops::cmul_grid(C0, s.Dk2.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n, s.Dk2.numel()); // C0 = Dk2 * F
ws.fft->inverse(ws.C0, ws.C2, ws.fft_ws); // C2 = IFFT2(Dk2 * F)
const Buf R = visco_acoustic2d_real_alias(ws.C0); // the product in C0 is dead
visco_ops::binary<visco_ops::MulOp>(R.data_ptr<float>(), s.Gd1.data_ptr<float>(), 1,
visco_acoustic2d_real_ptr(ws.C2), VISCO_REAL_STRIDE, n);
visco_ops::inplace<visco_ops::AddOp>(wf.u_next_t.template data_ptr<float>(), R.data_ptr<float>(), 1, n);
visco_ops::cmul_grid(C0, s.Dfrac.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n, ng); // C0 = Dfrac * F (R's bits are dead)
visco_ops::cmul_grid(C0, s.Dfrac.data_ptr<float>(), visco_acoustic2d_complex_ptr(ws.C1), n, s.Dfrac.numel()); // C0 = Dfrac * F (R's bits are dead)
ws.fft->inverse(ws.C0, ws.C2, ws.fft_ws); // C2 = IFFT2(Dfrac * F)
visco_ops::binary<visco_ops::MulOp>(R.data_ptr<float>(), s.Gd2.data_ptr<float>(), 1,
visco_acoustic2d_real_ptr(ws.C2), VISCO_REAL_STRIDE, n);
Expand Down
40 changes: 32 additions & 8 deletions src/sweep/equations/visco_acoustic.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,7 @@ def prepare_models(self, models):
vp, Q, omega = models
op = self.op
gamma = op.arctan(1.0 / Q) / np.pi
self._gbar = op.mean(gamma)
self._gbar = self._average_exponent(gamma)
if self.phase_shift:
csq_fac = op.cos(np.pi * gamma / 2)
vp_step = vp * csq_fac
Expand All @@ -290,10 +290,29 @@ def prepare_models(self, models):
A = op.zeros_like(vp)
return (vp_step, B1, B2, A)

def _average_exponent(self, gamma):
"""The frozen exponent ``gbar``: the mean of ``gamma`` over the model.

A shared model ``(nz, nx)`` gets one scalar. Per-shot batched models
(``(B, 1, nz, nx)``: the eager layout, and every impl='c' model) get
one mean per batch element, shaped ``(B, 1, 1, 1)``, so a shot's
filter exponent depends on its own Q only -- each mean is the same
reduction a single-shot run takes over that model. When all the
per-element means agree (a shared model the C path repeated across
the batch), the historical whole-batch scalar is kept bit-for-bit."""
op = self.op
if gamma.ndim <= 2 or gamma.shape[0] == 1:
return op.mean(gamma)
g = op.stack([op.mean(gamma[b]) for b in range(gamma.shape[0])])
if self.backend == 'torch' and bool((g == g[0]).all()):
return op.mean(gamma)
return g.reshape((gamma.shape[0],) + (1,) * (gamma.ndim - 1))

def _spectral_grids(self, k):
"""(D_k2, D_frac, D_loss) fractional-power grids on the runtime |k|
grid, using the average exponent ``self._gbar`` from the latest
:meth:`prepare_models` call. Cached per (gbar tensor, k shape): one
:meth:`prepare_models` call (one per batch element for per-shot
models, giving ``(B, 1, nz, nx)`` fractional grids). Cached per (gbar tensor, k shape): one
rebuild per forward on torch (prepare runs once outside the loop), one
per step inside a jax trace (where it stays trace-local). ``k = 0``
rows go through a safe branch so the pow's exponent-gradient stays
Expand Down Expand Up @@ -352,23 +371,28 @@ def c_eq_aux(self, prop):
from .base import init_wavenumbers
shape = tuple(int(v) for v in prop.shape_cuda)[-2:]
h = float(np.asarray(prop._grid_spacing).reshape(-1)[0])
gbar = float(self._gbar.detach().cpu())
key = (shape, h, str(prop.dev), round(gbar, 9),
# One exponent for a shared model; per-shot models with differing Q
# carry one per batch element, and the fractional grids become
# (B, nz, nx) stacks (the CUDA filter strides them per batch index).
gbars = [float(g) for g in self._gbar.detach().cpu().reshape(-1)]
key = (shape, h, str(prop.dev), tuple(round(g, 9) for g in gbars),
self.phase_shift, self.amplitude_damping)
cached = getattr(self, "_c_kmul_cache", None)
if cached is None or cached[0] != key:
k_np, _, _ = init_wavenumbers(shape, h)
def t(a):
return torch.as_tensor(
a.astype(np.float32), device=prop.dev).contiguous()
def frac(p):
g = [np.where(k_np > 0, np.power(k_np, 2.0*gbar + p, where=k_np > 0), 0.0)
for gbar in gbars]
return t(g[0] if len(g) == 1 else np.stack(g))
grids = []
if self.amplitude_damping:
grids.append(t(np.where(
k_np > 0, np.power(k_np, 2.0*gbar + 1.0, where=k_np > 0), 0.0)))
grids.append(frac(1.0))
if self.phase_shift:
grids.append(t(k_np * k_np))
grids.append(t(np.where(
k_np > 0, np.power(k_np, 2.0*gbar + 2.0, where=k_np > 0), 0.0)))
grids.append(frac(2.0))
self._c_kmul_cache = (key, tuple(grids))
return self._c_kmul_cache[1]

Expand Down
Loading
Loading