perf(eager): CPML on strips; boundary saving for more equations; async cpu and disk storage

[GeophyAI]

Six commits on the eager propagator: CPML memory on the absorbing strips, boundary saving for more equations, faster host storage and a new disk storage for the boundary ring, a notebook that compares eager and c, and a faster CPU test suite.

Changes

• CPML on the absorbing strips. Turned on by EagerOptions(cpml_strips=True); off by default.
  • The CPML memory fields live only on the absorbing strips.
  • Equations mark FieldSpec.cpml_axis and pass the field the derivative comes from to cpml_update / cpml_axis_update (equations/_cpml.py).
  • The eager propagator lays out the strips, boundary saving included. 17 equations are covered.
  • A one-time stderr notice now says when torch.compile is building the step.
• Boundary saving for more equations.
  • Elastic3D, DASMu, DASMu3D and ElasticVRR now split their step into two reversible halves, shared by func and interior_substeps. Eager boundary saving therefore covers them, free surface included.
  • Elastic 2D now refuses APM topography, as Elastic3D already did. Under bf16/int8 storage the looser self-check tolerance let a wrong gradient through.
• Refusals. A new equation flag, supports_boundary_saving_eager, is the eager/PropJax counterpart of supports_boundary_saving_c. ViscoAcoustic, AcousticVTI and AcousticTTI set it to False, so a boundary request raises. Their boundary-saving gradients came out 2–10% off.
• Ring storage.
  • Each step's ring is stored as one row.
  • storage='cpu' copies through pinned memory on a side stream and prefetches for the reverse pass.
  • storage='disk' is new for eager. It stages the ring in pinned slots, and background threads write and read four temp files. I/O errors raise.
  • pinned_memory, disk_dir, transfer_interval and ring_buffers now reach the eager backend; before, they were dropped silently.
• Docs.
  • Notebook 33 compares eager, eager with strips, and c on forward, full-storage and boundary-saving runs. It reports time, compile time, peak memory and gradient agreement.
  • Notebooks 07 and 19, the extending guide and the propagator guide are updated.
• CPU test suite. It had grown from about 4 to 17 minutes. Every new eager propagator compiled its step with Inductor, and the finite-difference tests build one per perturbation.
  • The CPU physics tests now run uncompiled.
  • test_compiled_stencil_slices compiles each case once and checks three things on it: no conv, one graph, and a match with the uncompiled step. It now also covers per-edge free surfaces, image topography and DASMu.
  • Locally the suite takes 6.5 minutes instead of 13.

Validation

• Default path (strips off). Records and gradients are bit-identical to dev, both uncompiled on CPU and compiled on GPU.
• Strips vs full grid. Bit-identical uncompiled for full, checkpoint and boundary-saving modes. This holds with and without a free surface, including per-edge surfaces and image-method topography. Compiled results agree within compile noise. The gradient matrix's eager-strips column passes 58/58.
• Boundary saving vs full tape, all models, compiled or not: DASMu 0.8–1.4e-4, DASMu3D 3–6e-5, ElasticVRR 0.8–3.2e-4.
• Storage. Every storage gives the same gradient as GPU storage, bit for bit: cpu, pageable cpu, and disk at several intervals and slot counts, in fp32 and int8.
• Free surface. A free-surface matrix over the elastic family passes: compile × strips × memory mode × storage.
• Tests. The eager, boundary-saving, memory-option, DAS and VRR suites pass.

On a V100, the cpu ring storage now runs at GPU-storage speed while keeping the ring off the device: Elastic3D 192³ goes from 11.8 s to 7.8 s, against 7.7 s with GPU storage.

🤖 Generated with Claude Code