m_

undo the view and make it a mask, this fuses the setitem with any pending compute too.

one behavior change is that for target not backed by a buffer (const and arange), rangeify makes output contiguous under the hood.
this is stricter better than raise and ask user to call contiguous, as that would no longer be fuse-able.

.github/workflows/test.yml

@@ -708,6 +708,8 @@ jobs:
         run: SKIP_SLOW_TEST=1 AMD_LLVM=0 pytest -n=auto test/backend/test_ops.py -k "test_sparse_categorical_crossentropy or test_tril or test_nonzero or test_softmax_argmax" --durations 20
       - name: Run RDNA4 emulator tests
         run: MOCKGPU_ARCH=rdna4 python -m pytest test/test_tiny.py -v --durations 20
+      - name: Run CDNA4 emulator tests
+        run: AMD_LLVM=1 MOCKGPU_ARCH=cdna4 python -m pytest test/test_tiny.py -v --durations 20
 
   testnvidia:
     strategy:

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/mmapeak/mmapeak.py

@@ -4,15 +4,16 @@ import os
 os.environ["AMD_AQL"] = "1"
 
 from tinygrad import Tensor, Device
+from tinygrad.helpers import getenv
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.amd.dsl import Reg, Inst, s, v
 
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
-NUM_WAVES = 2
+NUM_WAVES = 4
 FLOPS_PER_MATMUL = 16*16*16*2
-INTERNAL_LOOP = 1_000_00
+INTERNAL_LOOP = getenv("LOOP", 10_000)
 INSTRUCTIONS_PER_LOOP = 200
 
 def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> list[Inst]:
@@ -22,15 +23,6 @@ def repeat(insts:list[Inst], n:int, counter_sreg:Reg) -> list[Inst]:
   branch_inst = s_cbranch_scc1(simm16=-((loop_sz // 4) + 1) & 0xFFFF)
   return [s_mov_b32(counter_sreg, n)] + insts + [sub_inst, cmp_inst, branch_inst, s_endpgm()]
 
-def make_kernel(insts:list[Inst]):
-  def fxn(A:UOp) -> UOp:
-    threads = UOp.special(WAVE_SIZE * NUM_WAVES, "lidx0")
-    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
-    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo("mmapeak", estimates=Estimates(ops=FLOPs, mem=0)))
-    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-  return fxn
-
 def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs):
   if accum:
     inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[1]:vgprIndices[2]], 1, acc_cd=1, **kwargs)
@@ -39,7 +31,12 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
   else:
     inst = instruction(v[0:vgprIndices[0]], v[vgprIndices[1]:vgprIndices[2]], v[vgprIndices[3]:vgprIndices[4]], v[vgprIndices[5]])
   insts = repeat([inst for _ in range(INSTRUCTIONS_PER_LOOP)], n=INTERNAL_LOOP, counter_sreg=s[1])
-  fxn = make_kernel(insts)
+  def fxn(A:UOp) -> UOp:
+    threads = UOp.special(WAVE_SIZE * NUM_WAVES, "lidx0")
+    gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
+    FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
+    sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
   dummy = Tensor.zeros(1).contiguous().realize()
   out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
   ei = out.schedule()[-1].lower()

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/bench_emu.py

@@ -1,267 +0,0 @@
-#!/usr/bin/env python3
-"""Benchmark comparing Python vs Rust RDNA3 emulators on real tinygrad kernels."""
-import ctypes, time, os
-from pathlib import Path
-
-from tinygrad.renderer.amd.emu import run_asm as python_run_asm, decode_program
-from tinygrad.renderer.amd import decode_inst
-from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP, SOPPOp
-
-import tinygrad
-EXTRA_DIR = Path(tinygrad.__file__).parent.parent / "extra"
-REMU_PATH = EXTRA_DIR / "remu/target/release/libremu.so"
-if not REMU_PATH.exists():
-  REMU_PATH = EXTRA_DIR / "remu/target/release/libremu.dylib"
-
-def get_rust_remu():
-  """Load the Rust libremu shared library."""
-  if not REMU_PATH.exists(): return None
-  remu = ctypes.CDLL(str(REMU_PATH))
-  remu.run_asm.restype = ctypes.c_int32
-  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
-                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
-  return remu
-
-def count_instructions(kernel: bytes) -> int:
-  """Count instructions in a kernel."""
-  return len(decode_program(kernel))
-
-def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
-  """Allocate buffers and return args pointer + valid ranges."""
-  if init_data is None: init_data = {}
-  buffers = []
-  for i, size in enumerate(buf_sizes):
-    padded = ((size + 15) // 16) * 16 + 16
-    data = init_data.get(i, b'\x00' * padded)
-    data_list = list(data) + [0] * (padded - len(data))
-    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
-    buffers.append(buf)
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  return buffers, args, args_ptr, ranges
-
-def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
-  """Benchmark an emulator and return average time."""
-  gx, gy, gz = global_size
-  lx, ly, lz = local_size
-  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-  lib_ptr = ctypes.addressof(kernel_buf)
-
-  # Warmup
-  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
-
-  # Timed runs
-  times = []
-  for _ in range(iterations):
-    start = time.perf_counter()
-    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr, rsrc2)
-    end = time.perf_counter()
-    if result != 0:
-      print(f"  {name} returned error: {result}")
-      return None
-    times.append(end - start)
-
-  return sum(times) / len(times)
-
-def profile_instructions(kernel: bytes):
-  """Profile individual instruction compile times."""
-  from tinygrad.renderer.amd.emu import _get_runner, _canonical_runner_cache
-  from tinygrad.helpers import Context
-  _get_runner.cache_clear()
-  _canonical_runner_cache.clear()
-
-  results = []
-  i = 0
-  while i < len(kernel):
-    inst = decode_inst(kernel[i:])
-    if isinstance(inst, SOPP) and inst.op == SOPPOp.S_CODE_END: break
-    inst_bytes = bytes(kernel[i:i + inst.size() + 4])
-    try: inst_str = repr(inst)
-    except Exception: inst_str = f"<{type(inst).__name__}>"
-
-    # Time the full compile (sink + render + compile)
-    start = time.perf_counter()
-    with Context(CCACHE=0):
-      runner, is_new = _get_runner(inst_bytes)
-    compile_time = time.perf_counter() - start
-
-    results.append({
-      'inst_str': inst_str + ('' if is_new else ' [CACHED]'),
-      'compile_ms': compile_time * 1000 if is_new else 0,
-    })
-    i += inst.size()
-
-  return sorted(results, key=lambda x: x['compile_ms'], reverse=True)
-
-def benchmark_python_split(kernel: bytes, global_size, local_size, args_ptr, rsrc2: int, iterations: int = 5):
-  """Benchmark Python emulator with compile and execution times."""
-  from tinygrad.renderer.amd.emu import _get_runner, _canonical_runner_cache
-  from tinygrad.helpers import Context
-  _get_runner.cache_clear()
-  _canonical_runner_cache.clear()
-  decode_program.cache_clear()
-
-  # Measure compile time (decode_program builds sinks, renders, and compiles)
-  compile_start = time.perf_counter()
-  with Context(CCACHE=0):
-    program = decode_program(kernel)
-  compile_time = time.perf_counter() - compile_start
-  n_compiled = len(_canonical_runner_cache)
-
-  # Execution time
-  exec_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, rsrc2, iterations)
-  return compile_time, exec_time, len(program), n_compiled
-
-def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes], int] | None:
-  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data, rsrc2)."""
-  try:
-    from tinygrad import Tensor
-    from tinygrad.runtime.support.elf import elf_loader
-    from tinygrad.runtime.autogen import hsa
-    import numpy as np
-    np.random.seed(42)
-
-    ops = {
-      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
-      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
-      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
-      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
-      "reduce_sum": lambda: Tensor.empty(4096).sum(),
-      "reduce_max": lambda: Tensor.empty(4096).max(),
-      "softmax": lambda: Tensor.empty(256).softmax(),
-      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
-      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
-      "gelu": lambda: Tensor.empty(1024).gelu(),
-      "exp": lambda: Tensor.empty(1024).exp(),
-      "sin": lambda: Tensor.empty(1024).sin(),
-    }
-
-    if op_name not in ops: return None
-    out = ops[op_name]()
-    sched = out.schedule()
-
-    for ei in sched:
-      lowered = ei.lower()
-      if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        image = memoryview(bytearray(lib))
-        _, sections, _ = elf_loader(lib)
-        rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"), -1)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_sizes = [b.nbytes for b in lowered.bufs]
-            # Get initial data from numpy arrays if available
-            buf_data = {}
-            for i, buf in enumerate(lowered.bufs):
-              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
-                try: buf_data[i] = bytes(buf.base._buf)
-                except Exception: pass
-            # Extract rsrc2 from ELF (same as ops_amd.py)
-            group_segment_size = image[rodata_entry:rodata_entry+4].cast("I")[0]
-            lds_size = ((group_segment_size + 511) // 512) & 0x1FF
-            code = hsa.amd_kernel_code_t.from_buffer_copy(bytes(image[rodata_entry:rodata_entry+256]) + b'\x00'*256)
-            rsrc2 = code.compute_pgm_rsrc2 | (lds_size << 15)
-            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data, rsrc2)
-    return None
-  except Exception as e:
-    print(f"  Error getting kernel: {e}")
-    return None
-
-TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "sin", "gelu", "matmul_small"]
-
-def main():
-  import argparse
-  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
-  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
-  parser.add_argument("--profile", type=str, default=None, help="Profile instructions for a specific kernel (e.g. 'sin')")
-  parser.add_argument("--top", type=int, default=20, help="Number of top instructions to show in profile")
-  args = parser.parse_args()
-
-  # Profile mode: show individual instruction timing
-  if args.profile:
-    kernel_info = get_tinygrad_kernel(args.profile)
-    if kernel_info is None:
-      print(f"Failed to get kernel for '{args.profile}'")
-      return
-    kernel = kernel_info[0]
-    print(f"Profiling instructions for '{args.profile}' kernel...")
-    print("=" * 110)
-    results = profile_instructions(kernel)
-    print(f"{'Instruction':<90} {'Compile(ms)':>12}")
-    print("-" * 110)
-    for r in results[:args.top]:
-      inst = r['inst_str'][:87] + "..." if len(r['inst_str']) > 90 else r['inst_str']
-      print(f"{inst:<90} {r['compile_ms']:>12.3f}")
-    print("-" * 110)
-    total = sum(r['compile_ms'] for r in results)
-    print(f"{'TOTAL':<90} {total:>12.3f}")
-    return
-
-  rust_remu = get_rust_remu()
-  if rust_remu is None:
-    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
-    print("Running Python-only benchmarks...\n")
-
-  print("=" * 90)
-  print("RDNA3 Emulator Benchmark: Python vs Rust")
-  print("=" * 90)
-
-  results = []
-
-  print("\n[TINYGRAD KERNELS]")
-  print("-" * 90)
-
-  for op_name in TINYGRAD_TESTS:
-    print(f"\n{op_name}:", end=" ", flush=True)
-    kernel_info = get_tinygrad_kernel(op_name)
-    if kernel_info is None:
-      print("failed to compile")
-      continue
-
-    kernel, global_size, local_size, buf_sizes, buf_data, rsrc2 = kernel_info
-    buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
-
-    # Benchmark Python emulator (must be first to measure compile time before cache is populated)
-    py_compile, py_exec, n_insts, n_compiled = benchmark_python_split(kernel, global_size, local_size, args_ptr, rsrc2, args.iterations)
-
-    n_workgroups = global_size[0] * global_size[1] * global_size[2]
-    n_threads = local_size[0] * local_size[1] * local_size[2]
-    total_work = n_insts * n_workgroups * n_threads
-
-    print(f"{n_insts} insts ({n_compiled} unique) × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
-    rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size,
-                                   args_ptr, rsrc2, args.iterations) if rust_remu else None
-
-    if py_compile is not None:
-      py_exec_rate = total_work / py_exec / 1e6
-      print(f"  Compile:        {py_compile*1000:8.3f} ms  ({n_compiled} unique)")
-      print(f"  Exec:           {py_exec*1000:8.3f} ms  ({py_exec_rate:7.2f} M ops/s)")
-    if rust_time:
-      rust_rate = total_work / rust_time / 1e6
-      speedup = py_exec / rust_time if py_exec else 0
-      print(f"  Rust:           {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
-
-    results.append((op_name, n_insts, n_compiled, n_workgroups, py_compile, py_exec, rust_time))
-
-  # Summary table
-  print("\n" + "=" * 110)
-  print("SUMMARY")
-  print("=" * 110)
-  print(f"{'Name':<16} {'Insts':<6} {'Unique':<6} {'WGs':<5} {'Compile (ms)':<14} {'Exec (ms)':<12} {'Rust (ms)':<12} {'Speedup':<10}")
-  print("-" * 110)
-
-  for name, n_insts, n_compiled, n_wgs, py_compile, py_exec, rust_time in results:
-    compile_ms = f"{py_compile*1000:.3f}" if py_compile else "error"
-    exec_ms = f"{py_exec*1000:.3f}" if py_exec else "error"
-    if rust_time:
-      rust_ms = f"{rust_time*1000:.3f}"
-      speedup = f"{py_exec/rust_time:.1f}x" if py_exec else "N/A"
-    else:
-      rust_ms, speedup = "N/A", "N/A"
-    print(f"{name:<16} {n_insts:<6} {n_compiled:<6} {n_wgs:<5} {compile_ms:<14} {exec_ms:<12} {rust_ms:<12} {speedup:<10}")
-
-if __name__ == "__main__":
-  os.environ["AMD"] = "1"
-  main()

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

@@ -1,12 +1,15 @@
 # Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
 import unittest, ctypes
 from dataclasses import dataclass
+from pathlib import Path
 from tinygrad import Device
 
-from tinygrad.renderer.amd.emu import WaveState, decode_program, 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
-from test.amd.bench_emu import REMU_PATH
+import tinygrad
+REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.so"
+if not REMU_PATH.exists(): REMU_PATH = Path(tinygrad.__file__).parent.parent / "extra/remu/target/release/libremu.dylib"
 
 def set_valid_mem_ranges(ranges): pass  # emu2 doesn't need this
 
@@ -89,7 +92,7 @@ class RustEmulator:
 class PythonEmulator:
   def __init__(self):
     self.state: WaveState | None = None
-    self.program: dict | None = None
+    self.program: dict[int, tuple] = {}  # lazily populated: pc -> (name, fxn, globals)
     self.vmem_buf = None
     self.lds_buf = None
     self.kernel_buf = None  # Keep kernel bytes alive
@@ -99,27 +102,29 @@ class PythonEmulator:
     import ctypes
     from tinygrad.device import Buffer, BufferSpec
     from tinygrad.dtype import dtypes
-    # Store kernel in a ctypes buffer so generic instructions can read from vmem at actual PC address
+    # Store kernel in a ctypes buffer so _decode_at can read from memory at actual PC address
     self.kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
     self.lib_addr = ctypes.addressof(self.kernel_buf)
-    # Remap program dict to use actual addresses (like run_asm does)
-    program_raw = decode_program(kernel)
-    self.program = {self.lib_addr + offset: val for offset, val in program_raw.items()}
+    self.program = {}
     self.state = WaveState(n_lanes)
     self.state.pc = self.lib_addr  # Set PC to code base address
     self.vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
     self.lds_buf = Buffer('CPU', 65536 // 4, dtypes.uint32).ensure_allocated()
 
+  def _ensure_decoded(self, pc: int):
+    if pc not in self.program:
+      runner = _decode_at(pc, "rdna3")
+      self.program[pc] = (runner.p.function_name, runner._prg.fxn, runner.p.globals)
+
   def step(self) -> int:
     import ctypes
-    assert self.program is not None and self.state is not None
+    assert self.state is not None
     pc = self.state.pc
-    if pc == 0xFFFFFFFFFFFFFFFF or pc not in self.program: return -1
-    name, fxn, globals_list, _runner = self.program[pc]
-    if fxn is None: return 1  # unsupported instruction
+    if pc == 0xFFFFFFFFFFFFFFFF: return -1
+    self._ensure_decoded(pc)
+    name, fxn, globals_list = self.program[pc]
     buf_addrs = {0: self.state.sgpr_buf._buf.va_addr, 1: self.state.vgpr_buf._buf.va_addr,  # type: ignore[union-attr]
                  2: self.vmem_buf._buf.va_addr, 3: self.lds_buf._buf.va_addr}  # type: ignore[union-attr]
-    # Direct ctypes call - bypasses HCQ overhead
     fxn(*[ctypes.c_uint64(buf_addrs[g]) for g in globals_list], ctypes.c_int32(0))
     return -1 if self.state.pc == 0xFFFFFFFFFFFFFFFF else 0
 
@@ -140,7 +145,7 @@ class PythonEmulator:
                          exec_mask=sgpr[EXEC_LO.offset], sgpr=sgpr, vgpr=vgpr)
 
 def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
-                      local_size: tuple[int, int, int], program, max_steps: int, debug: bool, trace_len: int,
+                      local_size: tuple[int, int, int], max_steps: int, debug: bool, trace_len: int,
                       kernel_idx: int = 0, max_workgroups: int = 8) -> tuple[bool, str, int]:
   """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
   gx, gy, gz = global_size
@@ -181,9 +186,9 @@ def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: t
             rust_before = rust.get_snapshot()
             python_before = python.get_snapshot()
 
-            assert python.program is not None
-            inst_info = python.program.get(python.lib_addr + python_before.pc * 4)  # Convert word offset to actual address
-            inst_hex_name = inst_info[0] if inst_info else f"unknown at PC={python_before.pc}"
+            pc_addr = python.lib_addr + python_before.pc * 4  # Convert word offset to actual address
+            python._ensure_decoded(pc_addr)
+            inst_hex_name = python.program[pc_addr][0]
             # Decode the instruction to get mnemonic for sync_after checks
             try:
               # Format is mnemonic_hexbytes, e.g. v_exp_f32_e32_014b027e -> hex is 014b027e
@@ -310,12 +315,11 @@ def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int
     kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
     set_valid_mem_ranges(kernel_ranges)
 
-    program = decode_program(kernel.code)
     n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
 
     ok, msg, steps = run_single_kernel(
       kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
-      kernel.local_size, program, max_steps, debug, trace_len, ki
+      kernel.local_size, max_steps, debug, trace_len, ki
     )
     total_steps += steps
     if not ok:
@@ -341,9 +345,8 @@ def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list,
   ranges.add((args_ptr, ctypes.sizeof(args)))
   set_valid_mem_ranges(ranges)
 
-  program = decode_program(kernel)
   # Legacy wrapper assumes local_size = (n_lanes, 1, 1)
-  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), program, max_steps, debug, trace_len)
+  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, (n_lanes, 1, 1), max_steps, debug, trace_len)
   return ok, msg
 
 def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:

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/amd/test_rdna4_emu.py

@@ -1,96 +0,0 @@
-import unittest, ctypes
-from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
-from tinygrad.renderer.amd.dsl import v, s
-from tinygrad.renderer.amd.emu import WaveState, decode_program
-from tinygrad.device import Buffer, BufferSpec
-from tinygrad.dtype import dtypes
-
-class TestRDNA4Emu(unittest.TestCase):
-  def _run(self, insts: list, sgprs: dict[int, int] | None = None, vgprs: dict[tuple[int, int], int] | None = None) -> WaveState:
-    """Run instructions and return final WaveState."""
-    # Add S_ENDPGM if not present
-    if not any(isinstance(i, ir4.SOPP) and i.op == ir4.SOPPOp.S_ENDPGM for i in insts):
-      insts = list(insts) + [ir4.SOPP(ir4.SOPPOp.S_ENDPGM, simm=0)]
-
-    # Assemble and decode
-    code = b''.join(i.to_bytes() for i in insts)
-    code_buf = (ctypes.c_uint8 * len(code)).from_buffer_copy(code)
-    code_addr = ctypes.addressof(code_buf)
-    program_raw = decode_program(code, "rdna4")
-    program = {code_addr + offset: val for offset, val in program_raw.items()}
-
-    # Setup wave state
-    st = WaveState(n_lanes=1)
-    st.pc = code_addr
-    for idx, val in (sgprs or {}).items(): st._write_sgpr(idx, val)
-    for (reg, lane), val in (vgprs or {}).items(): st._write_vgpr(reg, lane, val)
-
-    # Setup vmem buffer with external_ptr=0 (maps to address 0, allows any pointer access)
-    vmem_buf = Buffer('CPU', 1 << 40, dtypes.uint32, options=BufferSpec(external_ptr=0)).ensure_allocated()
-
-    # Execute
-    c_bufs = [ctypes.c_uint64(st.sgpr_buf._buf.va_addr), ctypes.c_uint64(st.vgpr_buf._buf.va_addr),
-              ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(0), ctypes.c_uint64(0)]
-    for _ in range(100):
-      if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF or pc not in program: break
-      _, fxn, globals_list, _ = program[pc]
-      fxn(*[c_bufs[g] for g in globals_list])
-    return st
-
-  def test_vopd_dual_mov(self):
-    """Test VOPD with two V_DUAL_MOV_B32 operations: v[1]=s[1], v[2]=s[2]."""
-    insts = [ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-                      vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0])]
-    st = self._run(insts, sgprs={1: 0x40e00000, 2: 0x41100000})  # 7.0f, 9.0f
-    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
-    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
-
-  def test_vopd_dual_mov_after_other_vopd(self):
-    """Test VOPD reuse: first VOPD(v[3]=0, v[0]=?), then VOPD(v[1]=s[1], v[2]=s[2])."""
-    # This matches the BEAM kernel sequence that fails
-    insts = [
-      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),  # v[3]=0, v[0]=s[0]
-      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),  # v[1]=s[1], v[2]=s[2]
-    ]
-    st = self._run(insts, sgprs={0: 0x40a00000, 1: 0x40e00000, 2: 0x41100000})  # 5.0f, 7.0f, 9.0f
-    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
-    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
-
-  def test_vopd_with_s_add_f32_sequence(self):
-    """Test full BEAM kernel sequence: s_add_f32 then VOPD."""
-    # This is the exact sequence from the failing BEAM kernel
-    insts = [
-      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[0], ssrc0=s[0], ssrc1=s[8]),   # s[0] = s[0] + s[8]
-      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[1], ssrc0=s[1], ssrc1=s[9]),   # s[1] = s[1] + s[9]
-      ir4.SOP2(ir4.SOP2Op.S_ADD_F32, sdst=s[2], ssrc0=s[2], ssrc1=s[10]),  # s[2] = s[2] + s[10]
-      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-               vdstx=v[3], vdsty=v[0], srcx0=0, srcy0=s[0], vsrcx1=v[0], vsrcy1=v[0]),
-      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
-    ]
-    # Input: s[0:2] = [1,2,3], s[8:10] = [4,5,6]
-    # After s_add_f32: s[0:2] = [5,7,9]
-    st = self._run(insts, sgprs={0: 0x3f800000, 1: 0x40000000, 2: 0x40400000,  # 1.0, 2.0, 3.0
-                                  8: 0x40800000, 9: 0x40a00000, 10: 0x40c00000})  # 4.0, 5.0, 6.0
-    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
-    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
-
-  def test_s_mov_b32_then_vopd(self):
-    """Test s_mov_b32 followed by VOPD - simulates BEAM kernel sequence."""
-    # Use s_mov_b32 with SGPR source (copy from pre-initialized SGPRs)
-    # s[10:12] will have values set by test harness, copy to s[0:2], then VOPD to VGPRs
-    insts = [
-      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[0], ssrc0=s[10]),  # s[0] = s[10]
-      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[1], ssrc0=s[11]),  # s[1] = s[11]
-      ir4.SOP1(ir4.SOP1Op.S_MOV_B32, sdst=s[2], ssrc0=s[12]),  # s[2] = s[12]
-      ir4.VOPD(ir4.VOPDOp.V_DUAL_MOV_B32, ir4.VOPDOp.V_DUAL_MOV_B32,
-               vdstx=v[1], vdsty=v[2], srcx0=s[1], srcy0=s[2], vsrcx1=v[0], vsrcy1=v[0]),
-    ]
-    st = self._run(insts, sgprs={10: 0x40a00000, 11: 0x40e00000, 12: 0x41100000})  # 5.0, 7.0, 9.0
-    self.assertEqual(st._read_vgpr(1, 0), 0x40e00000)  # v[1] = 7.0
-    self.assertEqual(st._read_vgpr(2, 0), 0x41100000)  # v[2] = 9.0
-
-if __name__ == '__main__':
-  unittest.main()

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

@@ -1323,6 +1323,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/amddriver.py

@@ -90,9 +90,9 @@ class AMDDriver(VirtDriver):
   def _prepare_gpu(self, gpu_id):
     self.doorbells[gpu_id] = memoryview(bytearray(0x2000))
     self.gpus[gpu_id] = AMDGPU(gpu_id)
-    # IP versions: rdna3 = GC 11.0.0, NBIF 4.3.0; rdna4 = GC 12.0.0, NBIF 6.3.1
     ip_versions = {"rdna3": {"gc": (11, 0, 0), "sdma": (6, 0, 0), "nbif": (4, 3, 0)},
-                   "rdna4": {"gc": (12, 0, 0), "sdma": (6, 0, 0), "nbif": (6, 3, 1)}}[MOCKGPU_ARCH]
+                   "rdna4": {"gc": (12, 0, 0), "sdma": (6, 0, 0), "nbif": (6, 3, 1)},
+                   "cdna4": {"gc": (9, 5, 0), "sdma": (4, 4, 5), "nbif": (7, 9, 0)}}[MOCKGPU_ARCH]
     def ip_discovery_files(hwid, ver, base_addr):
       p = f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}/0'
       return [VirtFile(f'/sys/class/drm/renderD{gpu_id}/device/ip_discovery/die/0/{hwid}', functools.partial(DirFileDesc, child_names=['0'])),

test/mockgpu/amd/amdgpu.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import getbits, to_mv, getenv
 from tinygrad.runtime.support import c
 
 MOCKGPU_ARCH = getenv("MOCKGPU_ARCH", "rdna3")
-GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000}[MOCKGPU_ARCH]
+GFX_TARGET_VERSION = {"rdna3": 110000, "rdna4": 120000, "cdna4": 90500}[MOCKGPU_ARCH]
 import tinygrad.runtime.autogen.amd_gpu as amd_gpu, tinygrad.runtime.autogen.am.pm4_nv as pm4
 
 SDMA_MAX_COPY_SIZE = 0x400000
@@ -106,8 +106,8 @@ class PM4Executor(AMDQueue):
     return (self.rptr[0] - prev_rptr) + executed_in_ib
 
   def _exec_acquire_mem(self, n):
-    assert n == 6
-    for _ in range(7): self._next_dword() # TODO: implement
+    assert n in (5, 6)
+    for _ in range(n + 1): self._next_dword() # TODO: implement
 
   def _exec_release_mem(self, n):
     assert n == 6
@@ -184,6 +184,12 @@ class PM4Executor(AMDQueue):
     args_addr = self.gpu.regs[regCOMPUTE_USER_DATA_0] + (self.gpu.regs[regCOMPUTE_USER_DATA_0 + 1] << 32)
     lc = [self.gpu.regs[i] for i in range(regCOMPUTE_NUM_THREAD_X, regCOMPUTE_NUM_THREAD_X+3)]
     rsrc2 = self.gpu.regs[regCOMPUTE_PGM_RSRC2]
+    # Read all user data registers (hardware loads these directly into s[0:N])
+    user_sgpr_count = (rsrc2 >> 1) & 0x1F  # USER_SGPR_COUNT is bits 1:5
+    user_data = []
+    for i in range(user_sgpr_count):
+      try: user_data.append(self.gpu.regs[regCOMPUTE_USER_DATA_0 + i])
+      except KeyError: user_data.append(0)
 
     prg_sz = 0
     for st,sz in self.gpu.mapped_ranges:
@@ -197,11 +203,12 @@ class PM4Executor(AMDQueue):
     scratch_size = wavesize * 4  # This gives the scratch size per thread (lane)
 
     assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # Pass valid memory ranges, rsrc2, scratch_size and arch to Python emulator
+    # Pass valid memory ranges, rsrc2, scratch_size, arch, and user data registers to Python emulator
     if hasattr(remu, 'valid_mem_ranges'): remu.valid_mem_ranges = self.gpu.mapped_ranges
     if hasattr(remu, 'rsrc2'): remu.rsrc2 = rsrc2
     if hasattr(remu, 'scratch_size'): remu.scratch_size = scratch_size
     if hasattr(remu, 'arch'): remu.arch = self.gpu.arch
+    if hasattr(remu, 'user_data'): remu.user_data = user_data
     err = remu.run_asm(prg_addr, prg_sz, *gl, *lc, args_addr)
     if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
 
@@ -318,7 +325,7 @@ class AMDGPU(VirtGPU):
     self.regs = AMDGPURegisters()
     self.mapped_ranges = set()
     self.queues = []
-    self.arch = MOCKGPU_ARCH
+    self.arch = "cdna" if MOCKGPU_ARCH == "cdna4" else MOCKGPU_ARCH
 
   def map_range(self, vaddr, size): self.mapped_ranges.add((vaddr, size))
   def unmap_range(self, vaddr, size): self.mapped_ranges.remove((vaddr, size))
@@ -329,7 +336,7 @@ class AMDGPU(VirtGPU):
     self.queues.append(SDMAExecutor(self, base, size, rptr, wptr))
     return len(self.queues) - 1
 
-gpu_props = """cpu_cores_count 0
+_gpu_props_rdna = """cpu_cores_count 0
 simd_count 192
 mem_banks_count 1
 caches_count 206
@@ -367,3 +374,44 @@ sdma_fw_version 20
 unique_id 11673270660693242239
 num_xcc 1
 max_engine_clk_ccompute 2400"""
+
+_gpu_props_cdna = """cpu_cores_count 0
+simd_count 304
+mem_banks_count 1
+caches_count 206
+io_links_count 1
+p2p_links_count 5
+cpu_core_id_base 0
+simd_id_base 2147488032
+max_waves_per_simd 16
+lds_size_in_kb 128
+gds_size_in_kb 0
+num_gws 64
+wave_front_size 64
+array_count 16
+simd_arrays_per_engine 4
+cu_per_simd_array 19
+simd_per_cu 2
+max_slots_scratch_cu 32
+gfx_target_version {gfx_target_version}
+vendor_id 4098
+device_id 29772
+location_id 34304
+domain 0
+drm_render_minor {drm_render_minor}
+hive_id 0
+num_sdma_engines 2
+num_sdma_xgmi_engines 0
+num_sdma_queues_per_engine 6
+num_cp_queues 8
+max_engine_clk_fcompute 2100
+local_mem_size 0
+fw_version 2140
+capability 671588992
+debug_prop 1495
+sdma_fw_version 20
+unique_id 11673270660693242239
+num_xcc 1
+max_engine_clk_ccompute 2100"""
+
+gpu_props = _gpu_props_cdna if MOCKGPU_ARCH == "cdna4" else _gpu_props_rdna

test/mockgpu/amd/emu.py

@@ -7,7 +7,7 @@
 #   arg=4: scratch - per-lane scratch memory
 from __future__ import annotations
 import ctypes, functools, re, platform, subprocess, tempfile
-from typing import Any, Callable
+from typing import Callable
 
 # Set/restore DAZ+FTZ (denormals-are-zero + flush-to-zero) to match RDNA3 default float mode
 # x86: MXCSR bits DAZ(6)+FTZ(15), ARM64: FPCR bit FZ(24)
@@ -61,11 +61,13 @@ from tinygrad.engine.realize import get_runner
 from tinygrad.renderer.amd import decode_inst
 from tinygrad.runtime.autogen.amd.rdna3.str_pcode import PCODE as PCODE_RDNA3
 from tinygrad.runtime.autogen.amd.rdna4.str_pcode import PCODE as PCODE_RDNA4
+from tinygrad.runtime.autogen.amd.cdna.str_pcode import PCODE as PCODE_CDNA
 from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
 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
 
@@ -160,7 +162,7 @@ _pcode_fixes = {
 
 def _get_pcode_dict(op) -> dict:
   """Return the PCODE dictionary for the given opcode based on its architecture."""
-  return PCODE_RDNA4 if 'rdna4' in type(op).__module__ else PCODE_RDNA3
+  return PCODE_CDNA if 'cdna' in type(op).__module__ else PCODE_RDNA4 if 'rdna4' in type(op).__module__ else PCODE_RDNA3
 
 # Pcode parser
 @functools.cache
@@ -465,8 +467,8 @@ class _Ctx:
     pcode = get_pcode(op)
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
     if 'VCC' not in srcs: srcs['VCC'] = self.rsgpr_dyn(_c(vcc_reg))
-    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane,
-                 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0)})  # rounding mode: 0=RNE, RTZ constant
+    srcs.update({'EXEC': exec_mask, 'SCC': self.rsgpr_dyn(_c(SCC.offset)), 'laneId': lane, 'VDST': vdst_reg,
+                 'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0)})  # rounding mode constants
     _, assigns = parse_pcode(pcode, srcs)
 
     # For integer ops with clamp, compute overflow using wide arithmetic
@@ -543,10 +545,11 @@ class _Ctx:
 
 def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
   simm16 = ctx.inst_field_signed(type(inst).simm16).cast(dtypes.int16)
-  if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM):
+  if inst.op in (ir3.SOPPOp.S_ENDPGM, ir4.SOPPOp.S_ENDPGM, irc.SOPPOp.S_ENDPGM):
     return UOp.sink(ctx.wsgpr_dyn(_c(PC_LO_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)),
                           ctx.wsgpr_dyn(_c(PC_HI_IDX), UOp.const(dtypes.uint32, 0xFFFFFFFF)))
-  if inst.op in (ir3.SOPPOp.S_NOP, ir4.SOPPOp.S_NOP): return UOp.sink(*ctx.inc_pc())  # S_NOP is a no-op
+  # S_NOP and S_WAITCNT are no-ops in emulator (no pipeline/cache to wait on)
+  if inst.op in (ir3.SOPPOp.S_NOP, ir4.SOPPOp.S_NOP, irc.SOPPOp.S_NOP, irc.SOPPOp.S_WAITCNT): return UOp.sink(*ctx.inc_pc())
   # NOTE: we ignore SOPPs without PCODE
   if inst.op in _get_pcode_dict(inst.op):
     pcode = get_pcode(inst.op)
@@ -562,10 +565,7 @@ def _compile_sopp(inst: ir3.SOPP | ir4.SOPP, ctx: _Ctx) -> UOp:
 
 def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
   # Cache invalidation instructions are no-ops in the emulator (we don't model caches)
-  cache_inv_ops = [ir3.SMEMOp.S_GL1_INV, ir3.SMEMOp.S_DCACHE_INV, ir4.SMEMOp.S_DCACHE_INV]
-  if hasattr(ir4.SMEMOp, 'S_GL1_INV'): cache_inv_ops.append(ir4.SMEMOp.S_GL1_INV)
-  if inst.op in cache_inv_ops:
-    return UOp.sink(*ctx.inc_pc())
+  if '_INV' in inst.op.name: return UOp.sink(*ctx.inc_pc())
   # Dynamic sbase field (bits 5:0) - SGPR pair, field value * 2 = register offset
   sbase = ctx.inst_field(type(inst).sbase) * _c(2)
   # Dynamic sdata field (bits 12:6) - destination SGPR
@@ -573,34 +573,44 @@ def _compile_smem(inst: ir3.SMEM | ir4.SMEM, ctx: _Ctx) -> UOp:
   # RDNA4 uses 'ioffset', RDNA3 uses 'offset' - use type(inst) to get correct field
   offset_field = type(inst).ioffset if hasattr(type(inst), 'ioffset') else type(inst).offset  # type: ignore[union-attr]
   offset = ctx.inst_field_signed(offset_field)  # signed immediate
-  # Dynamic soffset field - SGPR for additional offset (NULL=124 reads as 0)
-  soffset = ctx.inst_field(type(inst).soffset)
-  addr = _u64(ctx.rsgpr_dyn(sbase), ctx.rsgpr_dyn(sbase + _c(1))) + offset.cast(dtypes.uint64) + ctx.rsgpr_dyn(soffset).cast(dtypes.uint64)
+  # Dynamic soffset field - SGPR for additional offset (NULL=124 reads as 0, CDNA soffset_en=0 means no soffset)
+  soffset_val = _c(0).cast(dtypes.uint64)
+  if not (isinstance(inst, irc.SMEM) and not inst.soffset_en):
+    soffset_val = ctx.inst_field(type(inst).soffset)
+    soffset_val = ctx.rsgpr_dyn(soffset_val).cast(dtypes.uint64)
+  addr = _u64(ctx.rsgpr_dyn(sbase), ctx.rsgpr_dyn(sbase + _c(1))) + offset.cast(dtypes.uint64) + soffset_val
   _SMEM_NDWORDS = {ir3.SMEMOp.S_LOAD_B32: 1, ir3.SMEMOp.S_LOAD_B64: 2, ir3.SMEMOp.S_LOAD_B128: 4,
     ir3.SMEMOp.S_LOAD_B256: 8, ir3.SMEMOp.S_LOAD_B512: 16, ir4.SMEMOp.S_LOAD_B32: 1, ir4.SMEMOp.S_LOAD_B64: 2,
-    ir4.SMEMOp.S_LOAD_B96: 3, ir4.SMEMOp.S_LOAD_B128: 4, ir4.SMEMOp.S_LOAD_B256: 8, ir4.SMEMOp.S_LOAD_B512: 16}
+    ir4.SMEMOp.S_LOAD_B96: 3, ir4.SMEMOp.S_LOAD_B128: 4, ir4.SMEMOp.S_LOAD_B256: 8, ir4.SMEMOp.S_LOAD_B512: 16,
+    irc.SMEMOp.S_LOAD_DWORD: 1, irc.SMEMOp.S_LOAD_DWORDX2: 2, irc.SMEMOp.S_LOAD_DWORDX4: 4,
+    irc.SMEMOp.S_LOAD_DWORDX8: 8, irc.SMEMOp.S_LOAD_DWORDX16: 16}
   ndwords = _SMEM_NDWORDS[inst.op]
   stores = [ctx.wsgpr_dyn(sdata_reg + _c(i), ctx.vmem.index((addr + UOp.const(dtypes.uint64, i * 4) >> UOp.const(dtypes.uint64, 2)).cast(dtypes.int)))
             for i in range(ndwords)]
   return UOp.sink(*stores, *ctx.inc_pc())
 
-def _compile_sop(inst: ir3.SOP1 | ir3.SOP2 | ir3.SOPC | ir3.SOPK | ir4.SOP1 | ir4.SOP2 | ir4.SOPC | ir4.SOPK, ctx: _Ctx) -> UOp:
+def _compile_sop(inst: ir3.SOP1|ir3.SOP2|ir3.SOPC|ir3.SOPK|ir4.SOP1|ir4.SOP2|ir4.SOPC|ir4.SOPK|irc.SOP1|irc.SOP2|irc.SOPC|irc.SOPK, ctx: _Ctx) -> UOp:
   bits = inst.canonical_op_bits
   literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None  # type: ignore[union-attr]
 
-  if isinstance(inst, (ir3.SOPK, ir4.SOPK)):
+  if isinstance(inst, (ir3.SOPK, ir4.SOPK, irc.SOPK)):
     sdst_off = ctx.inst_field(type(inst).sdst)
     simm16 = ctx.inst_field(type(inst).simm16)
     # Sign-extend simm16
     simm16_sext = simm16.cast(dtypes.int16).cast(dtypes.int32)
-    srcs = {'S0': ctx.rsgpr_dyn(sdst_off), 'SIMM16': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
+    # CDNA pcode uses S0 for the immediate in MOVK/MULK/ADDK/CMOVK (where RDNA uses SIMM16),
+    # but S0 = register for CMPK/SETREG. S1 is always the immediate for CDNA CMPK ops.
+    op_name = inst.op.name if hasattr(inst.op, 'name') else ''
+    s0_is_imm = isinstance(inst, irc.SOPK) and 'CMPK' not in op_name and 'SETREG' not in op_name
+    s0_val = simm16_sext if s0_is_imm else ctx.rsgpr_dyn(sdst_off)
+    srcs = {'S0': s0_val, 'SIMM16': simm16_sext, 'S1': simm16_sext, 'D0': ctx.rsgpr_dyn(sdst_off)}
     dst_off, dst_size = sdst_off, 1
-  elif isinstance(inst, (ir3.SOP1, ir4.SOP1)):
+  elif isinstance(inst, (ir3.SOP1, ir4.SOP1, irc.SOP1)):
     sdst_off = ctx.inst_field(type(inst).sdst)
     ssrc0_off = ctx.inst_field(type(inst).ssrc0)
     srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal)}
     dst_off, dst_size = sdst_off, bits['d'] // 32
-  elif isinstance(inst, (ir3.SOP2, ir4.SOP2)):
+  elif isinstance(inst, (ir3.SOP2, ir4.SOP2, irc.SOP2)):
     sdst_off = ctx.inst_field(type(inst).sdst)
     ssrc0_off = ctx.inst_field(type(inst).ssrc0)
     ssrc1_off = ctx.inst_field(type(inst).ssrc1)
@@ -608,7 +618,7 @@ def _compile_sop(inst: ir3.SOP1 | ir3.SOP2 | ir3.SOPC | ir3.SOPK | ir4.SOP1 | ir
             'S1': ctx.rsrc_dyn(ssrc1_off, None, bits['s1'], literal)}
     if literal is not None: srcs['SIMM32'] = literal
     dst_off, dst_size = sdst_off, bits['d'] // 32
-  elif isinstance(inst, (ir3.SOPC, ir4.SOPC)):
+  elif isinstance(inst, (ir3.SOPC, ir4.SOPC, irc.SOPC)):
     ssrc0_off = ctx.inst_field(type(inst).ssrc0)
     ssrc1_off = ctx.inst_field(type(inst).ssrc1)
     srcs = {'S0': ctx.rsrc_dyn(ssrc0_off, None, bits['s0'], literal),
@@ -619,7 +629,7 @@ def _compile_sop(inst: ir3.SOP1 | ir3.SOP2 | ir3.SOPC | ir3.SOPK | ir4.SOP1 | ir
 
   return ctx.compile_sop_pcode(inst.op, srcs, dst_off, dst_size)
 
-def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VOP1_SDST | ir4.VOP2, ctx: _Ctx) -> UOp:
+def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VOP1_SDST | ir4.VOP2 | irc.VOP1 | irc.VOP2, ctx: _Ctx) -> UOp:
   op_name = _op_name(inst)
   if op_name in ('V_READFIRSTLANE_B32_E32', 'V_PERMLANE64_B32_E32'): return ctx.compile_lane_pcode(inst.op, inst)
   lane, exec_mask, bits = ctx.range(), ctx.rsgpr_dyn(_c(EXEC_LO.offset)), inst.canonical_op_bits
@@ -628,7 +638,7 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VO
   write_hi_half = bits['d'] == 16 and (vdst_reg >= _c(128))
   if isinstance(write_hi_half, UOp): vdst_reg = write_hi_half.where(vdst_reg - _c(128), vdst_reg)
   elif write_hi_half: vdst_reg -= 128
-  if isinstance(inst, (ir3.VOP1, ir4.VOP1)):
+  if isinstance(inst, (ir3.VOP1, ir4.VOP1, irc.VOP1)):
     # Handle VOP1 hi-half source operand (src0 >= v[128] for 16-bit ops)
     src0_off = ctx.inst_field(type(inst).src0)
     s0 = ctx.rsrc_dyn(src0_off, lane, bits['s0'], literal)
@@ -654,12 +664,13 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VO
       s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
     srcs = {'S0': s0, 'S1': s1, 'D0': d0}
     if inst.op in (ir3.VOP2Op.V_FMAAK_F32_E32, ir3.VOP2Op.V_FMAMK_F32_E32, ir3.VOP2Op.V_FMAAK_F16_E32,
-                   ir3.VOP2Op.V_FMAMK_F16_E32):
+                   ir3.VOP2Op.V_FMAMK_F16_E32, irc.VOP2Op.V_FMAAK_F32_E32, irc.VOP2Op.V_FMAMK_F32_E32):
       assert literal is not None
       srcs['SIMM32'] = literal
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, opsel_dst_hi=write_hi_half)
 
-def _compile_vopc(inst: ir3.VOPC | ir3.VOP3 | ir4.VOPC | ir4.VOP3, ctx: _Ctx, opsel: int = 0, abs_bits: int = 0, neg_bits: int = 0) -> UOp:
+def _compile_vopc(inst: ir3.VOPC|ir3.VOP3|ir4.VOPC|ir4.VOP3|irc.VOPC|irc.VOP3, ctx: _Ctx,
+                  opsel: int = 0, abs_bits: int = 0, neg_bits: int = 0) -> UOp:
   exec_mask, op_name, bits = ctx.rsgpr_dyn(_c(EXEC_LO.offset)), _op_name(inst), inst.canonical_op_bits
   is_cmpx, is_vopc = 'CMPX' in op_name, hasattr(inst, 'vsrc1')  # is_vopc: e32 vs e64
 
@@ -707,7 +718,7 @@ def _compile_vopc(inst: ir3.VOPC | ir3.VOP3 | ir4.VOPC | ir4.VOP3, ctx: _Ctx, op
     stores = [ctx.wsgpr_dyn(dst_off, new_result)] if not is_vopc else [ctx.wsgpr_dyn(_c(VCC_LO.offset), new_result)]
   return UOp.sink(*stores, *ctx.inc_pc())
 
-def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3, ctx: _Ctx) -> UOp:
+def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   bits = inst.canonical_op_bits
   opsel, op_name = getattr(inst, 'opsel', 0) or 0, _op_name(inst)
@@ -741,13 +752,13 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3, ctx: _Ctx) -> UOp:
   src1 = _apply_src_mods(src1, 1, abs_bits, neg_bits, bits['s1'])
   src2 = _apply_src_mods(src2, 2, abs_bits, neg_bits, bits['s2'])
   srcs = {'S0': src0, 'S1': src1, 'S2': src2}
-  if inst.op in (ir3.VOP3Op.V_CNDMASK_B32_E64, ir3.VOP3Op.V_CNDMASK_B16) and src2 is not None: srcs['VCC'] = src2
+  if inst.op in (ir3.VOP3Op.V_CNDMASK_B32_E64, ir3.VOP3Op.V_CNDMASK_B16, irc.VOP3Op.V_CNDMASK_B32_E64) and src2 is not None: srcs['VCC'] = src2
   # FMAC instructions need D0 (accumulator) from destination register
   if 'FMAC' in op_name: srcs['D0'] = ctx.rvgpr_dyn(vdst_reg, lane)
   opsel_dst_hi = bool(opsel & 0b1000) and bits['d'] == 16
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, opsel_dst_hi=opsel_dst_hi, clmp=getattr(inst, 'clmp', 0))
 
-def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
+def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UOp:
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   bits, pcode, ops = inst.canonical_op_bits, get_pcode(inst.op), inst.canonical_operands
 
@@ -806,7 +817,7 @@ def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD, ctx: _Ctx) -> UOp:
   else:
     return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, sdst_reg=inst.sdst.offset)
 
-def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
+def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   op_name = _op_name(inst)
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   vdst_reg = ctx.inst_field(type(inst).vdst)
@@ -839,14 +850,15 @@ def _compile_wmma(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
     stores = [ctx.wvgpr_dyn(vdst_reg + _c(i // 32), UOp.const(dtypes.int, i % 32), mat_d[i].bitcast(dtypes.uint32), exec_mask) for i in range(256)]
   return UOp.sink(*stores, *ctx.inc_pc())
 
-def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
+def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P | irc.VOP3P, ctx: _Ctx) -> UOp:
   op_name = _op_name(inst)
   if 'WMMA' in op_name and ('16X16X16_F16' in op_name or '16X16X16_BF16' in op_name): return _compile_wmma(inst, ctx)
 
   lane = ctx.range()
   exec_mask = ctx.rsgpr_dyn(_c(EXEC_LO.offset))
   vdst_reg = ctx.inst_field(type(inst).vdst)
-  do_cast = any(x in op_name for x in ('F16', 'F32', 'BF16')) and 'IU' not in op_name
+  is_pk_f32 = 'PK' in op_name and 'F32' in op_name and 'MOV' not in op_name  # CDNA packed F32 ops
+  do_cast = any(x in op_name for x in ('F16', 'F32', 'BF16')) and 'IU' not in op_name and not is_pk_f32
   src0 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src0), lane, 16, do_cast=do_cast)
   src1 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src1), lane, 16, do_cast=do_cast)
   src2 = ctx.rsrc_dyn(ctx.inst_field(type(inst).src2), lane, 16, do_cast=do_cast)
@@ -854,7 +866,30 @@ def _compile_vop3p(inst: ir3.VOP3P | ir4.VOP3P, ctx: _Ctx) -> UOp:
   opsel_hi2 = getattr(inst, 'opsel_hi2', 1) if getattr(inst, 'opsel_hi2', 1) is not None else 1
   neg, neg_hi = getattr(inst, 'neg', 0) or 0, getattr(inst, 'neg_hi', 0) or 0
 
-  if 'FMA_MIX' in op_name:
+  if is_pk_f32:
+    # CDNA packed F32: read 32-bit sources, build 64-bit packed values using opsel.
+    # For VGPRs: opsel selects between v[reg] (0) and v[reg+1] (1) for each half.
+    # For SGPR pairs (off < 128): s[N] = lo float32, s[N+1] = hi float32.
+    # For inline constants (128 <= off < 256): broadcast same value to both halves.
+    src_offs = [ctx.inst_field(type(inst).src0), ctx.inst_field(type(inst).src1), ctx.inst_field(type(inst).src2)]
+    def build_pk_f32(src_lo: UOp, src_off: UOp, opsel_lo: int, opsel_hi_bit: int, neg_lo: int, neg_hi_bit: int) -> UOp:
+      is_vgpr = src_off >= _c(256)
+      vgpr_lo = ctx.rvgpr_dyn(src_off - _c(256), lane) if lane is not None else _c(0)
+      vgpr_hi = ctx.rvgpr_dyn(src_off - _c(256) + _c(1), lane) if lane is not None else _c(0)
+      # For SGPR pairs, opsel selects between s[N] (0) and s[N+1] (1); inline constants always broadcast.
+      is_sgpr_pair = src_off < _c(128)
+      sgpr_hi = ctx.rsgpr_dyn(src_off + _c(1), is_sgpr_pair)
+      scalar_lo_sel = src_lo if not opsel_lo else is_sgpr_pair.where(sgpr_hi, src_lo)
+      scalar_hi_sel = src_lo if not opsel_hi_bit else is_sgpr_pair.where(sgpr_hi, src_lo)
+      lo = is_vgpr.where(vgpr_hi if opsel_lo else vgpr_lo, scalar_lo_sel)
+      hi = is_vgpr.where(vgpr_hi if opsel_hi_bit else vgpr_lo, scalar_hi_sel)
+      if neg_lo: lo = lo ^ UOp.const(dtypes.uint32, 0x80000000)
+      if neg_hi_bit: hi = hi ^ UOp.const(dtypes.uint32, 0x80000000)
+      return _u64(lo, hi)
+    srcs = {'S0': build_pk_f32(src0, src_offs[0], opsel & 1, opsel_hi & 1, neg & 1, neg_hi & 1),
+            'S1': build_pk_f32(src1, src_offs[1], opsel & 2, opsel_hi & 2, neg & 2, neg_hi & 2),
+            'S2': build_pk_f32(src2, src_offs[2], opsel & 4, 1 if opsel_hi2 else 0, neg & 4, neg_hi & 4)}
+  elif 'FMA_MIX' in op_name:
     combined_opsel_hi = (opsel_hi & 0x3) | ((opsel_hi2 & 0x1) << 2)
     # For FMA_MIX: neg_hi is ABS (not neg!), neg is actual negation
     def apply_abs(v, bit, opsel_hi_bit, opsel_bit):
@@ -924,13 +959,18 @@ def _compile_vopd(inst: ir3.VOPD | ir4.VOPD, ctx: _Ctx) -> UOp:
       if dest.startswith('D0'): all_stores.append(ctx.wvgpr_dyn(vdst_reg, lane, _val_to_u32(val), exec_mask, after=srcy1))
   return UOp.sink(UOp.group(*all_stores).end(lane), *ctx.inc_pc())
 
-def _compile_mem_op(inst: ir3.DS | ir3.FLAT | ir3.GLOBAL | ir3.SCRATCH | ir4.DS | ir4.VFLAT | ir4.VGLOBAL | ir4.VSCRATCH, ctx: _Ctx) -> UOp:
+def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLAT|ir4.VGLOBAL|ir4.VSCRATCH
+                    |irc.DS|irc.FLAT|irc.GLOBAL|irc.SCRATCH, ctx: _Ctx) -> UOp:
   """Unified memory operation compiler for DS, FLAT, GLOBAL, SCRATCH."""
   exec_mask, op_name = ctx.rsgpr_dyn(_c(EXEC_LO.offset)), _op_name(inst)
   pcode = get_pcode(inst.op)
+  # CDNA pcode uses CalcGlobalAddr/CalcDsAddr to compute address from raw components, but make_addr already handles this.
+  # Strip the addr computation line and use pre-computed ADDR directly (rename 'addr' -> 'ADDR' in remaining pcode).
+  if isinstance(inst, (irc.GLOBAL, irc.FLAT, irc.SCRATCH, irc.DS)) and 'Calc' in pcode and 'Addr' in pcode:
+    pcode = re.sub(r'addr\s*=\s*Calc\w+Addr\([^)]*\)\s*;?\n?', '', pcode).replace('MEM[addr', 'MEM[ADDR')
 
-  is_lds = isinstance(inst, (ir3.DS, ir4.DS))
-  is_scratch = isinstance(inst, (ir3.SCRATCH, ir4.VSCRATCH))
+  is_lds = isinstance(inst, (ir3.DS, ir4.DS, irc.DS))
+  is_scratch = isinstance(inst, (ir3.SCRATCH, ir4.VSCRATCH, irc.SCRATCH))
   mem = ctx.lds if is_lds else ctx.scratch if is_scratch else ctx.vmem
   addr_shift = UOp.const(dtypes.uint32 if is_lds else dtypes.uint64, 2)
 
@@ -1038,7 +1078,7 @@ def _compile_mem_op(inst: ir3.DS | ir3.FLAT | ir3.GLOBAL | ir3.SCRATCH | ir4.DS
     if 'STORE' in op_name and data_bits_mem >= 64:
       vdata = vdata | (ctx.rvgpr_dyn(vdata_reg + _c(1), lane).cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
     srcs = {'ADDR': addr, 'VDATA': vdata, '_vmem': mem, '_active': active,
-            'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base}
+            'laneId': lane, 'v_addr': vaddr_base, 's_saddr': saddr_base, 'SADDR': saddr_base, 'OFFSET': offset}
     for i in range(data_bits_mem // 32):
       srcs[f'VDATA{i}'] = ctx.rvgpr_dyn(vdata_reg + _c(i), lane) if 'STORE' in op_name else UOp.const(dtypes.uint32, 0)
     return srcs
@@ -1075,7 +1115,7 @@ def _compile_mem_op(inst: ir3.DS | ir3.FLAT | ir3.GLOBAL | ir3.SCRATCH | ir4.DS
       return UOp.sink(*ended, *ctx.inc_pc())
 
   # Standard path: single lane range
-  writes_return_data = '_RTN' in op_name or (is_lds and op_name.startswith('DS_LOAD')) or bool(is_atomic and glc)
+  writes_return_data = '_RTN' in op_name or (is_lds and (op_name.startswith('DS_LOAD') or op_name.startswith('DS_READ'))) or bool(is_atomic and glc)
   lane = ctx.range()
   active = _lane_active(exec_mask, lane)
   pcode_vars, assigns = parse_pcode(pcode, make_srcs(lane))
@@ -1099,6 +1139,11 @@ _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
   ir4.VOP1: _compile_vop12, ir4.VOP1_SDST: _compile_vop12, ir4.VOP2: _compile_vop12, ir4.VOPC: _compile_vopc, ir4.VOP3: _compile_vop3,
   ir4.VOP3_SDST: _compile_vop3, ir4.VOP3SD: _compile_vop3sd, ir4.VOP3P: _compile_vop3p, ir4.VOPD: _compile_vopd,
   ir4.DS: _compile_mem_op, ir4.VFLAT: _compile_mem_op, ir4.VGLOBAL: _compile_mem_op, ir4.VSCRATCH: _compile_mem_op,
+  # CDNA instruction classes
+  irc.SOPP: _compile_sopp, irc.SMEM: _compile_smem, irc.SOP1: _compile_sop, irc.SOP2: _compile_sop, irc.SOPC: _compile_sop, irc.SOPK: _compile_sop,
+  irc.VOP1: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3,
+  irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p,
+  irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op,
 }
 
 # ═══════════════════════════════════════════════════════════════════════════════
@@ -1116,7 +1161,7 @@ def _get_runner(inst_bytes: bytes, arch: str = "rdna3"):
 
   # Check if instruction matches any cached canonical pattern
   for base, mask, size, runner in _canonical_runner_cache:
-    if inst_size == size and (inst_int & mask) == base: return runner, False
+    if inst_size == size and (inst_int & mask) == base: return runner
 
   # Look up handler by type, falling back to base classes for _LIT variants
   handler = _INST_HANDLERS.get(type(inst))
@@ -1136,30 +1181,17 @@ def _get_runner(inst_bytes: bytes, arch: str = "rdna3"):
   with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES=""):
     runner = get_runner('CPU', sink)
   _canonical_runner_cache.append((base, mask, size, runner))
-  return runner, True
+  return runner
 
-@functools.cache
-def decode_program(data: bytes, arch: str = "rdna3") -> dict[int, tuple[str, Callable, list[int], Any]]:
-  """Decode program to {pc: (name, fxn, globals, runner)}."""
-  result: dict[int, tuple[str, Callable, list[int], Any]] = {}
-  i = 0
-  while i < len(data):
-    inst = decode_inst(data[i:], arch)
-    if hasattr(inst, 'op') and inst.op in (ir3.SOPPOp.S_CODE_END, ir4.SOPPOp.S_CODE_END): break
-    try:
-      runner, is_new = _get_runner(bytes(data[i:i + inst.size() + 4]), arch)
-      if DEBUG >= 3:
-        try: inst_str = repr(inst)
-        except Exception: inst_str = f"<{type(inst).__name__} at PC={i}>"
-        msg = f"[emu] PC={i}: {inst_str}"
-        print(colored(msg, 'green') if is_new else msg)
-      result[i] = (runner.p.function_name, runner._prg.fxn, runner.p.globals, runner)
-    except Exception as e:
-      try: inst_str = repr(inst)
-      except Exception: inst_str = f"<{type(inst).__name__}>"
-      raise RuntimeError(f"[emu] Failed to compile PC={i} {inst_str}: {type(e).__name__}: {e}") from e
-    i += inst.size()
-  return result
+def _decode_at(pc: int, arch: str):
+  """Decode and compile instruction at absolute address pc. Returns CompiledRunner."""
+  inst_bytes = bytes((ctypes.c_char * 16).from_address(pc).raw)
+  inst = decode_inst(inst_bytes, arch)
+  try: return _get_runner(bytes(inst_bytes[:inst.size() + 4]), arch)
+  except Exception as e:
+    try: inst_str = repr(inst)
+    except Exception: inst_str = f"<{type(inst).__name__}>"
+    raise RuntimeError(f"[emu] Failed to compile {inst_str}: {type(e).__name__}: {e}") from e
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # WAVE STATE
@@ -1206,10 +1238,9 @@ class WaveState:
 # ═══════════════════════════════════════════════════════════════════════════════
 
 def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int, rsrc2: int = 0x19c,
-            scratch_size: int = 0, arch: str = "rdna3") -> int:
+            scratch_size: int = 0, arch: str = "rdna3", user_data: list[int]|None = None) -> int:
   """Execute AMD assembly program. scratch_size is private_segment_fixed_size from kernel descriptor (per-lane)."""
-  program_raw = decode_program(bytes((ctypes.c_char * lib_sz).from_address(lib).raw), arch)
-  program = {lib + offset: val for offset, val in program_raw.items()}  # Remap to actual addresses
+  program: dict[int, tuple[Callable, list[int]]] = {}  # lazily populated: pc -> (fxn, globals) extracted from runner
   lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT) * 512
   total_threads = lx * ly * lz
 
@@ -1226,8 +1257,12 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
           for wave_start in range(0, total_threads, WAVE_SIZE):
             n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState(min(WAVE_SIZE, total_threads - wave_start))
             st.pc = lib  # Set PC to code base address
-            st._write_sgpr(0, args_ptr & MASK32)
-            st._write_sgpr(1, (args_ptr >> 32) & MASK32)
+            # Initialize user SGPRs: hardware loads COMPUTE_USER_DATA registers directly into s[0:N]
+            if user_data:
+              for i, val in enumerate(user_data): st._write_sgpr(i, val)
+            else:
+              st._write_sgpr(0, args_ptr & MASK32)
+              st._write_sgpr(1, (args_ptr >> 32) & MASK32)
 
             # Workgroup IDs in SGPRs after user SGPRs
             sgpr_idx = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
@@ -1255,13 +1290,16 @@ def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int,
                       ctypes.c_uint64(vmem_buf._buf.va_addr), ctypes.c_uint64(lds_buf._buf.va_addr),
                       ctypes.c_uint64(scratch_buf._buf.va_addr if scratch_buf else 0)]
             for inst_count in range(1_000_000):
-              if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF or pc not in program: break
-              name, fxn, globals_list, _ = program[pc]
-              assert fxn is not None, f"[emu] No fxn for {name} at PC={pc}"
-              assert 4 not in globals_list or scratch_buf, f"SCRATCH instruction {name} but scratch_size=0"
-              if DEBUG >= 6:
-                inst = decode_inst(bytes((ctypes.c_char * 12).from_address(pc).raw), arch)
-                print(f"[emu] exec PC={pc:X}: {inst!r}")
+              if (pc := st.pc) == 0xFFFFFFFFFFFFFFFF: break
+              if pc not in program:
+                prev_len = len(_canonical_runner_cache)
+                runner = _decode_at(pc, arch)
+                program[pc] = (runner._prg.fxn, runner.p.globals)
+                if DEBUG >= 3:
+                  inst = decode_inst(bytes((ctypes.c_char * 16).from_address(pc).raw), arch)
+                  msg = f"[emu] PC={pc - lib}: {inst!r}"
+                  print(colored(msg, 'green') if len(_canonical_runner_cache) > prev_len else msg)
+              fxn, globals_list = program[pc]
               fxn(*[c_bufs[g] for g in globals_list])
             else: raise RuntimeError("exceeded 1M instructions, likely infinite loop")
   return 0

test/mockgpu/amd/pcode.py

@@ -40,7 +40,10 @@ def _bitreverse(v: UOp, bits: int) -> UOp:
 
 def _extract_bits(val: UOp, hi: int, lo: int) -> UOp:
   dt = dtypes.uint64 if val.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
-  return ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)
+  result = ((val >> _const(dt, lo)) if lo > 0 else val) & _const(val.dtype, (1 << (hi - lo + 1)) - 1)
+  # Downcast to uint32 when extracting <=32 bits from a 64-bit value, so .f32 bitcast works correctly
+  if dt == dtypes.uint64 and (hi - lo + 1) <= 32: result = result.cast(dtypes.uint32)
+  return result
 
 def _set_bit(old, pos, val):
   mask = _u32(1) << pos
@@ -554,7 +557,9 @@ class Parser:
       self.eat('LBRACKET')
       self.eat_val('laneId', 'IDENT')
       self.eat('RBRACKET')
-      result = (base >> _to_u32(self.vars['laneId'])) & _u32(1)
+      lane = self.vars['laneId']
+      shift = lane.cast(base.dtype) if base.dtype != dtypes.uint32 else _to_u32(lane)
+      result = (base >> shift) & _const(base.dtype, 1)
       if self.try_eat('DOT'):
         dt_name = self.eat('IDENT').val
         return result.cast(DTYPES.get(dt_name, dtypes.uint32))
@@ -806,6 +811,12 @@ def _subst_loop_var(line: str, loop_var: str, val: int) -> str:
 
 def _set_bits(old: UOp, val: UOp, width: int, offset: int) -> UOp:
   """Set bits [offset:offset+width) in old to val, masking and shifting appropriately."""
+  is64 = old.dtype in (dtypes.uint64, dtypes.int64) or offset + width > 32
+  if is64:
+    old = old.cast(dtypes.uint64) if old.dtype != dtypes.uint64 else old
+    mask = _u64(((1 << width) - 1) << offset)
+    v = (val.cast(dtypes.uint64) if val.dtype != dtypes.uint64 else val) & _u64((1 << width) - 1)
+    return (old & (mask ^ _u64(0xFFFFFFFFFFFFFFFF))) | (v << _u64(offset))
   mask = _u32(((1 << width) - 1) << offset)
   v = (val.cast(dtypes.uint32) if val.dtype != dtypes.uint32 else val) & _u32((1 << width) - 1)
   return (old & (mask ^ _u32(0xFFFFFFFF))) | (v << _u32(offset))

test/mockgpu/helpers.py

@@ -21,10 +21,11 @@ class PythonRemu:
   rsrc2: int = 0x19c  # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
   scratch_size: int = 0  # private_segment_fixed_size from kernel descriptor
   arch: str = "rdna3"  # Architecture: rdna3 or rdna4
+  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
-    return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size, self.arch)
+    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():
   # Use Python emulator only if PYTHON_REMU=1

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/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/renderer/amd/sqtt.py

@@ -670,7 +670,8 @@ if __name__ == "__main__":
     sys.exit(1)
   with open(sys.argv[1], "rb") as f:
     data = pickle.load(f)
+  prg_names = {e.tag: e.name for e in data if type(e).__name__ == "ProfileProgramEvent" and e.tag is not None}
   sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
   for i, event in enumerate(sqtt_events):
-    print(f"\n=== event {i} ===")
+    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
     print_packets(decode(event.blob))

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/am/amdev.py

@@ -193,7 +193,7 @@ class AMDev(PCIDevImplBase):
     if DEBUG >= 2: print(f"am {self.devfmt}: boot done")
 
   def init_sw(self, smi_dev=False):
-    self.smi_dev, self.is_err_state = smi_dev, False
+    self.smi_dev, self.is_err_state, self.has_aql_queue = smi_dev, False, False
 
     # Memory manager & firmware
     self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
@@ -226,7 +226,7 @@ class AMDev(PCIDevImplBase):
     self.reg("regSCRATCH_REG6").write(self.is_err_state) # set finalized state.
 
   def recover(self) -> bool:
-    if self.is_hive() or not self.is_err_state: return False # TODO: support mi300
+    if (self.has_aql_queue and self.is_hive()) or not self.is_err_state: return False # TODO: support aql queue recovery on hive
     if DEBUG >= 2: print(f"am {self.devfmt}: Start recovery")
     self.ih.interrupt_handler()
     self.gfx.reset_mec()

tinygrad/runtime/support/am/ip.py

@@ -291,6 +291,7 @@ class AM_GFX(AM_IP):
     self._enable_mec()
 
   def setup_ring(self, ring_addr:int, ring_size:int, rptr_addr:int, wptr_addr:int, eop_addr:int, eop_size:int, idx:int, aql:bool) -> tuple[int, int]:
+    self.adev.has_aql_queue |= aql
     pipe, queue, doorbell = idx // 4, idx % 4, am.AMDGPU_NAVI10_DOORBELL_MEC_RING0
     self._grbm_select(me=1, pipe=pipe, queue=queue, inst=0)
     restore_queue = aql and self.xccs > 1 and self.adev.partial_boot and (self.adev.regCP_HQD_ACTIVE.read(inst=0) & 1)

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

@@ -614,14 +614,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 +640,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 +771,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 +1216,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 +1331,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")
@@ -3879,10 +3904,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:
     """

tinygrad/viz/js/index.js

@@ -184,7 +184,7 @@ const WAVE_COLORS = {VALU:"#ffffc0", SALU:"#cef263", LOAD:"#ffc0c0", STORE:"#4fa
 const waveColor = (op) => {
   const cat = op.includes("VALU") || op === "VINTERP" ? "VALU" : op.includes("SALU") ? "SALU" : op.includes("VMEM") ? "VMEM"
             : op.includes("LOAD") || op === "SMEM" ? "LOAD" : op.includes("STORE") ? "STORE" : op;
-  ret = WAVE_COLORS[cat] ?? "#ffffff";
+  let ret = WAVE_COLORS[cat] ?? "#ffffff";
   if (op.includes("OTHER_") || op.includes("_ALT")) { ret = darkenHex(ret, 75) }
   if (op.includes("LDS_")) { ret = darkenHex(ret, 25) }
   return ret