test_module.py

# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import ctypes
import os
import pickle
import subprocess
import warnings
from pathlib import Path

import numpy as np
import pytest

import cuda.core
from cuda.core import Device, Kernel, Linker, LinkerOptions, ObjectCode, Program, ProgramOptions
from cuda.core._program import _can_load_generated_ptx
from cuda.core._utils.cuda_utils import CUDAError, driver, handle_return
from cuda.core._utils.version import binding_version, driver_version
from cuda.pathfinder import find_nvidia_binary_utility

try:
    import numba
except ImportError:
    numba = None

SAXPY_KERNEL = r"""
template<typename T>
__global__ void saxpy(const T a,
                    const T* x,
                    const T* y,
                    T* out,
                    size_t N) {
    const unsigned int tid = threadIdx.x + blockIdx.x * blockDim.x;
    for (size_t i=tid; i<N; i+=gridDim.x*blockDim.x) {
        out[tid] = a * x[tid] + y[tid];
    }
}
"""


def _is_nvfatbin_available():
    """Check if nvfatbin bindings are available."""
    try:
        from cuda.bindings import nvfatbin

        nvfatbin.version()
        return True
    except Exception:
        return False


nvfatbin_available = pytest.mark.skipif(not _is_nvfatbin_available(), reason="nvfatbin bindings not available")


@pytest.fixture(scope="module")
def cuda12_4_prerequisite_check():
    return binding_version() >= (12, 0, 0) and driver_version() >= (12, 4, 0)


@pytest.fixture(name="convert_path", params=[str, lambda p: p], ids=["str", "path"])
def convert_path_to_arg(request):
    return request.param


def test_kernel_attributes_init_disabled():
    with pytest.raises(RuntimeError, match=r"^KernelAttributes cannot be instantiated directly\."):
        cuda.core._module.KernelAttributes()  # Ensure back door is locked.


def test_kernel_attributes_per_device_view(get_saxpy_kernel_cubin):
    """kernel.attributes[device] returns a per-device view; values match."""
    kernel, _ = get_saxpy_kernel_cubin
    dev = Device()

    default_view = kernel.attributes
    int_view = kernel.attributes[dev.device_id]
    dev_view = kernel.attributes[dev]

    # Same value via every access path (default view = current device).
    assert default_view.num_regs == int_view.num_regs == dev_view.num_regs
    assert default_view.max_threads_per_block == int_view.max_threads_per_block

    # The bound views are distinct objects from the default view.
    assert int_view is not default_view
    assert int_view is not dev_view


def test_kernel_attributes_indexing_rejects_invalid_device(get_saxpy_kernel_cubin):
    """kernel.attributes[bad] raises through the Device(...) constructor."""
    kernel, _ = get_saxpy_kernel_cubin
    with pytest.raises((TypeError, ValueError, OverflowError)):
        kernel.attributes["not a device"]


def test_kernel_occupancy_init_disabled():
    with pytest.raises(RuntimeError, match=r"^KernelOccupancy cannot be instantiated directly\."):
        cuda.core._module.KernelOccupancy()  # Ensure back door is locked.


def test_kernel_init_disabled():
    with pytest.raises(RuntimeError, match=r"^Kernel objects cannot be instantiated directly\."):
        cuda.core._module.Kernel()  # Ensure back door is locked.


def test_object_code_init_disabled():
    with pytest.raises(RuntimeError, match=r"^ObjectCode objects cannot be instantiated directly\."):
        ObjectCode()  # Reject at front door.


@pytest.fixture
def get_saxpy_kernel_cubin(init_cuda):
    # prepare program
    prog = Program(SAXPY_KERNEL, code_type="c++")
    mod = prog.compile(
        "cubin",
        name_expressions=("saxpy<float>", "saxpy<double>"),
    )
    # run in single precision
    return mod.get_kernel("saxpy<float>"), mod


@pytest.fixture
def get_saxpy_kernel_ptx(init_cuda):
    # prepare program
    prog = Program(SAXPY_KERNEL, code_type="c++")
    mod = prog.compile(
        "ptx",
        name_expressions=("saxpy<float>", "saxpy<double>"),
    )
    ptx = mod.code
    return ptx, mod


@pytest.fixture
def get_saxpy_kernel_ltoir(init_cuda):
    # Create LTOIR code using link-time optimization
    prog = Program(SAXPY_KERNEL, code_type="c++", options=ProgramOptions(link_time_optimization=True))
    mod = prog.compile("ltoir", name_expressions=("saxpy<float>", "saxpy<double>"))
    return mod


@pytest.fixture
def get_saxpy_fatbin(init_cuda):
    from cuda.bindings import nvfatbin

    dev = Device()
    arch = dev.arch

    # Pick a second arch different from the current device
    second_arch = "75" if arch != "75" else "80"

    # Compile to cubin for current device arch
    prog = Program(SAXPY_KERNEL, code_type="c++", options=ProgramOptions(arch=f"sm_{arch}"))
    mod = prog.compile(
        "cubin",
        name_expressions=("saxpy<float>", "saxpy<double>"),
    )
    cubin = mod.code
    sym_map = mod.symbol_mapping

    # Compile to PTX targeting the second arch
    ptx_mod = Program(SAXPY_KERNEL, code_type="c++", options=ProgramOptions(arch=f"sm_{second_arch}")).compile(
        "ptx",
        name_expressions=("saxpy<float>", "saxpy<double>"),
    )
    ptx = ptx_mod.code

    # Create fatbin with both cubin + PTX
    handle = nvfatbin.create([], 0)
    nvfatbin.add_cubin(handle, cubin, len(cubin), arch, "saxpy")
    nvfatbin.add_ptx(handle, ptx, len(ptx), second_arch, "saxpy", f"-arch=sm_{second_arch}")
    fatbin_size = nvfatbin.size(handle)
    fatbin = bytearray(fatbin_size)
    nvfatbin.get(handle, fatbin)
    nvfatbin.destroy(handle)

    return bytes(fatbin), sym_map


@pytest.fixture(scope="module")
def get_saxpy_object():
    """Read the pre-built saxpy.o.

    In CI: produced by build stage into a test wheel file.
    In local dev: auto-built on demand if nvcc is available; if you edit
    saxpy.cu, remove the stale saxpy.o to force a rebuild.
    """
    binaries_dir = Path(__file__).parent / "test_binaries"
    obj_path = binaries_dir / "saxpy.o"

    if not obj_path.is_file():
        nvcc_path = find_nvidia_binary_utility("nvcc")
        if nvcc_path is None:
            pytest.skip(
                f"saxpy.o not found at {obj_path} and nvcc is unavailable. "
                "In CI this is downloaded from the build stage."
            )
        env = os.environ.copy()
        env["NVCC"] = nvcc_path
        subprocess.run(  # noqa: S603
            ["bash", str(binaries_dir / "build_test_binaries.sh")],  # noqa: S607
            check=True,
            env=env,
        )

    return obj_path.read_bytes()


def test_get_kernel(init_cuda):
    kernel = """extern "C" __global__ void ABC() { }"""

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        object_code = Program(kernel, "c++", options=ProgramOptions(relocatable_device_code=True)).compile("ptx")
        if any("The CUDA driver version is older than the backend version" in str(warning.message) for warning in w):
            pytest.skip("PTX version too new for current driver")

    # Verify lazy loading: get_kernel triggers module loading and returns a valid kernel
    kernel = object_code.get_kernel("ABC")
    assert object_code.handle is not None
    assert kernel.handle is not None
    # Also works with bytes
    assert object_code.get_kernel(b"ABC").handle is not None


@pytest.mark.parametrize(
    "attr, expected_type",
    [
        ("max_threads_per_block", int),
        ("shared_size_bytes", int),
        ("const_size_bytes", int),
        ("local_size_bytes", int),
        ("num_regs", int),
        ("ptx_version", int),
        ("binary_version", int),
        ("cache_mode_ca", bool),
        ("cluster_size_must_be_set", bool),
        ("max_dynamic_shared_size_bytes", int),
        ("preferred_shared_memory_carveout", int),
        ("required_cluster_width", int),
        ("required_cluster_height", int),
        ("required_cluster_depth", int),
        ("non_portable_cluster_size_allowed", bool),
        ("cluster_scheduling_policy_preference", int),
    ],
)
def test_read_only_kernel_attributes(get_saxpy_kernel_cubin, attr, expected_type):
    kernel, _ = get_saxpy_kernel_cubin

    # Default view: property access on the current-device view.
    value = getattr(kernel.attributes, attr)
    assert value is not None
    assert isinstance(value, expected_type), f"Expected {attr} to be of type {expected_type}, but got {type(value)}"

    # Per-device views via __getitem__: each device, both Device and ordinal forms.
    for device in Device.get_all_devices():
        value = getattr(kernel.attributes[device], attr)
        assert isinstance(value, expected_type), f"Expected {attr} to be of type {expected_type}, but got {type(value)}"
        value = getattr(kernel.attributes[device.device_id], attr)
        assert isinstance(value, expected_type), f"Expected {attr} to be of type {expected_type}, but got {type(value)}"


def test_object_code_load_ptx(get_saxpy_kernel_ptx):
    ptx, mod = get_saxpy_kernel_ptx
    sym_map = mod.symbol_mapping
    mod_obj = ObjectCode.from_ptx(ptx, symbol_mapping=sym_map)
    assert mod.code == ptx
    if not _can_load_generated_ptx():
        pytest.skip("PTX version too new for current driver")
    mod_obj.get_kernel("saxpy<double>")  # force loading


def test_object_code_load_ptx_from_file(get_saxpy_kernel_ptx, tmp_path, convert_path):
    ptx, mod = get_saxpy_kernel_ptx
    sym_map = mod.symbol_mapping
    assert isinstance(ptx, bytes)
    ptx_file = tmp_path / "test.ptx"
    ptx_file.write_bytes(ptx)
    arg = convert_path(ptx_file)
    mod_obj = ObjectCode.from_ptx(arg, symbol_mapping=sym_map)
    assert mod_obj.code == str(arg)
    assert mod_obj.code_type == "ptx"
    if not _can_load_generated_ptx():
        pytest.skip("PTX version too new for current driver")
    mod_obj.get_kernel("saxpy<double>")  # force loading


def test_object_code_load_cubin(get_saxpy_kernel_cubin):
    _, mod = get_saxpy_kernel_cubin
    cubin = mod.code
    sym_map = mod.symbol_mapping
    assert isinstance(cubin, bytes)
    mod = ObjectCode.from_cubin(cubin, symbol_mapping=sym_map)
    assert mod.code == cubin
    mod.get_kernel("saxpy<double>")  # force loading


def test_object_code_load_cubin_from_file(get_saxpy_kernel_cubin, tmp_path, convert_path):
    _, mod = get_saxpy_kernel_cubin
    cubin = mod.code
    sym_map = mod.symbol_mapping
    assert isinstance(cubin, bytes)
    cubin_file = tmp_path / "test.cubin"
    cubin_file.write_bytes(cubin)
    arg = convert_path(cubin_file)
    mod = ObjectCode.from_cubin(arg, symbol_mapping=sym_map)
    assert mod.code == str(arg)
    mod.get_kernel("saxpy<double>")  # force loading


def test_object_code_handle(get_saxpy_kernel_cubin):
    _, mod = get_saxpy_kernel_cubin
    assert mod.handle is not None


def test_object_code_load_ltoir(get_saxpy_kernel_ltoir):
    mod = get_saxpy_kernel_ltoir
    ltoir = mod.code
    sym_map = mod.symbol_mapping
    assert isinstance(ltoir, bytes)
    mod_obj = ObjectCode.from_ltoir(ltoir, symbol_mapping=sym_map)
    assert mod_obj.code == ltoir
    assert mod_obj.code_type == "ltoir"
    # ltoir doesn't support kernel retrieval directly as it's used for linking
    # Test that get_kernel fails for unsupported code type
    with pytest.raises(RuntimeError, match=r'Unsupported code type "ltoir"'):
        mod_obj.get_kernel("saxpy<float>")


def test_object_code_load_ltoir_from_file(get_saxpy_kernel_ltoir, tmp_path, convert_path):
    mod = get_saxpy_kernel_ltoir
    ltoir = mod.code
    sym_map = mod.symbol_mapping
    assert isinstance(ltoir, bytes)
    ltoir_file = tmp_path / "test.ltoir"
    ltoir_file.write_bytes(ltoir)
    arg = convert_path(ltoir_file)
    mod_obj = ObjectCode.from_ltoir(arg, symbol_mapping=sym_map)
    assert mod_obj.code == str(arg)
    assert mod_obj.code_type == "ltoir"
    # ltoir doesn't support kernel retrieval directly as it's used for linking


@nvfatbin_available
def test_object_code_load_fatbin(get_saxpy_fatbin):
    fatbin, sym_map = get_saxpy_fatbin
    assert isinstance(fatbin, bytes)
    mod_obj = ObjectCode.from_fatbin(fatbin, symbol_mapping=sym_map)
    assert mod_obj.code == fatbin
    assert mod_obj.code_type == "fatbin"
    mod_obj.get_kernel("saxpy<double>")  # force loading


@nvfatbin_available
def test_object_code_load_fatbin_from_file(get_saxpy_fatbin, tmp_path, convert_path):
    fatbin, sym_map = get_saxpy_fatbin
    assert isinstance(fatbin, bytes)
    fatbin_file = tmp_path / "test.fatbin"
    fatbin_file.write_bytes(fatbin)
    arg = convert_path(fatbin_file)
    mod_obj = ObjectCode.from_fatbin(arg, symbol_mapping=sym_map)
    assert mod_obj.code == str(arg)
    assert mod_obj.code_type == "fatbin"
    mod_obj.get_kernel("saxpy<double>")  # force loading


def test_object_code_load_object(get_saxpy_object):
    obj = get_saxpy_object
    assert isinstance(obj, bytes)
    mod_obj = ObjectCode.from_object(obj)
    assert mod_obj.code == obj
    assert mod_obj.code_type == "object"
    with pytest.raises(RuntimeError, match=r'Unsupported code type "object"'):
        mod_obj.get_kernel("saxpy<float>")


def test_object_code_load_object_from_file(get_saxpy_object, tmp_path):
    obj_file = tmp_path / "test.o"
    obj_file.write_bytes(get_saxpy_object)
    arg = str(obj_file)
    mod_obj = ObjectCode.from_object(arg)
    assert mod_obj.code == arg
    assert mod_obj.code_type == "object"


def test_object_code_load_object_with_linker(get_saxpy_object, init_cuda):
    arch = f"sm_{init_cuda.arch}"
    kernel_code = Program(
        r"""
        extern __device__ float saxpy_step(float a, float x, float y);
        extern "C" __global__ void linked_kernel(float a, float x, float y, float* out) {
            if (threadIdx.x == 0 && blockIdx.x == 0) *out = saxpy_step(a, x, y);
        }
        """,
        "c++",
        ProgramOptions(relocatable_device_code=True, arch=arch),
    ).compile("cubin")
    linked = Linker(
        kernel_code,
        ObjectCode.from_object(get_saxpy_object),
        options=LinkerOptions(arch=arch),
    ).link("cubin")
    kernel = linked.get_kernel("linked_kernel")

    stream = init_cuda.create_stream()
    host_buf = cuda.core.LegacyPinnedMemoryResource().allocate(4)
    result = np.from_dlpack(host_buf).view(np.float32)
    result[:] = 0.0
    dev_buf = init_cuda.memory_resource.allocate(4, stream=init_cuda.default_stream)

    cuda.core.launch(
        stream,
        cuda.core.LaunchConfig(grid=1, block=1),
        kernel,
        np.float32(2.0),
        np.float32(3.0),
        np.float32(4.0),
        dev_buf,
    )
    dev_buf.copy_to(host_buf, stream=stream)
    stream.sync()

    assert result[0] == 10.0


def test_saxpy_arguments(get_saxpy_kernel_cubin, cuda12_4_prerequisite_check):
    krn, _ = get_saxpy_kernel_cubin

    if cuda12_4_prerequisite_check:
        assert krn.num_arguments == 5
    else:
        with pytest.raises(NotImplementedError):
            _ = krn.num_arguments
        return

    arg_info = krn.arguments_info
    n_args = len(arg_info)
    assert n_args == krn.num_arguments

    class ExpectedStruct(ctypes.Structure):
        _fields_ = [
            ("a", ctypes.c_float),
            ("x", ctypes.POINTER(ctypes.c_float)),
            ("y", ctypes.POINTER(ctypes.c_float)),
            ("out", ctypes.POINTER(ctypes.c_float)),
            ("N", ctypes.c_size_t),
        ]

    offsets = [p.offset for p in arg_info]
    sizes = [p.size for p in arg_info]
    members = [getattr(ExpectedStruct, name) for name, _ in ExpectedStruct._fields_]
    expected_offsets = tuple(m.offset for m in members)
    assert all(actual == expected for actual, expected in zip(offsets, expected_offsets))
    expected_sizes = tuple(m.size for m in members)
    assert all(actual == expected for actual, expected in zip(sizes, expected_sizes))


@pytest.mark.parametrize("nargs", [0, 1, 2, 3, 16])
@pytest.mark.parametrize("c_type_name,c_type", [("int", ctypes.c_int), ("short", ctypes.c_short)], ids=["int", "short"])
def test_num_arguments(init_cuda, nargs, c_type_name, c_type, cuda12_4_prerequisite_check):
    if not cuda12_4_prerequisite_check:
        pytest.skip("Test requires CUDA 12")
    args_str = ", ".join([f"{c_type_name} p_{i}" for i in range(nargs)])
    src = f"__global__ void foo{nargs}({args_str}) {{ }}"
    prog = Program(src, code_type="c++")
    mod = prog.compile(
        "cubin",
        name_expressions=(f"foo{nargs}",),
    )
    krn = mod.get_kernel(f"foo{nargs}")
    assert krn.num_arguments == nargs

    class ExpectedStruct(ctypes.Structure):
        _fields_ = [(f"arg_{i}", c_type) for i in range(nargs)]

    members = tuple(getattr(ExpectedStruct, f"arg_{i}") for i in range(nargs))

    arg_info = krn.arguments_info
    assert all(actual.offset == expected.offset for actual, expected in zip(arg_info, members))
    assert all(actual.size == expected.size for actual, expected in zip(arg_info, members))


def test_num_args_error_handling(deinit_all_contexts_function, cuda12_4_prerequisite_check):
    if not cuda12_4_prerequisite_check:
        pytest.skip("Test requires CUDA 12")
    src = "__global__ void foo(int a) { }"
    prog = Program(src, code_type="c++")
    mod = prog.compile(
        "cubin",
        name_expressions=("foo",),
    )
    krn = mod.get_kernel("foo")
    # empty driver's context stack using function from conftest
    deinit_all_contexts_function()
    # with no current context, cuKernelGetParamInfo would report
    # exception which we expect to handle by raising
    with pytest.raises(CUDAError):
        # assignment resolves linter error "B018: useless expression"
        _ = krn.num_arguments


@pytest.mark.parametrize("block_size", [32, 64, 96, 120, 128, 256])
@pytest.mark.parametrize("smem_size_per_block", [0, 32, 4096])
def test_occupancy_max_active_block_per_multiprocessor(get_saxpy_kernel_cubin, block_size, smem_size_per_block):
    kernel, _ = get_saxpy_kernel_cubin
    dev_props = Device().properties
    assert block_size <= dev_props.max_threads_per_block
    assert smem_size_per_block <= dev_props.max_shared_memory_per_block
    num_blocks_per_sm = kernel.occupancy.max_active_blocks_per_multiprocessor(block_size, smem_size_per_block)
    assert isinstance(num_blocks_per_sm, int)
    assert num_blocks_per_sm > 0
    kernel_threads_per_sm = num_blocks_per_sm * block_size
    kernel_smem_size_per_sm = num_blocks_per_sm * smem_size_per_block
    assert kernel_threads_per_sm <= dev_props.max_threads_per_multiprocessor
    assert kernel_smem_size_per_sm <= dev_props.max_shared_memory_per_multiprocessor
    assert kernel.attributes.num_regs * num_blocks_per_sm <= dev_props.max_registers_per_multiprocessor


@pytest.mark.parametrize("block_size_limit", [32, 64, 96, 120, 128, 256, 0])
@pytest.mark.parametrize("smem_size_per_block", [0, 32, 4096])
def test_occupancy_max_potential_block_size_constant(get_saxpy_kernel_cubin, block_size_limit, smem_size_per_block):
    """Tests use case when shared memory needed is independent on the block size"""
    kernel, _ = get_saxpy_kernel_cubin
    dev_props = Device().properties
    assert block_size_limit <= dev_props.max_threads_per_block
    assert smem_size_per_block <= dev_props.max_shared_memory_per_block
    config_data = kernel.occupancy.max_potential_block_size(smem_size_per_block, block_size_limit)
    assert isinstance(config_data, tuple)
    assert len(config_data) == 2
    min_grid_size, max_block_size = config_data
    assert isinstance(min_grid_size, int)
    assert isinstance(max_block_size, int)
    assert min_grid_size > 0
    assert max_block_size > 0
    if block_size_limit > 0:
        assert max_block_size <= block_size_limit
    else:
        assert max_block_size <= dev_props.max_threads_per_block
    assert min_grid_size == config_data.min_grid_size
    assert max_block_size == config_data.max_block_size
    invalid_dsmem = Ellipsis
    with pytest.raises(TypeError):
        kernel.occupancy.max_potential_block_size(invalid_dsmem, block_size_limit)


@pytest.mark.skipif(numba is None, reason="Test requires numba to be installed")
@pytest.mark.parametrize("block_size_limit", [32, 64, 96, 120, 128, 277, 0])
def test_occupancy_max_potential_block_size_b2dsize(get_saxpy_kernel_cubin, block_size_limit):
    """Tests use case when shared memory needed depends on the block size"""
    kernel, _ = get_saxpy_kernel_cubin

    def shared_memory_needed(block_size: numba.intc) -> numba.size_t:
        "Size of dynamic shared memory needed by kernel of this block size"
        return 1024 * (block_size // 32)

    b2dsize_sig = numba.size_t(numba.intc)
    dsmem_needed_cfunc = numba.cfunc(b2dsize_sig)(shared_memory_needed)
    fn_ptr = ctypes.cast(dsmem_needed_cfunc.ctypes, ctypes.c_void_p).value
    b2dsize_fn = driver.CUoccupancyB2DSize(_ptr=fn_ptr)
    config_data = kernel.occupancy.max_potential_block_size(b2dsize_fn, block_size_limit)
    dev_props = Device().properties
    assert block_size_limit <= dev_props.max_threads_per_block
    min_grid_size, max_block_size = config_data
    assert isinstance(min_grid_size, int)
    assert isinstance(max_block_size, int)
    assert min_grid_size > 0
    assert max_block_size > 0
    if block_size_limit > 0:
        assert max_block_size <= block_size_limit
    else:
        assert max_block_size <= dev_props.max_threads_per_block


@pytest.mark.parametrize("num_blocks_per_sm, block_size", [(4, 32), (2, 64), (2, 96), (3, 120), (2, 128), (1, 256)])
def test_occupancy_available_dynamic_shared_memory_per_block(get_saxpy_kernel_cubin, num_blocks_per_sm, block_size):
    kernel, _ = get_saxpy_kernel_cubin
    dev_props = Device().properties
    assert block_size <= dev_props.max_threads_per_block
    assert num_blocks_per_sm * block_size <= dev_props.max_threads_per_multiprocessor
    smem_size = kernel.occupancy.available_dynamic_shared_memory_per_block(num_blocks_per_sm, block_size)
    assert smem_size <= dev_props.max_shared_memory_per_block
    assert num_blocks_per_sm * smem_size <= dev_props.max_shared_memory_per_multiprocessor


@pytest.mark.parametrize("cluster", [None, 2])
def test_occupancy_max_active_clusters(get_saxpy_kernel_cubin, cluster):
    kernel, _ = get_saxpy_kernel_cubin
    dev = Device()
    if dev.compute_capability < (9, 0):
        pytest.skip("Device with compute capability 90 or higher is required for cluster support")
    launch_config = cuda.core.LaunchConfig(grid=128, block=64, cluster=cluster)
    query_fn = kernel.occupancy.max_active_clusters
    with pytest.raises(TypeError, match=r"keyword-only argument"):
        query_fn(launch_config)
    max_active_clusters = query_fn(launch_config, stream=dev.default_stream)
    assert isinstance(max_active_clusters, int)
    assert max_active_clusters >= 0


def test_occupancy_max_potential_cluster_size(get_saxpy_kernel_cubin):
    kernel, _ = get_saxpy_kernel_cubin
    dev = Device()
    if dev.compute_capability < (9, 0):
        pytest.skip("Device with compute capability 90 or higher is required for cluster support")
    launch_config = cuda.core.LaunchConfig(grid=128, block=64)
    query_fn = kernel.occupancy.max_potential_cluster_size
    with pytest.raises(TypeError, match=r"keyword-only argument"):
        query_fn(launch_config)
    max_potential_cluster_size = query_fn(launch_config, stream=dev.default_stream)
    assert isinstance(max_potential_cluster_size, int)
    assert max_potential_cluster_size >= 0


def test_module_serialization_roundtrip(get_saxpy_kernel_cubin):
    _, objcode = get_saxpy_kernel_cubin
    result = pickle.loads(pickle.dumps(objcode))  # noqa: S301

    assert isinstance(result, ObjectCode)
    assert objcode.code == result.code
    assert objcode.symbol_mapping == result.symbol_mapping
    assert objcode.code_type == result.code_type


def test_kernel_from_handle(get_saxpy_kernel_cubin):
    """Test Kernel.from_handle() with a valid handle"""
    original_kernel, objcode = get_saxpy_kernel_cubin

    # Get the handle from the original kernel
    handle = int(original_kernel.handle)

    # Create a new Kernel from the handle
    kernel_from_handle = Kernel.from_handle(handle, objcode)
    assert isinstance(kernel_from_handle, Kernel)

    # Verify we can access kernel attributes
    max_threads = kernel_from_handle.attributes.max_threads_per_block
    assert isinstance(max_threads, int)
    assert max_threads > 0


def test_kernel_from_handle_no_module(get_saxpy_kernel_cubin):
    """Test Kernel.from_handle() without providing a module"""
    original_kernel, _ = get_saxpy_kernel_cubin

    # Get the handle from the original kernel
    handle = int(original_kernel.handle)

    # Create a new Kernel from the handle without a module
    # This is supported on CUDA 12+ backend (CUkernel)
    kernel_from_handle = Kernel.from_handle(handle)
    assert isinstance(kernel_from_handle, Kernel)

    # Verify we can still access kernel attributes
    max_threads = kernel_from_handle.attributes.max_threads_per_block
    assert isinstance(max_threads, int)
    assert max_threads > 0


@pytest.mark.parametrize(
    "invalid_value",
    [
        pytest.param("not_an_int", id="str"),
        pytest.param(2.71828, id="float"),
        pytest.param(None, id="None"),
        pytest.param({"handle": 123}, id="dict"),
        pytest.param([456], id="list"),
        pytest.param((789,), id="tuple"),
        pytest.param(3 + 4j, id="complex"),
        pytest.param(b"\xde\xad\xbe\xef", id="bytes"),
        pytest.param({999}, id="set"),
        pytest.param(object(), id="object"),
    ],
)
def test_kernel_from_handle_type_validation(invalid_value):
    """Test Kernel.from_handle() with wrong handle types"""
    with pytest.raises(TypeError):
        Kernel.from_handle(invalid_value)


def test_kernel_from_handle_invalid_module_type(get_saxpy_kernel_cubin):
    """Test Kernel.from_handle() with invalid module parameter"""
    original_kernel, _ = get_saxpy_kernel_cubin
    handle = int(original_kernel.handle)

    # Invalid module type (should fail type assertion in _from_obj)
    with pytest.raises((TypeError, AssertionError)):
        Kernel.from_handle(handle, mod="not_an_objectcode")

    with pytest.raises((TypeError, AssertionError)):
        Kernel.from_handle(handle, mod=12345)


def test_kernel_from_handle_multiple_instances(get_saxpy_kernel_cubin):
    """Test creating multiple Kernel instances from the same handle"""
    original_kernel, objcode = get_saxpy_kernel_cubin
    handle = int(original_kernel.handle)

    # Create multiple Kernel instances from the same handle
    kernel1 = Kernel.from_handle(handle, objcode)
    kernel2 = Kernel.from_handle(handle, objcode)
    kernel3 = Kernel.from_handle(handle, objcode)

    # All should be valid Kernel objects
    assert isinstance(kernel1, Kernel)
    assert isinstance(kernel2, Kernel)
    assert isinstance(kernel3, Kernel)

    # All should reference the same underlying CUDA kernel handle
    assert int(kernel1.handle) == int(kernel2.handle) == int(kernel3.handle) == handle


def test_kernel_from_handle_library_mismatch_warning(init_cuda):
    """Kernel.from_handle warns when caller-supplied module differs from the kernel's library."""
    prog1 = Program(SAXPY_KERNEL, code_type="c++")
    mod1 = prog1.compile("cubin", name_expressions=("saxpy<float>",))
    kernel = mod1.get_kernel("saxpy<float>")
    handle = int(kernel.handle)

    prog2 = Program(SAXPY_KERNEL, code_type="c++")
    mod2 = prog2.compile("cubin", name_expressions=("saxpy<float>",))
    mod2.get_kernel("saxpy<float>")

    with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter("always")
        k = Kernel.from_handle(handle, mod2)
        assert len(w) == 1
        assert "does not match" in str(w[0].message)

    assert k.attributes.max_threads_per_block > 0


def test_kernel_from_handle_foreign_kernel(init_cuda):
    """Kernel.from_handle with a driver-level kernel not created by cuda.core."""
    prog = Program(SAXPY_KERNEL, code_type="c++")
    mod = prog.compile("cubin", name_expressions=("saxpy<float>",))
    cubin = mod.code
    sym_map = mod.symbol_mapping

    cu_lib = handle_return(driver.cuLibraryLoadData(cubin, [], [], 0, [], [], 0))
    mangled = sym_map["saxpy<float>"]
    cu_kernel = handle_return(driver.cuLibraryGetKernel(cu_lib, mangled))
    handle = int(cu_kernel)

    k = Kernel.from_handle(handle)
    assert k.attributes.max_threads_per_block > 0


def test_kernel_keeps_library_alive(init_cuda):
    """Test that a Kernel keeps its underlying library alive after ObjectCode goes out of scope."""
    import gc

    import numpy as np

    def get_kernel_only():
        """Return a kernel, letting ObjectCode go out of scope."""
        code = """
        extern "C" __global__ void write_value(int* out) {
            if (threadIdx.x == 0 && blockIdx.x == 0) {
                *out = 42;
            }
        }
        """
        program = Program(code, "c++")
        object_code = program.compile("cubin")
        kernel = object_code.get_kernel("write_value")
        # ObjectCode goes out of scope here
        return kernel

    kernel = get_kernel_only()

    # Force GC to ensure ObjectCode destructor runs
    gc.collect()

    # Kernel should still be valid
    assert kernel.handle is not None
    assert kernel.num_arguments == 1

    # Actually launch the kernel to prove library is still loaded
    device = Device()
    stream = device.create_stream()

    # Allocate pinned host buffer and device buffer
    pinned_mr = cuda.core.LegacyPinnedMemoryResource()
    host_buf = pinned_mr.allocate(4)  # sizeof(int)
    result = np.from_dlpack(host_buf).view(np.int32)
    result[:] = 0

    dev_buf = device.memory_resource.allocate(4, stream=device.default_stream)

    # Launch kernel
    config = cuda.core.LaunchConfig(grid=1, block=1)
    cuda.core.launch(stream, config, kernel, dev_buf)

    # Copy result back to host
    dev_buf.copy_to(host_buf, stream=stream)
    stream.sync()

    assert result[0] == 42, f"Expected 42, got {result[0]}"