impl trait GpuEngine preparing for gpu accelerated FFT and multiexp

Cargo.toml

@@ -26,11 +26,13 @@ group = "0.11"
 static_assertions = "1.1.0"
 rand = "0.8"
 rand_core = { version = "0.6", default-features = false }
+ec-gpu = { version = "0.1.0", optional = true }
 
 [features]
 default = []
 asm = []
 prefetch = []
+gpu = ["ec-gpu"]
 
 [profile.bench]
 opt-level = 3

src/bn256/engine.rs

@@ -649,6 +649,20 @@ impl MultiMillerLoop for Bn256 {
     }
 }
 
+#[cfg(feature = "gpu")]
+impl ec_gpu::GpuEngine for Bn256 {
+    type Scalar = Fr;
+    type Fp = Fq;
+}
+
+#[cfg(feature = "gpu")]
+pub fn u64_to_u32(limbs: &[u64]) -> Vec<u32> {
+    limbs
+        .iter()
+        .flat_map(|limb| vec![(limb & 0xFFFF_FFFF) as u32, (limb >> 32) as u32])
+        .collect()
+}
+
 #[cfg(test)]
 use rand::SeedableRng;
 #[cfg(test)]

src/bn256/fq.rs

@@ -2,6 +2,8 @@
 use super::assembly::assembly_field;
 use super::common::common_field;
 use super::LegendreSymbol;
+#[cfg(feature = "gpu")]
+use super::engine::u64_to_u32;
 use crate::arithmetic::{adc, mac, sbb, BaseExt, FieldExt, Group};
 use core::convert::TryInto;
 use core::fmt;
@@ -274,6 +276,21 @@ impl ff::PrimeField for Fq {
     }
 }
 
+#[cfg(feature = "gpu")]
+impl ec_gpu::GpuField for Fq {
+    fn one() -> Vec<u32> {
+        u64_to_u32(&R.0[..])
+    }
+
+    fn r2() -> Vec<u32> {
+        u64_to_u32(&R2.0[..])
+    }
+
+    fn modulus() -> Vec<u32> {
+        u64_to_u32(&MODULUS.0[..])
+    }
+}
+
 #[cfg(test)]
 use ff::Field;
 #[cfg(test)]

src/bn256/fr.rs

@@ -2,6 +2,8 @@
 use super::assembly::assembly_field;
 use super::common::common_field;
 use super::LegendreSymbol;
+#[cfg(feature = "gpu")]
+use super::engine::u64_to_u32;
 use crate::arithmetic::{adc, mac, sbb, BaseExt, FieldExt, Group};
 use core::convert::TryInto;
 use core::fmt;
@@ -236,6 +238,21 @@ impl ff::PrimeField for Fr {
     }
 }
 
+#[cfg(feature = "gpu")]
+impl ec_gpu::GpuField for Fr {
+    fn one() -> Vec<u32> {
+        u64_to_u32(&R.0[..])
+    }
+
+    fn r2() -> Vec<u32> {
+        u64_to_u32(&R2.0[..])
+    }
+
+    fn modulus() -> Vec<u32> {
+        u64_to_u32(&MODULUS.0[..])
+    }
+}
+
 #[cfg(test)]
 use ff::Field;