From c4fb180b9709c0ccd2157345d20bc3b6e59199bc Mon Sep 17 00:00:00 2001
From: Jay White <JayWhite2357@users.noreply.github.com>
Date: Thu, 5 Sep 2024 10:15:16 -0400
Subject: [PATCH] perf: better parallelize multi-pairings (#129)

# Rationale for this change

Computing pairings is a significant portion of proof generation time.
This PR significantly reduces the cost of multi-pairings.

# What changes are included in this PR?

* Multi-pairings are parallelized naively rather than using the arkworks
version, which is surprisingly inefficient. It may be worth time
investigating why exactly this is more efficient.

# Are these changes tested?

Yes
---
 .../src/proof_primitive/dory/pairings.rs      | 69 +++++++++++--------
 1 file changed, 39 insertions(+), 30 deletions(-)

diff --git a/crates/proof-of-sql/src/proof_primitive/dory/pairings.rs b/crates/proof-of-sql/src/proof_primitive/dory/pairings.rs
index 2438c57f7..942fd7bef 100644
--- a/crates/proof-of-sql/src/proof_primitive/dory/pairings.rs
+++ b/crates/proof-of-sql/src/proof_primitive/dory/pairings.rs
@@ -1,4 +1,5 @@
-use ark_ec::pairing::{Pairing, PairingOutput};
+use ark_ec::pairing::{MillerLoopOutput, Pairing, PairingOutput};
+use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 #[tracing::instrument(level = "debug", skip_all)]
 // This is a wrapper around multi_pairing_impl simply because tracing doesn't work well with threading.
 pub fn pairing<P: Pairing>(
@@ -10,8 +11,8 @@ pub fn pairing<P: Pairing>(
 #[tracing::instrument(level = "debug", skip_all)]
 // This is a wrapper around multi_pairing_impl simply because tracing doesn't work well with threading.
 pub fn multi_pairing<P: Pairing>(
-    a: impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-    b: impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+    a: impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+    b: impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
 ) -> PairingOutput<P> {
     multi_pairing_impl(a, b)
 }
@@ -19,12 +20,12 @@ pub fn multi_pairing<P: Pairing>(
 // This is a wrapper around multi_pairing_2_impl simply because tracing doesn't work well with threading.
 pub fn multi_pairing_2<P: Pairing>(
     (a0, b0): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a1, b1): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
 ) -> (PairingOutput<P>, PairingOutput<P>) {
     multi_pairing_2_impl((a0, b0), (a1, b1))
@@ -33,20 +34,20 @@ pub fn multi_pairing_2<P: Pairing>(
 // This is a wrapper around multi_pairing_4_impl simply because tracing doesn't work well with threading.
 pub fn multi_pairing_4<P: Pairing>(
     (a0, b0): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a1, b1): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a2, b2): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a3, b3): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
 ) -> (
     PairingOutput<P>,
@@ -57,39 +58,47 @@ pub fn multi_pairing_4<P: Pairing>(
     multi_pairing_4_impl((a0, b0), (a1, b1), (a2, b2), (a3, b3))
 }
 fn multi_pairing_impl<P: Pairing>(
-    a: impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-    b: impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+    a: impl IntoIterator<Item = impl Into<P::G1Prepared> + Send>,
+    b: impl IntoIterator<Item = impl Into<P::G2Prepared> + Send>,
 ) -> PairingOutput<P> {
-    Pairing::multi_pairing(a, b)
+    let a: Vec<_> = a.into_iter().collect();
+    let b: Vec<_> = b.into_iter().collect();
+    Pairing::final_exponentiation(MillerLoopOutput(
+        a.into_par_iter()
+            .zip(b)
+            .map(|(x, y)| P::miller_loop(x, y).0)
+            .product(),
+    ))
+    .unwrap()
 }
 fn multi_pairing_2_impl<P: Pairing>(
     (a0, b0): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a1, b1): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
 ) -> (PairingOutput<P>, PairingOutput<P>) {
     rayon::join(|| multi_pairing_impl(a0, b0), || multi_pairing_impl(a1, b1))
 }
 fn multi_pairing_4_impl<P: Pairing>(
     (a0, b0): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a1, b1): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a2, b2): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
     (a3, b3): (
-        impl IntoIterator<Item = impl Into<P::G1Prepared>> + Send,
-        impl IntoIterator<Item = impl Into<P::G2Prepared>> + Send,
+        impl IntoIterator<Item = impl Into<P::G1Prepared> + Send> + Send,
+        impl IntoIterator<Item = impl Into<P::G2Prepared> + Send> + Send,
     ),
 ) -> (
     PairingOutput<P>,