Repository: numba/numba

Issue:
Support per-thread default streams in CUDA
Following the merge of https://github.com/numba/numba/pull/5136, it might be useful to have support for the per-thread default-stream option (reference: https://devblogs.nvidia.com/gpu-pro-tip-cuda-7-streams-simplify-concurrency/)

I think this wouldn't be a large effort to implement - opening this issue to note that it could be implemented and to gauge interest.

Are there any use cases this would unblock for anyone?


Discussion:
This is interesting. I checked the documentation and saw [this](https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html#stream-sync-behavior):
> The per-thread default stream is not a non-blocking stream and will synchronize with the legacy default stream if both are used in a program.

So I suppose the interplay between two libraries built with different `--default-stream` options, say CuPy (on `legacy`) and Numba (on `per-thread` with this proposal), is safe? Am I right that there will be a synchronization before a CuPy array on the `legacy` default stream being passed to a Numba kernel on the `per-thread` default stream?
> So I suppose the interplay between two libraries built with different `--default-stream` options, say CuPy (on `legacy`) and Numba (on `per-thread` with this proposal), is safe? Am I right that there will be a synchronization before a CuPy array on the `legacy` default stream being passed to a Numba kernel on the `per-thread` default stream?

When you say "a CuPy array on the `legacy` default stream", are you referring to a scenario like:

1. A `cuMemcpy<something>Async` with the array as the destination on the legacy default stream, then
2. The array immediately being passed to a kernel launched on a per-thread default stream?

If so, my understanding is that the interplay between them would be safe (though I would specifically be testing this scenario as part of implementing per-thread default stream support in Numba).

If I've misunderstood - could you explain a little more about the scenario you're thinking of please?
Yes, this is more or less what I had in mind, except that in CuPy a GPU array typically does not need memcpy to initialize, so it can be as simple as
```python
import cupy as cp
from numba import cuda

@cuda.jit
def some_func(x):
   # ...

a = cp.arange(0, 100).astype(cp.float64) # <- lives on "legacy" default stream
some_func[(1,), (100,)](a)               # <- may live on "per-thread" default stream
```
In this case, CuPy internally launches a ufunc kernel to initialize the array.

On the testing:
> though I would specifically be testing this scenario as part of implementing per-thread default stream support in Numba

It'd be nice to do this, but I am not sure how you could test a sync indeed happens as expected. If we just look at the output, even if there's no synchronization between the two kinds of default streams, the output could still be (non-deterministically) correct.
I was thinking about per-thread streams when I reviewed #5136, too. I have also been reviewing and experimenting with the legacy stream's implicit sync behavior.

I think the example from @leofang is safe:

```python
a = cp.arange(0, 100).astype(cp.float64) # <- lives on "legacy" default stream
some_func[(1,), (100,)](a)               # <- may live on "per-thread" default stream
```

because the legacy stream is implicitly synchronizing. The CUDA programming guide (at https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#default-stream)  says:

> Two commands from different streams cannot run concurrently if any one of the following operations is issued in-between them by the host thread:
> ...
> * any CUDA command to the NULL stream,

(The NULL stream is the legacy stream).

With that, we know that the per-thread default stream must run after any command to the legacy stream.

This article (https://devblogs.nvidia.com/gpu-pro-tip-cuda-7-streams-simplify-concurrency/) also says "Note that cudaStreamLegacy still synchronizes implicitly with the per-thread default streams if you happen to mix them in a program."

These lead me to believe that mixing the legacy and per-thread streams are safe.
Thanks for the discussions @sklam / @leofang - it seems like it should be safe. As an additional point of safety, I think that the per-thread default stream should be off by default, so that the default stream semantics are unchanged after adding the implementation. I'd imagine a config variable so that something like:

```
NUMBA_DEFAULT_STREAM="per-thread" python prog.py
```

or 

```python
numba.config.DEFAULT_STREAM="per-thread" # "legacy" being the default
```

would be used to enable it.


> It'd be nice to do this, but I am not sure how you could test a sync indeed happens as expected. If we just look at the output, even if there's no synchronization between the two kinds of default streams, the output could still be (non-deterministically) correct.

I'm not sure I could write an automated test for this - I'm imagining writing a small example that I'd be carefully checking runs of by hand under various circumstances (with some attempt to make the test prone to failure if the synchronisation doesn't happen / happens at the wrong point). Perhaps something could be done with events on the streams to check the order of things happening could be used, but I'd need to think a bit harder about it to come to a definite conclusion.
Making this a proper feature request, as it seems it will help with the use case mentioned in #6921 by @insertinterestingnamehere.
Quick note: the implementation would need (when PTDS is enabled) to use the appropriate driver API functions: those suffixed with `_ptds` (or `_ptsz` in some cases) instead of the variants with the `_v2` suffix (or no suffix, if there was no `_v2`-suffixed version).
A quick starter patch:

```diff
diff --git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py
index 77e4c5eae..a2bdb5789 100644
--- a/numba/cuda/cudadrv/driver.py
+++ b/numba/cuda/cudadrv/driver.py
@@ -309,17 +309,15 @@ class Driver(object):
         return functools.wraps(libfn)(wrapper)
 
     def _find_api(self, fname):
-        # Try version 2
-        try:
-            return getattr(self.lib, fname + "_v2")
-        except AttributeError:
-            pass
-
-        # Try regular
-        try:
-            return getattr(self.lib, fname)
-        except AttributeError:
-            pass
+        variants = ('_v2_ptds', '_v2_ptsz', '_ptds', '_ptsz', '_v2', '')
+        for variant in variants:
+            try:
+                full_name = f'{fname}{variant}'
+                api = getattr(self.lib, full_name)
+                print(f"\nUsing {full_name}\n")
+                return api
+            except AttributeError:
+                pass
 
         # Not found.
         # Delay missing function error to use
```

passes all tests:

```
$ python -m numba.runtests numba.cuda.tests

...

Ran 1140 tests in 87.612s

OK (skipped=10, expected failures=1)
```

A summary of the functions it reports can be produced:

```bash
grep "^Using" test_output.txt | sort | uniq | awk '{print $2}'
```

which gives:

```
cuCtxCreate_v2
cuCtxDestroy_v2
cuCtxGetCurrent
cuCtxGetDevice
cuCtxPopCurrent_v2
cuCtxPushCurrent_v2
cuCtxSynchronize
cuDeviceCanAccessPeer
cuDeviceComputeCapability
cuDeviceGet
cuDeviceGetAttribute
cuDeviceGetCount
cuDeviceGetName
cuDeviceGetUuid
cuDevicePrimaryCtxRelease_v2
cuDevicePrimaryCtxReset_v2
cuDevicePrimaryCtxRetain
cuEventCreate
cuEventDestroy_v2
cuEventElapsedTime
cuEventRecord_ptsz
cuEventSynchronize
cuFuncGetAttribute
cuInit
cuIpcCloseMemHandle
cuIpcGetMemHandle
cuIpcOpenMemHandle_v2
cuLaunchCooperativeKernel_ptsz
cuLaunchKernel_ptsz
cuLinkAddData_v2
cuLinkAddFile_v2
cuLinkComplete
cuLinkCreate_v2
cuLinkDestroy
cuMemAllocManaged
cuMemAlloc_v2
cuMemcpyDtoD_v2_ptds
cuMemcpyDtoHAsync_v2_ptsz
cuMemcpyDtoH_v2_ptds
cuMemcpyHtoDAsync_v2_ptsz
cuMemcpyHtoD_v2_ptds
cuMemFreeHost
cuMemFree_v2
cuMemGetAddressRange_v2
cuMemGetInfo_v2
cuMemHostAlloc
cuMemHostGetDevicePointer_v2
cuMemHostRegister_v2
cuMemHostUnregister
cuMemsetD8_v2_ptds
cuModuleGetFunction
cuModuleGetGlobal_v2
cuModuleLoadDataEx
cuModuleUnload
cuOccupancyMaxActiveBlocksPerMultiprocessor
cuOccupancyMaxPotentialBlockSize
cuPointerGetAttribute
cuProfilerStart
cuProfilerStop
cuStreamAddCallback_ptsz
cuStreamCreate
cuStreamDestroy_v2
cuStreamSynchronize_ptsz
cuStreamWaitEvent_ptsz
```

To finish the patch, we need:

- Validation that all the correct API functions are chosen - the above list of functions can be correlated with those defined using the `__CUDA_API_PTDS` and `__CUDA_API_PTSZ` macros in `cuda.h`.
- A config option for PTDS. On by default probably makes sense, assuming other CUDA Python libraries such as CuPy also have it on by default.
- Addition of some tests of PTDS behaviour.
- Tidying up the implementation of `_find_api` a little.
Adding this to the 0.54RC milestone as it seems that it could plausibly be implemented for the next release.
> Validation that all the correct API functions are chosen - the above list of functions can be correlated with those defined using the __CUDA_API_PTDS and __CUDA_API_PTSZ macros in cuda.h.

I've manually verified that the above implementation picks up all the correct APIs.
This is amazing! Thank you!

Resolving patch:
diff --git a/numba/core/config.py b/numba/core/config.py
--- a/numba/core/config.py
+++ b/numba/core/config.py
@@ -347,6 +347,10 @@ def avx_default():
         # Whether to generate verbose log messages when JIT linking
         CUDA_VERBOSE_JIT_LOG = _readenv("NUMBA_CUDA_VERBOSE_JIT_LOG", int, 1)
 
+        # Whether the default stream is the per-thread default stream
+        CUDA_PER_THREAD_DEFAULT_STREAM = _readenv(
+            "NUMBA_CUDA_PER_THREAD_DEFAULT_STREAM", int, 0)
+
         # Compute contiguity of device arrays using the relaxed strides
         # checking algorithm.
         NPY_RELAXED_STRIDES_CHECKING = _readenv(
diff --git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py
--- a/numba/cuda/cudadrv/driver.py
+++ b/numba/cuda/cudadrv/driver.py
@@ -309,17 +309,16 @@ def safe_cuda_api_call(*args):
         return functools.wraps(libfn)(wrapper)
 
     def _find_api(self, fname):
-        # Try version 2
-        try:
-            return getattr(self.lib, fname + "_v2")
-        except AttributeError:
-            pass
+        if config.CUDA_PER_THREAD_DEFAULT_STREAM:
+            variants = ('_v2_ptds', '_v2_ptsz', '_ptds', '_ptsz', '_v2', '')
+        else:
+            variants = ('_v2', '')
 
-        # Try regular
-        try:
-            return getattr(self.lib, fname)
-        except AttributeError:
-            pass
+        for variant in variants:
+            try:
+                return getattr(self.lib, f'{fname}{variant}')
+            except AttributeError:
+                pass
 
         # Not found.
         # Delay missing function error to use


Test patch:
diff --git a/numba/cuda/tests/cudadrv/test_ptds.py b/numba/cuda/tests/cudadrv/test_ptds.py
new file mode 100644
--- /dev/null
+++ b/numba/cuda/tests/cudadrv/test_ptds.py
@@ -0,0 +1,142 @@
+import multiprocessing as mp
+import logging
+import traceback
+from numba.cuda.testing import unittest, CUDATestCase
+from numba.cuda.testing import skip_on_cudasim
+
+
+def child_test():
+    from numba import cuda, int32, void
+    from numba.core import config
+    import io
+    import numpy as np
+    import threading
+
+    # Enable PTDS before we make any CUDA driver calls.  Enabling it first
+    # ensures that PTDS APIs are used because the CUDA driver looks up API
+    # functions on first use and memoizes them.
+    config.CUDA_PER_THREAD_DEFAULT_STREAM = 1
+
+    # Set up log capture for the Driver API so we can see what API calls were
+    # used.
+    logbuf = io.StringIO()
+    handler = logging.StreamHandler(logbuf)
+    cudadrv_logger = logging.getLogger('numba.cuda.cudadrv.driver')
+    cudadrv_logger.addHandler(handler)
+    cudadrv_logger.setLevel(logging.DEBUG)
+
+    # Set up data for our test, and copy over to the device
+    N = 2 ** 16
+    N_THREADS = 10
+    N_ADDITIONS = 4096
+
+    # Seed the RNG for repeatability
+    np.random.seed(1)
+    x = np.random.randint(low=0, high=1000, size=N, dtype=np.int32)
+    r = np.zeros_like(x)
+
+    # One input and output array for each thread
+    xs = [cuda.to_device(x) for _ in range(N_THREADS)]
+    rs = [cuda.to_device(r) for _ in range(N_THREADS)]
+
+    # Compute the grid size and get the [per-thread] default stream
+    n_threads = 256
+    n_blocks = N // n_threads
+    stream = cuda.default_stream()
+
+    # A simple multiplication-by-addition kernel. What it does exactly is not
+    # too important; only that we have a kernel that does something.
+    @cuda.jit(void(int32[::1], int32[::1]))
+    def f(r, x):
+        i = cuda.grid(1)
+
+        if i > len(r):
+            return
+
+        # Accumulate x into r
+        for j in range(N_ADDITIONS):
+            r[i] += x[i]
+
+    # This function will be used to launch the kernel from each thread on its
+    # own unique data.
+    def kernel_thread(n):
+        f[n_blocks, n_threads, stream](rs[n], xs[n])
+
+    # Create threads
+    threads = [threading.Thread(target=kernel_thread, args=(i,))
+               for i in range(N_THREADS)]
+
+    # Start all threads
+    for thread in threads:
+        thread.start()
+
+    # Wait for all threads to finish, to ensure that we don't synchronize with
+    # the device until all kernels are scheduled.
+    for thread in threads:
+        thread.join()
+
+    # Synchronize with the device
+    cuda.synchronize()
+
+    # Check output is as expected
+    expected = x * N_ADDITIONS
+    for i in range(N_THREADS):
+        np.testing.assert_equal(rs[i].copy_to_host(), expected)
+
+    # Return the driver log output to the calling process for checking
+    handler.flush()
+    return logbuf.getvalue()
+
+
+def child_test_wrapper(result_queue):
+    try:
+        output = child_test()
+        success = True
+    # Catch anything raised so it can be propagated
+    except: # noqa: E722
+        output = traceback.format_exc()
+        success = False
+
+    result_queue.put((success, output))
+
+
+@skip_on_cudasim('Streams not supported on the simulator')
+class TestPTDS(CUDATestCase):
+    def test_ptds(self):
+        # Run a test with PTDS enabled in a child process
+        ctx = mp.get_context('spawn')
+        result_queue = ctx.Queue()
+        proc = ctx.Process(target=child_test_wrapper, args=(result_queue,))
+        proc.start()
+        proc.join()
+        success, output = result_queue.get()
+
+        # Ensure the child process ran to completion before checking its output
+        if not success:
+            self.fail(output)
+
+        # Functions with a per-thread default stream variant that we expect to
+        # see in the output
+        ptds_functions = ('cuMemcpyHtoD_v2_ptds', 'cuLaunchKernel_ptsz',
+                          'cuMemcpyDtoH_v2_ptds')
+
+        for fn in ptds_functions:
+            with self.subTest(fn=fn, expected=True):
+                self.assertIn(fn, output)
+
+        # Non-PTDS versions of the functions that we should not see in the
+        # output:
+        legacy_functions = ('cuMemcpyHtoD_v2', 'cuLaunchKernel',
+                            'cuMemcpyDtoH_v2')
+
+        for fn in legacy_functions:
+            with self.subTest(fn=fn, expected=False):
+                # Ensure we only spot these function names appearing without a
+                # _ptds or _ptsz suffix by checking including the end of the
+                # line in the log
+                fn_at_end = f'{fn}\n'
+                self.assertNotIn(fn_at_end, output)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/numba/cuda/tests/cudapy/test_ipc.py b/numba/cuda/tests/cudapy/test_ipc.py
--- a/numba/cuda/tests/cudapy/test_ipc.py
+++ b/numba/cuda/tests/cudapy/test_ipc.py
@@ -10,8 +10,6 @@
                                 ContextResettingTestCase, ForeignArray)
 import unittest
 
-has_mp_get_context = hasattr(mp, 'get_context')
-
 
 def core_ipc_handle_test(the_work, result_queue):
     try:
@@ -79,7 +77,6 @@ def ipc_array_test(ipcarr, result_queue):
 
 
 @skip_under_cuda_memcheck('Hangs cuda-memcheck')
-@unittest.skipUnless(has_mp_get_context, "requires multiprocessing.get_context")
 @skip_on_cudasim('Ipc not available in CUDASIM')
 class TestIpcMemory(ContextResettingTestCase):
     def test_ipc_handle(self):
@@ -229,7 +226,6 @@ def staged_ipc_array_test(ipcarr, device_num, result_queue):
 
 
 @skip_under_cuda_memcheck('Hangs cuda-memcheck')
-@unittest.skipUnless(has_mp_get_context, "requires multiprocessing.get_context")
 @skip_on_cudasim('Ipc not available in CUDASIM')
 class TestIpcStaged(ContextResettingTestCase):
     def test_staged(self):