perf: CRP-2900 Use optimized G2 generator point multiplication (#219)

This saves 123M cycles during execution in a canister, as reported by
the canbench results.

The implementation is just a port of the existing algorithm in
https://github.com/dfinity/ic/blob/master/rs/crypto/internal/crypto_lib/bls12_381/type/src/lib.rs
used for multiplication by fixed points in the replica.

Author: randombit

Cargo.lock

@@ -13,4 +13,5 @@ candid = "0.10.0"
 getrandom = { version = "0.2", features = ["custom"] }
 hex = { workspace = true }
 ic-cdk = "0.18.5"
-ic-vetkeys = "0.3.0"
+ic-vetkeys = { path = "../ic_vetkeys" }
+

backend/rs/benchmarks/Cargo.toml

@@ -32,6 +32,180 @@ lazy_static::lazy_static! {
 const G1AFFINE_BYTES: usize = 48; // Size of compressed form
 const G2AFFINE_BYTES: usize = 96; // Size of compressed form
 
+struct G2PrecomputedTable {
+    tbl: Vec<G2Affine>,
+}
+
+impl G2PrecomputedTable {
+    /// The size of the windows
+    ///
+    /// This algorithm uses just `SUBGROUP_BITS/WINDOW_BITS` additions in
+    /// the online phase, at the cost of storing a table of size
+    /// `(SUBGROUP_BITS + WINDOW_BITS - 1)/WINDOW_BITS * (1 << WINDOW_BITS - 1)`
+    ///
+    /// This constant is configurable and can take values between 1 and 7
+    /// (inclusive)
+    ///
+    /// | WINDOW_BITS | TABLE_SIZE | online additions |
+    /// | ----------- | ---------- | ---------------- |
+    /// |           1 |       255  |              255 |
+    /// |           2 |       384  |              128 |
+    /// |           3 |       595  |               85 |
+    /// |           4 |       960  |               64 |
+    /// |           5 |      1581  |               51 |
+    /// |           6 |      2709  |               43 |
+    /// |           7 |      4699  |               37 |
+    ///
+    const WINDOW_BITS: usize = 4;
+
+    /// The bit length of the BLS12-381 subgroup
+    const SUBGROUP_BITS: usize = 255;
+
+    // A bitmask of all 1s that is WINDOW_BITS long
+    const WINDOW_MASK: u8 = (1 << Self::WINDOW_BITS) - 1;
+
+    // The total number of windows in a scalar
+    const WINDOWS: usize = Self::SUBGROUP_BITS.div_ceil(Self::WINDOW_BITS);
+
+    // We must select from 2^WINDOW_BITS elements in each table
+    // group. However one element of the table group is always the
+    // identity, and so can be omitted, which is the reason for the
+    // subtraction by 1 here.
+    const WINDOW_ELEMENTS: usize = (1 << Self::WINDOW_BITS) - 1;
+
+    // The total size of the table we will use
+    const TABLE_SIZE: usize = Self::WINDOW_ELEMENTS * Self::WINDOWS;
+
+    /// Precompute a table for fast multiplication
+    fn new(pt: &G2Affine) -> Self {
+        let mut ptbl = vec![ic_bls12_381::G2Projective::identity(); Self::TABLE_SIZE];
+
+        let mut accum = ic_bls12_381::G2Projective::from(pt);
+
+        for i in 0..Self::WINDOWS {
+            let tbl_i = &mut ptbl[Self::WINDOW_ELEMENTS * i..Self::WINDOW_ELEMENTS * (i + 1)];
+
+            tbl_i[0] = accum;
+            for j in 1..Self::WINDOW_ELEMENTS {
+                // Our table indexes are off by one due to the omitted
+                // identity element. So here we are checking if we are
+                // about to compute a point that is a doubling of a point
+                // we have previously computed. If so we can compute it
+                // using a (faster) doubling rather than using addition.
+
+                tbl_i[j] = if j % 2 == 1 {
+                    tbl_i[j / 2].double()
+                } else {
+                    tbl_i[j - 1] + tbl_i[0]
+                };
+            }
+
+            // move on to the next power
+            accum = tbl_i[Self::WINDOW_ELEMENTS / 2].double();
+        }
+
+        // batch convert the table to affine form, so we can use mixed addition
+        // in the online phase.
+        let mut tbl = vec![ic_bls12_381::G2Affine::identity(); Self::TABLE_SIZE];
+        ic_bls12_381::G2Projective::batch_normalize(&ptbl, &mut tbl);
+
+        Self { tbl }
+    }
+
+    /// Perform variable-time scalar multiplication using the precomputed table plus extra addition
+    fn mul_vartime(&self, scalar: &Scalar, extra_add: Option<&G2Affine>) -> ic_bls12_381::G2Affine {
+        let s = {
+            let mut s = scalar.to_bytes();
+            s.reverse(); // zkcrypto/bls12_381 uses little-endian
+            s
+        };
+
+        let mut accum = if let Some(add) = extra_add {
+            ic_bls12_381::G2Projective::from(add)
+        } else {
+            ic_bls12_381::G2Projective::identity()
+        };
+
+        for i in 0..Self::WINDOWS {
+            let tbl_for_i = &self.tbl[Self::WINDOW_ELEMENTS * i..Self::WINDOW_ELEMENTS * (i + 1)];
+
+            let b = Self::get_window(&s, Self::WINDOW_BITS * i);
+            if b > 0 {
+                accum += tbl_for_i[b as usize - 1];
+            }
+        }
+
+        G2Affine::from(accum)
+    }
+
+    /// Perform scalar multiplication using the precomputed table
+    fn mul(&self, scalar: &Scalar) -> ic_bls12_381::G2Affine {
+        let s = {
+            let mut s = scalar.to_bytes();
+            s.reverse(); // zkcrypto/bls12_381 uses little-endian
+            s
+        };
+
+        let mut accum = ic_bls12_381::G2Projective::identity();
+
+        for i in 0..Self::WINDOWS {
+            let tbl_for_i = &self.tbl[Self::WINDOW_ELEMENTS * i..Self::WINDOW_ELEMENTS * (i + 1)];
+
+            let b = Self::get_window(&s, Self::WINDOW_BITS * i);
+            accum += Self::ct_select(tbl_for_i, b as usize);
+        }
+
+        G2Affine::from(accum)
+    }
+
+    // Extract a WINDOW_BITS sized window out of s, depending on offset.
+    #[inline(always)]
+    fn get_window(s: &[u8], offset: usize) -> u8 {
+        const BITS_IN_BYTE: usize = 8;
+
+        let shift = offset % BITS_IN_BYTE;
+        let byte_offset = s.len() - 1 - (offset / BITS_IN_BYTE);
+
+        let w0 = s[byte_offset];
+
+        let single_byte_window = shift <= (BITS_IN_BYTE - Self::WINDOW_BITS) || byte_offset == 0;
+
+        let bits = if single_byte_window {
+            // If we can get the window out of single byte, do so
+            w0 >> shift
+        } else {
+            // Otherwise we must join two bytes and extract the result
+            let w1 = s[byte_offset - 1];
+            (w0 >> shift) | (w1 << (BITS_IN_BYTE - shift))
+        };
+
+        bits & Self::WINDOW_MASK
+    }
+
+    // Constant time table lookup
+    //
+    // This version is specifically adapted to this algorithm. If
+    // index is zero, then it returns the identity element. Otherwise
+    // it returns from[index-1].
+    #[inline(always)]
+    fn ct_select(from: &[ic_bls12_381::G2Affine], index: usize) -> ic_bls12_381::G2Affine {
+        use subtle::{ConditionallySelectable, ConstantTimeEq};
+
+        let mut val = ic_bls12_381::G2Affine::identity();
+
+        let index = index.wrapping_sub(1);
+        for (idx, v) in from.iter().enumerate() {
+            val.conditional_assign(v, usize::ct_eq(&idx, &index));
+        }
+
+        val
+    }
+}
+
+lazy_static::lazy_static! {
+    static ref G2_MUL_TABLE: G2PrecomputedTable = G2PrecomputedTable::new(&G2Affine::generator());
+}
+
 /// Derive a symmetric key using HKDF-SHA256
 fn hkdf(okm: &mut [u8], input: &[u8], domain_sep: &str) {
     let hk = hkdf::Hkdf::<sha2::Sha256>::new(None, input);
@@ -206,7 +380,7 @@ impl MasterPublicKey {
 
         let offset = hash_to_scalar_two_inputs(&self.serialize(), canister_id, dst);
 
-        let derived_key = G2Affine::from(self.point + G2Affine::generator() * offset);
+        let derived_key = G2_MUL_TABLE.mul_vartime(&offset, Some(&self.point));
         DerivedPublicKey { point: derived_key }
     }
 
@@ -290,7 +464,7 @@ impl DerivedPublicKey {
 
         let offset = hash_to_scalar_two_inputs(&self.serialize(), context, dst);
 
-        let derived_key = G2Affine::from(self.point + G2Affine::generator() * offset);
+        let derived_key = G2_MUL_TABLE.mul_vartime(&offset, Some(&self.point));
         Self { point: derived_key }
     }
 
@@ -733,7 +907,7 @@ impl IbeCiphertext {
 
         let tsig = ic_bls12_381::pairing(&pt, &dpk.point) * t;
 
-        let c1 = G2Affine::from(G2Affine::generator() * t);
+        let c1 = G2_MUL_TABLE.mul(&t);
         let c2 = Self::mask_seed(seed.value(), &tsig);
         let c3 = Self::mask_msg(msg, seed.value());
 
@@ -753,15 +927,15 @@ impl IbeCiphertext {
     ///
     /// Returns the plaintext, or Err if decryption failed
     pub fn decrypt(&self, vetkey: &VetKey) -> Result<Vec<u8>, String> {
-        let t = ic_bls12_381::pairing(vetkey.point(), &self.c1);
+        let tsig = ic_bls12_381::pairing(vetkey.point(), &self.c1);
 
-        let seed = Self::mask_seed(&self.c2, &t);
+        let seed = Self::mask_seed(&self.c2, &tsig);
 
         let msg = Self::mask_msg(&self.c3, &seed);
 
         let t = Self::hash_to_mask(&self.header, &seed, &msg);
 
-        let g_t = G2Affine::from(G2Affine::generator() * t);
+        let g_t = G2_MUL_TABLE.mul(&t);
 
         if self.c1 == g_t {
             Ok(msg)