openvm_circuit_primitives/bitwise_op_lookup/
mod.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    sync::{atomic::AtomicU32, Arc},
4};
5
6use openvm_circuit_primitives_derive::AlignedBorrow;
7use openvm_cpu_backend::CpuBackend;
8use openvm_stark_backend::{
9    interaction::InteractionBuilder,
10    p3_air::{Air, AirBuilder, BaseAir},
11    p3_field::{Field, PrimeCharacteristicRing},
12    p3_matrix::{dense::RowMajorMatrix, Matrix},
13    prover::AirProvingContext,
14    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
15};
16
17use crate::{Chip, ColumnsAir, StructReflection, StructReflectionHelper};
18
19mod bus;
20pub use bus::*;
21
22#[cfg(feature = "cuda")]
23mod cuda;
24#[cfg(feature = "cuda")]
25pub use cuda::*;
26
27#[cfg(test)]
28mod tests;
29
30#[derive(AlignedBorrow, StructReflection, Copy, Clone)]
31#[repr(C)]
32pub struct BitwiseOperationLookupCols<T, const NUM_BITS: usize> {
33    /// Binary decomposition of x (`x_bits[0]` is LSB, `x_bits[NUM_BITS-1]` is MSB)
34    pub x_bits: [T; NUM_BITS],
35    /// Binary decomposition of y (`y_bits[0]` is LSB, `y_bits[NUM_BITS-1]` is MSB)
36    pub y_bits: [T; NUM_BITS],
37    /// Number of range check operations requested for each (x, y) pair
38    pub mult_range: T,
39    /// Number of XOR operations requested for each (x, y) pair
40    pub mult_xor: T,
41}
42
43/// Number of multiplicity columns (mult_range and mult_xor)
44pub const NUM_BITWISE_OP_LOOKUP_MULT_COLS: usize = 2;
45
46#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
47#[columns_via(BitwiseOperationLookupCols<u8, NUM_BITS>)]
48pub struct BitwiseOperationLookupAir<const NUM_BITS: usize> {
49    pub bus: BitwiseOperationLookupBus,
50}
51
52impl<F: Field, const NUM_BITS: usize> BaseAirWithPublicValues<F>
53    for BitwiseOperationLookupAir<NUM_BITS>
54{
55}
56impl<F: Field, const NUM_BITS: usize> PartitionedBaseAir<F>
57    for BitwiseOperationLookupAir<NUM_BITS>
58{
59}
60impl<F: Field, const NUM_BITS: usize> BaseAir<F> for BitwiseOperationLookupAir<NUM_BITS> {
61    fn width(&self) -> usize {
62        BitwiseOperationLookupCols::<F, NUM_BITS>::width()
63    }
64}
65
66impl<AB: InteractionBuilder, const NUM_BITS: usize> Air<AB>
67    for BitwiseOperationLookupAir<NUM_BITS>
68{
69    fn eval(&self, builder: &mut AB) {
70        let main = builder.main();
71        let (local, next) = (
72            main.row_slice(0).expect("window should have two elements"),
73            main.row_slice(1).expect("window should have two elements"),
74        );
75        let local: &BitwiseOperationLookupCols<AB::Var, NUM_BITS> = (*local).borrow();
76        let next: &BitwiseOperationLookupCols<AB::Var, NUM_BITS> = (*next).borrow();
77
78        // 1. Binary constraints: ensure each bit is boolean
79        for i in 0..NUM_BITS {
80            builder.assert_bool(local.x_bits[i]);
81            builder.assert_bool(local.y_bits[i]);
82        }
83
84        // 2. Reconstruct x and y from their binary decompositions
85        // x = Σ(x_bits[i] * 2^i), y = Σ(y_bits[i] * 2^i)
86        let reconstruct = |bits: &[AB::Var; NUM_BITS]| {
87            bits.iter()
88                .enumerate()
89                .fold(AB::Expr::ZERO, |acc, (i, &bit)| {
90                    acc + bit * AB::Expr::from_usize(1 << i)
91                })
92        };
93        let x_reconstructed = reconstruct(&local.x_bits);
94        let y_reconstructed = reconstruct(&local.y_bits);
95
96        // 3. Compute z_xor algebraically from bits
97        // z_xor_bits[i] = x_bits[i] ^ y_bits[i] = x_bits[i] + y_bits[i] - 2 * x_bits[i] * y_bits[i]
98        // z_xor = Σ(z_xor_bits[i] * 2^i)
99        let z_xor_reconstructed = local
100            .x_bits
101            .iter()
102            .zip(local.y_bits.iter())
103            .enumerate()
104            .fold(AB::Expr::ZERO, |acc, (i, (&x_bit, &y_bit))| {
105                let xor_bit = x_bit + y_bit - AB::Expr::TWO * x_bit * y_bit;
106                acc + xor_bit * AB::Expr::from_usize(1 << i)
107            });
108
109        // 4. Combined index: idx = x * (2^NUM_BITS) + y
110        let combined_idx =
111            x_reconstructed.clone() * AB::Expr::from_usize(1 << NUM_BITS) + y_reconstructed.clone();
112        let next_combined_idx = reconstruct(&next.x_bits) * AB::Expr::from_usize(1 << NUM_BITS)
113            + reconstruct(&next.y_bits);
114
115        // 5. Constrain that combined index increments by 1 each row
116        builder
117            .when_transition()
118            .assert_one(next_combined_idx.clone() - combined_idx.clone());
119
120        // 6. Boundary constraints: first row has idx = 0, last row has idx = 2^(2*NUM_BITS) - 1
121        builder.when_first_row().assert_zero(combined_idx.clone());
122        builder.when_last_row().assert_eq(
123            combined_idx,
124            AB::Expr::from_usize((1 << (2 * NUM_BITS)) - 1),
125        );
126
127        // 7. Use reconstructed values for lookup bus interactions
128        self.bus
129            .receive(
130                x_reconstructed.clone(),
131                y_reconstructed.clone(),
132                AB::F::ZERO,
133                AB::F::ZERO,
134            )
135            .eval(builder, local.mult_range);
136        self.bus
137            .receive(
138                x_reconstructed,
139                y_reconstructed,
140                z_xor_reconstructed,
141                AB::F::ONE,
142            )
143            .eval(builder, local.mult_xor);
144    }
145}
146
147// Lookup chip for operations on size NUM_BITS integers. Uses gate-based constraints
148// with binary decomposition instead of preprocessed trace. Interactions are of form [x, y, z]
149// where z is either x ^ y for XOR or 0 for range check.
150
151pub struct BitwiseOperationLookupChip<const NUM_BITS: usize> {
152    pub air: BitwiseOperationLookupAir<NUM_BITS>,
153    pub count_range: Vec<AtomicU32>,
154    pub count_xor: Vec<AtomicU32>,
155}
156
157pub type SharedBitwiseOperationLookupChip<const NUM_BITS: usize> =
158    Arc<BitwiseOperationLookupChip<NUM_BITS>>;
159
160impl<const NUM_BITS: usize> BitwiseOperationLookupChip<NUM_BITS> {
161    pub fn new(bus: BitwiseOperationLookupBus) -> Self {
162        let num_rows = (1 << NUM_BITS) * (1 << NUM_BITS);
163        let count_range = (0..num_rows).map(|_| AtomicU32::new(0)).collect();
164        let count_xor = (0..num_rows).map(|_| AtomicU32::new(0)).collect();
165        Self {
166            air: BitwiseOperationLookupAir::new(bus),
167            count_range,
168            count_xor,
169        }
170    }
171
172    pub fn bus(&self) -> BitwiseOperationLookupBus {
173        self.air.bus
174    }
175
176    pub fn air_width(&self) -> usize {
177        BitwiseOperationLookupCols::<u8, NUM_BITS>::width()
178    }
179
180    pub fn request_range(&self, x: u32, y: u32) {
181        let upper_bound = 1 << NUM_BITS;
182        debug_assert!(x < upper_bound, "x out of range: {x} >= {upper_bound}");
183        debug_assert!(y < upper_bound, "y out of range: {y} >= {upper_bound}");
184        self.count_range[Self::idx(x, y)].fetch_add(1, std::sync::atomic::Ordering::Relaxed);
185    }
186
187    pub fn request_xor(&self, x: u32, y: u32) -> u32 {
188        let upper_bound = 1 << NUM_BITS;
189        debug_assert!(x < upper_bound, "x out of range: {x} >= {upper_bound}");
190        debug_assert!(y < upper_bound, "y out of range: {y} >= {upper_bound}");
191        self.count_xor[Self::idx(x, y)].fetch_add(1, std::sync::atomic::Ordering::Relaxed);
192        x ^ y
193    }
194
195    pub fn clear(&self) {
196        for i in 0..self.count_range.len() {
197            self.count_range[i].store(0, std::sync::atomic::Ordering::Relaxed);
198            self.count_xor[i].store(0, std::sync::atomic::Ordering::Relaxed);
199        }
200    }
201
202    /// Generates trace and resets all internal counters to 0.
203    pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
204        let num_cols = BitwiseOperationLookupCols::<F, NUM_BITS>::width();
205        let num_rows = (1 << NUM_BITS) * (1 << NUM_BITS);
206        let mut rows = F::zero_vec(num_rows * num_cols);
207
208        for (n, row) in rows.chunks_mut(num_cols).enumerate() {
209            let cols: &mut BitwiseOperationLookupCols<F, NUM_BITS> = row.borrow_mut();
210
211            // Compute x and y from row index: row n corresponds to (x, y) where
212            // x = n / (2^NUM_BITS), y = n % (2^NUM_BITS)
213            let x = (n / (1 << NUM_BITS)) as u32;
214            let y = (n % (1 << NUM_BITS)) as u32;
215
216            // Set x_bits and y_bits: decompose x and y into binary
217            for i in 0..NUM_BITS {
218                cols.x_bits[i] = F::from_u32((x >> i) & 1);
219                cols.y_bits[i] = F::from_u32((y >> i) & 1);
220            }
221
222            // Set multiplicities
223            cols.mult_range =
224                F::from_u32(self.count_range[n].swap(0, std::sync::atomic::Ordering::SeqCst));
225            cols.mult_xor =
226                F::from_u32(self.count_xor[n].swap(0, std::sync::atomic::Ordering::SeqCst));
227        }
228        RowMajorMatrix::new(rows, num_cols)
229    }
230
231    fn idx(x: u32, y: u32) -> usize {
232        (x * (1 << NUM_BITS) + y) as usize
233    }
234}
235
236impl<R, SC: StarkProtocolConfig, const NUM_BITS: usize> Chip<R, CpuBackend<SC>>
237    for BitwiseOperationLookupChip<NUM_BITS>
238{
239    /// Generates trace and resets all internal counters to 0.
240    fn generate_proving_ctx(&self, _: R) -> AirProvingContext<CpuBackend<SC>> {
241        let trace_row_maj = self.generate_trace::<Val<SC>>();
242        AirProvingContext::simple_no_pis(trace_row_maj)
243    }
244
245    fn constant_trace_height(&self) -> Option<usize> {
246        Some((1 << NUM_BITS) * (1 << NUM_BITS))
247    }
248}
openvm_circuit_primitives/bitwise_op_lookup/mod.rs

openvm_circuit_primitives/bitwise_op_lookup/
mod.rs
