files: from libcpp cimport vector
  from libcpp cimport unordered_map
  
  from cupy.cuda cimport device
  
  
  cdef class Chunk:
  
      cdef:
          readonly device.Device device
          readonly object mem
          readonly size_t ptr
          readonly Py_ssize_t offset
          readonly Py_ssize_t size
          public object stream_ptr
          public Chunk prev
          public Chunk next
  
  cdef class MemoryPointer:
  
      cdef:
          readonly device.Device device
          readonly object mem
          readonly size_t ptr
  
      cpdef copy_from_device(self, MemoryPointer src, Py_ssize_t size)
      cpdef copy_from_device_async(self, MemoryPointer src, size_t size,
                                   stream=?)
      cpdef copy_from_host(self, mem, size_t size)
      cpdef copy_from_host_async(self, mem, size_t size, stream=?)
      cpdef copy_from(self, mem, size_t size)
      cpdef copy_from_async(self, mem, size_t size, stream=?)
      cpdef copy_to_host(self, mem, size_t size)
      cpdef copy_to_host_async(self, mem, size_t size, stream=?)
      cpdef memset(self, int value, size_t size)
      cpdef memset_async(self, int value, size_t size, stream=?)
  
  
  cpdef MemoryPointer alloc(Py_ssize_t size)
  
  
  cpdef set_allocator(allocator=*)
  
  
  cdef class SingleDeviceMemoryPool:
  
      cdef:
          object _allocator
          dict _in_use
          dict _free
          object __weakref__
          object _weakref
          object _free_lock
          object _in_use_lock
          readonly Py_ssize_t _allocation_unit_size
          readonly int _device_id
          unordered_map.unordered_map[size_t, vector.vector[int]] _index
  
      cpdef MemoryPointer _alloc(self, Py_ssize_t size)
      cpdef MemoryPointer malloc(self, Py_ssize_t size)
      cpdef MemoryPointer _malloc(self, Py_ssize_t size)
      cpdef free(self, size_t ptr, Py_ssize_t size)
      cpdef free_all_blocks(self)
      cpdef free_all_free(self)
      cpdef n_free_blocks(self)
      cpdef used_bytes(self)
      cpdef free_bytes(self)
      cpdef total_bytes(self)
      cpdef Py_ssize_t _round_size(self, Py_ssize_t size)
      cpdef int _bin_index_from_size(self, Py_ssize_t size)
      cpdef list _arena(self, size_t stream_ptr)
      cdef vector.vector[int]* _arena_index(self, size_t stream_ptr)
      cpdef _append_to_free_list(self, Py_ssize_t size, chunk, size_t stream_ptr)
      cpdef bint _remove_from_free_list(self, Py_ssize_t size,
                                        chunk, size_t stream_ptr) except *
      cpdef tuple _split(self, Chunk chunk, Py_ssize_t size)
      cpdef Chunk _merge(self, Chunk head, Chunk remaining)
  
  cdef class MemoryPool:
  
      cdef:
          object _pools
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size)
      cpdef free_all_blocks(self)
      cpdef free_all_free(self)
      cpdef n_free_blocks(self)
      cpdef used_bytes(self)
      cpdef free_bytes(self)
      cpdef total_bytes(self)
  
files: cupy/cuda/memory.pxd
files: from libcpp cimport vector
  
  from cupy.cuda cimport device
  
  
  cdef class Chunk:
  
      cdef:
          readonly device.Device device
          readonly object mem
          readonly size_t ptr
          readonly Py_ssize_t offset
          readonly Py_ssize_t size
          public Chunk prev
          public Chunk next
  
  cdef class MemoryPointer:
  
      cdef:
          readonly device.Device device
          readonly object mem
          readonly size_t ptr
  
      cpdef copy_from_device(self, MemoryPointer src, Py_ssize_t size)
      cpdef copy_from_device_async(self, MemoryPointer src, size_t size, stream)
      cpdef copy_from_host(self, mem, size_t size)
      cpdef copy_from_host_async(self, mem, size_t size, stream)
      cpdef copy_from(self, mem, size_t size)
      cpdef copy_from_async(self, mem, size_t size, stream)
      cpdef copy_to_host(self, mem, size_t size)
      cpdef copy_to_host_async(self, mem, size_t size, stream)
      cpdef memset(self, int value, size_t size)
      cpdef memset_async(self, int value, size_t size, stream)
  
  
  cpdef MemoryPointer alloc(Py_ssize_t size)
  
  
  cpdef set_allocator(allocator=*)
  
  
  cdef class SingleDeviceMemoryPool:
  
      cdef:
          object _allocator
          dict _in_use
          list _free
          object __weakref__
          object _weakref
          object _free_lock
          object _in_use_lock
          readonly Py_ssize_t _allocation_unit_size
          readonly int _device_id
          vector.vector[int] _index
  
      cpdef MemoryPointer _alloc(self, Py_ssize_t size)
      cpdef MemoryPointer malloc(self, Py_ssize_t size)
      cpdef MemoryPointer _malloc(self, Py_ssize_t size)
      cpdef free(self, size_t ptr, Py_ssize_t size)
      cpdef free_all_blocks(self)
      cpdef free_all_free(self)
      cpdef n_free_blocks(self)
      cpdef used_bytes(self)
      cpdef free_bytes(self)
      cpdef total_bytes(self)
      cpdef Py_ssize_t _round_size(self, Py_ssize_t size)
      cpdef int _bin_index_from_size(self, Py_ssize_t size)
      cpdef _append_to_free_list(self, Py_ssize_t size, chunk)
      cpdef bint _remove_from_free_list(self, Py_ssize_t size, chunk) except *
      cpdef tuple _split(self, Chunk chunk, Py_ssize_t size)
      cpdef Chunk _merge(self, Chunk head, Chunk remaining)
  
  cdef class MemoryPool:
  
      cdef:
          object _pools
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size)
      cpdef free_all_blocks(self)
      cpdef free_all_free(self)
      cpdef n_free_blocks(self)
      cpdef used_bytes(self)
      cpdef free_bytes(self)
      cpdef total_bytes(self)
  
files: cupy/cuda/memory.pxd
files: # distutils: language = c++
  import collections
  import ctypes
  import gc
  import warnings
  import weakref
  
  from fastrlock cimport rlock
  from libcpp cimport algorithm
  
  from cupy.cuda import memory_hook
  from cupy.cuda import runtime
  from cupy.cuda import stream as stream_module
  
  from cupy.cuda cimport device
  from cupy.cuda cimport runtime
  
  
  class OutOfMemoryError(MemoryError):
  
      def __init__(self, size, total):
          msg = 'out of memory to allocate %d bytes ' \
                '(total %d bytes)' % (size, total)
          super(OutOfMemoryError, self).__init__(msg)
  
  
  class Memory(object):
  
      """Memory allocation on a CUDA device.
  
      This class provides an RAII interface of the CUDA memory allocation.
  
      Args:
          ~Memory.device (cupy.cuda.Device): Device whose memory the pointer
              refers to.
          ~Memory.size (int): Size of the memory allocation in bytes.
  
      """
  
      def __init__(self, Py_ssize_t size):
          self.size = size
          self.device = None
          self.ptr = 0
          if size > 0:
              self.device = device.Device()
              self.ptr = runtime.malloc(size)
  
      def __del__(self):
          if self.ptr:
              runtime.free(self.ptr)
  
      def __int__(self):
          """Returns the pointer value to the head of the allocation."""
          return self.ptr
  
  
  class ManagedMemory(Memory):
  
      """Managed memory (Unified memory) allocation on a CUDA device.
  
      This class provides an RAII interface of the CUDA managed memory
      allocation.
  
      Args:
          device (cupy.cuda.Device): Device whose memory the pointer refers to.
          size (int): Size of the memory allocation in bytes.
  
      """
  
      def __init__(self, Py_ssize_t size):
          self.size = size
          self.device = None
          self.ptr = 0
          if size > 0:
              self.device = device.Device()
              self.ptr = runtime.mallocManaged(size)
  
      def prefetch(self, stream):
          """(experimental) Prefetch memory.
  
          Args:
              stream (cupy.cuda.Stream): CUDA stream.
          """
          runtime.memPrefetchAsync(self.ptr, self.size, self.device.id,
                                   stream.ptr)
  
      def advise(self, int advise, device.Device device):
          """(experimental) Advise about the usage of this memory.
  
          Args:
              advics (int): Advise to be applied for this memory.
              device (cupy.cuda.Device): Device to apply the advice for.
  
          """
          runtime.memAdvise(self.ptr, self.size, advise, device.id)
  
  
  cdef set _peer_access_checked = set()
  
  
  cpdef _set_peer_access(int device, int peer):
      device_pair = device, peer
  
      if device_pair in _peer_access_checked:
          return
      cdef int can_access = runtime.deviceCanAccessPeer(device, peer)
      _peer_access_checked.add(device_pair)
      if not can_access:
          return
  
      cdef int current = runtime.getDevice()
      runtime.setDevice(device)
      try:
          runtime.deviceEnablePeerAccess(peer)
      finally:
          runtime.setDevice(current)
  
  cdef class Chunk:
  
      """A chunk points to a device memory.
  
      A chunk might be a splitted memory block from a larger allocation.
      The prev/next pointers contruct a doubly-linked list of memory addresses
      sorted by base address that must be contiguous.
  
      Args:
          mem (Memory): The device memory buffer.
          offset (int): An offset bytes from the head of the buffer.
          size (int): Chunk size in bytes.
          stream_ptr (size_t): Raw stream handle of cupy.cuda.Stream
  
      Attributes:
          device (cupy.cuda.Device): Device whose memory the pointer refers to.
          mem (Memory): The device memory buffer.
          ptr (size_t): Memory address.
          offset (int): An offset bytes from the head of the buffer.
          size (int): Chunk size in bytes.
          prev (Chunk): prev memory pointer if split from a larger allocation
          next (Chunk): next memory pointer if split from a larger allocation
          stream_ptr (size_t): Raw stream handle of cupy.cuda.Stream
      """
  
      def __init__(self, mem, Py_ssize_t offset, Py_ssize_t size, stream_ptr):
          assert mem.ptr > 0 or offset == 0
          self.mem = mem
          self.device = mem.device
          self.ptr = mem.ptr + offset
          self.offset = offset
          self.size = size
          self.stream_ptr = stream_ptr
          self.prev = None
          self.next = None
  
  cdef class MemoryPointer:
  
      """Pointer to a point on a device memory.
  
      An instance of this class holds a reference to the original memory buffer
      and a pointer to a place within this buffer.
  
      Args:
          mem (Memory): The device memory buffer.
          offset (int): An offset from the head of the buffer to the place this
              pointer refers.
  
      Attributes:
          ~MemoryPointer.device (cupy.cuda.Device): Device whose memory the
              pointer refers to.
          mem (Memory): The device memory buffer.
          ptr (size_t): Pointer to the place within the buffer.
      """
  
      def __init__(self, mem, Py_ssize_t offset):
          assert mem.ptr > 0 or offset == 0
          self.mem = mem
          self.device = mem.device
          self.ptr = mem.ptr + offset
  
      def __int__(self):
          """Returns the pointer value."""
          return self.ptr
  
      def __add__(x, y):
          """Adds an offset to the pointer."""
          cdef MemoryPointer self
          cdef Py_ssize_t offset
          if isinstance(x, MemoryPointer):
              self = x
              offset = <Py_ssize_t?>y
          else:
              self = <MemoryPointer?>y
              offset = <Py_ssize_t?>x
          assert self.ptr > 0 or offset == 0
          return MemoryPointer(self.mem,
                               self.ptr - self.mem.ptr + offset)
  
      def __iadd__(self, Py_ssize_t offset):
          """Adds an offset to the pointer in place."""
          assert self.ptr > 0 or offset == 0
          self.ptr += offset
          return self
  
      def __sub__(self, offset):
          """Subtracts an offset from the pointer."""
          return self + -offset
  
      def __isub__(self, Py_ssize_t offset):
          """Subtracts an offset from the pointer in place."""
          return self.__iadd__(-offset)
  
      cpdef copy_from_device(self, MemoryPointer src, Py_ssize_t size):
          """Copies a memory sequence from a (possibly different) device.
  
          Args:
              src (cupy.cuda.MemoryPointer): Source memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              _set_peer_access(src.device.id, self.device.id)
              runtime.memcpy(self.ptr, src.ptr, size,
                             runtime.memcpyDefault)
  
      cpdef copy_from_device_async(self, MemoryPointer src, size_t size,
                                   stream=None):
          """Copies a memory from a (possibly different) device asynchronously.
  
          Args:
              src (cupy.cuda.MemoryPointer): Source memory pointer.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
                  The default uses CUDA stream of the current context.
  
          """
          if stream is None:
              stream = stream_module.get_current_stream()
          if size > 0:
              _set_peer_access(src.device.id, self.device.id)
              runtime.memcpyAsync(self.ptr, src.ptr, size,
                                  runtime.memcpyDefault, stream.ptr)
  
      cpdef copy_from_host(self, mem, size_t size):
          """Copies a memory sequence from the host memory.
  
          Args:
              mem (ctypes.c_void_p): Source memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memcpy(self.ptr, mem.value, size,
                             runtime.memcpyHostToDevice)
  
      cpdef copy_from_host_async(self, mem, size_t size, stream=None):
          """Copies a memory sequence from the host memory asynchronously.
  
          Args:
              mem (ctypes.c_void_p): Source memory pointer. It must be a pinned
                  memory.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
                  The default uses CUDA stream of the current context.
  
          """
          if stream is None:
              stream = stream_module.get_current_stream()
          if size > 0:
              runtime.memcpyAsync(self.ptr, mem.value, size,
                                  runtime.memcpyHostToDevice, stream.ptr)
  
      cpdef copy_from(self, mem, size_t size):
          """Copies a memory sequence from a (possibly different) device or host.
  
          This function is a useful interface that selects appropriate one from
          :meth:`~cupy.cuda.MemoryPointer.copy_from_device` and
          :meth:`~cupy.cuda.MemoryPointer.copy_from_host`.
  
          Args:
              mem (ctypes.c_void_p or cupy.cuda.MemoryPointer): Source memory
                  pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if isinstance(mem, MemoryPointer):
              self.copy_from_device(mem, size)
          else:
              self.copy_from_host(mem, size)
  
      cpdef copy_from_async(self, mem, size_t size, stream=None):
          """Copies a memory sequence from an arbitrary place asynchronously.
  
          This function is a useful interface that selects appropriate one from
          :meth:`~cupy.cuda.MemoryPointer.copy_from_device_async` and
          :meth:`~cupy.cuda.MemoryPointer.copy_from_host_async`.
  
          Args:
              mem (ctypes.c_void_p or cupy.cuda.MemoryPointer): Source memory
                  pointer.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
                  The default uses CUDA stream of the current context.
  
          """
          if stream is None:
              stream = stream_module.get_current_stream()
          if isinstance(mem, MemoryPointer):
              self.copy_from_device_async(mem, size, stream)
          else:
              self.copy_from_host_async(mem, size, stream)
  
      cpdef copy_to_host(self, mem, size_t size):
          """Copies a memory sequence to the host memory.
  
          Args:
              mem (ctypes.c_void_p): Target memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memcpy(mem.value, self.ptr, size,
                             runtime.memcpyDeviceToHost)
  
      cpdef copy_to_host_async(self, mem, size_t size, stream=None):
          """Copies a memory sequence to the host memory asynchronously.
  
          Args:
              mem (ctypes.c_void_p): Target memory pointer. It must be a pinned
                  memory.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
                  The default uses CUDA stream of the current context.
  
          """
          if stream is None:
              stream = stream_module.get_current_stream()
          if size > 0:
              runtime.memcpyAsync(mem.value, self.ptr, size,
                                  runtime.memcpyDeviceToHost, stream.ptr)
  
      cpdef memset(self, int value, size_t size):
          """Fills a memory sequence by constant byte value.
  
          Args:
              value (int): Value to fill.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memset(self.ptr, value, size)
  
      cpdef memset_async(self, int value, size_t size, stream=None):
          """Fills a memory sequence by constant byte value asynchronously.
  
          Args:
              value (int): Value to fill.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
                  The default uses CUDA stream of the current context.
  
          """
          if stream is None:
              stream = stream_module.get_current_stream()
          if size > 0:
              runtime.memsetAsync(self.ptr, value, size, stream.ptr)
  
  
  cpdef MemoryPointer _malloc(Py_ssize_t size):
      mem = Memory(size)
      return MemoryPointer(mem, 0)
  
  
  cpdef MemoryPointer malloc_managed(Py_ssize_t size):
      """Allocate managed memory (unified memory).
  
      This method can be used as a CuPy memory allocator. The simplest way to
      use a managed memory as the default allocator is the following code::
  
          set_allocator(malloc_managed)
  
      The advantage using managed memory in CuPy is that device memory
      oversubscription is possible for GPUs that have a non-zero value for the
      device attribute cudaDevAttrConcurrentManagedAccess.
      CUDA >= 8.0 with GPUs later than or equal to Pascal is preferrable.
  
      Read more at: http://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#axzz4qygc1Ry1  # NOQA
  
      Args:
          size (int): Size of the memory allocation in bytes.
  
      Returns:
          ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
      """
      mem = ManagedMemory(size)
      return MemoryPointer(mem, 0)
  
  
  cdef object _current_allocator = _malloc
  
  
  cpdef MemoryPointer alloc(Py_ssize_t size):
      """Calls the current allocator.
  
      Use :func:`~cupy.cuda.set_allocator` to change the current allocator.
  
      Args:
          size (int): Size of the memory allocation.
  
      Returns:
          ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
  
      """
      return _current_allocator(size)
  
  
  cpdef set_allocator(allocator=_malloc):
      """Sets the current allocator.
  
      Args:
          allocator (function): CuPy memory allocator. It must have the same
              interface as the :func:`cupy.cuda.alloc` function, which takes the
              buffer size as an argument and returns the device buffer of that
              size.
  
      """
      global _current_allocator
      _current_allocator = allocator
  
  
  class PooledMemory(Memory):
  
      """Memory allocation for a memory pool.
  
      The instance of this class is created by memory pool allocator, so user
      should not instantiate it by hand.
  
      """
  
      def __init__(self, Chunk chunk, pool):
          self.device = chunk.device
          self.ptr = chunk.ptr
          self.size = chunk.size
          self.pool = pool
  
      def free(self):
          """Frees the memory buffer and returns it to the memory pool.
  
          This function actually does not free the buffer. It just returns the
          buffer to the memory pool for reuse.
  
          """
          cdef Py_ssize_t ptr
          ptr = self.ptr
          if ptr == 0:
              return
          self.ptr = 0
          pool = self.pool()
          if pool is None:
              return
  
          hooks = None
          # to avoid error at exit
          mh = memory_hook
          if mh is not None and mh.get_memory_hooks is not None:
              hooks = mh.get_memory_hooks()
          size = self.size
          if hooks:
              device_id = self.device.id
              pmem_id = id(self)
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.free_preprocess(device_id=device_id,
                                       mem_size=size,
                                       mem_ptr=ptr,
                                       pmem_id=pmem_id)
              try:
                  pool.free(ptr, size)
              finally:
                  for hook in hooks_values:
                      hook.free_postprocess(device_id=device_id,
                                            mem_size=size,
                                            mem_ptr=ptr,
                                            pmem_id=pmem_id)
          else:
              pool.free(ptr, size)
  
      __del__ = free
  
  
  cdef class SingleDeviceMemoryPool:
      """Memory pool implementation for single device.
  
      - The allocator attempts to find the smallest cached block that will fit
        the requested size. If the block is larger than the requested size,
        it may be split. If no block is found, the allocator will delegate to
        cudaMalloc.
      - If the cudaMalloc fails, the allocator will free all cached blocks that
        are not split and retry the allocation.
      """
  
      def __init__(self, allocator=_malloc):
          # cudaMalloc() is aligned to at least 512 bytes
          # cf. https://gist.github.com/sonots/41daaa6432b1c8b27ef782cd14064269
          self._allocation_unit_size = 512
          self._in_use = {}
          self._free = {}
          self._allocator = allocator
          self._weakref = weakref.ref(self)
          self._device_id = device.get_device_id()
          self._free_lock = rlock.create_fastrlock()
          self._in_use_lock = rlock.create_fastrlock()
  
      cpdef Py_ssize_t _round_size(self, Py_ssize_t size):
          """Round up the memory size to fit memory alignment of cudaMalloc."""
          unit = self._allocation_unit_size
          return ((size + unit - 1) // unit) * unit
  
      cpdef int _bin_index_from_size(self, Py_ssize_t size):
          """Get appropriate bins index from the memory size"""
          unit = self._allocation_unit_size
          return (size - 1) // unit
  
      cpdef list _arena(self, size_t stream_ptr):
          """Get appropriate arena (list of bins) of a given stream"""
          if stream_ptr not in self._free:
              self._free[stream_ptr] = []
          return self._free[stream_ptr]
  
      cdef vector.vector[int]* _arena_index(self, size_t stream_ptr):
          """Get appropriate arena sparse index of a given stream"""
          if self._index.count(stream_ptr) == 0:
              self._index[stream_ptr] = vector.vector[int]()
          return &self._index[stream_ptr]
  
      cpdef _append_to_free_list(self, Py_ssize_t size, chunk,
                                 size_t stream_ptr):
          cdef int index, bin_index
          cdef list arena
          cdef set free_list
          cdef vector.vector[int]* arena_index
  
          bin_index = self._bin_index_from_size(size)
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              arena = self._arena(stream_ptr)
              arena_index = self._arena_index(stream_ptr)
              index = algorithm.lower_bound(
                  arena_index.begin(), arena_index.end(),
                  bin_index) - arena_index.begin()
              size = <int>arena_index.size()
              if index < size and arena_index.at(index) == bin_index:
                  free_list = arena[index]
              else:
                  free_list = set()
                  arena_index.insert(arena_index.begin() + index, bin_index)
                  arena.insert(index, free_list)
              free_list.add(chunk)
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
      cpdef bint _remove_from_free_list(self, Py_ssize_t size, chunk,
                                        size_t stream_ptr) except *:
          cdef int index, bin_index
          cdef list arena
          cdef set free_list
          cdef vector.vector[int]* arena_index
  
          bin_index = self._bin_index_from_size(size)
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              arena = self._arena(stream_ptr)
              arena_index = self._arena_index(stream_ptr)
              index = algorithm.lower_bound(
                  arena_index.begin(), arena_index.end(),
                  bin_index) - arena_index.begin()
              if arena_index.at(index) != bin_index:
                  return False
              free_list = arena[index]
              if chunk in free_list:
                  free_list.remove(chunk)
                  return True
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return False
  
      cpdef tuple _split(self, Chunk chunk, Py_ssize_t size):
          """Split contiguous block of a larger allocation"""
          cdef Chunk head
          cdef Chunk remaining
          assert chunk.size >= size
          if chunk.size == size:
              return chunk, None
          head = Chunk(chunk.mem, chunk.offset, size, chunk.stream_ptr)
          remaining = Chunk(chunk.mem, chunk.offset + size, chunk.size - size,
                            chunk.stream_ptr)
          if chunk.prev is not None:
              head.prev = chunk.prev
              chunk.prev.next = head
          if chunk.next is not None:
              remaining.next = chunk.next
              chunk.next.prev = remaining
          head.next = remaining
          remaining.prev = head
          return head, remaining
  
      cpdef Chunk _merge(self, Chunk head, Chunk remaining):
          """Merge previously splitted block (chunk)"""
          assert head.stream_ptr == remaining.stream_ptr
          cdef Chunk merged
          size = head.size + remaining.size
          merged = Chunk(head.mem, head.offset, size, head.stream_ptr)
          if head.prev is not None:
              merged.prev = head.prev
              merged.prev.next = merged
          if remaining.next is not None:
              merged.next = remaining.next
              merged.next.prev = merged
          return merged
  
      cpdef MemoryPointer _alloc(self, Py_ssize_t rounded_size):
          hooks = memory_hook.get_memory_hooks()
          if hooks:
              memptr = None
              device_id = self._device_id
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.alloc_preprocess(device_id=device_id,
                                        mem_size=rounded_size)
              try:
                  memptr = self._allocator(rounded_size)
              finally:
                  for hook in hooks_values:
                      mem_ptr = memptr.ptr if memptr is not None else 0
                      hook.alloc_postprocess(device_id=device_id,
                                             mem_size=rounded_size,
                                             mem_ptr=mem_ptr)
              return memptr
          else:
              return self._allocator(rounded_size)
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size):
          rounded_size = self._round_size(size)
          hooks = memory_hook.get_memory_hooks()
          if hooks:
              memptr = None
              device_id = self._device_id
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.malloc_preprocess(device_id=device_id,
                                         size=size,
                                         mem_size=rounded_size)
              try:
                  memptr = self._malloc(rounded_size)
              finally:
                  if memptr is None:
                      mem_ptr = 0
                      pmem_id = 0
                  else:
                      mem_ptr = memptr.ptr
                      pmem_id = id(memptr.mem)
                  for hook in hooks_values:
                      hook.malloc_postprocess(device_id=device_id,
                                              size=size,
                                              mem_size=rounded_size,
                                              mem_ptr=mem_ptr,
                                              pmem_id=pmem_id)
              return memptr
          else:
              return self._malloc(rounded_size)
  
      cpdef MemoryPointer _malloc(self, Py_ssize_t size):
          cdef set free_list
          cdef Chunk chunk = None
          cdef Chunk remaining = None
          cdef int bin_index, index, length
  
          if size == 0:
              return MemoryPointer(Memory(0), 0)
  
          stream_ptr = stream_module.get_current_stream().ptr
  
          bin_index = self._bin_index_from_size(size)
          # find best-fit, or a smallest larger allocation
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              arena = self._arena(stream_ptr)
              arena_index = self._arena_index(stream_ptr)
              index = algorithm.lower_bound(
                  arena_index.begin(), arena_index.end(),
                  bin_index) - arena_index.begin()
              length = arena_index.size()
              for i in range(index, length):
                  free_list = arena[i]
                  if free_list:
                      chunk = free_list.pop()
                      break
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
          if chunk is not None:
              chunk, remaining = self._split(chunk, size)
          else:
              # cudaMalloc if a cache is not found
              try:
                  mem = self._alloc(size).mem
              except runtime.CUDARuntimeError as e:
                  if e.status != runtime.errorMemoryAllocation:
                      raise
                  self.free_all_blocks()
                  try:
                      mem = self._alloc(size).mem
                  except runtime.CUDARuntimeError as e:
                      if e.status != runtime.errorMemoryAllocation:
                          raise
                      gc.collect()
                      try:
                          mem = self._alloc(size).mem
                      except runtime.CUDARuntimeError as e:
                          if e.status != runtime.errorMemoryAllocation:
                              raise
                          else:
                              total = size + self.total_bytes()
                              raise OutOfMemoryError(size, total)
              chunk = Chunk(mem, 0, size, stream_ptr)
  
          assert chunk.stream_ptr == stream_ptr
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              self._in_use[chunk.ptr] = chunk
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          if remaining is not None:
              self._append_to_free_list(remaining.size, remaining, stream_ptr)
          pmem = PooledMemory(chunk, self._weakref)
          return MemoryPointer(pmem, 0)
  
      cpdef free(self, size_t ptr, Py_ssize_t size):
          cdef set free_list
          cdef Chunk chunk
  
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              chunk = self._in_use.pop(ptr, None)
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          if chunk is None:
              raise RuntimeError('Cannot free out-of-pool memory')
          stream_ptr = chunk.stream_ptr
  
          if chunk.next is not None:
              if self._remove_from_free_list(chunk.next.size, chunk.next,
                                             stream_ptr):
                  chunk = self._merge(chunk, chunk.next)
  
          if chunk.prev is not None:
              if self._remove_from_free_list(chunk.prev.size, chunk.prev,
                                             stream_ptr):
                  chunk = self._merge(chunk.prev, chunk)
  
          self._append_to_free_list(chunk.size, chunk, stream_ptr)
  
      cpdef free_all_blocks(self):
          cdef set free_list, keep_list
          cdef Chunk chunk
          # Free all **non-split** chunks
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for arena in self._free.itervalues():
                  for i in range(len(arena)):
                      free_list = arena[i]
                      keep_list = set()
                      for chunk in free_list:
                          if chunk.prev is not None or chunk.next is not None:
                              keep_list.add(chunk)
                      arena[i] = keep_list
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
      cpdef free_all_free(self):
          warnings.warn(
              'free_all_free is deprecated. Use free_all_blocks instead.',
              DeprecationWarning)
          self.free_all_blocks()
  
      cpdef n_free_blocks(self):
          cdef Py_ssize_t n = 0
          cdef set free_list
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for arena in self._free.itervalues():
                  for v in arena:
                      n += len(v)
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return n
  
      cpdef used_bytes(self):
          cdef Py_ssize_t size = 0
          cdef Chunk chunk
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              for chunk in self._in_use.itervalues():
                  size += chunk.size
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          return size
  
      cpdef free_bytes(self):
          cdef Py_ssize_t size = 0
          cdef set free_list
          cdef Chunk chunk
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for arena in self._free.itervalues():
                  for free_list in arena:
                      for chunk in free_list:
                          size += chunk.size
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return size
  
      cpdef total_bytes(self):
          return self.used_bytes() + self.free_bytes()
  
  
  cdef class MemoryPool(object):
  
      """Memory pool for all devices on the machine.
  
      A memory pool preserves any allocations even if they are freed by the user.
      Freed memory buffers are held by the memory pool as *free blocks*, and they
      are reused for further memory allocations of the same sizes. The allocated
      blocks are managed for each device, so one instance of this class can be
      used for multiple devices.
  
      .. note::
         When the allocation is skipped by reusing the pre-allocated block, it
         does not call ``cudaMalloc`` and therefore CPU-GPU synchronization does
         not occur. It makes interleaves of memory allocations and kernel
         invocations very fast.
  
      .. note::
         The memory pool holds allocated blocks without freeing as much as
         possible. It makes the program hold most of the device memory, which may
         make other CUDA programs running in parallel out-of-memory situation.
  
      Args:
          allocator (function): The base CuPy memory allocator. It is used for
              allocating new blocks when the blocks of the required size are all
              in use.
  
      """
  
      def __init__(self, allocator=_malloc):
          self._pools = collections.defaultdict(
              lambda: SingleDeviceMemoryPool(allocator))
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size):
          """Allocates the memory, from the pool if possible.
  
          This method can be used as a CuPy memory allocator. The simplest way to
          use a memory pool as the default allocator is the following code::
  
              set_allocator(MemoryPool().malloc)
  
          Also, the way to use a memory pool of Managed memory (Unified memory)
          as the default allocator is the following code::
  
              set_allocator(MemoryPool(malloc_managed).malloc)
  
          Args:
              size (int): Size of the memory buffer to allocate in bytes.
  
          Returns:
              ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
  
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.malloc(size)
  
      cpdef free_all_blocks(self):
          """Release free blocks."""
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          mp.free_all_blocks()
  
      cpdef free_all_free(self):
          """Release free blocks."""
          warnings.warn(
              'free_all_free is deprecated. Use free_all_blocks instead.',
              DeprecationWarning)
          self.free_all_blocks()
  
      cpdef n_free_blocks(self):
          """Count the total number of free blocks.
  
          Returns:
              int: The total number of free blocks.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.n_free_blocks()
  
      cpdef used_bytes(self):
          """Get the total number of bytes used.
  
          Returns:
              int: The total number of bytes used.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.used_bytes()
  
      cpdef free_bytes(self):
          """Get the total number of bytes acquired but not used in the pool.
  
          Returns:
              int: The total number of bytes acquired but not used in the pool.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.free_bytes()
  
      cpdef total_bytes(self):
          """Get the total number of bytes acquired in the pool.
  
          Returns:
              int: The total number of bytes acquired in the pool.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.total_bytes()
  
files: cupy/cuda/memory.pyx
files: # distutils: language = c++
  import collections
  import ctypes
  import gc
  import warnings
  import weakref
  
  from fastrlock cimport rlock
  from libcpp cimport algorithm
  
  from cupy.cuda import memory_hook
  from cupy.cuda import runtime
  
  from cupy.cuda cimport device
  from cupy.cuda cimport runtime
  
  
  class OutOfMemoryError(MemoryError):
  
      def __init__(self, size, total):
          msg = 'out of memory to allocate %d bytes ' \
                '(total %d bytes)' % (size, total)
          super(OutOfMemoryError, self).__init__(msg)
  
  
  class Memory(object):
  
      """Memory allocation on a CUDA device.
  
      This class provides an RAII interface of the CUDA memory allocation.
  
      Args:
          ~Memory.device (cupy.cuda.Device): Device whose memory the pointer
              refers to.
          ~Memory.size (int): Size of the memory allocation in bytes.
  
      """
  
      def __init__(self, Py_ssize_t size):
          self.size = size
          self.device = None
          self.ptr = 0
          if size > 0:
              self.device = device.Device()
              self.ptr = runtime.malloc(size)
  
      def __del__(self):
          if self.ptr:
              runtime.free(self.ptr)
  
      def __int__(self):
          """Returns the pointer value to the head of the allocation."""
          return self.ptr
  
  
  class ManagedMemory(Memory):
  
      """Managed memory (Unified memory) allocation on a CUDA device.
  
      This class provides an RAII interface of the CUDA managed memory
      allocation.
  
      Args:
          device (cupy.cuda.Device): Device whose memory the pointer refers to.
          size (int): Size of the memory allocation in bytes.
  
      """
  
      def __init__(self, Py_ssize_t size):
          self.size = size
          self.device = None
          self.ptr = 0
          if size > 0:
              self.device = device.Device()
              self.ptr = runtime.mallocManaged(size)
  
      def prefetch(self, stream):
          """(experimental) Prefetch memory.
  
          Args:
              stream (cupy.cuda.Stream): CUDA stream.
          """
          runtime.memPrefetchAsync(self.ptr, self.size, self.device.id,
                                   stream.ptr)
  
      def advise(self, int advise, device.Device device):
          """(experimental) Advise about the usage of this memory.
  
          Args:
              advics (int): Advise to be applied for this memory.
              device (cupy.cuda.Device): Device to apply the advice for.
  
          """
          runtime.memAdvise(self.ptr, self.size, advise, device.id)
  
  
  cdef set _peer_access_checked = set()
  
  
  cpdef _set_peer_access(int device, int peer):
      device_pair = device, peer
  
      if device_pair in _peer_access_checked:
          return
      cdef int can_access = runtime.deviceCanAccessPeer(device, peer)
      _peer_access_checked.add(device_pair)
      if not can_access:
          return
  
      cdef int current = runtime.getDevice()
      runtime.setDevice(device)
      try:
          runtime.deviceEnablePeerAccess(peer)
      finally:
          runtime.setDevice(current)
  
  cdef class Chunk:
  
      """A chunk points to a device memory.
  
      A chunk might be a splitted memory block from a larger allocation.
      The prev/next pointers contruct a doubly-linked list of memory addresses
      sorted by base address that must be contiguous.
  
      Args:
          mem (Memory): The device memory buffer.
          offset (int): An offset bytes from the head of the buffer.
          size (int): Chunk size in bytes.
  
      Attributes:
          device (cupy.cuda.Device): Device whose memory the pointer refers to.
          mem (Memory): The device memory buffer.
          ptr (int): Memory address.
          offset (int): An offset bytes from the head of the buffer.
          size (int): Chunk size in bytes.
          prev (Chunk): prev memory pointer if split from a larger allocation
          next (Chunk): next memory pointer if split from a larger allocation
      """
  
      def __init__(self, mem, Py_ssize_t offset, Py_ssize_t size):
          assert mem.ptr > 0 or offset == 0
          self.mem = mem
          self.device = mem.device
          self.ptr = mem.ptr + offset
          self.offset = offset
          self.size = size
          self.prev = None
          self.next = None
  
  cdef class MemoryPointer:
  
      """Pointer to a point on a device memory.
  
      An instance of this class holds a reference to the original memory buffer
      and a pointer to a place within this buffer.
  
      Args:
          mem (Memory): The device memory buffer.
          offset (int): An offset from the head of the buffer to the place this
              pointer refers.
  
      Attributes:
          ~MemoryPointer.device (cupy.cuda.Device): Device whose memory the
              pointer refers to.
          mem (Memory): The device memory buffer.
          ptr (int): Pointer to the place within the buffer.
      """
  
      def __init__(self, mem, Py_ssize_t offset):
          assert mem.ptr > 0 or offset == 0
          self.mem = mem
          self.device = mem.device
          self.ptr = mem.ptr + offset
  
      def __int__(self):
          """Returns the pointer value."""
          return self.ptr
  
      def __add__(x, y):
          """Adds an offset to the pointer."""
          cdef MemoryPointer self
          cdef Py_ssize_t offset
          if isinstance(x, MemoryPointer):
              self = x
              offset = <Py_ssize_t?>y
          else:
              self = <MemoryPointer?>y
              offset = <Py_ssize_t?>x
          assert self.ptr > 0 or offset == 0
          return MemoryPointer(self.mem,
                               self.ptr - self.mem.ptr + offset)
  
      def __iadd__(self, Py_ssize_t offset):
          """Adds an offset to the pointer in place."""
          assert self.ptr > 0 or offset == 0
          self.ptr += offset
          return self
  
      def __sub__(self, offset):
          """Subtracts an offset from the pointer."""
          return self + -offset
  
      def __isub__(self, Py_ssize_t offset):
          """Subtracts an offset from the pointer in place."""
          return self.__iadd__(-offset)
  
      cpdef copy_from_device(self, MemoryPointer src, Py_ssize_t size):
          """Copies a memory sequence from a (possibly different) device.
  
          Args:
              src (cupy.cuda.MemoryPointer): Source memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              _set_peer_access(src.device.id, self.device.id)
              runtime.memcpy(self.ptr, src.ptr, size,
                             runtime.memcpyDefault)
  
      cpdef copy_from_device_async(self, MemoryPointer src, size_t size, stream):
          """Copies a memory from a (possibly different) device asynchronously.
  
          Args:
              src (cupy.cuda.MemoryPointer): Source memory pointer.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
  
          """
          if size > 0:
              _set_peer_access(src.device.id, self.device.id)
              runtime.memcpyAsync(self.ptr, src.ptr, size,
                                  runtime.memcpyDefault, stream.ptr)
  
      cpdef copy_from_host(self, mem, size_t size):
          """Copies a memory sequence from the host memory.
  
          Args:
              mem (ctypes.c_void_p): Source memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memcpy(self.ptr, mem.value, size,
                             runtime.memcpyHostToDevice)
  
      cpdef copy_from_host_async(self, mem, size_t size, stream):
          """Copies a memory sequence from the host memory asynchronously.
  
          Args:
              mem (ctypes.c_void_p): Source memory pointer. It must be a pinned
                  memory.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
  
          """
          if size > 0:
              runtime.memcpyAsync(self.ptr, mem.value, size,
                                  runtime.memcpyHostToDevice, stream.ptr)
  
      cpdef copy_from(self, mem, size_t size):
          """Copies a memory sequence from a (possibly different) device or host.
  
          This function is a useful interface that selects appropriate one from
          :meth:`~cupy.cuda.MemoryPointer.copy_from_device` and
          :meth:`~cupy.cuda.MemoryPointer.copy_from_host`.
  
          Args:
              mem (ctypes.c_void_p or cupy.cuda.MemoryPointer): Source memory
                  pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if isinstance(mem, MemoryPointer):
              self.copy_from_device(mem, size)
          else:
              self.copy_from_host(mem, size)
  
      cpdef copy_from_async(self, mem, size_t size, stream):
          """Copies a memory sequence from an arbitrary place asynchronously.
  
          This function is a useful interface that selects appropriate one from
          :meth:`~cupy.cuda.MemoryPointer.copy_from_device_async` and
          :meth:`~cupy.cuda.MemoryPointer.copy_from_host_async`.
  
          Args:
              mem (ctypes.c_void_p or cupy.cuda.MemoryPointer): Source memory
                  pointer.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
  
          """
          if isinstance(mem, MemoryPointer):
              self.copy_from_device_async(mem, size, stream)
          else:
              self.copy_from_host_async(mem, size, stream)
  
      cpdef copy_to_host(self, mem, size_t size):
          """Copies a memory sequence to the host memory.
  
          Args:
              mem (ctypes.c_void_p): Target memory pointer.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memcpy(mem.value, self.ptr, size,
                             runtime.memcpyDeviceToHost)
  
      cpdef copy_to_host_async(self, mem, size_t size, stream):
          """Copies a memory sequence to the host memory asynchronously.
  
          Args:
              mem (ctypes.c_void_p): Target memory pointer. It must be a pinned
                  memory.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
  
          """
          if size > 0:
              runtime.memcpyAsync(mem.value, self.ptr, size,
                                  runtime.memcpyDeviceToHost, stream.ptr)
  
      cpdef memset(self, int value, size_t size):
          """Fills a memory sequence by constant byte value.
  
          Args:
              value (int): Value to fill.
              size (int): Size of the sequence in bytes.
  
          """
          if size > 0:
              runtime.memset(self.ptr, value, size)
  
      cpdef memset_async(self, int value, size_t size, stream):
          """Fills a memory sequence by constant byte value asynchronously.
  
          Args:
              value (int): Value to fill.
              size (int): Size of the sequence in bytes.
              stream (cupy.cuda.Stream): CUDA stream.
  
          """
          if size > 0:
              runtime.memsetAsync(self.ptr, value, size, stream.ptr)
  
  
  cpdef MemoryPointer _malloc(Py_ssize_t size):
      mem = Memory(size)
      return MemoryPointer(mem, 0)
  
  
  cpdef MemoryPointer malloc_managed(Py_ssize_t size):
      """Allocate managed memory (unified memory).
  
      This method can be used as a CuPy memory allocator. The simplest way to
      use a managed memory as the default allocator is the following code::
  
          set_allocator(malloc_managed)
  
      The advantage using managed memory in CuPy is that device memory
      oversubscription is possible for GPUs that have a non-zero value for the
      device attribute cudaDevAttrConcurrentManagedAccess.
      CUDA >= 8.0 with GPUs later than or equal to Pascal is preferrable.
  
      Read more at: http://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__MEMORY.html#axzz4qygc1Ry1  # NOQA
  
      Args:
          size (int): Size of the memory allocation in bytes.
  
      Returns:
          ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
      """
      mem = ManagedMemory(size)
      return MemoryPointer(mem, 0)
  
  
  cdef object _current_allocator = _malloc
  
  
  cpdef MemoryPointer alloc(Py_ssize_t size):
      """Calls the current allocator.
  
      Use :func:`~cupy.cuda.set_allocator` to change the current allocator.
  
      Args:
          size (int): Size of the memory allocation.
  
      Returns:
          ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
  
      """
      return _current_allocator(size)
  
  
  cpdef set_allocator(allocator=_malloc):
      """Sets the current allocator.
  
      Args:
          allocator (function): CuPy memory allocator. It must have the same
              interface as the :func:`cupy.cuda.alloc` function, which takes the
              buffer size as an argument and returns the device buffer of that
              size.
  
      """
      global _current_allocator
      _current_allocator = allocator
  
  
  class PooledMemory(Memory):
  
      """Memory allocation for a memory pool.
  
      The instance of this class is created by memory pool allocator, so user
      should not instantiate it by hand.
  
      """
  
      def __init__(self, Chunk chunk, pool):
          self.device = chunk.device
          self.ptr = chunk.ptr
          self.size = chunk.size
          self.pool = pool
  
      def free(self):
          """Frees the memory buffer and returns it to the memory pool.
  
          This function actually does not free the buffer. It just returns the
          buffer to the memory pool for reuse.
  
          """
          cdef Py_ssize_t ptr
          ptr = self.ptr
          if ptr == 0:
              return
          self.ptr = 0
          pool = self.pool()
          if pool is None:
              return
  
          hooks = None
          # to avoid error at exit
          mh = memory_hook
          if mh is not None and mh.get_memory_hooks is not None:
              hooks = mh.get_memory_hooks()
          size = self.size
          if hooks:
              device_id = self.device.id
              pmem_id = id(self)
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.free_preprocess(device_id=device_id,
                                       mem_size=size,
                                       mem_ptr=ptr,
                                       pmem_id=pmem_id)
              try:
                  pool.free(ptr, size)
              finally:
                  for hook in hooks_values:
                      hook.free_postprocess(device_id=device_id,
                                            mem_size=size,
                                            mem_ptr=ptr,
                                            pmem_id=pmem_id)
          else:
              pool.free(ptr, size)
  
      __del__ = free
  
  
  cdef class SingleDeviceMemoryPool:
      """Memory pool implementation for single device.
  
      - The allocator attempts to find the smallest cached block that will fit
        the requested size. If the block is larger than the requested size,
        it may be split. If no block is found, the allocator will delegate to
        cudaMalloc.
      - If the cudaMalloc fails, the allocator will free all cached blocks that
        are not split and retry the allocation.
      """
  
      def __init__(self, allocator=_malloc):
          # cudaMalloc() is aligned to at least 512 bytes
          # cf. https://gist.github.com/sonots/41daaa6432b1c8b27ef782cd14064269
          self._allocation_unit_size = 512
          self._in_use = {}
          self._free = []
          self._allocator = allocator
          self._weakref = weakref.ref(self)
          self._device_id = device.get_device_id()
          self._free_lock = rlock.create_fastrlock()
          self._in_use_lock = rlock.create_fastrlock()
  
      cpdef Py_ssize_t _round_size(self, Py_ssize_t size):
          """Round up the memory size to fit memory alignment of cudaMalloc."""
          unit = self._allocation_unit_size
          return ((size + unit - 1) // unit) * unit
  
      cpdef int _bin_index_from_size(self, Py_ssize_t size):
          """Get appropriate bins index from the memory size"""
          unit = self._allocation_unit_size
          return (size - 1) // unit
  
      cpdef _append_to_free_list(self, Py_ssize_t size, chunk):
          cdef int index, bin_index
          cdef set free_list
          bin_index = self._bin_index_from_size(size)
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              index = algorithm.lower_bound(
                  self._index.begin(), self._index.end(),
                  bin_index) - self._index.begin()
              if index < self._index.size() and self._index[index] == bin_index:
                  free_list = self._free[index]
              else:
                  free_list = set()
                  self._index.insert(
                      self._index.begin() + index, bin_index)
                  self._free.insert(index, free_list)
              free_list.add(chunk)
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
      cpdef bint _remove_from_free_list(self, Py_ssize_t size, chunk) except *:
          cdef int index, bin_index
          cdef set free_list
          bin_index = self._bin_index_from_size(size)
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              index = algorithm.lower_bound(
                  self._index.begin(), self._index.end(),
                  bin_index) - self._index.begin()
              if self._index[index] != bin_index:
                  return False
              free_list = self._free[index]
              if chunk in free_list:
                  free_list.remove(chunk)
                  return True
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return False
  
      cpdef tuple _split(self, Chunk chunk, Py_ssize_t size):
          """Split contiguous block of a larger allocation"""
          cdef Chunk head
          cdef Chunk remaining
          assert chunk.size >= size
          if chunk.size == size:
              return chunk, None
          head = Chunk(chunk.mem, chunk.offset, size)
          remaining = Chunk(chunk.mem, chunk.offset + size, chunk.size - size)
          if chunk.prev is not None:
              head.prev = chunk.prev
              chunk.prev.next = head
          if chunk.next is not None:
              remaining.next = chunk.next
              chunk.next.prev = remaining
          head.next = remaining
          remaining.prev = head
          return head, remaining
  
      cpdef Chunk _merge(self, Chunk head, Chunk remaining):
          """Merge previously splitted block (chunk)"""
          cdef Chunk merged
          size = head.size + remaining.size
          merged = Chunk(head.mem, head.offset, size)
          if head.prev is not None:
              merged.prev = head.prev
              merged.prev.next = merged
          if remaining.next is not None:
              merged.next = remaining.next
              merged.next.prev = merged
          return merged
  
      cpdef MemoryPointer _alloc(self, Py_ssize_t rounded_size):
          hooks = memory_hook.get_memory_hooks()
          if hooks:
              memptr = None
              device_id = self._device_id
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.alloc_preprocess(device_id=device_id,
                                        mem_size=rounded_size)
              try:
                  memptr = self._allocator(rounded_size)
              finally:
                  for hook in hooks_values:
                      mem_ptr = memptr.ptr if memptr is not None else 0
                      hook.alloc_postprocess(device_id=device_id,
                                             mem_size=rounded_size,
                                             mem_ptr=mem_ptr)
              return memptr
          else:
              return self._allocator(rounded_size)
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size):
          rounded_size = self._round_size(size)
          hooks = memory_hook.get_memory_hooks()
          if hooks:
              memptr = None
              device_id = self._device_id
              hooks_values = hooks.values()  # avoid six for performance
              for hook in hooks_values:
                  hook.malloc_preprocess(device_id=device_id,
                                         size=size,
                                         mem_size=rounded_size)
              try:
                  memptr = self._malloc(rounded_size)
              finally:
                  if memptr is None:
                      mem_ptr = 0
                      pmem_id = 0
                  else:
                      mem_ptr = memptr.ptr
                      pmem_id = id(memptr.mem)
                  for hook in hooks_values:
                      hook.malloc_postprocess(device_id=device_id,
                                              size=size,
                                              mem_size=rounded_size,
                                              mem_ptr=mem_ptr,
                                              pmem_id=pmem_id)
              return memptr
          else:
              return self._malloc(rounded_size)
  
      cpdef MemoryPointer _malloc(self, Py_ssize_t size):
          cdef set free_list
          cdef Chunk chunk = None
          cdef Chunk remaining = None
          cdef int bin_index, index, length
  
          if size == 0:
              return MemoryPointer(Memory(0), 0)
  
          # find best-fit, or a smallest larger allocation
          rlock.lock_fastrlock(self._free_lock, -1, True)
          bin_index = self._bin_index_from_size(size)
          try:
              index = algorithm.lower_bound(
                  self._index.begin(), self._index.end(),
                  bin_index) - self._index.begin()
              length = self._index.size()
              for i in range(index, length):
                  free_list = self._free[i]
                  if free_list:
                      chunk = free_list.pop()
                      break
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
          if chunk is not None:
              chunk, remaining = self._split(chunk, size)
          else:
              # cudaMalloc if a cache is not found
              try:
                  mem = self._alloc(size).mem
              except runtime.CUDARuntimeError as e:
                  if e.status != runtime.errorMemoryAllocation:
                      raise
                  self.free_all_blocks()
                  try:
                      mem = self._alloc(size).mem
                  except runtime.CUDARuntimeError as e:
                      if e.status != runtime.errorMemoryAllocation:
                          raise
                      gc.collect()
                      try:
                          mem = self._alloc(size).mem
                      except runtime.CUDARuntimeError as e:
                          if e.status != runtime.errorMemoryAllocation:
                              raise
                          else:
                              total = size + self.total_bytes()
                              raise OutOfMemoryError(size, total)
              chunk = Chunk(mem, 0, size)
  
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              self._in_use[chunk.ptr] = chunk
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          if remaining is not None:
              self._append_to_free_list(remaining.size, remaining)
          pmem = PooledMemory(chunk, self._weakref)
          return MemoryPointer(pmem, 0)
  
      cpdef free(self, size_t ptr, Py_ssize_t size):
          cdef set free_list
          cdef Chunk chunk
  
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              chunk = self._in_use.pop(ptr, None)
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          if chunk is None:
              raise RuntimeError('Cannot free out-of-pool memory')
  
          if chunk.next is not None:
              if self._remove_from_free_list(chunk.next.size, chunk.next):
                  chunk = self._merge(chunk, chunk.next)
  
          if chunk.prev is not None:
              if self._remove_from_free_list(chunk.prev.size, chunk.prev):
                  chunk = self._merge(chunk.prev, chunk)
  
          self._append_to_free_list(chunk.size, chunk)
  
      cpdef free_all_blocks(self):
          cdef set free_list, keep_list
          cdef Chunk chunk
          # Free all **non-split** chunks
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for i in range(len(self._free)):
                  free_list = self._free[i]
                  keep_list = set()
                  for chunk in free_list:
                      if chunk.prev is not None or chunk.next is not None:
                          keep_list.add(chunk)
                  self._free[i] = keep_list
          finally:
              rlock.unlock_fastrlock(self._free_lock)
  
      cpdef free_all_free(self):
          warnings.warn(
              'free_all_free is deprecated. Use free_all_blocks instead.',
              DeprecationWarning)
          self.free_all_blocks()
  
      cpdef n_free_blocks(self):
          cdef Py_ssize_t n = 0
          cdef set free_list
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for free_list in self._free:
                  n += len(free_list)
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return n
  
      cpdef used_bytes(self):
          cdef Py_ssize_t size = 0
          cdef Chunk chunk
          rlock.lock_fastrlock(self._in_use_lock, -1, True)
          try:
              for chunk in self._in_use.itervalues():
                  size += chunk.size
          finally:
              rlock.unlock_fastrlock(self._in_use_lock)
          return size
  
      cpdef free_bytes(self):
          cdef Py_ssize_t size = 0
          cdef set free_list
          cdef Chunk chunk
          rlock.lock_fastrlock(self._free_lock, -1, True)
          try:
              for free_list in self._free:
                  for chunk in free_list:
                      size += chunk.size
          finally:
              rlock.unlock_fastrlock(self._free_lock)
          return size
  
      cpdef total_bytes(self):
          return self.used_bytes() + self.free_bytes()
  
  
  cdef class MemoryPool(object):
  
      """Memory pool for all devices on the machine.
  
      A memory pool preserves any allocations even if they are freed by the user.
      Freed memory buffers are held by the memory pool as *free blocks*, and they
      are reused for further memory allocations of the same sizes. The allocated
      blocks are managed for each device, so one instance of this class can be
      used for multiple devices.
  
      .. note::
         When the allocation is skipped by reusing the pre-allocated block, it
         does not call ``cudaMalloc`` and therefore CPU-GPU synchronization does
         not occur. It makes interleaves of memory allocations and kernel
         invocations very fast.
  
      .. note::
         The memory pool holds allocated blocks without freeing as much as
         possible. It makes the program hold most of the device memory, which may
         make other CUDA programs running in parallel out-of-memory situation.
  
      Args:
          allocator (function): The base CuPy memory allocator. It is used for
              allocating new blocks when the blocks of the required size are all
              in use.
  
      """
  
      def __init__(self, allocator=_malloc):
          self._pools = collections.defaultdict(
              lambda: SingleDeviceMemoryPool(allocator))
  
      cpdef MemoryPointer malloc(self, Py_ssize_t size):
          """Allocates the memory, from the pool if possible.
  
          This method can be used as a CuPy memory allocator. The simplest way to
          use a memory pool as the default allocator is the following code::
  
              set_allocator(MemoryPool().malloc)
  
          Also, the way to use a memory pool of Managed memory (Unified memory)
          as the default allocator is the following code::
  
              set_allocator(MemoryPool(malloc_managed).malloc)
  
          Args:
              size (int): Size of the memory buffer to allocate in bytes.
  
          Returns:
              ~cupy.cuda.MemoryPointer: Pointer to the allocated buffer.
  
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.malloc(size)
  
      cpdef free_all_blocks(self):
          """Release free blocks."""
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          mp.free_all_blocks()
  
      cpdef free_all_free(self):
          """Release free blocks."""
          warnings.warn(
              'free_all_free is deprecated. Use free_all_blocks instead.',
              DeprecationWarning)
          self.free_all_blocks()
  
      cpdef n_free_blocks(self):
          """Count the total number of free blocks.
  
          Returns:
              int: The total number of free blocks.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.n_free_blocks()
  
      cpdef used_bytes(self):
          """Get the total number of bytes used.
  
          Returns:
              int: The total number of bytes used.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.used_bytes()
  
      cpdef free_bytes(self):
          """Get the total number of bytes acquired but not used in the pool.
  
          Returns:
              int: The total number of bytes acquired but not used in the pool.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.free_bytes()
  
      cpdef total_bytes(self):
          """Get the total number of bytes acquired in the pool.
  
          Returns:
              int: The total number of bytes acquired in the pool.
          """
          mp = <SingleDeviceMemoryPool>self._pools[device.get_device_id()]
          return mp.total_bytes()
  
files: cupy/cuda/memory.pyx
files: import ctypes
  import unittest
  
  import cupy.cuda
  from cupy.cuda import memory
  from cupy.cuda import stream as stream_module
  from cupy import testing
  
  
  class MockMemory(memory.Memory):
      cur_ptr = 1
  
      def __init__(self, size):
          self.ptr = MockMemory.cur_ptr
          MockMemory.cur_ptr += size
          self.size = size
          self.device = None
  
      def __del__(self):
          self.ptr = 0
          pass
  
  
  def mock_alloc(size):
      mem = MockMemory(size)
      return memory.MemoryPointer(mem, 0)
  
  # -----------------------------------------------------------------------------
  # Memory pointer
  
  
  @testing.gpu
  class TestMemoryPointer(unittest.TestCase):
  
      def test_int(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(1)
          self.assertEqual(pval, int(memptr))
  
      def test_add(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(8)
  
          memptr2 = memptr + 4
          self.assertIsInstance(memptr2, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr2))
  
          memptr3 = 4 + memptr
          self.assertIsInstance(memptr3, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr3))
  
          memptr += 4
          self.assertIsInstance(memptr, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr))
  
      def test_sub(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(8) + 4
  
          memptr2 = memptr - 4
          self.assertIsInstance(memptr2, memory.MemoryPointer)
          self.assertEqual(pval, int(memptr2))
  
          memptr -= 4
          self.assertIsInstance(memptr, memory.MemoryPointer)
          self.assertEqual(pval, int(memptr))
  
      def test_copy_to_and_from_host(self):
          a_gpu = memory.alloc(4)
          a_cpu = ctypes.c_int(100)
          a_gpu.copy_from(ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 4)
          b_cpu = ctypes.c_int()
          a_gpu.copy_to_host(
              ctypes.cast(ctypes.byref(b_cpu), ctypes.c_void_p), 4)
          self.assertEqual(b_cpu.value, a_cpu.value)
  
      def test_copy_from_device(self):
          a_gpu = memory.alloc(4)
          a_cpu = ctypes.c_int(100)
          a_gpu.copy_from(ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 4)
  
          b_gpu = memory.alloc(4)
          b_gpu.copy_from(a_gpu, 4)
          b_cpu = ctypes.c_int()
          b_gpu.copy_to_host(
              ctypes.cast(ctypes.byref(b_cpu), ctypes.c_void_p), 4)
          self.assertEqual(b_cpu.value, a_cpu.value)
  
      def test_memset(self):
          a_gpu = memory.alloc(4)
          a_gpu.memset(1, 4)
          a_cpu = ctypes.c_ubyte()
          for i in range(4):
              a_gpu.copy_to_host(
                  ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 1)
              self.assertEqual(a_cpu.value, 1)
              a_gpu += 1
  
  
  # -----------------------------------------------------------------------------
  # Memory pool
  
  @testing.gpu
  class TestSingleDeviceMemoryPool(unittest.TestCase):
  
      def setUp(self):
          self.pool = memory.SingleDeviceMemoryPool(allocator=mock_alloc)
          self.unit = self.pool._allocation_unit_size
          self.stream = stream_module.Stream()
          self.stream_ptr = self.stream.ptr
  
      def test_round_size(self):
          self.assertEqual(self.pool._round_size(self.unit - 1), self.unit)
          self.assertEqual(self.pool._round_size(self.unit), self.unit)
          self.assertEqual(self.pool._round_size(self.unit + 1), self.unit * 2)
  
      def test_bin_index_from_size(self):
          self.assertEqual(self.pool._bin_index_from_size(self.unit - 1), 0)
          self.assertEqual(self.pool._bin_index_from_size(self.unit), 0)
          self.assertEqual(self.pool._bin_index_from_size(self.unit + 1), 1)
  
      def test_split(self):
          mem = MockMemory(self.unit * 4)
          chunk = memory.Chunk(mem, 0, mem.size, self.stream_ptr)
          head, tail = self.pool._split(chunk, self.unit * 2)
          self.assertEqual(head.ptr, chunk.ptr)
          self.assertEqual(head.offset, 0)
          self.assertEqual(head.size, self.unit * 2)
          self.assertEqual(head.prev, None)
          self.assertEqual(head.next.ptr, tail.ptr)
          self.assertEqual(head.stream_ptr, self.stream_ptr)
          self.assertEqual(tail.ptr, chunk.ptr + self.unit * 2)
          self.assertEqual(tail.offset, self.unit * 2)
          self.assertEqual(tail.size, self.unit * 2)
          self.assertEqual(tail.prev.ptr, head.ptr)
          self.assertEqual(tail.next, None)
          self.assertEqual(tail.stream_ptr, self.stream_ptr)
  
          head_of_head, tail_of_head = self.pool._split(head, self.unit)
          self.assertEqual(head_of_head.ptr, chunk.ptr)
          self.assertEqual(head_of_head.offset, 0)
          self.assertEqual(head_of_head.size, self.unit)
          self.assertEqual(head_of_head.prev, None)
          self.assertEqual(head_of_head.next.ptr, tail_of_head.ptr)
          self.assertEqual(head_of_head.stream_ptr, self.stream_ptr)
          self.assertEqual(tail_of_head.ptr, chunk.ptr + self.unit)
          self.assertEqual(tail_of_head.offset, self.unit)
          self.assertEqual(tail_of_head.size, self.unit)
          self.assertEqual(tail_of_head.prev.ptr, head_of_head.ptr)
          self.assertEqual(tail_of_head.next.ptr, tail.ptr)
          self.assertEqual(tail_of_head.stream_ptr, self.stream_ptr)
  
          head_of_tail, tail_of_tail = self.pool._split(tail, self.unit)
          self.assertEqual(head_of_tail.ptr, chunk.ptr + self.unit * 2)
          self.assertEqual(head_of_tail.offset, self.unit * 2)
          self.assertEqual(head_of_tail.size, self.unit)
          self.assertEqual(head_of_tail.prev.ptr, tail_of_head.ptr)
          self.assertEqual(head_of_tail.next.ptr, tail_of_tail.ptr)
          self.assertEqual(head_of_tail.stream_ptr, self.stream_ptr)
          self.assertEqual(tail_of_tail.ptr, chunk.ptr + self.unit * 3)
          self.assertEqual(tail_of_tail.offset, self.unit * 3)
          self.assertEqual(tail_of_tail.size, self.unit)
          self.assertEqual(tail_of_tail.prev.ptr, head_of_tail.ptr)
          self.assertEqual(tail_of_tail.next, None)
          self.assertEqual(tail_of_tail.stream_ptr, self.stream_ptr)
  
      def test_merge(self):
          mem = MockMemory(self.unit * 4)
          chunk = memory.Chunk(mem, 0, mem.size, self.stream_ptr)
  
          head, tail = self.pool._split(chunk, self.unit * 2)
          head_ptr, tail_ptr = head.ptr, tail.ptr
          head_of_head, tail_of_head = self.pool._split(head, self.unit)
          head_of_tail, tail_of_tail = self.pool._split(tail, self.unit)
  
          merged_head = self.pool._merge(head_of_head, tail_of_head)
          self.assertEqual(merged_head.ptr, head.ptr)
          self.assertEqual(merged_head.offset, head.offset)
          self.assertEqual(merged_head.size, head.size)
          self.assertEqual(merged_head.prev, None)
          self.assertEqual(merged_head.next.ptr, tail_ptr)
          self.assertEqual(merged_head.stream_ptr, self.stream_ptr)
  
          merged_tail = self.pool._merge(head_of_tail, tail_of_tail)
          self.assertEqual(merged_tail.ptr, tail.ptr)
          self.assertEqual(merged_tail.offset, tail.offset)
          self.assertEqual(merged_tail.size, tail.size)
          self.assertEqual(merged_tail.prev.ptr, head_ptr)
          self.assertEqual(merged_tail.next, None)
          self.assertEqual(merged_tail.stream_ptr, self.stream_ptr)
  
          merged = self.pool._merge(merged_head, merged_tail)
          self.assertEqual(merged.ptr, chunk.ptr)
          self.assertEqual(merged.offset, chunk.offset)
          self.assertEqual(merged.size, chunk.size)
          self.assertEqual(merged.prev, None)
          self.assertEqual(merged.next, None)
          self.assertEqual(merged.stream_ptr, self.stream_ptr)
  
      def test_alloc(self):
          p1 = self.pool.malloc(self.unit * 4)
          p2 = self.pool.malloc(self.unit * 4)
          p3 = self.pool.malloc(self.unit * 8)
          self.assertNotEqual(p1.ptr, p2.ptr)
          self.assertNotEqual(p1.ptr, p3.ptr)
          self.assertNotEqual(p2.ptr, p3.ptr)
  
      def test_alloc_split(self):
          p = self.pool.malloc(self.unit * 4)
          ptr = p.ptr
          del p
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          self.assertEqual(ptr + self.unit * 2, tail.ptr)
  
      def test_free(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr1, p2.ptr)
  
      def test_free_stream(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          with self.stream:
              p2 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p2.ptr)
  
      def test_free_merge(self):
          p = self.pool.malloc(self.unit * 4)
          ptr = p.ptr
          del p
  
          # merge head into tail
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          del tail
          del head
          p = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr, p.ptr)
          del p
  
          # merge tail into head
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          del head
          del tail
          p = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr, p.ptr)
          del p
  
      def test_free_different_size(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          p2 = self.pool.malloc(self.unit * 8)
          self.assertNotEqual(ptr1, p2.ptr)
  
      def test_free_all_blocks(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          self.pool.free_all_blocks()
          p2 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p2.ptr)
          del p2
  
      def test_free_all_blocks_split(self):
          # do not free splitted blocks
          p = self.pool.malloc(self.unit * 4)
          del p
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          tailptr = tail.ptr
          del tail
          self.pool.free_all_blocks()
          p = self.pool.malloc(self.unit * 2)
          self.assertEqual(tailptr, p.ptr)
          del head
  
      def test_free_all_blocks_stream(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          with self.stream:
              p2 = self.pool.malloc(self.unit * 4)
              ptr2 = p2.ptr
              del p2
          self.pool.free_all_blocks()
          p3 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p3.ptr)
          self.assertNotEqual(ptr2, p3.ptr)
          with self.stream:
              p4 = self.pool.malloc(self.unit * 4)
              self.assertNotEqual(ptr1, p4.ptr)
              self.assertNotEqual(ptr2, p4.ptr)
  
      def test_free_all_free(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          with testing.assert_warns(DeprecationWarning):
              self.pool.free_all_free()
          p2 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p2.ptr)
  
      def test_used_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 2, self.pool.used_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 6, self.pool.used_bytes())
          del p2
          self.assertEqual(self.unit * 2, self.pool.used_bytes())
          del p1
          self.assertEqual(self.unit * 0, self.pool.used_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 1, self.pool.used_bytes())
          del p3
  
      def test_used_bytes_stream(self):
          p1 = self.pool.malloc(self.unit * 4)
          del p1
          with self.stream:
              p2 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 2, self.pool.used_bytes())
          del p2
  
      def test_free_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 0, self.pool.free_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 0, self.pool.free_bytes())
          del p2
          self.assertEqual(self.unit * 4, self.pool.free_bytes())
          del p1
          self.assertEqual(self.unit * 6, self.pool.free_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 5, self.pool.free_bytes())
          del p3
  
      def test_free_bytes_stream(self):
          p1 = self.pool.malloc(self.unit * 4)
          del p1
          with self.stream:
              p2 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 4, self.pool.free_bytes())
          del p2
  
      def test_total_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 2, self.pool.total_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p1
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p2
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p3
  
      def test_total_bytes_stream(self):
          p1 = self.pool.malloc(self.unit * 4)
          del p1
          with self.stream:
              p2 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p2
  
  
  @testing.parameterize(*testing.product({
      'allocator': [memory._malloc, memory.malloc_managed],
  }))
  @testing.gpu
  class TestMemoryPool(unittest.TestCase):
  
      def setUp(self):
          self.pool = memory.MemoryPool(self.allocator)
  
      def test_zero_size_alloc(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(0).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertNotIsInstance(mem, memory.PooledMemory)
  
      def test_double_free(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              mem.free()
              mem.free()
  
      def test_free_all_blocks(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertIsInstance(mem, memory.PooledMemory)
              self.assertEqual(self.pool.n_free_blocks(), 0)
              mem.free()
              self.assertEqual(self.pool.n_free_blocks(), 1)
              self.pool.free_all_blocks()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_blocks_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc.
              self.pool.free_all_blocks()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_free(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertIsInstance(mem, memory.PooledMemory)
              self.assertEqual(self.pool.n_free_blocks(), 0)
              mem.free()
              self.assertEqual(self.pool.n_free_blocks(), 1)
              with testing.assert_warns(DeprecationWarning):
                  self.pool.free_all_free()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_free_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc.
              with testing.assert_warns(DeprecationWarning):
                  self.pool.free_all_free()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_n_free_blocks_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc/free_all_free.
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_used_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.used_bytes())
  
      def test_free_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.free_bytes())
  
      def test_total_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.total_bytes())
  
files: tests/cupy_tests/cuda_tests/test_memory.py
files: import ctypes
  import unittest
  
  import cupy.cuda
  from cupy.cuda import memory
  from cupy import testing
  
  
  class MockMemory(memory.Memory):
      cur_ptr = 1
  
      def __init__(self, size):
          self.ptr = MockMemory.cur_ptr
          MockMemory.cur_ptr += size
          self.size = size
          self.device = None
  
      def __del__(self):
          self.ptr = 0
          pass
  
  
  def mock_alloc(size):
      mem = MockMemory(size)
      return memory.MemoryPointer(mem, 0)
  
  # -----------------------------------------------------------------------------
  # Memory pointer
  
  
  @testing.gpu
  class TestMemoryPointer(unittest.TestCase):
  
      def test_int(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(1)
          self.assertEqual(pval, int(memptr))
  
      def test_add(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(8)
  
          memptr2 = memptr + 4
          self.assertIsInstance(memptr2, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr2))
  
          memptr3 = 4 + memptr
          self.assertIsInstance(memptr3, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr3))
  
          memptr += 4
          self.assertIsInstance(memptr, memory.MemoryPointer)
          self.assertEqual(pval + 4, int(memptr))
  
      def test_sub(self):
          pval = MockMemory.cur_ptr
          memptr = mock_alloc(8) + 4
  
          memptr2 = memptr - 4
          self.assertIsInstance(memptr2, memory.MemoryPointer)
          self.assertEqual(pval, int(memptr2))
  
          memptr -= 4
          self.assertIsInstance(memptr, memory.MemoryPointer)
          self.assertEqual(pval, int(memptr))
  
      def test_copy_to_and_from_host(self):
          a_gpu = memory.alloc(4)
          a_cpu = ctypes.c_int(100)
          a_gpu.copy_from(ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 4)
          b_cpu = ctypes.c_int()
          a_gpu.copy_to_host(
              ctypes.cast(ctypes.byref(b_cpu), ctypes.c_void_p), 4)
          self.assertEqual(b_cpu.value, a_cpu.value)
  
      def test_copy_from_device(self):
          a_gpu = memory.alloc(4)
          a_cpu = ctypes.c_int(100)
          a_gpu.copy_from(ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 4)
  
          b_gpu = memory.alloc(4)
          b_gpu.copy_from(a_gpu, 4)
          b_cpu = ctypes.c_int()
          b_gpu.copy_to_host(
              ctypes.cast(ctypes.byref(b_cpu), ctypes.c_void_p), 4)
          self.assertEqual(b_cpu.value, a_cpu.value)
  
      def test_memset(self):
          a_gpu = memory.alloc(4)
          a_gpu.memset(1, 4)
          a_cpu = ctypes.c_ubyte()
          for i in range(4):
              a_gpu.copy_to_host(
                  ctypes.cast(ctypes.byref(a_cpu), ctypes.c_void_p), 1)
              self.assertEqual(a_cpu.value, 1)
              a_gpu += 1
  
  
  # -----------------------------------------------------------------------------
  # Memory pool
  
  @testing.gpu
  class TestSingleDeviceMemoryPool(unittest.TestCase):
  
      def setUp(self):
          self.pool = memory.SingleDeviceMemoryPool(allocator=mock_alloc)
          self.unit = self.pool._allocation_unit_size
  
      def test_round_size(self):
          self.assertEqual(self.pool._round_size(self.unit - 1), self.unit)
          self.assertEqual(self.pool._round_size(self.unit), self.unit)
          self.assertEqual(self.pool._round_size(self.unit + 1), self.unit * 2)
  
      def test_bin_index_from_size(self):
          self.assertEqual(self.pool._bin_index_from_size(self.unit - 1), 0)
          self.assertEqual(self.pool._bin_index_from_size(self.unit), 0)
          self.assertEqual(self.pool._bin_index_from_size(self.unit + 1), 1)
  
      def test_split(self):
          mem = MockMemory(self.unit * 4)
          chunk = memory.Chunk(mem, 0, mem.size)
          head, tail = self.pool._split(chunk, self.unit * 2)
          self.assertEqual(head.ptr,      chunk.ptr)
          self.assertEqual(head.offset,   0)
          self.assertEqual(head.size,     self.unit * 2)
          self.assertEqual(head.prev,     None)
          self.assertEqual(head.next.ptr, tail.ptr)
          self.assertEqual(tail.ptr,      chunk.ptr + self.unit * 2)
          self.assertEqual(tail.offset,   self.unit * 2)
          self.assertEqual(tail.size,     self.unit * 2)
          self.assertEqual(tail.prev.ptr, head.ptr)
          self.assertEqual(tail.next,     None)
  
          head_of_head, tail_of_head = self.pool._split(head, self.unit)
          self.assertEqual(head_of_head.ptr,      chunk.ptr)
          self.assertEqual(head_of_head.offset,   0)
          self.assertEqual(head_of_head.size,     self.unit)
          self.assertEqual(head_of_head.prev,     None)
          self.assertEqual(head_of_head.next.ptr, tail_of_head.ptr)
          self.assertEqual(tail_of_head.ptr,      chunk.ptr + self.unit)
          self.assertEqual(tail_of_head.offset,   self.unit)
          self.assertEqual(tail_of_head.size,     self.unit)
          self.assertEqual(tail_of_head.prev.ptr, head_of_head.ptr)
          self.assertEqual(tail_of_head.next.ptr, tail.ptr)
  
          head_of_tail, tail_of_tail = self.pool._split(tail, self.unit)
          self.assertEqual(head_of_tail.ptr,      chunk.ptr + self.unit * 2)
          self.assertEqual(head_of_tail.offset,   self.unit * 2)
          self.assertEqual(head_of_tail.size,     self.unit)
          self.assertEqual(head_of_tail.prev.ptr, tail_of_head.ptr)
          self.assertEqual(head_of_tail.next.ptr, tail_of_tail.ptr)
          self.assertEqual(tail_of_tail.ptr,      chunk.ptr + self.unit * 3)
          self.assertEqual(tail_of_tail.offset,   self.unit * 3)
          self.assertEqual(tail_of_tail.size,     self.unit)
          self.assertEqual(tail_of_tail.prev.ptr, head_of_tail.ptr)
          self.assertEqual(tail_of_tail.next,     None)
  
      def test_merge(self):
          mem = MockMemory(self.unit * 4)
          chunk = memory.Chunk(mem, 0, mem.size)
  
          head, tail = self.pool._split(chunk, self.unit * 2)
          head_ptr, tail_ptr = head.ptr, tail.ptr
          head_of_head, tail_of_head = self.pool._split(head, self.unit)
          head_of_tail, tail_of_tail = self.pool._split(tail, self.unit)
  
          merged_head = self.pool._merge(head_of_head, tail_of_head)
          self.assertEqual(merged_head.ptr,      head.ptr)
          self.assertEqual(merged_head.offset,   head.offset)
          self.assertEqual(merged_head.size,     head.size)
          self.assertEqual(merged_head.prev,     None)
          self.assertEqual(merged_head.next.ptr, tail_ptr)
  
          merged_tail = self.pool._merge(head_of_tail, tail_of_tail)
          self.assertEqual(merged_tail.ptr,      tail.ptr)
          self.assertEqual(merged_tail.offset,   tail.offset)
          self.assertEqual(merged_tail.size,     tail.size)
          self.assertEqual(merged_tail.prev.ptr, head_ptr)
          self.assertEqual(merged_tail.next,     None)
  
          merged = self.pool._merge(merged_head, merged_tail)
          self.assertEqual(merged.ptr,    chunk.ptr)
          self.assertEqual(merged.offset, chunk.offset)
          self.assertEqual(merged.size,   chunk.size)
          self.assertEqual(merged.prev,   None)
          self.assertEqual(merged.next,   None)
  
      def test_alloc(self):
          p1 = self.pool.malloc(self.unit * 4)
          p2 = self.pool.malloc(self.unit * 4)
          p3 = self.pool.malloc(self.unit * 8)
          self.assertNotEqual(p1.ptr, p2.ptr)
          self.assertNotEqual(p1.ptr, p3.ptr)
          self.assertNotEqual(p2.ptr, p3.ptr)
  
      def test_alloc_split(self):
          p = self.pool.malloc(self.unit * 4)
          ptr = p.ptr
          del p
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          self.assertEqual(ptr + self.unit * 2, tail.ptr)
  
      def test_free(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr1, p2.ptr)
  
      def test_free_merge(self):
          p = self.pool.malloc(self.unit * 4)
          ptr = p.ptr
          del p
  
          # merge head into tail
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          del tail
          del head
          p = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr, p.ptr)
          del p
  
          # merge tail into head
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          self.assertEqual(ptr, head.ptr)
          del head
          del tail
          p = self.pool.malloc(self.unit * 4)
          self.assertEqual(ptr, p.ptr)
          del p
  
      def test_free_different_size(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          p2 = self.pool.malloc(self.unit * 8)
          self.assertNotEqual(ptr1, p2.ptr)
  
      def test_free_all_blocks(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          self.pool.free_all_blocks()
          p2 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p2.ptr)
          del p2
  
          # do not free splitted blocks
          head = self.pool.malloc(self.unit * 2)
          tail = self.pool.malloc(self.unit * 2)
          tailptr = tail.ptr
          del tail
          self.pool.free_all_blocks()
          p = self.pool.malloc(self.unit * 2)
          self.assertEqual(tailptr, p.ptr)
          del head
  
      def test_free_all_free(self):
          p1 = self.pool.malloc(self.unit * 4)
          ptr1 = p1.ptr
          del p1
          with testing.assert_warns(DeprecationWarning):
              self.pool.free_all_free()
          p2 = self.pool.malloc(self.unit * 4)
          self.assertNotEqual(ptr1, p2.ptr)
  
      def test_used_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 2, self.pool.used_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 6, self.pool.used_bytes())
          del p2
          self.assertEqual(self.unit * 2, self.pool.used_bytes())
          del p1
          self.assertEqual(self.unit * 0, self.pool.used_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 1, self.pool.used_bytes())
          del p3
  
      def test_free_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 0, self.pool.free_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 0, self.pool.free_bytes())
          del p2
          self.assertEqual(self.unit * 4, self.pool.free_bytes())
          del p1
          self.assertEqual(self.unit * 6, self.pool.free_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 5, self.pool.free_bytes())
          del p3
  
      def test_total_bytes(self):
          p1 = self.pool.malloc(self.unit * 2)
          self.assertEqual(self.unit * 2, self.pool.total_bytes())
          p2 = self.pool.malloc(self.unit * 4)
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p1
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p2
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          p3 = self.pool.malloc(self.unit * 1)
          self.assertEqual(self.unit * 6, self.pool.total_bytes())
          del p3
  
  
  @testing.parameterize(*testing.product({
      'allocator': [memory._malloc, memory.malloc_managed],
  }))
  @testing.gpu
  class TestMemoryPool(unittest.TestCase):
  
      def setUp(self):
          self.pool = memory.MemoryPool(self.allocator)
  
      def test_zero_size_alloc(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(0).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertNotIsInstance(mem, memory.PooledMemory)
  
      def test_double_free(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              mem.free()
              mem.free()
  
      def test_free_all_blocks(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertIsInstance(mem, memory.PooledMemory)
              self.assertEqual(self.pool.n_free_blocks(), 0)
              mem.free()
              self.assertEqual(self.pool.n_free_blocks(), 1)
              self.pool.free_all_blocks()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_blocks_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc.
              self.pool.free_all_blocks()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_free(self):
          with cupy.cuda.Device(0):
              mem = self.pool.malloc(1).mem
              self.assertIsInstance(mem, memory.Memory)
              self.assertIsInstance(mem, memory.PooledMemory)
              self.assertEqual(self.pool.n_free_blocks(), 0)
              mem.free()
              self.assertEqual(self.pool.n_free_blocks(), 1)
              with testing.assert_warns(DeprecationWarning):
                  self.pool.free_all_free()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_free_all_free_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc.
              with testing.assert_warns(DeprecationWarning):
                  self.pool.free_all_free()
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_n_free_blocks_without_malloc(self):
          with cupy.cuda.Device(0):
              # call directly without malloc/free_all_free.
              self.assertEqual(self.pool.n_free_blocks(), 0)
  
      def test_used_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.used_bytes())
  
      def test_free_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.free_bytes())
  
      def test_total_bytes(self):
          with cupy.cuda.Device(0):
              self.assertEqual(0, self.pool.total_bytes())
  
files: tests/cupy_tests/cuda_tests/test_memory.py