diff --git a/examples/other_mnist/beautiful_mnist_torch.py b/examples/other_mnist/beautiful_mnist_torch.py index e1cef16d9c..227edec95a 100644 --- a/examples/other_mnist/beautiful_mnist_torch.py +++ b/examples/other_mnist/beautiful_mnist_torch.py @@ -1,5 +1,5 @@ -from tinygrad import dtypes, getenv -from tinygrad.helpers import trange, colored +from tinygrad import dtypes, getenv, Device +from tinygrad.helpers import trange, colored, DEBUG, temp from tinygrad.nn.datasets import mnist import torch from torch import nn, optim @@ -30,13 +30,15 @@ if __name__ == "__main__": import tinygrad.frontend.torch device = torch.device("tiny") else: - device = torch.device("mps") + device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu")) + if DEBUG >= 1: print(f"using torch backend {device}") X_train, Y_train, X_test, Y_test = mnist() X_train = torch.tensor(X_train.float().numpy(), device=device) Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=device) X_test = torch.tensor(X_test.float().numpy(), device=device) Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=device) + if getenv("TORCHVIZ"): torch.cuda.memory._record_memory_history() model = Model().to(device) optimizer = optim.Adam(model.parameters(), 1e-3) @@ -62,3 +64,6 @@ if __name__ == "__main__": if target := getenv("TARGET_EVAL_ACC_PCT", 0.0): if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green")) else: raise ValueError(colored(f"{test_acc=} < {target}", "red")) + if getenv("TORCHVIZ"): + torch.cuda.memory._dump_snapshot(fp:=temp("torchviz.pkl", append_user=True)) + print(f"saved torch memory snapshot to {fp}, view in https://pytorch.org/memory_viz") diff --git a/test/test_memory_planner.py b/test/test_memory_planner.py new file mode 100644 index 0000000000..dcd91569bd --- /dev/null +++ b/test/test_memory_planner.py @@ -0,0 +1,124 @@ +import unittest +from tinygrad import dtypes, Device +from tinygrad.device import Buffer +from tinygrad.engine.memory import _internal_memory_planner + +global_map = {} +def b(i, base=None, offset=0, pin=False, size=16): + global global_map + if i in global_map: return global_map[i] + global_map[i] = Buffer(Device.DEFAULT, size, dtypes.int8, base=global_map[base] if base is not None else None, offset=offset) + if pin: global_map[i].ref(1) + return global_map[i] + +def check_assign(buffers:list[list[Buffer]|tuple[Buffer, ...]]): + assigned = _internal_memory_planner(buffers, noopt_buffers=None) + + taken_parts = set() + first_appearance, last_appearance = {}, {} + for i,u in enumerate(buffers): + for buf in u: + if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue + if buf.base not in first_appearance: first_appearance[buf.base] = i + last_appearance[buf.base] = i + + for i,u in enumerate(buffers): + for buf in u: + if buf.is_allocated() or buf.base.is_allocated() or buf.lb_refcount > 0: continue + cur, base = assigned.get(buf, buf), assigned.get(buf.base, buf.base) + if buf._base is not None: + assert cur.base == base.base and cur.offset == buf.offset + base.offset, f"failed: {buf} {cur} {base} {buf.offset} {base.offset}" + else: + for part in taken_parts: + assert buf.base == part[3] or part[0] != cur.base or part[1] + part[2] <= cur.offset or part[1] >= cur.offset + buf.nbytes + if first_appearance[buf.base] == i: taken_parts.add((cur.base, cur.offset, buf.nbytes, buf.base)) + if last_appearance[buf.base] == i: taken_parts.remove((cur.base, cur.offset, buf.nbytes, buf.base)) + +class TestMemoryPlanner(unittest.TestCase): + def setUp(self): + global global_map + 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) + +if __name__ == "__main__": + unittest.main() diff --git a/tinygrad/viz/index.html b/tinygrad/viz/index.html index d4ebe19ccf..0533c543a5 100644 --- a/tinygrad/viz/index.html +++ b/tinygrad/viz/index.html @@ -15,11 +15,8 @@