Repository: numba/numba

Issue:
Add support for shared memory to the CUDA simulator
CUDA Python kernels which use shared memory will fail if the CUDA simulator is enabled.

Supporting this properly probably also requires solving the syncthreads situation, which is tricky.



Discussion:
(following up from the mailing list post)

It might be a good idea to indicate something to this effect in the docs, where shared mem is listed under the [supported features of the simulator](http://numba.pydata.org/numba-doc/dev/cuda/simulator.html).

What is the resolution required for the syncthreads situation? I didn't find an issue, but I am willing to try to help out.

Actually, upon inspection of the simulator code, I realized that it already handles syncthreads properly.  In that case, implementing shared memory is much easier.  @sklam: Ideas on how we would tackle this?

While trying to reproduce the issue with a simple code, I see that it occurs when I split code across modules, e.g. I've got two files:

`numba_cudasim_shared_inner.py`

``` python
from numba import cuda, float32
import math

def gen_inner(n, n_iter):
    @cuda.jit(device=True)
    def inner_fn(data):
        ti = cuda.threadIdx.x
        shm = cuda.shared.array((n, ), float32)
        shm[ti] = data[ti]
        for i in range(n_iter):
            shm[ti] = math.sin(shm[ti])
        data[ti] = shm[ti]
    return inner_fn
```

and `numba_cudasim_shared.py`

``` python
import numpy
from numba import cuda, float32

from numba_cudasim_shared_inner import gen_inner

def gen_outer(inner):
    @cuda.jit
    def outer_fn(data):
        inner(data)
    return outer_fn


test_data = (numpy.random.rand(32) - 0.5).astype(numpy.float32)
kernel = gen_outer(gen_inner(test_data.shape[0], 100))
kernel[(test_data.shape[0],), (1,)](test_data)
```

which produces a stack trace ending with 

```
  File "numba_cudasim_shared_inner.py", line 8, in inner_fn
    ti = cuda.threadIdx.x
AttributeError: tid=[0, 0, 0] ctaid=[0, 0, 0]: 'module' object has no attribute 'threadIdx'
```

If I place the shared memory declaration line first, in the inner kernel, I see `AttributeError: tid=[0, 0, 0] ctaid=[0, 0, 0]: 'module' object has no attribute 'shared'`.

So, this is in fact an issue with using kernels across multiple modules, because, if I define it all in the same file, it works fine.

The code in `numba_cudasim_shared_inner.py` and `numba_cudasim_shared.py` works when they are in a single file.   Shared memory is implemented and worked.

I am looking into why it has the attribute error when the files are splitted.


Resolving patch:
diff --git a/numba/cuda/simulator/kernel.py b/numba/cuda/simulator/kernel.py
--- a/numba/cuda/simulator/kernel.py
+++ b/numba/cuda/simulator/kernel.py
@@ -5,6 +5,36 @@
 import numpy as np
 import sys
 import threading
+from contextlib import contextmanager
+
+
+"""
+Global variable to keep track of the current "kernel context", i.e the
+FakeCUDAModule.  We only support one kernel launch at a time.
+No support for concurrent kernel launch.
+"""
+_kernel_context = None
+
+
+@contextmanager
+def _push_kernel_context(mod):
+    """
+    Push the current kernel context.
+    """
+    global _kernel_context
+    assert _kernel_context is None, "conrrent simulated kernel not supported"
+    _kernel_context = mod
+    try:
+        yield
+    finally:
+        _kernel_context = None
+
+
+def _get_kernel_context():
+    """
+    Get the current kernel context. This is usually done by a device function.
+    """
+    return _kernel_context
 
 
 class FakeCUDAKernel(object):
@@ -22,23 +52,25 @@ def __init__(self, fn, device, fastmath=False):
 
     def __call__(self, *args):
         if self._device:
-            return self.fn(*args)
+            with swapped_cuda_module(self.fn, _get_kernel_context()):
+                return self.fn(*args)
 
         fake_cuda_module = FakeCUDAModule(self.grid_dim, self.block_dim,
                                           self.dynshared_size)
-        # fake_args substitutes all numpy arrays for FakeCUDAArrays
-        # because they implement some semantics differently
-        def fake_arg(arg):
-            if isinstance(arg, np.ndarray) and arg.ndim > 0:
-                return to_device(arg)
-            return arg
-        fake_args = [fake_arg(arg) for arg in args]
-
-        with swapped_cuda_module(self.fn, fake_cuda_module):
-            # Execute one block at a time
-            for grid_point in np.ndindex(*self.grid_dim):
-                bm = BlockManager(self.fn, self.grid_dim, self.block_dim)
-                bm.run(grid_point, *fake_args)
+        with _push_kernel_context(fake_cuda_module):
+            # fake_args substitutes all numpy arrays for FakeCUDAArrays
+            # because they implement some semantics differently
+            def fake_arg(arg):
+                if isinstance(arg, np.ndarray) and arg.ndim > 0:
+                    return to_device(arg)
+                return arg
+            fake_args = [fake_arg(arg) for arg in args]
+
+            with swapped_cuda_module(self.fn, fake_cuda_module):
+                # Execute one block at a time
+                for grid_point in np.ndindex(*self.grid_dim):
+                    bm = BlockManager(self.fn, self.grid_dim, self.block_dim)
+                    bm.run(grid_point, *fake_args)
 
     def __getitem__(self, configuration):
         grid_dim = configuration[0]
diff --git a/numba/cuda/simulator/kernelapi.py b/numba/cuda/simulator/kernelapi.py
--- a/numba/cuda/simulator/kernelapi.py
+++ b/numba/cuda/simulator/kernelapi.py
@@ -10,7 +10,7 @@
 
 import numpy as np
 
-from numba import types, numpy_support
+from numba import numpy_support
 
 
 class Dim3(object):
@@ -207,11 +207,17 @@ def gridsize(self, n):
 
 @contextmanager
 def swapped_cuda_module(fn, fake_cuda_module):
+    from numba import cuda
+
     fn_globs = fn.__globals__
-    original_cuda, fn_globs['cuda'] = fn_globs['cuda'], fake_cuda_module
+    # get all globals that is the "cuda" module
+    orig = dict((k, v) for k, v in fn_globs.items() if v is cuda)
+    # build replacement dict
+    repl = dict((k, fake_cuda_module) for k, v in orig.items())
+    # replace
+    fn_globs.update(repl)
     try:
         yield
     finally:
-        fn_globs['cuda'] = original_cuda
-
-
+        # revert
+        fn_globs.update(orig)


Test patch:
diff --git a/numba/cuda/tests/__init__.py b/numba/cuda/tests/__init__.py
--- a/numba/cuda/tests/__init__.py
+++ b/numba/cuda/tests/__init__.py
@@ -3,12 +3,14 @@
 from numba import cuda
 from os.path import dirname, join
 
+
 def load_tests(loader, tests, pattern):
 
     suite = SerialSuite()
     this_dir = dirname(__file__)
     suite.addTests(load_testsuite(loader, join(this_dir, 'nocuda')))
     if cuda.is_available():
+        suite.addTests(load_testsuite(loader, join(this_dir, 'cudasim')))
         gpus = cuda.list_devices()
         if gpus and gpus[0].compute_capability >= (2, 0):
             suite.addTests(load_testsuite(loader, join(this_dir, 'cudadrv')))
diff --git a/numba/cuda/tests/cudasim/__init__.py b/numba/cuda/tests/cudasim/__init__.py
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudasim/__init__.py
@@ -0,0 +1,8 @@
+from numba.testing import SerialSuite
+from numba.testing import load_testsuite
+import os
+from numba import config
+
+
+def load_tests(loader, tests, pattern):
+    return SerialSuite(load_testsuite(loader, os.path.dirname(__file__)))
diff --git a/numba/cuda/tests/cudasim/support.py b/numba/cuda/tests/cudasim/support.py
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudasim/support.py
@@ -0,0 +1,6 @@
+from numba import cuda
+
+
+@cuda.jit(device=True)
+def cuda_module_in_device_function():
+    return cuda.threadIdx.x
diff --git a/numba/cuda/tests/cudasim/test_cudasim_issues.py b/numba/cuda/tests/cudasim/test_cudasim_issues.py
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudasim/test_cudasim_issues.py
@@ -0,0 +1,32 @@
+from __future__ import absolute_import, print_function, division
+
+import numpy as np
+
+from numba import unittest_support as unittest
+from numba import cuda
+
+
+class TestCudaSimIssues(unittest.TestCase):
+
+    def test_cuda_module_in_device_function(self):
+        """
+        Discovered in https://github.com/numba/numba/issues/1837.
+        When the `cuda` module is referenced in a device function,
+        it does not have the kernel API (e.g. cuda.threadIdx, cuda.shared)
+        """
+        from .support import cuda_module_in_device_function as inner
+
+        @cuda.jit
+        def outer(out):
+            tid = inner()
+            if tid < out.size:
+                out[tid] = tid
+
+        arr = np.zeros(10, dtype=np.int32)
+        outer[1, 11](arr)
+        expected = np.arange(arr.size, dtype=np.int32)
+        np.testing.assert_equal(expected, arr)
+
+
+if __name__ == '__main__':
+    unittest.main()