Fix example requirements

Author: mdboom

cuda_core/tests/example_tests/test_basic_examples.py

@@ -3,14 +3,15 @@
 
 # If we have subcategories of examples in the future, this file can be split along those lines
 
+import functools
 import glob
 import os
 import subprocess
 import sys
 
 import pytest
 
-from cuda.core import Device
+from cuda.core import Device, system
 
 
 def pytorch_installed():
@@ -22,16 +23,48 @@ def pytorch_installed():
         return False
 
 
+@functools.cache
+def cupy_installed():
+    try:
+        import cupy  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
+def cffi_installed():
+    try:
+        import cffi  # noqa: F401
+
+        return True
+    except ImportError:
+        return False
+
+
 def has_compute_capability_9_or_higher():
     dev = Device()
     arch = dev.compute_capability
     return arch > (9, 0)
 
 
+def has_multiple_devices():
+    return system.get_num_devices() >= 2
+
+
 REQUIREMENTS = {
-    "gl_interop_plasma.py": (lambda: False),  # requires a display
-    "pytorch_example.py": pytorch_installed,  # requires PyTorch
-    "thread_block_cluster.py": has_compute_capability_9_or_higher,  # requires CC 9.0+
+    "gl_interop_plasma.py": (lambda: False),
+    "pytorch_example.py": pytorch_installed,
+    "thread_block_cluster.py": has_compute_capability_9_or_higher,
+    "cuda_graphs.py": cupy_installed,
+    "jit_lto_fractal.py": cupy_installed,
+    "memory_ops.py": cupy_installed,
+    "saxpy.py": cupy_installed,
+    "simple_multi_gpu_example.py": lambda: has_multiple_devices() and cupy_installed(),
+    "strided_memory_view_cpu.py": cffi_installed,
+    "strided_memory_view_gpu.py": cupy_installed,
+    "tma_tensor_map.py": cupy_installed,
+    "vector_add.py": cupy_installed,
 }
 
 samples_path = os.path.join(os.path.dirname(__file__), "..", "..", "examples")