import numpy as np import pytest from numba import cuda import rmm import rmm._cuda.stream if sys.version_info < (3, 8): try: import pickle5 as pickle except ImportError: import pickle else: import pickle cuda.set_memory_manager(rmm.RMMNumbaManager) _driver_version = rmm._cuda.gpu.driverGetVersion() _runtime_version = rmm._cuda.gpu.runtimeGetVersion() _CUDAMALLOC_ASYNC_SUPPORTED = (_driver_version >= 11020) and (_runtime_version >= 11020) @pytest.fixture(scope="function", autouse=True) def rmm_auto_reinitialize(): # Run the test yield # Automatically reinitialize the current memory resource after running each # test
my_free(ptr) return finalizer # If NUMBA_CUDA_MEMORY_MANAGER is set to this module (e.g. # `NUMBA_CUDA_MEMORY_MANAGER=simple_emm_plugin`), then Numba will look at the # _numba_memory_manager global to determine what class to use for memory # management. # # This can be used to run the Numba test suite with the plugin, to verify that # the plugin is working correctly. For example, if the directory of this module # is on PYTHONPATH, then running: # # NUMBA_CUDA_MEMORY_MANAGER=simple_emm_plugin python -m numba.runtests \ # numba.cuda.tests # # will run all Numba CUDA tests with the plugin enabled. _numba_memory_manager = MyEMMPlugin if __name__ == '__main__': # Quick test of setting the memory manager and allocating/deleting an array cuda.set_memory_manager(MyEMMPlugin) ctx = cuda.current_context() print(f"Free before creating device array: {ctx.get_memory_info().free}") x = cuda.device_array(1000) print(f"Free after creating device array: {ctx.get_memory_info().free}") del x print(f"Free after freeing device array: {ctx.get_memory_info().free}")
def tearDown(self): # Set the memory manager back to the Numba internal one for subsequent # tests. cuda.close() cuda.set_memory_manager(cuda.cudadrv.driver.NumbaCUDAMemoryManager)
def test_bad_plugin_version(self): with self.assertRaises(RuntimeError) as raises: cuda.set_memory_manager(BadVersionEMMPlugin) self.assertIn("version 1 required", str(raises.exception))
def setUp(self): # Always start afresh with a new context and memory manager cuda.close() cuda.set_memory_manager(DeviceOnlyEMMPlugin)