perf: remove preprocessed trace bitwiseoplookupair [cpu & gpu tracegen] (#2366)

Resolves INT-5852.

Author: Tuanlinh12312

crates/circuits/primitives/cuda/src/bitwise_op_lookup.cu

@@ -1,5 +1,15 @@
 #include "fp.h"
 #include "launcher.cuh"
+#include "primitives/trace_access.h"
+
+constexpr uint32_t NUM_BITS = 8;
+
+template <typename T> struct BitwiseOperationLookupCols {
+    T x_bits[NUM_BITS];
+    T y_bits[NUM_BITS];
+    T mult_range;
+    T mult_xor;
+};
 
 __global__ void bitwise_op_lookup_tracegen(
     const uint32_t *count,
@@ -9,9 +19,26 @@ __global__ void bitwise_op_lookup_tracegen(
 ) {
     uint32_t row_idx = blockIdx.x * blockDim.x + threadIdx.x;
     if (row_idx < num_rows) {
-        trace[row_idx] = Fp(count[row_idx] + (cpu_count ? cpu_count[row_idx] : 0));
-        trace[row_idx + num_rows] =
-            Fp(count[row_idx + num_rows] + (cpu_count ? cpu_count[row_idx + num_rows] : 0));
+        uint32_t x = row_idx >> NUM_BITS;
+        uint32_t y = row_idx & ((1U << NUM_BITS) - 1);
+
+        Fp x_bits_array[NUM_BITS];
+        Fp y_bits_array[NUM_BITS];
+#pragma unroll
+        for (uint32_t i = 0; i < NUM_BITS; i++) {
+            x_bits_array[i] = Fp((x >> i) & 1);
+            y_bits_array[i] = Fp((y >> i) & 1);
+        }
+
+        uint32_t mult_range_val = count[row_idx] + (cpu_count ? cpu_count[row_idx] : 0);
+        uint32_t mult_xor_val = count[row_idx + num_rows] + 
+                                (cpu_count ? cpu_count[row_idx + num_rows] : 0);
+
+        RowSlice row(trace + row_idx, num_rows);
+        COL_WRITE_ARRAY(row, BitwiseOperationLookupCols, x_bits, x_bits_array);
+        COL_WRITE_ARRAY(row, BitwiseOperationLookupCols, y_bits, y_bits_array);
+        COL_WRITE_VALUE(row, BitwiseOperationLookupCols, mult_range, mult_range_val);
+        COL_WRITE_VALUE(row, BitwiseOperationLookupCols, mult_xor, mult_xor_val);
     }
 }
 

crates/circuits/primitives/src/bitwise_op_lookup/README.md

@@ -2,15 +2,21 @@
 
 XOR operation and range checking via lookup table
 
-This chip implements a lookup table approach for XOR operations and range checks for integers of size $`\texttt{NUM\_BITS}`$. The lookup table contains all possible combinations of $`x`$ and $`y`$ values (both in the range $`0..2^{\texttt{NUM\_BITS}}`$), along with their XOR result.
+This chip implements a lookup table approach for XOR operations and range checks for integers of size $`\texttt{NUM\_BITS}`$. The chip provides lookup table functionality for all possible combinations of $`x`$ and $`y`$ values (both in the range $`0..2^{\texttt{NUM\_BITS}}`$), enabling verification of XOR operations and range checks. In the trace, $x$ and $y$ are stored as binary decompositions (`x_bits` and `y_bits` arrays) rather than as full field elements.
 
-The lookup mechanism works through the Bus interface, with other circuits requesting lookups by incrementing multiplicity counters for the operations they need to perform. Each row in the lookup table corresponds to a specific $(x, y)$ pair.
+The lookup mechanism works through the Bus interface, with other circuits requesting lookups by incrementing multiplicity counters for the operations they need to perform. Each row in the trace corresponds to a specific $(x, y)$ pair.
 
-**Preprocessed Columns:**
-- `x`: Column containing the first input value ($0$ to $`2^{\texttt{NUM\_BITS}}-1`$)
-- `y`: Column containing the second input value ($0$ to $`2^{\texttt{NUM\_BITS}}-1`$)
-- `z_xor`: Column containing the XOR result of x and y ($x \oplus y$)
+The chip uses gate-based constraints to generate the trace columns instead of a preprocessed trace. The trace enumerates all valid $(x, y)$ pairs in order: row $n$ corresponds to $(x, y)$ where $x = \lfloor n / 2^{\texttt{NUM\_BITS}} \rfloor$ and $y = n \bmod 2^{\texttt{NUM\_BITS}}$. The enumeration order is: $(0, 0)$, $(0, 1)$, ..., $(0, 2^{\texttt{NUM\_BITS}}-1)$, $(1, 0)$, $(1, 1)$, ..., up to $(2^{\texttt{NUM\_BITS}}-1, 2^{\texttt{NUM\_BITS}}-1)$.
 
-**IO Columns:**
+**Columns:**
+- `x_bits[0..NUM_BITS-1]`: Binary decomposition of $x$ (where `x_bits[0]` is the least significant bit)
+- `y_bits[0..NUM_BITS-1]`: Binary decomposition of $y$ (where `y_bits[0]` is the least significant bit)
 - `mult_range`: Multiplicity column tracking the number of range check operations requested for each $(x, y)$ pair
 - `mult_xor`: Multiplicity column tracking the number of XOR operations requested for each $(x, y)$ pair
+
+The constraints enforce the enumeration pattern by:
+1. Ensuring each bit is binary (0 or 1) using `assert_bool` constraints
+2. Reconstructing $x$ and $y$ from their binary decompositions: $x = \sum_{i=0}^{\texttt{NUM\_BITS}-1} \texttt{x\_bits}[i] \cdot 2^i$
+3. Computing $z_{\texttt{xor}} = x \oplus y$ algebraically from bits: $z_{\texttt{xor}} = \sum_{i=0}^{\texttt{NUM\_BITS}-1} (\texttt{x\_bits}[i] + \texttt{y\_bits}[i] - 2 \cdot \texttt{x\_bits}[i] \cdot \texttt{y\_bits}[i]) \cdot 2^i$
+4. Constraining that the combined index $(x \cdot 2^{\texttt{NUM\_BITS}} + y)$ increments by 1 each row using transition constraints
+5. Enforcing boundary conditions: first row has index 0, last row has index $2^{2 \cdot \texttt{NUM\_BITS}} - 1$

crates/circuits/primitives/src/bitwise_op_lookup/cuda.rs

@@ -5,7 +5,9 @@ use openvm_cuda_common::{copy::MemCopyH2D as _, d_buffer::DeviceBuffer};
 use openvm_stark_backend::{prover::types::AirProvingContext, Chip};
 
 use crate::{
-    bitwise_op_lookup::{BitwiseOperationLookupChip, NUM_BITWISE_OP_LOOKUP_COLS},
+    bitwise_op_lookup::{
+        BitwiseOperationLookupChip, BitwiseOperationLookupCols, NUM_BITWISE_OP_LOOKUP_MULT_COLS,
+    },
     cuda_abi::bitwise_op_lookup::tracegen,
 };
 
@@ -22,7 +24,7 @@ impl<const NUM_BITS: usize> BitwiseOperationLookupChipGPU<NUM_BITS> {
     pub fn new() -> Self {
         // The first 2^(2 * NUM_BITS) indices are for range checking, the rest are for XOR
         let count = Arc::new(DeviceBuffer::<F>::with_capacity(
-            NUM_BITWISE_OP_LOOKUP_COLS * Self::num_rows(),
+            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
         ));
         count.fill_zero().unwrap();
         Self {
@@ -35,7 +37,7 @@ impl<const NUM_BITS: usize> BitwiseOperationLookupChipGPU<NUM_BITS> {
         assert_eq!(cpu_chip.count_range.len(), Self::num_rows());
         assert_eq!(cpu_chip.count_xor.len(), Self::num_rows());
         let count = Arc::new(DeviceBuffer::<F>::with_capacity(
-            NUM_BITWISE_OP_LOOKUP_COLS * Self::num_rows(),
+            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
         ));
         count.fill_zero().unwrap();
         Self {
@@ -53,8 +55,9 @@ impl<const NUM_BITS: usize> Default for BitwiseOperationLookupChipGPU<NUM_BITS>
 
 impl<RA, const NUM_BITS: usize> Chip<RA, GpuBackend> for BitwiseOperationLookupChipGPU<NUM_BITS> {
     fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<GpuBackend> {
+        let num_cols = BitwiseOperationLookupCols::<F, NUM_BITS>::width();
         debug_assert_eq!(
-            Self::num_rows() * NUM_BITWISE_OP_LOOKUP_COLS,
+            NUM_BITWISE_OP_LOOKUP_MULT_COLS * Self::num_rows(),
             self.count.len()
         );
         let cpu_count = self.cpu_chip.as_ref().map(|cpu_chip| {
@@ -69,7 +72,7 @@ impl<RA, const NUM_BITS: usize> Chip<RA, GpuBackend> for BitwiseOperationLookupC
         });
         // ATTENTION: we create a new buffer to copy `count` into because this chip is stateful and
         // `count` will be reused.
-        let trace = DeviceMatrix::<F>::with_capacity(Self::num_rows(), NUM_BITWISE_OP_LOOKUP_COLS);
+        let trace = DeviceMatrix::<F>::with_capacity(Self::num_rows(), num_cols);
         unsafe {
             tracegen(&self.count, &cpu_count, trace.buffer(), NUM_BITS as u32).unwrap();
         }

crates/circuits/primitives/src/bitwise_op_lookup/mod.rs

@@ -1,14 +1,13 @@
 use std::{
     borrow::{Borrow, BorrowMut},
-    mem::size_of,
     sync::{atomic::AtomicU32, Arc},
 };
 
 use openvm_circuit_primitives_derive::AlignedBorrow;
 use openvm_stark_backend::{
     config::{StarkGenericConfig, Val},
     interaction::InteractionBuilder,
-    p3_air::{Air, BaseAir, PairBuilder},
+    p3_air::{Air, AirBuilder, BaseAir},
     p3_field::{Field, FieldAlgebra},
     p3_matrix::{dense::RowMajorMatrix, Matrix},
     prover::{cpu::CpuBackend, types::AirProvingContext},
@@ -27,27 +26,21 @@ pub use cuda::*;
 #[cfg(test)]
 mod tests;
 
-#[derive(Default, AlignedBorrow, Copy, Clone)]
+#[derive(AlignedBorrow, Copy, Clone)]
 #[repr(C)]
-pub struct BitwiseOperationLookupCols<T> {
+pub struct BitwiseOperationLookupCols<T, const NUM_BITS: usize> {
+    /// Binary decomposition of x (x_bits[0] is LSB, x_bits[NUM_BITS-1] is MSB)
+    pub x_bits: [T; NUM_BITS],
+    /// Binary decomposition of y (y_bits[0] is LSB, y_bits[NUM_BITS-1] is MSB)
+    pub y_bits: [T; NUM_BITS],
     /// Number of range check operations requested for each (x, y) pair
     pub mult_range: T,
     /// Number of XOR operations requested for each (x, y) pair
     pub mult_xor: T,
 }
 
-#[derive(Default, AlignedBorrow, Copy, Clone)]
-#[repr(C)]
-pub struct BitwiseOperationLookupPreprocessedCols<T> {
-    pub x: T,
-    pub y: T,
-    /// XOR result of x and y (x ⊕ y)
-    pub z_xor: T,
-}
-
-pub const NUM_BITWISE_OP_LOOKUP_COLS: usize = size_of::<BitwiseOperationLookupCols<u8>>();
-pub const NUM_BITWISE_OP_LOOKUP_PREPROCESSED_COLS: usize =
-    size_of::<BitwiseOperationLookupPreprocessedCols<u8>>();
+/// Number of multiplicity columns (mult_range and mult_xor)
+pub const NUM_BITWISE_OP_LOOKUP_MULT_COLS: usize = 2;
 
 #[derive(Clone, Copy, Debug, derive_new::new)]
 pub struct BitwiseOperationLookupAir<const NUM_BITS: usize> {
@@ -64,52 +57,92 @@ impl<F: Field, const NUM_BITS: usize> PartitionedBaseAir<F>
 }
 impl<F: Field, const NUM_BITS: usize> BaseAir<F> for BitwiseOperationLookupAir<NUM_BITS> {
     fn width(&self) -> usize {
-        NUM_BITWISE_OP_LOOKUP_COLS
-    }
-
-    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
-        let rows: Vec<F> = (0..(1 << NUM_BITS))
-            .flat_map(|x: u32| {
-                (0..(1 << NUM_BITS)).flat_map(move |y: u32| {
-                    [
-                        F::from_canonical_u32(x),
-                        F::from_canonical_u32(y),
-                        F::from_canonical_u32(x ^ y),
-                    ]
-                })
-            })
-            .collect();
-        Some(RowMajorMatrix::new(
-            rows,
-            NUM_BITWISE_OP_LOOKUP_PREPROCESSED_COLS,
-        ))
+        BitwiseOperationLookupCols::<F, NUM_BITS>::width()
     }
 }
 
-impl<AB: InteractionBuilder + PairBuilder, const NUM_BITS: usize> Air<AB>
+impl<AB: InteractionBuilder, const NUM_BITS: usize> Air<AB>
     for BitwiseOperationLookupAir<NUM_BITS>
 {
     fn eval(&self, builder: &mut AB) {
-        let preprocessed = builder.preprocessed();
-        let prep_local = preprocessed.row_slice(0);
-        let prep_local: &BitwiseOperationLookupPreprocessedCols<AB::Var> = (*prep_local).borrow();
-
         let main = builder.main();
-        let local = main.row_slice(0);
-        let local: &BitwiseOperationLookupCols<AB::Var> = (*local).borrow();
+        let (local, next) = (main.row_slice(0), main.row_slice(1));
+        let local: &BitwiseOperationLookupCols<AB::Var, NUM_BITS> = (*local).borrow();
+        let next: &BitwiseOperationLookupCols<AB::Var, NUM_BITS> = (*next).borrow();
+
+        // 1. Binary constraints: ensure each bit is boolean
+        for i in 0..NUM_BITS {
+            builder.assert_bool(local.x_bits[i]);
+            builder.assert_bool(local.y_bits[i]);
+        }
+
+        // 2. Reconstruct x and y from their binary decompositions
+        // x = Σ(x_bits[i] * 2^i), y = Σ(y_bits[i] * 2^i)
+        let reconstruct = |bits: &[AB::Var; NUM_BITS]| {
+            bits.iter()
+                .enumerate()
+                .fold(AB::Expr::ZERO, |acc, (i, &bit)| {
+                    acc + bit * AB::Expr::from_canonical_usize(1 << i)
+                })
+        };
+        let x_reconstructed = reconstruct(&local.x_bits);
+        let y_reconstructed = reconstruct(&local.y_bits);
+
+        // 3. Compute z_xor algebraically from bits
+        // z_xor_bits[i] = x_bits[i] ^ y_bits[i] = x_bits[i] + y_bits[i] - 2 * x_bits[i] * y_bits[i]
+        // z_xor = Σ(z_xor_bits[i] * 2^i)
+        let z_xor_reconstructed = local
+            .x_bits
+            .iter()
+            .zip(local.y_bits.iter())
+            .enumerate()
+            .fold(AB::Expr::ZERO, |acc, (i, (&x_bit, &y_bit))| {
+                let xor_bit = x_bit + y_bit - AB::Expr::TWO * x_bit * y_bit;
+                acc + xor_bit * AB::Expr::from_canonical_usize(1 << i)
+            });
+
+        // 4. Combined index: idx = x * (2^NUM_BITS) + y
+        let combined_idx = x_reconstructed.clone() * AB::Expr::from_canonical_usize(1 << NUM_BITS)
+            + y_reconstructed.clone();
+        let next_combined_idx = reconstruct(&next.x_bits)
+            * AB::Expr::from_canonical_usize(1 << NUM_BITS)
+            + reconstruct(&next.y_bits);
+
+        // 5. Constrain that combined index increments by 1 each row
+        builder
+            .when_transition()
+            .assert_one(next_combined_idx.clone() - combined_idx.clone());
+
+        // 6. Boundary constraints: first row has idx = 0, last row has idx = 2^(2*NUM_BITS) - 1
+        builder.when_first_row().assert_zero(combined_idx.clone());
+        builder.when_last_row().assert_eq(
+            combined_idx,
+            AB::Expr::from_canonical_usize((1 << (2 * NUM_BITS)) - 1),
+        );
 
+        // 7. Use reconstructed values for lookup bus interactions
         self.bus
-            .receive(prep_local.x, prep_local.y, AB::F::ZERO, AB::F::ZERO)
+            .receive(
+                x_reconstructed.clone(),
+                y_reconstructed.clone(),
+                AB::F::ZERO,
+                AB::F::ZERO,
+            )
             .eval(builder, local.mult_range);
         self.bus
-            .receive(prep_local.x, prep_local.y, prep_local.z_xor, AB::F::ONE)
+            .receive(
+                x_reconstructed,
+                y_reconstructed,
+                z_xor_reconstructed,
+                AB::F::ONE,
+            )
             .eval(builder, local.mult_xor);
     }
 }
 
-// Lookup chip for operations on size NUM_BITS integers. Currently has pre-processed columns
-// for x ^ y and range check. Interactions are of form [x, y, z] where z is either x ^ y for
-// XOR or 0 for range check.
+// Lookup chip for operations on size NUM_BITS integers. Uses gate-based constraints
+// with binary decomposition instead of preprocessed trace. Interactions are of form [x, y, z]
+// where z is either x ^ y for XOR or 0 for range check.
 
 pub struct BitwiseOperationLookupChip<const NUM_BITS: usize> {
     pub air: BitwiseOperationLookupAir<NUM_BITS>,
@@ -137,7 +170,7 @@ impl<const NUM_BITS: usize> BitwiseOperationLookupChip<NUM_BITS> {
     }
 
     pub fn air_width(&self) -> usize {
-        NUM_BITWISE_OP_LOOKUP_COLS
+        BitwiseOperationLookupCols::<u8, NUM_BITS>::width()
     }
 
     pub fn request_range(&self, x: u32, y: u32) {
@@ -164,17 +197,33 @@ impl<const NUM_BITS: usize> BitwiseOperationLookupChip<NUM_BITS> {
 
     /// Generates trace and resets all internal counters to 0.
     pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
-        let mut rows = F::zero_vec(self.count_range.len() * NUM_BITWISE_OP_LOOKUP_COLS);
-        for (n, row) in rows.chunks_mut(NUM_BITWISE_OP_LOOKUP_COLS).enumerate() {
-            let cols: &mut BitwiseOperationLookupCols<F> = row.borrow_mut();
+        let num_cols = BitwiseOperationLookupCols::<F, NUM_BITS>::width();
+        let num_rows = (1 << NUM_BITS) * (1 << NUM_BITS);
+        let mut rows = F::zero_vec(num_rows * num_cols);
+
+        for (n, row) in rows.chunks_mut(num_cols).enumerate() {
+            let cols: &mut BitwiseOperationLookupCols<F, NUM_BITS> = row.borrow_mut();
+
+            // Compute x and y from row index: row n corresponds to (x, y) where
+            // x = n / (2^NUM_BITS), y = n % (2^NUM_BITS)
+            let x = (n / (1 << NUM_BITS)) as u32;
+            let y = (n % (1 << NUM_BITS)) as u32;
+
+            // Set x_bits and y_bits: decompose x and y into binary
+            for i in 0..NUM_BITS {
+                cols.x_bits[i] = F::from_canonical_u32((x >> i) & 1);
+                cols.y_bits[i] = F::from_canonical_u32((y >> i) & 1);
+            }
+
+            // Set multiplicities
             cols.mult_range = F::from_canonical_u32(
                 self.count_range[n].swap(0, std::sync::atomic::Ordering::SeqCst),
             );
             cols.mult_xor = F::from_canonical_u32(
                 self.count_xor[n].swap(0, std::sync::atomic::Ordering::SeqCst),
             );
         }
-        RowMajorMatrix::new(rows, NUM_BITWISE_OP_LOOKUP_COLS)
+        RowMajorMatrix::new(rows, num_cols)
     }
 
     fn idx(x: u32, y: u32) -> usize {
@@ -203,6 +252,6 @@ impl<const NUM_BITS: usize> ChipUsageGetter for BitwiseOperationLookupChip<NUM_B
         1 << (2 * NUM_BITS)
     }
     fn trace_width(&self) -> usize {
-        NUM_BITWISE_OP_LOOKUP_COLS
+        BitwiseOperationLookupCols::<u8, NUM_BITS>::width()
     }
 }