<commit_before>"""Device, context and memory management on PyCUDA."""

import atexit, copy_reg, os
import numpy
import pycuda.driver as drv
from pycuda.elementwise import ElementwiseKernel
from pycuda.reduction   import ReductionKernel
from pycuda import gpuarray
from pycuda import curandom
import pycuda.tools as cutools
from scikits.cuda import cublas
import scikits.cuda.misc as cumisc

# ------------------------------------------------------------------------------
# Basic types
# ------------------------------------------------------------------------------
from pycuda.driver   import Context, Device, Event, Stream
from pycuda.gpuarray import GPUArray

class IPCEvent(Event):
    def __init__(self):
        super(IPCEvent, self).__init__(
            drv.event_flags.INTERPROCESS | drv.event_flags.DISABLE_TIMING)

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
mem_alloc = None

_contexts = {}
_pools    = {}
_cublas_handles = {}
_pid      = None

_seed      = os.environ.get('CHAINER_SEED')
generator = None

def _seed_getter(N):
    if _seed is None:
        return curandom.seed_getter_uniform(N)
    return full((N,), _seed, numpy.int32)

def init(device=None):
    """Initialize CUDA global state.

    If you use chainer.cuda with multiprocessing, call init() for each process
    **after the fork** and before using chainer with PyCUDA.

    """
    global generator, _contexts, _pid

    pid = os.getpid()
    if _pid == pid:  # already initialized
        return

    drv.init()

    if device is None:  # use default device
        context = cutools.make_default_context()
        device  = Context.get_device()
    else:
        device  = Device(device)
        context = device.make_context()
    _contexts  = {device: context}

    generator = curandom.XORWOWRandomNumberGenerator(seed_getter=_seed_getter)

    _pid      = pid  # mark as initialized
    atexit.register(shutdown)

    cumisc.init(mem_alloc)
    _cublas_handles = {device: cumisc._global_cublas_handle}


def shutdown():
    """Finalize CUDA global state."""

    global _contexts, _cublas_handles, _pid, _pools

    pid = os.getpid()
    if _pid != pid:  # not initialized
        return

    for cublas_handle in _cublas_handles.itervalues():
        cublas.cublasDestroy(cublas_handle)
    _cublas_handles = {}

    cumisc.shutdown()

    _pools = {}

    for ctx in _contexts.itervalues():
        ctx.detach()
    _contexts = {}
    _pid      = None  # mark as uninitialized


def get_device(arg=None):
    """Get device from id or chainer array."""

    if arg is None:
        return Context.get_device()
    elif isinstance(arg, drv.Device):
        return arg
    elif isinstance(arg, GPUArray):
        return arg.gpudata.device
    return drv.Device(arg)


def use_device(arg, pop=True):
    """Switch context to use given device."""

    if pop:
        drv.Context.pop()

    device = get_device(arg)
    if device not in _contexts:
        _contexts[device] = device.make_context()
    else:
        _contexts[device].push()


class DeviceUser(object):
    """Device user for 'with' statement."""

    def __init__(self, arg):
        self.arg = arg

    def __enter__(self):
        use_device(self.arg, pop=False)

    def __exit__(self, typ, value, traceback):
        drv.Context.pop()

def using_device(arg):
    return DeviceUser(arg)


def mem_alloc(nbytes):
    """Allocate memory from memory pool corresponding to the current device."""
    global _pools<commit_msg>Add get_context<commit_after>"""Device, context and memory management on PyCUDA."""

import atexit, copy_reg, os
import numpy
import pycuda.driver as drv
from pycuda.elementwise import ElementwiseKernel
from pycuda.reduction   import ReductionKernel
from pycuda import gpuarray
from pycuda import curandom
import pycuda.tools as cutools
from scikits.cuda import cublas
import scikits.cuda.misc as cumisc

# ------------------------------------------------------------------------------
# Basic types
# ------------------------------------------------------------------------------
from pycuda.driver   import Context, Device, Event, Stream
from pycuda.gpuarray import GPUArray

class IPCEvent(Event):
    def __init__(self):
        super(IPCEvent, self).__init__(
            drv.event_flags.INTERPROCESS | drv.event_flags.DISABLE_TIMING)

# ------------------------------------------------------------------------------
# Global states
# ------------------------------------------------------------------------------
mem_alloc = None

_contexts = {}
_pools    = {}
_cublas_handles = {}
_pid      = None

_seed      = os.environ.get('CHAINER_SEED')
generator = None

def _seed_getter(N):
    if _seed is None:
        return curandom.seed_getter_uniform(N)
    return full((N,), _seed, numpy.int32)

def init(device=None):
    """Initialize CUDA global state.

    If you use chainer.cuda with multiprocessing, call init() for each process
    **after the fork** and before using chainer with PyCUDA.

    """
    global generator, _contexts, _pid

    pid = os.getpid()
    if _pid == pid:  # already initialized
        return

    drv.init()

    if device is None:  # use default device
        context = cutools.make_default_context()
        device  = Context.get_device()
    else:
        device  = Device(device)
        context = device.make_context()
    _contexts  = {device: context}

    generator = curandom.XORWOWRandomNumberGenerator(seed_getter=_seed_getter)

    _pid      = pid  # mark as initialized
    atexit.register(shutdown)

    cumisc.init(mem_alloc)
    _cublas_handles = {device: cumisc._global_cublas_handle}


def shutdown():
    """Finalize CUDA global state."""

    global _contexts, _cublas_handles, _pid, _pools

    pid = os.getpid()
    if _pid != pid:  # not initialized
        return

    for cublas_handle in _cublas_handles.itervalues():
        cublas.cublasDestroy(cublas_handle)
    _cublas_handles = {}

    cumisc.shutdown()

    _pools = {}

    for ctx in _contexts.itervalues():
        ctx.detach()
    _contexts = {}
    _pid      = None  # mark as uninitialized


def get_device(arg=None):
    """Get device from id or chainer array."""

    if arg is None:
        return Context.get_device()
    elif isinstance(arg, drv.Device):
        return arg
    elif isinstance(arg, GPUArray):
        return arg.gpudata.device
    return drv.Device(arg)


def use_device(arg, pop=True):
    """Switch context to use given device."""

    if pop:
        drv.Context.pop()

    device = get_device(arg)
    if device not in _contexts:
        _contexts[device] = device.make_context()
    else:
        _contexts[device].push()


class DeviceUser(object):
    """Device user for 'with' statement."""

    def __init__(self, arg):
        self.arg = arg

    def __enter__(self):
        use_device(self.arg, pop=False)

    def __exit__(self, typ, value, traceback):
        drv.Context.pop()

def using_device(arg):
    return DeviceUser(arg)


def get_context(device=None):
    """Get the context of given device.

    If device is not specified, it returns the current context.

    """
    return _contexts[get_device(device)]


def mem_alloc(nbytes):
    """Allocate memory from memory pool corresponding to the current device."""
    global _pools
{'fraction_non_alphanumeric': 0.17194875252865813, 'fraction_numerical': 0.0005394470667565745, 'mean_word_length': 2.217353579175705, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '603503', 'n_tokens_mistral': 2228, 'n_tokens_neox': 2164, 'n_words': 598}