Skip to content

[ROCm] Enable MSM on AMD GPUs by extending the HIP backend - #82

Open
jeffdaily wants to merge 1 commit into
supranational:mainfrom
jeffdaily:moat-port
Open

[ROCm] Enable MSM on AMD GPUs by extending the HIP backend#82
jeffdaily wants to merge 1 commit into
supranational:mainfrom
jeffdaily:moat-port

Conversation

@jeffdaily

Copy link
Copy Markdown

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

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

1 participant