Merge branch 'master' into rdna4_gemm

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

@@ -225,12 +225,14 @@ 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 (macOS)
+    - name: Setup AMD comgr+remu (macOS)
       if: inputs.amd == 'true' && runner.os == 'macOS'
       shell: bash
       run: |
@@ -238,6 +240,7 @@ 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/autogen.yml

@@ -58,7 +58,6 @@ jobs:
         python3 -c "from tinygrad.runtime.autogen import avcodec"
         python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
         python3 -c "from tinygrad.runtime.autogen import mlx5"
-        python3 -c "from tinygrad.runtime.autogen import ggml_common"
         REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
     - name: Check for differences
       run: |

.github/workflows/benchmark.yml

@@ -71,6 +71,10 @@ 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
@@ -534,8 +538,6 @@ 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
@@ -626,7 +628,7 @@ jobs:
     - name: IR3 openpilot compile3 0.11.0 driving_vision
       run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.11.0 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: DEBUG=2 openpilot compile3 0.10.1 driving_vision
@@ -634,7 +636,7 @@ jobs:
     - name: openpilot compile3 0.10.1 driving_vision
       run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
     - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
     - name: openpilot compile3 0.10.1 dmonitoring
       run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
     - name: benchmark MobileNetV2 on DSP
@@ -707,8 +709,6 @@ 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

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '19'
+  CACHE_VERSION: '18'
   CAPTURE_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: ${{ github.workspace }}
@@ -505,14 +505,12 @@ jobs:
         with:
           key: apps_llm
       - name: Test 1B LLM (llama)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
+        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.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.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster
+        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 (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.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
+        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
 
 # ****** Models Tests ******
 
@@ -643,7 +641,9 @@ jobs:
     runs-on: ubuntu-24.04
     timeout-minutes: 20
     env:
-      DEV: MOCKKFD+AMD
+      DEV: AMD
+      PYTHON_REMU: 1
+      MOCKGPU: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -669,7 +669,7 @@ jobs:
       - name: Run AMD renderer tests
         run: python -m pytest -n=auto test/amd/ --durations 20
       - name: Run AMD renderer tests (AMD:LLVM)
-        run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
+        run: DEV=AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
         run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
@@ -678,19 +678,20 @@ jobs:
         run: |
           PYTHONPATH=. DEV=NULL::gfx1100 python extra/mmapeak/mmapeak.py
           PYTHONPATH=. DEV=NULL::gfx1201 python3 -m pytest -n=auto test/testextra/test_tk.py test/backend/test_asm_gemm.py
-      - name: Run matmul on MOCKKFD
+      - name: Run matmul on MOCKGPU
         run: |
-          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_asm_matmul.py
-          PYTHONPATH="." DEV=MOCKKFD+AMD N=256 python3 extra/gemm/amd_copy_matmul.py
+          PYTHONPATH="." DEV=AMD MOCKGPU=1 N=256 python3 extra/gemm/amd_asm_matmul.py
+          PYTHONPATH="." DEV=AMD MOCKGPU=1 N=256 python3 extra/gemm/amd_copy_matmul.py
       - name: Run LLVM test
-        run: DEV=MOCKKFD+AMD:LLVM python test/device/test_amd_llvm.py
+        run: DEV=AMD:LLVM python test/device/test_amd_llvm.py
 
   testmockam:
     name: Linux (am)
     runs-on: ubuntu-24.04
     timeout-minutes: 15
     env:
-      DEV: MOCKPCI+AMD
+      DEV: PCI+AMD
+      MOCKGPU: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -702,13 +703,13 @@ jobs:
           amd: 'true'
       - name: Run test_tiny on MOCKAM
         run: python test/test_tiny.py
-      - name: Run test_tiny on MOCKUSB
-        run: GMMU=0 DEV=MOCKUSB+AMD python test/test_tiny.py
-      - name: Run test_hcq on MOCKPCI
+      - name: Run test_tiny on MOCKAM USB
+        run: GMMU=0 DEV=USB+AMD python test/test_tiny.py
+      - name: Run test_hcq on MOCKAM
         run: python -m pytest test/device/test_hcq.py
-      - name: Run disk copy tests on MOCKPCI
+      - name: Run disk copy tests on MOCKAM
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
-      - name: Run test_tiny on MOCKPCI Remote
+      - name: Run test_tiny on MOCKAM Remote
         run: |
           python extra/remote/serve.py 6667 &
           sleep 2
@@ -726,7 +727,8 @@ jobs:
     runs-on: ubuntu-22.04
     timeout-minutes: 15
     env:
-      DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
+      DEV: AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
+      MOCKGPU: 1
       SKIP_SLOW_TEST: 1
     steps:
       - name: Checkout Code
@@ -761,6 +763,7 @@ jobs:
     runs-on: ubuntu-22.04
     timeout-minutes: 20
     env:
+      MOCKGPU: 1
       FORWARD_ONLY: 1
     steps:
       - name: Checkout Code
@@ -773,7 +776,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCKNVK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=CUDA:PTX' || matrix.backend == 'nv' && 'DEV=NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
         run: |
           python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
@@ -858,19 +861,22 @@ jobs:
       run: DEV=METAL TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
       env:
-        DEV: MOCKKFD+AMD
+        MOCKGPU: 1
+        DEV: AMD
         FORWARD_ONLY: 1
       run: |
         python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
     - name: Run pytest (amd with llvm backend)
       env:
-        DEV: "MOCKKFD+AMD:LLVM"
+        MOCKGPU: 1
+        DEV: "AMD:LLVM"
         FORWARD_ONLY: 1
       run: |
         python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
     - name: Run pytest (ptx)
       env:
-        DEV: "MOCKNVK+NV:PTX"
+        MOCKGPU: 1
+        DEV: "NV:PTX"
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
@@ -945,8 +951,8 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
       - name: Run macOS-specific unit test
-        if: matrix.backend == 'llvm'
-        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated test/unit/test_cpu.py
+        if: matrix.backend == 'cpu'
+        run: python3 -m pytest test/unit/test_disk_tensor.py::TestDiskTensor::test_copy_to_cpu_not_truncated
 
 # ****** Windows Tests ******
 

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.schedule import ExecItem
+from tinygrad.engine.schedule import ExecItem
 schedule: List[ExecItem] = Tensor.schedule(l1, l2)
 
 print(f"The schedule contains {len(schedule)} items.")

docs/abstractions4.py

@@ -1,253 +0,0 @@
-# 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="." DEV=MOCKPCI+AMD
-
-from tinygrad import Tensor, Context, GlobalCounters, UOp, Device
-from tinygrad.helpers import DEV, 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 DEV.interface.startswith("MOCK") 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/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.
+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.
 
-::: tinygrad.schedule.ExecItem
+::: tinygrad.engine.schedule.ExecItem
 
 ## Lowering
 

docs/env_vars.md

@@ -34,8 +34,9 @@ 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]        | enable 2d specific optimizations
+IMAGE               | [1-2]      | 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.
@@ -64,8 +65,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 the applied optimizations at a kernel level
+DEBUG               | >= 3       | Outputs buffers used for each kernel (shape, dtype and strides) and the applied optimizations at a kernel level
 DEBUG               | >= 4       | Outputs the generated kernel code
-DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps
+DEBUG               | >= 5       | Displays the intermediate representation of the computation UOps (AST)
 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/nn.md

@@ -37,4 +37,4 @@
     options:
         show_signature: false
         separate_signature: false
-::: tinygrad.llm.gguf.gguf_load
+::: tinygrad.nn.state.gguf_load

docs/tinygpu.md

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

examples/compile_tensorflow.py

@@ -35,11 +35,12 @@ def compile_onnx_model(onnx_model):
   tinyonnx = TinyOnnx(onnx_model)
   the_input = Tensor.randn(1,32)
 
-  linear, output_bufs = jit_model(tinyonnx, the_input)
-  the_output = [tinyonnx.forward(the_input)]
+  run, special_names = jit_model(tinyonnx, the_input)
 
-  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
+  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
   prg = export_model_clang(functions, statements, bufs, {}, ["input0"], ["output0"])
+
+  the_output = run(the_input)
   cprog = ["#include <string.h>", "#include <stdio.h>", "#include <stdlib.h>"]
   cprog.append(prg)
 

examples/gpt2.py

@@ -5,9 +5,8 @@ with contextlib.suppress(ImportError): import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
 from tinygrad.uop.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
-from tinygrad.llm.gguf import gguf_load
 from tinygrad.nn import Embedding, Linear, LayerNorm
-from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
+from tinygrad.nn.state import gguf_load, torch_load, load_state_dict, get_state_dict
 from extra.bench_log import BenchEvent, WallTimeEvent
 
 MAX_CONTEXT = getenv("MAX_CONTEXT", 128)

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.nbytes() for x in get_parameters(llama.model))
+  param_bytes = sum(x.uop.size * x.dtype.itemsize 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

@@ -2,8 +2,7 @@ from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
-from tinygrad.llm.gguf import gguf_load
-from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
+from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
 from tinygrad.helpers import Profiling, Timing, DEBUG, colored, fetch, tqdm
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -325,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.nbytes() for x in get_parameters(model))
+  param_bytes = sum(x.uop.size * x.dtype.itemsize 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.schedule.memory import memory_planner
+from tinygrad.engine.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 // 2]  # in seconds
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 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 // 2]  # in seconds
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 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 // 2]  # in seconds
+          median_step_time = sorted(step_times)[(BENCHMARK + 1) // 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 // 2]  # in seconds
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 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:_}, "
@@ -1282,7 +1282,7 @@ def train_bert():
         previous_step = i
 
 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8, FP8_DTYPE
+  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad, FP8
   from examples.llama3 import MODEL_PARAMS
   from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
   from examples.mlperf.optim import GradAccClipAdamW
@@ -1416,9 +1416,9 @@ def train_llama3():
   optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                            eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)
 
+  # init grads
   for p in optim.params:
-    grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
-    p.grad = Tensor.zeros(p.shape, dtype=grad_dtype, device=p.device).contiguous()
+    p.grad = Tensor.zeros(p.shape, dtype=p.dtype, device=p.device).contiguous()
   grads = [p.grad for p in optim.params]
 
   scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
@@ -1433,17 +1433,6 @@ def train_llama3():
     load_state_dict(scheduler, safe_load(fn), realize=False)
 
   fp8_amax = [t for ts in model._fp8_amax.values() for t in ts] if FP8 else []
-  fp8_inv_scales = list(model._fp8_inv_scale.values()) if FP8 else []
-
-  if FP8:
-    from tinygrad.nn.state import get_state_dict
-    model_state = get_state_dict(model)
-    for wname in ["wqkv", "wo", "w13", "w2"]:
-      w = model_state[wname]
-      w._inv_scale = model._fp8_inv_scale[wname]
-      if optim.master_params:
-        idx = next(j for j, p in enumerate(optim.params) if p is w)
-        optim.master_params[idx].assign((optim.master_params[idx] * w._inv_scale.reshape(-1, *([1]*(w.ndim-1)))).contiguous())
 
   @TinyJit
   def minibatch(tokens:Tensor):
@@ -1468,7 +1457,7 @@ def train_llama3():
 
     lr_cpu = optim.lr.float().to("CPU")
     grad_norm_cpu = grad_norm.float().to("CPU")
-    Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
+    Tensor.realize(lr_cpu, grad_norm_cpu, *grads)
 
     return lr_cpu, grad_norm_cpu
 
@@ -1555,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 * (4.6e15 if FP8 else 2.3e15))) * 100
+      mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 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")
@@ -1588,9 +1577,8 @@ def train_llama3():
         safe_save(get_state_dict(scheduler), fn)
 
       if i == BENCHMARK:
-        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)
+        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]
+        estimated_total_minutes = int(median_step_time * (SAMPLES // GBS) / 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, functools
+import math, os
 if __name__ == "__main__":
   os.environ["DEFAULT_FLOAT"] = "bfloat16"
   os.environ["OPTIM_DTYPE"] = "bfloat16"
@@ -18,63 +18,37 @@ 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):
-  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)
+  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)
   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(), new_amax
+  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
 
-def matmul(x:Tensor, w:Tensor, fp8=FP8, amax_x:Tensor|None=None, w_inv_scale:Tensor|None=None,
-           x_fp8:Tensor|None=None, x_scale:Tensor|None=None, x_new_amax: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,)
-  assert w_inv_scale is not None, "fp8 matmul requires w_inv_scale (weights must be stored in fp8 with per-tensor scale)"
-  if x_fp8 is None: x_fp8, x_scale, x_new_amax = quantize_fp8(x, amax_state=amax_x)
-  if getenv("ASM_GEMM"):
+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)
+  combined_scale = x_scale * w_scale
+  if ASM_GEMM:
     from extra.gemm.cdna_asm_gemm import can_use_asm_gemm, asm_gemm
-    if can_use_asm_gemm(x_fp8, w.T): return asm_gemm(x_fp8, w.T, x_scale=x_scale, w_scale=w_inv_scale), x_new_amax, x_fp8, w
-  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale * w_inv_scale, x_new_amax, x_fp8, w
+    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
 
-def _rmsnorm_fwd(x_in:Tensor, eps:float) -> tuple[Tensor, Tensor]:
+def rmsnorm(x_in:Tensor, eps:float):
   x = x_in.float()
-  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))
+  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
+  return x.cast(x_in.dtype)
 
 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,
@@ -85,18 +59,22 @@ class FlatTransformer:
     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
-    self.hidden_dim = hidden_dim
 
     scaled_std = 0.02 / math.sqrt(2 * n_layers)
 
     # Attention
-    self._init_inv_scales = []  # populated by lin_per_layer when FP8
-    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
+    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.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim, std=scaled_std)
 
     # FeedForward
-    self.w13 = self.lin_per_layer(dim, hidden_dim * 2)
+    self.w1 = self.lin_per_layer(dim, hidden_dim)
     self.w2 = self.lin_per_layer(hidden_dim, dim, std=scaled_std)
+    self.w3 = self.lin_per_layer(dim, hidden_dim)
 
     self.norm_eps = norm_eps
     self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
@@ -111,100 +89,65 @@ class FlatTransformer:
 
     if FP8:
       def _amax(): return Tensor.full((), FP8_MAX).contiguous().requires_grad_(False)
-      names = ["xqkv", "xo", "x13", "x2"]
+      names = (["xqkv", "wqkv"] if WQKV else ["xq", "wq", "xk", "wk", "xv", "wv"]) + \
+              ["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}
-      # per-weight inv_scale: single (n_layers,) float32 tensor per weight (kernel reads float* pointers)
-      w_names = ["wqkv", "wo", "w13", "w2"]
-      self._fp8_inv_scale = {}
-      for wname, inv_scales in zip(w_names, self._init_inv_scales):
-        self._fp8_inv_scale[wname] = inv_scales.float().contiguous().requires_grad_(False)
-      del self._init_inv_scales
+      self._fp8_amax["xout"] = [_amax()]
+      self._fp8_amax["wout"] = [_amax()]
 
   def lin_per_layer(self, in_features:int, out_features:int, std:float=0.02):
-    if getenv("ZEROS"): w = Tensor.zeros(self.n_layers, out_features, in_features)
-    else: w = Tensor.normal(self.n_layers, out_features, in_features, mean=0.0, std=std)
-    if not FP8: return w
-    # per-layer scaled fp8 cast: fill the fp8 range for best precision
-    amax = w.abs().flatten(1).max(1).detach()
-    scale = FP8_MAX / (amax + 1e-8)
-    self._init_inv_scales.append((amax + 1e-8) / FP8_MAX)  # save for inv_scale init
-    return (w * scale.reshape(-1, 1, 1)).clamp(-FP8_MAX, FP8_MAX).cast(FP8_DTYPE)
+    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, wqkv:Tensor, wo:Tensor,
-                amax_xqkv=None, amax_xo=None, s_qkv=None, s_o=None):
+  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
     bsz, seqlen, _ = x.shape
-    new_amaxs, saves = [], []
 
-    x, rrms = rmsnorm(x, self.norm_eps)
-    saves.extend([x, rrms])
-    x = x * attention_norm
-
-    xqkv, *ret = matmul(x, wqkv, amax_x=amax_xqkv, w_inv_scale=s_qkv)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [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)
+    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)
 
     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)
-    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)
+    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)
 
-    out, *ret = matmul(attn, wo, amax_x=amax_xo, w_inv_scale=s_o)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [out])
-    return (out, *new_amaxs, *saves)
-
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                   amax_x13=None, amax_x2=None, s_13=None, s_2=None):
-    new_amaxs, saves = [], []
-
-    x, rrms = rmsnorm(x, self.norm_eps)
-    saves.extend([x, rrms])
-    x = x * ffn_norm
-
-    x_w13, *ret = matmul(x, w13, amax_x=amax_x13, w_inv_scale=s_13)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [x_w13])
-
-    if FP8 and getenv("FUSED_SILU_W13", 1):
-      from extra.amax.cast_amax import fused_quantize_fp8_w13
-      amax_s = amax_x2 if amax_x2 is not None else Tensor.full((), 1.0, dtype=dtypes.bfloat16, device=x_w13.device)
-      x2_fp8, x2_inv_scale, new_amax_x2 = fused_quantize_fp8_w13(x_w13, amax_s, FP8_DTYPE)
-      out, *ret = matmul(None, w2, w_inv_scale=s_2, x_fp8=x2_fp8, x_scale=x2_inv_scale, x_new_amax=new_amax_x2)
-    else:
-      x_w1, x_w3 = x_w13[..., :self.hidden_dim], x_w13[..., self.hidden_dim:]
-      out, *ret = matmul(x_w1.silu() * x_w3, w2, amax_x=amax_x2, w_inv_scale=s_2)
-    new_amaxs.extend(ret[:1])
-    saves.extend(ret[1:] + [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)
 
   @function(precompile=True, precompile_backward=True)
   def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w13:Tensor, w2:Tensor,
-                amax_xqkv=None, amax_xo=None,
-                amax_x13=None, amax_x2=None,
-                s_qkv=None, s_o=None, s_13=None, s_2=None):
-    attn, *attn_ret = self.attention(x, freqs_cis, attention_norm, wqkv, wo,
-                                     amax_xqkv=amax_xqkv, amax_xo=amax_xo,
-                                     s_qkv=s_qkv, s_o=s_o)
-    attn_amaxs, attn_saves = attn_ret[:2], attn_ret[2:]
-    h = x + attn
-    ffn, *ffn_ret = self.feed_forward(h, ffn_norm, w13, w2,
-                                      amax_x13=amax_x13, amax_x2=amax_x2,
-                                      s_13=s_13, s_2=s_2)
-    ffn_amaxs, ffn_saves = ffn_ret[:2], ffn_ret[2:]
-    h = h + ffn
-    return (h, *attn_amaxs, *ffn_amaxs, *attn_saves, *ffn_saves)
+                attention_norm: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_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)
 
   def shard(self, device:tuple[str, ...], mp:bool=False):
     from tinygrad.nn.state import get_parameters
@@ -212,44 +155,45 @@ 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
-      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
+      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.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
-      self.w13.shard_(device, axis=1).realize()             # (n_layers, hidden*2, dim) shard out
+      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
+      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
       self.attention_norm.shard_(device, axis=None).realize()
       self.ffn_norm.shard_(device, axis=None).realize()
       self.norm.weight.shard_(device, axis=None).realize()
       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)
-        for name in self._fp8_inv_scale:
-          self._fp8_inv_scale[name] = self._fp8_inv_scale[name].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
-    s = self._fp8_inv_scale if FP8 else None
     for i in range(self.n_layers):
-      amax_layer = {"amax_xqkv": a["xqkv"][i], "amax_xo": a["xo"][i],
-                    "amax_x13": a["x13"][i], "amax_x2": a["x2"][i]} if a else {}
-      scale_layer = {"s_qkv": s["wqkv"][i], "s_o": s["wo"][i],
-                     "s_13": s["w13"][i], "s_2": s["w2"][i]} if s else {}
-      h, *ret = self.run_layer(h, freqs_cis,
-                               self.attention_norm[i], self.wqkv[i], self.wo[i],
-                               self.ffn_norm[i], self.w13[i], self.w2[i],
-                               **amax_layer, **scale_layer)
-      if a:
-        amaxs = ret[:5]
-        amax_names = ["xqkv", "xo", "x13", "x2"]
-        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()
+      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_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()
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

examples/mlperf/models/llama.py

@@ -0,0 +1,80 @@
+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/optim.py

@@ -34,9 +34,7 @@ class GradAccClipAdamW(Optimizer):
     else:
       updates, extra = self._step([], grads)
     for i, tt in enumerate(self.params): tt.assign(self._apply_update(tt, updates[i], self.master_params[i] if self.master_params else None))
-    # collect inv_scale tensors attached to fp8 params (set by _apply_update)
-    fp8_inv_scales = [tt._inv_scale for tt in self.params if hasattr(tt, '_inv_scale')]
-    to_realize = extra+self.params+self.buffers+(self.master_params or [])+fp8_inv_scales
+    to_realize = extra+self.params+self.buffers+(self.master_params or [])
 
     Tensor.realize(*to_realize)
     return extra[-1]
@@ -79,12 +77,4 @@ class GradAccClipAdamW(Optimizer):
     new_w = w.detach() - up
     if master is not None: master.assign(new_w)
     if STOCHASTIC_ROUND and t.dtype == dtypes.bfloat16: return stochastic_round_bf16(new_w)
-    if t.dtype in dtypes.fp8s:
-      from examples.mlperf.models.flat_llama import FP8_MAX
-      amax = new_w.float().abs().flatten(1).max(1).detach()  # per-layer amax for (n_layers, out, in)
-      scale = FP8_MAX / (amax + 1e-8)
-      fp8_w = (new_w * scale.reshape(-1, *([1]*(new_w.ndim-1)))).clamp(-FP8_MAX, FP8_MAX).cast(t.dtype)
-      if hasattr(t, '_inv_scale'):
-        t._inv_scale.assign(((amax + 1e-8) / FP8_MAX).cast(t._inv_scale.dtype))
-      return fp8_w
     return new_w.cast(t.dtype)

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

@@ -2,6 +2,7 @@
 
 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
@@ -14,10 +15,9 @@ 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:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

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

@@ -2,6 +2,7 @@
 
 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
@@ -14,10 +15,9 @@ 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:-16} EVAL_BS=${EVAL_BS:-16} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
+export DP=${DP:-8} MP=${MP:-1} BS=${BS:-8} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-4}
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

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

@@ -2,5 +2,4 @@
 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
-SRC="AMD"; [[ $DEV == NULL* ]] && SRC="NULL"
-python -m tinygrad.viz.cli --profile -s "$SRC" --top 20
+extra/viz/cli.py --profile -s "${DEV:-AMD}"

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

@@ -4,6 +4,7 @@ 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
@@ -15,10 +16,9 @@ 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=16 EVAL_BS=16 GRADIENT_ACC_STEPS=2
+export DP=8 MP=1 BS=8 EVAL_BS=8 GRADIENT_ACC_STEPS=4
 export GBS=$((BS * GRADIENT_ACC_STEPS))
 
 export MODEL="llama3"

examples/openpilot/compile3.py

@@ -4,8 +4,7 @@ if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
-from tinygrad.uop.ops import Ops
-from tinygrad.engine.realize import get_runner
+from tinygrad.engine.realize import CompiledRunner
 from tinygrad.nn.onnx import OnnxRunner
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -36,11 +35,7 @@ def compile(onnx_file):
       ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)
-  # iterate kernel CALLs in the captured LINEAR UOp; toposort descends into batched graph CUSTOM_FUNCTIONs
-  kernel_asts = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
-  kernel_calls = [u for u in run_onnx_jit.captured.linear.toposort(gate=lambda x: x.op not in kernel_asts)
-                  if u.op is Ops.CALL and u.src[0].op in kernel_asts]
-  print(f"captured {len(kernel_calls)} kernels")
+  print(f"captured {len(run_onnx_jit.captured.jit_cache)} kernels")
   np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
@@ -48,14 +43,13 @@ def compile(onnx_file):
   kernel_count = 0
   read_image_count = 0
   gated_read_image_count = 0
-  for call in kernel_calls:
-    device = next(b.device for b in call.src[1:] if b.op is not Ops.BIND)
-    src = get_runner(device, call.src[0]).p.src
-    kernel_count += 1
-    read_image_count += src.count("read_image")
-    gated_read_image_count += src.count("?read_image")
-    for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', src)]:
-      if len(re.findall(fr'[\?\:]{v}\.[xyzw]', src)) > 0: gated_read_image_count += 1
+  for ei in run_onnx_jit.captured.jit_cache:
+    if isinstance(ei.prg, CompiledRunner):
+      kernel_count += 1
+      read_image_count += ei.prg.p.src.count("read_image")
+      gated_read_image_count += ei.prg.p.src.count("?read_image")
+      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
+        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
   print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
   if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
     assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"

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.nbytes() for x in get_parameters(transformer))
+  param_bytes = sum(x.uop.size * x.dtype.itemsize for x in get_parameters(transformer))
 
   outputted = args.prompt
   start_pos, toks = 0, tokenizer(outputted)["input_ids"]

examples/webgpu/stable_diffusion/compile.py

@@ -111,19 +111,19 @@ if __name__ == "__main__":
     return code
 
   def compile_step(model, step: Step):
-    linear, output_bufs = jit_model(step, *step.input)
-    functions, statements, bufs, _ = compile_net(linear, output_bufs)
+    run, special_names = jit_model(step, *step.input)
+    functions, statements, bufs, _ = compile_net(run, special_names)
     state = get_state_dict(model)
     weights = {id(x.uop.base.realized): name for name, x in state.items()}
     kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
     kernel_names = ', '.join([name for (name, _, _, _) in statements])
-    input_names = [f"input{i}" for i in range(len(step.input))]
-    output_names = [f"output{i}" for i in range(len(output_bufs))]
+    input_names = [name for _,name in special_names.items() if "input" in name]
+    output_names = [name for _,name in special_names.items() if "output" in name]
     input_buf_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
     output_buf_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
     kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
     exported_bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
-    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i in range(len(input_names))])
+    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
     input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new {input_buf_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,_ in enumerate(input_names)])
     return f"""\n    var {step.name} = function() {{
 
@@ -141,7 +141,7 @@ if __name__ == "__main__":
         const kernels = [{kernel_names}];
         const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
 
-        return async ({",".join([f'data{i}' for i in range(len(input_names))])}) => {{
+        return async ({",".join([f'data{i}' for i,(k,v) in enumerate(special_names.items()) if v != "output0"])}) => {{
             const commandEncoder = device.createCommandEncoder();
 
             {input_writer}

extra/amax/cast_amax.py

@@ -1,85 +0,0 @@
-import functools, pathlib
-from tinygrad import Tensor, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.renderer import Estimates
-from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
-
-FP8_MAX = 448.0
-NUM_WG, THREADS_PER_WG = 1024, 256
-
-def _compile(cpp_name:str, n_elems:int, hidden:int):
-  src = (pathlib.Path(__file__).parent/cpp_name).read_text()
-  defines = [f"-DN_ELEMS={n_elems}", f"-DHIDDEN={hidden}", f"-DNUM_WG={NUM_WG}", f"-DTHREADS_PER_WG={THREADS_PER_WG}"]
-  return src, HIPCCCompiler("gfx950", ["-std=c++20", "-ffast-math", *defines]).compile_cached(src)
-
-def _shard_shape(shape:tuple, axis:int, ndev:int) -> list:
-  s = list(shape); s[axis] //= ndev; return s
-
-@functools.cache
-def _custom_fused_bwd_w13(grad_xw13:UOp, xw13:UOp, grad_x2:UOp, amax_state:UOp, dname:str) -> UOp:
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  # read 2*N bf16 (xw13) + N bf16 (grad_x2) + 1 scalar; write 2*N bf16 (grad_xw13)
-  mem = n_elems * 2 * 5
-  sink = UOp.sink(grad_xw13.base, xw13.base, grad_x2.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_bwd_w13_{n_elems}", estimates=Estimates(ops=8*n_elems, mem=mem)))
-  src, lib = _compile("cast_amax_bwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-@functools.cache
-def _custom_fused_cast_amax_w13(fp8_out:UOp, amax_buf:UOp, xw13:UOp, amax_state:UOp, dname:str) -> UOp:
-  hidden = xw13.shape[2] // 2
-  n_elems = xw13.shape[0] * xw13.shape[1] * hidden
-  threads, workgroups = UOp.special(THREADS_PER_WG, "lidx0"), UOp.special(NUM_WG, "gidx0")
-  # read 2*N bf16 + 1 scalar, write N fp8 + NUM_WG bf16
-  mem = n_elems * 2 * 2 + n_elems + NUM_WG * 2
-  sink = UOp.sink(fp8_out.base, amax_buf.base, xw13.base, amax_state.base, threads, workgroups,
-                  arg=KernelInfo(f"fused_silu_mul_cast_amax_w13_{n_elems}", estimates=Estimates(ops=5*n_elems, mem=mem)))
-  src, lib = _compile("cast_amax_fwd_w13.cpp", n_elems, hidden)
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)),
-                               UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)))
-
-def _fused_quantize_bwd_w13(gradient:UOp, kernel:UOp):
-  # kernel.src[1:] is (fp8_out, amax_buf, xw13, amax_state); only xw13 needs a grad
-  _, _, xw13, amax_state = kernel.src[1:]
-  device = xw13.device
-  if isinstance(device, tuple):
-    axis, ndev = xw13.axis, len(device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    grad_xw13 = Tensor(Tensor.invalid(*_shard_shape(xw13.shape, axis, ndev), dtype=dtypes.bfloat16, device=device).uop.multi(axis), device=device)
-    dname = device[0].split(":")[0]
-  else:
-    grad_xw13 = Tensor.invalid(*xw13.shape, dtype=dtypes.bfloat16, device=device)
-    dname = device.split(":")[0] if isinstance(device, str) else device
-  grad_x2_t = Tensor(gradient, device=device).cast(dtypes.bfloat16)
-  fxn = functools.partial(_custom_fused_bwd_w13, dname=dname)
-  grad_xw13, *_ = Tensor.custom_kernel(grad_xw13, Tensor(xw13, device=device), grad_x2_t, Tensor(amax_state, device=device), fxn=fxn)
-  return (None, None, grad_xw13.uop, None)
-
-def fused_quantize_fp8_w13(xw13:Tensor, amax_state:Tensor, fp8_dtype) -> tuple[Tensor, Tensor, Tensor]:
-  # silu(xw1)*xw3 -> fp8 + amax over fused xw13 layout. Returns (fp8, inv_scale, new_amax).
-  assert xw13.dtype == dtypes.bfloat16, f"expected bf16, got {xw13.dtype}"
-  MBS, SEQ, H2 = xw13.shape
-  assert H2 % 2 == 0, f"w13 last-axis must be even, got {H2}"
-  HIDDEN = H2 // 2
-  if isinstance(xw13.device, tuple):
-    axis, ndev = xw13.uop.axis, len(xw13.device)
-    assert axis in (0, 1), f"unsupported sharding axis={axis}"
-    fp8_out = Tensor(Tensor.invalid(*_shard_shape((MBS, SEQ, HIDDEN), axis, ndev), dtype=fp8_dtype, device=xw13.device).uop.multi(axis), device=xw13.device)
-    amax_buf = Tensor(Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device).uop.multi(0), device=xw13.device)
-    dname = xw13.device[0].split(":")[0]
-  else:
-    fp8_out = Tensor.invalid(MBS, SEQ, HIDDEN, dtype=fp8_dtype, device=xw13.device)
-    amax_buf = Tensor.invalid(NUM_WG, dtype=dtypes.bfloat16, device=xw13.device)
-    dname = xw13.device.split(":")[0] if isinstance(xw13.device, str) else xw13.device
-  fxn = functools.partial(_custom_fused_cast_amax_w13, dname=dname)
-  fp8_out, amax_buf, *_ = Tensor.custom_kernel(fp8_out, amax_buf, xw13, amax_state, fxn=fxn, grad_fxn=_fused_quantize_bwd_w13)
-  # per-device scalar amax (no cross-device allreduce, matches _local_abs_max semantics)
-  if isinstance(amax_buf.device, tuple):
-    from examples.mlperf.models.flat_llama import _local_abs_max
-    new_amax = _local_abs_max(amax_buf).detach()
-  else: new_amax = amax_buf.max().detach()
-  inv_scale = (FP8_MAX / (amax_state + 1e-8)).float().reciprocal()
-  return fp8_out, inv_scale, new_amax

extra/amax/cast_amax_bwd_w13.cpp

@@ -1,68 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-
-#ifndef N_ELEMS
-#define N_ELEMS 234881024
-#endif
-#ifndef HIDDEN
-#define HIDDEN 14336
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC");
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_silu_mul_bwd_w13(
-    __hip_bfloat16*       __restrict__ grad_xw13_out,    // bf16, 2*N_ELEMS (interleaved layout)
-    const __hip_bfloat16* __restrict__ xw13,             // bf16, 2*N_ELEMS (interleaved)
-    const __hip_bfloat16* __restrict__ grad_x2,          // bf16, N_ELEMS
-    const __hip_bfloat16* __restrict__ amax_state)       // bf16 scalar
-{
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-  const int gid = wg * THREADS_PER_WG + tid;
-  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-
-  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    const int outer = base / HIDDEN;
-    const int inner = base % HIDDEN;
-    const int xw1_off = outer * 2 * HIDDEN + inner;
-    const int xw3_off = xw1_off + HIDDEN;
-
-    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
-    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
-    float4 g_raw  = *reinterpret_cast<const float4*>(&grad_x2[base]);
-
-    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
-    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
-    const __hip_bfloat16 *gv = reinterpret_cast<const __hip_bfloat16*>(&g_raw);
-
-    __hip_bfloat16 out1[VEC], out3[VEC];
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) {
-      const float f1 = static_cast<float>(x1[i]);
-      const float f3 = static_cast<float>(x3[i]);
-      const float fg = static_cast<float>(gv[i]);
-      const float sig = 1.0f / (1.0f + __expf(-f1));
-      const float silu = f1 * sig;
-      const float silu_prime = sig + silu * (1.0f - sig);
-      const float gs = fg * scale;
-      out1[i] = static_cast<__hip_bfloat16>(gs * silu_prime * f3);
-      out3[i] = static_cast<__hip_bfloat16>(gs * silu);
-    }
-
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw1_off]) = *reinterpret_cast<float4*>(out1);
-    *reinterpret_cast<float4*>(&grad_xw13_out[xw3_off]) = *reinterpret_cast<float4*>(out3);
-  }
-}

extra/amax/cast_amax_fwd_w13.cpp

@@ -1,79 +0,0 @@
-#include <hip/hip_runtime.h>
-#include <hip/hip_bf16.h>
-#include <hip/hip_fp8.h>
-
-#ifndef N_ELEMS
-#define N_ELEMS 234881024
-#endif
-#ifndef HIDDEN
-#define HIDDEN 14336
-#endif
-#ifndef NUM_WG
-#define NUM_WG 1024
-#endif
-#ifndef THREADS_PER_WG
-#define THREADS_PER_WG 256
-#endif
-
-constexpr int VEC = 8;
-constexpr float FP8_MAX = 448.0f;
-
-static_assert(N_ELEMS % VEC == 0, "N_ELEMS must be divisible by VEC");
-static_assert(HIDDEN % VEC == 0, "HIDDEN must be divisible by VEC (so VEC loads don't straddle block boundary)");
-
-extern "C" __global__ __launch_bounds__(THREADS_PER_WG) void
-fused_silu_mul_cast_amax_w13(
-    __hip_fp8_storage_t*  __restrict__ fp8_out,         // fp8, N_ELEMS
-    __hip_bfloat16*       __restrict__ amax_buf,        // bf16, NUM_WG (per-WG amaxes)
-    const __hip_bfloat16* __restrict__ xw13,            // bf16, 2*N_ELEMS
-    const __hip_bfloat16* __restrict__ amax_state)      // bf16 scalar
-{
-  __shared__ float sdata[THREADS_PER_WG];
-
-  const int tid = threadIdx.x;
-  const int wg  = blockIdx.x;
-  const int gid = wg * THREADS_PER_WG + tid;
-  const int stride_elems = NUM_WG * THREADS_PER_WG * VEC;
-
-  const float scale = FP8_MAX / (static_cast<float>(*amax_state) + 1e-8f);
-  float local_max = 0.0f;
-
-  // grid-stride over 8-element groups
-  for (int base = gid * VEC; base < N_ELEMS; base += stride_elems) {
-    // interleaved xw13 layout: xw1 and xw3 are not contiguous halves
-    const int outer = base / HIDDEN;
-    const int inner = base % HIDDEN;
-    const int xw1_off = outer * 2 * HIDDEN + inner;
-    const int xw3_off = xw1_off + HIDDEN;
-
-    float4 x1_raw = *reinterpret_cast<const float4*>(&xw13[xw1_off]);
-    float4 x3_raw = *reinterpret_cast<const float4*>(&xw13[xw3_off]);
-
-    const __hip_bfloat16 *x1 = reinterpret_cast<const __hip_bfloat16*>(&x1_raw);
-    const __hip_bfloat16 *x3 = reinterpret_cast<const __hip_bfloat16*>(&x3_raw);
-
-    __hip_fp8_storage_t out[VEC];
-    #pragma unroll
-    for (int i = 0; i < VEC; i++) {
-      const float f1 = static_cast<float>(x1[i]);
-      const float f3 = static_cast<float>(x3[i]);
-      const float silu = f1 / (1.0f + __expf(-f1));
-      const float x2 = silu * f3;
-      local_max = fmaxf(local_max, fabsf(x2));
-      const float x_scaled = fmaxf(-FP8_MAX, fminf(FP8_MAX, x2 * scale));
-      out[i] = __hip_cvt_float_to_fp8(x_scaled, __HIP_SATFINITE, __HIP_E4M3);
-    }
-
-    *reinterpret_cast<uint64_t*>(&fp8_out[base]) = *reinterpret_cast<uint64_t*>(out);
-  }
-
-  // LDS tree reduction: per-workgroup amax
-  sdata[tid] = local_max;
-  __syncthreads();
-  for (int s = THREADS_PER_WG / 2; s > 0; s >>= 1) {
-    if (tid < s) sdata[tid] = fmaxf(sdata[tid], sdata[tid + s]);
-    __syncthreads();
-  }
-
-  if (tid == 0) amax_buf[wg] = static_cast<__hip_bfloat16>(sdata[0]);
-}

extra/export_model.py

@@ -1,49 +1,47 @@
 from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType, dtypes
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Tensor
-from tinygrad.device import Device, Buffer
+from tinygrad.device import Device
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv, prod
-from tinygrad.uop.ops import Ops, UOp
-from tinygrad.codegen import get_program
+from tinygrad.helpers import Context, to_mv
+from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
 
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
 
-_KERNEL_ASTS = {Ops.SINK, Ops.PROGRAM, Ops.BEAM}
-def iter_kernel_calls(linear:UOp):
-  """Yield kernel CALLs from a LINEAR UOp. Toposort descends naturally into CUSTOM_FUNCTION graph batches; gate stops at kernel ASTs."""
-  return (u for u in linear.toposort(gate=lambda x: x.op not in _KERNEL_ASTS) if u.op is Ops.CALL and u.src[0].op in _KERNEL_ASTS)
+def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
+  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
+  canon, _seen = {}, {}
+  for ji in run.jit_cache:
+    for b in ji.bufs:
+      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
+  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}
 
-def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], List, Dict[str,Tuple[int,DType,int]], Dict[str,Buffer]]:
-  output_name = {id(b): f"output{i}" for i, b in enumerate(output_bufs)}
-  functions, bufs, bufs_to_save, statements, n = {}, {}, {}, [], 0
+  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
+  for ji in run.jit_cache:
+    fxn: ProgramSpec = ji.prg.p
+    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
+    cargs = []
+    for i,arg in enumerate(ji.bufs):
+      arg = canon[id(arg)]
+      key = id(arg)
+      if key not in bufs:
+        if key in special_names:
+          bufs[key] = (special_names[key], arg.size*arg.dtype.itemsize, arg.dtype, key)
+        else:
+          bufs[key] = (f"buf_{bufnum}", arg.size*arg.dtype.itemsize, arg.dtype, key)
+          bufnum += 1
+          if i > 0: bufs_to_save[bufs[key][0]] = arg   # if first usage of a buffer is not an output, and it's not a special name
+      cargs.append(bufs[key][0])
+    cargs += [var for var in fxn.vars if getattr(var, "op", None) is Ops.DEFINE_VAR] # symbolic vars; is it necessary or sufficient to check for DEFINE_VAR?
+    statements.append((fxn.function_name, cargs, fxn.global_size, fxn.local_size))
 
-  def name_of(bu:UOp, is_out:bool) -> str:
-    nonlocal n
-    if bu.op is Ops.PARAM: key, name, size = ("in", bu.arg), f"input{bu.arg}", prod(bu.shape)*bu.dtype.itemsize
-    else:
-      b = bu.buffer
-      key, size = (id(b.base), b.offset, b.size, b.dtype), b.size*b.dtype.itemsize
-      if key in bufs: return bufs[key][0]
-      if (name:=output_name.get(id(b))) is None:
-        name, n = f"buf_{n}", n+1
-        if not is_out: bufs_to_save[name] = b
-    bufs[key] = (name, size, bu.dtype, key)
-    return name
+  return functions, statements, {name:(size, dtype, key) for (name,size,dtype,key) in bufs.values()}, bufs_to_save
 
-  for call in iter_kernel_calls(linear):
-    arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
-    prg = get_program(call.src[0], Device[arg_uops[0].device].renderer)
-    functions[prg.function_name] = prg.src
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + [v for v in prg.vars if v.op is Ops.DEFINE_VAR]
-    statements.append((prg.function_name, cargs, prg.global_size, prg.local_size))
-
-  return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
-
-def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
+def jit_model(model, *args) -> Tuple[TinyJit,Dict[int,str]]:
   assert hasattr(model, "forward") or callable(model), "model needs a forward function"
   @TinyJit
   def run(*x):
@@ -52,10 +50,20 @@ def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
     out = [out] if isinstance(out, Tensor) else out
     return [o.realize() for o in out]
 
-  # run twice to trigger JIT capture
+  # twice to run the JIT
   for _ in range(2): the_output = run(*args)
-  assert run.captured is not None
-  return run.captured.linear, [o.uop.base.realized for o in the_output]
+  special_names = {}
+
+  # hack to put the inputs back
+  for (j,i),idx in run.input_replace.items():
+    realized_input = args[idx].uop.base.realized
+    run.jit_cache[j].bufs[i] = realized_input
+    special_names[id(realized_input)] = f'input{idx}'
+
+  # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
+  for i, output in enumerate(the_output):
+    special_names[id(output.uop.base.realized)] = f'output{i}'
+  return run, special_names
 
 def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,int,int]], bufs:Dict[str,Tuple[str,int,int]],
   bufs_to_save:Dict[str,Tensor], input_names:List[str], output_names:List[str], weight_names={}, model_name="model", symbolic_vars={}, wasm=False) -> str:
@@ -241,12 +249,12 @@ def export_model(model, target:str, *inputs, model_name: Optional[str] = "model"
   assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
 
   # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): linear, output_bufs = jit_model(model, *inputs)
-  functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
+  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
+  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
   state = get_state_dict(model)
   weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
-  input_names = [f"input{i}" for i in range(len(inputs))]
-  output_names = [f"output{i}" for i in range(len(output_bufs))]
+  input_names = [name for _,name in special_names.items() if "input" in name]
+  output_names = [name for _,name in special_names.items() if "output" in name]
 
   # handle symbolic variables; TODO: refactor to fix some of this stuff upstream in tinygrad
   symbolic_vars = OrderedDict()

extra/gemm/amd_copy_matmul.py

@@ -1,4 +1,4 @@
-from tinygrad import Device, UOp, getenv
+from tinygrad import UOp, getenv
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops
 from tinygrad.dtype import AddrSpace, dtypes
 
@@ -13,23 +13,18 @@ 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
 
@@ -76,10 +71,7 @@ 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.nbytes)
-b = MallocAllocator.alloc(nb.nbytes)
-c = MallocAllocator.alloc(nc.nbytes)
+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)
 
 MallocAllocator._copyin(b, flat_mv(nb.data))
 MallocAllocator._copyin(c, flat_mv(nc.data))

extra/gemm/cdna_asm_gemm.py

@@ -6,7 +6,7 @@ from tinygrad.renderer import Estimates
 from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
 from tinygrad.runtime.autogen.amd.cdna.ins import *
-from examples.mlperf.models.flat_llama import FP8_DTYPE, FP8_GRAD_DTYPE, quantize_fp8
+from examples.mlperf.models.flat_llama import FP8_DTYPE, FP8_GRAD_DTYPE, matmul, quantize_fp8
 
 # ** CDNA4 assembly gemm
 
@@ -2628,15 +2628,15 @@ def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
 # ** FP8 GEMM custom kernel
 
 @functools.cache
-def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, X_s:UOp, W_s:UOp, dname:str) -> UOp:
-  # A is (batch, M, K), B is (N, K) transposed, X_s is x_scale, W_s is w_scale — kernel multiplies by both
+def custom_hk_fp8_gemm(C:UOp, A:UOp, B:UOp, S:UOp, dname:str) -> UOp:
+  # A is (batch, M, K), B is (N, K) transposed, S is combined scale (scalar float)
   M, K = A.shape[0]*A.shape[1], A.shape[2]
   N, K2 = B.shape[(1 if B.ndim == 3 else 0):]
   assert K == K2, f"{A.shape} {B.shape}"
   block_size = 256
   threads = UOp.special(64 * 8, "lidx0")
   workgroups = UOp.special((M // block_size) * (N // block_size), "gidx0")
-  sink = UOp.sink(C.base, A.base, B.base, X_s.base, W_s.base, threads, workgroups,
+  sink = UOp.sink(C.base, A.base, B.base, S.base, threads, workgroups,
                   arg=KernelInfo(f"hk_fp8_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K, mem=(M*K+N*K)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
   kittens_path = pathlib.Path(__file__).parent.parent/"thunder"/"amd"
   src = (kittens_path/"gemm_fp8.cpp").read_text()
@@ -2698,34 +2698,31 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
 
 def custom_gemm_bw(gradient:UOp, kernel:UOp):
   inputs = kernel.src[1:]
-  # fp8 scaled gemm has 5 inputs (out, a, b, x_scale, w_scale), others have 3 (out, a, b)
-  if len(inputs) == 5:
-    out, a, b, s_x, s_w = inputs
-    a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
-    s_x_t, s_w_t = Tensor(s_x, device=a.device), Tensor(s_w, device=a.device)
+  # fp8 scaled gemm has 4 inputs (out, a, b, scale), others have 3 (out, a, b)
+  if len(inputs) == 4:
+    out, a, b, scale = inputs
+    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]]
-    g_fp8, g_scale, _ = quantize_fp8(g_t)
-    # dgrad: uses g_scale * x_scale * w_scale
-    grad_a = asm_gemm(g_fp8, b_t, x_scale=g_scale * s_x_t, w_scale=s_w_t)
-    # wgrad: no w_scale
-    _one = Tensor(1.0, dtype=dtypes.float, device=a.device)
-    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), x_scale=g_scale * s_x_t, w_scale=_one)
-    return (None, grad_a.uop, grad_b.uop, None, None)
+    # backward GEMMs in fp8 with scale applied inside kernel to prevent bf16 overflow
+    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)
+    # wgrad: g_fp8.T @ activation = (N, batch*seq) @ (batch*seq, K) → use permute to preserve sharding
+    grad_b = asm_gemm(g_fp8.permute(2, 0, 1).reshape(g_t.shape[-1], -1), a_t.reshape(-1, a_t.shape[-1]), combined_scale=bw_scale)
+    return (None, grad_a.uop, grad_b.uop, None)
   else:
     out, a, b = inputs
     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]]
-    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
+    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
     return (None, grad_a, grad_b)
 
 # ** main gemm function
 
-def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=None) -> Tensor:
+def asm_gemm(a:Tensor, b:Tensor, combined_scale:Tensor|None=None) -> Tensor:
   assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
   counters["used"] += 1
   unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
@@ -2757,12 +2754,10 @@ def asm_gemm(a:Tensor, b:Tensor, x_scale:Tensor|None=None, w_scale:Tensor|None=N
   renderer = Device[dname:=(a.device[0] if is_multi else a.device)].renderer
   dname, arch = dname.split(":")[0], renderer.target.arch
   if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    # fp8 gemm computes [email protected], kernel multiplies output by x_scale * w_scale before bf16 store
+    # fp8 gemm computes [email protected], with optional combined scale applied inside kernel before bf16 store
     if a.dtype == FP8_DTYPE:
-      _one = lambda: Tensor(1.0, dtype=dtypes.float, device=a.device)
-      xs = x_scale if x_scale is not None else _one()
-      ws = w_scale if w_scale is not None else _one()
-      out = Tensor.custom_kernel(out, a, b.T, xs, ws, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+      scale = combined_scale if combined_scale is not None else Tensor(1.0, dtype=dtypes.float, device=a.device)
+      out = Tensor.custom_kernel(out, a, b.T, scale, fxn=functools.partial(custom_hk_fp8_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
     else:
       out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
   else:

extra/gemm/tinygrad_nv_matmul.py

@@ -29,7 +29,7 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
   k.apply_opts(opts)
-  prg = get_program(k.ast.replace(arg=replace(k.ast.arg, opts_to_apply=tuple(k.applied_opts))), k.opts)
+  prg = get_program(k.ast, k.opts, k.applied_opts)
   new_src = prg.src
   # can mod source here
   prg = replace(prg, src=new_src)

extra/mlx_driver/loopback.py

@@ -6,8 +6,7 @@ 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 tinygrad.helpers import to_be32, to_be64
-from extra.mlx_driver.mlxdev import MLXDev, MLXQP
+from extra.mlx_driver.mlxdev import MLXDev, MLXQP, to_be
 
 BUF_SIZE = 0x1000
 MLX_PCI = getenv("MLX_PCI", "0000:41:00.0")
@@ -50,7 +49,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_be32(qp.rq_head)
+dev.dbr[qp.qp_dbr // 4] = to_be('I', qp.rq_head)
 
 # build send WQE in SQ from CPU (opcode 0x0a = SEND, ds_count=2)
 sq_head = qp.sq_head
@@ -61,7 +60,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_be64(int.from_bytes(bytes(wqe[0:8]), 'big'))
+doorbell_val = to_be('Q', 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
@@ -73,7 +72,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_be32(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_be('I', 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, Ops
+from tinygrad.uop.ops import AxisType
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.codegen.opt.postrange import Scheduler
 from tinygrad.codegen.opt.search import get_kernel_actions
@@ -85,7 +85,6 @@ 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)])
@@ -96,7 +95,6 @@ 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

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

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

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

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

extra/remu/src/helpers.rs

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

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

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

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

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

@@ -0,0 +1,155 @@
+# 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 library (note, that not default locations should be in LD_LIBRARY_PATH, so tinygrad can find it)."
+    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 "Press any key or symbol to continue..."
     read -n 1 -s
 
@@ -11,6 +11,11 @@ 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,7 +10,6 @@ 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__":
@@ -22,6 +21,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, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
+                     env={**os.environ, "AMD":"1", "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,8 +136,7 @@ def print_data(data:dict) -> None:
 
 def main() -> None:
   import tinygrad.viz.serve as viz
-  from tinygrad.uop.ops import RewriteTrace
-  data = viz.VizData()
+  viz.ctxs = []
 
   parser = argparse.ArgumentParser()
   parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
@@ -148,24 +147,24 @@ def main() -> None:
 
   with args.profile.open("rb") as f: profile = pickle.load(f)
 
-  viz.get_profile(profile, data=data)
+  viz.get_profile(profile)
 
   # List all kernels
   if args.kernel is None:
-    for c in data.ctxs:
+    for c in viz.ctxs:
       print(c["name"])
       for s in c["steps"]: print("  "+s["name"])
     return None
 
   # Find kernel trace
-  trace = next((c for c in data.ctxs if c["name"] == f"SQTT {args.kernel}"), None)
+  trace = next((c for c in viz.ctxs if c["name"] == f"Exec {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"])
-    ret = viz.get_render(data, s["query"])
-    print_data(ret)
+    data = viz.get_render(s["query"])
+    print_data(data)
     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), attn, l_vec
+  return attn.transpose(1, 2)
 
 @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/thunder/amd/gemm_fp8.cpp

@@ -93,7 +93,7 @@ constexpr int NUM_WARPS = 8;
 
 using G = kittens::group<NUM_WARPS>;
 
-__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *x_scale_ptr, float *w_scale_ptr) {
+__global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_ptr, fp8e4m3 *B_ptr, float *scale_ptr) {
     constexpr int M = GEMM_M, N = GEMM_N, K = GEMM_K;
 
     kittens::gl<fp8e4m3, 1, 1, M, K> A{A_ptr, nullptr, nullptr, nullptr, nullptr};
@@ -332,8 +332,8 @@ __global__ __launch_bounds__(512, 2) void hk_fp8_gemm(bf16 *C_ptr, fp8e4m3 *A_pt
         __builtin_amdgcn_s_barrier();
     }
 
-    // apply x_scale * w_scale before bf16 store to prevent overflow
-    float scale = *x_scale_ptr * *w_scale_ptr;
+    // apply combined scale (x_scale * w_scale) before bf16 store to prevent overflow
+    float scale = *scale_ptr;
     mul(cA, cA, scale);
     mul(cB, cB, scale);
     mul(cC, cC, scale);

extra/torch_backend/example.py

@@ -24,7 +24,7 @@ if __name__ == "__main__":
   kernel_count = GlobalCounters.kernel_count
   assert kernel_count > 0, "No kernels, test failed"
   # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 355
+  expected_kernels = 334
   expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
   if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
   assert kernel_count <= expected_kernels, f"{expectation}"

extra/torch_backend/test_kernel_fusion.py

@@ -23,7 +23,7 @@ class TestKernelFusionRegression(unittest.TestCase):
     def fn():
       x = torch.randn(128, 128, device=device)
       return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_relu_fusion(self):
     def fn():
@@ -31,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
       with torch.no_grad():
         return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 8)
 
   def test_batchnorm_fusion(self):
     def fn():
@@ -41,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.eval()
       with torch.no_grad():
         return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 11)
+    self._check_kernel_count(fn, 12)
 
   def test_reduce_fusion(self):
     def fn():
       x = torch.randn(64, 64, device=device)
       return (x * 2.0).sum()
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_matmul_elementwise_fusion(self):
     def fn():
       x = torch.randn(32, 32, device=device)
       w = torch.randn(32, 32, device=device)
       return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 9)
 
   def test_pooling_fusion(self):
     def fn():
       x = torch.randn(1, 8, 16, 16, device=device)
       return torch.nn.functional.max_pool2d(x * 2.0, 2)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 7)
 
   def test_residual_add_relu_fusion(self):
     def fn():
@@ -68,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       identity = torch.randn(1, 8, 16, 16, device=device)
       out = x + identity
       return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 9)
 
   def test_inplace_add_relu_fusion(self):
     def fn():
@@ -76,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       y = torch.randn(1, 16, 32, 32, device=device)
       x += y
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 9)
 
   def test_conv_bn_add_relu_fusion(self):
     def fn():
@@ -89,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
         out = bn(conv(x))
         out += identity
         return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 13)
+    self._check_kernel_count(fn, 14)
 
   def test_multiple_inplace_ops_fusion(self):
     def fn():
@@ -97,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       x += 1.0
       x *= 2.0
       return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 5)
+    self._check_kernel_count(fn, 6)
 
   def test_view_inplace_no_fusion_break(self):
     def fn():
@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       view = x[1:3]
       view += 1.0
       return x.sum()
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 10)
 
   def test_batchnorm_running_stats_update(self):
     def fn():
@@ -114,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       bn.train()
       with torch.no_grad():
         return bn(x)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 10)
 
   # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
   def test_mnist_training_fusion(self):
@@ -135,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
       loss.backward()
       optimizer.step()
       return loss
-    self._check_kernel_count(fn, 25)
+    self._check_kernel_count(fn, 26)
 
 if __name__ == "__main__":
   unittest.main()

extra/viz/README.md

@@ -0,0 +1,45 @@
+A command line tool for exploring the VIZ trace.
+
+# Lightweight tracing
+
+Supported on all backends.
+
+Flags: VIZ=-1 to only save the trace to a file, VIZ=1 also launches a web server.
+
+1. Set VIZ to -1 to save the trace.
+2. Use `extra/viz/cli.py` to inspect the trace files.
+
+## Inspect runtime profiling
+
+Use `extra/viz/cli.py --profile` to list all sources.
+
+List top slowest kernels on a source: `--profile -s "AMD"`
+List samples of a kernel on a source: `--profile -s "AMD" -i E_3 | head 4`
+
+## Inspect codegen and PatternMatcher
+
+Use `extra/viz/cli.py --rewrites` to list all sources.
+
+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
+
+Supported on AMD for RDNA3 and RDNA4 (best) and CDNA (developing).
+
+Flags: VIZ=-2 to save SQTT trace to a file. VIZ=2 also launches a web server. View other flags in tinygrad/runtime/ops_amd.py to configure SQTT as needed.
+
+Use `extra/viz/cli.py --profile | grep SQTT` to view all available SQTT traces.
+You can select a specific trace with --source, Example workflow:
+
+```bash
+# Run amd_asm_matmul with VIZ=-2 to capture the trace
+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
+
+# 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

@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+import argparse, pathlib, signal, sys, struct, json, itertools
+if hasattr(signal, "SIGPIPE"): signal.signal(signal.SIGPIPE, signal.SIG_DFL)
+from typing import Iterator
+from tinygrad.viz import serve as viz
+from tinygrad.uop.ops import RewriteTrace
+from tinygrad.helpers import temp, ansistrip, colored, time_to_str, ansilen, ProfilePointEvent, ProfileRangeEvent, TracingKey, unwrap
+
+# profile decoder used in CLI and tests
+def decode_profile(data:bytes) -> dict:
+  ret, off = data, 0
+  def u(fmt:str) -> tuple:
+    nonlocal off
+    vals = struct.unpack_from(fmt, ret, off)
+    off += struct.calcsize(fmt)
+    return vals
+  total_dur, global_peak, index_len, layout_len = u("<IQII")
+  strings, dtypes, markers = json.loads(ret[off:off+index_len]).values()
+  off += index_len
+  layout:dict[str, dict] = {}
+  # 0 means None, otherwise it's an enum value
+  def option(i:int) -> int|None: return None if i == 0 else i-1
+  for _ in range(layout_len):
+    klen = u("<B")[0]
+    k = ret[off:off+klen].decode()
+    off += klen
+    v:dict = {"events":[]}
+    layout[k] = v
+    event_type, event_count = u("<BI")
+    if event_type == 0:
+      for _ in range(event_count):
+        name, ref, key, st, dur, fmt = u("<IIIIfI")
+        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
+    else:
+      v["linear"] = u("<B")[0]
+      v["peak"] = u("<Q")[0]
+      for _ in range(event_count):
+        if v["linear"]:
+          ts, value = u("<IQ")
+          v["events"].append({"event":"freq", "ts":ts, "value":value})
+        else:
+          alloc, ts, key = u("<BII")
+          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}
+
+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')
+
+def main(args) -> None:
+  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
+  viz.ctxs = viz.get_rewrites(viz.trace)
+
+  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}")
+    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"]])
+    if args.src is None:
+      for k in profile["layout"]:
+        print(f"  {format_colored(k)}")
+      return None
+
+    # ** SQTT printer
+    data = get(profile["layout"], args.src)
+    if "SQTT" in args.src:
+      # modern terminals support 24-bit color
+      def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
+      print(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Delay':<4} {'Info'}")
+      print("-" * 100)
+      pc_map:dict[int, str] = {}
+      pkt_idxs:dict[str, itertools.count] = {}
+      dispatch_to_inst:dict[str, tuple[str, int]] = {}
+      inst_st:int|None = None
+      for e in viz.sqtt_timeline(*data):
+        if isinstance(e, ProfilePointEvent) and e.key == 'pcMap': pc_map = e.arg
+        if not isinstance(e, ProfileRangeEvent): continue
+        if inst_st is None: inst_st = int(e.st)
+        assert isinstance(e.name, TracingKey)
+        op_name, info = e.name.display_name, e.name.ret or ""
+        color = next((v for k,v in viz.wave_colors.items() if k in op_name), 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":
+          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"
+        if info.startswith("LINK:"):
+          inst, dispatch_st = dispatch_to_inst[info.replace("LINK:", "")]
+          phase, delay = "EXEC", int(e.st) - dispatch_st
+        if inst and phase: info = f"{phase:<8} {inst}"
+        unit = e.device.replace(" ", "-")
+        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
+
+    # ** Profiler printer
+    agg:dict[str, tuple[float, int]] = {}
+    total = 0
+    for e in data.get("events", []):
+      et = e["dur"] * 1e-6
+      if args.item is not None:
+        if ansistrip(e["name"]) == args.item:
+          ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None)
+          name = e["name"] + (" " * (46 - ansilen(e["name"])))
+          print(f"{format_colored(name)} {ptm}/{et*1e3:9.2f}ms  " + e.get("fmt", "").replace("\n", " | ") + "  ")
+      else:
+        t, c = agg.get(e["name"], (0.0, 0))
+        agg[e["name"]] = (t+et, c+1)
+        total += et
+    if agg and total > 0:
+      from tabulate import tabulate
+      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
+      rows = 20
+      table = [[format_colored(name), time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in items[:rows]]
+      if items[rows:]:
+        other_t = sum(t for _,(t,_) in items[rows:])
+        other_c = sum(c for _,(_,c) in items[rows:])
+        table.append(["Other", time_to_str(other_t, w=9), other_c, f"{(other_t/total*100.0):.2f}%"])
+      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
+    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")}
+  if args.src is None:
+    for k in rewrites: print(f"  {format_colored(k)}")
+    return None
+  steps = get(rewrites, args.src)
+  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"])
+    if isinstance(data.get("value"), Iterator):
+      for m in data["value"]:
+        if m.get("uop"): print(f"Input UOp:\n{m['uop']}")
+        if m.get("diff"):
+          loc = pathlib.Path(m["upat"][0][0])
+          print(f"Rewrite at {loc.parent.name}/{loc.name}:{m['upat'][0][1]}\n{m['upat'][1]}")
+          for line in m["diff"]:
+            print(line if args.no_color else colored(line, "red" if line.startswith("-") else "green" if line.startswith("+") else None))
+    if data.get("src") is not None: print(data["src"])
+
+def get_arg_parser() -> argparse.ArgumentParser:
+  parser = argparse.ArgumentParser(add_help=False)
+  g_mode = parser.add_argument_group("mode")
+  g_mode.add_argument("-p", "--profile", action="store_true", help="View profile")
+  g_mode.add_argument("-r", "--rewrites", action="store_true", help="View graph rewrites")
+  g_opts = parser.add_argument_group("optional args")
+  g_opts.add_argument("-s", "--src", type=str, default=None, metavar="NAME", help="Select a data source (default: list all sources)")
+  g_opts.add_argument("-i", "--item", type=str, default=None, metavar="NAME", help="Select an item within the source (default: list all items)")
+  g_opts.add_argument("--no-color", action="store_true", help="Turn off colored names")
+  g_opts.add_argument("--profile-path", type=pathlib.Path, metavar="PATH", help="Path to profile.pkl (optional file, default: latest profile)",
+                      default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  g_opts.add_argument("--rewrites-path", type=pathlib.Path, metavar="PATH", help="Path to rewrites.pkl (optional file, default: latest rewrites)",
+                      default=pathlib.Path(temp("rewrites.pkl", append_user=True)))
+  g_opts.add_argument("-h", "--help", action="help", help="show this help message and exit")
+  return parser
+
+if __name__ == "__main__":
+  args = get_arg_parser().parse_args()
+  if not args.profile and not args.rewrites:
+    get_arg_parser().print_help()
+    sys.exit(0)
+
+  try: main(args)
+  except KeyboardInterrupt: pass

pyproject.toml

@@ -19,11 +19,11 @@ build-backend = "setuptools.build_meta"
 include-package-data = true
 packages = [
   'tinygrad',
+  'tinygrad.apps',
   'tinygrad.codegen',
   'tinygrad.codegen.opt',
   'tinygrad.codegen.late',
   'tinygrad.engine',
-  'tinygrad.llm',
   'tinygrad.mixin',
   'tinygrad.nn',
   'tinygrad.renderer',
@@ -38,7 +38,6 @@ packages = [
   'tinygrad.runtime.graph',
   'tinygrad.runtime.support',
   'tinygrad.runtime.support.am',
-  'tinygrad.runtime.support.mlx',
   'tinygrad.runtime.support.nv',
   'tinygrad.schedule',
   'tinygrad.uop',
@@ -51,8 +50,8 @@ tinygrad = ["py.typed"]
 
 
 [project.optional-dependencies]
-# arm = ["unicorn"]
-# triton = ["triton-nightly>=2.1.0.dev20231014192330"]
+arm = ["unicorn"]
+triton = ["triton-nightly>=2.1.0.dev20231014192330"]
 linting = [
   "pylint",
   "mypy==1.19.1",
@@ -75,7 +74,7 @@ testing_minimal = [
   "hypothesis>=6.148.9",
   "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
 ]
-testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18", "capstone"]
+testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18"]
 testing = [
   "tinygrad[testing_unit]",
   "pillow",
@@ -93,6 +92,7 @@ testing = [
   "networkx",
   "nibabel",
   "bottle",
+  "capstone",
   "pycocotools",
   "boto3",
   "pandas",
@@ -112,9 +112,9 @@ docs = [
 [tool.mutmut]
 paths_to_mutate = ["tinygrad/"]
 do_not_mutate = [
+  "tinygrad/apps/*",
   "tinygrad/codegen/*",
   "tinygrad/engine/*",
-  "tinygrad/llm/*",
   "tinygrad/nn/*",
   "tinygrad/renderer/*",
   "tinygrad/runtime/*",

sz.py

@@ -56,7 +56,7 @@ def gen_diff(table_old, table_new):
 
 def display_diff(diff): return "+"+str(diff) if diff > 0 else str(diff)
 
-NONCORE_DIRS = {"tinygrad/llm", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
+NONCORE_DIRS = {"tinygrad/apps", "tinygrad/nn", "tinygrad/renderer", "tinygrad/runtime", "tinygrad/viz"}
 
 if __name__ == "__main__":
   if len(sys.argv) == 3:

test/amd/disasm.py

@@ -2,7 +2,6 @@
 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
@@ -839,11 +838,22 @@ def _disasm_vop1_sdwa(inst) -> str:
 
 def _decode_dpp(dpp: int) -> str:
   """Decode DPP control value to string."""
-  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}"
+  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}"
 
 def _disasm_vop1_dpp(inst) -> str:
   name = inst.op_name.lower().replace('_e32', '')

test/amd/helpers.py

@@ -12,19 +12,8 @@ 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

@@ -1,187 +0,0 @@
-"""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,6 +833,8 @@ 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

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

@@ -1,35 +0,0 @@
-"""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,17 +30,6 @@ 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,28 +20,6 @@ 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 = [
@@ -1614,7 +1592,8 @@ 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 saturates to 0 on RDNA3 hardware."""
+    """Clamp with NaN input should still produce NaN."""
+    import math
     quiet_nan = 0x7fc00000
     instructions = [
       s_mov_b32(s[0], quiet_nan),
@@ -1622,7 +1601,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.assertEqual(st.vgpr[0][1], 0)
+    self.assertTrue(math.isnan(i2f(st.vgpr[0][1])))
 
   def test_omod_ignored(self):
     """OMOD field is ignored on RDNA3 hardware."""
@@ -3626,30 +3605,32 @@ 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], v[255]),
+      v_mov_b32_e32(v[0], 0xCAFEBABE),  # source data
       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=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)
+    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)
 
   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], v[255]),
+      v_mov_b32_e32(v[0], 0x11111111),  # All lanes have this initially
       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)
-    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)
+    # 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)
 
 
 if __name__ == '__main__':

test/amd/test_compare_emulators.py

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

@@ -15,7 +15,7 @@ from extra.gemm.amd_asm_matmul import Kernel
 def custom_add_one(A:UOp) -> UOp:
   A = A.flatten()
   assert dtypes.is_float(A.dtype.base), f"buffer dtype must be float32, got {A.dtype}"
-  threads = UOp.special(A.numel(), "lidx0")
+  threads = UOp.special(A.size, "lidx0")
   insts = [
     s_load_b64(s[0:1], s[0:1], soffset=NULL),
     s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
@@ -27,13 +27,13 @@ def custom_add_one(A:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.numel()}", estimates=Estimates(ops=A.numel(), mem=A.numel()*4*2)))
+  sink = UOp.sink(A.base, threads, arg=KernelInfo(f"custom_add_one_{A.size}", estimates=Estimates(ops=A.size, mem=A.size*4*2)))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
 def custom_add_var(A:UOp, B:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
   assert A.dtype.base == dtypes.uint32, f"buffer dtype must be uint32, got {A.dtype}"
-  threads = UOp.special(A.numel(), "lidx0")
+  threads = UOp.special(A.size, "lidx0")
   var = UOp.variable("var", 0, 10)
   insts = [
     s_load_b128(s[4:7], s[0:1]),
@@ -46,7 +46,7 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
     global_store_b32(addr=v[0], data=v[1], saddr=s[4:5]),
     s_endpgm(),
   ]
-  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.numel()}"))
+  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
 
 def custom_wave_sync(A:UOp, arch:str) -> UOp:
@@ -104,21 +104,20 @@ 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], 4))
+  k.emit(r4.v_mov_b32_e32(v[1], 10))
   def emit_alt():
-    for i in range(2):
+    for i in range(4):
       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(2):
+    for _ in range(4):
       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(2):
+  for _ in range(4):
     emit_wmma()
     emit_alt()
   for _ in range(8): k.emit(s_nop(1))
@@ -132,7 +131,7 @@ def custom_handwritten(A:UOp, arch:str) -> UOp:
 
 def custom_data_deps(A:UOp, arch:str) -> UOp:
   A = A.flatten()
-  threads = UOp.special(A.numel(), "lidx0")
+  threads = UOp.special(A.size, "lidx0")
   k = Kernel(arch)
   k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
   k.emit(s_waitcnt_lgkmcnt(sdst=NULL, simm16=0))

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, GLOBALOp
+from tinygrad.runtime.autogen.amd.rdna3.enum import VOP1Op, VOP2Op, SOP2Op, DSOp
 
 def _srcs():
   """Create minimal source variables for pcode parsing."""
@@ -113,7 +113,6 @@ 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)."""
 
@@ -165,20 +164,6 @@ 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)
@@ -300,47 +285,6 @@ 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

@@ -37,7 +37,9 @@ dev.synchronize()
 '''
 
     env = os.environ.copy()
-    env["DEV"] = "MOCKKFD+AMD"
+    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,7 +1,6 @@
 #!/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
@@ -45,64 +44,6 @@ 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."""
@@ -116,6 +57,7 @@ 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_encoder.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """Tests for SQTT encoder: verifies the emulator produces correct SQTT traces for known kernels.
 
-Run with: DEV=MOCKKFD+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
+Run with: DEV=AMD MOCKGPU=1 python -m pytest test/amd/test_sqtt_encoder.py -v
 """
 import ctypes, unittest
 from tinygrad.helpers import Context

test/amd/test_sqtt_profiler.py

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

test/amd/test_sqttmap.py

@@ -1,5 +1,5 @@
 # test to compare every packet with the rocprof decoder
-import unittest, pickle
+import unittest, pickle, contextlib, io
 from typing import Iterator
 from pathlib import Path
 from tinygrad.helpers import DEBUG, getenv, temp, ansistrip
@@ -7,11 +7,17 @@ from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from tinygrad.viz.serve import sqtt_timeline
 from test.amd.disasm import disasm
-from test.null.test_viz import run_cli
 
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
 
+def run_cli(*cli_args) -> str:
+  from extra.viz.cli import main, get_arg_parser
+  args = get_arg_parser().parse_args(cli_args)
+  with contextlib.redirect_stdout(buf:=io.StringIO()):
+    main(args)
+  return buf.getvalue().strip()
+
 def rocprof_inst_traces_match(sqtt, prg, target):
   from tinygrad.viz.serve import amd_decode
   from extra.sqtt.roc import decode as roc_decode, InstExec
@@ -94,7 +100,9 @@ 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"} and not e.name.display_name.startswith("OTHER_"): insts += 1
+                                             "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
             else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
           self.assertEqual(execs, insts)
 
@@ -123,18 +131,7 @@ class TestSQTTMapBase(unittest.TestCase):
 
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"
 
-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 TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"
 
 class TestSQTTMapCDNA(TestSQTTMapBase):
   target = "gfx950"

test/backend/test_arange.py

@@ -4,7 +4,8 @@ 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.schedule import ExecItem
+from tinygrad.engine.schedule import ExecItem
+from tinygrad.uop.ops import Ops
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
 from test.helpers import needs_second_gpu
@@ -23,10 +24,6 @@ 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])
@@ -66,7 +63,7 @@ class TestIndexing(unittest.TestCase):
     print("*** indexing ***")
     with Context(NOOPT=1):
       GlobalCounters.reset()
-      rng = Tensor.arange(DSET, dtype=dtypes.int).reshape(1, 1, DSET, 1).expand(4, DDIM, DSET, 1)
+      rng = Tensor.ones(4, DDIM, DSET, dtype=dtypes.int)._cumalu(axis=-1, op=Ops.ADD, _include_initial=True).reshape(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

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
-from tinygrad.helpers import getenv, system, DEV
+from tinygrad.helpers import getenv, system
 from extra.gemm.cdna_asm_gemm import asm_gemm
 from test.helpers import needs_second_gpu
 from examples.mlperf.models.flat_llama import FP8_DTYPE
@@ -21,8 +21,9 @@ 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)
-  tst = asm_gemm(a, b)
-  tst.sum().backward()
+  with Context(ASM_GEMM=1):
+    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_()
@@ -31,8 +32,9 @@ 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)
-  ref = a_ref @ b_ref
-  ref.sum().backward()
+  with Context(ASM_GEMM=0):
+    ref = asm_gemm(a_ref, b_ref)
+    ref.sum().backward()
   Tensor.realize(ref, a_ref.grad, b_ref.grad)
 
   # no validation on the NULL device
@@ -131,15 +133,17 @@ class TestGemmLlama(unittest.TestCase):
   dtype = dtypes.bfloat16
 
   def setUp(self):
-    if not is_cdna4() or DEV.interface.startswith("MOCK"):
+    if not is_cdna4() or getenv("MOCKGPU"):
       self.skipTest("very slow on non mi350x")
 
-  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(self): (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 = asm_gemm(x, y)
+    z = 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_const_folding.py

@@ -1,7 +1,7 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.dtype import DTYPES_DICT
-from tinygrad.uop.ops import Ops, UOp
+from tinygrad.uop.ops import Ops
 from tinygrad.device import is_dtype_supported
 import numpy as np
 from test.helpers import not_support_multi_device
@@ -163,11 +163,6 @@ class TestMultiConstFolding(unittest.TestCase):
     np.testing.assert_equal((t ** one).numpy(), np.arange(16))
     np.testing.assert_equal((one ** t).numpy(), [1] * 16)
 
-class TestThreefryConstFolding(unittest.TestCase):
-  def test_threefry(self):
-    x = UOp.const(dtypes.uint64, 5, Device.DEFAULT, ()).threefry(UOp.const(dtypes.uint64, 10, Device.DEFAULT, ()))
-    self.assertIs(x.simplify().op, Ops.CONST)
-
 class TestTautologicalCompare(unittest.TestCase):
   # without const folding, these would have triggered -Wtautological-compare in clang
   def test_lt_false(self):

test/backend/test_custom_kernel.py

@@ -6,32 +6,32 @@ from tinygrad.uop.ops import KernelInfo, AxisType
 # **** kernels ****
 
 def custom_arange_kernel(C:UOp) -> UOp:
-  i = UOp.range(C.shape[0], 0)
-  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.shape[0]}"))
+  i = UOp.range(C.size, 0)
+  return C[i].store(i.cast(C.dtype.base)).end(i).sink(arg=KernelInfo(name=f"custom_arange_{C.size}"))
 
 def custom_eye_kernel(C:UOp) -> UOp:
   i = UOp.range(C.shape[0], 0)
   j = UOp.range(C.shape[1], 1)
-  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.numel()}"))
+  return C[i, j].store((i.eq(j)).cast(C.dtype.base)).end(i, j).sink(arg=KernelInfo(name=f"custom_eye_{C.size}"))
 
 def custom_add_one_kernel(B:UOp, A:UOp) -> UOp:
   A,B = A.flatten(), B.flatten()
-  assert B.numel() == A.numel()
-  i = UOp.range(A.numel(), 0)
-  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.numel()}"))
+  assert B.size == A.size
+  i = UOp.range(A.size, 0)
+  return B[i].store(A[i] + 1).end(i).sink(arg=KernelInfo(name=f"add_one_{A.size}"))
 
 def custom_elementwise_add_kernel(C:UOp, A:UOp, B:UOp) -> UOp:
   C,A,B = C.flatten(), A.flatten(), B.flatten()
-  i = UOp.range(C.numel(), 0)
-  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.numel()}")).simplify()
+  i = UOp.range(C.size, 0)
+  return C[i].store(A[i]+B[i]).end(i).sink(arg=KernelInfo(name=f"custom_add_kernel_{C.size}")).simplify()
 
 def custom_elementwise_addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
   C,D,A,B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
-  assert C.numel() == D.numel()
-  i = UOp.range(C.numel(), 0)
+  assert C.size == D.size
+  i = UOp.range(C.size, 0)
   store_c = C[i].store(A[i]+B[i])
   store_d = D[i].store(A[i]*B[i])
-  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.numel()}")).simplify()
+  return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name=f"custom_addmul_kernel_{C.size}")).simplify()
 
 def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
   assert A.shape[1] == B.shape[0]
@@ -52,7 +52,7 @@ def flip_contract_kernel(dest:UOp, src:UOp):
   j = UOp.range(dest.shape[1], 1, AxisType.UPCAST)
   vec = src[i, j].contract(j)
   store = UOp.group(*[dest[i, k].store(vec.gep(3-k)) for k in range(4)])
-  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.numel()}", opts_to_apply=()))
+  return store.end(i, j).sink(arg=KernelInfo(name=f"flip_contract_{dest.size}", opts_to_apply=()))
 
 def slice_sum_kernel(dest:UOp, src:UOp):
   G = UOp.range(src.shape[0], 0)
@@ -291,10 +291,10 @@ class TestCustomKernel(unittest.TestCase):
 
     def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
       o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
-      i = UOp.range(o1.numel(), 0)
+      i = UOp.range(o1.size, 0)
       store_o1 = o1[i].store(A[i]+B[i])
       store_o2 = o2[i].store(A[i]+B[i]+2)
-      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.numel()}")).simplify()
+      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.size}")).simplify()
 
     from tinygrad import function
     @function(precompile=True)

test/backend/test_dtype_alu.py

@@ -1,7 +1,7 @@
 import unittest, operator, math
 from tinygrad import Context, Tensor, dtypes, Device
 from tinygrad.dtype import DType, truncate, fp8_to_float
-from tinygrad.helpers import CI, EMULATED_DTYPES, DEV, getenv
+from tinygrad.helpers import CI, EMULATED_DTYPES, getenv
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
@@ -32,8 +32,7 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 #binary_operations.append(operator.truediv)
 
 # TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
-if ((DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"})
-    or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)):
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer):
   unary_operations.remove((Tensor.sin, np.sin))
   unary_operations.remove((Tensor.cos, np.cos))
 

test/backend/test_edgecases.py

@@ -27,10 +27,10 @@ import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
 from tinygrad.device import Device
-from tinygrad.helpers import DEV
+from tinygrad.helpers import getenv
 from tinygrad.renderer.nir import NIRRenderer
 
-MOCKGPU = DEV.interface.startswith("MOCK")
+MOCKGPU = getenv("MOCKGPU")
 
 class TestNaNEdgeCases(unittest.TestCase):
   # we don't need more of these. it's unclear if torch's behavior is desired here

test/backend/test_graph.py

@@ -2,12 +2,13 @@ import numpy as np
 import functools, unittest, ctypes
 
 from tinygrad.device import Device, Buffer
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import Context, from_mv
+from tinygrad.tensor import Tensor, _to_np_dtype
+from tinygrad.helpers import Context, dedup, from_mv
 from tinygrad.dtype import dtypes
 from tinygrad.engine.jit import MultiGraphRunner
-from tinygrad.schedule import linear_to_schedule
-from tinygrad.uop.ops import UOp, Ops, buffers
+from tinygrad.engine.realize import BufferXfer, get_runner, CompiledRunner
+from tinygrad.engine.schedule import ExecItem
+from tinygrad.uop.ops import UOp, Ops
 
 from test.helpers import needs_second_gpu
 
@@ -16,46 +17,77 @@ Tensor.manual_seed(1337)
 BUF_SIZE = 4096
 RUN_CNT = 5
 
-# cache AST by (device, num_inputs)
-cached_asts: dict[tuple[str, int], UOp] = {}
-def get_ast(device:str, num_inputs:int) -> UOp:
-  if (device, num_inputs) not in cached_asts:
+cached_prgs = {}
+def helper_exec_op(device, outbuf, inbufs):
+  if (device, len(inbufs)) not in cached_prgs:
     with Context(DEBUG=0):
-      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for _ in range(num_inputs)]
+      fst = [Tensor.randn(BUF_SIZE, dtype=dtypes.int).realize() for i in range(len(inbufs))]
       s = fst[0]
-      for i in range(1, num_inputs): s = s.bitwise_xor(fst[i])
-      cached_asts[(device, num_inputs)] = s.schedule()[-1].ast
-  return cached_asts[(device, num_inputs)]
+      for i in range(1, len(inbufs)): s = s.bitwise_xor(fst[i])
 
-def make_buffer(device, size=BUF_SIZE, fill=False):
-  buf = Buffer(device, size, dtypes.int).ensure_allocated()
+      si = s.schedule()[-1]
+      prg = get_runner(device, si.ast)
+    cached_prgs[(device, len(inbufs))] = prg
+
+  return ExecItem(UOp(Ops.NOOP), [outbuf] + inbufs, prg=cached_prgs[(device, len(inbufs))])
+
+def helper_copy_op(device, dest, src):
+  prg = BufferXfer(dest.nbytes, device, src.device)
+  return ExecItem(UOp(Ops.NOOP), [dest, src], prg=prg)
+
+def helper_alloc_rawbuffer(device, fill=False):
+  rawbuf = Buffer(device, BUF_SIZE, dtypes.int).ensure_allocated()
   if fill:
     with Context(DEBUG=0):
-      buf.copyin(Tensor(np.random.randint(-10000, 10000, size=size, dtype=np.int32)).realize().uop.base.realized.as_memoryview())
-  return buf
+      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+  return rawbuf
 
-def make_view(base, offset_elems, size_elems):
+def helper_create_offset_rawbuffer(base, offset=0):
+  x = Buffer(base.device, base.size-offset, base.dtype, base=base, offset=offset)
+  return x.ensure_allocated()
+
+def helper_alloc_rawbuffer_sized(device, size, fill=False):
+  rawbuf = Buffer(device, size, dtypes.int).ensure_allocated()
+  if fill:
+    with Context(DEBUG=0):
+      data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
+      rawbuf.copyin(Tensor(data).realize().uop.base.realized.as_memoryview())
+  return rawbuf
+
+def helper_make_view(base, offset_elems, size_elems):
   return Buffer(base.device, size_elems, base.dtype, base=base, offset=offset_elems * base.dtype.itemsize).ensure_allocated()
 
-def get_buf_uop(buf:Buffer, cache:dict[Buffer,UOp]) -> UOp:
-  if buf not in cache:
-    cache[buf] = u = UOp.new_buffer(buf.device, buf.size, buf.dtype)
-    buffers[u] = buf
-  return cache[buf]
-
-def make_graph(graph_cls, calls:list[UOp]):
-  linear = UOp(Ops.LINEAR, src=tuple(calls))
-  cf = UOp(Ops.CUSTOM_FUNCTION, dtypes.void, src=(linear,), arg="graph")
-  return graph_cls(cf, [])
-
-def run_schedule(calls:list[UOp]):
-  for ei in linear_to_schedule(UOp(Ops.LINEAR, src=tuple(calls))): ei.lower().run({})
-
-def zero_bufs(bufs):
-  for b in bufs:
-    mv = memoryview(bytearray(b.nbytes))
+def helper_run_jit(jis, bufs, out_buffers):
+  for rawbuf in out_buffers:
+    mv = memoryview(bytearray(rawbuf.size * rawbuf.dtype.itemsize))
     ctypes.memset(from_mv(mv), 0, len(mv))
-    b.copyin(mv)
+    rawbuf.copyin(mv)
+
+  for ei in jis: ei.run({}, jit=True)
+  return [rawbuf.as_memoryview() for rawbuf in bufs]
+
+def helper_test_graphs(graph_impl, graphs, runs=RUN_CNT):
+  reg_ji = []
+  bufs = []
+  out_buffers = set()
+  for graph in graphs:
+    for ji in graph:
+      out_buffers.update([ji.bufs[i] for i in (ji.prg.p.outs if isinstance(ji.prg, CompiledRunner) else [0])])
+      bufs += ji.bufs
+      reg_ji.append(ji)
+  bufs = dedup(bufs)
+
+  ground_thruth_bufs = helper_run_jit(reg_ji, bufs, out_buffers)
+  ground_truth_np = [np.frombuffer(x, _to_np_dtype(bufs[i].dtype)) for i,x in enumerate(ground_thruth_bufs)]
+
+  # Build graphs
+  gr_ji = [ExecItem(UOp(Ops.NOOP), [], prg=graph_impl(None, None, graph)) for graph in graphs]
+
+  for _ in range(runs):
+    test_bufs = helper_run_jit(gr_ji, bufs, out_buffers)
+    test_bufs_np = [np.frombuffer(x, _to_np_dtype(bufs[i].dtype)) for i,x in enumerate(test_bufs)]
+    for i in range(len(ground_thruth_bufs)): np.testing.assert_equal(ground_truth_np[i], test_bufs_np[i])
 
 @unittest.skipUnless(Device[Device.DEFAULT].graph is not None, "graph support required")
 class TestGraph(unittest.TestCase):
@@ -69,251 +101,236 @@ class TestGraph(unittest.TestCase):
 
   def test_order_2_writes_to_same_buf(self):
     d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(5)]
-    c: dict[Buffer,UOp] = {}
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
 
-    calls = [
-      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
+    graphs = [
+      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[0], [b0[3], b0[4]])]
     ]
 
-    zero_bufs([b[0]])
-    run_schedule(calls)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b[0]])
-      make_graph(Device[d0].graph, calls)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_order_read_write_same_buf(self):
     d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(5)]
-    c: dict[Buffer,UOp] = {}
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
 
-    calls = [
-      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[3],c), get_buf_uop(b[4],c), metadata=()),
+    graphs = [
+      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[1], [b0[3], b0[4]])]
     ]
 
-    zero_bufs([b[0], b[1]])
-    run_schedule(calls)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b[0], b[1]])
-      make_graph(Device[d0].graph, calls)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_order_write_read_same_buf(self):
     d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(5)]
-    c: dict[Buffer,UOp] = {}
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(5)]
 
-    calls = [
-      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), get_buf_uop(b[4],c), metadata=()),
+    graphs = [
+      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_exec_op(d0, b0[1], [b0[0], b0[4]])]
     ]
 
-    zero_bufs([b[0], b[1]])
-    run_schedule(calls)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b[0], b[1]])
-      make_graph(Device[d0].graph, calls)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_order_copy_writed(self):
     self.skip_if_not_multigraph()
-    d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(4)]
-    c: dict[Buffer,UOp] = {}
 
-    calls = [
-      get_ast(d0, 2).call(get_buf_uop(b[0],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
-      UOp(Ops.COPY).call(get_buf_uop(b[3],c), get_buf_uop(b[0],c), metadata=()),
+    d0 = Device.DEFAULT
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
+
+    graphs = [
+      [helper_exec_op(d0, b0[0], [b0[1], b0[2]]), helper_copy_op(d0, b0[3], b0[0])]
     ]
 
-    zero_bufs([b[0], b[3]])
-    run_schedule(calls)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b[0], b[3]])
-      make_graph(Device[d0].graph, calls)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_order_copy_then_read(self):
     self.skip_if_not_multigraph()
-    d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(4)]
-    c: dict[Buffer,UOp] = {}
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(b[1],c), get_buf_uop(b[0],c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=()),
+    d0 = Device.DEFAULT
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(4)]
+
+    graphs = [
+      [helper_copy_op(d0, b0[1], b0[0]), helper_exec_op(d0, b0[3], [b0[1], b0[2]])]
     ]
 
-    zero_bufs([b[1], b[3]])
-    run_schedule(calls)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b[1], b[3]])
-      make_graph(Device[d0].graph, calls)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_read_write_several_graphs(self):
     d0 = Device.DEFAULT
-    b = [make_buffer(d0, fill=True) for _ in range(8)]
-    c: dict[Buffer,UOp] = {}
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
 
-    calls1 = [get_ast(d0, 2).call(get_buf_uop(b[3],c), get_buf_uop(b[1],c), get_buf_uop(b[2],c), metadata=())]
-    calls2 = [get_ast(d0, 2).call(get_buf_uop(b[4],c), get_buf_uop(b[1],c), get_buf_uop(b[3],c), metadata=())]
-    calls3 = [get_ast(d0, 2).call(get_buf_uop(b[5],c), get_buf_uop(b[4],c), get_buf_uop(b[2],c), metadata=())]
+    graphs = [
+      [helper_exec_op(d0, b0[3], [b0[1], b0[2]])],
+      [helper_exec_op(d0, b0[4], [b0[1], b0[3]])],
+      [helper_exec_op(d0, b0[5], [b0[4], b0[2]])]
+    ]
 
-    out = [b[3], b[4], b[5]]
-    zero_bufs(out)
-    run_schedule(calls1 + calls2 + calls3)
-    expected = [np.frombuffer(x.as_memoryview(), np.int32).copy() for x in b]
+    helper_test_graphs(Device[d0].graph, graphs)
 
-    for _ in range(RUN_CNT):
-      zero_bufs(out)
-      make_graph(Device[d0].graph, calls1)([], {})
-      make_graph(Device[d0].graph, calls2)([], {})
-      make_graph(Device[d0].graph, calls3)([], {})
-      for i, buf in enumerate(b): np.testing.assert_equal(expected[i], np.frombuffer(buf.as_memoryview(), np.int32))
+    graphs = [
+      [helper_exec_op(d0, b0[3], [b0[1], b0[2]]), helper_exec_op(d0, b0[4], [b0[1], b0[2]]), helper_exec_op(d0, b0[5], [b0[1], b0[2]])],
+      [helper_exec_op(d0, b0[2], [b0[6], b0[7]])]
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
 
   @needs_second_gpu
   def test_copies_2_devs(self):
     self.skip_if_not_multigraph()
-    d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
-    b0 = [make_buffer(d0, fill=True) for _ in range(3)]
-    b1 = [make_buffer(d1, fill=True)]
-    c: dict[Buffer,UOp] = {}
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(b1[0],c), get_buf_uop(b0[0],c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b0[2],c), get_buf_uop(b0[0],c), get_buf_uop(b0[1],c), metadata=()),
+    d0, d1 = Device.DEFAULT, f"{Device.DEFAULT}:1"
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
+    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(1)]
+
+    graphs = [
+      [helper_copy_op(d0, b1[0], b0[0]), helper_exec_op(d0, b0[2], [b0[0], b0[1]])]
     ]
 
-    out = [b1[0], b0[2]]
-    zero_bufs(out)
-    run_schedule(calls)
-    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in b0 + b1}
+    helper_test_graphs(Device[d0].graph, graphs)
 
-    for _ in range(RUN_CNT):
-      zero_bufs(out)
-      make_graph(Device[d0].graph, calls)([], {})
-      for buf in b0 + b1: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+  @needs_second_gpu
+  def test_copies_after_graph_global(self):
+    self.skip_if_not_multigraph()
+
+    d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
+    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(6)]
+    b2 = [helper_alloc_rawbuffer(d2, fill=True) for _ in range(6)]
+    b3 = [helper_alloc_rawbuffer(d3, fill=True) for _ in range(6)]
+
+    graphs = [
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_exec_op(d0, b0[6], [b0[1], b0[2]]), helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
+      [helper_copy_op(d1, b0[2], b1[0])],
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_exec_op(d0, b0[6], [b0[1], b0[2]]), helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
+      [helper_copy_op(d3, b0[2], b3[0])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
+
+    graphs = [
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_copy_op(d0, b2[0], b0[2]), helper_copy_op(d0, b2[1], b0[5]),
+       helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
+      [helper_copy_op(d1, b0[2], b1[0])],
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]])],
+      [helper_copy_op(d3, b0[2], b3[0])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
+
+    graphs = [
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]]), helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]]), helper_copy_op(d0, b2[0], b0[2]), helper_copy_op(d0, b2[1], b0[5]),
+       helper_exec_op(d0, b0[7], [b0[0], b0[2]])],
+      [helper_copy_op(d1, b0[5], b1[0])],
+      [helper_copy_op(d3, b0[5], b3[0])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
+
+    graphs = [
+      [helper_copy_op(d1, b0[5], b1[0])],
+      [helper_copy_op(d3, b0[5], b3[0])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
+
+  @needs_second_gpu
+  def test_graph_after_copies_devs(self):
+    self.skip_if_not_multigraph()
+
+    d0, d1, d2, d3 = Device.DEFAULT, f"{Device.DEFAULT}:1", f"{Device.DEFAULT}:2", f"{Device.DEFAULT}:3"
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(8)]
+    b1 = [helper_alloc_rawbuffer(d1, fill=True) for _ in range(1)]
+    b2 = [helper_alloc_rawbuffer(d2, fill=True) for _ in range(2)]
+    b3 = [helper_alloc_rawbuffer(d3, fill=True) for _ in range(2)]
+
+    graphs = [
+      [helper_copy_op(d1, b0[0], b1[0])],
+      [helper_copy_op(d2, b0[1], b2[0]), helper_copy_op(d3, b0[2], b3[0])],
+      [helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
+
+    graphs = [
+      [helper_copy_op(d1, b0[0], b1[0])],
+      [helper_exec_op(d0, b0[2], [b0[0], b0[1]])],
+      [helper_copy_op(d2, b0[1], b2[0]), helper_copy_op(d3, b0[2], b3[0])],
+      [helper_exec_op(d0, b0[3], [b0[0], b0[2]]), helper_exec_op(d0, b0[4], [b0[3], b0[2]]),
+       helper_exec_op(d0, b0[5], [b0[0], b0[2]])],
+    ]
+
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_graph_offset_bufs(self):
     self.skip_if_not_multigraph()
+
     d0 = Device.DEFAULT
     if not hasattr(Device[d0].allocator, "_offset"): self.skipTest("device does not support _offset")
 
-    b0 = make_buffer(d0, fill=True)
-    b1 = make_view(b0, 0, b0.size)
-    b2 = make_view(b0, 0, b0.size)
-    c: dict[Buffer,UOp] = {}
+    b0 = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(1)]
+    b0 += [helper_create_offset_rawbuffer(b0[0]), helper_create_offset_rawbuffer(b0[0])]
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(b1,c), get_buf_uop(b0,c), get_buf_uop(b2,c), metadata=()),
+    graphs = [
+      [helper_copy_op(d0, b0[0], b0[2]), helper_exec_op(d0, b0[1], [b0[0], b0[2]])],
     ]
 
-    zero_bufs([b0])
-    run_schedule(calls)
-    expected = np.frombuffer(b0.as_memoryview(), np.int32).copy()
-
-    for _ in range(RUN_CNT):
-      zero_bufs([b0])
-      make_graph(Device[d0].graph, calls)([], {})
-      np.testing.assert_equal(expected, np.frombuffer(b0.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_partial_write_preserves_write_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
-    copy_src_full = make_buffer(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = make_buffer(d0, fill=True)
-    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
-    a, out = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
-    c: dict[Buffer,UOp] = {}
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    a, c = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
-      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(out,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
+    graphs = [
+      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, c, [v_hi, a])]
     ]
-
-    zero_bufs([base, out])
-    run_schedule(calls)
-    expected = {base: np.frombuffer(base.as_memoryview(), np.int32).copy(), out: np.frombuffer(out.as_memoryview(), np.int32).copy()}
-
-    for _ in range(RUN_CNT):
-      zero_bufs([base, out])
-      make_graph(Device[d0].graph, calls)([], {})
-      for buf in [base, out]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_partial_write_preserves_read_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = make_buffer(d0, BUF_SIZE * 2, fill=True)
-    copy_dst = make_buffer(d0, BUF_SIZE * 2, fill=True)
-    copy_src_lo = make_buffer(d0, fill=True)
-    v_lo, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE)
-    a, b = make_buffer(d0, fill=True), make_buffer(d0, fill=True)
-    c: dict[Buffer,UOp] = {}
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_dst = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 2, fill=True)
+    copy_src_lo = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    a, b = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(2)]
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(copy_dst,c), get_buf_uop(base,c), metadata=()),
-      UOp(Ops.COPY).call(get_buf_uop(v_lo,c), get_buf_uop(copy_src_lo,c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(v_hi,c), get_buf_uop(a,c), get_buf_uop(b,c), metadata=()),
+    graphs = [
+      [helper_copy_op(d0, copy_dst, base), helper_copy_op(d0, v_lo, copy_src_lo), helper_exec_op(d0, v_hi, [a, b])]
     ]
-
-    zero_bufs([copy_dst, base])
-    run_schedule(calls)
-    expected = {copy_dst: np.frombuffer(copy_dst.as_memoryview(), np.int32).copy(), base: np.frombuffer(base.as_memoryview(), np.int32).copy()}
-
-    for _ in range(RUN_CNT):
-      zero_bufs([copy_dst, base])
-      make_graph(Device[d0].graph, calls)([], {})
-      for buf in [copy_dst, base]: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
   def test_middle_write_splits_write_dep(self):
     self.skip_if_not_multigraph()
     self.skip_if_no_offset()
     d0 = Device.DEFAULT
 
-    base = make_buffer(d0, BUF_SIZE * 3, fill=True)
-    copy_src_full = make_buffer(d0, BUF_SIZE * 3, fill=True)
-    copy_src_mid = make_buffer(d0, fill=True)
-    v_lo, v_mid, v_hi = make_view(base, 0, BUF_SIZE), make_view(base, BUF_SIZE, BUF_SIZE), make_view(base, BUF_SIZE * 2, BUF_SIZE)
-    a, out1, out2 = make_buffer(d0, fill=True), make_buffer(d0, fill=True), make_buffer(d0, fill=True)
-    c: dict[Buffer,UOp] = {}
+    base = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
+    copy_src_full = helper_alloc_rawbuffer_sized(d0, BUF_SIZE * 3, fill=True)
+    copy_src_mid = helper_alloc_rawbuffer(d0, fill=True)
+    v_lo = helper_make_view(base, 0, BUF_SIZE)
+    v_mid = helper_make_view(base, BUF_SIZE, BUF_SIZE)
+    v_hi = helper_make_view(base, BUF_SIZE * 2, BUF_SIZE)
+    a, c, e = [helper_alloc_rawbuffer(d0, fill=True) for _ in range(3)]
 
-    calls = [
-      UOp(Ops.COPY).call(get_buf_uop(base,c), get_buf_uop(copy_src_full,c), metadata=()),
-      UOp(Ops.COPY).call(get_buf_uop(v_mid,c), get_buf_uop(copy_src_mid,c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(out1,c), get_buf_uop(v_lo,c), get_buf_uop(a,c), metadata=()),
-      get_ast(d0, 2).call(get_buf_uop(out2,c), get_buf_uop(v_hi,c), get_buf_uop(a,c), metadata=()),
+    graphs = [
+      [helper_copy_op(d0, base, copy_src_full), helper_copy_op(d0, v_mid, copy_src_mid),
+       helper_exec_op(d0, c, [v_lo, a]), helper_exec_op(d0, e, [v_hi, a])]
     ]
-
-    outs = [base, out1, out2]
-    zero_bufs(outs)
-    run_schedule(calls)
-    expected = {buf: np.frombuffer(buf.as_memoryview(), np.int32).copy() for buf in outs}
-
-    for _ in range(RUN_CNT):
-      zero_bufs(outs)
-      make_graph(Device[d0].graph, calls)([], {})
-      for buf in outs: np.testing.assert_equal(expected[buf], np.frombuffer(buf.as_memoryview(), np.int32))
+    helper_test_graphs(Device[d0].graph, graphs)
 
 if __name__ == '__main__':
   unittest.main()

test/backend/test_interop.py

@@ -3,12 +3,12 @@ import unittest
 import torch
 import numpy as np
 
-from tinygrad.helpers import CI, DEV
+from tinygrad.helpers import getenv, CI
 from tinygrad.tensor import Tensor
 from tinygrad.device import Device
 from tinygrad.dtype import _from_torch_dtype, _to_torch_dtype
 
-MOCKGPU = DEV.interface.startswith("MOCK")
+MOCKGPU = getenv("MOCKGPU")
 
 @unittest.skipIf(Device.DEFAULT not in ["METAL", "CUDA"] or MOCKGPU, f"no support on {Device.DEFAULT}")
 class TestInterop(unittest.TestCase):

test/backend/test_jit.py

@@ -1,15 +1,15 @@
 #!/usr/bin/env python
-import unittest
+import unittest, functools
 import numpy as np
 
 from hypothesis import given, settings, strategies as strat
-from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu
+from test.helpers import assert_jit_cache_len, not_support_multi_device, needs_second_gpu
 from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit, JitError, graph_class
+from tinygrad.engine.jit import TinyJit, JitError, GraphRunner, MultiGraphRunner, graph_class
+from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
 from tinygrad.device import Device
-from tinygrad.helpers import Context, JIT, DEV, GlobalCounters
+from tinygrad.helpers import Context, JIT, GlobalCounters, getenv
 from tinygrad.dtype import dtypes
-from tinygrad.uop.ops import Ops
 from extra.models.unet import ResBlock
 
 def _simple_test(add, extract=lambda x: x, N=10):
@@ -419,10 +419,10 @@ class TestJit(unittest.TestCase):
       if prev is not None: np.testing.assert_allclose(o, prev, atol=1e-4, rtol=1e-5)
       prev = o
 
+    graph_t = Device[Device.DEFAULT].graph.func if isinstance(Device[Device.DEFAULT].graph, functools.partial) else Device[Device.DEFAULT].graph
     # Checking that 2 graphs are inited.
-    assert len(jf.captured.linear.src) == 2
-    for si in jf.captured.linear.src:
-      assert call_is_graph(si)
+    assert isinstance(jf.jit_cache[0].prg, graph_t)
+    assert isinstance(jf.jit_cache[1].prg, graph_t)
 
   def test_jitted_clone(self):
     def f(a): return a.clone().realize()
@@ -583,7 +583,7 @@ class TestJitPrune(unittest.TestCase):
       a = Tensor.rand(16).realize()
       out = w2_prune(a)
       np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
-    assert_jit_cache_len(w2_prune, 1)
+    assert len(w2_prune.captured.jit_cache) == 1
 
   def test_prune_w_copy_correct(self):
     weights = Tensor.rand(16).realize()
@@ -617,7 +617,7 @@ class TestJitPrune(unittest.TestCase):
       out = w2_prune(a)
       np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
 
-    assert_jit_cache_len(w2_prune, 1)
+    assert len(w2_prune.captured.jit_cache) == 1, "prune should have removed the copy"
 
 class TestJitFree(unittest.TestCase):
   def test_free_intermediates(self):
@@ -688,9 +688,8 @@ class TestJitGraphSplit(unittest.TestCase):
     graph_t = graph_class(dev)
     if graph_t is None: return
 
-    got = f.captured.linear.src
+    got = f.jit_cache
     from tinygrad.runtime.graph.hcq import HCQGraph
-    from tinygrad.engine.jit import MultiGraphRunner
     if graph_t is HCQGraph:
       validate = hcqgraph
     elif issubclass(graph_t, MultiGraphRunner):
@@ -699,16 +698,16 @@ class TestJitGraphSplit(unittest.TestCase):
       validate = graph
 
     assert len(got) == len(validate), f"Expected {len(validate)} operations, got {len(got)}"
-    for expected, si in zip(validate, got):
-      ast = si.src[0]
+    for expected, got in zip(validate, got):
       if expected["type"] == "graph":
-        assert call_is_graph(si), f"Expected graph, got {ast.op}"
-        inner_cnt = len(ast.src[0].src)
-        assert inner_cnt == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {inner_cnt}"
+        assert isinstance(got.prg, GraphRunner), f"Expected GraphRunner, got {type(got.prg)}"
+        assert len(got.prg.jit_cache) == expected["cnt"], f"Expected {expected['cnt']} operations in graph, got {len(got.prg.jit_cache)}"
       elif expected["type"] == "comp":
-        assert ast.op in (Ops.SINK, Ops.PROGRAM, Ops.BEAM), f"Expected kernel, got {ast.op}"
-      elif expected["type"] in ("copy", "xfer"):
-        assert ast.op is Ops.COPY, f"Expected COPY, got {ast.op}"
+        assert isinstance(got.prg, CompiledRunner), f"Expected CompiledRunner, got {type(got.prg)}"
+      elif expected["type"] == "copy":
+        assert isinstance(got.prg, BufferCopy), f"Expected BufferCopy, got {type(got.prg)}"
+      elif expected["type"] == "xfer":
+        assert isinstance(got.prg, BufferXfer), f"Expected BufferXfer, got {type(got.prg)}"
 
   def ji_graph(self, cnt): return {"type": "graph", "cnt": cnt}
   def ji_comp(self): return {"type": "comp"}
@@ -813,7 +812,7 @@ class TestJitGraphSplit(unittest.TestCase):
       hcqgraph=[self.ji_graph(6)])
 
   @unittest.skip("this fails if you don't have SDMA or are using AMD_DISABLE_SDMA=1")
-  @unittest.skipIf(DEV.interface.startswith("MOCK"), "MockGPU does not support parallel copies")
+  @unittest.skipIf(getenv("MOCKGPU"), "MockGPU does not support parallel copies")
   def test_jit_multidev_copy(self):
     if Device.DEFAULT in {"CPU"}: raise unittest.SkipTest("CPU/LLVM is not a valid default device for this test (zero-copies)")
 

test/backend/test_linearizer.py

@@ -7,12 +7,11 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, CompiledRunner, get_program
-from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, DEV
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, getenv
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer
-from test.helpers import replace_opts
-MOCKGPU = DEV.interface.startswith("MOCK")
+MOCKGPU = getenv("MOCKGPU")
 
 from tinygrad.uop.ops import print_uops # noqa: F401 # pylint: disable=unused-import
 
@@ -45,7 +44,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int).cast(dtypes.char) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 1)
 
   @unittest.expectedFailure
@@ -53,7 +52,7 @@ class TestLinearizer(unittest.TestCase):
     tst = Tensor.ones(16, dtype=dtypes.int).contiguous().realize()
     out = tst.neg().cast(dtypes.char).cast(dtypes.int) * 2
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     self.assertEqual(len([x for x in uops if x.op is Ops.CAST]), 0)
 
   @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "broken on ptx")
@@ -63,7 +62,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.empty(16)
     out = img.conv2d(w, b)
     ast = helper_linearizer_opt(out)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     # slice at the last loop end
     uslice = [i for i,u in enumerate(uops) if u.op == Ops.END][-1]
     # only valid test if outermost range is the reduce
@@ -84,7 +83,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)).sum()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -92,7 +91,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor.randn(2, ).realize()
     out = a.reshape(2, 1).expand(2, 3).expand(2, 2, 3).sum()
     ast = helper_linearizer_opt(out, wanna_output=[np.broadcast_to(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)), (2, 2, 3)).sum()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -100,7 +99,7 @@ class TestLinearizer(unittest.TestCase):
     a = Tensor([2, 2]).realize()
     out = a.reshape(2, 1).pad(((1, 1), (1, 1)), value=2).sum()
     ast = helper_linearizer_opt(out, wanna_output=[24])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
     assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
@@ -113,7 +112,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(a.numpy()+b.numpy()[0]).sum()+b.numpy()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     # LOAD -> RANGE -> LOAD -> STORE
     assert len([x for x in uops[:ranges[0]] if x.op is Ops.LOAD]) == 1
@@ -123,7 +122,7 @@ class TestLinearizer(unittest.TestCase):
     b = Tensor.randn(1, 1).realize()
     out = (a.reshape(2, 1).expand(2, 3) + b[0]).sum() + b[0]
     ast = helper_linearizer_opt(out, wanna_output=[(np.broadcast_to(a.numpy().reshape(2, 1), (2, 3)) + b.numpy()[0]).sum() + b.numpy()])
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
     assert len(ranges) == 1 # NOTE: it collapses now
 
@@ -134,7 +133,7 @@ class TestLinearizer(unittest.TestCase):
     # these are of size 3 to avoid float4 coalesce
     r = a[:-1] + a[1:]
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_loads = len([uop for uop in uops if uop.op is Ops.LOAD])
     assert num_loads <= 4, "more load uops than needed"
     assert num_loads >= 4, "unexpected number of uops, maybe this test needs updating?"
@@ -146,7 +145,7 @@ class TestLinearizer(unittest.TestCase):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = a.expand([2]) + b.expand([2])
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops <= 1, "more alu uops than needed"
 
@@ -155,8 +154,8 @@ class TestLinearizer(unittest.TestCase):
     x, w = Tensor.randn((1,1,3)).realize(), Tensor.randn((1,1,2)).realize()
     r = Tensor.conv2d(x,w,padding=1).relu()
 
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]),
-                       renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer,
+                       opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0), Opt(op=OptOps.UNROLL, axis=0, arg=0)]).uops
     accs = [u for u in uops if u.op is Ops.DEFINE_REG]
     stores = [u for u in uops if u.op is Ops.STORE]
     assert len(accs) == 0  # it's removed now
@@ -168,7 +167,7 @@ class TestLinearizer(unittest.TestCase):
   @unittest.skipUnless(Device.DEFAULT == "CPU", "test only for CPU")
   def test_upcast_with_locals_cpu(self):
     out = Tensor.ones(64,64).contiguous() @ Tensor.ones(64,64).contiguous()
-    prg = get_program(replace_opts(out.schedule()[-1].ast, [Opt(OptOps.LOCAL, axis=0, arg=4)]), renderer=Device[Device.DEFAULT].renderer).uops
+    prg = get_program(out.schedule()[-1].ast, opts=[Opt(OptOps.LOCAL, axis=0, arg=4)]).uops
     self.assertEqual(len(prg.src.split("for")), 5)
 
   @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "test requires locals")
@@ -179,7 +178,7 @@ class TestLinearizer(unittest.TestCase):
     x, y = Tensor.rand(1,128), Tensor.rand(128, 128)
     r = (x@y).relu()
     opts_to_apply = [Opt(op=OptOps.GROUP, axis=0, arg=8), Opt(op=OptOps.LOCAL, axis=0, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=4)]
-    program = get_program(replace_opts(r.schedule()[-1].ast, opts_to_apply), renderer=Device[Device.DEFAULT].renderer)
+    program = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=opts_to_apply)
 
     stores = [u for u in program.uops if u.op is Ops.STORE and u.src[0].dtype.addrspace != AddrSpace.REG]
 
@@ -193,7 +192,7 @@ class TestLinearizer(unittest.TestCase):
   def test_zero_fold(self):
     a, b = Tensor.randn(1).realize(), Tensor.randn(1).realize()
     r = Tensor.stack(a, b)
-    uops = get_program(replace_opts(r.schedule()[-1].ast, [Opt(op=OptOps.UPCAST, axis=0, arg=0)]), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(r.schedule()[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=[Opt(op=OptOps.UPCAST, axis=0, arg=0)]).uops
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
@@ -203,14 +202,14 @@ class TestLinearizer(unittest.TestCase):
       if is_dtype_supported(tensor_dtype) and is_dtype_supported(acc_dtype):
         a = Tensor([1, 2, 3], dtype=tensor_dtype).sum()
         realized_ast = a.schedule()[-1].ast
-        program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+        program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
         local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
         assert local[0].dtype.base == acc_dtype
 
   def test_arg_acc_dtype(self):
     def helper_arg_acc_dtype(c: Tensor, expected_dtype:DType):
       realized_ast = c.schedule()[-1].ast
-      program = get_program(replace_opts(realized_ast, []), renderer=Device[Device.DEFAULT].renderer)
+      program = get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=[])
       local = [uop for uop in program.uops if uop.op is Ops.DEFINE_REG]
       self.assertEqual(local[0].dtype.base, expected_dtype)
 
@@ -238,7 +237,7 @@ class TestLinearizer(unittest.TestCase):
     opt = [Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4)]
     ast = helper_linearizer_opt(r, [opt])
     # the uops graph is DEFINE_REG -> 4x STORE 0.0 -> RANGE -> 4x ALU -> 4x STORE -> ENDRANGE
-    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
     begin_range = [i for i, x in enumerate(uops) if x.op is Ops.RANGE][-1]
     end_range = [i for i, x in enumerate(uops) if x.op is Ops.END][0]
     for i,u in enumerate(uops): print(i, u.op, [uops.index(s) for s in u.src], u.arg, u.dtype)
@@ -258,7 +257,7 @@ class TestLinearizer(unittest.TestCase):
     # shrink so that the dims do not collapse
     t = Tensor.ones(5, 6, 7).contiguous().realize().shrink(((0, 4), (0, 5), (0, 6)))
     ast = helper_linearizer_opt(t+1)
-    uops = get_program(replace_opts(ast, []), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=[]).uops
     idxs = dedup([uop for uop in uops if uop.op is Ops.SPECIAL])
     idxs = sorted(idxs, key=lambda uop: uop.arg)
     assert (idxs[0].arg, idxs[0].src[0].arg) == ('gidx0', 6), idxs[0]
@@ -291,7 +290,7 @@ class TestLinearizer(unittest.TestCase):
     sched_copy = sched[:]
     run_schedule(sched)
     np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])
-    program = get_program(replace_opts(sched_copy[-1].ast, []), renderer=Device[Device.DEFAULT].renderer)
+    program = get_program(sched_copy[-1].ast, renderer=Device[Device.DEFAULT].renderer, opts=())
     assert not any(u.op == Ops.WHERE for u in program.uops), "found where where where should be folded"
 
   def test_phi_simplification(self):
@@ -353,7 +352,7 @@ class TestLinearizer(unittest.TestCase):
             Opt(OptOps.UNROLL, 0, 4), Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UPCAST, 1, 2)] # upcast accs in both reduces
     ast = helper_linearizer_opt(out, opts=[opt])
     def get_recursive(uop): return set.union(set(uop.src), [uop], *[get_recursive(v) for v in uop.src])
-    uops = get_program(replace_opts(ast, opt), renderer=Device[Device.DEFAULT].renderer).uops
+    uops = get_program(ast, renderer=Device[Device.DEFAULT].renderer, opts=opt).uops
     local_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.DEFINE_LOCAL for x in get_recursive(u.src[0]))]
     global_stores = [u for u in uops if u.op is Ops.STORE and any(x.op is Ops.PARAM for x in get_recursive(u.src[0]))]
     barrier = [u for u in uops if u.op is Ops.BARRIER]
@@ -421,9 +420,7 @@ def _helper_linearizer_opt_ast(realized_ast:UOp, real_bufs:list[Buffer], opts=[]
   device = real_bufs[0].device
   wanna_output = [np.array(x).flatten() for x in wanna_output]
 
-  def get_prg(opts):
-    ast = realized_ast if opts is None else replace_opts(realized_ast, list(opts))
-    return CompiledRunner(replace(get_program(ast, renderer=Device[Device.DEFAULT].renderer), device=device))
+  def get_prg(opts): return CompiledRunner(replace(get_program(realized_ast, renderer=Device[Device.DEFAULT].renderer, opts=opts), device=device))
 
   def check_opt(opts):
     prg = get_prg(opts=opts)

test/backend/test_multitensor.py

@@ -1,13 +1,13 @@
-import unittest, random
+import unittest, functools, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import getenv, prod, Context
 from tinygrad.nn.state import get_parameters, get_state_dict
-from tinygrad.engine.realize import CompiledRunner, run_schedule
+from tinygrad.engine.realize import BufferCopy, CompiledRunner, run_schedule
 import numpy as np
 from hypothesis import given, strategies as strat, settings
-from test.helpers import not_support_multi_device, needs_second_gpu, slow, call_is_graph
+from test.helpers import not_support_multi_device, needs_second_gpu, slow
 
 settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
 settings.load_profile("my_profile")
@@ -305,8 +305,7 @@ class TestMultiTensor(unittest.TestCase):
     Xs = X.shard(device, shard_x)
     Ws = W.shard(device, shard_w)
     O = (Xs@Ws)
-    with np.errstate(all='ignore'):
-      np.testing.assert_allclose(X.numpy() @ W.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
+    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()
@@ -316,8 +315,7 @@ class TestMultiTensor(unittest.TestCase):
     W1s = W1.shard(device, shard_w)
     W2s = W2.shard(device, shard_w)
     O = (Xs@W1s)@W2s
-    with np.errstate(all='ignore'):
-      np.testing.assert_allclose((X.numpy() @ W1.numpy()) @ W2.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
+    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)
@@ -544,7 +542,7 @@ class TestMultiTensor(unittest.TestCase):
       b.shard_(devices_2)
       c = jf(a, b)
       np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
-    assert jf.captured is not None
+    assert len(jf.jit_cache) > 0
 
   def test_multi_tensor_jit_body(self):
     @TinyJit
@@ -558,7 +556,7 @@ class TestMultiTensor(unittest.TestCase):
     for _ in range(5):
       r = jf()
       np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
-    assert jf.captured is not None
+    assert len(jf.jit_cache) > 0
 
   def test_multitensor_jit_in_list(self):
     # test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
@@ -618,12 +616,15 @@ class TestMultiTensor(unittest.TestCase):
       o = jf(a, b, c, d).numpy()
       np.testing.assert_allclose(ref, o, atol=1e-4, rtol=1e-5)
 
+    graph_d0 = Device[d0].graph.func if isinstance(Device[d0].graph, functools.partial) else Device[d0].graph
+    graph_d1 = Device[d1].graph.func if isinstance(Device[d1].graph, functools.partial) else Device[d1].graph
     # Checking that 2 graphs per device, 1 copy and 1 last graph on device 1 are created.
-    sis = jf.captured.linear.src
-    assert len(sis) == 6
-    for si in (sis[0], sis[1], sis[2], sis[3], sis[5]):
-      assert call_is_graph(si)
-    assert sis[4].src[0].op is Ops.COPY
+    assert isinstance(jf.jit_cache[0].prg, graph_d0)
+    assert isinstance(jf.jit_cache[1].prg, graph_d0)
+    assert isinstance(jf.jit_cache[2].prg, graph_d1)
+    assert isinstance(jf.jit_cache[3].prg, graph_d1)
+    assert isinstance(jf.jit_cache[4].prg, BufferCopy)
+    assert isinstance(jf.jit_cache[5].prg, graph_d1)
 
   def test_bn_ast_on_devices(self):
     t = Tensor.empty((16, 64, 112, 112)).shard(devices_4, axis=0)

test/backend/test_ops.py

@@ -94,7 +94,7 @@ def prepare_test_op(low, high, shps, vals, forward_only=False):
 class TestOps(unittest.TestCase):
 
   def helper_test_exception(self, shps, torch_fxn, tinygrad_fxn=None, expected=None, forward_only=False, exact=False, vals=None, low=-1.5, high=1.5):
-    if DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV": self.skipTest('helper_test_exception fails in CI CUDA')
+    if getenv("MOCKGPU") and Device.DEFAULT == "NV": self.skipTest('helper_test_exception fails in CI CUDA')
     ts, tst = prepare_test_op(low, high, shps, vals, forward_only)
     if tinygrad_fxn is None:
       tinygrad_fxn = torch_fxn
@@ -281,17 +281,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: torch.arange(5.5, 175.5, 2.5), lambda: Tensor.arange(5.5, 175.5, 2.5), forward_only=True)
     helper_test_op([], lambda: torch.arange(-30.2, -0.3, 0.75), lambda: Tensor.arange(-30.2, -0.3, 0.75), forward_only=True)
     helper_test_op([], lambda: torch.arange(-50.3, -380.2, -2.25), lambda: Tensor.arange(-50.3, -380.2, -2.25), forward_only=True)
-    # boundary values that fit exactly in int8 (min=-128, max=127)
-    helper_test_op([], lambda: torch.arange(128, dtype=torch.int8), lambda: Tensor.arange(128, dtype=dtypes.int8), forward_only=True)
-    helper_test_op([], lambda: torch.arange(-128, 128, dtype=torch.int8), lambda: Tensor.arange(-128, 128, dtype=dtypes.int8), forward_only=True)
-    helper_test_op([], lambda: torch.arange(127, -129, -1, dtype=torch.int8),
-                   lambda: Tensor.arange(127, -129, -1, dtype=dtypes.int8), forward_only=True)
-    # overflow: tinygrad raises (torch silently wraps)
-    with self.assertRaises(OverflowError): Tensor.arange(2**33, dtype=dtypes.int)
-    with self.assertRaises(OverflowError): Tensor.arange(129, dtype=dtypes.int8)            # last=128 overflows
-    with self.assertRaises(OverflowError): Tensor.arange(-129, 128, dtype=dtypes.int8)      # start=-129 overflows
-    with self.assertRaises(OverflowError): Tensor.arange(128, 0, -1, dtype=dtypes.int8)     # start=128 overflows
-    with self.assertRaises(OverflowError): Tensor.arange(127, -130, -1, dtype=dtypes.int8)  # last=-129 overflows
 
   def test_arange_big(self):
     helper_test_op([], lambda: torch.arange(256, dtype=torch.int32), lambda: Tensor.arange(256), forward_only=True)
@@ -877,7 +866,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.sin())
     helper_test_op([()], lambda x: x.sin())
     # works on real CUDA but not CI
-    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.sin(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -886,7 +875,7 @@ class TestOps(unittest.TestCase):
   def test_cos(self):
     helper_test_op([(45,65)], lambda x: x.cos())
     helper_test_op([()], lambda x: x.cos())
-    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.cos(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -897,7 +886,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)
     helper_test_op([(45,65)], lambda x: x.tan(), low=-5, high=5)
     helper_test_op([()], lambda x: x.tan())
-    if not ((DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
+    if not ((getenv("MOCKGPU") and Device.DEFAULT == "NV") or Device.DEFAULT == "WEBGPU"):
       helper_test_op(None, lambda x: x.tan(), vals=[[math.nan, math.inf, -math.inf, 0.0]])
       helper_test_op(None, lambda x: x.tan(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                     atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
@@ -3310,7 +3299,7 @@ class TestOps(unittest.TestCase):
     helper_test_op([(32,10)], lambda x: x.masked_fill((x>0.1).detach(), -math.inf))
     helper_test_op([(32,10)], lambda x: x.masked_fill((x<0.1).detach(), -math.inf))
 
-  @unittest.skipIf((DEV.interface.startswith("MOCK") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
+  @unittest.skipIf((getenv("MOCKGPU") or Device.DEFAULT == "PYTHON"), "very slow on MOCKGPU because reduce does not fold")
   @unittest.skipIf(Device.DEFAULT == "WEBGPU", "webgpu runtime issue")
   @unittest.skipIf(Device.DEFAULT == "QCOM", "QCOM fails with: Resource deadlock avoided")
   def test_masked_select(self):

test/backend/test_opt_gemm.py

@@ -4,8 +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.schedule import ExecItem
-from test.helpers import replace_opts
+from tinygrad.engine.schedule import ExecItem
 
 class TestOptGemm(unittest.TestCase):
   @classmethod
@@ -13,14 +12,13 @@ class TestOptGemm(unittest.TestCase):
     N = 64
     cls.a = Tensor.randn(N, N).contiguous().realize()
     cls.b = Tensor.randn(N, N).contiguous().realize()
-    with np.errstate(all='ignore'):
-      cls.res = cls.a.T.numpy() @ cls.b.T.numpy()
+    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
     # TODO: this should be a generic test helper
     si = get_single_element(t.schedule())
-    run = CompiledRunner(get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer))
+    run = CompiledRunner(get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts))
     ExecItem(si.ast, list(si.bufs), prg=run).run()
     test = si.bufs[0].numpy().reshape(self.res.shape)
     np.testing.assert_allclose(self.res, test, atol=1e-4)

test/backend/test_pickle.py

@@ -125,13 +125,6 @@ class TestPickle(unittest.TestCase):
     out = add_fxn(x, y)
     np.testing.assert_equal(out.numpy(), 102)
 
-  def test_pickle_jit_no_del(self):
-    @TinyJit
-    def fn(x): return x + 1.0
-    for _ in range(3): fn(Tensor.randn(4))
-    loaded = pickle.loads(pickle.dumps(fn))
-    self.assertEqual(loaded(Tensor([1.0,2.0,3.0,4.0])).tolist(), [2.0,3.0,4.0,5.0])
-
   def test_pickle_context_var(self):
     v = ContextVar("test_var", 0)
     with Context(test_var=1):

test/backend/test_profiler.py

@@ -1,11 +1,11 @@
 import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
-from tinygrad.helpers import CI, DEV, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
+from tinygrad.helpers import CI, getenv, Context, ProfileRangeEvent, cpu_profile, cpu_events, ProfilePointEvent, dedup
 from tinygrad.device import Buffer, BufferSpec, Compiled, ProfileDeviceEvent, ProfileGraphEvent
 from tinygrad.runtime.support.hcq import HCQCompiled
 from tinygrad.engine.realize import get_runner
 
-MOCKGPU = DEV.interface.startswith("MOCK")
+MOCKGPU = getenv("MOCKGPU")
 def _dev_base(d):
   p = d.split(":")
   return p[0] if len(p) < 2 or not p[1].isdigit() else f"{p[0]}:{p[1]}"

test/backend/test_quantize_onnx.py

@@ -6,8 +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.schedule import ExecItem
-from test.helpers import replace_opts
+from tinygrad.engine.schedule import ExecItem
 
 N = 512
 
@@ -40,7 +39,7 @@ def create_gemm_model(model_path:str, batch_size=N, in_size=N, out_size=N, bias=
 
 def sexec(out:Tensor, opts:list[Opt], replace_src=None, run_count=3):
   si = out.schedule()[-1]
-  prg = get_program(replace_opts(si.ast, opts), renderer=Device[Device.DEFAULT].renderer)
+  prg = get_program(si.ast, renderer=Device[Device.DEFAULT].renderer, opts=opts)
   if replace_src is not None:
     old_name = prg.src.split("__attribute__((noinline)) void ")[1].split("(")[0]
     prg = replace(prg, src=replace_src + "/* DSP boilerplate */" + prg.src.split("/* DSP boilerplate */")[1].replace(old_name, "fxn"))

test/backend/test_randomness.py

@@ -360,7 +360,6 @@ class TestRandomness(unittest.TestCase):
     torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(1000)]
     self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
 
-  @unittest.skip("this test is flaky")
   def test_multinomial_counterexample(self):
     tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
     torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)

test/backend/test_schedule.py

@@ -12,20 +12,18 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType
 from tinygrad.uop.ops import UOp, Ops, UPat
 from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
-from tinygrad.engine.realize import CompiledRunner, run_schedule, run_linear
-from tinygrad.schedule import linear_to_schedule
+from tinygrad.engine.realize import CompiledRunner, run_schedule
 
 class KernelCountException(Exception): pass
 def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Tensor]|None=None, filter_sink=True):
   if to_prerealize:
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): Tensor.realize(*to_prerealize)
-  if isinstance(t, Tensor): linear, var_vals = t.linear_with_vars()
-  elif isinstance(t, list) and isinstance(t[0], Tensor): linear, var_vals = Tensor.linear_with_vars(*t)
+  if isinstance(t, Tensor): sched = t.schedule()
+  elif isinstance(t, list) and isinstance(t[0], Tensor): sched = Tensor.schedule(*t)
   else:
     assert isinstance(t, UOp), f"can't schedule {t}"
-    linear, var_vals = Tensor(t).linear_with_vars()
+    sched = Tensor(t).schedule()
   # test lowering all the ExecItems
-  sched = linear_to_schedule(linear)
   for si in sched: si.lower()
   kernel_cnt = len([si for si in sched if isinstance(si.prg, CompiledRunner) or not filter_sink])
   if kernel_cnt != allowed:
@@ -35,7 +33,7 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
         print("kernel", i+1)
         print(s.ast)
     raise KernelCountException(f"{kernel_cnt} != {allowed}")
-  return linear, var_vals
+  return sched
 
 def _realize_weights(m):
   for p in nn.state.get_parameters(m): p.realize()
@@ -93,17 +91,17 @@ class TestSchedule(unittest.TestCase):
   # all permutes, reshapes, expands and shrinks push through the reduce
   def test_arange_sum(self):
     a = Tensor.arange(6).reshape(3, 2).sum(axis=1)
-    run_linear(*check_schedule(a, 1))
+    run_schedule(check_schedule(a, 1))
     self.assertListEqual(a.tolist(), [1, 5, 9])
 
   def test_arange_sum_alt(self):
     a = (Tensor.arange(5).reshape(1,5).expand(6,5)*Tensor(2)).reshape(1,6,5).sum(axis=2)
-    run_linear(*check_schedule(a, 1))
+    run_schedule(check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), 20)
 
   def test_permute_arange(self):
     a = Tensor.arange(6).reshape(6, 1, 1).permute(2, 0, 1).sum(axis=1)
-    run_linear(*check_schedule(a, 1))
+    run_schedule(check_schedule(a, 1))
     self.assertListEqual(a.tolist(), [[15]])
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
@@ -111,13 +109,13 @@ class TestSchedule(unittest.TestCase):
   def test_expand_buffer_before_cast(self):
     a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
     b = a.cast(dtypes.half).expand((2, 4, 4))+2
-    run_linear(*check_schedule(b, 1))
+    run_schedule(check_schedule(b, 1))
     np.testing.assert_allclose(b.numpy(), np.broadcast_to(a.numpy().astype(np.float16), (2, 4, 4))+2, rtol=1e-3)
 
   def test_indexing_scalars_simple(self):
     X = Tensor.randn(2, 2).realize()
     xt = X[Tensor(1)][Tensor(0)]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])
 
   @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -137,46 +135,46 @@ class TestSchedule(unittest.TestCase):
     assume(a<x and b<y)
     X = Tensor.randn(x, y).realize()
     xt = X[Tensor(a)][Tensor(b)]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])
 
   def test_push_pads_elementwise(self):
     x = Tensor.full((4,4), 2.).contiguous().realize()
     y = Tensor.full((4,4), 4.).contiguous().realize()
     z = (x.reciprocal()*y).pad((None, (0,1),)).sum()
-    run_linear(*check_schedule(z, 1))
+    run_schedule(check_schedule(z, 1))
     self.assertEqual(z.item(), 32)
 
   def test_push_pads_contiguous(self):
     x = Tensor.full((4,1), 2.).contiguous()
     y = Tensor.full((4,4), 4.).contiguous()
     z = (x.reciprocal().expand(4,4)*y).pad((None, (0,1),)).sum()
-    run_linear(*check_schedule(z, 1, [x,y]))
+    run_schedule(check_schedule(z, 1, [x,y]))
     self.assertEqual(z.item(), 32)
 
   def test_constants_can_store(self):
     a = Tensor(2).contiguous()
-    run_linear(*check_schedule(a, 1))
+    run_schedule(check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), 2)
 
   def test_allow_push_permutes(self):
     a = Tensor.randn(10,10,10).realize()
     b = Tensor.randn(10,10,1).realize()
     c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    run_linear(*check_schedule(c, 1))
+    run_schedule(check_schedule(c, 1))
     np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy())
 
   def test_div_collapse_buffer(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     b = Tensor.full((4,), 2.0).contiguous().realize()
     expr = (a*b)/b
-    run_linear(*check_schedule(expr, 1))
+    run_schedule(check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 4.0))
 
   def test_div_collapse_const(self):
     a = Tensor.full((4,), 4.0).contiguous().realize()
     expr = a/a
-    run_linear(*check_schedule(expr, 1))
+    run_schedule(check_schedule(expr, 1))
     np.testing.assert_allclose(expr.numpy(), np.full((4,), 1.0))
 
   def test_div_collapse(self):
@@ -197,7 +195,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones(4).contiguous()
     b = Tensor.ones(4).contiguous()
     sched = check_schedule([a, b], 1)
-    run_linear(*sched)
+    run_schedule(sched)
     # a and b share the same underlying device memory
     self.assertIs(a.uop.realized, b.uop.realized)
 
@@ -206,7 +204,7 @@ class TestSchedule(unittest.TestCase):
     a = src.clone()
     b = src.clone()
     sched = check_schedule([a, b], 2, filter_sink=False)
-    run_linear(*sched)
+    run_schedule(sched)
     # a and b are assigned to the same device Buffer
     self.assertIsNot(a.uop.base.realized, b.uop.base.realized)
 
@@ -228,7 +226,7 @@ class TestSchedule(unittest.TestCase):
     ax = r.reshape(1)*2
     ay = r.reshape(1).shrink(((1,1),))*2
     out = ax+ay.pad(((1, 0),))
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     self.assertEqual(out.item(), 4.)
 
   def test_preserve_multistage_reduce(self):
@@ -236,7 +234,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(big_enough).realize()
     with Context(SPLIT_REDUCEOP=1):
       out = (x - x.max(keepdim=True)).max()
-      run_linear(*check_schedule(out, 4))
+      run_schedule(check_schedule(out, 4))
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(keepdims=True)).max())
 
   @unittest.skip("these two Tensors are the same")
@@ -246,7 +244,7 @@ class TestSchedule(unittest.TestCase):
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
 
   def test_example_matmul_contig(self):
@@ -255,7 +253,7 @@ class TestSchedule(unittest.TestCase):
     z = y.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.ones((64,64)))
 
   def test_example_matmul_same(self):
@@ -263,7 +261,7 @@ class TestSchedule(unittest.TestCase):
     z = x.matmul(x).sum()
     z.backward()
     out = x.grad.contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     # NOTE: the gradient flows twice
     np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))
 
@@ -277,7 +275,7 @@ class TestSchedule(unittest.TestCase):
     b_out = b.sum(1)
     b_out = b_out[:16]
     out = a_out + b_out + c
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4)
 
   def test_reduce_same_size(self):
@@ -286,7 +284,7 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum() + 2
     out1 = a.sum() + 4
     out2 = out0 * out1
-    run_linear(*check_schedule([out0, out1, out2], 3)) # TODO: 1?
+    run_schedule(check_schedule([out0, out1, out2], 3)) # TODO: 1?
     np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
@@ -297,7 +295,7 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum().exp2()
     # out1 has two paths to a.sum()
     out1 = a.sum() + out0
-    run_linear(*check_schedule([out0, out1], 2)) # TODO: 1?
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)
 
@@ -309,14 +307,13 @@ class TestSchedule(unittest.TestCase):
     b = (a + out0 + out1)
     out2 = b.sum().exp2()
     out3 = b.sum() + out2
-    # run_linear(*check_schedule([out0, out1, out2, out3], 1))
-    run_linear(*check_schedule([out0, out1, out2, out3], 4))
+    # run_schedule(check_schedule([out0, out1, out2, out3], 1))
+    run_schedule(check_schedule([out0, out1, out2, out3], 4))
     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)
-    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)
+    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)
@@ -325,8 +322,8 @@ class TestSchedule(unittest.TestCase):
     # b.sum() is not a descendant of the fused nodes
     out0 = a.sum() + b.sum() + 2
     out1 = a.sum() + b.sum() + 4
-    # run_linear(*check_schedule([out0, out1], 1))
-    run_linear(*check_schedule([out0, out1], 2))
+    # run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2))
     np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+b.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy().sum()+4, atol=1e-4, rtol=1e-4)
 
@@ -338,8 +335,8 @@ class TestSchedule(unittest.TestCase):
     # reduce node in the indirect path from r to out2
     out1 = (a - out0).max()
     out2 = r + out1
-    # run_linear(*check_schedule([r, out0, out1, out2], 1))
-    run_linear(*check_schedule([r, out0, out1, out2], 4))
+    # run_schedule(check_schedule([r, out0, out1, out2], 1))
+    run_schedule(check_schedule([r, out0, out1, out2], 4))
     np.testing.assert_allclose(r.numpy(), r_np:=a.numpy().sum(), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(r_np), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), out1_np:=(a.numpy() - out0_np).max(), atol=1e-4, rtol=1e-4)
@@ -352,8 +349,8 @@ class TestSchedule(unittest.TestCase):
     out0 = a.sum() + 4
     out1 = b.max() + out0*2
     out2 = a.sum() + out1
-    # run_linear(*check_schedule([out0, out1, out2], 1))
-    run_linear(*check_schedule([out0, out1, out2], 3))
+    # run_schedule(check_schedule([out0, out1, out2], 1))
+    run_schedule(check_schedule([out0, out1, out2], 3))
     np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out1.numpy(), out1_np:=b.numpy().max() + out0_np*2, atol=1e-4, rtol=1e-6)
     np.testing.assert_allclose(out2.numpy(), a.numpy().sum() + out1_np, atol=1e-4, rtol=1e-6)
@@ -367,8 +364,8 @@ class TestSchedule(unittest.TestCase):
     # e1 is in the indirect path from a.sum() to out1
     e = b + out0
     out1 = r + e[0][0][0]
-    # run_linear(*check_schedule([r, out0, out1, e], 3)) # 1 or 2 or 3? should be 1 (one reduce) but the different outputs might make it 3
-    run_linear(*check_schedule([r, out0, out1, e], 4))
+    # run_schedule(check_schedule([r, out0, out1, e], 3)) # 1 or 2 or 3? should be 1 (one reduce) but the different outputs might make it 3
+    run_schedule(check_schedule([r, out0, out1, e], 4))
     np.testing.assert_allclose(r.numpy(), r_np:=a.numpy().sum(), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(r_np), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=b.numpy() + out0_np, atol=1e-4, rtol=1e-4)
@@ -380,8 +377,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn((1, 16)).realize()
     out0 = a.sum() + 2
     out1 = a.sum() + b
-    # run_linear(*check_schedule([out0, out1], 2))
-    run_linear(*check_schedule([out0, out1], 3))
+    # run_schedule(check_schedule([out0, out1], 2))
+    run_schedule(check_schedule([out0, out1], 3))
     np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy(), atol=1e-4, rtol=1e-4)
 
@@ -389,7 +386,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.std(-1)
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_scaled_dot_product_attention_multireduce_fusion(self):
@@ -398,14 +395,14 @@ class TestSchedule(unittest.TestCase):
     k = Tensor.randn(32,8,16,8).realize()
     v = Tensor.randn(32,8,16,8).realize()
     out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_linear(*check_schedule(out, 4))
+    run_schedule(check_schedule(out, 4))
     if getenv("CHECK", 1):
       import torch
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
       np.testing.assert_allclose(out.numpy(), compare.numpy(), atol=1e-6, rtol=1e-3)
 
     out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_linear(*check_schedule(out, 4)) # TODO: should be 1?
+    run_schedule(check_schedule(out, 4)) # TODO: should be 1?
     if getenv("CHECK", 1):
       import torch
       compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -417,8 +414,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(4, 32).realize()
     c = Tensor.randn(4, 32).realize()
     out = (c * a.sum(-1, keepdim=True)).sum(-1) + (b * a.sum(-1, keepdim=True)).sum(-1) # a.sum has >1 children but should still fuse
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), \
       (c.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1) + (b.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1), atol=1e-4, rtol=1e-4)
 
@@ -426,16 +423,16 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()
     out = (a+a.sum(-1, keepdim=True)).sum(-1)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_reduce_expand_reduce_expand_fusion(self):
     Tensor.manual_seed(0)
     a = Tensor.randn(4, 32).realize()
     out = a+(a+a.sum(-1,keepdim=True)).sum(-1, keepdim=True)
-    # run_linear(*check_schedule(out, 2))
-    run_linear(*check_schedule(out, 3))
+    # run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 3))
     np.testing.assert_allclose(out.numpy(), \
       a.numpy()+(a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1,keepdims=True), atol=1e-4, rtol=1e-4)
 
@@ -444,8 +441,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(4, 32).realize()
     out0 = a+a.sum(-1, keepdim=True)
     out1 = out0.sum(-1)
-    # run_linear(*check_schedule(out, 2))
-    run_linear(*check_schedule([out0, out1], 3))
+    # run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule([out0, out1], 3))
     np.testing.assert_allclose(out0.numpy(), a.numpy()+a.numpy().sum(axis=-1,keepdims=True), atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(out1.numpy(), (a.numpy()+a.numpy().sum(axis=-1,keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
@@ -454,8 +451,8 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = (y + x.sum(axis=-1, keepdim=True)).sum(axis=-1)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (y.numpy() + x.numpy().sum(axis=-1, keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_simple_parallel(self):
@@ -463,15 +460,15 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = y.sum(axis=-1) + x.sum(axis=-1)
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), y.numpy().sum(axis=-1) + x.numpy().sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_sequential(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.std(-1)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_parallel(self):
@@ -479,16 +476,16 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = x.std(-1) + y.std(-1)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 3))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 3))
     np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1) + y.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_diffops_sequential(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = (x - x.max(-1, keepdim=True)).sum(-1)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).sum(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_diffops_parallel(self):
@@ -496,7 +493,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(4, 32).realize()
     y = Tensor.randn(4, 32).realize()
     out = x.sum(-1) + y.max(-1)
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), x.numpy().sum(axis=-1) + y.numpy().max(axis=-1), atol=1e-4, rtol=1e-4)
 
   def test_multireduce_fusion_sequential_and_parallel(self):
@@ -507,13 +504,13 @@ class TestSchedule(unittest.TestCase):
     out = [((x - mu).square().sum(-1)/x.shape[-1]).sqrt(), ((y - mu).square().sum(-1)/y.shape[-1]).sqrt()]
     np_mu = (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True) + \
       (y.numpy() - y.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True)
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 5))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 5))
     np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
     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 (same nesting context = should merge)
+    # two-stage cumsum + ops triggers parallel REDUCEs in one kernel that must share an END
     step, num_steps = 513, 10
     t = Tensor.arange(step).float().realize()
     phase = t.cumsum()
@@ -524,25 +521,19 @@ 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()
     c,d = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
     out = a@b + c@d
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), a.numpy()@b.numpy() + c.numpy()@d.numpy(), atol=1e-4, rtol=1e-4)
 
   def test_softmax_fusion(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 12, 64, 64).realize()
     out = x.softmax()
-    run_linear(*check_schedule(out, 3))
+    run_schedule(check_schedule(out, 3))
     expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
     np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)
 
@@ -553,8 +544,8 @@ class TestSchedule(unittest.TestCase):
     layer.bias = Tensor.randn(10,10).realize()
     x = Tensor.randn(20, 5, 10, 10).realize()
     out = layer(x)
-    # run_linear(*check_schedule(out, 2))
-    run_linear(*check_schedule(out, 3))
+    # run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 3))
     y = (x.numpy() - x.numpy().mean(layer.axis, keepdims=True))
     expected = y / np.sqrt((y*y).mean(layer.axis, keepdims=True) + layer.eps)
     np.testing.assert_allclose(out.numpy(), expected * layer.weight.numpy() + layer.bias.numpy(), atol=1e-4, rtol=1e-4)
@@ -569,7 +560,7 @@ class TestSchedule(unittest.TestCase):
     # schedule = check_schedule([b,c], 3)
     # self.assertIs(schedule[0].ast[0].src[0].arg, Ops.MUL)
     schedule = check_schedule([b,c], 4)
-    run_linear(*schedule)
+    run_schedule(schedule)
     np.testing.assert_allclose(b.numpy(), np_r.sum(0) + 8, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(c.numpy(), np_r.sum(1) + 12, atol=1e-4, rtol=1e-4)
 
@@ -581,7 +572,7 @@ class TestSchedule(unittest.TestCase):
     d = r.T * 4
     e = r * (d + a).sum(2)
     schedule = check_schedule([d, e], 3) # make sure it doesn't fuse
-    run_linear(*schedule)
+    run_schedule(schedule)
     np.testing.assert_allclose(d.numpy(), (a.numpy().sum(2) + b.numpy()).T * 4, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), (a.numpy().sum(2) + b.numpy()) * (d.numpy() + a.numpy()).sum(2), atol=1e-4, rtol=1e-4)
 
@@ -594,7 +585,7 @@ class TestSchedule(unittest.TestCase):
     r = a.sum(1) + c
     out = r[:4] * b + d.sum(1)[:4]
     schedule = check_schedule(out, 1)
-    run_linear(*schedule)
+    run_schedule(schedule)
     np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)
 
   def test_multireduce_midreduce_nochase(self):
@@ -602,7 +593,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(16, 16).realize()
     b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
     schedule = check_schedule(b, 1)
-    run_linear(*schedule)
+    run_schedule(schedule)
     np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)
 
   # pattern in test_transformer
@@ -612,8 +603,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(16, 16).realize()
     c = a.sum() + 2
     d = (a.sum() - b.sum()) * 4
-    # run_linear(*check_schedule([c, d], 1))
-    run_linear(*check_schedule([c, d], 2))
+    # run_schedule(check_schedule([c, d], 1))
+    run_schedule(check_schedule([c, d], 2))
     np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)
 
@@ -624,8 +615,8 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.randn(16, 16).realize()
     c = a.sum() + 2
     d = b.sum() - c
-    # run_linear(*check_schedule([c, d], 1))
-    run_linear(*check_schedule([c, d], 2))
+    # run_schedule(check_schedule([c, d], 1))
+    run_schedule(check_schedule([c, d], 2))
     np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), b.numpy().sum()-(a.numpy().sum()+2), atol=1e-4, rtol=1e-4)
 
@@ -638,8 +629,8 @@ class TestSchedule(unittest.TestCase):
     d = a.sum() * 2
     e = c * d
     f = b.sum() - e
-    # run_linear(*check_schedule([c, d, e, f], 1))
-    run_linear(*check_schedule([c, d, e, f], 4))
+    # run_schedule(check_schedule([c, d, e, f], 1))
+    run_schedule(check_schedule([c, d, e, f], 4))
     np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -653,8 +644,8 @@ class TestSchedule(unittest.TestCase):
     d = a.sum() * 2
     e = c * d
     f = (b - d).sum() - e
-    # run_linear(*check_schedule([c, d, e, f], 1))
-    run_linear(*check_schedule([c, d, e, f], 4))
+    # run_schedule(check_schedule([c, d, e, f], 1))
+    run_schedule(check_schedule([c, d, e, f], 4))
     np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
     np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -665,7 +656,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.rand(3, 4, 5).realize()
     b = Tensor.rand(3, 4, 5).realize()
     out = (a + b).pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(a.numpy()+b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)
 
   def test_multireduce_pad_reduce_safe(self):
@@ -673,7 +664,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(3, 4, 5).realize()
     b = Tensor.randn(3, 4, 5).realize()
     out = (a.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum(keepdim=True)+b.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()).contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(a.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(keepdims=True) + \
                                                    np.pad(b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-4, rtol=1e-4)
 
@@ -681,7 +672,7 @@ class TestSchedule(unittest.TestCase):
     Tensor.manual_seed(0)
     a = Tensor.rand(3, 4, 5).realize()
     out = a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)
 
   def test_multireduce_pad_reduce_unsafe(self):
@@ -689,8 +680,8 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.randn(3, 4, 5).abs().realize()
     b = Tensor.randn(3, 4, 5).abs().realize()
     out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    # run_linear(*check_schedule(out, 1))
-    run_linear(*check_schedule(out, 2))
+    # run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
                                                    b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)
 
@@ -698,13 +689,13 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones((3, )).contiguous().realize()
     b = Tensor.ones((3, )).contiguous().realize()
     out = (a + b).shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), [2, 0])
 
   def test_shrink_pad_unsafe(self):
     a = Tensor.ones((3, )).contiguous().realize()
     out = a.exp2().shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), [2, 0])
 
   def test_base_change_shrink_pad(self):
@@ -712,7 +703,7 @@ class TestSchedule(unittest.TestCase):
     b = a.exp2()
     c = b[:-1, :-1]
     d = c.pad(((0, 1), (0, 1))) * 2
-    run_linear(*check_schedule(d, 1))
+    run_schedule(check_schedule(d, 1))
     np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:-1, :-1], ((0, 1), (0, 1)))*2)
 
   def test_base_change_expand_pad(self):
@@ -720,7 +711,7 @@ class TestSchedule(unittest.TestCase):
     b = a.exp2()
     c = b[:, None, :]
     d = c.pad(((0, 0), (1, 1), (0, 0))) * 2
-    run_linear(*check_schedule(d, 1))
+    run_schedule(check_schedule(d, 1))
     np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)
 
   def test_fuse_arange_pad_replicate_mode(self):
@@ -728,7 +719,7 @@ class TestSchedule(unittest.TestCase):
     y = x.pad((-1,2,2,-1), mode="replicate")
     dx = y.sum().gradient(x)[0]
     sched = check_schedule(dx, 1)
-    run_linear(*sched)
+    run_schedule(sched)
     np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)
 
   # TODO like openpilot with imagef
@@ -737,7 +728,7 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     b = a.cast(dtypes.half).expand(2, 4, 4)
     c = b.cast(dtypes.int).expand(2, 2, 4, 4)
-    run_linear(*check_schedule(c, 1))
+    run_schedule(check_schedule(c, 1))
     np.testing.assert_equal(c.numpy(), np.ones(((2, 2, 4, 4)), dtype=np.int32))
 
   def test_base_change_pad_expand(self):
@@ -745,7 +736,7 @@ class TestSchedule(unittest.TestCase):
     b = Tensor.full((4, 4), 2.).contiguous().realize()
     c = (a + b).pad(((1, 1), (1, 1)))
     d = c.cast(dtypes.int).expand((2, 6, 6)) * 4
-    run_linear(*check_schedule(d, 1))
+    run_schedule(check_schedule(d, 1))
     c_np = np.pad((np.full((4, 4), 2., dtype=np.float32) + np.full((4, 4), 1., dtype=np.float32)), ((1, 1), (1, 1)), constant_values=0.0)
     np.testing.assert_equal(d.numpy(), np.broadcast_to(c_np.astype(np.half), (2, *c_np.shape)) * 4)
 
@@ -756,7 +747,7 @@ class TestSchedule(unittest.TestCase):
     p = P[0]
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
-    run_linear(*check_schedule(p, 3))
+    run_schedule(check_schedule(p, 4))  # TODO: this is high
     tiny_ret = p.numpy()
 
     P = np.ones((3, 3), dtype=np.float32)
@@ -838,7 +829,7 @@ class TestSchedule(unittest.TestCase):
 
   def _test_fusion(self, shapes, f, cnt):
     with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
-    run_linear(*check_schedule(compare:=f(*args), cnt))
+    run_schedule(check_schedule(compare:=f(*args), cnt))
     if getenv("COMPARE", 1):
       import torch
       good = f(*[torch.tensor(x.numpy()) for x in args])
@@ -882,9 +873,9 @@ class TestSchedule(unittest.TestCase):
   def test_cast_const_view(self):
     a = Tensor.ones((4, 4), dtype=dtypes.float32)
     casted_view = a.cast(dtypes.int32)
-    run_linear(*check_schedule(casted_view, 1))
+    run_schedule(check_schedule(casted_view, 1))
     realized_const_view = casted_view.contiguous()
-    run_linear(*check_schedule(realized_const_view, 0))
+    run_schedule(check_schedule(realized_const_view, 0))
     self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])
 
   @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
@@ -893,41 +884,41 @@ class TestSchedule(unittest.TestCase):
     assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
     a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
     casted_view = a.cast(dt2)
-    run_linear(*check_schedule(casted_view, 0))
+    run_schedule(check_schedule(casted_view, 0))
     realized_const_view = casted_view.contiguous()
-    run_linear(*check_schedule(realized_const_view, 1))
+    run_schedule(check_schedule(realized_const_view, 1))
     np.testing.assert_equal(realized_const_view.numpy(), [[0], [1], [0]])
 
   def test_simple_indexing(self):
     X = Tensor.randn(10, 10).realize()
     idxs = Tensor([0, 2]).realize()
     xt = X[idxs]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])
 
   def test_simple_indexing_alt(self):
     X = Tensor.arange(16).reshape(4, 4)
     xt = X[[1, 2], [-1, 2]]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), (np.arange(16).reshape(4, 4))[[1, 2], [-1, 2]])
 
   def test_advanced_indexing(self):
     X = Tensor.arange(10)+1
     xt = X[[0, -1]]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), (np.arange(10)+1)[[0, -1]])
 
   def test_advanced_indexing_alt(self):
     X = Tensor.arange(6).reshape(3, 2)+1
     xt = X[[Tensor([2]), Tensor([1])]]
-    run_linear(*check_schedule(xt, 1))
+    run_schedule(check_schedule(xt, 1))
     np.testing.assert_equal(xt.numpy(), 6)
 
   def test_push_through_reshape(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(10, 20).realize()
     out = x.argmax(1)
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), np.argmax(x.numpy(), 1))
 
   def test_arange_push_through_expand(self):
@@ -935,35 +926,35 @@ class TestSchedule(unittest.TestCase):
     a = Tensor.arange(4,)
     b = Tensor.randn(4, 4).realize()
     out = (a+b).sum()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (np.arange(4)+b.numpy()).sum(), atol=1e-5)
 
   def test_argmin(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmin(-1)
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_equal(out.numpy(), x.numpy().argmin(axis=-1))
 
   def test_argmax(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 32).realize()
     out = x.argmax(-1)
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_equal(out.numpy(), x.numpy().argmax(axis=-1))
 
   def test_arange_transposed(self):
     Tensor.manual_seed(0)
     x = Tensor.randint(4, 1).realize()
     a = ((Tensor.arange(4,)*x).T).sum()
-    run_linear(*check_schedule(a, 1))
+    run_schedule(check_schedule(a, 1))
     np.testing.assert_equal(a.numpy(), (np.arange(4)*x.numpy()).T.sum())
 
   def test_div_padded_arange(self):
     x = Tensor.full((2,2), 16)
     y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
     out = y.sum(axis=1)
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     self.assertListEqual(out.tolist(), [0, 12, 4, 0])
 
   def test_arange_transposed_descendants(self):
@@ -972,7 +963,7 @@ class TestSchedule(unittest.TestCase):
     a = (Tensor.arange(4,)*x).T
     b = Tensor.randint(4, 4).realize()
     out = (a+b).sum()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_equal(out.numpy(), ((np.arange(4)*x.numpy()).T+b.numpy()).sum())
 
   def test_arange_index(self):
@@ -980,7 +971,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10)
     out = (x + a[2]).sum()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_arange_index_contiguous(self):
@@ -988,7 +979,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10).contiguous()
     out = (x + a[2]).sum()
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_arange_index_child(self):
@@ -996,7 +987,7 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = Tensor.arange(10)+1
     out = (x + a[2]).sum()
-    run_linear(*check_schedule(out, 1))
+    run_schedule(check_schedule(out, 1))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   def test_user_contiguous(self):
@@ -1004,13 +995,13 @@ class TestSchedule(unittest.TestCase):
     x = Tensor.randn(5, 2).realize()
     a = (Tensor.arange(10)+1).contiguous()
     out = (x + a[2]).sum()
-    run_linear(*check_schedule(out, 2))
+    run_schedule(check_schedule(out, 2))
     np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)
 
   @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
   def test_arange_view_op(self):
     a = Tensor.arange(12).reshape(4, 3).shrink(((1, 2), (1, 3))).contiguous()
-    sched = run_linear(*check_schedule(a, 1))
+    sched = run_schedule(check_schedule(a, 1))
     self.assertIs(sched[1].ast.op, Ops.BUFFER_VIEW)
     np.testing.assert_equal(a.numpy(), [[4, 5]])
 
@@ -1019,16 +1010,16 @@ class TestSchedule(unittest.TestCase):
     from extra.models.llama import precompute_freqs_cis
     args = {"dim":32, "end":2048, "theta":10000}
     fused = precompute_freqs_cis(**args)
-    run_linear(*check_schedule(fused, 1))
+    run_schedule(check_schedule(fused, 1))
     if getenv("CHECK", 1):
       ref = precompute_freqs_cis(**args)
-      run_linear(*check_schedule(ref, 1))
+      run_schedule(check_schedule(ref, 1))
       np.testing.assert_equal(fused.numpy(), ref.numpy())
 
   def test_fuse_assign_contiguous(self):
     x = Tensor.zeros(4, 4, dtype=dtypes.int).contiguous().realize()
     a = Tensor.arange(8).reshape(4, 2)
-    run_linear(*check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
+    run_schedule(check_schedule(x.shrink((None, (0, 2))).assign(a.contiguous()), 2))
     np.testing.assert_equal(x.numpy(), [[0, 1, 0, 0], [2, 3, 0, 0], [4, 5, 0, 0], [6, 7, 0, 0]])
 
   def test_assign_non_contiguous_alt(self): self.test_assign_non_contiguous(alt=True)
@@ -1077,8 +1068,8 @@ class TestSchedule(unittest.TestCase):
     idx = Tensor([1,2,5,6], dtype=dtypes.int32)
     flat_base[idx] = Tensor([99,99,99,99])
     base.assign(flat_base.reshape(4, 4))
-    sched = check_schedule(base, 4)
-    run_linear(*sched)
+    sched = check_schedule(base, 6)  # TODO: this is high
+    run_schedule(sched)
     expected = list(range(16))
     for i, v in zip([1,2,5,6], [99,99,99,99]): expected[i] = v
     np.testing.assert_equal(base.reshape(16).numpy(), expected)
@@ -1087,7 +1078,7 @@ class TestSchedule(unittest.TestCase):
     X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
     Y = Tensor([1, 2]).realize()
     loss = X.sparse_categorical_crossentropy(Y)
-    run_linear(*check_schedule(loss, 3))
+    run_schedule(check_schedule(loss, 3))
     np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)
 
   def test_const_folding_alt(self):
@@ -1107,7 +1098,7 @@ class TestSchedule(unittest.TestCase):
     samples = Tensor.randint(BS:=getenv("BS", 512), high=cast(int,Y_train.shape[-1])).realize()
     yt = Tensor.randn(BS, 10).realize()
     loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-    run_linear(*check_schedule(loss, 4))
+    run_schedule(check_schedule(loss, 4))
     loss_fused = loss.numpy()
     loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
     np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -1117,7 +1108,7 @@ class TestSchedule(unittest.TestCase):
     r = (X+Tensor.arange(16).reshape(4, 4)).sum()
     out0 = r+2
     out1 = r+3
-    run_linear(*check_schedule([out0, out1], 2)) # TODO: 1?
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
     r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
     np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
     np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
@@ -1132,7 +1123,7 @@ class TestSchedule(unittest.TestCase):
     for shape in [(3, 3), (4, 4)]:
       a = Tensor.ones(*shape).contiguous().realize()
       a.assign(a / 1)
-      run_linear(*check_schedule(a, 0, filter_sink=False))
+      run_schedule(check_schedule(a, 0, filter_sink=False))
       self.assertListEqual(a.tolist(), [[1.]*shape[1]]*shape[0])
 
 class TestLimitBufs(unittest.TestCase):
@@ -1166,7 +1157,7 @@ class TestSwizzle(unittest.TestCase):
       a = Tensor.randint(32, 32).realize()
     r = (a+a).sum(1).sum(0)
     # double reduce collapses to a single reduce
-    run_linear(*check_schedule(r, 1))
+    run_schedule(check_schedule(r, 1))
     self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))
 
   def test_single_swizzle(self):
@@ -1176,7 +1167,7 @@ class TestSwizzle(unittest.TestCase):
       b = Tensor.ones((1, 1), dtype=a.dtype).contiguous().realize()
     # ADD(REDUCE(RESHAPE(LOAD)), LOAD) to ADD(REDUCE(RESHAPE(LOAD))), RESHAPE(LOAD)
     r = a.sum(0)+b
-    run_linear(*check_schedule(r, 1))
+    run_schedule(check_schedule(r, 1))
     self.assertEqual(r.numpy(), a.numpy().sum(0)+1)
 
   def test_double_swizzle_possible(self):
@@ -1186,7 +1177,7 @@ class TestSwizzle(unittest.TestCase):
       b = Tensor.randint(4,).realize()
     # parallel reduce!
     add = a.sum(0)+b.sum(0)
-    run_linear(*check_schedule(add, 1))
+    run_schedule(check_schedule(add, 1))
     self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))
 
   def test_swizzle_reduceop(self):
@@ -1194,7 +1185,7 @@ class TestSwizzle(unittest.TestCase):
     x = Tensor.randn(4,4).realize()
     y = Tensor.randn(4,4,4).realize()
     out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
-    run_linear(*check_schedule(out, 2)) # TODO: 1?
+    run_schedule(check_schedule(out, 2)) # TODO: 1?
     np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())
 
   def test_permute_rewrite(self):
@@ -1202,7 +1193,7 @@ class TestSwizzle(unittest.TestCase):
     y = Tensor.randn(4, 1, 16).realize()
     z = Tensor.randn(4, 4, 1).realize()
     t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
-    run_linear(*check_schedule(t, 2)) # TODO: 1?
+    run_schedule(check_schedule(t, 2)) # TODO: 1?
     t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
     np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)
 
@@ -1213,14 +1204,14 @@ class TestSwizzle(unittest.TestCase):
     a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
     b_reduce = b.sum(axis=(0,))
     t = a_reduce+b_reduce
-    run_linear(*check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
 
   def test_parallel_reduce_possible(self):
     Tensor.manual_seed(0)
     x = Tensor.randn(4, 2, 2).realize()
     y = Tensor.randn(4, 2, 2).realize()
     t = x.sum(axis=1)+y.sum(axis=1)
-    run_linear(*check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
     np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
 
   # kernels can only have 1 or n in each dim
@@ -1229,14 +1220,14 @@ class TestSwizzle(unittest.TestCase):
     x = Tensor.randn(4, 2, 2).realize()
     y = Tensor.randn(4, 3, 2).realize()
     t = x.sum(axis=1)+y.sum(axis=1)
-    run_linear(*check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
     np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)
 
   def test_unsafe_pad(self):
     x = Tensor.full((2,2), 1.0).contiguous()
     y = x*x.sum((1,)).reciprocal()
     t = y.pad(((0,1),None))
-    run_linear(*check_schedule(t, 3))
+    run_schedule(check_schedule(t, 3))
     np.testing.assert_equal(t.numpy(), [[0.5, 0.5], [0.5, 0.5], [0., 0.]])
 
 zero_pm = UPat(Ops.CONST, arg=0)
@@ -1247,7 +1238,7 @@ class TestView(unittest.TestCase):
     # all masked out, degrades to const 0
     b = a.pad(((0, 10), None))[10:]
     sched = check_schedule(b.contiguous(), 1)
-    run_linear(*sched)
+    run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 
   def test_mask_dim_1(self):
@@ -1256,7 +1247,7 @@ class TestView(unittest.TestCase):
     b = a.pad((None, (0, 10)))[:, 10:]
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
-    run_linear(*sched)
+    run_schedule(sched)
     np.testing.assert_equal(b.numpy(), 0)
 
   def test_partial_mask(self):
@@ -1265,7 +1256,7 @@ class TestView(unittest.TestCase):
     b = a.pad(((0, 5), None))[5:]
     assert b.shape == (10, 10)
     sched = check_schedule(b.contiguous(), 1)
-    run_linear(*sched)
+    run_schedule(sched)
     np.testing.assert_allclose(b.numpy(), np.pad(a.numpy(), ((0, 5), (0, 0)))[5:])
 
   # a*VIEW(x), where VIEW(x) = 0
@@ -1276,7 +1267,7 @@ class TestView(unittest.TestCase):
     bv = b.pad(((0, 2),))[-2:]
     # this becomes a late a*0
     late_mul = a*bv
-    run_linear(*check_schedule(late_mul, 2))
+    run_schedule(check_schedule(late_mul, 2))
     # the arange doesn't realize
     #self.assertIsNone(b.uop.base.realized)
     # mul doesn't realize
@@ -1299,14 +1290,14 @@ class TestView(unittest.TestCase):
     # NOTE: no longer checked
     # mul still collapses
     #self.assertIs(late_mul.uop.base.op, Ops.CONST)
-    run_linear(*s)
+    run_schedule(s)
     self.assertEqual(other_child.tolist(), [2, 3, 4])
 
 @unittest.skipIf(Device.DEFAULT == "CPU", "tests copy from another device to cpu")
 class TestCopyFolding(unittest.TestCase):
   def test_const_copy_is_free(self):
     b = Tensor(1).to("CPU") * 4
-    run_linear(*check_schedule(b, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     assert b.item() == 4
 
   def test_one_hot_with_copy(self):
@@ -1316,14 +1307,14 @@ class TestCopyFolding(unittest.TestCase):
 
   def test_const_copy_multi(self):
     x = Tensor.ones(1, device="CPU").to_(["CPU", "CPU:1"]) * 2
-    run_linear(*check_schedule(x, 2, filter_sink=False))
+    run_schedule(check_schedule(x, 2, filter_sink=False))
     self.assertEqual(x.item(), 2.0)
 
   def test_late_const_copy_folding(self):
     a = Tensor.arange(3).realize()
     zeros = Tensor.zeros(3).realize()
     b = (a*zeros).to("CPU") + 1
-    run_linear(*check_schedule(b, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     self.assertListEqual(b.tolist(), [1, 1, 1])
     self.assertEqual(b.device, "CPU")
 
@@ -1363,7 +1354,7 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.ones(4, 4).contiguous().realize()
     # use copy_to_device to bypass Tensor.to() shortcircuit and force a real same-device COPY in the graph
     a.assign(Tensor(a.uop.copy_to_device(a.device), a.device))
-    run_linear(*check_schedule(a, 2, filter_sink=False))
+    run_schedule(check_schedule(a, 2, filter_sink=False))
     self.assertListEqual(a.tolist(), [[1.]*4]*4)
 
   def test_clone(self):
@@ -1374,18 +1365,18 @@ class TestCopyFolding(unittest.TestCase):
     a = Tensor.arange(4)
     view = a.shrink(((0, 2),))
     b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     self.assertEqual(b.uop.base.buffer.size, 2)
-    self.assertEqual(b.uop.numel(), 2)
+    self.assertEqual(b.uop.size, 2)
     self.assertListEqual(b.tolist(), [0, 1])
 
   def test_expanded_copy(self):
     a = Tensor.arange(2)
     view = a.reshape(2, 1).expand(2, 2)
     b = view.clone()
-    run_linear(*check_schedule(b, 1, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
     self.assertEqual(b.uop.base.buffer.size, 4)
-    self.assertEqual(b.uop.numel(), 4)
+    self.assertEqual(b.uop.size, 4)
     self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
 
   def test_permuted_copy(self):

test/backend/test_subbuffer.py

@@ -1,7 +1,7 @@
 import unittest
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.device import Buffer
-from tinygrad.helpers import Context, DEV
+from tinygrad.helpers import Context, getenv
 from test.helpers import needs_second_gpu
 
 @unittest.skipUnless(hasattr(Device[Device.DEFAULT].allocator, "_offset"), "subbuffer not supported")
@@ -42,7 +42,7 @@ class TestSubBuffer(unittest.TestCase):
     assert out == [102, 103]
 
   @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or DEV.interface.startswith("MOCK"), "only NV, AMD, CUDA")
+  @unittest.skipIf(Device.DEFAULT not in {"CUDA", "NV", "AMD"} or getenv("MOCKGPU"), "only NV, AMD, CUDA")
   def test_subbuffer_transfer(self):
     t = Tensor.arange(0, 10, dtype=dtypes.uint8).realize()
     vt = t[2:5].contiguous().realize()

test/backend/test_transcendental.py

@@ -14,7 +14,7 @@ settings.load_profile("my_profile")
 
 class TestTranscendentalMath(unittest.TestCase):
   @unittest.skipUnless(is_dtype_supported(dtypes.float64), f"no float64 on {Device.DEFAULT}")
-  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
   @given(ht.float64, strat.sampled_from([(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.sin)]))
   def test_float64(self, x, op):
     if op[0] == Tensor.sin:
@@ -25,7 +25,7 @@ class TestTranscendentalMath(unittest.TestCase):
                                  op[1](np.array([x], dtype=_to_np_dtype(dtypes.float64))),
                                  atol=3e-2, rtol=1e-5)  # sin can have bigger atol for very big x
 
-  @unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
+  @unittest.skipIf(getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}, "crashed")
   @given(ht.float32, strat.sampled_from([(Tensor.exp, np.exp),(Tensor.log, np.log)] +
     ([(Tensor.sin, np.sin)] if is_dtype_supported(dtypes.ulong) else [])))
   def test_float32(self, x, op):
@@ -66,7 +66,7 @@ class TestFromFuzzer(unittest.TestCase):
     if not is_dtype_supported(dtype): return
     if dtype == dtypes.float64:
       # crashes in CI CUDA
-      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
+      if getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}: return
     def _test_value(n: float, unit: float=1.0):
       next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
       ulp = next_float - 1.0
@@ -88,7 +88,7 @@ class TestFromFuzzer(unittest.TestCase):
     if not is_dtype_supported(dtype): return
     if dtype == dtypes.float64:
       # crashes in CI CUDA
-      if DEV.interface.startswith("MOCK") and Device.DEFAULT in {"NV", "CUDA"}: return
+      if getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}: return
     def _test_value(n: float, unit: float=1.0):
       next_float = np.nextafter(1.0, 2.0, dtype=_to_np_dtype(dtype))
       ulp = next_float - 1.0

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.schedule import ExecItem
+from tinygrad.engine.schedule import ExecItem
 from tinygrad.device import is_dtype_supported
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer

test/device/test_hcq.py

@@ -1,7 +1,7 @@
 import unittest, ctypes, struct, os, random, numpy as np, time
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import mv_address, DEBUG, DEV
-from test.helpers import slow, replace_opts
+from tinygrad.helpers import getenv, mv_address, DEBUG, DEV
+from test.helpers import slow
 from tinygrad.device import Buffer, BufferSpec
 from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
 from tinygrad.runtime.autogen import libc
@@ -10,7 +10,7 @@ from tinygrad.engine.realize import get_runner, CompiledRunner, get_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad import Variable
 
-MOCKGPU = DEV.interface.startswith("MOCK")
+MOCKGPU = getenv("MOCKGPU")
 
 @unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run")
 class TestHCQ(unittest.TestCase):
@@ -76,7 +76,7 @@ class TestHCQ(unittest.TestCase):
         TestHCQ.d0.timeline_signal.wait(TestHCQ.d0.timeline_value)
         TestHCQ.d0.timeline_value += 1
 
-  @unittest.skipIf(Device.DEFAULT == "CPU" or (DEV.interface == "MOCKPCI" and DEV.device == "AMD"), "Can't handle async update on CPU/MOCKPCI device")
+  @unittest.skipIf(Device.DEFAULT in {"CPU"} or (DEV.interface == "PCI" and DEV.device == "AMD"), "Can't handle async update on CPU/MOCKAM device")
   def test_wait_late_set(self):
     for queue_type in [TestHCQ.d0.hw_compute_queue_t, TestHCQ.d0.hw_copy_queue_t]:
       if queue_type is None: continue
@@ -165,7 +165,7 @@ class TestHCQ(unittest.TestCase):
     b = a + 1
     si = b.schedule()[-1]
 
-    runner = CompiledRunner(get_program(replace_opts(si.ast, [Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]), TestHCQ.d0.renderer))
+    runner = CompiledRunner(get_program(si.ast, TestHCQ.d0.renderer, opts=[Opt(op=OptOps.LOCAL, axis=0, arg=3) for _ in range(3)]))
 
     zb = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
     zt = Buffer(Device.DEFAULT, 3 * 3 * 3, dtypes.int, options=BufferSpec(cpu_access=True, nolru=True)).ensure_allocated()
@@ -575,7 +575,7 @@ class TestHCQ(unittest.TestCase):
 
       np.testing.assert_equal(cpu_buffer.numpy(), local_buf.numpy(), "failed")
 
-  @unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "MOCKPCI"), "Emulate this on MOCKGPU to check the path in CI")
+  @unittest.skipUnless(MOCKGPU and not (DEV.device == "AMD" and DEV.interface == "PCI"), "Emulate this on MOCKGPU to check the path in CI")
   def test_on_device_hang(self):
     if not hasattr(self.d0, 'on_device_hang'): self.skipTest("device does not have on_device_hang")
 

test/external/external_llm_eval.py

@@ -1,4 +1,4 @@
-# eval for OpenAI API server
+# eval for tinygrad.apps.llm -- hits the server via OpenAI API
 # uses Meta's exact ARC-Challenge prompt template from lm-evaluation-harness llama3 tasks
 import argparse, re, pyarrow.parquet as pq
 from openai import OpenAI
@@ -8,7 +8,7 @@ LABEL = ["A", "B", "C", "D"]
 
 if __name__ == "__main__":
   parser = argparse.ArgumentParser()
-  parser.add_argument("--port", "-p", type=int, default=8000)
+  parser.add_argument("--port", "-p", type=int, default=11434)
   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_am_fault_recovery.py

@@ -1,6 +1,5 @@
 # ruff: noqa: F405
 import unittest, subprocess, os
-from tinygrad.helpers import DEV
 from tinygrad.runtime.autogen.amd.rdna3.ins import *  # noqa: F403
 from tinygrad.renderer.amd.dsl import s, v, Inst, NULL
 
@@ -28,7 +27,7 @@ _ILLEGAL_INST_ASM = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:
   ".rodata\n.p2align 6\n.amdhsa_kernel test\n.amdhsa_next_free_vgpr 8\n.amdhsa_next_free_sgpr 8\n" \
   ".amdhsa_wavefront_size32 1\n.amdhsa_user_sgpr_kernarg_segment_ptr 1\n.amdhsa_kernarg_size 8\n.end_amdhsa_kernel"
 
-@unittest.skipIf(DEV.device != "AMD" or not DEV.interface.startswith("MOCK"), "AMD with AM driver required")
+@unittest.skipIf(os.environ.get("AMD") != "1" or os.environ.get("MOCKGPU") == "1", "AMD with AM driver required")
 class TestAMFaultRecovery(unittest.TestCase):
   def _run_kernel(self, insts: list[Inst]) -> subprocess.CompletedProcess: return _run_asm(assemble_kernel(insts))