cubecl-core 0.11.0-pre.4

use cubecl_runtime::runtime::Runtime;
use std::{boxed::Box, println};

use alloc::string::{String, ToString};

use crate::{self as cubecl, as_bytes};
use cubecl::prelude::*;
use cubecl_runtime::server::{Handle, LaunchError, ResourceLimitError, ServerError};

#[derive(CubeLaunch, CubeType)]
pub struct ComptimeTag {
    array: Box<[f32]>,
    #[cube(comptime)]
    tag: String,
}

#[cube(launch)]
pub fn kernel_with_comptime_tag(mut output: ComptimeTag) {
    if UNIT_POS == 0 {
        if comptime![&output.tag == "zero"] {
            output.array[0] = f32::new(0f32);
        } else {
            output.array[0] = f32::new(1f32);
        }
    }
}

#[cube(launch)]
pub fn kernel_with_comptime(#[comptime] factor: u32) {
    let _ = factor;
}

#[cube(launch)]
pub fn kernel_with_generics<F: Float>(output: &mut [F]) {
    if UNIT_POS == 0 {
        output[0] = F::new(5f32);
    }
}

#[cube(launch)]
pub fn kernel_without_generics(output: &mut [f32]) {
    if UNIT_POS == 0 {
        output[0] = 5.0;
    }
}

#[cube(launch, address_type = "dynamic")]
pub fn kernel_dynamic_addressing(output: &mut [f32]) {
    if UNIT_POS == 0 {
        output[0] = 5.0;
    }
}

#[cube(launch)]
pub fn kernel_inplace(input: &[f32], output: &mut [f32]) {
    if UNIT_POS == 0 {
        output[0] = input[0] + 5.0;
    }
}

#[cube(launch)]
pub fn kernel_with_max_shared(
    output: &mut [u32],
    #[comptime] shared_size_1: usize,
    #[comptime] shared_size_2: usize,
) {
    let mut shared_1 = Shared::new_slice(shared_size_1);
    let mut shared_2 = Shared::new_slice(shared_size_2);
    if UNIT_POS < 8 {
        shared_1[shared_size_1 - UNIT_POS as usize - 1] = output[UNIT_POS as usize];
        shared_2[shared_size_2 - UNIT_POS as usize - 1] = output[8 - UNIT_POS as usize];
    }
    sync_cube();
    if UNIT_POS < 8 {
        let a = shared_1[shared_size_1 - UNIT_POS as usize - 2];
        let b = shared_2[shared_size_2 - UNIT_POS as usize - 1];
        output[UNIT_POS as usize] = a + b;
    }
}

#[cube(launch)]
pub fn kernel_resource_errors(output: &mut [u32], #[comptime] shared_size: usize) {
    let mut shared = Shared::new_slice(shared_size);
    // Add some dummy code to prevent smem from being optimized out
    shared[0] = 0;
    sync_cube();
    output[0] = shared[0];
}

pub fn test_kernel_with_comptime_tag<R: Runtime>(client: Client) {
    let handle = client.create_from_slice(f32::as_bytes(&[5.0]));
    let array_arg = unsafe { BufferArg::from_raw_parts(handle.clone(), 1) };

    kernel_with_comptime_tag::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        ComptimeTagLaunch::new(array_arg, "zero".to_string()),
    );

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

    assert_eq!(actual[0], f32::new(0.0));

    let handle = client.create_from_slice(f32::as_bytes(&[5.0]));
    let array_arg = unsafe { BufferArg::from_raw_parts(handle.clone(), 1) };

    kernel_with_comptime_tag::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        ComptimeTagLaunch::new(array_arg, "not_zero".to_string()),
    );

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

    assert_eq!(actual[0], f32::new(1.0));
}

pub fn test_kernel_comptime_name_disambiguation<R: Runtime>(client: Client) {
    let settings = KernelSettings::new(
        *CubeDim::new_single(),
        ExecutionMode::Checked,
        AddressType::U32,
    );

    let k1 = kernel_with_comptime::KernelWithComptime::new(
        settings.clone(),
        client.properties_shared(),
        client.target_properties_shared(),
        1,
    );
    let k2 = kernel_with_comptime::KernelWithComptime::new(
        settings,
        client.properties_shared(),
        client.target_properties_shared(),
        2,
    );

    let def1 = k1.define();
    let def2 = k2.define();

    assert_ne!(
        def1.settings.kernel_name, def2.settings.kernel_name,
        "Kernels with distinct comptime arguments must have distinct kernel names"
    );
    assert!(
        def1.settings
            .kernel_name
            .starts_with("kernel_with_comptime_")
    );
    assert!(
        def2.settings
            .kernel_name
            .starts_with("kernel_with_comptime_")
    );
}

pub fn test_kernel_with_generics<R: Runtime, F: Float + CubeElement>(client: Client) {
    let handle = client.create_from_slice(as_bytes![F: 0.0, 1.0]);

    kernel_with_generics::launch::<F>(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
    );

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

    assert_eq!(actual[0], F::new(5.0));
}

pub fn test_kernel_without_generics<R: Runtime>(client: Client) {
    let handle = client.create_from_slice(f32::as_bytes(&[0.0, 1.0]));

    kernel_without_generics::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
    );

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

    assert_eq!(actual[0], 5.0);
}

pub fn test_kernel_inplace<R: Runtime>(client: Client) {
    let handle = client.create_from_slice(f32::as_bytes(&[0.0, 1.0]));

    kernel_inplace::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
        BufferArg::alias(0, 2),
    );

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

    assert_eq!(actual[0], 5.0);
}

pub fn test_kernel_zero_cube_count<R: Runtime>(client: Client) {
    // A zero-element fill resolves to `Static(0, 0, 0)`. Launching it is a no-op.
    let handle = client.create_from_slice(f32::as_bytes(&[7.0, 8.0]));

    kernel_without_generics::launch(
        &client,
        CubeCount::Static(0, 0, 0),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
    );

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

    assert_eq!(actual, &[7.0, 8.0]);
}

pub fn test_kernel_dynamic_zero_cube_count<R: Runtime>(client: Client) {
    // The (0, 0, 0) count lives in a buffer, so the client guard can't see it.
    // Scoped to CUDA/HIP because wgpu can't bind a storage buffer for indirect
    // dispatch.
    let handle = client.create_from_slice(f32::as_bytes(&[7.0, 8.0]));
    let count = client.create_from_slice(u32::as_bytes(&[0, 0, 0]));

    kernel_without_generics::launch(
        &client,
        CubeCount::Dynamic(count.binding()),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
    );

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

    assert_eq!(actual, &[7.0, 8.0]);
}

pub fn test_kernel_max_shared<R: Runtime>(client: Client) {
    let total_shared_size = client.properties().hardware.max_shared_memory_size;

    let handle = client.create_from_slice(u32::as_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]));

    // Allocate 24kibi to a check buffer, and the rest to the second buffer
    let shared_size_1 = 24576 / size_of::<u32>();
    let shared_size_2 = (total_shared_size - 24576) / size_of::<u32>();

    kernel_with_max_shared::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 8) },
        shared_size_1,
        shared_size_2,
    );

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

    assert_eq!(actual, &[1, 9, 9, 9, 9, 9, 9, 1]);
}

/// The resource-limit error a launch was refused with, taken from the buffer
/// the launch was going to write.
///
/// A refused launch is not the flush's to report — it left the output
/// unwritten, and the claim on those bytes is what carries the reason. So the
/// read is the assertion: it must fail, and it must fail on the limit that
/// stopped the launch rather than on anything the runtime invented.
fn resource_error(client: &Client, out: Handle) -> ResourceLimitError {
    let err = client
        .read_one(out)
        .expect_err("a refused launch never wrote the buffer, so the read must fail");

    let ServerError::Several { mut errors, .. } = err else {
        panic!("a read of an unwritten buffer reports its failures, got: {err}");
    };
    match errors.remove(0) {
        ServerError::Unwritten { root, .. } => match *root {
            ServerError::Launch(LaunchError::TooManyResources(inner)) => inner,
            other => panic!("should be a resource error, is {other:?}"),
        },
        other => panic!("should be an unwritten-bytes report, is {other:?}"),
    }
}

pub fn test_shared_memory_error<R: Runtime>(client: Client) {
    // A CPU runtime emulates the cube model rather than dispatching it, so it
    // enforces none of these limits and refuses nothing.
    if client.properties().hardware.num_cpu_cores.is_some() {
        return;
    }

    let max_shared_size = client.properties().hardware.max_shared_memory_size;

    let shared_size = (max_shared_size + 1).div_ceil(size_of::<u32>());
    let requested_bytes = shared_size * size_of::<u32>();

    let handle = client.create_from_slice(u32::as_bytes(&[0]));
    kernel_resource_errors::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 1) },
        shared_size,
    );

    match resource_error(&client, handle) {
        ResourceLimitError::SharedMemory { requested, max, .. } => {
            assert_eq!(
                requested_bytes, requested,
                "Requested should be equal to requested size"
            );
            assert_eq!(
                max_shared_size, max,
                "Max should be equal to max shared size"
            );
        }
        other => panic!("should be a shared-memory limit, is {other:?}"),
    }
}

pub fn test_cube_dim_error<R: Runtime>(client: Client) {
    // A CPU runtime emulates the cube model rather than dispatching it, so it
    // enforces none of these limits and refuses nothing. Same reason the
    // shared-memory case skips it.
    if client.properties().hardware.num_cpu_cores.is_some() {
        return;
    }
    let max_cube_dim = client.properties().hardware.max_cube_dim;
    let max_units = client.properties().hardware.max_units_per_cube;

    // CPU has no limit, and + 1 will overflow
    if max_cube_dim.2 == u32::MAX {
        return;
    }

    let handle = client.create_from_slice(u32::as_bytes(&[0]));

    kernel_resource_errors::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_3d(1, 1, max_cube_dim.2 + 1),
        unsafe { BufferArg::from_raw_parts(handle.clone(), 1) },
        1,
    );
    match resource_error(&client, handle) {
        ResourceLimitError::CubeDim { requested, max, .. } => {
            assert_eq!((1, 1, max_cube_dim.2 + 1), requested);
            assert_eq!(max_cube_dim, max);
        }
        // Could also be valid
        ResourceLimitError::Units { requested, max, .. } if max_cube_dim.2 >= max_units => {
            assert_eq!(max_cube_dim.2 + 1, requested);
            assert_eq!(max_units, max);
        }
        other => panic!("should be a cube-dim limit, is {other:?}"),
    }
}

pub fn test_max_units_error<R: Runtime>(client: Client) {
    // A CPU runtime emulates the cube model rather than dispatching it, so it
    // enforces none of these limits and refuses nothing. Same reason the
    // shared-memory case skips it.
    if client.properties().hardware.num_cpu_cores.is_some() {
        return;
    }
    let max_cube_dim = client.properties().hardware.max_cube_dim;
    // CPU has no limit, and num_elems will overflow
    if max_cube_dim.2 == u32::MAX {
        return;
    }

    let max_units = client.properties().hardware.max_units_per_cube;
    let cube_dim: CubeDim = max_cube_dim.into();
    let requested_units = cube_dim.num_elems();

    let handle = client.create_from_slice(u32::as_bytes(&[0]));

    kernel_resource_errors::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        cube_dim,
        unsafe { BufferArg::from_raw_parts(handle.clone(), 1) },
        1,
    );

    match resource_error(&client, handle) {
        ResourceLimitError::Units { requested, max, .. } => {
            assert_eq!(requested_units, requested);
            assert_eq!(max_units, max);
        }
        other => panic!("should be a units limit, is {other:?}"),
    }
}

pub fn test_kernel_dynamic_addressing<R: Runtime>(client: Client, address_type: AddressType) {
    let handle = client.create_from_slice(f32::as_bytes(&[0.0, 1.0]));

    if !client.properties().supports_address(address_type) {
        println!("Skipping dynamic addressing kernel, no type support");
        return;
    }

    kernel_dynamic_addressing::launch(
        &client,
        CubeCount::Static(1, 1, 1),
        CubeDim::new_1d(1),
        address_type,
        unsafe { BufferArg::from_raw_parts(handle.clone(), 2) },
    );

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

    assert_eq!(actual[0], 5.0);
}

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

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

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

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

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

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

        #[ignore = "Seemingly flaky with CPU emulation"]
        #[$crate::runtime_tests::test_log::test]
        fn test_launch_with_max_shared() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_kernel_max_shared::<TestRuntime>(client);
        }
    };
}

/// Launch tests for backends that resolve a dynamic cube count to host (CUDA, HIP).
#[allow(missing_docs)]
#[macro_export]
macro_rules! testgen_launch_dynamic_count {
    () => {
        mod launch_dynamic_count {
            use super::*;

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

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

#[macro_export]
macro_rules! testgen_launch_untyped {
    () => {
        #[test]
        fn test_launch_dynamic_addressing_32() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_kernel_dynamic_addressing::<TestRuntime>(
                client.clone(),
                AddressType::U32,
            );
        }

        #[test]
        fn test_launch_dynamic_addressing_64() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_kernel_dynamic_addressing::<TestRuntime>(
                client,
                AddressType::U64,
            );
        }

        #[test]
        fn test_launch_shared_memory_error() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_shared_memory_error::<TestRuntime>(client);
        }

        #[test]
        fn test_launch_cube_dim_error() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_cube_dim_error::<TestRuntime>(client);
        }

        #[test]
        fn test_launch_units_error() {
            let client = TestRuntime::client(&Default::default());
            cubecl_core::runtime_tests::launch::test_max_units_error::<TestRuntime>(client);
        }
    };
}