commit - release 1

src/kernbench/policy/address/allocator.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from kernbench.policy.address.phyaddr import PhysAddr
+
+
+class AllocationError(Exception):
+    pass
+
+
+@dataclass(frozen=True)
+class AddressConfig:
+    sip_count: int
+    cubes_per_sip: int
+    pes_per_cube: int
+    hbm_bytes_per_cube: int
+    hbm_slices_per_cube: int
+    tcm_bytes_per_pe: int
+    tcm_scheduler_reserved_bytes: int
+    sram_bytes_per_cube: int
+
+    @property
+    def hbm_slice_bytes(self) -> int:
+        return self.hbm_bytes_per_cube // self.hbm_slices_per_cube
+
+    @property
+    def tcm_allocatable_bytes(self) -> int:
+        return self.tcm_bytes_per_pe - self.tcm_scheduler_reserved_bytes
+
+
+class PEMemAllocator:
+    def __init__(
+        self, rack_id: int, sip_id: int, cube_id: int, pe_id: int, cfg: AddressConfig,
+    ) -> None:
+        self._rack_id = rack_id
+        self._sip_id = sip_id
+        self._cube_id = cube_id
+        self._pe_id = pe_id
+        self._cfg = cfg
+        self._hbm_cursor = 0
+        self._tcm_cursor = 0
+
+    def alloc_hbm(self, nbytes: int) -> PhysAddr:
+        if self._hbm_cursor + nbytes > self._cfg.hbm_slice_bytes:
+            raise AllocationError(
+                f"HBM overflow: need {nbytes}, "
+                f"available {self._cfg.hbm_slice_bytes - self._hbm_cursor}"
+            )
+        pa = PhysAddr.pe_hbm_addr(
+            rack_id=self._rack_id, sip_id=self._sip_id, cube_id=self._cube_id,
+            pe_id=self._pe_id, pe_local_hbm_offset=self._hbm_cursor,
+            slice_size_bytes=self._cfg.hbm_slice_bytes,
+        )
+        self._hbm_cursor += nbytes
+        return pa
+
+    def alloc_tcm(self, nbytes: int) -> PhysAddr:
+        if self._tcm_cursor + nbytes > self._cfg.tcm_allocatable_bytes:
+            raise AllocationError(
+                f"TCM overflow: need {nbytes}, "
+                f"available {self._cfg.tcm_allocatable_bytes - self._tcm_cursor}"
+            )
+        pa = PhysAddr.pe_tcm_addr(
+            rack_id=self._rack_id, sip_id=self._sip_id, cube_id=self._cube_id,
+            pe_id=self._pe_id, tcm_offset=self._tcm_cursor,
+        )
+        self._tcm_cursor += nbytes
+        return pa
+
+    @property
+    def hbm_used(self) -> int:
+        return self._hbm_cursor
+
+    @property
+    def hbm_total(self) -> int:
+        return self._cfg.hbm_slice_bytes
+
+    @property
+    def tcm_used(self) -> int:
+        return self._tcm_cursor
+
+    @property
+    def tcm_total(self) -> int:
+        return self._cfg.tcm_allocatable_bytes