[ROCm] Enable MSM on AMD GPUs by extending the HIP backend - #82
Open
jeffdaily wants to merge 1 commit into
Open
Conversation
sppark already carries a HIP backend (util/cuda2hip.hpp, ff/mont_t.hip). This extends it to run multi-scalar multiplication (Pippenger) on AMD GPUs and adds a non-cooperative-launch fallback so the MSM path works across CDNA and RDNA.
Review the four changes in order:
1. Enable MSM (Pippenger) on AMD via HIP. Brings the EC point types (ec/affine_t.hpp, ec/jacobian_t.hpp, ec/xyzz_t.hpp), the field headers, and the Pippenger/sort/batch-addition kernels to the HIP backend, selecting the host (blst) field in the HIP host pass and the device mont_t field in the device pass to mirror CUDA's __CUDA_ARCH__ split. The G2 (fp2) MSM is not built on HIP because mont_t.hip does not yet implement the Montgomery modular inverse it needs; G1 MSM is fully supported.
2. Gate bn254 (alt_bn128) G1 MSM off on the HIP backend. That curve currently hangs the device, so both the build (build.rs) and the TU (#error) refuse the combination rather than ship a binary that wedges the GPU. bls12_381 and bls12_377 G1 MSM are supported on ROCm. bn254 on CUDA is unaffected.
3. Non-cooperative-launch fallback for the MSM sort/accumulate. Some GPUs and runtimes report cooperativeLaunch == 0: RDNA4 has no Global Wave Sync, and Windows ROCm under-reports cooperative launch on RDNA3. When cooperative launch is unavailable the radix sort and bucket accumulation run as a sequence of ordinary kernels whose kernel boundaries supply the grid-wide ordering the cooperative path got from this_grid().sync(); the split is derived identically on both paths so results match. The path is chosen at runtime from the device property (SPPARK_FORCE_NONCOOP forces it for testing), so a CUDA GPU that reports cooperativeLaunch == 1 takes the original cooperative path unchanged.
4. Remove the cooperativeLaunch term from the GPU-selection filter in util/all_gpus.cpp. This change is intentionally backend-agnostic (not guarded): a device that reports cooperativeLaunch == 0 is no longer filtered out of the usable-GPU list, because the fallback in change 3 now handles such devices. Modern CUDA GPUs report cooperativeLaunch == 1, so in practice they are unaffected and still take the cooperative path.
The HIP warp-shuffle compat shims emulate CUDA's logical 32-lane warp semantics on AMD regardless of the physical wavefront size, so the kernels are correct on both wave32 and wave64.
This work was authored with the assistance of Claude, an AI assistant by Anthropic.
Test Plan:
GPU validation of the MSM proof-of-concept (poc/msm-cuda, bls12_381) on AMD hardware, exercising both the cooperative and non-cooperative paths: gfx90a (CDNA2), gfx1100 (RDNA3) Linux, and gfx1101, gfx1201 Windows ROCm.
CUDA build is not regressed (compile-only, no NVIDIA GPU on the build host):
nvcc -c poc/msm-cuda/cuda/pippenger_inf.cu -x cu -std=c++17 -arch=sm_80 \
-DFEATURE_BLS12_381 -DTAKE_RESPONSIBILITY_FOR_ERROR_MESSAGE \
-I<sppark-root> -I<blst>/src -I<cccl-include> \
-Xcompiler -Wno-unused-function # CUDA 13.3, exit 0
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
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.
sppark already carries a HIP backend (util/cuda2hip.hpp, ff/mont_t.hip). This extends it to run multi-scalar multiplication (Pippenger) on AMD GPUs and adds a non-cooperative-launch fallback so the MSM path works across CDNA and RDNA.
Review the four changes in order:
Enable MSM (Pippenger) on AMD via HIP. Brings the EC point types (ec/affine_t.hpp, ec/jacobian_t.hpp, ec/xyzz_t.hpp), the field headers, and the Pippenger/sort/batch-addition kernels to the HIP backend, selecting the host (blst) field in the HIP host pass and the device mont_t field in the device pass to mirror CUDA's CUDA_ARCH split. The G2 (fp2) MSM is not built on HIP because mont_t.hip does not yet implement the Montgomery modular inverse it needs; G1 MSM is fully supported.
Gate bn254 (alt_bn128) G1 MSM off on the HIP backend. That curve currently hangs the device, so both the build (build.rs) and the TU (#error) refuse the combination rather than ship a binary that wedges the GPU. bls12_381 and bls12_377 G1 MSM are supported on ROCm. bn254 on CUDA is unaffected.
Non-cooperative-launch fallback for the MSM sort/accumulate. Some GPUs and runtimes report cooperativeLaunch == 0: RDNA4 has no Global Wave Sync, and Windows ROCm under-reports cooperative launch on RDNA3. When cooperative launch is unavailable the radix sort and bucket accumulation run as a sequence of ordinary kernels whose kernel boundaries supply the grid-wide ordering the cooperative path got from this_grid().sync(); the split is derived identically on both paths so results match. The path is chosen at runtime from the device property (SPPARK_FORCE_NONCOOP forces it for testing), so a CUDA GPU that reports cooperativeLaunch == 1 takes the original cooperative path unchanged.
Remove the cooperativeLaunch term from the GPU-selection filter in util/all_gpus.cpp. This change is intentionally backend-agnostic (not guarded): a device that reports cooperativeLaunch == 0 is no longer filtered out of the usable-GPU list, because the fallback in change 3 now handles such devices. Modern CUDA GPUs report cooperativeLaunch == 1, so in practice they are unaffected and still take the cooperative path.
The HIP warp-shuffle compat shims emulate CUDA's logical 32-lane warp semantics on AMD regardless of the physical wavefront size, so the kernels are correct on both wave32 and wave64.
This work was authored with the assistance of Claude, an AI assistant by Anthropic.
Test Plan:
GPU validation of the MSM proof-of-concept (poc/msm-cuda, bls12_381) on AMD hardware, exercising both the cooperative and non-cooperative paths: gfx90a (CDNA2), gfx1100 (RDNA3) Linux, and gfx1101, gfx1201 Windows ROCm.
CUDA build is not regressed (compile-only, no NVIDIA GPU on the build host):