cubecl-core 0.12.0-pre.2

use alloc::{vec, vec::Vec};
use cubecl_runtime::runtime::Runtime;

use crate::prelude::*;
use crate::{self as cubecl};
use cubecl_ir::features::{AtomicUsage, Plane};

#[cube(launch)]
/// First 32 elements should be 1, while last 32 elements may or may not be 1
fn kernel_test_sync_cube(buffer: &mut [u32], out: &mut [u32]) {
    let unit_pos = UNIT_POS as usize;
    buffer[unit_pos] = UNIT_POS;
    sync_cube();
    if unit_pos != 0 {
        out[unit_pos] = buffer[unit_pos - 1] + buffer[unit_pos];
    }
}

pub fn test_sync_cube<R: Runtime>(client: Client) {
    // Clamp the cube dim to the device's limit (e.g. the core count on CPU).
    let max_units = client.properties().hardware.max_units_per_cube;
    let dim_x = core::cmp::min(8, (max_units / 2).max(1));
    let units = (dim_x * 2) as usize;

    let handle = client.empty(32 * core::mem::size_of::<u32>());
    let test = client.empty(32 * core::mem::size_of::<u32>());

    kernel_test_sync_cube::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(dim_x, 2),
        unsafe { BufferArg::from_raw_parts(test, 32) },
        unsafe { BufferArg::from_raw_parts(handle.clone(), 32) },
    );

    let actual = client.read_one_unchecked(handle);
    let actual = u32::from_bytes(&actual);

    let expected: Vec<u32> = (0..units as i32)
        .map(|i| core::cmp::max(2 * i - 1, 0) as u32)
        .collect();

    assert_eq!(&actual[1..units], &expected[1..units]);
}

#[cube(launch)]
/// First 32 elements should be 1, while last 32 elements may or may not be 1
fn kernel_test_finished_sync_cube(buffer: &mut [u32], out: &mut [u32]) {
    let unit_pos = UNIT_POS as usize;
    buffer[unit_pos] = UNIT_POS;
    if UNIT_POS > 16 {
        terminate!();
    }
    sync_cube();
    sync_cube();
    if UNIT_POS != 0 {
        out[unit_pos] = buffer[unit_pos - 1] + buffer[unit_pos];
    }
    sync_cube();
}

pub fn test_finished_sync_cube<R: Runtime>(client: Client) {
    // Clamp the cube dim to the device's limit (e.g. the core count on CPU).
    let max_units = client.properties().hardware.max_units_per_cube;
    let dim_x = core::cmp::min(8, (max_units / 2).max(1));
    let checked = core::cmp::min(8, (dim_x * 2) as usize);

    let handle = client.empty(32 * core::mem::size_of::<u32>());
    let test = client.empty(32 * core::mem::size_of::<u32>());

    kernel_test_finished_sync_cube::launch(
        &client,
        CubeCount::Static(2, 1, 1),
        CubeDim::new_2d(dim_x, 2),
        unsafe { BufferArg::from_raw_parts(test, 32) },
        unsafe { BufferArg::from_raw_parts(handle.clone(), 32) },
    );

    let actual = client.read_one_unchecked(handle);
    let actual = u32::from_bytes(&actual);

    let expected: Vec<u32> = (0..checked as i32)
        .map(|i| core::cmp::max(2 * i - 1, 0) as u32)
        .collect();

    assert_eq!(&actual[1..checked], &expected[1..checked]);
}

#[cube(launch)]
/// First 32 elements should be 1, while last 32 elements may or may not be 1
fn kernel_test_sync_plane<F: Float>(out: &mut [F]) {
    let mut shared_memory = Shared::<F>::new();

    if UNIT_POS == 0 {
        *shared_memory = F::from_int(1);
    }

    sync_plane();

    out[UNIT_POS as usize] = *shared_memory;
}

pub fn test_sync_plane<R: Runtime>(client: Client) {
    if !client.features().plane.contains(Plane::Sync) {
        // We can't execute the test, skip.
        return;
    }

    let handle = client.empty(64 * core::mem::size_of::<f32>());

    kernel_test_sync_plane::launch::<f32>(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(32, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 64) },
    );

    let actual = client.read_one_unchecked(handle);
    let actual = f32::from_bytes(&actual);
    let expected = &[
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
        1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
    ];

    assert_eq!(&actual[0..32], expected);
}

/// Every unit reads its plane neighbour's slot in global memory before writing its own, so a
/// stale copy of the zero it read first is what a barrier ordering shared memory alone may hand
/// back; then it writes its slot, meets its plane at [`sync_plane_storage()`], and reads the
/// neighbour's slot again. `before` only keeps the early read alive; the host never reads it.
/// `plane_dim` is the plane size the device ran at, which decides who the neighbour is.
#[cube(launch)]
fn kernel_test_sync_plane_storage(
    slots: &mut [u32],
    before: &mut [u32],
    out: &mut [u32],
    plane_dim: &mut [u32],
) {
    let unit = ABSOLUTE_POS;
    if unit == 0 {
        plane_dim[0] = PLANE_DIM;
    }
    let first_of_plane = unit - UNIT_POS_PLANE as usize;
    let next = first_of_plane + ((UNIT_POS_PLANE + 1) % PLANE_DIM) as usize;
    before[unit] = slots[next];

    slots[unit] = unit as u32 + 1;
    sync_plane_storage();
    out[unit] = slots[next];
}

/// A plane's units read each other's global writes past [`sync_plane_storage()`], across many
/// cubes of two planes each.
pub fn test_sync_plane_storage<R: Runtime>(client: Client) {
    if !client.features().plane.contains(Plane::Sync) {
        // We can't execute the test, skip.
        return;
    }

    let cubes = 256u32;
    let plane_size = client.properties().hardware.plane_size_max;
    let units = core::cmp::min(
        2 * plane_size,
        client.properties().hardware.max_units_per_cube,
    );
    let count = (cubes * units) as usize;

    let slots = client.create_from_slice(u32::as_bytes(&vec![0u32; count]));
    let before = client.empty(count * core::mem::size_of::<u32>());
    let out = client.empty(count * core::mem::size_of::<u32>());
    let plane_dim = client.empty(core::mem::size_of::<u32>());

    kernel_test_sync_plane_storage::launch(
        &client,
        CubeCount::Static(cubes, 1, 1),
        CubeDim::new_1d(units),
        unsafe { BufferArg::from_raw_parts(slots, count) },
        unsafe { BufferArg::from_raw_parts(before, count) },
        unsafe { BufferArg::from_raw_parts(out.clone(), count) },
        unsafe { BufferArg::from_raw_parts(plane_dim.clone(), 1) },
    );

    let ran_at = u32::from_bytes(&client.read_one_unchecked(plane_dim))[0] as usize;
    let actual = client.read_one_unchecked(out);
    let actual = u32::from_bytes(&actual);
    let expected: Vec<u32> = (0..count)
        .map(|unit| {
            let first = unit - unit % ran_at;
            (first + (unit % ran_at + 1) % ran_at) as u32 + 1
        })
        .collect();
    assert_eq!(actual, &expected);
}

/// Each round every unit writes two slots of its plane's shared window, its own and the one a
/// plane further, meets its plane at [`sync_plane()`], and reads the two neighbouring slots at
/// twice its own position, the way a cooperative matrix load reads its cells; it meets its plane
/// again before the next round overwrites them. A barrier that orders a unit's first store but
/// not its second hands back the previous round's value.
#[cube(launch)]
fn kernel_test_sync_plane_two_stores(out: &mut [u32], #[comptime] units: usize) {
    let mut slots = Shared::<[u32]>::new_slice(comptime!(2 * units));
    let lane = UNIT_POS_PLANE as usize;
    let plane = PLANE_DIM as usize;
    let base = PLANE_POS as usize * 2 * plane;
    let unit = UNIT_POS as usize;
    let cube = CUBE_POS;
    for round in 0..ROUNDS {
        let value = (round * 2 * units) as u32;
        slots[base + lane] = value + (base + lane) as u32;
        slots[base + plane + lane] = value + (base + plane + lane) as u32;
        sync_plane();
        let at = ((cube * ROUNDS + round) * units + unit) * 2;
        out[at] = slots[base + 2 * lane];
        out[at + 1] = slots[base + 2 * lane + 1];
        sync_plane();
    }
}

const ROUNDS: usize = 64;

/// A plane's units read both of each other's shared writes past [`sync_plane()`], round after
/// round, across many cubes.
pub fn test_sync_plane_two_stores<R: Runtime>(client: Client) {
    if !client.features().plane.contains(Plane::Sync) {
        // We can't execute the test, skip.
        return;
    }

    let cubes = 1024usize;
    let units = client.properties().hardware.plane_size_max as usize;
    let count = cubes * ROUNDS * units * 2;
    let out = client.empty(count * core::mem::size_of::<u32>());
    let plane_dim = client.empty(core::mem::size_of::<u32>());

    kernel_test_plane_dim::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(units as u32),
        unsafe { BufferArg::from_raw_parts(plane_dim.clone(), 1) },
    );
    kernel_test_sync_plane_two_stores::launch(
        &client,
        CubeCount::Static(cubes as u32, 1, 1),
        CubeDim::new_1d(units as u32),
        unsafe { BufferArg::from_raw_parts(out.clone(), count) },
        units,
    );

    let plane = u32::from_bytes(&client.read_one_unchecked(plane_dim))[0] as usize;
    let actual = client.read_one_unchecked(out);
    let actual = u32::from_bytes(&actual);
    let mut expected = Vec::with_capacity(count);
    for _cube in 0..cubes {
        for round in 0..ROUNDS {
            let value = round * 2 * units;
            for unit in 0..units {
                let base = unit / plane * 2 * plane;
                let read = base + 2 * (unit % plane);
                expected.push((value + read) as u32);
                expected.push((value + read + 1) as u32);
            }
        }
    }
    let wrong = actual.iter().zip(&expected).filter(|(a, e)| a != e).count();
    assert_eq!(wrong, 0, "{wrong} of {count} reads saw a stale slot");
}

/// The plane size the device runs at.
#[cube(launch)]
fn kernel_test_plane_dim(out: &mut [u32]) {
    if UNIT_POS == 0 {
        out[0] = PLANE_DIM;
    }
}

#[cube(launch)]
/// All 64 elements should be 1
fn kernel_test_sync_cube_shared<F: Float>(out: &mut [F]) {
    let mut shared_memory = Shared::<F>::new();

    if UNIT_POS == 0 {
        *shared_memory = F::from_int(1);
    }

    sync_cube();

    out[UNIT_POS as usize] = *shared_memory;
}

pub fn test_sync_cube_shared<R: Runtime>(client: Client) {
    let max_cube_count = std::cmp::min(64, client.properties().hardware.max_units_per_cube);
    let handle = client.empty(max_cube_count as usize * core::mem::size_of::<f32>());

    kernel_test_sync_cube_shared::launch::<f32>(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(max_cube_count / 2, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), max_cube_count as usize) },
    );

    let actual = client.read_one_unchecked(handle);
    let actual = f32::from_bytes(&actual);
    let expected = vec![1.0; max_cube_count as usize];

    assert_eq!(&actual[0..max_cube_count as usize], expected);
}

#[cube(launch)]
/// Every round, the cube passes its units' values round through one shared array, frees it, meets,
/// and passes them round again through a second: each holds `len` values, more than half of what
/// a cube may hold, so they fit together only where the second takes the first's bytes. The round
/// opens on a barrier, so no unit writes the first again before every unit is done with the
/// second of the round before.
fn kernel_test_shared_freed_in_a_loop(out: &mut [u32], rounds: u32, #[comptime] len: usize) {
    let unit = UNIT_POS as usize;
    let next = (unit + 1) % CUBE_DIM as usize;
    for round in 0..rounds {
        sync_cube();
        let mut first = Shared::<[u32]>::new_slice(len);
        first[unit] = UNIT_POS + round;
        first[len - 1 - unit] = UNIT_POS + round;
        sync_cube();
        let passed = first[next] + first[len - 1 - next];
        unsafe { first.free() };
        sync_cube();
        let mut second = Shared::<[u32]>::new_slice(len);
        second[unit] = passed;
        second[len - 1 - unit] = passed;
        sync_cube();
        out[round as usize * CUBE_DIM as usize + unit] = second[next] + second[len - 1 - next];
    }
}

/// Shared memory freed inside a loop is taken by what the same round declares after it: two arrays
/// of five eighths of a cube's shared memory each fit, every round, and pass the right values.
///
/// Generated only for a runtime whose launch counts shared memory as the allocator lays it out
/// ([`testgen_shared_freed_in_a_loop`]): one that counts every declaration apart refuses the
/// launch, as a CPU runtime and wgpu do.
pub fn test_shared_freed_in_a_loop<R: Runtime>(client: Client) {
    let units = 32u32;
    let rounds = 3u32;
    let len = client.properties().hardware.max_shared_memory_size / size_of::<u32>() * 5 / 8;
    let count = (units * rounds) as usize;
    let handle = client.empty(count * size_of::<u32>());

    kernel_test_shared_freed_in_a_loop::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(units),
        unsafe { BufferArg::from_raw_parts(handle.clone(), count) },
        rounds,
        len,
    );

    let actual = client.read_one_unchecked(handle);
    let actual = u32::from_bytes(&actual);
    let expected: Vec<u32> = (0..rounds)
        .flat_map(|round| (0..units).map(move |unit| 4 * ((unit + 2) % units + round)))
        .collect();
    assert_eq!(actual, expected);
}

#[cube(launch)]
fn kernel_test_workgroup_uniform_load(out: &mut [u32]) {
    let mut count = Shared::new_slice(1usize);
    if UNIT_POS == 0 {
        count[0] = 3u32;
    }
    sync_cube();

    let n = workgroup_uniform_load(&count[0]);
    if n > 0 {
        sync_cube();
    }
    out[UNIT_POS as usize] = n;
}

pub fn test_workgroup_uniform_load<R: Runtime>(client: Client) {
    let max_cube_count = std::cmp::min(64, client.properties().hardware.max_units_per_cube);
    let handle = client.empty(max_cube_count as usize * core::mem::size_of::<u32>());

    kernel_test_workgroup_uniform_load::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(max_cube_count / 2, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), max_cube_count as usize) },
    );

    let actual = client.read_one_unchecked(handle);
    let expected = vec![3u32; max_cube_count as usize];
    assert_eq!(u32::from_bytes(&actual), &expected);
}

#[cube(launch)]
fn kernel_test_workgroup_uniform_load_atomic(out: &mut [u32]) {
    let count = Shared::<[Atomic<u32>]>::new_slice(1usize);
    if UNIT_POS == 0 {
        count[0].store(3u32);
    }
    sync_cube();

    let n = workgroup_uniform_load_atomic(&count[0]);
    if n > 0 {
        sync_cube();
    }
    out[UNIT_POS as usize] = n;
}

pub fn test_workgroup_uniform_load_atomic<R: Runtime>(client: Client) {
    let ty = Type::atomic(u32::elem_type_native());
    if !client
        .properties()
        .shared_atomic_type_usage(ty)
        .contains(AtomicUsage::LoadStore)
    {
        return;
    }

    let max_cube_count = std::cmp::min(64, client.properties().hardware.max_units_per_cube);
    let handle = client.empty(max_cube_count as usize * core::mem::size_of::<u32>());

    kernel_test_workgroup_uniform_load_atomic::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(max_cube_count / 2, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), max_cube_count as usize) },
    );

    let actual = client.read_one_unchecked(handle);
    let expected = vec![3u32; max_cube_count as usize];
    assert_eq!(u32::from_bytes(&actual), &expected);
}

#[cube(launch)]
fn kernel_test_workgroup_uniform_load_vec<N: Size>(out: &mut [Vector<f32, N>]) {
    let mut smem = Shared::new_slice(1usize);
    if UNIT_POS == 0 {
        smem[0] = Vector::new(7.0f32);
    }
    sync_cube();

    out[UNIT_POS as usize] = workgroup_uniform_load(&smem[0]);
}

pub fn test_workgroup_uniform_load_vec<R: Runtime>(client: Client) {
    let lanes = 4usize;
    let max_cube_count = std::cmp::min(64, client.properties().hardware.max_units_per_cube);
    let output = client.create_from_slice(f32::as_bytes(&vec![
        0.0f32;
        max_cube_count as usize * lanes
    ]));

    kernel_test_workgroup_uniform_load_vec::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(max_cube_count / 2, 2),
        lanes,
        unsafe { BufferArg::from_raw_parts(output.clone(), 64) },
    );

    let actual = client.read_one_unchecked(output);
    assert!(f32::from_bytes(&actual).iter().all(|&x| x == 7.0f32));
}

/// `workgroup_uniform_load` has to synchronise on its own: one unit publishes a
/// value and the rest read it back through the uniform load, with no explicit
/// `sync_cube` in between. The value is used as a loop bound, so a stale read
/// changes how much work the reader does rather than just its output.
#[cube(launch)]
fn kernel_test_workgroup_uniform_load_synchronizes(out: &mut [u32]) {
    let mut bound = Shared::new_slice(1usize);
    if UNIT_POS == 0 {
        // Derive the bound from real work so it cannot be constant-folded and
        // the store cannot be hoisted above the readers.
        let mut acc = 0u32;
        let mut i = 0u32;
        while i < 1024u32 {
            acc += i % 3u32;
            i += 1u32;
        }
        bound[0] = acc;
    }
    let n = workgroup_uniform_load(&bound[0]);
    let capped = min(n, 4096u32);
    let mut sum = 0u32;
    let mut k = 0u32;
    while k < capped {
        sum += k;
        k += 1u32;
    }
    out[UNIT_POS as usize] = sum;
}

fn expected_uniform_load_sum() -> u32 {
    let acc: u32 = (0..1024u32).map(|i| i % 3).sum();
    let capped = acc.min(4096);
    (0..capped).sum()
}

pub fn test_workgroup_uniform_load_synchronizes<R: Runtime>(client: Client) {
    let units = std::cmp::min(256, client.properties().hardware.max_units_per_cube);
    let handle = client.empty(units as usize * core::mem::size_of::<u32>());

    kernel_test_workgroup_uniform_load_synchronizes::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(units / 2, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), units as usize) },
    );

    let actual = client.read_one_unchecked(handle);
    let expected = vec![expected_uniform_load_sum(); units as usize];
    assert_eq!(u32::from_bytes(&actual), &expected);
}

/// Atomic counterpart of [`test_workgroup_uniform_load_synchronizes`].
#[cube(launch)]
fn kernel_test_workgroup_uniform_load_atomic_synchronizes(out: &mut [u32]) {
    let bound = Shared::<[Atomic<u32>]>::new_slice(1usize);
    if UNIT_POS == 0 {
        let mut acc = 0u32;
        let mut i = 0u32;
        while i < 1024u32 {
            acc += i % 3u32;
            i += 1u32;
        }
        Atomic::store(&bound[0], acc);
    }
    let n = workgroup_uniform_load_atomic(&bound[0]);
    let capped = min(n, 4096u32);
    let mut sum = 0u32;
    let mut k = 0u32;
    while k < capped {
        sum += k;
        k += 1u32;
    }
    out[UNIT_POS as usize] = sum;
}

pub fn test_workgroup_uniform_load_atomic_synchronizes<R: Runtime>(client: Client) {
    let units = std::cmp::min(256, client.properties().hardware.max_units_per_cube);
    let handle = client.empty(units as usize * core::mem::size_of::<u32>());

    kernel_test_workgroup_uniform_load_atomic_synchronizes::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_2d(units / 2, 2),
        unsafe { BufferArg::from_raw_parts(handle.clone(), units as usize) },
    );

    let actual = client.read_one_unchecked(handle);
    let expected = vec![expected_uniform_load_sum(); units as usize];
    assert_eq!(u32::from_bytes(&actual), &expected);
}

/// One cube reads what the others published, with no second dispatch: every cube reduces its
/// units to one partial, releases it with [`sync_storage`], and announces itself on a counter.
/// The cube whose arrival is the last acquires the rest and sums them.
///
/// Both halves of the scope are under test, and the shape is chosen so that each of them has to
/// work. The partial is written by *one* unit and is a reduction over all of them, so the
/// release has to cover a store the cube made together rather than one store per unit. The count
/// is taken by one unit and read by every one of them, which is the cube barrier. Nothing spins
/// — the cubes that are not last simply end — so this cannot hang on a device that does not run
/// them all at once.
///
/// It catches a lowering with no device release at all, which is what CUDA and Metal both had.
/// It does *not* catch one whose release sits on the wrong side of the barrier: on an M2 this
/// kernel answers correctly under that too, and the case that does not is a longer one, in
/// cubek's `the_last_cube_in_merges_the_others`.
#[cube(launch)]
fn kernel_test_sync_storage_across_cubes(
    partials: &mut [u32],
    counter: &mut [Atomic<u32>],
    out: &mut [u32],
    #[comptime] cubes: u32,
    #[comptime] units: u32,
) {
    let mut mine = Shared::<[u32]>::new_slice(units as usize);
    mine[UNIT_POS as usize] = CUBE_POS as u32 * units + UNIT_POS + 1;
    sync_cube();
    if UNIT_POS == 0 {
        let mut total = 0u32;
        let mut i = 0u32;
        while i < units {
            total += mine[i as usize];
            i += 1u32;
        }
        partials[CUBE_POS] = total;
    }

    let mut arrived = Shared::<u32>::new();
    // Release: the partial this cube just published is visible to whichever cube is last.
    sync_storage();
    if UNIT_POS == 0 {
        *arrived = counter[0].fetch_add(1u32);
    }
    // Acquire, and the cube half of the same scope: the count reaches every unit, and what the
    // cubes that arrived before published is visible to this one.
    sync_storage();

    if *arrived == cubes - 1 {
        // A stripe per unit, so the count has to have reached all of them.
        let mut sum = 0u32;
        let mut cube = UNIT_POS;
        while cube < cubes {
            sum += partials[cube as usize];
            cube += units;
        }
        out[UNIT_POS as usize] = sum;
    }
}

pub fn test_sync_storage_across_cubes<R: Runtime>(client: Client) {
    if !client.properties().features.device_memory_scope {
        // The runtime does not promise that one cube's writes reach another, so say so rather
        // than pass silently.
        std::println!("device memory scope not supported - skipped");
        return;
    }
    let ty = Type::atomic(u32::elem_type_native());
    if !client
        .properties()
        .atomic_type_usage(ty)
        .contains(AtomicUsage::Add)
    {
        std::println!("u32 atomic add not supported - skipped");
        return;
    }

    let cubes = 32u32;
    let units = core::cmp::min(32, client.properties().hardware.max_units_per_cube);

    let partials = client.empty(cubes as usize * core::mem::size_of::<u32>());
    let counter = client.create_from_slice(u32::as_bytes(&[0u32]));
    let out = client.create_from_slice(u32::as_bytes(&vec![0u32; units as usize]));

    kernel_test_sync_storage_across_cubes::launch(
        &client,
        CubeCount::Static(cubes, 1, 1),
        CubeDim::new_1d(units),
        unsafe { BufferArg::from_raw_parts(partials, cubes as usize) },
        unsafe { BufferArg::from_raw_parts(counter, 1) },
        unsafe { BufferArg::from_raw_parts(out.clone(), units as usize) },
        cubes,
        units,
    );

    let partial = |cube: u32| (0..units).map(|unit| cube * units + unit + 1).sum::<u32>();
    let expected: Vec<u32> = (0..units)
        .map(|unit| (unit..cubes).step_by(units as usize).map(partial).sum())
        .collect();

    let actual = client.read_one_unchecked(out);
    assert_eq!(u32::from_bytes(&actual), &expected);
}

/// Every cube reads every slot before writing its own, so the cubes that read early cache zeros in
/// whatever cache serves them, then writes its slot and arrives; the last to arrive sums every
/// slot. A storage barrier that orders the writes without making them visible to other cubes hands
/// it the zeros it cached: where [`kernel_test_sync_storage_across_cubes`] reads each slot once,
/// fresh, this reads over a stale copy. `before` only keeps the early reads alive; the host never
/// reads it.
#[cube(launch)]
fn kernel_test_sync_storage_over_a_stale_copy(
    slots: &mut [u32],
    counter: &mut [Atomic<u32>],
    before: &mut [u32],
    out: &mut [u32],
    #[comptime] cubes: u32,
) {
    // Every slot as it is before this cube writes, kept live so the read is made.
    let mut cached = 0u32;
    let mut slot = UNIT_POS;
    while slot < cubes {
        cached += slots[slot as usize];
        slot += CUBE_DIM;
    }
    before[CUBE_POS * CUBE_DIM as usize + UNIT_POS as usize] = cached;

    if UNIT_POS == 0 {
        slots[CUBE_POS] = CUBE_POS as u32 + 1;
    }
    let mut arrived = Shared::<u32>::new();
    // Release: this cube's slot is visible to whichever cube is last.
    sync_storage();
    if UNIT_POS == 0 {
        *arrived = counter[0].fetch_add(1u32);
    }
    // Acquire: every slot the cubes before published is visible here, over what this cube cached.
    sync_storage();

    if *arrived == cubes - 1 {
        let mut sum = 0u32;
        let mut slot = UNIT_POS;
        while slot < cubes {
            sum += slots[slot as usize];
            slot += CUBE_DIM;
        }
        out[UNIT_POS as usize] = sum;
    }
}

pub fn test_sync_storage_over_a_stale_copy<R: Runtime>(client: Client) {
    if !client.properties().features.device_memory_scope {
        std::println!("device memory scope not supported - skipped");
        return;
    }
    let ty = Type::atomic(u32::elem_type_native());
    if !client
        .properties()
        .atomic_type_usage(ty)
        .contains(AtomicUsage::Add)
    {
        std::println!("u32 atomic add not supported - skipped");
        return;
    }

    let cubes = 256u32;
    let units = core::cmp::min(64, client.properties().hardware.max_units_per_cube);

    let slots = client.create_from_slice(u32::as_bytes(&vec![0u32; cubes as usize]));
    let counter = client.create_from_slice(u32::as_bytes(&[0u32]));
    let before = client.empty((cubes * units) as usize * core::mem::size_of::<u32>());
    let out = client.create_from_slice(u32::as_bytes(&vec![0u32; units as usize]));

    kernel_test_sync_storage_over_a_stale_copy::launch(
        &client,
        CubeCount::Static(cubes, 1, 1),
        CubeDim::new_1d(units),
        unsafe { BufferArg::from_raw_parts(slots, cubes as usize) },
        unsafe { BufferArg::from_raw_parts(counter, 1) },
        unsafe { BufferArg::from_raw_parts(before, (cubes * units) as usize) },
        unsafe { BufferArg::from_raw_parts(out.clone(), units as usize) },
        cubes,
    );

    let expected: Vec<u32> = (0..units)
        .map(|unit| {
            (unit..cubes)
                .step_by(units as usize)
                .map(|slot| slot + 1)
                .sum()
        })
        .collect();
    let actual = client.read_one_unchecked(out);
    assert_eq!(u32::from_bytes(&actual), &expected);
}

#[macro_export]
macro_rules! testgen_sync_plane {
    () => {
        use super::*;

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_plane() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_plane::<TestRuntime>(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_plane_storage() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_plane_storage::<TestRuntime>(
                client,
            );
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_plane_two_stores() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_plane_two_stores::<TestRuntime>(
                client,
            );
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_cube() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_cube::<TestRuntime>(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_finished_sync_cube() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_finished_sync_cube::<TestRuntime>(
                client,
            );
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_storage_across_cubes() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_storage_across_cubes::<
                TestRuntime,
            >(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_storage_over_a_stale_copy() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_storage_over_a_stale_copy::<
                TestRuntime,
            >(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_sync_cube_shared() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_sync_cube_shared::<TestRuntime>(
                client,
            );
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_workgroup_uniform_load_synchronizes() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_workgroup_uniform_load_synchronizes::<
                TestRuntime,
            >(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_workgroup_uniform_load_atomic_synchronizes() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_workgroup_uniform_load_atomic_synchronizes::<
                TestRuntime,
            >(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_workgroup_uniform_load() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_workgroup_uniform_load::<TestRuntime>(
                client,
            );
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_workgroup_uniform_load_atomic() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_workgroup_uniform_load_atomic::<
                TestRuntime,
            >(client);
        }

        #[$crate::runtime_tests::test_log::test]
        fn test_workgroup_uniform_load_vec() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_workgroup_uniform_load_vec::<
                TestRuntime,
            >(client);
        }
    };
}

/// [`test_shared_freed_in_a_loop`], for a runtime whose launch counts shared memory as the
/// allocator lays it out.
#[allow(missing_docs)]
#[macro_export]
macro_rules! testgen_shared_freed_in_a_loop {
    () => {
        #[$crate::runtime_tests::test_log::test]
        fn test_shared_freed_in_a_loop() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::synchronization::test_shared_freed_in_a_loop::<TestRuntime>(
                client,
            );
        }
    };
}