Montgomery Multiplication for refactored metal EC backend (#23)

* feat: import mont_mul backend from https://github.com/geometryxyz/msl-secp256k1

* feat: conversion between bigint and arbitrary limb size

* test(bigint): add host test

* test(bigint): adapt from https://github.com/geometryxyz/msl-secp256k1

* refactor: add overflow detection and correct suitable bigint val for each cases

* chore: update path for contants.metal

* test(mont_mul): adapted mont mul tests from https://github.com/geometryxyz/msl-secp256k1

* feat(mont_mul): adapted utils function related to mont_mul from https://github.com/geometers/multiprecision

* feat: add limbs conversion from ark_ff bigint to arbitrary limbs in Vec<u32>

* chore: update path for utils for limb conversion

* fix(mont_mul): correct the conversion from arkworks' scalarfield to arbitrary limbs of Vec<u32>

* test(mont_mul): adapted benchmark functions from https://github.com/geometryxyz/msl-secp256k1

* feat(metal): add mont_mul cios

* test(mont_mul): add cios mont_mul test and benchmark

Author: moven0831

Cargo.lock

@@ -60,6 +60,7 @@ witness = { git = "https://github.com/philsippl/circom-witness-rs.git" }
 
 [dev-dependencies]
 serial_test = "3.0.0"
+stopwatch = "0.0.7"
 
 # [dependencies.rayon]
 # version = "1"

mopro-msm/Cargo.toml

@@ -12,7 +12,6 @@ pub trait ToLimbs {
 pub trait FromLimbs {
     fn from_u32_limbs(limbs: &[u32]) -> Self;
     fn from_limbs(limbs: &[u32], log_limb_size: u32) -> Self;
-    fn from_limbs(limbs: &[u32], log_limb_size: u32) -> Self;
     fn from_u128(num: u128) -> Self;
     fn from_u32(num: u32) -> Self;
 }
@@ -78,10 +77,6 @@ impl ToLimbs for Fq {
     fn to_limbs(&self, num_limbs: usize, log_limb_size: u32) -> Vec<u32> {
         self.0.to_limbs(num_limbs, log_limb_size)
     }
-
-    fn to_limbs(&self, num_limbs: usize, log_limb_size: u32) -> Vec<u32> {
-        self.0.to_limbs(num_limbs, log_limb_size)
-    }
 }
 
 impl FromLimbs for BigInteger256 {
@@ -134,35 +129,6 @@ impl FromLimbs for BigInteger256 {
 
         BigInteger256::new(result)
     }
-
-    fn from_limbs(limbs: &[u32], log_limb_size: u32) -> Self {
-        let mut result = [0u64; 4];
-        let limb_size = log_limb_size as usize;
-        let mut accumulated_bits = 0;
-        let mut current_u64 = 0u64;
-        let mut result_idx = 0;
-
-        for &limb in limbs {
-            // Add the current limb at the appropriate position
-            current_u64 |= (limb as u64) << accumulated_bits;
-            accumulated_bits += limb_size;
-
-            // If we've accumulated 64 bits or more, store the result
-            while accumulated_bits >= 64 && result_idx < 4 {
-                result[result_idx] = current_u64;
-                current_u64 = limb as u64 >> (limb_size - (accumulated_bits - 64));
-                accumulated_bits -= 64;
-                result_idx += 1;
-            }
-        }
-
-        // Handle any remaining bits
-        if accumulated_bits > 0 && result_idx < 4 {
-            result[result_idx] = current_u64;
-        }
-
-        BigInteger256::new(result)
-    }
 }
 
 impl FromLimbs for Fq {
@@ -193,9 +159,4 @@ impl FromLimbs for Fq {
         let bigint = BigInteger256::from_limbs(limbs, log_limb_size);
         Fq::new(mont_reduction::raw_reduction(bigint))
     }
-
-    fn from_limbs(limbs: &[u32], log_limb_size: u32) -> Self {
-        let bigint = BigInteger256::from_limbs(limbs, log_limb_size);
-        Fq::new(mont_reduction::raw_reduction(bigint))
-    }
 }

mopro-msm/src/msm/metal/abstraction/limbs_conversion.rs

@@ -1,4 +1,4 @@
-// source: https://github.com/geometryxyz/msl-secp256k1
+// adapted from: https://github.com/geometryxyz/msl-secp256k1
 
 using namespace metal;
 #include <metal_stdlib>
@@ -92,3 +92,76 @@ BigInt mont_mul_modified(
 
     return conditional_reduce(s, p);
 }
+
+/// The CIOS method for Montgomery multiplication from Tolga Acar's thesis:
+/// High-Speed Algorithms & Architectures For Number-Theoretic Cryptosystems
+/// https://www.proquest.com/openview/1018972f191afe55443658b28041c118/1
+BigInt mont_mul_cios(
+    BigInt x,
+    BigInt y,
+    BigInt p
+) {
+    uint t[NUM_LIMBS + 2] = {0};  // Extra space for carries
+    BigInt result;
+    
+    for (uint i = 0; i < NUM_LIMBS; i++) {
+        // Step 1: Multiply and add
+        uint c = 0;
+        for (uint j = 0; j < NUM_LIMBS; j++) {
+            uint r = t[j] + x.limbs[j] * y.limbs[i] + c;
+            c = r >> LOG_LIMB_SIZE;
+            t[j] = r & MASK;
+        }
+        uint r = t[NUM_LIMBS] + c;
+        t[NUM_LIMBS + 1] = r >> LOG_LIMB_SIZE;
+        t[NUM_LIMBS] = r & MASK;
+
+        // Step 2: Reduce
+        uint m = (t[0] * N0) & MASK;
+        r = t[0] + m * p.limbs[0];
+        c = r >> LOG_LIMB_SIZE;
+
+        for (uint j = 1; j < NUM_LIMBS; j++) {
+            r = t[j] + m * p.limbs[j] + c;
+            c = r >> LOG_LIMB_SIZE;
+            t[j - 1] = r & MASK;
+        }
+
+        r = t[NUM_LIMBS] + c;
+        c = r >> LOG_LIMB_SIZE;
+        t[NUM_LIMBS - 1] = r & MASK;
+        t[NUM_LIMBS] = t[NUM_LIMBS + 1] + c;
+    }
+
+    // Final reduction check
+    bool t_lt_p = false;
+    for (uint idx = 0; idx < NUM_LIMBS; idx++) {
+        uint i = NUM_LIMBS - 1 - idx;
+        if (t[i] < p.limbs[i]) {
+            t_lt_p = true;
+            break;
+        } else if (t[i] > p.limbs[i]) {
+            break;
+        }
+    }
+
+    if (t_lt_p) {
+        for (uint i = 0; i < NUM_LIMBS; i++) {
+            result.limbs[i] = t[i];
+        }
+    } else {
+        uint borrow = 0;
+        for (uint i = 0; i < NUM_LIMBS; i++) {
+            uint diff = t[i] - p.limbs[i] - borrow;
+            if (t[i] < (p.limbs[i] + borrow)) {
+                diff += (1 << LOG_LIMB_SIZE);
+                borrow = 1;
+            } else {
+                borrow = 0;
+            }
+            result.limbs[i] = diff;
+        }
+    }
+
+    return result;
+}

mopro-msm/src/msm/metal_msm/shader/mont_backend/mont.metal

@@ -0,0 +1,25 @@
+// source: https://github.com/geometryxyz/msl-secp256k1
+
+using namespace metal;
+#include <metal_stdlib>
+#include <metal_math>
+#include "mont.metal"
+
+kernel void run(
+    device BigInt* lhs [[ buffer(0) ]],
+    device BigInt* rhs [[ buffer(1) ]],
+    device BigInt* prime [[ buffer(2) ]],
+    device BigInt* result [[ buffer(3) ]],
+    uint gid [[ thread_position_in_grid ]]
+) {
+    BigInt a;
+    BigInt b;
+    BigInt p;
+    a.limbs = lhs->limbs;
+    b.limbs = rhs->limbs;
+    p.limbs = prime->limbs;
+
+    BigInt res = mont_mul_cios(a, b, p);
+    result->limbs = res.limbs;
+
+}

mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_cios.metal

@@ -0,0 +1,28 @@
+using namespace metal;
+#include <metal_stdlib>
+#include <metal_math>
+#include "mont.metal"
+
+kernel void run(
+    device BigInt* lhs [[ buffer(0) ]],
+    device BigInt* rhs [[ buffer(1) ]],
+    device BigInt* prime [[ buffer(2) ]],
+    device array<uint, 1>* cost [[ buffer(3) ]],
+    device BigInt* result [[ buffer(4) ]],
+    uint gid [[ thread_position_in_grid ]]
+) {
+    BigInt a;
+    BigInt b;
+    BigInt p;
+    a.limbs = lhs->limbs;
+    b.limbs = rhs->limbs;
+    p.limbs = prime->limbs;
+    array<uint, 1> cost_arr = *cost;
+
+    BigInt c = mont_mul_cios(a, a, p);
+    for (uint i = 1; i < cost_arr[0]; i ++) {
+        c = mont_mul_cios(c, a, p);
+    }
+    BigInt res = mont_mul_cios(c, b, p);
+    result->limbs = res.limbs;
+}

mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_cios_benchmarks.metal

@@ -0,0 +1,28 @@
+using namespace metal;
+#include <metal_stdlib>
+#include <metal_math>
+#include "mont.metal"
+
+kernel void run(
+    device BigInt* lhs [[ buffer(0) ]],
+    device BigInt* rhs [[ buffer(1) ]],
+    device BigInt* prime [[ buffer(2) ]],
+    device array<uint, 1>* cost [[ buffer(3) ]],
+    device BigInt* result [[ buffer(4) ]],
+    uint gid [[ thread_position_in_grid ]]
+) {
+    BigInt a;
+    BigInt b;
+    BigInt p;
+    a.limbs = lhs->limbs;
+    b.limbs = rhs->limbs;
+    p.limbs = prime->limbs;
+    array<uint, 1> cost_arr = *cost;
+
+    BigInt c = mont_mul_modified(a, a, p);
+    for (uint i = 1; i < cost_arr[0]; i ++) {
+        c = mont_mul_modified(c, a, p);
+    }
+    BigInt res = mont_mul_modified(c, b, p);
+    result->limbs = res.limbs;
+}

mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_modified_benchmarks.metal

@@ -0,0 +1,28 @@
+using namespace metal;
+#include <metal_stdlib>
+#include <metal_math>
+#include "mont.metal"
+
+kernel void run(
+    device BigInt* lhs [[ buffer(0) ]],
+    device BigInt* rhs [[ buffer(1) ]],
+    device BigInt* prime [[ buffer(2) ]],
+    device array<uint, 1>* cost [[ buffer(3) ]],
+    device BigInt* result [[ buffer(4) ]],
+    uint gid [[ thread_position_in_grid ]]
+) {
+    BigInt a;
+    BigInt b;
+    BigInt p;
+    a.limbs = lhs->limbs;
+    b.limbs = rhs->limbs;
+    p.limbs = prime->limbs;
+    array<uint, 1> cost_arr = *cost;
+
+    BigInt c = mont_mul_optimised(a, a, p);
+    for (uint i = 1; i < cost_arr[0]; i ++) {
+        c = mont_mul_optimised(c, a, p);
+    }
+    BigInt res = mont_mul_optimised(c, b, p);
+    result->limbs = res.limbs;
+}

mopro-msm/src/msm/metal_msm/shader/mont_backend/mont_mul_optimised_benchmarks.metal

@@ -1,10 +1,10 @@
 // adapted from: https://github.com/geometryxyz/msl-secp256k1
 
-use crate::msm::metal::abstraction::limbs_conversion::{FromLimbs, ToLimbs};
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
 use crate::msm::metal_msm::host::shader::{compile_metal, write_constants};
+use crate::msm::metal_msm::utils::limbs_conversion::{FromLimbs, ToLimbs};
 use ark_ff::{BigInt, BigInteger};
 use metal::*;
 

mopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_unsafe.rs

@@ -1,10 +1,10 @@
 // adapted from: https://github.com/geometryxyz/msl-secp256k1
 
-use crate::msm::metal::abstraction::limbs_conversion::{FromLimbs, ToLimbs};
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
 use crate::msm::metal_msm::host::shader::{compile_metal, write_constants};
+use crate::msm::metal_msm::utils::limbs_conversion::{FromLimbs, ToLimbs};
 use ark_ff::{BigInt, BigInteger};
 use metal::*;
 

mopro-msm/src/msm/metal_msm/tests/bigint/bigint_add_wide.rs

@@ -2,11 +2,11 @@
 
 use core::borrow;
 
-use crate::msm::metal::abstraction::limbs_conversion::{FromLimbs, ToLimbs};
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
 use crate::msm::metal_msm::host::shader::{compile_metal, write_constants};
+use crate::msm::metal_msm::utils::limbs_conversion::{FromLimbs, ToLimbs};
 use ark_ff::{BigInt, BigInteger};
 use metal::*;
 

mopro-msm/src/msm/metal_msm/tests/bigint/bigint_sub.rs

@@ -1,10 +1,10 @@
 // adapted from: https://github.com/geometryxyz/msl-secp256k1
 
-use crate::msm::metal::abstraction::limbs_conversion::{FromLimbs, ToLimbs};
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
 use crate::msm::metal_msm::host::shader::{compile_metal, write_constants};
+use crate::msm::metal_msm::utils::limbs_conversion::{FromLimbs, ToLimbs};
 use ark_bn254::Fr as ScalarField;
 use ark_ff::{BigInt, BigInteger, PrimeField};
 use metal::*;

mopro-msm/src/msm/metal_msm/tests/field/ff_add.rs

@@ -1,10 +1,10 @@
 // adapted from: https://github.com/geometryxyz/msl-secp256k1
 
-use crate::msm::metal::abstraction::limbs_conversion::{FromLimbs, ToLimbs};
 use crate::msm::metal_msm::host::gpu::{
     create_buffer, create_empty_buffer, get_default_device, read_buffer,
 };
 use crate::msm::metal_msm::host::shader::{compile_metal, write_constants};
+use crate::msm::metal_msm::utils::limbs_conversion::{FromLimbs, ToLimbs};
 use ark_bn254::Fr as ScalarField;
 use ark_ff::{BigInt, BigInteger, PrimeField};
 use metal::*;

mopro-msm/src/msm/metal_msm/tests/field/ff_sub.rs

@@ -2,3 +2,5 @@
 pub mod bigint;
 #[cfg(test)]
 pub mod field;
+#[cfg(test)]
+pub mod mont_backend;

mopro-msm/src/msm/metal_msm/tests/mod.rs

@@ -0,0 +1,8 @@
+#[cfg(test)]
+pub mod mont_benchmarks;
+#[cfg(test)]
+pub mod mont_mul_cios;
+#[cfg(test)]
+pub mod mont_mul_modified;
+#[cfg(test)]
+pub mod mont_mul_optimised;