| |
| |
|
|
| |
|
|
| import torch |
|
|
|
|
| class Base_IO_Buffer(object): |
|
|
| def __init__(self, pinned_tensor, dnvme_handle): |
| assert pinned_tensor.numel() % dnvme_handle.get_alignment() == 0 |
| self._dnvme_handle = dnvme_handle |
| self._pinned_tensor = pinned_tensor |
|
|
| def fill(self, src_tensor, src_offset): |
| pass |
|
|
| def drain(self, num_bytes, fd, file_offset): |
| pass |
|
|
| def is_empty(self): |
| pass |
|
|
| def is_full(self): |
| pass |
|
|
| def get_buffer(self): |
| pass |
|
|
| def get_offset(self): |
| pass |
|
|
| def get_aligned_num_bytes(self): |
| pass |
|
|
| def get_unaligned_num_bytes(self): |
| pass |
|
|
| def reset(self): |
| pass |
|
|
| def complete_ongoing_drain(self): |
| pass |
|
|
| def _drain(self, num_bytes, fd, file_offset, blocking=False): |
| assert num_bytes <= self.get_offset() |
| assert num_bytes % self._dnvme_handle.get_alignment() == 0 |
| buffer = self.get_buffer() |
| r = self._dnvme_handle.async_pwrite(torch.narrow(buffer, 0, 0, num_bytes), fd, file_offset) |
| assert 0 == r |
| if blocking: |
| assert 1 == self._dnvme_handle.wait() |
|
|
| @staticmethod |
| def fill_buffer(src_tensor, src_offset, buffer_tensor, buffer_offset): |
| src_bytes = src_tensor.numel() - src_offset |
| assert src_bytes > 0 |
|
|
| dst_bytes = buffer_tensor.numel() - buffer_offset |
| copy_bytes = min(src_bytes, dst_bytes) |
| assert (buffer_offset + copy_bytes) <= buffer_tensor.numel() |
|
|
| if copy_bytes > 0: |
| src_slice = torch.narrow(src_tensor, 0, src_offset, copy_bytes) |
| dst_slice = torch.narrow(buffer_tensor, 0, buffer_offset, copy_bytes) |
| dst_slice.data.copy_(src_slice.data) |
|
|
| return copy_bytes |
|
|
