cuwave 0.6.0__tar.gz → 0.7.0__tar.gz

{cuwave-0.6.0 → cuwave-0.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuwave
-Version: 0.6.0
+Version: 0.7.0
 Summary: GPU finite-difference wave solver with differentiable adjoints
 Author-email: Leon Herrmann <herrmann.leon@pm.me>
 License-Expression: MIT

{cuwave-0.6.0 → cuwave-0.7.0}/cuwave/sensitivity.py

@@ -637,11 +637,9 @@ def reconstruction_sensitivity(
         p = adjoint_step(P[m % 2], P[1 - m % 2], P[m % 2])  # lambda^n
         adjoint_excitation(p, signal, row)
         bc_step(p)
-        # inertia pairs lambda^n with the whole triplet, stiffness with its middle slot
-        gradient_step(*triplet(m), p)
+        return p
 
-    def backward(m, row, signal, source_rows, strip_rows=None):
-        adjoint(m, row, signal)
+    def reverse(m, row, source_rows, strip_rows=None):
         a, b, c = triplet(m)
         # read backwards, so the source rides two steps ahead of the state it rebuilds
         c = reverse_step(b, a, c)
@@ -650,6 +648,28 @@ def reconstruction_sensitivity(
             replay_strip(c, strip_rows, row)
         bc_step(c)
 
+    # the reverse and the adjoint march meet only in the gradient, so they may overlap
+    side = cp.cuda.Stream(non_blocking=True)
+    forked = cp.cuda.Event(disable_timing=True)
+    joined = cp.cuda.Event(disable_timing=True)
+    # past L2 the two would evict each other's fields
+    carried = U.nbytes + P.nbytes + sum(field.nbytes for field in grads.values())
+    overlap = carried <= cp.cuda.Device().attributes["L2CacheSize"]
+
+    # reverse step m on the side stream, beside adjoint step m + 1
+    def backward(m, row, signal, source_rows, strip_rows=None):
+        main = cp.cuda.get_current_stream()
+        lane = side if overlap else main
+        forked.record(main)
+        lane.wait_event(forked)
+        with lane:
+            reverse(m, row, source_rows, strip_rows)
+            joined.record(lane)
+        p = adjoint(m + 1, row, signal)
+        main.wait_event(joined)
+        # inertia pairs lambda^n with the whole triplet, stiffness with its middle slot
+        gradient_step(*triplet(m + 1), p)
+
     def leaf(k):
         s, e = starts[k], starts[k + 1]
         # row t - s + 2 is u^t, the two the segment starts from read off the slots
@@ -666,8 +686,9 @@ def reconstruction_sensitivity(
         if not s:
             # the forward rotation left the last three states live, which is the seed
             state["seed"] = float(cp.linalg.norm(triplet(m0)[2] * valid))
-        # row k of each is reverse step m0 + k, contiguous since the kernels index by row
-        windows = [Window(signal[m0 : m1 - 1])]
+        gradient_step(*triplet(m0), adjoint(m0, m0, signal))
+        # row k is reverse step m0 + k and adjoint step m0 + 1 + k, as kernels index rows
+        windows = [Window(signal[m0 + 1 : m1])]
         windows.append(Window(source.signal[s : e - 1][::-1].copy()))
         if num_strip:
             windows.append(Window(buffer[: e - s - 1][::-1].copy(), dtype=sim.dtype))
@@ -680,7 +701,6 @@ def reconstruction_sensitivity(
                     windows,
                     [U, P, *adjoint_memory],
                 )
-        adjoint(m1 - 1, m1 - 1, signal)
 
     for action, *segments in checkpoint_schedule(len(starts) - 1, checkpoints or 0):
         if action == "advance":

{cuwave-0.6.0 → cuwave-0.7.0}/cuwave/utils.py

@@ -428,7 +428,9 @@ def misfit_gradient(
         adjoint: which variant computes each shot, `reconstruction_sensitivity`
             where the stored history no longer fits.
         devices: GPUs to spread the shots over, every `len(devices)`-th on each, or
-            None for the current one.
+            None for the current one. A GPU listed twice runs two shots at once,
+            which pays only while the fields of both fit in its L2 cache together
+            and slows a larger grid down.
 
     Returns:
         (cost, gradient), the gradient a field over the padded grid, already

{cuwave-0.6.0 → cuwave-0.7.0}/cuwave.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuwave
-Version: 0.6.0
+Version: 0.7.0
 Summary: GPU finite-difference wave solver with differentiable adjoints
 Author-email: Leon Herrmann <herrmann.leon@pm.me>
 License-Expression: MIT

{cuwave-0.6.0 → cuwave-0.7.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "cuwave"
-version = "0.6.0"
+version = "0.7.0"
 description = "GPU finite-difference wave solver with differentiable adjoints"
 readme = "README.md"
 license = "MIT"