Add support for passing cupy arrays to "C" lang

Author: bouweandela

kernel_tuner/backends/compiler.py

@@ -21,6 +21,22 @@
     SkippableFailure,
 )
 
+try:
+    import cupy as cp
+except ImportError:
+    cp = None
+
+
+def is_cupy_array(array):
+    """Check if something is a cupy array."""
+    return cp is not None and isinstance(array, cp.ndarray)
+
+
+def get_array_module(*args):
+    """Return the array module for arguments."""
+    return np if cp is None else cp.get_array_module(*args)
+
+
 dtype_map = {
     "int8": C.c_int8,
     "int16": C.c_int16,
@@ -112,9 +128,9 @@ def ready_argument_list(self, arguments):
         ctype_args = [None for _ in arguments]
 
         for i, arg in enumerate(arguments):
-            if not isinstance(arg, (np.ndarray, np.number)):
+            if not (isinstance(arg, (np.ndarray, np.number)) or is_cupy_array(arg)):
                 raise TypeError(
-                    "Argument is not numpy ndarray or numpy scalar %s" % type(arg)
+                    f"Argument is not numpy or cupy ndarray or numpy scalar but a {type(arg)}"
                 )
             dtype_str = str(arg.dtype)
             if isinstance(arg, np.ndarray):
@@ -129,6 +145,8 @@ def ready_argument_list(self, arguments):
                     raise TypeError("unknown dtype for ndarray")
             elif isinstance(arg, np.generic):
                 data_ctypes = dtype_map[dtype_str](arg)
+            elif is_cupy_array(arg):
+                data_ctypes = C.c_void_p(arg.data.ptr)
             ctype_args[i] = Argument(numpy=arg, ctypes=data_ctypes)
         return ctype_args
 
@@ -326,7 +344,10 @@ def memset(self, allocation, value, size):
         :param size: The size of to the allocation unit in bytes
         :type size: int
         """
-        C.memset(allocation.ctypes, value, size)
+        if is_cupy_array(allocation.numpy):
+            cp.cuda.runtime.memset(allocation.numpy.data.ptr, value, size)
+        else:
+            C.memset(allocation.ctypes, value, size)
 
     def memcpy_dtoh(self, dest, src):
         """a simple memcpy copying from an Argument to a numpy array
@@ -337,18 +358,30 @@ def memcpy_dtoh(self, dest, src):
         :param src: An Argument for some memory allocation
         :type src: Argument
         """
-        dest[:] = src.numpy
+        if isinstance(dest, np.ndarray) and isinstance(src.numpy, cp.ndarray):
+            # Implicit conversion to a NumPy array is not allowed.
+            value = src.numpy.get()
+        else:
+            value = src.numpy
+        xp = get_array_module(dest)
+        dest[:] = xp.asarray(value)
 
     def memcpy_htod(self, dest, src):
         """a simple memcpy copying from a numpy array to an Argument
 
         :param dest: An Argument for some memory allocation
-        :type dst: Argument
+        :type dest: Argument
 
         :param src: A numpy array containing the source data
         :type src: np.ndarray
         """
-        dest.numpy[:] = src
+        if isinstance(dest.numpy, np.ndarray) and isinstance(src, cp.ndarray):
+            # Implicit conversion to a NumPy array is not allowed.
+            value = src.get()
+        else:
+            value = src
+        xp = get_array_module(dest.numpy)
+        dest.numpy[:] = xp.asarray(value)
 
     def cleanup_lib(self):
         """unload the previously loaded shared library"""

test/test_compiler_functions.py

@@ -11,11 +11,12 @@
     from unittest.mock import patch, Mock
 
 import kernel_tuner
-from kernel_tuner.backends.compiler import CompilerFunctions, Argument
+from kernel_tuner.backends.compiler import CompilerFunctions, Argument, is_cupy_array, get_array_module
 from kernel_tuner.core import KernelSource, KernelInstance
 from kernel_tuner import util
 
-from .context import skip_if_no_gfortran, skip_if_no_gcc, skip_if_no_openmp
+from .context import skip_if_no_gfortran, skip_if_no_gcc, skip_if_no_openmp, skip_if_no_cupy
+from .test_runners import env as cuda_env  # noqa: F401
 
 
 @skip_if_no_gcc
@@ -108,6 +109,29 @@ def test_ready_argument_list5():
     assert all(output[0].numpy == arg1)
 
 
+@skip_if_no_cupy
+def test_ready_argument_list6():
+    import cupy as cp
+
+    arg = cp.array([1, 2, 3], dtype=np.float32)
+    arguments = [arg]
+
+    cfunc = CompilerFunctions()
+    output = cfunc.ready_argument_list(arguments)
+    print(output)
+
+    assert len(output) == 1
+    assert output[0].numpy is arg
+    mem = cp.cuda.UnownedMemory(
+        ptr=output[0].ctypes.value,
+        size=int(arg.nbytes / arg.dtype.itemsize),
+        owner=None,
+    )
+    ptr = cp.cuda.MemoryPointer(mem, 0)
+    output_arg = cp.ndarray(shape=arg.shape, dtype=arg.dtype, memptr=ptr)
+    assert cp.all(output_arg == arg)
+
+
 @skip_if_no_gcc
 def test_byte_array_arguments():
     arg1 = np.array([1, 2, 3]).astype(np.int8)
@@ -206,8 +230,29 @@ def test_memset():
     assert all(x == np.zeros(4))
 
 
-@skip_if_no_gcc
+@skip_if_no_cupy
 def test_memcpy_dtoh():
+    import cupy as cp
+
+    a = [1, 2, 3, 4]
+    x = cp.asarray(a, dtype=np.float32)
+    x_c = C.c_void_p(x.data.ptr)
+    arg = Argument(numpy=x, ctypes=x_c)
+    output = np.zeros(len(x), dtype=x.dtype)
+
+    cfunc = CompilerFunctions()
+    cfunc.memcpy_dtoh(output, arg)
+
+    print(f"{type(x)=} {x=}")
+    print(f"{type(a)=} {a=}")
+    print(f"{type(output)=} {output=}")
+
+    assert all(output == a)
+    assert all(x.get() == a)
+
+
+@skip_if_no_gcc
+def test_memcpy_host_dtoh():
     a = [1, 2, 3, 4]
     x = np.array(a).astype(np.float32)
     x_c = x.ctypes.data_as(C.POINTER(C.c_float))
@@ -224,8 +269,44 @@ def test_memcpy_dtoh():
     assert all(x == a)
 
 
-@skip_if_no_gcc
+@skip_if_no_cupy
+def test_memcpy_device_dtoh():
+    import cupy as cp
+
+    a = [1, 2, 3, 4]
+    x = cp.asarray(a, dtype=np.float32)
+    x_c = C.c_void_p(x.data.ptr)
+    arg = Argument(numpy=x, ctypes=x_c)
+    output = cp.zeros_like(x)
+
+    cfunc = CompilerFunctions()
+    cfunc.memcpy_dtoh(output, arg)
+
+    print(f"{type(x)=} {x=}")
+    print(f"{type(a)=} {a=}")
+    print(f"{type(output)=} {output=}")
+
+    assert all(output.get() == a)
+    assert all(x.get() == a)
+
+
+@skip_if_no_cupy
 def test_memcpy_htod():
+    import cupy as cp
+
+    a = [1, 2, 3, 4]
+    src = np.array(a, dtype=np.float32)
+    x = cp.zeros(len(src), dtype=src.dtype)
+    x_c = C.c_void_p(x.data.ptr)
+    arg = Argument(numpy=x, ctypes=x_c)
+
+    cfunc = CompilerFunctions()
+    cfunc.memcpy_htod(arg, src)
+
+    assert all(arg.numpy.get() == a)
+
+
+def test_memcpy_host_htod():
     a = [1, 2, 3, 4]
     src = np.array(a).astype(np.float32)
     x = np.zeros_like(src)
@@ -238,6 +319,22 @@ def test_memcpy_htod():
     assert all(arg.numpy == a)
 
 
+@skip_if_no_cupy
+def test_memcpy_device_htod():
+    import cupy as cp
+
+    a = [1, 2, 3, 4]
+    src = cp.array(a, dtype=np.float32)
+    x = cp.zeros(len(src), dtype=src.dtype)
+    x_c = C.c_void_p(x.data.ptr)
+    arg = Argument(numpy=x, ctypes=x_c)
+
+    cfunc = CompilerFunctions()
+    cfunc.memcpy_htod(arg, src)
+
+    assert all(arg.numpy.get() == a)
+
+
 @skip_if_no_gfortran
 def test_complies_fortran_function_no_module():
     kernel_string = """
@@ -335,3 +432,58 @@ def test_benchmark(env):
     assert all(["nthreads" in result for result in results])
     assert all(["time" in result for result in results])
     assert all([result["time"] > 0.0 for result in results])
+
+
+@skip_if_no_cupy
+def test_is_cupy_array():
+    import cupy as cp
+
+    assert is_cupy_array(cp.array([1.0]))
+    assert not is_cupy_array(np.array([1.0]))
+
+
+def test_is_cupy_array_no_cupy():
+    assert not is_cupy_array(np.array([1.0]))
+
+
+@skip_if_no_cupy
+def test_get_array_module():
+    import cupy as cp
+
+    assert get_array_module(cp.array([1.0])) == cp
+    assert get_array_module(np.array([1.0])) == np
+
+
+@skip_if_no_cupy
+@skip_if_no_gcc
+def test_run_kernel():
+    import cupy as cp
+
+    kernel_string = """
+    __global__ void vector_add_kernel(float *c, const float *a, const float *b, int n) {
+        int i = blockIdx.x * block_size_x + threadIdx.x;
+        if (i<n) {
+            c[i] = a[i] + b[i];
+        }
+    }
+
+    extern "C" void vector_add(float *c, const float *a, const float *b, int n) {
+        dim3 dimGrid(n);
+        dim3 dimBlock(block_size_x);
+        vector_add_kernel<<<dimGrid, dimBlock>>>(c, a, b, n);
+    }
+    """
+    a = cp.asarray([1, 2.0], dtype=np.float32)
+    b = cp.asarray([3, 4.0], dtype=np.float32)
+    c = cp.zeros_like(b)
+    n = np.int32(len(c))
+
+    result = kernel_tuner.run_kernel(
+        kernel_name="vector_add",
+        kernel_source=kernel_string,
+        problem_size=n,
+        arguments=[c, a, b, n],
+        params={"block_size_x": 1},
+        lang="C",
+    )
+    assert cp.all((a + b) == c)