79 lines
2.2 KiB
Python
79 lines
2.2 KiB
Python
import unittest
|
|
from tinygrad.runtime.support.allocator import TLSFAllocator
|
|
|
|
class TestTLSFAllocator(unittest.TestCase):
|
|
def setUp(self):
|
|
self.allocator = TLSFAllocator(1024, block_size=16)
|
|
|
|
def test_basic_alloc_free(self):
|
|
addr1 = self.allocator.alloc(32)
|
|
self.assertEqual(addr1, 0)
|
|
|
|
addr2 = self.allocator.alloc(64)
|
|
self.assertEqual(addr2, 32)
|
|
|
|
self.allocator.free(addr1)
|
|
addr3 = self.allocator.alloc(32)
|
|
self.assertEqual(addr3, 0)
|
|
|
|
def test_block_size_alignment(self):
|
|
addr1 = self.allocator.alloc(20)
|
|
addr2 = self.allocator.alloc(35)
|
|
|
|
self.assertEqual(addr1 % 16, 0)
|
|
self.assertEqual(addr2 % 16, 0)
|
|
|
|
def test_merge_blocks(self):
|
|
addr1 = self.allocator.alloc(32)
|
|
addr2 = self.allocator.alloc(32)
|
|
self.allocator.alloc(32)
|
|
|
|
self.allocator.free(addr1)
|
|
self.allocator.free(addr2)
|
|
addr4 = self.allocator.alloc(64)
|
|
self.assertEqual(addr4, addr1)
|
|
|
|
def test_split_blocks(self):
|
|
addr1 = self.allocator.alloc(128)
|
|
self.allocator.free(addr1)
|
|
|
|
addr2 = self.allocator.alloc(32)
|
|
self.assertEqual(addr2, addr1)
|
|
|
|
addr3 = self.allocator.alloc(32)
|
|
self.assertEqual(addr3, addr1 + 32)
|
|
|
|
def test_out_of_memory(self):
|
|
with self.assertRaises(MemoryError):
|
|
self.allocator.alloc(2048)
|
|
|
|
def test_fragmentation_handling(self):
|
|
addrs = []
|
|
for _ in range(5):
|
|
addrs.append(self.allocator.alloc(32))
|
|
|
|
# Free alternate blocks
|
|
for i in range(0, len(addrs), 2):
|
|
self.allocator.free(addrs[i])
|
|
|
|
def test_custom_start_address(self):
|
|
allocator = TLSFAllocator(1024, start_addr=1000)
|
|
addr1 = allocator.alloc(32)
|
|
self.assertEqual(addr1, 1000)
|
|
|
|
addr2 = allocator.alloc(64)
|
|
self.assertEqual(addr2, 1032)
|
|
|
|
def test_block_tracking(self):
|
|
addr1 = self.allocator.alloc(32)
|
|
addr2 = self.allocator.alloc(64)
|
|
|
|
self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
|
|
self.assertTrue(addr2 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
|
|
|
|
self.allocator.free(addr1)
|
|
self.assertTrue(addr1 in [addr - self.allocator.start_addr for addr in self.allocator.blocks])
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|