//! PERF-2 P3-ii (`docs/perf2_campaign.md`): the multigrid V-cycle on the
//! CUDA device as the CG's preconditioner — the `mg_vcycle.cu` kernels of
//! the go/no-go benchmark, K = 1, with persistent device buffers per
//! operator. One CUDA context, stream and module per process (built on
//! first use); the CG, its residual and its stop stay f64 on the CPU.

use super::LevelExport;
use cudarc::driver::{
    CudaContext, CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtr, DeviceRepr,
    LaunchConfig, PushKernelArg, ValidAsZeroBits,
};
use cudarc::nvrtc::{CompileOptions, compile_ptx_with_opts};
use std::cell::Cell;
use std::collections::BTreeMap;
use std::sync::{Arc, Condvar, Mutex, OnceLock};

const KERNELS: &str = include_str!("../../../kernels/cuda/mg_vcycle.cu");

struct Runtime {
    _ctx: Arc<CudaContext>,
    stream: Arc<CudaStream>,
    _module: Arc<CudaModule>,
    f_half: CudaFunction,
    f_res: CudaFunction,
    f_restrict: CudaFunction,
    f_prolong: CudaFunction,
    f_zero: CudaFunction,
    f_coarsest: CudaFunction,
    // The lane-table kernels (P3-iii).
    l_half: CudaFunction,
    l_res: CudaFunction,
    l_restrict: CudaFunction,
    l_prolong: CudaFunction,
    l_zero: CudaFunction,
    l_coarsest: CudaFunction,
}

static RUNTIME: OnceLock<Runtime> = OnceLock::new();

fn runtime() -> &'static Runtime {
    RUNTIME.get_or_init(|| {
        let ctx = CudaContext::new(0).expect("CUDA context (device 0)");
        let stream = ctx.default_stream();
        let arch = std::env::var("RTX_CUDA_ARCH").unwrap_or_else(|_| "sm_120".to_string());
        let ptx = compile_ptx_with_opts(
            KERNELS,
            CompileOptions {
                arch: Some(Box::leak(arch.into_boxed_str())),
                ..Default::default()
            },
        )
        .expect("nvrtc: mg_vcycle.cu");
        let module = ctx.load_module(ptx).expect("mg_vcycle module");
        let f = |name: &str| module.load_function(name).expect(name);
        Runtime {
            f_half: f("mg_rb_half"),
            f_res: f("mg_residual"),
            f_restrict: f("mg_restrict"),
            f_prolong: f("mg_prolong"),
            f_zero: f("mg_zero"),
            f_coarsest: f("mg_coarsest"),
            l_half: f("ml_rb_half"),
            l_res: f("ml_residual"),
            l_restrict: f("ml_restrict"),
            l_prolong: f("ml_prolong"),
            l_zero: f("ml_zero"),
            l_coarsest: f("ml_coarsest"),
            _ctx: ctx,
            stream,
            _module: module,
        }
    })
}

struct DevLevel {
    n: usize,
    nx: i32,
    n_cells: usize,
    n_red: usize,
    n_black: usize,
    cells: CudaSlice<u32>,
    red: CudaSlice<u32>,
    black: CudaSlice<u32>,
    coarse_of: CudaSlice<u32>,
    children_ptr: CudaSlice<u32>,
    children_idx: CudaSlice<u32>,
    ae: CudaSlice<f32>,
    aw: CudaSlice<f32>,
    an: CudaSlice<f32>,
    as_: CudaSlice<f32>,
    ap: CudaSlice<f32>,
    b: CudaSlice<f32>,
    x: CudaSlice<f32>,
    r: CudaSlice<f32>,
}

/// One operator's hierarchy on the device (K = 1).
pub(super) struct DeviceVcycle {
    levels: Vec<DevLevel>,
    sweeps: usize,
    fine_cells: Vec<u32>,
    r_f32: Vec<f32>,
    z_f32: Vec<f32>,
}

fn cfg(n_items: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: ((n_items as u32).div_ceil(256).max(1), 1, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    }
}

impl DeviceVcycle {
    pub(super) fn new(levels: &[LevelExport], sweeps: usize) -> Self {
        let rt = runtime();
        let up_u = |v: &[u32]| -> CudaSlice<u32> {
            rt.stream
                .memcpy_stod(if v.is_empty() { &[0u32][..] } else { v })
                .expect("upload")
        };
        let up_f = |v: &[f32]| -> CudaSlice<f32> { rt.stream.memcpy_stod(v).expect("upload") };
        let dev: Vec<DevLevel> = levels
            .iter()
            .map(|l| {
                let n = l.nx * l.ny;
                DevLevel {
                    n,
                    nx: l.nx as i32,
                    n_cells: l.cells.len(),
                    n_red: l.red.len(),
                    n_black: l.black.len(),
                    cells: up_u(&l.cells),
                    red: up_u(&l.red),
                    black: up_u(&l.black),
                    coarse_of: up_u(&l.coarse_of),
                    children_ptr: up_u(&l.children_ptr),
                    children_idx: up_u(&l.children_idx),
                    ae: up_f(&l.ae),
                    aw: up_f(&l.aw),
                    an: up_f(&l.an),
                    as_: up_f(&l.as_),
                    ap: up_f(&l.ap),
                    b: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
                    x: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
                    r: rt.stream.alloc_zeros::<f32>(n).expect("alloc"),
                }
            })
            .collect();
        let n0 = dev[0].n;
        Self {
            levels: dev,
            sweeps,
            fine_cells: levels[0].cells.clone(),
            r_f32: vec![0.0; n0],
            z_f32: vec![0.0; n0],
        }
    }

    fn half(&mut self, l: usize, colour: u8) {
        let rt = runtime();
        let lv = &mut self.levels[l];
        let (list, n_list) = if colour == 0 {
            (&lv.red, lv.n_red)
        } else {
            (&lv.black, lv.n_black)
        };
        let (n_i, n_list_i) = (lv.n as i32, n_list as i32);
        unsafe {
            rt.stream
                .launch_builder(&rt.f_half)
                .arg(&n_list_i)
                .arg(list)
                .arg(&n_i)
                .arg(&lv.ae)
                .arg(&lv.aw)
                .arg(&lv.an)
                .arg(&lv.as_)
                .arg(&lv.ap)
                .arg(&lv.b)
                .arg(&mut lv.x)
                .arg(&lv.nx)
                .launch(cfg(n_list))
                .expect("mg_rb_half");
        }
    }

    fn smooth(&mut self, l: usize) {
        for _ in 0..self.sweeps {
            self.half(l, 0);
            self.half(l, 1);
            self.half(l, 1);
            self.half(l, 0);
        }
    }

    /// `z = M⁻¹ r` on the active cells (the same V-cycle as
    /// `Hierarchy::apply_preconditioner`, in f32 on the device).
    pub(super) fn apply(&mut self, r: &[f64], z: &mut [f64]) {
        let rt = runtime();
        let depth = self.levels.len();
        for (dst, &src) in self.r_f32.iter_mut().zip(r) {
            *dst = src as f32;
        }
        rt.stream
            .memcpy_htod(&self.r_f32, &mut self.levels[0].b)
            .expect("upload r");
        // Down.
        for l in 0..depth - 1 {
            {
                let lv = &mut self.levels[l];
                let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
                unsafe {
                    rt.stream
                        .launch_builder(&rt.f_zero)
                        .arg(&n_cells_i)
                        .arg(&lv.cells)
                        .arg(&n_i)
                        .arg(&mut lv.x)
                        .launch(cfg(lv.n_cells))
                        .expect("mg_zero");
                }
            }
            self.smooth(l);
            {
                let lv = &mut self.levels[l];
                let (n_i, n_cells_i) = (lv.n as i32, lv.n_cells as i32);
                unsafe {
                    rt.stream
                        .launch_builder(&rt.f_res)
                        .arg(&n_cells_i)
                        .arg(&lv.cells)
                        .arg(&n_i)
                        .arg(&lv.ae)
                        .arg(&lv.aw)
                        .arg(&lv.an)
                        .arg(&lv.as_)
                        .arg(&lv.ap)
                        .arg(&lv.b)
                        .arg(&lv.x)
                        .arg(&mut lv.r)
                        .arg(&lv.nx)
                        .launch(cfg(lv.n_cells))
                        .expect("mg_residual");
                }
            }
            let (fine, coarse) = self.levels.split_at_mut(l + 1);
            let (lf, lc) = (&fine[l], &mut coarse[0]);
            let (n_c_cells_i, n_f_i, n_c_i) = (lc.n_cells as i32, lf.n as i32, lc.n as i32);
            unsafe {
                rt.stream
                    .launch_builder(&rt.f_restrict)
                    .arg(&n_c_cells_i)
                    .arg(&lc.cells)
                    .arg(&lf.children_ptr)
                    .arg(&lf.children_idx)
                    .arg(&n_f_i)
                    .arg(&n_c_i)
                    .arg(&lf.r)
                    .arg(&mut lc.b)
                    .launch(cfg(lc.n_cells))
                    .expect("mg_restrict");
            }
        }
        // Coarsest.
        {
            let lv = &mut self.levels[depth - 1];
            let (k_i, n_cells_i, n_i, sw_i) = (1i32, lv.n_cells as i32, lv.n as i32, 50i32);
            let (n_red_i, n_black_i) = (lv.n_red as i32, lv.n_black as i32);
            unsafe {
                rt.stream
                    .launch_builder(&rt.f_coarsest)
                    .arg(&k_i)
                    .arg(&n_cells_i)
                    .arg(&lv.cells)
                    .arg(&n_red_i)
                    .arg(&lv.red)
                    .arg(&n_black_i)
                    .arg(&lv.black)
                    .arg(&n_i)
                    .arg(&lv.ae)
                    .arg(&lv.aw)
                    .arg(&lv.an)
                    .arg(&lv.as_)
                    .arg(&lv.ap)
                    .arg(&lv.b)
                    .arg(&mut lv.x)
                    .arg(&lv.nx)
                    .arg(&sw_i)
                    .launch(LaunchConfig {
                        grid_dim: (1, 1, 1),
                        block_dim: (32, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .expect("mg_coarsest");
            }
        }
        // Up.
        for l in (0..depth - 1).rev() {
            {
                let (fine, coarse) = self.levels.split_at_mut(l + 1);
                let (lf, lc) = (&mut fine[l], &coarse[0]);
                let (n_cells_i, n_f_i, n_c_i) = (lf.n_cells as i32, lf.n as i32, lc.n as i32);
                unsafe {
                    rt.stream
                        .launch_builder(&rt.f_prolong)
                        .arg(&n_cells_i)
                        .arg(&lf.cells)
                        .arg(&lf.coarse_of)
                        .arg(&n_f_i)
                        .arg(&n_c_i)
                        .arg(&mut lf.x)
                        .arg(&lc.x)
                        .launch(cfg(lf.n_cells))
                        .expect("mg_prolong");
                }
            }
            self.smooth(l);
        }
        rt.stream
            .memcpy_dtoh(&self.levels[0].x, &mut self.z_f32)
            .expect("download z");
        rt.stream.synchronize().expect("sync");
        for &idx in &self.fine_cells {
            z[idx as usize] = self.z_f32[idx as usize] as f64;
        }
    }
}

// ---------------------------------------------------------------------------
// PERF-2 P3-iii: the `MarchPlane` — K lanes' V-cycles in ONE batched launch
// set, and the rendezvous that gathers the lanes' preconditioner calls.
// ---------------------------------------------------------------------------

/// One lane's level as the `ml_*` kernels read it: 14 device pointers then
/// 6 ints (`n, nx, n_cells, n_red, n_black, pad`), mirroring `struct
/// LaneLevel` in `mg_vcycle.cu`.
#[repr(C)]
#[derive(Clone, Copy, Default)]
struct LaneLevel {
    ptrs: [u64; 14],
    ints: [i32; 6],
}
unsafe impl DeviceRepr for LaneLevel {}
unsafe impl ValidAsZeroBits for LaneLevel {}

impl DeviceVcycle {
    pub(super) fn depth(&self) -> usize {
        self.levels.len()
    }

    fn lane_level(&self, l: usize, stream: &CudaStream) -> LaneLevel {
        let lv = &self.levels[l];
        let pf = |s: &CudaSlice<f32>| s.device_ptr(stream).0;
        let pu = |s: &CudaSlice<u32>| s.device_ptr(stream).0;
        LaneLevel {
            ptrs: [
                pf(&lv.ae),
                pf(&lv.aw),
                pf(&lv.an),
                pf(&lv.as_),
                pf(&lv.ap),
                pf(&lv.b),
                pf(&lv.x),
                pf(&lv.r),
                pu(&lv.cells),
                pu(&lv.red),
                pu(&lv.black),
                pu(&lv.coarse_of),
                pu(&lv.children_ptr),
                pu(&lv.children_idx),
            ],
            ints: [
                lv.n as i32,
                lv.nx,
                lv.n_cells as i32,
                lv.n_red as i32,
                lv.n_black as i32,
                0,
            ],
        }
    }
}

/// One lane's preconditioner call: its operator on the device, the residual
/// in, the correction out.
pub(super) struct LaneJob<'a> {
    pub dev: &'a mut DeviceVcycle,
    pub r: &'a [f64],
    pub z: &'a mut [f64],
}

fn cfg_k(n_items: usize, k: usize) -> LaunchConfig {
    LaunchConfig {
        grid_dim: ((n_items as u32).div_ceil(256).max(1), k as u32, 1),
        block_dim: (256, 1, 1),
        shared_mem_bytes: 0,
    }
}

/// `z_k = M_k⁻¹ r_k` for every job, the lanes of equal hierarchy depth in
/// one batched launch set (grid.y = the lane). Each lane's result is the
/// K = 1 [`DeviceVcycle::apply`]'s, bit for bit (the same per-cell
/// arithmetic in the same order; gate G-P0).
pub(super) fn apply_batch(jobs: &mut [LaneJob<'_>]) {
    let rt = runtime();
    let stream = &rt.stream;
    let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
    for (i, j) in jobs.iter().enumerate() {
        groups.entry(j.dev.depth()).or_default().push(i);
    }
    for (depth, members) in groups {
        let k = members.len();
        // Residuals up.
        for &i in &members {
            let job = &mut jobs[i];
            for (dst, &src) in job.dev.r_f32.iter_mut().zip(job.r) {
                *dst = src as f32;
            }
            stream
                .memcpy_htod(&job.dev.r_f32, &mut job.dev.levels[0].b)
                .expect("upload r");
        }
        // The lane tables, one per level.
        let tables: Vec<CudaSlice<LaneLevel>> = (0..depth)
            .map(|l| {
                let rows: Vec<LaneLevel> = members
                    .iter()
                    .map(|&i| jobs[i].dev.lane_level(l, stream))
                    .collect();
                stream.memcpy_stod(&rows).expect("lane table")
            })
            .collect();
        let max_of = |f: &dyn Fn(&DevLevel) -> usize, l: usize| -> usize {
            members
                .iter()
                .map(|&i| f(&jobs[i].dev.levels[l]))
                .max()
                .unwrap_or(0)
        };
        let sweeps = jobs[members[0]].dev.sweeps;
        let smooth = |l: usize| {
            let n_red = max_of(&|lv| lv.n_red, l);
            let n_black = max_of(&|lv| lv.n_black, l);
            let half = |colour: i32| {
                let n = if colour == 0 { n_red } else { n_black };
                unsafe {
                    stream
                        .launch_builder(&rt.l_half)
                        .arg(&tables[l])
                        .arg(&colour)
                        .launch(cfg_k(n, k))
                        .expect("ml_rb_half");
                }
            };
            for _ in 0..sweeps {
                half(0);
                half(1);
                half(1);
                half(0);
            }
        };
        // Down.
        for l in 0..depth - 1 {
            let n_cells = max_of(&|lv| lv.n_cells, l);
            unsafe {
                stream
                    .launch_builder(&rt.l_zero)
                    .arg(&tables[l])
                    .launch(cfg_k(n_cells, k))
                    .expect("ml_zero");
            }
            smooth(l);
            unsafe {
                stream
                    .launch_builder(&rt.l_res)
                    .arg(&tables[l])
                    .launch(cfg_k(n_cells, k))
                    .expect("ml_residual");
            }
            let n_coarse = max_of(&|lv| lv.n_cells, l + 1);
            unsafe {
                stream
                    .launch_builder(&rt.l_restrict)
                    .arg(&tables[l])
                    .arg(&tables[l + 1])
                    .launch(cfg_k(n_coarse, k))
                    .expect("ml_restrict");
            }
        }
        // Coarsest.
        {
            let (k_i, sw_i) = (k as i32, 50i32);
            unsafe {
                stream
                    .launch_builder(&rt.l_coarsest)
                    .arg(&tables[depth - 1])
                    .arg(&k_i)
                    .arg(&sw_i)
                    .launch(LaunchConfig {
                        grid_dim: ((k as u32).div_ceil(32), 1, 1),
                        block_dim: (32, 1, 1),
                        shared_mem_bytes: 0,
                    })
                    .expect("ml_coarsest");
            }
        }
        // Up.
        for l in (0..depth - 1).rev() {
            let n_cells = max_of(&|lv| lv.n_cells, l);
            unsafe {
                stream
                    .launch_builder(&rt.l_prolong)
                    .arg(&tables[l])
                    .arg(&tables[l + 1])
                    .launch(cfg_k(n_cells, k))
                    .expect("ml_prolong");
            }
            smooth(l);
        }
        // Corrections down.
        for &i in &members {
            let job = &mut jobs[i];
            stream
                .memcpy_dtoh(&job.dev.levels[0].x, &mut job.dev.z_f32)
                .expect("download z");
        }
        stream.synchronize().expect("sync");
        for &i in &members {
            let job = &mut jobs[i];
            for &idx in &job.dev.fine_cells {
                job.z[idx as usize] = job.dev.z_f32[idx as usize] as f64;
            }
        }
    }
}

// --- The rendezvous ---------------------------------------------------------
//
// K lane threads each run an unchanged march. A lane inside a CG (between
// `cg_enter` and the guard's drop) that calls the preconditioner submits its
// job and blocks; when every lane currently inside a CG has submitted, the
// last arrival serves the whole batch and wakes the rest. A lane LEAVING its
// CG can complete a batch too (the others were waiting for it). Lanes
// outside a CG are not counted, so they never hold a batch up.

thread_local! {
    static LANE: Cell<Option<usize>> = const { Cell::new(None) };
}

/// Mark this thread as plane lane `lane` (`None` = a solo march: the K = 1
/// path, unchanged).
pub fn set_lane(lane: Option<usize>) {
    LANE.with(|l| l.set(lane));
}

pub fn lane() -> Option<usize> {
    LANE.with(|l| l.get())
}

struct PendingJob {
    dev: *mut DeviceVcycle,
    r: *const f64,
    r_len: usize,
    z: *mut f64,
    z_len: usize,
}
// The submitting thread blocks until its job is served, so the pointers
// outlive the batch.
unsafe impl Send for PendingJob {}

#[derive(Default)]
struct PlaneState {
    in_cg: usize,
    pending: Vec<PendingJob>,
    generation: u64,
    /// Batches served and lanes served (the plane's own bookkeeping).
    batches: u64,
    lanes_served: u64,
}

static PLANE: Mutex<PlaneState> = Mutex::new(PlaneState {
    in_cg: 0,
    pending: Vec::new(),
    generation: 0,
    batches: 0,
    lanes_served: 0,
});
static PLANE_CV: Condvar = Condvar::new();

fn plane_lock() -> std::sync::MutexGuard<'static, PlaneState> {
    PLANE.lock().unwrap_or_else(|e| e.into_inner())
}

fn serve(st: &mut PlaneState) {
    let pending = std::mem::take(&mut st.pending);
    let mut jobs: Vec<LaneJob<'_>> = pending
        .iter()
        .map(|p| unsafe {
            LaneJob {
                dev: &mut *p.dev,
                r: std::slice::from_raw_parts(p.r, p.r_len),
                z: std::slice::from_raw_parts_mut(p.z, p.z_len),
            }
        })
        .collect();
    apply_batch(&mut jobs);
    st.batches += 1;
    st.lanes_served += pending.len() as u64;
    st.generation += 1;
    PLANE_CV.notify_all();
}

/// Held by a lane for the duration of one CG solve.
pub(super) struct CgGuard(());

impl Drop for CgGuard {
    fn drop(&mut self) {
        let mut st = plane_lock();
        st.in_cg -= 1;
        if st.in_cg > 0 && st.pending.len() >= st.in_cg {
            serve(&mut st);
        }
    }
}

/// Enter a CG solve on this lane (`None` when this thread is not a lane).
pub(super) fn cg_enter() -> Option<CgGuard> {
    lane()?;
    plane_lock().in_cg += 1;
    Some(CgGuard(()))
}

/// This lane's preconditioner call inside the plane: submit, and either
/// serve the batch (last arrival) or wait for it.
pub(super) fn apply_lane(dev: &mut DeviceVcycle, r: &[f64], z: &mut [f64]) {
    let mut st = plane_lock();
    let my_generation = st.generation;
    st.pending.push(PendingJob {
        dev: std::ptr::from_mut(dev),
        r: r.as_ptr(),
        r_len: r.len(),
        z: z.as_mut_ptr(),
        z_len: z.len(),
    });
    if st.pending.len() >= st.in_cg {
        serve(&mut st);
        return;
    }
    while st.generation == my_generation {
        st = PLANE_CV.wait(st).unwrap_or_else(|e| e.into_inner());
    }
}

/// `(batches served, lane calls served)` so far in this process.
pub fn plane_counters() -> (u64, u64) {
    let st = plane_lock();
    (st.batches, st.lanes_served)
}

#[cfg(test)]
mod tests {
    //! G-P0 (`docs/perf2_campaign.md`, P3-iii): K distinct masked operators
    //! through the batched path equal each one's K = 1 device apply bit for
    //! bit, and the rendezvous serves uneven lanes exactly. GPU only:
    //! `RTX_CUDA_ARCH=sm_120 cargo test --release -p rtx-cfd --features cuda --lib device::tests`.
    use super::*;
    use crate::solvers::incompressible::poisson::{
        MgSmoother, MultigridParameters, PoissonProblem, export_hierarchy,
    };

    /// A masked channel operator with a hole at `(cx, cy)` and a seeded
    /// right-hand side (the go/no-go benchmark's construction).
    fn problem(nx: usize, ny: usize, cx: f64, cy: f64, seed: u64) -> PoissonProblem {
        let mut p = PoissonProblem::new(nx, ny);
        let (dx, dy, dt) = (2.5 / nx as f64, 0.41 / ny as f64, 3.24e-4);
        let (ae, an) = (dt * dy / dx, dt * dx / dy);
        let hole = |i: usize, j: usize| {
            let (x, y) = ((i as f64 + 0.5) * dx, (j as f64 + 0.5) * dy);
            (x - cx).powi(2) + (y - cy).powi(2) < 0.08 * 0.08
        };
        for j in 0..ny {
            for i in 0..nx {
                let idx = j * nx + i;
                if hole(i, j) {
                    p.active[idx] = false;
                    continue;
                }
                if i + 1 < nx && !hole(i + 1, j) {
                    p.ae[idx] = ae;
                }
                if i > 0 && !hole(i - 1, j) {
                    p.aw[idx] = ae;
                }
                if j + 1 < ny && !hole(i, j + 1) {
                    p.an[idx] = an;
                }
                if j > 0 && !hole(i, j - 1) {
                    p.as_[idx] = an;
                }
                if i + 1 == nx {
                    p.extra_diag[idx] = 2.0 * ae;
                }
            }
        }
        let mut state = seed | 1;
        for idx in 0..nx * ny {
            state ^= state << 13;
            state ^= state >> 7;
            state ^= state << 17;
            p.rhs[idx] = if p.active[idx] {
                1e-3 * ((state >> 11) as f64 / (1u64 << 53) as f64 - 0.5)
            } else {
                0.0
            };
        }
        p
    }

    fn params() -> MultigridParameters {
        MultigridParameters {
            smoother: MgSmoother::RedBlack,
            ..MultigridParameters::default()
        }
    }

    fn device_of(p: &PoissonProblem) -> DeviceVcycle {
        DeviceVcycle::new(
            &export_hierarchy(p, &params()),
            params().smoother_sweeps.max(1),
        )
    }

    fn bits(v: &[f64]) -> Vec<u64> {
        v.iter().map(|x| x.to_bits()).collect()
    }

    /// K operators with different holes: solo applies, then one batch.
    #[test]
    fn batched_lanes_equal_the_solo_apply_bit_for_bit() {
        let (nx, ny) = (186usize, 31usize);
        let probs: Vec<PoissonProblem> = [(0.2, 0.2, 3u64), (0.6, 0.15, 7), (1.1, 0.25, 11)]
            .iter()
            .map(|&(cx, cy, s)| problem(nx, ny, cx, cy, s))
            .collect();
        let n = nx * ny;
        let mut solo: Vec<Vec<f64>> = Vec::new();
        for p in &probs {
            let mut d = device_of(p);
            let mut z = vec![0.0; n];
            d.apply(&p.rhs, &mut z);
            solo.push(z);
        }
        assert!(
            solo[0].iter().any(|v| *v != 0.0),
            "the solo apply produced zeros"
        );
        // K = 1 through the batched path (the lane kernels vs the K = 1 kernels).
        {
            let mut d = device_of(&probs[0]);
            let mut z = vec![0.0; n];
            apply_batch(&mut [LaneJob {
                dev: &mut d,
                r: &probs[0].rhs,
                z: &mut z,
            }]);
            assert_eq!(
                bits(&z),
                bits(&solo[0]),
                "K = 1 through the lane kernels differs"
            );
        }
        // K = 3.
        let mut devs: Vec<DeviceVcycle> = probs.iter().map(device_of).collect();
        let mut zs: Vec<Vec<f64>> = vec![vec![0.0; n]; 3];
        {
            let mut jobs: Vec<LaneJob<'_>> = devs
                .iter_mut()
                .zip(probs.iter())
                .zip(zs.iter_mut())
                .map(|((dev, p), z)| LaneJob { dev, r: &p.rhs, z })
                .collect();
            apply_batch(&mut jobs);
        }
        for k in 0..3 {
            assert_eq!(
                bits(&zs[k]),
                bits(&solo[k]),
                "lane {k} differs from its solo apply"
            );
        }
        // A second batch on the same devices (persistent buffers) with the
        // lanes in another order and one lane left out.
        let mut z1 = vec![0.0; n];
        let mut z2 = vec![0.0; n];
        {
            let (a, b) = devs.split_at_mut(2);
            let mut jobs = [
                LaneJob {
                    dev: &mut b[0],
                    r: &probs[2].rhs,
                    z: &mut z2,
                },
                LaneJob {
                    dev: &mut a[1],
                    r: &probs[1].rhs,
                    z: &mut z1,
                },
            ];
            apply_batch(&mut jobs);
        }
        assert_eq!(bits(&z2), bits(&solo[2]));
        assert_eq!(bits(&z1), bits(&solo[1]));
    }

    /// Three lane threads with 2 / 4 / 5 preconditioner calls each inside
    /// one CG guard: every call's correction equals the solo apply of the
    /// same residual, the lanes leave at different times, and no thread
    /// hangs.
    #[test]
    fn the_rendezvous_serves_uneven_lanes_exactly() {
        let (nx, ny) = (124usize, 21usize);
        let n = nx * ny;
        let probs: Vec<PoissonProblem> = [(0.2, 0.2, 5u64), (0.7, 0.2, 9), (1.3, 0.22, 13)]
            .iter()
            .map(|&(cx, cy, s)| problem(nx, ny, cx, cy, s))
            .collect();
        let calls = [2usize, 4, 5];
        // Expected: the solo apply of rhs × (1 + c) for call c.
        let mut expected: Vec<Vec<Vec<u64>>> = Vec::new();
        for (k, p) in probs.iter().enumerate() {
            let mut d = device_of(p);
            let mut per_call = Vec::new();
            for c in 0..calls[k] {
                let r: Vec<f64> = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect();
                let mut z = vec![0.0; n];
                d.apply(&r, &mut z);
                per_call.push(bits(&z));
            }
            expected.push(per_call);
        }
        let (b0, l0) = plane_counters();
        let handles: Vec<_> = probs
            .into_iter()
            .enumerate()
            .map(|(k, p)| {
                let calls_k = calls[k];
                std::thread::spawn(move || {
                    set_lane(Some(k));
                    let mut d = device_of(&p);
                    let guard = cg_enter().expect("lane set");
                    let mut out = Vec::new();
                    for c in 0..calls_k {
                        let r: Vec<f64> = p.rhs.iter().map(|v| v * (1.0 + c as f64)).collect();
                        let mut z = vec![0.0; n];
                        apply_lane(&mut d, &r, &mut z);
                        out.push(bits(&z));
                    }
                    drop(guard);
                    out
                })
            })
            .collect();
        let results: Vec<Vec<Vec<u64>>> = handles
            .into_iter()
            .map(|h| h.join().expect("lane thread"))
            .collect();
        for k in 0..3 {
            assert_eq!(results[k].len(), calls[k]);
            for c in 0..calls[k] {
                assert_eq!(results[k][c], expected[k][c], "lane {k} call {c} differs");
            }
        }
        let (b1, l1) = plane_counters();
        assert_eq!(l1 - l0, 11, "every lane call served once");
        assert!(b1 - b0 >= 5 && b1 - b0 <= 11, "batches {}", b1 - b0);
    }
}