Repository: numba/numba

Issue:
Delayed CUDA deallocation breaks pinned/mapped context managers
- [X] I am using the latest released version of Numba (most recent is visible in
 the change log (https://github.com/numba/numba/blob/master/CHANGE_LOG).
- [X] I have included below a minimal working reproducer (if you are unsure how
 to write one see http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports).

Context managers cannot be used to repeatedly pin/map an existing array: the call to `cuMemHostUnregister` is delayed by the same mechanism as device memory deallocation, hence in many cases the memory will still be pinned on subsequent context manager invocations, raising `CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED`.

This fails:
```python
import numpy as np
from numba import cuda

arr = np.zeros(1)
with cuda.pinned(arr):
    pass
with cuda.pinned(arr):
    pass
```
```
...
CudaAPIError: [712] Call to cuMemHostRegister results in CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED
```
This works:
```python
import numpy as np
from numba import cuda

arr = np.zeros(1)
with cuda.pinned(arr):
    pass
cuda.current_context().deallocations.clear()
with cuda.pinned(arr):
    pass
```
Are there good reasons for routing finalizers that wrap `cuMemHostUnregister` through the deallocation queue instead of calling them immediately?

And what about `cuMemHostFree`, i.e., deallocation of memory that was allocated with `cuda.{pinned,mapped}_array`? Seems odd that a chunk of host memory has to wait in line together with objects in device memory to be freed, however I don't think I fully appreciate the implications of these events on asynchronous execution and system freezing in the case of corrupt contexts.


Discussion:
A cheap solution could be to always flush deallocations before the context manager yields, like this:
`numba/cuda/api.py, lines 238-252`
```python
# Page lock
@require_context
@contextlib.contextmanager
def pinned(*arylist):
    """A context manager for temporary pinning a sequence of host ndarrays.
    """
+   current_context().deallocations.clear()
    pmlist = []
    for ary in arylist:
        pm = current_context().mempin(ary, driver.host_pointer(ary),
                                      driver.host_memory_size(ary),
                                      mapped=False)
        pmlist.append(pm)
    yield
    del pmlist
```
However, this doesn't help if the context manager is called inside a `defer_cleanup()` context. That is, the following still fails:
```python
with cuda.pinned(arr):
    pass
with cuda.defer_cleanup():
    with cuda.pinned(arr):
        pass
```

---
Side note: To ensure cleanup in case of an exception within the `with` block, the last two lines should be wrapped in a try/finally statement:
```python
    try:
        yield
    finally:
        del pmlist
```
Thanks for the report, I can reproduce. I think your assessment of the problem is correct and evidently some thought needs to go into a suitable fix. Thanks for pointing out test cases and providing initial thoughts, v. useful.
I think you are right that at exit of `with cuda.pinned(arr):` the `arr` should have been unpinned without delay.  The same goes for `mapped`. 

The delayed device deallocation is needed to avoid breaking asynchronous execution because device  arrays can be created automatically and go out-of-scope in odd places.  As for `pinned` and `mapped`, there are explicitly created by users so it won't have the same problem of going out-of-scope unknowingly.  So, I don't think they need to have delayed cleanup.

Resolving patch:
diff --git a/numba/cuda/api.py b/numba/cuda/api.py
--- a/numba/cuda/api.py
+++ b/numba/cuda/api.py
@@ -248,7 +248,6 @@ def pinned(*arylist):
                                       mapped=False)
         pmlist.append(pm)
     yield
-    del pmlist
 
 
 @require_context
@@ -257,22 +256,27 @@ def mapped(*arylist, **kws):
     """A context manager for temporarily mapping a sequence of host ndarrays.
     """
     assert not kws or 'stream' in kws, "Only accept 'stream' as keyword."
-    pmlist = []
     stream = kws.get('stream', 0)
+    pmlist = []
+    devarylist = []
     for ary in arylist:
         pm = current_context().mempin(ary, driver.host_pointer(ary),
                                     driver.host_memory_size(ary),
                                     mapped=True)
         pmlist.append(pm)
-
-    devarylist = []
-    for ary, pm in zip(arylist, pmlist):
         devary = devicearray.from_array_like(ary, gpu_data=pm, stream=stream)
         devarylist.append(devary)
-    if len(devarylist) == 1:
-        yield devarylist[0]
-    else:
-        yield devarylist
+    try:
+        if len(devarylist) == 1:
+            yield devarylist[0]
+        else:
+            yield devarylist
+    finally:
+        # When exiting from `with cuda.mapped(*arrs) as mapped_arrs:`, the name
+        # `mapped_arrs` stays in scope, blocking automatic unmapping based on
+        # reference count. We therefore invoke the finalizer manually.
+        for pm in pmlist:
+            pm.free()
 
 
 def event(timing=True):
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
@@ -696,9 +696,8 @@ def allocator():
 
         self._attempt_allocation(allocator)
 
-        _memory_finalizer = _make_mem_finalizer(driver.cuMemFree, bytesize)
-        mem = AutoFreePointer(weakref.proxy(self), ptr, bytesize,
-                              _memory_finalizer(self, ptr))
+        finalizer = _alloc_finalizer(self, ptr, bytesize)
+        mem = AutoFreePointer(weakref.proxy(self), ptr, bytesize, finalizer)
         self.allocations[ptr.value] = mem
         return mem.own()
 
@@ -722,17 +721,14 @@ def allocator():
 
         owner = None
 
-        if mapped:
-            _hostalloc_finalizer = _make_mem_finalizer(driver.cuMemFreeHost,
-                                                       bytesize)
-            finalizer = _hostalloc_finalizer(self, pointer)
-            mem = MappedMemory(weakref.proxy(self), owner, pointer,
-                               bytesize, finalizer=finalizer)
+        finalizer = _hostalloc_finalizer(self, pointer, bytesize, mapped)
 
+        if mapped:
+            mem = MappedMemory(weakref.proxy(self), owner, pointer, bytesize,
+                               finalizer=finalizer)
             self.allocations[mem.handle.value] = mem
             return mem.own()
         else:
-            finalizer = _pinnedalloc_finalizer(self.deallocations, pointer)
             mem = PinnedMemory(weakref.proxy(self), owner, pointer, bytesize,
                                finalizer=finalizer)
             return mem
@@ -760,18 +756,16 @@ def allocator():
         else:
             allocator()
 
+        finalizer = _pin_finalizer(self, pointer, mapped)
+
         if mapped:
-            _mapped_finalizer = _make_mem_finalizer(driver.cuMemHostUnregister,
-                                                    size)
-            finalizer = _mapped_finalizer(self, pointer)
             mem = MappedMemory(weakref.proxy(self), owner, pointer, size,
                                finalizer=finalizer)
             self.allocations[mem.handle.value] = mem
             return mem.own()
         else:
             mem = PinnedMemory(weakref.proxy(self), owner, pointer, size,
-                               finalizer=_pinned_finalizer(self.deallocations,
-                                                           pointer))
+                               finalizer=finalizer)
             return mem
 
     def memunpin(self, pointer):
@@ -899,32 +893,60 @@ def load_module_image(context, image):
                   _module_finalizer(context, handle))
 
 
-def _make_mem_finalizer(dtor, bytesize):
-    def mem_finalize(context, handle):
-        allocations = context.allocations
-        deallocations = context.deallocations
+def _alloc_finalizer(context, handle, size):
+    allocations = context.allocations
+    deallocations = context.deallocations
 
-        def core():
-            if allocations:
-                del allocations[handle.value]
+    def core():
+        if allocations:
+            del allocations[handle.value]
+        deallocations.add_item(driver.cuMemFree, handle, size)
 
-            deallocations.add_item(dtor, handle, size=bytesize)
+    return core
 
-        return core
 
-    return mem_finalize
+def _hostalloc_finalizer(context, handle, size, mapped):
+    """
+    Finalize page-locked host memory allocated by `context.memhostalloc`.
 
+    This memory is managed by CUDA, and finalization entails deallocation. The
+    issues noted in `_pin_finalizer` are not relevant in this case, and the
+    finalization is placed in the `context.deallocations` queue along with
+    finalization of device objects.
+
+    """
+    allocations = context.allocations
+    deallocations = context.deallocations
+    if not mapped:
+        size = _SizeNotSet
 
-def _pinnedalloc_finalizer(deallocs, handle):
     def core():
-        deallocs.add_item(driver.cuMemFreeHost, handle)
+        if mapped and allocations:
+            del allocations[handle.value]
+        deallocations.add_item(driver.cuMemFreeHost, handle, size)
 
     return core
 
 
-def _pinned_finalizer(deallocs, handle):
+def _pin_finalizer(context, handle, mapped):
+    """
+    Finalize temporary page-locking of host memory by `context.mempin`.
+
+    This applies to memory not otherwise managed by CUDA. Page-locking can
+    be requested multiple times on the same memory, and must therefore be
+    lifted as soon as finalization is requested, otherwise subsequent calls to
+    `mempin` may fail with `CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED`, leading
+    to unexpected behavior for the context managers `cuda.{pinned,mapped}`.
+    This function therefore carries out finalization immediately, bypassing the
+    `context.deallocations` queue.
+
+    """
+    allocations = context.allocations
+
     def core():
-        deallocs.add_item(driver.cuMemHostUnregister, handle)
+        if mapped and allocations:
+            del allocations[handle.value]
+        driver.cuMemHostUnregister(handle)
 
     return core
 
@@ -1237,7 +1259,7 @@ def __init__(self, *args, **kwargs):
         self.refct -= 1
 
 
-class MappedMemory(MemoryPointer):
+class MappedMemory(AutoFreePointer):
     __cuda_memory__ = True
 
     def __init__(self, context, owner, hostpointer, size,


Test patch:
diff --git a/numba/cuda/tests/cudadrv/test_deallocations.py b/numba/cuda/tests/cudadrv/test_deallocations.py
--- a/numba/cuda/tests/cudadrv/test_deallocations.py
+++ b/numba/cuda/tests/cudadrv/test_deallocations.py
@@ -172,6 +172,61 @@ def test_device_memory(self):
         with self.check_ignored_exception(ctx):
             del mem
 
+    def test_pinned_contextmanager(self):
+        # Check that temporarily pinned memory is unregistered immediately,
+        # such that it can be re-pinned at any time
+        class PinnedException(Exception):
+            pass
+
+        arr = np.zeros(1)
+        ctx = cuda.current_context()
+        ctx.deallocations.clear()
+        with self.check_ignored_exception(ctx):
+            with cuda.pinned(arr):
+                pass
+            with cuda.pinned(arr):
+                pass
+            # Should also work inside a `defer_cleanup` block
+            with cuda.defer_cleanup():
+                with cuda.pinned(arr):
+                    pass
+                with cuda.pinned(arr):
+                    pass
+            # Should also work when breaking out of the block due to an exception
+            try:
+                with cuda.pinned(arr):
+                    raise PinnedException
+            except PinnedException:
+                with cuda.pinned(arr):
+                    pass
+
+    def test_mapped_contextmanager(self):
+        # Check that temporarily mapped memory is unregistered immediately,
+        # such that it can be re-mapped at any time
+        class MappedException(Exception):
+            pass
+
+        arr = np.zeros(1)
+        ctx = cuda.current_context()
+        ctx.deallocations.clear()
+        with self.check_ignored_exception(ctx):
+            with cuda.mapped(arr) as marr:
+                pass
+            with cuda.mapped(arr) as marr:
+                pass
+            # Should also work inside a `defer_cleanup` block
+            with cuda.defer_cleanup():
+                with cuda.mapped(arr) as marr:
+                    pass
+                with cuda.mapped(arr) as marr:
+                    pass
+            # Should also work when breaking out of the block due to an exception
+            try:
+                with cuda.mapped(arr) as marr:
+                    raise MappedException
+            except MappedException:
+                with cuda.mapped(arr) as marr:
+                    pass
 
 if __name__ == '__main__':
-    unittest.main()
\ No newline at end of file
+    unittest.main()