Support RDNA3.5 (gfx1151 / Strix Halo) builds: wave32 kernels + vendored-glm/HIP fixes

[warrenrross]

Addresses #8 (gfx1151 / RDNA3.5 support); relates to #4 (7900XTX — same wave32 root cause) and #3 (glm headers).

Ready-to-use artifacts mirroring this patch (for RDNA3.5 users + reproducers):

• Build recipe / Dockerfiles: https://github.com/warrenrross/gsplat-rocm-gfx1151
• Prebuilt image: docker pull warrenrross/gsplat-rocm:gfx1151-rocm72-trainer

Provenance. This PR was prepared by an AI coding agent working under human
direction and reviewed before submission. Every claim below was validated on real
RDNA3.5 / gfx1151 hardware (specs in Environment) — the agent ran the
build, the rasterize+backward smoke test, the numerical gradcheck, and an end-to-end
training job itself, and the reported numbers are from those runs. Happy to reshape
anything to fit the project's conventions.

Summary

ROCm/gsplat currently targets only CDNA Instinct GPUs (MI300X/MI325X), which are
wave64. On RDNA3.5 APUs (gfx1151, "Strix Halo" — AMD Ryzen AI Max), the build
fails for two unrelated reasons:

1. RDNA is wave32, but the kernels hardcode wavefront size 64
   (rocprim::warp_reduce<…,64>, rocprim_warpSum<64>, cg::tiled_partition<64>,
   (block_size+63)/64, LOGICAL_WARP_SIZE=64, manual __shfl offset=32).
   -mwavefrontsize64 does not help (rocPRIM hardcodes wave32 for RDNA regardless;
   forcing ROCPRIM_NAVI=0 just trips a DPP-broadcast assert).

   The dangerous part is not a build failure — it is a silent correctness failure.
   In RasterizeToPixels3DGSBwd.cu the bs64 kernel's DPP reduction is commented out
   (ROCm 6.4.1 compiler issue) and replaced by rocprim_warpSum<64>(...). On a wave32
   device that hardcoded-64 rocPRIM instantiation does not error: it hits
   rocPRIM's enable_if<VirtualWaveSize <= max_size()> guard
   (check_virtual_wave_size<64, size32>) and becomes a silent no-op, so the warp
   sum simply does not happen → wrong color/SH gradients with no error, no warning, no
   crash. The build succeeds, import gsplat succeeds, training runs and quietly
   converges to garbage. This is the headline reason the change is a correctness fix.

2. The HIP build's glm handling breaks during/after hipify: the vendored-glm
   include dir isn't passed to the HIP compile, hipify copies glm's .hpp headers but
   not their sibling .inl files, and hipify rewrites glm's CUDA_VERSION gate to an
   undefined TORCH_HIP_VERSION so glm #errors out.

This PR adds an RDNA3.5 build path: gate the wavefront-size literals on the target arch
(or, minimally, provide wave32 variants) and fix the HIP glm include/define plumbing.
Validated end-to-end on gfx1151 — builds cleanly (gsplat/csrc.so; import gsplat
works), passes a GPU rasterize+backward smoke test and a numerical gradcheck against the
fork's own torch reference (max abs error 2.0e-6), and completes a 7000-step
simple_trainer.py training run — inside
rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1 with
PYTORCH_ROCM_ARCH=gfx1151.

Environment

• GPU: AMD Radeon 8060S integrated GPU on a Ryzen AI Max ("Strix Halo") — gfx1151,
  RDNA3.5, wave32, unified memory (GTT).
• ROCm 7.2.1, PyTorch 2.9.1+rocm7.2.1, HIP 7.2 — image
  rocm/pytorch:rocm7.2.1_ubuntu24.04_py3.12_pytorch_release_2.9.1. (This patch keys the
  warp size off the __gfx9__ arch predicate rather than the deprecated
  __AMDGCN_WAVEFRONT_SIZE, so it builds on current ROCm — the older "build on 7.0"
  constraint no longer applies. Stock ROCm ≥7.2.1 also carries the gfx1151
  first-dispatch HSA fix, ROCm #5853.)
• GPU_ARCHS=gfx1151 PYTORCH_ROCM_ARCH=gfx1151.

Changes

1. setup.py — pass the vendored glm include to the HIP compile.
Add the vendored glm dir as the first HIP include_dirs entry (~L229):

osp.join(current_dir, "gsplat", "cuda", "csrc", "third_party", "glm")

Rationale: the fork needs its vendored glm (system libglm-dev 1.0.1 is
API-incompatible). Without the explicit include, the hipified tree can't find glm
headers.

2. setup.py — define -DTORCH_HIP_VERSION=8000 in hipcc_flags (~L211).
hipify rewrites glm's CUDA_VERSION check to TORCH_HIP_VERSION, which is undefined
when glm/simd/platform.h is parsed → evaluates 0 < 7000 →
#error "GLM requires CUDA 7.0 or higher". Defining it high satisfies the gate.

3. Build-doc / hipify note — complete the glm copy into the hipified tree.
hipify copies glm's .hpp but not its sibling .inl files, so the HIP compile fails
with fatal error: scalar_constants.inl: No such file or directory. After the first
(hipify) pass, top up the hipified glm:

cp -rf gsplat/cuda/csrc/third_party/glm/. gsplat/hip/csrc/third_party/glm/

Ideally this is folded into the hipify step in setup.py (copy the whole vendored glm
dir, not just hipify's header selection) so a one-pass build works. The two-pass
workaround (hipify, top-up glm, recompile) is documented for now.

4. Arch-keyed WARP_SIZE constant (the wave32 fix). Rather than a blind
64→32 flip, introduce a single compile-time WARP_SIZE constant in
gsplat/cuda/include/Common.cuh, selected by the __gfx9__ architecture predicate,
and thread it through every hardcoded site. The same source then compiles to a
correct wave64 object on CDNA and a correct wave32 object on RDNA — no regression to the
Instinct line.

#ifdef USE_ROCM
  #if defined(__gfx900__) || defined(__gfx906__) || defined(__gfx908__) ||      \
      defined(__gfx90a__) || defined(__gfx940__) || defined(__gfx941__) ||      \
      defined(__gfx942__) || defined(__gfx950__)
    constexpr int WARP_SIZE = 64;   // CDNA / GCN (gfx9) is wave64
  #else
    constexpr int WARP_SIZE = 32;   // RDNA (gfx10/11/12), incl. gfx1151, is wave32
  #endif
#else
  constexpr int WARP_SIZE = 32;     // NVIDIA
#endif

Why __gfx9__, not __AMDGCN_WAVEFRONT_SIZE? The latter is deprecated as of ROCm
7.0.2 and slated for removal, and warpSize is not a host-side constant (this value is
also read in host-side launch-sizing code). AMD's rocPRIM maintainers recommend the
__gfx9__-family predicate as the durable form
(ROCm/ROCm#4121). Only gfx9 (CDNA/GCN) is
wave64; every gfx10/11/12 (RDNA) target is wave32.

WARP_SIZE is threaded through (audit IDs F2–F7):
threadIdx.x % 64→% WARP_SIZE, for (i = 0; i < 64; ...)→i < WARP_SIZE,
offset = 64/2→offset = WARP_SIZE/2, LOGICAL_WARP_SIZE = 64→= WARP_SIZE,
(block_size + 63)/64→(block_size + WARP_SIZE - 1)/WARP_SIZE, across
Projection2DGSFused.cu, Projection2DGSPacked.cu, ProjectionEWA3DGSFused.cu,
ProjectionEWA3DGSPacked.cu, RasterizeToPixels2DGSBwd.cu,
RasterizeToPixels3DGSBwd.cu, RasterizeToPixelsFromWorld3DGSBwd.cu, Utils.cuh.

5. if constexpr guard for the wave64-only bs64 kernel (audit F1/F6). The 3DGS
backward dispatch is changed so the wave64-only bs64 kernel — the one whose
rocprim_warpSum<64> silently no-ops on wave32 — is never instantiated on a wave32
build:

auto KERNEL = rasterize_to_pixels_3dgs_bwd_kernel<CDIM, float>;
if constexpr (WARP_SIZE == 64) {
    if (block_size == 64)
        KERNEL = rasterize_bs64_to_pixels_3dgs_bwd_kernel<CDIM, float>;
}

This is what defuses the silent no-op: on wave32 the compiler never emits the bad
kernel, so there is no path to the corrupt gradient. The surviving 64 literals (lines
30, 47, 85, 86, 210 of RasterizeToPixels3DGSBwd.cu) are intentional — they are
inside the now-wave64-only bs64 kernel and must stay 64.

Edit the .cu/.cuh sources, not the generated .hip, so a fresh hipify
regenerates correct kernels.

Maintainer-friendliness note

The change is deliberately shaped to be a zero-regression add for the
officially-supported Instinct line: on a CDNA target the __gfx9__ branch resolves to
WARP_SIZE == 64 and the if constexpr keeps the existing bs64 kernel, so the
generated wave64 device code is byte-identical to today's. The wave32 path is purely
additive. If maintainers prefer a different surface, the same logic maps cleanly onto an
explicit #if defined(__GFX11__) || defined(__GFX12__) block or a
-DGSPLAT_WARP_SIZE= build flag — happy to reshape to fit the project's arch-dispatch
style.

One caveat to resolve before/at merge (host/device on CDNA): the __gfx9__ macros
are undefined in the host compile pass, so host-side WARP_SIZE resolves to 32. This
is benign for gfx1151-only builds (host 32 == device 32) but would be a host/device
mismatch on a CDNA build. For the upstream form, derive the host value from
PYTORCH_ROCM_ARCH / a -DGSPLAT_WARP_SIZE= flag so host and device agree on Instinct.

Test plan

• Build: PYTORCH_ROCM_ARCH=gfx1151 pip install --no-build-isolation -e . →
  gsplat/csrc.so produced, import gsplat succeeds. ✅
• Numerical correctness (the merge gate): gradcheck the HIP fwd+bwd against the fork's
  own torch reference (gsplat/cuda/_torch_impl.py, _torch_impl_2dgs.py) on a tiny
  fixed scene. Proves the wave32 reductions are correct and the silent-no-op path is
  gone. ✅ max abs error 2.0e-6 on the wave32 color/SH gradient (the path that
  silently no-op'd before the fix).
• Smoke: gsplat.rasterization(...) on a 100-Gaussian synthetic scene, 64×64 render +
  backward. ✅ RASTERIZE+BACKWARD OK.
• E2E: examples/simple_trainer.py default --max-steps 7000 on a real COLMAP project
  (4K photogrammetry). ✅ ran to completion — 2,658,384 Gaussians; PSNR 20.94 /
  SSIM 0.60 / LPIPS 0.355 over held-out views; no crash, stable memory.

Environment note (not part of this PR)

Early gfx1151 bring-up hit a first-GPU-op segfault in libhsa-runtime64.so
(ROCm #5853, a CWSR/VGPR-count mismatch) — purely a ROCm-runtime issue, unrelated to
these kernel changes. It is fixed in stock ROCm ≥7.2.1, which is why the validated
environment above uses the 7.2.1 base; no special env vars or workarounds are needed on
a current ROCm.

---

[bjoernellens1]

Independent confirmation from a second gfx1151 (Radeon 8060S / Strix Halo, RDNA3.5) device: applied the kernel-only portion of this fix (the arch-gated WARP_SIZE constant in Common.cuh plus the if constexpr (WARP_SIZE == 64) guard around the bs64 kernel dispatch in RasterizeToPixels3DGSBwd.cu, threaded through the same 9 files) on top of release/1.5.3b2, without the setup.py glm changes (we'd already fixed that separately).

Results on gfx1151, ROCm 7.2.2, PyTorch 2.10.0+rocm7.2.2:

• pip install --no-build-isolation . completes cleanly, extension imports.
• Forward+backward smoke test on random Gaussians at both tile_size=16 (block_size=256, default path) and tile_size=8 (block_size=64, exercises the bs64-guard path) produces finite, nonzero gradients with no crash.
• Central finite-difference gradcheck on colors/opacities agrees with the analytic backward gradients to within ~2.3% max relative error (float32, eps=1e-3) at both tile sizes — consistent with expected FD discretization noise, not a systematic error.

Before landing on this fix we also tried a naive blind 64→32 string substitution across the same hardcoded sites (no if constexpr guard). It builds and imports cleanly, which made it look correct at first — but it reproduces exactly the failure mode described in this PR's write-up: the bs64 kernel path's dpp_sclr_warpSum/dpp_sprd_warpSum dereference a NULL storage pointer and cause a real GPU memory-access fault during backward on gfx1151. That's strong independent confirmation that the if constexpr guard here isn't just a style choice — it's the difference between a working build and a silent-until-it-isn't crash.

Would be great to see this merged — it's the only fix we found (including a couple of other in-flight branches) that actually holds up under a gradcheck rather than just compiling.

[bjoernellens1]

Ran an independent validation pass of the kernel-side WARP_SIZE
arch-gating from this PR on real gfx1151 (Strix Halo) hardware — ported the
kernel-only changes (the WARP_SIZE threading through
Projection2DGSFused.cu/Projection2DGSPacked.cu,
ProjectionEWA3DGSFused.cu/ProjectionEWA3DGSPacked.cu,
RasterizeToPixels2DGSBwd.cu, RasterizeToPixels3DGSBwd.cu,
RasterizeToPixelsFromWorld3DGSBwd.cu, Common.cuh, Utils.cuh) onto a
downstream branch and ran the repo's test suite plus a separate
finite-difference gradcheck against it. Results, for what they're worth:

• The projection kernels (Projection2DGSFused/Packed,
  ProjectionEWA3DGSFused/Packed) pass the existing test suite against
  their torch reference implementations on gfx1151 — forward output
  verified correct with the wave32 gating in place.
• The backward rasterization kernels
  (RasterizeToPixels2DGSBwd/3DGSBwd/FromWorld3DGSBwd) were initially
  not exercised by the test suite — their torch reference path goes
  through nerfacc.render_weight_from_alpha, and stock nerfacc 0.5.2 has
  no ROCm/HIP detection at all, so its compiled extension stayed
  uninitialized and every test hitting that path failed before reaching a
  useful comparison (87 failures, all traced to this, zero attributable to
  gsplat kernels — see gsplat#5). That's a nerfacc-side gap, unrelated to
  this PR's kernel changes.
• Update: swapped in a working ROCm nerfacc build in place of stock
  nerfacc and reran the full suite. That alone recovered 86 of the 87
  failures: 200 passed / 1 skipped / 1 failed, up from 114/1/87. The
  backward rasterization kernels (including the wave32 WARP_SIZE gating
  from this PR) are now genuinely exercised end-to-end against their torch
  reference, not just covered by the narrower gradcheck below. The one
  remaining failure (test_projection_2dgs, a small numerical mismatch on
  6/1482 elements) is unrelated to nerfacc or to this PR's changes — not
  yet investigated.
• Also ran an independent finite-difference gradcheck (500 Gaussians, tile
  size 16/8, colors+opacities gradients) against RasterizeToPixels3DGSBwd
  directly — passed, max abs error ~2e-6.

Net: both the projection-kernel and backward-rasterization sides of the
wave32 gating now have real test-suite coverage on gfx1151, not just a
narrow gradcheck. Separately, while testing this I found and fixed an
unrelated pre-existing bug in RasterizeToPixels3DGSBwd.cu — a
single shared-storage rocprim::warp_reduce instance was being reused
across all 8 logical wave32 warps in a 256-thread block, corrupting the
backward loop bound (warp_bin_final); the two sibling backward kernel
files already used a race-free shuffle-based reduction, so this brings
the third file in line with that pattern. Filed as its own fix, mentioning
here since it's directly adjacent to this PR's WARP_SIZE work and touches
the same file.

Separately — noticed this PR also touches GLM/HIP plumbing (glm_path in
include_dirs + -DTORCH_HIP_VERSION=8000 to satisfy the hipified
CUDA_VERSION gate). I ran into the same underlying issue (hipify walking
and rewriting the vendored glm headers when they're reachable from
include_dirs) and ended up fixing it by keeping glm out of
include_dirs entirely and reaching it via explicit -I flags instead,
which avoids hipify touching glm's headers at all rather than working
around the rewrite after the fact. Wrote it up as a separate draft PR in
case it's useful here — happy to fold it into this PR directly if that's
easier to review than a second PR.