From e25f86721ddecaed62db1db1e6ba230a3d7bc0f2 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sun, 16 Aug 2026 21:31:53 -0400 Subject: [PATCH] more torch backend fixups (#17562) --- extra/torch_backend/backend.py | 48 +++++++++++----------------------- extra/torch_backend/test.py | 38 +++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 33 deletions(-) diff --git a/extra/torch_backend/backend.py b/extra/torch_backend/backend.py index fda5b0a5ae..71a983390e 100644 --- a/extra/torch_backend/backend.py +++ b/extra/torch_backend/backend.py @@ -4,7 +4,7 @@ # A006 Lambda argument `input` is shadowing a Python builtin from tinygrad import Tensor, dtypes, Device from tinygrad.uop.ops import Ops, GroupOp -from tinygrad.helpers import getenv, prod, strides_for_shape, argfix +from tinygrad.helpers import getenv, prod, strides_for_shape import torch.lib TORCH_DEBUG = getenv("TORCH_DEBUG") import torch, pathlib, operator, functools, weakref @@ -99,38 +99,11 @@ def _apply_view_ops(target, ops): for fn, args, kwargs in ops: target = fn(target, *args, **kwargs) return target -# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h -def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None: - if not (req := argfix(*args)): return None - new_shape, infer_idx = [], -1 - for i, s in enumerate(req): - if s is None: s = shape[i] if i < len(shape) else None - if not isinstance(s, int): return None - if s == -1: - if infer_idx != -1: return None - infer_idx = len(new_shape) - new_shape.append(s) - total = prod(shape) - if infer_idx != -1: - known = prod(x for x in new_shape if x != -1) - if known == 0: - if total != 0: return None - new_shape[infer_idx] = 0 - else: new_shape[infer_idx] = total // known - return tuple(new_shape) if prod(new_shape) == total else None - -# TODO: can we get rid of this? only for test_flatten_reshape_add +# a chain of reshapes (and detaches, which move nothing) is undone by reshaping the value back to the base def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool: if not (ops := _get_view_ops(view)): return False - shapes = [base.shape] - for fn, args, _ in ops: - if fn is Tensor.detach: continue # detach leaves every element where it was (a tracked view only on torch<2.10) - if fn is not Tensor.reshape: return False - if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False - shapes.append(next_shape) - if shapes[-1] != view.shape: return False - for s in reversed(shapes[:-1]): val = val.reshape(s) - base.assign(val) + if any(fn not in (Tensor.reshape, Tensor.detach) for fn, _, _ in ops): return False + base.assign(val.reshape(base.shape)) return True def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None: @@ -445,6 +418,7 @@ def _linalg_svd(self, full_matrices=False): from torch._decomp import get_decompositions decomps = [ aten.native_layer_norm_backward, + aten.native_group_norm_backward, aten.linalg_cross, aten.addmm, aten.addcmul, @@ -479,7 +453,13 @@ decomps = [ aten._softmax_backward_data, aten.embedding_dense_backward, aten.linalg_vector_norm, aten.binary_cross_entropy, aten.binary_cross_entropy_backward, + # the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize + aten.mse_loss, aten.mse_loss_backward, + aten.smooth_l1_loss, aten.smooth_l1_loss_backward, aten.upsample_nearest2d.out, + # NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel + aten.upsample_bicubic2d.out, + aten._adaptive_avg_pool2d, # activations aten.hardswish, aten.hardswish_backward, aten.hardtanh, aten.hardtanh_backward, @@ -661,8 +641,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{ "aten.acos": Tensor.acos, "aten.any": Tensor.any, "aten.bitwise_not": Tensor.bitwise_not, - "aten.argmax": Tensor.argmax, - "aten.argmin": Tensor.argmin, + # tinygrad indexes with int32, torch's arg reduces return int64 + "aten.argmax": lambda self, dim=None, keepdim=False: self.argmax(dim, keepdim).cast(dtypes.int64), + "aten.argmin": lambda self, dim=None, keepdim=False: self.argmin(dim, keepdim).cast(dtypes.int64), "aten.asinh": Tensor.asinh, "aten.mul": Tensor.mul, "aten.atanh": Tensor.atanh, @@ -688,6 +669,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{ self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None, "device": _from_torch_device(device) if device else None}.items() if v is not None}), "aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)), + "aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)), "aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)), "aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)), "aten.nonzero": Tensor.nonzero, diff --git a/extra/torch_backend/test.py b/extra/torch_backend/test.py index 74886a608d..69d7ba1798 100644 --- a/extra/torch_backend/test.py +++ b/extra/torch_backend/test.py @@ -166,6 +166,15 @@ class TestTorchBackend(unittest.TestCase): expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32) np.testing.assert_equal(y3.cpu().numpy(), expected) + def test_argmax_argmin(self): + a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4) + c = a.cpu() + for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)), + (torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices), + (torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]: + self.assertEqual(got.dtype, want.dtype) # torch's arg reduces are int64, tinygrad's are int32 + np.testing.assert_equal(got.cpu().numpy(), want.numpy()) + def test_isfinite(self): a = torch.ones(4, device=device) np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True]) @@ -373,6 +382,22 @@ class TestTorchBackend(unittest.TestCase): for bwd_eps in [1e-5, 0.3]: for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3) + def test_groupnorm_backward(self): + def run(dev): + x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_() + w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_() + torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward() + return x.grad.cpu().numpy(), w.grad.cpu().numpy() + for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3) + + def test_mse_smooth_l1_loss_backward(self): + def run(dev, loss): + x = torch.arange(4., device=dev).requires_grad_() + loss(x, torch.ones(4, device=dev)).backward() + return x.grad.cpu().numpy() + for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]: + np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6) + def test_batchnorm_unsqueeze(self): bn = torch.nn.BatchNorm2d(4).to(device) x = torch.randn(8, 4, 3, 3, device=device) @@ -817,6 +842,19 @@ class TestTorchBackend(unittest.TestCase): ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False) np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5) + def test_interpolate_bicubic_area(self): + a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4) + for mode, scale in [("bicubic", 2.0), ("area", 0.5)]: + ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode) + np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4) + + @unittest.expectedFailure + def test_interpolate_bicubic_backward(self): + # the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does + # aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel + x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_() + torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward() + @unittest.expectedFailure def test_interpolate_inexact_scale(self): # torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here