Merge branch 'master' into moveleftright

.github/workflows/test.yml

@@ -870,29 +870,29 @@ jobs:
     - name: Test ONNX Runner (WEBGPU)
       run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
 
-  #osxremote:
-  #  name: MacOS (remote metal)
-  #  runs-on: macos-15
-  #  timeout-minutes: 10
-  #  env:
-  #    REMOTE: 1
-  #    REMOTEDEV: METAL
-  #  steps:
-  #    - name: Checkout Code
-  #      uses: actions/checkout@v4
-  #    - name: Setup Environment
-  #      uses: ./.github/actions/setup-tinygrad
-  #      with:
-  #        key: macos-remote
-  #        deps: testing_minimal
-  #    - name: Check Device.DEFAULT and print some source
-  #      run: |
-  #        python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-  #        python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
-  #        DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-  #    - name: Run REMOTE=1 Test
-  #      run: |
-  #        python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+  osxremote:
+    name: MacOS (remote metal)
+    runs-on: macos-15
+    timeout-minutes: 10
+    env:
+      REMOTE: 1
+      REMOTEDEV: METAL
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: macos-remote
+          deps: testing_minimal
+      - name: Check Device.DEFAULT and print some source
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
 
   amdremote:
     name: Linux (remote)

setup.py

@@ -9,7 +9,7 @@ with open(directory / 'README.md', encoding='utf-8') as f:
 
 testing_minimal = [
   "numpy",
-  "torch==2.7.1",
+  "torch",
   "pytest",
   "pytest-xdist",
   "hypothesis",

test/test_jit.py

@@ -5,10 +5,9 @@ import numpy as np
 from hypothesis import given, settings, strategies as strat
 from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV
 from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
-from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
+from tinygrad.engine.jit import TinyJit
 from tinygrad.device import Device
-from tinygrad.helpers import Context, JIT, GlobalCounters, getenv
+from tinygrad.helpers import Context, JIT, GlobalCounters
 from tinygrad.dtype import dtypes
 from extra.models.unet import ResBlock
 
@@ -670,169 +669,5 @@ class TestJitFree(unittest.TestCase):
     out = fxn(Tensor([11,1,2,3,4]))
     self.assertEqual(out.item(), 13600)
 
-class TestJitGraphSplit(unittest.TestCase):
-  def compute(self, device, inp):
-    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
-    return (inp + 1.0).contiguous().realize()
-
-  def copy(self, device, to_device, inp):
-    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
-    return inp.to(to_device).realize()
-
-  def expect(self, f, *args, graph=None, multigraph=None, hcqgraph=None):
-    def _numpies(tpl): return tpl.numpy() if tpl.__class__ is Tensor else tuple([t.numpy() for t in tpl])
-
-    expected = _numpies(f(*args))
-    for i in range(4):
-      res = _numpies(f(*args))
-      np.testing.assert_allclose(res, expected, atol=1e-4, rtol=1e-5)
-
-    dev = Device[Device.DEFAULT]
-    graph_t = graph_class(dev)
-    if graph_t is None: return
-
-    got = f.jit_cache
-    from tinygrad.runtime.graph.hcq import HCQGraph
-    if graph_t is HCQGraph:
-      validate = hcqgraph
-    elif issubclass(graph_t, MultiGraphRunner):
-      validate = multigraph
-    else:
-      validate = graph
-
-    assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
-    for expected, got in zip(validate, got):
-      if expected["type"] == "graph":
-        assert isinstance(got.prg, GraphRunner), f"Expected GraphRunner, got {type(got.prg)}"
-        assert len(got.prg.jit_cache) == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {len(got.prg.jit_cache)}"
-      elif expected["type"] == "comp":
-        assert isinstance(got.prg, CompiledRunner), f"Expected CompiledRunner, got {type(got.prg)}"
-      elif expected["type"] == "copy":
-        assert isinstance(got.prg, BufferCopy), f"Expected BufferCopy, got {type(got.prg)}"
-      elif expected["type"] == "xfer":
-        assert isinstance(got.prg, BufferXfer), f"Expected BufferXfer, got {type(got.prg)}"
-
-  def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
-  def ji_comp(self): return {"type": "comp"}
-  def ji_copy(self): return {"type": "copy"}
-  def ji_xfer(self): return {"type": "xfer"}
-
-  def test_jit_split_simple(self):
-    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")
-
-    @TinyJit
-    def f(inp):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.compute(Device.DEFAULT, op1)
-      return op2
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    self.expect(f, inp,
-      graph=[self.ji_graph(3)],
-      multigraph=[self.ji_graph(3)],
-      hcqgraph=[self.ji_graph(3)])
-
-  def test_jit_cpu_simple(self):
-    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
-
-    @TinyJit
-    def f(inp, inp_cpu):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.compute("CPU", inp_cpu)
-      op3 = self.compute(Device.DEFAULT, op1)
-      return op2, op3
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
-    self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_comp(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(4)])
-
-  def test_jit_cpu_several(self):
-    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
-
-    @TinyJit
-    def f(inp, inp_cpu):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.compute("CPU", inp_cpu)
-      op3 = self.compute("CPU", op2)
-      op4 = self.compute(Device.DEFAULT, op1)
-      return op3, op4
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    inp_cpu = Tensor.randn(10, 10, device="CPU").realize()
-    self.expect(f, inp, inp_cpu,
-      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      hcqgraph=[self.ji_graph(5)])
-
-  def test_jit_multidev(self):
-    if Device.DEFAULT == "CPU": raise unittest.SkipTest("CPU is not a valid default device for this test")
-
-    try: Device[f"{Device.DEFAULT}:1"]
-    except Exception: raise unittest.SkipTest("no multidevice")
-
-    @TinyJit
-    def f(inp, inp_d1):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
-      op3 = self.compute(f"{Device.DEFAULT}:1", op2)
-      op4 = self.compute(Device.DEFAULT, op1)
-      return op3, op4
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
-    self.expect(f, inp, inp_d1,
-      graph=[self.ji_graph(2), self.ji_graph(2), self.ji_comp()],
-      multigraph=[self.ji_graph(5)],
-      hcqgraph=[self.ji_graph(5)])
-
-  def test_jit_multidev_xfer(self):
-    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
-
-    try: Device[f"{Device.DEFAULT}:1"]
-    except Exception: raise unittest.SkipTest("no multidevice")
-
-    @TinyJit
-    def f(inp, inp_d1):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.compute(f"{Device.DEFAULT}:1", inp_d1)
-      op3 = self.copy(f"{Device.DEFAULT}:1", Device.DEFAULT, op2)
-      op4 = self.compute(f"{Device.DEFAULT}:1", op2)
-      op5 = self.compute(Device.DEFAULT, op3)
-      return op1, op4, op5
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    inp_d1 = Tensor.randn(10, 10, device=f"{Device.DEFAULT}:1").realize()
-    self.expect(f, inp, inp_d1,
-      graph=[self.ji_graph(2), self.ji_comp(), self.ji_xfer(), self.ji_comp(), self.ji_comp()],
-      multigraph=[self.ji_graph(6)],
-      hcqgraph=[self.ji_graph(6)])
-
-  @unittest.skipIf(getenv("MOCKGPU"), "MockGPU does not support parallel copies")
-  def test_jit_multidev_copy(self):
-    if Device.DEFAULT in {"CPU", "LLVM"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
-    if Device.DEFAULT == "REMOTE": raise unittest.SkipTest("REMOTE gpu is broken")
-
-    @TinyJit
-    def f(inp):
-      op0 = self.compute(Device.DEFAULT, inp)
-      op1 = self.compute(Device.DEFAULT, op0)
-      op2 = self.copy(Device.DEFAULT, "CPU", op1)
-      op3 = self.compute("CPU", op2)
-      return op3
-
-    inp = Tensor.randn(10, 10, device=Device.DEFAULT).realize()
-    self.expect(f, inp,
-      graph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
-      multigraph=[self.ji_graph(2), self.ji_copy(), self.ji_comp()],
-      hcqgraph=[self.ji_graph(4)])
-
 if __name__ == '__main__':
   unittest.main()

test/test_nn.py

@@ -108,39 +108,105 @@ class TestNN(unittest.TestCase):
     _test_linear(Tensor.randn(BS, in_dim), in_dim, out_dim)
     _test_linear(Tensor.randn(BS, T, in_dim), in_dim, out_dim) # test with more dims
 
-  def _test_conv(self, tiny_conv, torch_conv, BS, C1, DIMS, C2, K, S, P, D=1):
+  def test_conv1d(self):
+    BS, C1, W = 4, 16, 224//4
+    C2, K, S, P = 64, 7, 2, 1
+
     # create in tinygrad
-    layer = tiny_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D)
+    layer = Conv1d(C1, C2, kernel_size=K, stride=S, padding=P)
 
     # create in torch
     with torch.no_grad():
-      torch_layer = torch_conv(C1, C2, kernel_size=K, stride=S, padding=P, dilation=D).eval()
+      torch_layer = torch.nn.Conv1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
       torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
       torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
 
     # test
-    x = Tensor.uniform(BS, C1, *DIMS)
+    x = Tensor.uniform(BS, C1, W)
     z = layer(x)
     torch_x = torch.tensor(x.numpy())
     torch_z = torch_layer(torch_x)
     np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
 
-  def test_conv1d(self): self._test_conv(Conv1d, torch.nn.Conv1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
-  def test_conv2d(self): self._test_conv(Conv2d, torch.nn.Conv2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
+  def test_conv2d(self):
+    BS, C1, H, W = 4, 16, 224//4, 224//4
+    C2, K, S, P = 64, 7, 2, 1
+
+    # create in tinygrad
+    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, H, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
 
   def test_conv1d_same_padding(self):
-    self._test_conv(Conv1d, torch.nn.Conv1d, BS=8, C1=3, DIMS=[32], C2=16, K=3, S=1, P='same')
+    BS, C1, W = 8, 3, 32
+    C2, K, S, P = 16, 3, 1, 'same'
+
+    # create in tinygrad
+    layer = Conv1d(C1, C2, kernel_size=K, stride=S, padding=P)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.Conv1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
+  def _run_conv2d_same_padding_test(self, BS, C1, C2, H, W, K, S, padding='same', D=1):
+    # create in tinygrad
+    layer = Conv2d(C1, C2, kernel_size=K, stride=S, padding=padding, dilation=D)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.Conv2d(C1, C2, kernel_size=K, stride=S, padding=padding, dilation=D).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, H, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
   def test_conv2d_same_padding_odd_input(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[29, 31], C2=32, K=5, S=1, P='same')
+    BS, C1, H, W = 16, 16, 29, 31
+    C2, K, S, P = 32, 5, 1, 'same'
+    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P)
+
   def test_conv2d_same_padding_large_kernel(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=16, DIMS=[28, 33], C2=32, K=9, S=1, P='same')
+    BS, C1, H, W = 16, 16, 28, 33
+    C2, K, S, P = 32, 9, 1, 'same'
+    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P)
+
   def test_conv2d_same_padding_with_dilation(self):
-    self._test_conv(Conv2d, torch.nn.Conv2d, BS=16, C1=3, DIMS=[28, 28], C2=32, K=3, S=1, P='same', D=3)
+    BS, C1, H, W = 16, 3, 28, 28
+    C2, K, S, P, D = 32, 3, 1, 'same', 3
+    self._run_conv2d_same_padding_test(BS, C1, C2, H, W, K, S, P, D)
 
   def test_conv2d_same_padding_invalid_stride(self):
-    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=2, padding='same')
+    C1, C2, K, S, P = 16, 32, 2, 2, 'same'
+    self.assertRaises(ValueError, Conv2d, C1, C2, kernel_size=K, stride=S, padding=P)
+
   def test_conv2d_same_padding_invalid_padding_str(self):
-    self.assertRaises(ValueError, Conv2d, in_channels=16, out_channels=32, kernel_size=2, stride=1, padding='not_same')
+    C1, C2, K, S, P = 16, 32, 2, 1, 'not_same'
+    self.assertRaises(ValueError, Conv2d, C1, C2, kernel_size=K, stride=S, padding=P)
 
   @unittest.skip("Takes too long to compile for Compiled backends")
   def test_conv2d_winograd(self):
@@ -163,13 +229,12 @@ class TestNN(unittest.TestCase):
     with Context(WINO=1):
       z = layer(x)
 
-    m = z.mean()
-    m.backward()
-
     torch_x = torch.tensor(x.numpy(), requires_grad=True)
     torch_z = torch_layer(torch_x)
     np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
 
+    m = z.mean()
+    m.backward()
     gw = layer.weight.grad.realize()
     gb = layer.bias.grad.realize()
     gx = x.grad.realize()
@@ -180,9 +245,44 @@ class TestNN(unittest.TestCase):
     np.testing.assert_allclose(gx.numpy(), torch_x.grad.numpy(), atol=5e-4, rtol=1e-5)
 
   def test_conv_transpose1d(self):
-    self._test_conv(ConvTranspose1d, torch.nn.ConvTranspose1d, BS=4, C1=16, DIMS=[224//4], C2=64, K=7, S=2, P=1)
+    BS, C1, W = 4, 16, 224//4
+    C2, K, S, P = 64, 7, 2, 1
+
+    # create in tinygrad
+    layer = ConvTranspose1d(C1, C2, kernel_size=K, stride=S, padding=P)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.ConvTranspose1d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
+
   def test_conv_transpose2d(self):
-    self._test_conv(ConvTranspose2d, torch.nn.ConvTranspose2d, BS=4, C1=16, DIMS=[224//4, 224//4], C2=64, K=7, S=2, P=1)
+    BS, C1, H, W = 4, 16, 224//4, 224//4
+    C2, K, S, P = 64, 7, 2, 1
+
+    # create in tinygrad
+    layer = ConvTranspose2d(C1, C2, kernel_size=K, stride=S, padding=P)
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.ConvTranspose2d(C1, C2, kernel_size=K, stride=S, padding=P).eval()
+      torch_layer.weight[:] = torch.tensor(layer.weight.numpy(), dtype=torch.float32)
+      torch_layer.bias[:] = torch.tensor(layer.bias.numpy(), dtype=torch.float32)
+
+    # test
+    x = Tensor.uniform(BS, C1, H, W)
+    z = layer(x)
+    torch_x = torch.tensor(x.numpy())
+    torch_z = torch_layer(torch_x)
+    np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-4, rtol=1e-5)
 
   def test_groupnorm(self):
     BS, H, W, C, G = 20, 10, 10, 6, 3

test/test_schedule.py

@@ -15,7 +15,8 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import PatternMatcher, UOp, Ops, GroupOp, UPat, graph_rewrite, track_rewrites
 from tinygrad.uop.symbolic import symbolic_simple
 from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.schedule.kernelize import merge_views, get_kernelize_map, Kernel
+from tinygrad.schedule.kernelize import get_kernelize_map, Kernel
+from tinygrad.opt.swizzler import merge_views
 from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
 from tinygrad.engine.realize import CompiledRunner, run_schedule, lower_schedule
 

test/test_softmax_fusion.py

@@ -86,18 +86,6 @@ class TestFuse(unittest.TestCase):
       return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
     self._test_fuse(embedding, a, atol=1e-5)
 
-  def test_attention_kernel_count(self):
-    wq = Tensor.empty(32, 32)
-    wk = Tensor.empty(32, 32)
-    wv = Tensor.empty(32, 32)
-    x = Tensor.empty(2, 100, 32)
-    q = (x @ wq).contiguous()
-    k = (x @ wk).contiguous()
-    v = (x @ wv).contiguous()
-    attn = q.scaled_dot_product_attention(k, v).fuse()
-    s = attn.schedule()
-    self.assertEqual(len(s), 4) # 3 matmul and 1 attention
-
   def test_flash_attention(self):
     BS = 4
     HEADS = 2

test/test_tensor.py

@@ -892,13 +892,13 @@ class TestIdxUpcast(unittest.TestCase):
 
   @unittest.skipUnless(is_dtype_supported(dtypes.long), "int64 is supported")
   def test_overflow_sym(self):
-    self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
+    self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 0, 2048).bind(32))
 
   def test_regular(self):
     self.do_op_then_assert(dtypes.int, 64, 64, 64)
 
   def test_regular_sym(self):
-    self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 1, 64).bind(32))
+    self.do_op_then_assert(dtypes.int, 2048, 2048, UOp.variable("dim3", 0, 64).bind(32))
 
   @unittest.skipIf(PTX, "PTX always convert Ops.INDEX to int64")
   def test_symfold(self):
@@ -910,7 +910,7 @@ class TestIdxUpcast(unittest.TestCase):
   @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
   def test_int64_unsupported_overflow_sym(self):
     with self.assertRaises(KeyError):
-      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 1, 2048).bind(32))
+      self.do_op_then_assert(dtypes.long, 2048, 2048, UOp.variable("dim3", 0, 2048).bind(32))
 
   @unittest.skipIf(is_dtype_supported(dtypes.long), "int64 is supported")
   def test_int64_unsupported_overflow(self):

tinygrad/codegen/__init__.py

@@ -16,7 +16,6 @@ from tinygrad.codegen.devectorizer import load_store_folding, load_store_indexin
   ReduceContext, correct_load_store, pm_render
 from tinygrad.codegen.optional import get_late_rewrite_patterns
 from tinygrad.codegen.linearize import block_create, pm_blockend_merge, block_merge, pm_finalize, BlockContext
-from tinygrad.opt import pm_optimize
 
 @dataclass
 class RewriteStep:
@@ -43,10 +42,6 @@ def get_rewrites_for_renderer(opts:Renderer, linearizer:bool=True) -> list[Rewri
 def _get_rewrites_for_renderer(opts:Renderer, linearizer:bool, _QUANTIZE, _DEVECTORIZE, _TRANSCENDENTAL) -> list[RewriteStep]:
   # ** lowerer (rewrite_shapetracker_with_index) **
   ret: list[RewriteStep] = []
-
-  # this is kernel.py
-  ret.append(RewriteStep(pm_optimize, ctx=lambda _: opts, name="optimize ast"))
-
   if _QUANTIZE and opts.device in {"CPU", "DSP"}: ret.append(RewriteStep(pm_quant, name="quantize"))
   ret.append(RewriteStep(pm_lowerer, get_index, name="lowerer", bottom_up=True))
 

tinygrad/engine/jit.py

@@ -52,14 +52,14 @@ def apply_graph_to_jit(jit_cache: list[ExecItem], input_rawbuffers: list[Buffer]
     can_be_graphed = ji_graph_dev is not None and ji_graph_dev.graph is not None and graph_class(ji_graph_dev).supports_exec_item([ji_graph_dev], ji)
 
     # Check if the current batch can be extended with this item.
-    can_share_graph = can_be_graphed and len(current_batch_devs) > 0 and \
-                      graph_class(current_batch_devs[0]).supports_exec_item(dedup(current_batch_devs + [ji_graph_dev]), ji)
+    new_batched_devs = dedup(current_batch_devs + [ji_graph_dev])
+    can_share_graph = can_be_graphed and len(current_batch_devs) > 0 and graph_class(current_batch_devs[0]).supports_exec_item(new_batched_devs, ji)
     can_extend_graph_batch = can_share_graph and (max_batch_size == 0 or len(current_batch) < max_batch_size)
 
     # Flush the current batch if any, since it can't be extended or is full.
     if not can_extend_graph_batch and len(current_batch) > 0: flush_batch()
     (current_batch if can_be_graphed else graphed_jit_cache).append(ji)
-    current_batch_devs = dedup(current_batch_devs + [ji_graph_dev]) if can_be_graphed else []
+    current_batch_devs = new_batched_devs if can_be_graphed else []
 
   if len(current_batch) > 0: flush_batch()
   return graphed_jit_cache

tinygrad/engine/realize.py

@@ -7,7 +7,9 @@ from tinygrad.uop.ops import Ops, PatternMatcher, UOp, UPat, Variable, sym_infer
 from tinygrad.device import Device, Buffer
 from tinygrad.renderer import Renderer, ProgramSpec, Estimates
 from tinygrad.engine.schedule import ScheduleItem
+from tinygrad.opt import get_optimized_ast
 from tinygrad.codegen import full_rewrite
+from tinygrad.uop.spec import type_verify
 
 # **************** Program Creation ****************
 
@@ -25,13 +27,16 @@ def get_program(ast:UOp, renderer:Renderer) -> ProgramSpec:
   """
 
   if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST")
+  modified_ast = get_optimized_ast(ast, renderer) if ast.arg is None or ast.arg.opts_to_apply is not None else ast
+  if __debug__: type_verify(list(modified_ast.toposort()))
 
   # linearize
   try:
-    uops = full_rewrite(ast, renderer)
+    uops = full_rewrite(modified_ast, renderer)
   except RuntimeError:
     print("***** LINEARIZE FAILURE *****")
     print(f"ast = {ast}")
+    print(f"opts = {modified_ast.arg.applied_opts}")
     raise
   assert uops[-1].op is Ops.SINK, "last uop must be sink"
 

tinygrad/opt/__init__.py

@@ -2,10 +2,9 @@
 
 from tinygrad.opt.kernel import Kernel
 from tinygrad.opt.heuristic import hand_coded_optimizations
-from tinygrad.uop.ops import UOp, PatternMatcher, UPat, Ops
+from tinygrad.uop.ops import UOp
 from tinygrad.helpers import NOOPT, BEAM, USE_TC, getenv
 from tinygrad.renderer import Renderer
-from tinygrad.uop.spec import type_verify
 
 def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
   """
@@ -28,11 +27,4 @@ def get_optimized_ast(ast:UOp, renderer:Renderer) -> UOp:
       kb = Kernel(ast, opts=renderer)
       rawbufs = bufs_from_lin(kb, allocate=False)
       k = beam_search(kb, rawbufs, BEAM.value, bool(getenv("BEAM_ESTIMATE", 1)))
-  ret = k.get_optimized_ast()
-  if __debug__: type_verify(list(ret.toposort()))
-  return ret
-
-pm_optimize = PatternMatcher([
-  (UPat(Ops.SINK, name="ast"), lambda ctx,ast:
-   get_optimized_ast(ast, ctx) if (ast.arg is None or ast.arg.opts_to_apply is not None) and ast.src[0].st is not None else None),
-])
+  return k.get_optimized_ast()

tinygrad/opt/kernel.py

@@ -14,7 +14,7 @@ from tinygrad.dtype import ImageDType, AddrSpace
 from tinygrad.helpers import all_same, colored, ansilen, dedup, prod, round_up, to_function_name, unwrap, argfix, DEBUG, TC_SELECT, TC_OPT, AMX
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import strides_for_shape, get_contraction
-from tinygrad.schedule.kernelize import view_left
+from tinygrad.opt.swizzler import view_left, view_right
 
 class OptOps(Enum):
   TC = auto(); UPCAST = auto(); UNROLL = auto(); LOCAL = auto() # noqa: E702
@@ -52,6 +52,8 @@ class TensorCoreOptions:
 class Kernel:
   def __init__(self, ast:UOp, opts:Renderer|None=None):
     assert ast.op is Ops.SINK, ast.op
+    ast = graph_rewrite(ast, view_left, name="Main View Left")
+    ast = graph_rewrite(ast, view_right, name="Main View Right")
     self.ast = ast
 
     self.opts = opts if opts is not None else Device[Device.DEFAULT].renderer
@@ -73,7 +75,7 @@ class Kernel:
       self.sts.append(unwrap(x.src[0].st))
 
     # add a shapetracker to the end to track the full shape, with 0 strides so it can merge
-    full_shape = ast.full_shape
+    full_shape = self.ast.full_shape
     self.sts.append(ShapeTracker.from_shape(full_shape, (0,)*len(full_shape)))
 
     # parameters for optimization

tinygrad/opt/swizzler.py

@@ -1,9 +1,11 @@
 from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, resolve, sint
-from tinygrad.helpers import all_same, prod, unwrap
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View, strides_for_shape, get_contraction_with_reduce
+from tinygrad.helpers import unwrap, prod, all_same
+from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.schedule.grouper import ALWAYS_CONTIGUOUS
 
+# **** swizzler
+
 merge_views = PatternMatcher([
   # merge adjacent views
   (UPat(Ops.VIEW, src=(UPat(Ops.VIEW, name="v1"),), name="v2"), lambda v1,v2: v1.replace(arg=v1.arg+v2.arg)),
@@ -17,6 +19,8 @@ merge_views = PatternMatcher([
   # only unmaksed VIEW on CONST replaces the ShapeTracker
   (UPat(Ops.VIEW, src=(UPat((Ops.CONST, Ops.DEFINE_VAR), name="x"),), name="view"),
    lambda x,view: x.replace(src=(x.src[0].replace(arg=x.st+view.st),)) if all(v.mask is None for v in (x.st+view.st).views) else None),
+  # VIEW on SINK is SINK
+  (UPat(Ops.VIEW, name="v").sink(), lambda v: v.src[0].sink()),
 ])
 
 def reduce_push_add_ones(src:UOp, r:UOp, view:UOp):
@@ -95,8 +99,10 @@ view_right = merge_views+PatternMatcher([
   # apply view after reduceops
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.VIEW, src=(UPat(GroupOp.All-ALWAYS_CONTIGUOUS, name="src"),), name="v"),), name="r"), reduceop_view_right),
   # apply view after elementwise ops
-  (UPat(GroupOp.All-{Ops.SINK, Ops.REDUCE_AXIS, Ops.LOAD, Ops.STORE}, name="root"), elementwise_view_right),
+  (UPat(GroupOp.All-{Ops.SINK, Ops.REDUCE_AXIS}, name="root"), elementwise_view_right),
   # merge axes for double reduce (invert of SPLIT_REDUCEOP=1)
   (UPat(Ops.REDUCE_AXIS, src=(UPat(Ops.REDUCE_AXIS, name="r1"),), name="r2"),
    lambda r1,r2: r1.replace(arg=(r1.arg[0], r2.arg[1]+r1.arg[1])) if r1.arg[0] is r2.arg[0] else None),
-])
+  # add VIEW to any DEFINE_GLOBAL that somehow lost its view
+  (UPat(Ops.STORE, src=(UPat(Ops.DEFINE_GLOBAL, name="d"),), name="x", allow_any_len=True), lambda d,x: x.replace(src=(d.view(d.st),)+x.src[1:])),
+])

tinygrad/runtime/support/nv/nvdev.py

@@ -66,7 +66,7 @@ class NVPageTableEntry:
     return self.read_fields(entry_id)[f'address{small}{sys}'] << 12
 
 class NVMemoryManager(MemoryManager):
-  va_allocator = TLSFAllocator((1 << 44), base=0x1000000000) # global for all devices.
+  va_allocator = TLSFAllocator((1 << 44), base=1 << 30) # global for all devices.
 
   def on_range_mapped(self): self.dev.NV_VIRTUAL_FUNCTION_PRIV_MMU_INVALIDATE.write((1 << 0) | (1 << 1) | (1 << 6) | (1 << 31))
 

tinygrad/schedule/kernelize.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, identity_element, resolve
+from tinygrad.uop.ops import UOp, Ops, GroupOp, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, identity_element, resolve, sint
 from tinygrad.uop.ops import track_rewrites, _substitute
 from tinygrad.uop.spec import type_verify, tensor_uop_spec
 from tinygrad.uop.symbolic import symbolic_simple
@@ -8,7 +8,6 @@ from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.schedule.multi import multi_pm
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.schedule.grouper import group_realizes, ALWAYS_CONTIGUOUS
-from tinygrad.opt.swizzler import merge_views, view_left, view_right, apply_swizzle, swizzle_reduceop
 
 # creation can recurse a lot
 import sys
@@ -150,11 +149,18 @@ create_kernels = PatternMatcher([
 
 # **** fix kernel AST
 
-early_buffer_ops = PatternMatcher([
+add_buffer_ops = PatternMatcher([
   # LOAD
-  (UPat(Ops.BUFFER, name="x"), lambda ctx,x: UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(x.size), (), ctx.index(x), tag=1)),
-  # no SINK for meta ops
+  (UPat(Ops.BUFFER, name="x"), lambda ctx,x: UOp.load(UOp(Ops.DEFINE_GLOBAL, x.dtype.ptr(x.size), (), ctx.index(x)).view(x.st),)),
+  # STORE (except for meta ops)
   (UPat(Ops.SINK, src=(UPat(Ops.CONTIGUOUS, src=(UPat(GroupOp.Meta, name="x"),),))), lambda x:x),
+  (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda ctx,sink:
+   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(ctx[i].size), (), i).view(s.st), s) for i,x in enumerate(sink.src)])),
+  # passthrough ASSIGN
+  (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
+  # VALID
+  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
+   lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
 ])
 
 def check_load_st(glbl:UOp, view:UOp):
@@ -168,16 +174,6 @@ def check_load_st(glbl:UOp, view:UOp):
                      +colored("   - a += a.T\n", "red")+colored("   + a += a.T.contiguous()", "green"))
 
 fix_kernel_ops = PatternMatcher([
-  # add the LOAD
-  (UPat(Ops.DEFINE_GLOBAL, name="x"), lambda x: x.replace(tag=None).view(x.st).load() if x.tag is not None else None),
-  # STORE (except for meta ops)
-  (UPat(Ops.SINK, src=UPat(GroupOp.All-{Ops.STORE}), name="sink"), lambda sink:
-   UOp.sink(*[UOp.store(UOp(Ops.DEFINE_GLOBAL, (s:=x.base).dtype.ptr(s.st.real_size()), (), i).view(s.st), s) for i,x in enumerate(sink.src)])),
-  # passthrough ASSIGN
-  (UPat(Ops.ASSIGN, name="x"), lambda x: x.src[1]),
-  # VALID
-  (UPat(Ops.VIEW, src=(UPat.cvar(),), name="self"),
-   lambda self: UOp.where(UOp(Ops.VALID, dtypes.bool, (UOp(Ops.VIEW, arg=self.st),)), self.const_like(self.base.arg), 0)),
   # remove CONTIGUOUS/DEVICE from kernel AST
   (UPat((Ops.CONTIGUOUS, Ops.MSELECT), src=(UPat.var("x"),)), lambda x: x),
   (UPat(Ops.VIEW, src=(UPat(Ops.DEVICE),), name="view"), lambda view: view.replace(src=())),
@@ -196,6 +192,8 @@ replace_globals = PatternMatcher([
 
 def fix_kernel_ast(k:UOp) -> UOp|None:
   if k.arg.ast.op in GroupOp.Meta or all(s.op is Ops.STORE for s in k.arg.ast.src): return None
+  # replace global memory ops with the BUFFER they write to
+  ast = graph_rewrite(k.arg.ast, replace_globals, bottom_up=True, name="replace globals")
   # replace buffer with define_global + add load/store last
   bufs = []
   for s in k.src:
@@ -203,15 +201,9 @@ def fix_kernel_ast(k:UOp) -> UOp|None:
     # traverse back through MSELECT and MSTACK. HACK: 0 branch of MSTACK only
     while s.op in {Ops.MSELECT, Ops.MSTACK}: s = s.src[0]
     bufs.append(s)
-  # replace global memory ops with the BUFFER they write to
-  ast = graph_rewrite(k.arg.ast, replace_globals, bottom_up=True, name="replace globals")
-  ast = graph_rewrite(ast, early_buffer_ops, bufs, bottom_up=True, name="replace buffer early")
+  ast = graph_rewrite(ast, add_buffer_ops+fix_kernel_ops, bufs, bottom_up=True, name="replace buffer")
   if ast.op is Ops.SINK and not all_same([x.device for x in k.src]):
     raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop.buffer for b in k.src)}")
-  # TODO: move these to codegen
-  ast = graph_rewrite(ast, view_left, name="Main View Left")
-  ast = graph_rewrite(ast, view_right, name="Main View Right")
-  ast = graph_rewrite(ast, view_left+fix_kernel_ops, bottom_up=True, name="replace buffer")
   return k.replace(arg=Kernel(ast, k.arg.metadata))
 
 create_ast = PatternMatcher([(UPat(Ops.KERNEL, name="k"), fix_kernel_ast),])
@@ -251,7 +243,7 @@ pm_fuse = PatternMatcher([
 def do_fusion(x:UOp):
   found_contiguous = {}
   def gate_contiguous(x):
-    if is_contiguous:=(x.op is Ops.CONTIGUOUS): found_contiguous[x] = x.replace(src=(UOp(Ops.VIEW, arg=x.st), UOp.unique()))
+    if is_contiguous:=(x.op is Ops.CONTIGUOUS): found_contiguous[x] = x.replace(src=(UOp(Ops.VIEW, arg=x.st),))
     return not is_contiguous
   x.toposort(gate=gate_contiguous)
   del gate_contiguous
@@ -289,7 +281,7 @@ def fuse_arange(root:UOp):
   return root.substitute(fuse_rep, name="fuse_arange") if fuse_rep else None
 
 do_fuse = PatternMatcher([
-  (UPat(Ops.FUSE, name="x"), do_fusion),
+  #(UPat(Ops.FUSE, name="x"), do_fusion),
   (UPat(Ops.REDUCE_AXIS, name="root"), fuse_arange),
 ])
 
@@ -324,6 +316,12 @@ finalize_contiguous = PatternMatcher([
 
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
+new_fixups = PatternMatcher([
+  (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).reshape(r.arg)),
+  # TODO: this should be BUFFER_VIEW
+  (UPat(Ops.COPY, src=(UPat(Ops.SHRINK, name="r"),UPat(name="d")), name="c"), lambda c,r,d: c.replace(src=(r.src[0],d)).shrink(r.arg)),
+])
+
 @track_rewrites(name=lambda sink,ret: f"Schedule {pluralize('Kernel',len([u for u in ret[sink].toposort() if u.op is Ops.KERNEL]))}")
 def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
   """
@@ -337,7 +335,7 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
   """
 
   # multi + merge_views + simplify
-  tensor_map = graph_rewrite_map(sink, multi_pm+do_fuse+merge_views+sym+replace_contiguous, ctx={}, name="merge_views")
+  tensor_map = graph_rewrite_map(sink, new_fixups+multi_pm+do_fuse+sym+replace_contiguous, ctx={}, name="merge_views")
 
   # display the cleaned up tensor graph
   if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Tensor Graph")

tinygrad/tensor.py

@@ -2234,7 +2234,7 @@ class Tensor(MathTrait):
     """
     def parse_formula(formula:str, *operands:Tensor):
       if "..." in (formula := formula.replace(" ", "")):
-        ell_chars, ell_longest = "".join(c for c in string.ascii_letters if c not in formula), 0
+        ell_chars, ell_longest = "".join(set(string.ascii_letters) - set(formula)), 0
         for i, inp in enumerate(filter(lambda x: "..." in x, inputs := formula.split("->")[0].split(","))):
           if (ell_count := max(operands[i].ndim, 1) - (len(inp) - len("..."))) > ell_longest: ell_longest = ell_count
           inputs[i] = inp.replace("...", ell_chars[-ell_count:])

tinygrad/uop/ops.py

@@ -154,8 +154,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
       sz = cast(PtrDType, self.dtype).size
       return ShapeTracker.from_shape((sz,)) if sz > 0 else None
 
-    # hack for PTX, CASTing the ptr loses the shape. even worse hack with tag
-    if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL and self.src[0].tag is None: return None
+    # hack for PTX, CASTing the ptr loses the shape
+    if self.op is Ops.CAST and self.src[0].op is Ops.DEFINE_GLOBAL: return None
 
     # otherwise we get the shape from sources
     if not (src_sts := [x.st for x in self.src if x.st is not None]): return None

tinygrad/viz/serve.py

@@ -132,8 +132,7 @@ def timeline_layout(events:list[tuple[int, int, float, DevEvent]]) -> dict:
     name, cat, info = e.name, None, None
     if (ref:=ref_map.get(name)) is not None:
       name = ctxs[ref]["name"]
-      # TODO: support symbolic by capturing var_vals in profile events
-      if isinstance(p:=contexts[0][ref].ret, ProgramSpec) and all(isinstance(es,int) for es in [p.estimates.ops, p.estimates.mem, p.estimates.lds]):
+      if isinstance(p:=contexts[0][ref].ret, ProgramSpec):
         info = f"{p.estimates.ops/(t:=dur*1e3):.2f} GFLOPS {p.estimates.mem/t:4.1f}|{p.estimates.lds/t:.1f} GB/s"
     elif isinstance(e.name, TracingKey):
       name, cat = e.name.display_name, e.name.cat