def test_getctype(): ffi = FFI() assert ffi.getctype("int**") == "int * *" # unicode literal assert type(ffi.getctype("int**")) is str
def test_getctype(): ffi = FFI() assert ffi.getctype("int**") == "int * *" # unicode literal assert type(ffi.getctype("int**")) is str
class CTC(object): """ """ def __init__(self, on_device='cpu', blank_label=0): libpath = get_ctc_lib() self.ffi = FFI() self.ffi.cdef(ctc_header()) self.ctclib = self.ffi.dlopen(libpath) supported_devices = ['cpu', 'gpu'] if on_device not in supported_devices: print("the requested device {} is not supported".format(on_device), file=sys.stderr) sys.exit(1) assign_device = 0 if on_device is 'cpu' else 1 self.options = self.ffi.new('ctcOptions*', { "loc": assign_device, "blank_label": blank_label })[0] self.size_in_bytes = self.ffi.new("size_t*") self.nout = None self.bsz = None def get_buf_size(self, ptr_to_buf): return self.ffi.sizeof( self.ffi.getctype(self.ffi.typeof(ptr_to_buf).item)) def buf_ref_from_array(self, arr): return self.ffi.from_buffer( self.ffi.buffer(self.ffi.cast('void*', arr.ptr), arr.nbytes)) def buf_ref_from_ptr(self, ptr, size): return self.ffi.from_buffer(self.ffi.buffer(ptr, size)) def get_gpu_workspace_size(self, lbl_lens, utt_lens, nout, bsz): self.nout = nout self.bsz = bsz _lbl_lens = self.ffi.cast("int*", lbl_lens.ravel().ctypes.data) _utt_lens = self.ffi.cast("int*", utt_lens.ravel().ctypes.data) status = self.ctclib.get_workspace_size(_lbl_lens, _utt_lens, self.nout, self.bsz, self.options, self.size_in_bytes) assert status is 0, "get_workspace_size() in warp-ctc failed" return self.size_in_bytes[0] def bind_to_gpu(self, acts, grads, lbls, lbl_lens, utt_lens, costs, workspace, scratch_size, stream): if stream is None: stream_ptr = self.ffi.cast('void*', 0) stream_buf_size = self.ffi.sizeof(self.ffi.new_handle(stream)) stream_buf = self.buf_ref_from_ptr(stream_ptr, stream_buf_size) else: stream_buf = self.ffi.cast("void*", stream.handle) self.options.stream = stream_buf flat_dims = np.prod(acts.shape) assert np.prod(grads.shape) == flat_dims acts_buf = self.ffi.cast("float*", self.buf_ref_from_array(acts)) grads_buf = self.ffi.cast("float*", self.buf_ref_from_array(grads)) costs_buf = self.ffi.cast("float*", self.buf_ref_from_array(costs)) warp_grads_buf_size = flat_dims * self.get_buf_size(grads_buf) warp_costs_buf_size = self.bsz * self.get_buf_size(costs_buf) warp_labels = self.ffi.cast("int*", lbls.ravel().ctypes.data) warp_label_lens = self.ffi.cast("int*", lbl_lens.ravel().ctypes.data) warp_input_lens = self.ffi.cast("int*", utt_lens.ravel().ctypes.data) workspace_buf = self.buf_ref_from_ptr( self.ffi.cast('void*', workspace), int(scratch_size)) ctc_status = self.ctclib.compute_ctc_loss(acts_buf, grads_buf, warp_labels, warp_label_lens, warp_input_lens, self.nout, self.bsz, costs_buf, workspace_buf, self.options) assert ctc_status is 0, "warp-ctc run failed" def bind_to_cpu(self, acts, lbls, utt_lens, lbl_lens, grads, costs, n_threads=1): self.options.num_threads = n_threads _, self.bsz, self.nout = acts.shape flat_dims = np.prod(acts.shape) assert np.prod(grads.shape) == flat_dims acts_buf = self.ffi.cast("float*", acts.ctypes.data) grads_buf = self.ffi.cast("float*", grads.ctypes.data) costs_buf = self.ffi.cast("float*", costs.ctypes.data) warp_grads_buf_size = flat_dims * self.get_buf_size(grads_buf) warp_costs_buf_size = self.bsz * self.get_buf_size(costs_buf) warp_labels = self.ffi.cast("int*", lbls.ravel().ctypes.data) warp_label_lens = self.ffi.cast("int*", lbl_lens.ravel().ctypes.data) warp_input_lens = self.ffi.cast("int*", utt_lens.ravel().ctypes.data) status = self.ctclib.get_workspace_size(warp_label_lens, warp_input_lens, self.nout, self.bsz, self.options, self.size_in_bytes) assert status is 0, "get_workspace_size() in warp-ctc failed" # TODO: workspace is a variable size buffer whose size is # determined during each call, so we can't initialize ahead # of time. Can we avoid this? workspace = self.ffi.new("char[]", self.size_in_bytes[0]) ctc_status = self.ctclib.compute_ctc_loss(acts_buf, grads_buf, warp_labels, warp_label_lens, warp_input_lens, self.nout, self.bsz, costs_buf, workspace, self.options) # transfer grads and costs back without copying self.ffi.memmove(grads, grads_buf, warp_grads_buf_size) grads = grads.reshape((acts.shape)) self.ffi.memmove(costs, costs_buf, warp_costs_buf_size) assert ctc_status is 0, "warp-ctc run failed"