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)
