A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying memory. It must have a unique type to support this behavior.
| 22 | |
| 23 | |
| 24 | class MemoryBuffer: |
| 25 | """ |
| 26 | A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying |
| 27 | memory. It must have a unique type to support this behavior. |
| 28 | """ |
| 29 | |
| 30 | def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype, source: Optional[torch.Tensor] = None): |
| 31 | self.numel = numel |
| 32 | self.numel_padded = numel_padded |
| 33 | self.dtype = dtype |
| 34 | if source is not None: |
| 35 | self.data = source |
| 36 | else: |
| 37 | self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device='cuda', requires_grad=False) |
| 38 | |
| 39 | def zero(self): |
| 40 | """Reset the buffer to zero.""" |
| 41 | self.data.zero_() |
| 42 | |
| 43 | def get(self, shape, start_index): |
| 44 | """Return a tensor with the input `shape` as a view into the |
| 45 | 1-D data starting at `start_index`.""" |
| 46 | end_index = start_index + shape.numel() |
| 47 | assert end_index <= self.numel, \ |
| 48 | 'requested tensor is out of the buffer range.' |
| 49 | buffer_tensor = self.data[start_index:end_index] |
| 50 | buffer_tensor = buffer_tensor.view(shape) |
| 51 | return buffer_tensor |
| 52 | |
| 53 | |
| 54 | def calc_padded_numel(shape: torch.Size, dtype: torch.dtype): |
no outgoing calls
no test coverage detected