Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Features: Orion implementation related peripheral building blocks #142

Merged
merged 6 commits into from
Nov 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions arith/gf2/src/gf2x128.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
#[cfg(target_arch = "x86_64")]
mod avx;
#[cfg(target_arch = "x86_64")]
pub type GF2x128 = avx::AVXGF2x128;

#[cfg(target_arch = "aarch64")]
mod neon;
#[cfg(target_arch = "aarch64")]
pub type GF2x128 = neon::NeonGF2x128;
354 changes: 354 additions & 0 deletions arith/gf2/src/gf2x128/avx.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,354 @@
use std::{
arch::x86_64::*,
mem::{transmute, zeroed},
ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};

use arith::{Field, FieldSerde, FieldSerdeResult, SimdField};

use crate::{GF2x64, GF2};

#[derive(Debug, Clone, Copy)]
pub struct AVXGF2x128 {
pub v: __m128i,
}

impl FieldSerde for AVXGF2x128 {
const SERIALIZED_SIZE: usize = 16;

#[inline(always)]
fn serialize_into<W: std::io::Write>(&self, mut writer: W) -> FieldSerdeResult<()> {
unsafe {
writer.write_all(transmute::<__m128i, [u8; Self::SERIALIZED_SIZE]>(self.v).as_ref())?
};
Ok(())
}

#[inline(always)]
fn deserialize_from<R: std::io::Read>(mut reader: R) -> FieldSerdeResult<Self> {
let mut u = [0u8; Self::SERIALIZED_SIZE];
reader.read_exact(&mut u)?;
unsafe {
Ok(AVXGF2x128 {
v: transmute::<[u8; Self::SERIALIZED_SIZE], __m128i>(u),
})
}
}

#[inline(always)]
fn try_deserialize_from_ecc_format<R: std::io::Read>(mut reader: R) -> FieldSerdeResult<Self> {
let mut u = [0u8; 32];
reader.read_exact(&mut u)?;
Ok(unsafe {
AVXGF2x128 {
v: transmute::<[u8; 16], __m128i>(u[..16].try_into().unwrap()),
}
})
}
}

impl Field for AVXGF2x128 {
const NAME: &'static str = "Galios Field 2 SIMD 128";

const SIZE: usize = 128 / 8;

const FIELD_SIZE: usize = 1; // in bits

const ZERO: Self = AVXGF2x128 {
v: unsafe { zeroed() },
};

const ONE: Self = AVXGF2x128 {
v: unsafe { transmute::<[u64; 2], __m128i>([!0u64, !0u64]) },
};

const INV_2: Self = AVXGF2x128 {
v: unsafe { zeroed() },
}; // should not be used

#[inline(always)]
fn zero() -> Self {
AVXGF2x128 {
v: unsafe { zeroed() },
}
}

#[inline(always)]
fn one() -> Self {
AVXGF2x128 {
v: unsafe { transmute::<[u64; 2], __m128i>([!0u64, !0u64]) },
}
}

#[inline(always)]
fn is_zero(&self) -> bool {
unsafe { transmute::<__m128i, [u8; 16]>(self.v) == [0; 16] }
}

#[inline(always)]
fn random_unsafe(mut rng: impl rand::RngCore) -> Self {
let mut u = [0u8; 16];
rng.fill_bytes(&mut u);
unsafe {
AVXGF2x128 {
v: *(u.as_ptr() as *const __m128i),
}
}
}

#[inline(always)]
fn random_bool(mut rng: impl rand::RngCore) -> Self {
let mut u = [0u8; 16];
rng.fill_bytes(&mut u);
unsafe {
AVXGF2x128 {
v: *(u.as_ptr() as *const __m128i),
}
}
}

#[inline(always)]
fn exp(&self, exponent: u128) -> Self {
if exponent == 0 {
return Self::one();
}
*self
}

#[inline(always)]
fn inv(&self) -> Option<Self> {
unimplemented!()
}

#[inline(always)]
fn as_u32_unchecked(&self) -> u32 {
unimplemented!("u32 for GF2x128 doesn't make sense")
}

#[inline(always)]
fn from_uniform_bytes(bytes: &[u8; 32]) -> Self {
unsafe {
AVXGF2x128 {
v: transmute::<[u8; 16], __m128i>(bytes[..16].try_into().unwrap()),
}
}
}
}

impl Default for AVXGF2x128 {
#[inline(always)]
fn default() -> Self {
Self::ZERO
}
}

impl PartialEq for AVXGF2x128 {
#[inline(always)]
fn eq(&self, other: &Self) -> bool {
unsafe { _mm_test_all_ones(_mm_cmpeq_epi8(self.v, other.v)) == 1 }
}
}

impl Mul<&AVXGF2x128> for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn mul(self, rhs: &AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_and_si128(self.v, rhs.v) },
}
}
}

impl Mul<AVXGF2x128> for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn mul(self, rhs: AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_and_si128(self.v, rhs.v) },
}
}
}

impl MulAssign<&AVXGF2x128> for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn mul_assign(&mut self, rhs: &AVXGF2x128) {
self.v = unsafe { _mm_and_si128(self.v, rhs.v) };
}
}

impl MulAssign<AVXGF2x128> for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn mul_assign(&mut self, rhs: AVXGF2x128) {
self.v = unsafe { _mm_and_si128(self.v, rhs.v) };
}
}

impl Sub for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn sub(self, rhs: AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_xor_si128(self.v, rhs.v) },
}
}
}

impl SubAssign for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn sub_assign(&mut self, rhs: AVXGF2x128) {
self.v = unsafe { _mm_xor_si128(self.v, rhs.v) };
}
}

impl Add for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn add(self, rhs: AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_xor_si128(self.v, rhs.v) },
}
}
}

impl AddAssign for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn add_assign(&mut self, rhs: AVXGF2x128) {
self.v = unsafe { _mm_xor_si128(self.v, rhs.v) };
}
}

impl Add<&AVXGF2x128> for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn add(self, rhs: &AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_xor_si128(self.v, rhs.v) },
}
}
}

impl AddAssign<&AVXGF2x128> for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn add_assign(&mut self, rhs: &AVXGF2x128) {
self.v = unsafe { _mm_xor_si128(self.v, rhs.v) };
}
}

impl Sub<&AVXGF2x128> for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn sub(self, rhs: &AVXGF2x128) -> AVXGF2x128 {
AVXGF2x128 {
v: unsafe { _mm_xor_si128(self.v, rhs.v) },
}
}
}

impl SubAssign<&AVXGF2x128> for AVXGF2x128 {
#[inline(always)]
#[allow(clippy::suspicious_op_assign_impl)]
fn sub_assign(&mut self, rhs: &AVXGF2x128) {
self.v = unsafe { _mm_xor_si128(self.v, rhs.v) };
}
}

impl<T: std::borrow::Borrow<AVXGF2x128>> std::iter::Sum<T> for AVXGF2x128 {
fn sum<I: Iterator<Item = T>>(iter: I) -> Self {
iter.fold(Self::zero(), |acc, item| acc + item.borrow())
}
}

impl<T: std::borrow::Borrow<AVXGF2x128>> std::iter::Product<T> for AVXGF2x128 {
fn product<I: Iterator<Item = T>>(iter: I) -> Self {
iter.fold(Self::one(), |acc, item| acc * item.borrow())
}
}

impl Neg for AVXGF2x128 {
type Output = AVXGF2x128;

#[inline(always)]
#[allow(clippy::suspicious_arithmetic_impl)]
fn neg(self) -> AVXGF2x128 {
AVXGF2x128 { v: self.v }
}
}

impl From<u32> for AVXGF2x128 {
#[inline(always)]
fn from(v: u32) -> Self {
assert!(v < 2);
if v == 0 {
AVXGF2x128::ZERO
} else {
AVXGF2x128::ONE
}
}
}

impl From<GF2> for AVXGF2x128 {
#[inline(always)]
fn from(v: GF2) -> Self {
assert!(v.v < 2);
if v.v == 0 {
AVXGF2x128::ZERO
} else {
AVXGF2x128::ONE
}
}
}

impl SimdField for AVXGF2x128 {
type Scalar = GF2;

const PACK_SIZE: usize = 128;

#[inline(always)]
fn scale(&self, challenge: &Self::Scalar) -> Self {
if challenge.v == 0 {
Self::ZERO
} else {
*self
}
}

#[inline(always)]
fn pack(base_vec: &[Self::Scalar]) -> Self {
assert_eq!(base_vec.len(), Self::PACK_SIZE);
let mut packed_to_gf2x64 = [GF2x64::ZERO; Self::PACK_SIZE / GF2x64::PACK_SIZE];
packed_to_gf2x64
.iter_mut()
.zip(base_vec.chunks(GF2x64::PACK_SIZE))
.for_each(|(gf2x64, pack)| *gf2x64 = GF2x64::pack(pack));

unsafe { transmute(packed_to_gf2x64) }
}

#[inline(always)]
fn unpack(&self) -> Vec<Self::Scalar> {
let packed_to_gf2x64: [GF2x64; Self::PACK_SIZE / GF2x64::PACK_SIZE] =
unsafe { transmute(*self) };

packed_to_gf2x64
.iter()
.flat_map(|packed| packed.unpack())
.collect()
}
}
Loading
Loading