Revert "use zip_extract and tar_extract in torch load (#14734)"

This reverts commit 9d9ef81608.

extra/nv_pma/decode.py

@@ -129,6 +129,14 @@ def decode_tpc_id(tpc_id:int) -> tuple[int, int, int]:
   # NOTE: valid only for ops_nv, cuda encoding is different
   return (tpc_id >> 5, (tpc_id >> 1) & 0xf, tpc_id & 1)
 
+def print_samples(samples:list[tuple[PMASample, int]]) -> None:
+  if not samples: return
+  base_pc = min(s.pc_offset for s, _ in samples)
+  for s, tpc_id in samples:
+    gpc, tpc, sm = decode_tpc_id(tpc_id)
+    stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
+    print(f"pc=0x{s.pc_offset - base_pc:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
+
 def print_packets(data:bytes, sm_version:int=0x800) -> None:
   record_size = 9 if sm_version >= 0x890 else 8
   tpc_state: dict[int, list[int]] = collections.defaultdict(list)
@@ -179,11 +187,7 @@ if __name__ == "__main__":
     print(f"\n{'='*60}\nDump {dump_idx} ({len(raw)} bytes, {len(raw)//32} packets)\n{'='*60}")
     if "--raw" in sys.argv: print_packets(raw, sm_ver)
     else:
-      samples = []
-      for s, tpc_id in decode(raw, sm_ver):
-        gpc, tpc, sm = decode_tpc_id(tpc_id)
-        stall_str = colored(f"{s.stall_reason.name:17}", STALL_COLORS.get(s.stall_reason, "white"))
-        print(f"pc=0x{s.pc_offset:06x} {stall_str} ev={s.stall_key:2d} active={s.active} wave={s.wave_id:2d} gpc={gpc} tpc={tpc} sm={sm}")
-        samples.append((s, tpc_id))
+      samples = list(decode(raw, sm_ver))
       print(f"\nDecoded {len(samples)} samples:")
+      print_samples(samples)
       print_aggregated(samples)

extra/viz/cli.py

@@ -1,6 +1,4 @@
 #!/usr/bin/env python3
-import os
-os.environ["VIZ"] = "0"
 import argparse, pathlib
 from typing import Iterator
 from tinygrad.viz import serve as viz

test/backend/test_dtype_alu.py

@@ -296,21 +296,18 @@ class TestDTypeALU(unittest.TestCase):
   @given(ht.int32, strat.sampled_from(dtypes_float+dtypes_int+dtypes_bool))
   def test_int32_cast(self, a, dtype): universal_test_cast(a, dtypes.int32, dtype)
 
-  @unittest.skip("relied on hacks")
   @given(strat.floats(width=32, min_value=1.0, max_value=254.0, allow_subnormal=False),
          strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned(self, a, float_dtype, unsigned_dtype):
     if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
     universal_test_cast(a, float_dtype, unsigned_dtype)
 
-  @unittest.skip("relied on hacks")
   @given(strat.floats(width=32, min_value=256.0, max_value=65000.0, allow_subnormal=False),
          strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_overflow(self, a, float_dtype, unsigned_dtype):
     if not is_dtype_supported(float_dtype): float_dtype = dtypes.float32
     universal_test_cast(a, float_dtype, unsigned_dtype)
 
-  @unittest.skip("relied on hacks")
   @given(strat.floats(width=32, min_value=-65000.0, max_value=-1.0, allow_subnormal=False),
          strat.sampled_from(dtypes_float), strat.sampled_from((dtypes.uint8, dtypes.uint16)))
   def test_float_cast_to_unsigned_underflow(self, a, float_dtype, unsigned_dtype):

test/backend/test_multitensor.py

@@ -1323,55 +1323,6 @@ class TestMultiAssign(unittest.TestCase):
       f(out, vi.bind(i))
     self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
 
-@unittest.skipIf(not_support_multi_device(), "need multi")
-class TestMultiSetitem(unittest.TestCase):
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
-
-  @needs_second_gpu
-  def setUp(self): pass
-
-  def _t(self, axis): return Tensor.arange(16).contiguous().realize().shard(self.device, axis=axis)
-
-  def test_setitem_scalar_axis0(self):
-    t = self._t(0)
-    t[1] = 99
-    self.assertListEqual(t.tolist(), [0,99,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
-
-  def test_setitem_scalar_axis_none(self):
-    t = self._t(None)
-    t[1] = 99
-    self.assertListEqual(t.tolist(), [0,99,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
-
-  def test_setitem_slice_cross_shard(self):
-    t = self._t(0)
-    t[2:6] = 99
-    self.assertListEqual(t.tolist(), [0,1,99,99,99,99,6,7,8,9,10,11,12,13,14,15])
-
-  def test_setitem_full_slice(self):
-    t = self._t(0)
-    t[:] = 42
-    self.assertListEqual(t.tolist(), [42]*16)
-
-  def test_setitem_stride(self):
-    t = self._t(0)
-    t[::4] = 0
-    self.assertListEqual(t.tolist(), [0,1,2,3,0,5,6,7,0,9,10,11,0,13,14,15])
-
-  def test_setitem_single_shard(self):
-    t = self._t(0)
-    t[13] = 99
-    self.assertListEqual(t.tolist(), [0,1,2,3,4,5,6,7,8,9,10,11,12,99,14,15])
-
-  def test_setitem_tensor_value_replicated(self):
-    t = self._t(0)
-    t[2:6] = Tensor([90, 91, 92, 93]).shard(self.device)
-    self.assertListEqual(t.tolist(), [0,1,90,91,92,93,6,7,8,9,10,11,12,13,14,15])
-
-  def test_setitem_tensor_value_sharded_aligned(self):
-    t = self._t(0)
-    t[::4] = Tensor([90, 91, 92, 93]).shard(self.device, axis=0)
-    self.assertListEqual(t.tolist(), [90,1,2,3,91,5,6,7,92,9,10,11,93,13,14,15])
-
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiTransformer(unittest.TestCase):
   @needs_second_gpu

test/backend/test_ops.py

@@ -3295,7 +3295,6 @@ class TestOps(unittest.TestCase):
 
 @unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}")
 class TestOpsUint8(unittest.TestCase):
-  @unittest.skip("relied on hacks")
   def test_cast(self):
     helper_test_op([(2,3,64,64)], lambda x: x.type(torch.uint8), lambda x: x.cast('uint8'), forward_only=True)
 

test/backend/test_outerworld_call.py

@@ -1,19 +0,0 @@
-import unittest
-from tinygrad import Tensor
-
-class TestOuterCall(unittest.TestCase):
-  def test_outer_call_assign(self):
-    a = Tensor.zeros(10,10).contiguous()
-    b = Tensor.ones(10,10).contiguous()
-    Tensor.realize(a,b)
-
-    pa = a.as_param(0)
-    pb = b.as_param(1)
-    out = Tensor.call(a, b, fxn=pa.assign(pa+pb))
-    out.realize()
-
-    print(a.numpy())
-    assert (a == 1).all().item()
-
-if __name__ == '__main__':
-  unittest.main()

test/backend/test_schedule.py

@@ -1082,14 +1082,6 @@ class TestSchedule(unittest.TestCase):
     new_uop = a.reshape(4,1).realize().uop
     assert new_uop.base.op is Ops.BUFFER
 
-  def test_self_assign_no_empty_kernel(self):
-    for shape in [(3, 3), (4, 4)]:
-      a = Tensor.ones(*shape).contiguous().realize()
-      a.assign(a / 1)
-      run_schedule(check_schedule(a, 0, filter_sink=False))
-      self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])
-
-class TestLimitBufs(unittest.TestCase):
   @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
   def test_limit_bufs_with_var(self):
     N = 31
@@ -1102,16 +1094,12 @@ class TestLimitBufs(unittest.TestCase):
     for X in range(1,N): root = root + bufs[X][vi] + bufs[X][vj]
     self.assertEqual(root.item(), N * 2)
 
-  def test_limit_bufs_arange_condition(self):
-    # WHERE with arange-based condition (pure index math, no device) and many buffer loads should not crash limit_bufs
-    with Context(MAX_KERNEL_BUFFERS=8):
-      N = 8
-      idx = Tensor.arange(N)
-      base = Tensor.zeros(N)
-      for i in range(4):
-        a, b = Tensor.rand(N).realize(), Tensor.rand(N).realize()
-        base = (idx >= i).where(a + b, base)
-      assert all(x > 0 for x in base.tolist())
+  def test_self_assign_no_empty_kernel(self):
+    for shape in [(3, 3), (4, 4)]:
+      a = Tensor.ones(*shape).contiguous().realize()
+      a.assign(a / 1)
+      run_schedule(check_schedule(a, 0, filter_sink=False))
+      self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])
 
 class TestSwizzle(unittest.TestCase):
   def test_swizzle_simple(self):

test/backend/test_setitem.py

@@ -36,6 +36,18 @@ class TestSetitem(unittest.TestCase):
     t[:3] *= 10
     self.assertListEqual(t.tolist(), [0, 10, 20, 3, 4, 5, 6, 7, 8, 9])
 
+  def test_setitem_into_unrealized(self):
+    t = Tensor.arange(4).reshape(2, 2)
+    t[1] = 5
+    np.testing.assert_allclose(t.numpy(), [[0, 1], [5, 5]])
+
+  def test_setitem_into_unrealized_sliced_compute(self):
+    # base computation contains SHRINK from prior slicing (like QR decomposition pattern)
+    a = Tensor.arange(6, dtype=dtypes.float).reshape(2, 3)
+    w = a[0] + a[1]  # unrealized ADD with SHRINK in graph: [3, 5, 7]
+    w[1] = 99
+    np.testing.assert_allclose(w.numpy(), [3, 99, 7])
+
   def test_setitem_fancy_on_unrealized_view(self):
     # fancy indexing setitem on unrealized SHRINK view (triggered infinite loop in graph_rewrite)
     base = Tensor.arange(20, dtype=dtypes.float).reshape(4, 5)
@@ -57,6 +69,31 @@ class TestSetitem(unittest.TestCase):
     t = Tensor.zeros(6, dtype=dtypes.float).contiguous().realize()
     with self.assertRaises(RuntimeError): t[2:4] = Tensor([1, 2], dtype=dtypes.int)
 
+  def test_setitem_into_empty(self):
+    t = Tensor.empty(4)
+    t[1] = 5
+    self.assertEqual(t[1].item(), 5)
+
+  def test_setitem_into_cont(self):
+    t = Tensor.ones(4)
+    with self.assertRaises(RuntimeError): t[1] = 5
+
+  def test_setitem_into_const_alu(self):
+    # TODO: this is not consistent
+    t = Tensor.ones(4) + Tensor.ones(4)
+    t[1] = 5
+    self.assertListEqual(t.tolist(), [2, 5, 2, 2])
+
+    t = Tensor.ones(4) + Tensor.ones(4)
+    t.realize()
+    with self.assertRaises(RuntimeError): t[1] = 5
+
+  def test_setitem_into_arange(self):
+    # NOTE: arange has no real buffer, but assigning to it is fine
+    t = Tensor.arange(4)
+    t[1] = 5
+    self.assertListEqual(t.tolist(), [0, 5, 2, 3])
+
   def test_setitem_chained_indexing(self):
     # N[i][j] must work the same as N[i, j]
     N1 = Tensor.zeros((3, 3)).contiguous().realize()

test/null/test_tensor_metadata.py

@@ -63,7 +63,6 @@ class TestTensorMetadata(unittest.TestCase):
     self.assertEqual(len(si.metadata), 3)
     self.assertEqual(set(m.name for m in si.metadata), {"relu", "sigmoid", "__mul__"})
 
-  @unittest.skip("flaky")
   def test_complex_backward(self):
     x = Tensor.rand(3, requires_grad=True).realize()
     y = Tensor.rand(3, requires_grad=True).realize()

test/null/test_tqdm.py

@@ -326,7 +326,7 @@ class TestProgressBar(unittest.TestCase):
     for _ in tinytqdm(range(10^7)): pass
     tinytqdm_time = time.perf_counter() - st
 
-    assert tinytqdm_time < 20 * tqdm_time
+    assert tinytqdm_time < 5 * tqdm_time
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_uops_stats.py

@@ -68,6 +68,28 @@ class TestMemoryCount(unittest.TestCase):
     _, mem = get_stats(a.assign(a+a))
     self.assertEqual(mem, 1024*1024*2)  # 1 read + 1 write
 
+  def test_setitem_slice_const(self):
+    t = Tensor.empty(100, dtype=dtypes.int).realize()
+    GlobalCounters.reset()
+    t[20:50] = 3
+    t.realize()
+    self.assertEqual(GlobalCounters.global_mem, 30*4)  # 30 elements written
+
+  def test_setitem_slice_tensor(self):
+    t = Tensor.empty(100, dtype=dtypes.int).realize()
+    v = Tensor.empty(30, dtype=dtypes.int).realize()
+    GlobalCounters.reset()
+    t[20:50] = v
+    t.realize()
+    self.assertEqual(GlobalCounters.global_mem, 30*4*2)  # 30 read + 30 written
+
+  def test_setitem_full(self):
+    t = Tensor.empty(100, dtype=dtypes.int).realize()
+    GlobalCounters.reset()
+    t[:] = 3
+    t.realize()
+    self.assertEqual(GlobalCounters.global_mem, 100*4)  # full buffer written
+
   @unittest.skipIf(Device.DEFAULT == "CPU", "test copy to CPU from other device")
   def test_copyout(self):
     a = Tensor.empty(32, dtype=dtypes.uint8).to("CPU")

test/null/test_viz.py

@@ -364,15 +364,6 @@ def load_profile(lst:list[ProfileEvent]) -> dict:
   return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
 
 class TestVizProfiler(BaseTestViz):
-  def test_transfer_uses_copy_device(self):
-    a = Tensor.ones(1, device="NULL").contiguous().realize()
-    a.to("NULL:1").realize()
-    range_events = [e for e in cpu_events if isinstance(e, ProfileRangeEvent)]
-    compute_events = [e for e in range_events if e.device == "NULL"]
-    copy_events = [e for e in range_events if e.device.endswith(":COPY")]
-    self.assertGreater(len(compute_events), 0, "expected compute events on base device")
-    self.assertGreater(len(copy_events), 0, "transfer must produce events with ':COPY' device suffix")
-
   def test_node(self):
     prof = [ProfileRangeEvent(device='NV', name='E_2', st=decimal.Decimal(1000), en=decimal.Decimal(1010)),
             ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000))]
@@ -583,7 +574,6 @@ class TestVizMemoryLayout(BaseTestViz):
     user_cnt = [len(b["arg"]["users"]) for b in buffers if b["arg"].get("users")]
     self.assertEqual(len(user_cnt), len(programs))
 
-  @unittest.skip("flaky")
   def test_inflight_buf(self):
     a = Tensor.empty(1, device="NULL")
     n = 4

test/unit/test_assign.py

@@ -756,26 +756,6 @@ class TestAssignOrdering(unittest.TestCase):
     self.assertEqual(buf[0:1, :].sum().item(), 4)
     self.assertEqual(buf[1:2, :].sum().item(), 8)
 
-  def test_multi_step_assign_read_write_same_buffer(self):
-    """Assign to m and param reading b, then update b, across multiple steps.
-    This is the optimizer bias-correction pattern from issue #13600: m accumulates,
-    param is updated using m/(1-b), and b is updated via *= after the reads."""
-    b = Tensor([0.5]).contiguous().realize()
-    m = Tensor([0.0]).contiguous().realize()
-    param = Tensor([1.0]).contiguous().realize()
-    for _ in range(10):
-      m.assign(0.9 * m + 0.1)
-      param.assign(param - m / (1 - b))
-      b *= 0.9
-      Tensor.realize(param, m, b)
-    # numpy reference
-    b_np, m_np, p_np = 0.5, 0.0, 1.0
-    for _ in range(10):
-      m_np = 0.9 * m_np + 0.1
-      p_np = p_np - m_np / (1 - b_np)
-      b_np *= 0.9
-    np.testing.assert_allclose(param.item(), p_np, atol=1e-5)
-
   def test_multiple_slice_assigns_then_read(self):
     """Multiple non-overlapping slice assigns then read."""
     buf = Tensor.zeros(4).contiguous().realize()

test/unit/test_disk_tensor.py

@@ -456,7 +456,6 @@ class TestDiskTensor(TempDirTestCase):
     np.testing.assert_equal(t1.numpy(), np.arange(128, dtype=np.uint8))
     np.testing.assert_equal(t2.numpy(), np.arange(64, dtype=np.uint8))
 
-  @unittest.skip("fails with setup_python_cap run")
   def test_disk_open_failure_state(self):
     from tinygrad.runtime.ops_disk import DiskDevice
     fn = pathlib.Path(self.tmp("dt_open_failure"))
@@ -477,7 +476,6 @@ class TestDiskTensor(TempDirTestCase):
     t2.to("CPU").realize()
     assert disk_device.size == 200
 
-  @unittest.skip("fails with setup_python_cap run")
   def test_disk_permission_error(self):
     fn = pathlib.Path(self.tmp("dt_permission"))
     fn.write_bytes(bytes(range(256)))

test/unit/test_indexing.py

@@ -1000,7 +1000,7 @@ def assert_backward_eq(tensor: Tensor, indexer):
 def get_set_tensor(indexed: Tensor, indexer):
   set_size = indexed[indexer].shape
   set_count = indexed[indexer].numel()
-  set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size).cast(indexed.dtype)
+  set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size) #.cast(dtypes.float64)
   return set_tensor
 
 @slow

test/unit/test_setitem_schedule.py

@@ -1,150 +0,0 @@
-import unittest
-from tinygrad import Tensor, dtypes, GlobalCounters
-
-class TestSetitemInto(unittest.TestCase):
-  def test_setitem_into_unrealized(self):
-    GlobalCounters.reset()
-    t = Tensor.arange(4, dtype=dtypes.int32).reshape(2, 2)
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 16)
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [[0, 1], [5, 5]])
-
-  def test_setitem_into_unrealized_sliced_compute(self):
-    # base computation contains SHRINK from prior slicing (like QR decomposition pattern)
-    GlobalCounters.reset()
-    a = Tensor.arange(8, dtype=dtypes.int32).reshape(2, 4)
-    w = a[0] + a[1]  # unrealized ADD with SHRINK in graph: [4, 6, 8, 10]
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    w[1] = 99
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    w.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4*4)
-    self.assertListEqual(w.tolist(), [4, 99, 8, 10])
-
-  def test_setitem_into_empty(self):
-    GlobalCounters.reset()
-    t = Tensor.empty(4, dtype=dtypes.int32)
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    # TODO: this can be just 4 if empty goes through is_realized setitem path
-    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(t[1].item(), 5)
-
-  def test_setitem_into_empty_alu(self):
-    GlobalCounters.reset()
-    t = Tensor.empty(4, dtype=dtypes.int32) + 1
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(t[1].item(), 5)
-
-  def test_setitem_into_tensor(self):
-    t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize()
-    GlobalCounters.reset()
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t[1].realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [1, 5, 3, 4])
-
-  def test_setitem_into_tensor_alu(self):
-    t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize() + 1
-    GlobalCounters.reset()
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t[1].realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [2, 5, 4, 5])
-
-  def test_setitem_into_cont(self):
-    GlobalCounters.reset()
-    t = Tensor.ones(4, dtype=dtypes.int32)
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4*4)
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [1, 5, 1, 1])
-
-  def test_setitem_into_const_alu(self):
-    GlobalCounters.reset()
-    t = Tensor.ones(4, dtype=dtypes.int32) + 1
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 4*4)
-    t[1].realize()
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [2, 5, 2, 2])
-
-  def test_setitem_into_arange(self):
-    # NOTE: arange has no real buffer, but assigning to it is fine
-    GlobalCounters.reset()
-    t = Tensor.arange(4, dtype=dtypes.int32)
-    t[1] = 5
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertListEqual(t.tolist(), [0, 5, 2, 3])
-
-  def test_setitem_slice_const(self):
-    t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
-    GlobalCounters.reset()
-    t[20:50] = 3
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 30*4)  # 30 elements written
-
-  def test_setitem_slice_tensor(self):
-    t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
-    v = Tensor.zeros(30, dtype=dtypes.int32).contiguous().realize()
-    GlobalCounters.reset()
-    t[20:50] = v
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 30*4*2)  # 30 read + 30 written
-
-  def test_setitem_full(self):
-    t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
-    GlobalCounters.reset()
-    t[:] = 3
-    self.assertEqual(GlobalCounters.kernel_count, 0)
-    t.realize()
-    self.assertEqual(GlobalCounters.kernel_count, 1)
-    self.assertEqual(GlobalCounters.global_mem, 100*4)  # full buffer written
-
-if __name__ == '__main__':
-  unittest.main()

tinygrad/engine/schedule.py

@@ -29,9 +29,7 @@ def create_schedule(sched_sink:UOp) -> tuple[list[ExecItem], UOp]:
       assert k.op in {Ops.CALL, Ops.END}, f"AFTER src[1] should be KERNEL or END, not {k.op}"
       in_degree.setdefault(k, 0)
       if k.op is Ops.END: assert k.src[0].op is Ops.CALL, f"END src[0] should be KERNEL, not {k.src[0].op}"
-      # WAR deps from rangeify are stored in AFTER src[2:]
-      kernel_deps = k.src[0].src[1:] if k.op is Ops.END else k.src[1:]
-      for s in kernel_deps + u.src[2:]:
+      for s in k.src[0].src[1:] if k.op is Ops.END else k.src[1:]:
         match (s := _unwrap_src(s)).op:
           case Ops.AFTER:
             children.setdefault(s.src[1], []).append(k)

tinygrad/helpers.py

@@ -182,7 +182,6 @@ CACHELEVEL, IGNORE_BEAM_CACHE, DEVECTORIZE = ContextVar("CACHELEVEL", 2), Contex
 VALIDATE_WITH_CPU, DISABLE_FAST_IDIV = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("DISABLE_FAST_IDIV", 0)
 CORRECT_DIVMOD_FOLDING, FUSE_OPTIM = ContextVar("CORRECT_DIVMOD_FOLDING", 0), ContextVar("FUSE_OPTIM", 0)
 ALLOW_DEVICE_USAGE, MAX_BUFFER_SIZE = ContextVar("ALLOW_DEVICE_USAGE", 1), ContextVar("MAX_BUFFER_SIZE", 0)
-MAX_KERNEL_BUFFERS = ContextVar("MAX_KERNEL_BUFFERS", 0)
 EMULATE, EMULATED_DTYPES = ContextVar("EMULATE", ""), ContextVar("EMULATED_DTYPES", "")
 CPU_COUNT = ContextVar("CPU_COUNT", max(1, len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else (os.cpu_count() or 1)))
 # Compilers

tinygrad/runtime/ops_metal.py

@@ -177,7 +177,7 @@ class MetalAllocator(LRUAllocator[MetalDevice]):
     # There is no real metal multidevice support for now, so transfer is used only for tests.
     src_dev.synchronize()
   def _cp_mv(self, dst, src, prof_desc):
-    with cpu_profile(prof_desc, f"{self.dev.device}:COPY"): dst[:] = src
+    with cpu_profile(prof_desc, self.dev.device): dst[:] = src
   def _as_buffer(self, src:MetalBuffer) -> memoryview:
     self.dev.synchronize()
     return to_mv(src.buf.contents(), src.size + src.offset)[src.offset:]

tinygrad/runtime/ops_null.py

@@ -24,7 +24,7 @@ class NullAllocator(Allocator['NullDevice']):
   def _copyout(self, dest:memoryview, src):
     if not NULL_ALLOW_COPYOUT: raise RuntimeError("no copyout on NULL")
   def _transfer(self, dest, src, sz:int, src_dev, dest_dev):
-    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{self.dev.device}:COPY"): pass
+    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", self.dev.device): pass
   def _offset(self, buf, offset:int, size:int): pass
 
 class NullGraph(MultiGraphRunner):

tinygrad/runtime/ops_qcom.py

@@ -329,7 +329,7 @@ class QCOMAllocator(HCQAllocatorBase):
     return self.dev._gpu_map(opts.external_ptr, size, image=opts.image) if opts.external_ptr else self.dev._gpu_alloc(size, image=opts.image)
 
   def _do_copy(self, src_addr, dest_addr, src_size, real_size, src_stride, dest_stride, prof_text, dest_off=0, src_off=0):
-    with cpu_profile(prof_text, f"{self.dev.device}:COPY"):
+    with cpu_profile(prof_text, self.dev.device):
       while src_off < src_size:
         ctypes.memmove(dest_addr+dest_off, src_addr+src_off, real_size)
         src_off, dest_off = src_off+src_stride, dest_off+dest_stride

tinygrad/runtime/support/hcq.py

@@ -516,7 +516,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
   def _copyin(self, dest:HCQBuffer, src:memoryview):
     if self.dev.hw_copy_queue_t is None:
       self.dev.synchronize()
-      with cpu_profile(f'TINY -> {self.dev.device}', f"{self.dev.device}:COPY"): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      with cpu_profile(f'TINY -> {self.dev.device}', self.dev.device): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
       return
 
     with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"TINY -> {self.dev.device}", enabled=PROFILE, dev_suff="SDMA:0"):
@@ -550,7 +550,7 @@ class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
     if self.dev.hw_copy_queue_t is None:
-      with cpu_profile(f'{self.dev.device} -> TINY', f"{self.dev.device}:COPY"): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      with cpu_profile(f'{self.dev.device} -> TINY', self.dev.device): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
       return
 
     with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=f"{self.dev.device} -> TINY", enabled=PROFILE, dev_suff="SDMA:0"):

tinygrad/schedule/rangeify.py

@@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes, PtrDType, ImageDType, AddrSpace
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, _substitute, KernelInfo, pm_gate_kernel_sink
 from tinygrad.uop.ops import graph_rewrite, identity_element, sint, AxisType, BottomUpGate, _remove_all_tags, range_str
 from tinygrad.uop.symbolic import symbolic
-from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS
+from tinygrad.helpers import argsort, prod, all_same, getenv, flatten, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ
 from tinygrad.helpers import PCONTIG, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -43,6 +43,7 @@ def assign_to_contiguous(assign:UOp, target:UOp, src:UOp):
   if target is not t and target.op_in_backward_slice_with_self(Ops.SHRINK):
     # base already realized: copy src only if it reads from the same buffer (overlapping read/write hazard)
     if t.op is Ops.CONTIGUOUS: return assign.replace(src=(target, src.contiguous())) if t in src.toposort() else None
+    if t.op is Ops.CONST: raise RuntimeError("setitem target must be a writable view backed by a buffer")
     mops: list[UOp] = []
     while target.op in GroupOp.Movement:
       mops.append(target)
@@ -312,7 +313,7 @@ DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8} # TODO: get from device?
 def limit_bufs(ctx:IndexingContext, root:UOp):
   if (device:=root._device) is None: return None # no device, index related calculations
   device = device if isinstance(device, str) else device[0].split(":")[0]
-  if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None
+  if not (MAX_BUFS:=getenv("MAX_KERNEL_BUFFERS", DEVICE_MAX_BUFS.get(device, 0))): return None
 
   bufs: set[UOp] = set()
   def gate_input(u:UOp):
@@ -324,7 +325,7 @@ def limit_bufs(ctx:IndexingContext, root:UOp):
   if len(bufs) > MAX_BUFS - 1: # NOTE: this -1 is for the output buffer
     srcs = []
     for s in root.src:
-      if s.op in GroupOp.Elementwise and s._device is not None:
+      if s.op in GroupOp.Elementwise:
         # Insert bufferize: all AxisType.REDUCE before bufferize are AxisType.LOOP
         orig_ranges, end_ranges = s.ranges, [x.replace(arg=(next(ctx.range_idx), AxisType.LOOP)) if x.op is Ops.RANGE else x for x in s.ranges]
         s = s.substitute(dict(zip(orig_ranges, end_ranges))).bufferize(*end_ranges, arg=BufferizeOpts(device=s.device)).index(*orig_ranges)
@@ -554,7 +555,7 @@ def tag_uop(ctx:tuple[list[UOp], set[UOp]], x:UOp):
   return x.replace(tag=(len(ctx[0])-1,))
 add_tags = pm_gate_kernel_sink+PatternMatcher([
   # don't tag BUFFERs, they are global
-  (UPat(GroupOp.All-{Ops.PARAM, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.END,
+  (UPat(GroupOp.All-{Ops.PARAM, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, Ops.CALL, Ops.END,
                      Ops.MSTACK, Ops.MSELECT, Ops.RANGE}.union(GroupOp.Movement), name="x"), tag_uop),
   (UPat({Ops.MSTACK, Ops.MSELECT}, name="x"), lambda ctx,x: None if all(s.op is Ops.PARAM for s in x.src) else tag_uop(ctx, x)),
 ])
@@ -601,15 +602,15 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
                         name="bufferize to store")
   tsink = graph_rewrite(tsink, pm_gate_kernel_sink+split_kernels, ctx=uop_list, bottom_up=True, name="split kernels")
 
-  # WAR deps: if kernel U reads buffer S, and S is also written by another kernel, S's write must wait for U to finish
-  afters = [u for u in tsink.toposort() if u.op is Ops.AFTER]
-  kernel_assign: dict[UOp, UOp] = {u.buf_uop:u for u in afters}
+  # if a kernel depends on a buffer, and that buffer is later assigned to, make the assign depend on the kernel's assign
+  kernel_assign: dict[UOp, UOp] = {}
   assign_rep: dict[UOp, UOp] = {}
-  for u in afters:
+  for u in tsink.toposort():
+    if u.op is not Ops.AFTER: continue
+    kernel_assign[u.buf_uop] = u
     for s in u.src[1].src:
       # TODO: this is probably broken for MSELECT/MSTACK
       if s.op not in {Ops.BUFFER, Ops.PARAM} or s is u.buf_uop or (a:=kernel_assign.get(s)) is None: continue
-      if a.src[1] is u.src[1]: continue  # same kernel (multi-output custom kernels)
       if any(x.op is Ops.AFTER and x.buf_uop is s for x in u.toposort()):
         raise RuntimeError(f"cycle detected in graph, kernel for {u.buf_uop} must either depend on AFTER or BUFFER")
       assign_rep[a] = kernel_assign[s] = a.replace(src=a.src+(u,))

tinygrad/tensor.py

@@ -614,15 +614,14 @@ class Tensor(OpMixin):
     print(t.numpy())
     ```
     """
-    dt = to_dtype(dtype or dtypes.default_float)
-    if not dtypes.is_float(dt): raise ValueError(f"rand only supports float dtypes, got {dt}")
+    if not dtypes.is_float(dtype := to_dtype(dtype or dtypes.default_float)): raise ValueError(f"rand only supports float dtypes, got {dtype}")
     if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
     if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
     device = cast(str, canonicalize_device(device))
 
     # if shape has 0, return zero tensor
-    if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dt, **kwargs)
-    num = ceildiv(numel * dt.itemsize, 4)
+    if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dtype, **kwargs)
+    num = ceildiv(numel * dtype.itemsize, 4)
 
     # generate per device seeds and rng counter if we haven't seen this device yet
     if device not in Tensor._device_seeds:
@@ -640,14 +639,14 @@ class Tensor(OpMixin):
     bits = Tensor._threefry_random_bits(Tensor._device_seeds[device], counts0, counts1)[:num]
 
     # bitcast to uint with same number of bits
-    _, nmant = dtypes.finfo(dt)
-    uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dt.itemsize]
+    _, nmant = dtypes.finfo(dtype)
+    uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dtype.itemsize]
     bits = bits.bitcast(uint_dtype)
     # only randomize the mantissa bits and set the exponent to 1
-    one = Tensor.ones_like(bits, device=bits.device, dtype=dt).bitcast(uint_dtype)
-    bits = bits.rshift(dt.bitsize - nmant).bitwise_or(one)
+    one = Tensor.ones_like(bits, device=bits.device, dtype=dtype).bitcast(uint_dtype)
+    bits = bits.rshift(dtype.bitsize - nmant).bitwise_or(one)
     # bitcast back to the original dtype and reshape
-    out = bits.bitcast(dt)[:numel].sub(1).reshape(shape).requires_grad_(kwargs.get("requires_grad"))
+    out = bits.bitcast(dtype)[:numel].sub(1).reshape(shape).requires_grad_(kwargs.get("requires_grad"))
     return out.contiguous() if contiguous else out
 
   # ***** creation helper functions *****
@@ -771,9 +770,8 @@ class Tensor(OpMixin):
     print(Tensor.eye(2, 4).numpy())
     ```
     """
-    m_ = n if m is None else m
-    if n < 0 or m_ < 0: raise ValueError(f"cannot have negative {n=}, {m_=}")
-    t = (Tensor.arange(n, device=device).unsqueeze(-1) == Tensor.arange(m_, device=device))
+    if n < 0 or ((m := n if m is None else m) < 0): raise ValueError(f"cannot have negative {n=}, {m=}")
+    t = (Tensor.arange(n, device=device).unsqueeze(-1) == Tensor.arange(m, device=device))
     return t.cast(dtype or dtypes.default_float).requires_grad_(requires_grad)
 
   def _multi_like(self, fxn, *args, **kwargs) -> Tensor:
@@ -1216,26 +1214,6 @@ class Tensor(OpMixin):
     x_dims = [p for p in indices_parsed if not isinstance(p['index'], sint)]
     x = x.reshape(tuple(p['size'] for p in x_dims))
 
-    # basic setitem: construct result with view region replaced by v using arange masks
-    if v is not None and not any(isinstance(p['index'], Tensor) for p in indices_parsed):
-      # broadcast v to getitem shape, reshape to self.ndim (squeeze None dims, unsqueeze int dims — all are size 1)
-      vb = v.cast(self.dtype)._broadcast_to(x.shape)
-      vb = vb.reshape(tuple(1 if isinstance(p['index'], sint) else p['size'] for p in indices_parsed if p['index'] is not None))
-      # undo movement ops per-dim and build boolean mask
-      per_dim = []
-      for d, m in enumerate(mops):
-        (s, e), st = m['boundary'], abs(m['stride'])
-        if st != 1 and vb.shape[d] > 1:  # un-stride: interleave with zeros
-          vb = vb.unsqueeze(d+1)
-          vb = vb.pad_to(tuple(st if j == d+1 else None for j in range(vb.ndim)))
-          vb = vb.reshape(vb.shape[:d] + (vb.shape[d]*vb.shape[d+1],) + vb.shape[d+2:])
-          vb = vb.shrink_to(tuple(e-s if j == d else None for j in range(self.ndim)))
-        idx = Tensor.arange(self.shape[d], device=self.device).reshape([1]*d + [self.shape[d]] + [1]*(self.ndim - d - 1))
-        per_dim.append((idx >= s) & (idx < e) & (((e-1-idx) if m['stride'] < 0 else (idx-s)) % st == 0))
-      vb = vb.flip(tuple(d for d, m in enumerate(mops) if m['stride'] < 0))
-      vb = vb.pad(tuple((m['boundary'][0], self.shape[d] - m['boundary'][1]) for d, m in enumerate(mops)))
-      return (functools.reduce(lambda a, b: a & b, per_dim) if per_dim else Tensor(True, dtype=dtypes.bool, device=self.device)).where(vb, self)
-
     # tensor indexing
     if tops := [(d, p) for d, p in enumerate(x_dims) if isinstance(p['index'], Tensor)]:
       dims, tensors, masks = [d for d, _ in tops], cast(list[Tensor], [p['index'] for _, p in tops]), []
@@ -1328,16 +1306,13 @@ class Tensor(OpMixin):
     idx = [indices] if (isinstance(indices, list) and all_int(indices)) or not isinstance(indices, (tuple, list)) else list(indices)
     is_disk = isinstance(self.device, str) and self.device.startswith("DISK")
     if any(isinstance(i, (Tensor, list, tuple)) for i in idx): # advanced setitem
-      if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
+      if isinstance(self.device, str) and self.device.startswith("DISK"): raise RuntimeError("advanced setitem is not supported for DISK tensors")
       if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
       self.assign(self._getitem(indices, v))
     elif is_disk or self.uop.is_realized: # basic setitem, self is realized. TODO: disk uop.base is a COPY and not realized
       self[indices].assign(v)
     else: # basic setitem, self is not realized
-      if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
-      # __iadd__/__isub__ on unrealized views creates a no-op ASSIGN; unwrap to get the computed value
-      if v.uop.op is Ops.ASSIGN: v = v._apply_uop(lambda x: x.src[1])
-      self.replace(self._getitem(indices, v))
+      self[indices].assign(v).realize()
 
   def __delitem__(self, indices) -> None:
     raise TypeError("Tensor does not support deleting items")
@@ -3904,7 +3879,10 @@ class Tensor(OpMixin):
     print(t.dtype, t.numpy())
     ```
     """
-    return self if self.dtype == (dt:=to_dtype(dtype)) else self._apply_uop(UOp.cast, dtype=dt)
+    if (dt:=to_dtype(dtype)) in {dtypes.uint8, dtypes.uint16} and dtypes.is_float(self.dtype):
+      # NOTE: values within the int32 range and outside the unsigned dtype range will cause values to wrap around
+      return self._apply_uop(UOp.cast, dtype=dtypes.int32)._apply_uop(UOp.cast, dtype=dt)
+    return self if self.dtype == dt else self._apply_uop(UOp.cast, dtype=dt)
 
   def bitcast(self, dtype:DTypeLike) -> Tensor:
     """