[jeffdaily]

This PR adds AMD GPU support to Caspar, so its generated kernels and runtime build and run on AMD GPUs through ROCm/HIP. The default NVIDIA CUDA build is unchanged. It is enabled with caslib.compile(out_dir, use_hip=True, hip_arch=...), or with -DUSE_HIP=ON in the generated build.

The CUDA runtime spellings Caspar emits and uses (cudaMalloc, cudaFuncSetAttribute, error handling and the rest of the runtime API, plus atomicAdd_block and the cub namespace) are mapped to their HIP equivalents in one compatibility header, source/runtime/cuda_to_hip.h. In a CUDA build the header includes the CUDA runtime, and also cooperative_groups and CUB when compiling CUDA, so host-only .cc files can include it. The header provides no device-side fallbacks of its own. The code generation templates emit this header, so the symbolic kernel definitions are unchanged.

The code generation gains a HIP path. CasparLibrary.compile() in code_generation/library.py takes use_hip and hip_arch, and the generated CMakeLists.txt has a USE_HIP branch that compiles the sources as HIP against hip and hipcub. The runtime sources get the matching HIP changes: shuffle-based tile reductions where HIP lacks cg::reduce, shared-memory atomics where it lacks cg::labeled_partition and coalesced-group shuffles (memops.cuh), a block-strided copy where it lacks cg::memcpy_async (shared_indices.cu), and a host-side pointer-attribute lookup in the pybind layer. A small Windows build fix (C++17, git path separator) is included for the all-clang ROCm toolchain.

The PR also fixes a solver bug that ROCm exposed: DoUpdateRFirst() did not zero r_kp1_norm2_tot before accumulating into it, so the first PCG step of each solver iteration added onto the previous iteration's value (or uninitialized memory on the first) before the PCG exit test read it. The bug is not specific to ROCm.

Tests

• symforce_caspar_codegen_test regenerates its checked-in kernels and solvers through the HIP-aware templates. It also checks in the generated CMakeLists.txt, with and without Python bindings, so changes to the HIP (or CUDA) build branch show up in review like other codegen changes.
• test/CMakeLists.txt builds the *_cuda_test.cc Caspar GPU tests with HIP when CMake finds a HIP compiler and no CUDA compiler. With -DCMAKE_HIP_ARCHITECTURES set the build needs no AMD GPU, so a ROCm-only machine or container compiles every checked-in Caspar kernel and solver. Running them needs an AMD GPU, as the CUDA tests need an NVIDIA one.

How to build for AMD GPUs

Pass use_hip=True and the target architecture when compiling a generated library:

caslib.generate(out_dir)
caslib.compile(out_dir, use_hip=True, hip_arch="gfx90a")

Set hip_arch to the target AMD GPU (for example gfx90a for CDNA2 or gfx1100 for RDNA3). Without it, CMake builds for the GPU it detects on the build machine. The build uses CMake's HIP language support, so it needs the ROCm Clang compiler (amdclang++; CMake does not accept the hipcc wrapper) and the hip and hipcub packages.

To build the Caspar GPU tests for ROCm:

cmake -B build -DCMAKE_HIP_ARCHITECTURES=gfx90a
cmake --build build --target bin.tests.caspar_padding_lane_cuda_test bin.tests.caspar_solver_cuda_test bin.tests.caspar_tridiagonal_cuda_test bin.tests.caspar_backend_coverage_cuda_test
ctest --test-dir build -R caspar_

Validation

The Caspar GPU tests (padding lane, solver, tridiagonal, backend coverage) and symforce_caspar_codegen_test pass on an AMD Instinct MI250X and an AMD Radeon PRO W7800 on Linux, and the four Caspar GPU tests also pass on an AMD Radeon 8060S on Windows. The Caspar test targets also compile and link with CUDA 12.8.

[aaron-skydio]

Thanks for the contribution, this looks interesting. I'm a little torn on whether we can commit to the maintenance burden of an AMD implementation here especially given the lack of tests. I think I'm going to leave this open and see how much interest there is from users in AMD support?

[jeffdaily]

At minimum I could add a build-only github action workflow. Would that be sufficient?

[bjoernellens1]

I've been building on top of this PR (running it on gfx1151 / ROCm 7.2) and hit a real data race in Caspar's solver runtime that's worth flagging here, since it's in the HIP-specific code path this PR introduces.

Bug: SumStore() in symforce/caspar/source/runtime/memops.cuh performs a two-stage warp-tile reduction into a shared-memory scratch buffer (inout_shared/shared_tmp) and returns without a trailing __syncthreads() after its last write. That buffer gets reused within the same kernel invocation — e.g. kernel_Point_alpha_numerator_denominator.cu calls SumStore twice back-to-back against the same buffer, and every kernel using the AddSum/WriteSum accessors calls SumStore immediately followed by SumFlushFinal. Only the 32 lanes in the meta_group_rank() == 0 tile do the final read of that buffer; the other threads finish SumStore right after the first-stage write, with nothing forcing them to wait for that tile's read before a subsequent write to the same buffer is issued.

On gfx1151 this reliably manifested as a real race: intermittent NaN/-nan in score_current/step_quality, runaway LM damping, and — the most convincing signal — run-to-run nondeterminism on the exact same compiled binary (no regeneration/recompile between runs). Five reruns of the identical binary produced five different reprojection RMS values (8.33 / 7.48 / 7.43 / 7.24 / 4.20 px), which rules out a build/environment issue and points squarely at a runtime race.

Suggested fix — a single trailing __syncthreads():

       if (group.thread_rank() == 0) {
         shared_tmp[offset] = tot;
       }
     }
+    __syncthreads();
   }

Verification: with the fix, the same determinism test (5 within-build reruns + 3 fresh-build reruns) produced the bit-identical, correct result (0.004022114544300473 px, all passing) — 8/8 vs. the prior 9/9 divergent failures. I later broadened this to a 22-run sweep across fresh RNG seeds and a 3x larger synthetic problem (12 cams / 64 points): 22/22 passed with no failures.

Honest caveat on mechanism: I initially attributed this to a write-after-read hazard across the SumStore→SumStore/SumFlushFinal call boundary, reasoned as fixed by the trailing barrier. On closer review that specific explanation doesn't fully hold up — both call sites that actually exhibit the divergence already have their own leading __syncthreads(), which should independently enforce safety there, making the trailing barrier look redundant at exactly the site where the divergence was observed. A later revert-and-retest (removing the barrier and rerunning the original determinism protocol) also failed to reliably reproduce the original 9/9 failure — an inconclusive result rather than a confirming one. So I can't fully defend the causal mechanism if pressed on it. What I can say with confidence: the fix is empirically robust (22/22 across broadened seed/size coverage, 0 failures), it's a real, defensible correctness improvement to SumStore's shared-buffer-reuse contract regardless of mechanism, and it costs nothing (one barrier, not gated to HIP since the missing-barrier hazard isn't platform-specific in principle, just harder to trigger on other scheduling regimes). I'd treat this as a legitimate hardening fix rather than a fully root-caused bug report.

I also independently checked FlushSumShared's HIP atomicAdd_block fallback (the largest HIP-specific rewrite in this file) against an independent NumPy host reference on a passing run — device output matched the host-computed reference to float32 precision (max relative error ~5e-6 for Cam nodes, ~6e-5 for Point nodes), so no evidence it's involved, though that only characterizes a clean run and doesn't rule out a rare mistimed interleaving.

Happy to open this as a proper PR against your branch if that's easier to review/merge than a diff in a comment — just let me know which you'd prefer.

[jeffdaily]

Note: this reply was drafted by an AI assistant.

Thanks for the thorough report -- five reruns with per-run RMS values and a 22/22 bit-identical confirmation with the fix in place is exactly the kind of evidence that makes this easy to act on. The diff in the comment is enough; no need to open a separate PR unless you'd prefer the authorship to be visible that way.

A barrier after the last shared-memory write is the standard idiom when the buffer can be reused by a later phase, so the change is sound on its face even with the causal chain not fully pinned down. I'll verify the mechanism against the surrounding code rather than take it on faith, fold the fix into the next update to this branch together with the rebase onto current main (the branch has drifted into conflict), credit you with a Co-authored-by trailer, and rerun the determinism check plus the test suite on our AMD hardware before pushing. Results will go here when that lands.

[aaron-skydio]

Yeah it seems like AMD support here would be useful to a bunch of people, so yeah I'm down to merge this

I'll want to have tests that at minimum cover the codegen part, and probably also the build, it's going to be harder to get something actually testing on an AMD gpu. As much of that as you're able to do would be helpful

[jeffdaily]

Note: this reply was drafted by an AI assistant.

Merged current main into the branch. The only conflict was an include in sort_indices.h.

On SumStore: after a768308 the trailing barrier from bjoernellens1's report isn't needed, so I took it back out. The NaN and run-to-run variation they saw fit #469. On an MI250X the padding-lane kernel without the register zeroing gives NaN at every padded size, with or without the barrier. With the zeroing, per-block sums are bit-identical over 200 runs whether or not the barrier is there, including the back-to-back SumStore kernel they pointed at. On a W7800 (wave32) that kernel matches an exact reference in every run, with or without the barrier. Thanks @bjoernellens1 for the report and the reproduction.

Tests: the *_cuda_test.cc tests now also build with HIP when CMake finds no CUDA compiler. With -DCMAKE_HIP_ARCHITECTURES set they build without an AMD GPU, so a ROCm container can compile every checked-in Caspar kernel and solver. symforce_caspar_codegen_test now also checks in the generated CMakeLists.txt, which is where the HIP build differs.

At the head I pushed before, a HIP build of a generated solver failed. The solver calls cudaGetDeviceCount (since the GPU device selection change, 2ee2978), and cuda_to_hip.h never mapped it. The new HIP build of the tests is what caught it, and it is mapped now. Running the tests on ROCm turned up one more bug: DoUpdateRFirst() never zeroed r_kp1_norm2_tot, so the first PCG step of each solver iteration added onto the previous iteration's value (or uninitialized memory on the first) before the PCG exit test read it. That's fixed in the template, and it affects CUDA too. A CUDA build of the tests also showed that host .cc files including the generated headers failed under g++, so cuda_to_hip.h now includes cooperative_groups and CUB only when compiling CUDA. I also removed some reduction helpers from cuda_to_hip.h that nothing used, and fixed the README's ROCm example, which called a function that doesn't exist.

Tested on an AMD Instinct MI250X and an AMD Radeon PRO W7800 on Linux (the Caspar GPU tests, the codegen tests and the rest of the test suite, apart from three tests that need tools those machines lack: rustfmt and a docs dependency) and an AMD Radeon 8060S on Windows (the Caspar GPU tests).

[aaron-skydio]

Thanks! I'll try and get this merged when I have time