Skip cupy test when cupy not installed

Signed-off-by: Ziheng Deng <zihengd@nvidia.com>

Author: gunnersdeng

test/conftest.py

@@ -218,3 +218,16 @@ def numba_cuda():
         pytest.xfail(f"Numba smoke test failed {result.returncode}. Skip.")
     import numba
     return numba.cuda
+
+
+def get_cupy_or_skip():
+    try:
+        import cupy as cupy
+    except ImportError:
+        pytest.skip("Cupy not installed. Skip test.")
+    return cupy
+
+
+@pytest.fixture(scope="session")
+def cupy():
+    return get_cupy_or_skip()

test/test_frontpage_example.py

@@ -6,7 +6,6 @@
 
 #example-begin
 import cuda.tile as ct
-import cupy
 
 TILE_SIZE = 16
 
@@ -18,18 +17,17 @@ def vector_add_kernel(a, b, result):
     b_tile = ct.load(b, index=(block_id,), shape=(TILE_SIZE,))
     result_tile = a_tile + b_tile
     ct.store(result, index=(block_id,), tile=result_tile)
-
-# Host-side function that launches the above kernel.
-def vector_add(a: cupy.ndarray, b: cupy.ndarray, result: cupy.ndarray):
-    assert a.shape == b.shape == result.shape
-    grid = (ct.cdiv(a.shape[0], TILE_SIZE), 1, 1)
-    ct.launch(cupy.cuda.get_current_stream(), grid, vector_add_kernel, (a, b, result))
 #example-end
 
-
 import numpy as np
 
-def test_vector_add():
+def test_vector_add(cupy):
+    # Host-side function that launches the above kernel.
+    def vector_add(a: cupy.ndarray, b: cupy.ndarray, result: cupy.ndarray):
+        assert a.shape == b.shape == result.shape
+        grid = (ct.cdiv(a.shape[0], TILE_SIZE), 1, 1)
+        ct.launch(cupy.cuda.get_current_stream(), grid, vector_add_kernel, (a, b, result))
+
     rng = cupy.random.default_rng()
     a = rng.random(128)
     b = rng.random(128)

test/test_full.py

@@ -5,14 +5,13 @@
 import numpy as np
 import pytest
 import torch
-import cupy as cp
 # Move cutile types to the top level?
 import cuda.tile as ct
 
 from pathlib import Path
 from math import ceil
 from util import assert_equal, jit_kernel
-from conftest import float_dtypes, int_dtypes, bool_dtypes, dtype_id
+from conftest import float_dtypes, int_dtypes, bool_dtypes, dtype_id, get_cupy_or_skip
 from cuda.tile._exception import TileTypeError
 from dataclasses import dataclass
 
@@ -65,7 +64,7 @@ def test_full_np_value_call(dtype, value, use_cupy, tmp_path: Path):
     dtype_str = torch_to_dtype_str[dtype].numpy
     if use_cupy:
         dtype_str = dtype_str.replace("np.", "cp.")
-        globals = {"cp": cp}
+        globals = {"cp": get_cupy_or_skip()}
     else:
         globals = {"np": np}
     value_str = str(value) if value is not None else ""
@@ -157,7 +156,7 @@ def test_full_np_dtype(value_dtype, use_cupy: bool, tmp_path: Path):
         dtype_str = dtype_str.replace("np.", "cp.")
         if value_str == "np.inf":
             value_str = "cp.inf"
-        globals = {"cp": cp}
+        globals = {"cp": get_cupy_or_skip()}
     else:
         globals = {"np": np}
     kernel = create_full_kernel("create_full_np_dtype", value_str, dtype_str,

test/test_readme_example.py

@@ -14,7 +14,7 @@
 file_path = os.path.realpath(__file__)
 
 
-def test_readme():
+def test_readme(cupy):
     readme_path = os.path.join(os.path.dirname(file_path), "..", "README.md")
     readme_txt = open(readme_path, 'r').read()
     header = "Example\n-------\n```python"

test/test_stream.py

@@ -4,10 +4,8 @@
 
 import torch
 from torch.testing import make_tensor
-import cupy
 import cuda.tile as ct
 from ctypes import c_void_p
-from cuda.bindings.driver import CUstream
 
 
 @ct.kernel
@@ -38,12 +36,12 @@ def test_torch_pass_stream_ptr():
 
 
 # -- Test CuPy Stream --
-def test_cupy_pass_stream():
+def test_cupy_pass_stream(cupy):
     stream = cupy.cuda.Stream()
     _test_stream(stream, stream.synchronize)
 
 
-def test_cupy_pass_stream_ptr():
+def test_cupy_pass_stream_ptr(cupy):
     stream = cupy.cuda.Stream()
     _test_stream(stream.ptr, stream.synchronize)
 
@@ -59,6 +57,7 @@ def test_numba_pass_stream_ptr(numba_cuda):
     handle = stream.handle
     # numba-cuda < 0.30: handle is ctypes c_void_p
     # numba-cuda >= 0.30: handle is cuda.bindings.driver.CUstream
+    from cuda.bindings.driver import CUstream
     if isinstance(handle, c_void_p):
         stream_ptr = handle.value
     elif isinstance(handle, CUstream):