Repository: cupy/cupy

Issue:
Add RawKernel.compile() method
As discussed in #1657, a `.compile()` method is desired. I implemented it and added two tests.

A few words on this implementation: while it functions as requested, upon reflection I feel the `compile()` method should belong to the module level instead of the `RawKernel` class. The later implies that a CUDA kernel, instead of a module, is returned and used. However, I don't see any other way to do so without breaking backward compatibility, so this is a design choice I leave for you to worry in the future. Should you have any comments or suggestions, I'm more than happy to discuss and make changes. 

Thanks.
documenting cupy.cuda.function.Module?
I was working with an Nvidia engineer @laytonjb to help a group of scientists migrate their Python codebase to GPU, and he suggested to use cupy. After some experiments I realized that cupy provides almost identical functionalities as in `pycuda.driver.module_from_file`, namely to load precompiled cubin (CUDA binary) and grab the kernels therein. We wonder why this great feature is not documented at all (I hope we didn't miss anything!). IMHO this is a huge attraction for pycuda users (as I am), and for anyone who needs the flexibility of occasionally working with low-level CUDA kernels. Several issues related to JIT compilation (such as #1258, #1398, more recently #1655, etc) could've been less urgent if this were documented in the first place.

For people who are looking for this feature, below is the steps I found that worked perfectly for us : 
suppose we have a file named `cupy_mod.cu` which is defined as follows
```cuda
extern "C"{ //avoid C++ name mangling 
//C=A*B, so Ay=Bx
__global__ void mat_mul(double * A, double * B, double * C, int Ax, int Bx, int By) {
   /* implementation goes here */
   }

/* other kernels */
}
```
then the steps to take is 
1. compile the .cu file to .cubin (CUDA binary) with `nvcc -arch=sm_XX -cubin -o cupy_mod.cubin cupy_mod.cu`
2. load it in python
```python
import cupy as cp

# create a Module object in python
mod = cp.cuda.function.Module()

# load the cubin
mod.load_file("/path/to/cupy_mod.cubin")

# fetch the kernel to make it a Python function object
mat_mul_cp = mod.get_function("mat_mul")

# declare A, B, C as 2D cupy arrays of dtype cp.float64
# be sure they are C contiguous arrays!
# ...omitted...

# call the function with a tuple of grid size, a tuple of block size, and a tuple of all arguments required by the kernel
# if the kernel requires shared memory, append `shared_mem=n_bytes` to the function call
mat_mul_cp( ((A.shape[1]+128-1)//128, 0, 0), (128, 0, 0), (A, B, C, cp.int32(A.shape[0]), cp.int32(B.shape[0]), cp.int32(B.shape[1])))
```


Discussion:
Thanks @leofang! Will take a look.
Please block this PR by #1941. If the backend is changeable, `RawKernel.compile()` will need to take an extra argument `backend='nvrtc'`. 
Can I take up on this issue
Sorry for not responding to this issue.
We now provide `RawKernel`. Does that satisfy your requirement?
https://docs-cupy.chainer.org/en/latest/tutorial/kernel.html#raw-kernels
Hi @kmaehashi, thanks for your reply. I'm aware of `RawKernel` and I saw it was referenced in several issues, but there're two reasons that we found directly accessing the functionality of `cupy.cuda.function.Module` is more useful. 

First, we have a legacy codebase written in CUDA and glued to Python by PyCUDA, so the source codes for Python and CUDA are separated. We find this is easier to maintain because we can version-control and syntax-highlight the CUDA code easily, which would not be feasible as far as I know if the CUDA kernels were hard-wired as strings in the Python code.

Second, for deploying our Python + CUDA (PyCUDA) + MPI (mpi4py) code in a HPC cluster with high-speed interconnect and job scheduler, we've seen on-the-fly compiling caused mysterious crashes that we were unable to pinpoint, at least it happened with PyCUDA, so pre-compilation to cubin on each node before production runs is highly desired. (Well, I know there's a caching option for `RawKernel`, so a way to circumvent this point is to do a "trial run" first to compile. It's just a matter of taste for how things can be done :P) 


> First, we have a legacy codebase written in CUDA and glued to Python by PyCUDA, so the source codes for Python and CUDA are separated. We find this is easier to maintain because we can version-control and syntax-highlight the CUDA code easily, which would not be feasible as far as I know if the CUDA kernels were hard-wired as strings in the Python code.

I agree with separating CUDA source file from Python code is better. In fact our example does that; read CUDA source file and pass it to `RawKernel`. https://github.com/cupy/cupy/tree/master/examples/gemm
In this way, you don't have to manage `cubin`s separately.

> Second, for deploying our Python + CUDA (PyCUDA) + MPI (mpi4py) code in a HPC cluster with high-speed interconnect and job scheduler, we've seen on-the-fly compiling caused mysterious crashes that we were unable to pinpoint, at least it happened with PyCUDA, so pre-compilation to cubin on each node before production runs is highly desired.

We run Python + CuPy + MPI (mpi4py) in our HPC cluster everyday (our motivation is to do [distributed deep learning](https://docs.chainer.org/en/latest/chainermn/index.html)), but haven't heard of such "mysterious crashes". One difference between PyCUDA and CuPy regarding compilation is that PyCUDA uses `NVCC` (external command) whereas CuPy directly uses `NVRTC` (C/C++ library behind `NVCC` command). (I don't think PyCUDA's approach is unstable, though.)

> (Well, I know there's a caching option for RawKernel, so a way to circumvent this point is to do a "trial run" first to compile. It's just a matter of taste for how things can be done :P)

Absolutely. You can just run the code with small dataset instead of writing precompilation code.

In conclusion, I'd like to suggest you to try `RawKernel`, but I'm also interested in the use-case of HPC cluster users migrating from pure CUDA to CUDA + Python.
Will it make sense for you if we provide `k = RawKernel(...); k.compile()` which just run compilation without actually running it, or you still have some other use-cases to load external cubin file?
@kmaehashi sorry for my late reply. I thought I posted it last Friday but apparently it was not sent out...

> I agree with separating CUDA source file from Python code is better. In fact our example does that; read CUDA source file and pass it to `RawKernel`. https://github.com/cupy/cupy/tree/master/examples/gemm
> In this way, you don't have to manage `cubin`s separately.
> 
This is indeed a very nice way of integrating CUDA source codes with Python. Thanks for pointing me to this example.

> > Second, for deploying our Python + CUDA (PyCUDA) + MPI (mpi4py) code in a HPC cluster with high-speed interconnect and job scheduler, we've seen on-the-fly compiling caused mysterious crashes that we were unable to pinpoint, at least it happened with PyCUDA, so pre-compilation to cubin on each node before production runs is highly desired.
> 
> We run Python + CuPy + MPI (mpi4py) in our HPC cluster everyday (our motivation is to do [distributed deep learning](https://docs.chainer.org/en/latest/chainermn/index.html)), but haven't heard of such "mysterious crashes". One difference between PyCUDA and CuPy regarding compilation is that PyCUDA uses `NVCC` (external command) whereas CuPy directly uses `NVRTC` (C/C++ library behind `NVCC` command). (I don't think PyCUDA's approach is unstable, though.)

If this works then perhaps it'd solve the problem for us. We suspected that a bad interplay among MPI, InfiniBand driver and Slurm when managing a spawned process for `NVCC` caused the crashes we saw, although we ran out of resources to verify this guess. I looked up quickly the `NVRTC` doc, and if I understand it correctly it does not spawn any new process when compiling the CUDA codes. Another motivation for us to move to CuPy!

> > (Well, I know there's a caching option for RawKernel, so a way to circumvent this point is to do a "trial run" first to compile. It's just a matter of taste for how things can be done :P)
> 
> Absolutely. You can just run the code with small dataset instead of writing precompilation code.
> 
> In conclusion, I'd like to suggest you to try `RawKernel`, but I'm also interested in the use-case of HPC cluster users migrating from pure CUDA to CUDA + Python.
> Will it make sense for you if we provide `k = RawKernel(...); k.compile()` which just run compilation without actually running it, or you still have some other use-cases to load external cubin file?

I see that `RawKernel` specifies a kernel in a `.cu` file by name. If a `.compile()` method is to be implemented, I'd expect that it can compile the entire `.cu` file, which may contain more than one kernel. But to support this may require some API change and thus break backward compatibility for `RawKernel`? What do you think?

> If this works then perhaps it'd solve the problem for us. We suspected that a bad interplay among MPI, InfiniBand driver and Slurm when managing a spawned process for NVCC caused the crashes we saw, although we ran out of resources to verify this guess. I looked up quickly the NVRTC doc, and if I understand it correctly it does not spawn any new process when compiling the CUDA codes. Another motivation for us to move to CuPy!

We're also using MPI over InfiniBand but haven't experienced issues regarding compilation when using CuPy. As you see in the doc, yes, NVRTC does not spawn any processes.

> I see that RawKernel specifies a kernel in a .cu file by name. If a .compile() method is to be implemented, I'd expect that it can compile the entire .cu file, which may contain more than one kernel. But to support this may require some API change and thus break backward compatibility for RawKernel? What do you think?

Do you have multiple kernels in one `.cu` file that are difficult to separate?
I think we can provide such feature without breaking compatibility, like:

```py
kernels = cp.RawKernel.compile(code, ['kern1', 'kern2'])
kernels['kern1'](grid, block, args)
```

or, more like PyCUDA,

```py
module = cp.RawKernel.compile(code)
module.get_function('kern1')(grid, block, args)
```
> Do you have multiple kernels in one `.cu` file that are difficult to separate?
>
 
Yes I think in our case it's hard to separate. We have a series of kernels in one `.cu` file, each of which is supposed to be compiled twice to handle single and double precisions. A set of macros is used for controlling this, and in our old way we use a shell script to pass the marco values to `nvcc`. In the end we generate two `.cubin` for the two precision versions, and we load the corresponding one programmatically given user input. As a result, in a heterogeneous environment with N different generations of CUDA devices, we need 2*N `.cubin` files.

> I think we can provide such feature without breaking compatibility, like:
> 
> ```python
> kernels = cp.RawKernel.compile(code, ['kern1', 'kern2'])
> kernels['kern1'](grid, block, args)
> ```
> 
> or, more like PyCUDA,
> 
> ```python
> module = cp.RawKernel.compile(code)
> module.get_function('kern1')(grid, block, args)
> ```

The second approach may be the way to go in my opinion. We just need to add a part of code to ` cp.RawKernel.compile` to handle the one-argument (`code`) scenario. I could try to do a PR for this if you'd like me to. It's a good chance for me to learn the CuPy codebase.
Ah, I think we don't need a list of kernel names even if the former interface.

```py
module = cp.RawKernel.compile(code)
module['kern1'](grid, block, args)
```

I'd like to hear opinions from @okuta if this interface is okay.
I disussed with @okuta and concluded that 

```py
module = cp.RawKernel.compile(code)
module.get_function('kern1')(grid, block, args)
```

is better as it is similar to PyCUDA's interface.
@leofang Could you try implementing it?
@kmaehashi ok I'll try
Hi, has the compile method been implemented now?
It's in #1889 but is not merged to master yet.

Resolving patch:
diff --git a/cupy/__init__.py b/cupy/__init__.py
--- a/cupy/__init__.py
+++ b/cupy/__init__.py
@@ -655,6 +655,7 @@ def isscalar(num):
 
 from cupy.core import ElementwiseKernel  # NOQA
 from cupy.core import RawKernel  # NOQA
+from cupy.core import RawModule  # NOQA
 from cupy.core import ReductionKernel  # NOQA
 
 # -----------------------------------------------------------------------------
diff --git a/cupy/core/__init__.py b/cupy/core/__init__.py
--- a/cupy/core/__init__.py
+++ b/cupy/core/__init__.py
@@ -62,3 +62,4 @@
 from cupy.core.internal import complete_slice  # NOQA
 from cupy.core.internal import get_size  # NOQA
 from cupy.core.raw import RawKernel  # NOQA
+from cupy.core.raw import RawModule  # NOQA


Test patch:
diff --git a/tests/cupy_tests/core_tests/test_raw.py b/tests/cupy_tests/core_tests/test_raw.py
--- a/tests/cupy_tests/core_tests/test_raw.py
+++ b/tests/cupy_tests/core_tests/test_raw.py
@@ -1,9 +1,10 @@
+import pytest
 import unittest
 
 import cupy
 
 
-_test_source = r'''
+_test_source1 = r'''
 extern "C" __global__
 void test_sum(const float* x1, const float* x2, float* y) {
     int tid = blockDim.x * blockIdx.x + threadIdx.x;
@@ -11,18 +12,91 @@
 }
 '''
 
+# test compiling and invoking multiple kernels in one single .cubin
+_test_source2 = r'''
+extern "C"{
+
+__global__ void test_sum(const float* x1, const float* x2, float* y, \
+                         unsigned int N)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N)
+    {
+        y[tid] = x1[tid] + x2[tid];
+    }
+}
+
+__global__ void test_multiply(const float* x1, const float* x2, float* y, \
+                              unsigned int N)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N)
+    {
+        y[tid] = x1[tid] * x2[tid];
+    }
+}
+
+}
+'''
+
+# test C macros
+_test_source3 = r'''
+#ifndef PRECISION
+    #define PRECISION 2
+#endif
+
+#if PRECISION == 2
+    #define TYPE double
+#elif PRECISION == 1
+    #define TYPE float
+#else
+    #error precision not supported
+#endif
+
+extern "C"{
+
+__global__ void test_sum(const TYPE* x1, const TYPE* x2, TYPE* y, \
+                         unsigned int N)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N)
+    {
+        y[tid] = x1[tid] + x2[tid];
+    }
+}
+
+__global__ void test_multiply(const TYPE* x1, const TYPE* x2, TYPE* y, \
+                              unsigned int N)
+{
+    unsigned int tid = blockDim.x * blockIdx.x + threadIdx.x;
+    if (tid < N)
+    {
+        y[tid] = x1[tid] * x2[tid];
+    }
+}
+
+}
+'''
+
 
 class TestRaw(unittest.TestCase):
 
     def setUp(self):
-        self.kern = cupy.RawKernel(_test_source, 'test_sum')
+        self.kern = cupy.RawKernel(_test_source1, 'test_sum')
+        self.mod2 = cupy.RawModule(_test_source2)
+        self.mod3 = cupy.RawModule(_test_source3, ("-DPRECISION=2",))
+
+    def _helper(self, kernel, dtype):
+        N = 10
+        x1 = cupy.arange(N**2, dtype=dtype).reshape(N, N)
+        x2 = cupy.ones((N, N), dtype=dtype)
+        y = cupy.zeros((N, N), dtype=dtype)
+        kernel((N,), (N,), (x1, x2, y, N**2))
+        return x1, x2, y
 
     def test_basic(self):
-        x1 = cupy.arange(100, dtype=cupy.float32).reshape(10, 10)
-        x2 = cupy.ones((10, 10), dtype=cupy.float32)
-        y = cupy.zeros((10, 10), dtype=cupy.float32)
-        self.kern((10,), (10,), (x1, x2, y))
-        assert (y == x1 + x2).all()
+        x1, x2, y = self._helper(self.kern, cupy.float32)
+        assert cupy.allclose(y, x1 + x2)
 
     def test_kernel_attributes(self):
         attrs = self.kern.attributes
@@ -40,3 +114,47 @@ def test_kernel_attributes(self):
         assert self.kern.num_regs > 0
         assert self.kern.max_threads_per_block > 0
         assert self.kern.shared_size_bytes == 0
+
+    def test_module(self):
+        module = self.mod2
+        ker_sum = module.get_function('test_sum')
+        ker_times = module.get_function('test_multiply')
+
+        x1, x2, y = self._helper(ker_sum, cupy.float32)
+        assert cupy.allclose(y, x1 + x2)
+
+        x1, x2, y = self._helper(ker_times, cupy.float32)
+        assert cupy.allclose(y, x1 * x2)
+
+    def test_compiler_flag(self):
+        module = self.mod3
+        ker_sum = module.get_function('test_sum')
+        ker_times = module.get_function('test_multiply')
+
+        x1, x2, y = self._helper(ker_sum, cupy.float64)
+        assert cupy.allclose(y, x1 + x2)
+
+        x1, x2, y = self._helper(ker_times, cupy.float64)
+        assert cupy.allclose(y, x1 * x2)
+
+    def test_invalid_compiler_flag(self):
+        with pytest.raises(cupy.cuda.compiler.CompileException) as ex:
+            cupy.RawModule(_test_source3, ("-DPRECISION=3",))
+        assert 'precision not supported' in str(ex)
+
+    def test_module_load_failure(self):
+        # in principle this test is better done in test_driver.py, but
+        # this error is more likely to appear when using RawModule, so
+        # let us do it here
+        import os
+        with pytest.raises(cupy.cuda.driver.CUDADriverError) as ex:
+            cupy.RawModule(os.path.expanduser("~/this_does_not_exist.cubin"))
+        assert 'CUDA_ERROR_FILE_NOT_FOUND' in str(ex)
+
+    def test_get_function_failure(self):
+        # in principle this test is better done in test_driver.py, but
+        # 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")
+        assert 'CUDA_ERROR_NOT_FOUND' in str(ex)