Repository: cupy/cupy
Pull request: Support cuComplex.h in `cupy.RawKernel` and `cupy.RawModule`
Closes #1866. Closes #2111.

Previously the support of complex numbers in `cupy.RawKernel` (and thus `cupy.RawModule`) was poorly documented, see https://github.com/cupy/cupy/pull/1866#issuecomment-457791608 and #2111 for example. This PR aims to solve this problem once and for all. This PR supersedes #1866.

The support of the APIs from `cuComplex.h` is enabled by prepending a set of C macros that replace those APIs by their Thrust counterparts. Thus, users can just copy and paste their `cuComplex`-based codes to `cupy.RawKernel` or `cupy.RawModule`, and make them work without any modification. Note that the complex mathematical functions are supported without any additional macros.

For backward compatibility, the newly added option `enable_cuComplex` is set to `False` by default.

attn: @grlee77 

Problem statement:
Use complex number in cupy.RawKernel
Hi!

I wonder what is the native/suggested way to use complex numbers in Cupy.RawKernel?

I was using `#include<cuComplex.h>` until I got an error on my newly installed cupy on another computer. I scrambled through your code and it seems to me that the preferred way is to use `#include<cupy/complex/complex.h>`. Is this the preferred library to use for cupy?

Hints from issue discussion:
See #1866 for some discussion of this and an example
@grlee77 Fantastic. Hope that PR got merged soon.

Commit message: apply review from https://github.com/cupy/cupy/pull/2551#discussion_r344032104

Patch to review:
diff --git a/cupy/core/carray.pxi b/cupy/core/carray.pxi
index fdc57ba2299..f42a8b3a63f 100644
--- a/cupy/core/carray.pxi
+++ b/cupy/core/carray.pxi
@@ -1,4 +1,5 @@
 import os
+import re
 
 from cupy import cuda
 
@@ -122,9 +123,21 @@ cpdef str _get_header_source():
     return _header_source
 
 
+cdef inline str _convert_cuComplex_to_Thrust(str source):
+    source = re.sub(r'<cuComplex.h>',
+                    r'<cupy/cuComplex_bridge.h>  // translate_cucomplex',
+                    source)
+    return source
+
+
 cpdef function.Module compile_with_cache(
         str source, tuple options=(), arch=None, cachd_dir=None,
-        prepend_cupy_headers=True, backend='nvrtc'):
+        prepend_cupy_headers=True, backend='nvrtc', translate_cucomplex=False):
+    if translate_cucomplex:
+        source = _convert_cuComplex_to_Thrust(source)
+        _cupy_header_list.append('cupy/cuComplex_bridge.h')
+        prepend_cupy_headers=True
+
     if prepend_cupy_headers:
         source = _cupy_header + source
     extra_source = _get_header_source()
diff --git a/cupy/core/core.pxd b/cupy/core/core.pxd
index 0d640e629f2..aaae93969f3 100644
--- a/cupy/core/core.pxd
+++ b/cupy/core/core.pxd
@@ -99,7 +99,7 @@ cpdef ndarray asfortranarray(ndarray a, dtype=*)
 
 cpdef Module compile_with_cache(str source, tuple options=*, arch=*,
                                 cachd_dir=*, prepend_cupy_headers=*,
-                                backend=*)
+                                backend=*, translate_cucomplex=*)
 
 
 # TODO(niboshi): Move to _routines_creation.pyx
diff --git a/cupy/core/include/cupy/cuComplex_bridge.h b/cupy/core/include/cupy/cuComplex_bridge.h
new file mode 100644
index 00000000000..6e572768017
--- /dev/null
+++ b/cupy/core/include/cupy/cuComplex_bridge.h
@@ -0,0 +1,34 @@
+/* 
+   This header is to support the "translate_cucomplex" option
+   that turns cuComplex function calls to their Thrust counterparts.
+*/ 
+
+/* ------------------- single complex ------------------- */
+#define cuFloatComplex                complex<float>
+#define cuComplex                     complex<float>
+#define cuCrealf                      real
+#define cuCimagf                      imag
+#define make_cuFloatComplex(A, B)     complex<float>(A, B)
+#define make_cuComplex(A, B)          complex<float>(A, B)
+#define cuConjf                       conj
+#define cuCaddf(A, B)                 (A + B)
+#define cuCsubf(A, B)                 (A - B)
+#define cuCmulf(A, B)                 (A * B)
+#define cuCdivf(A, B)                 (A / B)
+#define cuCabsf                       abs
+#define cuComplexDoubleToFloat        complex<float>
+#define cuCfmaf(A, B, C)              (A * B + C)
+
+/* ------------------- double complex ------------------- */
+#define cuDoubleComplex               complex<double>
+#define cuCreal                       real
+#define cuCimag                       imag
+#define make_cuDoubleComplex(A, B)    complex<double>(A, B)
+#define cuConj                        conj
+#define cuCadd(A, B)                  (A + B)
+#define cuCsub(A, B)                  (A - B)
+#define cuCmul(A, B)                  (A * B)
+#define cuCdiv(A, B)                  (A / B)
+#define cuCabs                        abs
+#define cuComplexFloatToDouble        complex<double>
+#define cuCfma(A, B, C)               (A * B + C)
diff --git a/cupy/core/raw.pxd b/cupy/core/raw.pxd
index 9d7b0d8a72c..a836cfcb31b 100644
--- a/cupy/core/raw.pxd
+++ b/cupy/core/raw.pxd
@@ -6,6 +6,7 @@ cdef class RawKernel:
         readonly tuple options
         object _kernel
         readonly str backend
+        bint cuComplex
 
 
 cdef class RawModule:
@@ -17,3 +18,4 @@ cdef class RawModule:
         dict kernels
         readonly str backend
         object module
+        bint cuComplex
diff --git a/cupy/core/raw.pyx b/cupy/core/raw.pyx
index 0a7f705d653..162de97caec 100644
--- a/cupy/core/raw.pyx
+++ b/cupy/core/raw.pyx
@@ -24,9 +24,12 @@ cdef class RawKernel:
         options (str): Compiler options passed to NVRTC. For details, see
             https://docs.nvidia.com/cuda/nvrtc/index.html#group__options.
         backend (str): Either `nvrtc` or `nvcc`. Defaults to `nvrtc`
+        translate_cucomplex (bool): Whether the CUDA source includes the
+            "cuComplex.h" header or not. Defaults to False.
     """
 
-    def __init__(self, code, name, options=(), backend='nvrtc'):
+    def __init__(self, code, name, options=(), backend='nvrtc',
+                 translate_cucomplex=False):
         if isinstance(code, six.binary_type):
             code = code.decode('UTF-8')
         if isinstance(name, six.binary_type):
@@ -38,6 +41,7 @@ cdef class RawKernel:
         self.options = options
         self._kernel = None
         self.backend = backend
+        self.cuComplex = translate_cucomplex
 
     def __call__(self, grid, block, args, **kwargs):
         """__call__(self, grid, block, args, *, shared_mem=0)
@@ -60,7 +64,7 @@ cdef class RawKernel:
     def kernel(self):
         if self._kernel is None:
             self._kernel = _get_raw_kernel(self.code, self.name, self.options,
-                                           self.backend)
+                                           self.backend, self.cuComplex)
         return self._kernel
 
     @property
@@ -183,9 +187,11 @@ cdef class RawKernel:
 
 
 @cupy.util.memoize(for_each_device=True)
-def _get_raw_kernel(code, name, options=(), backend='nvrtc'):
+def _get_raw_kernel(code, name, options=(), backend='nvrtc',
+                    translate_cucomplex=False):
     module = cupy.core.core.compile_with_cache(
-        code, options, prepend_cupy_headers=False, backend=backend)
+        code, options, prepend_cupy_headers=False, backend=backend,
+        translate_cucomplex=translate_cucomplex)
     return module.get_function(name)
 
 
@@ -210,11 +216,14 @@ cdef class RawModule:
             needed. For details, see
             https://docs.nvidia.com/cuda/nvrtc/index.html#group__options.
         backend (str): Either `nvrtc` or `nvcc`. Defaults to `nvrtc`
+        translate_cucomplex (bool): Whether the CUDA source includes the
+            "cuComplex.h" header or not. Defaults to False.
 
     .. note::
         Each kernel in ``RawModule`` possesses independent function attributes.
     """
-    def __init__(self, code_or_path, options=(), backend='nvrtc'):
+    def __init__(self, code_or_path, options=(), backend='nvrtc',
+                 translate_cucomplex=False):
         if isinstance(code_or_path, six.binary_type):
             code_or_path = code_or_path.decode('UTF-8')
         if isinstance(backend, six.binary_type):
@@ -231,7 +240,8 @@ cdef class RawModule:
 
         if self.code is not None:
             self.module = cupy.core.core.compile_with_cache(
-                code, options, prepend_cupy_headers=False, backend=backend)
+                code, options, prepend_cupy_headers=False, backend=backend,
+                translate_cucomplex=translate_cucomplex)
             self.backend = backend
         elif self.cubin_path is not None:
             self.module = Module()
@@ -240,6 +250,7 @@ cdef class RawModule:
 
         self.options = options
         self.kernels = {}
+        self.cuComplex = translate_cucomplex
 
     def get_function(self, name):
         """Retrieve a CUDA kernel by its name from the module.
@@ -253,7 +264,8 @@ cdef class RawModule:
         if name in self.kernels:
             return self.kernels[name]
         else:
-            ker = RawKernel(None, name, self.options)
+            ker = RawKernel(None, name, self.options, self.backend,
+                            self.cuComplex)
             ker._kernel = self.module.get_function(name)
             self.kernels[name] = ker
             return ker
diff --git a/docs/source/tutorial/kernel.rst b/docs/source/tutorial/kernel.rst
index 9edea30a699..ad02e55d941 100644
--- a/docs/source/tutorial/kernel.rst
+++ b/docs/source/tutorial/kernel.rst
@@ -197,6 +197,33 @@ In other words, you have control over grid size, block size, shared memory size
           [30., 32., 34., 36., 38.],
           [40., 42., 44., 46., 48.]], dtype=float32)
 
+Raw kernels operating on complex-valued arrays can be created as well:
+
+.. doctest::
+
+   >>> complex_kernel = cp.RawKernel(r'''
+   ... #include <cupy/complex.cuh>
+   ... extern "C" __global__
+   ... void my_func(const complex<float>* x1, const complex<float>* x2,
+   ...              complex<float>* y, float a) {
+   ...     int tid = blockDim.x * blockIdx.x + threadIdx.x;
+   ...     y[tid] = x1[tid] + a * x2[tid];
+   ... }
+   ... ''', 'my_func')
+   >>> x1 = cupy.arange(25, dtype=cupy.complex64).reshape(5, 5)
+   >>> x2 = 1j*cupy.arange(25, dtype=cupy.complex64).reshape(5, 5)
+   >>> y = cupy.zeros((5, 5), dtype=cupy.complex64)
+   >>> complex_kernel((5,), (5,), (x1, x2, y, cupy.float32(2.0)))  # grid, block and arguments
+   >>> y
+   array([[ 0. +0.j,  1. +2.j,  2. +4.j,  3. +6.j,  4. +8.j],
+          [ 5.+10.j,  6.+12.j,  7.+14.j,  8.+16.j,  9.+18.j],
+          [10.+20.j, 11.+22.j, 12.+24.j, 13.+26.j, 14.+28.j],
+          [15.+30.j, 16.+32.j, 17.+34.j, 18.+36.j, 19.+38.j],
+          [20.+40.j, 21.+42.j, 22.+44.j, 23.+46.j, 24.+48.j]],
+         dtype=complex64)
+
+Note that while we encourage the usage of ``complex<T>`` types for complex numbers (available by including ``<cupy/complex.cuh>`` as shown above), for CUDA codes already written using functions from ``cuComplex.h`` there is no need to make the conversion yourself: just set the option ``translate_cucomplex=True`` when creating a :class:`~cupy.RawKernel` instance.
+
 The CUDA kernel attributes can be retrieved by either accessing the :attr:`~cupy.RawKernel.attributes` dictionary,
 or by accessing the :class:`~cupy.RawKernel` object's attributes directly; the latter can also be used to set certain
 attributes:
@@ -278,6 +305,8 @@ For dealing a large raw CUDA source or loading an existing CUDA binary, the :cla
     >>> ker_times((N,), (N,), (x1, x2, y, N**2)) # y = x1 * x2
     >>> assert cp.allclose(y, x1 * x2)
 
+The instruction above for using complex numbers in :class:`~cupy.RawKernel` also applies to :class:`~cupy.RawModule`.
+
 CuPy also supports the Texture Reference API. A handle to the texture reference in a module can be retrieved by name via :meth:`~cupy.RawModule.get_texref`. Then, you need to pass it to :class:`~cupy.cuda.texture.TextureReference`, along with a resource descriptor and texture descriptor, for binding the reference to the array. (The interface of :class:`~cupy.cuda.texture.TextureReference` is meant to mimic that of :class:`~cupy.cuda.texture.TextureObject` to help users make transition to the latter, since as of CUDA Toolkit 10.1 the former is marked as deprecated.)
 
 
diff --git a/tests/cupy_tests/core_tests/test_raw.py b/tests/cupy_tests/core_tests/test_raw.py
index af4344cef89..63bf8188a14 100644
--- a/tests/cupy_tests/core_tests/test_raw.py
+++ b/tests/cupy_tests/core_tests/test_raw.py
@@ -83,6 +83,67 @@
 }
 '''
 
+_test_cuComplex = r'''
+#include <cuComplex.h>
+#define N 100
+
+extern "C"{
+/* ------------------- double complex ------------------- */
+
+__global__ void test_div(cuDoubleComplex* arr1, cuDoubleComplex* arr2,
+                         cuDoubleComplex* out) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N) {
+        out[tid] = cuCdiv(arr1[tid], arr2[tid]);
+    }
+}
+
+__global__ void test_fma(cuDoubleComplex* A, cuDoubleComplex* B,
+                         cuDoubleComplex* C, cuDoubleComplex* out) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N) {
+        out[tid] = cuCfma(A[tid], B[tid], C[tid]);
+    }
+}
+
+__global__ void test_make(cuDoubleComplex* arr) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    cuDoubleComplex out = make_cuDoubleComplex(1.8, 2.9);
+    if (tid < N) {
+        arr[tid] = make_cuDoubleComplex(cuCreal(out), -3.*cuCimag(out));
+    }
+}
+
+/* ------------------- single complex ------------------- */
+
+__global__ void test_divf(cuComplex* arr1, cuComplex* arr2,
+                          cuComplex* out) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N) {
+        out[tid] = cuCdivf(arr1[tid], arr2[tid]);
+    }
+}
+
+__global__ void test_fmaf(cuFloatComplex* A, cuFloatComplex* B,
+                          cuFloatComplex* C, cuFloatComplex* out) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N) {
+        out[tid] = cuCfmaf(A[tid], B[tid], C[tid]);
+    }
+}
+
+__global__ void test_makef(cuComplex* arr) {
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    cuComplex out = make_cuFloatComplex(1.8, 2.9);
+    if (tid < N) {
+        arr[tid] = make_cuFloatComplex(cuCrealf(out), -3.*cuCimagf(out));
+    }
+}
+
+}
+'''
+
+
 if 'CUPY_CACHE_DIR' in os.environ:
     _old_cache_dir = os.environ['CUPY_CACHE_DIR']
     _is_cache_env_var_set = True
@@ -192,3 +253,58 @@ def test_get_function_failure(self):
         with pytest.raises(cupy.cuda.driver.CUDADriverError) as ex:
             self.mod2.get_function("no_such_kernel")
         assert 'CUDA_ERROR_NOT_FOUND' in str(ex.value)
+
+    def test_cuFloatComplex(self):
+        N = 100
+        block = 32
+        grid = (N + block - 1) // block
+        dtype = cupy.complex64
+
+        mod = cupy.RawModule(_test_cuComplex, translate_cucomplex=True)
+        a = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        a = a.astype(dtype)
+        b = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        b = b.astype(dtype)
+        c = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        c = c.astype(dtype)
+        out = cupy.zeros((N,), dtype=dtype)
+
+        ker = mod.get_function('test_divf')
+        ker((grid,), (block,), (a, b, out))
+        assert (out == a/b).all()
+
+        ker = mod.get_function('test_fmaf')
+        ker((grid,), (block,), (a, b, c, out))
+        assert (out == a * b + c).all()
+
+        ker = mod.get_function('test_makef')
+        ker((grid,), (block,), (out,))
+        # because of precision issue, the (A==B).all() semantics would fail
+        assert cupy.allclose(out, 1.8 - 1j * 8.7)
+
+    def test_cuDoubleComplex(self):
+        N = 100
+        block = 32
+        grid = (N + block - 1) // block
+        dtype = cupy.complex128
+
+        mod = cupy.RawModule(_test_cuComplex, translate_cucomplex=True)
+        a = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        a = a.astype(dtype)
+        b = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        b = b.astype(dtype)
+        c = cupy.random.random((N,)) + 1j*cupy.random.random((N,))
+        c = c.astype(dtype)
+        out = cupy.zeros((N,), dtype=dtype)
+
+        ker = mod.get_function('test_div')
+        ker((grid,), (block,), (a, b, out))
+        assert (out == a/b).all()
+
+        ker = mod.get_function('test_fma')
+        ker((grid,), (block,), (a, b, c, out))
+        assert (out == a * b + c).all()
+
+        ker = mod.get_function('test_make')
+        ker((grid,), (block,), (out,))
+        assert (out == 1.8 - 1j * 8.7).all()


Code review comments:
diff_hunk: @@ -6,6 +6,7 @@ cdef class RawKernel:
         readonly tuple options
         object _kernel
         readonly str backend
+        bint cuComplex
path: cupy/core/raw.pxd
text: @user1:
nitpick: use `translate_cucomplex` for consistency and PEP8

@author:
done

diff_hunk: @@ -122,9 +123,21 @@ cpdef str _get_header_source():
     return _header_source
 
 
+cdef inline str _convert_cuComplex_to_Thrust(str source):
+    source = re.sub(r'<cuComplex.h>',
path: cupy/core/carray.pxi
text: @user1:
Ah, what I meant by "regex that matches the whole line" was like this:

```py
# module level (to precompile the regex)
cucomplex_include_tokens = ['', '#', 'include', '<', r'cuComplex\.h', '>']
cucomplex_include_pattern = re.compile(r'\s*'.join(cucomplex_include_tokens))

# in _convert_cuComplex_to_Thrust
lines = []
for line in source.splitlines(keepends=True):
    if cucomplex_include_pattern.match(line):
        lines += '#include <cupy/cuComplex_bridge.h>  // translate_cucomplex'
    else:
        lines += line
return ''.join(lines)
```

@author:
I see. Sorry I know nothing about regex. While your regex code looks fine to me (I just googled a few regex tutorials), for some reason it can't detect the following cases:
```python
>>> cucomplex_include_pattern.match(r'#include <cuComplex.h>')  # ok
<_sre.SRE_Match object; span=(0, 22), match='#include <cuComplex.h>'>
>>> cucomplex_include_pattern.match(r'//#include <cuComplex.h>')  # miss!
>>> cucomplex_include_pattern.match(r'// #include <cuComplex.h>')  # miss!
>>> cucomplex_include_pattern.match(r'// # include <cuComplex.h> ')  #miss!
```
@user1 could you help me diagnose please? 

@user1:
```py
>>> cucomplex_include_pattern.match(r'//#include <cuComplex.h>')  # miss!
>>> cucomplex_include_pattern.match(r'// #include <cuComplex.h>')  # miss!
>>> cucomplex_include_pattern.match(r'// # include <cuComplex.h> ')  #miss!
```

This behavior is what I intended. They are commented out so we don't have to replace it.

@author:
Oh oh, I see! I misunderstood what you meant. Will make a change before going to bed :)

@author:
@user1 This is added now.