openvm-cuda-backend 2.0.0

//! Tests for the CUDA backend.
//!
//! Shared tests are generated by the [`openvm_backend_tests::backend_test_suite!`] macro using
//! the reference engine/backend.
//! GPU-specific tests (monomial vs DAG, stacked reduction) remain here.
//! Additional GPU-specific test cases for parameterized shared tests are added below.

use std::sync::Arc;

use itertools::Itertools;
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use openvm_cuda_common::copy::{MemCopyD2H, MemCopyH2D};
use openvm_cuda_common::{
    common::get_device,
    stream::{CudaStream, GpuDeviceCtx, StreamGuard},
};
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use openvm_stark_backend::{
    hasher::{Hasher as CpuMerkleHasher, MerkleHasher, MultiFieldHasher},
    multi_field_packing::pack_f_to_sf,
    p3_symmetric::TruncatedPermutation,
};
use openvm_stark_backend::{
    prover::{
        stacked_pcs::stacked_commit,
        stacked_reduction::{prove_stacked_opening_reduction, StackedReductionCpu},
        DeviceDataTransporter, MatrixDimensions, MultiRapProver,
    },
    test_utils::{default_test_params_small, FibFixture, InteractionsFixture11, TestFixture},
    verifier::stacked_reduction::{verify_stacked_reduction, StackedReductionError},
    FiatShamirTranscript, StarkEngine, StarkProtocolConfig,
};
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use openvm_stark_sdk::config::baby_bear_bn254_poseidon2::Bn254Scalar;
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use openvm_stark_sdk::config::bn254_poseidon2::{
    default_bn254_poseidon2_width2, default_bn254_poseidon2_width3,
};
use openvm_stark_sdk::{
    config::baby_bear_poseidon2::{
        default_duplex_sponge, BabyBearPoseidon2RefEngine, DuplexSponge,
    },
    utils::setup_tracing_with_log_level,
};
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use p3_field::BasedVectorSpace;
use p3_field::{PrimeCharacteristicRing, TwoAdicField};
#[cfg(feature = "baby-bear-bn254-poseidon2")]
use p3_symmetric::Permutation;
use test_case::test_case;
use tracing::{debug, Level};

#[cfg(feature = "baby-bear-bn254-poseidon2")]
use crate::cuda::bn254_merkle_tree::Bn254Digest;
use crate::{
    base::DeviceMatrix,
    cuda::{batch_ntt_small::batch_ntt_small, logup_zerocheck::frac_matrix_vertically_repeat},
    merkle_tree::MerkleTreeGpu,
    prelude::{EF, F, SC},
    sponge::DuplexSpongeGpu,
    BabyBearPoseidon2GpuEngine, GpuBackend,
};

type RefEngine = BabyBearPoseidon2RefEngine<DuplexSponge>;
type Engine = RefEngine;

fn test_ctx() -> GpuDeviceCtx {
    GpuDeviceCtx {
        device_id: get_device().unwrap() as u32,
        stream: StreamGuard::new(CudaStream::new_non_blocking().unwrap()),
    }
}

// ===========================================================================
// Shared test suite (engine-generic + WHIR)
// ===========================================================================

openvm_backend_tests::backend_test_suite!(Engine);

// ===========================================================================
// BN254 Poseidon2 engine end-to-end test
// ===========================================================================

/// End-to-end prove + verify using the BabyBear-BN254 Poseidon2 GPU engine.
///
/// Exercises the full BN254 pipeline: Merkle commitment, MultiField32Challenger
/// transcript, BN254 sponge grinding kernel, and proof verification.
#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test_case(2, 10  ; "standard")]
#[test_case(2, 1   ; "log_trace_degree_1_lt_l_skip")]
#[test_case(2, 0   ; "log_trace_degree_0_lt_l_skip")]
fn test_bn254_fib_air_roundtrip(l_skip: usize, log_trace_degree: usize) {
    use openvm_stark_backend::{
        test_utils::FibFixture, SystemParams, WhirConfig, WhirParams, WhirProximityStrategy,
    };
    use openvm_stark_sdk::config::log_up_params::log_up_security_params_baby_bear_100_bits;

    setup_tracing_with_log_level(Level::DEBUG);

    let n_stack = 8;
    let w_stack = 8;
    let k_whir = 4;
    let whir_params = WhirParams {
        k: k_whir,
        log_final_poly_len: k_whir,
        query_phase_pow_bits: 1,
        proximity: WhirProximityStrategy::UniqueDecoding,
        folding_pow_bits: 1,
        mu_pow_bits: 1,
    };
    let log_blowup = 1;
    let whir = WhirConfig::new(log_blowup, l_skip + n_stack, whir_params, 80);
    let params = SystemParams {
        l_skip,
        n_stack,
        w_stack,
        log_blowup,
        whir,
        logup: log_up_security_params_baby_bear_100_bits(0.0),
        max_constraint_degree: 3,
    };

    let fib = FibFixture::new(0, 1, 1 << log_trace_degree);

    let engine = crate::BabyBearBn254Poseidon2GpuEngine::new(params);
    let (pk, vk) = fib.keygen(&engine);
    let proof = fib.prove(&engine, &pk);
    engine
        .verify(&vk, &proof)
        .expect("BN254 verification failed");
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
fn bn254_host_merkle_layers(
    matrix: &[F],
    height: usize,
    width: usize,
    rows_per_query: usize,
) -> Vec<Vec<Bn254Digest>> {
    let perm = default_bn254_poseidon2_width3();
    let hasher = CpuMerkleHasher::new(
        MultiFieldHasher::<F, Bn254Scalar, _, 3, 16, 1>::new(perm.clone()),
        TruncatedPermutation::new(default_bn254_poseidon2_width2()),
    );

    let query_stride = height / rows_per_query;
    let mut row_buf = vec![F::ZERO; width];
    let mut leaf_hashes = Vec::with_capacity(rows_per_query);
    let mut query_digest_layer = Vec::with_capacity(query_stride);
    for query_idx in 0..query_stride {
        leaf_hashes.clear();
        for row_offset in 0..rows_per_query {
            let row_idx = row_offset * query_stride + query_idx;
            for col_idx in 0..width {
                row_buf[col_idx] = matrix[col_idx * height + row_idx];
            }
            leaf_hashes.push(hasher.hash_slice(&row_buf));
        }
        query_digest_layer.push(hasher.tree_compress(leaf_hashes.clone()));
    }

    let mut digest_layers = vec![query_digest_layer];
    while digest_layers.last().unwrap().len() > 1 {
        let prev_layer = digest_layers.last().unwrap();
        let layer = prev_layer
            .chunks_exact(2)
            .map(|pair| hasher.compress(pair[0], pair[1]))
            .collect_vec();
        digest_layers.push(layer);
    }
    digest_layers
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
fn bn254_host_merkle_layers_ext(
    matrix: &[EF],
    height: usize,
    width: usize,
    rows_per_query: usize,
) -> Vec<Vec<Bn254Digest>> {
    let perm = default_bn254_poseidon2_width3();
    let hasher = CpuMerkleHasher::new(
        MultiFieldHasher::<F, Bn254Scalar, _, 3, 16, 1>::new(perm.clone()),
        TruncatedPermutation::new(default_bn254_poseidon2_width2()),
    );

    let query_stride = height / rows_per_query;
    let mut row_buf = Vec::with_capacity(width * 4);
    let mut leaf_hashes = Vec::with_capacity(rows_per_query);
    let mut query_digest_layer = Vec::with_capacity(query_stride);
    for query_idx in 0..query_stride {
        leaf_hashes.clear();
        for row_offset in 0..rows_per_query {
            row_buf.clear();
            let row_idx = row_offset * query_stride + query_idx;
            for col_idx in 0..width {
                row_buf.extend_from_slice(
                    matrix[col_idx * height + row_idx].as_basis_coefficients_slice(),
                );
            }
            leaf_hashes.push(hasher.hash_slice(&row_buf));
        }
        query_digest_layer.push(hasher.tree_compress(leaf_hashes.clone()));
    }

    let mut digest_layers = vec![query_digest_layer];
    while digest_layers.last().unwrap().len() > 1 {
        let prev_layer = digest_layers.last().unwrap();
        let layer = prev_layer
            .chunks_exact(2)
            .map(|pair| hasher.compress(pair[0], pair[1]))
            .collect_vec();
        digest_layers.push(layer);
    }
    digest_layers
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
fn bn254_host_merkle_proofs(
    host_layers: &[Vec<Bn254Digest>],
    query_indices: &[usize],
) -> Vec<Vec<Bn254Digest>> {
    let proof_depth = host_layers.len() - 1;
    query_indices
        .iter()
        .map(|&index| {
            (0..proof_depth)
                .map(|layer_idx| host_layers[layer_idx][(index >> layer_idx) ^ 1])
                .collect_vec()
        })
        .collect_vec()
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
fn bn254_row_hash_emulated(vals: &[F]) -> Bn254Digest {
    let perm = default_bn254_poseidon2_width3();
    let mut state = [Bn254Scalar::ZERO; 3];
    let mut buf = [F::ZERO; 16];
    let mut cnt = 0usize;

    for &value in vals {
        buf[cnt] = value;
        cnt += 1;
        if cnt == 16 {
            state[0] = pack_f_to_sf(&buf[..8]);
            state[1] = pack_f_to_sf(&buf[8..16]);
            perm.permute_mut(&mut state);
            cnt = 0;
        }
    }

    if cnt > 0 {
        state[0] = pack_f_to_sf(&buf[..cnt.min(8)]);
        if cnt > 8 {
            state[1] = pack_f_to_sf(&buf[8..cnt.min(16)]);
        }
        perm.permute_mut(&mut state);
    }

    [state[0]]
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test]
fn test_bn254_merkle_gpu_matches_host_large_matrix() {
    let ctx = test_ctx();
    let height = 1 << 16;
    let width = 19;
    let rows_per_query = 1 << 4;
    let host_matrix = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(37).wrapping_add((i >> 7) as u32)))
        .collect_vec();

    let host_layers = bn254_host_merkle_layers(&host_matrix, height, width, rows_per_query);
    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );
    let gpu_tree = MerkleTreeGpu::<F, Bn254Digest>::new_with_hash::<
        crate::hash_scheme::Bn254Poseidon2MerkleHash,
    >(device_matrix, rows_per_query, true, &ctx)
    .unwrap();
    let gpu_layers = gpu_tree
        .digest_layers
        .iter()
        .map(|layer| layer.to_host_on(&ctx).unwrap())
        .collect_vec();

    for (layer_idx, (gpu_layer, host_layer)) in
        gpu_layers.iter().zip(host_layers.iter()).enumerate()
    {
        assert_eq!(
            gpu_layer.len(),
            host_layer.len(),
            "layer {layer_idx} length mismatch"
        );
        if let Some((digest_idx, (gpu_digest, host_digest))) = gpu_layer
            .iter()
            .zip(host_layer.iter())
            .enumerate()
            .find(|(_, (gpu_digest, host_digest))| gpu_digest != host_digest)
        {
            panic!(
                "layer {layer_idx} digest {digest_idx} mismatch: gpu={gpu_digest:?} host={host_digest:?}"
            );
        }
    }
    assert_eq!(
        gpu_tree.root(),
        *host_layers.last().unwrap().first().unwrap()
    );
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test]
fn test_bn254_merkle_proof_queries_gpu_match_host() {
    let ctx = test_ctx();
    let height = 1 << 12;
    let width = 19;
    let rows_per_query = 1;
    let query_indices = [0, 1, 7, 42, 255, 1023];
    let host_matrix_a = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(17).wrapping_add((i >> 4) as u32)))
        .collect_vec();
    let host_matrix_b = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(29).wrapping_add((i >> 6) as u32)))
        .collect_vec();

    let host_layers_a = bn254_host_merkle_layers(&host_matrix_a, height, width, rows_per_query);
    let host_layers_b = bn254_host_merkle_layers(&host_matrix_b, height, width, rows_per_query);
    let tree_a = MerkleTreeGpu::<F, Bn254Digest>::new_with_hash::<
        crate::hash_scheme::Bn254Poseidon2MerkleHash,
    >(
        DeviceMatrix::new(
            Arc::new(host_matrix_a.to_device_on(&ctx).unwrap()),
            height,
            width,
        ),
        rows_per_query,
        true,
        &ctx,
    )
    .unwrap();
    let tree_b = MerkleTreeGpu::<F, Bn254Digest>::new_with_hash::<
        crate::hash_scheme::Bn254Poseidon2MerkleHash,
    >(
        DeviceMatrix::new(
            Arc::new(host_matrix_b.to_device_on(&ctx).unwrap()),
            height,
            width,
        ),
        rows_per_query,
        true,
        &ctx,
    )
    .unwrap();

    let gpu_proofs = MerkleTreeGpu::<F, Bn254Digest>::batch_query_merkle_proofs(
        &[&tree_a, &tree_b],
        &query_indices,
        &ctx,
    )
    .unwrap();

    assert_eq!(
        gpu_proofs[0],
        bn254_host_merkle_proofs(&host_layers_a, &query_indices)
    );
    assert_eq!(
        gpu_proofs[1],
        bn254_host_merkle_proofs(&host_layers_b, &query_indices)
    );
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test]
fn test_bn254_row_hash_gpu_matches_host_multi_block_rows() {
    let ctx = test_ctx();
    let height = 1 << 12;
    let width = 19;
    let rows_per_query = 1;
    let host_matrix = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(53).wrapping_add((i >> 5) as u32)))
        .collect_vec();

    let host_layers = bn254_host_merkle_layers(&host_matrix, height, width, rows_per_query);
    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );
    let gpu_tree = MerkleTreeGpu::<F, Bn254Digest>::new_with_hash::<
        crate::hash_scheme::Bn254Poseidon2MerkleHash,
    >(device_matrix, rows_per_query, true, &ctx)
    .unwrap();
    let gpu_layer0 = gpu_tree.digest_layers[0].to_host_on(&ctx).unwrap();

    if let Some((digest_idx, (gpu_digest, host_digest))) = gpu_layer0
        .iter()
        .zip(host_layers[0].iter())
        .enumerate()
        .find(|(_, (gpu_digest, host_digest))| gpu_digest != host_digest)
    {
        panic!("row-hash digest {digest_idx} mismatch: gpu={gpu_digest:?} host={host_digest:?}");
    }
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test]
fn test_bn254_row_hash_ext_gpu_matches_host_multi_block_rows() {
    let ctx = test_ctx();
    let height = 1 << 11;
    let width = 5;
    let rows_per_query = 1;
    let host_matrix = (0..width * height)
        .map(|i| {
            EF::from_basis_coefficients_fn(|j| {
                F::from_u32(
                    (i as u32)
                        .wrapping_mul(43)
                        .wrapping_add((j as u32) * 11 + (i >> 3) as u32),
                )
            })
        })
        .collect_vec();

    let host_layers = bn254_host_merkle_layers_ext(&host_matrix, height, width, rows_per_query);
    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );
    let gpu_tree = MerkleTreeGpu::<EF, Bn254Digest>::new_with_hash::<
        crate::hash_scheme::Bn254Poseidon2MerkleHash,
    >(device_matrix, rows_per_query, true, &ctx)
    .unwrap();
    let gpu_layer0 = gpu_tree.digest_layers[0].to_host_on(&ctx).unwrap();

    if let Some((digest_idx, (gpu_digest, host_digest))) = gpu_layer0
        .iter()
        .zip(host_layers[0].iter())
        .enumerate()
        .find(|(_, (gpu_digest, host_digest))| gpu_digest != host_digest)
    {
        panic!(
            "row-hash-ext digest {digest_idx} mismatch: gpu={gpu_digest:?} host={host_digest:?}"
        );
    }
}

#[cfg(feature = "baby-bear-bn254-poseidon2")]
#[test]
fn test_bn254_row_hash_emulation_matches_host_multi_block_rows() {
    let perm = default_bn254_poseidon2_width3();
    let host_hasher = CpuMerkleHasher::new(
        MultiFieldHasher::<F, Bn254Scalar, _, 3, 16, 1>::new(perm),
        TruncatedPermutation::new(default_bn254_poseidon2_width2()),
    );
    let row = (0..19)
        .map(|i| F::from_u32((i as u32).wrapping_mul(53).wrapping_add((i >> 1) as u32)))
        .collect_vec();

    assert_eq!(bn254_row_hash_emulated(&row), host_hasher.hash_slice(&row));
}

#[test]
fn test_merkle_gpu_supports_512_rows_per_query() {
    use openvm_cuda_common::copy::MemCopyH2D;

    let ctx = test_ctx();
    let height = 1 << 9;
    let width = 7;
    let rows_per_query = 1 << 9;
    let host_matrix = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(31).wrapping_add((i >> 2) as u32)))
        .collect_vec();

    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );
    let tree = MerkleTreeGpu::<F, crate::prelude::Digest>::new_with_hash::<
        crate::hash_scheme::Poseidon2MerkleHash,
    >(device_matrix, rows_per_query, true, &ctx)
    .expect("rows_per_query=512 should be supported");

    assert_eq!(tree.query_stride(), 1);
    assert_eq!(tree.proof_depth(), 0);
}

#[test]
fn test_merkle_batch_query_zero_work_returns_empty() {
    use openvm_cuda_common::copy::MemCopyH2D;

    let ctx = test_ctx();
    let height = 1 << 3;
    let width = 4;
    let rows_per_query = height;
    let host_matrix = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(19).wrapping_add(7)))
        .collect_vec();

    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );
    let tree = MerkleTreeGpu::<F, crate::prelude::Digest>::new_with_hash::<
        crate::hash_scheme::Poseidon2MerkleHash,
    >(device_matrix, rows_per_query, true, &ctx)
    .unwrap();

    let proofs =
        MerkleTreeGpu::<F, crate::prelude::Digest>::batch_query_merkle_proofs(&[&tree], &[0], &ctx)
            .unwrap();
    assert_eq!(proofs.len(), 1);
    assert_eq!(proofs[0].len(), 1);
    assert!(proofs[0][0].is_empty());

    let empty_queries =
        MerkleTreeGpu::<F, crate::prelude::Digest>::batch_query_merkle_proofs(&[&tree], &[], &ctx)
            .unwrap();
    assert_eq!(empty_queries.len(), 1);
    assert!(empty_queries[0].is_empty());
}

#[test]
fn test_merkle_batch_open_rows_empty_queries_returns_empty() {
    use openvm_cuda_common::copy::MemCopyH2D;

    let ctx = test_ctx();
    let height = 1 << 3;
    let width = 4;
    let host_matrix = (0..width * height)
        .map(|i| F::from_u32((i as u32).wrapping_mul(23).wrapping_add(11)))
        .collect_vec();
    let device_matrix = DeviceMatrix::new(
        Arc::new(host_matrix.to_device_on(&ctx).unwrap()),
        height,
        width,
    );

    let opened_rows = MerkleTreeGpu::<F, crate::prelude::Digest>::batch_open_rows(
        &[&device_matrix],
        &[],
        1,
        1,
        &ctx,
    )
    .unwrap();
    assert_eq!(opened_rows.len(), 1);
    assert!(opened_rows[0].is_empty());
}

// ===========================================================================
// GPU-specific tests (not shared)
// ===========================================================================

#[test]
fn test_interactions_roundtrip_with_l_skip_zero() {
    use openvm_stark_backend::test_utils::test_system_params_small;

    setup_tracing_with_log_level(Level::DEBUG);

    let engine = BabyBearPoseidon2GpuEngine::new(test_system_params_small(0, 8, 3));
    let fixture = InteractionsFixture11;
    let (vk, proof) = fixture.keygen_and_prove(&engine);
    engine
        .verify(&vk, &proof)
        .expect("l_skip=0 interactions roundtrip should verify");
}

#[test]
fn test_gpu_l_skip_10_is_rejected() {
    use crate::cuda::batch_ntt_small::validate_gpu_l_skip;

    assert!(validate_gpu_l_skip(9).is_ok());
    assert!(validate_gpu_l_skip(10).is_err());
}

#[test_case(1 ; "l_skip_1")]
#[test_case(4 ; "l_skip_4")]
#[test_case(9 ; "l_skip_9")]
fn test_batch_ntt_small_partial_last_block_roundtrip(l_skip: usize) {
    use openvm_cuda_common::copy::{MemCopyD2H, MemCopyH2D};

    setup_tracing_with_log_level(Level::DEBUG);
    let gpu_ctx = test_ctx();

    let block_size = 1usize << l_skip;
    let cnt_blocks = (1024usize >> l_skip) + 1;
    let original = (0..(cnt_blocks * block_size))
        .map(|i| F::from_u32((i as u32).wrapping_mul(17).wrapping_add(5)))
        .collect_vec();
    let mut d_values = original.to_device_on(&gpu_ctx).unwrap();

    unsafe {
        batch_ntt_small(
            &mut d_values,
            l_skip,
            cnt_blocks,
            false,
            gpu_ctx.stream.as_raw(),
        )
        .unwrap();
        batch_ntt_small(
            &mut d_values,
            l_skip,
            cnt_blocks,
            true,
            gpu_ctx.stream.as_raw(),
        )
        .unwrap();
    }

    assert_eq!(d_values.to_host_on(&gpu_ctx).unwrap(), original);
}

/// GPU integration regression for multi-warp skip-domain sizes.
///
/// These cases exercise `ntt_coset_interpolate` through the normal proving path with
/// a mixture of cached, preprocessed, interaction, and regular traces. `l_skip = 6`
/// is the first shared-memory NTT size; `l_skip = 9` exercises a larger supported
/// skip domain while staying under the current round0 CUDA resource limit.
#[test_case(6 ; "l_skip_6")]
#[test_case(9 ; "l_skip_9")]
fn test_mixture_fixture_gpu_roundtrip_large_l_skip(l_skip: usize) {
    use openvm_stark_backend::test_utils::{test_system_params_small, MixtureFixture};

    setup_tracing_with_log_level(Level::DEBUG);

    let n_stack = 13 - l_skip;
    let engine = BabyBearPoseidon2GpuEngine::new(test_system_params_small(l_skip, n_stack, 3));
    let fixture = MixtureFixture::standard(10, engine.config().clone());
    let (vk, proof) = fixture.keygen_and_prove(&engine);
    engine.verify(&vk, &proof).unwrap_or_else(|err| {
        panic!("l_skip={l_skip} MixtureFixture proof should verify on the GPU backend: {err:?}")
    });
}

#[test]
fn test_frac_matrix_vertically_repeat_guards_tail_rows() {
    use openvm_cuda_common::{
        copy::{MemCopyD2H, MemCopyH2D},
        d_buffer::DeviceBuffer,
    };
    use openvm_stark_backend::prover::fractional_sumcheck_gkr::Frac;

    setup_tracing_with_log_level(Level::DEBUG);
    let gpu_ctx = test_ctx();

    let width = 3usize;
    let height = 769usize;
    let lifted_height = 1538usize;
    let padded_height = 2048usize;
    let tail_rows = padded_height - lifted_height;
    let output_len = width * lifted_height + tail_rows;
    let canary = Frac::new(EF::from_u32(777), EF::from_u32(999));

    let input = (0..(width * height))
        .map(|i| Frac::new(EF::from_u32(i as u32 + 1), EF::from_u32(i as u32 + 1001)))
        .collect_vec();
    let mut expected = vec![canary; output_len];
    for col in 0..width {
        for row in 0..lifted_height {
            expected[col * lifted_height + row] = input[col * height + (row % height)];
        }
    }

    let d_input = input.to_device_on(&gpu_ctx).unwrap();
    let mut d_output = DeviceBuffer::<Frac<EF>>::with_capacity_on(output_len, &gpu_ctx);
    vec![canary; output_len]
        .copy_to_on(&mut d_output, &gpu_ctx)
        .unwrap();

    unsafe {
        frac_matrix_vertically_repeat(
            d_output.as_mut_ptr(),
            d_input.as_ptr(),
            width as u32,
            lifted_height as u32,
            height as u32,
            gpu_ctx.stream.as_raw(),
        )
        .unwrap();
    }

    let output = d_output.to_host_on(&gpu_ctx).unwrap();
    assert_eq!(output.len(), expected.len());
    for (idx, (got, want)) in output.iter().zip(expected.iter()).enumerate() {
        assert_eq!(got.p, want.p, "numerator mismatch at index {idx}");
        assert_eq!(got.q, want.q, "denominator mismatch at index {idx}");
    }
}

#[test_case(9)]
#[test_case(2 ; "when log_height equals l_skip")]
#[test_case(1 ; "when log_height less than l_skip")]
#[test_case(0 ; "when log_height is zero")]
fn test_stacked_opening_reduction(
    log_trace_degree: usize,
) -> Result<(), StackedReductionError<EF>> {
    setup_tracing_with_log_level(Level::DEBUG);

    let gpu_engine = BabyBearPoseidon2GpuEngine::new(default_test_params_small());
    let params = gpu_engine.config().params().clone();

    let engine = BabyBearPoseidon2RefEngine::<DuplexSponge>::new(params.clone());
    let fib = FibFixture::new(0, 1, 1 << log_trace_degree);
    let (pk, _vk) = fib.keygen(&engine);
    let pk = engine.device().transport_pk_to_device(&pk);
    let mut ctx = fib.generate_proving_ctx();

    ctx.sort_for_stacking();

    let (_, common_main_pcs_data) = {
        stacked_commit(
            engine.config().hasher(),
            params.l_skip,
            params.n_stack,
            params.log_blowup,
            params.k_whir(),
            &ctx.common_main_traces()
                .map(|(_, trace)| trace)
                .collect_vec(),
        )
        .unwrap()
    };

    let omega_skip = F::two_adic_generator(params.l_skip);
    let omega_skip_pows = omega_skip.powers().take(1 << params.l_skip).collect_vec();

    let device = engine.device();
    let ((_, batch_proof), r) = device
        .prove_rap_constraints(
            &mut DuplexSpongeGpu::default(),
            &pk,
            &ctx,
            &common_main_pcs_data,
        )
        .unwrap();

    let need_rot = pk.per_air[ctx.per_trace[0].0].vk.params.need_rot;
    let need_rot_per_commit = vec![vec![need_rot]];
    let (stacking_proof, _) = prove_stacked_opening_reduction::<SC, _, _, _, StackedReductionCpu<SC>>(
        device,
        &mut DuplexSpongeGpu::default(),
        params.n_stack,
        vec![&common_main_pcs_data],
        need_rot_per_commit.clone(),
        &r,
    );

    debug!(?batch_proof.column_openings);

    let u_prism = verify_stacked_reduction(
        &mut default_duplex_sponge(),
        &stacking_proof,
        &[common_main_pcs_data.layout],
        &need_rot_per_commit,
        params.l_skip,
        params.n_stack,
        &batch_proof.column_openings,
        &r,
        &omega_skip_pows,
    )?;
    assert_eq!(u_prism.len(), params.n_stack + 1);
    Ok(())
}

/// Tests that monomial-based and DAG-based zerocheck evaluation paths produce identical results.
#[test]
fn test_monomial_vs_dag_equivalence() {
    use openvm_cuda_common::copy::{MemCopyD2H, MemCopyH2D};
    use openvm_stark_backend::{
        poly_common::eval_eq_uni_at_one, test_utils::prove_up_to_batch_constraints,
    };
    use p3_util::log2_strict_usize;

    use crate::{
        cuda::logup_zerocheck::{fold_selectors_round0, interpolate_columns_gpu, MainMatrixPtrs},
        logup_zerocheck::{
            batch_mle::{TraceCtx, ZerocheckMleBatchBuilder},
            batch_mle_monomial::{compute_lambda_combinations, ZerocheckMonomialBatch},
            fold_ple::fold_ple_evals_rotate,
        },
        poly::EqEvalSegments,
        prelude::EF,
    };

    setup_tracing_with_log_level(Level::DEBUG);
    let gpu_ctx = test_ctx();

    let log_trace_degree = 5;
    let threshold = 32u32;

    let gpu_engine = BabyBearPoseidon2GpuEngine::new(default_test_params_small());
    let device = gpu_engine.device();
    let params = gpu_engine.params();
    let l_skip = params.l_skip;

    let fib = FibFixture::new(0, 1, 1 << log_trace_degree);
    let (pk, _vk) = fib.keygen(&gpu_engine);
    let pk = <_ as DeviceDataTransporter<SC, GpuBackend>>::transport_pk_to_device(device, &pk);

    let ctx_for_challenges =
        <_ as DeviceDataTransporter<SC, GpuBackend>>::transport_proving_ctx_to_device(
            device,
            &fib.generate_proving_ctx(),
        )
        .into_sorted();
    let mut prover_sponge = DuplexSpongeGpu::default();
    let ((_, _), r) =
        prove_up_to_batch_constraints(&gpu_engine, &mut prover_sponge, &pk, ctx_for_challenges);

    let proving_ctx =
        <_ as DeviceDataTransporter<SC, GpuBackend>>::transport_proving_ctx_to_device(
            device,
            &fib.generate_proving_ctx(),
        )
        .into_sorted();

    let height = proving_ctx.per_trace[0].1.common_main.height();
    let n_calc = log2_strict_usize(height).saturating_sub(l_skip);
    let xi_len = l_skip + n_calc + 1;

    let mut xi: Vec<EF> = Vec::with_capacity(xi_len);
    for _ in 0..l_skip {
        xi.push(prover_sponge.sample_ext());
    }
    xi.extend_from_slice(&r);
    while xi.len() < xi_len {
        xi.push(prover_sponge.sample_ext());
    }
    assert!(xi.len() > l_skip, "xi vector must have enough elements");

    let omega_skip = F::two_adic_generator(l_skip);
    let omega_skip_pows: Vec<F> = omega_skip.powers().take(1 << l_skip).collect();
    let d_omega_skip_pows = omega_skip_pows.to_device_on(&gpu_ctx).unwrap();

    let (air_idx, air_ctx) = &proving_ctx.per_trace[0];
    let height = air_ctx.common_main.height();
    let n = log2_strict_usize(height) as isize - l_skip as isize;
    let n_lift = n.max(0) as usize;

    let eq_xis = EqEvalSegments::new(&xi[l_skip..], &gpu_ctx).expect("failed to compute eq_xis");

    let sel_height = 1 << n_lift;
    let mut sel_cols = F::zero_vec(3 * sel_height);
    sel_cols[sel_height..2 * sel_height - 1].fill(F::ONE);
    sel_cols[0] = F::ONE;
    sel_cols[2 * sel_height + sel_height - 1] = F::ONE;
    let d_sels_base = sel_cols.to_device_on(&gpu_ctx).unwrap();

    let (l, r_fold) = if n.is_negative() {
        (
            l_skip.wrapping_add_signed(n),
            r[0].exp_power_of_2(-n as usize),
        )
    } else {
        (l_skip, r[0])
    };
    let omega = F::two_adic_generator(l);
    let is_first = eval_eq_uni_at_one(l, r_fold);
    let is_last = eval_eq_uni_at_one(l, r_fold * omega);
    let d_sels_folded = openvm_cuda_common::d_buffer::DeviceBuffer::<EF>::with_capacity_on(
        sel_height * 3,
        &gpu_ctx,
    );
    unsafe {
        fold_selectors_round0(
            d_sels_folded.as_mut_ptr(),
            d_sels_base.as_ptr(),
            is_first,
            is_last,
            sel_height,
            gpu_ctx.stream.as_raw(),
        )
        .unwrap();
    }

    let inv_lagrange_denoms_r0 =
        crate::utils::compute_barycentric_inv_lagrange_denoms(l_skip, &omega_skip_pows, r[0]);
    let d_inv_lagrange_denoms_r0 = inv_lagrange_denoms_r0.to_device_on(&gpu_ctx).unwrap();

    let mat_folded = fold_ple_evals_rotate(
        l_skip,
        &d_omega_skip_pows,
        &air_ctx.common_main,
        &d_inv_lagrange_denoms_r0,
        true,
        &gpu_ctx,
    )
    .unwrap();

    let lambda = prover_sponge.sample_ext();
    let air_pk = &pk.per_air[*air_idx];
    let max_num_constraints = air_pk
        .vk
        .symbolic_constraints
        .constraints
        .constraint_idx
        .len();
    let h_lambda_pows: Vec<EF> = lambda.powers().take(max_num_constraints).collect();
    let d_lambda_pows = h_lambda_pows.to_device_on(&gpu_ctx).unwrap();

    let d_public_values = if air_ctx.public_values.is_empty() {
        openvm_cuda_common::d_buffer::DeviceBuffer::new()
    } else {
        air_ctx.public_values.to_device_on(&gpu_ctx).unwrap()
    };

    let dag = &air_pk.vk.symbolic_constraints;
    let has_constraints = dag.constraints.num_constraints() > 0;
    assert!(has_constraints, "FibFixture should have constraints");

    let has_monomials = air_pk
        .other_data
        .zerocheck_monomials
        .as_ref()
        .map(|m| m.num_monomials > 0)
        .unwrap_or(false);
    assert!(
        has_monomials,
        "Proving key should have expanded monomials for monomial path"
    );

    let s_deg = params.max_constraint_degree as usize + 1;

    for test_round in 1..=n_lift.min(3) {
        let n_round = n_lift.saturating_sub(test_round - 1);
        let test_height = 1 << n_round;
        let num_y = (test_height / 2) as u32;

        if num_y == 0 || num_y > threshold {
            continue;
        }

        debug!(test_round, num_y, %threshold, "testing monomial vs DAG equivalence");

        let has_interactions = false;
        let mut columns: Vec<*const EF> = Vec::new();
        columns.push(eq_xis.get_ptr(n_round));
        for col in 0..3 {
            columns.push(d_sels_folded.as_ptr().wrapping_add(col * sel_height));
        }
        for col in 0..mat_folded.width() {
            columns.push(
                mat_folded
                    .buffer()
                    .as_ptr()
                    .wrapping_add(col * mat_folded.height()),
            );
        }

        let interpolated = crate::base::DeviceMatrix::<EF>::with_capacity_on(
            s_deg * num_y as usize,
            columns.len(),
            &gpu_ctx,
        );
        let d_columns = columns.to_device_on(&gpu_ctx).unwrap();
        unsafe {
            interpolate_columns_gpu(
                interpolated.buffer(),
                &d_columns,
                s_deg,
                num_y as usize,
                gpu_ctx.stream.as_raw(),
            )
            .expect("failed to interpolate columns");
        }

        let interpolated_height = interpolated.height();
        let eq_xi_ptr = eq_xis.get_ptr(n_round);
        let sels_ptr = interpolated
            .buffer()
            .as_ptr()
            .wrapping_add(interpolated_height);

        let main_ptrs = [MainMatrixPtrs {
            data: interpolated
                .buffer()
                .as_ptr()
                .wrapping_add(4 * interpolated_height),
            air_width: mat_folded.width() as u32 / 2,
        }];
        let main_ptrs_dev = main_ptrs.to_device_on(&gpu_ctx).unwrap();

        let trace_ctx = TraceCtx {
            trace_idx: 0,
            air_idx: *air_idx,
            n_lift,
            num_y,
            has_constraints: true,
            has_interactions,
            norm_factor: F::ONE,
            eq_xi_ptr,
            sels_ptr,
            prep_ptr: MainMatrixPtrs {
                data: std::ptr::null(),
                air_width: 0,
            },
            main_ptrs_dev,
            public_ptr: d_public_values.as_ptr(),
            eq_3bs_ptr: std::ptr::null(),
        };

        let dag_builder =
            ZerocheckMleBatchBuilder::new(std::iter::once(&trace_ctx), &pk, s_deg as u32, &gpu_ctx)
                .unwrap();
        let dag_output = dag_builder.evaluate(&d_lambda_pows, s_deg as u32).unwrap();
        let dag_results: Vec<EF> = dag_output.to_host_on(&gpu_ctx).expect("copy DAG output");

        let lambda_comb = compute_lambda_combinations(&pk, 0, &d_lambda_pows, &gpu_ctx).unwrap();
        let mono_batch = ZerocheckMonomialBatch::new(
            std::iter::once(&trace_ctx),
            &pk,
            &[&lambda_comb],
            &gpu_ctx,
        )
        .unwrap();
        let mono_output = mono_batch.evaluate(s_deg as u32).unwrap();
        let mono_results: Vec<EF> = mono_output
            .to_host_on(&gpu_ctx)
            .expect("copy monomial output");

        assert_eq!(
            dag_results.len(),
            mono_results.len(),
            "Output lengths should match"
        );
        for (i, (dag_val, mono_val)) in dag_results.iter().zip(mono_results.iter()).enumerate() {
            assert_eq!(
                dag_val, mono_val,
                "Mismatch at index {i} for num_y={num_y}: DAG={dag_val:?}, monomial={mono_val:?}"
            );
        }
        debug!(
            num_y,
            num_results = dag_results.len(),
            "monomial vs DAG equivalence verified"
        );
    }
}