Add ROCm/HIP support for AMD GPUs (Linux + Windows) - #970
Conversation
Completes ROCm/HIP support for gsplat's differentiable Gaussian-splatting rasterization (PyTorch's hipify translates the .cu at build time). The project carried only build-config ROCm awareness; this fills the gaps so the rasterizer builds and is GPU-validated on AMD Linux (gfx90a MI250X wave64, gfx1100 W7800 wave32) and Windows (gfx1151 Radeon 8060S wave32), with every change behind USE_ROCM / torch.version.hip and the CUDA path unchanged. The substantive work was HIP Cooperative Groups: ROCm's CG (7.2.1) has no cg::reduce and no cg::labeled_partition. cg::reduce over a 32-thread tile becomes a butterfly shfl_xor all-reduce; cg::labeled_partition is rebuilt from warp.match_any(label) -- a masked all-reduce over same-label lanes plus one atomicAdd per label by the lowest-set-bit lane. Matching CUDA's one-atomic-per-label granularity (not a per-lane fan-out) is load-bearing: the larger atomic count changes float accumulation order and pushed the v_viewmats gradient past gsplat's own test tolerance. The 32-lane tiles stay within their 32 lanes on a 64-lane wavefront (shfl_xor/match_any use width 32), so no wave64 lane-count rework is needed. Other gaps: ROCm -ffast-math is more aggressive than CUDA's and perturbed ill-conditioned projection-covariance gradients past tolerance, so FAST_MATH defaults off on ROCm (opt-in FAST_MATH=1); torch's hipify mangles the bundled GLM, worked around by monkeypatching hipify to skip the GLM dir (GLM 1.0.2 detects __HIP__ natively); plus std::min/max host shims, cuda::std->std, namespace cub = hipcub (radix sort uses begin_bit=0), the hipFuncSetAttribute const void* cast, and __CUDA_ARCH__ || __HIP_DEVICE_COMPILE__ guards. 3DGUT on ROCm: the 3DGUT path needs cuda::std::optional, which the ROCm stack does not ship. The header-only ROCm/libhipcxx port supplies the cuda::std namespace, so 3DGUT builds and runs on ROCm: clone github.com/ROCm/libhipcxx and point LIBHIPCXX_INCLUDE at its include/ (build.py then defaults BUILD_3DGUT on; unset, it stays off so a plain pip install does not hard-error). The two RasterizeToPixelsFromWorld3DGS kernels also needed the hipFuncSetAttribute const-void* cast. Windows (gfx1151) enablement, all guarded to win32 + torch.version.hip so Linux and CUDA are byte-identical: build.py dropped the nvcc-only -allow-unsupported-compiler and the MSVC -Xcompiler /Zc:preprocessor and /openmp from the device flags (amdclang is a gcc-style driver: it rejects the first and treats the MSVC flags as input files; it uses -fopenmp instead). The host .cpp op-wrappers include c10/cuda/CUDAGuard.h, which on ROCm pulls in the HIP runtime headers whose GCC __attribute__ syntax MSVC cl.exe cannot parse; build.py routes each wrapper through hipcc/amdclang by handing torch a byte-identical .cu shim next to the source (same dir, so every relative #include resolves identically; the wrappers carry no device code, so the device pass compiles nothing). distributed.py now imports torch.distributed.nn lazily: the Windows ROCm torch wheel ships without a c10d backend, and the eager import failed at module load; collective helpers are multi-GPU-only and unused single-process. Also adds "AMD GPUs (ROCm)" install notes to the README and the Windows install guide (docs/INSTALL_WIN.md). Test Plan: ``` HIP_VISIBLE_DEVICES=<n> PYTORCH_ROCM_ARCH=<arch> pip install -e . --no-build-isolation HIP_VISIBLE_DEVICES=<n> python -m pytest tests/test_basic.py # 3DGUT: clone ROCm/libhipcxx, then LIBHIPCXX_INCLUDE=<clone>/include BUILD_3DGUT=1 BUILD_3DGS=1 BUILD_2DGS=1 BUILD_ADAM=1 BUILD_RELOC=1 BUILD_LOSSES=1 \ HIP_VISIBLE_DEVICES=<n> PYTORCH_ROCM_ARCH=<arch> pip install -e . --no-build-isolation # Windows gfx1151: TheRock torch[device-gfx1151] wheel, Python on a space-free path, # DISTUTILS_USE_SDK=1, HIP_DEVICE_LIB_PATH=<rocm>/lib/llvm/amdgcn/bitcode, MSVC link.exe ahead of MSYS on PATH. ``` Validated on gfx90a (MI250X, ROCm 7.2.1, PyTorch 2.13 ROCm), gfx1100 (W7800, ROCm 7.2.1), and gfx1151 (Radeon 8060S, TheRock ROCm 7.x, PyTorch 2.12 ROCm). tests/test_basic.py core 3DGS/2DGS passes 108 (HIP vs a pure-torch reference, forward + autograd.grad); an independent compositing oracle agrees with a backward finite-difference check; forward is bit-exact deterministic. With libhipcxx the 3DGUT/lidar/UT and distortion/ftheta suites pass. (test_rasterize_to_pixels_eval3d still needs nerfacc, which is CUDA-only -- a separate dependency gap, not a kernel issue. On gfx1151 one UT-projection gaussian's derived pixel radius lands 12 vs 10 px over a heuristic atol, 1/769252 -- an FP-rounding boundary, not a kernel defect.) Authored with Claude.
|
@jeffdaily and I are working together inside AMD to consolidate this PR with my #957 |
|
We used this PR's head ( The main failure/fix clusters are:
This is not a rebased alternative PR, and I am not claiming a full current-upstream suite pass for the stack. It is concrete gfx1100 failure/fix evidence intended to help the ongoing consolidation with #957. I split two independent current- |
|
I opened the gfx1100 follow-up as a stacked draft PR directly against the head branch of this ROCm port: jeffdaily#1 This keeps the review diff limited to the 24 follow-up commits instead of duplicating #970 against upstream |
Completes ROCm/HIP support for gsplat's differentiable Gaussian-splatting rasterization (PyTorch's hipify translates the .cu at build time). The project carried only build-config ROCm awareness; this fills the gaps so the rasterizer builds and is GPU-validated on AMD Linux (gfx90a MI250X wave64, gfx1100 W7800 wave32) and Windows (gfx1151 Radeon 8060S wave32), with every change behind USE_ROCM / torch.version.hip and the CUDA path unchanged.
The substantive work was HIP Cooperative Groups: ROCm's CG (7.2.1) has no cg::reduce and no cg::labeled_partition. cg::reduce over a 32-thread tile becomes a butterfly shfl_xor all-reduce; cg::labeled_partition is rebuilt from warp.match_any(label) -- a masked all-reduce over same-label lanes plus one atomicAdd per label by the lowest-set-bit lane. Matching CUDA's one-atomic-per-label granularity (not a per-lane fan-out) is load-bearing: the larger atomic count changes float accumulation order and pushed the v_viewmats gradient past gsplat's own test tolerance. The 32-lane tiles stay within their 32 lanes on a 64-lane wavefront (shfl_xor/match_any use width 32), so no wave64 lane-count rework is needed.
Other gaps: ROCm -ffast-math is more aggressive than CUDA's and perturbed ill-conditioned projection-covariance gradients past tolerance, so FAST_MATH defaults off on ROCm (opt-in FAST_MATH=1); torch's hipify mangles the bundled GLM, worked around by monkeypatching hipify to skip the GLM dir (GLM 1.0.2 detects HIP natively); plus std::min/max host shims, cuda::std->std, namespace cub = hipcub (radix sort uses begin_bit=0), the hipFuncSetAttribute const void* cast, and CUDA_ARCH || HIP_DEVICE_COMPILE guards.
3DGUT on ROCm: the 3DGUT path needs cuda::std::optional, which the ROCm stack does not ship. The header-only ROCm/libhipcxx port supplies the cuda::std namespace, so 3DGUT builds and runs on ROCm: clone github.com/ROCm/libhipcxx and point LIBHIPCXX_INCLUDE at its include/ (build.py then defaults BUILD_3DGUT on; unset, it stays off so a plain pip install does not hard-error). The two RasterizeToPixelsFromWorld3DGS kernels also needed the hipFuncSetAttribute const-void* cast.
Windows (gfx1151) enablement, all guarded to win32 + torch.version.hip so Linux and CUDA are byte-identical: build.py dropped the nvcc-only -allow-unsupported-compiler and the MSVC -Xcompiler /Zc:preprocessor and /openmp from the device flags (amdclang is a gcc-style driver: it rejects the first and treats the MSVC flags as input files; it uses -fopenmp instead). The host .cpp op-wrappers include c10/cuda/CUDAGuard.h, which on ROCm pulls in the HIP runtime headers whose GCC attribute syntax MSVC cl.exe cannot parse; build.py routes each wrapper through hipcc/amdclang by handing torch a byte-identical .cu shim next to the source (same dir, so every relative #include resolves identically; the wrappers carry no device code, so the device pass compiles nothing). distributed.py now imports torch.distributed.nn lazily: the Windows ROCm torch wheel ships without a c10d backend, and the eager import failed at module load; collective helpers are multi-GPU-only and unused single-process.
Also adds "AMD GPUs (ROCm)" install notes to the README and the Windows install guide (docs/INSTALL_WIN.md).
Test Plan:
Validated on gfx90a (MI250X, ROCm 7.2.1, PyTorch 2.13 ROCm), gfx1100 (W7800, ROCm 7.2.1), and gfx1151 (Radeon 8060S, TheRock ROCm 7.x, PyTorch 2.12 ROCm). tests/test_basic.py core 3DGS/2DGS passes 108 (HIP vs a pure-torch reference, forward + autograd.grad); an independent compositing oracle agrees with a backward finite-difference check; forward is bit-exact deterministic. With libhipcxx the 3DGUT/lidar/UT and distortion/ftheta suites pass. (test_rasterize_to_pixels_eval3d still needs nerfacc, which is CUDA-only -- a separate dependency gap, not a kernel issue. On gfx1151 one UT-projection gaussian's derived pixel radius lands 12 vs 10 px over a heuristic atol, 1/769252 -- an FP-rounding boundary, not a kernel defect.)
Authored with Claude.
This port was prepared with the assistance of Claude (AI), reviewed before submission.