perf: remove 15 cols from poseidon2 wide (#1397)

crates/recursion/core-v2/src/chips/poseidon2_wide/air.rs

@@ -55,17 +55,17 @@ where
         // For now, include only memory constraints.
         (0..WIDTH).for_each(|i| {
             builder.send_single(
-                prep_local.memory_preprocessed[i].addr,
+                prep_local.input[i],
                 local_row.external_rounds_state()[0][i],
-                prep_local.memory_preprocessed[i].mult,
+                prep_local.is_real_neg,
             )
         });
 
         (0..WIDTH).for_each(|i| {
             builder.send_single(
-                prep_local.memory_preprocessed[i + WIDTH].addr,
+                prep_local.output[i].addr,
                 local_row.perm_output()[i],
-                prep_local.memory_preprocessed[i + WIDTH].mult,
+                prep_local.output[i].mult,
             )
         });
 

crates/recursion/core-v2/src/chips/poseidon2_wide/columns/preprocessed.rs

@@ -1,9 +1,14 @@
 use sp1_derive::AlignedBorrow;
 
-use crate::chips::{mem::MemoryAccessCols, poseidon2_wide::WIDTH};
+use crate::{
+    chips::{mem::MemoryAccessCols, poseidon2_wide::WIDTH},
+    Address,
+};
 
 #[derive(AlignedBorrow, Clone, Copy, Debug)]
 #[repr(C)]
 pub struct Poseidon2PreprocessedCols<T: Copy> {
-    pub memory_preprocessed: [MemoryAccessCols<T>; 2 * WIDTH],
+    pub input: [Address<T>; WIDTH],
+    pub output: [MemoryAccessCols<T>; WIDTH],
+    pub is_real_neg: T,
 }

crates/recursion/core-v2/src/chips/poseidon2_wide/trace.rs

@@ -48,9 +48,6 @@ impl<F: PrimeField32, const DEGREE: usize> MachineAir<F> for Poseidon2WideChip<D
         let num_columns = <Self as BaseAir<F>>::width(self);
         let mut values = vec![F::zero(); padded_nb_rows * num_columns];
 
-        let mut dummy_row = vec![F::zero(); num_columns];
-        self.populate_perm([F::zero(); WIDTH], None, &mut dummy_row);
-
         let populate_len = events.len() * num_columns;
         let (values_pop, values_dummy) = values.split_at_mut(populate_len);
         join(
@@ -62,6 +59,8 @@ impl<F: PrimeField32, const DEGREE: usize> MachineAir<F> for Poseidon2WideChip<D
                 )
             },
             || {
+                let mut dummy_row = vec![F::zero(); num_columns];
+                self.populate_perm([F::zero(); WIDTH], None, &mut dummy_row);
                 values_dummy
                     .par_chunks_mut(num_columns)
                     .for_each(|row| row.copy_from_slice(&dummy_row))
@@ -108,20 +107,16 @@ impl<F: PrimeField32, const DEGREE: usize> MachineAir<F> for Poseidon2WideChip<D
             .par_chunks_mut(PREPROCESSED_POSEIDON2_WIDTH)
             .zip_eq(instrs)
             .for_each(|(row, instr)| {
-                let cols: &mut Poseidon2PreprocessedCols<_> = row.borrow_mut();
-
                 // Set the memory columns. We read once, at the first iteration,
                 // and write once, at the last iteration.
-                cols.memory_preprocessed = std::array::from_fn(|j| {
-                    if j < WIDTH {
-                        MemoryAccessCols { addr: instr.addrs.input[j], mult: F::neg_one() }
-                    } else {
-                        MemoryAccessCols {
-                            addr: instr.addrs.output[j - WIDTH],
-                            mult: instr.mults[j - WIDTH],
-                        }
-                    }
-                });
+                *row.borrow_mut() = Poseidon2PreprocessedCols {
+                    input: instr.addrs.input,
+                    output: std::array::from_fn(|j| MemoryAccessCols {
+                        addr: instr.addrs.output[j],
+                        mult: instr.mults[j],
+                    }),
+                    is_real_neg: F::neg_one(),
+                }
             });
         Some(RowMajorMatrix::new(values, PREPROCESSED_POSEIDON2_WIDTH))
     }