Fix NRT stats for cuda python (#168)

* Fix NRT stats when cuda-python bindings in use

The NRT stats implementation did not take the difference in ctypes
pointers between the ctypes and cuda-python bindings into account.

* NRT tests: Add a test for single stats

This functionality was previously not exercised by the test suite.

Author: gmarkall

numba_cuda/numba/cuda/runtime/nrt.py

@@ -5,7 +5,8 @@
 
 from numba import cuda, config
 from numba.core.runtime.nrt import _nrt_mstats
-from numba.cuda.cudadrv.driver import Linker, driver, launch_kernel
+from numba.cuda.cudadrv.driver import (Linker, driver, launch_kernel,
+                                       USE_NV_BINDING)
 from numba.cuda.cudadrv import devices
 from numba.cuda.api import get_current_device
 from numba.cuda.utils import _readenv
@@ -128,6 +129,18 @@ def _single_thread_launch(self, module, stream, name, params=()):
             cooperative=False
         )
 
+    def _ctypes_pointer(self, array):
+        """
+        Given an array, return a ctypes pointer to the data suitable for
+        passing to ``launch_kernel``.
+        """
+        ptr = array.device_ctypes_pointer
+
+        if USE_NV_BINDING:
+            ptr = ctypes.c_void_p(int(ptr))
+
+        return ptr
+
     def ensure_initialized(self, stream=None):
         """
         If memsys is not initialized, initialize memsys
@@ -174,12 +187,13 @@ def memsys_stats_enabled(self, stream=None):
         context
         """
         enabled_ar = cuda.managed_array(1, np.uint8)
+        enabled_ptr = self._ctypes_pointer(enabled_ar)
 
         self._single_thread_launch(
             self._memsys_module,
             stream,
             "NRT_MemSys_stats_enabled",
-            (enabled_ar.device_ctypes_pointer,)
+            (enabled_ptr,)
         )
 
         cuda.synchronize()
@@ -198,12 +212,13 @@ def _copy_memsys_to_host(self, stream):
         ])
 
         stats_for_read = cuda.managed_array(1, dt)
+        stats_ptr = self._ctypes_pointer(stats_for_read)
 
         self._single_thread_launch(
             self._memsys_module,
             stream,
             "NRT_MemSys_read",
-            [stats_for_read.device_ctypes_pointer]
+            [stats_ptr]
         )
         cuda.synchronize()
 
@@ -231,11 +246,13 @@ def _get_single_stat(self, stat, stream=None):
         Get a single stat from the memsys
         """
         got = cuda.managed_array(1, np.uint64)
+        got_ptr = self._ctypes_pointer(got)
+
         self._single_thread_launch(
             self._memsys_module,
             stream,
             f"NRT_MemSys_read_{stat}",
-            [got.device_ctypes_pointer]
+            [got_ptr]
         )
 
         cuda.synchronize()
@@ -294,11 +311,13 @@ def set_memsys_to_module(self, module, stream=None):
             raise RuntimeError(
                 "Please allocate NRT Memsys first before setting to module.")
 
+        memsys_ptr = self._ctypes_pointer(self._memsys)
+
         self._single_thread_launch(
             module,
             stream,
             "NRT_MemSys_set",
-            [self._memsys.device_ctypes_pointer,]
+            [memsys_ptr]
         )
 
     @_alloc_init_guard

numba_cuda/numba/cuda/tests/nrt/test_nrt.py

@@ -229,6 +229,38 @@ def test_nrt_explicit_stats_query_raises_exception_when_disabled(self):
                     stats_func()
                 self.assertIn("NRT stats are disabled.", str(raises.exception))
 
+    def test_read_one_stat(self):
+        @cuda.jit
+        def foo():
+            tmp = np.ones(3)
+            arr = np.arange(5 * tmp[0]) # noqa: F841
+            return None
+
+        with (
+            override_config('CUDA_ENABLE_NRT', True),
+            override_config('CUDA_NRT_STATS', True)
+        ):
+
+            # Switch on stats
+            rtsys.memsys_enable_stats()
+
+            # Launch the kernel a couple of times to increase stats
+            foo[1, 1]()
+            foo[1, 1]()
+
+            # Get stats struct and individual stats
+            stats = rtsys.get_allocation_stats()
+            stats_alloc = rtsys.memsys_get_stats_alloc()
+            stats_mi_alloc = rtsys.memsys_get_stats_mi_alloc()
+            stats_free = rtsys.memsys_get_stats_free()
+            stats_mi_free = rtsys.memsys_get_stats_mi_free()
+
+            # Check individual stats match stats struct
+            self.assertEqual(stats.alloc, stats_alloc)
+            self.assertEqual(stats.mi_alloc, stats_mi_alloc)
+            self.assertEqual(stats.free, stats_free)
+            self.assertEqual(stats.mi_free, stats_mi_free)
+
 
 if __name__ == '__main__':
     unittest.main()