diff --git a/src/field/mod.rs b/src/field/mod.rs index bb02a27d..4011767d 100644 --- a/src/field/mod.rs +++ b/src/field/mod.rs @@ -4,6 +4,8 @@ pub mod extension_field; pub mod fft; pub mod field_types; pub(crate) mod interpolation; +pub(crate) mod packable; +pub(crate) mod packed_field; #[cfg(test)] mod field_testing; diff --git a/src/field/packable.rs b/src/field/packable.rs new file mode 100644 index 00000000..50a26e50 --- /dev/null +++ b/src/field/packable.rs @@ -0,0 +1,14 @@ +use crate::field::crandall_field::CrandallField; +use crate::field::field_types::Field; +use crate::field::packed_field::{PackedField, Singleton}; + +/// Points us to the default packing for a particular field. There may me multiple choices of +/// PackedField for a particular Field (e.g. Singleton works for all fields), but this is the +/// recommended one. The recommended packing varies by target_arch and target_feature. +pub trait Packable: Field { + type PackedType: PackedField; +} + +impl Packable for F { + default type PackedType = Singleton; +} diff --git a/src/field/packed_field.rs b/src/field/packed_field.rs new file mode 100644 index 00000000..ec5258b7 --- /dev/null +++ b/src/field/packed_field.rs @@ -0,0 +1,211 @@ +use std::fmt; +use std::fmt::{Debug, Formatter}; +use std::iter::{Product, Sum}; +use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; + +use crate::field::field_types::Field; + +pub trait PackedField: + 'static + + Add + + Add + + AddAssign + + AddAssign + + Copy + + Debug + + Default + // TODO: Implementing Div sounds like a pain so it's a worry for later. + + Mul + + Mul + + MulAssign + + MulAssign + + Neg + + Product + + Send + + Sub + + Sub + + SubAssign + + SubAssign + + Sum + + Sync +{ + type FieldType: Field; + + const LOG2_WIDTH: usize; + const WIDTH: usize = 1 << Self::LOG2_WIDTH; + + fn square(&self) -> Self; + + fn zero() -> Self; + fn one() -> Self; + + fn broadcast(x: Self::FieldType) -> Self; + + fn new_from_slice(arr: &[Self::FieldType]) -> Self; + fn to_vec(&self) -> Vec; + + fn interleave(&self, other: Self, r: usize) -> (Self, Self); + + fn pack_slice(buf: &[Self::FieldType]) -> &[Self] { + assert!( + buf.len() % Self::WIDTH == 0, + "Slice length (got {}) must be a multiple of packed field width ({}).", + buf.len(), + Self::WIDTH + ); + let buf_ptr = buf.as_ptr().cast::(); + // assert!(buf_ptr.align_offset(align_of::()) == 0, "Not aligned"); + let n = buf.len() / Self::WIDTH; + unsafe { std::slice::from_raw_parts(buf_ptr, n) } + } + fn pack_slice_mut(buf: &mut [Self::FieldType]) -> &mut [Self] { + assert!( + buf.len() % Self::WIDTH == 0, + "Slice length (got {}) must be a multiple of packed field width ({}).", + buf.len(), + Self::WIDTH + ); + let buf_ptr = buf.as_mut_ptr().cast::(); + // assert!(buf_ptr.align_offset(align_of::()) == 0, "Not aligned"); + let n = buf.len() / Self::WIDTH; + unsafe { std::slice::from_raw_parts_mut(buf_ptr, n) } + } +} + +#[derive(Copy, Clone)] +#[repr(transparent)] +pub struct Singleton(pub F); + +impl Add for Singleton { + type Output = Self; + fn add(self, rhs: Self) -> Self { + Self(self.0 + rhs.0) + } +} +impl Add for Singleton { + type Output = Self; + fn add(self, rhs: F) -> Self { + self + Self::broadcast(rhs) + } +} +impl AddAssign for Singleton { + fn add_assign(&mut self, rhs: Self) { + *self = *self + rhs; + } +} +impl AddAssign for Singleton { + fn add_assign(&mut self, rhs: F) { + *self = *self + rhs; + } +} + +impl Debug for Singleton { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "({:?})", self.0) + } +} + +impl Default for Singleton { + fn default() -> Self { + Self::zero() + } +} + +impl Mul for Singleton { + type Output = Self; + fn mul(self, rhs: Self) -> Self { + Self(self.0 * rhs.0) + } +} +impl Mul for Singleton { + type Output = Self; + fn mul(self, rhs: F) -> Self { + self * Self::broadcast(rhs) + } +} +impl MulAssign for Singleton { + fn mul_assign(&mut self, rhs: Self) { + *self = *self * rhs; + } +} +impl MulAssign for Singleton { + fn mul_assign(&mut self, rhs: F) { + *self = *self * rhs; + } +} + +impl Neg for Singleton { + type Output = Self; + fn neg(self) -> Self { + Self(-self.0) + } +} + +impl Product for Singleton { + fn product>(iter: I) -> Self { + Self(iter.map(|x| x.0).product()) + } +} + +impl PackedField for Singleton { + const LOG2_WIDTH: usize = 0; + type FieldType = F; + + fn square(&self) -> Self { + *self * *self + } + + fn zero() -> Self { + Self::broadcast(F::ZERO) + } + fn one() -> Self { + Self::broadcast(F::ONE) + } + + fn broadcast(x: F) -> Self { + Self(x) + } + + fn new_from_slice(arr: &[Self::FieldType]) -> Self { + Self(arr[0]) + } + fn to_vec(&self) -> Vec { + vec![self.0] + } + + fn interleave(&self, other: Self, r: usize) -> (Self, Self) { + match r { + 0 => (*self, other), // This is a no-op whenever r == LOG2_WIDTH. + _ => panic!("r cannot be more than LOG2_WIDTH"), + } + } +} + +impl Sub for Singleton { + type Output = Self; + fn sub(self, rhs: Self) -> Self { + Self(self.0 - rhs.0) + } +} +impl Sub for Singleton { + type Output = Self; + fn sub(self, rhs: F) -> Self { + self - Self::broadcast(rhs) + } +} +impl SubAssign for Singleton { + fn sub_assign(&mut self, rhs: Self) { + *self = *self - rhs; + } +} +impl SubAssign for Singleton { + fn sub_assign(&mut self, rhs: F) { + *self = *self - rhs; + } +} + +impl Sum for Singleton { + fn sum>(iter: I) -> Self { + Self(iter.map(|x| x.0).sum()) + } +} diff --git a/src/lib.rs b/src/lib.rs index 45ee1fee..45928fb9 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ #![feature(destructuring_assignment)] #![feature(generic_const_exprs)] +#![feature(specialization)] pub mod field; pub mod fri;