openvm-cuda-backend 2.0.0

//! Batched monomial-based MLE evaluation for zerocheck and logup.
//!
//! This module provides batch evaluators for monomial evaluations across multiple AIRs,
//! enabling efficient GPU kernel launches that process multiple traces in a single launch.

use openvm_cuda_common::{
    copy::MemCopyH2D,
    d_buffer::DeviceBuffer,
    error::{CudaError, MemCopyError},
    stream::GpuDeviceCtx,
};
use openvm_stark_backend::prover::{fractional_sumcheck_gkr::Frac, DeviceMultiStarkProvingKey};
use p3_field::PrimeCharacteristicRing;
use tracing::debug;

use crate::{
    cuda::logup_zerocheck::{
        logup_monomial_batched, precompute_lambda_combinations,
        precompute_logup_denom_combinations, precompute_logup_numer_combinations,
        zerocheck_monomial_batched, zerocheck_monomial_par_y_batched, BlockCtx, EvalCoreCtx,
        LogupMonomialCommonCtx, LogupMonomialCtx, MonomialAirCtx,
    },
    error::KernelError,
    gpu_backend::GenericGpuBackend,
    hash_scheme::GpuHashScheme,
    logup_zerocheck::{batch_mle::TraceCtx, block_ctxs::build_block_ctxs},
    prelude::EF,
};

const THREADS_PER_BLOCK: u32 = 256;

/// Returns true if the trace can use the monomial evaluation path.
///
/// A trace is eligible if it has constraints and the AIR has expanded monomials.
pub(crate) fn trace_has_monomials<HS: GpuHashScheme>(
    trace: &TraceCtx,
    pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
) -> bool {
    trace.has_constraints
        && pk.per_air[trace.air_idx]
            .other_data
            .zerocheck_monomials
            .as_ref()
            .map(|m| m.num_monomials > 0)
            .unwrap_or(false)
}

/// Get the number of monomials for a trace. Returns 0 if the trace has no monomials.
pub(crate) fn get_num_monomials<HS: GpuHashScheme>(
    trace: &TraceCtx,
    pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
) -> u32 {
    pk.per_air[trace.air_idx]
        .other_data
        .zerocheck_monomials
        .as_ref()
        .map(|m| m.num_monomials)
        .unwrap_or(0)
}

/// Get the rules_len for a trace's zerocheck DAG.
pub(crate) fn get_zerocheck_rules_len<HS: GpuHashScheme>(
    trace: &TraceCtx,
    pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
) -> usize {
    pk.per_air[trace.air_idx]
        .other_data
        .zerocheck_mle
        .inner
        .d_rules
        .len()
}

/// Precompute lambda combinations for a single AIR's monomials.
///
/// Returns a buffer of length `num_monomials` where each element is
/// `sum_l(coefficient_l * lambda_pows[constraint_idx_l])` for that monomial.
///
/// The AIR must have nonempty monomials.
pub(crate) fn compute_lambda_combinations<HS: GpuHashScheme>(
    pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
    air_idx: usize,
    lambda_pows: &DeviceBuffer<EF>,
    device_ctx: &GpuDeviceCtx,
) -> Result<DeviceBuffer<EF>, CudaError> {
    let monomials = pk.per_air[air_idx]
        .other_data
        .zerocheck_monomials
        .as_ref()
        .expect("AIR must have monomials");
    let mut buf =
        DeviceBuffer::<EF>::with_capacity_on(monomials.num_monomials as usize, device_ctx);
    unsafe {
        precompute_lambda_combinations(
            &mut buf,
            monomials.d_headers.as_ptr(),
            monomials.d_lambda_terms.as_ptr(),
            lambda_pows,
            monomials.num_monomials,
            device_ctx.stream.as_raw(),
        )?;
    }
    Ok(buf)
}

/// Batch evaluator for monomial-based zerocheck MLE evaluation.
///
/// Pre-builds GPU contexts for all traces, then evaluates in a single kernel launch.
///
/// The caller must filter traces using [`trace_has_monomials`] before constructing.
/// The batch must contain at least one trace.
///
/// The struct holds references to `TraceCtx` which guarantees the underlying
/// device buffers (including `main_ptrs_dev`) remain valid for the struct's lifetime.
pub(crate) struct ZerocheckMonomialBatch<'a> {
    traces: Vec<&'a TraceCtx>,
    block_ctxs: DeviceBuffer<BlockCtx>,
    air_ctxs: DeviceBuffer<MonomialAirCtx>,
    air_offsets: DeviceBuffer<u32>,
    /// Cheap clone: just `(device_id, Arc<CudaStream>)`.
    device_ctx: GpuDeviceCtx,
}

impl<'a> ZerocheckMonomialBatch<'a> {
    /// Creates a new batch from an iterator of traces.
    ///
    /// `lambda_combinations` must contain one buffer per trace (in iteration order),
    /// each precomputed via [`compute_lambda_combinations`].
    ///
    /// # Panics
    ///
    /// Panics if `traces` is empty or if `lambda_combinations` length doesn't match.
    pub fn new<HS: GpuHashScheme>(
        traces: impl IntoIterator<Item = &'a TraceCtx>,
        pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
        lambda_combinations: &[&DeviceBuffer<EF>],
        device_ctx: &GpuDeviceCtx,
    ) -> Result<Self, MemCopyError> {
        let traces: Vec<_> = traces.into_iter().collect();
        assert!(
            !traces.is_empty(),
            "ZerocheckMonomialBatch requires at least one trace"
        );
        assert_eq!(
            traces.len(),
            lambda_combinations.len(),
            "lambda_combinations must have one buffer per trace"
        );

        let threads_per_block = THREADS_PER_BLOCK;

        // Block layout: each AIR contributes `mono_blocks * num_y` blocks, where
        // `mono_blocks = ceil(num_monomials / tpb)`. `local_block_idx_x` runs `0..total_blocks`
        // and the kernel decodes it as `y_int * mono_blocks + mono_block_idx`.
        let (block_ctxs_h, air_offsets) = build_block_ctxs(traces.iter().map(|t| {
            let monomials = pk.per_air[t.air_idx]
                .other_data
                .zerocheck_monomials
                .as_ref()
                .unwrap();
            let mono_blocks = monomials.num_monomials.div_ceil(threads_per_block);
            mono_blocks * t.num_y
        }));

        // Build MonomialAirCtx for each trace
        let air_ctxs_h: Vec<MonomialAirCtx> = traces
            .iter()
            .zip(lambda_combinations)
            .map(|(t, lc)| {
                let monomials = pk.per_air[t.air_idx]
                    .other_data
                    .zerocheck_monomials
                    .as_ref()
                    .unwrap();

                let eval_ctx = EvalCoreCtx {
                    d_selectors: t.sels_ptr,
                    d_preprocessed: t.prep_ptr,
                    d_main: t.main_ptrs_dev.as_ptr(),
                    d_public: t.public_ptr,
                };

                MonomialAirCtx {
                    d_headers: monomials.d_headers.as_ptr(),
                    d_variables: monomials.d_variables.as_ptr(),
                    d_lambda_combinations: lc.as_ptr(),
                    num_monomials: monomials.num_monomials,
                    eval_ctx,
                    d_eq_xi: t.eq_xi_ptr,
                    num_y: t.num_y,
                }
            })
            .collect();

        // Upload to device
        let block_ctxs = block_ctxs_h.to_device_on(device_ctx)?;
        let air_ctxs = air_ctxs_h.to_device_on(device_ctx)?;
        let air_offsets = air_offsets.to_device_on(device_ctx)?;

        debug!(
            num_airs = traces.len(),
            num_blocks = block_ctxs.len(),
            "ZerocheckMonomialBatch created"
        );

        Ok(Self {
            traces,
            block_ctxs,
            air_ctxs,
            air_offsets,
            device_ctx: device_ctx.clone(),
        })
    }

    /// Returns the trace indices in order.
    pub fn trace_indices(&self) -> impl Iterator<Item = usize> + '_ {
        self.traces.iter().map(|t| t.trace_idx)
    }

    /// Evaluates the batch and returns the output device buffer.
    ///
    /// The buffer contains `num_airs * num_x` elements, laid out as
    /// `[air0_x0, air0_x1, ..., air1_x0, air1_x1, ...]`.
    /// See [`crate::logup_zerocheck`] module docs for async-free/peak memory behavior.
    pub fn evaluate(&self, num_x: u32) -> Result<DeviceBuffer<EF>, KernelError> {
        let num_blocks = self.block_ctxs.len();
        let num_airs = self.air_ctxs.len();

        debug!(
            %num_blocks,
            %num_x,
            %num_airs,
            "zerocheck_monomial_batched"
        );

        let mut tmp_sums =
            DeviceBuffer::<EF>::with_capacity_on(num_blocks * num_x as usize, &self.device_ctx);
        let mut output =
            DeviceBuffer::<EF>::with_capacity_on(num_airs * num_x as usize, &self.device_ctx);

        debug_assert_eq!(
            self.air_offsets.len(),
            num_airs + 1,
            "air_offsets must have num_airs + 1 elements"
        );
        // SAFETY: All device pointers in block_ctxs and air_ctxs were constructed from
        // valid DeviceBuffers that outlive this call (TraceCtx references, pk monomial data,
        // lambda_combinations). The air_offsets buffer has length num_airs + 1 as required.
        unsafe {
            zerocheck_monomial_batched(
                &mut tmp_sums,
                &mut output,
                &self.block_ctxs,
                &self.air_ctxs,
                &self.air_offsets,
                num_blocks as u32,
                num_x,
                num_airs as u32,
                THREADS_PER_BLOCK,
                self.device_ctx.stream.as_raw(),
            )?;
        }

        Ok(output)
    }
}

// Constants for par-y kernel
const THREADS_PER_BLOCK_PAR_Y: u32 = 128;
const DEFAULT_MAX_MONOMIALS_PER_THREAD: u32 = 64;
const WAVES_TARGET: u32 = 4;

/// Batch evaluator for monomial-based zerocheck MLE evaluation, parallelizing over y_int.
///
/// This variant is optimized for traces with high `num_y`: each thread handles one y_int
/// and loops over a chunk of monomials. The chunk size is auto-tuned based on SM count.
///
/// The caller must filter traces using [`trace_has_monomials`] before constructing.
/// The batch must contain at least one trace.
pub(crate) struct ZerocheckMonomialParYBatch<'a> {
    traces: Vec<&'a TraceCtx>,
    block_ctxs: DeviceBuffer<BlockCtx>,
    air_ctxs: DeviceBuffer<MonomialAirCtx>,
    air_offsets: DeviceBuffer<u32>,
    num_blocks: u32,
    chunk_size: u32,
    /// Cheap clone: just `(device_id, Arc<CudaStream>)`.
    device_ctx: GpuDeviceCtx,
}

impl<'a> ZerocheckMonomialParYBatch<'a> {
    /// Creates a new batch from an iterator of traces.
    ///
    /// `lambda_combinations` must contain one buffer per trace (in iteration order),
    /// each precomputed via [`compute_lambda_combinations`].
    ///
    /// The `sm_count` and `num_x` parameters are used to auto-tune the chunk size
    /// for optimal SM utilization.
    ///
    /// # Panics
    ///
    /// Panics if `traces` is empty or if `lambda_combinations` length doesn't match.
    #[allow(clippy::too_many_arguments)]
    pub fn new<HS: GpuHashScheme>(
        traces: impl IntoIterator<Item = &'a TraceCtx>,
        pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
        lambda_combinations: &[&DeviceBuffer<EF>],
        sm_count: u32,
        num_x: u32,
        max_monomials_per_thread: Option<u32>,
        device_ctx: &GpuDeviceCtx,
    ) -> Result<Self, MemCopyError> {
        let traces: Vec<_> = traces.into_iter().collect();
        assert!(
            !traces.is_empty(),
            "ZerocheckMonomialParYBatch requires at least one trace"
        );
        assert_eq!(
            traces.len(),
            lambda_combinations.len(),
            "lambda_combinations must have one buffer per trace"
        );

        let threads_per_block = THREADS_PER_BLOCK_PAR_Y;
        let max_mono_per_thread =
            max_monomials_per_thread.unwrap_or(DEFAULT_MAX_MONOMIALS_PER_THREAD);

        // First pass: collect per-AIR info
        let mut per_air_info: Vec<(u32, u32)> = Vec::new(); // (y_blocks, num_monomials)
        let mut max_monomials = 0u32;

        for t in traces.iter() {
            let monomials = pk.per_air[t.air_idx]
                .other_data
                .zerocheck_monomials
                .as_ref()
                .expect("AIR with constraints must have monomials");

            let y_blocks = t.num_y.div_ceil(threads_per_block);
            max_monomials = max_monomials.max(monomials.num_monomials);
            per_air_info.push((y_blocks, monomials.num_monomials));
        }

        // Determine chunk_size based on SM utilization and cap
        // We want: total_blocks * num_x >= sm_count * WAVES_TARGET
        // total_blocks = sum of (y_blocks * air_mono_chunks) per AIR
        // where air_mono_chunks = ceil(num_monomials / chunk_size)
        //
        // Also: chunk_size <= max_mono_per_thread
        //
        // Start with chunk_size = max_mono_per_thread and adjust if needed
        let target_blocks = sm_count * WAVES_TARGET;

        // Initial estimate: use max_mono_per_thread as chunk_size
        let mut chunk_size = max_mono_per_thread;
        loop {
            let total_blocks: u32 = per_air_info
                .iter()
                .map(|(y_blocks, num_mono)| {
                    let air_mono_chunks = num_mono.div_ceil(chunk_size);
                    y_blocks * air_mono_chunks
                })
                .sum();

            if total_blocks * num_x >= target_blocks || chunk_size <= 1 {
                break;
            }
            // Need more blocks: reduce chunk_size
            chunk_size = (chunk_size / 2).max(1);
        }

        // Block layout: each AIR contributes `y_blocks * air_mono_chunks` blocks.
        // `local_block_idx_x` runs `0..total_blocks` and the kernel decodes it as
        // `y_block * air_mono_chunks + mono_chunk`.
        let (block_ctxs_h, air_offsets) =
            build_block_ctxs(per_air_info.iter().map(|(y_blocks, num_mono)| {
                let air_mono_chunks = num_mono.div_ceil(chunk_size);
                y_blocks * air_mono_chunks
            }));

        let num_blocks = block_ctxs_h.len() as u32;

        // Build MonomialAirCtx for each trace
        let air_ctxs_h: Vec<MonomialAirCtx> = traces
            .iter()
            .zip(lambda_combinations)
            .map(|(t, lc)| {
                let monomials = pk.per_air[t.air_idx]
                    .other_data
                    .zerocheck_monomials
                    .as_ref()
                    .unwrap();

                let eval_ctx = EvalCoreCtx {
                    d_selectors: t.sels_ptr,
                    d_preprocessed: t.prep_ptr,
                    d_main: t.main_ptrs_dev.as_ptr(),
                    d_public: t.public_ptr,
                };

                MonomialAirCtx {
                    d_headers: monomials.d_headers.as_ptr(),
                    d_variables: monomials.d_variables.as_ptr(),
                    d_lambda_combinations: lc.as_ptr(),
                    num_monomials: monomials.num_monomials,
                    eval_ctx,
                    d_eq_xi: t.eq_xi_ptr,
                    num_y: t.num_y,
                }
            })
            .collect();

        // Upload to device
        let block_ctxs = block_ctxs_h.to_device_on(device_ctx)?;
        let air_ctxs = air_ctxs_h.to_device_on(device_ctx)?;
        let air_offsets = air_offsets.to_device_on(device_ctx)?;

        debug!(
            num_airs = traces.len(),
            num_blocks, chunk_size, max_monomials, "ZerocheckMonomialParYBatch created"
        );

        Ok(Self {
            traces,
            block_ctxs,
            air_ctxs,
            air_offsets,
            num_blocks,
            chunk_size,
            device_ctx: device_ctx.clone(),
        })
    }

    /// Returns the trace indices in order.
    pub fn trace_indices(&self) -> impl Iterator<Item = usize> + '_ {
        self.traces.iter().map(|t| t.trace_idx)
    }

    /// Evaluates the batch and returns the output device buffer.
    ///
    /// The buffer contains `num_airs * num_x` elements, laid out as
    /// `[air0_x0, air0_x1, ..., air1_x0, air1_x1, ...]`.
    /// See [`crate::logup_zerocheck`] module docs for async-free/peak memory behavior.
    pub fn evaluate(&self, num_x: u32) -> Result<DeviceBuffer<EF>, KernelError> {
        let num_airs = self.air_ctxs.len();

        debug!(
            num_blocks = %self.num_blocks,
            %num_x,
            %num_airs,
            chunk_size = %self.chunk_size,
            "zerocheck_monomial_par_y_batched"
        );

        let mut tmp_sums = DeviceBuffer::<EF>::with_capacity_on(
            self.num_blocks as usize * num_x as usize,
            &self.device_ctx,
        );
        let mut output =
            DeviceBuffer::<EF>::with_capacity_on(num_airs * num_x as usize, &self.device_ctx);

        debug_assert_eq!(
            self.air_offsets.len(),
            num_airs + 1,
            "air_offsets must have num_airs + 1 elements"
        );
        // SAFETY: All device pointers in block_ctxs and air_ctxs were constructed from
        // valid DeviceBuffers that outlive this call (TraceCtx references, pk monomial data,
        // lambda_combinations). The air_offsets buffer has length num_airs + 1 as required.
        unsafe {
            zerocheck_monomial_par_y_batched(
                &mut tmp_sums,
                &mut output,
                &self.block_ctxs,
                &self.air_ctxs,
                &self.air_offsets,
                self.num_blocks,
                num_x,
                num_airs as u32,
                self.chunk_size,
                THREADS_PER_BLOCK_PAR_Y,
                self.device_ctx.stream.as_raw(),
            )?;
        }

        Ok(output)
    }
}

// ============================================================================
// LOGUP MONOMIAL EVALUATION
// ============================================================================

/// Precomputed logup combinations for a single AIR.
pub struct LogupCombinations {
    pub d_numer_combinations: DeviceBuffer<EF>,
    pub d_denom_combinations: DeviceBuffer<EF>,
    pub bus_term_sum: EF,
}

/// Precompute logup combinations for a single AIR's interaction monomials.
///
/// The AIR must have nonempty interaction monomials.
#[allow(clippy::too_many_arguments)]
pub(crate) fn compute_logup_combinations<HS: GpuHashScheme>(
    pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
    air_idx: usize,
    d_beta_pows: &DeviceBuffer<EF>,
    d_eq_3bs: &DeviceBuffer<EF>,
    eq_3bs_host: &[EF],
    beta_pows_host: &[EF],
    device_ctx: &GpuDeviceCtx,
) -> Result<LogupCombinations, CudaError> {
    let monomials = pk.per_air[air_idx]
        .other_data
        .interaction_monomials
        .as_ref()
        .expect("AIR must have interaction monomials");

    let stream = device_ctx.stream.as_raw();
    // Precompute numerator combinations: sum_i(coeff_i * eq_3bs[interaction_idx_i])
    let mut d_numer_combinations = if monomials.num_numer_monomials > 0 {
        DeviceBuffer::<EF>::with_capacity_on(monomials.num_numer_monomials as usize, device_ctx)
    } else {
        DeviceBuffer::new()
    };
    if monomials.num_numer_monomials > 0 {
        unsafe {
            precompute_logup_numer_combinations(
                &mut d_numer_combinations,
                monomials.d_numer_headers.as_ptr(),
                monomials.d_numer_terms.as_ptr(),
                d_eq_3bs,
                monomials.num_numer_monomials,
                stream,
            )?;
        }
    }

    // Precompute denominator combinations: sum_i(coeff_i * beta_pows[field_idx_i] *
    // eq_3bs[interaction_idx_i])
    let mut d_denom_combinations = if monomials.num_denom_monomials > 0 {
        DeviceBuffer::<EF>::with_capacity_on(monomials.num_denom_monomials as usize, device_ctx)
    } else {
        DeviceBuffer::new()
    };
    if monomials.num_denom_monomials > 0 {
        unsafe {
            precompute_logup_denom_combinations(
                &mut d_denom_combinations,
                monomials.d_denom_headers.as_ptr(),
                monomials.d_denom_terms.as_ptr(),
                d_beta_pows,
                d_eq_3bs,
                monomials.num_denom_monomials,
                stream,
            )?;
        }
    }

    // Compute bus_term_sum on CPU: sum_i(beta_pows[message_len_i] * (bus_idx[i]+1) * eq_3bs[i])
    let interactions = &pk.per_air[air_idx].vk.symbolic_constraints.interactions;
    debug_assert_eq!(
        interactions.len(),
        eq_3bs_host.len(),
        "interaction count must match eq_3bs"
    );
    let mut bus_term_sum = EF::ZERO;
    for (i, interaction) in interactions.iter().enumerate() {
        let beta_len = beta_pows_host[interaction.message.len()];
        let bus_idx = interaction.bus_index as u32;
        bus_term_sum += beta_len * EF::from_u32(bus_idx + 1) * eq_3bs_host[i];
    }

    Ok(LogupCombinations {
        d_numer_combinations,
        d_denom_combinations,
        bus_term_sum,
    })
}

const THREADS_PER_BLOCK_LOGUP: u32 = 128;

/// Batch evaluator for logup monomial MLE evaluation.
///
/// Each block evaluates a monomial chunk for a y_int, producing a FracExt output
/// compatible with standard reduction.
pub(crate) struct LogupMonomialBatch<'a> {
    traces: Vec<&'a TraceCtx>,
    block_ctxs: DeviceBuffer<BlockCtx>,
    common_ctxs: DeviceBuffer<LogupMonomialCommonCtx>,
    numer_ctxs: DeviceBuffer<LogupMonomialCtx>,
    denom_ctxs: DeviceBuffer<LogupMonomialCtx>,
    air_offsets: DeviceBuffer<u32>,
    num_blocks: u32,
    device_ctx: GpuDeviceCtx,
}

impl<'a> LogupMonomialBatch<'a> {
    /// Creates a new batch from an iterator of traces.
    ///
    /// `logup_combinations` must contain one `LogupCombinations` per trace (in iteration order),
    /// each precomputed via [`compute_logup_combinations`].
    ///
    /// # Panics
    ///
    /// Panics if `traces` is empty or if `logup_combinations` length doesn't match.
    pub fn new<HS: GpuHashScheme>(
        traces: impl IntoIterator<Item = &'a TraceCtx>,
        pk: &DeviceMultiStarkProvingKey<GenericGpuBackend<HS>>,
        logup_combinations: &[&LogupCombinations],
        device_ctx: &GpuDeviceCtx,
    ) -> Result<Self, MemCopyError> {
        let traces: Vec<_> = traces.into_iter().collect();
        assert!(
            !traces.is_empty(),
            "LogupMonomialBatch requires at least one trace"
        );
        assert_eq!(
            traces.len(),
            logup_combinations.len(),
            "logup_combinations must have one entry per trace"
        );

        let threads_per_block = THREADS_PER_BLOCK_LOGUP;

        // Block layout: each AIR contributes `num_y * mono_blocks` blocks, where `mono_blocks`
        // covers the larger of numerator/denominator monomial counts.
        // `local_block_idx_x` runs `0..total_blocks` and the kernel decodes it as
        // `y_int * mono_blocks + mono_block`.
        let (block_ctxs_h, air_offsets) = build_block_ctxs(traces.iter().map(|t| {
            let monomials = pk.per_air[t.air_idx]
                .other_data
                .interaction_monomials
                .as_ref()
                .unwrap();
            let max_monomials = monomials
                .num_numer_monomials
                .max(monomials.num_denom_monomials);
            let mono_blocks = max_monomials.div_ceil(threads_per_block).max(1);
            t.num_y * mono_blocks
        }));

        let num_blocks = block_ctxs_h.len() as u32;

        // Build logup monomial ctxs for each trace
        let common_ctxs_h: Vec<LogupMonomialCommonCtx> = traces
            .iter()
            .zip(logup_combinations)
            .map(|(t, lc)| {
                let monomials = pk.per_air[t.air_idx]
                    .other_data
                    .interaction_monomials
                    .as_ref()
                    .unwrap();
                let max_monomials = monomials
                    .num_numer_monomials
                    .max(monomials.num_denom_monomials);
                let mono_blocks = max_monomials.div_ceil(threads_per_block).max(1);

                let eval_ctx = EvalCoreCtx {
                    d_selectors: t.sels_ptr,
                    d_preprocessed: t.prep_ptr,
                    d_main: t.main_ptrs_dev.as_ptr(),
                    d_public: t.public_ptr,
                };

                LogupMonomialCommonCtx {
                    eval_ctx,
                    d_eq_xi: t.eq_xi_ptr,
                    bus_term_sum: lc.bus_term_sum,
                    num_y: t.num_y,
                    mono_blocks,
                }
            })
            .collect();
        let numer_ctxs_h: Vec<LogupMonomialCtx> = traces
            .iter()
            .zip(logup_combinations)
            .map(|(t, lc)| {
                let monomials = pk.per_air[t.air_idx]
                    .other_data
                    .interaction_monomials
                    .as_ref()
                    .unwrap();
                LogupMonomialCtx {
                    d_headers: monomials.d_numer_headers.as_ptr(),
                    d_variables: monomials.d_numer_variables.as_ptr(),
                    d_combinations: lc.d_numer_combinations.as_ptr(),
                    num_monomials: monomials.num_numer_monomials,
                }
            })
            .collect();
        let denom_ctxs_h: Vec<LogupMonomialCtx> = traces
            .iter()
            .zip(logup_combinations)
            .map(|(t, lc)| {
                let monomials = pk.per_air[t.air_idx]
                    .other_data
                    .interaction_monomials
                    .as_ref()
                    .unwrap();
                LogupMonomialCtx {
                    d_headers: monomials.d_denom_headers.as_ptr(),
                    d_variables: monomials.d_denom_variables.as_ptr(),
                    d_combinations: lc.d_denom_combinations.as_ptr(),
                    num_monomials: monomials.num_denom_monomials,
                }
            })
            .collect();

        // Upload to device
        let block_ctxs = block_ctxs_h.to_device_on(device_ctx)?;
        let common_ctxs = common_ctxs_h.to_device_on(device_ctx)?;
        let numer_ctxs = numer_ctxs_h.to_device_on(device_ctx)?;
        let denom_ctxs = denom_ctxs_h.to_device_on(device_ctx)?;
        let air_offsets = air_offsets.to_device_on(device_ctx)?;

        debug!(
            num_airs = traces.len(),
            num_blocks, "LogupMonomialBatch created"
        );

        Ok(Self {
            traces,
            block_ctxs,
            common_ctxs,
            numer_ctxs,
            denom_ctxs,
            air_offsets,
            num_blocks,
            device_ctx: device_ctx.clone(),
        })
    }

    /// Returns the trace indices in order.
    pub fn trace_indices(&self) -> impl Iterator<Item = usize> + '_ {
        self.traces.iter().map(|t| t.trace_idx)
    }

    /// Evaluates the batch and returns the output device buffer.
    ///
    /// The buffer contains `num_airs * num_x` FracExt elements, laid out as
    /// `[air0_x0, air0_x1, ..., air1_x0, air1_x1, ...]`.
    /// See [`crate::logup_zerocheck`] module docs for async-free/peak memory behavior.
    pub fn evaluate(&self, num_x: u32) -> Result<DeviceBuffer<Frac<EF>>, KernelError> {
        let num_airs = self.common_ctxs.len();

        debug!(
            num_blocks = %self.num_blocks,
            %num_x,
            %num_airs,
            "logup_monomial_batched"
        );

        let mut tmp_sums = DeviceBuffer::<Frac<EF>>::with_capacity_on(
            self.num_blocks as usize * num_x as usize,
            &self.device_ctx,
        );
        let mut output =
            DeviceBuffer::<Frac<EF>>::with_capacity_on(num_airs * num_x as usize, &self.device_ctx);

        debug_assert_eq!(
            self.air_offsets.len(),
            num_airs + 1,
            "air_offsets must have num_airs + 1 elements"
        );

        // SAFETY: All device pointers were constructed from valid DeviceBuffers that outlive this
        // call.
        unsafe {
            logup_monomial_batched(
                &mut tmp_sums,
                &mut output,
                &self.block_ctxs,
                &self.common_ctxs,
                &self.numer_ctxs,
                &self.denom_ctxs,
                &self.air_offsets,
                self.num_blocks,
                num_x,
                num_airs as u32,
                THREADS_PER_BLOCK_LOGUP,
                self.device_ctx.stream.as_raw(),
            )?;
        }

        Ok(output)
    }
}