mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-09 22:13:43 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
257 lines
8.0 KiB
Python
257 lines
8.0 KiB
Python
import unittest
|
|
from tinygrad import dtypes
|
|
from tinygrad.uop.ops import UOp, Ops
|
|
from tinygrad.schedule.memory import memory_plan_rewrite
|
|
|
|
global_map = {}
|
|
held_bufs: set[UOp] = set()
|
|
def b(i, base=None, offset=0, pin=False, size=16):
|
|
global global_map
|
|
if i in global_map: return global_map[i]
|
|
if base is not None:
|
|
global_map[i] = global_map[base]
|
|
return global_map[i]
|
|
global_map[i] = UOp.new_buffer("NULL", size, dtypes.int8)
|
|
if pin: held_bufs.add(global_map[i])
|
|
return global_map[i]
|
|
|
|
def _make_linear(buffer_lists, copies=None):
|
|
copy_pairs = {frozenset((id(dst), id(src))) for dst, src in copies} if copies else set()
|
|
calls = []
|
|
for bufs in buffer_lists:
|
|
is_copy = len(bufs) == 2 and frozenset((id(bufs[0]), id(bufs[1]))) in copy_pairs
|
|
if is_copy:
|
|
src0 = bufs[0].copy_to_device(bufs[1].device)
|
|
else:
|
|
src0 = UOp(Ops.SINK, src=tuple(bufs))
|
|
calls.append(src0.call(*bufs))
|
|
return UOp(Ops.LINEAR, src=tuple(calls))
|
|
|
|
def _get_planned_view(buf:UOp) -> tuple[UOp, int, int]|None:
|
|
view = buf.src[0] if buf.op is Ops.BITCAST else buf
|
|
if view.op is not Ops.SHRINK or view.src[0].op is not Ops.BUFFER: return None
|
|
return (arena:=view.src[0]), view.src[1].val * arena.dtype.itemsize, view.src[2].val * arena.dtype.itemsize
|
|
|
|
def _get_arena(buf, linear, result):
|
|
for orig_si, new_si in zip(linear.src, result.src):
|
|
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
|
if orig is buf and (planned:=_get_planned_view(new)) is not None: return planned[0]
|
|
return None
|
|
|
|
def check_assign(buffer_lists, copies=None):
|
|
linear = _make_linear(buffer_lists, copies)
|
|
result = memory_plan_rewrite(linear, held_bufs)
|
|
|
|
# build mapping: original buf -> (arena, offset_bytes, nbytes) from the result
|
|
replace_map: dict[int, tuple[UOp, int, int]] = {}
|
|
for orig_si, new_si in zip(linear.src, result.src):
|
|
for orig, new in zip(orig_si.src[1:], new_si.src[1:]):
|
|
if (planned:=_get_planned_view(new)) is not None and id(orig) not in replace_map:
|
|
replace_map[id(orig)] = planned
|
|
|
|
# verify pinned buffers are not planned
|
|
for buf in held_bufs:
|
|
assert id(buf) not in replace_map, "pinned buffer was planned"
|
|
|
|
# compute lifetimes
|
|
first_appearance, last_appearance = {}, {}
|
|
for i, bufs in enumerate(buffer_lists):
|
|
for buf in bufs:
|
|
if buf in held_bufs: continue
|
|
if id(buf) not in first_appearance: first_appearance[id(buf)] = i
|
|
last_appearance[id(buf)] = i
|
|
|
|
# verify non-overlapping: no two live buffers share the same arena region
|
|
taken_parts: set[tuple[int, int, int, int]] = set() # (id(arena), offset, nbytes, id(buf))
|
|
for i, bufs in enumerate(buffer_lists):
|
|
for buf in bufs:
|
|
if buf in held_bufs or id(buf) not in replace_map: continue
|
|
arena, off, nb = replace_map[id(buf)]
|
|
for part in taken_parts:
|
|
assert id(buf) == part[3] or part[0] != id(arena) or part[1] + part[2] <= off or part[1] >= off + nb, \
|
|
f"overlap at step {i}: [{off}, {off+nb}) conflicts with [{part[1]}, {part[1]+part[2]})"
|
|
if first_appearance.get(id(buf)) == i: taken_parts.add((id(arena), off, nb, id(buf)))
|
|
if last_appearance.get(id(buf)) == i: taken_parts.discard((id(arena), off, nb, id(buf)))
|
|
|
|
class TestMemoryPlanner(unittest.TestCase):
|
|
def setUp(self):
|
|
global global_map
|
|
held_bufs.clear()
|
|
global_map = {}
|
|
|
|
def test_simple_buffer(self):
|
|
bs = [
|
|
[b(0), b(1), b(2)],
|
|
[b(1), b(2), b(3)],
|
|
[b(4), b(3)],
|
|
[b(5), b(2)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_simple_pinned(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1), b(2, pin=True)],
|
|
[b(1), b(2), b(3)],
|
|
[b(4), b(3)],
|
|
[b(5), b(2)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_all_pinned(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1, pin=True)],
|
|
[b(1), b(2, pin=True)],
|
|
[b(4, pin=True), b(3, pin=True)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_simple_buffer_offset(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
|
|
[b(1), b(2), b(3, base=0, offset=1, size=8)],
|
|
[b(4), b(3)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_buffer_offset(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1, base=0, offset=1, size=8), b(2)],
|
|
[b(1), b(2), b(3, base=0, offset=1, size=8)],
|
|
[b(4), b(3)],
|
|
[b(5, base=2, offset=2, size=8), b(3)],
|
|
[b(6), b(5), b(0)],
|
|
[b(7), b(8, pin=True)],
|
|
[b(8), b(9, base=2, offset=2, size=8)],
|
|
[b(9), b(3), b(5)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_buffer_offset2(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1), b(2)],
|
|
[b(1), b(2), b(3)],
|
|
[b(4), b(3)],
|
|
[b(5), b(3)],
|
|
[b(6), b(5), b(0)],
|
|
[b(7), b(8, pin=True)],
|
|
[b(8), b(9)],
|
|
[b(9), b(3), b(5)],
|
|
[b(11), b(0)],
|
|
[b(11), b(10), b(5)],
|
|
[b(12), b(11), b(0)],
|
|
[b(6), b(12), b(7)],
|
|
[b(13), b(6), b(11)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_all_offsets_of_one(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1)],
|
|
[b(3, base=1, offset=0, size=8), b(2, base=0, offset=0, size=8)],
|
|
[b(5, base=1, offset=8, size=8), b(4, base=0, offset=8, size=8)],
|
|
[b(7, base=1, offset=4, size=8), b(6, base=0, offset=4, size=8)],
|
|
|
|
[b(4), b(5), b(2)],
|
|
[b(3), b(7)],
|
|
[b(10), b(6), b(7)],
|
|
[b(11), b(3), b(2)],
|
|
[b(12), b(5), b(4), b(3), b(2)],
|
|
[b(13), b(6), b(12), b(7)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_very_small_buffers(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1, size=32)],
|
|
[b(3, size=4), b(4, size=6)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_very_big_buffers(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1, size=34359738368000)],
|
|
[b(3, size=1 << 128), b(4, size=1 << 64)],
|
|
]
|
|
check_assign(bs)
|
|
|
|
def test_copy_bufs_separate_from_compute(self):
|
|
bs = [
|
|
[b(0), b(1)],
|
|
[b(1), b(2)],
|
|
[b(3), b(2)],
|
|
]
|
|
linear = _make_linear(bs, copies=[(b(1), b(0))])
|
|
result = memory_plan_rewrite(linear)
|
|
r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
|
|
assert r1_arena is not None and r2_arena is not None
|
|
assert r1_arena is not r2_arena
|
|
|
|
def test_copy_bufs_reuse_among_copies(self):
|
|
bs = [
|
|
[b(0), b(1)],
|
|
[b(2), b(1)],
|
|
[b(3), b(2)],
|
|
]
|
|
linear = _make_linear(bs, copies=[(b(1), b(0)), (b(2), b(1))])
|
|
result = memory_plan_rewrite(linear)
|
|
r1_arena, r2_arena = _get_arena(b(1), linear, result), _get_arena(b(2), linear, result)
|
|
assert r1_arena is not None and r2_arena is not None
|
|
assert r1_arena is r2_arena
|
|
|
|
def test_compute_bufs_reuse_among_compute(self):
|
|
bs = [
|
|
[b(0), b(1)],
|
|
[b(2), b(1)],
|
|
[b(3), b(2)],
|
|
[b(4), b(3)],
|
|
]
|
|
linear = _make_linear(bs, copies=[(b(1), b(0))])
|
|
result = memory_plan_rewrite(linear)
|
|
r2_arena, r3_arena = _get_arena(b(2), linear, result), _get_arena(b(3), linear, result)
|
|
assert r2_arena is not None and r3_arena is not None
|
|
assert r2_arena is r3_arena
|
|
|
|
def test_copy_and_compute_no_cross_reuse(self):
|
|
bs = [
|
|
[b(0), b(1)],
|
|
[b(2), b(1)],
|
|
[b(3), b(2)],
|
|
]
|
|
linear = _make_linear(bs, copies=[(b(2), b(1))])
|
|
result = memory_plan_rewrite(linear)
|
|
r0_arena, r2_arena = _get_arena(b(0), linear, result), _get_arena(b(2), linear, result)
|
|
assert r0_arena is not None and r2_arena is not None
|
|
assert r0_arena is not r2_arena
|
|
|
|
def test_multiple_copy_bufs_with_offsets(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1), b(2)],
|
|
[b(3, base=0, offset=1, size=8), b(1), b(2)],
|
|
[b(4), b(3)],
|
|
[b(5), b(4)],
|
|
]
|
|
check_assign(bs, copies=[(b(1), b(0)), (b(2), b(0))])
|
|
|
|
def test_copy_bufs_pinned_mixed(self):
|
|
bs = [
|
|
[b(0, pin=True), b(1), b(2)],
|
|
[b(1), b(3), b(2)],
|
|
[b(4), b(3)],
|
|
[b(5), b(4), b(0)],
|
|
]
|
|
check_assign(bs, copies=[(b(1), b(0)), (b(3), b(1))])
|
|
|
|
def test_deferred_copy_frees_chain(self):
|
|
bs = []
|
|
copies = []
|
|
for i in range(6):
|
|
copy_buf, compute_buf = b(i * 2 + 1), b(i * 2 + 2)
|
|
bs.append([copy_buf, b(0, pin=True)])
|
|
bs.append([compute_buf, copy_buf])
|
|
copies.append((copy_buf, b(0, pin=True)))
|
|
bs.append([b(100, pin=True)])
|
|
check_assign(bs, copies=copies)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|