add 3 more tests, augment API with single getters

Author: isVoid

numba_cuda/numba/cuda/dispatcher.py

@@ -365,7 +365,8 @@ def launch(self, args, griddim, blockdim, stream=0, sharedmem=0):
         rtsys.ensure_allocated(stream_handle)
         rtsys.set_memsys_to_module(cufunc.module, stream_handle)
         rtsys.ensure_initialized(stream_handle)
-        rtsys.memsys_enable_stats(stream_handle)
+        if config.CUDA_NRT_STATS:
+            rtsys.memsys_enable_stats(stream_handle)
 
         # Invoke kernel
         driver.launch_kernel(cufunc.handle,

numba_cuda/numba/cuda/runtime/memsys.cu

@@ -25,6 +25,30 @@ extern "C" __global__ void NRT_MemSys_read(uint64_t *managed_memsys)
   managed_memsys[3] = TheMSys->stats.mi_free;
 }
 
+extern "C" __global__ void NRT_MemSys_read_alloc(uint64_t *managed_memsys)
+{
+  detail::check_memsys();
+  managed_memsys[0] = TheMSys->stats.alloc;
+}
+
+extern "C" __global__ void NRT_MemSys_read_free(uint64_t *managed_memsys)
+{
+  detail::check_memsys();
+  managed_memsys[0] = TheMSys->stats.free;
+}
+
+extern "C" __global__ void NRT_MemSys_read_mi_alloc(uint64_t *managed_memsys)
+{
+  detail::check_memsys();
+  managed_memsys[0] = TheMSys->stats.mi_alloc;
+}
+
+extern "C" __global__ void NRT_MemSys_read_mi_free(uint64_t *managed_memsys)
+{
+  detail::check_memsys();
+  managed_memsys[0] = TheMSys->stats.mi_free;
+}
+
 extern "C" __global__ void NRT_MemSys_init(void)
 {
   detail::check_memsys();

numba_cuda/numba/cuda/runtime/nrt.py

@@ -163,6 +163,9 @@ def _copy_memsys_to_host(self, stream):
 
     @_alloc_init_guard
     def get_allocation_stats(self, stream=None):
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
         memsys = self._copy_memsys_to_host(stream)
         return _nrt_mstats(
             alloc=memsys["alloc"],
@@ -171,6 +174,51 @@ def get_allocation_stats(self, stream=None):
             mi_free=memsys["mi_free"]
         )
 
+    @_alloc_init_guard
+    def _get_single_stat(self, stat, stream=None):
+        got = cuda.managed_array(1, np.uint64)
+        self._single_thread_launch(
+            self._memsys_module,
+            stream,
+            f"NRT_MemSys_read_{stat}",
+            [got.device_ctypes_pointer]
+        )
+
+        cuda.synchronize()
+        return got[0]
+
+    @_alloc_init_guard
+    def memsys_get_stats_alloc(self, stream=None):
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("alloc")
+
+    @_alloc_init_guard
+    def memsys_get_stats_free(self, stream=None):
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("free")
+
+    @_alloc_init_guard
+    def memsys_get_stats_mi_alloc(self, stream=None):
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("mi_alloc")
+
+    @_alloc_init_guard
+    def memsys_get_stats_mi_free(self, stream=None):
+        enabled = self.memsys_stats_enabled(stream)
+        if not enabled:
+            raise RuntimeError("NRT stats are disabled.")
+
+        return self._get_single_stat("mi_free")
+
     def set_memsys_to_module(self, module, stream=None):
         if self._memsys is None:
             raise RuntimeError(
@@ -186,6 +234,7 @@ def set_memsys_to_module(self, module, stream=None):
 
     @_alloc_init_guard
     def print_memsys(self, stream=None):
+        """Print the current statistics of memsys, for debugging purpose."""
         cuda.synchronize()
         self._single_thread_launch(
             self._memsys_module,

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

@@ -6,7 +6,7 @@
 from unittest.mock import patch
 from numba.cuda.testing import CUDATestCase
 
-from numba.cuda.tests.nrt.mock_numpy import cuda_empty
+from numba.cuda.tests.nrt.mock_numpy import cuda_empty, cuda_ones, cuda_arange
 from numba.tests.support import run_in_subprocess
 
 from numba import cuda
@@ -160,6 +160,71 @@ def test_stats_env_var_default_off(self):
         env.pop('NUMBA_CUDA_NRT_STATS', None)
         self.check_env_var_off(env)
 
+    def test_stats_status_toggle(self):
+
+        @cuda.jit
+        def foo():
+            tmp = cuda_ones(3)
+            arr = cuda_arange(5 * tmp[0]) # noqa: F841
+            return None
+
+        # Switch on stats
+        rtsys.memsys_enable_stats()
+        # check the stats are on
+        self.assertTrue(rtsys.memsys_stats_enabled())
+
+        for i in range(2):
+            # capture the stats state
+            stats_1 = rtsys.get_allocation_stats()
+            # Switch off stats
+            rtsys.memsys_disable_stats()
+            # check the stats are off
+            self.assertFalse(rtsys.memsys_stats_enabled())
+            # run something that would move the counters were they enabled
+            with patch('numba.config.CUDA_ENABLE_NRT', True, create=True):
+                foo[1, 1]()
+            # Switch on stats
+            rtsys.memsys_enable_stats()
+            # check the stats are on
+            self.assertTrue(rtsys.memsys_stats_enabled())
+            # capture the stats state (should not have changed)
+            stats_2 = rtsys.get_allocation_stats()
+            # run something that will move the counters
+            with patch('numba.config.CUDA_ENABLE_NRT', True, create=True):
+                foo[1, 1]()
+            # capture the stats state (should have changed)
+            stats_3 = rtsys.get_allocation_stats()
+            # check stats_1 == stats_2
+            self.assertEqual(stats_1, stats_2)
+            # check stats_2 < stats_3
+            self.assertLess(stats_2, stats_3)
+
+    def test_rtsys_stats_query_raises_exception_when_disabled(self):
+        # Checks that the standard rtsys.get_allocation_stats() query raises
+        # when stats counters are turned off.
+
+        rtsys.memsys_disable_stats()
+        self.assertFalse(rtsys.memsys_stats_enabled())
+
+        with self.assertRaises(RuntimeError) as raises:
+            rtsys.get_allocation_stats()
+
+        self.assertIn("NRT stats are disabled.", str(raises.exception))
+
+    def test_nrt_explicit_stats_query_raises_exception_when_disabled(self):
+        # Checks the various memsys_get_stats functions raise if queried when
+        # the stats counters are disabled.
+        method_variations = ('alloc', 'free', 'mi_alloc', 'mi_free')
+        for meth in method_variations:
+            stats_func = getattr(rtsys, f'memsys_get_stats_{meth}')
+            with self.subTest(stats_func=stats_func):
+                # Turn stats off
+                rtsys.memsys_disable_stats()
+                self.assertFalse(rtsys.memsys_stats_enabled())
+                with self.assertRaises(RuntimeError) as raises:
+                    stats_func()
+                self.assertIn("NRT stats are disabled.", str(raises.exception))
+
 
 if __name__ == '__main__':
     unittest.main()