diff --git a/test/test_subbuffer.py b/test/test_subbuffer.py index 9c9a34d5e7..c469a57423 100644 --- a/test/test_subbuffer.py +++ b/test/test_subbuffer.py @@ -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() diff --git a/tinygrad/device.py b/tinygrad/device.py index c912eb21a6..61fd207d4f 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -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"" @@ -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: