From 896593c40f1884ac2f2b7f976cc508f9242769ed Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Wed, 30 Nov 2022 16:11:11 -0800 Subject: [PATCH 1/7] Add derived impl of squaring --- ff-macros/src/montgomery/mod.rs | 10 +++- ff-macros/src/montgomery/square.rs | 95 ++++++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 1 deletion(-) create mode 100644 ff-macros/src/montgomery/square.rs diff --git a/ff-macros/src/montgomery/mod.rs b/ff-macros/src/montgomery/mod.rs index 6d95f8caf..e268db700 100644 --- a/ff-macros/src/montgomery/mod.rs +++ b/ff-macros/src/montgomery/mod.rs @@ -2,6 +2,7 @@ use std::str::FromStr; use num_bigint::BigUint; use num_traits::One; +use quote::format_ident; mod biginteger; use biginteger::*; @@ -9,8 +10,10 @@ use biginteger::*; mod mul; use mul::*; +mod square; +use square::*; + mod sum_of_products; -use quote::format_ident; use sum_of_products::*; use crate::utils; @@ -68,6 +71,7 @@ pub fn mont_config_helper( let sub_with_borrow = sub_with_borrow_impl(limbs); let subtract_modulus = subtract_modulus_impl(&modulus); let mul_assign = mul_assign_impl(can_use_no_carry_mul_opt, limbs, &modulus_limbs); + let square_in_place = square_in_place_impl(can_use_no_carry_mul_opt, limbs, &modulus_limbs); let sum_of_products = sum_of_products_impl(limbs, &modulus_limbs); let mixed_radix = if let Some(large_subgroup_generator) = large_subgroup_generator { @@ -135,6 +139,10 @@ pub fn mont_config_helper( fn mul_assign(a: &mut F, b: &F) { #mul_assign } + #[inline(always)] + fn square_in_place(a: &mut F) { + #square_in_place + } fn sum_of_products( a: &[F; M], diff --git a/ff-macros/src/montgomery/square.rs b/ff-macros/src/montgomery/square.rs new file mode 100644 index 000000000..25510272d --- /dev/null +++ b/ff-macros/src/montgomery/square.rs @@ -0,0 +1,95 @@ +use quote::quote; + +pub(super) fn square_in_place_impl( + can_use_no_carry_mul_opt: bool, + num_limbs: usize, + modulus_limbs: &[u64], +) -> proc_macro2::TokenStream { + let mut body = proc_macro2::TokenStream::new(); + let mut default = proc_macro2::TokenStream::new(); + + let modulus_0 = modulus_limbs[0]; + let double_num_limbs = 2 * num_limbs; + default.extend(quote! { + let mut r = [0u64; #double_num_limbs]; + let mut carry = 0; + }); + for i in 0..(num_limbs - 1) { + for j in (i + 1)..num_limbs { + default.extend(quote! { + r[#i + #j] = fa::mac_with_carry(r[#i + #j], (a.0).0[#i], (a.0).0[#j], &mut carry); + }) + } + default.extend(quote! { + r[#num_limbs + #i] = carry; + carry = 0; + }); + } + default.extend(quote! { r[#double_num_limbs - 1] = r[#double_num_limbs - 2] >> 63; }); + for i in 2..(double_num_limbs - 1) { + default.extend(quote! { + r[#double_num_limbs - #i] = + (r[#double_num_limbs - #i] << 1) | (r[#double_num_limbs - (#i + 1)] >> 63); + }); + } + default.extend(quote! { r[1] <<= 1; }); + for i in 0..num_limbs { + default.extend(quote! { + r[2 * #i] = fa::mac_with_carry(r[2 * #i], (a.0).0[#i], (a.0).0[#i], &mut carry); + carry = fa::adc(&mut r[2 * #i + 1], 0, carry); + }); + } + // Montgomery reduction + default.extend(quote! { let mut carry2 = 0; }); + for i in 0..num_limbs { + default.extend(quote! { + let k = r[#i].wrapping_mul(Self::INV); + let mut carry = 0; + fa::mac_discard(r[#i], k, #modulus_0, &mut carry); + }); + for j in 1..num_limbs { + let idx = j + i; + let modulus_j = modulus_limbs[j]; + default.extend(quote! { + r[#idx] = fa::mac_with_carry(r[#idx], k, #modulus_j, &mut carry); + }); + } + default.extend(quote! { carry2 = fa::adc(&mut r[#num_limbs + #i], carry, carry2); }); + } + default.extend(quote! { (a.0).0.copy_from_slice(&r[#num_limbs..]); }); + + if num_limbs == 1 { + // We default to multiplying with `a` using the `Mul` impl + // for the N == 1 case + return quote!({ + *a *= *a; + }); + } else if (2..=6).contains(&num_limbs) && can_use_no_carry_mul_opt { + body.extend(quote!({ + if cfg!(all( + feature = "asm", + target_feature = "bmi2", + target_feature = "adx", + target_arch = "x86_64" + )) { + #[cfg( + all( + feature = "asm", + target_feature = "bmi2", + target_feature = "adx", + target_arch = "x86_64" + ) + )] + #[allow(unsafe_code, unused_mut)] + ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); + } else { + #default + } + })) + } else { + body.extend(quote!({ #default })) + } + + body.extend(quote!(__subtract_modulus(a);)); + body +} From 4ac5b6d8800d50e9f0821e70a4f4939fff8f690d Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Thu, 8 Dec 2022 18:08:54 -0800 Subject: [PATCH 2/7] Update ff-macros/src/montgomery/mod.rs --- ff-macros/src/montgomery/mod.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/ff-macros/src/montgomery/mod.rs b/ff-macros/src/montgomery/mod.rs index d752a400c..80694094c 100644 --- a/ff-macros/src/montgomery/mod.rs +++ b/ff-macros/src/montgomery/mod.rs @@ -3,7 +3,6 @@ use std::str::FromStr; use num_bigint::BigUint; use num_traits::One; -use quote::format_ident; mod biginteger; use biginteger::*; From 352567204bea37ff2ff07043a58fcd95f3c2445d Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Fri, 9 Dec 2022 15:57:24 -0800 Subject: [PATCH 3/7] Fix macro to use `subtract_modulus_with_carry` --- ff-macros/src/montgomery/mod.rs | 7 +++- ff-macros/src/montgomery/square.rs | 16 ++++++--- test-templates/src/fields.rs | 52 +++++++++++++++--------------- 3 files changed, 44 insertions(+), 31 deletions(-) diff --git a/ff-macros/src/montgomery/mod.rs b/ff-macros/src/montgomery/mod.rs index 80694094c..a086399c1 100644 --- a/ff-macros/src/montgomery/mod.rs +++ b/ff-macros/src/montgomery/mod.rs @@ -80,7 +80,12 @@ pub fn mont_config_helper( &modulus_limbs, modulus_has_spare_bit, ); - let square_in_place = square_in_place_impl(can_use_no_carry_mul_opt, limbs, &modulus_limbs); + let square_in_place = square_in_place_impl( + can_use_no_carry_mul_opt, + limbs, + &modulus_limbs, + modulus_has_spare_bit, + ); let sum_of_products = sum_of_products_impl(limbs, &modulus_limbs); let mixed_radix = if let Some(large_subgroup_generator) = large_subgroup_generator { diff --git a/ff-macros/src/montgomery/square.rs b/ff-macros/src/montgomery/square.rs index 25510272d..fa5978a98 100644 --- a/ff-macros/src/montgomery/square.rs +++ b/ff-macros/src/montgomery/square.rs @@ -4,6 +4,7 @@ pub(super) fn square_in_place_impl( can_use_no_carry_mul_opt: bool, num_limbs: usize, modulus_limbs: &[u64], + modulus_has_spare_bit: bool, ) -> proc_macro2::TokenStream { let mut body = proc_macro2::TokenStream::new(); let mut default = proc_macro2::TokenStream::new(); @@ -33,6 +34,7 @@ pub(super) fn square_in_place_impl( }); } default.extend(quote! { r[1] <<= 1; }); + for i in 0..num_limbs { default.extend(quote! { r[2 * #i] = fa::mac_with_carry(r[2 * #i], (a.0).0[#i], (a.0).0[#i], &mut carry); @@ -82,14 +84,20 @@ pub(super) fn square_in_place_impl( )] #[allow(unsafe_code, unused_mut)] ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); + __subtract_modulus(a); } else { #default + __subtract_modulus(a); } - })) + })); + return body; } else { - body.extend(quote!({ #default })) + body.extend(quote!( #default )); + if modulus_has_spare_bit { + body.extend(quote!(__subtract_modulus(a);)); + } else { + body.extend(quote!(__subtract_modulus_with_carry(a, carry2 != 0);)); + } } - - body.extend(quote!(__subtract_modulus(a);)); body } diff --git a/test-templates/src/fields.rs b/test-templates/src/fields.rs index f249a8395..dbe7938fd 100644 --- a/test-templates/src/fields.rs +++ b/test-templates/src/fields.rs @@ -209,55 +209,55 @@ macro_rules! __test_field { let mut rng = test_rng(); let zero = <$field>::zero(); let one = <$field>::one(); - assert_eq!(one.inverse().unwrap(), one); - assert!(one.is_one()); + assert_eq!(one.inverse().unwrap(), one, "One inverse failed"); + assert!(one.is_one(), "One is not one"); - assert!(<$field>::ONE.is_one()); - assert_eq!(<$field>::ONE, one); + assert!(<$field>::ONE.is_one(), "One constant is not one"); + assert_eq!(<$field>::ONE, one, "One constant is incorrect"); for _ in 0..ITERATIONS { // Associativity let a = <$field>::rand(&mut rng); let b = <$field>::rand(&mut rng); let c = <$field>::rand(&mut rng); - assert_eq!((a * b) * c, a * (b * c)); + assert_eq!((a * b) * c, a * (b * c), "Associativity failed"); // Commutativity - assert_eq!(a * b, b * a); + assert_eq!(a * b, b * a, "Commutativity failed"); // Identity - assert_eq!(one * a, a); - assert_eq!(one * b, b); - assert_eq!(one * c, c); + assert_eq!(one * a, a, "Identity mul failed"); + assert_eq!(one * b, b, "Identity mul failed"); + assert_eq!(one * c, c, "Identity mul failed"); - assert_eq!(zero * a, zero); - assert_eq!(zero * b, zero); - assert_eq!(zero * c, zero); + assert_eq!(zero * a, zero, "Mul by zero failed"); + assert_eq!(zero * b, zero, "Mul by zero failed"); + assert_eq!(zero * c, zero, "Mul by zero failed"); // Inverses - assert_eq!(a * a.inverse().unwrap(), one); - assert_eq!(b * b.inverse().unwrap(), one); - assert_eq!(c * c.inverse().unwrap(), one); + assert_eq!(a * a.inverse().unwrap(), one, "Mul by inverse failed"); + assert_eq!(b * b.inverse().unwrap(), one, "Mul by inverse failed"); + assert_eq!(c * c.inverse().unwrap(), one, "Mul by inverse failed"); // Associativity and commutativity simultaneously let t0 = (a * b) * c; let t1 = (a * c) * b; let t2 = (b * c) * a; - assert_eq!(t0, t1); - assert_eq!(t1, t2); + assert_eq!(t0, t1, "Associativity + commutativity failed"); + assert_eq!(t1, t2, "Associativity + commutativity failed"); // Squaring - assert_eq!(a * a, a.square()); - assert_eq!(b * b, b.square()); - assert_eq!(c * c, c.square()); + assert_eq!(a * a, a.square(), "Squaring failed"); + assert_eq!(b * b, b.square(), "Squaring failed"); + assert_eq!(c * c, c.square(), "Squaring failed"); // Distributivity - assert_eq!(a * (b + c), a * b + a * c); - assert_eq!(b * (a + c), b * a + b * c); - assert_eq!(c * (a + b), c * a + c * b); - assert_eq!((a + b).square(), a.square() + b.square() + a * b.double()); - assert_eq!((b + c).square(), c.square() + b.square() + c * b.double()); - assert_eq!((c + a).square(), a.square() + c.square() + a * c.double()); + assert_eq!(a * (b + c), a * b + a * c, "Distributivity failed"); + assert_eq!(b * (a + c), b * a + b * c, "Distributivity failed"); + assert_eq!(c * (a + b), c * a + c * b, "Distributivity failed"); + assert_eq!((a + b).square(), a.square() + b.square() + a * b.double(), "Distributivity for square failed"); + assert_eq!((b + c).square(), c.square() + b.square() + c * b.double(), "Distributivity for square failed"); + assert_eq!((c + a).square(), a.square() + c.square() + a * c.double(), "Distributivity for square failed"); } } From b41c011ef8f19e5ebd44824aabd772c7319a372d Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Fri, 9 Dec 2022 18:15:26 -0800 Subject: [PATCH 4/7] Tweak --- ff-macros/src/montgomery/square.rs | 8 +++++--- ff/src/fields/models/fp/mod.rs | 1 + 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/ff-macros/src/montgomery/square.rs b/ff-macros/src/montgomery/square.rs index fa5978a98..e2eb9ab9d 100644 --- a/ff-macros/src/montgomery/square.rs +++ b/ff-macros/src/montgomery/square.rs @@ -58,7 +58,7 @@ pub(super) fn square_in_place_impl( } default.extend(quote! { carry2 = fa::adc(&mut r[#num_limbs + #i], carry, carry2); }); } - default.extend(quote! { (a.0).0.copy_from_slice(&r[#num_limbs..]); }); + default.extend(quote! { (a.0).0 = r[#num_limbs..].try_into().unwrap(); }); if num_limbs == 1 { // We default to multiplying with `a` using the `Mul` impl @@ -83,8 +83,10 @@ pub(super) fn square_in_place_impl( ) )] #[allow(unsafe_code, unused_mut)] - ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); - __subtract_modulus(a); + { + ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); + __subtract_modulus(a); + } } else { #default __subtract_modulus(a); diff --git a/ff/src/fields/models/fp/mod.rs b/ff/src/fields/models/fp/mod.rs index ceb108aee..7ddf59b98 100644 --- a/ff/src/fields/models/fp/mod.rs +++ b/ff/src/fields/models/fp/mod.rs @@ -131,6 +131,7 @@ pub type Fp768

= Fp; pub type Fp832

= Fp; impl, const N: usize> Fp { + #[doc(hidden)] #[inline] pub fn is_geq_modulus(&self) -> bool { self.0 >= P::MODULUS From 141e4bcb57fdc986d462aa1f78189e61090ea651 Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Fri, 9 Dec 2022 18:17:30 -0800 Subject: [PATCH 5/7] Use macro backtrace in nightly --- .github/workflows/ci.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5d4656af6..39c56113f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -116,7 +116,9 @@ jobs: command: test args: "--workspace \ --package ark-test-curves \ - --all-features" + --all-features + -- -Z macro-backtrace + " if: matrix.rust == 'nightly' test_assembly: From 47d643c5b1821424478314614ecd7395148a47b3 Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Fri, 9 Dec 2022 18:45:32 -0800 Subject: [PATCH 6/7] Fix --- ff-macros/src/montgomery/mul.rs | 4 +--- ff-macros/src/montgomery/square.rs | 25 ++++++++++++------------- 2 files changed, 13 insertions(+), 16 deletions(-) diff --git a/ff-macros/src/montgomery/mul.rs b/ff-macros/src/montgomery/mul.rs index 63e58cca6..23050e66d 100644 --- a/ff-macros/src/montgomery/mul.rs +++ b/ff-macros/src/montgomery/mul.rs @@ -61,9 +61,7 @@ pub(super) fn mul_assign_impl( } })) } else { - body.extend(quote!({ - #default - })) + body.extend(quote!({ #default })) } body.extend(quote!(__subtract_modulus(a);)); } else { diff --git a/ff-macros/src/montgomery/square.rs b/ff-macros/src/montgomery/square.rs index e2eb9ab9d..eb9a6cef4 100644 --- a/ff-macros/src/montgomery/square.rs +++ b/ff-macros/src/montgomery/square.rs @@ -17,8 +17,9 @@ pub(super) fn square_in_place_impl( }); for i in 0..(num_limbs - 1) { for j in (i + 1)..num_limbs { + let idx = i + j; default.extend(quote! { - r[#i + #j] = fa::mac_with_carry(r[#i + #j], (a.0).0[#i], (a.0).0[#j], &mut carry); + r[#idx] = fa::mac_with_carry(r[#idx], (a.0).0[#i], (a.0).0[#j], &mut carry); }) } default.extend(quote! { @@ -28,17 +29,16 @@ pub(super) fn square_in_place_impl( } default.extend(quote! { r[#double_num_limbs - 1] = r[#double_num_limbs - 2] >> 63; }); for i in 2..(double_num_limbs - 1) { - default.extend(quote! { - r[#double_num_limbs - #i] = - (r[#double_num_limbs - #i] << 1) | (r[#double_num_limbs - (#i + 1)] >> 63); - }); + let idx = double_num_limbs - i; + default.extend(quote! { r[#idx] = (r[#idx] << 1) | (r[#idx - 1] >> 63); }); } default.extend(quote! { r[1] <<= 1; }); for i in 0..num_limbs { + let idx = 2 * i; default.extend(quote! { - r[2 * #i] = fa::mac_with_carry(r[2 * #i], (a.0).0[#i], (a.0).0[#i], &mut carry); - carry = fa::adc(&mut r[2 * #i + 1], 0, carry); + r[#idx] = fa::mac_with_carry(r[#idx], (a.0).0[#i], (a.0).0[#i], &mut carry); + carry = fa::adc(&mut r[#idx + 1], 0, carry); }); } // Montgomery reduction @@ -63,9 +63,9 @@ pub(super) fn square_in_place_impl( if num_limbs == 1 { // We default to multiplying with `a` using the `Mul` impl // for the N == 1 case - return quote!({ + quote!({ *a *= *a; - }); + }) } else if (2..=6).contains(&num_limbs) && can_use_no_carry_mul_opt { body.extend(quote!({ if cfg!(all( @@ -85,14 +85,13 @@ pub(super) fn square_in_place_impl( #[allow(unsafe_code, unused_mut)] { ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); - __subtract_modulus(a); } } else { #default - __subtract_modulus(a); } })); - return body; + body.extend(quote!(__subtract_modulus(a);)); + body } else { body.extend(quote!( #default )); if modulus_has_spare_bit { @@ -100,6 +99,6 @@ pub(super) fn square_in_place_impl( } else { body.extend(quote!(__subtract_modulus_with_carry(a, carry2 != 0);)); } + body } - body } From 8fbdc5b5ce2db163b87f4632cca56c0f73c3993d Mon Sep 17 00:00:00 2001 From: Pratyush Mishra Date: Fri, 9 Dec 2022 18:49:24 -0800 Subject: [PATCH 7/7] Clean up assembly branches --- ff-macros/src/montgomery/mul.rs | 12 +++++++++++- ff-macros/src/montgomery/square.rs | 12 +++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/ff-macros/src/montgomery/mul.rs b/ff-macros/src/montgomery/mul.rs index 23050e66d..8f6859aea 100644 --- a/ff-macros/src/montgomery/mul.rs +++ b/ff-macros/src/montgomery/mul.rs @@ -57,7 +57,17 @@ pub(super) fn mul_assign_impl( #[allow(unsafe_code, unused_mut)] ark_ff::x86_64_asm_mul!(#num_limbs, (a.0).0, (b.0).0); } else { - #default + #[cfg( + not(all( + feature = "asm", + target_feature = "bmi2", + target_feature = "adx", + target_arch = "x86_64" + )) + )] + { + #default + } } })) } else { diff --git a/ff-macros/src/montgomery/square.rs b/ff-macros/src/montgomery/square.rs index eb9a6cef4..df4c866e5 100644 --- a/ff-macros/src/montgomery/square.rs +++ b/ff-macros/src/montgomery/square.rs @@ -87,7 +87,17 @@ pub(super) fn square_in_place_impl( ark_ff::x86_64_asm_square!(#num_limbs, (a.0).0); } } else { - #default + #[cfg( + not(all( + feature = "asm", + target_feature = "bmi2", + target_feature = "adx", + target_arch = "x86_64" + )) + )] + { + #default + } } })); body.extend(quote!(__subtract_modulus(a);));