torch compile

examples/other_mnist/beautiful_mnist_torch.py

@@ -28,6 +28,7 @@ def forward(self, x):
 if __name__ == "__main__":
   if getenv("TINY_BACKEND"):
     import tinygrad.nn.torch  # noqa: F401
+    import extra.torch_backend.compile  # noqa: F401
     device = torch.device("tiny")
   else:
     device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
@@ -43,7 +44,7 @@ def forward(self, x):
   optimizer = optim.Adam(model.parameters(), 1e-3)
 
   loss_fn = nn.CrossEntropyLoss()
-  #@torch.compile
+  @torch.compile
   def step(samples):
     X,Y = X_train[samples], Y_train[samples]
     out = model(X)

extra/torch_backend/backend.py

@@ -55,22 +55,36 @@ def device_count(self): return getenv("GPUS", 1) # TODO: device count in tiny?
 torch.utils.generate_methods_for_privateuse1_backend()
 aten = torch.ops.aten
 
+# a tiny tensor has no Storage at all, but dynamo fakifies its inputs through untyped_storage(). hand it a meta storage,
+# torch's own "a storage with no data", which is what MetaConverter rebuilds from it anyway. dynamo compares storage
+# identity to decide if two tensors alias, so it is cached on the tinygrad Tensor: t, t.detach() and nn.Parameter(t)
+# are different torch tensors over one tinygrad Tensor and do alias
+_untyped_storage = torch.Tensor.untyped_storage
+def untyped_storage(self:torch.Tensor):
+  # only a tiny tensor is storage-less by our doing: sparse and nested ones are torch's, and keep torch's own error
+  if torch._C._has_storage(self) or self.device.type != "tiny": return _untyped_storage(self)
+  if getattr(x:=unwrap(self), "_meta_storage", None) is None: x._meta_storage = torch.UntypedStorage(self.nbytes, device="meta")
+  return x._meta_storage
+torch.Tensor.untyped_storage = untyped_storage
+
 # track view relationships for in place operations
 def canonical_base(view: Tensor): return getattr(view, "_view_base", view)
 def derived_views(base: Tensor): return [t for tref in getattr(base, "_views", set()) if (t:=tref()) is not None]
 def unwrap_args(args, kwargs):
   return [unwrap(x) if isinstance(x, torch.Tensor) else x for x in args], {k:unwrap(v) if isinstance(v, torch.Tensor) else v for k,v in kwargs.items()}
+# record that `view` is `base` with `ops` applied, so a later in place write to either goes through the base
+def register_view(base: Tensor, view: Tensor, ops) -> Tensor:
+  view._view_base, view._view_ops = base, ops
+  base._views = getattr(base, "_views", set())
+  base._views.add(weakref.ref(view))
+  return view
+
 def wrap_view_op(fn):
   @functools.wraps(fn)
   def _wrap(*args, **kwargs):
     args, kwargs = unwrap_args(args, kwargs)
     ret = fn(*args, **kwargs)
-    base = canonical_base(args[0])
-    ret._view_base = base
-    base._views = getattr(base, "_views", set())
-    base._views.add(weakref.ref(ret))
-    ret._view_ops = _get_view_ops(args[0]) + [(fn, args[1:], kwargs)]
-    return wrap(ret)
+    return wrap(register_view(canonical_base(args[0]), ret, _get_view_ops(args[0]) + [(fn, args[1:], kwargs)]))
   return _wrap
 
 # NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
@@ -349,6 +363,8 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
 from torch._decomp import get_decompositions
 decomps = [
   aten.native_layer_norm_backward,
+  # functionalizing a training BatchNorm under torch.compile gives this instead of aten.native_batch_norm
+  aten._native_batch_norm_legit_functional,
   aten.native_group_norm_backward,
   aten.linalg_cross,
   aten.addmm,

extra/torch_backend/compile.py

@@ -0,0 +1,52 @@
+# a dynamo backend that lowers the AOT-autograd FX graph into a TinyJit, so a compiled function replays kernels
+# instead of dispatching op-by-op. importing this makes a bare torch.compile use it
+import torch
+from torch._dynamo.backends.registry import register_backend, set_default_backend
+from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_func
+from tinygrad import Tensor, TinyJit
+from extra.torch_backend.backend import wrap, unwrap, canonical_base, register_view, _get_view_ops
+
+# TinyJit is static shape, and a dim dynamo marked dynamic arrives as a SymInt with no Tensor to bind. NOTE: process wide
+torch._dynamo.config.automatic_dynamic_shapes = False
+
+def _jit_input(x:torch.Tensor) -> Tensor:
+  # a JIT input has to be a real buffer: .tiny() brings other devices over, .clone() gives a deviceless const a device
+  t = unwrap(x.tiny())
+  return t if t.device is not None else t.replace(t.clone(t._torch_device))
+
+def _copy_out(outs:list[Tensor|None]) -> list[Tensor|None]:
+  # copy every output, recording a returned view as a view of its base's copy. sharing one buffer instead would hand the
+  # next graph two JIT inputs on it, which jit.py rejects. one realize for all: per-output is a graph rewrite each
+  ret:list[Tensor|None] = [None if t is None else t.clone() for t in outs]
+  copies = {t: c for t, c in zip(outs, ret) if t is not None and t is canonical_base(t)}
+  for t, c in zip(outs, ret):
+    if t is None or t is (b:=canonical_base(t)): continue
+    if (base_copy:=copies.get(b)) is not None: register_view(base_copy, c, _get_view_ops(t))
+  if (real:=[t for t in ret if t is not None]): Tensor.realize(*real)
+  return ret
+
+def _tiny_compiler(gm:torch.fx.GraphModule, sample_inputs):
+  # dynamo emits op-free graphs at training graph breaks and TinyJit cannot capture one, so run it eagerly
+  if not any(n.op in ("call_function", "call_method", "call_module") for n in gm.graph.nodes):
+    return make_boxed_func(torch.compiler.disable(gm))
+  # the JIT hands back the buffers it captured: copy inside so a passthrough output is written on this call, copy outside
+  # so a retained output survives the next. a backward graph has a None output per input that wanted no grad
+  @TinyJit
+  def jitted(*args:Tensor): return _copy_out([x if x is None else unwrap(x) for x in gm(*[wrap(a) for a in args])])
+  # this runs under an active dynamo, which would otherwise trace the backend itself
+  @torch.compiler.disable
+  def run(*args:torch.Tensor):
+    return [t if t is None else wrap(t) for t in _copy_out(jitted(*map(_jit_input, args)))]
+  return make_boxed_func(run)
+
+# mode= and options= are inductor knobs torch hands to any named backend, drop them like torch's own aot_autograd does
+@register_backend
+def tiny(gm:torch.fx.GraphModule, sample_inputs, **_ignored):
+  # a non-Tensor input is a SymInt for a dynamic dim. NOTE: never format a tiny tensor in here, repr() of one runs a
+  # float64 reduction that METAL cannot compile
+  if not all(isinstance(x, torch.Tensor) for x in sample_inputs):
+    raise RuntimeError(f"the tiny backend needs static shapes, got {[type(x).__name__ for x in sample_inputs]}")
+  return aot_module_simplified(gm, sample_inputs, fw_compiler=_tiny_compiler)
+
+# torch.compile defaults to inductor, which cannot run a tiny tensor. NOTE: process wide, like the config above
+set_default_backend("tiny")

extra/torch_backend/torch_tests.py

@@ -3,8 +3,9 @@
 import torch
 import extra.torch_backend.backend
 
-from torch.testing._internal.common_utils import TestCase, is_privateuse1_backend_available
-assert is_privateuse1_backend_available() and torch._C._get_privateuse1_backend_name() == "tiny"
+from torch.testing._internal.common_utils import TestCase
+# NOTE: this is what torch's own is_privateuse1_backend_available did, inlined: it went private in torch 2.13
+assert torch.tiny.is_available() and torch._C._get_privateuse1_backend_name() == "tiny"
 from torch.testing._internal.common_device_type import ops, onlyOn, instantiate_device_type_tests
 from torch.testing._internal.common_methods_invocations import unary_ufuncs, binary_ufuncs, reduction_ops, shape_funcs
 

pyproject.toml

@@ -71,7 +71,7 @@ linting = [
 # ]
 testing_minimal = [
   "numpy",
-  "torch==2.9.1",
+  "torch==2.13.0",
   "pytest",
   "pytest-xdist",
   "pytest-split",