perf(cuda): compile every kernel once (wheel 93.1 -> 42.4 MB)

[GeophyAI]

Why

The 0.3.3 wheel is 93.1 MB, 7 MB under PyPI's 100 MB per-file limit. In its cu12 core, 65.4 of the 111.8 MB of SASS were duplicate copies of the same kernels:

• A kernel template launched from a header (a driver hook or a LAUNCH_* macro) is instantiated, for five stencil orders × every arch of the fat binary, in every translation unit that includes the header, whether that unit calls it or not. That means an equation's forward.cu and backward.cu, and every equation sharing a kernel family (Elastic / DASMu / ViscoElastic, Elastic3D / DASMu3D, Acoustic / ViscoAcoustic). Each 2-D elastic kernel had six copies.
• A static __global__ kernel defined in a header is compiled into every unit that includes it, referenced or not.

What

• cuda/launch/by_order.cuh: a table of a kernel template's five stencil-order specializations (2/4/6/8/-1).
  • Each family's kernels.cu defines its tables; it is the only translation unit that names the specializations.
  • The order-dispatch macros in kernels.cuh launch through the table, so call sites are unchanged.
  • The table's pointer type spells out the kernel's parameter types and must match exactly (a mismatch does not compile). The launch is the one nvcc emits for kernel<N><<<>>>.
  • Not extern template: nvcc 12.9 warns that extern __global__ templates stop being supported in whole-program mode once -static-global-template-stub defaults to true (#20279), and nvcc 13 requires a named specialization to be defined in the unit that names it (#20280).
• Kernels defined in headers move to their kernels.cu, behind a declaration in the header: air-clear, the 3-D carrier capture, LSRTM compute_v2_lambda, the VRZ utt correction and kappa kernels, and the TTI-2nd rho correction.
• visco_acoustic2d's visco_ops launchers move to a new spectral_ops.cu.
• Two call sites relied on a kernel's default arguments, which a function pointer does not carry: the das_mu2d / das_mu3d NOPML carriers. They now pass the same nullptrs explicitly.
• Six macros that nothing invokes (VRZ *_ADJOINT / BUILD_VRZ_ADJOINT_FIELDS*, elastic3d *_ADJOINT) are left as they were: a table would instantiate kernels the core never contained.
• Docs:
  • The add-an-equation checklist in docs/dev/cuda_drivers.md states the convention.
  • The sizes in docs/dev/release_cores.md are updated.
  • CHANGELOG entry.

Sizes (manylinux_2_28 images, the same arch lists as 0.3.3):

	0.3.3	this PR
cu12 core	157.3 MB	70.8 MB
cu13 core	189.9 MB	85.8 MB
wheel	93.1 MB	42.4 MB

Tested

• Machine code (both cores built in the same images as 0.3.3): every kernel × arch SASS, compared as encoded instruction words, and every PTX entry are byte-identical to 0.3.3's.
  • cu12: 4254 SASS + 709 PTX. cu13: 4963 SASS + 709 PTX.
  • The same 709 kernels as 0.3.3, each now once per arch instead of 1380 copies per arch.
• Gradient consistency matrix (test/backend_gradient_matrix.py --backends eager/gpu c-cuda --scale cuda-suite, 29 cases): the 0.3.3 core's log and this core's are identical line for line, apart from timestamps.
  • Both runs give pass=56, fail=12.
  • The 12 failures are all present in 0.3.3: 9 recursive-checkpoint cells TTI / VTI do not offer, DASZhao3D boundary saving (unsupported), and LSRTM3D's source-wavelet gradient (full and chunk), which gradient_cases.py keeps failing on purpose.
• Old-core vs new-core A/B on an RTX 6000 Ada:
  • Coverage: 35 cases (the matrix's 29 plus ElasticVRR, ViscoAcoustic, ViscoElastic and the 2-D / 3-D APM topography paths); spatial orders 2/4/6/8/10, so every table slot; full / boundary saving / chunk / recursive checkpointing; 3470 records and gradients.
  • Every array the 0.3.3 core reproduces bitwise from run to run is bitwise equal.
  • The 0.3.3 core does not reproduce some arrays bitwise from run to run: DASZhao 2-D / 3-D in every mode; acoustic3d, LSRTM3D, VRZ3D and ViscoAcoustic in the checkpoint modes; acoustic2d at order 10 in the checkpoint modes. Over 8 + 8 interleaved runs these differ old-vs-new exactly as they differ old-vs-old.
• Full suite on this core: 1895 passed, 44 skipped, the same as 0.3.3.
• Timing unchanged for acoustic2d, ViscoAcoustic (120 runs each: median 203.2 vs 202.8 ms), Elastic and Elastic3D.
• Not covered:
  • The cu13 core only runs on the A100s (KW60443's driver cannot run CUDA 13); its machine code is byte-identical to 0.3.3's.
  • Multi-GPU domain decomposition beyond the single-GPU two-tile tests.

🤖 Generated with Claude Code