Clean up testsuite

cuda_core/tests/conftest.py

@@ -5,12 +5,29 @@
 # this software. Any use, reproduction, disclosure, or distribution of
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
+try:
+    from cuda.bindings import driver
+except ImportError:
+    from cuda import cuda as driver
 
-from cuda.core.experimental._device import Device
+from cuda.core.experimental import Device
+from cuda.core.experimental import _device
+from cuda.core.experimental._utils import handle_return
 import pytest
 
-@pytest.fixture(scope="module")
+@pytest.fixture(scope="function")
 def init_cuda():
     device = Device()
     device.set_current()
-
+    yield
+    _device_unset_current()
+
+def _device_unset_current():
+    handle_return(driver.cuCtxPopCurrent())
+    with _device._tls_lock:
+        del _device._tls.devices
+
+@pytest.fixture(scope="function")
+def deinit_cuda():
+    yield
+    _device_unset_current()

cuda_core/tests/example_tests/test_basic_examples.py

@@ -20,6 +20,6 @@
     'example', sample_files
 )
 class TestExamples:
-    def test_example(self, example):
+    def test_example(self, example, deinit_cuda):
         filename = os.path.basename(example)
         run_example(samples_path, example)

cuda_core/tests/example_tests/utils.py

@@ -6,7 +6,6 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda import cuda
 import gc
 import os
 import sys

cuda_core/tests/test_device.py

@@ -6,12 +6,20 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda import cuda, cudart
-from cuda.core.experimental._device import Device
-from cuda.core.experimental._utils import handle_return, ComputeCapability, CUDAError, \
-                             precondition
-import pytest
+try:
+    from cuda.bindings import driver, runtime
+except ImportError:
+    from cuda import cuda as driver
+    from cuda import cudart as runtime
 
+from cuda.core.experimental import Device
+from cuda.core.experimental._utils import handle_return, ComputeCapability
+
+def test_device_set_current(deinit_cuda):
+    device = Device()
+    device.set_current()
+    assert handle_return(driver.cuCtxGetCurrent()) is not None
+
 def test_device_repr():
     device = Device(0)
     assert str(device).startswith('<Device 0')
@@ -24,44 +32,40 @@ def test_device_alloc(init_cuda):
     assert buffer.size == 1024
     assert buffer.device_id == 0
 
-def test_device_set_current():
-    device = Device()
-    device.set_current()
-
-def test_device_create_stream():
+def test_device_create_stream(init_cuda):
     device = Device()
     stream = device.create_stream()
     assert stream is not None
     assert stream.handle
 
 def test_pci_bus_id():
     device = Device()
-    bus_id = handle_return(cudart.cudaDeviceGetPCIBusId(13, device.device_id))
+    bus_id = handle_return(runtime.cudaDeviceGetPCIBusId(13, device.device_id))
     assert device.pci_bus_id == bus_id[:12].decode()
 
 def test_uuid():
     device = Device()
-    driver_ver = handle_return(cuda.cuDriverGetVersion())
+    driver_ver = handle_return(driver.cuDriverGetVersion())
     if driver_ver >= 11040:
-        uuid = handle_return(cuda.cuDeviceGetUuid_v2(device.device_id))
+        uuid = handle_return(driver.cuDeviceGetUuid_v2(device.device_id))
     else:
-        uuid = handle_return(cuda.cuDeviceGetUuid(device.device_id))
+        uuid = handle_return(driver.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()
-    name = handle_return(cuda.cuDeviceGetName(128, device.device_id))
+    name = handle_return(driver.cuDeviceGetName(128, device.device_id))
     name = name.split(b'\0')[0]
     assert device.name == name.decode()
 
 def test_compute_capability():
     device = Device()
-    major = handle_return(cudart.cudaDeviceGetAttribute(
-        cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, device.device_id))
-    minor = handle_return(cudart.cudaDeviceGetAttribute(
-        cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, device.device_id))
+    major = handle_return(runtime.cudaDeviceGetAttribute(
+        runtime.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, device.device_id))
+    minor = handle_return(runtime.cudaDeviceGetAttribute(
+        runtime.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, device.device_id))
     expected_cc = ComputeCapability(major, minor)
     assert device.compute_capability == expected_cc
 

cuda_core/tests/test_event.py

@@ -6,34 +6,38 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda import cuda
-from cuda.core.experimental._event import  EventOptions, Event
-from cuda.core.experimental._utils import handle_return
-from cuda.core.experimental._device import Device
+from cuda.core.experimental import Device, EventOptions
 import pytest
 
-def test_is_timing_disabled():
-    options = EventOptions(enable_timing=False)
-    event = Event._init(options)
-    assert event.is_timing_disabled == True
+@pytest.mark.parametrize("enable_timing", [True, False, None])
+def test_timing(init_cuda, enable_timing):
+    options = EventOptions(enable_timing=enable_timing)
+    stream = Device().create_stream()
+    event = stream.record(options=options)
+    assert event.is_timing_disabled == (not enable_timing if enable_timing is not None else True)
+
 
-def test_is_sync_busy_waited():
-    options = EventOptions(busy_waited_sync=True)
-    event = Event._init(options)
+def test_is_sync_busy_waited(init_cuda):
+    options = EventOptions(enable_timing=False, busy_waited_sync=True)
+    stream = Device().create_stream()
+    event = stream.record(options=options)
     assert event.is_sync_busy_waited == True
 
-def test_sync():
-    options = EventOptions()
-    event = Event._init(options)
+    options = EventOptions(enable_timing=False)
+    stream = Device().create_stream()
+    event = stream.record(options=options)
+    assert event.is_sync_busy_waited == False
+
+def test_sync(init_cuda):
+    options = EventOptions(enable_timing=False)
+    stream = Device().create_stream()
+    event = stream.record(options=options)
     event.sync()
     assert event.is_done == True
 
-def test_is_done():
-    options = EventOptions()
-    event = Event._init(options)
+def test_is_done(init_cuda):
+    options = EventOptions(enable_timing=False)
+    stream = Device().create_stream()
+    event = stream.record(options=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

@@ -6,14 +6,10 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda import cuda
-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.core.experimental import Device, Stream, LaunchConfig
 import pytest
 
-def test_launch_config_init():
+def test_launch_config_init(init_cuda):
     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)
@@ -50,7 +46,7 @@ def test_launch_config_invalid_values():
     with pytest.raises(ValueError):
         LaunchConfig(grid=(1, 1, 1), block=(0, 1))
 
-def test_launch_config_stream():
+def test_launch_config_stream(init_cuda):
     stream = Device().create_stream()
     config = LaunchConfig(grid=(1, 1, 1), block=(1, 1, 1), stream=stream, shmem_size=0)
     assert config.stream == stream

cuda_core/tests/test_memory.py

@@ -6,12 +6,15 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda import cuda
+try:
+    from cuda.bindings import driver
+except ImportError:
+    from cuda import cuda as driver
+
+from cuda.core.experimental import Device
 from cuda.core.experimental._memory import Buffer, MemoryResource
-from cuda.core.experimental._device import Device
 from cuda.core.experimental._utils import handle_return
 import ctypes
-import pytest
 
 class DummyDeviceMemoryResource(MemoryResource):
     def __init__(self, device):
@@ -66,11 +69,11 @@ def __init__(self, device):
         self.device = device
 
     def allocate(self, size, stream=None) -> Buffer:
-        ptr = handle_return(cuda.cuMemAllocManaged(size, cuda.CUmemAttach_flags.CU_MEM_ATTACH_GLOBAL.value))
+        ptr = handle_return(driver.cuMemAllocManaged(size, driver.CUmemAttach_flags.CU_MEM_ATTACH_GLOBAL.value))
         return Buffer(ptr=ptr, size=size, mr=self)
 
     def deallocate(self, ptr, size, stream=None):
-        handle_return(cuda.cuMemFree(ptr))
+        handle_return(driver.cuMemFree(ptr))
 
     @property
     def is_device_accessible(self) -> bool:
@@ -89,11 +92,11 @@ def __init__(self, device):
         self.device = device
 
     def allocate(self, size, stream=None) -> Buffer:
-        ptr = handle_return(cuda.cuMemAllocHost(size))
+        ptr = handle_return(driver.cuMemAllocHost(size))
         return Buffer(ptr=ptr, size=size, mr=self)
 
     def deallocate(self, ptr, size, stream=None):
-        handle_return(cuda.cuMemFreeHost(ptr))
+        handle_return(driver.cuMemFreeHost(ptr))
 
     @property
     def is_device_accessible(self) -> bool:

cuda_core/tests/test_module.py

@@ -6,10 +6,7 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-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
+from cuda.core.experimental._module import ObjectCode
 import pytest
 import importlib
 

cuda_core/tests/test_program.py

@@ -6,9 +6,8 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda.core.experimental._program import Program
+from cuda.core.experimental import Program
 from cuda.core.experimental._module import ObjectCode, Kernel
-from cuda.core.experimental._device import Device
 import pytest
 
 def test_program_init_valid_code_type():

cuda_core/tests/test_stream.py

@@ -6,68 +6,69 @@
 # this software and related documentation outside the terms of the EULA
 # is strictly prohibited.
 
-from cuda.core.experimental._stream import Stream, StreamOptions, LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, default_stream
-from cuda.core.experimental._event import Event, EventOptions
-from cuda.core.experimental._device import Device
+from cuda.core.experimental import Device, Stream, StreamOptions
+from cuda.core.experimental._stream import LEGACY_DEFAULT_STREAM, PER_THREAD_DEFAULT_STREAM, default_stream
+from cuda.core.experimental._event import Event
 import pytest
 
 def test_stream_init():
     with pytest.raises(NotImplementedError):
         Stream()
 
-def test_stream_init_with_options():
-    stream = Stream._init(options=StreamOptions(nonblocking=True, priority=0))
+def test_stream_init_with_options(init_cuda):
+    stream = Device().create_stream(options=StreamOptions(nonblocking=True, priority=0))
     assert stream.is_nonblocking is True
     assert stream.priority == 0
 
-def test_stream_handle():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_handle(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     assert isinstance(stream.handle, int)
 
-def test_stream_is_nonblocking():
-    stream = Stream._init(options=StreamOptions(nonblocking=True))
+def test_stream_is_nonblocking(init_cuda):
+    stream = Device().create_stream(options=StreamOptions(nonblocking=True))
     assert stream.is_nonblocking is True
 
-def test_stream_priority():
-    stream = Stream._init(options=StreamOptions(priority=0))
+def test_stream_priority(init_cuda):
+    stream = Device().create_stream(options=StreamOptions(priority=0))
     assert stream.priority == 0
-    stream = Stream._init(options=StreamOptions(priority=-1))
+    stream = Device().create_stream(options=StreamOptions(priority=-1))
     assert stream.priority == -1
     with pytest.raises(ValueError):
-        stream = Stream._init(options=StreamOptions(priority=1))
+        stream = Device().create_stream(options=StreamOptions(priority=1))
 
-def test_stream_sync():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_sync(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     stream.sync()  # Should not raise any exceptions
 
-def test_stream_record():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_record(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     event = stream.record()
     assert isinstance(event, Event)
 
-def test_stream_record_invalid_event():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_record_invalid_event(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     with pytest.raises(TypeError):
         stream.record(event="invalid_event")
 
-def test_stream_wait_event():
-    stream = Stream._init(options=StreamOptions())
-    event = Event._init()
-    stream.record(event)
-    stream.wait(event)  # Should not raise any exceptions
+def test_stream_wait_event(init_cuda):
+    s1 = Device().create_stream()
+    s2 = Device().create_stream()
+    e1 = s1.record()
+    s2.wait(e1)  # Should not raise any exceptions
+    s2.sync()
 
-def test_stream_wait_invalid_event():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_wait_invalid_event(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     with pytest.raises(ValueError):
         stream.wait(event_or_stream="invalid_event")
 
-def test_stream_device():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_device(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     device = stream.device
     assert isinstance(device, Device)
 
-def test_stream_context():
-    stream = Stream._init(options=StreamOptions())
+def test_stream_context(init_cuda):
+    stream = Device().create_stream(options=StreamOptions())
     context = stream.context
     assert context is not None