Repository: cupy/cupy

Issue:
Ability to pass structured types by value as kernel parameters
Currently, it is not possible to pass structures **by value** as kernel arguments. Only some fixed size numpy dtype arguments have been hardcoded yet in `cupy/cuda/function.pyx`. This implies that a cupy kernel launch only succeeds with numpy.generic arguments which are scalar types with itemsize 1, 2, 4, 8 or 16 bytes. This allows us to pass directly by value np.int8 up to np.complex128 numpy data types as kernel arguments. I would like to extend pass-by-value to any numpy dtype.

I am using RawKernels extensively and I would like to be able to pass cuda builtin vector types such as `float3` as kernel parameters. This would make multidimensional code in physics much clearer for example (ie. replace many scalar kernel arguments by vectorized arguments or static arrays of scalars). Depending on device compute capability we can now pass up to 4kB of parameters to cuda kernels, and doing it in 16 bytes chunks is not very efficient.

The definition of the structure layout using numpy.dtype should be the responsibility of the user. This off course implies defining the struct with the same offsets, alignment and itemsize between host and device. `pyopencl` and `pycuda` offer utility functions to convert a device side struct to a numpy dtype with matching layout, but this feature would be overkill for now.

Please note that **this feature is not related to general structured array support in cupy**. It is however already possible to pass a pointer to any array of structures to a RawKernel: simply cast any numpy array to np.int8 data type and let cupy handle np.int8 cuda arrays. Within the kernel, the pointer will be cast to the kernel argument type anyway. This off course only works if host and device structures layout match exactly.

Here is a simple example to illustrate what I would be able to do:
```python
import numpy as np
import cupy as cp

code = r"""
    template <typename T>
    __global__ void affect(T* const __restrict__ output,
                         const T arg0, const T arg1,
                         const size_t nelements) {
        ptrdiff_t i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= nelements)
            return;
        output[i] = (i%2 == 0 ? arg0 : arg1);
    }
"""
module = cp.RawModule(code=code,
                      name_expressions=('affect<float3>',),
                      options=('-std=c++11',))
kernel = module.get_function('affect<float3>')

float3 = np.dtype( { 'names': ['x', 'y', 'z'],
                     'formats': [np.float32]*3 } )

arg0 = np.random.rand(3).astype(np.float32).view(float3)
arg1 = np.random.rand(3).astype(np.float32).view(float3)

N = 512
h_output = np.empty(dtype=float3, shape=N)
d_output = cp.asarray(h_output.view(np.float32))

grid = (N,)
block = (256,)
args = (d_output, arg0, arg1, np.uint64(N))
kernel(grid, block, args)

h_output[0::2] = arg0
h_output[1::2] = arg1

np.testing.assert_array_equal(h_output,
                              cp.asnumpy(d_output).view(float3))
```





Discussion:

Resolving patch:
diff --git a/examples/custom_struct/builtin_vectors.py b/examples/custom_struct/builtin_vectors.py
new file mode 100644
--- /dev/null
+++ b/examples/custom_struct/builtin_vectors.py
@@ -0,0 +1,48 @@
+import sys
+import numpy
+import cupy
+
+code = '''
+__device__ double3 operator+(const double3& lhs, const double3& rhs) {
+    return make_double3(lhs.x + rhs.x,
+                        lhs.y + rhs.y,
+                        lhs.z + rhs.z);
+}
+
+extern "C" __global__ void sum_kernel(const double3* lhs,
+                                            double3  rhs,
+                                            double3* out) {
+  int i = threadIdx.x;
+  out[i] = lhs[i] + rhs;
+}
+'''
+
+double3 = numpy.dtype(
+    {
+        'names': ['x', 'y', 'z'],
+        'formats': [numpy.float64]*3
+    }
+)
+
+
+def main():
+    N = 8
+
+    # The kernel computes out = lhs+rhs where lhs and rhs are double3 vectors.
+    # lhs is an array of N such vectors and rhs is double3 kernel parameter.
+
+    lhs = cupy.random.rand(3*N, dtype=numpy.float64).reshape(N, 3)
+    rhs = numpy.random.rand(3).astype(numpy.float64)
+    out = cupy.empty_like(lhs)
+
+    kernel = cupy.RawKernel(code, 'sum_kernel')
+    args = (lhs, rhs.view(double3), out)
+    kernel((1,), (N,), args)
+
+    expected = lhs + cupy.asarray(rhs[None, :])
+    cupy.testing.assert_array_equal(expected, out)
+    print("Kernel output matches expected value.")
+
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/examples/custom_struct/complex_struct.py b/examples/custom_struct/complex_struct.py
new file mode 100644
--- /dev/null
+++ b/examples/custom_struct/complex_struct.py
@@ -0,0 +1,132 @@
+import sys
+import numpy
+import cupy
+
+struct_definition = '''
+struct complex_struct {
+    int4 a;
+    char b;
+    double c[2];
+    short1 d;
+    unsigned long long int e[3];
+};
+'''
+
+struct_layout_code = '''
+{struct_definition}
+
+extern "C" __global__ void get_struct_layout(
+                                unsigned long long *itemsize,
+                                unsigned long long *sizes,
+                                unsigned long long *offsets) {{
+    const complex_struct* ptr = NULL;
+
+    itemsize[0] = sizeof(complex_struct);
+
+    sizes[0] = sizeof(ptr->a);
+    sizes[1] = sizeof(ptr->b);
+    sizes[2] = sizeof(ptr->c);
+    sizes[3] = sizeof(ptr->d);
+    sizes[4] = sizeof(ptr->e);
+
+    offsets[0] = (unsigned long long)&ptr->a;
+    offsets[1] = (unsigned long long)&ptr->b;
+    offsets[2] = (unsigned long long)&ptr->c;
+    offsets[3] = (unsigned long long)&ptr->d;
+    offsets[4] = (unsigned long long)&ptr->e;
+}}
+'''.format(struct_definition=struct_definition)
+
+
+kernel_code = '''
+{struct_definition}
+
+extern "C" __global__ void test_kernel(const complex_struct s,
+                                       double* out) {{
+    int i = threadIdx.x;
+    double sum = 0.0;
+    sum += s.a.x + s.a.y + s.a.z + s.a.w;
+    sum += s.b;
+    sum += s.c[0] + s.c[1];
+    sum += s.d.x;
+    sum += s.e[0] + s.e[1] + s.e[2];
+    out[i] = i * sum;
+}}
+'''.format(struct_definition=struct_definition)
+
+
+def make_packed(basetype, N, itemsize):
+    # A small utility function to make packed structs
+    # Can represent simple packed vectors such as float4 or double[3].
+    assert 0 < N <= 4, N
+    names = list('xyzw')[:N]
+    formats = [basetype]*N
+    return numpy.dtype(dict(names=names,
+                            formats=formats,
+                            itemsize=itemsize))
+
+
+def main():
+    # This program demonstrate how to build a hostside
+    # representation of device structure 'complex_struct'
+    # defined in variable 'struct_definition' that can be
+    # used as a RawKernel argument.
+
+    # First step is to determine structure memory layout
+    #  itemsize -> overall struct size
+    #  sizes    -> individual struct member sizes, determined with sizeof
+    #  offsets  -> individual struct member offsets, determined with offsetof
+    # Results (in terms of bytes) are copied to host after kernel launch.
+    # Note that 'complex_struct' has 5 members named a, b, c, d and e.
+    itemsize = cupy.ndarray(shape=(1,), dtype=numpy.uint64)
+    sizes = cupy.ndarray(shape=(5,), dtype=numpy.uint64)
+    offsets = cupy.ndarray(shape=(5,), dtype=numpy.uint64)
+
+    kernel = cupy.RawKernel(struct_layout_code, 'get_struct_layout')
+    kernel((1,), (1,), (itemsize, sizes, offsets))
+
+    (itemsize, sizes, offsets) = map(cupy.asnumpy, (itemsize, sizes, offsets))
+    print("Overall structure itemsize: {} bytes".format(itemsize.item()))
+    print("Structure members itemsize: {}".format(sizes))
+    print("Structure members offsets: {}".format(offsets))
+
+    # Second step: build a numpy dtype for each struct member
+    atype = make_packed(numpy.int32,   4, sizes[0])
+    btype = make_packed(numpy.int8,    1, sizes[1])
+    ctype = make_packed(numpy.float64, 2, sizes[2])
+    dtype = make_packed(numpy.int16,   1, sizes[3])
+    etype = make_packed(numpy.uint64,  3, sizes[4])
+
+    # Third step: create the complex struct representation with
+    #  the right offsets
+    names = list('abcde')
+    formats = [atype, btype, ctype, dtype, etype]
+    complex_struct = numpy.dtype(dict(names=names,
+                                      formats=formats,
+                                      offsets=offsets,
+                                      itemsize=itemsize.item()))
+
+    # Build a complex_struct kernel argument
+    s = numpy.empty(shape=(1,), dtype=complex_struct)
+    s['a'] = numpy.arange(0, 4).astype(numpy.int32).view(atype)
+    s['b'] = numpy.arange(4, 5).astype(numpy.int8).view(btype)
+    s['c'] = numpy.arange(5, 7).astype(numpy.float64).view(ctype)
+    s['d'] = numpy.arange(7, 8).astype(numpy.int16).view(dtype)
+    s['e'] = numpy.arange(8, 11).astype(numpy.uint64).view(etype)
+    print("Complex structure value:\n  {}".format(s))
+
+    # Setup test kernel
+    N = 8
+    out = cupy.empty(shape=(N,), dtype=numpy.float64)
+    kernel = cupy.RawKernel(kernel_code, 'test_kernel')
+    kernel((1,), (N,), (s, out))
+
+    # the sum of all members of our complex struct instance is 55.0
+    expected = cupy.arange(N) * 55.0
+
+    cupy.testing.assert_array_almost_equal(expected, out)
+    print("Kernel output matches expected value.")
+
+
+if __name__ == '__main__':
+    sys.exit(main())
diff --git a/examples/custom_struct/packed_matrix.py b/examples/custom_struct/packed_matrix.py
new file mode 100644
--- /dev/null
+++ b/examples/custom_struct/packed_matrix.py
@@ -0,0 +1,91 @@
+import sys
+import numpy
+import cupy
+
+code = '''
+template<typename T>
+struct Matrix {
+    T value[4][4];
+
+    __device__ T& operator() (int i, int j) {
+        return this->value[i][j];
+    }
+
+    __device__ const T& operator() (int i, int j) const {
+        return this->value[i][j];
+    }
+};
+
+template<typename T>
+__device__ Matrix<T> operator+ (const Matrix<T>& lhs, const Matrix<T>& rhs) {
+    Matrix<T> res;
+    for (int i = 0; i<4; i++) {
+        for (int j = 0; j<4; j++) {
+            res(i,j) = lhs(i,j) + rhs(i,j);
+        }
+    }
+    return res;
+}
+
+template<typename T>
+__device__ Matrix<T> operator* (const Matrix<T>& lhs, const Matrix<T>& rhs) {
+    Matrix<T> res;
+    for (int i = 0; i<4; i++) {
+        for (int j = 0; j<4; j++) {
+            res(i,j) = T(0);
+            for (int k = 0; k<4; k++) {
+                res(i,j) += lhs(i,k) * rhs(k,j);
+            }
+        }
+    }
+    return res;
+}
+
+template<typename T>
+__global__ void kernel(const Matrix<T>* A,
+                       const Matrix<T>* B,
+                       const Matrix<T>  C,
+                             Matrix<T>* out) {
+  int i = threadIdx.x;
+  out[i] = A[i] * B[i] + C;
+}
+'''
+
+
+def main():
+    N = 8
+    module = cupy.RawModule(code=code, options=('-std=c++11',),
+                            name_expressions=('kernel<float>',
+                                              'kernel<double>'))
+
+    # The kernel computes out = A*B+C where A, B and C are 4x4 matrices.
+    # A and B are arrays of N such matrices and C is a matrix kernel parameter.
+
+    for (ctype, dtype) in zip(('float', 'double'),
+                              (numpy.float32, numpy.float64)):
+
+        A = cupy.random.rand(16*N, dtype=dtype).reshape(N, 4, 4)
+        B = cupy.random.rand(16*N, dtype=dtype).reshape(N, 4, 4)
+        C = numpy.random.rand(16).astype(dtype).reshape(4, 4)
+        out = cupy.empty_like(A)
+
+        Matrix = numpy.dtype(
+            {
+                'names': ['value'],
+                'formats': [(dtype, (4, 4))]
+            }
+        )
+
+        kernel = module.get_function('kernel<{}>'.format(ctype))
+        args = (A, B, C.ravel().view(Matrix), out)
+        kernel((1,), (N,), args)
+
+        expected = cupy.matmul(A, B) + cupy.asarray(C[None, :, :])
+
+        cupy.testing.assert_array_almost_equal(expected, out)
+        print("Kernel output matches expected value for "
+              "type '{}'.".format(ctype))
+
+
+if __name__ == '__main__':
+    sys.exit(main())


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_function.py b/tests/cupy_tests/core_tests/test_function.py
--- a/tests/cupy_tests/core_tests/test_function.py
+++ b/tests/cupy_tests/core_tests/test_function.py
@@ -1,14 +1,19 @@
 import unittest
+import pytest
 
 import numpy
 
 import cupy
+from cupy._core import core
 from cupy.cuda import compiler
+from cupy.cuda import runtime
 from cupy import testing
 
 
 def _compile_func(kernel_name, code):
-    mod = compiler.compile_with_cache(code)
+    # workaround for hipRTC
+    extra_source = core._get_header_source() if runtime.is_hip else None
+    mod = compiler.compile_with_cache(code, extra_source=extra_source)
     return mod.get_function(kernel_name)
 
 
@@ -54,3 +59,160 @@ def test_numpy_scalar(self):
 
         expected = a_cpu + b
         testing.assert_array_equal(x, expected)
+
+    def test_numpy_dtype(self):
+        code = '''
+extern "C" __global__ void test_kernel(const double* a, double3 b, double* x) {
+  int i = threadIdx.x;
+  x[i] = a[i] + b.x + b.y + b.z;
+}
+'''
+
+        double3 = numpy.dtype({'names': ['x', 'y', 'z'],
+                               'formats': [numpy.float64]*3})
+        a_cpu = numpy.arange(24, dtype=numpy.float64)
+        a = cupy.array(a_cpu)
+        b = numpy.random.rand(3).view(double3)
+        x = cupy.empty_like(a)
+
+        func = _compile_func('test_kernel', code)
+
+        func.linear_launch(a.size, (a, b, x))
+
+        expected = a_cpu + b['x'] + b['y'] + b['z']
+        testing.assert_array_equal(x, expected)
+
+    def test_static_array(self):
+        code = '''
+struct double5 {
+    double value[5];
+    __device__ const double& operator[](size_t i) const { return value[i]; }
+};
+
+extern "C" __global__ void test_kernel(const double* a, double5 b, double* x) {
+  int i = threadIdx.x;
+  x[i] = a[i] + b[0] + b[1] + b[2] + b[3] + b[4];
+}
+'''
+
+        a_cpu = numpy.arange(24, dtype=numpy.float64)
+        a = cupy.array(a_cpu)
+        x = cupy.empty_like(a)
+
+        func = _compile_func('test_kernel', code)
+
+        # We cannot pass np.ndarray kernel arguments of size > 1
+        b = numpy.arange(5).astype(numpy.float64)
+        with pytest.raises(TypeError):
+            func.linear_launch(a.size, (a, b, x))
+
+        double5 = numpy.dtype({
+            'names': ['dummy'],
+            'formats': [(numpy.float64, (5,))]
+        })
+        func.linear_launch(a.size, (a, b.view(double5), x))
+
+        expected = a_cpu + b.sum()
+        testing.assert_array_equal(x, expected)
+
+    def test_custom_user_struct(self):
+        struct_definition = '''
+struct custom_user_struct {
+    int4 a;
+    char b;
+    double c[2];
+    short1 d;
+    unsigned long long int e[3];
+};
+'''
+
+        # first step is to determine struct memory layout
+        struct_layout_code = '''
+{struct_definition}
+extern "C" __global__ void get_struct_layout(
+                                unsigned long long *itemsize,
+                                unsigned long long *sizes,
+                                unsigned long long *offsets) {{
+    const custom_user_struct* ptr = NULL;
+
+    itemsize[0] = sizeof(custom_user_struct);
+
+    sizes[0] = sizeof(ptr->a);
+    sizes[1] = sizeof(ptr->b);
+    sizes[2] = sizeof(ptr->c);
+    sizes[3] = sizeof(ptr->d);
+    sizes[4] = sizeof(ptr->e);
+
+    offsets[0] = (unsigned long long)&ptr->a;
+    offsets[1] = (unsigned long long)&ptr->b;
+    offsets[2] = (unsigned long long)&ptr->c;
+    offsets[3] = (unsigned long long)&ptr->d;
+    offsets[4] = (unsigned long long)&ptr->e;
+}}
+'''.format(struct_definition=struct_definition)
+
+        itemsize = cupy.ndarray(shape=(1,), dtype=numpy.uint64)
+        sizes = cupy.ndarray(shape=(5,), dtype=numpy.uint64)
+        offsets = cupy.ndarray(shape=(5,), dtype=numpy.uint64)
+        func = _compile_func('get_struct_layout', struct_layout_code)
+        func.linear_launch(1, (itemsize, sizes, offsets))
+
+        # Build structure data type recursively
+        names = list('abcde')
+        itemsize = cupy.asnumpy(itemsize).item()
+        sizes = cupy.asnumpy(sizes).tolist()
+        offsets = cupy.asnumpy(offsets).tolist()
+
+        def make_packed(basetype, N, itemsize):
+            assert 0 < N <= 4, N
+            names = list('xyzw')[:N]
+            formats = [basetype]*N
+            return numpy.dtype(dict(names=names,
+                                    formats=formats,
+                                    itemsize=itemsize))
+
+        # structure member data types
+        int4 = make_packed(numpy.int32, 4, sizes[0])
+        char = make_packed(numpy.int8, 1, sizes[1])
+        double2 = make_packed(numpy.float64, 2, sizes[2])
+        short1 = make_packed(numpy.int16, 1, sizes[3])
+        ulong3 = make_packed(numpy.uint64, 3, sizes[4])
+
+        formats = [int4, char, double2, short1, ulong3]
+        struct_dtype = numpy.dtype(dict(names=names,
+                                        formats=formats,
+                                        offsets=offsets,
+                                        itemsize=itemsize))
+
+        s = numpy.empty(shape=(1,), dtype=struct_dtype)
+        s['a'] = numpy.arange(0, 4).astype(numpy.int32).view(int4)
+        s['b'] = numpy.arange(4, 5).astype(numpy.int8).view(char)
+        s['c'] = numpy.arange(5, 7).astype(numpy.float64).view(double2)
+        s['d'] = numpy.arange(7, 8).astype(numpy.int16).view(short1)
+        s['e'] = numpy.arange(8, 11).astype(numpy.uint64).view(ulong3)
+
+        # test kernel code
+        code = '''
+{struct_definition}
+extern "C" __global__ void test_kernel(const double* a,
+                                       custom_user_struct s,
+                                       double* x) {{
+    int i = threadIdx.x;
+    double sum = s.a.x + s.a.y + s.a.z + s.a.w;
+    sum += s.b;
+    sum += s.c[0] + s.c[1];
+    sum += s.d.x;
+    sum += s.e[0] + s.e[1] + s.e[2];
+    x[i] = a[i] + sum;
+}}
+'''.format(struct_definition=struct_definition)
+
+        a_cpu = numpy.arange(24, dtype=numpy.float64)
+        a = cupy.array(a_cpu)
+        x = cupy.empty_like(a)
+
+        func = _compile_func('test_kernel', code)
+        func.linear_launch(a.size, (a, s, x))
+
+        expected = a_cpu + 55.0
+        testing.assert_array_equal(x, expected)
diff --git a/tests/example_tests/test_custom_struct.py b/tests/example_tests/test_custom_struct.py
new file mode 100644
--- /dev/null
+++ b/tests/example_tests/test_custom_struct.py
@@ -0,0 +1,34 @@
+import re
+import unittest
+
+from example_tests import example_test
+
+
+class TestCustomStruct(unittest.TestCase):
+
+    def test_builtin_vectors(self):
+        output = example_test.run_example('custom_struct/builtin_vectors.py')
+        assert re.match(
+            r"Kernel output matches expected value.",
+            output.decode('utf-8'),
+        )
+
+    def test_packed_matrix(self):
+        output = example_test.run_example('custom_struct/packed_matrix.py')
+        assert re.match(
+            r"Kernel output matches expected value for type 'float'.\r?\n"
+            r"Kernel output matches expected value for type 'double'.",
+            output.decode('utf-8'),
+        )
+
+    def test_complex_struct(self):
+        output = example_test.run_example('custom_struct/complex_struct.py')
+        assert re.match(
+            r"Overall structure itemsize: \d+ bytes\r?\n"
+            r"Structure members itemsize: \[(\s*\d+){5}]\r?\n"
+            r"Structure members offsets: \[(\s*\d+){5}]\r?\n"
+            r"Complex structure value:\r?\n"
+            r"\s+\[.*\]\r?\n"
+            r"Kernel output matches expected value.",
+            output.decode('utf-8'),
+        )