Repository: cupy/cupy
Pull request: Support separate compilation in `RawKernel`
Resolves #2324. The example there can be run with this PR.

The main change of this PR is two-fold:
1. add wrappers for `cuLinkAddFile` at various levels (driver & function),
2. link to CUDA device runtime (cudadevrt) if requested.

***UPDATE:*** With this PR the separate compilation becomes possible. The dynamical parallelism is just one of the cases in which separate compilation is required, thus the PR title is changed to better reflect this.

**~~Two questions for CuPy core devs~~**: 
1. ~~Do you know if `libcudadevrt.a` always lives in `/usr/local/cuda/lib/`, or it can reside in a different place from the CUDA installation?~~
2. ~~Is there an official way to get the CUDA library path against which CuPy is built and linked?~~

~~In this PR I added an additional env variable `CUPY_CUDADEVRT_PATH` to resolve the above issues, but perhaps we don't need that.~~

~~Tests and docs will follow after the above questions are addressed.~~ 

Thank you.

Problem statement:
Dynamic Parallelism 
I am interested in utilizing Dynamic Parallelism in my cupy projects. 
```cpp
extern "C"
{

__global__ void test_kernel_inner()
{
    printf("  Hello inner world!\n");

}

__global__ void test_kernel()
{
    printf("Hello outer world!\n");
    test_kernel_inner<<<2, 1>>>();
}

}
```
Trying to launch the above  cuda code with the following commands

```py
TestKernel = cp.RawKernel(Test, 'test_kernel', ('-rdc=true',))
TestKernel((1,),(1,),())
```

I get the error
```
  File "/home/palmerss/PycharmProjects/TestProject/Test.py", line 23, in <module>
    TestKernel((1,),(1,),())
  File "cupy/core/raw.pyx", line 51, in cupy.core.raw.RawKernel.__call__
  File "cupy/util.pyx", line 55, in cupy.util.memoize.decorator.ret
  File "cupy/core/raw.pyx", line 57, in cupy.core.raw._get_raw_kernel
  File "cupy/core/carray.pxi", line 125, in cupy.core.core.compile_with_cache
  File "cupy/core/carray.pxi", line 166, in cupy.core.core.compile_with_cache
  File "/home/palmerss/anaconda3/envs/TestProject/lib/python3.7/site-packages/cupy/cuda/compiler.py", line 168, in compile_with_cache
    cubin = ls.complete()
  File "cupy/cuda/function.pyx", line 216, in cupy.cuda.function.LinkState.complete
  File "cupy/cuda/function.pyx", line 217, in cupy.cuda.function.LinkState.complete
  File "cupy/cuda/driver.pyx", line 161, in cupy.cuda.driver.linkComplete
  File "cupy/cuda/driver.pyx", line 82, in cupy.cuda.driver.check_status
cupy.cuda.driver.CUDADriverError: CUDA_ERROR_UNKNOWN: unknown error
```

Hints from issue discussion:
ref: https://stackoverflow.com/questions/57152396/dynamic-parallelism-in-cupy
@palmerss could you test #2426? With that your example worked in my environment.
I will have time to test within the next week or two! Sorry for the delay.
@leofang This seems to be working correctly :) thanks!
@palmerss Thank you for testing!

Commit message: Fixed path on Win

Patch to review:
diff --git a/cupy/cuda/compiler.py b/cupy/cuda/compiler.py
index d2a40fd4351..27ddbc7cb67 100644
--- a/cupy/cuda/compiler.py
+++ b/cupy/cuda/compiler.py
@@ -16,6 +16,8 @@
 
 _nvrtc_version = None
 _nvrtc_max_compute_capability = None
+_win32 = (sys.platform == 'win32')
+_osx = (sys.platform == 'darwin')
 
 
 class NVCCException(Exception):
@@ -67,6 +69,27 @@ def _get_arch():
                _nvrtc_max_compute_capability)
 
 
+def _check_cudadevrt_needed(options):
+    require_cudadevrt = False
+    for option in options:
+        if option in ('--device-c', '-dc', '-rdc=true',
+                      '--relocatable-device-code=true'):
+            require_cudadevrt = True
+            break
+    return require_cudadevrt
+
+
+def _remove_rdc_option(options):
+    new_options = []
+    for option in options:
+        if option in ('--device-c', '-dc', '-rdc=true',
+                      '--relocatable-device-code=true'):
+            continue
+        else:
+            new_options.append(option)
+    return tuple(new_options)
+
+
 class TemporaryDirectory(object):
     def __enter__(self):
         self.path = tempfile.mkdtemp()
@@ -117,7 +140,8 @@ def compile_using_nvrtc(source, options=(), arch=None, filename='kern.cu'):
 
 
 def compile_using_nvcc(source, options=(), arch=None,
-                       filename='kern.cu', code_type='cubin'):
+                       filename='kern.cu', code_type='cubin',
+                       separate_compilation=False):
     if not arch:
         arch = _get_arch()
 
@@ -125,7 +149,7 @@ def compile_using_nvcc(source, options=(), arch=None,
         raise ValueError('Invalid code_type %s. Should be cubin or ptx')
 
     arch_str = '-gencode=arch=compute_{cc},code=sm_{cc}'.format(cc=arch)
-    cmd = ['nvcc', '--%s' % code_type, arch_str] + list(options)
+    cmd = ['nvcc', arch_str]
 
     with TemporaryDirectory() as root_dir:
         first_part = filename.split('.')[0]
@@ -137,19 +161,55 @@ def compile_using_nvcc(source, options=(), arch=None,
         with open(cu_path, 'w') as cu_file:
             cu_file.write(source)
 
-        cmd.append(cu_path)
-
-        try:
-            _run_nvcc(cmd, root_dir)
-        except NVCCException as e:
-            cex = CompileException(str(e), source, cu_path, options, 'nvcc')
-
-            dump = _get_bool_env_variable(
-                'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
-            if dump:
-                cex.dump(sys.stderr)
-
-            raise cex
+        if not separate_compilation:  # majority cases
+            cmd.append('--%s' % code_type)
+            cmd += list(options)
+            cmd.append(cu_path)
+
+            try:
+                _run_nvcc(cmd, root_dir)
+            except NVCCException as e:
+                cex = CompileException(str(e), source, cu_path, options,
+                                       'nvcc')
+
+                dump = _get_bool_env_variable(
+                    'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
+                if dump:
+                    cex.dump(sys.stderr)
+
+                raise cex
+        else:  # two steps: compile to object and device-link
+            cmd_partial = cmd.copy()
+            cmd_partial.append('--cubin')
+
+            obj = path + '.o'
+            options_obj = options
+            options_obj += ('-o', obj)
+            cmd += list(options_obj)
+            cmd.append(cu_path)
+
+            try:
+                _run_nvcc(cmd, root_dir)
+            except NVCCException as e:
+                cex = CompileException(str(e), source, cu_path, options,
+                                       'nvcc')
+
+                dump = _get_bool_env_variable(
+                    'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
+                if dump:
+                    cex.dump(sys.stderr)
+
+                raise cex
+
+            options = _remove_rdc_option(options)
+            options += ('--device-link', obj, '-o', path + '.cubin')
+            cmd = cmd_partial + list(options)
+
+            try:
+                _run_nvcc(cmd, root_dir)
+            except NVCCException as e:
+                cex = CompileException(str(e), '', '', options, 'nvcc')
+                raise cex
 
         if code_type == 'ptx':
             with open(result_path, 'rb') as ptx_file:
@@ -263,10 +323,29 @@ def _compile_with_cache_cuda(source, options, arch, cache_dir,
         ptx = compile_using_nvrtc(source, options, arch, name + '.cu')
         ls = function.LinkState()
         ls.add_ptr_data(ptx, 'cupy.ptx')
+        # for separate compilation
+        if _check_cudadevrt_needed(options):
+            global _win32, _osx
+            # defer import to here to avoid circular dependency
+            from cupy.cuda import get_cuda_path
+            cudadevrt = get_cuda_path()
+            if _win32:
+                # rely on os.altsep
+                cudadevrt += '/lib/x64/cudadevrt.lib'
+            elif _osx:
+                cudadevrt += '/lib/libcudadevrt.a'
+            else:  # linux
+                cudadevrt += '/lib64/libcudadevrt.a'
+            cudadevrt = os.path.normpath(cudadevrt)
+            if not os.path.isfile(cudadevrt):
+                raise RuntimeError('Relocatable PTX code is requested, but '
+                                   'cudadevrt is not found.')
+            ls.add_ptr_file(cudadevrt)
         cubin = ls.complete()
     elif backend == 'nvcc':
+        rdc = True if _check_cudadevrt_needed(options) else False
         cubin = compile_using_nvcc(source, options, arch, name + '.cu',
-                                   code_type='cubin')
+                                   code_type='cubin', separate_compilation=rdc)
     else:
         raise ValueError('Invalid backend %s' % backend)
 
diff --git a/cupy/cuda/cupy_cuda.h b/cupy/cuda/cupy_cuda.h
index d7568c2d73b..dc56c305705 100644
--- a/cupy/cuda/cupy_cuda.h
+++ b/cupy/cuda/cupy_cuda.h
@@ -98,6 +98,10 @@ CUresult cuLinkAddData(...) {
     return CUDA_SUCCESS;
 }
 
+CUresult cuLinkAddFile(...) {
+    return CUDA_SUCCESS;
+}
+
 CUresult cuLinkComplete(...) {
     return CUDA_SUCCESS;
 }
diff --git a/cupy/cuda/driver.pxd b/cupy/cuda/driver.pxd
index c4f611aaf1e..0f3d080879d 100644
--- a/cupy/cuda/driver.pxd
+++ b/cupy/cuda/driver.pxd
@@ -106,6 +106,7 @@ cpdef ctxDestroy(intptr_t ctx)
 
 cpdef intptr_t linkCreate() except? 0
 cpdef linkAddData(intptr_t state, int input_type, bytes data, unicode name)
+cpdef linkAddFile(intptr_t state, int input_type, unicode path)
 cpdef bytes linkComplete(intptr_t state)
 cpdef linkDestroy(intptr_t state)
 cpdef intptr_t moduleLoad(str filename) except? 0
diff --git a/cupy/cuda/driver.pyx b/cupy/cuda/driver.pyx
index 40249f5ebd1..8ce455524db 100644
--- a/cupy/cuda/driver.pyx
+++ b/cupy/cuda/driver.pyx
@@ -39,6 +39,9 @@ cdef extern from 'cupy_cuda.h' nogil:
     int cuLinkAddData(LinkState state, CUjitInputType type, void* data,
                       size_t size, const char* name, unsigned int  numOptions,
                       CUjit_option* options, void** optionValues)
+    int cuLinkAddFile(LinkState state, CUjitInputType type, const char* path,
+                      unsigned int numOptions, CUjit_option* options, void**
+                      optionValues)
     int cuLinkComplete(LinkState state, void** cubinOut, size_t* sizeOut)
     int cuLinkDestroy(LinkState state)
     int cuModuleLoad(Module* module, char* fname)
@@ -195,6 +198,15 @@ cpdef linkAddData(intptr_t state, int input_type, bytes data, unicode name):
     check_status(status)
 
 
+cpdef linkAddFile(intptr_t state, int input_type, unicode path):
+    cdef bytes b_path = path.encode()
+    cdef const char* b_path_ptr = b_path
+    with nogil:
+        status = cuLinkAddFile(<LinkState>state, <CUjitInputType>input_type,
+                               b_path_ptr, 0, <CUjit_option*>0, <void**>0)
+    check_status(status)
+
+
 cpdef bytes linkComplete(intptr_t state):
     cdef void* cubinOut
     cdef size_t sizeOut
diff --git a/cupy/cuda/function.pxd b/cupy/cuda/function.pxd
index b73dc054768..ef8b78b2943 100644
--- a/cupy/cuda/function.pxd
+++ b/cupy/cuda/function.pxd
@@ -33,4 +33,5 @@ cdef class LinkState:
         public intptr_t ptr
 
     cpdef add_ptr_data(self, unicode data, unicode name)
+    cpdef add_ptr_file(self, unicode path)
     cpdef bytes complete(self)
diff --git a/cupy/cuda/function.pyx b/cupy/cuda/function.pyx
index 43fa3d538f5..99217310cdb 100644
--- a/cupy/cuda/function.pyx
+++ b/cupy/cuda/function.pyx
@@ -231,6 +231,9 @@ cdef class LinkState:
         cdef bytes data_byte = data.encode()
         driver.linkAddData(self.ptr, driver.CU_JIT_INPUT_PTX, data_byte, name)
 
+    cpdef add_ptr_file(self, unicode path):
+        driver.linkAddFile(self.ptr, driver.CU_JIT_INPUT_LIBRARY, path)
+
     cpdef bytes complete(self):
         cubin = driver.linkComplete(self.ptr)
         return cubin
diff --git a/docs/source/tutorial/kernel.rst b/docs/source/tutorial/kernel.rst
index 9edea30a699..356a8ccf939 100644
--- a/docs/source/tutorial/kernel.rst
+++ b/docs/source/tutorial/kernel.rst
@@ -201,6 +201,10 @@ The CUDA kernel attributes can be retrieved by either accessing the :attr:`~cupy
 or by accessing the :class:`~cupy.RawKernel` object's attributes directly; the latter can also be used to set certain
 attributes:
 
+Dynamical parallelism is supported by :class:`~cupy.RawKernel`. You just need to provide the linking flag (such as ``-dc``) to :class:`~cupy.RawKernel`'s ``options`` arugment. The static CUDA device runtime library (``cudadevrt``) is automatically discovered by CuPy. For further detail, see `CUDA Toolkit's documentation`_.
+
+.. _CUDA Toolkit's documentation: https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#compiling-and-linking
+
 .. doctest::
 
    >>> add_kernel = cp.RawKernel(r'''
@@ -212,7 +216,7 @@ attributes:
    ... ''', 'my_add')
    >>> add_kernel.attributes  # doctest: +SKIP
    {'max_threads_per_block': 1024, 'shared_size_bytes': 0, 'const_size_bytes': 0, 'local_size_bytes': 0, 'num_regs': 10, 'ptx_version': 70, 'binary_version': 70, 'cache_mode_ca': 0, 'max_dynamic_shared_size_bytes': 49152, 'preferred_shared_memory_carveout': -1}
-   >>> add_kernel.max_dynamic_shared_size_bytes
+   >>> add_kernel.max_dynamic_shared_size_bytes  # doctest: +SKIP
    49152
    >>> add_kernel.max_dynamic_shared_size_bytes = 50000  # set a new value for the attribute  # doctest: +SKIP
    >>> add_kernel.max_dynamic_shared_size_bytes  # doctest: +SKIP
diff --git a/tests/cupy_tests/core_tests/test_raw.py b/tests/cupy_tests/core_tests/test_raw.py
index af4344cef89..5885274672a 100644
--- a/tests/cupy_tests/core_tests/test_raw.py
+++ b/tests/cupy_tests/core_tests/test_raw.py
@@ -83,6 +83,29 @@
 }
 '''
 
+# dynamic parallelism
+_test_source4 = r'''
+extern "C"{
+
+__global__ void test_kernel_inner(float *arr, int N)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+
+    if (tid < N)
+        arr[tid] = 1.0;
+}
+
+__global__ void test_kernel(float *arr, int N, int inner_blk)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+
+    if (tid < N/inner_blk)
+        test_kernel_inner<<<1, inner_blk>>>(arr+tid*inner_blk, inner_blk);
+}
+
+}
+'''
+
 if 'CUPY_CACHE_DIR' in os.environ:
     _old_cache_dir = os.environ['CUPY_CACHE_DIR']
     _is_cache_env_var_set = True
@@ -190,5 +213,29 @@ def test_get_function_failure(self):
         # this error is more likely to appear when using RawModule, so
         # let us do it here
         with pytest.raises(cupy.cuda.driver.CUDADriverError) as ex:
-            self.mod2.get_function("no_such_kernel")
+            self.mod2.get_function('no_such_kernel')
         assert 'CUDA_ERROR_NOT_FOUND' in str(ex.value)
+
+    def test_dynamical_parallelism(self):
+        ker = cupy.RawKernel(_test_source4, 'test_kernel', options=('-dc',),
+                             backend=self.backend)
+        N = 169
+        inner_chunk = 13
+        x = cupy.zeros((N,), dtype=cupy.float32)
+        ker((1,), (N//inner_chunk,), (x, N, inner_chunk))
+        assert (x == 1.0).all()
+
+    def test_dynamical_parallelism_compile_failure(self):
+        # no option for separate compilation is given should cause an error
+        ker = cupy.RawKernel(_test_source4, 'test_kernel',
+                             backend=self.backend)
+        N = 10
+        inner_chunk = 2
+        x = cupy.zeros((N,), dtype=cupy.float32)
+        if self.backend == 'nvrtc':
+            # raised when calling ls.complete()
+            with pytest.raises(cupy.cuda.driver.CUDADriverError):
+                ker((1,), (N//inner_chunk,), (x, N, inner_chunk))
+        else:  # nvcc
+            with pytest.raises(cupy.cuda.compiler.CompileException):
+                ker((1,), (N//inner_chunk,), (x, N, inner_chunk))


Code review comments:
diff_hunk: @@ -67,6 +69,27 @@ def _get_arch():
                _nvrtc_max_compute_capability)
 
 
+def _check_cudadevrt_needed(options):
+    require_cudadevrt = False
+    for option in options:
+        if option in ('--device-c', '-dc', '-rdc=true',
path: cupy/cuda/compiler.py
text: @user1:
Could you define the option list as module variable (e.g. `_rdc_flags`) so that it can be reused in `_remove_rdc_option`?

diff_hunk: @@ -263,10 +323,29 @@ def _compile_with_cache_cuda(source, options, arch, cache_dir,
         ptx = compile_using_nvrtc(source, options, arch, name + '.cu')
         ls = function.LinkState()
         ls.add_ptr_data(ptx, 'cupy.ptx')
+        # for separate compilation
+        if _check_cudadevrt_needed(options):
+            global _win32, _osx
+            # defer import to here to avoid circular dependency
+            from cupy.cuda import get_cuda_path
+            cudadevrt = get_cuda_path()
+            if _win32:
+                # rely on os.altsep
+                cudadevrt += '/lib/x64/cudadevrt.lib'
+            elif _osx:
+                cudadevrt += '/lib/libcudadevrt.a'
+            else:  # linux
+                cudadevrt += '/lib64/libcudadevrt.a'
path: cupy/cuda/compiler.py
text: @user1:
Use `lib64` first, and if it does not exist, use `lib`.
c.f. https://github.com/cupy/cupy/blob/b9db374f4d0ccee358a998d7b283248f77cef351/install/build.py#L115-L118

@author:
Ah OK. Thanks. I wasn't aware that `lib64` is possible on Mac. 

@author:
I guess I parsed your comment wrong. Now I fixed it.

diff_hunk: @@ -263,10 +323,29 @@ def _compile_with_cache_cuda(source, options, arch, cache_dir,
         ptx = compile_using_nvrtc(source, options, arch, name + '.cu')
         ls = function.LinkState()
         ls.add_ptr_data(ptx, 'cupy.ptx')
+        # for separate compilation
+        if _check_cudadevrt_needed(options):
+            global _win32, _osx
+            # defer import to here to avoid circular dependency
+            from cupy.cuda import get_cuda_path
+            cudadevrt = get_cuda_path()
+            if _win32:
+                # rely on os.altsep
+                cudadevrt += '/lib/x64/cudadevrt.lib'
+            elif _osx:
+                cudadevrt += '/lib/libcudadevrt.a'
+            else:  # linux
+                cudadevrt += '/lib64/libcudadevrt.a'
+            cudadevrt = os.path.normpath(cudadevrt)
path: cupy/cuda/compiler.py
text: @user1:
Please avoid using `normpath` as CUDA_PATH may contain `..` (it is not very common setup but not explicitly prohibited).

@author:
I made the requested change, but wouldn't `..` be correctly handled by `normpath`? 

@author:
```python
>>> import os.path
>>> os.path.normpath('/home/leofang/cupy/../')
'/home/leofang'
>>>
```

@user1:
If the parent directory is a symlink, it will be dereferenced by `normpath`.
c.f. https://stackoverflow.com/questions/34865153/os-path-normpath-and-symbolic-links

@author:
Thanks, @user1!

diff_hunk: @@ -137,19 +161,55 @@ def compile_using_nvcc(source, options=(), arch=None,
         with open(cu_path, 'w') as cu_file:
             cu_file.write(source)
 
-        cmd.append(cu_path)
-
-        try:
-            _run_nvcc(cmd, root_dir)
-        except NVCCException as e:
-            cex = CompileException(str(e), source, cu_path, options, 'nvcc')
-
-            dump = _get_bool_env_variable(
-                'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
-            if dump:
-                cex.dump(sys.stderr)
-
-            raise cex
+        if not separate_compilation:  # majority cases
+            cmd.append('--%s' % code_type)
+            cmd += list(options)
+            cmd.append(cu_path)
+
+            try:
+                _run_nvcc(cmd, root_dir)
+            except NVCCException as e:
+                cex = CompileException(str(e), source, cu_path, options,
+                                       'nvcc')
+
+                dump = _get_bool_env_variable(
+                    'CUPY_DUMP_CUDA_SOURCE_ON_ERROR', False)
+                if dump:
+                    cex.dump(sys.stderr)
+
+                raise cex
+        else:  # two steps: compile to object and device-link
+            cmd_partial = cmd.copy()
+            cmd_partial.append('--cubin')
+
+            obj = path + '.o'
+            options_obj = options
+            options_obj += ('-o', obj)
+            cmd += list(options_obj)
path: cupy/cuda/compiler.py
text: @user1:
To make the code more explicit (options is immutable):
```py
cmd += list(options + ('-o', obj))
```

diff_hunk: @@ -16,6 +16,8 @@
 
 _nvrtc_version = None
 _nvrtc_max_compute_capability = None
+_win32 = (sys.platform == 'win32')
path: cupy/cuda/compiler.py
text: @user1:
Use `sys.platform.startswith`.
https://docs.python.org/3/library/sys.html#sys.platform

diff_hunk: @@ -67,6 +69,27 @@ def _get_arch():
                _nvrtc_max_compute_capability)
 
 
+def _check_cudadevrt_needed(options):
+    require_cudadevrt = False
path: cupy/cuda/compiler.py
text: @user1:
More simply:
```py
def _check_cudadevrt_needed(options):
    return any(o for o in options if o in _rdc_flags)
```

@author:
That's more pythonic!

diff_hunk: @@ -67,6 +69,27 @@ def _get_arch():
                _nvrtc_max_compute_capability)
 
 
+def _check_cudadevrt_needed(options):
+    require_cudadevrt = False
+    for option in options:
+        if option in ('--device-c', '-dc', '-rdc=true',
+                      '--relocatable-device-code=true'):
+            require_cudadevrt = True
+            break
+    return require_cudadevrt
+
+
+def _remove_rdc_option(options):
+    new_options = []
path: cupy/cuda/compiler.py
text: @user1:
More simply:
```py
def _remove_rdc_option(options):
    return tuple(o for o in options if o not in _rdc_flags)
```

@author:
I like one liners! Thanks a lot.