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

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.
 

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<class, class H, const H*> 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<class, class H, const H*> 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<class, class H, const H*> 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];

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

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<field_t, field_h, a4>{ 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<class affine_t>
     __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<class affine_t>
     __device__ void uadd(const affine_t& p2, bool subtract = false)
     {
@@ -542,12 +546,12 @@ class xyzz_t {
     }
 #else
     template<class affine_t>
-    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

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,17 +63,17 @@ 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<typename... Ts> constexpr fp_t(Ts... a)  : fp_mont{a...} {}
 };
 typedef mont_t<254, device::ALT_BN128_r, device::ALT_BN128_m0,
                     device::ALT_BN128_rRR, device::ALT_BN128_rone,
                     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<typename... Ts> 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 <blst_t.hpp>
 
 # if defined(__GNUC__) && !defined(__clang__)

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,17 +67,17 @@ 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<typename... Ts> constexpr fp_t(Ts... a)  : fp_mont{a...} {}
 };
 typedef mont_t<253, device::BLS12_377_r, device::BLS12_377_m0,
                     device::BLS12_377_rRR, device::BLS12_377_rone,
                     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<typename... Ts> 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 <blst_t.hpp>
 
 # if defined(__GNUC__) && !defined(__clang__)

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,17 +67,17 @@ 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<typename... Ts> constexpr fp_t(Ts... a)  : fp_mont{a...} {}
 };
 typedef mont_t<255, device::BLS12_381_r, device::BLS12_381_m0,
                     device::BLS12_381_rRR, device::BLS12_381_rone,
                     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<typename... Ts> 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 <blst_t.hpp>
 
 # if defined(__GNUC__) && !defined(__clang__)