File size: 1,324 Bytes
9c6594c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
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
|