accelerate to_affine conversion

Author: austinabell

p256/src/arithmetic/field.rs

@@ -63,12 +63,10 @@ primeorder::impl_mont_field_element!(
 
 impl FieldElement {
     #[cfg(all(target_os = "zkvm", target_arch = "riscv32"))]
-    #[inline(never)]
     pub(crate) fn from_words_le(fe: [u32; 8]) -> CtOption<Self> {
-        // use elliptic_curve::bigint::Encoding;
-        // println!("r2: {:0X?}", fe_from_montgomery(R_2.as_words()));
-
         let fe = FieldElement256::new_unchecked(fe);
+
+        // Convert to montgomery form with aR mod p
         let mut mont = FieldElement256::default();
         fe.mul_unchecked(&R_2_LE, &mut mont);
 
@@ -81,16 +79,24 @@ impl FieldElement {
         CtOption::new(Self(uint), is_within_modulus)
     }
 
+    #[cfg(all(target_os = "zkvm", target_arch = "riscv32"))]
+    pub(crate) fn to_words_le(&self) -> [u32; 8] {
+        use crate::elliptic_curve::bigint::Encoding;
+        // NOTE: this from mont conversion could be accelerated, but it's very little cycles.
+        let canonical = self.to_canonical();
+        let input = canonical.to_le_bytes();
+        let array = bytemuck::cast::<_, [u32; 8]>(input);
+
+        array
+    }
+
     /// Returns the multiplicative inverse of self, if self is non-zero.
-    #[inline(never)]
     pub fn invert(&self) -> CtOption<Self> {
         #[cfg(all(target_os = "zkvm", target_arch = "riscv32"))]
         {
             use crate::elliptic_curve::bigint::Encoding;
 
-            let canonical = self.to_canonical();
-            let input = canonical.to_le_bytes();
-            let input_words = bytemuck::cast::<_, [u32; 8]>(input);
+            let input_words = self.to_words_le();
             let mut output = [0u32; 8];
             risc0_bigint2::field::modinv_256_unchecked(
                 &input_words,

p256/src/arithmetic/field/field32.rs

@@ -9,7 +9,6 @@ use crate::arithmetic::util::{adc, mac, sbb};
 pub type Fe = [u32; 8];
 
 /// Translate a field element out of the Montgomery domain.
-#[inline]
 pub const fn fe_from_montgomery(w: &Fe) -> Fe {
     let w = fe32_to_fe64(w);
     montgomery_reduce(&[w[0], w[1], w[2], w[3], 0, 0, 0, 0])

primeorder/src/field.rs

@@ -456,7 +456,6 @@ macro_rules! impl_field_op {
         impl ::core::ops::$op for $fe {
             type Output = $fe;
 
-            #[inline]
             fn $op_fn(self, rhs: $fe) -> $fe {
                 $fe($func(self.as_ref(), rhs.as_ref()).into())
             }
@@ -465,7 +464,6 @@ macro_rules! impl_field_op {
         impl ::core::ops::$op<&$fe> for $fe {
             type Output = $fe;
 
-            #[inline]
             fn $op_fn(self, rhs: &$fe) -> $fe {
                 $fe($func(self.as_ref(), rhs.as_ref()).into())
             }
@@ -474,7 +472,6 @@ macro_rules! impl_field_op {
         impl ::core::ops::$op<&$fe> for &$fe {
             type Output = $fe;
 
-            #[inline]
             fn $op_fn(self, rhs: &$fe) -> $fe {
                 $fe($func(self.as_ref(), rhs.as_ref()).into())
             }

primeorder/src/projective.rs

@@ -57,6 +57,37 @@ where
 
     /// Returns the affine representation of this point, or `None` if it is the identity.
     pub fn to_affine(&self) -> AffinePoint<C> {
+        #[cfg(all(target_os = "zkvm", target_arch = "riscv32"))]
+        {
+            use crate::risc0::felt_to_u32_words_le;
+            let z = felt_to_u32_words_le::<C>(&self.z);
+            let mut z_inv = [0u32; 8];
+            risc0_bigint2::field::modinv_256_unchecked(&z, &C::PRIME_LE_WORDS, &mut z_inv);
+
+            let mut buffer = [0u32; 8];
+            let x_buffer = felt_to_u32_words_le::<C>(&self.x);
+            let y_buffer = felt_to_u32_words_le::<C>(&self.y);
+            risc0_bigint2::field::modmul_256_unchecked(
+                &x_buffer,
+                &z_inv,
+                &C::PRIME_LE_WORDS,
+                &mut buffer,
+            );
+
+            let x = C::from_u32_words_le(buffer);
+
+            risc0_bigint2::field::modmul_256_unchecked(
+                &y_buffer,
+                &z_inv,
+                &C::PRIME_LE_WORDS,
+                &mut buffer,
+            );
+            let y = C::from_u32_words_le(buffer);
+            return x
+                .and_then(|x| y.map(|y| AffinePoint { x, y, infinity: 0 }))
+                .unwrap_or(AffinePoint::IDENTITY);
+        }
+
         self.z
             .invert()
             .map(|zinv| AffinePoint {

primeorder/src/risc0.rs

@@ -175,27 +175,26 @@ where
     }
 }
 
-// TODO remove inline
-#[inline(never)]
+pub(crate) fn felt_to_u32_words_le<C>(data: &C::FieldElement) -> [u32; 8]
+where
+    C: PrimeCurveParams,
+{
+    // TODO alignment
+    let mut data_bytes: [u8; 32] = data.to_repr().as_slice().try_into().unwrap();
+    data_bytes.reverse();
+    bytemuck::cast::<_, [u32; 8]>(data_bytes)
+}
+
 pub(crate) fn projective_to_affine<C>(p: &ProjectivePoint<C>) -> ec::AffinePoint<8, C>
 where
     C: PrimeCurveParams,
 {
     let aff = p.to_affine();
-    let x_bytes = aff.x.to_repr();
-    let y_bytes = aff.y.to_repr();
-    let mut x_bytes_arr: [u8; 32] = x_bytes.as_slice().try_into().unwrap();
-    let mut y_bytes_arr: [u8; 32] = y_bytes.as_slice().try_into().unwrap();
-    x_bytes_arr.reverse();
-    y_bytes_arr.reverse();
-    // TODO make more alignment safe
-    let x = bytemuck::cast::<_, [u32; 8]>(x_bytes_arr);
-    let y = bytemuck::cast::<_, [u32; 8]>(y_bytes_arr);
+    let x = felt_to_u32_words_le::<C>(&aff.x);
+    let y = felt_to_u32_words_le::<C>(&aff.y);
     ec::AffinePoint::new_unchecked(x, y)
 }
 
-// TODO remove inline
-#[inline(never)]
 pub(crate) fn affine_to_projective<C>(affine: &ec::AffinePoint<8, C>) -> ProjectivePoint<C>
 where
     C: PrimeCurveParams,
@@ -216,11 +215,11 @@ where
     }
 }
 
-#[inline(never)]
 pub(crate) fn scalar_to_words<C>(s: &Scalar<C>) -> [u32; 8]
 where
     C: PrimeCurveParams,
 {
+    // TODO alignment
     let mut bytes: [u8; 32] = s.to_repr().as_slice().try_into().unwrap();
     // U256 is big endian, need to flip to little endian.
     bytes.reverse();