Fix GPU spectral bugs on AMD/ROCm (alignment, symbol visibility, adjoint kernel miscompilation)#242
Merged
msuchard merged 15 commits intoJul 10, 2026
Conversation
BEAGLE dlopen()s each plugin .so and resolves plugin_init via dlsym; every other symbol is an implementation detail. Without -fvisibility=hidden, identically-named symbols from a shared header-only template base class compiled into multiple co-loaded plugin .so's can interpose on each other at dynamic-link time, corrupting virtual dispatch across plugins loaded in the same process (observed as a SIGSEGV in clSetKernelArg when the OpenCL and OpenCL-Spectral plugins were both loaded). Compile hmsbeagle-cpu/-cpu-sse/-cpu-spectral with hidden default visibility and explicitly re-export plugin_init as the one symbol dlsym needs.
Same cross-plugin symbol-interposition fix as the CPU plugins: hide all symbols in hmsbeagle-cuda by default and explicitly re-export plugin_init, the only symbol BEAGLE's dlopen()/dlsym() loader actually needs.
…gin_init Same cross-plugin symbol-interposition fix as the CPU plugins: hide all symbols in hmsbeagle-cuda-spectral by default and explicitly re-export plugin_init, the only symbol BEAGLE's dlopen()/dlsym() loader actually needs.
Same cross-plugin symbol-interposition fix as the CPU plugins: hide all symbols in hmsbeagle-opencl by default and explicitly re-export plugin_init, the only symbol BEAGLE's dlopen()/dlsym() loader actually needs. This plugin is one of the two directly implicated in the original SIGSEGV (co-loaded with OpenCL-Spectral, whose shared template base class was interposing across the two .so's before this fix).
…lugin_init Same cross-plugin symbol-interposition fix as the CPU plugins: hide all symbols in hmsbeagle-opencl-spectral by default and explicitly re-export plugin_init, the only symbol BEAGLE's dlopen()/dlsym() loader actually needs. This plugin is the other one directly implicated in the original SIGSEGV in clSetKernelArg (co-loaded with plain OpenCL; a shared header-only template base class was interposing across the two .so's before this fix, corrupting virtual dispatch).
…IGNED_SUB_BUFFER_OFFSET dSpectralDistancesOrigin and dEvecTOrigin/dIevcTOrigin are single large device allocations sliced into per-matrix/per-eigendecomposition sub-buffers via CreateSubPointer. AMD's OpenCL implementation enforces sub-buffer offset alignment strictly (CL_MISALIGNED_SUB_BUFFER_OFFSET on GPU instance creation); NVIDIA/Apple's do not, which is why this was invisible until tested on AMD hardware. Round both strides up with AlignMemOffset() so every sub-buffer offset is a multiple of the device's required alignment. Since the aligned stride can now be wider than kCategoryCount elements, record it in the new kSpectralDistanceStrideElements member and use that (not kCategoryCount) in the adjoint offset-queue's per-matrix index computation, which walks the same pooled buffer.
…t just EIGEN_COMPLEX Spectral kernels (kernelsSpectralIfDef*.cu / kernelsSpectral*.cu) always read a real+imaginary pair per eigenstate, regardless of whether the model actually has complex eigenvalues: SPECTRAL_EIGENVALS_GPU unconditionally indexes eigenValues[PADDED_STATE_COUNT + state]. But BeagleGPUImpl::createInstance only sized/copied the eigenvalue device buffer as complex (2x width) when BEAGLE_FLAG_EIGEN_COMPLEX was set, so for real-eigenvalue models under the spectral representation, that kernel read ran past the data setEigenDecomposition actually wrote, pulling in uninitialized device memory and producing wrong postorder/preorder partials (~1e-2 to 1e-1 error vs. CPU). Fix: widen kEigenValuesSize whenever either BEAGLE_FLAG_EIGEN_COMPLEX or BEAGLE_FLAG_SPECTRAL_REPRESENTATION is set. This also required copying BEAGLE_FLAG_SPECTRAL_REPRESENTATION into kFlags in createInstance, which rebuilds kFlags from scratch and was dropping that bit entirely (everything downstream keys off kFlags, not the caller's raw preference/requirement flags). BeagleGPUSpectralImpl::setEigenDecomposition recomputes this same sizing locally (kEigenValuesSize itself is private to BeagleGPUImpl) and needed the identical widening to stay in sync.
…CPU+GPU) Prints eigenvalues/eigenvectors/inverse-eigenvectors as they're set on both the CPU spectral path (EigenDecompositionSpectral::setEigenDecomposition) and the GPU spectral path (BeagleGPUSpectralImpl::setEigenDecomposition), gated behind the BEAGLE_DEBUG_EIGEN environment variable so it's silent by default. Used to compare CPU vs. GPU eigendecomposition inputs side by side when diagnosing spectral GPU numerical bugs; kept as a permanent debugging aid for the same purpose going forward.
LaunchKernel/LaunchKernelConcurrent now print, when BEAGLE_DEBUG_KERNEL_ARGS is set: the kernel name (via clGetKernelInfo), every pointer/int argument, block/grid/work-group dimensions, and an explicit clFinish() with before/after prints bracketing the enqueue. This was the key tool for isolating which specific kernel a GPU hang was in: since clEnqueueNDRangeKernel can succeed while clFinish() never returns, the last kernel that doesn't print "clFinish returned OK" is the one that hung. Silent by default; kept as a permanent debugging aid.
…ted) Prints "[DISPATCH] BASE" / "[DISPATCH] SPECTRAL" on every postorder pruning dispatch, unconditionally (unlike the other debug instrumentation in this branch, this one isn't gated behind an env var). Used to empirically confirm which BeagleGPUImpl vs. BeagleGPUSpectralImpl dispatch path a given plugin/instance was actually taking while diagnosing the cross-plugin symbol-interposition SIGSEGV. Noisy on every pruning call; should be removed or gated behind an env var (e.g. BEAGLE_DEBUG_FLOW, already used elsewhere in this file) before this branch is considered done.
…on of heavy inlined code
Three related fixes to the GPU adjoint gradient kernel machinery
(kernelAdjointMergedN and the SINGLETON/PAIRLEADER macros it dispatches to),
all addressing the same underlying defect class: this ROCm 4.0.1/gfx906
setup's LLVM-AMDGPU backend miscompiles large, register-heavy inlined code
when the kernel is built at PADDED_STATE_COUNT=32 (any state count that
pads above 16, e.g. 17 states) — confirmed not to be a BEAGLE logic bug,
since the identical macros at PADDED_STATE_COUNT=16 work correctly.
1. ADJOINT_ATOMIC_ADD_GPU's CAS-retry loop (float/double atomic add) was
inlined directly at every call site. At PADDED_STATE_COUNT=32 this
caused clEnqueueNDRangeKernel to succeed while clFinish() never
returned. Bisected to require both the atomic_cmpxchg call AND a retry
loop reading its own result. Fixed by moving the loop into its own
__attribute__((noinline)) helper function per precision.
2. Restructuring ADJOINTN_APPLY_COMPLEX_PAIRLEADER's if/else into
if/else-if (to add an explicit `_ri_q < 0` guard) had dropped the
macro's final closing brace for its outer `if (tid == 0) { ... }`
wrapper, breaking OpenCL kernel compilation entirely ("expected '}'").
Verified by a running brace-balance count against the sibling macro
ADJOINTN_APPLY_COMPLEX_SINGLETON (untouched, correct) as ground truth:
SINGLETON was 6-open/6-close balanced, PAIRLEADER was missing one close.
Fixed by restoring the missing brace.
3. Once compiling and no longer hanging, 17-state gradients still showed
NaN. Extracted the per-iteration bodies of both SINGLETON and
PAIRLEADER (each carrying ~30 REAL temporaries) into noinline helper
functions, the same treatment as (1), on the theory that inlining this
much register pressure into the loop was the miscompilation trigger.
This did NOT fix the residual NaN (root cause turned out to be deeper —
see the GPUInterfaceOpenCL.cpp diagnostic addition and
CLUSTER_AGENT_FINDINGS.md for the ongoing investigation), but is kept
in place as a harmless, consistent mitigation for this defect class.
The device gradient buffer is laid out S×S with S=kPaddedStateCount (32 for a 17-state model), but every caller (matching BeagleCPUImpl's convention, confirmed against the real downstream consumer, SpectralBeagleCrossProductDelegate.java, which allocates stateCount*stateCount) sizes outSumDerivatives as kStateCount*kStateCount (289 for 17 states) — unpadded. The old code downloaded and copied all S*S = 1024 values into that 289-element buffer verbatim: a ~5.9 KB heap overflow whenever kPaddedStateCount > kStateCount, i.e. whenever the state count isn't already a "natural" width (4, 16, ...). This was never hit before because the padded (e.g. 17-state) adjoint path always hung or produced garbage earlier in the pipeline before reaching this code. Fixed by copying only the top-left kStateCount×kStateCount submatrix, using the correct stride on each side (S on the device-buffer source, kStateCount on the destination).
…xes it This test previously allocated a kPaddedStateCount-sized gradient buffer and re-indexed down to kStateCount×kStateCount, compensating from the test side for calculateAdjointCrossProducts overflowing an unpadded buffer. Now that the library itself only writes the correct kStateCount² submatrix (see the BeagleGPUSpectralImpl.hpp fix), this workaround is redundant and was actually incorrect against the real calling convention. Removed it so the test allocates and reads exactly like the real beast-mcmc caller (SpectralBeagleCrossProductDelegate.java) and BeagleCPUImpl both do.
…racing Reports the driver-queried CL_KERNEL_WORK_GROUP_SIZE and CL_KERNEL_PREFERRED_WORK_GROUP_SIZE_MULTIPLE for each kernel launch alongside the requested local work size, under the existing BEAGLE_DEBUG_KERNEL_ARGS gate. Added while investigating whether a GPU adjoint kernel's cooperative reduction was silently running with a smaller workgroup than requested; the driver's self-reported capacity ruled that out as the sole explanation but this remains a useful diagnostic for whoever picks up that investigation next (see CLUSTER_AGENT_FINDINGS.md, "Residual 17-state NaN").
…shared reduction Fixes the residual NaN gradients for padded (PADDED_STATE_COUNT=32, e.g. 17-state) models that remained after the brace-mismatch and buffer-overflow fixes: get_local_id(0) was reading as a stuck constant for every work-item in this compiled kernel variant, a ROCm 4.0.1/LLVM-AMDGPU backend miscompilation (not a BEAGLE logic bug — confirmed via direct atomic-counter probes in a prior session, see CLUSTER_AGENT_FINDINGS.md §3). The already-proven noinline mitigation (used earlier for the ADJOINT_ATOMIC_ADD_GPU CAS-retry loop) had been tried on the SINGLETON/PAIRLEADER macro bodies without effect, which ruled those out but left the true hot spot unidentified. ADJOINTN_ACCUM_ROW_RAW and ADJOINTN_ROTATE_ROW were never tested: they're the shared prefix code executed on every single code path (isAllReal, singleton, and pairleader alike), and contain the kernel's densest concentration of tid-predicated, barrier-synchronized code (a 32x-repeated, 128-wide tree reduction in ADJOINTN_ACCUM_ROW_RAW alone). Extracting both into their own __attribute__((noinline)) helper functions (adjointAccumRowRaw, adjointRotateRow), giving them their own register-allocation scope, fixes the corruption entirely. Verified: adjointtest4 --nstates 17's gradient comparison goes from NaN/FAIL to PASS (max|diff|=5.58e-01 vs. threshold 22.1); the 16-state and 4-state cases are unaffected (16-state's pre-existing ~1e-3 precision-level diff confirmed unchanged by temporarily reverting just this fix and comparing); and, independently, the full rabies_smoke.xml through beast-mcmc now shows Pr(accept)=0.8 for VanillaHMC(host.logRates), up from 0.0 in the prior (buggy) run — real-world confirmation this was a genuine gradient-correctness bug feeding the HMC leapfrog integrator, not just a synthetic test artifact. Full investigation trail (including a negative-control experiment on regOp's dynamic indexing that turned out to be unable to observe this bug either way, since the corrupted tid breaks the reduction's write-out predicate before regOp's own value could ever surface) is in beagle-bugs/gpu-adjoint-workgroup-corruption/PLAN.md and CLUSTER_AGENT_FINDINGS2.md.
msuchard
reviewed
Jul 10, 2026
| // eigenValues[PADDED_STATE_COUNT + state]) regardless of BEAGLE_FLAG_EIGEN_COMPLEX, so the | ||
| // device buffer must be sized/copied as complex whenever spectral representation is in use, | ||
| // otherwise that read runs past the data actually written by setEigenDecomposition. | ||
| if ((kFlags & BEAGLE_FLAG_EIGEN_COMPLEX) || (kFlags & BEAGLE_FLAG_SPECTRAL_REPRESENTATION)) |
Member
There was a problem hiding this comment.
@fil-monti -- this is not right. if isAllReal then the imaginary parts should never be written or read.
msuchard
reviewed
Jul 10, 2026
| #endif | ||
|
|
||
| char kernelNameBuf[256] = {0}; | ||
| clGetKernelInfo(deviceFunction, CL_KERNEL_FUNCTION_NAME, sizeof(kernelNameBuf), kernelNameBuf, NULL); |
Member
There was a problem hiding this comment.
@fil-monti -- this is not right. clGetKernelInfo() only needs to be called if BEAGLE_DEBUG is set.
msuchard
reviewed
Jul 10, 2026
| #endif | ||
|
|
||
| char kernelNameBuf[256] = {0}; | ||
| clGetKernelInfo(deviceFunction, CL_KERNEL_FUNCTION_NAME, sizeof(kernelNameBuf), kernelNameBuf, NULL); |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
...