def __init__(self):
self.gpus = _DeviceList()
# A thread local stack
self.context_stack = servicelib.TLStack()
# Remember the main thread
# Only the main thread can *actually* destroy
self._mainthread = threading.current_thread()
# Avoid mutation of runtime state in multithreaded programs
self._lock = threading.RLock()
def __init__(self):
self.gpus = _DeviceList()
# A thread local stack
self.context_stack = servicelib.TLStack()
# Remembers all context
# key: context.handle.value
# value: Context
self._contexts = {}
# Device's context map
# key: Device.id
# value: Context
self._devctx = {}
# Remember the main thread
# Only the main thread can *actually* destroy
self._mainthread = threading.current_thread()
# Avoid mutation of runtime state in multithreaded programs
self._lock = threading.RLock()
def __exit__(self, exc_type, exc_val, exc_tb):
assert _get_device() is self
self.context.pop()
_gpustack.pop()
def reset(self):
if self._context:
self._context.reset()
self._context = None
def get_gpu(i):
return gpus[i]
_gpustack = servicelib.TLStack()
def _get_device(devnum=0):
"""Get the current device or use a device by device number.
"""
if not _gpustack:
_gpustack.push(get_gpu(devnum))
return _gpustack.top
def get_context(devnum=0):
"""Get the current device or use a device by device number, and
return the CUDA context.
"""
return _get_device(devnum=devnum).context