Skip to content

Add AMD GPU (ROCm/HIP) support to the Caspar backend - #465

Open
jeffdaily wants to merge 2 commits into
symforce-org:mainfrom
AMD-Ecosystem:moat-port
Open

Add AMD GPU (ROCm/HIP) support to the Caspar backend#465
jeffdaily wants to merge 2 commits into
symforce-org:mainfrom
AMD-Ecosystem:moat-port

Conversation

@jeffdaily

Copy link
Copy Markdown

This PR adds AMD GPU support to Caspar so its generated kernels and runtime build and run on AMD GPUs through ROCm/HIP, while keeping the default NVIDIA CUDA build unchanged. It is enabled with compile_caspar_library(..., use_hip=True, hip_arch=...) (or -DUSE_HIP=ON in the generated build); when off, the build is exactly as before.

The CUDA spellings Caspar emits and uses -- cudaMalloc, __syncthreads, cooperative-group reductions, CUB primitives, and the runtime API -- are mapped to their HIP equivalents through a small compatibility header (source/runtime/cuda_to_hip.h). On an NVIDIA build the header is a transparent passthrough; on a ROCm build it aliases the cuda* symbols to hip* and supplies device-side fallbacks where HIP lacks a cooperative-groups primitive (cg::reduce, cg::labeled_partition). Because the mapping lives in one header and the codegen templates emit it, the symbolic kernel definitions are unchanged.

The code generation gains a HIP path: code_generation/library.py takes use_hip/hip_arch, and the Jinja build-file and kernel templates emit the HIP-correct includes and the USE_HIP CMake option. The runtime sources get the corresponding HIP-compat (shared-memory atomics, the reduction fallback, 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.

How to build for AMD GPUs

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

compile_caspar_library(caslib, output_dir, use_hip=True, hip_arch="gfx90a")

Set hip_arch to the target AMD GPU (for example gfx90a for CDNA2 or gfx1100 for RDNA3). The ROCm build needs a HIP-enabled compiler (hipcc/amdclang++) and the hip and hipcub packages.

Validation

The full code-generation pipeline (generate, HIP-compile, execute on GPU, verify numerical output) was exercised on real AMD hardware: Linux CDNA2 (gfx90a) and RDNA3 (gfx1100), and Windows RDNA4 (gfx1201). The generated kernels (cooperative-group reductions, shared-memory atomics, scatter/gather) produce results matching the CUDA path. The default NVIDIA CUDA build is unchanged.

Adds AMD GPU support to Caspar so its generated kernels and runtime also
build and run on AMD GPUs via HIP/ROCm, while leaving the default NVIDIA
CUDA build unchanged (enabled with use_hip=True / -DUSE_HIP=ON; off by
default).

Review in this order:

1. symforce/caspar/source/runtime/cuda_to_hip.h (new): a compatibility
   header that maps the CUDA spellings Caspar emits and uses (cudaMalloc,
   __syncthreads, cooperative-group reductions, CUB primitives, the
   runtime API) onto their HIP equivalents, and supplies device-side
   fallbacks where HIP lacks a cg:: primitive (reduce, labeled_partition).

2. code_generation/library.py and source/templates/*.jinja: the codegen
   gains a use_hip/hip_arch path; the generated build file and kernel
   templates emit the compat include and HIP-correct spellings, so the
   symbolic kernel definitions are unchanged.

3. source/runtime/*.cu, memops.cuh, pybind_array_tools.{cc,h}: the
   runtime HIP-compat (shared-memory atomics, the cooperative-group
   reduction fallback, and a host-side pointer-attribute lookup).

4. A small Windows build fix (C++17, git path separator) for the
   all-clang ROCm toolchain.

5. README: how to build a generated library for AMD GPUs.

Authored with assistance from Claude.

Test Plan:

The full code-generation pipeline (generate -> HIP-compile -> execute on
GPU -> verify numerical output) was exercised on real AMD hardware,
Linux CDNA2 (gfx90a) and RDNA3 (gfx1100) and Windows RDNA4 (gfx1201):

  compile_caspar_library(caslib, out_dir, use_hip=True, hip_arch="gfx90a")
  # generated kernel executes on GPU; output matches the CUDA path.

The default NVIDIA CUDA build (use_hip=False) is unchanged.
The generated CMake template pinned CMAKE_HIP_ARCHITECTURES to gfx90a when
unset, but the pin ran after project(... LANGUAGES ... HIP), which already
enables the HIP language and resolves the architecture. The block was dead
code that could only mislead, or become a live footgun if file order ever
changed.

project(... LANGUAGES ... HIP) already honors an explicit
-DCMAKE_HIP_ARCHITECTURES, otherwise auto-detects the host GPU, and errors
on a no-GPU build host. Removing the pin lets that single mechanism decide
the arch, so a user on a non-gfx90a card no longer risks a silently
mistargeted default. The hip_arch argument still threads through as
-DCMAKE_HIP_ARCHITECTURES.

This work was authored with the assistance of the Claude AI assistant.
@aaron-skydio

Copy link
Copy Markdown
Member

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

Copy link
Copy Markdown
Author

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

bjoernellens1 pushed a commit to bjoernellens1/symforce that referenced this pull request Aug 4, 2026
@bjoernellens1

Copy link
Copy Markdown

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 SumStoreSumStore/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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants