|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
from .base_io_buffer import Base_IO_Buffer |
|
|
|
|
|
|
|
|
class Single_IO_Buffer(Base_IO_Buffer): |
|
|
|
|
|
def __init__(self, pinned_tensor, dnvme_handle): |
|
|
super(Single_IO_Buffer, self).__init__(pinned_tensor, dnvme_handle) |
|
|
self._pinned_offset = 0 |
|
|
|
|
|
def fill(self, src_tensor, src_offset): |
|
|
copy_bytes = Base_IO_Buffer.fill_buffer(src_tensor, src_offset, self._pinned_tensor, self._pinned_offset) |
|
|
self._pinned_offset += copy_bytes |
|
|
return copy_bytes |
|
|
|
|
|
def drain(self, num_bytes, fd, file_offset): |
|
|
self._drain(num_bytes, fd, file_offset, blocking=True) |
|
|
self._pinned_offset = 0 |
|
|
|
|
|
def get_buffer(self): |
|
|
return self._pinned_tensor |
|
|
|
|
|
def get_offset(self): |
|
|
return self._pinned_offset |
|
|
|
|
|
def get_aligned_num_bytes(self): |
|
|
aligned_size = self._dnvme_handle.get_alignment() |
|
|
return (self._pinned_offset // aligned_size) * aligned_size |
|
|
|
|
|
def get_unaligned_num_bytes(self): |
|
|
return self._pinned_offset % self._dnvme_handle.get_alignment() |
|
|
|
|
|
def is_full(self): |
|
|
return self._pinned_offset == self._pinned_tensor.numel() |
|
|
|
|
|
def is_empty(self): |
|
|
return self._pinned_offset == 0 |
|
|
|
|
|
def reset(self): |
|
|
self._pinned_offset = 0 |
|
|
|
