update tests, couple TODOs left

Author: ksimpson-work

cuda_core/tests/test_context.py

@@ -1,4 +1,4 @@
-from cuda.core._context import Context
+from cuda.core.experimental._context import Context
 
 def test_context_initialization():
     try:

cuda_core/tests/test_device.py

@@ -1,5 +1,53 @@
-from cuda.core._device import Device
+from cuda import cuda, cudart
+from cuda.core.experimental._device import Device
+from cuda.core.experimental._utils import handle_return, ComputeCapability, CUDAError, \
+                             precondition
 
 def test_device_initialization():
     device = Device()
-    assert device is not None
+    assert device is not None
+
+def test_device_repr():
+    device = Device()
+    assert str(device).startswith('<Device 0')
+
+def test_device_set_current():
+    device = Device()
+    device.set_current()
+
+def test_device_create_stream():
+    device = Device()
+    stream = device.create_stream()
+    assert stream is not None
+
+
+def test_pci_bus_id():
+    device = Device(0)
+    bus_id = handle_return(cudart.cudaDeviceGetPCIBusId(13, device.device_id))
+    assert device.pci_bus_id == bus_id[:12].decode()
+
+def test_uuid():
+    device = Device(0)
+    driver_ver = handle_return(cuda.cuDriverGetVersion())
+    if driver_ver >= 11040:
+        uuid = handle_return(cuda.cuDeviceGetUuid_v2(device.device_id))
+    else:
+        uuid = handle_return(cuda.cuDeviceGetUuid(device.device_id))
+    uuid = uuid.bytes.hex()
+    expected_uuid = f"{uuid[:8]}-{uuid[8:12]}-{uuid[12:16]}-{uuid[16:20]}-{uuid[20:]}"
+    assert device.uuid == expected_uuid
+
+def test_name():
+    device = Device(0)
+    name = handle_return(cuda.cuDeviceGetName(128, device.device_id))
+    name = name.split(b'\0')[0]
+    assert device.name == name.decode()
+
+def test_compute_capability():
+    device = Device(0)
+    major = handle_return(cudart.cudaDeviceGetAttribute(
+        cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, device.device_id))
+    minor = handle_return(cudart.cudaDeviceGetAttribute(
+        cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, device.device_id))
+    expected_cc = ComputeCapability(major, minor)
+    assert device.compute_capability == expected_cc

cuda_core/tests/test_event.py

@@ -1,9 +1,38 @@
-from cuda.core._event import Event
+from cuda import cuda
+from cuda.core.experimental._event import  EventOptions, Event
+from cuda.core.experimental._utils import handle_return
 
-def test_event_initialization():
+def test_is_timing_disabled():
+    options = EventOptions(enable_timing=False)
+    event = Event._init(options)
+    assert event.is_timing_disabled == True
+
+def test_is_sync_busy_waited():
+    options = EventOptions(busy_waited_sync=True)
+    event = Event._init(options)
+    assert event.is_sync_busy_waited == True
+
+def test_is_ipc_supported():
+    options = EventOptions(support_ipc=True)
     try:
-        event = Event()
-    except NotImplementedError as e:
+        event = Event._init(options)
+    except NotImplementedError:
         assert True
     else:
-        assert False, "Expected NotImplementedError was not raised"
+        assert False
+
+def test_sync():
+    options = EventOptions()
+    event = Event._init(options)
+    event.sync()
+    assert event.is_done == True
+
+def test_is_done():
+    options = EventOptions()
+    event = Event._init(options)
+    assert event.is_done == True
+
+def test_handle():
+    options = EventOptions()
+    event = Event._init(options)
+    assert isinstance(event.handle, int)

cuda_core/tests/test_launcher.py

@@ -1,8 +1,77 @@
-from cuda.core._launcher import LaunchConfig
+from cuda.core.experimental._launcher import LaunchConfig
+from cuda.core.experimental._stream import Stream
+from cuda.core.experimental._device import Device
+from cuda.core.experimental._utils import handle_return
+from cuda import cuda
 
-def test_launch_initialization():
+def test_launch_config_init():
     config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream=None, shmem_size=0)
-    
     assert config.grid == (1, 1, 1)
     assert config.block == (1, 1, 1)
+    assert config.stream is None
+    assert config.shmem_size == 0
+
+    config = LaunchConfig(grid=(2, 2, 2), block=(2, 2, 2), stream=Device().create_stream(), shmem_size=1024)
+    assert config.grid == (2, 2, 2)
+    assert config.block == (2, 2, 2)
+    assert isinstance(config.stream, Stream)
+    assert config.shmem_size == 1024
+
+def test_launch_config_cast_to_3_tuple():
+    config = LaunchConfig(grid=1, block=1)
+    assert config._cast_to_3_tuple(1) == (1, 1, 1)
+    assert config._cast_to_3_tuple((1, 2)) == (1, 2, 1)
+    assert config._cast_to_3_tuple((1, 2, 3)) == (1, 2, 3)
+
+    # Edge cases
+    assert config._cast_to_3_tuple(999) == (999, 1, 1)
+    assert config._cast_to_3_tuple((999, 888)) == (999, 888, 1)
+    assert config._cast_to_3_tuple((999, 888, 777)) == (999, 888, 777)
+
+def test_launch_config_invalid_values():
+    try:
+        LaunchConfig(grid=0, block=1)
+    except ValueError:
+        assert True
+    else:
+        assert False
+
+    try:
+        LaunchConfig(grid=(0, 1), block=1)
+    except ValueError:
+        assert True
+    else:
+        assert False
+
+    try:
+        LaunchConfig(grid=(1, 1, 1), block=0)
+    except ValueError:
+        assert True
+    else:
+        assert False
+
+    try:
+        LaunchConfig(grid=(1, 1, 1), block=(0, 1))
+    except ValueError:
+        assert True
+    else:
+        assert False
+
+def test_launch_config_stream():
+    stream = Device().create_stream()
+    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream=stream, shmem_size=0)
+    assert config.stream == stream
+
+    try:
+        LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream="invalid_stream", shmem_size=0)
+    except ValueError:
+        assert True
+    else:
+        assert False
+
+def test_launch_config_shmem_size():
+    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream=None, shmem_size=2048)
+    assert config.shmem_size == 2048
+
+    config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream=None)
     assert config.shmem_size == 0

cuda_core/tests/test_memory.py

@@ -1,14 +1,20 @@
-from cuda.core._memory import Buffer, MemoryResource
+# FILE: test_memory.py
+
+from cuda.core.experimental._memory import Buffer, MemoryResource
+from cuda.core.experimental._device import Device
+from cuda import cuda
+from cuda.core.experimental._utils import handle_return
 
 class DummyMemoryResource(MemoryResource):
     def __init__(self):
         pass
 
     def allocate(self, size, stream=None) -> Buffer:
-        pass
+        ptr = handle_return(cuda.cuMemAlloc(size))
+        return Buffer(ptr=ptr, size=size, mr=self)
 
     def deallocate(self, ptr, size, stream=None):
-        pass
+        handle_return(cuda.cuMemFree(ptr))
 
     @property
     def is_device_accessible(self) -> bool:
@@ -24,7 +30,42 @@ def device_id(self) -> int:
 
 def test_buffer_initialization():
     dummy_mr = DummyMemoryResource()
-    buffer = Buffer(ptr=1234, size=1024, mr=dummy_mr)
-    assert buffer._ptr == 1234
-    assert buffer._size == 1024
-    assert buffer._mr == dummy_mr
+    buffer = dummy_mr.allocate(size=1024)
+    assert buffer.handle != 0
+    assert buffer.size == 1024
+    assert buffer.memory_resource == dummy_mr
+    assert buffer.is_device_accessible == True
+    assert buffer.is_host_accessible == True
+    assert buffer.device_id == 0
+    dummy_mr.deallocate(buffer.handle, buffer.size)
+
+def test_buffer_copy_to():
+    dummy_mr = DummyMemoryResource()
+    src_buffer = dummy_mr.allocate(size=1024)
+    dst_buffer = dummy_mr.allocate(size=1024)
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+    src_buffer.copy_to(dst_buffer, stream=stream)
+    # Assuming cuMemcpyAsync is correctly called, we can't directly check the result here
+    dummy_mr.deallocate(src_buffer.handle, src_buffer.size)
+    dummy_mr.deallocate(dst_buffer.handle, dst_buffer.size)
+
+def test_buffer_copy_from():
+    dummy_mr = DummyMemoryResource()
+    src_buffer = dummy_mr.allocate(size=1024)
+    dst_buffer = dummy_mr.allocate(size=1024)
+    device = Device()
+    device.set_current()
+    stream = device.create_stream()
+    dst_buffer.copy_from(src_buffer, stream=stream)
+    # Assuming cuMemcpyAsync is correctly called, we can't directly check the result here
+    dummy_mr.deallocate(src_buffer.handle, src_buffer.size)
+    dummy_mr.deallocate(dst_buffer.handle, dst_buffer.size)
+
+def test_buffer_close():
+    dummy_mr = DummyMemoryResource()
+    buffer = dummy_mr.allocate(size=1024)
+    buffer.close()
+    assert buffer.handle == 0
+    assert buffer.memory_resource == None

cuda_core/tests/test_module.py

@@ -1,11 +1,35 @@
-from cuda.core._module import ObjectCode
-
-def test_module_initialization():
-    module_code = b"dummy_code"
-    code_type = "ptx"
-    module = ObjectCode(module=module_code, code_type=code_type)
-    assert module._handle is not None
-    assert module._code_type == code_type
-    assert module._module == module_code
-    assert module._loader is not None
-    assert module._sym_map == {}
+import pytest
+from cuda import cuda
+from cuda.core.experimental._device import Device
+from cuda.core.experimental._module import Kernel, ObjectCode
+from cuda.core.experimental._utils import handle_return
+
+@pytest.fixture(scope='module')
+def init_cuda():
+    Device().set_current()
+
+def test_object_code_initialization():
+    # Test with supported code types
+    for code_type in ["cubin", "ptx", "fatbin"]:
+        module_data = b"dummy_data"
+        obj_code = ObjectCode(module_data, code_type)
+        assert obj_code._code_type == code_type
+        assert obj_code._module == module_data
+        assert obj_code._handle is not None
+
+    # Test with unsupported code type
+    with pytest.raises(ValueError):
+        ObjectCode(b"dummy_data", "unsupported_code_type")
+
+#TODO add ObjectCode tests which provide the appropriate data for cuLibraryLoadFromFile
+def test_object_code_initialization_with_str():
+    assert True
+
+def test_object_code_initialization_with_jit_options():
+    assert True
+
+def test_object_code_get_kernel():
+    assert True
+
+def test_kernel_from_obj():
+    assert True

cuda_core/tests/test_program.py

@@ -1,7 +1,50 @@
-from cuda.core._program import Program
+import pytest
+from cuda import nvrtc
+from cuda.core.experimental._program import Program
+from cuda.core.experimental._module import ObjectCode, Kernel
 
-def test_program_initialization():
-    code = "__device__ int test_func() { return 0; }"
+def test_program_init_valid_code_type():
+    code = "extern \"C\" __global__ void my_kernel() {}"
     program = Program(code, "c++")
-    assert program._handle is not None
-    assert program._backend == "nvrtc"
+    assert program.backend == "nvrtc"
+    assert program.handle is not None
+
+def test_program_init_invalid_code_type():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    with pytest.raises(NotImplementedError):
+        Program(code, "python")
+
+def test_program_init_invalid_code_format():
+    code = 12345
+    with pytest.raises(TypeError):
+        Program(code, "c++")
+
+def test_program_compile_valid_target_type():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    program = Program(code, "c++")
+    object_code = program.compile("ptx")
+    kernel = object_code.get_kernel("my_kernel")
+    assert isinstance(object_code, ObjectCode)
+    assert isinstance(kernel, Kernel)
+
+def test_program_compile_invalid_target_type():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    program = Program(code, "c++")
+    with pytest.raises(NotImplementedError):
+        program.compile("invalid_target")
+
+def test_program_backend_property():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    program = Program(code, "c++")
+    assert program.backend == "nvrtc"
+
+def test_program_handle_property():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    program = Program(code, "c++")
+    assert program.handle is not None
+
+def test_program_close():
+    code = "extern \"C\" __global__ void my_kernel() {}"
+    program = Program(code, "c++")
+    program.close()
+    assert program.handle is None