IQ.Pilot Prebuilt Release @ ab07000

tinygrad_repo/test/external/external_test_tlsf.py

@@ -0,0 +1,78 @@
+import unittest
+from tinygrad.runtime.support.memory 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()