sched docs

fixes: #15698

.github/actions/setup-tinygrad/action.yml

@@ -225,14 +225,12 @@ runs:
       if: inputs.amd == 'true' && runner.os == 'Linux'
       shell: bash
       run: |
-        cargo build --release --manifest-path ./extra/remu/Cargo.toml
-        sudo ln -sf ${{ github.workspace }}/extra/remu/target/release/libremu.so /usr/local/lib/libremu.so
         sudo tee --append /etc/ld.so.conf.d/rocm.conf <<'EOF'
           /opt/rocm/lib
           /opt/rocm/lib64
         EOF
         sudo ldconfig
-    - name: Setup AMD comgr+remu (macOS)
+    - name: Setup AMD comgr (macOS)
       if: inputs.amd == 'true' && runner.os == 'macOS'
       shell: bash
       run: |
@@ -240,7 +238,6 @@ runs:
         curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
           jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
           sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
-        cargo build --release --manifest-path ./extra/remu/Cargo.toml
 
     # **** gpuocelot ****
 

.github/workflows/benchmark.yml

@@ -71,10 +71,6 @@ jobs:
         uv venv /tmp/tinygrad_pytest_ci
         source /tmp/tinygrad_pytest_ci/bin/activate
         uv pip install .[testing]
-    - name: setup other stuff
-      run: |
-        mkdir -p extra/remu/target/release/
-        ln -s ~/tinygrad/extra/remu/target/release/libremu.so extra/remu/target/release/libremu.so
     - name: setup staging db
       run: |
         echo "CACHEDB=/tmp/pytest-db-ci.db" >> $GITHUB_ENV
@@ -538,6 +534,8 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
+    - name: Test GPU crash recovery
+      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
     - name: Train MNIST
       run: time PYTHONPATH=. DEV=AMD TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
     - name: Run 10 CIFAR training steps
@@ -709,6 +707,8 @@ jobs:
       run: time DEBUG=3 DEV=AMD AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Test driver warm start time
       run: time DEBUG=3 DEV=AMD python3 test/test_tiny.py TestTiny.test_plus
+    - name: Test GPU crash recovery
+      run: DEV=AMD python3 -m pytest -rA test/external/external_test_gpu_crash.py
     # Fails on 9070
     # - name: Test tensor cores
     #   run: |

.github/workflows/test.yml

@@ -508,6 +508,8 @@ jobs:
         run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
       - name: Test 1B LLM (llama q4)
         run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
+      - name: Test 1B LLM (qwen3.5)
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster
       - name: Test 1B LLM (qwen)
         # NOTE: qwen is dumb and only knows about female chickens
         run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
@@ -642,7 +644,6 @@ jobs:
     timeout-minutes: 20
     env:
       DEV: AMD
-      PYTHON_REMU: 1
       MOCKGPU: 1
     steps:
       - name: Checkout Code

docs/abstractions3.py

@@ -38,7 +38,7 @@ optim.schedule_step()   # this will step the optimizer without running realize
 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph
 
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 schedule: List[ExecItem] = Tensor.schedule(l1, l2)
 
 print(f"The schedule contains {len(schedule)} items.")

docs/abstractions4.py

@@ -0,0 +1,253 @@
+# tinygrad allows you to write kernels at many different abstractions levels.
+# This is for RDNA3, but if you don't have one you can run with the emulator
+# PYTHONPATH="." MOCKGPU=1 DEV=AMD
+
+from tinygrad import Tensor, Context, GlobalCounters, UOp, Device
+from tinygrad.helpers import DEBUG, getenv
+from tinygrad.uop.ops import AxisType, KernelInfo, Ops
+from tinygrad.dtype import AddrSpace, dtypes
+from tinygrad.runtime.autogen.amd.rdna3.ins import *
+
+def eval_harness(name, tensor, fxn, check=None):
+  print(f"***** {name}")
+  GlobalCounters.reset()
+  with Context(DEBUG=max(DEBUG.value, 2)): out = fxn(tensor).item()
+  assert check is None or abs(out - check) < abs(check) * 1e-3, f"out was wrong {out}, expected {check}, off by {out/check}x"
+  print(f"computed in {GlobalCounters.time_sum_s*1000:.2f} ms, {(a.nbytes()/1e9)/GlobalCounters.time_sum_s:.2f} GB/s")
+  return out
+
+SZ = 256*1024 if getenv("MOCKGPU") else 1024*1024*1024
+
+def example_2_hip(a:Tensor, correct):
+  GLOBALS = 1024
+  THREADS = 256
+  def hip_reduce_sum(out:UOp, buf:UOp) -> UOp:
+    assert SZ % (GLOBALS * THREADS) == 0
+    CHUNK = SZ // (GLOBALS * THREADS)
+    # NOTE: tinygrad doesn't populate HIP hidden kernargs, so blockDim.x/gridDim.x read as 0.
+    # We hardcode block/grid sizes as constexpr to avoid any dependency on those builtins.
+    code = f"""
+    #include <hip/hip_runtime.h>
+    constexpr unsigned int BLOCK = {THREADS};
+    constexpr unsigned int CHUNK = {CHUNK};
+    extern "C" __global__ void hip_reduce_sum_kernel(float* __restrict__ block_sums, const float* __restrict__ x) {{
+      __shared__ float sdata[BLOCK];
+
+      unsigned int tid = threadIdx.x;
+      unsigned int gid = blockIdx.x * BLOCK + tid;
+
+      // Each thread sums CHUNK consecutive elements from its own region
+      float sum = 0.0f;
+      const float* base = x + gid * CHUNK;
+      #pragma unroll 16
+      for (unsigned int k = 0; k < CHUNK; k++) {{
+        sum += base[k];
+      }}
+
+      sdata[tid] = sum;
+      __syncthreads();
+
+      // Block reduction in shared memory
+      for (unsigned int s = BLOCK / 2; s > 0; s >>= 1) {{
+        if (tid < s) {{
+          sdata[tid] += sdata[tid + s];
+        }}
+        __syncthreads();
+      }}
+
+      // One partial sum per block
+      if (tid == 0) {{
+        block_sums[blockIdx.x] = sdata[0];
+      }}
+    }}"""
+
+    # TODO: remove the need for the compiler here, you should just be able to remove Ops.BINARY
+    from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
+    lib = HIPCCCompiler(Device[Device.DEFAULT].renderer.target.arch, []).compile_cached(code)
+    # the sink specifies the GLOBAL and LOCAL sizes, along with the input buffers and name
+    sink = UOp.sink(UOp.special(GLOBALS, 'gidx0'), UOp.special(THREADS, 'lidx0'), out, buf,
+                    arg=KernelInfo(name="hip_reduce_sum_kernel"))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
+  eval_harness("HIP kernel", a, lambda x: Tensor.empty(GLOBALS).custom_kernel(x, fxn=hip_reduce_sum)[0].sum(), check=correct)
+
+def example_3_custom_uop(a:Tensor, correct):
+  # This GPU has 32 CUs, keep them all busy
+  CU_COUNT = 32
+  def custom_sum(out:UOp, buf:UOp) -> UOp:
+    LCLS = 256
+    buf = buf.reshape(CU_COUNT, -1, LCLS)
+
+    glbl = UOp.range(CU_COUNT, 0, AxisType.GLOBAL)
+    lane = UOp.range(LCLS, 1, AxisType.LOCAL)
+
+    # accumulate the globals into a per lane accumulator
+    reduce_loop = UOp.range(buf.shape[1], 2, AxisType.REDUCE)
+    acc = UOp.placeholder((1,), dtypes.float, slot=6, addrspace=AddrSpace.REG)
+    acc = acc.after(acc.store(0))
+    acc = acc.after(acc[0].store(acc.after(reduce_loop)[0] + buf[glbl, reduce_loop, lane]).end(reduce_loop))
+
+    # store all the per lane accumulators to LOCAL
+    local_accs = UOp.placeholder((LCLS,), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
+    local_accs = local_accs.after(local_accs[lane].store(acc[0]).barrier())
+
+    # accumulate LOCALs into a single per CU accumulator
+    late_reduce_loop = UOp.range(LCLS, 3, AxisType.REDUCE)
+    acc2 = UOp.placeholder((1,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
+    acc2 = acc2.after(acc2.store(0))
+    acc2 = acc2.after(acc2[0].store(acc2.after(late_reduce_loop)[0] + local_accs[late_reduce_loop]).end(late_reduce_loop))[0]
+
+    # store (NOTE: since the address doesn't depend on the warp, this will be automatically gated)
+    return out[glbl].store(acc2).end(lane, glbl).sink(arg=KernelInfo(opts_to_apply=()))
+
+  eval_harness("custom UOp kernel", a, lambda x: Tensor.empty(CU_COUNT).custom_kernel(x, fxn=custom_sum)[0].sum(), check=correct)
+
+def example_5_custom_assembly(a:Tensor, correct):
+  # Kernel class copied from amd_asm_matmul
+  class Kernel:
+    def __init__(self, arch='gfx1100'): self.instructions, self.labels, self.pos, self.arch = [], {}, 0, arch
+    def label(self, name): self.labels[name] = self.pos
+    def emit(self, inst, target=None):
+      self.instructions.append(inst)
+      inst._target, inst._pos = target, self.pos
+      self.pos += inst.size()
+      return inst
+    def waitcnt(self, lgkm=None, vm=None):
+      # Wait for memory operations. lgkm=N waits until N lgkm ops remain, vm=N waits until N vmem ops remain.
+      vmcnt, lgkmcnt, expcnt = vm if vm is not None else 63, lgkm if lgkm is not None else 63, 7
+      waitcnt = (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
+      self.emit(s_waitcnt(simm16=waitcnt))
+    def finalize(self, sink:UOp) -> UOp:
+      for inst in self.instructions:
+        if inst._target is None: continue
+        offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
+        if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
+        inst.simm16 = offset_dwords
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=Device.DEFAULT),
+                                   UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
+
+  CU_COUNT = 32
+  LANES = 64
+  def asm_sum(out:UOp, buf:UOp) -> UOp:
+    V_LANE_ID = 0             # lane_id set on startup
+    S_WORKGROUP_X = 2         # workgroup_id_x
+    S_LOOP_CTR = 3
+    k = Kernel()
+    # mul lane id by 16 for offsets (4 for float, 4 for b128)
+    k.emit(v_mul_lo_u32(v[0], v[V_LANE_ID], 16))
+    k.emit(v_add_nc_u32_e32(v[1], 4096, v[0]))
+    k.emit(v_add_nc_u32_e32(v[2], 4096, v[1]))
+    k.emit(v_add_nc_u32_e32(v[3], 4096, v[2]))
+    # load both addresses
+    k.emit(s_load_b128(sdata=s[4:7], sbase=s[0:1], offset=0x0, soffset=NULL))
+    k.waitcnt(lgkm=0)
+    # offset buffer pointer by workgroup_id_x * chunk_size_bytes
+    k.emit(s_mul_i32(s[S_LOOP_CTR], s[S_WORKGROUP_X], buf.numel()*4//CU_COUNT))
+    k.emit(s_add_u32(s[6], s[6], s[S_LOOP_CTR]))
+    k.emit(s_addc_u32(s[7], s[7], 0))
+    # zero the accumulators
+    k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[4], vdsty=v[5], srcx0=0, srcy0=0))
+    k.emit(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[6], vdsty=v[7], srcx0=0, srcy0=0))
+
+    def emit_loads(base_vreg, reg_len):
+      assert reg_len%4 == 0
+      k.emit(s_clause(simm16=(reg_len//4)-1))
+      for i in range(reg_len//4):
+        offset = i*LANES*16
+        assert offset < 16384
+        k.emit(global_load_b128(vdst=v[base_vreg+i*4:base_vreg+i*4+3], addr=v[offset//4096], saddr=s[6:7], offset=offset%4096))
+      k.emit(s_add_u32(s[6], s[6], reg_len * LANES * 4))
+      k.emit(s_addc_u32(s[7], s[7], 0))
+
+    def tree_reduce_to_4567(base_vreg, reg_len):
+      assert reg_len%4 == 0
+      reg_len //= 4
+      while reg_len > 1:
+        half = reg_len // 2
+        for j in range(half):
+          a, b = base_vreg + j*4, base_vreg + (j+half)*4
+          # v[a+0](bank0) += v[b+2](bank2), v[a+1](bank1) += v[b+3](bank3) — src0 and src1 on different banks
+          k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[a], vdsty=v[a+1], srcx0=v[a], vsrcx1=v[b+2], srcy0=v[a+1], vsrcy1=v[b+3]))
+          # v[a+2](bank2) += v[b+0](bank0), v[a+3](bank3) += v[b+1](bank1) — src0 and src1 on different banks
+          k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[a+2], vdsty=v[a+3], srcx0=v[a+2], vsrcx1=v[b], srcy0=v[a+3], vsrcy1=v[b+1]))
+        reg_len = half
+      k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[4], vdsty=v[5], srcx0=v[4], vsrcx1=v[base_vreg], srcy0=v[5], vsrcy1=v[base_vreg+1]))
+      k.emit(VOPD(VOPDOp.V_DUAL_ADD_F32, VOPDOp.V_DUAL_ADD_F32, vdstx=v[6], vdsty=v[7], srcx0=v[6], vsrcx1=v[base_vreg+2], srcy0=v[7], vsrcy1=v[base_vreg+3]))
+
+    BASE_REG = 8
+    LOAD_UNROLL = 64
+    INNER_UNROLL = 2
+
+    assert buf.numel() % (CU_COUNT*LANES*LOAD_UNROLL*INNER_UNROLL) == 0
+    total_batches = buf.numel()//(CU_COUNT*LANES*LOAD_UNROLL*INNER_UNROLL)
+    k.emit(s_mov_b32(s[S_LOOP_CTR], total_batches-1))
+
+    k.label('LOOP')
+    for _ in range(INNER_UNROLL):
+      emit_loads(BASE_REG, reg_len=LOAD_UNROLL)
+      k.waitcnt(vm=0)
+      tree_reduce_to_4567(BASE_REG, reg_len=LOAD_UNROLL)
+    k.emit(s_sub_u32(s[S_LOOP_CTR], s[S_LOOP_CTR], 1))
+    k.emit(s_cbranch_scc0(), target='LOOP')
+
+    # add into v[4]
+    k.emit(v_add_f32_e32(v[4], v[4], v[5]))
+    k.emit(v_add_f32_e32(v[6], v[6], v[7]))
+    k.emit(v_add_f32_e32(v[4], v[4], v[6]))
+
+    # warp shuffle into v[4] on lane 0 using DPP row_shl within each 16-lane row
+    for shift in [1, 2, 4, 8]:
+      k.emit(v_add_f32_e32(v[4], DPP, v[4], vsrc0=v[4], dpp=0x100 | shift, row_mask=0xf, bank_mask=0xf, bc=1))
+    # combine rows: get lane 16's value to lane 0 via permlanex16
+    k.emit(v_permlanex16_b32(v[5], v[4], 0, 0))
+    k.emit(v_add_f32_e32(v[4], v[4], v[5]))
+
+    # atomic store (only on lane 0)
+    k.emit(s_mov_b32(EXEC_LO, 1))
+    k.emit(v_mov_b32_e32(v[0], 0))
+    k.emit(global_atomic_add_f32(addr=v[0], saddr=s[4:5], data=v[4]))
+
+    k.emit(s_sendmsg(simm16=3))  # DEALLOC_VGPRS
+    k.emit(s_endpgm())
+    return k.finalize(UOp.sink(UOp.special(CU_COUNT, 'gidx0'), UOp.special(LANES, 'lidx0'), out, buf, arg=KernelInfo(name="asm_reduce")))
+
+  out = Tensor.zeros(1,).contiguous().realize()
+  eval_harness("RDNA3 assembly kernel", a, lambda x: out.custom_kernel(x, fxn=asm_sum)[0], check=correct)
+
+if __name__ == "__main__":
+  examples = [int(x) for x in getenv("EXAMPLES", "1,2,3,4,5").split(",")]
+
+  correct = None
+  # First define a Tensor and realize it. We will focus on a 1GB sum kernel on RDNA3
+  a = (Tensor.randn(SZ) if getenv("RAND") else Tensor.ones(SZ)).contiguous().realize()
+
+  if 1 in examples:
+    # *****
+    # This is the high level tinygrad way.
+    # Note that this is split into multiple kernels for speed.
+    correct = eval_harness("basic kernel", a, lambda x: x.sum())
+
+  if 2 in examples:
+    # *****
+    # You can import kernels from CUDA/HIP/Metal.
+    # ChatGPT is great at writing these Kernel
+    example_2_hip(a, correct)
+
+  if 3 in examples:
+    # *****
+    # Now we get to the lower abstraction layers of tinygrad.
+    # You can write a kernel in UOps, and it's 2.5x faster than normal.
+    example_3_custom_uop(a, correct)
+
+  if 4 in examples:
+    # *****
+    # You can also BEAM search stock tinygrad for a faster kernel.
+    # This does even better than all the kernels to date in this simple case.
+    with Context(BEAM=2):
+      eval_harness("BEAMed kernel", a, lambda x: x.sum(), check=correct)
+
+  if 5 in examples:
+    # *****
+    # If you really want to go crazy with speed, you can code in assembly.
+    # There's not too much to gain here over BEAM, but it's a few percent faster.
+    example_5_custom_assembly(a, correct)

docs/developer/developer.md

@@ -17,9 +17,9 @@ The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not al
 
 ## Scheduling
 
-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/schedule/__init__.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
 
-::: tinygrad.engine.schedule.ExecItem
+::: tinygrad.schedule.ExecItem
 
 ## Lowering
 

docs/env_vars.md

@@ -34,9 +34,8 @@ DEBUG               | [1-7]      | enable debugging output (operations, timings,
 DEV                 | [AMD, NV, ...] | enable a specific backend, see [below](#dev-variable)
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
-IMAGE               | [1-2]      | enable 2d specific optimizations
+IMAGE               | [1]        | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -65,8 +64,8 @@ Variable | Value | Description
 ---|---|---
 DEBUG               | >= 1       | Enables debugging and lists devices being used
 DEBUG               | >= 2       | Provides performance metrics for operations, including timing, memory usage, bandwidth for each kernel execution
-DEBUG               | >= 3       | Outputs buffers used for each kernel (shape, dtype and strides) and the applied optimizations at a kernel level
+DEBUG               | >= 3       | Outputs the applied optimizations at a kernel level
 DEBUG               | >= 4       | Outputs the generated kernel code
-DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps (AST)
+DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps
 DEBUG               | >= 6       | Displays the intermediate representation of the computation UOps in a linearized manner, detailing the operation sequence
 DEBUG               | >= 7       | Outputs the assembly code generated for the target hardware

docs/tinygpu.md

@@ -57,3 +57,5 @@ export PATH="$HOME/.local/bin:$PATH"
 ```bash
 DEV={AMD|NV} python3 tinygrad/apps/llm.py
 ```
+
+**Note:** Use `JITBEAM=2` to search for faster kernels (one-time search cost, results cached).

examples/llama.py

@@ -445,7 +445,7 @@ After you are done speaking, output [EOS]. You are not Chad.
   print(f"using LLaMA{LLAMA_SUFFIX}-{args.size} model")
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
   llama = LLaMa.build(MODEL_PATH, TOKENIZER_PATH, model_gen=args.gen, model_size=args.size, quantize=args.quantize, device=device)
-  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(llama.model))
+  param_bytes = sum(x.nbytes() for x in get_parameters(llama.model))
 
   outputted = pre_prompt if chatbot else args.prompt
   start_pos, toks = 0, [llama.tokenizer.bos_id()] + llama.tokenizer.encode(outputted)

examples/llama3.py

@@ -324,7 +324,7 @@ if __name__ == "__main__":
 
   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(args.shard)) if args.shard > 1 else Device.DEFAULT
   model = build_transformer(args.model, model_size=args.size, quantize=args.quantize, device=device)
-  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(model))
+  param_bytes = sum(x.nbytes() for x in get_parameters(model))
 
   if not args.no_api and not args.benchmark:
     from bottle import Bottle, request, response, HTTPResponse, abort, static_file

examples/llm.c/export.py

@@ -5,7 +5,7 @@ from tinygrad import Device, nn, Tensor, dtypes
 from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import DEV, dedup, flatten, getenv, GlobalCounters, to_function_name
 from tinygrad.engine.realize import get_kernel
-from tinygrad.engine.memory import memory_planner
+from tinygrad.schedule.memory import memory_planner
 from tinygrad.uop.ops import Ops
 
 DEV.value = "CPU"

examples/mlperf/model_train.py

@@ -246,7 +246,7 @@ def train_resnet():
 
       if i == BENCHMARK:
         assert not math.isnan(loss)
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * steps_in_train_epoch * epochs / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
@@ -593,7 +593,7 @@ def train_retinanet():
 
       if i == BENCHMARK:
         assert not math.isnan(loss)
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * steps_in_train_epoch * EPOCHS / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
@@ -868,7 +868,7 @@ def train_unet3d():
         i += 1
 
         if i == BENCHMARK:
-          median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+          median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
           estimated_total_minutes = int(median_step_time * SAMPLES_PER_EPOCH * NUM_EPOCHS / 60)
           print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
           if (TRAIN_BEAM or EVAL_BEAM) and epoch == start_epoch: break
@@ -1167,7 +1167,7 @@ def train_bert():
       i += 1
 
       if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+        median_step_time = sorted(step_times)[BENCHMARK // 2]  # in seconds
         estimated_total_minutes = int(median_step_time * train_steps / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
@@ -1544,7 +1544,7 @@ def train_llama3():
 
       mem_gb = GlobalCounters.mem_used / 1e9
       gflops = GlobalCounters.global_ops / 1e9 / dev_time
-      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 2.3e15)) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * (4.6e15 if FP8 else 2.3e15))) * 100
       tqdm.write(
           f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
           f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
@@ -1577,8 +1577,9 @@ def train_llama3():
         safe_save(get_state_dict(scheduler), fn)
 
       if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
-        estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 60)
+        median_step_time = sorted(step_times)[BENCHMARK // 2]
+        estimated_steps = 200_000 // GBS if getenv("LLAMA3_SIZE", "8B") == "8B" else MAX_STEPS
+        estimated_total_minutes = int(median_step_time * estimated_steps / 60)
         print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
         print(f"epoch global_ops: {GlobalCounters.global_ops:_}, "
               f"epoch global_mem: {GlobalCounters.global_mem:_}")

examples/mlperf/models/flat_llama.py

@@ -1,4 +1,4 @@
-import math, os
+import math, os, functools
 if __name__ == "__main__":
   os.environ["DEFAULT_FLOAT"] = "bfloat16"
   os.environ["OPTIM_DTYPE"] = "bfloat16"
@@ -18,37 +18,63 @@ from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
 
 FP8 = getenv("FP8", 0)
-WQKV = getenv("WQKV", 0)
 
 FP8_DTYPE = dtypes.fp8e4m3
 FP8_GRAD_DTYPE = dtypes.fp8e5m2
 FP8_MAX = 448.0
 
+# per-device abs max without allreduce (matches TE delayed scaling behavior)
+@functools.cache
+def _local_abs_max_fxn(x_p, device):
+  x = Tensor(x_p, device=device)
+  inner = Tensor(x.uop.src[0]) if x.uop.op is Ops.MULTI else x
+  return (inner.abs().max(),)
+
+def _local_abs_max(x:Tensor) -> Tensor:
+  param = x.as_param(0)
+  fxn = _local_abs_max_fxn(param.uop, x.device)
+  return Tensor(fxn[0].uop.call(x.uop).gettuple(0))
+
 def quantize_fp8(x:Tensor, amax_state:Tensor|None=None):
-  if amax_state is not None:
-    scale = FP8_MAX / (amax_state + 1e-8)
-    amax_state.assign(x.abs().max().detach())
-  else:
-    scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
+  new_amax = (_local_abs_max(x) if isinstance(x.device, tuple) else x.abs().max()).detach()
+  scale = FP8_MAX / ((amax_state if amax_state is not None else new_amax) + 1e-8)
   x_scaled = x * scale
   x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
-  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
+  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal(), new_amax
 
-def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, amax_w:Tensor|None=None) -> Tensor:
-  if not fp8: return x @ w.T
-  from tinygrad.helpers import ASM_GEMM
-  x_fp8, x_scale = quantize_fp8(x, amax_state=amax_x)
-  w_fp8, w_scale = quantize_fp8(w, amax_state=amax_w)
+def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, amax_w:Tensor|None=None) -> tuple[Tensor,...]:
+  if not fp8:
+    if getenv("ASM_GEMM"):
+      from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
+      if can_use_asm_gemm(x, w.T): return (asm_gemm(x, w.T),)
+    return (x @ w.T,)
+  x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
+  w_fp8, w_scale, w_new_amax = quantize_fp8(w, amax_state=amax_w)
   combined_scale = x_scale * w_scale
-  if ASM_GEMM:
+  if getenv("ASM_GEMM"):
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w_fp8.T): return asm_gemm(x_fp8, w_fp8.T, combined_scale=combined_scale)
-  return x_fp8.dot(w_fp8.T, dtype=dtypes.float) * combined_scale
+    if can_use_asm_gemm(x_fp8, w_fp8.T): return asm_gemm(x_fp8, w_fp8.T, combined_scale=combined_scale), x_new_amax, w_new_amax, x_fp8, w_fp8
+  return x_fp8.dot(w_fp8.T, dtype=dtypes.float) * combined_scale, x_new_amax, w_new_amax, x_fp8, w_fp8
 
-def rmsnorm(x_in:Tensor, eps:float):
+def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
   x = x_in.float()
-  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return x.cast(x_in.dtype)
+  rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return (x * rrms).cast(x_in.dtype), rrms
+
+@functools.cache
+def _rmsnorm_fwd_fxn(x_in_p, eps, device):
+  return _rmsnorm_fwd(Tensor(x_in_p, device=device), eps)
+
+def _rmsnorm_bwd(grad:UOp, call:UOp) -> tuple:
+  x_normed = Tensor(call.gettuple(0)).float()
+  do_float = Tensor(grad).float()
+  d_x = Tensor(call.gettuple(1)) * (do_float - x_normed * (do_float * x_normed).mean(-1, keepdim=True))
+  return (d_x.cast(call.src[1].dtype).uop,)
+
+def rmsnorm(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+  fxn = _rmsnorm_fwd_fxn(x_in.as_param(0).uop, eps, x_in.device)
+  call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, grad_fxn=_rmsnorm_bwd)
+  return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
 
 class FlatTransformer:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
@@ -63,12 +89,7 @@ class FlatTransformer:
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
-    if WQKV:
-      self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    else:
-      self.wq = self.lin_per_layer(dim, self.n_heads * self.head_dim)
-      self.wk = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
-      self.wv = self.lin_per_layer(dim, self.n_kv_heads * self.head_dim)
+    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
     self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
 
     # FeedForward
@@ -89,8 +110,7 @@ class FlatTransformer:
 
     if FP8:
       def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
-      names = (["xqkv", "wqkv"] if WQKV else ["xq", "wq", "xk", "wk", "xv", "wv"]) + \
-              ["xo", "wo", "x1", "w1", "x2", "w2", "x3", "w3"]
+      names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x2", "w2", "x3", "w3"]
       # _fp8_amax[name][layer_idx] = scalar amax tensor
       self._fp8_amax = {name: [_amax() for _ in range(n_layers)] for name in names}
       self._fp8_amax["xout"] = [_amax()]
@@ -100,54 +120,73 @@ class FlatTransformer:
     if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features)
     return Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
 
-  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wo:Tensor, wqkv:Tensor|None=None,
-                wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None,
-                amax_xqkv=None, amax_wqkv=None, amax_xq=None, amax_wq=None, amax_xk=None, amax_wk=None,
-                amax_xv=None, amax_wv=None, amax_xo=None, amax_wo=None):
-    x = rmsnorm(x, self.norm_eps) * attention_norm
+  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
+                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None):
     bsz, seqlen, _ = x.shape
+    new_amaxs, saves = [], []
 
-    if wqkv is not None:
-      xqkv = matmul(x, wqkv, amax_x=amax_xqkv, amax_w=amax_wqkv)
-      xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-      xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-    else:
-      assert wq is not None and wk is not None and wv is not None
-      xq = matmul(x, wq, amax_x=amax_xq, amax_w=amax_wq).reshape(bsz, seqlen, self.n_heads, self.head_dim)
-      xk = matmul(x, wk, amax_x=amax_xk, amax_w=amax_wk).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-      xv = matmul(x, wv, amax_x=amax_xv, amax_w=amax_wv).reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * attention_norm
+
+    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, amax_w=amax_wqkv)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [xqkv])
+    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
+    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
+    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
+    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
 
     xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
     if FP8: xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16)
     xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-    attn = attn.reshape(bsz, seqlen, -1)
-    return matmul(attn, wo, amax_x=amax_xo, amax_w=amax_wo)
+    if getenv("HK_FLASH_ATTENTION"):
+      from extra.thunder.amd.fa import flash_attention
+      attn, *save = flash_attention(xq, xk, xv, is_causal=True)
+      saves.extend(save)
+    else:
+      attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True)
+    attn = attn.transpose(1, 2).reshape(bsz, seqlen, -1)
+
+    out, *ret = matmul(attn, wo, amax_x=amax_xo, amax_w=amax_wo)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [out])
+    return (out, *new_amaxs, *saves)
 
   def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
                    amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
-    x = rmsnorm(x, self.norm_eps) * ffn_norm
-    x_w1 = matmul(x, w1, amax_x=amax_x1, amax_w=amax_w1).silu()
-    x_w3 = matmul(x.contiguous_backward(), w3, amax_x=amax_x3, amax_w=amax_w3)
-    return matmul(x_w1 * x_w3, w2, amax_x=amax_x2, amax_w=amax_w2)
+    new_amaxs, saves = [], []
+
+    x, rrms = rmsnorm(x, self.norm_eps)
+    saves.extend([x, rrms])
+    x = x * ffn_norm
+
+    x_w1, *ret = matmul(x, w1, amax_x=amax_x1, amax_w=amax_w1)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [x_w1])
+    x_w3, *ret = matmul(x.contiguous_backward(), w3, amax_x=amax_x3, amax_w=amax_w3)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [x_w3])
+    out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, amax_w=amax_w2)
+    new_amaxs.extend(ret[:2])
+    saves.extend(ret[2:] + [out])
+    return (out, *new_amaxs, *saves)
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wo:Tensor,
+                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
                 ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor,
-                wqkv:Tensor|None=None, wq:Tensor|None=None, wk:Tensor|None=None, wv:Tensor|None=None,
-                amax_xqkv=None, amax_wqkv=None, amax_xq=None, amax_wq=None, amax_xk=None, amax_wk=None,
-                amax_xv=None, amax_wv=None, amax_xo=None, amax_wo=None,
+                amax_xqkv=None, amax_wqkv=None, amax_xo=None, amax_wo=None,
                 amax_x1=None, amax_w1=None, amax_x2=None, amax_w2=None, amax_x3=None, amax_w3=None):
-    h = x + self.attention(x, freqs_cis, attention_norm, wo, wqkv=wqkv, wq=wq, wk=wk, wv=wv,
-                           amax_xqkv=amax_xqkv, amax_wqkv=amax_wqkv, amax_xq=amax_xq, amax_wq=amax_wq,
-                           amax_xk=amax_xk, amax_wk=amax_wk, amax_xv=amax_xv, amax_wv=amax_wv,
-                           amax_xo=amax_xo, amax_wo=amax_wo)
-    return h + self.feed_forward(h, ffn_norm, w1, w2, w3,
-                                 amax_x1=amax_x1, amax_w1=amax_w1, amax_x2=amax_x2, amax_w2=amax_w2,
-                                 amax_x3=amax_x3, amax_w3=amax_w3)
+    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
+                                     amax_xqkv=amax_xqkv, amax_wqkv=amax_wqkv, amax_xo=amax_xo, amax_wo=amax_wo)
+    attn_amaxs, attn_saves = attn_ret[:4], attn_ret[4:]
+    h = x + attn
+    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w1, w2, w3,
+                                      amax_x1=amax_x1, amax_w1=amax_w1, amax_x2=amax_x2, amax_w2=amax_w2, amax_x3=amax_x3, amax_w3=amax_w3)
+    ffn_amaxs, ffn_saves = ffn_ret[:6], ffn_ret[6:]
+    h = h + ffn
+    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -155,12 +194,7 @@ class FlatTransformer:
       for v in get_parameters(self): v.shard_(device, axis=None)
     else:
       # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      if WQKV:
-        self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      else:
-        self.wq.shard_(device, axis=1).realize()            # (n_layers, n_heads*head_dim, dim) shard out
-        self.wk.shard_(device, axis=1).realize()            # (n_layers, n_kv_heads*head_dim, dim) shard out
-        self.wv.shard_(device, axis=1).realize()            # (n_layers, n_kv_heads*head_dim, dim) shard out
+      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
       self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
       self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
       self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
@@ -171,29 +205,32 @@ class FlatTransformer:
       self.tok_embeddings.weight.shard_(device, axis=0).realize()
       self.output.shard_(device, axis=1).realize()
       self.freqs_cis.shard_(device, axis=None).realize()
+      if FP8:
+        for name in self._fp8_amax:
+          for i in range(len(self._fp8_amax[name])):
+            self._fp8_amax[name][i] = self._fp8_amax[name][i].to(device).contiguous().requires_grad_(False)
 
   def __call__(self, tokens:Tensor):
     h = self.tok_embeddings(tokens)
     freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
     a = self._fp8_amax if FP8 else None
     for i in range(self.n_layers):
-      if WQKV:
-        attn_kwargs = {"wqkv": self.wqkv[i]}
-        amax_attn = {"amax_xqkv": a["xqkv"][i], "amax_wqkv": a["wqkv"][i]} if a else {}
-      else:
-        attn_kwargs = {"wq": self.wq[i], "wk": self.wk[i], "wv": self.wv[i]}
-        amax_attn = {"amax_xq": a["xq"][i], "amax_wq": a["wq"][i],
-                     "amax_xk": a["xk"][i], "amax_wk": a["wk"][i],
-                     "amax_xv": a["xv"][i], "amax_wv": a["wv"][i]} if a else {}
-      amax_layer = {"amax_xo": a["xo"][i], "amax_wo": a["wo"][i],
+      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_wqkv": a["wqkv"][i],
+                    "amax_xo": a["xo"][i], "amax_wo": a["wo"][i],
                     "amax_x1": a["x1"][i], "amax_w1": a["w1"][i],
                     "amax_x2": a["x2"][i], "amax_w2": a["w2"][i],
                     "amax_x3": a["x3"][i], "amax_w3": a["w3"][i]} if a else {}
-      h = self.run_layer(h, freqs_cis,
-                         self.attention_norm[i], self.wo[i],
-                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i],
-                         **attn_kwargs, **amax_attn, **amax_layer)
-    logits = (self.norm(h).contiguous().contiguous_backward() @ self.output[0].T).contiguous_backward()
+      h, *ret = self.run_layer(h, freqs_cis,
+                               self.attention_norm[i], self.wqkv[i], self.wo[i],
+                               self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i],
+                               **amax_layer)
+      if a:
+        amaxs = ret[:10]
+        amax_names = ["xqkv", "wqkv", "xo", "wo", "x1", "w1", "x3", "w3", "x2", "w2"]
+        for name, new_val in zip(amax_names, amaxs):
+          a[name][i].assign(new_val)
+
+    logits = matmul(self.norm(h).contiguous().contiguous_backward(), self.output[0], fp8=False)[0].contiguous_backward()
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

examples/mlperf/models/llama.py

@@ -1,80 +0,0 @@
-from tinygrad import Tensor, nn
-from tinygrad.helpers import getenv
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-class Attention:
-  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    if getenv("WQKV"):
-      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
-    else:
-      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-
-    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
-    if getenv("WQKV"):
-      xqkv = self.wqkv(x)
-      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-    else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-
-    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
-    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
-    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    bsz, seqlen, _, _ = xq.shape
-
-    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-
-    attn = attn.reshape(bsz, seqlen, -1)
-    return self.wo(attn)
-
-class FeedForward:
-  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.w1 = linear(dim, hidden_dim, bias=False)
-    self.w2 = linear(hidden_dim, dim, bias=False)
-    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
-
-  def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x)
-    return self.w2(w1 * w3)
-
-class TransformerBlock:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
-    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
-    self.feed_forward = FeedForward(dim, hidden_dim, linear)
-    self.attention_norm = nn.RMSNorm(dim, norm_eps)
-    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor):
-    h = x + self.attention(self.attention_norm(x), freqs_cis)
-    return h + self.feed_forward(self.ffn_norm(h))
-
-class Transformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = embedding(vocab_size, dim)
-    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for layer in self.layers: h = layer(h, freqs_cis)
-    logits = self.output(self.norm(h))
-    return logits

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh

@@ -2,7 +2,6 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -15,6 +14,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh

@@ -2,7 +2,6 @@
 
 export PYTHONPATH="."
 export DEV=${DEV:-AMD}
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -15,6 +14,7 @@ export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
 export MASTER_WEIGHTS=${MASTER_WEIGHTS:-1}
 export FP8=${FP8:-1}
+export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/profile.sh

@@ -2,4 +2,5 @@
 export BENCHMARK=5
 export EVAL_BS=0
 VIZ=${VIZ:--1} FULL_LAYERS=1 DEBUG=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_beam.sh
-extra/viz/cli.py --profile -s "${DEV:-AMD}"
+SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
+extra/viz/cli.py --profile -s "$SRC"

examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/run_and_time.sh

@@ -4,7 +4,6 @@ set -o pipefail  # Make pipeline fail if any command fails
 
 export PYTHONPATH="."
 export DEV=AMD
-export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
 export DEVICE_IN_FUNCTION_BUG=1
@@ -16,6 +15,7 @@ export ASM_GEMM=1
 export WQKV=1
 export MASTER_WEIGHTS=1
 export FP8=1
+export ALLREDUCE_CAST=1
 
 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=8 MP=1 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=4

examples/qwq.py

@@ -66,7 +66,7 @@ if __name__ == "__main__":
   model_path = Path(args.weights) if args.weights else download_weights(model_info["total_num_weights"])
   transformer = load_model(model_path, model_info["model_params"])
   tokenizer = AutoTokenizer.from_pretrained(model_info["tokenizer"])
-  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(transformer))
+  param_bytes = sum(x.nbytes() for x in get_parameters(transformer))
 
   outputted = args.prompt
   start_pos, toks = 0, tokenizer(outputted)["input_ids"]

extra/gemm/amd_copy_matmul.py

@@ -1,4 +1,4 @@
-from tinygrad import UOp, getenv
+from tinygrad import Device, UOp, getenv
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 
@@ -13,18 +13,23 @@ assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
 
 use_wmma = getenv("WMMA")
 if use_wmma:
+  is_rdna4 = Device[Device.DEFAULT].renderer.target.arch.startswith("gfx12")
+
   WAVES_M, WAVES_N = 2, 2
   LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-  UNROLL_M, UNROLL_N = 1, 1
 
   # wmma params
   WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
   WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
+  UNROLL_M, UNROLL_N = (WMMA_ACC, 1) if is_rdna4 else (1, 1)
 else:
   WAVES_M, WAVES_N = 4, 1
   LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
   UNROLL_M, UNROLL_N = 4, 4
 
+# total lanes must be the warp size
+assert LANES_PER_WAVE_M*LANES_PER_WAVE_N == WARP_SIZE
+
 # WARP_SIZE * total waves
 THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
 
@@ -71,7 +76,10 @@ def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
     acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
     a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
     b_frag = B_local.reshape(WAVES_N, TN, WMMA_N, BLOCK_K // WMMA_K, WMMA_K)[wave_n, tile_n, lane_n, k]
-
+    if is_rdna4:
+      # NOTE: since this is part of K, these 2 can be anywhere in the frags and long as a and b match
+      a_frag = a_frag.reshape(2, 8)[lane_m, :]
+      b_frag = b_frag.reshape(2, 8)[lane_m, :]
     wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
     acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
   else:

extra/gemm/amx.py

@@ -44,9 +44,9 @@ nc = np.random.randn(N, N).astype(np.float32)
 
 ns = nb.reshape(-1, 32).sum(axis=0)
 
-a = MallocAllocator.alloc(na.size * np.dtype(np.float32).itemsize)
-b = MallocAllocator.alloc(nb.size * np.dtype(np.float32).itemsize)
-c = MallocAllocator.alloc(nc.size * np.dtype(np.float32).itemsize)
+a = MallocAllocator.alloc(na.nbytes)
+b = MallocAllocator.alloc(nb.nbytes)
+c = MallocAllocator.alloc(nc.nbytes)
 
 MallocAllocator._copyin(b, flat_mv(nb.data))
 MallocAllocator._copyin(c, flat_mv(nc.data))

extra/gemm/cdna_asm_gemm.py

@@ -2704,7 +2704,7 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     a_t, b_t, g_t, s_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device), Tensor(scale, device=a.device)
     g_t = g_t[:a.shape[0]]
     # backward GEMMs in fp8 with scale applied inside kernel to prevent bf16 overflow
-    g_fp8, g_scale = quantize_fp8(g_t)
+    g_fp8, g_scale, _ = quantize_fp8(g_t)
     bw_scale = g_scale * s_t
     # dgrad: g_fp8 @ weight (asm_gemm computes a@b)
     grad_a = asm_gemm(g_fp8, b_t, combined_scale=bw_scale)
@@ -2716,8 +2716,11 @@ def custom_gemm_bw(gradient:UOp, kernel:UOp):
     assert all_same([gradient.device, a.device, b.device, out.device])
     a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
     g_t = g_t[:a.shape[0]]
-    grad_a = (g_t @ b_t.T).uop
-    grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+    if can_use_asm_gemm(g_t, b_t.T): grad_a = asm_gemm(g_t, b_t.T).uop
+    else: grad_a = (g_t @ b_t.T).uop
+    a_t_flat, g_t_flat = a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1), g_t.reshape(-1, g_t.shape[-1])
+    if can_use_asm_gemm(a_t_flat, g_t_flat): grad_b = asm_gemm(a_t_flat, g_t_flat).uop
+    else: grad_b = (a_t_flat @ g_t_flat).uop
     return (None, grad_a, grad_b)
 
 # ** main gemm function

extra/gemm/rdna4_asm_matmul.py

@@ -0,0 +1,245 @@
+# RDNA4 128x128 GEMM using WMMA — optimized DS scheduling
+import numpy as np
+from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import getenv, colored
+from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.engine.realize import Estimates
+from tinygrad.renderer.amd.dsl import s, v, VCC_LO, NULL, src, ttmp
+from tinygrad.runtime.autogen.amd.rdna4.ins import *
+
+BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 16
+TILES_M, TILES_N = 4, 4
+THREADS, ELEM = 128, 2
+LDS_A_ROW = BLOCK_K*ELEM   # 32
+LDS_B_ROW = BLOCK_N*ELEM  # 256
+LDS_A_SIZE = BLOCK_M * LDS_A_ROW  # 4096
+LDS_B_SIZE = BLOCK_K * LDS_B_ROW  # 4096
+LDS_SIZE = LDS_A_SIZE + LDS_B_SIZE  # 8192
+LDS_B_OFF = LDS_A_SIZE
+ACC, DA, DB, FA, FB, ET = 60, 188, 196, 204, 44, 10
+
+def build_kernel(N, arch='gfx1200'):
+  assert N % BLOCK_M == 0 and N >= 256
+  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  I, L, B = [], {}, []
+  def e(i): I.append(i); return i
+  def label(n): L[n] = sum(i.size() for i in I)
+  def br(i, t): B.append((len(I)-1, t))
+
+  e(s_load_b128(sdata=s[4:7], sbase=s[0:1], ioffset=0, soffset=NULL))
+  e(s_load_b64(sdata=s[8:9], sbase=s[0:1], ioffset=0x10, soffset=NULL))
+  e(s_wait_kmcnt(simm16=0))
+  e(s_mov_b32(s[10], ttmp[9])); e(s_and_b32(s[11], ttmp[7], 0xFFFF))
+  e(s_lshl_b32(s[10], s[10], 7)); e(s_lshl_b32(s[11], s[11], 7))
+  e(s_mov_b32(s[12], N)); e(s_lshl_b32(s[13], s[12], 1))
+  e(s_mul_i32(s[14], s[12], BLOCK_K*ELEM))
+  e(s_add_co_i32(s[17], s[12], -2*BLOCK_K))  # loop bound
+
+  e(v_and_b32_e32(v[1], 31, v[0])); e(v_lshrrev_b32_e32(v[2], 5, v[0]))
+  e(v_and_b32_e32(v[3], 1, v[2])); e(v_lshrrev_b32_e32(v[2], 1, v[2]))
+
+  e(v_lshlrev_b32_e32(v[4], 5, v[0]))
+  # B store: transposed layout for stride-32 reads. addr = LDS_B_OFF + (tid%8)*512 + (tid/8)*32
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[5], 9, v[48]))   # (tid%8)*512
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))  # (tid/8)*32
+  e(v_add_nc_u32_e32(v[5], v[5], v[48])); e(v_add_nc_u32_e32(v[5], LDS_B_OFF, v[5]))
+
+  e(v_add_nc_u32_e32(v[48], s[11], v[0]))
+  e(v_mul_lo_u32(v[6], v[48], N*ELEM)); e(v_mov_b32_e32(v[7], 0))
+  e(v_lshrrev_b32_e32(v[48], 3, v[0])); e(v_mul_lo_u32(v[8], v[48], N*ELEM))
+  e(v_and_b32_e32(v[48], 7, v[0])); e(v_lshlrev_b32_e32(v[48], 5, v[48]))
+  e(v_add_nc_u32_e32(v[8], v[8], v[48]))
+  e(s_mul_i32(s[15], s[10], ELEM)); e(v_add_nc_u32_e32(v[8], s[15], v[8]))
+  e(v_mov_b32_e32(v[9], 0))
+
+  # LDS read addrs with padded strides (eliminates bank conflicts)
+  # A: (lane%16)*LDS_A_ROW + (lane/16)*16 + wave_m*64*LDS_A_ROW
+  # B: (lane%16)*LDS_B_ROW + (lane/16)*16 + wave_n*64*ELEM + LDS_B_OFF
+  LLA, LLB = 40, 43
+  e(v_and_b32_e32(v[50], 15, v[1])); e(v_lshrrev_b32_e32(v[51], 4, v[1]))
+  e(v_lshlrev_b32_e32(v[LLA], 5, v[50]))        # (lane%16) * 32
+  e(v_lshlrev_b32_e32(v[51], 4, v[51]))        # (lane/16) * 16
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[51]))
+  e(v_lshlrev_b32_e32(v[52], 11, v[2]))         # wave_m * 2048
+  e(v_add_nc_u32_e32(v[LLA], v[LLA], v[52]))
+  # B read: transposed layout. addr = LDS_B_OFF + (lane%16)*32 + (lane/16)*16 + wave_n*2*512
+  # wave_n selects column panels: wave_n*2 panels (each panel=16 cols, wave_n covers 64 cols = 4 panels)
+  # But wave_n*2*512 = wave_n*1024. Hmm, wave_n covers cols [wave_n*64 : (wave_n+1)*64].
+  # Each panel = 16 cols = 512 bytes. wave_n*64/16 = wave_n*4 panels. Offset = wave_n*4*512 = wave_n*2048.
+  e(v_lshlrev_b32_e32(v[LLB], 5, v[50]))         # (lane%16) * 32 (stride 32!)
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[51]))     # + (lane/16)*16
+  e(v_lshlrev_b32_e32(v[52], 11, v[3]))           # wave_n * 2048
+  e(v_add_nc_u32_e32(v[LLB], v[LLB], v[52]))
+  e(v_add_nc_u32_e32(v[LLB], LDS_B_OFF, v[LLB]))
+
+  for i in range(0, 128, 2):
+    e(VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[ACC+i], vdsty=v[ACC+i+1], srcx0=0, srcy0=0))
+  e(s_mov_b32(s[16], 0))
+
+  if not NO_GLOBAL:
+    for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+    for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+    e(s_wait_loadcnt(simm16=0))
+  if not NO_DS:
+    for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+  if not NO_GLOBAL:
+    e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+    e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+
+  # =============================================================================
+  def emit_iter_body(load_set='AB'):
+    if not NO_DS:
+      e(s_wait_dscnt(simm16=0))
+      e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+    if not NO_GLOBAL:
+      if 'A' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DA+i*4:DA+i*4+3], vaddr=v[6:7], saddr=s[4:5], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[6], BLOCK_K*ELEM, v[6]))
+      if 'B' in load_set:
+        for i in range(2): e(global_load_b128(vdst=v[DB+i*4:DB+i*4+3], vaddr=v[8:9], saddr=s[6:7], ioffset=i*16))
+        e(v_add_nc_u32_e32(v[8], s[14], v[8]))
+    if not NO_DS:
+      # Issue 6 loads: A[0:3] + B[0] + B[1]. B[2:3] interleaved with WMMAs.
+      for tm in range(TILES_M):
+        aoff = tm * 16 * LDS_A_ROW
+        e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+      e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+      e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+      e(s_wait_dscnt(simm16=0))  # wait for 6 loads (no stall!)
+    if not NO_ALU:
+      # B[0] WMMAs — issue B[2] during compute
+      if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+0)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+      # B[1] WMMAs — issue B[3] during compute
+      if not NO_DS:
+        e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+1)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+      # B[2] WMMAs — B[2] loaded during B[0] WMMAs (~100 cycles ago)
+      if not NO_DS: e(s_wait_dscnt(simm16=1))  # B[2] done, B[3] may still be loading
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+2)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+      # B[3] WMMAs
+      if not NO_DS: e(s_wait_dscnt(simm16=0))
+      for tm in range(TILES_M):
+        ac = ACC + (tm*TILES_N+3)*8
+        e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+    if not NO_GLOBAL and not NO_DS: e(s_wait_loadcnt(simm16=0))
+    if not NO_DS:
+      for i in range(2): e(ds_store_b128(addr=v[4], data0=v[DA+i*4:DA+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+      for i in range(2): e(ds_store_b128(addr=v[5], data0=v[DB+i*4:DB+i*4+3], offset0=(i*16)&0xFF, offset1=(i*16)>>8))
+    e(s_add_co_i32(s[16], s[16], BLOCK_K))
+
+  label('LOOP')
+  emit_iter_body(load_set='A')
+  emit_iter_body(load_set='B')
+  e(s_cmp_lt_i32(s[16], s[17])); e(s_cbranch_scc1(simm16=0)); br(I[-1], 'LOOP')
+
+  emit_iter_body(load_set='AB')  # tail with prefetch
+
+  # Final iteration: no prefetch, no ds_store needed
+  if not NO_DS:
+    e(s_wait_dscnt(simm16=0))
+    e(s_barrier_signal(ssrc0=src[193])); e(s_barrier_wait(simm16=0xFFFF))
+  if not NO_DS:
+    for tm in range(TILES_M):
+      aoff = tm * 16 * LDS_A_ROW
+      e(ds_load_b128(vdst=v[FA+tm*4:FA+tm*4+3], addr=v[LLA], offset0=aoff&0xFF, offset1=aoff>>8))
+    e(ds_load_b128(vdst=v[FB:FB+3], addr=v[LLB], offset0=0, offset1=0))
+    e(ds_load_b128(vdst=v[FB+4:FB+7], addr=v[LLB], offset0=0, offset1=2))
+    e(s_wait_dscnt(simm16=0))
+  if not NO_ALU:
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+8:FB+11], addr=v[LLB], offset0=0, offset1=4))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+0)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB:FB+3], src2=v[ac:ac+7]))
+    if not NO_DS: e(ds_load_b128(vdst=v[FB+12:FB+15], addr=v[LLB], offset0=0, offset1=6))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+1)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+4:FB+7], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=1))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+2)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+8:FB+11], src2=v[ac:ac+7]))
+    if not NO_DS: e(s_wait_dscnt(simm16=0))
+    for tm in range(TILES_M):
+      ac = ACC + (tm*TILES_N+3)*8
+      e(v_wmma_f32_16x16x16_f16(vdst=v[ac:ac+7], src0=v[FA+tm*4:FA+tm*4+3], src1=v[FB+12:FB+15], src2=v[ac:ac+7]))
+
+  label('EPILOGUE')
+  e(v_and_b32_e32(v[ET], 15, v[1]))
+  e(v_lshrrev_b32_e32(v[ET+1], 4, v[1])); e(v_lshlrev_b32_e32(v[ET+1], 3, v[ET+1]))
+  e(v_lshlrev_b32_e32(v[ET+2], 6, v[2])); e(v_add_nc_u32_e32(v[ET+2], s[11], v[ET+2]))
+  e(v_lshlrev_b32_e32(v[ET+3], 6, v[3])); e(v_add_nc_u32_e32(v[ET+3], s[10], v[ET+3]))
+  e(v_add_nc_u32_e32(v[ET+3], v[ET+3], v[ET])); e(v_mov_b32_e32(v[ET+5], 0))
+
+  for tm in range(TILES_M):
+    for tn in range(TILES_N):
+      ac = ACC + (tm*TILES_N+tn)*8; r_off, c_off = tm*16, tn*16
+      e(v_add_nc_u32_e32(v[ET+6], r_off, v[ET+2])); e(v_add_nc_u32_e32(v[ET+6], v[ET+1], v[ET+6]))
+      e(v_mul_lo_u32(v[ET+4], v[ET+6], s[12])); e(v_add_nc_u32_e32(v[ET+4], v[ET+4], v[ET+3]))
+      if c_off: e(v_add_nc_u32_e32(v[ET+4], c_off, v[ET+4]))
+      e(v_lshlrev_b32_e32(v[ET+4], 1, v[ET+4]))
+      for elem in range(8):
+        e(v_cvt_f16_f32_e32(v[ET+7], v[ac+elem]))
+        e(global_store_b16(vaddr=v[ET+4:ET+5], vsrc=v[ET+7], saddr=s[8:9]))
+        if elem < 7: e(v_add_nc_u32_e32(v[ET+4], s[13], v[ET+4]))
+
+  e(s_wait_storecnt(simm16=0)); e(s_sendmsg(simm16=3)); e(s_endpgm())
+
+  for idx, target in B:
+    off = (L[target] - sum(i.size() for i in I[:idx+1])) // 4
+    assert -32768 <= off <= 32767; I[idx].simm16 = off
+  return I
+
+N = getenv("N", 4096)
+
+def test_matmul():
+  dev = Device[Device.DEFAULT]
+  arch = getattr(dev.renderer, 'arch', 'gfx1200')
+  print(f"Device arch: {arch}")
+  insts = build_kernel(N, arch)
+
+  rng = np.random.default_rng(42)
+  a = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  b = Tensor(rng.random((N, N), dtype=np.float32).astype(np.float16))
+  c = Tensor.empty(N, N, dtype=dtypes.half)
+  Tensor.realize(a, b, c)
+
+  grid, local = (N//BLOCK_N, N//BLOCK_M, 1), (THREADS, 1, 1)
+  print(f"Grid: {grid}, Local: {local}")
+
+  dname = Device.DEFAULT
+  def asm_kernel(A, B, C):
+    gidxs = [UOp.special(n, f"gidx{i}") for i,n in enumerate(grid)]
+    lidxs = [UOp.special(THREADS, "lidx0")]
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=max(LDS_SIZE, 65536//getenv("LIMIT_OCC",2)), addrspace=AddrSpace.LOCAL), (), 'lds')
+    sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
+                    arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+  c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
+  ei = c.schedule()[0].lower()
+
+  ets = []
+  with Context(DEBUG=2):
+    for _ in range(getenv("CNT", 5)): ets.append(ei.run(wait=True))
+  print(f"REAL TFLOPS {N*N*N*2 / min(ets) * 1e-12:.2f}")
+
+  if getenv("VERIFY", 1):
+    GlobalCounters.reset()
+    c_np = c.float().numpy()
+    a_np, b_np = a.float().numpy(), b.float().numpy()
+    ref = a_np @ b_np
+    err = np.sqrt(np.mean((c_np - ref)**2)) / np.sqrt(np.mean(ref**2))
+    print(f"relative RMSE {err:.6f}")
+    if err != err or err > 0.05: raise RuntimeError(f"matmul is wrong! RMSE={err}")
+
+if __name__ == "__main__":
+  test_matmul()

extra/mlx_driver/loopback.py

@@ -6,7 +6,8 @@ from tinygrad.device import Device, BufferSpec
 from tinygrad.runtime.support.system import PCIDevice
 from tinygrad.runtime.support.memory import AddrSpace
 from tinygrad.runtime.ops_amd import AMDComputeQueue
-from extra.mlx_driver.mlxdev import MLXDev, MLXQP, to_be
+from tinygrad.helpers import to_be32, to_be64
+from extra.mlx_driver.mlxdev import MLXDev, MLXQP
 
 BUF_SIZE = 0x1000
 MLX_PCI = getenv("MLX_PCI", "0000:41:00.0")
@@ -49,7 +50,7 @@ rq_wqe = qp.qp_buf.view((qp.rq_head & rq_mask) * 16, 16)
 rq_wqe[:] = struct.pack('>IIQ', len(test_msg), dev.mkey, dst_paddr)
 qp.rq_head += 1
 # ring recv doorbell from CPU (DBR offset 0 = recv counter)
-dev.dbr[qp.qp_dbr // 4] = to_be('I', qp.rq_head)
+dev.dbr[qp.qp_dbr // 4] = to_be32(qp.rq_head)
 
 # build send WQE in SQ from CPU (opcode 0x0a = SEND, ds_count=2)
 sq_head = qp.sq_head
@@ -60,7 +61,7 @@ wqe[0:8] = struct.pack('>II', (sq_head << 8) | 0x0a, (qp.qp_info['qpn'] << 8) |
 wqe[11] = 0x08  # CE: signal completion
 wqe[16:32] = struct.pack('>IIQ', len(test_msg), dev.mkey, src_paddr)
 qp.sq_head += 1
-doorbell_val = to_be('Q', int.from_bytes(bytes(wqe[0:8]), 'big'))
+doorbell_val = to_be64(int.from_bytes(bytes(wqe[0:8]), 'big'))
 
 # map MLX5 UAR and DBR into GPU VA
 uar_paddr = dev.pci_dev.bar_info(0)[0] + dev.uar * 0x1000
@@ -72,7 +73,7 @@ print(f"UAR gpu_va=0x{uar_gpu_va:x} DBR gpu_va=0x{dbr_gpu_va:x}")
 q = AMDComputeQueue(gpu)
 q.wait(gpu.timeline_signal, gpu.timeline_value - 1)
 # write DBR (32-bit sq_head) - send doorbell at qp_dbr + 4
-q.release_mem(dbr_gpu_va + qp.qp_dbr + 4, to_be('I', qp.sq_head), q.pm4.data_sel__mec_release_mem__send_32_bit_low,
+q.release_mem(dbr_gpu_va + qp.qp_dbr + 4, to_be32(qp.sq_head), q.pm4.data_sel__mec_release_mem__send_32_bit_low,
               q.pm4.int_sel__mec_release_mem__none)
 # write UAR doorbell (64-bit)
 q.release_mem(uar_gpu_va + 0x800, doorbell_val, q.pm4.data_sel__mec_release_mem__send_64_bit_data,

extra/optimization/test_beam_search.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad.helpers import BEAM, Timing, CI, prod
 from tinygrad import Variable, Device, Tensor
 from tinygrad.nn import Conv2d
-from tinygrad.uop.ops import AxisType
+from tinygrad.uop.ops import AxisType, Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.opt.postrange import Scheduler
 from tinygrad.codegen.opt.search import get_kernel_actions
@@ -85,6 +85,7 @@ class TestBeamSearch(unittest.TestCase):
     size = max(tc.dims[0], tc.dims[1]) * 8
     a, b = Tensor.rand(size, size, dtype=tc.dtype_in), Tensor.rand(size, size, dtype=tc.dtype_in)
     ast = a.matmul(b, dtype=tc.dtype_out).schedule()[-1].ast
+    if ast.op is Ops.BEAM: ast = ast.src[0]
     s = Scheduler(ast, Device[Device.DEFAULT].renderer)
     s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
     up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
@@ -95,6 +96,7 @@ class TestBeamSearch(unittest.TestCase):
   def test_max_up(self):
     a = Tensor.rand(16, 16)
     ast = a.schedule()[-1].ast
+    if ast.op is Ops.BEAM: ast = ast.src[0]
     s = Scheduler(ast, Device[Device.DEFAULT].renderer)
     for max_up in (2, 4):
       actions = get_kernel_actions(s, include_0=False, max_up=max_up)

extra/remu/Cargo.lock

@@ -1,66 +0,0 @@
-# This file is automatically @generated by Cargo.
-# It is not intended for manual editing.
-version = 4
-
-[[package]]
-name = "autocfg"
-version = "1.1.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
-
-[[package]]
-name = "cfg-if"
-version = "1.0.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
-
-[[package]]
-name = "crunchy"
-version = "0.2.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "7a81dae078cea95a014a339291cec439d2f232ebe854a9d672b796c6afafa9b7"
-
-[[package]]
-name = "float-cmp"
-version = "0.9.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "98de4bbd547a563b716d8dfa9aad1cb19bfab00f4fa09a6a4ed21dbcf44ce9c4"
-dependencies = [
- "num-traits",
-]
-
-[[package]]
-name = "half"
-version = "2.3.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "bc52e53916c08643f1b56ec082790d1e86a32e58dc5268f897f313fbae7b4872"
-dependencies = [
- "cfg-if",
- "crunchy",
- "num-traits",
-]
-
-[[package]]
-name = "libm"
-version = "0.2.8"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ec2a862134d2a7d32d7983ddcdd1c4923530833c9f2ea1a44fc5fa473989058"
-
-[[package]]
-name = "num-traits"
-version = "0.2.17"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c"
-dependencies = [
- "autocfg",
- "libm",
-]
-
-[[package]]
-name = "remu"
-version = "0.1.0"
-dependencies = [
- "float-cmp",
- "half",
- "num-traits",
-]

extra/remu/Cargo.toml

@@ -1,15 +0,0 @@
-[package]
-name = "remu"
-version = "0.1.0"
-edition = "2021"
-rust-version = "1.80.0"
-
-[lib]
-crate-type = ["cdylib"]
-
-[dependencies]
-half = { version = "2.3.1", features = ["num-traits"] }
-num-traits = "0.2.17"
-
-[dev-dependencies]
-float-cmp = "0.9.0"

extra/remu/README.md

@@ -1,80 +0,0 @@
-## Intro
-
-Remu is an RDNA3 emulator built to test correctness of RDNA3 code. It is used in [tinygrad's AMD CI](https://github.com/tinygrad/tinygrad).
-
-Most of the common instructions are implemented, but some formats like IMG are not supported.
-
-Remu is only for testing correctness of program output, it is not a cycle accurate simulator.
-
-## Build Locally
-
-Remu is written in Rust. Make sure you have [Cargo](https://doc.rust-lang.org/cargo/getting-started/installation.html).
-
-To build the project, run:
-
-```bash
-cargo build --release --manifest-path ./extra/remu/Cargo.toml
-```
-
-This will produce a binary in the `extra/remu/target/release` directory.
-
-## Usage with tinygrad
-
-The latest binaries are released in https://github.com/Qazalin/remu/releases. Alternatively, you can [build locally](#build-locally).
-
-Tinygrad does not yet output RDNA3 kernels directly. You can either install comgr or use `DEV=AMD:LLVM` (default) if you have [LLVM@19](https://github.com/tinygrad/tinygrad/blob/e2ed673c946c8f1774d816c75e52a994c2dd8a88/.github/actions/setup-tinygrad/action.yml#L208).
-
-`PYTHONPATH="." MOCKGPU=1 DEV=AMD python test/test_tiny.py TestTiny.test_plus` runs an emulated RDNA3 kernel with Remu.
-
-Add `DEBUG=6` to see Remu's logs.
-
-### DEBUG output
-
-Remu runs each thread one at a time in a nested for loop, see lib.rs. The DEBUG output prints information about the current thread.
-
-The DEBUG output has 3 sections:
-
-```
-<------------ 1 ----------> <--- 2 ---> <--------------------------------------- 3 ------------------------------------------>
-[0   0   0  ] [0   0   0  ] 0  F4080100 SMEM { op: 2, sdata: 4, sbase: 0, offset: 0, soffset: 124, glc: false, dlc: false }
-```
-
-#### Section 1: Grid info
-
-`[gid.x, gid.y, gid.z], [lid.x, lid.y, lid.z]` of the current thread.
-
-#### Section 2: Wave info
-
-`<lane> <instruction hex>`
-
-RDNA3 divides threads into chunks of 32. Each thread is assigned to a "lane" from 0-31.
-
-In Remu, even though all threads run one at a time, each 32 thread chunk (a wave) shares state like SGPR, VGPR, LDS, EXEC mask, etc.
-Remu can simulate up to one wave sync instruction.
-For more details, see work_group.rs.
-
-Section 2 can have a green or gray color.
-
-Green = The thread is actively executing the instruction.
-
-Gray = The thread has been "turned off" by the EXEC mask, it skips execution of some instructions. (refer to "EXECute Mask" on [page 23](https://www.amd.com/content/dam/amd/en/documents/radeon-tech-docs/instruction-set-architectures/rdna3-shader-instruction-set-architecture-feb-2023_0.pdf#page=23) of ISA docs for more details.)
-
-To see the colors in action, try running `DEBUG=6 PYTHONPATH="." MOCKGPU=1 DEV=AMD python test/test_ops.py TestOps.test_arange_big`. See how only lane 0 writes to global memory:
-```
-[255 0   0  ] [0   0   0  ] 0  DC6A0000 FLAT { op: 26, offset: 0, dlc: false, glc: false, slc: false, seg: 2, addr: 8, data: 0, saddr: 0, sve: false, vdst: 0 }
-[255 0   0  ] [1   0   0  ] 1  DC6A0000
-[255 0   0  ] [2   0   0  ] 2  DC6A0000
-[255 0   0  ] [3   0   0  ] 3  DC6A0000
-[255 0   0  ] [3   0   0  ] 4  DC6A0000
-```
-
-#### Section 3: Decoded Instruction
-
-This prints the instruction type and all the parsed bitfields.
-
-Remu output vs llvm-objdump:
-
-```
-s_load_b64 s[0:1], s[0:1], 0x10                            // 00000000160C: F4040000 F8000010
-SMEM { op: 1, sdata: 0, sbase: 0, offset: 16, soffset: 124, glc: false, dlc: false }
-```

extra/remu/rustfmt.toml

@@ -1 +0,0 @@
-max_width = 150

extra/remu/src/helpers.rs

@@ -1,162 +0,0 @@
-use half::f16;
-use num_traits::{float::FloatCore, PrimInt, Unsigned, clamp};
-
-pub fn bits<T>(word: T, hi: usize, lo: usize) -> T where T: PrimInt + Unsigned {
-    assert!(hi >= lo);
-    let width = hi - lo + 1;
-    (word >> lo) & ((T::one() << width) - T::one())
-}
-
-pub fn nth(val: u32, pos: usize) -> u32 {
-    (val >> (31 - pos as u32)) & 1
-}
-pub fn f16_lo(val: u32) -> f16 {
-    f16::from_bits((val & 0xffff) as u16)
-}
-pub fn f16_hi(val: u32) -> f16 {
-    f16::from_bits(((val >> 16) & 0xffff) as u16)
-}
-
-pub fn sign_ext(num: u64, bits: usize) -> i64 {
-    let mut value = num;
-    let is_negative = (value >> (bits - 1)) & 1 != 0;
-    if is_negative {
-        value |= !0 << bits;
-    }
-    value as i64
-}
-
-pub trait IEEEClass<T> {
-    fn exponent(&self) -> T;
-}
-impl IEEEClass<u32> for f32 {
-    fn exponent(&self) -> u32 {
-        (self.to_bits() & 0b01111111100000000000000000000000) >> 23
-    }
-}
-impl IEEEClass<u16> for f16 {
-    fn exponent(&self) -> u16 {
-        (self.to_bits() & 0b0111110000000000) >> 10
-    }
-}
-impl IEEEClass<u64> for f64 {
-    fn exponent(&self) -> u64 {
-        (self.to_bits() & 0b0111111111110000000000000000000000000000000000000000000000000000) >> 52
-    }
-}
-
-pub trait VOPModifier<T> {
-    fn negate(&self, pos: usize, modifier: usize) -> T;
-    fn absolute(&self, pos: usize, modifier: usize) -> T;
-    fn clmp(&self, cm: bool) -> T;
-}
-impl<T> VOPModifier<T> for T
-where
-    T: FloatCore,
-{
-    fn negate(&self, pos: usize, modifier: usize) -> T {
-        match (modifier >> pos) & 1 {
-            1 => -*self,
-            _ => *self,
-        }
-    }
-    fn absolute(&self, pos: usize, modifier: usize) -> T {
-        match (modifier >> pos) & 1 {
-            1 => self.abs(),
-            _ => *self,
-        }
-    }
-    fn clmp(&self, cm:bool) -> T {
-        if !cm { return *self }
-        let r = clamp(*self, T::zero(), T::one());
-        if r == T::zero() { T::zero() } else { r }
-    }
-}
-
-pub fn extract_mantissa(x: f64) -> f64 {
-    if x.is_infinite() || x.is_nan() {
-        return x;
-    }
-    let bits = x.to_bits();
-    let mantissa_mask: u64 = 0x000FFFFFFFFFFFFF;
-    let bias: u64 = 1023;
-    let normalized_mantissa_bits = (bits & mantissa_mask) | ((bias - 1) << 52);
-    return f64::from_bits(normalized_mantissa_bits);
-}
-pub fn ldexp(x: f64, exp: i32) -> f64 {
-    x * 2f64.powi(exp)
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    #[test]
-    fn test_extract_mantissa() {
-        assert_eq!(extract_mantissa(2.0f64), 0.5);
-    }
-
-    #[test]
-    fn test_normal_exponent() {
-        assert_eq!(2.5f32.exponent(), 128);
-        assert_eq!(1.17549435e-38f32.exponent(), 1);
-        assert_eq!(f32::INFINITY.exponent(), 255);
-        assert_eq!(f32::NEG_INFINITY.exponent(), 255);
-    }
-
-    #[test]
-    fn test_denormal_exponent() {
-        assert_eq!(1.0e-40f32.exponent(), 0);
-        assert_eq!(1.0e-42f32.exponent(), 0);
-        assert_eq!(1.0e-44f32.exponent(), 0);
-        assert_eq!((1.17549435e-38f32 / 2.0).exponent(), 0);
-    }
-
-    #[test]
-    fn test_normal_exponent_f16() {
-        assert_eq!(f16::from_f32(3.14f32).exponent(), 16);
-        assert_eq!(f16::NEG_INFINITY.exponent(), 31);
-        assert_eq!(f16::INFINITY.exponent(), 31);
-    }
-
-    #[test]
-    fn test_neg() {
-        assert_eq!(0.3_f32.negate(0, 0b001), -0.3_f32);
-        assert_eq!(0.3_f32.negate(1, 0b010), -0.3_f32);
-        assert_eq!(0.3_f32.negate(2, 0b100), -0.3_f32);
-        assert_eq!(0.3_f32.negate(0, 0b110), 0.3_f32);
-        assert_eq!(0.3_f32.negate(1, 0b010), -0.3_f32);
-        assert_eq!(0.0_f32.negate(0, 0b001).to_bits(), (-0.0f32).to_bits());
-        assert_eq!((-0.0_f32).negate(0, 0b001).to_bits(), 0);
-    }
-
-    #[test]
-    fn test_sign_ext() {
-        assert_eq!(sign_ext(0b000000000000000101000, 21), 40);
-        assert_eq!(sign_ext(0b111111111111111011000, 21), -40);
-        assert_eq!(sign_ext(0b000000000000000000000, 21), 0);
-        assert_eq!(sign_ext(0b111111111111111111111, 21), -1);
-        assert_eq!(sign_ext(0b111000000000000000000, 21), -262144);
-        assert_eq!(sign_ext(0b000111111111111111111, 21), 262143);
-        assert_eq!(sign_ext(7608, 13), -584);
-    }
-}
-
-use std::sync::LazyLock;
-pub static DEBUG: LazyLock<bool> = LazyLock::new(|| std::env::var("DEBUG").map(|v| v.parse::<usize>().unwrap_or(0) >= 6).unwrap_or(false));
-
-pub fn colored(st:&str, color:&str) -> String {
-    let ansi_code = match color {
-        "green" => format!("\x1b[{};2;39;176;139m", 38),
-        "gray" => format!("\x1b[{};2;169;169;169m", 38),
-        _ => format!("\x1b[{};2;255;255;255m", 38),
-    };
-    format!("{}{}{}", ansi_code, st, "\x1b[0m")
-}
-
-#[macro_export]
-macro_rules! todo_instr {
-    ($x:expr) => {{
-        println!("{:08X}", $x);
-        Err(1)
-    }};
-}

extra/remu/src/lib.rs

@@ -1,77 +0,0 @@
-use crate::state::StateSnapshot;
-use crate::work_group::{WaveContext, WorkGroup};
-use std::os::raw::c_char;
-use std::slice;
-mod helpers;
-mod rdna3;
-mod state;
-mod thread;
-mod work_group;
-
-#[no_mangle]
-pub extern "C" fn run_asm(lib: *const c_char, lib_sz: u32, gx: u32, gy: u32, gz: u32, lx: u32, ly: u32, lz: u32, args_ptr: *const u64) -> i32 {
-    if lib.is_null() || (lib_sz % 4) != 0 {
-        panic!("Pointer is null or length is not properly aligned to 4 bytes");
-    }
-    let kernel = unsafe { slice::from_raw_parts(lib as *const u32, (lib_sz / 4) as usize).to_vec() };
-    let dispatch_dim = match (gy != 1, gz != 1) {
-        (true, true) => 3,
-        (true, false) => 2,
-        _ => 1,
-    };
-    for gx in 0..gx {
-        for gy in 0..gy {
-            for gz in 0..gz {
-                let mut wg = WorkGroup::new(dispatch_dim, [gx, gy, gz], [lx, ly, lz], &kernel, args_ptr);
-                if let Err(err) = wg.exec_waves() {
-                    return err;
-                }
-            }
-        }
-    }
-    0
-}
-
-// FFI functions for single-stepping comparison tests
-
-#[no_mangle]
-pub extern "C" fn wave_create(lib: *const c_char, lib_sz: u32, n_lanes: u32) -> *mut WaveContext {
-    if lib.is_null() || (lib_sz % 4) != 0 { return std::ptr::null_mut(); }
-    let kernel = unsafe { slice::from_raw_parts(lib as *const u32, (lib_sz / 4) as usize).to_vec() };
-    Box::into_raw(Box::new(WaveContext::new(kernel, n_lanes as usize)))
-}
-
-#[no_mangle]
-pub extern "C" fn wave_step(ctx: *mut WaveContext) -> i32 {
-    if ctx.is_null() { return -99; }
-    unsafe { (*ctx).step() }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_get_snapshot(ctx: *const WaveContext, out: *mut StateSnapshot) {
-    if ctx.is_null() || out.is_null() { return; }
-    unsafe { *out = (*ctx).get_snapshot(); }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_set_sgpr(ctx: *mut WaveContext, idx: u32, val: u32) {
-    if ctx.is_null() || idx >= 128 { return; }
-    unsafe { (*ctx).scalar_reg[idx as usize] = val; }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_set_vgpr(ctx: *mut WaveContext, lane: u32, idx: u32, val: u32) {
-    if ctx.is_null() || lane >= 32 || idx >= 256 { return; }
-    unsafe { (*ctx).vec_reg.get_lane_mut(lane as usize)[idx as usize] = val; }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_init_lds(ctx: *mut WaveContext, size: u32) {
-    if ctx.is_null() { return; }
-    unsafe { (*ctx).lds.data.resize(size as usize, 0); }
-}
-
-#[no_mangle]
-pub extern "C" fn wave_free(ctx: *mut WaveContext) {
-    if !ctx.is_null() { unsafe { drop(Box::from_raw(ctx)); } }
-}

extra/remu/src/rdna3.rs

@@ -1,223 +0,0 @@
-use crate::helpers::{bits, sign_ext};
-
-#[derive(Debug, PartialEq)]
-pub enum Instruction {
-    SOP2 { op: u8, ssrc0: u8, ssrc1: u8, sdst: u8 },
-    SOP1 { op: u8, ssrc0: u8, sdst: u8 },
-    SOPK { op: u8, simm16: i16, sdst: u8 },
-    SOPP { op: u8, simm16: i16 },
-    SOPC { op: u8, ssrc0: u8, ssrc1: u8 },
-
-    SMEM { op: u8, sdata: u8, sbase: u8, offset: i32, soffset: u8, glc: bool, dlc: bool },
-
-    VOP1 { op: u8, vdst: u8, src: u16 },
-    VOP2 { op: u8, vdst: u8, vsrc: u8, src: u16 },
-    VOPC { op: u8, vsrc: u8, src: u16 },
-    VOP3 { op: u32, opsel: u8, cm: bool, abs: u8, vdst: u8, neg: u8, omod: u8, src2: u16, src1: u16, src0: u16 },
-    VOP3SD { op: u32, cm: bool, sdst: u8, vdst: u8, neg: u8, omod: u8, src2: u16, src1: u16, src0: u16 },
-    VOP3P { op: u8, vdst: u8, neg_hi: u8, opsel: u8, opsel_hi: u8, opsel_hi2: bool, cm: bool, src2: u16, src1: u16, src0: u16, neg: u8 },
-    VOPD { opx: u8, opy: u8, vdstx: u8, vdsty: u8, vsrcx1: u8, vsrcy1: u8, srcx0: u16, srcy0: u16 },
-
-    DS { op: u8, gds: bool, offset1: u8, offset0: u8, vdst: u8, data1: u8, data0: u8, addr: u8 },
-
-    FLAT { op: u8, offset: u16, dlc: bool, glc: bool, slc: bool, seg: u8, addr: u8, data: u8, saddr: u8, sve: bool, vdst: u8 }
-}
-
-const VOP3SD_OPS: [u32; 7] = [764, 765, 766, 767, 768, 769, 770];
-
-pub fn decode(word:u32, word1:Option<&u32>) -> Instruction {
-    match bits(word, 31, 30) {
-        0b11 => {
-            let word = (*word1.unwrap() as u64) << 32 | (word as u64);
-            match bits(word, 29, 26) {
-                0b1101 => {
-                    let sbase = (bits(word, 5, 0) as u8) << 1;
-                    let sdata = bits(word, 12, 6) as u8;
-                    let dlc = bits(word, 13, 13) != 0;
-                    let glc = bits(word, 14, 14) != 0;
-                    let op = bits(word, 25, 18) as u8;
-                    let offset = sign_ext(bits(word, 52, 32), 21) as i32;
-                    let soffset = bits(word, 63, 57) as u8;
-                    Instruction::SMEM { sbase, sdata, dlc, glc, op, offset, soffset }
-                }
-                0b0101 => {
-                    let op = bits(word, 25, 16) as u32;
-                    let vdst = bits(word, 7, 0) as u8;
-                    let cm = bits(word, 15, 15) != 0;
-                    let src0 = bits(word, 40, 32) as u16;
-                    let src1 = bits(word, 49, 41) as u16;
-                    let src2 = bits(word, 58, 50) as u16;
-                    let omod = bits(word, 60, 59) as u8;
-                    let neg = bits(word, 63, 61) as u8;
-                     if VOP3SD_OPS.contains(&op) {
-                         let sdst = bits(word, 14, 8) as u8;
-                         Instruction::VOP3SD { op, vdst, sdst, cm, src0, src1, src2, omod, neg }
-                     } else {
-                        let abs = bits(word, 10, 8) as u8;
-                        let opsel = bits(word, 14, 11) as u8;
-                        Instruction::VOP3 { opsel, cm, abs, vdst, neg, omod, src2, src1, src0, op }
-                     }
-                }
-                0b0011 => {
-                    let op = bits(word, 22, 16) as u8;
-                    let vdst = bits(word, 7, 0) as u8;
-                    let neg_hi = bits(word, 10, 8) as u8;
-                    let opsel = bits(word, 13, 11) as u8;
-                    let opsel_hi2 = bits(word, 14, 14) != 0;
-                    let cm = bits(word, 15, 15) != 0;
-                    let src0 = bits(word, 40, 32) as u16;
-                    let src1 = bits(word, 49, 41) as u16;
-                    let src2 = bits(word, 58, 50) as u16;
-                    let opsel_hi = bits(word, 60, 59) as u8;
-                    let neg = bits(word, 63, 61) as u8;
-                    Instruction::VOP3P { op, vdst, neg_hi, opsel, opsel_hi, opsel_hi2, cm, src0, src1, src2, neg }
-                }
-                0b0110 => {
-                    let offset0 = bits(word, 7, 0) as u8;
-                    let offset1 = bits(word, 15, 8) as u8;
-                    let gds = bits(word, 17, 17) != 0;
-                    let op = bits(word, 25, 18) as u8;
-                    let addr = bits(word, 39, 32) as u8;
-                    let data0 = bits(word, 47, 40) as u8;
-                    let data1 = bits(word, 55, 48) as u8;
-                    let vdst = bits(word, 63, 56) as u8;
-                    Instruction::DS { op, gds, offset1, offset0, vdst, data1, data0, addr }
-                }
-                0b0111 => {
-                    let offset = bits(word, 12, 0) as u16;
-                    let dlc = bits(word, 13, 13) != 0;
-                    let glc = bits(word, 14, 14) != 0;
-                    let slc = bits(word, 15, 15) != 0;
-                    let seg = bits(word, 17, 16) as u8;
-                    let op = bits(word, 24, 18) as u8;
-                    let addr = bits(word, 39, 32) as u8;
-                    let data = bits(word, 47, 40) as u8;
-                    let saddr = bits(word, 54, 48) as u8;
-                    let sve = bits(word, 55, 55) != 0;
-                    let vdst = bits(word, 63, 56) as u8;
-                    Instruction::FLAT { offset, dlc, glc, slc, seg, op, addr, data, saddr, sve, vdst }
-                },
-                0b0010 => {
-                    let srcx0 = bits(word, 8, 0) as u16;
-                    let vsrcx1 = bits(word, 16, 9) as u8;
-                    let opy = bits(word, 21, 17) as u8;
-                    let opx = bits(word, 25, 22) as u8;
-                    let srcy0 = bits(word, 40, 32) as u16;
-                    let vsrcy1 = bits(word, 48, 41) as u8;
-                    let vdsty = bits(word, 55, 49) as u8;
-                    let vdstx = bits(word, 63, 56) as u8;
-                    Instruction::VOPD { opx, opy, vdstx, vdsty, vsrcx1, vsrcy1, srcx0, srcy0 }
-                }
-                _ => todo!(),
-            }
-        }
-        0b10 => {
-            let ssrc0 = bits(word, 7, 0) as u8;
-            let ssrc1 = bits(word, 15, 8) as u8;
-            let simm16 = word as i16;
-            let sdst = bits(word, 22, 16) as u8;
-            match bits(word, 29, 23) {
-                0b1111101 => Instruction::SOP1 { ssrc0, sdst, op: bits(word, 15, 8) as u8 },
-                0b1111110 => Instruction::SOPC { ssrc0, ssrc1, op: bits(word, 22, 16) as u8 },
-                0b1111111 => Instruction::SOPP { simm16, op: bits(word, 22, 16) as u8 },
-                _ => {
-                    match bits(word, 29, 28) {
-                        0b11 => Instruction::SOPK { simm16, sdst, op: bits(word, 27, 23) as u8 },
-                        _ => Instruction::SOP2 { ssrc0, ssrc1, sdst, op: bits(word, 29, 23) as u8 }
-                    }
-                }
-            }
-        }
-        _ => {
-            let vdst = bits(word, 24, 17) as u8;
-            let src = bits(word, 8, 0) as u16;
-            let vsrc = bits(word, 16, 9) as u8;
-            match bits(word, 30, 25) {
-                0b111110 => Instruction::VOPC { vsrc, src, op: bits(word, 24, 17) as u8 },
-                0b111111 => Instruction::VOP1 { vdst, src, op: vsrc },
-                _ => Instruction::VOP2 { vdst, vsrc, src, op: bits(word, 30, 25) as u8 },
-            }
-        },
-    }
-}
-
-#[cfg(test)]
-mod test_rdna3 {
-    use super::*;
-
-    use std::process::{Stdio, Command};
-    use std::io::{Result, Write};
-
-    const LLVM_ARGS: &[&str; 3] = &["--arch=amdgcn", "--mcpu=gfx1100", "--triple=amdgcn-amd-amdhsa"];
-    const OFFSET_PRG: usize = 16;
-    const NULL: u8 = 124;
-
-    fn llvm_assemble(asm: &str) -> Result<Vec<u8>> {
-        let mut proc = Command::new("llvm-mc").args(LLVM_ARGS).args(["-filetype=obj", "-o", "-"]).stdin(Stdio::piped()).stdout(Stdio::piped()).spawn()?;
-        proc.stdin.as_mut().unwrap().write_all(asm.as_bytes())?;
-        let out = proc.wait_with_output()?;
-        match out.status.success() {
-            true => Ok(out.stdout),
-            false => Err(std::io::Error::new(std::io::ErrorKind::Other, "llvm-mc err")),
-        }
-    }
-
-    fn llvm_disassemble(code: &Vec<u8>) -> Result<String> {
-        let mut proc = Command::new("llvm-objdump").args(LLVM_ARGS).args(["--disassemble", "-"]).stdin(Stdio::piped()).stdout(Stdio::piped()).spawn()?;
-        proc.stdin.as_mut().unwrap().write_all(code)?;
-        let out = proc.wait_with_output()?;
-        match out.status.success() {
-            true => Ok(String::from_utf8(out.stdout).unwrap()),
-            false => Err(std::io::Error::new(std::io::ErrorKind::Other, "llvm-objdump err")),
-        }
-    }
-
-    fn test_decode(asm: &str) -> Instruction {
-        let lib = llvm_assemble(asm).unwrap();
-        println!("{}", llvm_disassemble(&lib).unwrap());
-        let stream: Vec<u32> = lib.chunks_exact(4).map(|chunk| u32::from_le_bytes(chunk.try_into().unwrap())).skip(OFFSET_PRG).collect();
-        decode(stream[0], stream.get(1))
-    }
-
-    #[test]
-    fn test_decode_smem() {
-        assert_eq!(test_decode("s_load_b128 s[4:7], s[0:1], null"), Instruction::SMEM { op: 2, sdata: 4, sbase: 0, offset: 0, soffset: NULL, glc: false, dlc: false });
-        assert_eq!(test_decode("s_load_b32 s10, s[0:1], 0xc"), Instruction::SMEM { op: 0, sdata: 10, sbase: 0, offset: 0xc, soffset: NULL, glc: false, dlc: false });
-        assert_eq!(test_decode("s_load_b32 s0, s[4:5], s6"), Instruction::SMEM { op: 0, sdata: 0, sbase: 4, offset: 0, soffset: 6, glc: false, dlc: false });
-        assert_eq!(test_decode("s_load_b32 s0, s[4:5], glc dlc"), Instruction::SMEM { op: 0, sdata: 0, sbase: 4, offset: 0, soffset: NULL, glc: true, dlc: true });
-        assert_eq!(test_decode("s_load_b32 s0, s[4:5], glc"), Instruction::SMEM { op: 0, sdata: 0, sbase: 4, offset: 0, soffset: NULL, glc: true, dlc: false });
-        assert_eq!(test_decode("s_load_b32 s0, s[4:5], -20"), Instruction::SMEM { op: 0, sdata: 0, sbase: 4, offset: -20, soffset: NULL, glc: false, dlc: false });
-        assert_eq!(test_decode("s_load_b32 s0, s[4:5], -1048576"), Instruction::SMEM { op: 0, sdata: 0, sbase: 4, offset: -1048576, soffset: NULL, glc: false, dlc: false });
-    }
-
-    #[test]
-    fn test_decode_salu() {
-        assert_eq!(test_decode("s_add_u32 s1 s2 s3"), Instruction::SOP2 { op: 0, ssrc0: 2, ssrc1: 3, sdst: 1 });
-        assert_eq!(test_decode("s_add_u32 vcc_hi exec_lo vcc_lo"), Instruction::SOP2 { op: 0, ssrc0: 126, ssrc1: 106, sdst: 107 });
-        assert_eq!(test_decode("s_mov_b32 s1 -0.5"), Instruction::SOP1 { op: 0, ssrc0: 241, sdst: 1 });
-        assert_eq!(test_decode("s_cmpk_eq_i32 s0 -30"), Instruction::SOPK { op: 3, sdst: 0, simm16: -30 });
-        assert_eq!(test_decode("s_cmpk_eq_u32 s0 65535"), Instruction::SOPK { op: 9, sdst: 0, simm16: -1 });
-        assert_eq!(test_decode("s_cmp_ge_i32 s1 s2"), Instruction::SOPC { op: 3, ssrc0: 1, ssrc1: 2 });
-    }
-
-    #[test]
-    fn test_decode_valu_e32() {
-        assert_eq!(test_decode("v_mov_b32 v0, v0"), Instruction::VOP1 { op: 1, vdst: 0, src: 256 });
-        assert_eq!(test_decode("v_mov_b32 v0, s0"), Instruction::VOP1 { op: 1, vdst: 0, src: 0 });
-        assert_eq!(test_decode("v_cmp_t_f32 v1, v0"), Instruction::VOPC { op: 31, vsrc: 0, src: 257 });
-    }
-
-    #[test]
-    fn test_decode_valu_e64() {
-        assert_eq!(test_decode("v_log_f32_e64 v2, |v0|"), Instruction::VOP3 { op: 423, vdst: 2, src0: 256, src1: 0, src2: 0, abs: 0b001, neg: 0, opsel: 0, omod: 0, cm: false });
-        assert_eq!(test_decode("v_div_scale_f32 v2, s1, v0, v1, v2"), Instruction::VOP3SD { op: 764, cm: false, vdst: 2, sdst: 1, src0: 256, src1: 257, src2: 258, omod: 0, neg: 0 });
-        assert_eq!(test_decode("v_pk_add_i16 v1, v0, v2"), Instruction::VOP3P { op: 2, vdst: 1, neg_hi: 0, opsel: 0, opsel_hi: 3, opsel_hi2: true, cm: false, src2: 0, src1: 258, src0: 256, neg: 0 });
-    }
-
-    #[test]
-    fn test_decode_ds() {
-        assert_eq!(test_decode("ds_add_u32 v2, v4 offset:16"), Instruction::DS { op: 0, gds: false, offset1: 0, offset0: 0x10, vdst: 0, data1: 0, data0: 4, addr: 2 });
-        assert_eq!(test_decode("ds_store_b32 v0, v1, offset: 0x04 gds"), Instruction::DS { op: 13, gds: true, offset1: 0, offset0: 0x04, vdst: 0, data1: 0, data0: 1, addr: 0 });
-        assert_eq!(test_decode("ds_load_u8 v1, v0 offset:16"), Instruction::DS { op: 58, gds: false, offset1: 0, offset0: 16, vdst: 1, data1: 0, data0: 0, addr: 0 });
-    }
-}

extra/remu/src/state.rs

@@ -1,272 +0,0 @@
-use std::ops::{Index, IndexMut};
-
-pub trait Register {
-    fn read64(&self, idx: usize) -> u64;
-    fn write64(&mut self, idx: usize, addr: u64);
-}
-impl<T> Register for T where T: Index<usize, Output = u32> + IndexMut<usize> {
-    fn read64(&self, idx: usize) -> u64 {
-        let lsb = self[idx] as u64;
-        let msb = self[idx + 1] as u64;
-        (msb << 32) | lsb
-    }
-
-    fn write64(&mut self, idx: usize, value: u64) {
-        self[idx] = (value & 0xffffffff) as u32;
-        self[idx + 1] = ((value & (0xffffffff << 32)) >> 32) as u32;
-    }
-}
-
-#[derive(Debug, Clone)]
-pub struct VGPR {
-    values: [[u32; 256]; 32],
-    pub default_lane: Option<usize>,
-}
-impl Index<usize> for VGPR {
-    type Output = u32;
-    fn index(&self, index: usize) -> &Self::Output {
-        &self.values[self.default_lane.unwrap()][index]
-    }
-}
-impl IndexMut<usize> for VGPR {
-    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
-        &mut self.values[self.default_lane.unwrap()][index]
-    }
-}
-impl VGPR {
-    pub fn new() -> Self {
-        VGPR {
-            values: [[0; 256]; 32],
-            default_lane: None,
-        }
-    }
-    pub fn get_lane(&self, lane: usize) -> [u32; 256] {
-        *self.values.get(lane).unwrap()
-    }
-    pub fn get_lane_mut(&mut self, lane: usize) -> &mut [u32; 256] {
-        self.values.get_mut(lane).unwrap()
-    }
-}
-
-pub trait Value {
-    fn mut_hi16(&mut self, val: u16);
-    fn mut_lo16(&mut self, val: u16);
-}
-impl Value for u32 {
-    fn mut_hi16(&mut self, val: u16) {
-        *self = ((val as u32) << 16) | (*self as u16 as u32);
-    }
-    fn mut_lo16(&mut self, val: u16) {
-        *self = ((((*self & (0xffff << 16)) >> 16) as u32) << 16) | val as u32;
-    }
-}
-
-#[derive(Debug, Clone, Copy)]
-pub struct WaveValue {
-    pub value: u32,
-    pub warp_size: usize,
-    pub default_lane: Option<usize>,
-    pub mutations: Option<[bool; 32]>,
-}
-impl WaveValue {
-    pub fn new(value: u32, warp_size: usize) -> Self {
-        Self {
-            value,
-            warp_size,
-            default_lane: None,
-            mutations: None,
-        }
-    }
-    pub fn read(&self) -> bool {
-        (self.value >> self.default_lane.unwrap()) & 1 == 1
-    }
-    pub fn set_lane(&mut self, value: bool) {
-        if self.mutations.is_none() {
-            self.mutations = Some([false; 32])
-        }
-        self.mutations.as_mut().unwrap()[self.default_lane.unwrap()] = value;
-    }
-    pub fn apply_muts(&mut self) {
-        self.value = 0;
-        for lane in 0..self.warp_size {
-            if self.mutations.unwrap()[lane] {
-                self.value |= 1 << lane;
-            }
-        }
-    }
-}
-
-/// C-compatible state snapshot for FFI - used for comparing emulator states
-#[repr(C)]
-#[derive(Clone, Debug)]
-pub struct StateSnapshot {
-    pub pc: u32,
-    pub scc: u32,
-    pub vcc: u32,
-    pub exec_mask: u32,
-    pub sgpr: [u32; 128],
-    pub vgpr: [[u32; 256]; 32],
-}
-
-impl StateSnapshot {
-    pub fn new() -> Self {
-        Self { pc: 0, scc: 0, vcc: 0, exec_mask: 0, sgpr: [0; 128], vgpr: [[0; 256]; 32] }
-    }
-}
-
-#[derive(Clone, Debug)]
-pub struct VecDataStore {
-    pub data: Vec<u8>,
-}
-
-impl VecDataStore {
-    pub fn new() -> Self {
-        Self { data: Vec::new() }
-    }
-    pub fn write(&mut self, addr: usize, val: u32) {
-        if addr + 4 >= self.data.len() {
-            self.data.resize(self.data.len() + addr + 5, 0);
-        }
-        self.data[addr..addr + 4].iter_mut().enumerate().for_each(|(i, x)| {
-            *x = val.to_le_bytes()[i];
-        });
-    }
-    pub fn write64(&mut self, addr: usize, val: u64) {
-        self.write(addr, (val & 0xffffffff) as u32);
-        self.write(addr + 4, ((val & (0xffffffff << 32)) >> 32) as u32);
-    }
-    pub fn read(&self, addr: usize) -> u32 {
-        let mut bytes: [u8; 4] = [0; 4];
-        bytes.copy_from_slice(&self.data[addr + 0..addr + 4]);
-        u32::from_le_bytes(bytes)
-    }
-    pub fn read64(&mut self, addr: usize) -> u64 {
-        let lsb = self.read(addr);
-        let msb = self.read(addr + 4);
-        ((msb as u64) << 32) | lsb as u64
-    }
-}
-
-#[cfg(test)]
-mod test_state {
-    use super::*;
-
-    #[test]
-    fn test_wave_value() {
-        let mut val = WaveValue::new(0b11000000000000011111111111101110, 32);
-        val.default_lane = Some(0);
-        assert!(!val.read());
-        val.default_lane = Some(31);
-        assert!(val.read());
-    }
-
-    #[test]
-    fn test_wave_value_small() {
-        let mut val = WaveValue::new(0, 1);
-        val.default_lane = Some(0);
-        assert!(!val.read());
-        assert_eq!(val.value, 0);
-        val.set_lane(true);
-        val.apply_muts();
-        assert!(val.read());
-        assert_eq!(val.value, 1);
-    }
-
-    #[test]
-    fn test_wave_value_small_alt() {
-        let mut val = WaveValue::new(0, 2);
-        val.default_lane = Some(0);
-        assert!(!val.read());
-        assert_eq!(val.value, 0);
-        val.set_lane(true);
-        val.apply_muts();
-        assert!(val.read());
-        assert_eq!(val.value, 1);
-    }
-
-    #[test]
-    fn test_wave_value_exec() {
-        let warp_size = 32;
-        let val = WaveValue::new(u32::MAX, warp_size);
-        assert_eq!(val.value, u32::MAX);
-        let warp_size = 3;
-        let val = WaveValue::new((1 << warp_size) - 1, warp_size);
-        assert_eq!(val.value, 7)
-    }
-
-    #[test]
-    fn test_wave_value_toggle_one() {
-        let warp_size = 2;
-        let mut val = WaveValue::new(0b11, warp_size);
-        // 0
-        val.default_lane = Some(0);
-        val.set_lane(false);
-        // 1
-        val.default_lane = Some(1);
-        val.set_lane(true);
-        val.apply_muts();
-        assert_eq!(val.value, 2);
-    }
-
-    #[test]
-    fn test_wave_value_mutate_small() {
-        let mut val = WaveValue::new(0, 2);
-        val.default_lane = Some(0);
-        assert!(!val.read());
-        assert_eq!(val.value, 0);
-        val.set_lane(true);
-        val.apply_muts();
-        assert!(val.read());
-        assert_eq!(val.value, 1);
-    }
-
-    #[test]
-    fn test_wave_value_mutations() {
-        let mut val = WaveValue::new(0b10001, 32);
-        val.default_lane = Some(0);
-        val.set_lane(false);
-        assert!(val.mutations.unwrap().iter().all(|x| !x));
-        val.default_lane = Some(1);
-        val.set_lane(true);
-        assert_eq!(val.value, 0b10001);
-        assert_eq!(
-            val.mutations,
-            Some([
-                false, true, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false, false,
-                false, false, false, false, false, false, false, false, false, false, false, false, false,
-            ])
-        );
-
-        val.apply_muts();
-        assert_eq!(val.value, 0b10);
-    }
-
-    #[test]
-    fn test_write16() {
-        let mut vgpr = VGPR::new();
-        vgpr.default_lane = Some(0);
-        vgpr[0] = 0b11100000000000001111111111111111;
-        vgpr[0].mut_lo16(0b1011101111111110);
-        assert_eq!(vgpr[0], 0b11100000000000001011101111111110);
-    }
-
-    #[test]
-    fn test_write16hi() {
-        let mut vgpr = VGPR::new();
-        vgpr.default_lane = Some(0);
-        vgpr[0] = 0b11100000000000001111111111111111;
-        vgpr[0].mut_hi16(0b1011101111111110);
-        assert_eq!(vgpr[0], 0b10111011111111101111111111111111);
-    }
-
-    #[test]
-    fn test_vgpr() {
-        let mut vgpr = VGPR::new();
-        vgpr.default_lane = Some(0);
-        vgpr[0] = 42;
-        vgpr.default_lane = Some(10);
-        vgpr[0] = 10;
-        assert_eq!(vgpr.get_lane(0)[0], 42);
-        assert_eq!(vgpr.get_lane(10)[0], 10);
-    }
-}

extra/remu/src/work_group.rs

@@ -1,323 +0,0 @@
-use crate::helpers::{colored, DEBUG};
-use crate::state::{Register, StateSnapshot, VecDataStore, WaveValue, VGPR};
-use crate::thread::{Thread, END_PRG, SGPR_COUNT};
-use std::collections::HashMap;
-
-pub const WAVE_SIZE: usize = 32;
-
-pub struct WorkGroup<'a> {
-    dispatch_dim: u32,
-    id: [u32; 3],
-    lds: VecDataStore,
-    kernel: &'a Vec<u32>,
-    kernel_args: *const u64,
-    launch_bounds: [u32; 3],
-    wave_state: HashMap<usize, WaveState>,
-}
-
-#[derive(Debug, Clone)]
-struct WaveState {
-    scalar_reg: [u32; SGPR_COUNT],
-    scc: u32,
-    vcc: WaveValue,
-    exec: WaveValue,
-    vec_reg: VGPR,
-    pc: usize,
-    sds: HashMap<usize, VecDataStore>,
-}
-
-const SYNCS: [u32; 4] = [0xBF89FC07, 0xBC7C0000, 0xBF890007, 0xbFB60003];
-const S_BARRIER: u32 = 0xBFBD0000;
-
-/// Context for single-stepping through a wave - holds all mutable state
-pub struct WaveContext {
-    pub kernel: Vec<u32>,
-    pub scalar_reg: [u32; SGPR_COUNT],
-    pub scc: u32,
-    pub pc: usize,
-    pub vec_reg: VGPR,
-    pub vcc: WaveValue,
-    pub exec: WaveValue,
-    pub lds: VecDataStore,
-    pub sds: HashMap<usize, VecDataStore>,
-    pub n_lanes: usize,
-}
-
-impl WaveContext {
-    pub fn new(kernel: Vec<u32>, n_lanes: usize) -> Self {
-        let active = (!0u32).wrapping_shr(32 - (n_lanes as u32));
-        Self {
-            kernel,
-            scalar_reg: [0; SGPR_COUNT],
-            scc: 0,
-            pc: 0,
-            vec_reg: VGPR::new(),
-            vcc: WaveValue::new(0, n_lanes),
-            exec: WaveValue::new(active, n_lanes),
-            lds: VecDataStore::new(),
-            sds: (0..=31).map(|i| (i, VecDataStore::new())).collect(),
-            n_lanes,
-        }
-    }
-
-    /// Execute a single instruction. Returns: 0=continue, -1=endpgm, -2=barrier, 1=done (pc past program), negative=error
-    pub fn step(&mut self) -> i32 {
-        if self.pc >= self.kernel.len() { return 1; }
-        if self.kernel[self.pc] == END_PRG { return -1; }
-        if self.kernel[self.pc] == S_BARRIER { self.pc += 1; return -2; }
-        // Skip sync/nop instructions
-        if SYNCS.contains(&self.kernel[self.pc]) || self.kernel[self.pc] >> 20 == 0xbf8 || self.kernel[self.pc] == 0x7E000000 {
-            self.pc += 1;
-            return 0;
-        }
-
-        let mut sgpr_co = None;
-        for lane_id in 0..self.n_lanes {
-            self.vec_reg.default_lane = Some(lane_id);
-            self.vcc.default_lane = Some(lane_id);
-            self.exec.default_lane = Some(lane_id);
-            let mut thread = Thread {
-                scalar_reg: &mut self.scalar_reg,
-                scc: &mut self.scc,
-                vec_reg: &mut self.vec_reg,
-                vcc: &mut self.vcc,
-                exec: &mut self.exec,
-                lds: &mut self.lds,
-                sds: &mut self.sds.get_mut(&lane_id).unwrap(),
-                pc_offset: 0,
-                stream: self.kernel[self.pc..].to_vec(),
-                scalar: false,
-                simm: None,
-                warp_size: self.n_lanes,
-                sgpr_co: &mut sgpr_co,
-            };
-            if let Err(e) = thread.interpret() { return e; }
-            if thread.scalar {
-                self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-                break;
-            }
-            if lane_id == self.n_lanes - 1 {
-                self.pc = ((self.pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-            }
-        }
-        if self.vcc.mutations.is_some() { self.vcc.apply_muts(); self.vcc.mutations = None; }
-        if self.exec.mutations.is_some() { self.exec.apply_muts(); self.exec.mutations = None; }
-        if let Some((idx, mut wv)) = sgpr_co.take() { wv.apply_muts(); self.scalar_reg[idx] = wv.value; }
-        0
-    }
-
-    pub fn get_snapshot(&self) -> StateSnapshot {
-        let mut snap = StateSnapshot::new();
-        snap.pc = self.pc as u32;
-        snap.scc = self.scc;
-        snap.vcc = self.vcc.value;
-        snap.exec_mask = self.exec.value;
-        snap.sgpr = self.scalar_reg;
-        for lane in 0..32 { snap.vgpr[lane] = self.vec_reg.get_lane(lane); }
-        snap
-    }
-}
-
-impl<'a> WorkGroup<'a> {
-    pub fn new(dispatch_dim: u32, id: [u32; 3], launch_bounds: [u32; 3], kernel: &'a Vec<u32>, kernel_args: *const u64) -> Self {
-        Self { dispatch_dim, id, kernel, launch_bounds, kernel_args, lds: VecDataStore::new(), wave_state: HashMap::new() }
-    }
-
-    pub fn exec_waves(&mut self) -> Result<(), i32> {
-        let mut threads = vec![];
-        for z in 0..self.launch_bounds[2] {
-            for y in 0..self.launch_bounds[1] {
-                for x in 0..self.launch_bounds[0] {
-                    threads.push([x, y, z])
-                }
-            }
-        }
-        let waves = threads.chunks(WAVE_SIZE).collect::<Vec<_>>();
-
-        let mut sync = false;
-        for (i, x) in self.kernel.iter().enumerate() {
-            if i != 0 && *x == S_BARRIER {
-                sync = true;
-                break;
-            }
-        }
-
-        for _ in 0..=(sync as usize) {
-            for w in waves.iter().enumerate() {
-                self.exec_wave(w)?
-            }
-        }
-        Ok(())
-    }
-
-    fn exec_wave(&mut self, (wave_id, threads): (usize, &&[[u32; 3]])) -> Result<(), i32> {
-        let (mut scalar_reg, mut scc, mut pc, mut vec_reg, mut vcc, mut exec, mut sds) = match self.wave_state.get(&wave_id) {
-          None => {
-            let mut scalar_reg = [0; SGPR_COUNT];
-            scalar_reg.write64(0, self.kernel_args as u64);
-
-            let [gx, gy, gz] = self.id;
-            match self.dispatch_dim {
-              3 => (scalar_reg[13], scalar_reg[14], scalar_reg[15]) = (gx, gy, gz),
-              2 => (scalar_reg[14], scalar_reg[15]) = (gx, gy),
-              _ => scalar_reg[15] = gx,
-            }
-
-            let mut vec_reg = VGPR::new();
-            for (t, [x, y, z]) in threads.iter().enumerate() {
-              vec_reg.get_lane_mut(t)[0] = match &self.launch_bounds {
-                [_, 1, 1] => *x,
-                _ => (z << 20) | (y << 10) | x,
-              }
-            }
-
-            let vcc = WaveValue::new(0, threads.len());
-            let active = (!0u32).wrapping_shr(32 - (threads.len() as u32));
-            let exec = WaveValue::new(active, threads.len());
-
-            let sds = (0..=31).map(|i| (i, VecDataStore::new())).collect();
-            (scalar_reg, 0, 0, vec_reg, vcc, exec, sds)
-          }
-
-          Some(val) => {
-            let val = val.clone();
-            (val.scalar_reg, val.scc, val.pc, val.vec_reg, val.vcc, val.exec, val.sds)
-          }
-        };
-
-        loop {
-            if self.kernel[pc] == END_PRG {
-                break Ok(());
-            }
-            if self.kernel[pc] == S_BARRIER && self.wave_state.get(&wave_id).is_none() {
-                self.wave_state.insert(wave_id, WaveState { scalar_reg, scc, vec_reg, vcc, exec, pc, sds });
-                break Ok(());
-            }
-            if self.kernel[pc] == S_BARRIER || SYNCS.contains(&self.kernel[pc]) || self.kernel[pc] >> 20 == 0xbf8 || self.kernel[pc] == 0x7E000000 {
-                pc += 1;
-                continue;
-            }
-
-            let mut sgpr_co = None;
-            for (lane_id, [x, y, z]) in threads.iter().enumerate() {
-                vec_reg.default_lane = Some(lane_id);
-                vcc.default_lane = Some(lane_id);
-                exec.default_lane = Some(lane_id);
-                if *DEBUG {
-                    let lane = format!("{:<2} {:08X} ", lane_id, self.kernel[pc]);
-                    let state = match exec.read() {
-                        true => "green",
-                        false => "gray",
-                    };
-                    let [id0, id1, id2] = self.id;
-                    print!("[{id0:<3} {id1:<3} {id2:<3}] [{x:<3} {y:<3} {z:<3}] {}", colored(&lane, state));
-                }
-                let mut thread = Thread {
-                    scalar_reg: &mut scalar_reg,
-                    scc: &mut scc,
-                    vec_reg: &mut vec_reg,
-                    vcc: &mut vcc,
-                    exec: &mut exec,
-                    lds: &mut self.lds,
-                    sds: &mut sds.get_mut(&lane_id).unwrap(),
-                    pc_offset: 0,
-                    stream: self.kernel[pc..self.kernel.len()].to_vec(),
-                    scalar: false,
-                    simm: None,
-                    warp_size: threads.len(),
-                    sgpr_co: &mut sgpr_co,
-                };
-                thread.interpret()?;
-                if *DEBUG {
-                    println!();
-                }
-                if thread.scalar {
-                    pc = ((pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-                    break;
-                }
-                if lane_id == threads.len() - 1 {
-                    pc = ((pc as isize) + 1 + (thread.pc_offset as isize)) as usize;
-                }
-            }
-
-            if vcc.mutations.is_some() {
-                vcc.apply_muts();
-                vcc.mutations = None;
-            }
-            if exec.mutations.is_some() {
-                exec.apply_muts();
-                exec.mutations = None;
-            }
-            if let Some((idx, mut wv)) = sgpr_co.take() {
-                wv.apply_muts();
-                scalar_reg[idx] = wv.value;
-            }
-        }
-    }
-}
-
-#[cfg(test)]
-mod test_workgroup {
-    use super::*;
-
-    // TODO: make this generic by adding the assembler
-    fn global_store_sgpr(addr: u64, instructions: Vec<u32>, src: u32) -> Vec<u32> {
-        [
-            instructions,
-            vec![
-                0x7E020200 + src,
-                0x7E0402FF,
-                addr as u32,
-                0x7E0602FF,
-                (addr >> 32) as u32,
-                0xDC6A0000,
-                0x007C0102,
-            ],
-            vec![END_PRG],
-        ]
-        .concat()
-    }
-
-    #[test]
-    fn test_wave_value_state_vcc() {
-        let mut ret: u32 = 0;
-        let kernel = vec![
-            0xBEEA00FF,
-            0b11111111111111111111111111111111, // initial vcc state
-            0x7E140282,
-            0x7C94010A, // cmp blockDim.x == 2
-        ];
-        let addr = (&mut ret as *mut u32) as u64;
-        let kernel = global_store_sgpr(addr, kernel, 106);
-        let mut wg = WorkGroup::new(1, [0, 0, 0], [3, 1, 1], &kernel, [addr].as_ptr());
-        wg.exec_waves().unwrap();
-        assert_eq!(ret, 0b100);
-    }
-
-    #[test]
-    fn test_wave_value_state_exec() {
-        let mut ret: u32 = 0;
-        let kernel = vec![
-            0xBEFE00FF,
-            0b11111111111111111111111111111111,
-            0x7E140282,
-            0x7D9C010A, // cmpx blockDim.x <= 2
-        ];
-        let addr = (&mut ret as *mut u32) as u64;
-        let kernel = global_store_sgpr(addr, kernel, 126);
-        let mut wg = WorkGroup::new(1, [0, 0, 0], [4, 1, 1], &kernel, [addr].as_ptr());
-        wg.exec_waves().unwrap();
-        assert_eq!(ret, 0b0111);
-    }
-
-    #[test]
-    fn test_wave_value_sgpr_co() {
-        let mut ret: u32 = 0;
-        let kernel = vec![0xBE8D00FF, 0x7FFFFFFF, 0x7E1402FF, u32::MAX, 0xD700000A, 0x0002010A];
-        let addr = (&mut ret as *mut u32) as u64;
-        let kernel = global_store_sgpr(addr, kernel, 0);
-        let mut wg = WorkGroup::new(1, [0, 0, 0], [5, 1, 1], &kernel, [addr].as_ptr());
-        wg.exec_waves().unwrap();
-        assert_eq!(ret, 0b11110);
-    }
-}

extra/remu/test/hwtest.py

@@ -1,155 +0,0 @@
-# ruff: noqa: F405, F403
-# allow define from star imports
-
-import numpy as np
-import unittest
-import subprocess, struct, math, functools
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import getenv
-
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.renderer.amd.asm import waitcnt
-
-from test.testextra.test_cfg_viz import asm_kernel
-
-def get_output(asm:list, n_threads:int=1, vdst:VGPR=v[1]):
-  out = Tensor([0]*n_threads, dtype=dtypes.uint32).realize()
-  insts = [
-    s_load_b64(s[0:1], s[0:1], NULL),
-    *asm,
-    v_lshlrev_b32_e32(v[0], 2, v[0]),
-    s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
-    #global_store_b32(v[0], v[1], s[0:1]),
-    global_store_b32(addr=v[0], data=vdst, saddr=s[0:1]),
-    s_endpgm()
-  ]
-  out = Tensor.custom_kernel(out, fxn=functools.partial(asm_kernel, name="test", insts=insts, device=out.device, n_threads=n_threads))[0]
-  out.realize()
-  return out.tolist()
-
-def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
-def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
-def f32_to_bits(x:float) -> int: return struct.unpack('<I', struct.pack('<f', x))[0]
-
-@unittest.skipUnless(Device.DEFAULT == "AMD", "tests RDNA3")
-class TestHW(unittest.TestCase):
-  def setUp(self):
-    if getenv("MOCKGPU"): subprocess.run(["cargo", "build", "--release", "--manifest-path", "./extra/remu/Cargo.toml"], check=True)
-
-  def test_simple_v_mov(self):
-    out = get_output([
-      v_mov_b32_e32(v[1], 2),
-    ])
-    self.assertEqual(out, [2])
-
-  def test_simple_s_mov(self):
-    out = get_output([
-      s_mov_b32(s[7], 0x7fffffff),
-      v_mov_b32_e32(v[1], s[7]),
-    ])
-    self.assertEqual(out, [0x7fffffff])
-
-  def test_exec_mov(self):
-    out = get_output([
-      v_mov_b32_e32(v[1], 42),
-      s_mov_b32(EXEC_LO, 0b10),
-      v_mov_b32_e32(v[1], 10),
-      s_mov_b32(EXEC_LO, 0b11),
-    ], n_threads=2)
-    np.testing.assert_equal(out, [42, 10])
-
-  def test_exec_cmp_vopc(self):
-    out = get_output([
-      s_mov_b32(VCC_LO, 0), # reset vcc
-      v_mov_b32_e32(v[1], 42),
-      v_mov_b32_e32(v[2], 10),
-      s_mov_b32(EXEC_LO, 0b01),
-      v_cmp_ne_u32_e32(v[1], v[2]),
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], VCC_LO),
-    ], n_threads=2)[0]
-    np.testing.assert_equal(out, 1)
-
-  def test_exec_cmpx_vop3(self):
-    out = get_output([
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], 42),
-      v_mov_b32_e32(v[2], 10),
-      s_mov_b32(EXEC_LO, 0b01),
-      v_cmpx_ne_u32_e32(v[1], v[2]),
-      s_mov_b32(s[10], EXEC_LO),
-      s_mov_b32(EXEC_LO, 0b11),
-      v_mov_b32_e32(v[1], s[10]),
-    ], n_threads=2)[0]
-    np.testing.assert_equal(out & 0b11, 0b01)
-
-  def test_fmac_vop3_modifier(self):
-    init_state = [
-      v_mov_b32_e32(a:=v[1], f16_to_bits(4.0)),
-      v_mov_b32_e32(b:=v[2], f16_to_bits(3.0)),
-      v_mov_b32_e32(c:=v[3], f16_to_bits(2.0)),
-    ]
-    def run_fmac(a, b): return get_output(init_state+[v_fmac_f16_e64(c, a, b)], vdst=c)[0]
-    self.assertEqual(run_fmac(a, b), f16_to_bits(14.0))
-    self.assertEqual(run_fmac(a, -b), f16_to_bits(-10.0))
-    self.assertEqual(run_fmac(-a, -b), f16_to_bits(14.0))
-
-  def test_s_abs_i32(self):
-    def check(x, y, dst=s[10], scc=0):
-      for reg,val in [(dst, y), (SCC, scc)]:
-        self.assertEqual(get_output([
-          s_mov_b32(dst, x),
-          s_abs_i32(dst, dst),
-          v_mov_b32_e32(v[1], reg)
-        ])[0], val)
-
-    check(0x00000001, 0x00000001, scc=1)
-    check(0x7fffffff, 0x7fffffff, scc=1)
-    check(0x80000000, 0x80000000, scc=1)
-    check(0x80000001, 0x7fffffff, scc=1)
-    check(0x80000002, 0x7ffffffe, scc=1)
-    check(0xffffffff, 0x00000001, scc=1)
-    check(0, 0, scc=0)
-
-  def test_v_rcp_f32_neg_vop3(self):
-    def v_neg_rcp_f32(x:float, y:float):
-      out = get_output([
-        v_mov_b32_e32(v[2], f32_to_bits(x)),
-        v_rcp_f32_e64(v[2], -v[2]),
-      ], vdst=v[2])[0]
-      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
-
-    v_neg_rcp_f32(math.inf, -0.0)
-    v_neg_rcp_f32(-math.inf, 0.0)
-    v_neg_rcp_f32(0.0, -math.inf)
-    v_neg_rcp_f32(-0.0, math.inf)
-    v_neg_rcp_f32(-2.0, 0.5)
-    v_neg_rcp_f32(2.0, -0.5)
-
-  def test_v_cndmask_b32_neg(self):
-    def v_neg(x:float, y:float):
-      out = get_output([
-        v_mov_b32_e32(v[1], f32_to_bits(x)),
-        s_mov_b32(s[10], 1),
-        v_cndmask_b32_e64(v[1], v[1], -v[1], s[10]),
-      ])[0]
-      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
-
-    v_neg(-0.0, 0.0)
-    v_neg(0.0, -0.0)
-    v_neg(2.0, -2.0)
-    v_neg(math.inf, -math.inf)
-    v_neg(-math.inf, math.inf)
-
-  @unittest.skip("how does VOPD work in the dsl")
-  def test_v_subrev_wrap(self):
-    out = get_output([
-      #v_dual_mov_b32(v[1], 0xffffffff, v[2], 0x0),
-      #v_dual_mov_b32(vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
-      #VOPD(opx=VOPDOp.V_DUAL_MOV_B32, opy=VOPDOp.V_DUAL_MOV_B32, vdstx=v[1], srcx=0xffffffff, vdsty=v[2], srcy=0x0),
-      v_subrev_co_u32(v[2], VCC_LO, v[2], v[1]),
-    ], vdst=v[2])[0]
-    self.assertEqual(out, 0xffff_ffff)
-
-if __name__ == "__main__":
-  unittest.main()

extra/setup_mock_amd_osx.sh

@@ -3,7 +3,7 @@ INSTALL_PATH="${1:-/opt/homebrew/lib}"
 if [ ! -d "$INSTALL_PATH" ]; then
     USER=$(whoami)
     echo "No path $INSTALL_PATH. Will create. Might need your password..."
-    echo "You can stop now and provide any location as an argument where you want to save the libs (note, that not default locations should be in LD_LIBRARY_PATH, so tinygrad can find the libs)."
+    echo "You can stop now and provide any location as an argument where you want to save the library (note, that not default locations should be in LD_LIBRARY_PATH, so tinygrad can find it)."
     echo "Press any key or symbol to continue..."
     read -n 1 -s
 
@@ -11,11 +11,6 @@ if [ ! -d "$INSTALL_PATH" ]; then
     sudo chown -R "$USER":staff "$INSTALL_PATH"
 fi
 
-# Download libremu.dylib
-curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-    jq -r '.assets[] | select(.name == "libremu.dylib").browser_download_url' | \
-    xargs curl -L -o $INSTALL_PATH/libremu.dylib
-
 # Download libamd_comgr.dylib
 curl -s https://api.github.com/repos/tinygrad/amdcomgr_dylib/releases/latest | \
     jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \

extra/sqtt/examples/generate_examples.py

@@ -10,6 +10,7 @@ EXAMPLES = {
   "plus":"test/test_tiny.py TestTiny.test_plus",
   "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
   "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
+  "handwritten":"test/amd/test_custom_kernel.py TestCustomKernel.test_handwritten",
 }
 
 if __name__ == "__main__":
@@ -21,6 +22,6 @@ if __name__ == "__main__":
     for i in range(2):
       # AM_RESET=1 gets a clear trace, does not work on mi300 machines
       subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
-                     env={**os.environ, "AMD":"1", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+                     env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
       PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
       print(f"saved SQTT trace to {dest}")

extra/sqtt/roc.py

@@ -136,7 +136,8 @@ def print_data(data:dict) -> None:
 
 def main() -> None:
   import tinygrad.viz.serve as viz
-  viz.ctxs = []
+  from tinygrad.uop.ops import RewriteTrace
+  data = viz.VizData()
 
   parser = argparse.ArgumentParser()
   parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
@@ -147,24 +148,24 @@ def main() -> None:
 
   with args.profile.open("rb") as f: profile = pickle.load(f)
 
-  viz.get_profile(profile)
+  viz.get_profile(profile, data=data)
 
   # List all kernels
   if args.kernel is None:
-    for c in viz.ctxs:
+    for c in data.ctxs:
       print(c["name"])
       for s in c["steps"]: print("  "+s["name"])
     return None
 
   # Find kernel trace
-  trace = next((c for c in viz.ctxs if c["name"] == f"Exec {args.kernel}"), None)
+  trace = next((c for c in data.ctxs if c["name"] == f"SQTT {args.kernel}"), None)
   if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
   n = 0
   for s in trace["steps"]:
     if "PKTS" in s["name"]: continue
     print(s["name"])
-    data = viz.get_render(s["query"])
-    print_data(data)
+    ret = viz.get_render(data, s["query"])
+    print_data(ret)
     n += 1
     if n > args.n: break
 

extra/thunder/amd/fa.py

@@ -81,7 +81,7 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
 
   attn, l_vec = Tensor.custom_kernel(attn, l_vec, xq, xk, xv, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D), grad_fxn=grad)[:2]
 
-  return attn.transpose(1, 2)
+  return attn.transpose(1, 2), attn, l_vec
 
 @functools.cache
 def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):

extra/viz/README.md

@@ -24,7 +24,7 @@ List all codegen steps for a kernel: `--rewrites -s E_3`
 Get source code: `--rewrites -s E_3 -i "View Source"`
 Inspect a graph rewrite: `--rewrites -s E_3 -i "initial symbolic"`
 
-# SQTT tracing
+## SQTT tracing
 
 Supported on AMD for RDNA3 and RDNA4 (best) and CDNA (developing).
 
@@ -38,8 +38,12 @@ You can select a specific trace with --source, Example workflow:
 VIZ=-2 python extra/gemm/amd_asm_matmul.py
 
 # View barriers
-extra/viz/cli.py --profile -s "SQTT kernel PKTS SE:0" | rg BARRIER | head -10
+extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | rg BARRIER | head -10
+
+# Get bank conflicts from performance counters
+
+python extra/viz/cli.py -p -s "kernel PMC" -i "SQC_LDS_BANK_CONFLICT"
 
 # Find the EXEC corresponding to a DISPATCH at cycle 410
-extra/viz/cli.py --profile -s "SQTT kernel PKTS SE:0" | awk '/EXEC/ && $1 - $5 == 410'
+extra/viz/cli.py --profile -s "kernel SQTT SE:0 PKTS" | awk '/EXEC/ && $1 - $5 == 410'
 ```

extra/viz/cli.py

@@ -47,20 +47,22 @@ def decode_profile(data:bytes) -> dict:
 def get(data:dict, key:str):
   for k,v in data.items():
     if ansistrip(k) == key: return v
-  raise RuntimeError(f'item "{key}" not found in list')
+  import difflib
+  match = difflib.get_close_matches(key, [ansistrip(k) for k in data], n=1, cutoff=0.6)
+  raise RuntimeError(f'item "{key}" not found in list'+(f", did you mean {match[0]!r}?" if match else ''))
 
 def main(args) -> None:
-  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
-  viz.ctxs = viz.get_rewrites(viz.trace)
+  data = viz.VizData(viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {})))
+  viz.load_rewrites(data)
 
   def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
 
   if args.profile:
     events:list = viz.load_pickle(args.profile_path, default=[])
-    if (profile_bytes:=viz.get_profile(events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
+    if (profile_bytes:=viz.get_profile(data, events)) is None: raise RuntimeError(f"empty profile in {args.profile_path}")
     profile = decode_profile(profile_bytes)
-    profile["layout"].update([(f'{c["name"]} {s["name"]}', s["data"]) for c in viz.ctxs if c["name"].startswith("SQTT") for s in c["steps"]
-                              if "PKTS" in s["name"]])
+    profile["layout"].update([(f'{c["name"][5:]}{" SQTT" if s["name"].endswith("PKTS") else ""} {s["name"]}', s["data"]) for c in data.ctxs
+                              if c["name"].startswith("SQTT") for s in c["steps"] if s["name"].endswith(("PMC", "PKTS"))])
     if args.src is None:
       for k in profile["layout"]:
         print(f"  {format_colored(k)}")
@@ -87,7 +89,7 @@ def main(args) -> None:
         op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
         phase, delay = None, 0
         idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
-        if e.device.startswith("WAVE") or e.device == "OTHER_SIMD":
+        if e.device.startswith("WAVE"):
           inst = f"0x{(pc:=int(info.replace('PC:', ''))):05x} {pc_map[pc]}" if info else f"{'':7} {op_name}"
           dispatch_to_inst[f"{e.device}-{idx}"] = (inst, int(e.st))
           phase = "DISPATCH"
@@ -99,6 +101,20 @@ def main(args) -> None:
         print(f"{int(e.st)-inst_st:<12} {unit:<20} {op_str}{' '*(22-ansilen(op_str))} {int(unwrap(e.en)-e.st):<4} {str(delay or ''):<4} {info}")
       return None
 
+    # ** PMC printer
+    if "PMC" in args.src:
+      table = viz.unpack_pmc(data[0])
+      cols = table["cols"]
+      rows:list = []
+      for r in table["rows"]:
+        if args.item is None: rows.append(r[:2])
+        elif args.item == r[0]:
+          rows = r[2]["rows"] if len(r) > 2 else [r[:2]]
+          cols = r[2]["cols"] if len(r) > 2 else cols
+      from tabulate import tabulate
+      print(tabulate(rows, headers=cols, tablefmt="github"))
+      return None
+
     # ** Profiler printer
     agg:dict[str, tuple[float, int]] = {}
     total = 0
@@ -126,7 +142,7 @@ def main(args) -> None:
     return None
 
   # ** Graph rewrites printer
-  rewrites = {c["name"]:{s["name"]:s for s in c["steps"]} for c in viz.ctxs if c.get("steps")}
+  rewrites = {c["name"]:{s["name"]:s for s in c["steps"]} for c in data.ctxs if c.get("steps")}
   if args.src is None:
     for k in rewrites: print(f"  {format_colored(k)}")
     return None
@@ -134,7 +150,7 @@ def main(args) -> None:
   if args.item is None:
     for k,v in steps.items(): print(" "*v["depth"]+k+(f" - {v['match_count']}" if v.get('match_count', 0) else ''))
   else:
-    data = viz.get_render(get(steps, args.item)["query"])
+    data = viz.get_render(data, get(steps, args.item)["query"])
     if isinstance(data.get("value"), Iterator):
       for m in data["value"]:
         if m.get("uop"): print(f"Input UOp:\n{m['uop']}")

test/amd/disasm.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 import re
 from typing import Callable
+from test.amd.helpers import decode_dpp16
 from tinygrad.renderer.amd.dsl import Inst, Reg
 
 # Special register mappings for disassembly
@@ -838,22 +839,11 @@ def _disasm_vop1_sdwa(inst) -> str:
 
 def _decode_dpp(dpp: int) -> str:
   """Decode DPP control value to string."""
-  if dpp < 0x100: return f"quad_perm:[{dpp&3},{(dpp>>2)&3},{(dpp>>4)&3},{(dpp>>6)&3}]"
-  if 0x100 <= dpp <= 0x10f: return f"row_shl:{dpp & 0xf}"
-  if 0x110 <= dpp <= 0x11f: return f"row_shr:{dpp & 0xf}"
-  if 0x120 <= dpp <= 0x12f: return f"row_ror:{dpp & 0xf}"
-  if dpp == 0x130: return "wave_shl:1"
-  if dpp == 0x134: return "wave_rol:1"
-  if dpp == 0x138: return "wave_shr:1"
-  if dpp == 0x13c: return "wave_ror:1"
-  if dpp == 0x140: return "row_mirror"
-  if dpp == 0x141: return "row_half_mirror"
-  if dpp == 0x142: return "row_bcast:15"
-  if dpp == 0x143: return "row_bcast:31"
-  if 0x150 <= dpp <= 0x15f: return f"row_newbcast:{dpp & 0xf}"
-  if 0x160 <= dpp <= 0x16f: return f"row_share:{dpp & 0xf}"
-  if 0x170 <= dpp <= 0x17f: return f"row_xmask:{dpp & 0xf}"
-  return f"dpp:{dpp:#x}"
+  op, arg = decode_dpp16(dpp)
+  if op == "quad_perm": return f"quad_perm:[{','.join(str(x) for x in arg)}]"
+  if op in ("row_mirror", "row_half_mirror"): return op
+  if op == "dpp": return f"dpp:{arg:#x}"
+  return f"{op}:{arg}"
 
 def _disasm_vop1_dpp(inst) -> str:
   name = inst.op_name.lower().replace('_e32', '')

test/amd/helpers.py

@@ -12,8 +12,19 @@ ARCH_TO_TARGET:dict[str, list[str]] = {
 
 TARGET_TO_ARCH:dict[str, str] = {t:arch for arch,targets in ARCH_TO_TARGET.items() for t in targets}
 
+_DPP16_RANGE_OPS = {0x100: "row_shl", 0x110: "row_shr", 0x120: "row_ror", 0x150: "row_newbcast", 0x160: "row_share", 0x170: "row_xmask"}
+_DPP16_EXACT_OPS = {0x130: ("wave_shl", 1), 0x134: ("wave_rol", 1), 0x138: ("wave_shr", 1), 0x13c: ("wave_ror", 1),
+                    0x140: ("row_mirror", 0), 0x141: ("row_half_mirror", 0), 0x142: ("row_bcast", 15), 0x143: ("row_bcast", 31)}
+
 def get_target(arch:str) -> str: return ARCH_TO_TARGET[arch][0]
 
+def decode_dpp16(dpp: int) -> tuple[str, int | tuple[int, int, int, int]]:
+  """Decode a DPP16 control word into a symbolic operation and argument."""
+  if dpp < 0x100: return "quad_perm", ((dpp >> 0) & 0x3, (dpp >> 2) & 0x3, (dpp >> 4) & 0x3, (dpp >> 6) & 0x3)
+  if dpp in _DPP16_EXACT_OPS: return _DPP16_EXACT_OPS[dpp]
+  if (base := dpp & 0x1f0) in _DPP16_RANGE_OPS: return _DPP16_RANGE_OPS[base], dpp & 0xf
+  return "dpp", dpp
+
 def get_mattr(arch:str) -> str:
   return {"rdna3":"+real-true16,+wavefrontsize32", "rdna4":"+real-true16,+wavefrontsize32", "cdna":"+wavefrontsize64"}[arch]
 

test/amd/hw/test_dpp.py

@@ -0,0 +1,187 @@
+"""Tests for DPP16 source swizzles.
+
+These instructions trap in the default wave32 hw helper, so this file uses a
+minimal wave64 lane-store harness and compares emulator vs hardware directly
+when USE_HW=1.
+"""
+import ctypes, unittest
+from tinygrad.runtime.autogen.amd.rdna3.ins import *
+from tinygrad.helpers import flat_mv
+from test.amd.hw.helpers import USE_HW, assemble
+from test.mockgpu.amd.emu import run_asm
+
+WAVE64 = 64
+
+def _wave64_code(instructions: list, out_reg: int = 1) -> bytes:
+  return assemble([
+    s_mov_b32(s[80], s[0]),
+    s_mov_b32(s[81], s[1]),
+    v_mov_b32_e32(v[255], v[0]),
+    *instructions,
+    s_load_b64(s[92:93], s[80:81], 0, soffset=NULL),
+    s_waitcnt(0),
+    v_lshlrev_b32_e32(v[240], 2, v[255]),
+    global_store_b32(addr=v[240], data=v[out_reg], saddr=s[92:93], offset=0),
+    s_endpgm(),
+  ])
+
+def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
+  out_buf = (ctypes.c_uint32 * WAVE64)(*([0] * WAVE64))
+  args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
+  code = _wave64_code(instructions, out_reg)
+  kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
+  rsrc2 = 0x19c | (128 << 15)
+  scratch_size = 0x10000
+  result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, WAVE64, 1, 1, ctypes.addressof(args), rsrc2, scratch_size)
+  assert result == 0, f"run_asm failed with {result}"
+  return list(out_buf)
+
+def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+  dev = Device["AMD"]
+  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
+  code = _wave64_code(instructions, out_reg)
+  byte_str = ', '.join(f'0x{b:02x}' for b in code)
+  asm_src = f""".text
+.globl test
+.p2align 8
+.type test,@function
+test:
+.byte {byte_str}
+
+.rodata
+.p2align 6
+.amdhsa_kernel test
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 96
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_kernarg_size 8
+  .amdhsa_group_segment_fixed_size 65536
+  .amdhsa_private_segment_fixed_size 65536
+  .amdhsa_enable_private_segment 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: test
+    .symbol: test.kd
+    .kernarg_segment_size: 8
+    .group_segment_fixed_size: 65536
+    .private_segment_fixed_size: 65536
+    .kernarg_segment_align: 8
+    .wavefront_size: 64
+    .sgpr_count: 96
+    .vgpr_count: 256
+    .max_flat_workgroup_size: 1024
+...
+.end_amdgpu_metadata
+"""
+  lib = compiler.compile(asm_src)
+  prg = AMDProgram(dev, "test", lib)  # type: ignore[arg-type]
+  out_gpu = dev.allocator.alloc(WAVE64 * 4)
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
+  out = bytearray(WAVE64 * 4)
+  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
+  return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(WAVE64)]
+
+def run_wave64(instructions: list, out_reg: int = 1) -> list[int]:
+  emu = _run_wave64_emu(instructions, out_reg)
+  if not USE_HW: return emu
+  hw = _run_wave64_hw(instructions, out_reg)
+  if emu != hw:
+    diffs = [f"lane {i}: emu=0x{e:08x} hw=0x{h:08x}" for i, (e, h) in enumerate(zip(emu, hw)) if e != h]
+    raise AssertionError("Emulator vs Hardware mismatch:\n" + '\n'.join(diffs[:16]))
+  return hw
+
+class TestDPP16(unittest.TestCase):
+  def _run_copy(self, dpp: int, *, row_mask: int = 0xf, bank_mask: int = 0xf, bc: int = 1, dst_seed: int | None = None) -> list[int]:
+    instructions = [
+      v_mul_u32_u24_e32(v[0], 10, v[255]),
+      v_add_nc_u32_e32(v[0], 3, v[0]),
+    ]
+    if dst_seed is not None: instructions.append(v_mov_b32_e32(v[1], dst_seed))
+    instructions += [v_mov_b32_e32(v[2], 0), v_or_b32_e32(v[1], DPP, v[2], vsrc0=v[0], dpp=dpp, row_mask=row_mask, bank_mask=bank_mask, bc=bc)]
+    return run_wave64(instructions)
+
+  def test_quad_perm_reverse(self):
+    out = self._run_copy(0x1b)
+    self.assertEqual(out[0], 33)
+    self.assertEqual(out[1], 23)
+    self.assertEqual(out[2], 13)
+    self.assertEqual(out[3], 3)
+    self.assertEqual(out[4], 73)
+
+  def test_row_shl(self):
+    out = self._run_copy(0x101)
+    self.assertEqual(out[0], 13)
+    self.assertEqual(out[7], 83)
+    self.assertEqual(out[14], 153)
+    self.assertEqual(out[15], 0)
+    self.assertEqual(out[16], 173)
+
+  def test_row_shr(self):
+    out = self._run_copy(0x111)
+    self.assertEqual(out[0], 0)
+    self.assertEqual(out[1], 3)
+    self.assertEqual(out[8], 73)
+    self.assertEqual(out[15], 143)
+    self.assertEqual(out[16], 0)
+    self.assertEqual(out[17], 163)
+
+  def test_row_ror(self):
+    out = self._run_copy(0x121)
+    self.assertEqual(out[0], 153)
+    self.assertEqual(out[1], 3)
+    self.assertEqual(out[15], 143)
+    self.assertEqual(out[16], 313)
+
+  def test_row_mirror(self):
+    out = self._run_copy(0x140)
+    self.assertEqual(out[0], 153)
+    self.assertEqual(out[5], 103)
+    self.assertEqual(out[8], 73)
+    self.assertEqual(out[16], 313)
+
+  def test_row_half_mirror(self):
+    out = self._run_copy(0x141)
+    self.assertEqual(out[0], 73)
+    self.assertEqual(out[7], 3)
+    self.assertEqual(out[8], 153)
+    self.assertEqual(out[15], 83)
+    self.assertEqual(out[16], 233)
+
+  def test_row_mask(self):
+    out = self._run_copy(0x101, row_mask=0x5, dst_seed=0xDEADBEEF)
+    self.assertEqual(out[0], 13)
+    self.assertEqual(out[15], 0)
+    self.assertEqual(out[16], 0xDEADBEEF)
+    self.assertEqual(out[32], 333)
+    self.assertEqual(out[47], 0)
+    self.assertEqual(out[48], 0xDEADBEEF)
+
+  def test_bank_mask(self):
+    out = self._run_copy(0x101, bank_mask=0x5, dst_seed=0xDEADBEEF)
+    self.assertEqual(out[0], 13)
+    self.assertEqual(out[3], 43)
+    self.assertEqual(out[4], 0xDEADBEEF)
+    self.assertEqual(out[8], 93)
+    self.assertEqual(out[12], 0xDEADBEEF)
+
+class TestVOPCDPP16(unittest.TestCase):
+  def test_row_bcast15_materializes_vcc(self):
+    out = run_wave64([
+      v_mov_b32_e32(v[0], v[255]),
+      v_cmp_eq_u32_e32(DPP, v[0], vsrc0=v[0], dpp=0x142, row_mask=0xf, bank_mask=0xf, bc=1),
+      v_mov_b32_e32(v[2], 0),
+      v_mov_b32_e32(v[3], 1),
+      v_cndmask_b32_e32(v[1], v[2], v[3]),
+    ])
+    for lane in (0, 16, 32, 48): self.assertEqual(out[lane], 1)
+    for lane in (1, 15, 31, 47, 63): self.assertEqual(out[lane], 0)

test/amd/hw/test_ds.py

@@ -833,8 +833,6 @@ class TestDsPermute(unittest.TestCase):
       src_lane = lane ^ 1
       expected = src_lane + 100
       self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
-
-
 class TestDSSubDword(unittest.TestCase):
   """Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
 

test/amd/hw/test_rdna4_permlane_var.py

@@ -0,0 +1,129 @@
+"""RDNA4 V_PERMLANE16_VAR_B32 / V_PERMLANEX16_VAR_B32 coverage.
+
+Exercises the generated pcode path end-to-end in the emulator and compares against
+real RDNA4 hardware when USE_HW=1.
+"""
+import ctypes, unittest
+import tinygrad.runtime.autogen.amd.rdna4.ins as r4
+from tinygrad.helpers import flat_mv
+from tinygrad.renderer.amd.dsl import NULL
+from test.amd.hw.helpers import USE_HW, assemble
+from test.mockgpu.amd.emu import run_asm
+
+LANES = 32
+
+def _code(instructions: list, out_reg: int = 2) -> bytes:
+  return assemble([
+    r4.s_mov_b32(r4.s[80], r4.s[0]),
+    r4.s_mov_b32(r4.s[81], r4.s[1]),
+    r4.v_mov_b32_e32(r4.v[255], r4.v[0]),
+    *instructions,
+    r4.s_load_b64(r4.s[92:93], r4.s[80:81], soffset=NULL),
+    r4.s_wait_kmcnt(simm16=0),
+    r4.v_lshlrev_b32_e32(r4.v[240], 2, r4.v[255]),
+    r4.v_mov_b32_e32(r4.v[241], 0),
+    r4.global_store_b32(vaddr=r4.v[240:241], saddr=r4.s[92:93], vsrc=r4.v[out_reg]),
+    r4.s_endpgm(),
+  ])
+
+def _run_emu(instructions: list, out_reg: int = 2) -> list[int]:
+  out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES))
+  args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
+  code = _code(instructions, out_reg)
+  kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
+  result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args), arch='rdna4')
+  assert result == 0, f"run_asm failed with {result}"
+  return list(out_buf)
+
+def _run_hw(instructions: list, out_reg: int = 2) -> list[int]:
+  from tinygrad.device import Device
+  from tinygrad.runtime.ops_amd import AMDProgram
+  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+  dev = Device['AMD']
+  if not dev.arch.startswith('gfx12'): raise unittest.SkipTest('requires RDNA4 hardware')
+  compiler = HIPCompiler(dev.arch)
+  code = _code(instructions, out_reg)
+  byte_str = ', '.join(f'0x{b:02x}' for b in code)
+  asm_src = f""".text
+.globl test
+.p2align 8
+.type test,@function
+test:
+.byte {byte_str}
+
+.rodata
+.p2align 6
+.amdhsa_kernel test
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 96
+  .amdhsa_wavefront_size32 1
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_kernarg_size 8
+  .amdhsa_group_segment_fixed_size 65536
+  .amdhsa_private_segment_fixed_size 65536
+  .amdhsa_enable_private_segment 1
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.version:
+  - 1
+  - 0
+amdhsa.kernels:
+  - .name: test
+    .symbol: test.kd
+    .kernarg_segment_size: 8
+    .group_segment_fixed_size: 65536
+    .private_segment_fixed_size: 65536
+    .kernarg_segment_align: 8
+    .wavefront_size: 32
+    .sgpr_count: 96
+    .vgpr_count: 256
+    .max_flat_workgroup_size: 1024
+...
+.end_amdgpu_metadata
+"""
+  lib = compiler.compile(asm_src)
+  prg = AMDProgram(dev, 'test', lib)
+  out_gpu = dev.allocator.alloc(LANES * 4)
+  prg(out_gpu, global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
+  out = bytearray(LANES * 4)
+  dev.allocator._copyout(flat_mv(memoryview(out)), out_gpu)
+  return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(LANES)]
+
+def run_rdna4(instructions: list, out_reg: int = 2) -> list[int]:
+  emu = _run_emu(instructions, out_reg)
+  if not USE_HW: return emu
+  hw = _run_hw(instructions, out_reg)
+  if emu != hw:
+    diffs = [f"lane {i}: emu=0x{e:08x} hw=0x{h:08x}" for i, (e, h) in enumerate(zip(emu, hw)) if e != h]
+    raise AssertionError("Emulator vs Hardware mismatch:\n" + '\n'.join(diffs[:16]))
+  return hw
+
+class TestPermlaneVarRDNA4(unittest.TestCase):
+  def test_v_permlane16_var_b32_reverse(self):
+    out = run_rdna4([
+      r4.v_mov_b32_e32(r4.v[0], r4.v[255]),
+      r4.v_xor_b32_e32(r4.v[1], 15, r4.v[255]),
+      r4.v_permlane16_var_b32(r4.v[2], r4.v[0], r4.v[1]),
+    ])
+    self.assertEqual(out[0], 15)
+    self.assertEqual(out[5], 10)
+    self.assertEqual(out[15], 0)
+    self.assertEqual(out[16], 31)
+    self.assertEqual(out[21], 26)
+    self.assertEqual(out[31], 16)
+
+  def test_v_permlanex16_var_b32_cross_row(self):
+    out = run_rdna4([
+      r4.v_mov_b32_e32(r4.v[0], r4.v[255]),
+      r4.v_mov_b32_e32(r4.v[1], r4.v[255]),
+      r4.v_permlanex16_var_b32(r4.v[2], r4.v[0], r4.v[1]),
+    ])
+    self.assertEqual(out[0], 16)
+    self.assertEqual(out[5], 21)
+    self.assertEqual(out[15], 31)
+    self.assertEqual(out[16], 0)
+    self.assertEqual(out[21], 5)
+    self.assertEqual(out[31], 15)

test/amd/hw/test_vinterp.py

@@ -0,0 +1,35 @@
+"""Tests for VINTERP instructions."""
+import unittest
+from test.amd.hw.helpers import *
+
+class TestVInterp(unittest.TestCase):
+  def test_v_interp_p10_f32(self):
+    instructions = [
+      v_mov_b32_e32(v[10], v[255]),
+      v_cvt_f32_u32_e32(v[1], v[10]),
+      s_mov_b32(s[0], f2i(100.0)),
+      v_add_f32_e32(v[1], s[0], v[1]),
+      v_cvt_f32_u32_e32(v[3], v[10]),
+      s_mov_b32(s[1], f2i(10.0)),
+      v_add_f32_e32(v[3], s[1], v[3]),
+      s_mov_b32(s[2], f2i(2.0)),
+      v_interp_p10_f32(v[4], v[1], s[2], v[3]),
+    ]
+    st = run_program(instructions, n_lanes=8)
+    for lane in range(4): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 212.0, places=5)
+    for lane in range(4, 8): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 224.0, places=5)
+
+  def test_v_interp_p10_f16_f32(self):
+    instructions = [
+      v_mov_b32_e32(v[10], v[255]),
+      v_cvt_f32_u32_e32(v[11], v[10]),
+      v_cvt_f16_f32_e32(v[1], v[11]),
+      s_mov_b32(s[0], f2i(10.0)),
+      v_add_f32_e32(v[12], s[0], v[11]),
+      v_cvt_f16_f32_e32(v[3], v[12]),
+      s_mov_b32(s[1], f2i(2.0)),
+      v_interp_p10_f16_f32(v[4], v[1], s[1], v[3]),
+    ]
+    st = run_program(instructions, n_lanes=8)
+    for lane in range(4): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 12.0, places=5)
+    for lane in range(4, 8): self.assertAlmostEqual(i2f(st.vgpr[lane][4]), 24.0, places=5)

test/amd/hw/test_vop2.py

@@ -30,6 +30,17 @@ class TestBasicArithmetic(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertAlmostEqual(i2f(st.vgpr[0][2]), 8.0, places=5)
 
+  def test_v_add_f32_dpp_row_shl(self):
+    """V_ADD_F32 DPP row_shl swizzles src0 before the add."""
+    instructions = [
+      v_cvt_f32_u32_e32(v[0], v[255]),
+      v_add_f32_e32(v[1], DPP, v[0], vsrc0=v[0], dpp=0x101, row_mask=0xf, bank_mask=0xf, bc=1),
+    ]
+    st = run_program(instructions, n_lanes=16)
+    self.assertAlmostEqual(i2f(st.vgpr[0][1]), 1.0, places=5)
+    self.assertAlmostEqual(i2f(st.vgpr[1][1]), 3.0, places=5)
+    self.assertAlmostEqual(i2f(st.vgpr[14][1]), 29.0, places=5)
+
   def test_v_fmac_f32(self):
     """V_FMAC_F32: d = d + a*b using inline constants."""
     instructions = [

test/amd/hw/test_vop3.py

@@ -20,6 +20,28 @@ class TestFMA(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertAlmostEqual(i2f(st.vgpr[0][3]), 9.0, places=5)
 
+  def test_v_mullit_f32_basic(self):
+    """V_MULLIT_F32 multiplies when the guard input is valid."""
+    instructions = [
+      v_mov_b32_e32(v[0], f2i(2.0)),
+      v_mov_b32_e32(v[1], f2i(3.0)),
+      v_mov_b32_e32(v[2], f2i(1.0)),
+      v_mullit_f32(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertAlmostEqual(i2f(st.vgpr[0][3]), 6.0, places=5)
+
+  def test_v_mullit_f32_invalid_guard(self):
+    """V_MULLIT_F32 returns -MAX_FLOAT_F32 when the guard input is non-positive."""
+    instructions = [
+      v_mov_b32_e32(v[0], f2i(2.0)),
+      v_mov_b32_e32(v[1], f2i(3.0)),
+      v_mov_b32_e32(v[2], f2i(0.0)),
+      v_mullit_f32(v[3], v[0], v[1], v[2]),
+    ]
+    st = run_program(instructions, n_lanes=1)
+    self.assertEqual(st.vgpr[0][3], 0xFF7FFFFF)
+
   def test_v_fma_f32_negative(self):
     """V_FMA_F32 with negative multiplier."""
     instructions = [
@@ -1592,8 +1614,7 @@ class TestModifierInteractions(unittest.TestCase):
     self.assertEqual(st.vgpr[0][2], 0x80000000, "-|(-0.0)| = -0.0")
 
   def test_clamp_with_nan(self):
-    """Clamp with NaN input should still produce NaN."""
-    import math
+    """Clamp with NaN input saturates to 0 on RDNA3 hardware."""
     quiet_nan = 0x7fc00000
     instructions = [
       s_mov_b32(s[0], quiet_nan),
@@ -1601,7 +1622,7 @@ class TestModifierInteractions(unittest.TestCase):
       VOP3(VOP3Op.V_ADD_F32, vdst=v[1], src0=v[0], src1=0.0, clmp=1),
     ]
     st = run_program(instructions, n_lanes=1)
-    self.assertTrue(math.isnan(i2f(st.vgpr[0][1])))
+    self.assertEqual(st.vgpr[0][1], 0)
 
   def test_omod_ignored(self):
     """OMOD field is ignored on RDNA3 hardware."""
@@ -3605,32 +3626,30 @@ class TestPermlane(unittest.TestCase):
     """V_PERMLANE16_B32 broadcast lane 0 to all lanes in row."""
     # lanesel = all zeros -> all positions read from lane 0 within row
     instructions = [
-      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source data
+      v_mov_b32_e32(v[0], v[255]),
       s_mov_b32(s[0], 0),                # lanesel low = 0 (all read lane 0)
       s_mov_b32(s[1], 0),                # lanesel high = 0
       v_permlane16_b32(v[1], v[0], s[0], s[1]),
     ]
-    st = run_program(instructions, n_lanes=4)
-    # All lanes read from lane 0 of their row
-    for lane in range(4):
-      self.assertEqual(st.vgpr[lane][1], 0xCAFEBABE)
+    st = run_program(instructions, n_lanes=32)
+    for lane in range(16): self.assertEqual(st.vgpr[lane][1], 0)
+    for lane in range(16, 32): self.assertEqual(st.vgpr[lane][1], 16)
 
   def test_v_permlanex16_b32_identity(self):
     """V_PERMLANEX16_B32 cross-row read with identity selection."""
-    # In wave32: row 0 (lanes 0-15) reads from row 1 (lanes 16-31) and vice versa
-    # With single lane in row 0, it reads from lane 0 of row 1 (lane 16)
-    # But lane 16 doesn't exist in 1-lane test, so use 32 lanes
     instructions = [
-      v_mov_b32_e32(v[0], 0x11111111),  # All lanes have this initially
+      v_mov_b32_e32(v[0], v[255]),
       s_mov_b32(s[0], 0x76543210),       # lanesel low
       s_mov_b32(s[1], 0xFEDCBA98),       # lanesel high
       v_permlanex16_b32(v[1], v[0], s[0], s[1]),
     ]
     st = run_program(instructions, n_lanes=32)
-    # Lane 0 in row 0 reads from lane 0 of row 1 (lane 16)
-    self.assertEqual(st.vgpr[0][1], 0x11111111)
-    # Lane 16 in row 1 reads from lane 0 of row 0 (lane 0)
-    self.assertEqual(st.vgpr[16][1], 0x11111111)
+    self.assertEqual(st.vgpr[0][1], 16)
+    self.assertEqual(st.vgpr[5][1], 21)
+    self.assertEqual(st.vgpr[15][1], 31)
+    self.assertEqual(st.vgpr[16][1], 0)
+    self.assertEqual(st.vgpr[21][1], 5)
+    self.assertEqual(st.vgpr[31][1], 15)
 
 
 if __name__ == '__main__':

test/amd/test_compare_emulators.py

@@ -1,528 +0,0 @@
-# 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 test.mockgpu.amd.emu import WaveState, _decode_at, WAVE_SIZE, VCC_LO, EXEC_LO, SCC
-from tinygrad.renderer.amd import decode_inst
-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
-
-def _is_f32_nan(bits: int) -> bool:
-  """Check if 32-bit value is a NaN (exponent all 1s, mantissa non-zero)."""
-  return (bits & 0x7f800000) == 0x7f800000 and (bits & 0x007fffff) != 0
-
-def _vals_equal(a: int, b: int) -> bool:
-  """Compare two 32-bit values, treating all NaN bit patterns as equal."""
-  if a == b: return True
-  return _is_f32_nan(a) and _is_f32_nan(b)
-
-@dataclass
-class KernelSnapshot:
-  code: bytes
-  src: str
-  global_size: tuple[int, int, int]
-  local_size: tuple[int, int, int]
-  buf_idxs: list[int]  # indices into shared buffer pool
-  buf_sizes: list[int]  # sizes for each buffer index
-
-@dataclass
-class StateSnapshot:
-  pc: int
-  scc: int
-  vcc: int
-  exec_mask: int
-  sgpr: list[int]
-  vgpr: list[list[int]]
-
-  def diff(self, other: 'StateSnapshot', n_lanes: int, arrow: str = " vs ") -> list[str]:
-    """Return list of differences between two states."""
-    diffs = []
-    if self.pc != other.pc: diffs.append(f"pc: {self.pc}{arrow}{other.pc}")
-    if self.scc != other.scc: diffs.append(f"scc: {self.scc}{arrow}{other.scc}")
-    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x}{arrow}0x{other.vcc:08x}")
-    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x}{arrow}0x{other.exec_mask:08x}")
-    for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
-      # Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
-      if i in (106, 107, 126, 127): continue
-      if not _vals_equal(a, b): diffs.append(f"sgpr[{i}]: 0x{a:08x}{arrow}0x{b:08x}")
-    for lane in range(n_lanes):
-      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
-        if not _vals_equal(a, b): diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x}{arrow}0x{b:08x}")
-    return diffs
-
-class CStateSnapshot(ctypes.Structure):
-  _fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
-              ("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
-
-  def to_snapshot(self) -> StateSnapshot:
-    return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
-                         sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
-
-class RustEmulator:
-  def __init__(self):
-    self.lib = ctypes.CDLL(str(REMU_PATH))
-    self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_create.restype = ctypes.c_void_p
-    self.lib.wave_step.argtypes = [ctypes.c_void_p]
-    self.lib.wave_step.restype = ctypes.c_int32
-    self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
-    self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
-    self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
-    self.lib.wave_free.argtypes = [ctypes.c_void_p]
-    self.ctx = None
-
-  def create(self, kernel: bytes, n_lanes: int):
-    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
-    self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
-    self._kernel_buf = kernel_buf
-
-  def step(self) -> int: return self.lib.wave_step(self.ctx)
-  def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
-  def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
-  def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
-
-  def get_snapshot(self) -> StateSnapshot:
-    snap = CStateSnapshot()
-    self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
-    return snap.to_snapshot()
-
-  def free(self):
-    if self.ctx:
-      self.lib.wave_free(self.ctx)
-      self.ctx = None
-
-class PythonEmulator:
-  def __init__(self):
-    self.state: WaveState | 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
-    self.lib_addr = 0  # Base address of kernel code
-
-  def create(self, kernel: bytes, n_lanes: int):
-    import ctypes
-    from tinygrad.device import Buffer, BufferSpec
-    from tinygrad.dtype import dtypes
-    # 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)
-    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.state is not None
-    pc = self.state.pc
-    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]
-    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
-
-  def set_sgpr(self, idx: int, val: int):
-    assert self.state is not None
-    self.state._write_sgpr(idx, val)
-  def set_vgpr(self, lane: int, idx: int, val: int):
-    assert self.state is not None
-    self.state._write_vgpr(idx, lane, val)
-
-  def get_snapshot(self) -> StateSnapshot:
-    assert self.state is not None
-    sgpr = [self.state._read_sgpr(i) for i in range(128)]
-    vgpr = [[self.state._read_vgpr(reg, lane) for reg in range(256)] for lane in range(WAVE_SIZE)]
-    # Convert actual PC address to word offset for comparison with Rust emulator
-    pc_offset = (self.state.pc - self.lib_addr) // 4 if self.state.pc != 0xFFFFFFFFFFFFFFFF else 0xFFFFFFFFFFFFFFFF
-    return StateSnapshot(pc=pc_offset, scc=self.state._read_sgpr(SCC.offset), vcc=sgpr[VCC_LO.offset],
-                         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], 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
-  lx, ly, lz = local_size
-  total_steps = 0
-  wg_count = 0
-
-  for gidz in range(gz):
-    for gidy in range(gy):
-      for gidx in range(gx):
-        if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
-        wg_count += 1
-        rust = RustEmulator()
-        python = PythonEmulator()
-        rust.create(kernel, n_lanes)
-        python.create(kernel, n_lanes)
-
-        # Initialize LDS (64KB, standard size for AMD GPUs)
-        rust.init_lds(65536)
-
-        for emu in (rust, python):
-          emu.set_sgpr(0, args_ptr & 0xffffffff)
-          emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
-          emu.set_sgpr(13, gidx)
-          emu.set_sgpr(14, gidy)
-          emu.set_sgpr(15, gidz)
-          # Initialize v[0] with packed workitem IDs for each lane
-          for lane in range(n_lanes):
-            tid = lane
-            z, y, x = tid // (lx * ly), (tid // lx) % ly, tid % lx
-            emu.set_vgpr(lane, 0, (z << 20) | (y << 10) | x)
-
-        step = 0
-        trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
-        prev_sync_after = False  # Track if previous instruction had known Rust bugs
-        try:
-          while step < max_steps:
-            rust_before = rust.get_snapshot()
-            python_before = python.get_snapshot()
-
-            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
-              parts = inst_hex_name.rsplit('_', 1)
-              inst_bytes_hex = parts[1] if len(parts) == 2 else ""
-              inst_bytes = bytes.fromhex(inst_bytes_hex) if inst_bytes_hex else b''
-              decoded = decode_inst(inst_bytes) if inst_bytes else None
-              inst_mnemonic = repr(decoded).split('(')[0] if decoded else ""
-            except Exception:
-              inst_mnemonic = ""
-            # For generic instructions, use function name for sync_after check
-            if not inst_mnemonic: inst_mnemonic = inst_hex_name
-            inst_str = inst_hex_name
-            trace.append((step, python_before.pc, inst_str, rust_before, python_before))
-            if len(trace) > trace_len: trace.pop(0)
-
-            if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
-
-            # Instructions with known Rust emulator bugs or precision differences - sync Python to Rust after execution
-            # v_div_scale/v_div_fixup: Rust has different VCC handling
-            # v_cvt_f16_f32: Rust clears high 16 bits, but hardware (and Python) preserves them
-            # s_add_i32/s_sub_i32: Rust has incorrect SCC overflow detection
-            # v_exp_f32/v_log_f32/v_ldexp_f32: precision differences in transcendental functions
-            # s_delay_alu: Rust handles differently
-            # v_add_co_ci_u32/v_sub_co_ci_u32/v_subrev_co_ci_u32: Rust preserves inactive VCC bits, but hardware clears all bits
-            sync_after = any(x in inst_mnemonic.lower() for x in ('v_div_scale', 'v_div_fixup', 'v_cvt_f16_f32', 's_add_i32', 's_sub_i32',
-                                                                   'v_exp_f32', 'v_log_f32', 'v_ldexp_f32', 's_delay_alu',
-                                                                   'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'))
-            # Skip comparison if previous instruction had known Rust bugs (states were synced but may still differ slightly)
-            diffs = rust_before.diff(python_before, n_lanes) if not prev_sync_after else []
-            if diffs:
-              trace_lines = []
-              for idx, (s, pc, d, rb, pb) in enumerate(trace):
-                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
-                if idx < len(trace) - 1:
-                  next_rb, next_pb = trace[idx + 1][3:5]
-                  rust_diffs = rb.diff(next_rb, n_lanes, "->")
-                  python_diffs = pb.diff(next_pb, n_lanes, "->")
-                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
-                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append("             python: (no changes)")
-                else:
-                  # Last traced instruction - compare with current state
-                  rust_diffs = rb.diff(rust_before, n_lanes, "->")
-                  python_diffs = pb.diff(python_before, n_lanes, "->")
-                  if rust_diffs: trace_lines.append(f"             rust:   {', '.join(rust_diffs[:5])}")
-                  if python_diffs: trace_lines.append(f"             python: {', '.join(python_diffs[:5])}")
-                  elif rust_diffs: trace_lines.append("             python: (no changes)")
-              trace_str = "\n".join(trace_lines)
-              msg = f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ (rust vs python):\n  "
-              msg += "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}"
-              return False, msg, total_steps
-
-            rust_result = rust.step()
-            python_result = python.step()
-
-            if rust_result != python_result:
-              # Rust returns 1 for unsupported instructions - skip test
-              if rust_result == 1 and python_result == 0:
-                raise unittest.SkipTest(f"Rust emulator doesn't support instruction: {inst_str}")
-              trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
-              msg = (f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: "
-                     f"rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}")
-              return False, msg, total_steps
-
-            # Sync Python state to Rust after instructions with known Rust emulator differences
-            if sync_after:
-              rust_after = rust.get_snapshot()
-              for i in range(128): python.set_sgpr(i, rust_after.sgpr[i])
-              for lane in range(n_lanes):
-                for i in range(256): python.set_vgpr(lane, i, rust_after.vgpr[lane][i])
-              assert python.state is not None
-              # Convert Rust's word-based PC to Python's actual address
-              python.state.pc = python.lib_addr + rust_after.pc * 4
-              python.state._write_sgpr(SCC.offset, rust_after.scc)
-              python.state._write_sgpr(VCC_LO.offset, rust_after.vcc)
-              python.state._write_sgpr(EXEC_LO.offset, rust_after.exec_mask)
-            prev_sync_after = sync_after
-
-            if rust_result == -1:
-              total_steps += step + 1
-              break
-            if rust_result == 1:
-              total_steps += step + 1
-              break
-            if rust_result < 0 and rust_result != -2:
-              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
-
-            step += 1
-          else:
-            return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
-        finally:
-          rust.free()
-
-  return True, f"Completed {gx*gy*gz} workgroups", total_steps
-
-def compare_emulators_multi_kernel(kernels: list[KernelSnapshot], buf_pool: dict[int, int], max_steps: int = 1000,
-                                    debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
-  """Run all kernels through both emulators with shared buffer pool."""
-  if buf_data is None: buf_data = {}
-
-  # Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
-  buf_id_to_ptr: dict[int, int] = {}
-  buffers = []
-  for buf_id, size in buf_pool.items():
-    padded_size = ((size + 15) // 16) * 16 + 16  # round up to 16 bytes + extra padding
-    # Initialize with data from COPY if available
-    init_data = buf_data.get(buf_id, b'\x00' * padded_size)
-    init_list = list(init_data) + [0] * (padded_size - len(init_data))
-    buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
-    buffers.append((buf, padded_size))
-    buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
-
-  # Set up valid memory ranges
-  ranges = {(ctypes.addressof(b), size) for b, size in buffers}
-
-  total_steps = 0
-  for ki, kernel in enumerate(kernels):
-    # Create args array for this kernel's buffers
-    args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
-    args_ptr = ctypes.addressof(args)
-
-    # Update valid ranges to include this args array
-    kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
-    set_valid_mem_ranges(kernel_ranges)
-
-    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, max_steps, debug, trace_len, ki
-    )
-    total_steps += steps
-    if not ok:
-      return False, msg
-
-  return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
-
-def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
-                                   global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
-  """Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
-  # Allocate buffers
-  buffers = []
-  for size in buf_sizes:
-    buf = (ctypes.c_uint8 * size)(*[0] * size)
-    buffers.append(buf)
-
-  # Create args array with buffer pointers
-  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
-  args_ptr = ctypes.addressof(args)
-
-  # Set up valid memory ranges for Python emulator
-  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
-  ranges.add((args_ptr, ctypes.sizeof(args)))
-  set_valid_mem_ranges(ranges)
-
-  # 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), max_steps, debug, trace_len)
-  return ok, msg
-
-def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
-  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
-  from tinygrad import Tensor
-  from tinygrad.runtime.support.elf import elf_loader
-
-  out = op_fn(Tensor)
-  sched = out.schedule()
-  kernels = []
-  buf_pool: dict[int, int] = {}  # buffer id -> size
-  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
-
-  for ei in sched:
-    lowered = ei.lower()
-    if ei.ast.op.name == 'COPY':
-      # Handle COPY: extract source data to initialize destination buffer
-      if len(lowered.bufs) >= 2:
-        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
-        dst_id = id(dst_buf)
-        if dst_id not in buf_pool:
-          buf_pool[dst_id] = dst_buf.nbytes
-        # Get source data if it's from numpy/CPU
-        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
-          src_data = bytes(src_buf.base._buf)
-          buf_data[dst_id] = src_data
-    elif ei.ast.op.name == 'SINK':
-      if lowered.prg and lowered.prg.p.lib:
-        lib = bytes(lowered.prg.p.lib)
-        _, sections, _ = elf_loader(lib)
-        for sec in sections:
-          if sec.name == '.text':
-            buf_idxs = []
-            buf_sizes = []
-            for b in lowered.bufs:
-              buf_id = id(b)
-              if buf_id not in buf_pool:
-                buf_pool[buf_id] = b.nbytes
-              buf_idxs.append(buf_id)
-              buf_sizes.append(b.nbytes)
-            kernels.append(KernelSnapshot(
-              code=bytes(sec.content),
-              src=lowered.prg.p.src,
-              global_size=tuple(lowered.prg.p.global_size),
-              local_size=tuple(lowered.prg.p.local_size),
-              buf_idxs=buf_idxs,
-              buf_sizes=buf_sizes
-            ))
-  if not kernels: raise RuntimeError("No kernel found")
-  return kernels, buf_pool, buf_data
-
-def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
-  """Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
-  kernels, _, _ = get_kernels_from_tinygrad(op_fn)
-  k = kernels[-1]
-  return k.code, k.global_size, k.local_size, k.buf_sizes
-
-@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
-class TestTinygradKernels(unittest.TestCase):
-  """Compare emulators on real tinygrad-compiled kernels."""
-
-  def _test_kernel(self, op_fn, max_steps=10000):
-    kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
-    ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
-    self.assertTrue(ok, msg)
-
-  # Basic ops - consolidated tests covering key instruction patterns
-  def test_unary_ops(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu().exp().log().sqrt().reciprocal())
-  def test_binary_ops(self): self._test_kernel(lambda T: (T([1.0, 2.0]) + T([3.0, 4.0])) * T([0.5, 0.5]) - T([1.0, 1.0]))
-  def test_trig(self): self._test_kernel(lambda T: T([0.1, 1.0, 3.14, -1.0]*8).sin() + T([0.1, 1.0, 3.14, -1.0]*8).cos())
-  def test_compare(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
-  def test_bitwise(self): self._test_kernel(lambda T: (T([0xF0, 0x0F, 0xFF]*11).int() & T([0x0F, 0x0F, 0x00]*11).int()) | T([1]*33).int())
-  def test_int_ops(self): self._test_kernel(lambda T: ((T.empty(64).int() + T.empty(64).int()) * T.empty(64).int()).float())
-
-  # Reductions
-  def test_reduce(self): self._test_kernel(lambda T: T.empty(64).sum() + T.empty(64).max())
-  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
-
-  # Matmul
-  def test_gemm(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=100000)
-  @unittest.skip("Rust emulator crashes on this kernel (assertion failure in thread.rs)")
-  def test_gemm_fp16(self): self._test_kernel(lambda T: T.empty(16, 16).half() @ T.empty(16, 16).half(), max_steps=100000)
-
-  # Complex ops
-  def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
-  def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
-
-  # Memory patterns
-  def test_memory(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous() + T.empty(4, 1).expand(4, 4))
-
-  # Cast ops
-  def test_cast(self): self._test_kernel(lambda T: T.empty(32).half().float() + T.empty(32).int().float())
-
-  # Pooling - regression for VCC wave32 mode
-  def test_pool2d(self):
-    self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4)) + T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4)))
-
-  # Convolution
-  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 2, 8, 8).conv2d(T.empty(2, 2, 3, 3)), max_steps=50000)
-
-  # Regression tests
-  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
-  def test_interpolate(self): self._test_kernel(lambda T: T.empty(1,2,16,16).relu().cast('uint8').interpolate((8,8), mode="linear"))
-  def test_index_int64(self):
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
-  def test_gelu(self): self._test_kernel(lambda T: T.empty(32, 32).gelu())
-  def test_exp(self): self._test_kernel(lambda T: T.empty(1024).exp())
-  def test_cross_entropy(self):
-    import numpy as np
-    np.random.seed(0)
-    classes = np.random.randint(0, 10, (16,), dtype=np.int32).tolist()
-    x_np = np.random.randn(16, 10).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(16,10) + 0).cross_entropy((T(classes).int().reshape(16) + 0)))
-  def test_isinf(self): self._test_kernel(lambda T: T([float('-inf'), 0., float('inf'), 1.1]*8).isinf())
-  def test_sin_f64(self):
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T([2.0], dtype=dtypes.float64).sin())
-
-  def test_sin_large_f32(self):
-    """Test sin with large values that trigger Payne-Hanek range reduction."""
-    # Values around 859240 trigger the Payne-Hanek algorithm
-    # This tests the integer multiply-high instructions used in range reduction
-    self._test_kernel(lambda T: T([859240.0, 1000000.0, 100594688.0]).sin())
-
-  def test_clip_zero_one(self):
-    """Test clip(0, 1) - regression for binary_crossentropy failure."""
-    import numpy as np
-    np.random.seed(0)
-    x_np = np.random.uniform(-2, 2, (32, 10)).astype(np.float32).tolist()
-    self._test_kernel(lambda T: T(x_np).clip(0, 1))
-
-  def test_mod_int64(self):
-    """Test int64 modulo, especially edge cases like 1 % -1."""
-    from tinygrad import dtypes
-    self._test_kernel(lambda T: T([1, 10, -10, 7], dtype=dtypes.int64) % T([-1, 3, 3, -3], dtype=dtypes.int64))
-
-  def test_expand_flatten_sum(self):
-    """Test flatten of expanded tensor followed by sum.
-
-    Bug: flatten() of an expanded tensor produces wrong results for certain sizes.
-    Sizes that are multiples of 32 work (32, 48, 64), but sizes like 33, 49, 50 fail.
-    This breaks masked_select and nonzero operations.
-    """
-    import numpy as np
-    np.random.seed(0)
-    x_np = np.random.uniform(-2, 2, (33,)).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).unsqueeze(-1).expand(33, 3).flatten().sum())
-
-  @unittest.skip("slow and broken with AMD:LLVM")
-  def test_nonzero(self):
-    """Test nonzero operation - counts and gathers indices of non-zero elements."""
-    import numpy as np
-    np.random.seed(42)
-    x_np = np.random.rand(10, 5, 3).astype(np.float32)
-    self._test_kernel(lambda T: (T(x_np.tolist()) > 0.5).nonzero())
-
-  @unittest.skip("Precision differences in v_exp/v_log accumulate across kernels, causing memory divergence")
-  def test_softmax_argmax_fused(self):
-    """Test fused softmax+argmax - tracks exp2 precision issue.
-
-    The fused kernel recomputes softmax inline and Python emulator's exp2 polynomial
-    has up to 1 ULP error vs native exp2f, causing accumulated differences.
-    """
-    import torch
-    torch.manual_seed(0)
-    x_np = torch.rand(4, 10).numpy()
-    self._test_kernel(lambda T: T(x_np.tolist()).softmax(1).argmax())
-
-if __name__ == "__main__":
-  unittest.main()

test/amd/test_custom_kernel.py

@@ -104,20 +104,21 @@ def custom_handwritten(A:UOp, arch:str) -> UOp:
   lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
   k = Kernel(arch)
   k.emit(r4.s_nop(0))
-  k.emit(r4.v_mov_b32_e32(v[1], 10))
+  k.emit(r4.v_mov_b32_e32(v[1], 4))
   def emit_alt():
-    for i in range(4):
+    for i in range(2):
       k.emit(r4.v_mov_b32_e32(v[20+i], 4.0))
       k.emit(r4.v_rcp_f32_e32(v[22+i], v[20+i]))
       k.emit(r4.s_mov_b32(s[20+i], i))
+      k.emit(r4.s_mul_i32(s[14+i], s[12+i], 32))
   def emit_wmma():
-    for _ in range(4):
+    for _ in range(2):
       k.emit(r4.v_wmma_f32_16x16x16_f16(v[0:7], v[8:11], v[8:11], 1))
   k.label("start")
   k.emit(s_mov_b32(s[1], 10))
   k.label("loop")
   # wmma should've overlapped here if it was a different unit?
-  for _ in range(4):
+  for _ in range(2):
     emit_wmma()
     emit_alt()
   for _ in range(8): k.emit(s_nop(1))

test/amd/test_emu2_pcode.py

@@ -7,7 +7,7 @@ from tinygrad.uop.ops import UOp, Ops
 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
+from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp, GLOBALOp
 
 def _srcs():
   """Create minimal source variables for pcode parsing."""
@@ -113,6 +113,7 @@ class TestParseExpr(unittest.TestCase):
     result = parse_expr('cond ? a : b', vrs)
     self.assertEqual(result.op, Ops.WHERE)
 
+
 class TestForLoopParsing(unittest.TestCase):
   """Test for loop parsing (CLZ/CTZ patterns)."""
 
@@ -164,6 +165,20 @@ class TestForLoopParsing(unittest.TestCase):
 class TestDSPcodePatterns(unittest.TestCase):
   """Test DS instruction pcode patterns."""
 
+  def test_global_atomic_add_f32_parsing(self):
+    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
+    vmem = UOp(Ops.PARAM, dtypes.uint32.ptr(1024), arg=2)
+    srcs = {
+      'ADDR': UOp.const(dtypes.uint64, 0),
+      'DATA': UOp.const(dtypes.uint32, 0x3f800000),
+      '_vmem': vmem,
+    }
+
+    _, assigns = parse_pcode(PCODE[GLOBALOp.GLOBAL_ATOMIC_ADD_F32], srcs)
+    mem_write = next(val for dest, val in assigns if dest == 'MEM[ADDR].f32')
+    self.assertEqual(mem_write[1].op, Ops.ADD)  # type: ignore[index]
+    self.assertEqual(mem_write[1].dtype, dtypes.float32)  # type: ignore[index]
+
   def test_ds_load_b32_pcode(self):
     """Test DS_LOAD_B32 pcode is parseable."""
     pcode = PCODE.get(DSOp.DS_LOAD_B32)
@@ -285,6 +300,47 @@ class TestConditionalParsing(unittest.TestCase):
     # Result should be a WHERE (ternary becomes WHERE)
     self.assertEqual(val.op, Ops.WHERE)
 
+class TestConcatWidthParsing(unittest.TestCase):
+  """Test that bit extracts keep the right width for concat/unary ops."""
+
+  def test_permlanex16_altrow_concat(self):
+    for row, expected in [(0, 1), (1, 0), (2, 3), (3, 2)]:
+      parsed = parse_expr('{ row[1], ~row[0] }', {'row': UOp.const(dtypes.uint32, row)})
+      self.assertEqual(parsed.simplify().arg, expected)
+
+  def test_permlane64_altlane_concat(self):
+    for lane, expected in [(0, 32), (1, 33), (31, 63), (32, 0), (63, 31)]:
+      parsed = parse_expr('{ ~lane[5], lane[4:0] }', {'lane': UOp.const(dtypes.uint32, lane)})
+      self.assertEqual(parsed.simplify().arg, expected)
+
+  def test_permlane64_wave64_pcode_indices(self):
+    vgpr = UOp(Ops.PARAM, dtypes.uint32.ptr(256), arg=0)
+    srcs = {
+      'SRC0': UOp.const(dtypes.uint32, 0),
+      'VDST': UOp.const(dtypes.uint32, 1),
+      'EXEC_LO': UOp.const(dtypes.uint32, 0xFFFFFFFF),
+      'EXEC': UOp.const(dtypes.uint64, 0xFFFFFFFFFFFFFFFF),
+      '_vgpr': vgpr,
+      '_wave_size': 64,
+      'S0': UOp.const(dtypes.uint32, 0),
+      'S1': UOp.const(dtypes.uint32, 0),
+      'S2': UOp.const(dtypes.uint32, 0),
+    }
+
+    def load_idx(v: UOp) -> int:
+      simp = v.simplify()
+      self.assertEqual(simp.op, Ops.LOAD)
+      self.assertEqual(simp.src[0].op, Ops.INDEX)
+      idx = simp.src[0].src[1].simplify()
+      self.assertEqual(idx.op, Ops.CONST)
+      return idx.arg
+
+    _, assigns = parse_pcode(PCODE[VOP1Op.V_PERMLANE64_B32_E32], srcs)
+    self.assertEqual(len(assigns), 64)
+    for lane, (dst_idx, src_idx) in {0: (64, 32), 31: (95, 63), 32: (96, 0), 63: (127, 31)}.items():
+      self.assertEqual(assigns[lane][1][0].simplify().arg, dst_idx)  # type: ignore[index]
+      self.assertEqual(load_idx(assigns[lane][1][1]), src_idx)  # type: ignore[index]
+
 class TestAllPcode(unittest.TestCase):
   """Test that all pcode from all architectures can be parsed."""
 

test/amd/test_mockgpu_invalid.py

@@ -39,7 +39,6 @@ dev.synchronize()
     env = os.environ.copy()
     env["AMD"] = "1"
     env["MOCKGPU"] = "1"
-    env["PYTHON_REMU"] = "1"
     env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
 
     st = time.perf_counter()

test/amd/test_roundtrip.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
 import unittest, io, sys, re
+from dataclasses import dataclass
 from tinygrad import Device
 from tinygrad.renderer.amd import detect_format
 from test.amd.helpers import llvm_assemble, llvm_disasm, get_target, get_mattr
@@ -44,6 +45,64 @@ def compile_and_disasm_batch(instrs: list[str], arch: str = 'rdna3') -> list[str
   code = b''.join(llvm_assemble(instrs, mcpu, mattr))
   return llvm_disasm(code, mcpu, mattr)[:len(instrs)]
 
+@dataclass
+class KernelSnapshot:
+  code: bytes
+  src: str
+  global_size: tuple[int, int, int]
+  local_size: tuple[int, int, int]
+  buf_idxs: list[int]  # indices into shared buffer pool
+  buf_sizes: list[int]  # sizes for each buffer index
+
+def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, int], dict[int, bytes]]:
+  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
+  from tinygrad import Tensor
+  from tinygrad.runtime.support.elf import elf_loader
+
+  out = op_fn(Tensor)
+  sched = out.schedule()
+  kernels = []
+  buf_pool: dict[int, int] = {}  # buffer id -> size
+  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
+
+  for ei in sched:
+    lowered = ei.lower()
+    if ei.ast.op.name == 'COPY':
+      # Handle COPY: extract source data to initialize destination buffer
+      if len(lowered.bufs) >= 2:
+        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
+        dst_id = id(dst_buf)
+        if dst_id not in buf_pool:
+          buf_pool[dst_id] = dst_buf.nbytes
+        # Get source data if it's from numpy/CPU
+        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
+          src_data = bytes(src_buf.base._buf)
+          buf_data[dst_id] = src_data
+    elif ei.ast.op.name == 'SINK':
+      if lowered.prg and lowered.prg.p.lib:
+        lib = bytes(lowered.prg.p.lib)
+        _, sections, _ = elf_loader(lib)
+        for sec in sections:
+          if sec.name == '.text':
+            buf_idxs = []
+            buf_sizes = []
+            for b in lowered.bufs:
+              buf_id = id(b)
+              if buf_id not in buf_pool:
+                buf_pool[buf_id] = b.nbytes
+              buf_idxs.append(buf_id)
+              buf_sizes.append(b.nbytes)
+            kernels.append(KernelSnapshot(
+              code=bytes(sec.content),
+              src=lowered.prg.p.src,
+              global_size=tuple(lowered.prg.p.global_size),
+              local_size=tuple(lowered.prg.p.local_size),
+              buf_idxs=buf_idxs,
+              buf_sizes=buf_sizes
+            ))
+  if not kernels: raise RuntimeError("No kernel found")
+  return kernels, buf_pool, buf_data
+
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestTinygradKernelRoundtrip(unittest.TestCase):
   """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
@@ -57,7 +116,6 @@ class TestTinygradKernelRoundtrip(unittest.TestCase):
     """
     arch = self.arch
 
-    from test.amd.test_compare_emulators import get_kernels_from_tinygrad
     from tinygrad.runtime.support.elf import elf_loader
     from tinygrad.runtime.support.compiler_amd import HIPCompiler, AMDLLVMCompiler
     from tinygrad.helpers import DEV

test/amd/test_sqtt_profiler.py

@@ -1,15 +1,16 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit
 from tinygrad.device import Compiled, ProfileProgramEvent, ProfileDeviceEvent
-from tinygrad.viz.serve import load_amd_counters
+from tinygrad.viz.serve import load_amd_counters, VizData
 
 @contextlib.contextmanager
 def save_sqtt():
-  yield (ret:=[])
+  data = VizData()
+  yield data.ctxs
   Device[Device.DEFAULT].synchronize()
   Device[Device.DEFAULT]._at_profile_finalize()
-  load_amd_counters(ret, Compiled.profile_events)
-  ret[:] = [r for r in ret if r["name"].startswith("SQTT")]
+  load_amd_counters(data, Compiled.profile_events)
+  data.ctxs[:] = [r for r in data.ctxs if r["name"].startswith("SQTT")]
 
 @unittest.skipUnless(Device.DEFAULT == "AMD", "only runs on AMD")
 class TestSQTTProfiler(unittest.TestCase):

test/amd/test_sqttmap.py

@@ -100,9 +100,7 @@ class TestSQTTMapBase(unittest.TestCase):
             elif "WAVE" in e.device:
               # sopk/immediates don't get ALU/MEM EXEC
               if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
-                                             "WAVEEND", "WAVERDY"}: insts += 1
-            # OTHER_ is its own stream, it's the INST from other SIMDs that share the same EXEC.
-            elif e.device.startswith("OTHER"): continue
+                                             "WAVEEND", "WAVERDY"} and not e.name.display_name.startswith("OTHER_"): insts += 1
             else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
           self.assertEqual(execs, insts)
 
@@ -131,7 +129,18 @@ class TestSQTTMapBase(unittest.TestCase):
 
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
-class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
+class TestSQTTMapRDNA4(TestSQTTMapBase):
+  target = "gfx1200"
+
+  @unittest.expectedFailure
+  def test_rdna4_wmma(self):
+    events, kernels, target = self.examples["profile_handwritten_run_0"]
+    row_ends = {}
+    for e in sqtt_timeline(events[0].blob, list(kernels.values())[0].lib, target):
+      if type(e).__name__ != "ProfileRangeEvent" or e.device != "ALUEXEC:0 WMMA": continue
+      if (et:=row_ends.get(e.device)) is not None and e.st < et:
+        raise RuntimeError(f"WMMA exec overlaps in {e.device}: {e.st} {et}.")
+      row_ends[e.device] = e.en
 
 class TestSQTTMapCDNA(TestSQTTMapBase):
   target = "gfx950"

test/backend/test_arange.py

@@ -4,8 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
-from tinygrad.uop.ops import Ops
+from tinygrad.schedule import ExecItem
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu
@@ -24,6 +23,10 @@ class TestArange(unittest.TestCase):
     self.assertEqual(self._get_flops(Tensor.arange(256), np.arange(256)), 0)
     self.assertEqual(self._get_flops(Tensor.arange(2560), np.arange(2560)), 0)
 
+  @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
+  def test_arange_cumsum(self):
+    np.testing.assert_equal(Tensor.arange(513).cumsum(0).numpy(), np.arange(513).cumsum())
+
   def test_arange_cat(self):
     t = Tensor.arange(2, dtype=dtypes.int)+Tensor([3])
     self.assertEqual(t.cat(t).tolist(), [3, 4, 3, 4])
@@ -63,7 +66,7 @@ class TestIndexing(unittest.TestCase):
     print("*** indexing ***")
     with Context(NOOPT=1):
       GlobalCounters.reset()
-      rng = Tensor.ones(4, DDIM, DSET, dtype=dtypes.int)._cumalu(axis=-1, op=Ops.ADD, _include_initial=True).reshape(4, DDIM, DSET, 1)
+      rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
       idxs = idxs.reshape(4,1,1,1).expand(4, DDIM, DSET, 1)
       reshape_dataset = dataset.T.reshape(1, DDIM, DSET, 1).expand(4, DDIM, DSET, 1)
       full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, DDIM, DSET, 1))

test/backend/test_asm_gemm.py

@@ -21,9 +21,8 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
 
   a, b = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
   if multi: a, b = a.shard(devs, axis=a_shard), b.shard(devs, axis=b_shard)
-  with Context(ASM_GEMM=1):
-    tst = asm_gemm(a, b)
-    tst.sum().backward()
+  tst = asm_gemm(a, b)
+  tst.sum().backward()
   Tensor.realize(tst, a.grad, b.grad)
 
   a_ref, b_ref = a_rand.clone().requires_grad_(), b_rand.clone().requires_grad_()
@@ -32,9 +31,8 @@ def run_asm_gemm(a_shape, b_shape, dtype=dtypes.float16, a_shard=None, b_shard=N
     a_ref = a_ref.cast(dtypes.bfloat16)
     b_ref = b_ref.cast(dtypes.bfloat16)
   if multi: a_ref, b_ref = a_ref.shard(devs, axis=a_shard), b_ref.shard(devs, axis=b_shard)
-  with Context(ASM_GEMM=0):
-    ref = asm_gemm(a_ref, b_ref)
-    ref.sum().backward()
+  ref = a_ref @ b_ref
+  ref.sum().backward()
   Tensor.realize(ref, a_ref.grad, b_ref.grad)
 
   # no validation on the NULL device
@@ -136,14 +134,12 @@ class TestGemmLlama(unittest.TestCase):
     if not is_cdna4() or getenv("MOCKGPU"):
       self.skipTest("very slow on non mi350x")
 
-  @Context(ASM_GEMM=1)
-  def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype)@Tensor.empty(N, N, dtype=self.dtype)).realize()
+  def test_empty(self): asm_gemm(Tensor.empty(N:=getenv("N", 4096), N, dtype=self.dtype), Tensor.empty(N, N, dtype=self.dtype)).realize()
 
-  @Context(ASM_GEMM=1)
   def test_empty_bw(self):
     x = Tensor.empty(1, N:=getenv("N", 4096), N, dtype=self.dtype, requires_grad=True)
     y = Tensor.empty((N, N), dtype=self.dtype, requires_grad=True)
-    z = x @ y
+    z = asm_gemm(x, y)
     z.sum().backward()
     Tensor.realize(z, x.grad, y.grad)
     # FP8 forward output is bf16, gradients use fp8e5m2 (aka bf8)

test/backend/test_graph.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import Context, dedup, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
 from tinygrad.engine.realize import BufferXfer, get_runner, CompiledRunner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.uop.ops import UOp, Ops
 
 from test.helpers import needs_second_gpu
@@ -60,7 +60,7 @@ def helper_make_view(base, offset_elems, size_elems):
 
 def helper_run_jit(jis, bufs, out_buffers):
   for rawbuf in out_buffers:
-    mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
+    mv = memoryview(bytearray(rawbuf.nbytes))
     ctypes.memset(from_mv(mv), 0, len(mv))
     rawbuf.copyin(mv)
 

test/backend/test_multitensor.py

@@ -305,7 +305,8 @@ class TestMultiTensor(unittest.TestCase):
     Xs = X.shard(device, shard_x)
     Ws = W.shard(device, shard_w)
     O = (Xs@Ws)
-    np.testing.assert_allclose(X.numpy() @ W.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
+    with np.errstate(all='ignore'):
+      np.testing.assert_allclose(X.numpy() @ W.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
 
   def _test_double_matmul_shard_axis(self, shard_x, shard_w, device):
     X = Tensor.kaiming_uniform(N, N).realize()
@@ -315,7 +316,8 @@ class TestMultiTensor(unittest.TestCase):
     W1s = W1.shard(device, shard_w)
     W2s = W2.shard(device, shard_w)
     O = (Xs@W1s)@W2s
-    np.testing.assert_allclose((X.numpy() @ W1.numpy()) @ W2.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
+    with np.errstate(all='ignore'):
+      np.testing.assert_allclose((X.numpy() @ W1.numpy()) @ W2.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
 
   def test_matmul_shard_none(self): return self._test_matmul_shard_axis(None, None, devices_2)
   def test_matmul_shard_X_0(self): return self._test_matmul_shard_axis(0, None, devices_2)

test/backend/test_opt_gemm.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, Device
 from tinygrad.helpers import get_single_element
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 
 class TestOptGemm(unittest.TestCase):
   @classmethod
@@ -12,7 +12,8 @@ class TestOptGemm(unittest.TestCase):
     N = 64
     cls.a = Tensor.randn(N, N).contiguous().realize()
     cls.b = Tensor.randn(N, N).contiguous().realize()
-    cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
+    with np.errstate(all='ignore'):
+      cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
 
   def _test_gemm_unrolled_permute_l(self, opts=[]):
     t = self.a.T @ self.b.T

test/backend/test_quantize_onnx.py

@@ -6,7 +6,7 @@ from tinygrad import Tensor, Context, Device, dtypes
 from tinygrad.uop.ops import Ops
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 
 N = 512
 

test/backend/test_schedule.py

@@ -312,8 +312,9 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out0.numpy(), np_out0:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), np_out1:=a.numpy().sum()+np_out0, atol=1e-4, rtol=1e-4)
     np_b = (a.numpy() + np_out0 + np_out1)
-    np.testing.assert_allclose(out2.numpy(), np_out2:=np.exp2(np_b.sum()), atol=1e-4, rtol=1e-4)
-    np.testing.assert_allclose(out3.numpy(), np_b.sum()+np_out2, atol=1e-4, rtol=1e-4)
+    with np.errstate(over='ignore'):
+      np.testing.assert_allclose(out2.numpy(), np_out2:=np.exp2(np_b.sum()), atol=1e-4, rtol=1e-4)
+      np.testing.assert_allclose(out3.numpy(), np_b.sum()+np_out2, atol=1e-4, rtol=1e-4)
 
   def test_reduce_ext_reduce_child(self):
     Tensor.manual_seed(0)
@@ -510,7 +511,7 @@ class TestSchedule(unittest.TestCase):
     np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)
 
   def test_cumsum_parallel_reduce_fused(self):
-    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END
+    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END (same nesting context = should merge)
     step, num_steps = 513, 10
     t = Tensor.arange(step).float().realize()
     phase = t.cumsum()
@@ -521,6 +522,12 @@ class TestSchedule(unittest.TestCase):
     expected = (expected * np.array([1,0,0,1,0,0,0,0,1,0]).reshape(num_steps, 1)).flatten()
     np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
 
+  @unittest.skipIf(Device.DEFAULT == "CL", "TODO: fails on CI CL")
+  def test_reduce_different_nesting_depth(self):
+    # two REDUCEs sharing the same RANGE at different nesting depths must NOT merge
+    x = Tensor.arange(768).reshape(3, 256).float()
+    np.testing.assert_allclose((x.sum(axis=1) + x.sum(axis=1).sum()).numpy(), x.numpy().sum(axis=1) + x.numpy().sum(axis=1).sum())
+
   def test_multimatmul_fusion(self):
     Tensor.manual_seed(0)
     a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()

test/backend/test_uops.py

@@ -8,7 +8,7 @@ from tinygrad.device import Buffer, Device
 from tinygrad.uop.ops import Ops, UOp, KernelInfo, AxisType
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.engine.realize import CompiledRunner, get_program, get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer

test/external/external_llm_eval.py

@@ -8,7 +8,7 @@ LABEL = ["A", "B", "C", "D"]
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--port", "-p", type=int, default=11434)
+  parser.add_argument("--port", "-p", type=int, default=8000)
   parser.add_argument("--limit", "-L", type=int, default=None)
   parser.add_argument("--max_tokens", "-T", type=int, default=4096)
   parser.add_argument("--offset", "-O", type=int, default=0)

test/external/external_test_amd.py

@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Device, Tensor
-from tinygrad.engine.schedule import create_schedule
+from tinygrad.schedule import create_schedule
 from tinygrad.runtime.ops_amd import AMDDevice
 
 class TestAMD(unittest.TestCase):

test/external/external_test_hip_compile.py

@@ -2,7 +2,7 @@ import time, unittest
 from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
-from tinygrad.engine.schedule import create_schedule
+from tinygrad.schedule import create_schedule
 from tinygrad.codegen.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):

test/external/external_test_opt.py

@@ -7,7 +7,7 @@ from tinygrad import GlobalCounters, Tensor, Device
 from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing, run_schedule
-from tinygrad.engine.schedule import linear_to_schedule
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.tensor import _to_np_dtype
 
 class CLCache:

test/external/external_uop_gc.py

@@ -1,6 +1,6 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op, _apply_reshape
 from tinygrad.uop.divandmod import fold_divmod_general

test/external/fuzz_graph.py

@@ -5,7 +5,7 @@ from tinygrad.helpers import Context, getenv, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import BufferXfer, get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.jit import apply_graph_to_jit
 
@@ -85,7 +85,7 @@ def run_jit(jis, all_buffers, input_buffers, var_vals):
   with Context(DEBUG=0):
     for rawbuf in all_buffers:
       if rawbuf in input_buffers: continue
-      mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
+      mv = memoryview(bytearray(rawbuf.nbytes))
       ctypes.memset(from_mv(mv), 0, len(mv))
       rawbuf.copyin(mv)
 

test/mockgpu/amd/README

@@ -4,7 +4,7 @@ Test with `pytest -n12 test/amd/`
 `DEV=AMD:LLVM pytest -n12 test/amd/`
 
 * dsl.py -- helpers for the autogen instruction classes in `__init__.py`. should be standalone with init
-* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `DEV=AMD MOCKGPU=1 PYTHON_REMU=1`
+* test/mockgpu/amd/emu.py -- an emulator for RDNA that runs in tinygrad with `DEV=AMD MOCKGPU=1`
 * generate.py -- extract assembly format + instruction pseudocode from AMD XML + PDF
 * test/mockgpu/amd/pcode.py -- pseudocode to UOp transformation
 * sqtt.py -- SQTT parser
@@ -20,20 +20,19 @@ test_llvm.py tests asm/disasm on the LLVM tests, confirming it behaves the same
 
 tinygrad's dtype tests should pass with and without LLVM. they run in about 12 seconds.
 
-`DEV=AMD PYTHON_REMU=1 MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
-`DEV=AMD:LLVM PYTHON_REMU=1 MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`DEV=AMD MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
+`DEV=AMD:LLVM MOCKGPU=1 pytest -n=12 test/backend/test_dtype_alu.py test/backend/test_dtype.py`
 
 The ops tests also pass, but they are very slow, so you should run them one at a time.
 
-`SKIP_SLOW_TEST=1 DEV=AMD PYTHON_REMU=1 MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
-`SKIP_SLOW_TEST=1 DEV=AMD:LLVM PYTHON_REMU=1 MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 DEV=AMD MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
+`SKIP_SLOW_TEST=1 DEV=AMD:LLVM MOCKGPU=1 pytest -n=12 test/backend/test_ops.py`
 
 When something is caught by main tinygrad tests, a local regression test should be added to `test/amd`.
 While working with tinygrad, you can dump the assembly with `DEBUG=7`. These tests all pass on real hardware
-If a test is failing with `DEV=AMD PYTHON_REMU=1 MOCKGPU=1` it's because an instruction is emulated incorrectly.
+If a test is failing with `DEV=AMD MOCKGPU=1` it's because an instruction is emulated incorrectly.
 You can test without `MOCKGPU=1` to test on real hardware, if it works on real hardware there's a bug in the emulator.
 IMPORTANT: if a test is failing in the emulator, it's an instruction bug. Use DEBUG=7, get the instructions, and debug.
 
 Currently, only RDNA3 is well supported, but when finished, this will support RDNA3+RDNA4+CDNA in ~3000 lines.
 Get line count with `cloc --by-file tinygrad/renderer/amd/*.py`
-

test/mockgpu/amd/amdgpu.py

@@ -1,7 +1,7 @@
 import ctypes, time
 from dataclasses import replace
 from test.mockgpu.gpu import VirtGPU
-from test.mockgpu.helpers import _try_dlopen_remu
+from test.mockgpu.helpers import PythonRemu
 from tinygrad.helpers import getbits, to_mv, getenv, DEV
 from tinygrad.runtime.support import c
 
@@ -41,7 +41,7 @@ WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
 WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
 
-remu = _try_dlopen_remu()
+remu = PythonRemu()
 
 def create_sdma_packets():
   # TODO: clean up this, if we want to keep it
@@ -212,13 +212,13 @@ class PM4Executor(AMDQueue):
     scratch_size = wavesize * (16 if self.gpu.arch == "cdna" else 4)  # per-thread scratch size in bytes
 
     assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # 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)
+    # Pass valid memory ranges, rsrc2, scratch_size, arch, and user data registers to the emulator
+    remu.valid_mem_ranges = self.gpu.mapped_ranges
+    remu.rsrc2 = rsrc2
+    remu.scratch_size = scratch_size
+    remu.arch = self.gpu.arch
+    remu.user_data = user_data
+    err = remu.run_asm(prg_addr, prg_sz, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr)
     if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
 
   def _exec_indirect_buffer(self, n):

test/mockgpu/amd/emu.py

@@ -67,6 +67,7 @@ from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
 from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
+from test.amd.helpers import decode_dpp16
 from test.mockgpu.amd.pcode import parse_block, _FUNCS, _set_bits, _val_to_bits
 
 MASK32 = 0xFFFFFFFF
@@ -233,7 +234,6 @@ VOPD_TO_VOP2 = {
   ir4.VOPDOp.V_DUAL_DOT2ACC_F32_F16: ir3.VOP2Op.V_DOT2ACC_F32_F16_E32,
 }
 def _wave_size(arch: str) -> int: return 64 if arch.startswith("cdna") else 32
-WAVE_SIZE = 32  # default wave size for RDNA (exported for test_compare_emulators)
 # Special registers stored after inline constants (256-259)
 PC_LO_IDX, PC_HI_IDX, SCRATCH_STRIDE_IDX = 256, 257, 259
 # SGPR buffer: 0-127 = SGPRs, 128-255 = inline constants, 256-259 = special registers
@@ -346,7 +346,7 @@ def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[d
   # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
   lines: list[str] = []
   for l in raw_lines:
-    if lines and lines[-1].endswith('&&'): lines[-1] = lines[-1] + ' ' + l
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
     else: lines.append(l)
   _, final, _ = parse_block(lines, 0, env, assigns=assigns)
   sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
@@ -639,11 +639,13 @@ class _Ctx:
     src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), _c(0))  # VGPR index or 0
     src1_off = self.inst_field(type(inst).src1) if hasattr(type(inst), 'src1') else None
     src2_off = self.inst_field(type(inst).src2) if hasattr(type(inst), 'src2') else None
+    src1_reg = (src1_off >= _c(256)).where(src1_off - _c(256), src1_off) if src1_off is not None else _c(0)
+    src2_reg = (src2_off >= _c(256)).where(src2_off - _c(256), src2_off) if src2_off is not None else _c(0)
     exec_val = self.rexec()
     exec_lo = exec_val.cast(dtypes.uint32) if exec_val.dtype == dtypes.uint64 else exec_val
     srcs = {
       'SRC0': src0_reg, 'VDST': vdst_off, 'EXEC_LO': exec_lo, 'EXEC': exec_val if exec_val.dtype == dtypes.uint64 else exec_val.cast(dtypes.uint64),
-      '_vgpr': self.vgpr, '_wave_size': self.wave_size,
+      '_vgpr': self.vgpr, '_wave_size': self.wave_size, 'SRC1': src1_reg, 'SRC2': src2_reg,
       'S0': self.rsrc_dyn(src0_off, _c(0, dtypes.int)) if 'WRITELANE' in op_name else src0_reg,
       'S1': self.rsrc_dyn(src1_off, _c(0, dtypes.int)) if src1_off is not None else _c(0),
       'S2': self.rsrc_dyn(src2_off, _c(0, dtypes.int)) if src2_off is not None else _c(0),
@@ -663,10 +665,11 @@ class _Ctx:
     vcc_reg = sdst_reg if sdst_reg is not None else VCC_LO.offset
     if 'VCC' not in srcs: srcs['VCC'] = self.rmask(_c(vcc_reg))
     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), '_vgpr': self.vgpr, '_wave_size': self.wave_size,
-                 # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults)
-                 'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
-                 'WORD0': _c(0), 'WORD1': _c(1)})  # rounding mode and SDWA constants
+                  'ROUND_MODE': _c(0), 'ROUND_TOWARD_ZERO': _c(0), 'ROUND_NEAREST_EVEN': _c(0), '_vgpr': self.vgpr, '_wave_size': self.wave_size,
+                  'MAX_FLOAT_F32': UOp.const(dtypes.float32, 3.4028234663852886e38),
+                  # CDNA SDWA byte/word select constants (E32 always uses BYTE0/WORD0 defaults)
+                  'SDWA_SRC0_SEL': _c(0), 'BYTE0': _c(0), 'BYTE1': _c(1), 'BYTE2': _c(2), 'BYTE3': _c(3),
+                  'WORD0': _c(0), 'WORD1': _c(1)})  # rounding mode and SDWA constants
     _, assigns = parse_pcode(pcode, srcs)
 
     # For integer ops with clamp, compute overflow using wide arithmetic
@@ -716,6 +719,10 @@ class _Ctx:
         new_vcc = _set_lane_bit(old_vcc, lane, val, exec_mask)
         raw_stores.extend([('vcc', s) for s in self.wmask(_c(VCC_LO.offset), new_vcc)])
       elif dest.startswith('D0'):
+        dest_suffix = re.match(r'D0\.(\w+)', dest)
+        if dest_suffix is not None:
+          target_dt = {'u16': dtypes.uint16, 'i16': dtypes.int16, 'f16': dtypes.half}.get(dest_suffix.group(1))
+          if target_dt is not None and val.dtype != target_dt: val = val.cast(target_dt)
         if (slice_match := re.match(r'D0\[(\d+)\s*:\s*(\d+)\]', dest)):
           d0_hi_bit, d0_lo_bit = int(slice_match.group(1)), int(slice_match.group(2))
           if d0_hi_bit != 31 or d0_lo_bit != 0:
@@ -728,7 +735,8 @@ class _Ctx:
         # For integer ops with clamp, use pre-computed saturated value; for floats, clamp to [0,1]
         if int_saturate is not None: val = int_saturate
         elif clmp and val.dtype in (dtypes.float32, dtypes.half, dtypes.float64):
-          val = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
+          clamped = val.maximum(UOp.const(val.dtype, 0.0)).minimum(UOp.const(val.dtype, 1.0))
+          val = _FUNCS['isNAN'](val).where(UOp.const(val.dtype, 0.0), clamped)
         if val.dtype in (dtypes.uint64, dtypes.int64, dtypes.float64):
           lo, hi = _split64(val)
           raw_stores.extend([('vgpr', self.wvgpr_dyn(vdst_reg, lane, lo, exec_mask)),
@@ -916,6 +924,49 @@ def _sdwa_write(old: UOp, val: UOp, dst_sel: UOp, dst_unused: UOp) -> UOp:
   # For PAD and SEXT, unused bits are zero (PAD) or sign-extended (SEXT). For DWORD, just return val.
   return dst_sel.eq(_c(6)).where(val, dst_unused.eq(_c(2)).where(preserved, placed))
 
+def _dpp_quad_sel(quad_lane: UOp, sels: tuple[int, int, int, int]) -> UOp:
+  sel = _c(sels[0], dtypes.int)
+  for i, src in enumerate(sels[1:], start=1): sel = quad_lane.eq(_c(i, dtypes.int)).where(_c(src, dtypes.int), sel)
+  return sel
+
+def _dpp16_ctrl(lane: UOp, dpp: int, row_mask: int, bank_mask: int, wave_size: int) -> tuple[UOp, UOp, UOp]:
+  """Return (src_lane, row/bank enabled, in-bounds) for a DPP16 swizzle."""
+  lane_i = lane.cast(dtypes.int)
+  row_base, lane_in_row = lane_i & _c(~15, dtypes.int), lane_i & _c(15, dtypes.int)
+  row = lane_i // _c(16, dtypes.int)
+  bank = lane_in_row >> _c(2, dtypes.int)
+  enabled = (((_c(row_mask) >> row.cast(dtypes.uint32)) & _c(1)).ne(_c(0)) &
+             (((_c(bank_mask) >> bank.cast(dtypes.uint32)) & _c(1)).ne(_c(0))))
+  op, arg = decode_dpp16(dpp)
+  src_lane, valid = lane_i, UOp.const(dtypes.bool, True)
+
+  if op == 'quad_perm':
+    assert isinstance(arg, tuple)
+    src_lane = (lane_i & _c(~3, dtypes.int)) + _dpp_quad_sel(lane_i & _c(3, dtypes.int), arg)
+  else:
+    assert isinstance(arg, int)
+    if op == 'row_shl': src_lane, valid = row_base + lane_in_row + _c(arg, dtypes.int), lane_in_row <= _c(15 - arg, dtypes.int)
+    elif op == 'row_shr': src_lane, valid = row_base + lane_in_row - _c(arg, dtypes.int), lane_in_row >= _c(arg, dtypes.int)
+    elif op == 'row_ror': src_lane = row_base + ((lane_in_row - _c(arg, dtypes.int)) & _c(15, dtypes.int))
+    elif op == 'row_mirror': src_lane = row_base + (_c(15, dtypes.int) - lane_in_row)
+    elif op == 'row_half_mirror': src_lane = row_base + ((lane_in_row & _c(8, dtypes.int)) | (_c(7, dtypes.int) - (lane_in_row & _c(7, dtypes.int))))
+    elif op == 'row_bcast': src_lane = row_base
+    elif op == 'wave_shl': src_lane, valid = lane_i + _c(arg, dtypes.int), lane_i < _c(wave_size - arg, dtypes.int)
+    elif op == 'wave_rol': src_lane = (lane_i + _c(arg, dtypes.int)) % _c(wave_size, dtypes.int)
+    elif op == 'wave_shr': src_lane, valid = lane_i - _c(arg, dtypes.int), lane_i >= _c(arg, dtypes.int)
+    elif op == 'wave_ror': src_lane = (lane_i - _c(arg, dtypes.int)) % _c(wave_size, dtypes.int)
+    else: raise NotImplementedError(f"DPP16 control {dpp:#x} ({op}:{arg}) not implemented in emulator")
+  return src_lane, enabled, valid
+
+def _load_dpp16_src0(ctx: _Ctx, inst, lane: UOp, fallback: UOp) -> UOp:
+  """Load a DPP16-swizzled src0 value from vsrc0."""
+  src_lane, enabled, valid = _dpp16_ctrl(lane, getattr(inst, 'dpp', 0) or 0, getattr(inst, 'row_mask', 0xf) or 0xf,
+                                         getattr(inst, 'bank_mask', 0xf) or 0xf, ctx.wave_size)
+  safe_src_lane = (enabled & valid).where(src_lane, _c(0, dtypes.int))
+  swizzled = ctx.rvgpr_dyn(ctx.inst_field(type(inst).vsrc0), safe_src_lane)
+  invalid = UOp.const(fallback.dtype, 0) if getattr(inst, 'bc', 0) else fallback
+  return enabled.where(valid.where(swizzled, invalid), fallback)
+
 def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc.VOPC_SDWA_SDST, ctx: _Ctx) -> UOp:
   """Compile CDNA SDWA (Sub-Dword Access) VOP1/VOP2/VOPC instructions."""
   is_vopc = isinstance(inst, irc.VOPC_SDWA_SDST)
@@ -999,33 +1050,43 @@ def _compile_sdwa(inst: irc.VOP1_SDWA | irc.VOP2_SDWA | irc.VOP2_SDWA_SDST | irc
     return UOp.sink(UOp.sink(*stores).end(lane), *ctx.inc_pc())
   return UOp.sink(*ctx.inc_pc())
 
-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:
+def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP1_DPP16 | ir3.VOP2 | ir3.VOP2_DPP16 |
+                   ir4.VOP1 | ir4.VOP1_SDST | ir4.VOP1_DPP16 | ir4.VOP2 | ir4.VOP2_DPP16 |
+                   irc.VOP1 | irc.VOP1_DPP16 | irc.VOP2 | irc.VOP2_DPP16, 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)
   # v_accvgpr_mov_b32: ACCVGPR[vdst] = ACCVGPR[src0] (VOP1 encoding, no pcode)
   if 'ACCVGPR_MOV' in op_name:
     lane, exec_mask = ctx.range(), ctx.rexec()
     vdst_reg = ctx.inst_field(type(inst).vdst)  # VGPRField: raw ACCVGPR index (0-255)
-    src0_off = ctx.inst_field(type(inst).src0)  # SrcField: raw 256 + ACCVGPR index
-    val = ctx.raccvgpr_dyn(src0_off - _c(256), lane)
+    acc_src0_off = ctx.inst_field(type(inst).src0)  # SrcField: raw 256 + ACCVGPR index
+    val = ctx.raccvgpr_dyn(acc_src0_off - _c(256), lane)
     return UOp.sink(ctx.waccvgpr_dyn(vdst_reg, lane, val, exec_mask).end(lane), *ctx.inc_pc())
   lane, exec_mask, bits = ctx.range(), ctx.rexec(), inst.canonical_op_bits
   literal = ctx.inst_field(type(inst).literal) if hasattr(type(inst), 'literal') else None  # type: ignore[union-attr]
   is_f64 = 'F64' in op_name and 'B64' not in op_name
+  is_float = any(x in op_name for x in ('F16', 'F32', 'F64'))
+  is_dpp16 = hasattr(type(inst), 'dpp') and hasattr(type(inst), 'vsrc0')
   vdst_reg = ctx.inst_field(type(inst).vdst)
   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
+  src0_off: UOp | None = None
   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, is_f64)
-    if bits['s0'] == 16:
+    d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))
+    if is_dpp16:
+      s0 = _load_dpp16_src0(ctx, inst, lane, d0)
+    else:
+      src0_off = ctx.inst_field(type(inst).src0)
+      s0 = ctx.rsrc_dyn(src0_off, lane, bits['s0'], literal, is_f64)
+    if bits['s0'] == 16 and not is_dpp16:
       src0_hi = src0_off >= _c(384)
       # Only compute hi-half when src0_off >= 384, use guarded index to prevent OOB access
       src0_reg = src0_hi.where(src0_off - _c(384), _c(0))
       s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
-    d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))
+    if is_dpp16 and is_float:
+      s0 = _apply_src_mods(s0, 0, 1 if getattr(inst, 'src0_abs', 0) else 0, 1 if getattr(inst, 'src0_neg', 0) else 0, bits['s0'])
     srcs:dict[str, UOp | int] = {'S0': s0, 'D0': d0}
   else:
     vsrc1_reg = ctx.inst_field(type(inst).vsrc1)
@@ -1038,13 +1099,19 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VO
       s1 = _cond_hi16(vsrc1_hi, ctx.rvgpr_dyn(vsrc1_actual, lane))
       d0 = _cond_hi16(write_hi_half, ctx.rvgpr_dyn(vdst_reg, lane))  # FMAC/FMAMK hi-half dest needs hi-half accumulator
     # Handle VOP2 hi-half src0 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, is_f64)
-    if bits['s0'] == 16:
+    if is_dpp16:
+      s0 = _load_dpp16_src0(ctx, inst, lane, d0)
+    else:
+      src0_off = ctx.inst_field(type(inst).src0)
+      s0 = ctx.rsrc_dyn(src0_off, lane, bits['s0'], literal, is_f64)
+    if bits['s0'] == 16 and not is_dpp16:
       src0_hi = src0_off >= _c(384)
       # Only compute hi-half when src0_off >= 384, use guarded index to prevent OOB access
       src0_reg = src0_hi.where(src0_off - _c(384), _c(0))
       s0 = src0_hi.where(_hi16(ctx.rvgpr_dyn(src0_reg, lane)), s0)
+    if is_dpp16 and is_float:
+      s0 = _apply_src_mods(s0, 0, 1 if getattr(inst, 'src0_abs', 0) else 0, 1 if getattr(inst, 'src0_neg', 0) else 0, bits['s0'])
+      s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1'])
     srcs = {'S0': s0, 'S1': s1, 'D0': d0}
     # FMAAK_(DTYPE)_E32 series
     if 'V_FMAA' in _op_name(inst) or 'V_FMAM' in _op_name(inst):
@@ -1052,10 +1119,11 @@ def _compile_vop12(inst: ir3.VOP1 | ir3.VOP1_SDST | ir3.VOP2 | ir4.VOP1 | ir4.VO
       srcs['SIMM32'] = literal
   return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask, opsel_dst_hi=write_hi_half, src0_off=src0_off)
 
-def _compile_vopc(inst: ir3.VOPC|ir3.VOP3|ir4.VOPC|ir4.VOP3|irc.VOPC|irc.VOP3, ctx: _Ctx,
+def _compile_vopc(inst: ir3.VOPC|ir3.VOPC_DPP16|ir3.VOP3|ir4.VOPC|ir4.VOPC_DPP16|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.rexec(), _op_name(inst), inst.canonical_op_bits
   is_cmpx, is_vopc = 'CMPX' in op_name, hasattr(inst, 'vsrc1')  # is_vopc: e32 vs e64
+  is_dpp16 = hasattr(type(inst), 'dpp') and hasattr(type(inst), 'vsrc0')
 
   # Handle both VOPC (vsrc1) and VOP3 (src1) instruction formats - read operands dynamically
   if is_vopc:
@@ -1078,11 +1146,14 @@ def _compile_vopc(inst: ir3.VOPC|ir3.VOP3|ir4.VOPC|ir4.VOP3|irc.VOPC|irc.VOP3, c
   is_float, is_f64, pcode = any(x in op_name for x in ('_F32', '_F64', '_F16')), '_F64' in op_name, get_pcode(inst.op)
   def get_cmp_bit(lane) -> UOp:
     lc = lane.cast(dtypes.int) if isinstance(lane, UOp) else _c(lane, dtypes.int)
-    s0 = ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64)
+    s0 = _load_dpp16_src0(ctx, inst, lc, _c(0)) if is_dpp16 else ctx.rsrc_dyn(src0_off, lc, bits['s0'], literal, is_f64)
     s1 = _cond_hi16(vsrc1_hi, ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)) if bits['s0'] == 16 \
       else ctx.rsrc_dyn(src1_off, lc, bits['s1'], literal, is_f64)
     if bits['s0'] == 16 and opsel: s0, s1 = _apply_opsel(s0, 0, opsel), _apply_opsel(s1, 1, opsel)
     if is_float:
+      if is_dpp16:
+        s0 = _apply_src_mods(s0, 0, 1 if getattr(inst, 'src0_abs', 0) else 0, 1 if getattr(inst, 'src0_neg', 0) else 0, bits['s0'])
+        s1 = _apply_src_mods(s1, 0, 1 if getattr(inst, 'src1_abs', 0) else 0, 1 if getattr(inst, 'src1_neg', 0) else 0, bits['s1'])
       s0 = _apply_src_mods(s0, 0, abs_bits, neg_bits, bits['s0'])
       s1 = _apply_src_mods(s1, 1, abs_bits, neg_bits, bits['s1'])
     for dest, val in parse_pcode(pcode, {'S0': s0, 'S1': s1, 'laneId': lc, 'D0': UOp.const(dtypes.uint64, 0)})[1]:
@@ -1177,6 +1248,19 @@ def _compile_vop3(inst: ir3.VOP3 | ir4.VOP3 | irc.VOP3, ctx: _Ctx) -> UOp:
   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_vinterp(inst: ir3.VINTERP | ir4.VINTERP, ctx: _Ctx) -> UOp:
+  lane, exec_mask = ctx.range(), ctx.rexec()
+  inst_type = type(inst)
+  vdst_reg = ctx.inst_field(inst_type.vdst)
+  src0_off, src1_off, src2_off = ctx.inst_field(inst_type.src0), ctx.inst_field(inst_type.src1), ctx.inst_field(inst_type.src2)
+  src0_reg = (src0_off >= _c(256)).where(src0_off - _c(256), src0_off)
+  src2_reg = (src2_off >= _c(256)).where(src2_off - _c(256), src2_off)
+  srcs = {
+    'SRC0': src0_reg, 'SRC2': src2_reg,
+    'S0': ctx.rsrc_dyn(src0_off, lane), 'S1': ctx.rsrc_dyn(src1_off, lane), 'S2': ctx.rsrc_dyn(src2_off, lane),
+  }
+  return ctx.compile_vop_pcode(inst.op, srcs, lane, vdst_reg, exec_mask)
+
 def _compile_vop3sd(inst: ir3.VOP3SD | ir4.VOP3SD | irc.VOP3SD, ctx: _Ctx) -> UOp:
   exec_mask = ctx.rexec()
   bits, pcode, ops = inst.canonical_op_bits, get_pcode(inst.op), inst.canonical_operands
@@ -1788,7 +1872,7 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
                 'DATA2': _u64(ctx.rvgpr_dyn(data1_reg, lane), ctx.rvgpr_dyn(data1_reg + _c(1), lane)) if has_data1 else UOp.const(dtypes.uint64, 0)}
       # RDNA3 uses ADDR/OFFSET, RDNA4 uses vgpr_a/offset (lowercase) + CalcDsAddr function
       return {'ADDR': addr, 'ADDR_BASE': addr, 'OFFSET': offset, 'OFFSET0': offset0, 'OFFSET1': offset1, '_lds': mem, 'laneId': lane,
-              'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, **data}
+              'vgpr_a': ctx.rvgpr_dyn(addr_reg, lane), 'offset': offset, 'offset0': offset0, 'offset1': offset1, **data}
     active = _lane_active(exec_mask, lane)
     # saddr < 124 means valid SGPR pair, otherwise use 0 (NULL means no saddr contribution)
     use_saddr = (saddr_reg < _c(124)) if saddr_reg is not None else UOp.const(dtypes.bool, False)
@@ -1937,17 +2021,20 @@ def _compile_mubuf(inst: irc.MUBUF, ctx: _Ctx) -> UOp:
 # Dispatch table: instruction type -> handler function
 _INST_HANDLERS: dict[type, Callable[..., UOp]] = {
   ir3.SOPP: _compile_sopp, ir3.SMEM: _compile_smem, ir3.SOP1: _compile_sop, ir3.SOP2: _compile_sop, ir3.SOPC: _compile_sop, ir3.SOPK: _compile_sop,
-  ir3.VOP1: _compile_vop12, ir3.VOP1_SDST: _compile_vop12, ir3.VOP2: _compile_vop12, ir3.VOPC: _compile_vopc, ir3.VOP3: _compile_vop3,
+  ir3.VOP1: _compile_vop12, ir3.VOP1_SDST: _compile_vop12, ir3.VOP1_DPP16: _compile_vop12, ir3.VOP2: _compile_vop12, ir3.VOP2_DPP16: _compile_vop12,
+  ir3.VOPC: _compile_vopc, ir3.VOPC_DPP16: _compile_vopc, ir3.VOP3: _compile_vop3, ir3.VINTERP: _compile_vinterp,
   ir3.VOP3_SDST: _compile_vop3, ir3.VOP3SD: _compile_vop3sd, ir3.VOP3P: _compile_vop3p, ir3.VOPD: _compile_vopd,
   ir3.DS: _compile_mem_op, ir3.FLAT: _compile_mem_op, ir3.GLOBAL: _compile_mem_op, ir3.SCRATCH: _compile_mem_op,
   # RDNA4 instruction classes
   ir4.SOPP: _compile_sopp, ir4.SMEM: _compile_smem, ir4.SOP1: _compile_sop, ir4.SOP2: _compile_sop, ir4.SOPC: _compile_sop, ir4.SOPK: _compile_sop,
-  ir4.VOP1: _compile_vop12, ir4.VOP1_SDST: _compile_vop12, ir4.VOP2: _compile_vop12, ir4.VOPC: _compile_vopc, ir4.VOP3: _compile_vop3,
+  ir4.VOP1: _compile_vop12, ir4.VOP1_SDST: _compile_vop12, ir4.VOP1_DPP16: _compile_vop12, ir4.VOP2: _compile_vop12, ir4.VOP2_DPP16: _compile_vop12,
+  ir4.VOPC: _compile_vopc, ir4.VOPC_DPP16: _compile_vopc, ir4.VOP3: _compile_vop3, ir4.VINTERP: _compile_vinterp,
   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.VOP1: _compile_vop12, irc.VOP1_DPP16: _compile_vop12, irc.VOP2: _compile_vop12, irc.VOP2_DPP16: _compile_vop12,
+  irc.VOPC: _compile_vopc, irc.VOP3: _compile_vop3,
   irc.VOP3_SDST: _compile_vop3, irc.VOP3SD: _compile_vop3sd, irc.VOP3P: _compile_vop3p,
   irc.VOP1_SDWA: _compile_sdwa, irc.VOP2_SDWA: _compile_sdwa, irc.VOP2_SDWA_SDST: _compile_sdwa, irc.VOPC_SDWA_SDST: _compile_sdwa,
   irc.DS: _compile_mem_op, irc.FLAT: _compile_mem_op, irc.GLOBAL: _compile_mem_op, irc.SCRATCH: _compile_mem_op,

test/mockgpu/amd/pcode.py

@@ -4,7 +4,7 @@ from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.uop.decompositions import f2f
 
-# Type alias for vars dict: stores UOps for variables and tuples for lambda definitions
+# Type alias for vars dict: stores UOps and tuples for lambda definitions
 VarVal = UOp | tuple[str, list[str], str]
 
 def _const(dt, v): return UOp.const(dt, v)
@@ -50,6 +50,22 @@ def _extract_bits(val: UOp, hi: int, lo: int) -> UOp:
   if result.dtype != target_dt: result = result.cast(target_dt)
   return result
 
+def _expr_bits(v: UOp) -> int:
+  if v.dtype == dtypes.bool: return 1
+  if v.op in (Ops.AND, Ops.XOR):
+    widths: list[int] = []
+    for src in v.src:
+      if src.op == Ops.CONST and isinstance(src.arg, int) and src.arg > 0 and (src.arg & (src.arg + 1)) == 0:
+        widths.append(src.arg.bit_length())
+    if widths: return max(widths)
+  return v.dtype.bitsize
+
+def _countbits(v: UOp) -> UOp:
+  dt = dtypes.uint64 if _expr_bits(v) > 32 or v.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
+  vv, out = v.cast(dt), _u32(0)
+  for i in range(_expr_bits(v)): out = out + ((vv >> _const(dt, i)) & _const(dt, 1)).cast(dtypes.uint32)
+  return out
+
 def _set_bit(old, pos, val):
   mask = _u32(1) << pos
   return (old & (mask ^ _u32(0xFFFFFFFF))) | ((val.cast(dtypes.uint32) & _u32(1)) << pos)
@@ -335,6 +351,7 @@ _FUNCS: dict[str, Callable[..., UOp]] = {
   # System NOPs - these are scheduling hints, no effect on emulation
   'MIN': lambda a, b: (a < b).where(a, b),
   's_nop': lambda a: _u32(0),
+  'countbits': _countbits,
   # Address calculation for memory operations
   'CalcDsAddr': lambda a, o, *r: a.cast(dtypes.uint32) + o.cast(dtypes.uint32),
   'CalcGlobalAddr': lambda v, s, *r: v.cast(dtypes.uint64) + s.cast(dtypes.uint64),
@@ -389,7 +406,7 @@ def tokenize(s: str) -> list[Token]:
     if c.isspace():
       i += 1
       continue
-    if i + 1 < n and s[i:i+2] in ('+=', '-='):
+    if i + 1 < n and s[i:i+2] in ('+=', '-=', '|=', '&=', '^='):
       tokens.append(Token('ASSIGN_OP', s[i:i+2]))
       i += 2
       continue
@@ -503,7 +520,7 @@ class Parser:
   def unary(self) -> UOp:
     if self.try_eat_val('~', 'OP'):
       inner = self.unary()
-      return inner ^ _const(inner.dtype, (1 << (inner.dtype.itemsize * 8)) - 1)
+      return inner ^ _const(inner.dtype, (1 << _expr_bits(inner)) - 1)
     if self.try_eat_val('!', 'OP'):
       inner = self.unary()
       return inner.eq(_const(inner.dtype, 0))
@@ -539,7 +556,10 @@ class Parser:
       self.eat('COMMA')
       lo = self.parse()
       self.eat('RBRACE')
-      return (hi.cast(dt:=_BITS_DT.get((s:=lo.dtype.bitsize) * 2, dtypes.uint64)) << _const(dt, s)) | lo.cast(dt)
+      lo_bits, hi_bits = _expr_bits(lo), _expr_bits(hi)
+      total_bits = lo_bits + hi_bits
+      dt = _BITS_DT.get(total_bits, dtypes.uint32 if total_bits <= 32 else dtypes.uint64)
+      return (hi.cast(dt) << _const(dt, lo_bits)) | lo.cast(dt)
     if self.at('NUM'):
       num = self.eat('NUM').val
       if self.try_eat('QUOTE'):
@@ -576,8 +596,8 @@ class Parser:
       if name == 'OVERFLOW_F32': return _const(dtypes.uint32, 0x7F7FFFFF).bitcast(dtypes.float32)
       if name == 'UNDERFLOW_F64': return _const(dtypes.uint64, 1).bitcast(dtypes.float64)
       if name == 'OVERFLOW_F64': return _const(dtypes.uint64, 0x7FEFFFFFFFFFFFFF).bitcast(dtypes.float64)
-      if name == 'WAVE32': return _const(dtypes.bool, self.vars.get('_wave_size', 32) <= 32)
-      if name == 'WAVE64': return _const(dtypes.bool, self.vars.get('_wave_size', 32) > 32)
+      if name.lower() == 'wave32': return _const(dtypes.bool, self.vars.get('_wave_size', 32) <= 32)
+      if name.lower() == 'wave64': return _const(dtypes.bool, self.vars.get('_wave_size', 32) > 32)
       if name == 'WAVE_MODE' and self.try_eat('DOT') and self.try_eat_val('IEEE', 'IDENT'): return _u32(1)
       if self.try_eat('LBRACE'):
         idx = self.eat('NUM').val
@@ -685,7 +705,7 @@ class Parser:
       dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
       base_cast = base.cast(dt) if base.dtype != dt else base
       result = ((base_cast >> _const(dt, idx)) & _const(dt, 1))
-      return _cast_to(result, dt_suffix) if dt_suffix else result
+      return _cast_to(result, dt_suffix) if dt_suffix else result.cast(dtypes.bool)
     if var_name:
       idx_u32 = _to_u32(first)
       elems = [(i, self.vars[f'{var_name}@{i}']) for i in range(256) if f'{var_name}@{i}' in self.vars]
@@ -699,7 +719,7 @@ class Parser:
     dt = dtypes.uint64 if base.dtype in (dtypes.uint64, dtypes.int64) else dtypes.uint32
     base_cast = base.cast(dt) if base.dtype != dt else base
     result = (base_cast >> first.cast(dt)) & _const(dt, 1)
-    return _cast_to(result, dt_suffix) if dt_suffix else result
+    return _cast_to(result, dt_suffix) if dt_suffix else result.cast(dtypes.bool)
 
   def _handle_brace_index(self, base) -> UOp:
     self.eat('LBRACE')
@@ -845,7 +865,7 @@ class Parser:
         hi = mem.index(safe_idx_hi, *gate)
         combined = val.cast(dtypes.uint64) | (hi.cast(dtypes.uint64) << UOp.const(dtypes.uint64, 32))
         val = is_unaligned.where((combined >> (byte_off.cast(dtypes.uint64) * UOp.const(dtypes.uint64, 8))).cast(dtypes.uint32), val)
-    return val
+    return _cast_to(val, dt)
 
   def _coerce_cmp(self, l: UOp, r: UOp) -> tuple[UOp, UOp]:
     if l.dtype != r.dtype:
@@ -1044,14 +1064,12 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       elif j < len(toks) and toks[j].type == 'EQUALS': j += 1
       rhs = parse_tokens(toks[j:], env, funcs)
       if compound_op:
-        mem = env.get('_vmem') if '_vmem' in env else env.get('_lds')
-        if isinstance(mem, UOp):
-          adt = dtypes.uint64 if addr.dtype == dtypes.uint64 else dtypes.uint32
-          idx = (addr >> _const(adt, 2)).cast(dtypes.int)
-          old = mem.index(idx)
-          if dt in (dtypes.uint64, dtypes.int64, dtypes.float64):
-            old = old.cast(dtypes.uint64) | (mem.index(((addr + _const(adt, 4)) >> _const(adt, 2)).cast(dtypes.int)).cast(dtypes.uint64) << _u64(32))
-          rhs = (old + rhs) if compound_op == '+=' else (old - rhs)
+        old = Parser([Token('EOF', '')], env, funcs)._handle_mem_load(addr, dt)
+        if compound_op == '+=': rhs = old + rhs
+        elif compound_op == '-=': rhs = old - rhs
+        elif compound_op == '|=': rhs = old | rhs
+        elif compound_op == '&=': rhs = old & rhs
+        elif compound_op == '^=': rhs = old ^ rhs
       if assigns is not None: assigns.append((f'MEM[{_tok_str(addr_toks)}].{dt_name}', (addr, rhs)))
       i += 1
       continue
@@ -1188,7 +1206,11 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       old = block_assigns.get(var, env.get(var, _u32(0)))
       rhs = parse_tokens(toks[assign_op+1:], env, funcs)
       if rhs.dtype != old.dtype: rhs = rhs.cast(old.dtype)
-      block_assigns[var] = env[var] = (old + rhs) if toks[assign_op].val == '+=' else (old - rhs)
+      if toks[assign_op].val == '+=': block_assigns[var] = env[var] = old + rhs
+      elif toks[assign_op].val == '-=': block_assigns[var] = env[var] = old - rhs
+      elif toks[assign_op].val == '|=': block_assigns[var] = env[var] = old | rhs
+      elif toks[assign_op].val == '&=': block_assigns[var] = env[var] = old & rhs
+      elif toks[assign_op].val == '^=': block_assigns[var] = env[var] = old ^ rhs
       i += 1
       continue
 
@@ -1335,4 +1357,3 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
-

test/mockgpu/helpers.py

@@ -1,5 +1,4 @@
 import ctypes, ctypes.util
-from tinygrad.helpers import getenv
 
 def _try_dlopen_gpuocelot():
   GPUOCELOT_PATHS = [ctypes.util.find_library("gpuocelot")] if ctypes.util.find_library("gpuocelot") is not None else []
@@ -16,7 +15,7 @@ def _try_dlopen_gpuocelot():
   return None
 
 class PythonRemu:
-  """Python RDNA3/RDNA4 emulator wrapper that matches the libremu.so interface."""
+  """Python RDNA3/RDNA4 emulator wrapper used by mockgpu."""
   valid_mem_ranges: set[tuple[int, int]] = set()
   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
@@ -26,20 +25,3 @@ class PythonRemu:
   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 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
-  if int(getenv("PYTHON_REMU", "1")):
-    return PythonRemu()
-  REMU_PATHS = ["extra/remu/target/release/libremu.so", "libremu.so", "/usr/local/lib/libremu.so",
-               "extra/remu/target/release/libremu.dylib", "libremu.dylib", "/usr/local/lib/libremu.dylib", "/opt/homebrew/lib/libremu.dylib"]
-  for path in REMU_PATHS:
-    try:
-      remu = ctypes.CDLL(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]
-    except OSError: pass
-    else: return remu
-  print("Could not find libremu.so")
-  return None

test/null/test_autogen.py

@@ -1,7 +1,6 @@
 import ctypes, struct, subprocess, tempfile, unittest
-from typing import Annotated
 from tinygrad.helpers import OSX, WIN
-from tinygrad.runtime.support.c import DLL, record, init_records
+from tinygrad.runtime.support.c import DLL, record, Field
 from tinygrad.runtime.support import c
 from tinygrad.runtime.support.autogen import gen
 
@@ -12,15 +11,56 @@ class TestC(unittest.TestCase):
       subprocess.check_output(('clang', '-x', 'c', '-fPIC', '-shared', '-', '-o', f.name), input=src.encode())
       return DLL("test", f.name)
 
+  def test_struct_array_init(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 12
+      a = Field(ctypes.c_int * 3, 0)
+
+    f = Foo((1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+    f = Foo((ctypes.c_int * 3)(1,2,3))
+    assert f.a[0] == 1
+    assert f.a[1] == 2
+    assert f.a[2] == 3
+
+  def test_field_ranges(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 2
+      s = Field(ctypes.c_int8, 0)
+      u = Field(ctypes.c_uint8, 1)
+
+    f = Foo()
+    f.s = -1
+    f.u = -1
+    assert f.s == -1
+    assert f.u == 255
+
+  # this syntax is inherited from ctypes, but it seems a bit nonsensical?
+  def test_voidp_none(self):
+    @record
+    class Foo(c.Struct):
+      SIZE = 8
+      p = Field(ctypes.c_void_p, 0)
+
+    f = Foo(None)
+    assert f.p is None
+    f.p = ctypes.c_void_p(0xDEADBEEF)
+    assert f.p == 0xDEADBEEF
+    f.p = None
+    assert f.p is None
+
   def test_packed_struct(self):
     @record
-    class Baz:
+    class Baz(c.Struct):
       SIZE = 8
-      a: Annotated[ctypes.c_uint, 0, 30]
-      b: Annotated[ctypes.c_uint, 3, 30, 6]
-      c: Annotated[ctypes.c_uint, 7, 2, 4]
-      d: Annotated[ctypes.c_uint, 7, 2, 6]
-    init_records()
+      a = Field(ctypes.c_uint, 0, 30)
+      b = Field(ctypes.c_uint, 3, 30, 6)
+      c = Field(ctypes.c_uint, 7, 2, 4)
+      d = Field(ctypes.c_uint, 7, 2, 6)
 
     b = Baz(0x3AAADEAD, 0xBEEF, 1, 0)
     assert b.a == 0x3AAADEAD
@@ -36,13 +76,12 @@ class TestC(unittest.TestCase):
 
   def test_packed_struct_interop(self):
     @record
-    class Baz:
+    class Baz(c.Struct):
       SIZE = 8
-      a: Annotated[ctypes.c_int, 0, 30]
-      b: Annotated[ctypes.c_int, 3, 30, 6]
-      c: Annotated[ctypes.c_int, 7, 2, 4]
-      d: Annotated[ctypes.c_int, 7, 2, 6]
-    init_records()
+      a = Field(ctypes.c_int, 0, 30)
+      b = Field(ctypes.c_int, 3, 30, 6)
+      c = Field(ctypes.c_int, 7, 2, 4)
+      d = Field(ctypes.c_int, 7, 2, 6)
 
     src = '''
       struct __attribute__((packed)) baz {
@@ -58,24 +97,23 @@ class TestC(unittest.TestCase):
     '''
     dll = self.compile(src)
     b = Baz(0xAA000, 0x00BB0, 0, 1)
-    @dll.bind
+    @dll.bind(ctypes.c_int, Baz)
     def test(x:Baz) -> ctypes.c_int: ...
     self.assertEqual(test(b), b.a + b.b + b.c + b.d)
 
   # https://github.com/python/cpython/issues/90914
   def test_bitfield_interop(self):
     @record
-    class Baz:
+    class Baz(c.Struct):
       SIZE = 1
-      a: Annotated[ctypes.c_bool, 0, 1, 0]
-      b: Annotated[ctypes.c_bool, 0, 1, 1]
-      c: Annotated[ctypes.c_bool, 0, 1, 2]
-      d: Annotated[ctypes.c_bool, 0, 1, 3]
-      e: Annotated[ctypes.c_bool, 0, 1, 4]
-      f: Annotated[ctypes.c_bool, 0, 1, 5]
-      g: Annotated[ctypes.c_bool, 0, 1, 6]
-      h: Annotated[ctypes.c_bool, 0, 1, 7]
-    init_records()
+      a = Field(ctypes.c_bool, 0, 1, 0)
+      b = Field(ctypes.c_bool, 0, 1, 1)
+      c = Field(ctypes.c_bool, 0, 1, 2)
+      d = Field(ctypes.c_bool, 0, 1, 3)
+      e = Field(ctypes.c_bool, 0, 1, 4)
+      f = Field(ctypes.c_bool, 0, 1, 5)
+      g = Field(ctypes.c_bool, 0, 1, 6)
+      h = Field(ctypes.c_bool, 0, 1, 7)
     src = '''#include <stdbool.h>
       struct baz {
         bool a:1, b:1, c:1, d:1, e:1, f:1, g:1, h:1;
@@ -86,23 +124,22 @@ class TestC(unittest.TestCase):
       }
     '''
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(ctypes.c_int, Baz)
     def test(x:Baz) -> ctypes.c_int: ...
     for i in range(8): self.assertEqual(test(Baz(*(j==i for j in range(8)))), i==2)
 
   def test_struct_interop(self):
     @record
-    class Baz:
+    class Baz(c.Struct):
       SIZE = 32
-      a: Annotated[ctypes.c_int, 0]
-      b: Annotated[ctypes.c_int, 4]
-      c: Annotated[ctypes.c_int, 8]
-      d: Annotated[ctypes.c_int, 12]
-      e: Annotated[ctypes.c_int, 16]
-      f: Annotated[ctypes.c_int, 20]
-      g: Annotated[ctypes.c_int, 24]
-      h: Annotated[ctypes.c_int, 28]
-    init_records()
+      a = Field(ctypes.c_int, 0)
+      b = Field(ctypes.c_int, 4)
+      c = Field(ctypes.c_int, 8)
+      d = Field(ctypes.c_int, 12)
+      e = Field(ctypes.c_int, 16)
+      f = Field(ctypes.c_int, 20)
+      g = Field(ctypes.c_int, 24)
+      h = Field(ctypes.c_int, 28)
     src = '''#include <stdio.h>
       struct baz {
         int a, b, c, d, e, f, g, h;
@@ -113,16 +150,15 @@ class TestC(unittest.TestCase):
       }
     '''
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(Baz, Baz)
     def test(x:Baz) -> Baz: ...
     self.assertEqual(bytes(test(Baz(*range(8)))), struct.pack("8i", *range(7, -1, -1)))
 
   def test_aos_interop(self):
     @record
-    class Item:
+    class Item(c.Struct):
       SIZE = 4
-      val: Annotated[ctypes.c_int, 0]
-    init_records()
+      val = Field(ctypes.c_int, 0)
     src = """
     struct item { int val; };
       int test(struct item arr[3]) {
@@ -132,16 +168,15 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(ctypes.c_int, Item * 3)
     def test(arr:(Item * 3)) -> ctypes.c_int: ...
     self.assertEqual(test((Item * 3)(Item(10), Item(20), Item(30))), 60)
 
   def test_soa_interop(self):
     @record
-    class Row:
+    class Row(c.Struct):
       SIZE = 16
-      data: Annotated[ctypes.c_int * 3, 0]
-    init_records()
+      data = Field(ctypes.c_int * 3, 0)
     src = """
       struct row { int data[3]; };
       struct row test(struct row x) {
@@ -149,7 +184,7 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(Row, Row)
     def test(x:Row) -> Row: ...
     r = test(Row((ctypes.c_int * 3)(10, 20, 30)))
     self.assertIsInstance(r, Row)
@@ -159,10 +194,9 @@ class TestC(unittest.TestCase):
 
   def test_soa_ptr_interop(self):
     @record
-    class Row:
+    class Row(c.Struct):
       SIZE = 8
-      data: Annotated[c.POINTER[ctypes.c_int], 0]
-    init_records()
+      data = Field(c.POINTER[ctypes.c_int], 0)
     src = """
       struct row { int *data; };
       int test(struct row x) {
@@ -170,21 +204,20 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(ctypes.c_int, Row)
     def test(x:Row) -> ctypes.c_int: ...
     assert test(Row((ctypes.c_int * 3)(10, 20, 30))) == 60
 
   def test_nested_struct_interop(self):
     @record
-    class Inner:
+    class Inner(c.Struct):
       SIZE = 4
-      a: Annotated[ctypes.c_int, 0]
+      a = Field(ctypes.c_int, 0)
     @record
-    class Outer:
+    class Outer(c.Struct):
       SIZE = 8
-      inner: Annotated[Inner, 0]
-      b: Annotated[ctypes.c_int, 4]
-    init_records()
+      inner = Field(Inner, 0)
+      b = Field(ctypes.c_int, 4)
     src = """
       struct i { int a; };
       struct o { struct i i; int b; };
@@ -193,7 +226,7 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(Outer, Outer)
     def test(x:Outer) -> Outer: ...
     o = test(Outer(Inner(10), 20))
     self.assertEqual(o.inner.a, 20)
@@ -201,11 +234,10 @@ class TestC(unittest.TestCase):
 
   def test_struct_pointer_interop(self):
     @record
-    class Foo:
+    class Foo(c.Struct):
       SIZE = 8
-      a: Annotated[ctypes.c_int, 0]
-      b: Annotated[ctypes.c_int, 4]
-    init_records()
+      a = Field(ctypes.c_int, 0)
+      b = Field(ctypes.c_int, 4)
     src = """
       struct foo { int a, b; };
       struct foo *test(struct foo *f) {
@@ -216,7 +248,7 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(ctypes.POINTER(Foo), ctypes.POINTER(Foo))
     def test(f:ctypes.POINTER(Foo)) -> ctypes.POINTER(Foo): ...
     inp = ctypes.pointer(Foo(10, 20))
     out = test(inp)
@@ -228,16 +260,15 @@ class TestC(unittest.TestCase):
     # Mimics how mesa.struct_lp_build_tgsi_params.mask is used
     from tinygrad.runtime.support.c import POINTER
     @record
-    class Inner:
+    class Inner(c.Struct):
       SIZE = 8
-      value: Annotated[ctypes.c_int, 0]
-      flag: Annotated[ctypes.c_int, 4]
+      value = Field(ctypes.c_int, 0)
+      flag = Field(ctypes.c_int, 4)
     @record
-    class Outer:
+    class Outer(c.Struct):
       SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      inner_ptr: Annotated[POINTER[Inner], 8]
-    init_records()
+      x = Field(ctypes.c_int, 0)
+      inner_ptr = Field(POINTER[Inner], 8)
 
     src = """
       struct inner { int value; int flag; };
@@ -247,7 +278,7 @@ class TestC(unittest.TestCase):
       }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(ctypes.c_int, ctypes.POINTER(Inner))
     def test(p:POINTER[Inner]) -> ctypes.c_int: ...
 
     inner = Inner(value=42, flag=10)
@@ -261,17 +292,16 @@ class TestC(unittest.TestCase):
     # This causes the pointed-to object to be garbage collected, leading to use-after-free.
     from tinygrad.runtime.support.c import POINTER
     @record
-    class MaskContext:
+    class MaskContext(c.Struct):
       SIZE = 16
-      value: Annotated[ctypes.c_int, 0]
-      initialized: Annotated[ctypes.c_int, 4]
-      ptr: Annotated[ctypes.c_void_p, 8]
+      value = Field(ctypes.c_int, 0)
+      initialized = Field(ctypes.c_int, 4)
+      ptr = Field(ctypes.c_void_p, 8)
     @record
-    class Params:
+    class Params(c.Struct):
       SIZE = 16
-      x: Annotated[ctypes.c_int, 0]
-      mask: Annotated[POINTER[MaskContext], 8]
-    init_records()
+      x = Field(ctypes.c_int, 0)
+      mask = Field(POINTER[MaskContext], 8)
 
     src = """
       struct mask_ctx { int value; int initialized; void *ptr; };
@@ -279,9 +309,9 @@ class TestC(unittest.TestCase):
       int mask_end(struct mask_ctx *m) { return m->value + m->initialized; }
     """
     dll = self.compile(src)
-    @dll.bind
+    @dll.bind(None, ctypes.POINTER(MaskContext), ctypes.c_int)
     def mask_begin(m:POINTER[MaskContext], val:ctypes.c_int) -> None: ...
-    @dll.bind
+    @dll.bind(ctypes.c_int, ctypes.POINTER(MaskContext))
     def mask_end(m:POINTER[MaskContext]) -> ctypes.c_int: ...
 
     # When MaskContext() is created inline, it gets garbage collected after the pointer
@@ -399,6 +429,10 @@ typedef struct
     self.assertTrue(hasattr(rect, 'height'))
     self.assertTrue(hasattr(rect, 'color'))
 
+    p2 = Point(10, 20)
+    self.assertEqual(p2.x, 10)
+    self.assertEqual(p2.y, 20)
+
   def test_struct_ordering(self):
     namespace = self.run_gen("""
     struct A;

test/null/test_device.py

@@ -116,7 +116,9 @@ class TestDevVar(unittest.TestCase):
     for d, t in [("AMD", Target(device="AMD", renderer="")), ("AMD:LLVM", Target(device="AMD", renderer="LLVM")),
                  (":LLVM", Target(device="", renderer="LLVM")), ("AMD::gfx1100", Target(device="AMD", arch="gfx1100")),
                  ("AMD:LLVM:gfx1100", Target(device="AMD", renderer="LLVM", arch="gfx1100")), ("::gfx1100", Target(arch="gfx1100")),
-                 ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB"))]:
+                 ("USB+", Target(interface="USB")), ("USB+AMD", Target(device="AMD", interface="USB")),
+                 ("PCI:0+AMD", Target(device="AMD", interface="PCI", indices="0")), (":0+AMD", Target(device="AMD", indices="0")),
+                 ("PCI:0,1+AMD", Target(device="AMD", interface="PCI", indices="0,1"))]:
       with Context(DEV=d):
         self.assertEqual(DEV.value, t)
         self.assertEqual(str(DEV.value), d)

test/null/test_llm_server.py

@@ -10,7 +10,9 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_tok.role = Mock(return_value=[100, 101])
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
+    cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
     cls.mock_tok.end_turn = Mock(return_value=[998])
+    cls.mock_tok.preset = "llama3"
 
     cls.mock_model = Mock()
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))

test/null/test_llm_tokenizer.py

@@ -46,5 +46,26 @@ class TestLLMTokenizer(unittest.TestCase):
   def test_llama_repeat(self): self._test_coding(self.llama_tok, "00000000000000000", [ 931, 931, 931, 931, 931, 410 ])
   def test_llama_pat(self): self._test_coding(self.llama_tok, "today\n  \n", [ 31213, 14211 ])
 
+  def test_tekken_from_gguf_kv(self):
+    kv = {
+      "tokenizer.ggml.tokens": ["<unk>", "<s>", "</s>", "[INST]", "[/INST]", "hello"],
+      "tokenizer.ggml.token_type": [3, 3, 3, 3, 3, 1],
+      "tokenizer.ggml.pre": "tekken",
+    }
+    tok = SimpleTokenizer.from_gguf_kv(kv)
+    self.assertEqual(tok.role("user"), [3])
+    self.assertEqual(tok.encode("hello"), [5])
+    self.assertEqual(tok.end_turn(2), [4])
+    self.assertEqual(tok.role("assistant"), [])
+
+  def test_stream_decoder(self):
+    """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""
+    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]
+    be = {b: chr(b) for b in bs} | {b: chr(256+i) for i,b in enumerate(b for b in range(256) if b not in bs)}
+    token_bytes = {25677: b'\x20\xf0\x9f\x98', 138: b'\x8a'}  # ' ' + 3/4 emoji | 1/4 emoji (qwen3.5)
+    tok = SimpleTokenizer({"".join(be[b] for b in v): k for k, v in token_bytes.items()}, {})
+    dec = tok.stream_decoder()
+    self.assertEqual(dec(25677) + dec(138) + dec(), " 😊")
+
 if __name__ == '__main__':
   unittest.main()

test/null/test_memory_planner.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import dtypes
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.memory import memory_plan_rewrite
+from tinygrad.schedule.memory import memory_plan_rewrite
 
 global_map = {}
 held_bufs: set[UOp] = set()

test/null/test_multitensor.py

@@ -1,6 +1,7 @@
 import gc, unittest
 from tinygrad import Tensor, GlobalCounters, dtypes
 from tinygrad.engine.jit import TinyJit
+from tinygrad.helpers import Context
 
 class TestMultiRamUsage(unittest.TestCase):
   def setUp(self):
@@ -139,6 +140,19 @@ class TestMultiRamUsage(unittest.TestCase):
     mem_4 = run_layers(4)
     self.assertEqual(mem_2, mem_4, f"graph memory should not grow with layers: 2 layers={mem_2}, 4 layers={mem_4}")
 
+  def test_allreduce_cast_dtype_memory(self):
+    N = 32
+    devices_2 = ("NULL:1", "NULL:2")
+    mem = {}
+    for allreduce_cast in (0, 1):
+      GlobalCounters.reset()
+      with Context(ALLREDUCE_CAST=allreduce_cast, SCACHE=0):
+        x = Tensor.empty((N, N), dtype=dtypes.bfloat16, device="NULL:1").shard(devices_2, axis=0)
+        x.sum(0).realize()
+      mem[allreduce_cast] = GlobalCounters.global_mem
+    # with ALLREDUCE_CAST, allreduce copies happen in bf16 (2 bytes) instead of fp32 (4 bytes)
+    self.assertLess(mem[1], mem[0])
+
 class TestMultiAxis(unittest.TestCase):
   def test_reshape_shard_invalid(self):
     devices = ("NULL:0", "NULL:1")

test/null/test_schedule.py

@@ -1,7 +1,7 @@
 # schedule tests that pass on NULL backend (no copyout needed)
 import gc, unittest, time
 from tinygrad import nn, dtypes, Device, Tensor
-from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
+from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo
 from tinygrad.helpers import DEBUG, GlobalCounters, Context
 from tinygrad.engine.realize import CompiledRunner, run_schedule
 
@@ -143,6 +143,36 @@ class TestSimpleSchedule(unittest.TestCase):
     self.assertEqual(len(Tensor.schedule(a1, a2)), 1)
 
 class TestSchedule(unittest.TestCase):
+  def test_create_schedule_handles_multi_kernel_after_and_after_deps(self):
+    def named_copy(name:str):
+      def fxn(out:UOp, src:UOp) -> UOp:
+        i = UOp.range(src.size, 0)
+        return out[i].store(src[i]).end(i).sink(arg=KernelInfo(name=name))
+      return fxn
+
+    src = Tensor.zeros(4, dtype=dtypes.float).contiguous().realize()
+    dep = Tensor.zeros(4, dtype=dtypes.float).contiguous().realize()
+    out = Tensor.zeros(4, dtype=dtypes.float).contiguous().realize()
+    ones = Tensor.ones(4, dtype=dtypes.float).contiguous().realize()
+    twos = Tensor.full((4,), 2.0, dtype=dtypes.float).contiguous().realize()
+    threes = Tensor.full((4,), 3.0, dtype=dtypes.float).contiguous().realize()
+
+    ka = Tensor.custom_kernel(src, ones, fxn=named_copy("ka"))[0]
+    kb = Tensor.custom_kernel(src, twos, fxn=named_copy("kb"))[0]
+    src_after = Tensor(src.uop.after(*ka.uop.src[1:], *kb.uop.src[1:]))
+
+    kd = Tensor.custom_kernel(dep, threes, fxn=named_copy("kd"))[0]
+    kc = Tensor.custom_kernel(out, src_after, fxn=named_copy("kc"))[0]
+    out_after = Tensor(kc.uop.src[0].after(*kc.uop.src[1:], kd.uop))
+
+    schedule = out_after.schedule()
+    names = [si.ast.arg.name for si in schedule]
+    self.assertEqual(set(names), {"ka", "kb", "kc", "kd"})
+    self.assertEqual(names[-1], "kc")
+    self.assertLess(names.index("ka"), names.index("kc"))
+    self.assertLess(names.index("kb"), names.index("kc"))
+    self.assertLess(names.index("kd"), names.index("kc"))
+
   @unittest.skipIf(Device.DEFAULT == "CPU", "devices must mismatch")
   def test_error_on_device_mismatch(self):
     a = Tensor.empty(10)

test/null/test_schedule_cache.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Variable, Context
 from tinygrad.helpers import cpu_events
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 
 def schedule_one():
   Tensor([1]).schedule()

test/null/test_tensor_metadata.py

@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, dtypes
 from tinygrad.tensor import _METADATA
 from tinygrad.engine.realize import capturing
-from tinygrad.engine.schedule import linear_to_schedule
+from tinygrad.schedule import linear_to_schedule
 from tinygrad.helpers import Context
 
 @unittest.skip("tensor metadata is no longer supported")

test/null/test_tensor_uop_mixin.py

@@ -0,0 +1,74 @@
+import math, unittest
+from tinygrad import Tensor
+
+def _t(*shape):
+  return Tensor.arange(math.prod(shape)).reshape(*shape)
+
+# Tensor().func().uop should be the same as UOp.func()
+def _check(tc: unittest.TestCase, t: Tensor, fn):
+  tc.assertIs(fn(t).uop, fn(t.uop), f"\ntensor.uop = {fn(t).uop}\nuop = {fn(t.uop)}")
+
+class TestTensorUOpGetitem(unittest.TestCase):
+  # ---- pure slice patterns ----
+  def test_slice_full(self):           _check(self, _t(4), lambda x: x[slice(None)])
+  def test_slice_positive(self):       _check(self, _t(8), lambda x: x[1:5])
+  def test_slice_open_start(self):     _check(self, _t(8), lambda x: x[:5])
+  def test_slice_open_stop(self):      _check(self, _t(8), lambda x: x[3:])
+  def test_slice_negative_start(self): _check(self, _t(8), lambda x: x[-3:])
+  def test_slice_negative_stop(self):  _check(self, _t(8), lambda x: x[:-2])
+  def test_slice_both_negative(self):  _check(self, _t(8), lambda x: x[-5:-1])
+
+  # ---- slice with stride ----
+  def test_slice_stride(self):                  _check(self, _t(6), lambda x: x[::2])
+  def test_slice_start_stop_stride(self):       _check(self, _t(6), lambda x: x[1:5:2])
+  def test_slice_reverse(self):                 _check(self, _t(6), lambda x: x[::-1])
+  def test_slice_singleton_negative_step(self): _check(self, _t(8), lambda x: x[3:2:-1])
+
+  # ---- empty / out-of-bounds slice ----
+  def test_slice_empty(self):    _check(self, _t(6), lambda x: x[3:1])
+  def test_slice_oob_stop(self): _check(self, _t(6), lambda x: x[0:100])
+
+  # ---- single int (reduces a dim) ----
+  def test_int_positive(self): _check(self, _t(8), lambda x: x[3])
+  def test_int_negative(self): _check(self, _t(8), lambda x: x[-1])
+
+  # ---- ellipsis ----
+  def test_ellipsis_only(self):       _check(self, _t(2, 3, 4), lambda x: x[...])
+  def test_ellipsis_then_int(self):   _check(self, _t(2, 3, 4), lambda x: x[..., -1])
+  def test_ellipsis_then_slice(self): _check(self, _t(2, 3, 4), lambda x: x[..., 1:3])
+  def test_ellipsis_then_none(self):  _check(self, _t(2, 3), lambda x: x[..., None])
+
+  # ---- None (unsqueeze) ----
+  def test_none_front(self):    _check(self, _t(4), lambda x: x[None])
+  def test_none_back(self):     _check(self, _t(4), lambda x: x[:, None])
+  def test_none_middle(self):   _check(self, _t(2, 3), lambda x: x[:, None, :])
+  def test_multiple_none(self): _check(self, _t(2, 3), lambda x: x[None, :, None])
+
+  # ---- mixed multi-dim ----
+  def test_int_then_slice(self):    _check(self, _t(2, 3), lambda x: x[1, :])
+  def test_multi_int(self):         _check(self, _t(2, 3, 4), lambda x: x[1, 2])
+  def test_mixed_slice_int(self):   _check(self, _t(2, 3, 4), lambda x: x[0:2, -1, 1:3])
+  def test_mixed_slice_slice(self): _check(self, _t(3, 4, 5), lambda x: x[1:3, :, 0:2])
+  def test_high_rank_combo(self):   _check(self, _t(4, 5, 6), lambda x: x[1:3, :, -1, None])
+
+class TestTensorUOpCumalu(unittest.TestCase):
+  def test_cumsum_1d(self):       _check(self, _t(5), lambda x: x.cumsum())
+  def test_cumsum_2d(self):       _check(self, _t(3, 4), lambda x: x.cumsum(1))
+  def test_cumsum_non_last(self): _check(self, _t(3, 4), lambda x: x.cumsum(0))
+  def test_cumsum_large(self):    _check(self, _t(600), lambda x: x.cumsum())  # exercises _split_cumalu
+  def test_cumprod(self):         _check(self, _t(4), lambda x: x.cumprod(0))
+
+class TestTensorUOpCat(unittest.TestCase):
+  def test_cat_dim0(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=0))
+  def test_cat_dim1(self):     _check(self, _t(2, 3), lambda x: x.cat(x, dim=1))
+  def test_cat_3tensors(self): _check(self, _t(2, 3), lambda x: x.cat(x, x, dim=0))
+  def test_cat_neg_dim(self):  _check(self, _t(2, 3, 4), lambda x: x.cat(x, dim=-1))
+
+class TestTensorUOpStack(unittest.TestCase):
+  def test_stack_dim0(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=0))
+  def test_stack_dim1(self):     _check(self, _t(2, 3), lambda x: x.stack(x, dim=1))
+  def test_stack_3tensors(self): _check(self, _t(2, 3), lambda x: x.stack(x, x, dim=0))
+  def test_stack_new_last(self): _check(self, _t(2, 3), lambda x: x.stack(x, dim=-1))
+
+if __name__ == "__main__":
+  unittest.main()

test/null/test_viz.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import VIZ, cpu_profile, ProfilePointEvent, unwrap
 from tinygrad.device import Buffer
 
 from tinygrad.uop.ops import tracked_keys, tracked_ctxs, uop_fields, active_rewrites, active_group, _name_cnt, RewriteTrace
-from tinygrad.viz.serve import get_rewrites, get_full_rewrite, uop_to_json
+from tinygrad.viz.serve import load_rewrites, get_full_rewrite, uop_to_json, VizData
 
 @track_rewrites(name=True)
 def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=None) -> UOp:
@@ -21,19 +21,19 @@ def exec_rewrite(sink:UOp, pm_lst:list[PatternMatcher], names:None|list[str]=Non
 # small container class for the viz server module
 class VizTrace:
   # loader init
-  def __init__(self): self._trace:RewriteTrace|None = None
+  def __init__(self): self._data:VizData|None = None
   @property
-  def trace(self) -> RewriteTrace: return unwrap(self._trace)
-  def set_trace(self) -> None:
-    self._trace = RewriteTrace(tracked_keys.copy(), tracked_ctxs.copy(), uop_fields.copy())
-    import tinygrad.viz.serve as serve_module
-    serve_module.trace = self._trace
+  def data(self) -> VizData: return unwrap(self._data)
+  def set_data(self) -> None:
+    data = VizData(RewriteTrace(tracked_keys.copy(), tracked_ctxs.copy(), uop_fields.copy()))
+    load_rewrites(data)
+    self._data = data
   # the API
-  def list_items(self) -> list[dict]: return get_rewrites(self.trace)
+  def list_items(self) -> list[dict]:
+    return self.data.ctxs
   def get_details(self, rewrite_idx:int, step:int) -> Generator[dict, None, None]:
-    lst = self.list_items()
-    assert len(lst) > rewrite_idx, f"only loaded {len(lst)} traces, expecting at least {rewrite_idx}"
-    return get_full_rewrite(self.trace.rewrites[rewrite_idx][step])
+    assert len(self.data.trace.rewrites) > rewrite_idx, f"only loaded {len(self.data.trace.rewrites)} traces, expecting at least {rewrite_idx}"
+    return get_full_rewrite(self.data, self.data.trace.rewrites[rewrite_idx][step])
 
 @contextlib.contextmanager
 def save_viz():
@@ -52,7 +52,7 @@ def save_viz():
   try:
     yield viz
   finally:
-    viz.set_trace()
+    viz.set_data()
     TRACK_MATCH_STATS.value = prev_tms
     PROFILE.value = prev_profile
     VIZ.value = prev_viz
@@ -194,7 +194,7 @@ class TestViz(unittest.TestCase):
     class TestStruct:
       colored_field: str
     a = UOp(Ops.CUSTOM, arg=TestStruct(colored("xyz", "magenta")+colored("12345", "blue")))
-    a2 = uop_to_json(a)[id(a)]
+    a2 = uop_to_json(VizData(), a)[id(a)]
     self.assertEqual(ansistrip(a2["label"]), f"CUSTOM\n{TestStruct.__qualname__}(colored_field='xyz12345')")
 
   def test_colored_label_multiline(self):
@@ -217,11 +217,11 @@ class TestViz(unittest.TestCase):
       # use smaller stack limit for faster test (default is 250000)
       with Context(REWRITE_STACK_LIMIT=100): self.assertRaises(RuntimeError, exec_rewrite, a, [pm])
     graphs = flatten(x["graph"].values() for x in viz.get_details(0, 0))
-    self.assertEqual(graphs[0], uop_to_json(a)[id(a)])
-    self.assertEqual(graphs[1], uop_to_json(b)[id(b)])
+    self.assertEqual(graphs[0], uop_to_json(VizData(), a)[id(a)])
+    self.assertEqual(graphs[1], uop_to_json(VizData(), b)[id(b)])
     # fallback to NOOP with the error message
     nop = UOp(Ops.NOOP, arg="infinite loop in fixed_point_rewrite")
-    self.assertEqual(graphs[2], uop_to_json(nop)[id(nop)])
+    self.assertEqual(graphs[2], uop_to_json(VizData(), nop)[id(nop)])
 
   def test_const_node_visibility(self):
     with save_viz() as viz:
@@ -241,7 +241,7 @@ class TestViz(unittest.TestCase):
     c = UOp.const(dtypes.float, 1.0, device="CPU", shape=(3,4))  # creates CONST->RESHAPE->EXPAND chain
     a = UOp(Ops.DEFINE_VAR, dtypes.float, arg=("a", 0.0, 10.0))
     alu = a + c
-    graph = uop_to_json(alu)
+    graph = uop_to_json(VizData(), alu)
     # the RESHAPE and EXPAND nodes from the const should not appear in the graph
     labels = {v["label"].split("\n")[0] for v in graph.values()}
     self.assertNotIn("RESHAPE", labels)
@@ -335,7 +335,7 @@ class TestVizIntegration(unittest.TestCase):
       prg = get_program(ast, Device[Device.DEFAULT].renderer)
     lst = viz.list_items()
     self.assertEqual(len(lst), 3)
-    self.assertEqual(lst[0]["name"], "Process 1 Buffer n1")
+    self.assertEqual(lst[0]["name"], "Callify 1 Buffer n1")
     self.assertEqual(lst[1]["name"], "Schedule 1 Kernel n1")
     self.assertEqual(lst[2]["name"], prg.name)
 
@@ -417,7 +417,7 @@ from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent, ProfileGraphE
 from tinygrad.viz.serve import get_profile
 from extra.viz.cli import decode_profile
 
-def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(lst))
+def load_profile(lst:list[ProfileEvent]) -> dict: return decode_profile(get_profile(VizData(), lst))
 
 class TestVizProfiler(unittest.TestCase):
   def test_transfer_uses_copy_device(self):
@@ -470,7 +470,7 @@ class TestVizProfiler(unittest.TestCase):
     j = load_profile(prof)
     event = j['layout']['NV:SDMA:0']['events'][0]
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s")
+    self.assertEqual(event['fmt'], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
 
   def test_graph(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -512,7 +512,7 @@ class TestVizProfiler(unittest.TestCase):
     j = load_profile(prof)
     sdma_events = j['layout']['NV:1:SDMA:0']['events']
     gbs = sz/(dur*1e-6)*1e-9
-    self.assertEqual(sdma_events[0]['fmt'], f"{gbs:.0f} GB/s")
+    self.assertEqual(sdma_events[0]["fmt"], f"{gbs:.0f} GB/s\n{sz/1e6:.0f} MB")
 
   def test_block_ordering(self):
     prof = [ProfileDeviceEvent(device='NV', tdiff=decimal.Decimal(-1000)),
@@ -563,7 +563,7 @@ class TestVizProfiler(unittest.TestCase):
     step = 10
     n_events = 1_000
     prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
-    sz = len(get_profile(prof))
+    sz = len(get_profile(VizData(), prof))
     self.assertLessEqual(sz/n_events, 26)
 
   def test_calltrace(self):
@@ -586,7 +586,7 @@ class TestVizProfiler(unittest.TestCase):
     step = decimal.Decimal(dur_mins*60*1e6//n_events)
     prof = [ProfileRangeEvent("CPU", name="k_test", st=decimal.Decimal(ts:=i*step), en=decimal.Decimal(ts)+step) for i in range(n_events)]
     with self.assertRaisesRegex(ValueError, "timestamp out of range"):
-      get_profile(prof)
+      get_profile(VizData(), prof)
 
   def test_python_marker(self):
     with save_viz():

test/testextra/test_tk.py

@@ -3,7 +3,7 @@ import unittest, math, time
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.engine.realize import get_runner
-from tinygrad.engine.schedule import ExecItem
+from tinygrad.schedule import ExecItem
 from tinygrad.engine.jit import TinyJit
 import numpy as np
 

test/unit/test_allreduce.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad import Tensor
+from tinygrad import Tensor, dtypes
 from tinygrad.helpers import Context
 from tinygrad.uop.ops import Ops
 
@@ -25,5 +25,28 @@ class TestRingAllReduce(unittest.TestCase):
       out = t.sum(0)
       self.assertListEqual(out.tolist(), [4]*N*100)
 
+class TestAllreduceCast(unittest.TestCase):
+  def _get_copy_dtypes(self, dtype, allreduce_cast):
+    ds = tuple(f"CPU:{i}" for i in range(2))
+    with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
+      t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
+      schedules = t.sum(0).schedule_with_vars()[0]
+      return {si.bufs[0].dtype.scalar() for si in schedules if si.ast.op is Ops.COPY}
+
+  def test_allreduce_cast_bf16(self):
+    # with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32
+    self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=1))
+    self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=0))
+
+  def test_allreduce_cast_half(self):
+    self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=1))
+    self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=0))
+
+  def test_allreduce_cast_float32_noop(self):
+    # float32 should not be affected by ALLREDUCE_CAST (no promotion happens)
+    dtypes_on = self._get_copy_dtypes(dtypes.float, allreduce_cast=1)
+    dtypes_off = self._get_copy_dtypes(dtypes.float, allreduce_cast=0)
+    self.assertEqual(dtypes_on, dtypes_off)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_assign.py

@@ -947,5 +947,38 @@ class TestPartialAssignToSharedBuffer(unittest.TestCase):
     for v, s in zip(views, shapes):
       np.testing.assert_allclose(v.numpy(), np.ones(s))
 
+
+class TestAfterCachePatterns(unittest.TestCase):
+  def test_double_store_after(self):
+    a = Tensor.zeros(10).contiguous()
+    b = Tensor.zeros(10).contiguous()
+    c = Tensor.ones(10).contiguous()
+    Tensor.realize(a, b, c)
+
+    a_store = a.uop.store(c.uop)
+    b_store = b.uop.store(c.uop)
+
+    with self.assertRaises(AssertionError):
+      a = Tensor(a.uop.after(a_store, b_store))
+      a.realize()
+      np.testing.assert_array_equal(a.numpy(), 1)
+      np.testing.assert_array_equal(b.numpy(), 1)
+
+  def test_double_store_after_different_sizes(self):
+    full = Tensor.zeros(2).contiguous()
+    head = Tensor.zeros(1).contiguous()
+    full_src = Tensor([1, 2], dtype=dtypes.float).contiguous()
+    head_src = Tensor([3], dtype=dtypes.float).contiguous()
+    Tensor.realize(full, head, full_src, head_src)
+
+    full_store = full.uop.store(full_src.uop)
+    head_store = head.uop.store(head_src.uop)
+
+    with self.assertRaises(AssertionError):
+      head = Tensor(head.uop.after(head_store, full_store))
+      head.realize()
+      np.testing.assert_array_equal(head.numpy(), [3])
+      np.testing.assert_array_equal(full.numpy(), [1, 2])
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_attention.py

@@ -1,7 +1,10 @@
 import unittest
 import numpy as np
 from tinygrad import Tensor, dtypes
-from tinygrad.apps.llm import apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk
+from tinygrad.apps.llm import (
+  GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
+  apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
+)
 
 def apply_rope(x:Tensor, start_pos:int):
   B, H, T, Hd = x.shape
@@ -18,6 +21,158 @@ class TestAttention(unittest.TestCase):
     self.assertGreater((result - apply_rope(x, 5)).abs().max().item(), 1e-6)
     with self.assertRaises(AssertionError): apply_rope(Tensor.randn(1, 1, 4, 7, dtype=dtypes.float32), 0)
 
+  def test_partial_rope_in_attention(self):
+    dim, rope_dim, seqlen = 8, 4, 3
+    config = TransformerConfig(num_blocks=1, dim=dim, hidden_dim=16, n_heads=1, n_kv_heads=1,
+                               norm_eps=1e-5, vocab_size=32, head_dim=dim, rope_theta=10000.0,
+                               rope_dim=rope_dim, v_head_dim=dim, max_context=8)
+    block = TransformerBlock(config)
+
+    x = Tensor.randn(1, seqlen, dim, dtype=dtypes.float32)
+    x_norm = block.attn_norm(x)
+    k = block.attn_k(x_norm).reshape(1, seqlen, 1, dim).transpose(1, 2)
+
+    precompute_freqs_cis.cache_clear()
+    block.cache_kv = Tensor.empty(2, 1, 1, config.max_context, max(dim, config.v_head_dim), device=x.device)
+    block.freqs_cis = precompute_freqs_cis(rope_dim, config.max_context, config.rope_theta)
+    block._attention(x_norm, 0).realize()
+
+    expected = apply_rope_new(k[..., :rope_dim], block.freqs_cis[:seqlen]).cat(k[..., rope_dim:], dim=-1)
+    np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
+
+class TestGatedDeltaNetBlock(unittest.TestCase):
+  def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
+    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
+
+  def _make_config(self, **kwargs):
+    return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
+                                 "rope_dim":4, "v_head_dim":4, "max_context":4, "full_attention_interval":2,
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
+
+  def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
+    block = GatedDeltaNetBlock(config, config.ssm)
+    block.attn_norm.weight = self._tensor_linspace(0.8, 1.2, (config.dim,))
+    block.attn_qkv.weight = self._tensor_linspace(-0.15, 0.2, (block.conv_channels, config.dim))
+    block.attn_gate.weight = self._tensor_linspace(-0.1, 0.15, (config.ssm.inner_size, config.dim))
+    block.ssm_alpha.weight = self._tensor_linspace(-0.08, 0.12, (block.num_v_heads, config.dim))
+    block.ssm_beta.weight = self._tensor_linspace(-0.12, 0.07, (block.num_v_heads, config.dim))
+    block.ssm_conv1d["weight"] = self._tensor_linspace(-0.05, 0.05, (block.conv_channels, block.ssm_conv_kernel))
+    block.ssm_dt["bias"] = self._tensor_linspace(-0.1, 0.1, (block.num_v_heads,))
+    block.ssm_a = self._tensor_linspace(-0.1, -0.05, (block.num_v_heads,))
+    block.ssm_norm.weight = self._tensor_linspace(0.9, 1.1, (block.head_v_dim,))
+    block.ssm_out.weight = self._tensor_linspace(-0.2, 0.18, (config.dim, config.ssm.inner_size))
+    return block
+
+  def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
+    x_norm = block.attn_norm(x)
+    block._init_state(x_norm)
+    return block._attention(x_norm, start_pos).realize().numpy()
+
+  def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
+    conv_flat = (block.ssm_conv_kernel - 1) * block.conv_channels
+    cache = block.delta_cache.numpy()
+    conv_state = cache[:, :conv_flat].reshape(cache.shape[0], block.ssm_conv_kernel - 1, block.conv_channels)
+    recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
+    return conv_state, recurrent_state
+
+  def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
+    return x.astype(np.float32) @ weight.T.astype(np.float32)
+
+  def _rms_norm_np(self, x:np.ndarray, weight:np.ndarray, eps:float) -> np.ndarray:
+    x_float = x.astype(np.float32)
+    return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
+
+  def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
+    return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
+
+  def _softplus_np(self, x:np.ndarray) -> np.ndarray:
+    return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
+
+  def _silu_np(self, x:np.ndarray) -> np.ndarray:
+    return x / (1.0 + np.exp(-x))
+
+  def _naive_attention(self, block:GatedDeltaNetBlock, x:Tensor):
+    x_np = x.numpy().astype(np.float32)
+    B, T, _ = x_np.shape
+    conv_state = np.zeros((B, block.ssm_conv_kernel - 1, block.conv_channels), dtype=np.float32)
+    recurrent_state = np.zeros((B, block.num_v_heads, block.head_v_dim, block.head_v_dim), dtype=np.float32)
+    conv_weight = block.ssm_conv1d["weight"].numpy().astype(np.float32).T[None, :, :]
+    qkv_weight = block.attn_qkv.weight.numpy().astype(np.float32)
+    gate_weight = block.attn_gate.weight.numpy().astype(np.float32)
+    alpha_weight = block.ssm_alpha.weight.numpy().astype(np.float32)
+    beta_weight = block.ssm_beta.weight.numpy().astype(np.float32)
+    out_weight = block.ssm_out.weight.numpy().astype(np.float32)
+    dt_bias = block.ssm_dt["bias"].numpy().astype(np.float32)
+    ssm_a = block.ssm_a.numpy().astype(np.float32)
+    attn_norm_weight = block.attn_norm.weight.numpy().astype(np.float32)
+    ssm_norm_weight = block.ssm_norm.weight.numpy().astype(np.float32)
+    outputs, conv_states, recurrent_states = [], [], []
+
+    for t in range(T):
+      x_norm = self._rms_norm_np(x_np[:, t:t+1, :], attn_norm_weight, block.attn_norm.eps)
+      x_half = x_norm.astype(np.float16)
+      out_gate = self._linear_np(x_half, gate_weight).reshape(B, 1, block.num_v_heads, block.head_v_dim)
+      beta = 1.0 / (1.0 + np.exp(-self._linear_np(x_half, beta_weight))).reshape(B, block.num_v_heads, 1, 1)
+      alpha = np.exp((self._softplus_np(self._linear_np(x_half, alpha_weight) + dt_bias)).reshape(B, block.num_v_heads, 1, 1) *
+                     ssm_a.reshape(1, block.num_v_heads, 1, 1))
+      conv_window = np.concatenate([conv_state, self._linear_np(x_half, qkv_weight)], axis=1)
+      conv_out = self._silu_np((conv_window * conv_weight).sum(axis=1))
+      q, k, v = np.split(conv_out, [block.q_dim, 2 * block.q_dim], axis=-1)
+      q = self._normalize_np(q.reshape(B, block.num_k_heads, block.head_k_dim))
+      k = self._normalize_np(k.reshape(B, block.num_k_heads, block.head_k_dim))
+      v = v.reshape(B, block.num_v_heads, block.head_v_dim)
+      if block.num_v_heads != block.num_k_heads:
+        k_repeat = block.num_v_heads // block.num_k_heads
+        q = np.repeat(q[:, None, :, :], k_repeat, axis=1).reshape(B, block.num_v_heads, block.head_k_dim)
+        k = np.repeat(k[:, None, :, :], k_repeat, axis=1).reshape(B, block.num_v_heads, block.head_k_dim)
+      q, k, v = (q * (block.head_k_dim ** -0.5))[..., None], k[..., None], v[..., None]
+      recurrent_state = recurrent_state * alpha
+      recurrent_state = recurrent_state + np.matmul((v - np.matmul(recurrent_state, k)) * beta, np.swapaxes(k, -1, -2))
+      core_attn_out = np.matmul(recurrent_state, q).squeeze(-1).reshape(B, 1, block.num_v_heads, block.head_v_dim)
+      core_attn_out = self._rms_norm_np(core_attn_out, ssm_norm_weight, block.ssm_norm.eps)
+      out = self._linear_np((core_attn_out * self._silu_np(out_gate)).reshape(B, 1, -1).astype(np.float16), out_weight)
+      conv_state = conv_window[:, 1:, :]
+      outputs.append(out)
+      conv_states.append(conv_state.copy())
+      recurrent_states.append(recurrent_state.copy())
+
+    return outputs, conv_states, recurrent_states
+
+  def test_gatedeltanet_reference_and_reset(self):
+    config = self._make_config(max_context=3)
+    block = self._make_block(config)
+    x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)
+
+    expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
+
+    for step in range(x.shape[1]):
+      out = self._run_attention(block, x[:, step:step+1], step)
+      conv_state, recurrent_state = self._cache_views(block)
+      np.testing.assert_allclose(out, expected_outs[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet output mismatch at step {step}")
+      np.testing.assert_allclose(conv_state, expected_conv[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet conv cache mismatch at step {step}")
+      np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet recurrent cache mismatch at step {step}")
+
+    warmup = Tensor.linspace(-0.5, 0.5, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
+    prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
+
+    for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
+    Tensor.realize(*block._state_reset_ops())
+    expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)
+
+    for step in range(prompt.shape[1]):
+      out = self._run_attention(block, prompt[:, step:step+1], step)
+      conv_state, recurrent_state = self._cache_views(block)
+      np.testing.assert_allclose(out, expected_outs[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet reset output mismatch at step {step}")
+      np.testing.assert_allclose(conv_state, expected_conv[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet reset conv cache mismatch at step {step}")
+      np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
+                                 err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
+
 class TestPairwiseTopk(unittest.TestCase):
   def test_basic_topk(self):
     x = Tensor([[[1.0, 3.0, 2.0, 5.0, 4.0]]])

test/unit/test_dtype_spec.py

@@ -64,7 +64,8 @@ class TestTypeSpec(unittest.TestCase):
     tested = 0
     for dtype_str, dtype in [
       ("bool", dtypes.bool), ("int8", dtypes.int8), ("int", dtypes.int), ("uint32", dtypes.uint32), ("float32", dtypes.float32)]:
-      np.testing.assert_equal(Tensor(n, dtype=dtype_str).numpy(), Tensor(n, dtype=dtype).numpy())
+      with np.errstate(invalid='ignore'):
+        np.testing.assert_equal(Tensor(n, dtype=dtype_str).numpy(), Tensor(n, dtype=dtype).numpy())
       np.testing.assert_equal(Tensor(n).cast(dtype_str).numpy(), Tensor(n).cast(dtype).numpy())
       if dtype.itemsize == 4:
         np.testing.assert_equal(Tensor(n).bitcast(dtype_str).numpy(), Tensor(n).bitcast(dtype).numpy())

test/unit/test_function.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 from tinygrad.function import function
 from tinygrad import Tensor, GlobalCounters
-from tinygrad.uop.ops import UOp
+from tinygrad.uop.ops import UOp, KernelInfo
 
 class TestFunction(unittest.TestCase):
   def test_simple(self):
@@ -398,6 +398,74 @@ class TestFunctionTuple(unittest.TestCase):
     np.testing.assert_allclose(x.grad.numpy(), [1., 1., 1.])
     np.testing.assert_allclose(y.grad.numpy(), [1., 1., 1.])
 
+  def test_grad_unused_tuple_output_recursive(self):
+    # only one output is used
+    @function(precompile=True, precompile_backward=True)
+    def f(x:Tensor, w:Tensor):
+      a = x @ w
+      b = (x @ w) * 2  # shares x@w with a
+      return (a, b)
+
+    x = Tensor([[1., 2.], [3., 4.]], requires_grad=True).contiguous()
+    w = Tensor([[1., 0.], [0., 1.]], requires_grad=True).contiguous()
+    Tensor.realize(x, w)
+    t1, _ = f(x, w)
+    t1.sum().backward()
+    Tensor.realize(x.grad, w.grad)
+    # only t1 = x @ w flows to loss; dL/dw = x.T @ ones(2,2)
+    np.testing.assert_allclose(w.grad.numpy(), np.array([[1., 2.], [3., 4.]]).T @ np.ones((2, 2)))
+    np.testing.assert_allclose(x.grad.numpy(), np.ones((2, 2)) @ np.array([[1., 0.], [0., 1.]]).T)
+
+  def test_custom_kernel_save_unused_output(self):
+    def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
+      i = UOp.range(A.size, 0)
+      j = UOp.range(D.size, 1)
+      store_c = C[i].store(A[i] * 2.0).end(i)
+      store_d = D[j].store(A[j]).end(j)
+      return UOp.group(store_c, store_d).sink(arg=KernelInfo(name="my_kernel"))
+
+    def my_grad(d_c:UOp, call:UOp):
+      a_input = call.src[3]
+      return (None, None, (Tensor(d_c) * 2.0 + Tensor(a_input) * 0).uop)
+
+    @function(precompile=True, precompile_backward=True)
+    def f(a:Tensor):
+      c = Tensor.invalid(*a.shape, dtype=a.dtype, device=a.device)
+      d = Tensor.invalid(3, dtype=a.dtype, device=a.device)
+      c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
+      return c, d
+
+    a = Tensor([1., 2., 3., 4.], requires_grad=True).contiguous()
+    Tensor.realize(a)
+    c, _ = f(a)
+    c.sum().backward()
+    Tensor.realize(a.grad)
+    np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
+
+  def test_custom_kernel_both_outputs_used(self):
+    def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
+      i = UOp.range(A.size, 0)
+      store_c = C[i].store(A[i] * 2.0)
+      store_d = D[i].store(A[i] * 3.0)
+      return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="my_kernel"))
+
+    def my_grad(d_combined:UOp, call:UOp):
+      return (None, None, Tensor(d_combined).uop)
+
+    @function(precompile=True, precompile_backward=True)
+    def f(a:Tensor):
+      c = Tensor.invalid(*a.shape, dtype=a.dtype, device=a.device)
+      d = Tensor.invalid(*a.shape, dtype=a.dtype, device=a.device)
+      c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
+      return (c, d)
+
+    a = Tensor([1., 2., 3., 4.], requires_grad=True).contiguous()
+    Tensor.realize(a)
+    c, d = f(a)
+    (c.sum() + d.sum()).backward()  # dL/da = (1 + 1) since grad_fxn passes d_combined through
+    Tensor.realize(a.grad)
+    np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
+
 class TestFunctionGrad(unittest.TestCase):
   def test_function_grad_ops(self, precompile=False, precompile_backward=False):
     N = 64

test/unit/test_gguf.py

@@ -30,11 +30,15 @@ class TestGGUF(unittest.TestCase):
 
   def test_dequantization_q4_0(self): self._test_dequantization(GGMLQuantizationType.Q4_0)
   def test_dequantization_q4_1(self): self._test_dequantization(GGMLQuantizationType.Q4_1)
+  def test_dequantization_q5_0(self): self._test_dequantization(GGMLQuantizationType.Q5_0)
+  def test_dequantization_q5_1(self): self._test_dequantization(GGMLQuantizationType.Q5_1)
   def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
   def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
   def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
   def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
   def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
+  def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
   def test_dequantization_mxfp4_old(self):
     def encode(nibbles, E):
       packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
@@ -77,6 +81,13 @@ class TestGGUF(unittest.TestCase):
     out = ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value)
     np.testing.assert_equal(out.numpy(), expected)
 
+  def test_dequantization_q1_0(self):
+    # Q1_0: 2 bytes fp16 scale + 16 bytes (128 1-bit values)
+    block = np.frombuffer(np.float16(2.0).tobytes() + np.packbits(np.random.choice([0, 1], size=128)).tobytes(), dtype=np.uint8).copy()
+    expected = np.float16(2.0) * (np.unpackbits(block[2:], bitorder="little").astype(np.int8) * 2 - 1)
+    # TODO: replace 41 with GGMLQuantizationType.Q1_0.value on next gguf-py release
+    np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 128, 41).numpy().flatten(), expected)
+
   def test_expected_failure_unknown_type(self):
     with self.assertRaises(ValueError):
       ggml_data_to_tensor(Tensor.empty(512, dtype=dtypes.uint8), 256, 1337)
@@ -120,17 +131,21 @@ class TestGGUF(unittest.TestCase):
 class TestGGUFGEMV(unittest.TestCase):
   def _test_gguf_gemv(self, qtype: GGMLQuantizationType):
     block_size, type_size = GGML_QUANT_SIZES[qtype]
-    rows, cols = 8192, 2048
+    rows, cols = (1024, 512) if qtype == GGMLQuantizationType.BF16 else (8192, 2048)
     n_blocks = rows * cols // block_size
     rng = np.random.default_rng(42)
-    # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
-    q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
-    scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
-    if qtype == GGMLQuantizationType.Q8_0: q_data[:, :2] = scales[:, :2]                    # d at offset 0
-    elif qtype in (GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K): q_data[:, :4] = scales[:, :4] # d, dmin at offset 0
-    elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]                 # d at end
-    elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
-    q_data = q_data.flatten()
+    if qtype == GGMLQuantizationType.BF16:
+      q_data = (rng.standard_normal(rows * cols).astype(np.float32).view(np.uint32) >> 16).astype(np.uint16).view(np.uint8)
+    else:
+      # generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
+      q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
+      scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
+      if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0): q_data[:, :2] = scales[:, :2]  # d at offset 0
+      elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
+        q_data[:, :4] = scales[:, :4]  # d, m/dmin at offset 0
+      elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2]               # d at end
+      elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
+      q_data = q_data.flatten()
     ref = dequantize(q_data, qtype).reshape(rows, cols)
 
     # build a minimal gguf in memory: header + 1 tensor info + aligned data
@@ -147,16 +162,20 @@ class TestGGUFGEMV(unittest.TestCase):
     _, tensors = gguf_load(Tensor(np.frombuffer(buf, dtype=np.uint8)).to(None))
 
     x = rng.standard_normal(cols).astype(np.float32)
-    np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
-    # can only expect the weights to be identical if we really support float16 (ie. not decompositions)
-    if is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
+    with np.errstate(all='ignore'):
+      np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
+    if qtype == GGMLQuantizationType.BF16 or is_dtype_supported(dtypes.half): np.testing.assert_equal(tensors["weight"].numpy(), ref)
     assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"
 
   def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
+  def test_gguf_gemv_q5_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_0)
+  def test_gguf_gemv_q5_1(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_1)
   def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
   def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
   def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
   def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
+  @unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "Backend must support bfloat16")
+  def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
 
 if __name__ == '__main__':
   unittest.main()

test/unit/test_llm_mla.py

@@ -0,0 +1,72 @@
+import unittest
+import numpy as np
+from tinygrad import Tensor
+from tinygrad.apps.llm import Transformer, TransformerConfig, apply_rope
+
+class TestMLA(unittest.TestCase):
+  def _make_config(self, **kwargs):
+    return TransformerConfig(**{
+      "num_blocks": 1, "dim": 64, "hidden_dim": 128, "n_heads": 4, "n_kv_heads": 1,
+      "norm_eps": 1e-5, "vocab_size": 100, "head_dim": 16, "rope_theta": 10000.0, "rope_dim": 8, "max_context": 32,
+      "kv_lora_rank": 16, "v_head_dim": 8,
+    } | kwargs)
+
+  def test_mla_attention_matches_naive(self):
+    config = self._make_config(max_context=16)
+    from tinygrad.apps.llm import MLATransformerBlock, precompute_freqs_cis
+
+    block = MLATransformerBlock(config)
+    c = config
+    B, T = 1, 4
+    q_nope_head_dim = c.head_dim - c.rope_dim
+
+    x = Tensor.randn(B, T, c.dim)
+    x_norm = block.attn_norm(x)
+
+    # --- Our absorbed implementation ---
+    q = block.attn_q(x_norm).reshape(B, T, c.n_heads, c.head_dim).transpose(1, 2)
+    q_nope, q_rope = q[..., :q_nope_head_dim], q[..., q_nope_head_dim:]
+    freqs = precompute_freqs_cis(c.rope_dim, 16, c.rope_theta)
+    q_rope = apply_rope(q_rope, freqs[0:T])
+
+    kv_a = block.attn_kv_a_mqa(x_norm)
+    c_kv = block.attn_kv_a_norm(kv_a[..., :c.kv_lora_rank])
+    k_rope = kv_a[..., c.kv_lora_rank:].reshape(B, T, 1, c.rope_dim).transpose(1, 2)
+    k_rope = apply_rope(k_rope, freqs[0:T])
+
+    # --- Naive (non-absorbed): expand K and V, do standard attention ---
+    k_nope_naive = c_kv.unsqueeze(1) @ block.attn_k_b["weight"]  # (B, H, T, nope)
+    k_naive = k_nope_naive.cat(k_rope.expand(-1, c.n_heads, -1, -1), dim=-1)  # (B, H, T, nope+rope)
+    v_naive = c_kv.unsqueeze(1) @ block.attn_v_b["weight"].transpose(-1, -2)  # (B, H, T, v_dim)
+
+    q_naive = q_nope.cat(q_rope, dim=-1)
+    scale = 1.0 / c.head_dim ** 0.5
+    scores_naive = (q_naive @ k_naive.transpose(-1, -2)) * scale
+    # causal mask
+    mask = Tensor.full((1, 1, T, T), float("-inf")).triu(1)
+    attn_naive = (scores_naive + mask).softmax(-1) @ v_naive  # (B, H, T, v_dim)
+    out_naive = block.attn_output(attn_naive.transpose(1, 2).reshape(B, T, -1))
+
+    # --- Absorbed: q_nope @ wk_b^T, then dot with compressed kv ---
+    q_nope_abs = q_nope @ block.attn_k_b["weight"].transpose(-1, -2)  # (B, H, T, lora)
+    q_abs = q_nope_abs.cat(q_rope, dim=-1)  # (B, H, T, lora+rope)
+    k_abs = c_kv.reshape(B, 1, T, c.kv_lora_rank).cat(k_rope.reshape(B, 1, T, c.rope_dim), dim=-1)
+    scores_abs = (q_abs @ k_abs.transpose(-1, -2)) * scale
+    attn_abs = (scores_abs + mask).softmax(-1)
+    # attn @ v_compressed @ wv_b
+    v_compressed = c_kv.reshape(B, 1, T, c.kv_lora_rank)
+    attn_abs_out = (attn_abs @ v_compressed) @ block.attn_v_b["weight"].transpose(-1, -2)
+    out_abs = block.attn_output(attn_abs_out.transpose(1, 2).reshape(B, T, -1))
+
+    # Compare
+    naive_np = out_naive.realize().numpy()
+    abs_np = out_abs.realize().numpy()
+    np.testing.assert_allclose(naive_np, abs_np, atol=1e-4, rtol=1e-4,
+      err_msg="Absorbed MLA should match naive MLA")
+
+  def test_shared_expert_gate_optional(self):
+    from tinygrad import nn
+    model = Transformer(self._make_config(num_experts=4, num_experts_per_tok=2, shared_expert_dim=32, shared_expert_gate=False))
+    self.assertNotIn('blk.0.ffn_gate_inp_shexp.weight', nn.state.get_state_dict(model))
+    out = model.blk[0]._feed_forward(Tensor.randn(1, 4, model.blk[0].config.dim))
+    self.assertEqual(out.shape, (1, 4, model.blk[0].config.dim))

test/unit/test_llm_moe.py

@@ -1,12 +1,15 @@
 import unittest
 import numpy as np
+from dataclasses import replace
 from tinygrad import Tensor
 from tinygrad.apps.llm import TransformerBlock, TransformerConfig
 
 def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
-  return TransformerConfig(num_blocks=1, dim=dim, hidden_dim=hidden, n_heads=n_heads, n_kv_heads=n_heads,
-                           norm_eps=1e-5, vocab_size=100, head_dim=dim//n_heads, rope_theta=10000, max_context=16,
-                           num_experts=num_experts, num_experts_per_tok=num_experts_per_tok)
+  return TransformerConfig(
+    num_blocks=1, dim=dim, hidden_dim=hidden, n_heads=n_heads, n_kv_heads=n_heads,
+    norm_eps=1e-5, vocab_size=100, head_dim=dim//n_heads, rope_theta=10000,
+    rope_dim=dim//n_heads, v_head_dim=dim//n_heads, max_context=16,
+    num_experts=num_experts, num_experts_per_tok=num_experts_per_tok)
 
 class TestMoEFeedForward(unittest.TestCase):
   def test_moe_feed_forward(self):
@@ -24,10 +27,10 @@ class TestMoEFeedForward(unittest.TestCase):
 
     # input of ones -> after norm still ~ones -> experts 0,2 selected -> weighted sum of silu outputs
     h = Tensor.ones(1, 1, dim)
-    out = block._feed_forward(h)
+    out = block._feed_forward(block.ffn_norm(h))
 
-    # expected: residual + moe_output ≈ 1 + avg(silu(1), silu(3))
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    # expected moe_output ≈ avg(silu(1), silu(3))
+    expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
     np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
 
   def test_moe_feed_forward_batched(self):
@@ -45,17 +48,16 @@ class TestMoEFeedForward(unittest.TestCase):
 
     # test with BS=2, T=3
     h = Tensor.ones(2, 3, dim)
-    out = block._feed_forward(h)
+    out = block._feed_forward(block.ffn_norm(h))
 
     # all outputs should match the BS=1 expected value
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
     np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
 
   def test_moe_feed_forward_norm_topk_prob(self):
     dim, hidden, n_heads = 8, 16, 2
     num_experts, k = 4, 2
 
-    from dataclasses import replace
     block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k), norm_topk_prob=True))
 
     block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
@@ -65,10 +67,34 @@ class TestMoEFeedForward(unittest.TestCase):
     block.ffn_norm.weight = Tensor.ones(dim)
 
     h = Tensor.ones(1, 1, dim)
-    out = block._feed_forward(h)
+    out = block._feed_forward(block.ffn_norm(h))
 
-    expected = 1 + (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
     np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
 
+  def test_moe_feed_forward_shared_expert(self):
+    dim, hidden, n_heads = 8, 16, 2
+    num_experts, k = 4, 2
+
+    block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k), shared_expert_dim=dim))
+
+    block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
+    block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
+    block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
+    block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T
+    block.ffn_gate_shexp.weight = Tensor.eye(dim) * 2
+    block.ffn_up_shexp.weight = Tensor.eye(dim)
+    block.ffn_down_shexp.weight = Tensor.eye(dim)
+    block.ffn_gate_inp_shexp["weight"] = Tensor.zeros(dim)
+    block.ffn_norm.weight = Tensor.ones(dim)
+
+    h = Tensor.ones(1, 1, dim)
+    out = block._feed_forward(block.ffn_norm(h))
+
+    moe_expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
+    shared_expected = Tensor([2.0]).silu().item() * 0.5
+    expected = moe_expected + shared_expected
+    np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
+
 if __name__ == '__main__':
   unittest.main()

test/unit/test_llm_server.py

@@ -1,11 +1,11 @@
 import unittest
 from unittest.mock import patch
 from tinygrad import Tensor, UOp
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 from tinygrad.apps.llm import Transformer, TransformerConfig
 
 TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
-                          norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, max_context=32)
+                           norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
 
 class TestTransformerGenerate(unittest.TestCase):
   def test_kv_cache_reuse(self):

test/unit/test_schedule_cache.py

@@ -2,7 +2,7 @@ import unittest
 import functools
 from tinygrad import Tensor, Variable, UOp
 from tinygrad.uop.ops import KernelInfo
-from tinygrad.engine.schedule import schedule_cache
+from tinygrad.schedule import schedule_cache
 
 def custom_set0_kernel(A:UOp, num:int) -> UOp:
   return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))