openvm_benchmarks_fields/

lib.rs

//! Extension Field Benchmark
//!
//! Benchmarks for field arithmetic operations on GPU.
//! Measures compute throughput by doing many ops per element to amortize memory access.
//!
//! Throughput is reported in Gops/s (Giga-operations per second = billion ops/sec).

use std::{ffi::c_void, sync::OnceLock, time::Instant};

use openvm_cuda_common::{
    copy::MemCopyH2D,
    d_buffer::DeviceBuffer,
    stream::{cudaStream_t, GpuDeviceCtx},
};

pub fn bench_ctx() -> &'static GpuDeviceCtx {
    static CTX: OnceLock<GpuDeviceCtx> = OnceLock::new();
    CTX.get_or_init(|| GpuDeviceCtx::for_current_device().expect("benchmark CUDA context"))
}

// ============================================================================
// FFI Bindings
// ============================================================================

// Benchmark kernels
mod ffi {
    use super::*;

    #[link(name = "ext_field_bench")]
    extern "C" {
        pub(crate) fn init_fp(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // BabyBear quartic extension (simple implementation)
        pub fn init_fp4(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp4(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp4(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp4(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // BabyBear quartic extension (optimized bb31_4_t)
        pub(crate) fn init_fpext(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fpext(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fpext(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fpext(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        pub fn init_fp5(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp5(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp5(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp5(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        pub fn init_fp6(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp6(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp6(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp6(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        pub fn init_fp2x3(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp2x3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp2x3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp2x3(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        pub fn init_fp3x2(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_fp3x2(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_fp3x2(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_fp3x2(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // KoalaBear base field
        pub fn init_kb(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_kb(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_kb(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_kb(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // KoalaBear quintic extension (x^5 + x + 4)
        pub fn init_kb5(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_kb5(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_kb5(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_kb5(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // KoalaBear sextic extension (x^6 + x^3 + 1)
        pub fn init_kb6(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_kb6(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_kb6(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_kb6(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // KoalaBear 2×3 tower (u²=3, v³=1+u)
        pub fn init_kb2x3(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_kb2x3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_kb2x3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_kb2x3(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // KoalaBear 3×2 tower (w³=-w-4, z²=3)
        pub fn init_kb3x2(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_kb3x2(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_kb3x2(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_kb3x2(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // Goldilocks base field (64-bit prime: 2^64 - 2^32 + 1)
        pub(crate) fn init_gl(
            out: *mut c_void,
            raw_data: *const u64,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_gl(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_gl(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_gl(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;

        // Goldilocks cubic extension (X³ - X - 1)
        pub(crate) fn init_gl3(
            out: *mut c_void,
            raw_data: *const u64,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn add_gl3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn mul_gl3(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn inv_gl3(
            out: *mut c_void,
            a: *const c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
    }

    // Poseidon2 benchmark kernels
    #[link(name = "ext_field_bench")]
    extern "C" {
        pub(crate) fn init_poseidon2_bb(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn run_poseidon2_bb(
            states: *mut c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn init_poseidon2_kb(
            out: *mut c_void,
            raw_data: *const u32,
            n: usize,
            stream: cudaStream_t,
        ) -> i32;
        pub(crate) fn run_poseidon2_kb(
            states: *mut c_void,
            n: usize,
            reps: i32,
            stream: cudaStream_t,
        ) -> i32;
    }
}

macro_rules! wrap_init_u32 {
    ($vis:vis $name:ident) => {
        /// Launches the benchmark init kernel on the benchmark stream.
        ///
        /// # Safety
        /// The caller must provide valid device pointers for `out` and `raw_data`,
        /// and `out` must have capacity for `n` field elements of the target type.
        #[inline]
        $vis unsafe extern "C" fn $name(out: *mut c_void, raw_data: *const u32, n: usize) -> i32 {
            ffi::$name(out, raw_data, n, bench_ctx().stream.as_raw())
        }
    };
}

macro_rules! wrap_init_u64 {
    ($vis:vis $name:ident) => {
        /// Launches the benchmark init kernel on the benchmark stream.
        ///
        /// # Safety
        /// The caller must provide valid device pointers for `out` and `raw_data`,
        /// and `out` must have capacity for `n` field elements of the target type.
        #[inline]
        $vis unsafe extern "C" fn $name(out: *mut c_void, raw_data: *const u64, n: usize) -> i32 {
            ffi::$name(out, raw_data, n, bench_ctx().stream.as_raw())
        }
    };
}

macro_rules! wrap_binary {
    ($name:ident) => {
        #[inline]
        unsafe extern "C" fn $name(
            out: *mut c_void,
            a: *const c_void,
            b: *const c_void,
            n: usize,
            reps: i32,
        ) -> i32 {
            ffi::$name(out, a, b, n, reps, bench_ctx().stream.as_raw())
        }
    };
}

macro_rules! wrap_unary {
    ($name:ident) => {
        #[inline]
        unsafe extern "C" fn $name(out: *mut c_void, a: *const c_void, n: usize, reps: i32) -> i32 {
            ffi::$name(out, a, n, reps, bench_ctx().stream.as_raw())
        }
    };
}

wrap_init_u32!(pub init_fp);
wrap_binary!(add_fp);
wrap_binary!(mul_fp);
wrap_unary!(inv_fp);

wrap_init_u32!(pub init_fp4);
wrap_binary!(add_fp4);
wrap_binary!(mul_fp4);
wrap_unary!(inv_fp4);

wrap_init_u32!(init_fpext);
wrap_binary!(add_fpext);
wrap_binary!(mul_fpext);
wrap_unary!(inv_fpext);

wrap_init_u32!(pub init_fp5);
wrap_binary!(add_fp5);
wrap_binary!(mul_fp5);
wrap_unary!(inv_fp5);

wrap_init_u32!(pub init_fp6);
wrap_binary!(add_fp6);
wrap_binary!(mul_fp6);
wrap_unary!(inv_fp6);

wrap_init_u32!(pub init_fp2x3);
wrap_binary!(add_fp2x3);
wrap_binary!(mul_fp2x3);
wrap_unary!(inv_fp2x3);

wrap_init_u32!(pub init_fp3x2);
wrap_binary!(add_fp3x2);
wrap_binary!(mul_fp3x2);
wrap_unary!(inv_fp3x2);

wrap_init_u32!(pub init_kb);
wrap_binary!(add_kb);
wrap_binary!(mul_kb);
wrap_unary!(inv_kb);

wrap_init_u32!(pub init_kb5);
wrap_binary!(add_kb5);
wrap_binary!(mul_kb5);
wrap_unary!(inv_kb5);

wrap_init_u32!(pub init_kb6);
wrap_binary!(add_kb6);
wrap_binary!(mul_kb6);
wrap_unary!(inv_kb6);

wrap_init_u32!(pub init_kb2x3);
wrap_binary!(add_kb2x3);
wrap_binary!(mul_kb2x3);
wrap_unary!(inv_kb2x3);

wrap_init_u32!(pub init_kb3x2);
wrap_binary!(add_kb3x2);
wrap_binary!(mul_kb3x2);
wrap_unary!(inv_kb3x2);

wrap_init_u64!(init_gl);
wrap_binary!(add_gl);
wrap_binary!(mul_gl);
wrap_unary!(inv_gl);

wrap_init_u64!(init_gl3);
wrap_binary!(add_gl3);
wrap_binary!(mul_gl3);
wrap_unary!(inv_gl3);

wrap_init_u32!(init_poseidon2_bb);

#[inline]
unsafe extern "C" fn run_poseidon2_bb(states: *mut c_void, n: usize, reps: i32) -> i32 {
    ffi::run_poseidon2_bb(states, n, reps, bench_ctx().stream.as_raw())
}

wrap_init_u32!(init_poseidon2_kb);

#[inline]
unsafe extern "C" fn run_poseidon2_kb(states: *mut c_void, n: usize, reps: i32) -> i32 {
    ffi::run_poseidon2_kb(states, n, reps, bench_ctx().stream.as_raw())
}

/// Check CUDA return code, panic on error
pub fn cuda_check(code: i32) {
    assert!(code == 0, "CUDA error: {}", code);
}

/// Sync and check
pub fn sync() {
    bench_ctx().stream.synchronize().expect("sync failed");
}

// ============================================================================
// Configuration & Results
// ============================================================================

pub struct BenchConfig {
    pub num_elements: usize,
    pub warmup_iters: usize,
    pub bench_iters: usize,
    pub ops_per_element: i32,
}

impl Default for BenchConfig {
    fn default() -> Self {
        Self {
            num_elements: 1 << 22, // 4M elements
            warmup_iters: 3,
            bench_iters: 10,
            ops_per_element: 100,
        }
    }
}

#[derive(Debug, Clone)]
pub struct OpResult {
    pub avg_time_ms: f64,
    pub throughput_gops: f64,
}

#[derive(Clone)]
pub struct FieldBenchResult {
    pub field_name: String,
    pub bits: usize, // bits of provable security
    pub u32s_per_element: usize,
    pub init: OpResult,
    pub add: OpResult,
    pub mul: OpResult,
    pub inv: OpResult,
}

/// Print a table of benchmark results
pub fn print_benchmark_tables(results: &[FieldBenchResult], baseline: &FieldBenchResult) {
    // Find max field name length for alignment
    let max_name_len = results
        .iter()
        .map(|r| r.field_name.len())
        .max()
        .unwrap_or(10);

    // Table 1: Time (ms)
    println!("### Time (ms)");
    println!();
    print!("| {:width$} | bits |", "Field", width = max_name_len);
    println!("    init |     add |     mul |     inv |");
    print!("|{:-<width$}--|-----:|", "", width = max_name_len);
    println!("--------:|--------:|--------:|--------:|");
    for r in results {
        print!(
            "| {:width$} | {:>4} |",
            r.field_name,
            r.bits,
            width = max_name_len
        );
        println!(
            " {:>7.3} | {:>7.3} | {:>7.3} | {:>7.3} |",
            r.init.avg_time_ms, r.add.avg_time_ms, r.mul.avg_time_ms, r.inv.avg_time_ms
        );
    }
    println!();

    // Table 2: Throughput (Gops/s)
    println!("### Throughput (Gops/s)");
    println!();
    print!("| {:width$} | bits |", "Field", width = max_name_len);
    println!("    init |     add |     mul |     inv |");
    print!("|{:-<width$}--|-----:|", "", width = max_name_len);
    println!("--------:|--------:|--------:|--------:|");
    for r in results {
        print!(
            "| {:width$} | {:>4} |",
            r.field_name,
            r.bits,
            width = max_name_len
        );
        println!(
            " {:>7.1} | {:>7.1} | {:>7.1} | {:>7.1} |",
            r.init.throughput_gops,
            r.add.throughput_gops,
            r.mul.throughput_gops,
            r.inv.throughput_gops
        );
    }
    println!();

    // Table 3: Relative to baseline (xN)
    println!("### Relative to {} (xN)", baseline.field_name);
    println!();
    print!("| {:width$} | bits |", "Field", width = max_name_len);
    println!("    init |     add |     mul |     inv |");
    print!("|{:-<width$}--|-----:|", "", width = max_name_len);
    println!("--------:|--------:|--------:|--------:|");
    for r in results {
        let init_ratio = baseline.init.throughput_gops / r.init.throughput_gops;
        let add_ratio = baseline.add.throughput_gops / r.add.throughput_gops;
        let mul_ratio = baseline.mul.throughput_gops / r.mul.throughput_gops;
        let inv_ratio = baseline.inv.throughput_gops / r.inv.throughput_gops;

        print!(
            "| {:width$} | {:>4} |",
            r.field_name,
            r.bits,
            width = max_name_len
        );
        println!(
            " {:>7.1} | {:>7.1} | {:>7.1} | {:>7.1} |",
            init_ratio, add_ratio, mul_ratio, inv_ratio
        );
    }
    println!();
}

// ============================================================================
// Benchmark Implementation
// ============================================================================

pub fn random_u32s(count: usize, seed: u64) -> Vec<u32> {
    let mut rng = seed;
    (0..count)
        .map(|_| {
            rng = rng
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            (rng >> 32) as u32
        })
        .collect()
}

pub fn random_u64s(count: usize, seed: u64) -> Vec<u64> {
    let mut rng = seed;
    (0..count)
        .map(|_| {
            rng = rng
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            rng
        })
        .collect()
}

pub fn measure<F: FnMut()>(config: &BenchConfig, total_ops: u64, mut f: F) -> OpResult {
    // Warmup
    for _ in 0..config.warmup_iters {
        f();
    }
    sync();

    // Timed
    let start = Instant::now();
    for _ in 0..config.bench_iters {
        f();
    }
    sync();
    let elapsed = start.elapsed();

    let avg_time_ms = elapsed.as_secs_f64() * 1000.0 / config.bench_iters as f64;
    let throughput_gops =
        (total_ops * config.bench_iters as u64) as f64 / elapsed.as_secs_f64() / 1e9;

    OpResult {
        avg_time_ms,
        throughput_gops,
    }
}

pub fn bench_fp(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Only 3 device buffers: a, b, out (init works in-place)
    let d_a = random_u32s(n, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n, bench_ctx());

    // Benchmark init (in-place: raw u32 -> Fp) - also initializes d_a
    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp".into(),
        bits: 31,
        u32s_per_element: 1,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fp4(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    let d_a = random_u32s(n * 4, 11111).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 4, 22222).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 4, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp4(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp4(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp4(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp4(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp4(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp4".into(),
        bits: 124,
        u32s_per_element: 4,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fpext(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Only 3 device buffers: a, b, out (init works in-place)
    let d_a = random_u32s(n * 4, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 4, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 4, bench_ctx());

    // Benchmark init (in-place: raw u32s -> FpExt) - also initializes d_a
    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fpext(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fpext(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fpext(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fpext(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fpext(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "FpExt".into(),
        bits: 124,
        u32s_per_element: 4,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fp5(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Only 3 device buffers: a, b, out (init works in-place)
    let d_a = random_u32s(n * 5, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 5, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 5, bench_ctx());

    // Benchmark init (in-place: raw u32s -> Fp5) - also initializes d_a
    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp5(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp5(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp5(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp5(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp5(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp5".into(),
        bits: 155,
        u32s_per_element: 5,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fp6(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Only 3 device buffers: a, b, out (init works in-place)
    let d_a = random_u32s(n * 6, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    // Benchmark init (in-place: raw u32s -> Fp6) - also initializes d_a
    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp6(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp6(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp6(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp6(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp6(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp6".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fp2x3(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    let d_a = random_u32s(n * 6, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp2x3(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp2x3(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp2x3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp2x3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp2x3(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp2x3".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_fp3x2(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    let d_a = random_u32s(n * 6, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_fp3x2(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_fp3x2(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_fp3x2(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_fp3x2(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_fp3x2(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Fp3x2".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_kb(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Only 3 device buffers: a, b, out (init works in-place)
    let d_a = random_u32s(n, 12345).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n, 67890).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n, bench_ctx());

    // Benchmark init (in-place: raw u32 -> Kb) - also initializes d_a
    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_kb(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_kb(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_kb(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_kb(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_kb(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Kb".into(),
        bits: 31,
        u32s_per_element: 1,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_kb5(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Kb5 has 5 u32s per element
    let d_a = random_u32s(n * 5, 11111).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 5, 22222).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 5, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_kb5(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_kb5(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_kb5(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_kb5(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_kb5(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Kb5".into(),
        bits: 155,
        u32s_per_element: 5,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_kb6(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Kb6 has 6 u32s per element
    let d_a = random_u32s(n * 6, 33333).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 44444).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_kb6(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_kb6(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_kb6(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_kb6(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_kb6(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Kb6".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_kb2x3(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Kb2x3 has 6 u32s per element
    let d_a = random_u32s(n * 6, 55555).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 66666).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_kb2x3(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_kb2x3(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_kb2x3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_kb2x3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_kb2x3(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Kb2x3".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_kb3x2(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Kb3x2 has 6 u32s per element
    let d_a = random_u32s(n * 6, 77777).to_device_on(bench_ctx()).unwrap();
    let d_b = random_u32s(n * 6, 88888).to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u32>::with_capacity_on(n * 6, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_kb3x2(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_kb3x2(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_kb3x2(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_kb3x2(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_kb3x2(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    FieldBenchResult {
        field_name: "Kb3x2".into(),
        bits: 186,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_gl(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Goldilocks uses u64 elements (8 bytes)
    let h_a = random_u64s(n, 99999);
    let h_b = random_u64s(n, 88888);

    let d_a = h_a.to_device_on(bench_ctx()).unwrap();
    let d_b = h_b.to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u64>::with_capacity_on(n, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_gl(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_gl(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_gl(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_gl(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_gl(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    // u32s_per_element = 2 for Goldilocks (64-bit)
    FieldBenchResult {
        field_name: "Gl".into(),
        bits: 64,
        u32s_per_element: 2,
        init,
        add,
        mul,
        inv,
    }
}

pub fn bench_gl3(config: &BenchConfig) -> FieldBenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // Gl3 uses 3 x u64 elements (24 bytes)
    let h_a = random_u64s(n * 3, 77777);
    let h_b = random_u64s(n * 3, 66666);

    let d_a = h_a.to_device_on(bench_ctx()).unwrap();
    let d_b = h_b.to_device_on(bench_ctx()).unwrap();
    let d_out = DeviceBuffer::<u64>::with_capacity_on(n * 3, bench_ctx());

    let init = measure(config, n as u64, || {
        cuda_check(unsafe { init_gl3(d_a.as_mut_raw_ptr(), d_a.as_ptr(), n) });
    });
    cuda_check(unsafe { init_gl3(d_b.as_mut_raw_ptr(), d_b.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;

    let add = measure(config, ops, || {
        cuda_check(unsafe {
            add_gl3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let mul = measure(config, ops, || {
        cuda_check(unsafe {
            mul_gl3(
                d_out.as_mut_raw_ptr(),
                d_a.as_raw_ptr(),
                d_b.as_raw_ptr(),
                n,
                reps,
            )
        });
    });

    let inv = measure(config, ops, || {
        cuda_check(unsafe { inv_gl3(d_out.as_mut_raw_ptr(), d_a.as_raw_ptr(), n, reps) });
    });

    // u32s_per_element = 6 for Gl3 (3 x 64-bit = 24 bytes)
    FieldBenchResult {
        field_name: "Gl3".into(),
        bits: 192,
        u32s_per_element: 6,
        init,
        add,
        mul,
        inv,
    }
}

// ============================================================================
// Poseidon2 Benchmark
// ============================================================================

#[derive(Debug, Clone)]
pub struct Poseidon2BenchResult {
    pub name: String,
    pub avg_time_ms: f64,
    pub throughput_gops: f64,
}

pub fn bench_poseidon2_bb(config: &BenchConfig) -> Poseidon2BenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // 16 u32s per poseidon2 state
    let d_states = random_u32s(n * 16, 55555)
        .to_device_on(bench_ctx())
        .unwrap();

    // Initialize: raw u32 -> Fp field elements
    cuda_check(unsafe { init_poseidon2_bb(d_states.as_mut_raw_ptr(), d_states.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;
    let result = measure(config, ops, || {
        cuda_check(unsafe { run_poseidon2_bb(d_states.as_mut_raw_ptr(), n, reps) });
    });

    Poseidon2BenchResult {
        name: "BB Poseidon2".into(),
        avg_time_ms: result.avg_time_ms,
        throughput_gops: result.throughput_gops,
    }
}

pub fn bench_poseidon2_kb(config: &BenchConfig) -> Poseidon2BenchResult {
    let n = config.num_elements;
    let reps = config.ops_per_element;

    // 16 u32s per poseidon2 state
    let d_states = random_u32s(n * 16, 66666)
        .to_device_on(bench_ctx())
        .unwrap();

    // Initialize: raw u32 -> Kb field elements
    cuda_check(unsafe { init_poseidon2_kb(d_states.as_mut_raw_ptr(), d_states.as_ptr(), n) });
    sync();

    let ops = n as u64 * reps as u64;
    let result = measure(config, ops, || {
        cuda_check(unsafe { run_poseidon2_kb(d_states.as_mut_raw_ptr(), n, reps) });
    });

    Poseidon2BenchResult {
        name: "KB Poseidon2".into(),
        avg_time_ms: result.avg_time_ms,
        throughput_gops: result.throughput_gops,
    }
}

pub fn run_all_benchmarks(config: &BenchConfig) {
    println!("=== Extension Field Benchmark ===");
    println!();
    println!("Elements: {}", config.num_elements);
    println!("Ops per element: {}", config.ops_per_element);
    println!(
        "Warmup: {}, Bench iters: {}",
        config.warmup_iters, config.bench_iters
    );
    println!();

    // Collect all BabyBear results
    let fp = bench_fp(config);
    let fpext = bench_fpext(config);
    let fp5 = bench_fp5(config);
    let fp6 = bench_fp6(config);
    let fp2x3 = bench_fp2x3(config);
    let fp3x2 = bench_fp3x2(config);

    // Collect all KoalaBear results
    let kb = bench_kb(config);
    let kb5 = bench_kb5(config);
    let kb6 = bench_kb6(config);
    let kb2x3 = bench_kb2x3(config);
    let kb3x2 = bench_kb3x2(config);

    // Collect Goldilocks results
    let gl = bench_gl(config);
    let gl3 = bench_gl3(config);

    // Combine all results and sort by bits (then by name for consistent ordering)
    let mut all_results = vec![
        fp.clone(),
        fpext,
        fp5,
        fp6,
        fp2x3,
        fp3x2,
        kb,
        kb5,
        kb6,
        kb2x3,
        kb3x2,
        gl,
        gl3,
    ];
    all_results.sort_by(|a, b| {
        a.bits
            .cmp(&b.bits)
            .then_with(|| a.field_name.cmp(&b.field_name))
    });
    print_benchmark_tables(&all_results, &fp);

    // Poseidon2 benchmarks
    println!("### Poseidon2 Permutation");
    println!();
    let p2_bb = bench_poseidon2_bb(config);
    let p2_kb = bench_poseidon2_kb(config);
    println!(
        "| {:15} | {:>10} | {:>12} |",
        "Permutation", "Time (ms)", "Gops/s"
    );
    println!("|{:-<17}|{:-<12}|{:-<14}|", "", "", "");
    println!(
        "| {:15} | {:>10.3} | {:>12.1} |",
        p2_bb.name, p2_bb.avg_time_ms, p2_bb.throughput_gops
    );
    println!(
        "| {:15} | {:>10.3} | {:>12.1} |",
        p2_kb.name, p2_kb.avg_time_ms, p2_kb.throughput_gops
    );
    if p2_bb.throughput_gops > 0.0 {
        println!();
        println!(
            "KB/BB speedup: {:.2}x",
            p2_kb.throughput_gops / p2_bb.throughput_gops
        );
    }
    println!();
}