forked from tinygrad/tinygrad
support copyin/copyout/is_allocated for subbuffers (#10869)
* support copyin/copyout/is_allocated for subbuffers * simple * clean up * rm underlying_buf * add function is_initialized * add tests * better test_subbuffer_copy_in_out * fix allocator --------- Co-authored-by: b1tg <[email protected]> Co-authored-by: George Hotz <[email protected]>
This commit is contained in:
b1tg
@@ -9,6 +9,7 @@ class TestSubBuffer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.buf = Buffer(Device.DEFAULT, 10, dtypes.uint8).ensure_allocated()
|
||||
self.buf.copyin(memoryview(bytearray(range(10))))
|
||||
self.buf_unalloc = Buffer(Device.DEFAULT, 10, dtypes.uint8)
|
||||
|
||||
def test_subbuffer(self):
|
||||
vbuf = self.buf.view(2, dtypes.uint8, offset=3).ensure_allocated()
|
||||
@@ -64,5 +65,120 @@ class TestSubBuffer(unittest.TestCase):
|
||||
tst = vbuf.as_buffer().tolist()
|
||||
assert tst == [13, 14]
|
||||
|
||||
def test_subbuffer_is_allocated(self):
|
||||
buf = self.buf_unalloc
|
||||
sub_buf = buf.view(3, dtypes.uint8, offset=4)
|
||||
self.assertFalse(buf.is_allocated())
|
||||
self.assertFalse(buf.is_initialized())
|
||||
self.assertFalse(sub_buf.is_allocated())
|
||||
self.assertFalse(sub_buf.is_initialized())
|
||||
|
||||
# base buffer alloc
|
||||
buf.allocate()
|
||||
self.assertTrue(buf.is_allocated())
|
||||
self.assertTrue(buf.is_initialized())
|
||||
self.assertTrue(sub_buf.is_allocated())
|
||||
self.assertFalse(sub_buf.is_initialized())
|
||||
|
||||
# sub buffer alloc
|
||||
sub_buf.allocate()
|
||||
self.assertTrue(sub_buf.is_initialized())
|
||||
|
||||
# sub buffer dealloc
|
||||
sub_buf.deallocate()
|
||||
self.assertTrue(buf.is_allocated())
|
||||
self.assertTrue(buf.is_initialized())
|
||||
self.assertTrue(sub_buf.is_allocated())
|
||||
self.assertFalse(sub_buf.is_initialized())
|
||||
|
||||
# base buffer dealloc
|
||||
buf.deallocate()
|
||||
self.assertFalse(buf.is_allocated())
|
||||
self.assertFalse(buf.is_initialized())
|
||||
self.assertFalse(sub_buf.is_allocated())
|
||||
self.assertFalse(sub_buf.is_initialized())
|
||||
|
||||
# sub buffer alloc
|
||||
sub_buf.ensure_allocated()
|
||||
self.assertTrue(buf.is_allocated())
|
||||
self.assertTrue(buf.is_initialized())
|
||||
self.assertTrue(sub_buf.is_allocated())
|
||||
self.assertTrue(sub_buf.is_initialized())
|
||||
|
||||
def test_subbuffer_copy_in_out(self):
|
||||
sub_buf = self.buf.view(3, dtypes.uint8, offset=3).ensure_allocated() # [3:6]
|
||||
data_out_sub = bytearray([0]*3)
|
||||
sub_buf.copyout(memoryview(data_out_sub))
|
||||
assert data_out_sub == bytearray(range(3, 6))
|
||||
sub_buf.copyin(memoryview(bytearray(range(3))))
|
||||
assert sub_buf.as_buffer().tolist() == list(range(3))
|
||||
assert self.buf.as_buffer().tolist()[3:6] == list(range(3))
|
||||
sub_buf.copyout(memoryview(data_out_sub))
|
||||
assert data_out_sub == bytearray(range(3))
|
||||
data_out_base = bytearray([0]*10)
|
||||
self.buf.copyout(memoryview(data_out_base))
|
||||
assert data_out_base[0:3] == bytearray(range(0, 3))
|
||||
assert data_out_base[3:6] == data_out_sub
|
||||
assert data_out_base[6:10] == bytearray(range(6, 10))
|
||||
|
||||
def test_subbuffer_copy_in_out_view_of_view(self):
|
||||
view1 = self.buf.view(7, dtypes.uint8, offset=2).ensure_allocated() # [2:9]
|
||||
view2 = view1.view(3, dtypes.uint8, offset=2).ensure_allocated() # [4:7]
|
||||
self.assertTrue(view1.is_allocated())
|
||||
self.assertTrue(view2.is_allocated())
|
||||
|
||||
data_in = bytearray([7, 8, 9])
|
||||
view2.copyin(memoryview(data_in))
|
||||
data_out_v2 = bytearray([0]*3)
|
||||
view2.copyout(memoryview(data_out_v2))
|
||||
assert data_in == data_out_v2
|
||||
|
||||
expected_base_data = bytearray([0]*10)
|
||||
expected_base_data = memoryview(bytearray(range(10)))
|
||||
expected_base_data[4:7] = data_in
|
||||
|
||||
data_out_base = bytearray([0]*10)
|
||||
self.buf.copyout(memoryview(data_out_base))
|
||||
assert expected_base_data == data_out_base
|
||||
|
||||
def test_subbuffer_alloc(self):
|
||||
sub_buf = self.buf.view(4, dtypes.int8, offset=3)
|
||||
sub_buf.allocate()
|
||||
sub_buf.copyin(memoryview(bytearray(range(10, 14))))
|
||||
assert self.buf.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
|
||||
|
||||
sub_buf = self.buf_unalloc.view(4, dtypes.int8, offset=3)
|
||||
sub_buf.allocate()
|
||||
sub_buf.copyin(memoryview(bytearray(range(10, 14))))
|
||||
assert self.buf_unalloc.as_buffer().tolist()[3:7] == sub_buf.as_buffer().tolist()
|
||||
|
||||
def test_subbuffer_dealloc(self):
|
||||
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
|
||||
sub_buf.deallocate()
|
||||
assert self.buf.as_buffer().tolist() == list(range(10))
|
||||
|
||||
def test_subbuffer_double_dealloc(self):
|
||||
sub_buf = self.buf.view(3, dtypes.uint8, offset=4).ensure_allocated()
|
||||
self.buf.deallocate()
|
||||
with self.assertRaises(AssertionError):
|
||||
self.buf.deallocate()
|
||||
sub_buf.deallocate()
|
||||
with self.assertRaises(AssertionError):
|
||||
sub_buf.deallocate()
|
||||
|
||||
def test_subbuffer_uaf(self):
|
||||
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
|
||||
assert self.buf.as_buffer().tolist(), list(range(10))
|
||||
sub_buf.deallocate()
|
||||
with self.assertRaises(AssertionError):
|
||||
sub_buf.as_buffer().tolist()
|
||||
assert self.buf.as_buffer().tolist(), list(range(10))
|
||||
|
||||
sub_buf = self.buf.view(4, dtypes.int8, offset=3).ensure_allocated()
|
||||
assert sub_buf.as_buffer().tolist(), list(range(3, 7))
|
||||
self.buf.deallocate()
|
||||
with self.assertRaises(AssertionError):
|
||||
sub_buf.as_buffer().tolist()
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
|
||||
+10
-7
@@ -130,10 +130,13 @@ class Buffer:
|
||||
def ref(self, cnt):
|
||||
self.base._uop_refcount += cnt
|
||||
return self
|
||||
def is_allocated(self) -> bool: return hasattr(self, '_buf')
|
||||
def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_allocated() else self
|
||||
# check if the underlying buffer is allocated and the current buffer/view is initialized
|
||||
def is_initialized(self) -> bool: return self.is_allocated() and hasattr(self, '_buf')
|
||||
# check if the underlying buffer is allocated, possibly from the base object
|
||||
def is_allocated(self) -> bool: return self.base.is_allocated() if self._base is not None else hasattr(self, '_buf')
|
||||
def ensure_allocated(self) -> Buffer: return self.allocate() if not self.is_initialized() else self
|
||||
def allocate(self, opaque=None, external_ptr=None) -> Buffer:
|
||||
assert not self.is_allocated(), "can't allocate already allocated buffer"
|
||||
assert not self.is_initialized(), "can't allocate already allocated buffer"
|
||||
if DEBUG >= 7: print(f"buffer: allocate {self.nbytes} bytes on {self.device}")
|
||||
if (mbs:=getenv("MAX_BUFFER_SIZE", 0)) > 0 and self.size > mbs: raise RuntimeError(f"buffer of size {self.size/1e6:.2f}M is too large")
|
||||
self.allocator:Allocator = Device[self.device].allocator
|
||||
@@ -149,7 +152,7 @@ class Buffer:
|
||||
if not self.device.startswith("DISK"): GlobalCounters.mem_used += self.nbytes
|
||||
return self
|
||||
def deallocate(self):
|
||||
assert self.is_allocated(), "buffer must be allocated to deallocate"
|
||||
assert hasattr(self, '_buf'), "buffer must be allocated to deallocate"
|
||||
if DEBUG is not None and DEBUG >= 7: print(f"buffer: deallocate {self.nbytes} bytes on {self.device}")
|
||||
if self._base is None and (self.options is None or self.options.external_ptr is None):
|
||||
if GlobalCounters is not None and not self.device.startswith("DISK"): GlobalCounters.mem_used -= self.nbytes
|
||||
@@ -167,7 +170,7 @@ class Buffer:
|
||||
return self.__class__, (self.device, self.size, self.dtype, None, self.options, buf, self.uop_refcount)
|
||||
@property
|
||||
def nbytes(self): return self.size*self.dtype.itemsize
|
||||
def __del__(self): (not self.is_allocated()) or self.deallocate()
|
||||
def __del__(self): (not hasattr(self, '_buf')) or self.deallocate()
|
||||
def __repr__(self):
|
||||
return f"<buf real:{self.is_allocated()} device:{self.device} size:{self.size} dtype:{self.dtype}" + \
|
||||
(f" offset:{self.offset}" if self._base is not None else "") + (f" {self.options=}" if self.options is not None else "") + ">"
|
||||
@@ -188,13 +191,13 @@ class Buffer:
|
||||
def copyin(self, mv:memoryview):
|
||||
mv = flat_mv(mv)
|
||||
assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
|
||||
assert self.is_allocated(), "can't copyin to unallocated buffer"
|
||||
assert self.is_initialized(), "can't copyin to unallocated buffer"
|
||||
self.allocator._copyin(self._buf, mv)
|
||||
return self
|
||||
def copyout(self, mv:memoryview) -> memoryview:
|
||||
mv = flat_mv(mv)
|
||||
assert len(mv) == self.nbytes, f"size mismatch, {len(mv)=} != {self.dtype=} {self.size=}"
|
||||
assert self.is_allocated(), "can't copyout unallocated buffer"
|
||||
assert self.is_initialized(), "can't copyout unallocated buffer"
|
||||
self.allocator._copyout(mv, self._buf)
|
||||
return mv
|
||||
def view(self, size:int, dtype:DType, offset:int) -> Buffer:
|
||||
|
||||
Reference in New Issue
Block a user