[ROCm] Add AMD GPU (HIP) support to the RPUCuda backend#770
Open
jeffdaily wants to merge 3 commits into
Open
Conversation
Collaborator
|
Hello @jeffdaily! Thanks a lot for this PR! We are looking into this and try to get access to an AMD GPU to test it out. Please can you add your signature to the commits so the DCO check can pass? Thanks again! |
This adds an optional USE_HIP build of the RPUCuda CUDA backend so the analog tile simulator runs on AMD GPUs through ROCm, alongside the existing USE_CUDA path. Every HIP change is gated behind USE_HIP, so we have made every effort to leave the CUDA build compiling the same sources it did before. The build stays on the project's standalone scikit-build/CMake CUDA toolchain (PyTorch is used only for headers and linkage), so no torch cpp_extension hipify is involved. Review order: start with src/rpucuda/cuda/cuda_to_hip.h, the single HIP-aware header. It is included once from cuda_util.h, ahead of the CUDA runtime / cuBLAS / cuRAND includes and guarded by USE_HIP, and aliases the cuda*/cublas*/ curand* spellings the backend uses to their hip*/hipBLAS/hipRAND equivalents so the CUDA-spelled device sources compile unchanged. cuBLAS maps to hipBLAS (GEMM/GEMV/GER/scal/copy/nrm2, handle, host and device pointer modes, op and status enums) and cuRAND to hipRAND (host generator plus device per-thread state). In CUDA, curandState, curandState_t and curandStateXORWOW are one XORWOW type; hipRAND declares hiprandState and hiprandStateXORWOW_t as distinct structs, so all three spellings are mapped to one hipRAND type to keep the backend's template instantiations consistent (mapping them apart drops a CudaArray instantiation and breaks the module load). The one semantically delicate piece is bit_line_maker.cu. Its stochastic pulse train is a warp-size-coupled serialized bit format: a 32-lane ballot is packed into 32-bit words and the pulsed-weight-updater kernels read those words back by bit position, so the on-device word layout must stay independent of the physical wavefront width. A 64-wide CDNA wavefront holds two independent 32-lane logical warps, and a native 64-lane ballot would merge them into one 64-bit value and corrupt the format. The compat header therefore redefines the ballot/shuffle sync intrinsics as width-32 logical-warp operations: the ballot takes the full wavefront ballot and selects only the calling lane's own 32-lane subgroup as a 32-bit word (shift 0 on wave32; lanes 32-63 select the high word on wave64), and the shuffles use width 32. The word width and the consumer stay unchanged, so a wave32 and a wave64 device emit the same 32-bit-per-32-lane stream the CUDA build produces. CMake gains an option(USE_HIP) and a dependencies_hip.cmake mirroring the CUDA path: it enables the HIP language, marks the .cu sources LANGUAGE HIP, links hipBLAS/hipRAND/hipCUB, takes the architecture list from CMAKE_HIP_ARCHITECTURES (defaulting to gfx90a only when unset, never hardcoded), keeps RPU_USE_CUDA defined so the device tile code and the CudaAnalogTile binding stay built, and forces interprocedural optimization off (HIP device linking does not finalize LTO, which would yield an empty module init). rpu_cub.h selects hipCUB and its hipcub namespace. The pybind translation units are compiled with the HIP toolchain because they inherit the GPU library's offload-arch usage requirement, and rpu_base_tiles_cuda.cpp uses the lightweight c10 HIP stream header instead of the full ATen CUDA context header, which would drag in hipSOLVER/hipSPARSE headers not part of this build. RPU_CXX_STANDARD is a new cache variable (default 17) so the standard can track newer PyTorch headers that require C++20; that also exposed and fixed an ill-formed template-id constructor name in rpu_linearstep_device.h. The pybind layer hand-rolls an at::cuda::getCurrentCUDAStream shim in rpu_base_tiles_cuda.cpp because this build defines its own USE_HIP and never runs PyTorch source-hipify on that translation unit. Which underlying c10 symbol the shim must call depends on the torch hipify generation, not the ROCm version: hipify v2 keeps c10::cuda::getCurrentCUDAStream as the public stream API (c10::hip::getCurrentHIPStream stays guarded by USE_ROCM, which this build does not define), while hipify v1 removes it so the hip-spelled symbol must be used. CMake detects the hipify version from the build's own torch and defines TORCH_HIPIFY_V2 when it is >= 2.0.0; the shim selects its branch from that define, and detection failure leaves it undefined, matching the safe v1 default. Building with amdclang++ on Windows needs two further fixes: guard the MSVC-only /O2 flag with if(MSVC) (Clang-on-Windows is not MSVC) while keeping the C++ standard settings for all Windows builds, and change the forward declaration of PulsedUpdateMetaParameter from class to struct to match its definition (struct and class mangle identically on ELF but differently on the MSVC ABI used by amdclang++, which otherwise leaves virtual-method symbols unresolved at link). The ROCm/HIP build is documented in the developer install guide alongside the CUDA build. This work was authored with the assistance of Claude, an AI assistant. Test Plan: Built and validated on AMD GPUs across CDNA and RDNA, exercising the analog tiles on device against the CPU tiles (tolerance-based comparisons). Linux, AMD Instinct MI250X (gfx90a, CDNA2, wave64), ROCm 7.2.1: ``` USE_HIP=ON USE_CUDA=0 python setup.py build_ext -j16 --inplace \ -DUSE_HIP=ON -DUSE_CUDA=OFF -DRPU_CXX_STANDARD=20 \ -DCMAKE_HIP_ARCHITECTURES=gfx90a \ -DCMAKE_HIP_COMPILER=/opt/rocm/llvm/bin/clang++ \ -DCMAKE_PREFIX_PATH="$TORCH_CMAKE;/opt/rocm" \ -DRPU_BLAS=OpenBLAS -DRPU_USE_TORCH_BUFFERS=OFF -DCMAKE_BUILD_TYPE=Release HIP_VISIBLE_DEVICES=0 PYTHONPATH=src python -m pytest -q \ tests/test_specific_tiles.py tests/test_simulator_tiles.py \ tests/test_bindings_tiles.py tests/test_torch_tiles.py \ tests/test_inference_tiles.py tests/test_layers_linear.py \ tests/test_layers_convolution.py ``` test_specific_tiles (the bit-line-maker plus pulsed-update path) passes 18/18; the tile/binding/torch/inference/layer suites pass (one stochastic weight-programming test is order-sensitive and passes in isolation). A fat binary for "gfx90a;gfx1100" builds and emits both code objects for bit_line_maker.cu, confirming the ballot rewrite is wave-width agnostic. The same suites also pass on a Linux gfx1100 (RDNA3, wave32) device. Windows, AMD Radeon RX 9070 XT (gfx1201, RDNA4, wave32), ROCm 7.14: ``` cmake -DUSE_HIP=ON -DUSE_CUDA=OFF -DRPU_CXX_STANDARD=20 \ -DCMAKE_HIP_ARCHITECTURES=gfx1201 -DRPU_BLAS=OpenBLAS -DBUILD_TEST=OFF ninja rpu_base ``` test_specific_tiles 18/18; the tile/binding/torch/inference/layer suites pass (over 900 tests), with one order-sensitive stochastic test that passes in isolation. Signed-off-by: Jeff Daily <jeff.daily@amd.com>
Author
|
Thanks @PabloCarmona! I've added my Signed-off-by to the commit, so the DCO check passes now. |
The top-level CMakeLists set CMAKE_HIP_ARCHITECTURES to gfx90a when it was unset, and that block runs before cmake/dependencies_hip.cmake reaches enable_language(HIP). Pinning ahead of the enable preempts CMake's own host-GPU detection, so a builder on a non-gfx90a AMD GPU who did not pass -DCMAKE_HIP_ARCHITECTURES would silently get gfx90a code objects that fail to load at runtime on their card. Drop the pin and let enable_language(HIP) resolve the architecture: it honors an explicit -DCMAKE_HIP_ARCHITECTURES, otherwise auto-detects the host GPU, and otherwise errors so the builder must choose. The RPU_GPU and test targets continue to read CMAKE_HIP_ARCHITECTURES for their per-target HIP_ARCHITECTURES property, now from the resolved value. Authored with assistance from Claude. Signed-off-by: Jeff Daily <jeff.daily@amd.com>
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.
Summary
This adds an optional
USE_HIPbuild of the RPUCuda backend so the analog tile simulator runs on AMD GPUs through ROCm, alongside the existingUSE_CUDApath, which is left byte-identical. The build stays on the project's standalone scikit-build/CMake CUDA toolchain (PyTorch is used only for headers and linkage), so no torchcpp_extensionhipify is involved.Build it with
-DUSE_HIP=ON(mutually exclusive withUSE_CUDA) and select the GPU with-DCMAKE_HIP_ARCHITECTURES=<arch>(e.g.gfx90a,gfx1100). The ROCm/HIP build is documented in the developer install guide alongside the CUDA build.What is implemented
src/rpucuda/cuda/cuda_to_hip.h(the single HIP-aware header, included once fromcuda_util.hunderUSE_HIP): aliases thecuda*/cublas*/curand*spellings tohip*/ hipBLAS / hipRAND so the CUDA-spelled device sources compile unchanged.curandState,curandState_tandcurandStateXORWOW(one XORWOW type in CUDA) are all mapped to one hipRAND type, since hipRAND declares them as distinct structs and splitting them would drop aCudaArrayinstantiation.bit_line_maker.cuwarp-size handling: its stochastic pulse train packs a 32-lane ballot into 32-bit words, so the on-device layout must be independent of the physical wavefront width. The compat header redefines the ballot/shuffle sync intrinsics as width-32 logical-warp operations (a wave64 CDNA device selects the calling lane's own 32-lane subgroup), so wave32 and wave64 emit the same stream the CUDA build produces.option(USE_HIP)plus adependencies_hip.cmakemirroring the CUDA path (HIP language,.cuasLANGUAGE HIP, hipBLAS/hipRAND/hipCUB, arch fromCMAKE_HIP_ARCHITECTURES, IPO off since HIP device linking does not finalize LTO). A newRPU_CXX_STANDARDcache variable (default 17) lets the standard track newer PyTorch headers.rpu_base_tiles_cuda.cpp: keyed on the torch hipify generation (not the ROCm version) via a CMake-detectedTORCH_HIPIFY_V2define, since whichc10stream symbol is public depends on the hipify generation./O2flag withif(MSVC), and fix aclassvsstructforward-declaration mismatch that mangles differently on the MSVC ABI.Validation
The analog tiles are exercised on device against the CPU tiles (tolerance-based), via the pytest tile/binding/torch/inference/layer suites:
test_specific_tiles(the bit-line-maker + pulsed-update path) passes 18/18 on each, and the tile/binding/torch/inference/layer suites pass (one order-sensitive stochastic weight-programming test passes in isolation). A fat binary forgfx90a;gfx1100emits both code objects forbit_line_maker.cu, confirming the ballot rewrite is wave-width agnostic.The CUDA/NVIDIA path is unchanged and not affected by this build.
This work was authored with the assistance of Claude, an AI assistant.