warriors/trisha/wgpu/accelerator.rs

//! GPU-accelerated Tip5 batch hashing via wgpu.
//!
//! Implements `triton_vm::gpu::GpuAccelerator` to offload Tip5 hash_varlen
//! to the GPU during STARK proving. The prover calls `hash_varlen_batch`
//! to hash all LDE table rows into Merkle tree leaves — the dominant cost
//! in Merkle tree construction.
//!
//! Uses the `tip5.wgsl` compute shader with Goldilocks field arithmetic.

use std::collections::HashMap;
use std::sync::Mutex;

use wgpu;
use wgpu::util::DeviceExt;

use triton_vm::gpu::GpuAccelerator;
use twenty_first::math::traits::PrimitiveRootOfUnity;
use twenty_first::prelude::*;

const DIGEST_LEN: usize = 5;
const WORKGROUP_SIZE: u32 = 64;

/// Tip5 constants for GPU upload.
struct Tip5Constants {
    /// 256-entry split-and-lookup S-box table (as u32 for GPU).
    lookup_table: [u32; 256],
    /// MDS first column as canonical Goldilocks field elements (lo, hi pairs).
    mds_column: [[u32; 2]; 16],
    /// Round constants as canonical Goldilocks field elements (lo, hi pairs).
    round_constants: [[u32; 2]; 80],
}

impl Tip5Constants {
    fn load() -> Self {
        use twenty_first::tip5::LOOKUP_TABLE;
        use twenty_first::tip5::MDS_MATRIX_FIRST_COLUMN;
        use twenty_first::tip5::ROUND_CONSTANTS;

        let lookup_table: [u32; 256] = {
            let mut table = [0u32; 256];
            for (i, &v) in LOOKUP_TABLE.iter().enumerate() {
                table[i] = v as u32;
            }
            table
        };

        let mds_column: [[u32; 2]; 16] = {
            let mut col = [[0u32; 2]; 16];
            for (i, &v) in MDS_MATRIX_FIRST_COLUMN.iter().enumerate() {
                // MDS coefficients are small positive i64 values.
                // Convert to Montgomery representation via BFieldElement::new().
                let bfe = BFieldElement::new(v as u64);
                let raw = bfe.raw_u64();
                col[i] = [raw as u32, (raw >> 32) as u32];
            }
            col
        };

        let round_constants: [[u32; 2]; 80] = {
            let mut rc = [[0u32; 2]; 80];
            for (i, &bfe) in ROUND_CONSTANTS.iter().enumerate() {
                // Round constants are already BFieldElements in Montgomery form.
                let raw = bfe.raw_u64();
                rc[i] = [raw as u32, (raw >> 32) as u32];
            }
            rc
        };

        Tip5Constants {
            lookup_table,
            mds_column,
            round_constants,
        }
    }
}

/// wgpu-based GPU accelerator implementing triton-vm's GpuAccelerator.
///
/// Accelerates:
/// - Tip5 batch hashing (Merkle tree leaf construction)
/// - NTT/iNTT (polynomial interpolation during proving)
/// - Merkle tree construction (Tip5 hash_pair)
/// - FRI fold (split-and-fold with XFE arithmetic)
pub struct WgpuTip5Accelerator {
    device: wgpu::Device,
    queue: wgpu::Queue,
    // Tip5 pipeline + constants
    tip5_pipeline: wgpu::ComputePipeline,
    lookup_buf: wgpu::Buffer,
    mds_buf: wgpu::Buffer,
    rc_buf: wgpu::Buffer,
    // Merkle tree pipelines
    _hash_pair_pipeline: wgpu::ComputePipeline,
    hash_pair_flat_pipeline: wgpu::ComputePipeline,
    // NTT pipelines
    ntt_butterfly_pipeline: wgpu::ComputePipeline,
    ntt_normalize_pipeline: wgpu::ComputePipeline,
    // FRI fold pipeline
    fri_fold_pipeline: wgpu::ComputePipeline,
    // GEMV pipelines
    gemv_bfe_pipeline: wgpu::ComputePipeline,
    gemv_xfe_pipeline: wgpu::ComputePipeline,
    // Mining pipeline
    mine_pipeline: wgpu::ComputePipeline,
    // Persistent buffers
    twiddle_cache: Mutex<HashMap<(usize, bool), wgpu::Buffer>>,
    two_inverse_buf: wgpu::Buffer,
    // VRAM budget and staging pool
    vram_budget: u64,
    staging_pool: Mutex<Vec<(u64, wgpu::Buffer)>>,
}

impl WgpuTip5Accelerator {
    /// Create a new accelerator using an existing wgpu device and queue.
    pub fn new(device: wgpu::Device, queue: wgpu::Queue) -> Self {
        let goldilocks_src = include_str!("shaders/goldilocks.wgsl");

        // Compile Tip5 shader
        let tip5_src = include_str!("shaders/tip5.wgsl");
        let tip5_full = format!("{}\n{}", goldilocks_src, tip5_src);
        let tip5_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("tip5"),
            source: wgpu::ShaderSource::Wgsl(tip5_full.into()),
        });
        let tip5_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("tip5_hash_rows"),
            layout: None,
            module: &tip5_module,
            entry_point: Some("hash_rows"),
            compilation_options: Default::default(),
            cache: None,
        });
        let hash_pair_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("tip5_hash_pair"),
            layout: None,
            module: &tip5_module,
            entry_point: Some("hash_pair"),
            compilation_options: Default::default(),
            cache: None,
        });
        let hash_pair_flat_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("tip5_hash_pair_flat"),
                layout: None,
                module: &tip5_module,
                entry_point: Some("hash_pair_flat"),
                compilation_options: Default::default(),
                cache: None,
            });

        // Compile FRI fold shader
        let fri_src = include_str!("shaders/fri.wgsl");
        let fri_full = format!("{}\n{}", goldilocks_src, fri_src);
        let fri_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("fri"),
            source: wgpu::ShaderSource::Wgsl(fri_full.into()),
        });
        let fri_fold_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("fri_fold_round"),
            layout: None,
            module: &fri_module,
            entry_point: Some("fri_fold_round"),
            compilation_options: Default::default(),
            cache: None,
        });

        // Compile NTT shader
        let ntt_src = include_str!("shaders/ntt.wgsl");
        let ntt_full = format!("{}\n{}", goldilocks_src, ntt_src);
        let ntt_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("ntt"),
            source: wgpu::ShaderSource::Wgsl(ntt_full.into()),
        });
        let ntt_butterfly_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("ntt_butterfly"),
                layout: None,
                module: &ntt_module,
                entry_point: Some("ntt_butterfly"),
                compilation_options: Default::default(),
                cache: None,
            });
        let ntt_normalize_pipeline =
            device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
                label: Some("ntt_normalize"),
                layout: None,
                module: &ntt_module,
                entry_point: Some("ntt_normalize"),
                compilation_options: Default::default(),
                cache: None,
            });

        // Compile GEMV shader
        let gemv_src = include_str!("shaders/gemv.wgsl");
        let gemv_full = format!("{}\n{}", goldilocks_src, gemv_src);
        let gemv_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("gemv"),
            source: wgpu::ShaderSource::Wgsl(gemv_full.into()),
        });
        let gemv_bfe_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("gemv_bfe"),
            layout: None,
            module: &gemv_module,
            entry_point: Some("gemv_bfe"),
            compilation_options: Default::default(),
            cache: None,
        });
        let gemv_xfe_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("gemv_xfe"),
            layout: None,
            module: &gemv_module,
            entry_point: Some("gemv_xfe"),
            compilation_options: Default::default(),
            cache: None,
        });

        // Compile mining shader
        let mine_src = include_str!("shaders/mine.wgsl");
        let mine_full = format!("{}\n{}", goldilocks_src, mine_src);
        let mine_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("mine"),
            source: wgpu::ShaderSource::Wgsl(mine_full.into()),
        });
        let mine_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("mine"),
            layout: None,
            module: &mine_module,
            entry_point: Some("mine"),
            compilation_options: Default::default(),
            cache: None,
        });

        // Upload Tip5 constants
        let constants = Tip5Constants::load();

        let lookup_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("tip5_lookup"),
            contents: bytemuck::cast_slice(&constants.lookup_table),
            usage: wgpu::BufferUsages::STORAGE,
        });

        let mds_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("tip5_mds"),
            contents: bytemuck::cast_slice(&constants.mds_column),
            usage: wgpu::BufferUsages::STORAGE,
        });

        let rc_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("tip5_round_constants"),
            contents: bytemuck::cast_slice(&constants.round_constants),
            usage: wgpu::BufferUsages::STORAGE,
        });

        // Precompute two_inverse = XFE(2).inverse() — used by every FRI fold
        let two_inv = XFieldElement::new([
            BFieldElement::new(2),
            BFieldElement::new(0),
            BFieldElement::new(0),
        ])
        .inverse();
        let two_inv_data: [[u32; 2]; 3] = {
            let mut data = [[0u32; 2]; 3];
            for (i, bfe) in two_inv.coefficients.iter().enumerate() {
                let raw = bfe.raw_u64();
                data[i] = [raw as u32, (raw >> 32) as u32];
            }
            data
        };
        let two_inverse_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("fri_two_inv_persistent"),
            contents: bytemuck::cast_slice(&two_inv_data),
            usage: wgpu::BufferUsages::STORAGE,
        });

        // Compute VRAM budget: use 75% of max_buffer_size to leave headroom.
        // max_storage_buffer_binding_size is the hard per-binding limit.
        let limits = device.limits();
        let per_binding = limits.max_storage_buffer_binding_size as u64;
        let per_buffer = limits.max_buffer_size;
        let effective_max = per_binding.min(per_buffer);
        let vram_budget = (effective_max * 3) / 4; // 75%
        eprintln!(
            "GPU: VRAM budget {:.0} MB (max_buffer {} MB, max_binding {} MB)",
            vram_budget as f64 / (1024.0 * 1024.0),
            per_buffer as f64 / (1024.0 * 1024.0),
            per_binding as f64 / (1024.0 * 1024.0),
        );

        WgpuTip5Accelerator {
            device,
            queue,
            tip5_pipeline,
            _hash_pair_pipeline: hash_pair_pipeline,
            hash_pair_flat_pipeline,
            lookup_buf,
            mds_buf,
            rc_buf,
            ntt_butterfly_pipeline,
            ntt_normalize_pipeline,
            fri_fold_pipeline,
            gemv_bfe_pipeline,
            gemv_xfe_pipeline,
            mine_pipeline,
            twiddle_cache: Mutex::new(HashMap::new()),
            two_inverse_buf,
            vram_budget,
            staging_pool: Mutex::new(Vec::new()),
        }
    }
}

impl WgpuTip5Accelerator {
    /// Fused GPU NTT: all butterfly layers + optional normalization in a single
    /// command encoder submission. Eliminates per-layer CPU↔GPU sync.
    ///
    /// `omega` is the twiddle root: omega.inverse() for iNTT, omega for forward NTT.
    /// `normalize` adds a final multiply-by-1/n pass (iNTT only).
    /// Ensure twiddle factor buffer exists in cache for given domain size and direction.
    fn ensure_twiddle_cached(&self, n: usize, omega: BFieldElement, is_inverse: bool) {
        let mut cache = self.twiddle_cache.lock().unwrap();
        let key = (n, is_inverse);
        if cache.contains_key(&key) {
            return;
        }
        let mut twiddles = vec![[0u32; 2]; n / 2];
        let mut w = BFieldElement::new(1);
        for tw in twiddles.iter_mut() {
            let raw = w.raw_u64();
            *tw = [raw as u32, (raw >> 32) as u32];
            w *= omega;
        }
        let buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("ntt_twiddles_cached"),
                contents: bytemuck::cast_slice(&twiddles),
                usage: wgpu::BufferUsages::STORAGE,
            });
        cache.insert(key, buf);
    }

    /// Acquire a staging buffer (MAP_READ ￿ COPY_DST) from the pool, or create one.
    fn acquire_staging(&self, size: u64) -> wgpu::Buffer {
        let mut pool = self.staging_pool.lock().unwrap();
        // Find smallest buffer that fits
        if let Some(idx) = pool
            .iter()
            .enumerate()
            .filter(￿(_, (s, _))￿ *s >= size)
            .min_by_key(￿(_, (s, _))￿ *s)
            .map(￿(i, _)￿ i)
        {
            return pool.remove(idx).1;
        }
        drop(pool);
        self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("staging_pooled"),
            size,
            usage: wgpu::BufferUsages::MAP_READ ￿ wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        })
    }

    /// Return a staging buffer to the pool for reuse.
    fn release_staging(&self, size: u64, buf: wgpu::Buffer) {
        let mut pool = self.staging_pool.lock().unwrap();
        if pool.len() < 8 {
            pool.push((size, buf));
        }
        // else: drop the buffer (pool full)
    }

    /// GPU inner path for hash_varlen_batch — assumes inputs fit in VRAM.
    fn hash_varlen_batch_gpu(&self, inputs: &[&[BFieldElement]], row_len: usize) -> Vec<Digest> {
        let num_rows = inputs.len() as u32;

        // Flatten input: convert BFieldElement → Montgomery u64 → [u32; 2]
        let input_data: Vec<[u32; 2]> = inputs
            .iter()
            .flat_map(|row| {
                row.iter().map(|bfe| {
                    let raw = bfe.raw_u64();
                    [raw as u32, (raw >> 32) as u32]
                })
            })
            .collect();

        let input_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("tip5_input"),
                contents: bytemuck::cast_slice(&input_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let output_size = (num_rows as usize * DIGEST_LEN * 8) as u64;
        let output_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("tip5_output"),
            size: output_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let params = [num_rows, row_len as u32, 0u32, 0u32];
        let params_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("tip5_params"),
                contents: bytemuck::cast_slice(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        let bind_group_layout = self.tip5_pipeline.get_bind_group_layout(0);
        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("tip5_bind_group"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.lookup_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: self.mds_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: self.rc_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: input_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: output_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: params_buf.as_entire_binding(),
                },
            ],
        });

        let workgroups = (num_rows + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("tip5_encoder"),
            });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("tip5_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.tip5_pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(workgroups, 1, 1);
        }

        let staging_buf = self.acquire_staging(output_size);
        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging_buf, 0, output_size);
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("GPU readback channel closed")
            .expect("GPU readback failed");

        let data = slice.get_mapped_range();
        let output_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);

        let digests: Vec<Digest> = (0..num_rows as usize)
            .map(|i| {
                let base = i * DIGEST_LEN;
                let mut elements = [BFieldElement::new(0); DIGEST_LEN];
                for j in 0..DIGEST_LEN {
                    let [lo, hi] = output_pairs[base + j];
                    let raw = (hi as u64) << 32 | lo as u64;
                    elements[j] = BFieldElement::from_raw_u64(raw);
                }
                Digest::new(elements)
            })
            .collect();

        drop(data);
        staging_buf.unmap();
        self.release_staging(output_size, staging_buf);

        digests
    }

    fn gpu_ntt_core(&self, column: &mut [BFieldElement], omega: BFieldElement, normalize: bool) {
        let n = column.len();
        let n_u32 = n as u32;
        let log_n = n.trailing_zeros();

        // Bit-reversal permutation (CPU, before GPU butterflies)
        for k in 0..n {
            let rev_k = ((k as u32).reverse_bits() >> (32 - log_n)) as usize;
            if k < rev_k {
                column.swap(k, rev_k);
            }
        }

        // Upload column data (new each call — data changes)
        let col_data: Vec<[u32; 2]> = column
            .iter()
            .map(|bfe| {
                let raw = bfe.raw_u64();
                [raw as u32, (raw >> 32) as u32]
            })
            .collect();

        let data_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("ntt_data"),
                contents: bytemuck::cast_slice(&col_data),
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            });

        // Twiddle factors: cached per (domain_size, direction).
        // ensure_twiddle_cached populates the cache; we hold the lock for bind group creation.
        self.ensure_twiddle_cached(n, omega, normalize);
        let cache = self.twiddle_cache.lock().unwrap();
        let twiddle_buf = cache.get(&(n, normalize)).unwrap();

        // Pre-create all per-layer params buffers so bind groups can reference them
        let layer_params: Vec<wgpu::Buffer> = (0..log_n)
            .map(|layer| {
                let params = [n_u32, layer, 0u32, 0u32];
                self.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("ntt_params"),
                        contents: bytemuck::cast_slice(&params),
                        usage: wgpu::BufferUsages::UNIFORM,
                    })
            })
            .collect();

        let bg_layout = self.ntt_butterfly_pipeline.get_bind_group_layout(0);
        let layer_bind_groups: Vec<wgpu::BindGroup> = layer_params
            .iter()
            .map(|params_buf| {
                self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("ntt_bg"),
                    layout: &bg_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: data_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: twiddle_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: params_buf.as_entire_binding(),
                        },
                    ],
                })
            })
            .collect();

        let n_butterflies = n_u32 / 2;
        let workgroups = (n_butterflies + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;

        // Single command encoder for ALL butterfly layers + normalization + readback
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("ntt_fused_encoder"),
            });

        // Dispatch all butterfly layers in one compute pass.
        // wgpu guarantees sequential execution within a pass — layer N sees
        // layer N-1's writes to the data buffer.
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("ntt_butterfly_fused"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.ntt_butterfly_pipeline);
            for bg in &layer_bind_groups {
                pass.set_bind_group(0, bg, &[]);
                pass.dispatch_workgroups(workgroups, 1, 1);
            }
        }

        // Optional normalization pass (iNTT: multiply all by 1/n)
        if normalize {
            let n_inv = BFieldElement::new(n as u64).inverse();
            let n_inv_raw = n_inv.raw_u64();
            let norm_twiddle = [[n_inv_raw as u32, (n_inv_raw >> 32) as u32]];
            let norm_tw_buf = self
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("ntt_norm_tw"),
                    contents: bytemuck::cast_slice(&norm_twiddle),
                    usage: wgpu::BufferUsages::STORAGE,
                });
            let norm_params = [n_u32, 0u32, 0u32, 0u32];
            let norm_params_buf =
                self.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("ntt_norm_params"),
                        contents: bytemuck::cast_slice(&norm_params),
                        usage: wgpu::BufferUsages::UNIFORM,
                    });
            let norm_bg_layout = self.ntt_normalize_pipeline.get_bind_group_layout(0);
            let norm_bg = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("ntt_norm_bg"),
                layout: &norm_bg_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: data_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: norm_tw_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: norm_params_buf.as_entire_binding(),
                    },
                ],
            });
            let norm_workgroups = (n_u32 + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
            {
                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("ntt_normalize"),
                    timestamp_writes: None,
                });
                pass.set_pipeline(&self.ntt_normalize_pipeline);
                pass.set_bind_group(0, &norm_bg, &[]);
                pass.dispatch_workgroups(norm_workgroups, 1, 1);
            }
        }

        // Readback in the same submission
        let buf_size = (n * 8) as u64;
        let staging_buf = self.acquire_staging(buf_size);
        encoder.copy_buffer_to_buffer(&data_buf, 0, &staging_buf, 0, buf_size);

        // Single submit — entire NTT runs on GPU without CPU intervention
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("NTT readback channel closed")
            .expect("NTT readback failed");

        let data = slice.get_mapped_range();
        let result_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);
        for (i, &[lo, hi]) in result_pairs.iter().enumerate() {
            let raw = (hi as u64) << 32 | lo as u64;
            column[i] = BFieldElement::from_raw_u64(raw);
        }
        drop(data);
        staging_buf.unmap();
        self.release_staging(buf_size, staging_buf);
    }

    /// GPU nonce mining: search for nonce such that Tip5(message ++ nonce) < target.
    ///
    /// `message` is the base message in Montgomery form BFieldElements.
    /// `target` is the difficulty target as a canonical u64.
    /// `max_attempts` limits total nonces tried.
    /// Returns `(nonce_raw_u64, digest_mont_values, attempts)` or None.
    pub fn mine(
        &self,
        message: &[BFieldElement],
        target: u64,
        max_attempts: u64,
    ) -> Option<(u64, Vec<BFieldElement>, u64)> {
        let msg_len = message.len() as u32;
        let batch_size: u64 = 256 * 1024; // 256 workgroups × 256 threads

        // Upload base message
        let msg_data: Vec<[u32; 2]> = message
            .iter()
            .map(|bfe| {
                let raw = bfe.raw_u64();
                [raw as u32, (raw >> 32) as u32]
            })
            .collect();
        // Ensure at least 1 element for empty message
        let msg_upload = if msg_data.is_empty() {
            vec![[0u32; 2]]
        } else {
            msg_data
        };
        let msg_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mine_message"),
                contents: bytemuck::cast_slice(&msg_upload),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let target_lo = target as u32;
        let target_hi = (target >> 32) as u32;

        // Result buffer: [found, nonce_lo, nonce_hi, d0_lo, d0_hi, ..., d4_lo, d4_hi] = 13 u32s
        let result_size = 13u64 * 4;
        let result_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("mine_result"),
            size: result_size,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });

        // Target uniform
        let target_data = [target_hi, 0u32, 0u32, 0u32];
        let target_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("mine_target"),
                contents: bytemuck::cast_slice(&target_data),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        let mut attempts: u64 = 0;
        while attempts < max_attempts {
            // Clear result buffer
            let zero_data = [0u32; 13];
            let zero_buf = self
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("mine_zero"),
                    contents: bytemuck::cast_slice(&zero_data),
                    usage: wgpu::BufferUsages::COPY_SRC,
                });
            let mut clear_enc =
                self.device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some("mine_clear"),
                    });
            clear_enc.copy_buffer_to_buffer(&zero_buf, 0, &result_buf, 0, result_size);
            self.queue.submit(std::iter::once(clear_enc.finish()));

            // Nonce offset in Montgomery form
            let nonce_bfe = BFieldElement::new(attempts);
            let nonce_raw = nonce_bfe.raw_u64();
            let params_data = [
                msg_len,
                nonce_raw as u32,
                (nonce_raw >> 32) as u32,
                target_lo,
            ];
            let params_buf = self
                .device
                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                    label: Some("mine_params"),
                    contents: bytemuck::cast_slice(&params_data),
                    usage: wgpu::BufferUsages::UNIFORM,
                });

            let bg_layout = self.mine_pipeline.get_bind_group_layout(0);
            let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("mine_bg"),
                layout: &bg_layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: msg_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: self.lookup_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 2,
                        resource: self.mds_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 3,
                        resource: self.rc_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 4,
                        resource: params_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 5,
                        resource: target_buf.as_entire_binding(),
                    },
                    wgpu::BindGroupEntry {
                        binding: 6,
                        resource: result_buf.as_entire_binding(),
                    },
                ],
            });

            let workgroups = (batch_size as u32) / 256;
            let mut encoder = self
                .device
                .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                    label: Some("mine_encoder"),
                });
            {
                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("mine_pass"),
                    timestamp_writes: None,
                });
                pass.set_pipeline(&self.mine_pipeline);
                pass.set_bind_group(0, &bind_group, &[]);
                pass.dispatch_workgroups(workgroups, 1, 1);
            }

            // Readback result
            let staging_buf = self.acquire_staging(result_size);
            encoder.copy_buffer_to_buffer(&result_buf, 0, &staging_buf, 0, result_size);
            self.queue.submit(std::iter::once(encoder.finish()));

            let slice = staging_buf.slice(..);
            let (tx, rx) = std::sync::mpsc::channel();
            slice.map_async(wgpu::MapMode::Read, move |r| {
                tx.send(r).unwrap();
            });
            self.device.poll(wgpu::Maintain::Wait);
            rx.recv()
                .expect("mine readback channel closed")
                .expect("mine readback failed");

            let data = slice.get_mapped_range();
            let result_u32: &[u32] = bytemuck::cast_slice(&data);

            if result_u32[0] != 0 {
                // Found! Extract nonce and digest
                let nonce_lo = result_u32[1];
                let nonce_hi = result_u32[2];
                let nonce_mont_raw = (nonce_hi as u64) << 32 | nonce_lo as u64;
                let nonce_bfe = BFieldElement::from_raw_u64(nonce_mont_raw);
                let nonce_val = nonce_bfe.value();

                let mut digest = Vec::with_capacity(5);
                for i in 0..5 {
                    let lo = result_u32[3 + i * 2];
                    let hi = result_u32[3 + i * 2 + 1];
                    let raw = (hi as u64) << 32 | lo as u64;
                    digest.push(BFieldElement::from_raw_u64(raw));
                }

                drop(data);
                staging_buf.unmap();
                self.release_staging(result_size, staging_buf);
                return Some((nonce_val, digest, attempts + batch_size));
            }

            drop(data);
            staging_buf.unmap();
            self.release_staging(result_size, staging_buf);
            attempts += batch_size;
        }

        None
    }
}

impl GpuAccelerator for WgpuTip5Accelerator {
    fn name(&self) -> &str {
        "wgpu-tip5"
    }

    fn hash_varlen_batch(&self, inputs: &[&[BFieldElement]]) -> Vec<Digest> {
        if inputs.is_empty() {
            return vec![];
        }

        // All rows in a single proving call have the same length (LDE table rows).
        // Verify this assumption and use uniform row length for GPU dispatch.
        let row_len = inputs[0].len();
        let uniform = inputs.iter().all(|r| r.len() == row_len);

        if !uniform || row_len == 0 {
            // Fall back to CPU for non-uniform or empty rows.
            return inputs
                .iter()
                .map(|input| Tip5::hash_varlen(input))
                .collect();
        }

        // Check if total GPU memory exceeds budget; if so, process in chunks.
        let bytes_per_row = row_len * 8; // input
        let bytes_per_digest = DIGEST_LEN * 8; // output
        let total_input_bytes = inputs.len() * bytes_per_row;
        let total_output_bytes = inputs.len() * bytes_per_digest;
        let total_bytes = (total_input_bytes + total_output_bytes) as u64;

        if total_bytes > self.vram_budget {
            // Chunk: fit both input + output within budget
            let per_row_total = (bytes_per_row + bytes_per_digest) as u64;
            let chunk_rows = (self.vram_budget / per_row_total).max(1) as usize;
            let mut all_digests = Vec::with_capacity(inputs.len());
            for chunk_start in (0..inputs.len()).step_by(chunk_rows) {
                let chunk_end = (chunk_start + chunk_rows).min(inputs.len());
                let chunk = &inputs[chunk_start..chunk_end];
                let chunk_digests = self.hash_varlen_batch_gpu(chunk, row_len);
                all_digests.extend(chunk_digests);
            }
            return all_digests;
        }

        self.hash_varlen_batch_gpu(inputs, row_len)
    }

    fn intt_bfe(&self, column: &mut [BFieldElement]) {
        let n = column.len();
        if n <= 1 || !n.is_power_of_two() {
            twenty_first::math::ntt::intt(column);
            return;
        }
        // CPU fallback for small domains or oversized columns
        if n < 1024 || (n as u64 * 8) > self.vram_budget {
            twenty_first::math::ntt::intt(column);
            return;
        }

        let omega = BFieldElement::primitive_root_of_unity(n as u64).unwrap();
        self.gpu_ntt_core(column, omega.inverse(), true);
    }

    fn intt_xfe(&self, column: &mut [XFieldElement]) {
        let n = column.len();
        if n <= 1 || !n.is_power_of_two() {
            twenty_first::math::ntt::intt(column);
            return;
        }

        // CPU fallback for small domains or oversized columns (3 BFE columns)
        if n < 1024 || (n as u64 * 8) > self.vram_budget {
            twenty_first::math::ntt::intt(column);
            return;
        }

        // Deinterleave: split XFE column into 3 independent BFE columns.
        // iNTT is linear over BFE, so iNTT(xfe_col) = reassemble(iNTT(c0), iNTT(c1), iNTT(c2)).
        let mut c0 = vec![BFieldElement::new(0); n];
        let mut c1 = vec![BFieldElement::new(0); n];
        let mut c2 = vec![BFieldElement::new(0); n];
        for (i, xfe) in column.iter().enumerate() {
            c0[i] = xfe.coefficients[0];
            c1[i] = xfe.coefficients[1];
            c2[i] = xfe.coefficients[2];
        }

        // Run 3 independent BFE iNTTs on GPU
        self.intt_bfe(&mut c0);
        self.intt_bfe(&mut c1);
        self.intt_bfe(&mut c2);

        // Reinterleave back into XFE column
        for (i, xfe) in column.iter_mut().enumerate() {
            xfe.coefficients[0] = c0[i];
            xfe.coefficients[1] = c1[i];
            xfe.coefficients[2] = c2[i];
        }
    }

    fn ntt_bfe(&self, column: &mut [BFieldElement]) {
        let n = column.len();
        if n <= 1 || !n.is_power_of_two() {
            twenty_first::math::ntt::ntt(column);
            return;
        }
        if n < 1024 || (n as u64 * 8) > self.vram_budget {
            twenty_first::math::ntt::ntt(column);
            return;
        }

        let omega = BFieldElement::primitive_root_of_unity(n as u64).unwrap();
        self.gpu_ntt_core(column, omega, false);
    }

    fn ntt_xfe(&self, column: &mut [XFieldElement]) {
        let n = column.len();
        if n <= 1 || !n.is_power_of_two() {
            twenty_first::math::ntt::ntt(column);
            return;
        }

        if n < 1024 || (n as u64 * 8) > self.vram_budget {
            twenty_first::math::ntt::ntt(column);
            return;
        }

        // Deinterleave: split XFE column into 3 independent BFE columns.
        // NTT is linear over BFE, so NTT(xfe_col) = reassemble(NTT(c0), NTT(c1), NTT(c2)).
        let mut c0 = vec![BFieldElement::new(0); n];
        let mut c1 = vec![BFieldElement::new(0); n];
        let mut c2 = vec![BFieldElement::new(0); n];
        for (i, xfe) in column.iter().enumerate() {
            c0[i] = xfe.coefficients[0];
            c1[i] = xfe.coefficients[1];
            c2[i] = xfe.coefficients[2];
        }

        // Run 3 independent BFE forward NTTs on GPU
        self.ntt_bfe(&mut c0);
        self.ntt_bfe(&mut c1);
        self.ntt_bfe(&mut c2);

        // Reinterleave back into XFE column
        for (i, xfe) in column.iter_mut().enumerate() {
            xfe.coefficients[0] = c0[i];
            xfe.coefficients[1] = c1[i];
            xfe.coefficients[2] = c2[i];
        }
    }

    fn gemv_bfe(
        &self,
        matrix: &[BFieldElement],
        nrows: usize,
        ncols: usize,
        weights: &[XFieldElement],
    ) -> Vec<XFieldElement> {
        let cpu_gemv_bfe = |mat: &[BFieldElement],
                            nr: usize,
                            nc: usize,
                            w: &[XFieldElement]|
         -> Vec<XFieldElement> {
            (0..nr)
                .map(|i| {
                    let row = &mat[i * nc..(i + 1) * nc];
                    row.iter()
                        .zip(w.iter())
                        .map(|(&m, &wt)| wt * m)
                        .fold(XFieldElement::new([BFieldElement::new(0); 3]), |acc, x| {
                            acc + x
                        })
                })
                .collect()
        };

        if nrows < 256 || ncols < 4 {
            return cpu_gemv_bfe(matrix, nrows, ncols, weights);
        }

        // VRAM guard: chunk rows if matrix doesn't fit
        let matrix_bytes = (nrows * ncols * 8) as u64;
        if matrix_bytes > self.vram_budget {
            let rows_per_chunk = (self.vram_budget / (ncols as u64 * 8)).max(256) as usize;
            let mut result = Vec::with_capacity(nrows);
            for chunk_start in (0..nrows).step_by(rows_per_chunk) {
                let chunk_end = (chunk_start + rows_per_chunk).min(nrows);
                let chunk_nrows = chunk_end - chunk_start;
                let chunk_matrix = &matrix[chunk_start * ncols..chunk_end * ncols];
                if chunk_nrows < 256 {
                    result.extend(cpu_gemv_bfe(chunk_matrix, chunk_nrows, ncols, weights));
                } else {
                    result.extend(self.gemv_bfe(chunk_matrix, chunk_nrows, ncols, weights));
                }
            }
            return result;
        }

        // Flatten matrix: BFE → vec2<u32>
        let matrix_data: Vec<[u32; 2]> = matrix
            .iter()
            .map(|bfe| {
                let raw = bfe.raw_u64();
                [raw as u32, (raw >> 32) as u32]
            })
            .collect();

        // Flatten weights: XFE → 3 × vec2<u32>
        let weights_data: Vec<[u32; 2]> = weights
            .iter()
            .flat_map(|xfe| {
                xfe.coefficients.iter().map(|bfe| {
                    let raw = bfe.raw_u64();
                    [raw as u32, (raw >> 32) as u32]
                })
            })
            .collect();

        let matrix_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_matrix"),
                contents: bytemuck::cast_slice(&matrix_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let weights_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_weights"),
                contents: bytemuck::cast_slice(&weights_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        // Output: nrows XFEs = nrows * 3 * 8 bytes
        let output_size = (nrows * 3 * 8) as u64;
        let output_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("gemv_output"),
            size: output_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let params = [nrows as u32, ncols as u32, 0u32, 0u32];
        let params_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_params"),
                contents: bytemuck::cast_slice(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        let bind_group_layout = self.gemv_bfe_pipeline.get_bind_group_layout(0);
        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("gemv_bfe_bg"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: matrix_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: weights_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: output_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: params_buf.as_entire_binding(),
                },
            ],
        });

        let workgroups = (nrows as u32 + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("gemv_bfe_encoder"),
            });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("gemv_bfe_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.gemv_bfe_pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(workgroups, 1, 1);
        }

        // Read back results
        let staging_buf = self.acquire_staging(output_size);
        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging_buf, 0, output_size);
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("GEMV BFE readback channel closed")
            .expect("GEMV BFE readback failed");

        let data = slice.get_mapped_range();
        let result_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);

        let result: Vec<XFieldElement> = (0..nrows)
            .map(|i| {
                let base = i * 3;
                let mut coeffs = [BFieldElement::new(0); 3];
                for j in 0..3 {
                    let [lo, hi] = result_pairs[base + j];
                    let raw = (hi as u64) << 32 | lo as u64;
                    coeffs[j] = BFieldElement::from_raw_u64(raw);
                }
                XFieldElement::new(coeffs)
            })
            .collect();

        drop(data);
        staging_buf.unmap();
        self.release_staging(output_size, staging_buf);

        result
    }

    fn gemv_xfe(
        &self,
        matrix: &[XFieldElement],
        nrows: usize,
        ncols: usize,
        weights: &[XFieldElement],
    ) -> Vec<XFieldElement> {
        let cpu_gemv_xfe = |mat: &[XFieldElement],
                            nr: usize,
                            nc: usize,
                            w: &[XFieldElement]|
         -> Vec<XFieldElement> {
            (0..nr)
                .map(|i| {
                    let row = &mat[i * nc..(i + 1) * nc];
                    row.iter()
                        .zip(w.iter())
                        .map(|(&m, &wt)| m * wt)
                        .fold(XFieldElement::new([BFieldElement::new(0); 3]), |acc, x| {
                            acc + x
                        })
                })
                .collect()
        };

        if nrows < 256 || ncols < 4 {
            return cpu_gemv_xfe(matrix, nrows, ncols, weights);
        }

        // VRAM guard: chunk rows if XFE matrix doesn't fit (24 bytes per element)
        let matrix_bytes = (nrows * ncols * 24) as u64;
        if matrix_bytes > self.vram_budget {
            let rows_per_chunk = (self.vram_budget / (ncols as u64 * 24)).max(256) as usize;
            let mut result = Vec::with_capacity(nrows);
            for chunk_start in (0..nrows).step_by(rows_per_chunk) {
                let chunk_end = (chunk_start + rows_per_chunk).min(nrows);
                let chunk_nrows = chunk_end - chunk_start;
                let chunk_matrix = &matrix[chunk_start * ncols..chunk_end * ncols];
                if chunk_nrows < 256 {
                    result.extend(cpu_gemv_xfe(chunk_matrix, chunk_nrows, ncols, weights));
                } else {
                    result.extend(self.gemv_xfe(chunk_matrix, chunk_nrows, ncols, weights));
                }
            }
            return result;
        }

        // Flatten matrix: XFE → 3 × vec2<u32>
        let matrix_data: Vec<[u32; 2]> = matrix
            .iter()
            .flat_map(|xfe| {
                xfe.coefficients.iter().map(|bfe| {
                    let raw = bfe.raw_u64();
                    [raw as u32, (raw >> 32) as u32]
                })
            })
            .collect();

        // Flatten weights: XFE → 3 × vec2<u32>
        let weights_data: Vec<[u32; 2]> = weights
            .iter()
            .flat_map(|xfe| {
                xfe.coefficients.iter().map(|bfe| {
                    let raw = bfe.raw_u64();
                    [raw as u32, (raw >> 32) as u32]
                })
            })
            .collect();

        let matrix_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_xfe_matrix"),
                contents: bytemuck::cast_slice(&matrix_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let weights_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_xfe_weights"),
                contents: bytemuck::cast_slice(&weights_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let output_size = (nrows * 3 * 8) as u64;
        let output_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("gemv_xfe_output"),
            size: output_size,
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let params = [nrows as u32, ncols as u32, 0u32, 0u32];
        let params_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("gemv_xfe_params"),
                contents: bytemuck::cast_slice(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        let bind_group_layout = self.gemv_xfe_pipeline.get_bind_group_layout(0);
        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("gemv_xfe_bg"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: matrix_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: weights_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: output_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: params_buf.as_entire_binding(),
                },
            ],
        });

        let workgroups = (nrows as u32 + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("gemv_xfe_encoder"),
            });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("gemv_xfe_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.gemv_xfe_pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(workgroups, 1, 1);
        }

        let staging_buf = self.acquire_staging(output_size);
        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging_buf, 0, output_size);
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move |result| {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("GEMV XFE readback channel closed")
            .expect("GEMV XFE readback failed");

        let data = slice.get_mapped_range();
        let result_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);

        let result: Vec<XFieldElement> = (0..nrows)
            .map(|i| {
                let base = i * 3;
                let mut coeffs = [BFieldElement::new(0); 3];
                for j in 0..3 {
                    let [lo, hi] = result_pairs[base + j];
                    let raw = (hi as u64) << 32 | lo as u64;
                    coeffs[j] = BFieldElement::from_raw_u64(raw);
                }
                XFieldElement::new(coeffs)
            })
            .collect();

        drop(data);
        staging_buf.unmap();
        self.release_staging(output_size, staging_buf);

        result
    }

    fn fri_fold(
        &self,
        codeword: &[XFieldElement],
        domain_point_inverses: &[BFieldElement],
        folding_challenge: XFieldElement,
    ) -> Vec<XFieldElement> {
        let n = codeword.len();
        let half_n = n / 2;

        // CPU fallback for small, non-power-of-2, or oversized codewords
        // Codeword needs n*24 bytes (XFE) + n/2*8 (domain inverses)
        let codeword_bytes = (n as u64) * 24 + (half_n as u64) * 8;
        if half_n == 0 || !n.is_power_of_two() || half_n < 512 || codeword_bytes > self.vram_budget
        {
            let one = XFieldElement::new([
                BFieldElement::new(1),
                BFieldElement::new(0),
                BFieldElement::new(0),
            ]);
            let two_inverse = XFieldElement::new([
                BFieldElement::new(2),
                BFieldElement::new(0),
                BFieldElement::new(0),
            ])
            .inverse();
            return (0..half_n)
                .map(|i| {
                    let scaled_offset_inv = folding_challenge * domain_point_inverses[i];
                    let left_summand = (one + scaled_offset_inv) * codeword[i];
                    let right_summand = (one - scaled_offset_inv) * codeword[n / 2 + i];
                    (left_summand + right_summand) * two_inverse
                })
                .collect();
        }

        // Flatten codeword: each XFE = 3 BFE = 3 × vec2<u32>
        let codeword_data: Vec<[u32; 2]> = codeword
            .iter()
            .flat_map(|xfe| {
                xfe.coefficients.iter().map(|bfe| {
                    let raw = bfe.raw_u64();
                    [raw as u32, (raw >> 32) as u32]
                })
            })
            .collect();

        // Domain point inverses: BFE values
        let dinv_data: Vec<[u32; 2]> = domain_point_inverses
            .iter()
            .map(|bfe| {
                let raw = bfe.raw_u64();
                [raw as u32, (raw >> 32) as u32]
            })
            .collect();

        // Folding challenge: 3 BFE
        let challenge_data: [[u32; 2]; 3] = {
            let mut data = [[0u32; 2]; 3];
            for (i, bfe) in folding_challenge.coefficients.iter().enumerate() {
                let raw = bfe.raw_u64();
                data[i] = [raw as u32, (raw >> 32) as u32];
            }
            data
        };

        // Create GPU buffers (two_inverse uses persistent buffer)
        let codeword_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("fri_codeword"),
                contents: bytemuck::cast_slice(&codeword_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let dinv_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("fri_domain_inv"),
                contents: bytemuck::cast_slice(&dinv_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let challenge_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("fri_challenge"),
                contents: bytemuck::cast_slice(&challenge_data),
                usage: wgpu::BufferUsages::STORAGE,
            });

        let output_size = (half_n * 3 * 8) as u64; // half_n XFEs × 3 BFEs × 8 bytes
        let output_buf = self.device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("fri_folded"),
            size: output_size,
            usage: wgpu::BufferUsages::STORAGE ￿ wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        });

        let params = [half_n as u32, 0u32, 0u32, 0u32];
        let params_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("fri_params"),
                contents: bytemuck::cast_slice(&params),
                usage: wgpu::BufferUsages::UNIFORM,
            });

        // Create bind group
        let bind_group_layout = self.fri_fold_pipeline.get_bind_group_layout(0);
        let bind_group = self.device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some("fri_bg"),
            layout: &bind_group_layout,
            entries: &[
                wgpu::BindGroupEntry {
                    binding: 0,
                    resource: codeword_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 1,
                    resource: dinv_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 2,
                    resource: challenge_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 3,
                    resource: self.two_inverse_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 4,
                    resource: output_buf.as_entire_binding(),
                },
                wgpu::BindGroupEntry {
                    binding: 5,
                    resource: params_buf.as_entire_binding(),
                },
            ],
        });

        // Dispatch
        let workgroups = (half_n as u32 + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("fri_encoder"),
            });
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                label: Some("fri_fold_pass"),
                timestamp_writes: None,
            });
            pass.set_pipeline(&self.fri_fold_pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(workgroups, 1, 1);
        }

        // Read back results
        let staging_buf = self.acquire_staging(output_size);
        encoder.copy_buffer_to_buffer(&output_buf, 0, &staging_buf, 0, output_size);
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move ￿result￿ {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("FRI readback channel closed")
            .expect("FRI readback failed");

        let data = slice.get_mapped_range();
        let result_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);

        let folded: Vec<XFieldElement> = (0..half_n)
            .map(|i| {
                let base = i * 3;
                let mut coeffs = [BFieldElement::new(0); 3];
                for j in 0..3 {
                    let [lo, hi] = result_pairs[base + j];
                    let raw = (hi as u64) << 32 | lo as u64;
                    coeffs[j] = BFieldElement::from_raw_u64(raw);
                }
                XFieldElement::new(coeffs)
            })
            .collect();

        drop(data);
        staging_buf.unmap();
        self.release_staging(output_size, staging_buf);

        folded
    }

    fn merkle_tree(&self, leaves: &[Digest]) -> twenty_first::util_types::merkle_tree::MerkleTree {
        use twenty_first::util_types::merkle_tree::MerkleTree;

        let n = leaves.len();
        if n == 0 || !n.is_power_of_two() {
            return MerkleTree::par_new(leaves).unwrap();
        }
        // CPU fallback for small trees or if flat buffer exceeds VRAM
        let buf_bytes = (2 * n * 5 * 8) as u64;
        if n < 512 || buf_bytes > self.vram_budget {
            return MerkleTree::par_new(leaves).unwrap();
        }

        // Flat node array: 2n nodes × 5 digest elements. Index 0 unused, root at 1,
        // leaves at [n..2n). Entire tree lives on GPU — no per-level readback.
        let num_nodes = 2 * n;
        let digest_len = 5;

        let mut nodes_flat: Vec<[u32; 2]> = vec![[0u32; 2]; num_nodes * digest_len];
        for (i, leaf) in leaves.iter().enumerate() {
            let base = (n + i) * digest_len;
            for (j, &bfe) in leaf.0.iter().enumerate() {
                let raw = bfe.raw_u64();
                nodes_flat[base + j] = [raw as u32, (raw >> 32) as u32];
            }
        }

        let buf_size = (num_nodes * digest_len * 8) as u64;
        let nodes_buf = self
            .device
            .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                label: Some("merkle_nodes"),
                contents: bytemuck::cast_slice(&nodes_flat),
                usage: wgpu::BufferUsages::STORAGE ￿ wgpu::BufferUsages::COPY_SRC,
            });

        // Pre-create params buffers for each level
        let num_levels = (n as f64).log2() as u32;
        let mut level_params: Vec<wgpu::Buffer> = Vec::with_capacity(num_levels as usize);
        let mut level_size = n / 2;
        let mut parent_start = n / 2;
        while level_size >= 1 {
            let params = [level_size as u32, parent_start as u32, 0u32, 0u32];
            level_params.push(
                self.device
                    .create_buffer_init(&wgpu::util::BufferInitDescriptor {
                        label: Some("merkle_flat_params"),
                        contents: bytemuck::cast_slice(&params),
                        usage: wgpu::BufferUsages::UNIFORM,
                    }),
            );
            parent_start /= 2;
            level_size /= 2;
        }

        let bg_layout = self.hash_pair_flat_pipeline.get_bind_group_layout(0);
        let level_bind_groups: Vec<wgpu::BindGroup> = level_params
            .iter()
            .map(￿params_buf￿ {
                self.device.create_bind_group(&wgpu::BindGroupDescriptor {
                    label: Some("merkle_flat_bg"),
                    layout: &bg_layout,
                    entries: &[
                        wgpu::BindGroupEntry {
                            binding: 0,
                            resource: self.lookup_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 1,
                            resource: self.mds_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 2,
                            resource: self.rc_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 9,
                            resource: nodes_buf.as_entire_binding(),
                        },
                        wgpu::BindGroupEntry {
                            binding: 10,
                            resource: params_buf.as_entire_binding(),
                        },
                    ],
                })
            })
            .collect();

        // Single command encoder: all levels dispatched as separate compute passes
        // (each pass gets an implicit barrier for prior writes to nodes_buf)
        let mut encoder = self
            .device
            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                label: Some("merkle_fused_encoder"),
            });

        let mut dispatch_size = n / 2;
        for bg in &level_bind_groups {
            let workgroups = (dispatch_size as u32 + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE;
            {
                let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("merkle_level"),
                    timestamp_writes: None,
                });
                pass.set_pipeline(&self.hash_pair_flat_pipeline);
                pass.set_bind_group(0, bg, &[]);
                pass.dispatch_workgroups(workgroups, 1, 1);
            }
            dispatch_size /= 2;
        }

        // Single readback of entire tree
        let staging_buf = self.acquire_staging(buf_size);
        encoder.copy_buffer_to_buffer(&nodes_buf, 0, &staging_buf, 0, buf_size);
        self.queue.submit(std::iter::once(encoder.finish()));

        let slice = staging_buf.slice(..);
        let (tx, rx) = std::sync::mpsc::channel();
        slice.map_async(wgpu::MapMode::Read, move ￿result￿ {
            tx.send(result).unwrap();
        });
        self.device.poll(wgpu::Maintain::Wait);
        rx.recv()
            .expect("Merkle readback channel closed")
            .expect("Merkle readback failed");

        let data = slice.get_mapped_range();
        let result_pairs: &[[u32; 2]] = bytemuck::cast_slice(&data);

        let nodes: Vec<Digest> = (0..num_nodes)
            .map(|i| {
                let base = i * digest_len;
                let mut elements = [BFieldElement::new(0); 5];
                for j in 0..5 {
                    let [lo, hi] = result_pairs[base + j];
                    let raw = (hi as u64) << 32 | lo as u64;
                    elements[j] = BFieldElement::from_raw_u64(raw);
                }
                Digest::new(elements)
            })
            .collect();

        drop(data);
        staging_buf.unmap();
        self.release_staging(buf_size, staging_buf);

        MerkleTree::from_nodes(nodes)
    }
}