diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs index 0c544a4a5..28286f395 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -1,7 +1,6 @@ -#[cfg(target_arch = "x86_64")] -pub mod avx; pub mod component; pub mod constraint_eval; +pub mod simd; pub mod trace_gen; #[cfg(test)] diff --git a/crates/prover/src/examples/wide_fibonacci/avx.rs b/crates/prover/src/examples/wide_fibonacci/simd.rs similarity index 81% rename from crates/prover/src/examples/wide_fibonacci/avx.rs rename to crates/prover/src/examples/wide_fibonacci/simd.rs index 50fa953da..ef9722941 100644 --- a/crates/prover/src/examples/wide_fibonacci/avx.rs +++ b/crates/prover/src/examples/wide_fibonacci/simd.rs @@ -5,8 +5,10 @@ use tracing::{span, Level}; use super::component::{WideFibAir, WideFibComponent}; use crate::core::air::accumulation::DomainEvaluationAccumulator; use crate::core::air::{AirProver, ComponentProver, ComponentTrace}; -use crate::core::backend::avx512::qm31::PackedSecureField; -use crate::core::backend::avx512::{AVX512Backend, BaseFieldVec, PackedBaseField, VECS_LOG_SIZE}; +use crate::core::backend::simd::column::BaseFieldVec; +use crate::core::backend::simd::m31::{PackedBaseField, LOG_N_LANES}; +use crate::core::backend::simd::qm31::PackedSecureField; +use crate::core::backend::simd::SimdBackend; use crate::core::backend::{Col, Column, ColumnOps}; use crate::core::constraints::coset_vanishing; use crate::core::fields::m31::BaseField; @@ -16,20 +18,20 @@ use crate::core::poly::BitReversedOrder; use crate::core::ColumnVec; use crate::examples::wide_fibonacci::component::N_COLUMNS; -impl AirProver for WideFibAir { - fn prover_components(&self) -> Vec<&dyn ComponentProver> { +impl AirProver for WideFibAir { + fn prover_components(&self) -> Vec<&dyn ComponentProver> { vec![&self.component] } } pub fn gen_trace( log_size: usize, -) -> ColumnVec> { - assert!(log_size >= VECS_LOG_SIZE); +) -> ColumnVec> { + assert!(log_size >= LOG_N_LANES as usize); let mut trace = (0..N_COLUMNS) - .map(|_| Col::::zeros(1 << log_size)) + .map(|_| Col::::zeros(1 << log_size)) .collect_vec(); - for vec_index in 0..(1 << (log_size - VECS_LOG_SIZE)) { + for vec_index in 0..(1 << (log_size - LOG_N_LANES as usize)) { let mut a = PackedBaseField::one(); let mut b = PackedBaseField::from_array(std::array::from_fn(|i| { BaseField::from_u32_unchecked((vec_index * 16 + i) as u32) @@ -44,15 +46,15 @@ pub fn gen_trace( let domain = CanonicCoset::new(log_size as u32).circle_domain(); trace .into_iter() - .map(|eval| CircleEvaluation::::new(domain, eval)) + .map(|eval| CircleEvaluation::::new(domain, eval)) .collect_vec() } -impl ComponentProver for WideFibComponent { +impl ComponentProver for WideFibComponent { fn evaluate_constraint_quotients_on_domain( &self, - trace: &ComponentTrace<'_, AVX512Backend>, - evaluation_accumulator: &mut DomainEvaluationAccumulator, + trace: &ComponentTrace<'_, SimdBackend>, + evaluation_accumulator: &mut DomainEvaluationAccumulator, ) { assert_eq!(trace.polys.len(), self.n_columns()); // TODO(spapini): Steal evaluation from commitment. @@ -65,9 +67,9 @@ impl ComponentProver for WideFibComponent { let zero_domain = CanonicCoset::new(self.log_column_size()).coset; let mut denoms = BaseFieldVec::from_iter(eval_domain.iter().map(|p| coset_vanishing(zero_domain, p))); - >::bit_reverse_column(&mut denoms); + >::bit_reverse_column(&mut denoms); let mut denom_inverses = BaseFieldVec::zeros(denoms.len()); - >::batch_inverse(&denoms, &mut denom_inverses); + >::batch_inverse(&denoms, &mut denom_inverses); span.exit(); let _span = span!(Level::INFO, "Constraint pointwise eval").entered(); @@ -76,7 +78,7 @@ impl ComponentProver for WideFibComponent { let [accum] = evaluation_accumulator.columns([(constraint_log_degree_bound, self.n_columns() - 1)]); - for vec_row in 0..(1 << (eval_domain.log_size() - VECS_LOG_SIZE as u32)) { + for vec_row in 0..(1 << (eval_domain.log_size() - LOG_N_LANES)) { // Numerator. let a = trace_eval[0].data[vec_row]; let mut row_res = @@ -110,18 +112,17 @@ impl ComponentProver for WideFibComponent { } } -#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] #[cfg(test)] mod tests { use tracing::{span, Level}; + use super::{gen_trace, WideFibAir}; use crate::core::channel::{Blake2sChannel, Channel}; use crate::core::fields::m31::BaseField; use crate::core::fields::IntoSlice; use crate::core::prover::{prove, verify}; use crate::core::vcs::blake2_hash::Blake2sHasher; use crate::core::vcs::hasher::Hasher; - use crate::examples::wide_fibonacci::avx::{gen_trace, WideFibAir}; use crate::examples::wide_fibonacci::component::{WideFibComponent, LOG_N_COLUMNS}; #[test_log::test]