From 9bd0fc44239cc0b9bf9e1bc17208cb13fa021ae1 Mon Sep 17 00:00:00 2001 From: Andrew Milson Date: Tue, 30 Apr 2024 23:33:26 -0400 Subject: [PATCH] Implement FieldOps for SIMD backend --- crates/prover/src/core/backend/simd/column.rs | 159 ++++++++++++++++++ crates/prover/src/core/backend/simd/m31.rs | 9 +- crates/prover/src/core/backend/simd/mod.rs | 37 ++++ crates/prover/src/core/backend/simd/qm31.rs | 7 + 4 files changed, 211 insertions(+), 1 deletion(-) create mode 100644 crates/prover/src/core/backend/simd/column.rs diff --git a/crates/prover/src/core/backend/simd/column.rs b/crates/prover/src/core/backend/simd/column.rs new file mode 100644 index 000000000..202fb7a3f --- /dev/null +++ b/crates/prover/src/core/backend/simd/column.rs @@ -0,0 +1,159 @@ +use bytemuck::{cast_slice, cast_slice_mut, Zeroable}; +use itertools::Itertools; +use num_traits::Zero; + +use super::m31::{PackedBaseField, N_LANES}; +use super::qm31::PackedSecureField; +use crate::core::backend::Column; +use crate::core::fields::m31::BaseField; +use crate::core::fields::qm31::SecureField; + +#[derive(Clone, Debug)] +pub struct BaseFieldVec { + pub data: Vec, + pub length: usize, +} + +impl AsRef<[BaseField]> for BaseFieldVec { + fn as_ref(&self) -> &[BaseField] { + &cast_slice(&self.data)[..self.length] + } +} + +impl AsMut<[BaseField]> for BaseFieldVec { + fn as_mut(&mut self) -> &mut [BaseField] { + &mut cast_slice_mut(&mut self.data)[..self.length] + } +} + +impl Column for BaseFieldVec { + fn zeros(length: usize) -> Self { + let data = vec![PackedBaseField::zeroed(); length.div_ceil(N_LANES)]; + Self { data, length } + } + + fn to_cpu(&self) -> Vec { + self.as_ref().to_vec() + } + + fn len(&self) -> usize { + self.length + } + + fn at(&self, index: usize) -> BaseField { + self.data[index / N_LANES].to_array()[index % N_LANES] + } +} + +impl FromIterator for BaseFieldVec { + fn from_iter>(iter: I) -> Self { + let mut chunks = iter.into_iter().array_chunks(); + let mut data = (&mut chunks).map(PackedBaseField::from_array).collect_vec(); + let mut length = data.len() * N_LANES; + + if let Some(remainder) = chunks.into_remainder() { + if !remainder.is_empty() { + length += remainder.len(); + let mut last = [BaseField::zero(); N_LANES]; + last[..remainder.len()].copy_from_slice(remainder.as_slice()); + data.push(PackedBaseField::from_array(last)); + } + } + + Self { data, length } + } +} + +#[derive(Clone, Debug)] +pub struct SecureFieldVec { + pub data: Vec, + pub length: usize, +} + +impl Column for SecureFieldVec { + fn zeros(length: usize) -> Self { + Self { + data: vec![PackedSecureField::zeroed(); length.div_ceil(N_LANES)], + length, + } + } + + fn to_cpu(&self) -> Vec { + self.data + .iter() + .flat_map(|x| x.to_array()) + .take(self.length) + .collect() + } + + fn len(&self) -> usize { + self.length + } + + fn at(&self, index: usize) -> SecureField { + self.data[index / N_LANES].to_array()[index % N_LANES] + } +} + +impl FromIterator for SecureFieldVec { + fn from_iter>(iter: I) -> Self { + let mut chunks = iter.into_iter().array_chunks(); + let mut data = (&mut chunks) + .map(PackedSecureField::from_array) + .collect_vec(); + let mut length = data.len() * N_LANES; + + if let Some(remainder) = chunks.into_remainder() { + if !remainder.is_empty() { + length += remainder.len(); + let mut last = [SecureField::zero(); N_LANES]; + last[..remainder.len()].copy_from_slice(remainder.as_slice()); + data.push(PackedSecureField::from_array(last)); + } + } + + Self { data, length } + } +} + +impl FromIterator for SecureFieldVec { + fn from_iter>(iter: I) -> Self { + let data = (&mut iter.into_iter()).collect_vec(); + let length = data.len() * N_LANES; + + Self { data, length } + } +} + +#[cfg(test)] +mod tests { + use std::array; + + use rand::rngs::SmallRng; + use rand::{Rng, SeedableRng}; + + use super::BaseFieldVec; + use crate::core::backend::simd::column::SecureFieldVec; + use crate::core::backend::Column; + use crate::core::fields::m31::BaseField; + use crate::core::fields::qm31::SecureField; + + #[test] + fn base_field_vec_from_iter_works() { + let values: [BaseField; 30] = array::from_fn(BaseField::from); + + let res = values.into_iter().collect::(); + + assert_eq!(res.to_cpu(), values); + } + + #[test] + fn secure_field_vec_from_iter_works() { + let mut rng = SmallRng::seed_from_u64(0); + let values: [SecureField; 30] = rng.gen(); + + let res = values.into_iter().collect::(); + + assert_eq!(res.to_cpu(), values); + } +} diff --git a/crates/prover/src/core/backend/simd/m31.rs b/crates/prover/src/core/backend/simd/m31.rs index d3494b298..257787fbc 100644 --- a/crates/prover/src/core/backend/simd/m31.rs +++ b/crates/prover/src/core/backend/simd/m31.rs @@ -6,6 +6,7 @@ use std::simd::{u32x16, Simd, Swizzle}; use bytemuck::{Pod, Zeroable}; use num_traits::{One, Zero}; +use rand::distributions::{Distribution, Standard}; use crate::core::backend::simd::utils::{LoEvensInterleaveHiEvens, LoOddsInterleaveHiOdds}; use crate::core::fields::m31::{BaseField, M31, P}; @@ -60,7 +61,7 @@ impl PackedBaseField { self.to_array().into_iter().sum() } - /// Doubles each element. + /// Doubles each element in the vector. pub fn double(self) -> Self { // TODO: Make more optimal. self + self @@ -212,6 +213,12 @@ impl From<[BaseField; N_LANES]> for PackedBaseField { } } +impl Distribution for Standard { + fn sample(&self, rng: &mut R) -> PackedBaseField { + PackedBaseField::from_array(rng.gen()) + } +} + #[cfg(target_arch = "aarch64")] fn _mul_neon(a: PackedBaseField, b: PackedBaseField) -> PackedBaseField { use core::arch::aarch64::{int32x2_t, vqdmull_s32}; diff --git a/crates/prover/src/core/backend/simd/mod.rs b/crates/prover/src/core/backend/simd/mod.rs index ff94a9f5b..757474307 100644 --- a/crates/prover/src/core/backend/simd/mod.rs +++ b/crates/prover/src/core/backend/simd/mod.rs @@ -1,7 +1,44 @@ +use self::column::{BaseFieldVec, SecureFieldVec}; +use self::m31::PackedBaseField; +use self::qm31::PackedSecureField; +use super::ColumnOps; +use crate::core::fields::m31::BaseField; +use crate::core::fields::qm31::SecureField; +use crate::core::fields::{FieldExpOps, FieldOps}; + pub mod cm31; +pub mod column; pub mod m31; pub mod qm31; mod utils; #[derive(Copy, Clone, Debug)] pub struct SimdBackend; + +impl ColumnOps for SimdBackend { + type Column = BaseFieldVec; + + fn bit_reverse_column(_column: &mut Self::Column) { + todo!() + } +} + +impl FieldOps for SimdBackend { + fn batch_inverse(column: &Self::Column, dst: &mut Self::Column) { + PackedBaseField::batch_inverse(&column.data, &mut dst.data); + } +} + +impl ColumnOps for SimdBackend { + type Column = SecureFieldVec; + + fn bit_reverse_column(_column: &mut Self::Column) { + todo!() + } +} + +impl FieldOps for SimdBackend { + fn batch_inverse(column: &Self::Column, dst: &mut Self::Column) { + PackedSecureField::batch_inverse(&column.data, &mut dst.data); + } +} diff --git a/crates/prover/src/core/backend/simd/qm31.rs b/crates/prover/src/core/backend/simd/qm31.rs index 459fed27f..8e4c92f88 100644 --- a/crates/prover/src/core/backend/simd/qm31.rs +++ b/crates/prover/src/core/backend/simd/qm31.rs @@ -4,6 +4,7 @@ use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; use bytemuck::{Pod, Zeroable}; use num_traits::{One, Zero}; +use rand::distributions::{Distribution, Standard}; use super::cm31::PackedCM31; use super::m31::{PackedBaseField, N_LANES}; @@ -224,6 +225,12 @@ impl Neg for PackedSecureField { } } +impl Distribution for Standard { + fn sample(&self, rng: &mut R) -> PackedSecureField { + PackedSecureField::from_array(rng.gen()) + } +} + #[cfg(test)] mod tests { use std::array;