DeepFaceLive/xlib/avecl/_internal/backend/Buffer.py
2021-10-22 19:31:35 +04:00

111 lines
3.8 KiB
Python

from typing import Iterable, Union
import numpy as np
from . import OpenCL as CL
class Buffer:
__slots__ = ['_device','_cl_mem','_size','_on_initialize']
def __init__(self, device : 'Device', size : int, on_initialize = None):
"""
represents physical buffer associated with physical device
device Device
size int
"""
Buffer._object_count += 1
self._device = device
self._size = size
self._cl_mem = None
self._on_initialize = on_initialize
def __del__(self):
#print('Buffer.__del__')
Buffer._object_count -= 1
self.free_cl_mem()
def get_device(self) -> 'Device': return self._device
def get_size(self) -> int: return self._size
def has_cl_mem(self) -> bool: return self._cl_mem is not None
def get_cl_mem(self) -> CL.cl_mem:
if self._cl_mem is None:
self._cl_mem = self._device._cl_mem_pool_alloc(self._size)
if self._on_initialize is not None:
self._on_initialize()
return self._cl_mem
def free_cl_mem(self):
if self._cl_mem is not None:
self._device._cl_mem_pool_free(self._cl_mem)
self._cl_mem = None
def set(self, value : Union['Buffer', np.ndarray]):
"""
Parameters
value Buffer copy data from other Buffer.
np.ndarray copies values from ndarray
to Buffer's memory
"""
if isinstance(value, Buffer):
if self != value:
if self._size != value._size:
raise Exception(f'Unable to copy from Buffer with {value._size} size to buffer with {self._size} size.')
if self._device == value._device:
CL.clEnqueueCopyBuffer(self._device._get_ctx_q(), value.get_cl_mem(), self.get_cl_mem(), 0, 0, self._size, 0, None, None)
else:
# Transfer between devices will cause low performance
raise NotImplementedError()
else:
if not isinstance(value, np.ndarray):
raise ValueError (f'Invalid type {value.__class__}. Must be np.ndarray.')
if value.nbytes != self._size:
raise ValueError(f'Value size {value.nbytes} does not match Buffer size {self._size}.')
if not value.flags.contiguous:
value = value.reshape(-1)
if not value.flags.contiguous:
raise ValueError ("Unable to write from non-contiguous np array.")
ev = CL.cl_event()
clr = CL.clEnqueueWriteBuffer(self._device._get_ctx_q(), self.get_cl_mem(), False, 0, value.nbytes, value.ctypes.data_as(CL.c_void_p), 0, None, ev)
if clr != CL.CLERROR.SUCCESS:
raise Exception(f'clEnqueueWriteBuffer error: {clr}')
CL.clWaitForEvents(1, ( CL.cl_event * 1 )(ev) )
CL.clReleaseEvent(ev)
def np(self, shape : Iterable, dtype : np.dtype, out=None):
"""
Returns data of buffer as np.ndarray with specified shape and dtype
"""
if out is None:
out = np.empty (shape, dtype)
if out.nbytes != self._size:
raise ValueError(f'Unable to represent Buffer with size {self._size} as shape {shape} with dtype {dtype}')
clr = CL.clEnqueueReadBuffer(self._device._get_ctx_q(), self.get_cl_mem(), True, 0, self._size, out.ctypes.data, 0, None, None)
if clr != CL.CLERROR.SUCCESS:
raise Exception(f'clEnqueueReadBuffer error: {clr}')
return out
def __str__(self):
return f'Buffer [{self._size} bytes][{f"{self._cl_mem.value}" if self._cl_mem is not None else "unallocated"}] on {str(self._device)}'
def __repr__(self):
return self.__str__()
_object_count = 0