Repository: cupy/cupy

Issue:
jit.rawkernel decorated function cannot be used on anothor device
When I trying to make the kernel function decorated by jit.rawkernel() to work on another device, it shows `"CUDA_ERROR_INVALID_HANDLE: invalid resource handle", though it can be successfully executed on the first device.

The following code can demonstrate it more clearly..
  
* Code to reproduce
``` python
import cupy as cp
import cupyx as cpx
from torch.utils.dlpack import to_dlpack
from torch.utils.dlpack import from_dlpack
import numpy as np
import torch

@cpx.jit.rawkernel()
def add_one_kernel(var0, out, m, n):
    tid = cpx.jit.threadIdx.x + cpx.jit.blockIdx.x * cpx.jit.blockDim.x
    if tid < n:
        for i in range(m):
            out[i, tid] = var0[i, tid] + 1

def AddOne_cupy(var0, window, valid):
    var0 = cp.fromDlpack(to_dlpack(var0))
    with cp.cuda.Device(var0.device.id):
        out = cp.zeros_like(var0)
        add_one_kernel[64, 64](var0, out, var0.shape[0], var0.shape[1])
        out = from_dlpack(out.toDlpack())
        out[:window, :] = np.nan
        out[~valid] = np.nan
    print(f"success on device {var0.device.id}")
    return out

cp.show_config()
### speed test
var0 = torch.randn(2500, 4000).cuda()
valid = torch.ones(var0.shape).bool().cuda()
window=20
AddOne_cupy(var0, window, valid)

var1 = torch.randn(2500, 4000).cuda(1)
valid1 = torch.ones(var0.shape).bool().cuda(1)
AddOne_cupy(var1, window, valid1)

```
**"success on device 0"** is printed! but fail to execute on the second device
* Error messages, stack traces, or logs
```
Traceback (most recent call last):
  File "test2.py", line 38, in <module>
    AddOne_cupy(var1, window, valid1, stream1)
  File "test2.py", line 19, in AddOne_cupy
    add_one_kernel[64, 64](var0, out, var0.shape[0], var0.shape[1])
  File "/opt/anaconda3/lib/python3.7/site-packages/cupyx/jit/_interface.py", line 123, in <lambda>
    return lambda *args, **kwargs: self(grid, block, args, **kwargs)
  File "/opt/anaconda3/lib/python3.7/site-packages/cupyx/jit/_interface.py", line 111, in __call__
    kern(grid, block, args, shared_mem, stream, enable_cooperative_groups)
  File "cupy/cuda/function.pyx", line 201, in cupy.cuda.function.Function.__call__
  File "cupy/cuda/function.pyx", line 183, in cupy.cuda.function._launch
  File "cupy_backends/cuda/api/driver.pyx", line 306, in cupy_backends.cuda.api.driver.launchKernel
  File "cupy_backends/cuda/api/driver.pyx", line 125, in cupy_backends.cuda.api.driver.check_status
cupy_backends.cuda.api.driver.CUDADriverError: CUDA_ERROR_INVALID_HANDLE: invalid resource handle

```

* Conditions (you can just paste the output of `python -c 'import cupy; cupy.show_config()'`)
  - CuPy version: 9.4.0
  - OS/Platform:  Linux-4.18.0-147.8.1.el8_1.x86_64-x86_64-with-centos-8.1.1911-Core
  - CUDA version: 11010
  - cuDNN/NCCL version: 8005


Discussion:
Thanks for the report! Seems we should cache compiled kernel per device:
https://github.com/cupy/cupy/blob/8d9e5aebe8d20faf1809ed11a0d7852be6fdfba6/cupyx/jit/_interface.py#L108-L109

Is your GPU model different between the first one and the second one?
> Thanks for the report! Seems we should cache compiled kernel per device:
> https://github.com/cupy/cupy/blob/8d9e5aebe8d20faf1809ed11a0d7852be6fdfba6/cupyx/jit/_interface.py#L108-L109
> 
> Is your GPU model different between the first one and the second one?

They are same.. 
I simplified it to the above example code. It is quite simple, adding one to every element of an array...
> Thanks for the report! Seems we should cache compiled kernel per device:
> https://github.com/cupy/cupy/blob/8d9e5aebe8d20faf1809ed11a0d7852be6fdfba6/cupyx/jit/_interface.py#L108-L109
> 
> Is your GPU model different between the first one and the second one?

The printed cache code is as follow:
```
{(<cupyx.jit._cuda_types.CArray object at 0x7fa5c07287d0>, <cupyx.jit._cuda_types.CArray object at 0x7fa5c0728910>, <cupyx.jit._cuda_types.Scalar object at 0x7fa5c0728950>, <cupyx.jit._cuda_types.Scalar object at 0x7fa5c07288d0>): '\nextern "C" __global__ void add_one_kernel(CArray<float, 2, true, true> var0, CArray<float, 2, true, true> out, int m, int n) {\n  unsigned int tid;\n  int i;\n  tid = (threadIdx.x + (blockIdx.x * blockDim.x));\n  if ((tid < (unsigned int)(n))) {\n    for (int __it = 0, __stop = m, __step = 1; __step >= 0 ? __it < __stop : __it > __stop; __it += __step) {\n      i = __it;\n      out._indexing(thrust::make_tuple(i, tid)) = (var0._indexing(thrust::make_tuple(i, tid)) + 1.0f);\n    }\n  }\n  else {\n  }\n}'}
```

The device info:
![image](https://user-images.githubusercontent.com/46940676/135019752-2f15f0d1-3c1e-4c41-b6aa-b4f8a2401897.png)

> Is your GPU model different between the first one and the second one?

Ah, the device architecture is not related here. `self._cache` holds functions loaded by `cuModuleGetFunction` so it should be loaded for each device. We have to fix it to hold a cache for each device anyways.

Resolving patch:
diff --git a/cupyx/jit/_interface.py b/cupyx/jit/_interface.py
--- a/cupyx/jit/_interface.py
+++ b/cupyx/jit/_interface.py
@@ -87,16 +87,21 @@ def __call__(
                 raise TypeError(f'{type(x)} is not supported for RawKernel')
             in_types.append(t)
         in_types = tuple(in_types)
+        device_id = cupy.cuda.get_device_id()
 
-        kern = self._cache.get(in_types)
+        kern = self._cache.get((in_types, device_id))
         if kern is None:
-            result = _compile.transpile(
-                self._func,
-                ['extern "C"', '__global__'],
-                self._mode,
-                in_types,
-                _cuda_types.void,
-            )
+            result = self._cached_codes.get(in_types)
+            if result is None:
+                result = _compile.transpile(
+                    self._func,
+                    ['extern "C"', '__global__'],
+                    self._mode,
+                    in_types,
+                    _cuda_types.void,
+                )
+                self._cached_codes[in_types] = result
+
             fname = result.func_name
             # workaround for hipRTC: as of ROCm 4.1.0 hipRTC still does not
             # recognize "-D", so we have to compile using hipcc...
@@ -106,8 +111,7 @@ def __call__(
                 options=('-DCUPY_JIT_MODE', '--std=c++11'),
                 backend=backend)
             kern = module.get_function(fname)
-            self._cache[in_types] = kern
-            self._cached_codes[in_types] = result.code
+            self._cache[(in_types, device_id)] = kern
 
         kern(grid, block, args, shared_mem, stream, enable_cooperative_groups)
 
@@ -134,7 +138,7 @@ def cached_codes(self):
             warnings.warn(
                 'No codes are cached because compilation is deferred until '
                 'the first function call.')
-        return self._cached_codes
+        return [result.code for result in self._cached_codes]
 
     @property
     def cached_code(self):


Test patch:
diff --git a/tests/cupyx_tests/jit_tests/test_raw.py b/tests/cupyx_tests/jit_tests/test_raw.py
--- a/tests/cupyx_tests/jit_tests/test_raw.py
+++ b/tests/cupyx_tests/jit_tests/test_raw.py
@@ -646,3 +646,21 @@ def f(arr):
         x = cupy.zeros((1,), dtype=cupy.uint32)
         f((1,), (1,), (x,))
         assert x == N
+
+    @testing.multi_gpu(2)
+    def test_device_cache(self):
+        @jit.rawkernel()
+        def f(x, y):
+            tid = jit.threadIdx.x + jit.blockDim.x * jit.blockIdx.x
+            y[tid] = x[tid]
+
+        with device.Device(0):
+            x = testing.shaped_random((30,), dtype=numpy.int32, seed=0)
+            y = testing.shaped_random((30,), dtype=numpy.int32, seed=1)
+            f((5,), (6,), (x, y))
+            assert bool((x == y).all())
+        with device.Device(1):
+            x = testing.shaped_random((30,), dtype=numpy.int32, seed=2)
+            y = testing.shaped_random((30,), dtype=numpy.int32, seed=3)
+            f((5,), (6,), (x, y))
+            assert bool((x == y).all())