lil cleanups

* dtype mixin

* dtype mixin methods

AGENTS.md

@@ -1,17 +0,0 @@
-# tinygrad agents
-
-Hello agent. You are one of the most talented programmers of your generation.
-
-You are looking forward to putting those talents to use to improve tinygrad.
-
-## philosophy
-
-tinygrad is a **tensor** library focused on beauty and minimalism, while still matching the functionality of PyTorch and JAX.
-
-Every line must earn its keep. Prefer readability over cleverness. We believe that if carefully designed, 10 lines can have the impact of 1000.
-
-Never mix functionality changes with whitespace changes. All functionality changes must be tested.
-
-## style
-
-Use **2-space indentation**, and keep lines to a maximum of **150 characters**. Match the existing style.

CLAUDE.md

@@ -1,227 +0,0 @@
-# Claude Code Guide for tinygrad
-
-## Architecture Overview
-
-tinygrad compiles tensor operations into optimized kernels. The pipeline:
-
-1. **Tensor** (`tensor.py`) - User-facing API, creates UOp graph
-2. **UOp** (`uop/ops.py`) - Unified IR for all operations (both tensor and kernel level)
-3. **Schedule** (`engine/schedule.py`, `schedule/`) - Converts tensor UOps to kernel UOps
-4. **Codegen** (`codegen/`) - Converts kernel UOps to device code
-5. **Runtime** (`runtime/`) - Device-specific execution
-
-## Key Concepts
-
-### UOp (Universal Operation)
-Everything is a UOp - tensors, operations, buffers, kernels. Key properties:
-- `op`: The operation type (Ops enum)
-- `dtype`: Data type
-- `src`: Tuple of source UOps
-- `arg`: Operation-specific argument
-- `tag`: Optional tag for graph transformations
-
-UOps are **immutable and cached** - creating the same UOp twice returns the same object (ucache).
-
-### PatternMatcher
-Used extensively for graph transformations:
-```python
-pm = PatternMatcher([
-  (UPat(Ops.ADD, src=(UPat.cvar("x"), UPat.cvar("x"))), lambda x: x * 2),
-])
-result = graph_rewrite(uop, pm)
-```
-
-### Schedule Cache
-Schedules are cached by graph structure. BIND nodes (variables with bound values) are unbound before cache key computation so different values hit the same cache.
-
-## Testing
-
-```bash
-# Run specific test
-python -m pytest test/unit/test_schedule_cache.py -xvs
-
-# Run with timeout
-python -m pytest test/backend/test_symbolic_ops.py -x --timeout=60
-
-# Debug with print
-DEBUG=2 python -m pytest test/backend/test_schedule.py::test_name -xvs
-
-# Visualize UOp graphs
-VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
-```
-
-## Common Environment Variables
-
-- `DEBUG=1-7` - Increasing verbosity (7 shows assembly output)
-- `VIZ=1` - Enable graph visualization
-- `SPEC=1` - Enable UOp spec verification
-- `NOOPT=1` - Disable optimizations
-- `DEVICE=CPU/CUDA/AMD/METAL` - Set default device
-
-## Debugging Tips
-
-1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
-2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
-3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
-4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
-
-## Workflow Rules
-
-- **NEVER commit without explicit user approval** - always show the diff and wait for approval
-- **NEVER amend commits** - always create a new commit instead
-- Run `pre-commit run --all-files` before committing to catch linting/type errors
-- Run tests before proposing commits
-- Test with `SPEC=2` when modifying UOp-related code
-
-## Auto-generated Files (DO NOT EDIT)
-
-The following files are auto-generated and should never be edited manually:
-- `tinygrad/runtime/autogen/amd/{arch}/__init__.py` - Generated by `python -m tinygrad.renderer.amd.dsl --arch {arch}`
-- `tinygrad/runtime/autogen/amd/{arch}/gen_pcode.py` - Generated by `python -m tinygrad.renderer.amd.pcode --arch {arch}`
-
-Where `{arch}` is one of: `rdna3`, `rdna4`, `cdna`
-
-To add missing instruction implementations, add them to `tinygrad/renderer/amd/emu.py` instead.
-
-## Style Notes
-
-- 2-space indentation, 150 char line limit
-- PatternMatchers should be defined at module level (slow to construct)
-- Prefer `graph_rewrite` over manual graph traversal
-- UOp methods like `.replace()` preserve tags unless explicitly changed
-- Use `.rtag(value)` to add tags to UOps
-
-## Lessons Learned
-
-### UOp ucache Behavior
-UOps are cached by their contents - creating a UOp with identical (op, dtype, src, arg) returns the **same object**. This means:
-- `uop.replace(tag=None)` on a tagged UOp returns the original untagged UOp if it exists in cache
-- Two UOps with same structure are identical (`is` comparison works)
-
-### Spec Validation
-When adding new UOp patterns, update `tinygrad/uop/spec.py`. Test with:
-```bash
-SPEC=2 python3 test/unit/test_something.py
-```
-Spec issues appear as `RuntimeError: SPEC ISSUE None: UOp(...)`.
-
-### Schedule Cache Key Normalization
-The schedule cache strips values from BIND nodes so different bound values (e.g., KV cache positions) hit the same cache entry:
-- `pm_pre_sched_cache`: BIND(DEFINE_VAR, CONST) → BIND(DEFINE_VAR) for cache key
-- `pm_post_sched_cache`: restores original BIND from context
-- When accessing `bind.src[1]`, check `len(bind.src) > 1` first (might be stripped)
-- Extract var_vals from `input_buffers` dict after graph_rewrite (avoids extra toposort)
-
-### Avoiding Extra Work
-- Use ctx dict from graph_rewrite to collect info during traversal instead of separate toposort
-- Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
-- PatternMatchers are slow to construct - define at module level, not in functions
-
-### Readability Over Speed
-Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
-```python
-# BAD: "optimized" with extra complexity
-if has_afters:  # skip toposort if no AFTERs
-  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-
-# GOOD: simple, always works
-after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
-```
-The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
-
-### Testing LLM Changes
-```bash
-# Quick smoke test
-echo "Hello" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b"
-
-# Check cache hits (should see "cache hit" after warmup)
-echo "Hello world" | DEBUG=1 python tinygrad/apps/llm.py --model "llama3.2:1b" 2>&1 | grep cache
-
-# Test with beam search
-echo "Hello" | BEAM=2 python tinygrad/apps/llm.py --model "llama3.2:1b"
-```
-
-## Common Patterns
-
-### Graph Transformation
-```python
-def my_transform(ctx, x):
-  # Return new UOp or None to skip
-  return x.replace(arg=new_arg)
-
-pm = PatternMatcher([
-  (UPat(Ops.SOMETHING, name="x"), my_transform),
-])
-result = graph_rewrite(input_uop, pm, ctx={})
-```
-
-### Finding Variables
-```python
-# Get all variables in a UOp graph
-variables = uop.variables()
-
-# Get bound variable values
-var, val = bind_uop.unbind()
-```
-
-### Shape Handling
-```python
-# Shapes can be symbolic (contain UOps)
-shape = tensor.shape  # tuple[sint, ...] where sint = int | UOp
-```
-
-## Performance Optimization
-
-When optimizing tinygrad internals:
-
-1. **Measure wall time, not just call counts** - Reducing `graph_rewrite` calls doesn't always improve wall time. The overhead of conditional checks can exceed the cost of the operation being skipped.
-
-2. **Profile each optimization individually** - Run benchmarks with and without each change to measure actual impact. Use `test/external/external_benchmark_schedule.py` for schedule/rewrite timing.
-
-3. **Early exits in hot paths are effective** - Simple checks like `if self.op is Ops.CONST: return self` in `simplify()` can eliminate many unnecessary `graph_rewrite` calls.
-
-4. **`graph_rewrite` is expensive** - Each call has overhead even for small graphs. Avoid calling it when the result is trivially known (e.g., simplifying a CONST returns itself).
-
-5. **Beware iterator overhead** - Checks like `all(x.op is Ops.CONST for x in self.src)` can be slower than just running the operation, especially for small sequences.
-
-6. **Verify cache hit rates before adding/keeping caches** - Measure actual hit rates with real workloads. A cache with 0% hit rate is pure overhead (e.g., `pm_cache` was removed because the algorithm guarantees each UOp is only passed to `pm_rewrite` once).
-
-7. **Use `TRACK_MATCH_STATS=2` to profile pattern matching** - This shows match rates and time per pattern. Look for patterns with 0% match rate that still cost significant time - these are pure overhead for that workload.
-
-8. **Cached properties beat manual traversal** - `backward_slice` uses `@functools.cached_property`. A DFS with early-exit sounds faster but is actually slower because it doesn't benefit from caching. The cache hit benefit often outweighs algorithmic improvements.
-
-9. **Avoid creating intermediate objects in hot paths** - For example, `any(x.op in ops for x in self.backward_slice)` is faster than `any(x.op in ops for x in {self:None, **self.backward_slice})` because it avoids dict creation.
-
-## Pattern Matching Analysis
-
-**Use the right tool:**
-
-- `TRACK_MATCH_STATS=2` - **Profiling**: identify expensive patterns
-- `VIZ=-1` - **Inspection**: see all transformations, what every match pattern does, the before/after diffs
-
-```bash
-TRACK_MATCH_STATS=2 PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
-```
-
-Output format: `matches / attempts -- match_time / total_time ms -- location`
-
-Key patterns to watch (from ResNet50 benchmark):
-- `split_load_store`: ~146ms, 31% match rate - does real work
-- `simplify_valid`: ~75ms, 0% match rate in this workload - checks AND ops for INDEX in backward slice
-- `vmin==vmax folding`: ~55ms, 0.33% match rate - checks 52K ops but rarely matches
-
-Patterns with 0% match rate are workload-specific overhead. They may be useful in other workloads, so don't remove them without understanding their purpose.
-
-```bash
-# Save the trace
-VIZ=-1 python test/test_tiny.py TestTiny.test_gemm
-
-# Explore it
-./extra/viz/cli.py --help
-```
-
-## AMD Performance Counter Profiling
-
-Set VIZ to `-2` to save performance counters traces for the AMD backend.
-
-Use the CLI in `./extra/sqtt/roc.py` to explore the trace.

extra/nv_pma/decode.py

@@ -129,14 +129,6 @@ 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)
@@ -187,7 +179,11 @@ 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 = list(decode(raw, sm_ver))
+      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))
       print(f"\nDecoded {len(samples)} samples:")
-      print_samples(samples)
       print_aggregated(samples)

extra/viz/cli.py

@@ -1,4 +1,6 @@
 #!/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/amd/hw/helpers.py

@@ -6,7 +6,7 @@ Set USE_HW=1 to run on both emulator and hardware, comparing results.
 import ctypes, math, os, struct
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
 
-from tinygrad.renderer.amd.emu import run_asm
+from test.mockgpu.amd.emu import run_asm
 from tinygrad.renderer.amd.dsl import NULL, SCC, VCC_LO, VCC_HI, EXEC_LO, EXEC_HI, M0
 
 def _i32(f: float) -> int: return struct.unpack('<I', struct.pack('<f', f))[0]
@@ -75,7 +75,7 @@ def i642f(i: int) -> float: return struct.unpack('<d', struct.pack('<Q', i))[0]
 def assemble(instructions: list) -> bytes:
   return b''.join(inst.to_bytes() for inst in instructions)
 
-# Simple WaveState class for test output parsing (mirrors emu.py interface for tests)
+# Simple WaveState class for test output parsing (mirrors test/mockgpu/amd/emu.py interface for tests)
 class WaveState:
   def __init__(self):
     self.vgpr = [[0] * 256 for _ in range(32)]  # vgpr[lane][reg]

test/amd/test_compare_emulators.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 from pathlib import Path
 from tinygrad import Device
 
-from tinygrad.renderer.amd.emu import WaveState, _decode_at, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
+from test.mockgpu.amd.emu import WaveState, _decode_at, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
 from tinygrad.renderer.amd import decode_inst
 from test.amd.helpers import KernelInfo
 import tinygrad

test/amd/test_emu2_pcode.py

@@ -4,8 +4,8 @@ from collections import defaultdict
 from tinygrad.helpers import DEBUG
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.renderer.amd.emu import parse_pcode
-from tinygrad.renderer.amd.pcode import parse_expr
+from test.mockgpu.amd.emu import parse_pcode
+from test.mockgpu.amd.pcode import parse_expr
 from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE
 from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp
 

test/backend/test_dtype.py

@@ -251,7 +251,7 @@ class TestEmulatedHalf(TestHalfDType):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="half"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()
@@ -355,7 +355,7 @@ class TestEmulatedInt64DType(TestInt64DType):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()
@@ -371,7 +371,7 @@ class TestEmulatedUInt64DType(TestUint64DType):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="long"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()
@@ -385,7 +385,7 @@ class TestEmulatedBFloat16Type(TestBFloat16Type):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="bfloat16"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()
@@ -397,7 +397,7 @@ class TestEmulatedFp8e4m3(TestFp8e4m3):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e4m3"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()
@@ -409,7 +409,7 @@ class TestEmulatedFp8e5m2(TestFp8e5m2):
   def setUpClass(cls):
     cls.stack = contextlib.ExitStack()
     cls.stack.enter_context(Context(EMULATED_DTYPES="fp8e5m2"))
-    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10, allow_subnormal=False)
 
   @classmethod
   def tearDownClass(cls): cls.stack.close()

test/backend/test_dtype_alu.py

@@ -296,18 +296,21 @@ 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

@@ -465,6 +465,19 @@ class TestMultiTensor(unittest.TestCase):
     y_shard = norm_sharded(x_sharded).realize()
     np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
 
+  def test_sdpa_causal_shard_batch(self):
+    B, H, T, D = 4, 2, 10, 16
+    q = Tensor.rand(B, H, T, D)
+    k = Tensor.rand(B, H, T, D)
+    v = Tensor.rand(B, H, T, D)
+    q_shard = q.shard(devices_2, axis=0)
+    k_shard = k.shard(devices_2, axis=0)
+    v_shard = v.shard(devices_2, axis=0)
+    Tensor.realize(q, k, v, q_shard, k_shard, v_shard)
+    y = Tensor.scaled_dot_product_attention(q, k, v, is_causal=True).realize()
+    y_shard = Tensor.scaled_dot_product_attention(q_shard, k_shard, v_shard, is_causal=True).realize()
+    np.testing.assert_allclose(y_shard.numpy(), y.numpy(), atol=1e-6, rtol=1e-6)
+
   # NOTE: this is failing on LLVM CI, no idea why. Works locally.
   @slow
   def test_data_parallel_resnet(self):
@@ -1323,6 +1336,55 @@ 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,6 +3295,7 @@ 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

@@ -0,0 +1,19 @@
+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

@@ -1018,7 +1018,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.arange(16).contiguous().realize()
     GlobalCounters.reset()
     a[4] = 3
-    # TODO: update when this becomes lazy
+    self.assertEqual(GlobalCounters.kernel_count, 0)
+    a.realize()
     self.assertEqual(GlobalCounters.kernel_count, 1)
     self.assertListEqual(a.tolist(), [0, 1, 2, 3, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15])
 
@@ -1081,6 +1082,14 @@ 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
@@ -1093,12 +1102,16 @@ class TestSchedule(unittest.TestCase):
     for X in range(1,N): root = root + bufs[X][vi] + bufs[X][vj]
     self.assertEqual(root.item(), N * 2)
 
-  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])
+  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())
 
 class TestSwizzle(unittest.TestCase):
   def test_swizzle_simple(self):

test/backend/test_setitem.py

@@ -36,18 +36,6 @@ 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)
@@ -69,10 +57,6 @@ 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_noncontiguous(self):
-    t = Tensor.ones(4)
-    with self.assertRaises(RuntimeError): t[1] = 5
-
   def test_setitem_chained_indexing(self):
     # N[i][j] must work the same as N[i, j]
     N1 = Tensor.zeros((3, 3)).contiguous().realize()
@@ -162,6 +146,8 @@ class TestSetitem(unittest.TestCase):
     @TinyJit
     def f(t:Tensor, a:Tensor):
       t[2:4, 3:5] = a
+      # NOTE: without return t or an explicit realize, it's lazy and not captured
+      return t
 
     for i in range(1, 6):
       t = Tensor.zeros(6, 6).contiguous().realize()

test/external/external_test_hive_reset.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+import subprocess, sys
+from tinygrad.helpers import getenv
+
+LOOPS = getenv("LOOPS", 10)
+BROKEN = getenv("BROKEN", 0)
+
+BROKEN_KERNEL_SCRIPT = """
+from tinygrad.device import Device
+from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
+from tinygrad.runtime.support.compiler_amd import compile_hip
+dev = Device["AMD"]
+assert isinstance(dev, AMDDevice) and dev.is_am(), "Need AM driver (not KFD)"
+broken_src = '''
+extern "C" __attribute__((global)) void broken(int* dummy) {
+  volatile int* bad_ptr = (volatile int*)0xDEAD00000000ULL;
+  *bad_ptr = 0x42;
+}
+'''
+broken_lib = compile_hip(broken_src, dev.arch)
+broken_prg = AMDProgram(dev, "broken", broken_lib)
+buf = dev.allocator.alloc(64)
+try:
+  broken_prg(buf, global_size=(1,1,1), local_size=(1,1,1), wait=True)
+  print("  ERROR: Kernel did not fault!")
+except RuntimeError as e:
+  print(f"  Got expected error: {e}")
+"""
+
+for i in range(LOOPS):
+  print(f"=== Running hive_reset.py ({i+1}/{LOOPS}) ===")
+  subprocess.run([sys.executable, "extra/amdpci/hive_reset.py"], check=True)
+  print("=== hive_reset complete ===")
+
+  if BROKEN:
+    print(f"=== Running broken kernel ({i+1}/{LOOPS}) ===")
+    ret = subprocess.run([sys.executable, "-c", BROKEN_KERNEL_SCRIPT])
+    print(f"=== broken kernel exited with code {ret.returncode} ===")
+
+  print(f"=== Running test_tiny.py ({i+1}/{LOOPS}) ===")
+  ret = subprocess.run([sys.executable, "test/test_tiny.py", "TestTiny.test_plus"])
+  print(f"=== test_tiny.py exited with code {ret.returncode} ===")

test/helpers.py

@@ -41,14 +41,18 @@ def assert_jit_cache_len(fxn, expected_len):
     assert type(fxn.jit_cache[0].prg).__name__.endswith('Graph')
     assert len(fxn.jit_cache[0].prg.jit_cache) == expected_len
 
-def rand_for_dtype(dt:DType, size:int):
+def rand_for_dtype(dt:DType, size:int, allow_subnormal=True):
   if dtypes.is_unsigned(dt):
     return np.random.randint(0, 100, size=size, dtype=_to_np_dtype(dt))
   elif dtypes.is_int(dt):
     return np.random.randint(-100, 100, size=size, dtype=_to_np_dtype(dt))
   elif dt == dtypes.bool:
     return np.random.choice([True, False], size=size)
-  return np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
+  ret = np.random.uniform(-10, 10, size=size).astype(_to_np_dtype(dt))
+  if not allow_subnormal:
+    min_normal = 2.0 ** (2 - (1 << (dtypes.finfo(dt)[0] - 1)))
+    ret = np.where(np.abs(ret) < min_normal, 0, ret)
+  return ret
 
 def timeit(fxn:Callable[..., T], *args, **kwargs) -> tuple[T, float]:
   st = time.perf_counter_ns()

test/mockgpu/amd/README

@@ -4,12 +4,12 @@ Test with `pytest -n12 test/amd/`
 `AMD_LLVM=1 pytest -n12 test/amd/`
 
 * dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
+* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `AMD=1 MOCKGPU=1 PYTHON_REMU=1`
 * generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
-* pcode.py -- pseudocode to UOp transformation
+* test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
 * sqtt.py -- SQTT parser
 
-The code should be as readable and deduplicated as possible. asm and emu shouldn't be required for dsl.
+The code should be as readable and deduplicated as possible. emu (in test/mockgpu/amd/) shouldn't be required for dsl.
 
 The autogen folder is autogenerated from the AMD PDFs with `python3 -m tinygrad.renderer.amd.pdf --arch all`
 

test/mockgpu/amd/emu.py

@@ -67,7 +67,7 @@ from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
 from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
-from tinygrad.renderer.amd.pcode import parse_block, _FUNCS
+from test.mockgpu.amd.pcode import parse_block, _FUNCS
 
 MASK32 = 0xFFFFFFFF
 

test/mockgpu/helpers.py

@@ -24,7 +24,7 @@ class PythonRemu:
   user_data: list[int] = []  # All COMPUTE_USER_DATA registers (loaded into s[0:N])
 
   def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
-    from tinygrad.renderer.amd.emu import run_asm
+    from test.mockgpu.amd.emu import run_asm
     return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size, self.arch, self.user_data)
 
 def _try_dlopen_remu():

test/null/test_pattern_matcher.py

@@ -203,7 +203,7 @@ class TestPatternMatcher(unittest.TestCase):
 
   def _assert_eq_upat(self, a:UPat, b:UPat):
     assert (sorted(map(str,a.op)) if a.op else [] == (sorted(map(str,b.op)) if b.op else []))
-    assert (sorted(a.dtype) if a.dtype else [] == (sorted(b.dtype) if b.dtype else []))
+    assert (sorted(a.match_dtype) if a.match_dtype else [] == (sorted(b.match_dtype) if b.match_dtype else []))
     assert (a.name, type(a.src)) == (b.name, type(b.src))
     def simple_src(u:UPat):
       if u.src is None: return []

test/null/test_tensor_metadata.py

@@ -1,6 +1,7 @@
 import unittest
 from tinygrad import Tensor, dtypes
 from tinygrad.tensor import _METADATA
+from tinygrad.engine.realize import capturing
 from tinygrad.helpers import Context
 
 class TestTensorMetadata(unittest.TestCase):
@@ -62,6 +63,7 @@ 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()
@@ -90,5 +92,25 @@ class TestTensorMetadata(unittest.TestCase):
       si = out.schedule()[-1]
       self.assertEqual(si.metadata, ())
 
+  def _has_metadata(self, h, name):
+    items = []
+    capturing.append(type("", (), {"add": lambda _, ei: items.append(ei)})())
+    try: h.realize()
+    finally: capturing.clear()
+    return any(m.name == name for ei in items for m in ei.metadata)
+
+  def test_metadata_survives_realize_pending_assign(self):
+    shared = Tensor.rand(4)
+    c = Tensor.zeros(8).contiguous().realize()
+    c[:4].assign(shared)
+    self.assertTrue(self._has_metadata(c[:4].relu(), "relu"))
+
+  @unittest.expectedFailure
+  def test_metadata_lost_realize_pending_assign(self):
+    shared = Tensor.rand(4)
+    c = Tensor.zeros(8).contiguous().realize()
+    c[:4].assign(shared)
+    self.assertTrue(self._has_metadata((c[:4] + shared).relu(), "relu"))
+
 if __name__ == '__main__':
   unittest.main()

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 < 5 * tqdm_time
+    assert tinytqdm_time < 20 * tqdm_time
 
 if __name__ == '__main__':
   unittest.main()

test/null/test_uops_stats.py

@@ -68,28 +68,6 @@ 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,6 +364,15 @@ 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))]
@@ -574,6 +583,7 @@ 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/testextra/test_tk.py

@@ -5,21 +5,18 @@ from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import get_runner
 from tinygrad.engine.schedule import ExecItem
 from tinygrad.engine.jit import TinyJit
-from tinygrad.helpers import CI
 import numpy as np
 
 from extra.thunder.tiny.tk import WARP_THREADS
 from extra.thunder.tiny.tk.kernel import Kernel
 from extra.thunder.tiny.tk.tiles import ST_16X32, RT_16X32, RT_16X16, TileLayout
 
-@unittest.skipIf(CI or Device.DEFAULT not in ["AMD"], "only amd")
 class TestTK(unittest.TestCase):
   def setUp(self):
-    arch = Device["AMD"].arch
+    arch = getattr(Device[Device.DEFAULT].renderer, "arch", "")
     if not arch.startswith("gfx9"):
       self.skipTest(f"arch {arch} not supported")
 
-  @unittest.skipIf(CI, "no wmma in ci")
   def test_simple_matmul(self):
     N = 8192
     BLOCK_SIZE = 64
@@ -73,7 +70,6 @@ class TestTK(unittest.TestCase):
 
     np.testing.assert_allclose(c.numpy(), ref.numpy())
 
-  @unittest.skipIf(CI, "no wmma in ci")
   def test_simple_matmul_transposed(self):
     N = 8192
     BLOCK_N, BLOCK_M, BLOCK_K = 64, 64, 128

test/unit/test_assign.py

@@ -756,6 +756,26 @@ 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,6 +456,7 @@ 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"))
@@ -476,6 +477,7 @@ 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(dtypes.float64)
+  set_tensor = Tensor.randint(set_count, high=set_count).reshape(set_size).cast(indexed.dtype)
   return set_tensor
 
 @slow

test/unit/test_setitem_schedule.py

@@ -0,0 +1,150 @@
+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,7 +29,9 @@ 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}"
-      for s in k.src[0].src[1:] if k.op is Ops.END else k.src[1:]:
+      # 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:]:
         match (s := _unwrap_src(s)).op:
           case Ops.AFTER:
             children.setdefault(s.src[1], []).append(k)

tinygrad/helpers.py

@@ -182,6 +182,7 @@ 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/mixin/dtype.py

@@ -0,0 +1,96 @@
+from typing import Self
+from tinygrad.dtype import DType, dtypes
+
+class DTypeMixin:
+  @property
+  def dtype(self) -> DType: raise NotImplementedError
+
+  def cast(self, dtype:DType) -> Self: raise NotImplementedError
+
+  def element_size(self) -> int:
+    """
+    Returns the size in bytes of an individual element in the tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([5], dtype=dtypes.int16)
+    print(t.element_size())
+    ```
+    """
+    return self.dtype.itemsize
+
+  def is_floating_point(self) -> bool:
+    """
+    Returns `True` if the tensor contains floating point types, i.e. is one of `dtypes.float64`, `dtypes.float32`,
+    `dtypes.float16`, `dtypes.bfloat16`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([8, 9], dtype=dtypes.float32)
+    print(t.is_floating_point())
+    ```
+    """
+    return dtypes.is_float(self.dtype)
+
+  def float(self) -> Self:
+    """
+    Convenience method to cast `self` to a `float32` Tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([-1, 2, 3], dtype=dtypes.int32)
+    print(t.dtype, t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = t.float()
+    print(t.dtype, t.numpy())
+    ```
+    """
+    return self.cast(dtypes.float32)
+
+  def half(self) -> Self:
+    """
+    Convenience method to cast `self` to a `float16` Tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([-1, 2, 3], dtype=dtypes.int32)
+    print(t.dtype, t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = t.half()
+    print(t.dtype, t.numpy())
+    ```
+    """
+    return self.cast(dtypes.float16)
+
+  def int(self) -> Self:
+    """
+    Convenience method to cast `self` to a `int32` Tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([-1.5, -0.5, 0.0, 0.5, 1.5])
+    print(t.dtype, t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = t.int()
+    print(t.dtype, t.numpy())
+    ```
+    """
+    return self.cast(dtypes.int32)
+
+  def bool(self) -> Self:
+    """
+    Convenience method to cast `self` to a `bool` Tensor.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = Tensor([-1, 0, 1])
+    print(t.dtype, t.numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    t = t.bool()
+    print(t.dtype, t.numpy())
+    ```
+    """
+    return self.cast(dtypes.bool)
+
+  def bfloat16(self) -> Self: return self.cast(dtypes.bfloat16)
+  def double(self) -> Self: return self.cast(dtypes.double)
+  def long(self) -> Self: return self.cast(dtypes.long)
+  def short(self) -> Self: return self.cast(dtypes.short)

tinygrad/mixin/math.py

@@ -2,9 +2,10 @@ import math
 from typing import Self
 from tinygrad.uop import Ops
 from tinygrad.dtype import dtypes, ConstType
+from tinygrad.mixin.dtype import DTypeMixin
 
 
-class MathMixin:
+class MathMixin(DTypeMixin):
   # required to implement
   def alu(self, op: Ops, *src: Self) -> Self:
     raise NotImplementedError
@@ -23,16 +24,11 @@ class MathMixin:
     return self.ne(True)
 
   def neg(self) -> Self:
-    if (dtype := getattr(self, "dtype")) is None:
-      raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
-    return self.logical_not() if dtype.scalar() == dtypes.bool else self * (-1)
+    return self.logical_not() if self.dtype.scalar() == dtypes.bool else self * (-1)
 
   def _check_dtype(self) -> None:
-    if (dtype := getattr(self, "dtype")) is not None:
-      if isinstance(dtype, tuple):
-        dtype = dtype[0]
-      if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)):
-        raise RuntimeError(f"{dtype} is not supported")
+    if not (dtypes.is_bool(self.dtype) or dtypes.is_int(self.dtype)):
+      raise RuntimeError(f"{self.dtype} is not supported")
 
   def add(self, x: Self | ConstType, reverse: bool = False) -> Self:
     """
@@ -199,10 +195,10 @@ class MathMixin:
     return self.mod(x, True)
 
   def __lt__(self, x: Self | ConstType) -> Self:
-    return self.alu(Ops.CMPLT, self.ufix(x))
+    return self._binop(Ops.CMPLT, x, False)
 
   def __gt__(self, x: Self | ConstType) -> Self:
-    return self.ufix(x).alu(Ops.CMPLT, self)
+    return self._binop(Ops.CMPLT, x, True)
 
   def __ge__(self, x: Self | ConstType) -> Self:
     return (self < x).logical_not()
@@ -211,7 +207,7 @@ class MathMixin:
     return (self > x).logical_not()
 
   def ne(self, x: Self | ConstType) -> Self:
-    return self.alu(Ops.CMPNE, self.ufix(x))
+    return self._binop(Ops.CMPNE, x, False)
 
   def eq(self, x: Self | ConstType) -> Self:
     return self.ne(x).logical_not()
@@ -240,7 +236,17 @@ class MathMixin:
     return self.rshift(x, True)
 
   def maximum(self, x: Self | ConstType) -> Self:
-    return self.alu(Ops.MAX, self.ufix(x))
+    """
+    Computes element-wise maximum of `self` and `x`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-1, 2, 3]).maximum(1).numpy())
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-1, 2, 3]).maximum(Tensor([-4, -2, 9])).numpy())
+    ```
+    """
+    return self._binop(Ops.MAX, x, False)
 
   def minimum(self, x: Self | ConstType) -> Self:
     return -(-self).maximum(-self.ufix(x))
@@ -514,3 +520,73 @@ class MathMixin:
     ```
     """
     return self.sqrt().reciprocal()
+
+  def log(self) -> Self:
+    """
+    Computes the natural logarithm element-wise.
+
+    See: https://en.wikipedia.org/wiki/Logarithm
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1., 2., 4., 8.]).log().numpy())
+    ```
+    """
+    return self.log2()*math.log(2)
+
+  def log10(self) -> Self:
+    """
+    Computes the base-10 logarithm element-wise.
+
+    See: https://en.wikipedia.org/wiki/Logarithm
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1., 2., 4., 8.]).log10().numpy())
+    ```
+    """
+    return self.log2()*math.log10(2)
+
+  def atanh(self) -> Self:
+    """
+    Applies the Inverse Hyperbolic Tangent (atanh) function element-wise.
+
+    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#atanh
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-0.9, -0.6, -0.3, 0., 0.3, 0.6, 0.9]).atanh().numpy())
+    ```
+    """
+    return ((1 + self)/(1 - self)).log() / 2
+
+  def asinh(self) -> Self:
+    """
+    Applies the Inverse Hyperbolic Sine (asinh) function element-wise.
+
+    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#asinh
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy())
+    ```
+    """
+    return (self + (self.square() + 1).sqrt()).log()
+
+  def acosh(self) -> Self:
+    """
+    Applies the Inverse Hyperbolic Cosine (acosh) function element-wise.
+
+    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#acosh
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).acosh().numpy())
+    ```
+    """
+    return (self + (self.square() - 1).sqrt()).log()
+
+  def round(self) -> Self:
+    """
+    Rounds the tensor element-wise with rounding half to even.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).round().numpy())
+    ```
+    """
+    return ((self > 0).eq((b := self.trunc() / 2.0).trunc().eq(b))).where((self - 0.5).ceil(), (self + 0.5).floor())

tinygrad/mixin/movement.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING, Self
 from tinygrad.uop import Ops
-from tinygrad.helpers import prod, argfix, flatten, dedup, make_tuple, ceildiv
+from tinygrad.helpers import prod, argfix, argsort, flatten, dedup, make_tuple, ceildiv
 from tinygrad.uop.ops import resolve, smax
 
 if TYPE_CHECKING:
@@ -373,3 +373,27 @@ class MovementMixin:
     x = x.shrink_to(noop + flatten((k, o, 1) for k, o in zip(k_, o_))).reshape(noop + flatten((k, o) for k, o in zip(k_, o_)))
     # permute to move reduce to the end
     return x.permute(*range(len(noop)), *[len(noop) + i * 2 + 1 for i in range(len(i_))], *[len(noop) + i * 2 for i in range(len(i_))])
+
+  def unfold(self, dim:int, size, step:int) -> Self:
+    """
+    Unfolds the tensor along dimension `dim` into overlapping windows.
+
+    Each window has length `size` and begins every `step` elements of `self`.
+    Returns the input tensor with dimension `dim` replaced by dims `(n_windows, size)`
+    where `n_windows = (self.shape[dim] - size) // step + 1`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    unfolded = Tensor.arange(8).unfold(0,2,2)
+    print("\\n".join([repr(x.numpy()) for x in unfolded]))
+    ```
+    ```python exec="true" source="above" session="tensor" result="python"
+    unfolded = Tensor.arange(27).reshape(3,3,3).unfold(-1,2,3)
+    print("\\n".join([repr(x.numpy()) for x in unfolded]))
+    ```
+    """
+    if size < 0: raise RuntimeError(f'size must be >= 0 but got {size=}')
+    if step <= 0: raise RuntimeError(f'step must be > 0 but got {step=}')
+    if size > self.shape[dim]: raise RuntimeError(f'maximum size for tensor at dimension {dim} is {self.shape[dim]} but size is {size}')
+    dim = self._resolve_dim(dim)
+    perm_to_last = tuple(i for i in range(self.ndim) if i != dim) + (dim,)
+    return self.permute(perm_to_last)._pool((size,), step).permute(argsort(perm_to_last) + (self.ndim,))

tinygrad/renderer/amd/elf.py

@@ -9,98 +9,11 @@ from tinygrad.renderer.amd.dsl import Reg, FixedBitField
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_code_end # same encoding as RDNA4
 from tinygrad.runtime.autogen.amd.cdna.ins import s_nop as s_nop_cdna
 
-def put(dst:bytearray, off:int, data:bytes) -> None:
-  end = off + len(data)
-  if end > len(dst): raise ValueError("write past end of buffer")
-  dst[off:end] = data
-
-def create_elf(prg:bytes, kd:dict, arch:str) -> bytes:
-  is_cdna, is_rdna4 = arch == "cdna", arch == "rdna4"
-  padding_inst = (s_nop_cdna(0) if is_cdna else s_code_end()).to_bytes()
-  text = prg + padding_inst * ((hsa.AMD_ISA_ALIGN_BYTES - len(prg) % hsa.AMD_ISA_ALIGN_BYTES) % hsa.AMD_ISA_ALIGN_BYTES)
-  text_offset = round_up(ctypes.sizeof(libc.Elf64_Ehdr), hsa.AMD_ISA_ALIGN_BYTES)
-  rodata_offset = text_offset + len(text)
-
-  # ** pack rodata object
-  desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t()
-  desc.group_segment_fixed_size = kd.get("group_segment_fixed_size", 0)
-  desc.private_segment_fixed_size = kd.get("private_segment_fixed_size", 0)
-  desc.kernarg_size = kd.get("kernarg_size", 0)
-  desc.kernel_code_entry_byte_offset = text_offset-rodata_offset
-  # rsrc1
-  vgpr_granule = max(0, (kd["next_free_vgpr"] + 7) // 8 - 1)
-  # CDNA: add 6 for VCC(2) + FLAT_SCRATCH(2) + XNACK_MASK(2)
-  # next_free_sgpr is unused in RDNA
-  # NOTE: CU mode is the default, it seems faster and simpler
-  sgpr_granule = max(0, ceildiv(kd["next_free_sgpr"] + 6, 8) - 1) if is_cdna else 0
-  desc.compute_pgm_rsrc1 = (vgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT_SHIFT |
-                            sgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT |
-                            kd.get("float_round_mode_32", 0) << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_ROUND_MODE_32_SHIFT |
-                            kd.get("float_round_mode_16_64", 0) << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_ROUND_MODE_16_64_SHIFT |
-                            kd.get("float_denorm_mode_32", 0) << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_32_SHIFT |
-                            kd.get("float_denorm_mode_16_64", 3) << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_16_64_SHIFT |
-                            kd.get("dx10_clamp", 0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_DX10_CLAMP_SHIFT |
-                            kd.get("ieee_mode", 0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_IEEE_MODE_SHIFT |
-                            kd.get("fp16_overflow", 0) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX9_PLUS_FP16_OVFL_SHIFT |
-                            (0 if is_cdna else kd.get("workgroup_processor_mode", 0)) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_WGP_MODE_SHIFT |
-                            (0 if is_cdna else kd.get("memory_ordered", 1)) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_MEM_ORDERED_SHIFT |
-                            (0 if is_cdna else kd.get("forward_progress", 0)) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_FWD_PROGRESS_SHIFT)
-  # rsrc2
-  desc.compute_pgm_rsrc2 = (kd.get("enable_private_segment", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT_SHIFT |
-                            kd.get("user_sgpr_count", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_USER_SGPR_COUNT_SHIFT |
-                            kd.get("system_sgpr_workgroup_id_x", 1) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X_SHIFT |
-                            kd.get("system_sgpr_workgroup_id_y", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y_SHIFT |
-                            kd.get("system_sgpr_workgroup_id_z", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Z_SHIFT |
-                            kd.get("system_sgpr_workgroup_info", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_INFO_SHIFT |
-                            kd.get("system_vgpr_workitem_id", 0) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_VGPR_WORKITEM_ID_SHIFT)
-  # rsrc3
-  if is_cdna:
-    amdhsa_accum_offset = ((kd.get("accum_offset", 4) // 4) - 1) & amdgpu_kd.COMPUTE_PGM_RSRC3_GFX90A_ACCUM_OFFSET
-    desc.compute_pgm_rsrc3 = amdhsa_accum_offset << amdgpu_kd.COMPUTE_PGM_RSRC3_GFX90A_ACCUM_OFFSET_SHIFT
-  else:
-    desc.compute_pgm_rsrc3 = kd.get("shared_vgpr_count", 0) << amdgpu_kd.COMPUTE_PGM_RSRC3_GFX10_GFX11_SHARED_VGPR_COUNT_SHIFT
-  # kernel code properties
-  desc.kernel_code_properties = (kd.get("user_sgpr_dispatch_ptr", 0) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_PTR_SHIFT |
-                                 kd.get("user_sgpr_queue_ptr", 0) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_QUEUE_PTR_SHIFT |
-                                 kd.get("user_sgpr_kernarg_segment_ptr", 0) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR_SHIFT |
-                                 kd.get("user_sgpr_dispatch_id", 0) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_ID_SHIFT |
-                                 kd.get("user_sgpr_private_segment_size",0) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_SIZE_SHIFT |
-                                 kd.get("wavefront_size32", 0 if is_cdna else 1) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32_SHIFT |
-                                 kd.get("uses_dynamic_stack", 0) << amdgpu_kd.KERNEL_CODE_PROPERTY_USES_DYNAMIC_STACK_SHIFT)
-  rodata = bytes(desc)
-
-  # ** pack elf sections
-  sh_names:list[int] = []
-  strtab = bytearray(b"\x00")
-  for name in [".text", ".rodata", ".strtab"]:
-    sh_names.append(len(strtab))
-    strtab += name.encode("ascii") + b"\x00"
-
-  rodata_offset = round_up(text_offset+(text_size:=len(text)), hsa.AMD_KERNEL_CODE_ALIGN_BYTES)
-  strtab_offset = rodata_offset+(rodata_size:=len(rodata))
-  shdr_offset   = strtab_offset+(strtab_size:=len(strtab))
-
-  sections = [(libc.SHT_PROGBITS, libc.SHF_ALLOC | libc.SHF_EXECINSTR, text_offset, text_offset, text_size),
-              (libc.SHT_PROGBITS, libc.SHF_ALLOC, rodata_offset, rodata_offset, rodata_size),
-              (libc.SHT_STRTAB, 0, 0, strtab_offset, strtab_size)]
-  shdrs = (libc.Elf64_Shdr * len(sections))()
-  for i,s in enumerate(sections): shdrs[i] = libc.Elf64_Shdr(sh_names[i], *s)
-
-  ehdr = libc.Elf64_Ehdr()
-  ehdr.e_shoff, ehdr.e_shnum, ehdr.e_shstrndx = shdr_offset, len(sections), 2
-
-  elf = bytearray(shdr_offset + ctypes.sizeof(shdrs))
-  put(elf, 0, bytes(ehdr))
-  put(elf, text_offset, text)
-  put(elf, rodata_offset, rodata)
-  put(elf, strtab_offset, strtab)
-  put(elf, shdr_offset, bytes(shdrs))
-  return bytes(elf)
-
 _arch_map = {"gfx9": "cdna", "gfx10": "rdna3", "gfx11": "rdna3", "gfx12": "rdna4"}
 def do_assemble_amd(ctx, prg:UOp, lin:UOp) -> UOp:
   insts = [u.arg for u in lin.src]
-  # scan for max vgpr/sgpr
+
+  # ** scan for max vgpr/sgpr
   max_vgpr, max_sgpr = 0, 0
   for inst in insts:
     for name, field in inst._fields:
@@ -109,7 +22,8 @@ def do_assemble_amd(ctx, prg:UOp, lin:UOp) -> UOp:
       if not isinstance(val, Reg): continue
       if 256 <= val.offset < 512: max_vgpr = max(max_vgpr, (val.offset - 256) + val.sz)
       elif val.offset < 106: max_sgpr = max(max_sgpr, val.offset + val.sz)
-  # scan sink for metadata
+
+  # ** scan sink for metadata
   sink, n_bufs, n_vars, lds_size, gids = prg.src[0], 0, 0, 0, set()
   for u in sink.toposort():
     if u.op is Ops.PARAM: n_bufs += 1
@@ -119,9 +33,65 @@ def do_assemble_amd(ctx, prg:UOp, lin:UOp) -> UOp:
   src = "\n".join(str(inst) for inst in insts)
   code_bytes = b"".join(inst.to_bytes() for inst in insts)
   arch = next(v for k, v in _arch_map.items() if ctx.arch.startswith(k))
-  kd = {"kernarg_size":n_bufs*8+n_vars*4, "group_segment_fixed_size":lds_size,
-        "user_sgpr_kernarg_segment_ptr":1, "user_sgpr_count":2,
-        "system_sgpr_workgroup_id_x":int(0 in gids), "system_sgpr_workgroup_id_y":int(1 in gids), "system_sgpr_workgroup_id_z":int(2 in gids),
-        "next_free_vgpr":round_up(max_vgpr, 8), "next_free_sgpr":round_up(max_sgpr, 8)}
-  binary = create_elf(code_bytes, kd, arch)
+  is_cdna, is_rdna4 = arch == "cdna", arch == "rdna4"
+
+  # ** pad text to ISA alignment
+  padding_inst = (s_nop_cdna(0) if is_cdna else s_code_end()).to_bytes()
+  text = code_bytes + padding_inst * ((hsa.AMD_ISA_ALIGN_BYTES - len(code_bytes) % hsa.AMD_ISA_ALIGN_BYTES) % hsa.AMD_ISA_ALIGN_BYTES)
+  text_offset = round_up(ctypes.sizeof(libc.Elf64_Ehdr), hsa.AMD_ISA_ALIGN_BYTES)
+
+  # ** pack kernel descriptor (rodata)
+  next_free_vgpr, next_free_sgpr = round_up(max_vgpr, 8), round_up(max_sgpr, 8)
+  vgpr_granule = max(0, (next_free_vgpr + 7) // 8 - 1)
+  # CDNA: add 6 for VCC(2) + FLAT_SCRATCH(2) + XNACK_MASK(2), next_free_sgpr is unused in RDNA.
+  sgpr_granule = max(0, ceildiv(next_free_sgpr + 6, 8) - 1) if is_cdna else 0
+  desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t()
+  desc.group_segment_fixed_size = lds_size
+  desc.kernarg_size = n_bufs * 8 + n_vars * 4
+  desc.kernel_code_entry_byte_offset = -len(text)
+
+  # https://llvm.org/docs/AMDGPUUsage.html#amdgpu-amdhsa-compute-pgm-rsrc1-gfx6-gfx12-table
+  # NOTE: CU mode is the default
+  desc.compute_pgm_rsrc1 = (vgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WORKITEM_VGPR_COUNT_SHIFT |
+                            sgpr_granule << amdgpu_kd.COMPUTE_PGM_RSRC1_GRANULATED_WAVEFRONT_SGPR_COUNT_SHIFT |
+                            3 << amdgpu_kd.COMPUTE_PGM_RSRC1_FLOAT_DENORM_MODE_16_64_SHIFT |
+                            (0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_DX10_CLAMP_SHIFT |
+                            (0 if is_rdna4 else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX6_GFX11_ENABLE_IEEE_MODE_SHIFT |
+                            (0 if is_cdna else 1) << amdgpu_kd.COMPUTE_PGM_RSRC1_GFX10_PLUS_MEM_ORDERED_SHIFT)
+  desc.compute_pgm_rsrc2 = (2 << amdgpu_kd.COMPUTE_PGM_RSRC2_USER_SGPR_COUNT_SHIFT |
+                            int(0 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X_SHIFT |
+                            int(1 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y_SHIFT |
+                            int(2 in gids) << amdgpu_kd.COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Z_SHIFT)
+  desc.kernel_code_properties = (1 << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR_SHIFT |
+                                 (0 if is_cdna else 1) << amdgpu_kd.KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32_SHIFT)
+  rodata = bytes(desc)
+
+  # ** pack ELF
+  sh_names:list[int] = []
+  strtab = bytearray(b"\x00")
+  for name in [".text", ".rodata", ".strtab"]:
+    sh_names.append(len(strtab))
+    strtab += name.encode("ascii") + b"\x00"
+
+  rodata_offset = round_up(text_offset + (text_size := len(text)), hsa.AMD_KERNEL_CODE_ALIGN_BYTES)
+  strtab_offset = rodata_offset + (rodata_size := len(rodata))
+  shdr_offset   = strtab_offset + (strtab_size := len(strtab))
+
+  sections = [(libc.SHT_PROGBITS, libc.SHF_ALLOC | libc.SHF_EXECINSTR, text_offset, text_offset, text_size),
+              (libc.SHT_PROGBITS, libc.SHF_ALLOC, rodata_offset, rodata_offset, rodata_size),
+              (libc.SHT_STRTAB, 0, 0, strtab_offset, strtab_size)]
+  shdrs = (libc.Elf64_Shdr * len(sections))()
+  for i, s in enumerate(sections): shdrs[i] = libc.Elf64_Shdr(sh_names[i], *s)
+
+  ehdr = libc.Elf64_Ehdr()
+  ehdr.e_shoff, ehdr.e_shnum, ehdr.e_shstrndx = shdr_offset, len(sections), 2
+
+  elf = bytearray(shdr_offset + ctypes.sizeof(shdrs))
+  elf[0:ctypes.sizeof(ehdr)] = bytes(ehdr)
+  elf[text_offset:text_offset+text_size] = text
+  elf[rodata_offset:rodata_offset+rodata_size] = rodata
+  elf[strtab_offset:strtab_offset+strtab_size] = strtab
+  elf[shdr_offset:shdr_offset+ctypes.sizeof(shdrs)] = bytes(shdrs)
+  binary = bytes(elf)
+
   return prg.replace(src=prg.src[:3]+(UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))

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, self.dev.device): dst[:] = src
+    with cpu_profile(prof_desc, f"{self.dev.device}:COPY"): 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}", self.dev.device): pass
+    with cpu_profile(f"{src_dev.device} -> {dest_dev.device}", f"{self.dev.device}:COPY"): pass
   def _offset(self, buf, offset:int, size:int): pass
 
 class NullGraph(MultiGraphRunner):

tinygrad/runtime/ops_nv.py

@@ -753,7 +753,7 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.iface.rm_control(self.profiler, nv_gpu.NVB0CC_CTRL_CMD_POWER_REQUEST_FEATURES, power_params)
 
     self.pma_buf = self.iface.alloc(getenv("PMA_BUFFER_SIZE", 512) << 20, uncached=True, cpu_cached=True, cpu_access=True)
-    self.pma_bytes = self.iface.alloc(0x1000, uncached=True, cpu_cached=True, read_only=True)
+    self.pma_bytes = self.iface.alloc(0x1000, uncached=True, cpu_cached=True, cpu_access=True, read_only=True)
     self.pma_rptr = 0
 
     pma_stream = nv_gpu.struct_NVB0CC_CTRL_ALLOC_PMA_STREAM_PARAMS(hMemPmaBuffer=self.pma_buf.meta.hMemory,

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, self.dev.device):
+    with cpu_profile(prof_text, f"{self.dev.device}:COPY"):
       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}', self.dev.device): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      with cpu_profile(f'TINY -> {self.dev.device}', f"{self.dev.device}:COPY"): 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', self.dev.device): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      with cpu_profile(f'{self.dev.device} -> TINY', f"{self.dev.device}:COPY"): 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
+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 PCONTIG, partition, get_single_element
 from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify
 from tinygrad.codegen.opt import Opt
@@ -43,7 +43,6 @@ 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)
@@ -313,7 +312,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:=getenv("MAX_KERNEL_BUFFERS", DEVICE_MAX_BUFS.get(device, 0))): return None
+  if not (MAX_BUFS:=MAX_KERNEL_BUFFERS.value or DEVICE_MAX_BUFS.get(device, 0)): return None
 
   bufs: set[UOp] = set()
   def gate_input(u:UOp):
@@ -325,7 +324,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:
+      if s.op in GroupOp.Elementwise and s._device is not None:
         # 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)
@@ -555,7 +554,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.CALL, Ops.END,
+  (UPat(GroupOp.All-{Ops.PARAM, Ops.CONST, Ops.DEVICE, Ops.UNIQUE, Ops.LUNIQUE, Ops.DEFINE_VAR, Ops.BIND, 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)),
 ])
@@ -602,15 +601,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")
 
-  # 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] = {}
+  # 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}
   assign_rep: dict[UOp, UOp] = {}
-  for u in tsink.toposort():
-    if u.op is not Ops.AFTER: continue
-    kernel_assign[u.buf_uop] = u
+  for u in afters:
     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

@@ -189,12 +189,10 @@ class Tensor(OpMixin):
     all_tensors[weakref.ref(ret)] = None
     return ret
 
-  def _apply_broadcasted_uop(self, fxn:Callable, x:Tensor|ConstType, reverse=False) -> Tensor:
-    lhs,rhs = self._broadcasted(x, reverse)
-    return lhs._apply_uop(fxn, rhs)
-
   # _binop and alu are used by MathMixin
-  def _binop(self, op, x, reverse): return self._apply_broadcasted_uop(lambda *u: UOp.alu(u[0], op, *u[1:]), x, reverse)
+  def _binop(self, op, x, reverse):
+    lhs,rhs = self._broadcasted(x, reverse)
+    return lhs._apply_uop(lambda *u: u[0].alu(op, *u[1:]), rhs)
   def alu(self, op: Ops, *src: Tensor) -> Tensor: return self._apply_uop(lambda *u: u[0].alu(op, *u[1:]), *src)
 
   def requires_grad_(self, requires_grad=True) -> Tensor:
@@ -614,14 +612,15 @@ class Tensor(OpMixin):
     print(t.numpy())
     ```
     """
-    if not dtypes.is_float(dtype := to_dtype(dtype or dtypes.default_float)): raise ValueError(f"rand only supports float dtypes, got {dtype}")
+    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 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=dtype, **kwargs)
-    num = ceildiv(numel * dtype.itemsize, 4)
+    if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dt, **kwargs)
+    num = ceildiv(numel * dt.itemsize, 4)
 
     # generate per device seeds and rng counter if we haven't seen this device yet
     if device not in Tensor._device_seeds:
@@ -639,14 +638,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(dtype)
-    uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dtype.itemsize]
+    _, nmant = dtypes.finfo(dt)
+    uint_dtype = {1: dtypes.uint8, 2: dtypes.uint16, 4: dtypes.uint32, 8: dtypes.uint64}[dt.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=dtype).bitcast(uint_dtype)
-    bits = bits.rshift(dtype.bitsize - nmant).bitwise_or(one)
+    one = Tensor.ones_like(bits, device=bits.device, dtype=dt).bitcast(uint_dtype)
+    bits = bits.rshift(dt.bitsize - nmant).bitwise_or(one)
     # bitcast back to the original dtype and reshape
-    out = bits.bitcast(dtype)[:numel].sub(1).reshape(shape).requires_grad_(kwargs.get("requires_grad"))
+    out = bits.bitcast(dt)[:numel].sub(1).reshape(shape).requires_grad_(kwargs.get("requires_grad"))
     return out.contiguous() if contiguous else out
 
   # ***** creation helper functions *****
@@ -770,8 +769,9 @@ class Tensor(OpMixin):
     print(Tensor.eye(2, 4).numpy())
     ```
     """
-    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))
+    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))
     return t.cast(dtype or dtypes.default_float).requires_grad_(requires_grad)
 
   def _multi_like(self, fxn, *args, **kwargs) -> Tensor:
@@ -1214,6 +1214,26 @@ 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]), []
@@ -1309,10 +1329,13 @@ class Tensor(OpMixin):
       if is_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))
-    else: # basic setitem
-      if is_disk: self[indices].assign(v)
-      else:
-        self[indices].assign(v).realize()
+    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))
 
   def __delitem__(self, indices) -> None:
     raise TypeError("Tensor does not support deleting items")
@@ -1422,30 +1445,6 @@ class Tensor(OpMixin):
     assert chunks > 0, f"expect chunks to be greater than 0, got: {chunks}"
     return list(self.split(ceildiv(dim_sz, chunks) if dim_sz else [0]*chunks, dim=dim))
 
-  def unfold(self, dim:int, size:sint, step:int) -> Tensor:
-    """
-    Unfolds the tensor along dimension `dim` into overlapping windows.
-
-    Each window has length `size` and begins every `step` elements of `self`.
-    Returns the input tensor with dimension `dim` replaced by dims `(n_windows, size)`
-    where `n_windows = (self.shape[dim] - size) // step + 1`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    unfolded = Tensor.arange(8).unfold(0,2,2)
-    print("\\n".join([repr(x.numpy()) for x in unfolded]))
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    unfolded = Tensor.arange(27).reshape(3,3,3).unfold(-1,2,3)
-    print("\\n".join([repr(x.numpy()) for x in unfolded]))
-    ```
-    """
-    if size < 0: raise RuntimeError(f'size must be >= 0 but got {size=}')
-    if step <= 0: raise RuntimeError(f'step must be > 0 but got {step=}')
-    if size > self.shape[dim]: raise RuntimeError(f'maximum size for tensor at dimension {dim} is {self.shape[dim]} but size is {size}')
-    dim = self._resolve_dim(dim)
-    perm_to_last = tuple(i for i in range(self.ndim) if i != dim) + (dim,)
-    return self.permute(perm_to_last)._pool((size,), step).permute(argsort(perm_to_last) + (self.ndim,))
-
   def meshgrid(self:Tensor, *args:Tensor, indexing:str="ij") -> tuple[Tensor, ...]:
     """
     Generates coordinate matrices from coordinate vectors.
@@ -2821,7 +2820,7 @@ class Tensor(OpMixin):
     print(Tensor([False, True]).logical_not().numpy())
     ```
     """
-    return self.cast(dtypes.bool)._apply_broadcasted_uop(UOp.ne, True)
+    return self.cast(dtypes.bool).ne(True)
 
   def neg(self) -> Tensor:
     """
@@ -2845,30 +2844,6 @@ class Tensor(OpMixin):
     """
     return self._apply_uop(UOp.contiguous_backward)
 
-  def log(self) -> Tensor:
-    """
-    Computes the natural logarithm element-wise.
-
-    See: https://en.wikipedia.org/wiki/Logarithm
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([1., 2., 4., 8.]).log().numpy())
-    ```
-    """
-    return self.log2()*math.log(2)
-
-  def log10(self) -> Tensor:
-    """
-    Computes the base-10 logarithm element-wise.
-
-    See: https://en.wikipedia.org/wiki/Logarithm
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([1., 2., 4., 8.]).log10().numpy())
-    ```
-    """
-    return self.log2()*math.log10(2)
-
   def log2(self) -> Tensor:
     """
     Computes the base-2 logarithm element-wise.
@@ -2996,16 +2971,6 @@ class Tensor(OpMixin):
 
   # ***** math functions *****
 
-  def round(self: Tensor) -> Tensor:
-    """
-    Rounds the tensor element-wise with rounding half to even.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-3.5, -2.5, -1.5, -0.5, 0.5, 1.5, 2.5, 3.5]).round().numpy())
-    ```
-    """
-    return ((self > 0) == ((b := self.trunc() / 2.0).trunc() == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
-
   def lerp(self, end:Tensor, weight:Tensor|float) -> Tensor:
     """
     Linearly interpolates between `self` and `end` by `weight`.
@@ -3111,42 +3076,6 @@ class Tensor(OpMixin):
     """
     return (self.exp() + self.neg().exp()) / 2
 
-  def atanh(self) -> Tensor:
-    """
-    Applies the Inverse Hyperbolic Tangent (atanh) function element-wise.
-
-    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#atanh
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-0.9, -0.6, -0.3, 0., 0.3, 0.6, 0.9]).atanh().numpy())
-    ```
-    """
-    return ((1 + self)/(1 - self)).log() / 2
-
-  def asinh(self) -> Tensor:
-    """
-    Applies the Inverse Hyperbolic Sine (asinh) function element-wise.
-
-    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#asinh
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).asinh().numpy())
-    ```
-    """
-    return (self + (self.square() + 1).sqrt()).log()
-
-  def acosh(self) -> Tensor:
-    """
-    Applies the Inverse Hyperbolic Cosine (acosh) function element-wise.
-
-    - Described: https://en.wikipedia.org/wiki/Inverse_hyperbolic_functions#acosh
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-3., -2., -1., 0., 1., 2., 3.]).acosh().numpy())
-    ```
-    """
-    return (self + (self.square() - 1).sqrt()).log()
-
   def erf(self) -> Tensor:
     """
     Applies error function element-wise.
@@ -3266,7 +3195,7 @@ class Tensor(OpMixin):
       numerator, denominator = numerator.cast(dt), denominator.cast(dt)
       if rounding_mode == "trunc": return numerator.idiv(denominator)
       if rounding_mode == "floor":
-        truncate_div, truncate_mod = numerator.idiv(denominator), numerator._apply_broadcasted_uop(UOp.mod, denominator)
+        truncate_div, truncate_mod = numerator.idiv(denominator), numerator._binop(Ops.MOD, denominator, False)
         opposite_sign = ((numerator>0)&(denominator<0)) | ((numerator<0)&(denominator>0))
         return (opposite_sign&(truncate_mod!=0)).where(truncate_div-1, truncate_div)
     if rounding_mode == "trunc": return d.trunc().cast(output_dtype)
@@ -3348,19 +3277,6 @@ class Tensor(OpMixin):
     # NOTE: pow(int, float) -> int
     return ret.round().cast(self.dtype) if not reverse and not dtypes.is_float(self.dtype) and dtypes.is_float(exponent.dtype) else ret
 
-  def maximum(self, x:Tensor|ConstType) -> Tensor:
-    """
-    Computes element-wise maximum of `self` and `x`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-1, 2, 3]).maximum(1).numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([-1, 2, 3]).maximum(Tensor([-4, -2, 9])).numpy())
-    ```
-    """
-    return self._apply_broadcasted_uop(UOp.maximum, x)
-
   def minimum(self, x:Tensor|ConstType) -> Tensor:
     """
     Computes element-wise minimum of `self` and `x`.
@@ -3445,10 +3361,6 @@ class Tensor(OpMixin):
   def __ilshift__(self, x) -> Tensor: return self.assign(self.lshift(x)) # type: ignore[misc]
   def __irshift__(self, x) -> Tensor: return self.assign(self.rshift(x)) # type: ignore[misc]
 
-  def __lt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, False)
-  def __gt__(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.__lt__, x, True)
-  def ne(self, x) -> Tensor: return self._apply_broadcasted_uop(UOp.ne, x, False)
-
   def __eq__(self, x) -> Tensor: return self.eq(x)                      # type: ignore[override]
 
   # ***** encoding/decoding ops *****
@@ -3608,7 +3520,7 @@ class Tensor(OpMixin):
     # handle attention mask
     if is_causal:
       if attn_mask is not None: raise RuntimeError("cannot set attn_mask when is_causal=True")
-      attn_mask = qk.ones_like(requires_grad=False, device=self.device, dtype=dtypes.bool).tril()
+      attn_mask = qk.ones_like(requires_grad=False, dtype=dtypes.bool).tril()
     if attn_mask is not None:
       if attn_mask.dtype == dtypes.bool: attn_mask = attn_mask.where(0, -float("inf"))
       qk = qk + attn_mask
@@ -3812,17 +3724,6 @@ class Tensor(OpMixin):
 
   # ***** Tensor Properties *****
 
-  def element_size(self) -> int:
-    """
-    Returns the size in bytes of an individual element in the tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([5], dtype=dtypes.int16)
-    print(t.element_size())
-    ```
-    """
-    return self.dtype.itemsize
-
   def nbytes(self) -> int:
     """
     Returns the total number of bytes of all elements in the tensor.
@@ -3834,18 +3735,6 @@ class Tensor(OpMixin):
     """
     return int(self.numel()) * self.element_size()
 
-  def is_floating_point(self) -> bool:
-    """
-    Returns `True` if the tensor contains floating point types, i.e. is one of `dtypes.float64`, `dtypes.float32`,
-    `dtypes.float16`, `dtypes.bfloat16`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([8, 9], dtype=dtypes.float32)
-    print(t.is_floating_point())
-    ```
-    """
-    return dtypes.is_float(self.dtype)
-
   def size(self, dim:int|None=None) -> sint|tuple[sint, ...]:
     """
     Returns the size of the tensor. If `dim` is specified, return the length along dimension `dim`. Otherwise return the shape of the tensor.
@@ -3879,10 +3768,7 @@ class Tensor(OpMixin):
     print(t.dtype, t.numpy())
     ```
     """
-    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)
+    return self if self.dtype == (dt:=to_dtype(dtype)) else self._apply_uop(UOp.cast, dtype=dt)
 
   def bitcast(self, dtype:DTypeLike) -> Tensor:
     """
@@ -3915,71 +3801,6 @@ class Tensor(OpMixin):
       return Tensor.stack(*(tmp>>8*i*ns for i in range(os//ns)), dim=-1).flatten(-2).cast(new_uint).bitcast(dtype)
     return self._apply_uop(UOp.bitcast, dtype=dt) if self.dtype != dt else self
 
-  def float(self) -> Tensor:
-    """
-    Convenience method to cast `self` to a `float32` Tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([-1, 2, 3], dtype=dtypes.int32)
-    print(t.dtype, t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = t.float()
-    print(t.dtype, t.numpy())
-    ```
-    """
-    return self.cast(dtypes.float32)
-
-  def half(self) -> Tensor:
-    """
-    Convenience method to cast `self` to a `float16` Tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([-1, 2, 3], dtype=dtypes.int32)
-    print(t.dtype, t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = t.half()
-    print(t.dtype, t.numpy())
-    ```
-    """
-    return self.cast(dtypes.float16)
-
-  def int(self) -> Tensor:
-    """
-    Convenience method to cast `self` to a `int32` Tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([-1.5, -0.5, 0.0, 0.5, 1.5])
-    print(t.dtype, t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = t.int()
-    print(t.dtype, t.numpy())
-    ```
-    """
-    return self.cast(dtypes.int32)
-
-  def bool(self) -> Tensor:
-    """
-    Convenience method to cast `self` to a `bool` Tensor.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = Tensor([-1, 0, 1])
-    print(t.dtype, t.numpy())
-    ```
-    ```python exec="true" source="above" session="tensor" result="python"
-    t = t.bool()
-    print(t.dtype, t.numpy())
-    ```
-    """
-    return self.cast(dtypes.bool)
-
-  def bfloat16(self) -> Tensor: return self.cast(dtypes.bfloat16)
-  def double(self) -> Tensor: return self.cast(dtypes.double)
-  def long(self) -> Tensor: return self.cast(dtypes.long)
-  def short(self) -> Tensor: return self.cast(dtypes.short)
-
   # *** image Tensor function replacements ***
 
   def image_dot(self, w:Tensor, dtype:DTypeLike|None=None) -> Tensor:

tinygrad/uop/ops.py

@@ -893,13 +893,13 @@ def get_location() -> tuple[str, int]:
   return frm.f_code.co_filename, frm.f_lineno
 
 class UPat(OpMixin):
-  __slots__ = ("op", "dtype", "arg", "name", "src", "is_any")
+  __slots__ = ("op", "match_dtype", "arg", "name", "src", "is_any")
   def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None,
                src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,
                name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False):
     assert op is None or isinstance(op, (Ops, tuple, set)), "op must be Ops or tuple of Ops"
     self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op)
-    self.dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype)
+    self.match_dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype)
     self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject
     self.src: Any = None
     self.is_any = is_any
@@ -922,9 +922,14 @@ class UPat(OpMixin):
       upat_match = [src] if isinstance(src, UPat) else ([] if src is None else self.src[0])
       self.early_reject = {pp.op[0] for pp in upat_match if pp.op is not None and len(pp.op) == 1}
 
+  @property
+  def dtype(self) -> DType: return self.match_dtype[0] if self.match_dtype is not None else dtypes.void
+
+  def _check_dtype(self) -> None: pass
+
   def __reduce__(self):
-    return UPat, (self.op, self.dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location)
-  def named(self, name:str): return UPat(self.op, self.dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject)
+    return UPat, (self.op, self.match_dtype, self._in_src, self.arg, self.name, not self.strict_length, self.custom_early_reject, self.location)
+  def named(self, name:str): return UPat(self.op, self.match_dtype, self._in_src, self.arg, name, not self.strict_length, self.custom_early_reject)
 
   @staticmethod
   def any(*src): return UPat(src=src, is_any=True)
@@ -948,23 +953,23 @@ class UPat(OpMixin):
   # copied from UOp
   def sink(self, *srcs:UPat|None, **kwargs): return UPat(Ops.SINK, dtypes.void, (self,)+tuple([x for x in srcs if x is not None]), **kwargs)
   def index(self, idx:UPat, valid:UPat|None=None, **kwargs):
-    return UPat(Ops.INDEX, self.dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
+    return UPat(Ops.INDEX, self.match_dtype, (self,idx,valid) if valid is not None else (self,idx), **kwargs)
   def cast(self, dtype=None, **kwargs): return UPat(Ops.CAST, dtype, (self,), **kwargs)
   def bitcast(self, dtype=None): return UPat(Ops.BITCAST, dtype, (self,))
   def gep(self, i:int|None=None, **kwargs): return UPat(Ops.GEP, None, (self,), (i,) if i is not None else None, **kwargs)
   def load(self, *src:UPat, **kwargs): return UPat(Ops.LOAD, src=(self,)+src, **kwargs)
-  def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, self.dtype, (self,)+src, **kwargs)
-  def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.dtype, (self,x), **kwargs)
-  def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.dtype, src=(self,)+src, **kwargs)
-  def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.dtype, src=self, **kwargs)
-  def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.dtype, src=(self,)+args, **kwargs)
-  def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.dtype, (self,)+src, **kwargs)
-  def end(self, *src:UPat, **kwargs): return UPat(Ops.END, self.dtype, (self,)+src, **kwargs)
+  def store(self, *src:UPat, **kwargs): return UPat(Ops.STORE, self.match_dtype, (self,)+src, **kwargs)
+  def assign(self, x:UPat, **kwargs): return UPat(Ops.ASSIGN, self.match_dtype, (self,x), **kwargs)
+  def reduce(self, *src:UPat, **kwargs): return UPat(Ops.REDUCE, self.match_dtype, src=(self,)+src, **kwargs)
+  def broadcast(self, **kwargs): return UPat(Ops.VECTORIZE, self.match_dtype, src=self, **kwargs)
+  def contiguous(self, *args, **kwargs): return UPat(Ops.CONTIGUOUS, dtype=self.match_dtype, src=(self,)+args, **kwargs)
+  def after(self, *src:UPat, **kwargs): return UPat(Ops.AFTER, self.match_dtype, (self,)+src, **kwargs)
+  def end(self, *src:UPat, **kwargs): return UPat(Ops.END, self.match_dtype, (self,)+src, **kwargs)
 
-  def const_like(self, b:ConstLike): return UPat.const(self.dtype, cast(ConstType, b))
+  def const_like(self, b:ConstLike): return UPat.const(self.match_dtype, cast(ConstType, b))
   def alu(self, op:Ops, *src:UPat):
     asrc = (self,)+src
-    return UPat(op, dtypes.bool if op in {Ops.CMPLT, Ops.CMPNE} else asrc[-1].dtype, list(asrc) if op in GroupOp.Commutative else asrc)
+    return UPat(op, dtypes.bool if op in {Ops.CMPLT, Ops.CMPNE} else asrc[-1].match_dtype, list(asrc) if op in GroupOp.Commutative else asrc)
 
   def match(self:UPat, uop:UOp, store:dict[str, UOp]) -> list[dict[str, UOp]]:
     if self.is_any:
@@ -972,7 +977,7 @@ class UPat(OpMixin):
       return flatten([x for x in matches if x is not None])
     if (self.op is not None and uop.op not in self.op) or \
        (self.name is not None and store.setdefault(self.name, uop) is not uop) or \
-       (self.dtype is not None and uop.dtype not in self.dtype and uop.dtype.scalar() not in self.dtype) or \
+       (self.match_dtype is not None and uop.dtype not in self.match_dtype and uop.dtype.scalar() not in self.match_dtype) or \
        (self.arg is not None and self.arg != uop.arg) or \
        (len(uop.src) < self.required_len) or \
        (self.strict_length and len(uop.src) != self.required_len): return []

tinygrad/uop/upat.py

@@ -22,10 +22,10 @@ def _get_clause(self:UPat, base:UOp, depth=0) -> UOp:
   if self.strict_length or self.required_len > 0:
     and_clause.append(UOp(Ops.CUSTOM, src=(base,), arg=("len({0}.src)"+(" == " if self.strict_length else " >= ")+str(self.required_len))))
   if self.name is not None: and_clause.append(UOp(Ops.STORE, src=(UOp(Ops.DEFINE_VAR, arg=self.name), base)))
-  if self.dtype is not None:
-    if len(self.dtype) > 1:
-      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})"))
-    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.dtype[0])), arg="({0}.dtype == {1} or {0}.dtype._scalar == {1})"))
+  if self.match_dtype is not None:
+    if len(self.match_dtype) > 1:
+      and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=tuple(self.match_dtype))), arg="({0}.dtype in {1} or {0}.dtype._scalar in {1})"))
+    else: and_clause.append(UOp(Ops.CUSTOM, src=(base, UOp(Ops.BIND, arg=self.match_dtype[0])), arg="({0}.dtype == {1} or {0}.dtype._scalar == {1})"))
   if self.src is not None:
     # single match
     if len(self.src) == 1 and isinstance(self.src[0], tuple):