diff --git a/README.md b/README.md index ca164879..2c9c201b 100644 --- a/README.md +++ b/README.md @@ -23,7 +23,7 @@ The goal of the sppark library is to provide foundational components for applica ## Platform and Language Compatibility -This library primarily supports x86_64 with Nvidia's Volta+ GPU hardware platforms on Linux and Windows operating systems. A limited support for AMD's RDNA and CDNA GPUs is provided. Non-GPU portions can be utilized even on ARM64, and additionally on Mac. +This library primarily supports x86_64 with Nvidia's Volta+ GPU hardware platforms on Linux and Windows operating systems. A limited support for AMD's RDNA and CDNA GPUs is provided through ROCm/HIP, covering the MSM and NTT primitives. Non-GPU portions can be utilized even on ARM64, and additionally on Mac. We show how to interface with Rust and Go. Caveat lector. Achieving highest possible GPU performance requires interfacing with target language memory management, possibly its async facilities, and might even require changes to object's data layout. These are hard to generalize and consequently are also a matter of discussion, likely on a case-by-case basis. diff --git a/ec/affine_t.hpp b/ec/affine_t.hpp index 9aab63ae..82827043 100644 --- a/ec/affine_t.hpp +++ b/ec/affine_t.hpp @@ -5,7 +5,7 @@ #ifndef __SPPARK_EC_AFFINE_T_HPP__ #define __SPPARK_EC_AFFINE_T_HPP__ -#ifndef __CUDACC__ +#if !defined(__CUDACC__) && !defined(__HIPCC__) # undef __host__ # define __host__ # undef __device__ @@ -27,7 +27,7 @@ template friend class xyzz_t; inline __host__ __device__ Affine_t(const field_t& x, const field_t& y) : X(x), Y(y) {} -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) inline __device__ bool is_inf() const { return (bool)(X.is_zero(Y)); } #else @@ -52,7 +52,7 @@ template friend class xyzz_t; friend inline bool operator!=(const Affine_t& p1, const point_t& p2) { return p2 != p1; } -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) class mem_t { field_h X, Y; @@ -91,7 +91,7 @@ template friend class xyzz_t; return affine_t{czero(X, inf), czero(Y, inf)}; } -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) class mem_t { field_h X, Y; int inf[sizeof(field_t)%32 ? (sizeof(field_t)%16 ? 2 : 4) : 8]; diff --git a/ec/jacobian_t.hpp b/ec/jacobian_t.hpp index 16ad15a8..e4efc3f8 100644 --- a/ec/jacobian_t.hpp +++ b/ec/jacobian_t.hpp @@ -7,7 +7,7 @@ #include "affine_t.hpp" -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) # pragma nv_diag_suppress 284 // NULL reference is not allowed #endif @@ -39,7 +39,7 @@ class jacobian_t { return affine_t{xa, ya}; } -#ifdef __CUDACC__ // mask a warning +#if defined(__CUDACC__) || defined(__HIPCC__) // mask a warning inline jacobian_t& operator=(const affine_t& a) { X = a.X; @@ -387,7 +387,7 @@ class jacobian_t { */ void add(const jacobian_t& p2) { -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) jacobian_t p1 = *this; #else jacobian_t &p1 = *this; @@ -479,7 +479,7 @@ class jacobian_t { void add(const affine_t& p2) { -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) jacobian_t p1 = *this; #else jacobian_t &p1 = *this; @@ -586,7 +586,7 @@ class jacobian_t { { return !p1.eq(p2); } }; -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) # pragma nv_diag_default 284 #endif #endif diff --git a/ec/xyzz_t.hpp b/ec/xyzz_t.hpp index 16b03359..09efff74 100644 --- a/ec/xyzz_t.hpp +++ b/ec/xyzz_t.hpp @@ -7,7 +7,7 @@ #include "affine_t.hpp" -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) # pragma nv_diag_suppress 284 // NULL reference is not allowed #endif @@ -25,7 +25,7 @@ class xyzz_t { ZZZ(field_t::one(is_inf)), ZZ(ZZZ) {} -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) class mem_t { friend class xyzz_t; field_h X, Y, ZZZ, ZZ; @@ -89,7 +89,7 @@ class xyzz_t { { return jacobian_t{ X*ZZ, Y*ZZZ, ZZ }; } #endif -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) inline __device__ bool is_inf() const { return (bool)(ZZZ.is_zero(ZZ)); } #else @@ -99,12 +99,16 @@ class xyzz_t { inline __host__ __device__ void inf() { ZZZ.zero(); ZZ.zero(); } inline __host__ __device__ void cneg(bool neg) { ZZZ.cneg(neg); } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) static inline __device__ void prefetch(const xyzz_t* p_) { const unsigned char* p = (const unsigned char*)p_; for (size_t i = 0; i < sizeof(*p_); i += 128) +#if defined(__HIP_DEVICE_COMPILE__) + __builtin_prefetch(p+i, 0, 1); +#else asm("prefetch.global.L2 [%0];" :: "l"(p+i)); +#endif } #endif @@ -123,7 +127,7 @@ class xyzz_t { return; } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) xyzz_t p31 = *this; #else xyzz_t& p31 = *this; @@ -171,7 +175,7 @@ class xyzz_t { M = p31.X^2; M = M + M + M; /* M = 3*X1^2[+a*ZZ1^2] */ if (a4 != nullptr) { -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) U = *a4; U *= p31.ZZ^2; #else @@ -194,12 +198,12 @@ class xyzz_t { } else { /* X1==X2 && Y1==-Y2 */\ p31.inf(); /* set |p3| to infinity */\ } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) *this = p31; #endif } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) __device__ void uadd(const xyzz_t& p2) { xyzz_t p31 = *this; @@ -339,7 +343,7 @@ class xyzz_t { *this = p31; } #else - inline void uadd(const xyzz_t& p2) { add(p2); } + __host__ __device__ void uadd(const xyzz_t& p2) { add(p2); } #endif /* @@ -351,7 +355,7 @@ class xyzz_t { template __host__ __device__ void add(const affine_t& p2, bool subtract = false) { -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) xyzz_t p31 = *this; #else xyzz_t& p31 = *this; @@ -403,7 +407,7 @@ class xyzz_t { M = p2.X^2; M = M + M + M; /* M = 3*X1^2[+a] */ if (a4 != nullptr) { -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) M += (U = *a4); #else M += *a4; @@ -423,12 +427,12 @@ class xyzz_t { p31.inf(); /* set |p3| to infinity */ } } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) *this = p31; #endif } -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) template __device__ void uadd(const affine_t& p2, bool subtract = false) { @@ -542,12 +546,12 @@ class xyzz_t { } #else template - inline void uadd(const affine_t& p2, bool subtract = false) + __host__ __device__ void uadd(const affine_t& p2, bool subtract = false) { add(p2, subtract); } #endif }; -#ifdef __CUDACC__ +#if defined(__CUDACC__) || defined(__HIPCC__) # pragma nv_diag_default 284 #endif #endif diff --git a/ff/alt_bn128.hpp b/ff/alt_bn128.hpp index 72207a3a..5544a555 100644 --- a/ff/alt_bn128.hpp +++ b/ff/alt_bn128.hpp @@ -46,10 +46,12 @@ namespace device { }; static __device__ __constant__ const uint32_t ALT_BN128_m0 = 0xefffffff; } -# if defined(__CUDA_ARCH__) || defined(__HIPCC__) // device-side field types +# if defined(__CUDA_ARCH__) || \ + (defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) || \ + (!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // device-side field types # if defined(__CUDA_ARCH__) # include "mont_t.cuh" -# elif defined(__HIPCC__) +# else # include "mont_t.hip" typedef uint64_t vec256[4]; # endif @@ -61,8 +63,8 @@ typedef mont_t<254, device::ALT_BN128_P, device::ALT_BN128_M0, device::ALT_BN128_Px4> fp_mont; struct fp_t : public fp_mont { using mem_t = fp_t; - __device__ __forceinline__ fp_t() {} - __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} + __host__ __device__ __forceinline__ fp_t() {} + __host__ __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} template constexpr fp_t(Ts... a) : fp_mont{a...} {} }; typedef mont_t<254, device::ALT_BN128_r, device::ALT_BN128_m0, @@ -70,8 +72,8 @@ typedef mont_t<254, device::ALT_BN128_r, device::ALT_BN128_m0, device::ALT_BN128_rx4> fr_mont; struct fr_t : public fr_mont { using mem_t = fr_t; - __device__ __forceinline__ fr_t() {} - __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} + __host__ __device__ __forceinline__ fr_t() {} + __host__ __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} template constexpr fr_t(Ts... a) : fr_mont{a...} {} # ifdef __HIPCC__ __host__ __forceinline__ fr_t(vec256 a) : fr_mont(a) {} @@ -83,7 +85,9 @@ struct fr_t : public fr_mont { # endif #endif -#if !defined(__CUDA_ARCH__) && !defined(__HIPCC__) // host-side field types +#if !defined(__CUDA_ARCH__) && \ + !(defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) && \ + !(!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // host-side field types # include # if defined(__GNUC__) && !defined(__clang__) diff --git a/ff/bls12-377.hpp b/ff/bls12-377.hpp index 9b010a33..6a96a67a 100644 --- a/ff/bls12-377.hpp +++ b/ff/bls12-377.hpp @@ -50,10 +50,12 @@ namespace device { }; static __device__ __constant__ /*const*/ uint32_t BLS12_377_m0 = 0xffffffff; } -# if defined(__CUDA_ARCH__) || defined(__HIPCC__) // device-side field types +# if defined(__CUDA_ARCH__) || \ + (defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) || \ + (!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // device-side field types # if defined(__CUDA_ARCH__) # include "mont_t.cuh" -# elif defined(__HIPCC__) +# else # include "mont_t.hip" typedef uint64_t vec256[4]; # endif @@ -65,8 +67,8 @@ typedef mont_t<377, device::BLS12_377_P, device::BLS12_377_M0, device::BLS12_381_Px128> fp_mont; struct fp_t : public fp_mont { using mem_t = fp_t; - __device__ __forceinline__ fp_t() {} - __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} + __host__ __device__ __forceinline__ fp_t() {} + __host__ __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} template constexpr fp_t(Ts... a) : fp_mont{a...} {} }; typedef mont_t<253, device::BLS12_377_r, device::BLS12_377_m0, @@ -74,8 +76,8 @@ typedef mont_t<253, device::BLS12_377_r, device::BLS12_377_m0, device::BLS12_377_rx8> fr_mont; struct fr_t : public fr_mont { using mem_t = fr_t; - __device__ __forceinline__ fr_t() {} - __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} + __host__ __device__ __forceinline__ fr_t() {} + __host__ __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} template constexpr fr_t(Ts... a) : fr_mont{a...} {} # ifdef __HIPCC__ __host__ __forceinline__ fr_t(vec256 a) : fr_mont(a) {} @@ -87,7 +89,9 @@ struct fr_t : public fr_mont { # endif #endif -#if !defined(__CUDA_ARCH__) && !defined(__HIPCC__) // host-side field types +#if !defined(__CUDA_ARCH__) && \ + !(defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) && \ + !(!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // host-side field types # include # if defined(__GNUC__) && !defined(__clang__) diff --git a/ff/bls12-381.hpp b/ff/bls12-381.hpp index 52e0135d..5779855f 100644 --- a/ff/bls12-381.hpp +++ b/ff/bls12-381.hpp @@ -50,10 +50,12 @@ namespace device { }; static __device__ __constant__ /*const*/ uint32_t BLS12_381_m0 = 0xffffffff; } -# if defined(__CUDA_ARCH__) || defined(__HIPCC__) // device-side field types +# if defined(__CUDA_ARCH__) || \ + (defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) || \ + (!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // device-side field types # if defined(__CUDA_ARCH__) # include "mont_t.cuh" -# elif defined(__HIPCC__) +# else # include "mont_t.hip" typedef uint64_t vec256[4]; # endif @@ -65,8 +67,8 @@ typedef mont_t<381, device::BLS12_381_P, device::BLS12_381_M0, device::BLS12_381_Px8> fp_mont; struct fp_t : public fp_mont { using mem_t = fp_t; - __device__ __forceinline__ fp_t() {} - __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} + __host__ __device__ __forceinline__ fp_t() {} + __host__ __device__ __forceinline__ fp_t(const fp_mont& a) : fp_mont(a) {} template constexpr fp_t(Ts... a) : fp_mont{a...} {} }; typedef mont_t<255, device::BLS12_381_r, device::BLS12_381_m0, @@ -74,8 +76,8 @@ typedef mont_t<255, device::BLS12_381_r, device::BLS12_381_m0, device::BLS12_381_rx2> fr_mont; struct fr_t : public fr_mont { using mem_t = fr_t; - __device__ __forceinline__ fr_t() {} - __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} + __host__ __device__ __forceinline__ fr_t() {} + __host__ __device__ __forceinline__ fr_t(const fr_mont& a) : fr_mont(a) {} template constexpr fr_t(Ts... a) : fr_mont{a...} {} # ifdef __HIPCC__ __host__ __forceinline__ fr_t(vec256 a) : fr_mont(a) {} @@ -87,7 +89,9 @@ struct fr_t : public fr_mont { # endif #endif -#if !defined(__CUDA_ARCH__) && !defined(__HIPCC__) // host-side field types +#if !defined(__CUDA_ARCH__) && \ + !(defined(SPPARK_HIP_HOST_FIELD) && defined(__HIP_DEVICE_COMPILE__)) && \ + !(!defined(SPPARK_HIP_HOST_FIELD) && defined(__HIPCC__)) // host-side field types # include # if defined(__GNUC__) && !defined(__clang__) diff --git a/ff/mont_t.hip b/ff/mont_t.hip index a515f787..fa049e4d 100644 --- a/ff/mont_t.hip +++ b/ff/mont_t.hip @@ -1,6 +1,10 @@ // Copyright Supranational LLC +// Copyright (c) 2026 Advanced Micro Devices, Inc. // Licensed under the Apache License, Version 2.0, see LICENSE for details. // SPDX-License-Identifier: Apache-2.0 +// +// ROCm/HIP host+device field arithmetic and MSM support extensions +// by Jeff Daily . #if defined(__HIPCC__) && !defined(__SPPARK_FF_MONT_T_HIP__) # define __SPPARK_FF_MONT_T_HIP__ @@ -30,7 +34,14 @@ # define __MONT_T_XSTR(x) __MONT_T_STR(x) # define S_OP(op) "s_" #op "_b" __MONT_T_XSTR(__AMDGCN_WAVEFRONT_SIZE) " " -# define inline __device__ __forceinline__ +// The MSM driver runs its post-kernel point accumulation (collect/integrate) +// on the CPU using this same field type, so the field must be callable from +// host as well as device. clang's HIP front end type-checks host functions in +// the device pass and vice versa, so every method is __host__ __device__ with +// the GCN-asm fast path guarded by __HIP_DEVICE_COMPILE__ and a portable C++ +// fallback for the host pass. (nvcc tolerates the device-only spelling because +// it discards unreferenced host code from its device pass; clang does not.) +# define inline __host__ __device__ __forceinline__ // // To instantiate declare modulus as __device__ __constant___ const and @@ -44,7 +55,7 @@ template> 32; + } + carry = (cond_t)c; +#endif + final_sub(carry); + } + + return *this; + } + friend inline mont_t operator<<(mont_t a, unsigned l) + { return a <<= l; } + + inline mont_t& operator>>=(unsigned r) + { + while (r--) { + uint32_t mask = 0 - (val[0] & 1); // all ones if odd + uint64_t c = 0; + uint32_t tmp[n+1]; + for (size_t i=0; i> 32; + } + tmp[n] = (N % 32 == 0) ? (uint32_t)c : 0; + + for (size_t i=0; i> 1) | (tmp[i+1] << 31); + } + + return *this; + } + friend inline mont_t operator>>(mont_t a, unsigned r) + { return a >>= r; } + inline mont_t& operator+=(const mont_t& b) { cond_t carry; +#if defined(__HIP_DEVICE_COMPILE__) asm("v_add_co_u32 %0, %1, %0, %2" : "+v"(val[0]), "=s"(carry) : "v"(b[0])); for (size_t i=1; i> 32; + } + carry = (cond_t)c; +#endif final_sub(carry); return *this; @@ -114,6 +186,7 @@ public: cond_t borrow, carry; uint32_t tmp[n]; +#if defined(__HIP_DEVICE_COMPILE__) asm("v_sub_co_u32 %0, %1, %0, %2" : "+v"(val[0]), "=s"(borrow) : "v"(b[0])); for (size_t i=1; i> 32) & 1; + } + borrow = (cond_t)(0 - bw); // all-ones mask when a borrow occurred + uint64_t c = 0; + for (size_t i=0; i> 32; + } +#endif return *this; } @@ -163,8 +252,14 @@ public: if (N%32 != 0) { ret[n-1] = u[1] + top; } else { +#if defined(__HIP_DEVICE_COMPILE__) asm( v_addc_co_u32 "%0, %1, %2, %3, %1" : "=v"(ret[n-1]), "+s"(carry) : "v"(u[1]), "v"(top)); +#else + uint64_t s = (uint64_t)u[1] + top + (carry & 1); + ret[n-1] = (uint32_t)s; + carry = (cond_t)(s >> 32); +#endif } if (++j == n) @@ -180,8 +275,14 @@ public: if (N%32 != 0) { top = u[1]; } else { +#if defined(__HIP_DEVICE_COMPILE__) asm( v_addc_co_u32 "%0, %1, %2, 0, %1" : "=v"(top), "+s"(carry) : "v"(u[1])); +#else + uint64_t s = (uint64_t)u[1] + (carry & 1); + top = (uint32_t)s; + carry = (cond_t)(s >> 32); +#endif } } @@ -220,6 +321,14 @@ public: static inline const mont_t& one() { return *reinterpret_cast(ONE); } + static inline mont_t one(int or_zero) + { + mont_t ret; + for (size_t i = 0; i < n; i++) + ret[i] = or_zero ? 0 : ONE[i]; + return ret; + } + inline void zero() { if (n%4 == 0) { @@ -236,35 +345,139 @@ public: friend inline mont_t czero(const mont_t& a, int set_z) { mont_t ret; - cond_t cond; +#if defined(__HIP_DEVICE_COMPILE__) + cond_t cond; asm("v_cmp_ne_u32 %0, %1, 0" : "=s"(cond) : "v"(set_z)); for (size_t i=0; i> 32) & 1; + } + + if (flag && nz) + for (size_t i = 0; i < n; i++) + val[i] = tmp[i]; + + return *this; + } + static inline mont_t cneg(mont_t a, bool flag) + { return a.cneg(flag); } + + inline mont_t operator-() const + { return cneg(*this, true); } + + // make the value "positive" and return the original "sign" + inline bool abs() + { + uint32_t tmp[n]; + uint64_t borrow = 0; + + for (size_t i = 0; i < n; i++) { + uint64_t d = (uint64_t)MOD[i] - val[i] - borrow; + tmp[i] = (uint32_t)d; + borrow = (d >> 32) & 1; + } + + bool sign = tmp[n-1] < val[n-1]; + + if (sign) + for (size_t i = 0; i < n; i++) + val[i] = tmp[i]; + + return sign; } private: @@ -273,6 +486,7 @@ private: cond_t borrow; uint32_t tmp[n]; +#if defined(__HIP_DEVICE_COMPILE__) asm("v_sub_co_u32 %0, %1, %2, %3" : "=v"(tmp[0]), "=s"(borrow) : "v"(val[0]), "v"(MOD[0])); for (size_t i=1; i> 32) & 1; + } + borrow = (cond_t)bw; + // subtract MOD when there was an incoming carry OR val >= MOD (no borrow) + if ((carry & 1) || !borrow) + for (size_t i=0; i -#include +#if defined(__HIPCC__) +# include +#else +# include +# include +#endif #include #include @@ -23,13 +27,14 @@ #endif template + class bucket_h = typename bucket_t::mem_t, + class affine_t = typename bucket_t::affine_t> __device__ __forceinline__ static void add(bucket_h ret[], const affine_h points[], uint32_t npoints, const uint32_t bitmap[], const uint32_t refmap[], bool accumulate, uint32_t sid) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) static __device__ uint32_t streams[BATCH_ADD_NSTREAMS]; uint32_t& current = streams[sid % BATCH_ADD_NSTREAMS]; @@ -38,8 +43,12 @@ static void add(bucket_h ret[], const affine_h points[], uint32_t npoints, const uint32_t tid = (threadIdx.x + blockDim.x*blockIdx.x) / degree; const uint32_t xid = tid % warp_sz; +#if defined(__HIPCC__) + uint32_t laneid = threadIdx.x % WARP_SZ; +#else uint32_t laneid; asm("mov.u32 %0, %laneid;" : "=r"(laneid)); +#endif bucket_t acc; acc.inf(); @@ -104,18 +113,23 @@ static void add(bucket_h ret[], const affine_h points[], uint32_t npoints, acc.uadd(down); // .add() triggers spills ... in .shfl_down() } +#if defined(__HIP_DEVICE_COMPILE__) + SPPARK_GRID_SYNC(); +#elif defined(__CUDA_ARCH__) cooperative_groups::this_grid().sync(); +#endif if (xid == 0) ret[tid/warp_sz] = acc; if (threadIdx.x + blockIdx.x == 0) current = 0; +#endif } template + class bucket_h = typename bucket_t::mem_t, + class affine_t = typename bucket_t::affine_t> __launch_bounds__(BATCH_ADD_BLOCK_SIZE) __global__ void batch_addition(bucket_h ret[], const affine_h points[], uint32_t npoints, const uint32_t bitmap[], bool accumulate = false, @@ -123,8 +137,8 @@ void batch_addition(bucket_h ret[], const affine_h points[], uint32_t npoints, { add(ret, points, npoints, bitmap, nullptr, accumulate, sid); } template + class bucket_h = typename bucket_t::mem_t, + class affine_t = typename bucket_t::affine_t> __launch_bounds__(BATCH_ADD_BLOCK_SIZE) __global__ void batch_diff(bucket_h ret[], const affine_h points[], uint32_t npoints, const uint32_t bitmap[], const uint32_t refmap[], @@ -132,12 +146,13 @@ void batch_diff(bucket_h ret[], const affine_h points[], uint32_t npoints, { add(ret, points, npoints, bitmap, refmap, accumulate, sid); } template + class bucket_h = typename bucket_t::mem_t, + class affine_t = typename bucket_t::affine_t> __launch_bounds__(BATCH_ADD_BLOCK_SIZE) __global__ void batch_addition(bucket_h ret[], const affine_h points[], size_t npoints, const uint32_t digits[], const uint32_t& ndigits) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) const uint32_t degree = bucket_t::degree; const uint32_t warp_sz = WARP_SZ / degree; const uint32_t tid = (threadIdx.x + blockDim.x*blockIdx.x) / degree; @@ -165,6 +180,7 @@ void batch_addition(bucket_h ret[], const affine_h points[], size_t npoints, if (xid == 0) ret[tid/warp_sz] = acc; +#endif } template diff --git a/msm/pippenger.cuh b/msm/pippenger.cuh index 2d9e2106..d53f55e9 100644 --- a/msm/pippenger.cuh +++ b/msm/pippenger.cuh @@ -1,12 +1,17 @@ // Copyright Supranational LLC +// Copyright (c) 2026 Advanced Micro Devices, Inc. // Licensed under the Apache License, Version 2.0, see LICENSE for details. // SPDX-License-Identifier: Apache-2.0 #ifndef __SPPARK_MSM_PIPPENGER_CUH__ #define __SPPARK_MSM_PIPPENGER_CUH__ -#include -#include +#if defined(__HIPCC__) +# include +#else +# include +# include +#endif #include #include @@ -28,7 +33,7 @@ * Break down |scalars| to signed |wbits|-wide digits. */ -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) // Transposed scalar_t template class scalar_T { @@ -76,7 +81,7 @@ void breakdown(vec2d_t digits, const scalar_t scalars[], size_t len, { assert(len <= (1U<<31) && wbits < 32); -#ifdef __CUDA_ARCH__ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) extern __shared__ scalar_T xchange[]; const uint32_t tid = threadIdx.x; const uint32_t tix = threadIdx.x + blockIdx.x*blockDim.x; @@ -142,22 +147,44 @@ void breakdown(vec2d_t digits, const scalar_t scalars[], size_t len, # error "invalid MSM_NSTREAMS" #endif +// Persistent per-stream work-queue counter for accumulate(). Hoisted to namespace +// scope (it was a function-local static __device__) so the non-cooperative path can +// zero it from a separate reset kernel before each launch. The cooperative path +// resets it in-kernel after a grid.sync, exactly as before. +static __device__ uint32_t accumulate_streams[MSM_NSTREAMS]; + +// Zero one accumulate work-queue counter. Used only by the non-cooperative path, +// which cannot reset the counter in-kernel (that needs a grid-wide barrier). +// Templated so the definition has vague linkage in a header. +template +__global__ void reset_accumulate_counter(uint32_t sid) +{ +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) + accumulate_streams[sid % MSM_NSTREAMS] = 0; +#endif +} + template + bool COOP = true, + class bucket_h = typename bucket_t::mem_t, + class affine_t = typename bucket_t::affine_t> __launch_bounds__(ACCUMULATE_NTHREADS) __global__ void accumulate(bucket_h buckets_[], uint32_t nwins, uint32_t wbits, /*const*/ affine_h points_[], const vec2d_t digits, const vec2d_t histogram, uint32_t sid = 0) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) vec2d_t buckets{buckets_, 1U<<--wbits}; const affine_h* points = points_; - static __device__ uint32_t streams[MSM_NSTREAMS]; - uint32_t& current = streams[sid % MSM_NSTREAMS]; + uint32_t& current = accumulate_streams[sid % MSM_NSTREAMS]; +#if defined(__HIPCC__) + uint32_t laneid = threadIdx.x % WARP_SZ; +#else uint32_t laneid; asm("mov.u32 %0, %laneid;" : "=r"(laneid)); +#endif const uint32_t degree = bucket_t::degree; const uint32_t warp_sz = WARP_SZ / degree; const uint32_t lane_id = laneid / degree; @@ -182,10 +209,16 @@ void accumulate(bucket_h buckets_[], uint32_t nwins, uint32_t wbits, uint32_t idx, len = h[0]; +#if defined(__HIPCC__) + idx = __shfl_up_sync(0xffffffff, len, degree, WARP_SZ); + if ((threadIdx.x % WARP_SZ) < degree) + idx = 0; +#else asm("{ .reg.pred %did;" " shfl.sync.up.b32 %0|%did, %1, %2, 0, 0xffffffff;" " @!%did mov.b32 %0, 0;" "}" : "=r"(idx) : "r"(len), "r"(degree)); +#endif if (lane_id == 0 && x != 0) idx = h[-1]; @@ -216,16 +249,28 @@ void accumulate(bucket_h buckets_[], uint32_t nwins, uint32_t wbits, x = __shfl_sync(0xffffffff, x, 0) + lane_id; } - cooperative_groups::this_grid().sync(); - - if (threadIdx.x + blockIdx.x == 0) - current = 0; + // The grid-wide barrier exists only so every block finishes draining the work + // queue before block 0 resets |current| for the next launch. On GPUs without + // cooperative launch (COOP == false) the host resets the counter with + // reset_accumulate_counter before each launch instead, and the following + // integrate kernel's launch boundary fences the bucket writes either way. + if constexpr (COOP) { +#if defined(__HIP_DEVICE_COMPILE__) + SPPARK_GRID_SYNC(); +#elif defined(__CUDA_ARCH__) + cooperative_groups::this_grid().sync(); +#endif + if (threadIdx.x + blockIdx.x == 0) + current = 0; + } +#endif } -template +template __launch_bounds__(256) __global__ void integrate(bucket_h buckets_[], uint32_t nwins, uint32_t wbits, uint32_t nbits) { +#if defined(__CUDA_ARCH__) || defined(__HIP_DEVICE_COMPILE__) const uint32_t degree = bucket_t::degree; uint32_t Nthrbits = 31 - __clz(blockDim.x / degree); @@ -293,6 +338,7 @@ void integrate(bucket_h buckets_[], uint32_t nwins, uint32_t wbits, uint32_t nbi buckets[bid][2*tid] = p; buckets[bid][2*tid+1] = acc; +#endif } #undef asm @@ -317,14 +363,15 @@ void breakdown(vec2d_t digits, const scalar_t scalars[], #endif #include +#include #include #include #include template + class affine_h = typename affine_t::mem_t, + class bucket_h = typename bucket_t::mem_t> class msm_t { const gpu_t& gpu; size_t npoints; @@ -333,6 +380,7 @@ class msm_t { affine_h *d_points; scalar_t *d_scalars; vec2d_t d_hist; + bool coop; // device supports cooperative launch (Global Wave Sync) template using vec_t = slice_t; @@ -352,6 +400,15 @@ public: size_t ffi_affine_sz = sizeof(affine_t), int device_id = -1) : gpu(select_gpu(device_id)), d_points(nullptr), d_scalars(nullptr) { + // Cooperative launch (hipLaunchCooperativeKernel / cudaLaunchCooperative- + // Kernel) needs Global Wave Sync, which some GPUs/runtimes lack (RDNA4 has + // no GWS; Windows ROCm under-reports it on RDNA3). When absent, the MSM + // sort/accumulate kernels are driven through the non-cooperative fallback. + // SPPARK_FORCE_NONCOOP forces that path for testing on capable GPUs. + coop = gpu.props().cooperativeLaunch != 0; + if (getenv("SPPARK_FORCE_NONCOOP")) + coop = false; + npoints = (np+WARP_SZ-1) & ((size_t)0-WARP_SZ); wbits = 17; @@ -395,6 +452,66 @@ public: } private: + // Launch one (1- or 2-window) radix sort. With cooperative launch this is the + // single persistent `sort` kernel; without it, the two-level pass is split into + // sort_upper_count -> sort_upper_scatter -> sort_lower (same grid/block, kernel + // boundaries standing in for the grid-wide barriers), and the small single-block + // case reuses `sort` with gridDim.x == 1 (it has no barrier). + void launch_sort(const stream_t& s, dim3 grid, size_t shared_sz, + vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t kwbits, uint32_t lsbits0, uint32_t lsbits1) + { + if (coop) { + s.launch_coop(sort, {grid, SORT_BLOCKDIM, shared_sz}, + inouts, len, win, temps, histograms, + kwbits, lsbits0, lsbits1); + return; + } + + int lg = lg2(grid.x); + bool two_level = (kwbits > DIGIT_BITS) || + (lg && kwbits > (uint32_t)lg + 1); + if (two_level) { + sort_upper_count<<>>( + inouts, len, win, temps, histograms, kwbits, lsbits0, lsbits1); + CUDA_OK(cudaGetLastError()); + sort_upper_scatter<<>>( + inouts, len, win, temps, histograms, kwbits, lsbits0, lsbits1); + CUDA_OK(cudaGetLastError()); + sort_lower<<>>( + inouts, len, win, temps, histograms, kwbits, lsbits0, lsbits1); + CUDA_OK(cudaGetLastError()); + } else { + grid.x = 1; + sort<<>>( + inouts, len, win, temps, histograms, kwbits, lsbits0, lsbits1); + CUDA_OK(cudaGetLastError()); + } + } + + // Launch the bucket accumulation. With cooperative launch it is the persistent + // `accumulate` kernel (resets its work queue in-kernel after a grid.sync); + // without it, the work-queue counter is zeroed by a tiny reset kernel beforehand + // and the COOP=false instantiation drops the trailing barrier/reset. + void launch_accumulate(const stream_t& s, uint32_t sid, bucket_h* buckets_, + affine_h* points_, vec2d_t d_digits) + { + if (coop) { + s.launch_coop(accumulate, + {gpu.sm_count(), 0}, + buckets_, nwins, wbits, points_, d_digits, d_hist, sid); + return; + } + + reset_accumulate_counter<><<<1, 1, 0, s>>>(sid); + CUDA_OK(cudaGetLastError()); + accumulate + <<>>( + buckets_, nwins, wbits, points_, d_digits, d_hist, sid); + CUDA_OK(cudaGetLastError()); + } + void digits(const scalar_t d_scalars[], size_t len, vec2d_t& d_digits, vec2d_t&d_temps, bool mont) { @@ -432,14 +549,14 @@ private: uint32_t top = scalar_t::bit_length() - wbits * (nwins-1); uint32_t win; for (win = 0; win < nwins-1; win += 2) { - gpu[2].launch_coop(sort, {{grid_size, 2}, SORT_BLOCKDIM, shared_sz}, - d_digits, len, win, d_temps, d_hist, - wbits-1, wbits-1, win == nwins-2 ? top-1 : wbits-1); + launch_sort(gpu[2], {grid_size, 2}, shared_sz, + d_digits, len, win, d_temps, d_hist, + wbits-1, wbits-1, win == nwins-2 ? top-1 : wbits-1); } if (win < nwins) { - gpu[2].launch_coop(sort, {{grid_size, 1}, SORT_BLOCKDIM, shared_sz}, - d_digits, len, win, d_temps, d_hist, - wbits-1, top-1, 0u); + launch_sort(gpu[2], {grid_size, 1}, shared_sz, + d_digits, len, win, d_temps, d_hist, + wbits-1, top-1, 0u); } #endif } @@ -512,10 +629,8 @@ public: ); CUDA_OK(cudaGetLastError()); - gpu[i&1].launch_coop(accumulate, - {gpu.sm_count(), 0}, - d_buckets, nwins, wbits, &d_points[d_off], d_digits, d_hist, i&1 - ); + launch_accumulate(gpu[i&1], i&1, d_buckets, &d_points[d_off], + d_digits); gpu[i&1].record(ev); integrate<< template + class affine_t = typename bucket_t::affine_t> static void mult_pippenger(point_t& ret, const affine_t points[], size_t npoints, const scalar_t _scalars[], bool mont, thread_pool_t* da_pool = nullptr) @@ -350,7 +350,7 @@ static void mult_pippenger(point_t& ret, const affine_t points[], size_t npoints } template + class affine_t = typename bucket_t::affine_t> static void mult_pippenger(point_t& ret, const std::vector& points, const std::vector& scalars, bool mont, thread_pool_t* da_pool = nullptr) @@ -363,7 +363,7 @@ static void mult_pippenger(point_t& ret, const std::vector& points, #include template + class affine_t = typename bucket_t::affine_t> static void mult_pippenger(point_t& ret, slice_t points, slice_t scalars, bool mont, thread_pool_t* da_pool = nullptr) diff --git a/msm/sort.cuh b/msm/sort.cuh index 0b2928ce..e834efee 100644 --- a/msm/sort.cuh +++ b/msm/sort.cuh @@ -1,4 +1,5 @@ // Copyright Supranational LLC +// Copyright (c) 2026 Advanced Micro Devices, Inc. // Licensed under the Apache License, Version 2.0, see LICENSE for details. // SPDX-License-Identifier: Apache-2.0 @@ -22,6 +23,27 @@ __global__ void sort(vec2d_t inouts, size_t len, uint32_t win, vec2d_t temps, vec2d_t histograms, uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1); +// Non-cooperative fallback kernels for GPUs/runtimes without cooperative launch +// (cooperativeLaunch == 0, e.g. RDNA4 has no Global Wave Sync and Windows ROCm +// under-reports it on RDNA3). The cooperative `sort` above runs the cross-block +// top-level radix pass as one persistent kernel with grid-wide barriers; these +// split that pass at its barriers into ordinary kernels whose boundaries provide +// the same grid-wide ordering. Launch sort_upper_count, then sort_upper_scatter, +// then sort_lower (same grid/block as the cooperative launch) for the two-level +// case; the single-block case reuses `sort` with gridDim.x == 1. +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_upper_count(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1); +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_upper_scatter(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1); +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_lower(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1); + #ifndef __MSM_SORT_DONT_IMPLEMENT__ #ifndef WARP_SZ @@ -42,9 +64,13 @@ uint32_t pack(uint32_t a, uint32_t mask, uint32_t b) { uint32_t ret; +#if defined(__HIPCC__) + ret = (a & ~mask) | (mask & b); +#else asm("lop3.b32 %0, %1, %2, %3, 0xb8;" // a & ~mask | mask & b : "=r"(ret) : "r"(a), "r"(mask), "r"(b)); +#endif return ret; } @@ -52,12 +78,22 @@ uint32_t pack(uint32_t a, uint32_t mask, uint32_t b) __device__ __forceinline__ uint32_t sum_up(uint32_t sum, const uint32_t limit = WARP_SZ) { +#if defined(__HIPCC__) + const uint32_t laneid = threadIdx.x % WARP_SZ; + #pragma unroll + for (uint32_t off = 1; off < limit; off <<= 1) { + uint32_t v = __shfl_up_sync(0xffffffff, sum, off, WARP_SZ); + if (laneid >= off) + sum += v; + } +#else #pragma unroll for (uint32_t off = 1; off < limit; off <<= 1) asm("{ .reg.b32 %v; .reg.pred %did;" " shfl.sync.up.b32 %v|%did, %0, %1, 0, 0xffffffff;" " @%did add.u32 %0, %0, %v;" "}" : "+r"(sum) : "r"(off)); +#endif return sum; } @@ -117,22 +153,35 @@ void scatter(uint2 dst[], const uint32_t src[], uint32_t base, uint32_t len, } } -__device__ -static void upper_sort(uint2 dst[], const uint32_t src[], uint32_t len, - uint32_t lsbits, uint32_t bits, uint32_t digit, - uint32_t histogram[]) +// Partition |len| elements across the grid (one contiguous slice per block). +// Pure function of len, gridDim.x, blockIdx.x, so each phase recomputes it +// identically whether the phases run as one cooperative kernel or several. +__device__ __forceinline__ +static void upper_partition(uint32_t len, uint32_t& base, uint32_t& slice) { uint32_t grid_div = 31 - __clz(gridDim.x); uint32_t grid_rem = (1<> grid_div; // / gridDim.x; - uint32_t rem = len & grid_rem; // % gridDim.x; - uint32_t base; + slice = len >> grid_div; // / gridDim.x; + uint32_t rem = len & grid_rem; // % gridDim.x; if (blockIdx.x < rem) base = ++slice * blockIdx.x; else base = slice * blockIdx.x + rem; +} + +// Phase 1 of the cross-block top-level radix pass: count this block's digits and +// publish its per-digit histogram to global memory. All cross-block state lives in +// |histogram|, so a grid-wide barrier (cooperative path) or a kernel boundary +// (fallback path) after this is sufficient before phase 2. +__device__ __forceinline__ +static void upper_count_emit(const uint32_t src[], uint32_t len, + uint32_t lsbits, uint32_t bits, uint32_t digit, + uint32_t histogram[]) +{ + uint32_t base, slice; + upper_partition(len, base, slice); const uint32_t mask = (1<dst by +// the resulting per-block offsets, and (block 0) finalize the global digit bases. +// The shared counters[] are recomputed from |histogram| here, so this needs no +// shared state carried over from phase 1. +__device__ __forceinline__ +static void upper_scan_scatter(uint2 dst[], const uint32_t src[], uint32_t len, + uint32_t lsbits, uint32_t bits, uint32_t digit, + uint32_t histogram[]) +{ + uint32_t base, slice; + upper_partition(len, base, slice); + + const uint32_t mask = (1<= 1) + h.x = v; + } +#else asm("{ .reg.b32 %v; .reg.pred %did;"); asm("shfl.sync.up.b32 %v|%did, %0, 1, 0, 0xffffffff;" :: "r"(sum)); asm("@%did mov.b32 %0, %v;" : "+r"(h.x)); asm("}"); +#endif h.y = __shfl_down_sync(0xffffffff, sum, gridDim.x-1) - h.x; if (blockIdx.x == 0 && sub_laneid == 0 && warp_off < 1< DIGIT_BITS || (lg_gridDim && wbits > lg_gridDim+1)) { - uint32_t top_bits = wbits / 2; - uint32_t low_bits = wbits - top_bits; - - if (low_bits < lg_gridDim+1) { - low_bits = lg_gridDim+1; - top_bits = wbits - low_bits; - } + uint32_t top_bits, low_bits; + upper_bits(wbits, top_bits, low_bits); upper_sort(temp, inout, len, lsbits, top_bits, low_bits, histogram); - histogram += blockIdx.x< inouts, size_t len, uint32_t win, wbits, blockIdx.y==0 ? lsbits0 : lsbits1); } +// Non-cooperative fallback for the two-level (cross-block) sort_row path. The host +// launches these three kernels in order, with the same grid/block as the +// cooperative `sort`, when the device lacks cooperative launch. Each kernel +// boundary supplies the grid-wide ordering the cooperative path got from +// grid.sync. The single-block sort_row path needs no barrier and reuses `sort` +// with gridDim.x == 1, so it has no fallback kernel here. +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_upper_count(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1) +{ + win += blockIdx.y; + uint32_t lsbits = blockIdx.y==0 ? lsbits0 : lsbits1; + uint32_t top_bits, low_bits; + upper_bits(wbits, top_bits, low_bits); + upper_count_emit(inouts[win], len, lsbits, top_bits, low_bits, histograms[win]); +} + +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_upper_scatter(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1) +{ + win += blockIdx.y; + uint32_t lsbits = blockIdx.y==0 ? lsbits0 : lsbits1; + uint32_t top_bits, low_bits; + upper_bits(wbits, top_bits, low_bits); + upper_scan_scatter(temps[blockIdx.y], inouts[win], len, lsbits, + top_bits, low_bits, histograms[win]); +} + +__launch_bounds__(SORT_BLOCKDIM) +__global__ void sort_lower(vec2d_t inouts, size_t len, uint32_t win, + vec2d_t temps, vec2d_t histograms, + uint32_t wbits, uint32_t lsbits0, uint32_t lsbits1) +{ + win += blockIdx.y; + uint32_t top_bits, low_bits; + upper_bits(wbits, top_bits, low_bits); + lower_sort_loop(inouts[win], temps[blockIdx.y], histograms[win], + top_bits, low_bits); +} + # undef asm #endif #endif diff --git a/poc/msm-cuda/benches/msm.rs b/poc/msm-cuda/benches/msm.rs index 35fbfa56..a539ea51 100644 --- a/poc/msm-cuda/benches/msm.rs +++ b/poc/msm-cuda/benches/msm.rs @@ -40,10 +40,12 @@ fn criterion_benchmark(c: &mut Criterion) { group.finish(); } -#[cfg(not(any(feature = "bls12_381", feature = "bls12_377")))] +// G2 (fp2) MSM is not yet available on the ROCm/HIP backend, so only the G1 +// benchmark group is registered there. +#[cfg(any(feature = "rocm", not(any(feature = "bls12_381", feature = "bls12_377"))))] criterion_group!(benches, criterion_benchmark); -#[cfg(any(feature = "bls12_381", feature = "bls12_377"))] +#[cfg(all(not(feature = "rocm"), any(feature = "bls12_381", feature = "bls12_377")))] fn criterion_benchmark_fp2(c: &mut Criterion) { let bench_npow = std::env::var("BENCH_NPOW").unwrap_or("23".to_string()); let npoints_npow = i32::from_str(&bench_npow).unwrap(); @@ -68,7 +70,7 @@ fn criterion_benchmark_fp2(c: &mut Criterion) { group.finish(); } -#[cfg(any(feature = "bls12_381", feature = "bls12_377"))] +#[cfg(all(not(feature = "rocm"), any(feature = "bls12_381", feature = "bls12_377")))] criterion_group!(benches, criterion_benchmark, criterion_benchmark_fp2); criterion_main!(benches); diff --git a/poc/msm-cuda/build.rs b/poc/msm-cuda/build.rs index daccf927..5cd3ff64 100644 --- a/poc/msm-cuda/build.rs +++ b/poc/msm-cuda/build.rs @@ -90,24 +90,58 @@ fn main() { if cfg!(target_os = "windows") && !cfg!(target_env = "msvc") { return; } - // Detect if there is CUDA compiler and engage "cuda" feature accordingly + // Detect a CUDA (nvcc) or ROCm (hipcc) compiler and compile the GPU MSM + // accordingly. nvcc is preferred when both are present; set NVCC=off to + // force the ROCm path. The sppark build dependency auto-detects the same + // toolchain and exports DEP_SPPARK_TARGET, which sppark::build::ccmd() + // reads to return the matching cc::Build (CUDA or ROCm). + println!("cargo:rerun-if-env-changed=NVCC"); let nvcc = match env::var("NVCC") { Ok(var) => which::which(var), Err(_) => which::which("nvcc"), }; + println!("cargo:rerun-if-env-changed=HIPCC"); + let hipcc = match env::var("HIPCC") { + Ok(var) => which::which(var), + Err(_) => which::which("hipcc"), + }; + + let backend = if nvcc.is_ok() { + Some("cuda") + } else if hipcc.is_ok() { + Some("rocm") + } else { + None + }; - if nvcc.is_ok() { - let mut nvcc: cc::Build = sppark::build::ccmd(); - if cfg!(feature = "quiet") { - nvcc.flag("-diag-suppress=177"); // bug in the warning system. + if let Some(backend) = backend { + // bn254 (alt_bn128) G1 MSM is not yet supported on the ROCm/HIP backend: + // the kernel hangs the GPU for this curve (see the project notes for the + // deferred bn254-on-ROCm investigation). Refuse the combination up front + // so a user never builds a binary that wedges the device. bls12_381 and + // bls12_377 G1 MSM are fully supported on ROCm. + if backend == "rocm" && cfg!(feature = "bn254") { + panic!( + "the bn254 curve is not yet supported on the ROCm/HIP MSM backend; \ + use bls12_381 or bls12_377, or the CUDA backend for bn254" + ); + } + let mut ccmd: cc::Build = sppark::build::ccmd(); + if backend == "cuda" && cfg!(feature = "quiet") { + ccmd.flag("-diag-suppress=177"); // bug in the warning system. } - nvcc.define(curve, None); + ccmd.define(curve, None); if let Some(def) = cc_opt { - nvcc.define(def, None); + ccmd.define(def, None); + } + if backend == "rocm" { + // The MSM kernels pass CUDA's 32-bit warp mask; suppress ROCm 7's + // native 64-bit-mask *_sync builtins so the compat polyfills apply. + ccmd.define("SPPARK_DISABLE_NATIVE_WARP_SYNC", None); } - nvcc.file("cuda/pippenger_inf.cu").compile("blst_cuda_msm"); + ccmd.file("cuda/pippenger_inf.cu").compile("blst_cuda_msm"); - println!("cargo:rustc-cfg=feature=\"cuda\""); + println!("cargo:rustc-cfg=feature=\"{}\"", backend); println!("cargo:rerun-if-changed=cuda"); println!("cargo:rerun-if-env-changed=CXXFLAGS"); } diff --git a/poc/msm-cuda/cuda/pippenger_inf.cu b/poc/msm-cuda/cuda/pippenger_inf.cu index d6a2c00f..3247389a 100644 --- a/poc/msm-cuda/cuda/pippenger_inf.cu +++ b/poc/msm-cuda/cuda/pippenger_inf.cu @@ -2,14 +2,53 @@ // Licensed under the Apache License, Version 2.0, see LICENSE for details. // SPDX-License-Identifier: Apache-2.0 -#include +#if defined(__HIPCC__) +# include +// The MSM driver runs its post-kernel point accumulation on the CPU, which needs +// host-callable field arithmetic. Select the host (blst) field in the HIP host +// pass and the device mont_t field only in the device pass (mirroring CUDA's +// __CUDA_ARCH__ split). NTT, which never does host field math, keeps the device +// field in both passes (upstream default). +# define SPPARK_HIP_HOST_FIELD +#else +# include +#endif + +// The G2 (fp2 extension field) MSM relies on the Montgomery modular inverse +// (vt_inverse_mod_x) provided only by the CUDA mont_t.cuh device field; the +// ROCm/HIP mont_t.hip does not implement it yet, so G2 MSM is not built on +// HIP. G1 MSM is fully supported. See the project notes for the deferred G2 +// work. +#if !defined(__HIPCC__) +# define SPPARK_MSM_FP2 +#endif + +// bn254 (alt_bn128) G1 MSM is not yet supported on the ROCm/HIP backend: the +// kernel hangs the GPU for this curve. Refuse to build it rather than emit a +// binary that wedges the device. bls12_381 and bls12_377 G1 MSM are supported. +// See the project notes for the deferred bn254-on-ROCm investigation. +#if defined(__HIPCC__) && defined(FEATURE_BN254) +# error "bn254 G1 MSM is not yet supported on the ROCm/HIP backend" +#endif #if defined(FEATURE_BLS12_381) -# include +# if defined(SPPARK_MSM_FP2) +# include +# else +# include +# endif #elif defined(FEATURE_BLS12_377) -# include +# if defined(SPPARK_MSM_FP2) +# include +# else +# include +# endif #elif defined(FEATURE_BN254) -# include +# if defined(SPPARK_MSM_FP2) +# include +# else +# include +# endif #else # error "no FEATURE" #endif @@ -33,7 +72,8 @@ RustError::by_value mult_pippenger_inf(point_t* out, const affine_t points[], return mult_pippenger(out, points, npoints, scalars, false, ffi_affine_sz); } -#if defined(FEATURE_BLS12_381) || defined(FEATURE_BLS12_377) || defined(FEATURE_BN254) +#if defined(SPPARK_MSM_FP2) && \ + (defined(FEATURE_BLS12_381) || defined(FEATURE_BLS12_377) || defined(FEATURE_BN254)) typedef jacobian_t point_fp2_t; typedef xyzz_t bucket_fp2_t; typedef bucket_fp2_t::affine_inf_t affine_fp2_t; diff --git a/poc/msm-cuda/src/lib.rs b/poc/msm-cuda/src/lib.rs index 1c1cae22..12e4200c 100644 --- a/poc/msm-cuda/src/lib.rs +++ b/poc/msm-cuda/src/lib.rs @@ -80,7 +80,11 @@ pub fn multi_scalar_mult_arkworks( ret } -#[cfg(any(feature = "bls12_381", feature = "bls12_377", feature = "bn254"))] +// G2 (fp2) MSM is not yet available on the ROCm/HIP backend (see cuda/pippenger_inf.cu). +#[cfg(all( + not(feature = "rocm"), + any(feature = "bls12_381", feature = "bls12_377", feature = "bn254") +))] pub fn multi_scalar_mult_fp2_arkworks( points: &[G], scalars: &[::BigInt], diff --git a/poc/msm-cuda/tests/msm.rs b/poc/msm-cuda/tests/msm.rs index 766648b1..8a4fa91d 100644 --- a/poc/msm-cuda/tests/msm.rs +++ b/poc/msm-cuda/tests/msm.rs @@ -38,7 +38,37 @@ fn msm_correctness() { assert_eq!(msm_result, arkworks_result); } -#[cfg(any(feature = "bls12_381", feature = "bls12_377", feature = "bn254"))] +// Check the MSM against arkworks across a range of sizes in one run. On GPUs that +// lack cooperative launch this drives the non-cooperative sort/accumulate fallback; +// on GPUs that have it, set SPPARK_FORCE_NONCOOP=1 to exercise the same fallback and +// confirm it agrees with both the cooperative path and arkworks. +#[test] +fn msm_correctness_sizes() { + for npow in [4usize, 8, 10, 12, 14, 16] { + let npoints = 1usize << npow; + + let (points, scalars) = + util::generate_points_scalars::(npoints); + + let msm_result = multi_scalar_mult_arkworks(points.as_slice(), unsafe { + std::mem::transmute::<&[_], &[BigInteger256]>(scalars.as_slice()) + }) + .into_affine(); + + let arkworks_result = + VariableBaseMSM::multi_scalar_mul(points.as_slice(), unsafe { + std::mem::transmute::<&[_], &[BigInteger256]>(scalars.as_slice()) + }) + .into_affine(); + + assert_eq!(msm_result, arkworks_result, "npow = {}", npow); + } +} + +#[cfg(all( + not(feature = "rocm"), + any(feature = "bls12_381", feature = "bls12_377", feature = "bn254") +))] #[test] fn msm_fp2_correctness() { let test_npow = std::env::var("TEST_NPOW").unwrap_or("14".to_string()); diff --git a/util/all_gpus.cpp b/util/all_gpus.cpp index 4a6158f2..976fedc6 100644 --- a/util/all_gpus.cpp +++ b/util/all_gpus.cpp @@ -19,7 +19,7 @@ class gpus_t { for (int id = 0; id < n; id++) { cudaDeviceProp prop; if (cudaGetDeviceProperties(&prop, id) == cudaSuccess && - prop.major >= PROP_MAJOR_MIN && prop.cooperativeLaunch) { + prop.major >= PROP_MAJOR_MIN) { (void)cudaSetDevice(id); gpus.push_back(new gpu_t(gpus.size(), id, prop)); } diff --git a/util/cuda2hip.hpp b/util/cuda2hip.hpp index 36a0ca1f..ef19c6e0 100644 --- a/util/cuda2hip.hpp +++ b/util/cuda2hip.hpp @@ -13,6 +13,16 @@ # endif #endif +// ROCm 7's native __shfl_*_sync / __ballot_sync builtins require a 64-bit lane +// mask and static_assert on a 32-bit one. The MSM kernels pass CUDA's 32-bit +// 0xffffffff mask with logical-32-lane-warp semantics, so those translation +// units suppress the native builtins (-DSPPARK_DISABLE_NATIVE_WARP_SYNC) and use +// the CUDA-compatible 32-bit-mask polyfills below. NTT does not use the *_sync +// shuffles but does rely on the native __syncwarp memory ordering, so it keeps +// the native builtins. +#if defined(SPPARK_DISABLE_NATIVE_WARP_SYNC) +# define HIP_DISABLE_WARP_SYNC_BUILTINS +#endif #include #ifdef NDEBUG # define assert(e) (void)(e) @@ -121,10 +131,22 @@ cudaLaunchCooperativeKernel(const T* func, dim3 gridDim, dim3 blockDim, stream); } -#if HIP_VERSION_MAJOR < 7 +#if defined(SPPARK_DISABLE_NATIVE_WARP_SYNC) || HIP_VERSION_MAJOR < 7 static inline __device__ void __syncwarp() { __builtin_amdgcn_wave_barrier(); } #endif +/* + * Grid-wide barrier for a cooperative launch. The full ROCm + * instantiates host-side helpers that reference + * __device__-only __ockl_* builtins, so it fails to compile when pulled into a + * translation unit that also carries host code (the MSM driver). __ockl_grid_sync + * is the underlying builtin and is all the kernels need. NTT includes the real + * header directly and is unaffected; this shim is keyed off SPPARK_GRID_SYNC so + * it only appears where a source explicitly opts in. + */ +extern "C" __device__ __attribute__((convergent)) void __ockl_grid_sync(void); +#define SPPARK_GRID_SYNC() __ockl_grid_sync() + /* * To match CUDA, the 3-argument polyfills below are designed to produce * a result as if the wavefront size is 32 irregardless of its actual size.