Merge branch 'master' into clone_up_front

* no tags on contiguous

* fixes

* fix test

* simple fix, and new test

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

@@ -230,9 +230,12 @@ runs:
         sudo chown -R $USER:$USER /var/cache/apt/archives/
 
     - name: Add clang to PATH (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
+      if: runner.os == 'Linux'
       shell: bash
-      run: echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
+      run: |
+        if [ -d /usr/lib/llvm-20/bin ]; then
+          echo "/usr/lib/llvm-20/bin" >> "$GITHUB_PATH"
+        fi
 
     # **** AMD ****
     - name: Setup AMD (Linux)

.github/workflows/benchmark.yml

@@ -97,7 +97,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -137,7 +137,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -185,7 +185,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -227,7 +227,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -272,7 +272,7 @@ jobs:
         shell: bash -e -o pipefail {0}
     env:
       DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
     if: github.repository_owner == 'tinygrad'
     steps:
     - name: Checkout Code
@@ -319,7 +319,7 @@ jobs:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 10
+    timeout-minutes: 11
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -436,13 +436,7 @@ jobs:
     - name: UsbGPU tiny tests
       run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
     - name: UsbGPU copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    - name: UsbGPU (USB4/TB) install script
-      run: sh extra/setup_tinygpu_osx.sh
-    - name: UsbGPU (USB4/TB) boot time
-      run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py
 
   testcomma:
     strategy:
@@ -542,7 +536,7 @@ jobs:
   testcommausbgpubenchmark:
     name: UsbGPU Benchmark (comma)
     runs-on: [self-hosted, Linux, comma4]
-    timeout-minutes: 10
+    timeout-minutes: 14
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -562,7 +556,7 @@ jobs:
     - name: openpilot run_pickle big_driving_supercombo
       run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
     - name: Test copy speeds
-      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py
 
   driverbenchmarks:
     name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -629,16 +623,6 @@ jobs:
     - name: Run 10 MLPerf Bert training steps (1 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
       run: BENCHMARK_LOG=bert_10steps CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 python3 test/test_tiny.py
-        if [[ "${{ matrix.dev }}" == "AMD" ]]; then
-          DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD_AQL=1 python3 test/test_tiny.py
-        fi
-        pkill -f 'extra/remote/serve.py' || true
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 

.github/workflows/platform.yml

@@ -70,17 +70,19 @@ jobs:
     - name: Run pytest (amd)
       env:
         DEV: MOCKKFD+AMD
+        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
+        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py --durations=20
     - name: Run pytest (ptx)
       env:
         DEV: "MOCK+NV:PTX"
+        HCQ_RUNTIME_DEV: PYTHON
         FORWARD_ONLY: 1
         # TODO: failing due to library loading error
         CAPTURE_PROCESS_REPLAY: 0
       run: |
-        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py \
+        python3 -m pytest -n=auto test/device/test_hcq2.py test/test_tiny.py \
           test/testextra/test_hevc.py::TestHevc::test_hevc_decode_compile --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay

.github/workflows/test.yml

@@ -253,7 +253,7 @@ jobs:
         deps: testing_unit
         llvm: 'true'
     - name: Test SPEC=2
-      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py -k "not test_setitem_big" -k "not test_conv2d_ceildiv_edge_case" --splits 2 --group ${{ matrix.group }}
+      run: SPEC=2 pytest --maxfail=10 -n auto --durations=30 test/unit test/backend test/opt --ignore test/backend/test_custom_kernel.py --ignore test/unit/test_hashing.py --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -478,6 +478,7 @@ jobs:
     timeout-minutes: 20
     env:
       DEV: MOCKKFD+AMD
+      HCQ_RUNTIME_DEV: PYTHON
     steps:
       - name: Checkout Code
         uses: actions/checkout@v6
@@ -504,7 +505,7 @@ jobs:
       - name: Run AMD renderer tests (AMD:LLVM)
         run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
       - name: Run SQTT profiling tests
-        run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
+        run: SQTT_BUFFER_SIZE=16 VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
       - name: Run AMD emulated tests on NULL backend
         env:
           AMD: 0
@@ -544,14 +545,6 @@ jobs:
         run: python -m pytest test/device/test_hcq2.py
       - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
-      - name: Run test_tiny on MOCKPCI Remote
-        env:
-          HCQ2: 0
-        run: |
-          python extra/remote/serve.py 6667 &
-          sleep 2
-          REMOTE=127.0.0.1:6667 python test/test_tiny.py
-          REMOTE=127.0.0.1:6667 python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk; kill %1
 
   testamd:
     strategy:
@@ -619,7 +612,7 @@ jobs:
           cuda: 'true'
           ocelot: 'true'
       - name: Set env
-        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'ptx' && 'DEV=MOCK+CUDA:PTX' || matrix.backend == 'nv' && 'DEV=MOCK+NV\nSKIP_SLOW_TEST=1\nHCQ_RUNTIME_DEV=PYTHON' }}" >> $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"

AGENTS.md

@@ -5,3 +5,4 @@
 - Run `python -m ruff check .` to lint
 - Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
 - Do not do amend commits. Always do a new commit if a force push to origin would be required.
+- tinygrad has user space PCI drivers for AMD and NVIDIA GPUs. Do not insert the unneeded kernel modules.

conftest.py

@@ -2,7 +2,7 @@ import os, pytest, signal, threading
 
 @pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 120)), os.kill, args=(os.getpid(), signal.SIGABRT))
   t.start()
   try: yield
   finally:

docs/developer/developer.md

@@ -40,7 +40,3 @@ Then we render the UOps into code with a `Renderer`, then we compile the code to
 Runtimes are responsible for device-specific interactions. They handle tasks such as initializing devices, allocating memory, loading/launching programs, and more. You can find more information about the runtimes API on the [runtime overview page](runtime.md).
 
 All runtime implementations can be found in the [runtime directory](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime).
-
-### HCQ Compatible Runtimes
-
-HCQ API is a lower-level API for defining runtimes. Interaction with HCQ-compatible devices occurs at a lower level, with commands issued directly to hardware queues. Some examples of such backends are [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) and [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py), which are userspace drivers for NVIDIA and AMD devices respectively. You can find more information about the API on [HCQ overview page](hcq.md)

docs/developer/hcq.md

@@ -1,128 +0,0 @@
-# HCQ Compatible Runtime
-
-## Overview
-
-The main aspect of HCQ-compatible runtimes is how they interact with devices. In HCQ, all interactions with devices occur in a hardware-friendly manner using [command queues](#command-queues). This approach allows commands to be issued directly to devices, bypassing runtime overhead such as HIP or CUDA. Additionally, by using the HCQ API, these runtimes can benefit from various optimizations and features, including [HCQGraph](#hcqgraph) and built-in profiling capabilities.
-
-### Command Queues
-
-To interact with devices you create a `HWQueue`. Some methods are required, like timestamp and synchronization methods like [signal](#tinygrad.runtime.support.hcq.HWQueue.signal) and [wait](#tinygrad.runtime.support.hcq.HWQueue.wait), while others are dependent on it being a compute or copy queue.
-
-For example, the following Python code enqueues a wait, execute, and signal command on the HCQ-compatible device:
-```python
-HWQueue().wait(signal_to_wait, value_to_wait) \
-         .exec(program, args_state, global_dims, local_dims) \
-         .signal(signal_to_fire, value_to_fire) \
-         .submit(your_device)
-```
-
-Each runtime should implement the required functions that are defined in the `HWQueue` classes.
-
-::: tinygrad.runtime.support.hcq.HWQueue
-    options:
-        members: [
-            "signal",
-            "wait",
-            "timestamp",
-            "bind",
-            "submit",
-            "memory_barrier",
-            "exec",
-            "copy",
-        ]
-        show_source: false
-
-### HCQ Compatible Device
-
-The `HCQCompiled` class defines the API for HCQ-compatible devices. This class serves as an abstract base class that device-specific implementations should inherit from and implement.
-
-::: tinygrad.runtime.support.hcq.HCQCompiled
-    options:
-        show_source: false
-
-#### Signals
-
-Signals are device-dependent structures used for synchronization and timing in HCQ-compatible devices. They should be designed to record both a `value` and a `timestamp` within the same signal. HCQ-compatible backend implementations should use `HCQSignal` as a base class.
-
-::: tinygrad.runtime.support.hcq.HCQSignal
-    options:
-        members: [value, timestamp, wait]
-        show_source: false
-
-The following Python code demonstrates the usage of signals:
-
-```python
-signal = your_device.new_signal(value=0)
-
-HWQueue().timestamp(signal) \
-         .signal(signal, value_to_fire) \
-         .submit(your_device)
-
-signal.wait(value_to_fire)
-signaled_value = signal.value # should be the same as `value_to_fire`
-timestamp = signal.timestamp
-```
-
-##### Synchronization signals
-
-Each HCQ-compatible device must allocate two signals for global synchronization purposes. These signals are passed to the `HCQCompiled` base class during initialization: an active timeline signal `self.timeline_signal` and a shadow timeline signal `self._shadow_timeline_signal` which helps to handle signal value overflow issues. You can find more about synchronization in the [synchronization section](#synchronization)
-
-### HCQ Compatible Allocator
-
-The `HCQAllocator` base class simplifies allocator logic by leveraging [command queues](#command-queues) abstractions. This class efficiently handles copy and transfer operations, leaving only the alloc and free functions to be implemented by individual backends.
-
-::: tinygrad.runtime.support.hcq.HCQAllocator
-    options:
-        members: [
-            "_alloc",
-            "_free",
-        ]
-        show_source: false
-
-#### HCQ Allocator Result Protocol
-
-Backends must adhere to the `HCQBuffer` protocol when returning allocation results.
-
-::: tinygrad.runtime.support.hcq.HCQBuffer
-    options:
-        members: true
-        show_source: false
-
-### HCQ Compatible Program
-
-`HCQProgram` is a base class for defining programs compatible with HCQ-enabled devices. It provides a flexible framework for handling different argument layouts (see `HCQArgsState`).
-
-::: tinygrad.runtime.support.hcq.HCQProgram
-    options:
-        members: true
-        show_source: false
-
-#### Arguments State
-
-`HCQArgsState` is a base class for managing the argument state for HCQ programs. Backend implementations should create a subclass of `HCQArgsState` to manage arguments for the given program.
-
-::: tinygrad.runtime.support.hcq.HCQArgsState
-    options:
-        members: true
-        show_source: false
-
-**Lifetime**: The `HCQArgsState` is passed to `HWQueue.exec` and is guaranteed not to be freed until `HWQueue.submit` for the same queue is called.
-
-### Synchronization
-
-HCQ-compatible devices use a global timeline signal for synchronizing all operations. This mechanism ensures proper ordering and completion of tasks across the device. By convention, `self.timeline_value` points to the next value to signal. So, to wait for all previous operations on the device to complete, wait for `self.timeline_value - 1` value. The following Python code demonstrates the typical usage of signals to synchronize execution to other operations on the device:
-
-```python
-HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
-         .exec(...)
-         .signal(your_device.timeline_signal, your_device.next_timeline()) \
-         .submit(your_device)
-
-# Optionally wait for execution
-your_device.timeline_signal.wait(your_device.timeline_value - 1)
-```
-
-## HCQGraph
-
-[HCQGraph](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/graph/hcq.py) is a core feature that implements `GraphRunner` for HCQ-compatible devices. `HCQGraph` builds static `HWQueue` for all operations per device. To optimize enqueue time, only the necessary parts of the queues are updated for each run using the symbolic variables, avoiding a complete rebuild.
-Optionally, queues can implement a `bind` API, which allows further optimization by eliminating the need to copy the queues into the device ring.

examples/gpt2.py

@@ -57,7 +57,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
     h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).contiguous()
+    return (h + self.mlp(self.ln_2(h))).clone()
 
 class Transformer:
   def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):

examples/mlperf/models/gpt_oss.py

@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
 from tinygrad.uop.ops import Ops, UOp
 from extra.models.llama import apply_rotary_emb
 from extra.llama_kernels.rmsnorm import rmsnorm
-from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
+from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm, mx_pack
 from extra.gemm.moe_gemm import grouped_mx_gemm
 from extra.gemm.moe_routing import route, dispatch, combine, router_mfma
 
@@ -305,10 +305,25 @@ class GPTOSS:
       h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
 
     h_normed = self.norm(h)
-    pad = (-self.dim) % 256
-    h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
-    if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
-    else: logits = h_normed @ self.output.T
+
+    if getenv("FP8_LMHEAD", 0) and ASM_GEMM:
+      pad = (-self.dim) % 256
+      h2 = h_normed.reshape(-1, self.dim).pad(((0, 0), (0, pad)))
+      w2 = self.output.pad(((0, 0), (0, pad)))
+      hq, he8, hsi = quantize_mxfp8(h2)
+      oq, oe8, _ = quantize_mxfp8(w2)
+      if hsi is not None and can_use_asm_gemm(hq, oq.T):
+        logits = asm_gemm(hq, oq.T, mx=True, mx_scales=(hsi, he8, mx_pack(oe8), oe8), mx_w_stored=False)
+        logits = logits.reshape(bsz, seqlen, self.vocab_size).cast(dtypes.bfloat16)
+      else:
+        logits = h_normed @ self.output.T
+    elif ASM_GEMM:
+      pad = (-self.dim) % 256
+      h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
+      logits = asm_gemm(h_padded, w_padded.T) if can_use_asm_gemm(h_padded, w_padded.T) and getenv("VOCAB_ASM", 1) else h_normed @ self.output.T
+    else:
+      logits = h_normed @ self.output.T
+
     return logits
 
 def _get_pads(uop:UOp) -> list[UOp]:

extra/amdflash/README.md

@@ -25,11 +25,9 @@ The paths are state-dependent and are not interchangeable:
   Use it only when an empty or corrupt flash has stalled the PSP PBL. Healthy
   autonomous boot gates this engine; the usual gated status is
   `ROM_SW_STATUS=0x04000800`.
-* **`fw_live.py probe`, `ifwi-step`, and `ifwi-all`** use the early PSP
-  boot-firmware mailbox. They must run after autonomous PSP boot but before a
-  host driver or `AMDev` initializes SOS.
-* **`fw_live.py live-flash`** boots `AMDev`, stages an image in trained VRAM,
-  and invokes the Linux PSP v13 live-update command sequence.
+* **`fw_live.py probe`** queries the early PSP boot-firmware mailbox.
+* Firmware-mediated write commands are retained for protocol documentation but
+  are disabled because an exact stock reflash did not validate safely.
 * **`fw_live.py dump`** reads an exact 2 MiB raw image through
   `ROM_INDEX/ROM_DATA`. It refuses devices where the raw SMUIO controller is
   unavailable; the NBIO SOC15 function-ROM aperture is not a physical SPI
@@ -78,55 +76,41 @@ Navi31 ROM_SW details used by the implementation:
 
 ## Firmware-mediated access
 
-Query the early mailbox without changing flash:
+The read-only commands are:
 
 ```sh
 python3 extra/amdflash/fw_live.py probe
-```
-
-Stream one exact item only when its type matches the firmware request:
-
-```sh
-python3 extra/amdflash/fw_live.py stream 0x37 vbios-item.bin --yes
-```
-
-For a complete 2 MiB IFWI, either perform one requested step or follow requests
-until firmware reports a `0x2xx` completion state:
-
-```sh
-python3 extra/amdflash/fw_live.py ifwi-step full-ifwi.bin --yes
-python3 extra/amdflash/fw_live.py ifwi-all full-ifwi.bin --yes
-```
-
-A completion state requires a hard reset; it is not another item request. The
-resolver supports recovery metadata types `0x01`-`0x08` and `0x80`-`0x89`,
-including the firmware-selected inactive partition. This path is signature
-enforcing and intentionally refuses mismatched item types.
-
-The early protocol is:
-
-* `START_TRANSFER`: `(size << 8) | item_type`
-* `DATA_TRANSFER`: one little-endian image dword per mailbox command
-* `END_TRANSFER`: `(bytes_sent << 8) | item_type`
-
-Each dword requires a firmware acknowledgement, so large partition transfers
-are slow through a USB-PCIe bridge.
-
-The fully initialized PSP path and the healthy-state dump are:
-
-```sh
-python3 extra/amdflash/fw_live.py live-flash signed-update.bin --yes
 python3 extra/amdflash/fw_live.py dump current-spi.bin
 ```
 
-The PSP validates live-update inputs and may reject an image even when its size
-and alignment are valid. `dump` produces exactly `0x200000` bytes, requires the
-raw IFWI magic at offset zero, rejects mirrored 1 MiB apertures, and restores
-the ROM controller/index state before writing output.
+`dump` produces exactly `0x200000` bytes, requires the raw IFWI magic at offset
+zero, rejects mirrored 1 MiB apertures, and restores the ROM controller/index
+state before writing output.
+
+The validated early-firmware sequence is available as:
+
+```sh
+python3 extra/amdflash/fw_live.py --transport usb ifwi-all full-ifwi.bin --yes
+```
+
+It resolves at most Navi31's configured 19 items, streams the item associated
+with terminal phase `0x2xx`, and then stops. PSP selects the destination
+partition; item `0x08` always comes from the payload referenced by the first
+ISH descriptor, matching AMDVBFlash. A hard power cycle is required afterward.
+
+A successful PSP update is not a byte-identical raw rewrite. On the validated
+stock test, both A/B payloads matched the source exactly, PSP selected and
+booted the updated B partition, and firmware changed only its update cookie,
+B descriptor counter/checksum, and generated metadata near `0x1ef000`.
+
+The `stream`, `ifwi-step`, and `live-flash` commands remain disabled. Testing
+showed that the PSP live path parses a raw stock IFWI but fails with status
+`0xC` (`PSP Write To SPI Error`) after writing an `$AMDVBFL` cookie. Use the
+verified ROM_SW path for recovery.
 
 ## Safety
 
-All erase, program, and firmware-streaming commands require `--yes`. Read-only
-commands still touch controller and mailbox registers but do not issue SPI
-program/erase or PSP transfer-start commands. Preserve a known-good full dump
-outside the repository.
+ROM_SW erase/program and `ifwi-all` commands require `--yes`; other
+firmware-streaming commands are disabled. Read-only commands still touch controller and mailbox registers but
+do not issue SPI program/erase or PSP transfer-start commands. Preserve a
+known-good full dump outside the repository.

extra/amdflash/fw_live.py

@@ -19,12 +19,24 @@ GET_BOOT_PARTITION, GET_FB_STATE, GET_TRANSFER_TYPE = 0x01, 0x06, 0x07
 START_TRANSFER, DATA_TRANSFER, END_TRANSFER = 0x08, 0x09, 0x0A
 SPI_GET_MODEL_ID = 0x0B
 LIVE_ADDR_LO, LIVE_ADDR_HI, LIVE_UPDATE = 0x02, 0x03, 0x04
+PSP_ERRORS = {
+  0x01: "generic error", 0x02: "out of bounds", 0x03: "invalid parameter",
+  0x04: "off-chip boot error", 0x05: "address not set", 0x06: "parse off-chip error",
+  0x07: "address map error", 0x08: "parse on-chip error", 0x09: "full update error",
+  0x0A: "partition update error", 0x0B: "map on-chip error", 0x0C: "write to SPI error",
+  0x0D: "signature validation error", 0x0E: "invalid command", 0x0F: "signature not found",
+  0x10: "state machine not initialized", 0x11: "state machine transfer error",
+  0x12: "initialization error",
+}
 
 
 class PSPFlashMailbox:
   def __init__(self, pci_dev): self.mmio = MMIO(pci_dev)
 
   def command(self, command: int, data: int | None = None, *, timeout: float = 10.0) -> tuple[int, int]:
+    status = self.mmio.read32(COMMAND)
+    if not status & 0x80000000:
+      raise RuntimeError(f"PSP mailbox is not ready before command {command:#x}: status={status:#010x}")
     if data is not None: self.mmio.write32(COMMAND_DATA, data)
     self.mmio.write32(COMMAND, command << 16)
     self.mmio.write32(DOORBELL, 1)
@@ -36,8 +48,8 @@ class PSPFlashMailbox:
   def require(self, command: int, data: int | None = None, *, timeout: float = 10.0, name: str = '') -> int:
     error, response = self.command(command, data, timeout=timeout)
     if error:
-      hint = " (update interface is gated after SOS initialization)" if error == 0xA else ""
-      raise RuntimeError(f"PSP {name or hex(command)} failed: error={error:#x}{hint}")
+      detail = PSP_ERRORS.get(error, "unknown error")
+      raise RuntimeError(f"PSP {name or hex(command)} failed: error={error:#x} ({detail})")
     return response
 
   def probe(self) -> dict[str, tuple[int, int]]:
@@ -47,16 +59,14 @@ class PSPFlashMailbox:
       result[name] = self.command(command)
     return result
 
-  def stream(self, payload: bytes, item_type: int):
+  def stream(self, payload: bytes, item_type: int, transfer_type: int | None = None):
     if not payload: raise ValueError("payload is empty")
     if len(payload) > 0xFFFFFF: raise ValueError("payload exceeds the mailbox's 24-bit size field")
     if len(payload) & 3: raise ValueError("payload size must be divisible by four")
     if not 0 <= item_type <= 0xff: raise ValueError("item type must fit in eight bits")
-    transfer_type = self.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
+    if transfer_type is None: transfer_type = self.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
     requested = transfer_type & 0xff
     print(f"firmware transfer_type={transfer_type:#x}", flush=True)
-    if transfer_type & 0x200:
-      raise RuntimeError(f"firmware reports completion state {transfer_type:#x}; hard reset required")
     if requested != item_type:
       raise RuntimeError(f"firmware requests item {requested:#x}, not {item_type:#x}")
     self.require(START_TRANSFER, (len(payload) << 8) | item_type, name="START_TRANSFER")
@@ -79,9 +89,8 @@ class PSPFlashMailbox:
     print(f"stream complete: type={item_type:#x} size={sent:#x} elapsed={time.monotonic()-started:.1f}s")
 
 
-def resolve_ifwi_item(image: bytes, item_type: int, active: int) -> tuple[int, bytes]:
+def resolve_ifwi_item(image: bytes, item_type: int) -> tuple[int, bytes]:
   """Resolve AMDVBFlash recovery-layout item types to exact IFWI bytes."""
-  if active not in (1, 2): raise ValueError(f"unexpected one-based active partition {active}")
   if item_type == 0x01: offset, size = 0, 0x54
   elif item_type in (0x02, 0x03):
     offset = 0x2000 if item_type == 0x02 else 0x3000
@@ -107,11 +116,11 @@ def resolve_ifwi_item(image: bytes, item_type: int, active: int) -> tuple[int, b
     _, _, size, offset = match[0]
   elif item_type == 0x89: offset, size = 0x1f0000, 0x100
   elif item_type == 0x08:
-    # Firmware requests the inactive partition. A 0x2xx transfer state after
-    # this item is completion/reset-required status, not another item instance.
-    descriptor = 0x13000 if active == 1 else 0x12000
-    offset = struct.unpack_from('<I', image, descriptor + 0x10)[0]
-    size = struct.unpack_from('<I', image, descriptor + 0x18)[0]
+    # AMDVBFlash's GetPartitionDetails follows the first ISH entry (firmware ID
+    # 0x13c) and streams its payload. PSP, not the host resolver, selects the
+    # destination partition.
+    offset = struct.unpack_from('<I', image, 0x12000 + 0x10)[0]
+    size = struct.unpack_from('<I', image, 0x12000 + 0x18)[0]
   else:
     raise ValueError(f"IFWI resolver does not yet support requested item {item_type:#x}")
   payload = image[offset:offset+size]
@@ -139,6 +148,10 @@ class LivePSPFlash:
 def open_mailbox(args): return PSPFlashMailbox(open_gpu(args.device, args.transport))
 
 
+def reject_unvalidated_firmware_write():
+  raise RuntimeError("firmware writes are disabled: stock reflash validation failed; use romless.py for recovery")
+
+
 def cmd_probe(args):
   result = open_mailbox(args).probe()
   for name, (error, response) in result.items(): print(f"{name}: error={error:#x} response={response:#x}")
@@ -147,21 +160,21 @@ def cmd_probe(args):
 
 def cmd_stream(args):
   if not args.yes: raise RuntimeError("refusing to stream without --yes")
+  reject_unvalidated_firmware_write()
   payload = Path(args.image).read_bytes()
   open_mailbox(args).stream(payload, args.item_type)
 
 
 def cmd_ifwi_step(args):
   if not args.yes: raise RuntimeError("refusing to stream without --yes")
+  reject_unvalidated_firmware_write()
   image = Path(args.ifwi).read_bytes()
   if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
   mailbox = open_mailbox(args)
-  active = mailbox.require(GET_BOOT_PARTITION, name="GET_BOOT_PARTITION")
   state = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
   request = state & 0xff
-  if state & 0x200: raise RuntimeError(f"firmware reports completion state {state:#x}; reset instead of streaming another item")
-  _, payload = resolve_ifwi_item(image, request, active)
-  mailbox.stream(payload, request)
+  _, payload = resolve_ifwi_item(image, request)
+  mailbox.stream(payload, request, transfer_type=state)
   next_request = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
   print(f"next firmware transfer_type={next_request:#x}")
 
@@ -171,22 +184,24 @@ def cmd_ifwi_all(args):
   image = Path(args.ifwi).read_bytes()
   if len(image) != 0x200000: raise ValueError("Navi31 IFWI image must be exactly 2 MiB")
   mailbox = open_mailbox(args)
-  active = mailbox.require(GET_BOOT_PARTITION, name="GET_BOOT_PARTITION")
   current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  for step in range(32):
-    if current & 0x200:
-      print(f"IFWI update complete: firmware state={current:#x}; hard reset required")
+  for step in range(19):  # Navi31 ROMItemCount from AMDVBFlash ASICDetails.xml
+    request, phase = current & 0xff, current >> 8
+    print(f"IFWI step {step}: state={current:#x} item={request:#x} phase={phase}", flush=True)
+    _, payload = resolve_ifwi_item(image, request)
+    mailbox.stream(payload, request, transfer_type=current)
+    # AMDVBFlash tests the high byte belonging to the item just streamed. Phase
+    # 2 terminates the loop only after that item has completed successfully.
+    if phase == 2:
+      print(f"IFWI stream complete after terminal state {current:#x}; hard power cycle required")
       return
-    request = current & 0xff
-    print(f"IFWI step {step}: state={current:#x} item={request:#x}", flush=True)
-    _, payload = resolve_ifwi_item(image, request, active)
-    mailbox.stream(payload, request)
     current = mailbox.require(GET_TRANSFER_TYPE, name="GET_TRANSFER_TYPE")
-  raise RuntimeError(f"IFWI request cycle did not close after 32 items (state={current:#x})")
+  raise RuntimeError(f"IFWI stream did not reach terminal phase after 19 items (state={current:#x})")
 
 
 def cmd_live_flash(args):
   if not args.yes: raise RuntimeError("refusing to flash without --yes")
+  reject_unvalidated_firmware_write()
   image = Path(args.ifwi).read_bytes()
   if not image or len(image) > 16 * 1024 * 1024 or len(image) & 3:
     raise ValueError("live PSP image must be non-empty, 4-byte aligned, and at most 16 MiB")

extra/amdpci/am_smi.py

@@ -66,6 +66,8 @@ class AMSMI(AMDev):
   def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
     self.pcibus, self.devfmt = pcibus, pcibus
     self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
+    self.is_vf = bool(self.mmio[am.mmRCC_IOV_FUNC_IDENTIFIER] & 1)
+    self.vf_rlc_gated:list[tuple[int, int]] = []
     self.pci_state = self.read_pci_state()
     if self.pci_state == "D0": self._init_from_d0()
 

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.engine.realize import compile_linear
+from tinygrad.engine.realize import lower_and_compile
 from tinygrad.codegen.opt import OptOps
 
 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
@@ -39,7 +39,7 @@ if __name__ == "__main__":
     c = a.matmul(b, dtype=acc_dtype).realize()
 
   if getenv("SHOULD_USE_TC"):
-    linear = compile_linear(a.matmul(b, dtype=acc_dtype).schedule_linear())
+    linear = lower_and_compile(a.matmul(b, dtype=acc_dtype).schedule_linear())
     call = get_single_element(list(linear.src))
     applied_opts = call.src[0].src[0].arg.applied_opts
     assert any(opt.op is OptOps.TC for opt in applied_opts), f"TC not triggered, {applied_opts}"

extra/hcq1/graph.py

@@ -1,7 +1,8 @@
 import collections, time
 from typing import Any, cast
 from tinygrad.helpers import round_up, PROFILE, ALL2ALL, merge_dicts, getenv, suppress_finalizing, TracingKey, unwrap
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQSignal, HCQBuffer, HWQueue, HCQArgsState, BumpAllocator, MMIOInterface
+from extra.hcq1.hcq import HCQCompiled, HCQAllocator, HCQSignal, HWQueue, HCQArgsState
+from tinygrad.runtime.support.hcq import HCQBuffer, BumpAllocator, MMIOInterface
 from tinygrad.device import Buffer, BufferSpec, Compiled, Device, MultiBuffer, ProfileGraphEntry, ProfileGraphEvent
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, Variable
@@ -102,7 +103,7 @@ class HCQGraph(MultiGraphRunner):
       elif is_rdma:
         enqueue_queue = self.comp_queues[enqueue_dev]
         rdma_key = (cast(HCQCompiled, Device[bufs[0].device]).rdma_dev(), enqueue_dev.rdma_dev())
-        from tinygrad.runtime.ops_rdma import RDMACopyQueue
+        from extra.hcq1.ops_rdma import RDMACopyQueue
         self.rdma_queues.setdefault(rdma_key, RDMACopyQueue(enqueue_dev.rdma_dev()))
       else:
         assert (enqueue_dev.hw_copy_queue_t is not None), "device must implement a copy queue"

extra/hcq1/hcq.py

@@ -0,0 +1,551 @@
+from __future__ import annotations
+from typing import cast, Callable, Type, TypeVar, Generic, Any
+import contextlib, decimal, statistics, time, ctypes, array, collections, itertools
+from tinygrad.helpers import PROFILE, getenv, from_mv, cpu_profile, ProfileRangeEvent, unwrap
+from tinygrad.helpers import suppress_finalizing, TracingKey
+from tinygrad.device import BufferSpec, Compiled, LRUAllocator, ProfileDeviceEvent, ProfileProgramEvent, Program, TinyELF
+from tinygrad.uop.ops import sym_infer, sint, UOp
+from tinygrad.runtime.support.memory import BumpAllocator, MMIOInterface
+from tinygrad.renderer import Renderer
+from tinygrad.runtime.support.hcq import HCQBuffer
+
+SignalType = TypeVar('SignalType', bound='HCQSignal')
+HCQDeviceType = TypeVar('HCQDeviceType', bound='HCQCompiled')
+ProgramType = TypeVar('ProgramType', bound='HCQProgram')
+ArgsStateType = TypeVar('ArgsStateType', bound='HCQArgsState')
+
+class HWQueue(Generic[SignalType, HCQDeviceType, ProgramType, ArgsStateType]):
+  """
+  A base class for hardware command queues in the HCQ (Hardware Command Queue) API.
+  """
+
+  def __init__(self):
+    self._q:Any = []
+    self.binded_device:HCQDeviceType|None = None
+    self.q_sints:list[tuple[int, int]] = []
+    self.mv_sints:list[tuple[MMIOInterface, int, int, int|None]] = []
+    self.syms:list[sint] = []
+    self._prev_resolved_syms:list[int|None] = []
+
+  def _new_sym(self, sym:sint) -> int:
+    if sym not in self.syms:
+      self.syms.append(sym)
+      self._prev_resolved_syms.append(None)
+    return self.syms.index(sym)
+
+  def q(self, *values):
+    """
+    Enqueues values in the queue.
+
+    Args:
+      values: The values to enqueue in the queue.
+    """
+
+    for v in values:
+      if isinstance(v, UOp):
+        self.q_sints.append((len(self._q), self._new_sym(v)))
+        self._q.append(0xbadc0ded)
+      else: self._q.append(v)
+
+  # *** common commands  ***
+
+  def timestamp(self, signal:SignalType):
+    """
+    Enqueues a timestamp command which records the current time in a signal after all previously enqueued commands are completed.
+
+    Args:
+      signal: The signal to store the timestamp
+    """
+
+  def signal(self, signal:SignalType, value:sint):
+    """
+    Enqueues a signal command which sets the signal to the given value, ensuring all previous operations are completed.
+
+    Args:
+      signal: The signal to set
+      value: The value to set the signal to
+    """
+
+  def wait(self, signal:SignalType, value:sint):
+    """
+    Enqueues a wait command which halts execution until the signal is greater than or equal to a specific value.
+
+    Args:
+      signal: The signal to wait on
+      value: The value to wait for
+    """
+
+  # *** commands for compute queues ***
+
+  def memory_barrier(self):
+    """
+    Enqueues a memory barrier command to ensure memory coherence between agents. Only on compute queues.
+    """
+
+  def exec(self, prg:ProgramType, args_state:ArgsStateType, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
+    """
+    Enqueues an execution command for a kernel program. Only on compute queues.
+
+    Args:
+      prg: The program to execute
+      args_state: The args state to execute program with
+      global_size: The global work size
+      local_size: The local work size
+    """
+
+  def write(self, b:HCQBuffer, val:sint, b64:bool=False):
+    """
+    Enqueues a command to write a value to a buffer address after all previously enqueued commands are completed.
+
+    Args:
+      b: The buffer to write to
+      val: The value to write
+      b64: If True, write a 64-bit value; otherwise write 32-bit
+    """
+    raise NotImplementedError("write not implemented")
+
+  def poll_bit(self, b:HCQBuffer, val:sint, mask:int):
+    """
+    Enqueues a poll command which halts execution until (mem[b] & mask) == val.
+    val must be 0 or mask (i.e. checks if masked bits are all clear or all set).
+
+    Args:
+      b: The buffer to poll
+      val: The expected value after masking (0 or mask)
+      mask: The bit mask to test
+    """
+    raise NotImplementedError("poll_bit not implemented")
+
+  # *** commands for copy queues ***
+
+  def copy(self, dest:HCQBuffer, src:HCQBuffer, copy_size:int):
+    """
+    Enqueues a copy command to transfer data. Only on copy queues.
+
+    Args:
+      dest: The destination buffer of the copy
+      src: The source buffer of the copy
+      copy_size: The size of data to copy
+    """
+
+  # *** submit and bind commands  ***
+
+  def bind(self, dev:HCQDeviceType):
+    """
+    Associates the queue with a specific device for optimized execution.
+
+    This optional method allows backend implementations to tailor the queue for efficient use on the given device. When implemented, it can eliminate
+    the need to copy queues into the device, thereby enhancing performance.
+
+    Args:
+      dev: The target device for queue optimization.
+
+    Note:
+      Implementing this method is optional but recommended for performance gains.
+    """
+
+  def bind_args_state(self, args_state:ArgsStateType):
+    for vals, mem, fmt in args_state.bind_data: self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt)
+
+  def bind_sints(self, *vals:sint, mem:MMIOInterface, struct_t:Type[ctypes.Structure], start_field:str, fmt, mask:int|None=None):
+    self.bind_sints_to_mem(*vals, mem=mem, fmt=fmt, mask=mask, offset=getattr(struct_t, start_field).offset)
+
+  def bind_sints_to_mem(self, *vals:sint, mem:MMIOInterface, fmt, mask:int|None=None, offset:int=0):
+    mv = mem.view(offset=offset, size=len(vals)*8, fmt=fmt)
+    for i, val in enumerate(vals):
+      if isinstance(val, int): mv[i] = val if mask is None else ((mv[i] & ~mask) | val)
+      else: self.mv_sints.append((mv, i, self._new_sym(val), mask))
+
+  def _apply_var_vals(self, var_vals:dict[str, int]):
+    resolved_syms: list[int|None] = [sym_infer(sym, var_vals) for sym in self.syms]
+
+    for off, sym_idx in self.q_sints:
+      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
+      self._q[off] = resolved_syms[sym_idx]
+
+    for mv, off, sym_idx, mask in self.mv_sints:
+      if self._prev_resolved_syms[sym_idx] == resolved_syms[sym_idx]: continue
+      mv[off] = resolved_syms[sym_idx] if mask is None else ((mv[off] & ~mask) | resolved_syms[sym_idx])
+
+    self._prev_resolved_syms = resolved_syms
+
+  def submit(self, dev:HCQDeviceType, var_vals:dict[str, int]|None=None):
+    """
+    Submits the command queue to a specific device for execution.
+
+    Args:
+      dev: The device to submit the queue to
+    """
+
+    if var_vals is not None: self._apply_var_vals(var_vals)
+    self._submit(dev)
+    return self
+  def _submit(self, dev:HCQDeviceType): raise NotImplementedError("need _submit")
+
+class HCQSignal(Generic[HCQDeviceType]):
+  def __init__(self, base_buf:HCQBuffer, value:int=0, owner:HCQDeviceType|None=None, is_timeline:bool=False, timestamp_divider=1000, virt=False):
+    self.base_buf, self.owner, self.is_timeline = base_buf, owner, is_timeline
+    self.should_return = isinstance(self.base_buf.va_addr, int) and self.owner is not None and not virt
+    self.timestamp_divider:decimal.Decimal = decimal.Decimal(timestamp_divider)
+    if isinstance(self.base_buf.va_addr, int) and not virt: self.value = value
+
+  def __del__(self):
+    if self.should_return: HCQCompiled.signal_pool[unwrap(self.owner).peer_group].append(self.base_buf)
+
+  @property
+  def value_addr(self) -> sint: return self.base_buf.va_addr
+
+  @property
+  def timestamp_addr(self) -> sint: return self.base_buf.va_addr + 8
+
+  @property
+  def value(self) -> int: return self.base_buf.cpu_view().view(0, 8, 'Q')[0]
+
+  @value.setter
+  def value(self, new_value:int): self.base_buf.cpu_view().view(0, 8, 'Q')[0] = new_value
+
+  @property
+  def timestamp(self) -> decimal.Decimal:
+    """
+    Get the timestamp field of the signal.
+
+    This property provides read-only access to the signal's timestamp.
+
+    Returns:
+      The timestamp in microseconds.
+    """
+    return self.base_buf.cpu_view().view(8, 8, 'Q')[0] / self.timestamp_divider
+
+  def _sleep(self, time_spent_since_last_sleep_ms:int):
+    """
+    Optional function which can implement sleep functionality for the signal.
+    Raises RuntimeError if a fault is detected.
+    """
+
+  def wait(self, value:int, timeout:int|None=None):
+    """
+    Waits the signal is greater than or equal to a specific value.
+
+    Args:
+      value: The value to wait for.
+      timeout: Maximum time to wait in milliseconds. Defaults to 30s.
+    """
+    timeout = timeout or getenv("HCQDEV_WAIT_TIMEOUT_MS", 30000)
+    start_time = int(time.perf_counter() * 1000)
+    while (not_passed:=(prev_value:=self.value) < value) and (cur_time:=int(time.perf_counter() * 1000)) - start_time < timeout:
+      self._sleep(cur_time - start_time)
+      if self.value != prev_value: start_time = int(time.perf_counter() * 1000) # progress was made, reset timer
+    if not_passed and self.value < value: raise RuntimeError(f"Wait timeout: {timeout} ms! (the signal is not set to {value}, but {self.value})")
+
+@contextlib.contextmanager
+def hcq_profile(dev:HCQCompiled, enabled, desc, queue_type:Callable[[], HWQueue]|None=None, queue:HWQueue|None=None, dev_suff:str|None=None,
+                profile_key:bytes|None=None):
+  st, en = (dev.new_signal(), dev.new_signal()) if enabled else (None, None)
+  assert queue is not None or queue_type is not None, "Either queue or queue_type must be provided"
+
+  if enabled and queue is not None: queue.timestamp(st)
+  elif enabled and queue_type is not None:
+    queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(st).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
+
+  try: yield (st, en)
+  finally:
+    if enabled and queue is not None: queue.timestamp(en)
+    elif enabled and queue_type is not None:
+      queue_type().wait(dev.timeline_signal, dev.timeline_value - 1).timestamp(en).signal(dev.timeline_signal, dev.next_timeline()).submit(dev)
+
+    if enabled and PROFILE: dev.sig_prof_records.append((unwrap(st), unwrap(en), desc, f"{dev.device}:{dev_suff}" if dev_suff else dev.device,
+                                                         profile_key))
+
+class HCQArgsState(Generic[ProgramType]):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=()):
+    self.buf, self.prg, self.bufs, self.vals = buf, prg, bufs, vals
+    self.bind_data:list[tuple[tuple[sint, ...], MMIOInterface, str]] = []
+
+  def bind_sints_to_buf(self, *vals:sint, buf:HCQBuffer, fmt, offset=0): self.bind_data.append((vals, buf.cpu_view().view(offset=offset), fmt))
+
+class CLikeArgsState(HCQArgsState[ProgramType]):
+  def __init__(self, buf:HCQBuffer, prg:ProgramType, bufs:tuple[HCQBuffer, ...], vals:tuple[sint|None, ...]=(), prefix:list[int]|None=None):
+    super().__init__(buf, prg, bufs, vals=vals)
+
+    if prefix is not None: self.buf.cpu_view().view(size=len(prefix) * 4, fmt='I')[:] = array.array('I', prefix)
+
+    self.bind_sints_to_buf(*[b.va_addr for b in bufs], buf=self.buf, fmt='Q', offset=len(prefix or []) * 4)
+    for v,(val_offset,dt) in zip(vals, TinyELF.iter_sig(prg.signature[-len(vals):], len(bufs) * 8)):
+      assert v is not None
+      self.bind_sints_to_buf(v, buf=self.buf, fmt=dt.fmt, offset=len(prefix or []) * 4 + val_offset)
+
+class HCQProgram(Program[HCQDeviceType]):
+  def __init__(self, args_state_t:Type[HCQArgsState], dev:HCQDeviceType, obj:TinyELF, kernargs_alloc_size:int, base:int|None=None):
+    self.args_state_t, self.dev, self.name, self.signature, self.kernargs_alloc_size = args_state_t, dev, obj.name, obj.signature, kernargs_alloc_size
+    self.profile_key = obj.profile_key
+    self.prof_prg_counter = next(self.dev.prof_prg_counter)
+    if PROFILE: Compiled.profile_events += [ProfileProgramEvent(dev.device, obj.name, obj.lib, base, self.prof_prg_counter, self.profile_key)]
+
+  @staticmethod
+  def _fini(dev, buf, spec): dev.allocator.free(buf, buf.size, spec)
+
+  def fill_kernargs(self, bufs:tuple[HCQBuffer, ...], vals:tuple[int|None, ...]=(), kernargs:HCQBuffer|None=None) -> HCQArgsState:
+    """
+    Fills arguments for the kernel, optionally allocating space from the device if `kernargs_ptr` is not provided.
+    Args:
+      bufs: Buffers to be written to kernel arguments.
+      vals: Values to be written to kernel arguments.
+      kernargs_ptr: Optional pointer to pre-allocated kernel arguments memory.
+    Returns:
+      Arguments state with the given buffers and values set for the program.
+    """
+    argsbuf = kernargs or self.dev.kernargs_buf.offset(offset=self.dev.kernargs_offset_allocator.alloc(self.kernargs_alloc_size, 8),
+                                                       size=self.kernargs_alloc_size)
+    return self.args_state_t(argsbuf, self, bufs, vals=vals)
+
+  def __call__(self, *bufs:HCQBuffer, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1),
+               vals:tuple[int|None, ...]=(), wait:bool=False, timeout:int|None=None) -> float|None:
+    """
+    Enqueues the program for execution with the given arguments and dimensions.
+
+    Args:
+      bufs: Buffer arguments to execute the kernel with.
+      global_size: Specifies the global work size for kernel execution (equivalent to CUDA's grid size).
+      local_size: Specifies the local work size for kernel execution (equivalent to CUDA's block size).
+      vals: Value arguments to execute the kernel with.
+      wait: If True, waits for the kernel to complete execution.
+
+    Returns:
+      Execution time of the kernel if 'wait' is True, otherwise None.
+    """
+
+    kernargs = self.fill_kernargs(bufs, vals)
+    q = unwrap(self.dev.hw_compute_queue_t)().wait(self.dev.timeline_signal, self.dev.timeline_value - 1).memory_barrier()
+
+    self.dev.prof_exec_counter += 1
+    with hcq_profile(self.dev, queue=q, desc=self.name, enabled=wait or PROFILE, profile_key=self.profile_key) as (sig_st, sig_en):
+      q.exec(self, kernargs, global_size, local_size)
+
+    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+
+    if wait: self.dev.synchronize(timeout=timeout)
+    return (float(sig_en.timestamp - sig_st.timestamp) / 1e6) if wait else None
+
+class HCQCompiled(Compiled, Generic[SignalType]):
+  """
+  A base class for devices compatible with the HCQ (Hardware Command Queue) API.
+  """
+  peer_groups: dict[str, list[HCQCompiled]] = collections.defaultdict(list)
+  signal_pages: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
+  signal_pool: dict[str, list[HCQBuffer]] = collections.defaultdict(list) # per peer group
+  cpu_devices: list[HCQCompiled] = []
+
+  def __init__(self, device:str, allocator:HCQAllocatorBase, compilers:list[type[Renderer]], runtime:type[Program]|None,
+               signal_t:Type[SignalType]|None=None, comp_queue_t:Callable[..., HWQueue]|None=None, copy_queue_t:Callable[..., HWQueue]|None=None,
+               kernargs_size=(16 << 20), sigalloc_size=0x1000, can_recover:bool=False, arch=None):
+    from extra.hcq1.graph import HCQGraph
+    super().__init__(device, allocator, compilers, runtime, HCQGraph, arch=arch)
+
+    self.peer_group = getattr(getattr(self, 'iface', None), 'peer_group', device.split(":")[0])
+    HCQCompiled.peer_groups[self.peer_group].append(self)
+
+    self.signal_t, self.hw_compute_queue_t, self.hw_copy_queue_t = signal_t, comp_queue_t, copy_queue_t
+
+    self.timeline_value:int = 1
+    self.sig_prof_records:list[tuple[HCQSignal, HCQSignal, str|TracingKey, str, bytes|None]] = []
+    self.prof_exec_counter:int = 0
+    self.prof_prg_counter = itertools.count(0)
+
+    if signal_t is not None:
+      # Map signals if any
+      for sig_page in HCQCompiled.signal_pages[self.peer_group]: cast(HCQAllocator, self.allocator)._map(sig_page)
+
+      self.sigalloc_size = sigalloc_size
+      self.timeline_signal, self._shadow_timeline_signal = self.new_signal(value=0, is_timeline=True), self.new_signal(value=0, is_timeline=True)
+
+    if comp_queue_t is not None:
+      self.kernargs_buf:HCQBuffer = self.allocator.alloc(kernargs_size, BufferSpec(cpu_access=True))
+      self.kernargs_offset_allocator:BumpAllocator = BumpAllocator(self.kernargs_buf.size, wrap=True)
+
+    self.can_recover = can_recover # Whether the device can recover from faults or timeouts
+    self.error_state:Exception|None = None # Exception if error is unrecoverable and sync will always fail
+
+    if self._is_cpu(): HCQCompiled.cpu_devices.append(self)
+
+  def synchronize(self, timeout:int|None=None):
+    if self.error_state is not None: raise self.error_state
+    if not hasattr(self, 'timeline_signal'): return
+
+    # If we have any work on CPU devices, need to synchronize them. This is just an optimization to release GIL allowing to finish faster.
+    if not self._is_cpu():
+      for dev in HCQCompiled.cpu_devices: dev.synchronize()
+
+    try: self.timeline_signal.wait(self.timeline_value - 1, timeout=timeout if timeout is not None and self.can_recover else None)
+    except RuntimeError as e:
+      self.error_state = e
+      if hasattr(self, 'on_device_hang'): self.on_device_hang()
+      raise e
+
+    if self.timeline_value > (1 << 31): self._wrap_timeline_signal()
+    if PROFILE:
+      Compiled.profile_events += [ProfileRangeEvent(dev, name, st.timestamp, en.timestamp, pk) for st,en,name,dev,pk in self.sig_prof_records]
+      self.sig_prof_records = []
+
+  def next_timeline(self):
+    self.timeline_value += 1
+    return self.timeline_value - 1
+
+  def new_signal(self, **kwargs) -> SignalType:
+    assert self.signal_t is not None, "Device does not support signals"
+    if not HCQCompiled.signal_pool[pg:=self.peer_group]:
+      HCQCompiled.signal_pages[pg].append(alc:=self.allocator.alloc(self.sigalloc_size, BufferSpec(host=True, uncached=True, cpu_access=True)))
+      HCQCompiled.signal_pool[pg] += [alc.offset(offset=off, size=16) for off in range(0, alc.size, 16)]
+      for dev in HCQCompiled.peer_groups[pg]: cast(HCQAllocator, dev.allocator)._map(alc)
+    return self.signal_t(base_buf=HCQCompiled.signal_pool[pg].pop(), owner=self, **kwargs)
+
+  def device_props(self) -> dict[str,Any]: return {} # to be overridden if needed. dict keys are backend dependent.
+
+  def hw_compute_queues(self) -> list[tuple[str|None, Callable[[], HWQueue]]]:
+    return [(None, self.hw_compute_queue_t)] if self.hw_compute_queue_t is not None else []
+  def hw_copy_queues(self) -> list[tuple[str, Callable[[], HWQueue]]]:
+    return [("SDMA:0", self.hw_copy_queue_t)] if self.hw_copy_queue_t is not None else []
+
+  def _at_profile_finalize(self):
+    self.synchronize() # Expect device to be synchronizes
+
+    def _sync(d:HCQCompiled, q_t:Callable[[], HWQueue]):
+      q_t().timestamp(d.timeline_signal).signal(d.timeline_signal, d.next_timeline()).submit(d)
+      st = time.perf_counter_ns()
+      d.timeline_signal.wait(d.timeline_value - 1)  # average of the two
+      et = time.perf_counter_ns()
+      return (decimal.Decimal(et+st) / 2000) - d.timeline_signal.timestamp
+
+    for prefix, q_t in self.hw_compute_queues() + self.hw_copy_queues():
+      devname = f"{self.device}:{prefix}" if prefix else self.device
+      Compiled.profile_events += [ProfileDeviceEvent(devname, statistics.median([_sync(self, q_t) for _ in range(40)]), props=self.device_props())]
+
+  def _wrap_timeline_signal(self):
+    self.timeline_signal, self._shadow_timeline_signal, self.timeline_value = self._shadow_timeline_signal, self.timeline_signal, 1
+    self.timeline_signal.value = 0
+    cast(HCQAllocatorBase, self.allocator).b_timeline = [0] * len(cast(HCQAllocatorBase, self.allocator).b)
+
+  def _realloc(self, oldbuf:HCQBuffer|None, new_size:int, options:BufferSpec|None=None, force=False) -> tuple[HCQBuffer, bool]:
+    if oldbuf is not None: self.allocator.free(oldbuf, oldbuf.size, options=options)
+    try: buf, realloced = self.allocator.alloc(new_size, options=options), True
+    except MemoryError:
+      if force: raise
+      buf, realloced = self.allocator.alloc(oldbuf.size if oldbuf is not None else new_size, options=options), False
+    return buf, realloced
+
+  def _is_cpu(self) -> bool: return hasattr(self, 'device') and self.device.split(":")[0] == "CPU"
+
+  def rdma_dev(self):
+    from extra.hcq1.ops_rdma import get_rdma_device
+    for i in itertools.count():
+      if (dev:=next((d for d in HCQCompiled.peer_groups[self.peer_group] if type(d).__name__ == 'RDMADevice'), None)): return dev
+      try: get_rdma_device(i)
+      except IndexError: raise RuntimeError(f"No RDMA found for peer group '{self.peer_group}'")
+
+  def finalize(self):
+    try: self.synchronize() # Try to finalize device in any case.
+    except RuntimeError as e: print(f"{self.device} synchronization failed before finalizing: {e}")
+    super().finalize()
+
+class HCQAllocatorBase(LRUAllocator[HCQDeviceType], Generic[HCQDeviceType]):
+  """
+  A base allocator class compatible with the HCQ (Hardware Command Queue) API.
+
+  This class implements basic copy operations following the HCQ API, utilizing both types of `HWQueue`.
+  """
+
+  def __init__(self, dev:HCQDeviceType, batch_size:int=(2 << 20), batch_cnt:int=32, copy_bufs=None, **kwargs):
+    super().__init__(dev, **kwargs)
+    self.b = copy_bufs or [self._alloc(batch_size, BufferSpec(host=True)) for _ in range(batch_cnt)]
+    self.b_timeline, self.b_next = [0] * len(self.b), 0
+
+  def _map(self, buf:HCQBuffer) -> HCQBuffer:
+    if self.dev in buf.mapped_devs: return buf
+    if buf.owner is None: raise RuntimeError(f"map failed: buffer {buf.va_addr} has no owner, it's a virtual buffer")
+    if not hasattr(self, '_do_map'): raise NotImplementedError("map failed: no method implemented")
+
+    # Since it's unified memory space, any buffer mapping is valid for all devices after successful map.
+    # Devices can save mappings and internal metadata as a new buffer.
+    if (mb:=self._do_map(buf)) is not None: buf.mappings[self.dev] = mb
+    buf.mapped_devs.append(self.dev)
+    return buf
+
+  @suppress_finalizing
+  def _free(self, buf:HCQBuffer, options:BufferSpec|None=None):
+    for dev in buf.mapped_devs: dev.synchronize()
+    for d, mb in buf.mappings.items(): d.allocator._unmap(mb)
+    if hasattr(self, '_do_free'): self._do_free(buf, options)
+
+  def _unmap(self, mb): self.dev.iface.free(mb)
+
+  def _offset(self, buf, size:int, offset:int) -> HCQBuffer: return buf.offset(offset=offset, size=size)
+
+class HCQAllocator(HCQAllocatorBase, Generic[HCQDeviceType]):
+  def _copyin(self, dest:HCQBuffer, src:memoryview):
+    if self.dev.hw_copy_queue_t is None:
+      self.dev.synchronize()
+      with cpu_profile(f'TINY -> {self.dev.device}', f"{self.dev.device}:COPY"): ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
+      return
+
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"TINY -> {self.dev.device}", ret=src.nbytes), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for i in range(0, src.nbytes, self.b[0].size):
+        self.b_next = (self.b_next + 1) % len(self.b)
+        self.dev.timeline_signal.wait(self.b_timeline[self.b_next])
+
+        lsize = min(self.b[self.b_next].size, src.nbytes - i)
+        self.b[self.b_next].cpu_view().view(size=lsize, fmt='B')[:] = src.cast('B')[i:i+lsize]
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(dest.offset(i), self.b[self.b_next], lsize) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.b_timeline[self.b_next] = self.dev.timeline_value - 1
+
+  def copy_from_disk(self, dest:HCQBuffer, src, size):
+    def _get_temp_buf():
+      # Check if the next buffer is safe to be used (its signal has passed) and reserve it.
+      if self.b_timeline[(self.b_next + 1) % len(self.b)] <= self.dev.timeline_signal.value:
+        self.b_timeline[(self.b_next + 1) % len(self.b)], self.b_next = (1 << 64), (self.b_next + 1) % len(self.b)
+        return (self.b[self.b_next].cpu_view(), self.b_next)
+      return None
+
+    assert self.dev.hw_copy_queue_t is not None
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"DISK -> {self.dev.device}", ret=size), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for (batch_info, dst_off, src_off, copy_size) in src.device.allocator._copyout_sharded(src, size, _get_temp_buf, seg_len=self.b[0].size,
+                                                                                             use_ioring=type(self.b[0].cpu_view()) is MMIOInterface):
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(dest.offset(dst_off), self.b[batch_info[1]].offset(src_off), copy_size) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.b_timeline[batch_info[1]] = self.dev.timeline_value - 1
+
+  def _copyout(self, dest:memoryview, src:HCQBuffer):
+    self.dev.synchronize()
+    if self.dev.hw_copy_queue_t is None:
+      with cpu_profile(f'{self.dev.device} -> TINY', f"{self.dev.device}:COPY"): ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+      return
+
+    with hcq_profile(self.dev, queue_type=self.dev.hw_copy_queue_t, desc=TracingKey(f"{self.dev.device} -> TINY", ret=dest.nbytes), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      for i in range(0, dest.nbytes, cp_size:=self.b[0].size):
+        self.dev.hw_copy_queue_t().wait(self.dev.timeline_signal, self.dev.timeline_value - 1) \
+                                  .copy(self.b[0], src.offset(i), lsize:=min(cp_size, dest.nbytes-i)) \
+                                  .signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+        self.dev.timeline_signal.wait(self.dev.timeline_value - 1)
+        dest.cast('B')[i:i+lsize] = self.b[0].cpu_view().view(size=lsize, fmt='B')[:]
+
+  def _transfer(self, dest:HCQBuffer, src:HCQBuffer, sz:int, src_dev:HCQDeviceType, dest_dev:HCQDeviceType):
+    if src_dev.peer_group != dest_dev.peer_group: return src_dev.rdma_dev().allocator._transfer(dest, src, sz, src_dev, dest_dev)
+
+    cast(HCQAllocator, src_dev.allocator)._map(dest)
+
+    assert src_dev.hw_copy_queue_t is not None
+    with hcq_profile(src_dev, queue_type=src_dev.hw_copy_queue_t, desc=TracingKey(f"{src_dev.device} -> {dest_dev.device}", ret=sz), enabled=PROFILE,
+                     dev_suff="SDMA:0"):
+      src_dev.hw_copy_queue_t().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
+                               .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
+                               .copy(dest, src, sz) \
+                               .signal(src_dev.timeline_signal, src_dev.next_timeline()).submit(src_dev)
+
+    if src_dev != dest_dev:
+      unwrap(dest_dev.hw_compute_queue_t)().wait(src_dev.timeline_signal, src_dev.timeline_value - 1) \
+                                           .wait(dest_dev.timeline_signal, dest_dev.timeline_value - 1) \
+                                           .signal(dest_dev.timeline_signal, dest_dev.next_timeline()).submit(dest_dev)

extra/hcq1/ops_rdma.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
-import mmap, struct, functools
+import mmap, struct, functools, atexit
 from typing import cast
 from tinygrad.uop.ops import sint
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue, HCQBuffer, FileIOInterface
+from extra.hcq1.hcq import HCQCompiled, HCQAllocatorBase, HCQAllocator, HWQueue
+from tinygrad.runtime.support.hcq import HCQBuffer, FileIOInterface
 from tinygrad.runtime.support.system import System, PCIIfaceBase, PCIAllocationMeta
 from tinygrad.runtime.support.memory import VirtMapping, AddrSpace
 from tinygrad.runtime.support.mlx.mlxdev import MLXDev, MLXQP
@@ -103,3 +104,9 @@ class RDMADevice(HCQCompiled):
   def __init__(self, device:str=""):
     self.iface = MLXIface(self, int(device.split(":")[1]) if ":" in device else 0)
     super().__init__(device, RDMAAllocator(self), [], None, signal_t=None)
+
+@functools.cache
+def get_rdma_device(index:int) -> RDMADevice:
+  dev = RDMADevice(f"RDMA:{index}")
+  atexit.register(dev.finalize)
+  return dev

extra/hcq1/remote.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+import os, mmap, array, functools, contextlib, itertools, struct, socket, subprocess, time, enum, atexit
+from tinygrad.helpers import getenv, temp, ceildiv, unwrap, fetch, system, _ensure_downloads_dir, DEBUG, flatten
+from tinygrad.runtime.support.hcq import FileIOInterface, MMIOInterface
+from tinygrad.runtime.support.system import PCIDevice, System
+
+class RemoteCmd(enum.IntEnum):
+  PROBE,MAP_BAR,MAP_SYSMEM_FD,CFG_READ,CFG_WRITE,RESET,MMIO_READ,MMIO_WRITE,MAP_SYSMEM,SYSMEM_READ,SYSMEM_WRITE,RESIZE_BAR,PING = range(13)
+
+class RemoteMMIOInterface(MMIOInterface):
+  def __init__(self, dev:RemotePCIDevice, residx:int, nbytes:int, fmt='B', off=0, rd_cmd=RemoteCmd.MMIO_READ, wr_cmd=RemoteCmd.MMIO_WRITE):
+    self.dev, self.residx, self.nbytes, self.fmt, self.off, self.el_sz = dev, residx, nbytes, fmt, off, struct.calcsize(fmt)
+    self.rd_cmd, self.wr_cmd = rd_cmd, wr_cmd
+
+  def __getitem__(self, index):
+    sl = index if isinstance(index, slice) else slice(index, index + 1)
+    start, stop = (sl.start or 0) * self.el_sz, (sl.stop or len(self)) * self.el_sz
+    data = self.dev._bulk_read(self.rd_cmd, self.residx, self.off + start, stop - start)
+    result = data if self.fmt == 'B' else list(struct.unpack(f'<{(stop - start) // self.el_sz}{self.fmt}', data))
+    return result if isinstance(index, slice) else result[0]
+
+  def __setitem__(self, index, val):
+    start = (index.start or 0) * self.el_sz if isinstance(index, slice) else index * self.el_sz
+    data = (val if self.fmt == 'B' else struct.pack(f'<{len(val)}{self.fmt}', *val)) if isinstance(index, slice) else struct.pack(f'<{self.fmt}', val)
+    self.dev._bulk_write(self.wr_cmd, self.residx, self.off + start, data)
+
+  def view(self, offset:int=0, size:int|None=None, fmt=None):
+    return RemoteMMIOInterface(self.dev, self.residx, size or (self.nbytes - offset), fmt or self.fmt, self.off + offset, self.rd_cmd, self.wr_cmd)
+
+class RemotePCIDevice(PCIDevice):
+  _bulk_sent:int = 0
+  _bulk_recv:int = 0
+  _rpc_count:int = 0
+  _start_time:float = 0.0
+
+  @staticmethod
+  @functools.cache
+  def remote_sock(host:str, port:int) -> socket.socket:
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
+    sock.settimeout(getenv("REMOTE_TIMEOUT", 3))
+    sock.connect((host, port))
+    sock.settimeout(None)
+    if DEBUG >= 1 and RemotePCIDevice._start_time == 0.0:
+      RemotePCIDevice._start_time = time.perf_counter()
+      def _print_stats():
+        dt = time.perf_counter() - RemotePCIDevice._start_time
+        sent_mb, recv_mb = RemotePCIDevice._bulk_sent / 1e6, RemotePCIDevice._bulk_recv / 1e6
+        print(f"remote: sent {sent_mb:,.2f} MB ({sent_mb/dt:,.2f} MB/s), recv {recv_mb:,.2f} MB ({recv_mb/dt:,.2f} MB/s), "
+              f"{RemotePCIDevice._rpc_count:,} roundtrips in {dt:.2f}s")
+      atexit.register(_print_stats)
+    return sock
+
+  @staticmethod
+  @functools.cache
+  def remote_list(vendor:int, devices:tuple[tuple[int, tuple[int, ...]], ...], base_class:int|None) -> list[tuple[socket.socket, str]]:
+    payload = array.array('I', itertools.chain.from_iterable((m, d) for m, ds in devices for d in ds)).tobytes()
+    def q(r:str) -> list[tuple[socket.socket, str]]:
+      sock = RemotePCIDevice.remote_sock((host:=r.strip().split(":")[0]), (port:=int(r.strip().split(":")[1]) if ":" in r else 6667))
+      data_len, _, _, _ = RemotePCIDevice._rpc(sock, 0, RemoteCmd.PROBE, base_class or 0, len(payload), vendor, payload=payload)
+      return [(sock, f"remote:{host}:{port}:{d}") for d in RemotePCIDevice._recvall(sock, data_len).decode().split('\n')]
+    return flatten([q(r) for r in getenv("REMOTE", "").split(",") if r.strip()])
+
+  @staticmethod
+  def _recvall(sock:socket.socket, n:int) -> bytes:
+    data = b''
+    while len(data) < n and (chunk:=sock.recv(n - len(data))): data += chunk
+    if len(data) < n: raise RuntimeError("Connection closed")
+    return data
+
+  @staticmethod
+  def _rpc(sock:socket.socket, dev_id:int, cmd:int, *args:int, bar:int=0, readout_size:int=0, payload:bytes=b'', has_fd=False):
+    sock.sendall(struct.pack('<BIIQQQ', cmd, dev_id, bar, *(*args, 0, 0, 0)[:3]) + payload)
+    if has_fd:
+      msg, anc, _, _ = sock.recvmsg(17, socket.CMSG_LEN(4))
+      fd = struct.unpack('<i', anc[0][2][:4])[0]
+    else: msg, fd = RemotePCIDevice._recvall(sock, 17), None
+    if (resp:=struct.unpack('<BQQ', msg))[0] != 0:
+      raise RuntimeError(f"RPC failed: {RemotePCIDevice._recvall(sock, resp[1]).decode('utf-8') if resp[1] > 0 else 'unknown error'}")
+    RemotePCIDevice._rpc_count += 1
+    return (resp[1], resp[2]) + ((RemotePCIDevice._recvall(sock, readout_size) if readout_size > 0 else None),) + (fd,)
+
+  def __init__(self, devpref:str, pcibus:str, sock:socket.socket):
+    self.sock, self.pcibus, self.dev_id = sock, pcibus, int(pcibus.split(':')[-1]) if ':' in pcibus else 0
+    self.peer_group = sock.getpeername()[0]
+    for buft in [socket.SO_SNDBUF, socket.SO_RCVBUF]: self.sock.setsockopt(socket.SOL_SOCKET, buft, 64 << 20)
+
+    self.lock_fd = System.flock_acquire(f"{devpref.lower()}_{pcibus.lower()}.lock")
+
+  def _bulk_read(self, cmd:int, idx:int, offset:int, size:int) -> bytes:
+    RemotePCIDevice._bulk_recv += size
+    return unwrap(self._rpc(self.sock, self.dev_id, cmd, offset, size, bar=idx, readout_size=size)[2])
+  def _bulk_write(self, cmd:int, idx:int, offset:int, data:bytes):
+    RemotePCIDevice._bulk_sent += len(data)
+    self.sock.sendall(struct.pack('<BIIQQQ', cmd, self.dev_id, idx, offset, len(data), 0) + data)
+
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    paddrs_len, handle, _, _ = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM, size, int(contiguous))
+    paddrs = list(struct.unpack(f'<{paddrs_len // 8}Q', self._recvall(self.sock, paddrs_len)))
+    return RemoteMMIOInterface(self, handle, size, fmt='B', rd_cmd=RemoteCmd.SYSMEM_READ, wr_cmd=RemoteCmd.SYSMEM_WRITE), paddrs
+
+  def reset(self): self._rpc(self.sock, self.dev_id, RemoteCmd.RESET)
+  def read_config(self, offset:int, size:int): return self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_READ, offset, size)[0]
+  def write_config(self, offset:int, value:int, size:int): self._rpc(self.sock, self.dev_id, RemoteCmd.CFG_WRITE, offset, size, value)
+
+  @functools.cache
+  def bar_info(self, bar_idx:int) -> tuple[int, int]: return self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_BAR, bar=bar_idx)[:2]
+  def map_bar(self, bar:int, off:int=0, addr:int=0, size:int|None=None, fmt='B') -> MMIOInterface:
+    return RemoteMMIOInterface(self, bar, size or self.bar_info(bar)[1], fmt).view(off, size, fmt)
+  def resize_bar(self, bar_idx:int): self._rpc(self.sock, self.dev_id, RemoteCmd.RESIZE_BAR, bar=bar_idx)
+
+class APLRemotePCIDevice(RemotePCIDevice):
+  APP_PATH = "/Applications/TinyGPU.app/Contents/MacOS/TinyGPU"
+
+  @classmethod
+  def ensure_app(cls):
+    commit = "c0d024f9ff0e1dc8fdf217f255da7101d91e8323"
+    app_name = f"TinyGPU_{commit}.zip"
+    if (_ensure_downloads_dir() / app_name).is_file() and os.path.exists(cls.APP_PATH): return
+    print("Downloading TinyGPU.app...")
+    with contextlib.suppress(RuntimeError): system("pkill -f TinyGPU")
+    system(f"ditto -xk {fetch(f'https://github.com/tinygrad/tinygpu_releases/raw/{commit}/TinyGPU.zip', name=app_name)} /Applications")
+    print(system(f"{cls.APP_PATH} install"))
+
+  def __init__(self, devpref:str, pcibus:str):
+    self.ensure_app()
+    sock_path, sock = getenv("APL_REMOTE_SOCK", temp("tinygpu.sock")), socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+    for i in range(100):
+      with contextlib.suppress(ConnectionRefusedError, FileNotFoundError):
+        sock.connect(sock_path)
+        break
+      if i == 0: subprocess.Popen([self.APP_PATH, "server", sock_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+      time.sleep(0.05)
+    else: raise RuntimeError(f"Failed to connect to TinyGPU server at {sock_path}.")
+    super().__init__(devpref, "usb4", sock=sock)
+
+  def alloc_sysmem(self, size:int, vaddr:int=0, contiguous:bool=False) -> tuple[MMIOInterface, list[int]]:
+    mapped_size, _, _, fd = self._rpc(self.sock, self.dev_id, RemoteCmd.MAP_SYSMEM_FD, size, int(contiguous), has_fd=True)
+    memview = MMIOInterface(FileIOInterface(fd=fd).mmap(0, mapped_size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, 0), mapped_size, fmt='B')
+
+    # paddrs are returned as (paddr, size) pairs until a (paddr=0, size=0) terminator in the beginning of the mapping.
+    paddrs_raw = list(itertools.takewhile(lambda p: p[1] != 0, zip(memview.view(fmt='Q')[0::2], memview.view(fmt='Q')[1::2])))
+    return memview, [p + i for p, sz in paddrs_raw for i in range(0, sz, 0x1000)][:ceildiv(size, 0x1000)]

extra/hcq1/test_hcq.py

@@ -3,7 +3,8 @@ from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import mv_address, DEBUG, DEV
 from test.helpers import slow, replace_opts
 from tinygrad.device import Buffer, BufferSpec
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQBuffer
+from extra.hcq1.hcq import HCQCompiled
+from tinygrad.runtime.support.hcq import HCQBuffer
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.system import PCIIfaceBase
 from tinygrad.engine.realize import get_runtime

extra/hcq1/test_hcq_graph.py

@@ -3,8 +3,8 @@ from tinygrad import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes
-from tinygrad.runtime.graph.hcq import HCQGraph
-from tinygrad.runtime.support.hcq import HCQCompiled
+from extra.hcq1.graph import HCQGraph
+from extra.hcq1.hcq import HCQCompiled
 from tinygrad.runtime.support.usb import USBMMIOInterface
 from test.mockgpu.usb import MockUSB
 

extra/hcq2/ops_amd2.py

@@ -1,716 +0,0 @@
-from __future__ import annotations
-from typing import cast
-import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
-assert sys.platform != 'win32'
-from dataclasses import dataclass
-from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, to_name
-from tinygrad.uop.ops import sint, UOp
-from tinygrad.device import BufferSpec, Buffer
-from tinygrad.dtype import dtypes
-from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
-from tinygrad.helpers import ceildiv, unwrap, pluralize
-from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
-from tinygrad.renderer.llvmir import AMDLLVMRenderer
-from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
-from tinygrad.runtime.autogen.am import am
-from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
-from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
-from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
-from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
-from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
-from tinygrad.runtime.ops_amd import SQTT, PMC
-from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
-if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401 # pylint: disable=unused-import
-
-from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
-from tinygrad.uop.ops import Ops, UPat, PatternMatcher
-
-# *****************
-# PM4
-
-def _queue_args(hq:HWQueue, q) -> list[UOp]: # the ring and its pointers, tagged {name}_{queue} like the device's bufferize rules
-  shapes = [("ring", (q.ring.size,), q.ring.dtype)] + [(n, (1,), dtypes.uint64) for n in ("write_ptr", "doorbell", "put_value")]
-  return [UOp.placeholder(s, d, 0, device=hq.devs, volatile=True, tag=to_name(n, hq.queue)) for n, s, d in shapes]
-
-def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
-
-class AMDComputeQueue(HWQueue):
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
-    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))),
-     lambda ctx, dst, val: ctx.signal(dst, val)),
-  ])
-
-  def __init__(self, ctx, submit):
-    super().__init__(ctx, submit)
-    self.pm4, self.gc, self.soc, self.nbio, self.target = self.dev.pm4, self.dev.gc, self.dev.soc, self.dev.nbio, self.dev.target
-
-  def pkt3(self, cmd, *vals): self.q(self.pm4.PACKET3(cmd, _dw(vals) - 1), *vals)
-
-  def wreg(self, reg:AMDReg, *args:sint, **kwargs:int):
-    if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
-    if self.pm4.PACKET3_SET_SH_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_SH_REG_END:
-      set_packet, set_packet_start = self.pm4.PACKET3_SET_SH_REG, self.pm4.PACKET3_SET_SH_REG_START
-    elif self.pm4.PACKET3_SET_UCONFIG_REG_START <= reg.addr[0] < self.pm4.PACKET3_SET_UCONFIG_REG_START + 2**16-1:
-      set_packet, set_packet_start = self.pm4.PACKET3_SET_UCONFIG_REG, self.pm4.PACKET3_SET_UCONFIG_REG_START
-    else: raise RuntimeError(f'Cannot set {reg.name} ({reg.addr[0]}) via pm4 packet')
-    self.pkt3(set_packet, reg.addr[0] - set_packet_start, *(args or (reg.encode(**kwargs),)))
-
-  def wait_reg_mem(self, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
-    wrm_info_dw = self.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | self.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
-                | self.pm4.WAIT_REG_MEM_FUNCTION(op) | self.pm4.WAIT_REG_MEM_ENGINE(0)
-    self.pkt3(self.pm4.PACKET3_WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4)
-
-  def acquire_mem(self, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
-    if self.target[0] != 9:
-      cache_flags_dw = self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLI_INV(gli) \
-                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_INV(glm) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLM_WB(glm) \
-                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_INV(glk) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLK_WB(glk) \
-                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GLV_INV(glv) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL1_INV(gl1) \
-                     | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_INV(gl2) | self.pm4.PACKET3_ACQUIRE_MEM_GCR_CNTL_GL2_WB(gl2)
-      return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, 0, *data64_le(sz), *data64_le(addr), 0, cache_flags_dw)
-    cp_coher_cntl = self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_ICACHE_ACTION_ENA(gli) | \
-                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_SH_KCACHE_ACTION_ENA(glk) | \
-                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_ACTION_ENA(gl2) | \
-                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TCL1_ACTION_ENA(gl1) | \
-                    self.pm4.PACKET3_ACQUIRE_MEM_CP_COHER_CNTL_TC_WB_ACTION_ENA(gl2)
-    return self.pkt3(self.pm4.PACKET3_ACQUIRE_MEM, cp_coher_cntl, *data64_le(sz), *data64_le(addr), 0x0000000A)
-
-  def release_mem(self, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache_flush=False):
-    if self.target[0] != 9:
-      cache_flags_dw = 0 if not cache_flush else (self.pm4.PACKET3_RELEASE_MEM_GCR_GLV_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL1_INV \
-                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_WB \
-                     | self.pm4.PACKET3_RELEASE_MEM_GCR_GLM_INV | self.pm4.PACKET3_RELEASE_MEM_GCR_GL2_WB | self.pm4.PACKET3_RELEASE_MEM_GCR_SEQ)
-      event_dw = self.pm4.PACKET3_RELEASE_MEM_EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) \
-               | self.pm4.PACKET3_RELEASE_MEM_EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
-      memsel_dw = self.pm4.PACKET3_RELEASE_MEM_DATA_SEL(data_sel) | self.pm4.PACKET3_RELEASE_MEM_INT_SEL(int_sel) \
-                | self.pm4.PACKET3_RELEASE_MEM_DST_SEL(0)
-    else:
-      cache_flags_dw = 0 if not cache_flush else (self.pm4.EOP_TC_WB_ACTION_EN | self.pm4.EOP_TC_NC_ACTION_EN)
-      event_dw = self.pm4.EVENT_TYPE(self.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | \
-                 self.pm4.EVENT_INDEX(self.pm4.event_index__mec_release_mem__end_of_pipe)
-      memsel_dw = self.pm4.DATA_SEL(data_sel) | self.pm4.INT_SEL(int_sel)
-      ctxid = 0
-    addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64)
-    val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64)
-    self.pkt3(self.pm4.PACKET3_RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid)
-
-  def memory_barrier(self):
-    pf = '' if self.nbio.version[0] == 2 else '0' if self.nbio.version[:2] != (7, 11) else '1'
-    self.wait_reg_mem(reg=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
-                      reg_done=getattr(self.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff)
-    self.acquire_mem()
-
-  def exec(self, call:UOp, prg:UOp):
-    data, lib = amd_build_program(self.dev, prg, self.devs)
-    info = prg.arg
-
-    # kernargs: a nested blob linear inside a getaddr, packed into the tail of the cmdbuf
-    ka_words = [get_call_arg_uops(call)[gi].getaddr(self.devs) for gi in info.globals] + \
-               [b.ccast(v.dtype) for v, b in zip(info.vars, get_call_var_uops(call, prg))] # a bound value is a bare const, the var has the width
-    pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
-    assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
-    ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4))
-
-    prog_addr = lib.getaddr(self.devs) + data.entry_point_offset
-    scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=self.devs).rtag("scratch").getaddr(self.devs)
-    args_addr = ka.getaddr(self.devs)
-
-    user_regs:list = []
-    if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000]
-    if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size]
-    user_regs += [args_addr]
-
-    dispatch_init = self.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
-      **({'cs_w32_en': int(data.wave32)} if self.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
-    self.acquire_mem(gli=0, gl2=0)
-    self.wreg(self.gc.regCOMPUTE_PGM_LO, prog_addr >> 8)
-    self.wreg(self.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2)
-    self.wreg(self.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3)
-    self.wreg(self.gc.regCOMPUTE_TMPRING_SIZE, self.dev.tmpring_size(data.private_segment_size))
-    for xcc_id in range(self.dev.xccs):
-      self.wreg(self.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // self.dev.xccs * xcc_id) >> 8)
-    self.wreg(self.gc.regCOMPUTE_RESTART_X, 0, 0, 0)
-    self.wreg(self.gc.regCOMPUTE_USER_DATA_0, *user_regs)
-    self.wreg(self.gc.regCOMPUTE_RESOURCE_LIMITS, self.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH")))
-    self.wreg(self.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0)
-    self.pkt3(self.pm4.PACKET3_DISPATCH_DIRECT, *info.global_size, dispatch_init)
-    self.pkt3(self.pm4.PACKET3_EVENT_WRITE, self.pm4.EVENT_TYPE(self.soc.CS_PARTIAL_FLUSH) | self.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))
-
-  def wait(self, signal:UOp, value:UOp): self.wait_reg_mem(value.cast(dtypes.uint32), mem=signal.getaddr(self.devs))
-
-  def timestamp(self, signal:UOp):
-    self.release_mem(signal.getaddr(self.devs), 0, self.pm4.data_sel__mec_release_mem__send_gpu_clock_counter,
-                     self.pm4.int_sel__mec_release_mem__none)
-
-  def signal(self, signal:UOp, value:UOp):
-    self.release_mem(signal.getaddr(self.devs), value, self.pm4.data_sel__mec_release_mem__send_32_bit_low,
-                     self.pm4.int_sel__mec_release_mem__send_interrupt_after_write_confirm, cache_flush=True)
-
-  def submit(self, cmdbuf:UOp) -> UOp:
-    q = self.dev.compute_queue
-
-    ring, wptr, doorbell, put = _queue_args(self, q)
-
-    size_dw = cmdbuf.max_numel() // 4
-    p = put.index(0).load()
-    i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf,))
-    copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
-    next_put = p + size_dw
-    flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
-    return doorbell.after(flush).index(0).store(next_put)
-
-# *****************
-# SDMA
-
-class AMDSDMAQueue(HWQueue):
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), lambda ctx, call: ctx.copy(call)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ()),
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
-    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))),
-     lambda ctx, dst, val: ctx.signal(dst, val)),
-  ])
-
-  def __init__(self, ctx, submit):
-    super().__init__(ctx, submit)
-    self.sdma, self.target, self.max_copy_size = self.dev.sdma, self.dev.target, self.dev.max_copy_size
-
-  def copy(self, call:UOp):
-    sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
-    hdr = self.sdma.SDMA_OP_COPY | self.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_COPY_LINEAR)
-    for off in range(0, sz, self.max_copy_size):
-      self.q(hdr, self.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, self.max_copy_size)-1), 0,
-             *(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(self.devs), call.src[1].getaddr(self.devs))))
-
-  def wait(self, signal:UOp, value:UOp):
-    op = self.sdma.SDMA_OP_POLL_REGMEM | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
-       | self.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
-    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), 0xffffffff,
-           self.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | self.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))
-
-  def timestamp(self, signal:UOp):
-    self.q(self.sdma.SDMA_OP_TIMESTAMP | self.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(self.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL),
-           signal.getaddr(self.devs))
-
-  def signal(self, signal:UOp, value:UOp): # a fence packet then a trap
-    op = self.sdma.SDMA_OP_FENCE | (self.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if self.target[0] != 9 else 0)
-    self.q(op, signal.getaddr(self.devs), value.cast(dtypes.uint32), self.sdma.SDMA_OP_TRAP, 0)
-
-  def submit(self, cmdbuf:UOp) -> UOp:
-    # sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail
-    q = unwrap(self.dev.sdma_queue(int(self.queue.split(":")[1])))
-
-    ring, wptr, doorbell, put = _queue_args(self, q)
-
-    rs, size_dw = q.ring.size, cmdbuf.max_numel() // 4
-    put_b = put.index(0).load()
-    tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
-    fits = (size_dw <= rs - tail).cast(dtypes.int)
-    start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
-    zi = UOp.range(zero_amt, 10, dtype=dtypes.int, src=(cmdbuf,))
-    zero_tail = ring.index(tail + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
-    i = UOp.range(size_dw, 11, dtype=dtypes.int, src=(cmdbuf,))
-    copy = ring.index(start_dw + i).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
-    next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype)
-    flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
-    return doorbell.after(flush).index(0).store(next_put)
-
-@dataclass(frozen=True)
-class AMDProgramData:
-  entry_point_offset:int; rsrc1:int; rsrc2:int; rsrc3:int; wave32:bool
-  private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
-  enable_dispatch_ptr:int; enable_private_segment_sgpr:int
-
-_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
-def amd_build_program(dev, prg:UOp, devs:tuple[str, ...]) -> tuple[AMDProgramData, UOp]:
-  # the image parses once per lib, each device set gets its own program buffer of it
-  if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, devs))) is None:
-    data, image = _amd_program_image(dev, lib)
-    buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
-    cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
-  return cached
-
-@functools.cache
-def _amd_program_image(dev, lib:bytes) -> tuple[AMDProgramData, bytes]:
-  image, sections, relocs = elf_loader(lib)
-  rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
-  for off, sym, typ, addent in relocs:
-    assert typ == 5, f"unknown AMD reloc {typ}"  # R_AMDGPU_REL64
-    image[off:off+8] = struct.pack('<q', sym - off + addent)
-  desc = amdgpu_kd.llvm_amdhsa_kernel_descriptor_t.from_buffer_copy(bytes(image[rodata:rodata+ctypes.sizeof(amdgpu_kd.llvm_amdhsa_kernel_descriptor_t)]))
-  if (lds:=((desc.group_segment_fixed_size+511)//512)&0x1FF) > (dev.iface.props['lds_size_in_kb']*1024)//512:
-    raise RuntimeError("Too many resources requested: group_segment_size")
-  edp = desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_DISPATCH_PTR
-
-  data = AMDProgramData(entry_point_offset=rodata + desc.kernel_code_entry_byte_offset,
-    rsrc1=desc.compute_pgm_rsrc1 | ((1<<20) if dev.target[0]==11 else 0),  # priv=1 on gfx11 for cwsr
-    rsrc2=desc.compute_pgm_rsrc2 | (lds<<15), rsrc3=desc.compute_pgm_rsrc3,
-    wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
-    kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
-    enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
-  return data, bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
-
-class AMDAllocator(HCQAllocator['AMDDevice']):
-  def __init__(self, dev:AMDDevice):
-    super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
-
-  def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
-    return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
-
-  def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
-
-  def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
-
-  def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
-
-@dataclass
-class AMDQueueDesc:
-  ring: Buffer; read_ptr: Buffer; write_ptr: Buffer; doorbell: Buffer; put_value: Buffer  # noqa: E702
-  eop_buffer: Buffer|None = None; cwsr_buffer: Buffer|None = None; params: tuple|None = None  # noqa: E702
-
-class KFDIface:
-  kfd:FileIOInterface|None = None
-  event_page:HCQBuffer|None = None
-  gpus:list[FileIOInterface] = []
-  count:int = 0
-
-  def _is_usable_gpu(self, gpu_id):
-    with contextlib.suppress(OSError): return int(gpu_id.read()) != 0
-    return False
-
-  def __init__(self, dev, device_id):
-    self.dev = dev
-
-    kfd_topo_path = "/sys/devices/virtual/kfd/kfd/topology/nodes"
-
-    # Initialize KFD interface during first run
-    if KFDIface.kfd is None:
-      KFDIface.kfd = FileIOInterface("/dev/kfd", os.O_RDWR)
-      gpus = [g for g in FileIOInterface(kfd_topo_path).listdir() if self._is_usable_gpu(FileIOInterface(f"{kfd_topo_path}/{g}/gpu_id"))]
-      KFDIface.gpus = hcq_filter_visible_devices(sorted(gpus, key=lambda x: int(x.split('/')[-1])), "AMD")
-      KFDIface.count = len(KFDIface.gpus)
-
-    if device_id >= len(KFDIface.gpus): raise RuntimeError(f"No device found for {device_id}. Requesting more devices than the system has?")
-
-    self.gpu_id = int(FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/gpu_id").read())
-    self.props = {(p:=l.split())[0]: int(p[1]) for l in FileIOInterface(f"{kfd_topo_path}/{KFDIface.gpus[device_id]}/properties").read().splitlines()}
-    self.dev_sysfs_path = f"/sys/class/drm/renderD{self.props['drm_render_minor']}/device"
-    ip_base = f"{self.dev_sysfs_path}/ip_discovery/die/0"
-    id2ip = {am.GC_HWID: am.GC_HWIP, am.SDMA0_HWID: am.SDMA0_HWIP, am.NBIF_HWID: am.NBIF_HWIP}
-    ip_hw = [(id2ip[int(hwid)], int(hwid)) for hwid in FileIOInterface(ip_base).listdir() if hwid.isnumeric() and int(hwid) in id2ip]
-    self.ip_versions = {ip:tuple(int(FileIOInterface(f'{ip_base}/{hw}/0/{part}').read()) for part in ['major','minor','revision']) for ip,hw in ip_hw}
-    self.drm_fd = FileIOInterface(f"/dev/dri/renderD{self.props['drm_render_minor']}", os.O_RDWR)
-
-    self.kfd_ver = ((ver_st:=kfd.AMDKFD_IOC_GET_VERSION(KFDIface.kfd)).major_version, ver_st.minor_version)
-    kfd.AMDKFD_IOC_ACQUIRE_VM(KFDIface.kfd, drm_fd=self.drm_fd.fd, gpu_id=self.gpu_id)
-    if self.kfd_ver >= (1,14): kfd.AMDKFD_IOC_RUNTIME_ENABLE(KFDIface.kfd, mode_mask=0)
-
-    # Set these for our device.
-    if KFDIface.event_page is None:
-      KFDIface.event_page = self.alloc(0x8000, uncached=True)
-      kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_page_offset=KFDIface.event_page.meta.handle)
-    else: self.map(KFDIface.event_page)
-
-    # Event to wait for queues completion
-    self.dev.queue_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_SIGNAL, auto_reset=1)
-    self.dev.queue_event_mailbox_ptr = KFDIface.event_page.va_addr + self.dev.queue_event.event_slot_index * 8
-
-    # OS events to collect memory and hardware faults
-    self.mem_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_MEMORY)
-    self.hw_fault_event = kfd.AMDKFD_IOC_CREATE_EVENT(KFDIface.kfd, event_type=kfd.KFD_IOC_EVENT_HW_EXCEPTION)
-
-    self.queue_event_arr = (kfd.struct_kfd_event_data * 3)(kfd.struct_kfd_event_data(event_id=self.dev.queue_event.event_id),
-      kfd.struct_kfd_event_data(event_id=self.mem_fault_event.event_id), kfd.struct_kfd_event_data(event_id=self.hw_fault_event.event_id))
-    self.queue_event_arr_ptr = ctypes.addressof(self.queue_event_arr)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, cpu_addr=None) -> HCQBuffer:
-    flags = kfd.KFD_IOC_ALLOC_MEM_FLAGS_WRITABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_EXECUTABLE | kfd.KFD_IOC_ALLOC_MEM_FLAGS_NO_SUBSTITUTE
-
-    if uncached: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED | kfd.KFD_IOC_ALLOC_MEM_FLAGS_GTT
-    else: flags |= (kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR if host else kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM)
-
-    # Make mapped cpu address to be uncachable
-    if cpu_addr is not None: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_COHERENT | kfd.KFD_IOC_ALLOC_MEM_FLAGS_UNCACHED
-
-    if cpu_access or host: flags |= kfd.KFD_IOC_ALLOC_MEM_FLAGS_PUBLIC
-
-    if flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR:
-      buf = addr = cpu_addr or FileIOInterface.anon_mmap(0, size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | mmap.MAP_ANONYMOUS, 0)
-    else: buf, addr = 0, FileIOInterface.anon_mmap(0, size, 0, mmap.MAP_PRIVATE | mmap.MAP_ANONYMOUS | MAP_NORESERVE, 0)
-
-    try: mem = kfd.AMDKFD_IOC_ALLOC_MEMORY_OF_GPU(self.kfd, va_addr=addr, size=size, gpu_id=self.gpu_id, flags=flags, mmap_offset=buf)
-    except OSError as e:
-      if e.errno == errno.EINVAL and (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_VRAM) and cpu_access:
-        raise MemoryError("Cannot allocate host-visible VRAM. Ensure the resizable BAR option is enabled on your system.") from e
-      if e.errno == errno.ENOMEM: raise MemoryError(f"Cannot allocate {size} bytes: no memory is available.") from e
-      raise
-
-    if not (flags & kfd.KFD_IOC_ALLOC_MEM_FLAGS_USERPTR):
-      buf = self.drm_fd.mmap(mem.va_addr, mem.size, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED | MAP_FIXED, mem.mmap_offset)
-      assert addr == buf == mem.va_addr
-
-    view = MMIOInterface(mem.va_addr, mem.size, fmt='B') if cpu_access or host else None
-    self.map(hcqbuf:=HCQBuffer(mem.va_addr, mem.size, meta=mem, view=view, owner=self.dev))
-    return hcqbuf
-
-  def free(self, mem):
-    self._unmap(mem)
-    if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
-    kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def unmap(self, mem):
-    self._unmap(mem)
-    if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
-
-  def _unmap(self, mem):
-    gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
-    assert stm.n_success == 1
-
-  def map(self, mem):
-    if mem.owner is not None and mem.owner._is_cpu():
-      mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
-      mapped._owns_kfd_handle = True
-      return mapped
-
-    c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
-    stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
-    assert stm.n_success == 1
-    return HCQBuffer(mem.va_addr, mem.size, meta=mem.meta, owner=mem.owner)
-
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
-                   xcc_id=0, idx=0):
-    queue = kfd.AMDKFD_IOC_CREATE_QUEUE(KFDIface.kfd, ring_base_address=ring._buf.va_addr, ring_size=ring._buf.size, gpu_id=self.gpu_id,
-      queue_type=queue_type, queue_percentage=kfd.KFD_MAX_QUEUE_PERCENTAGE|(xcc_id<<8), queue_priority=getenv("AMD_KFD_QUEUE_PRIORITY", 7),
-      eop_buffer_address=eop_buffer._buf.va_addr if eop_buffer else 0, eop_buffer_size=eop_buffer._buf.size if eop_buffer else 0,
-      ctl_stack_size=ctl_stack_size, ctx_save_restore_address=cwsr_buffer._buf.va_addr if cwsr_buffer else 0, ctx_save_restore_size=ctx_save_restore_size,
-      write_pointer_address=gart._buf.va_addr+wptr, read_pointer_address=gart._buf.va_addr+rptr+8*xcc_id)
-
-    if not hasattr(self, 'doorbells'):
-      self.doorbells_base = queue.doorbell_offset & (~0x1fff) # doorbell is two pages
-      self.doorbells = cast(FileIOInterface, KFDIface.kfd).mmap(0, 0x2000, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, self.doorbells_base)
-
-    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
-    doorbell = Buffer("CPU", 1, dtypes.uint64,
-      options=BufferSpec(external_ptr=self.doorbells + queue.doorbell_offset - self.doorbells_base), preallocate=True)
-    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr+8*xcc_id).ensure_allocated(),
-      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer)
-
-  def sleep(self, tm:int):
-    kfd.AMDKFD_IOC_WAIT_EVENTS(KFDIface.kfd, events_ptr=self.queue_event_arr_ptr, num_events=3, wait_for_all=0, timeout=tm)
-    if self.queue_event_arr[1].memory_exception_data.gpu_id or self.queue_event_arr[2].hw_exception_data.gpu_id: self.on_device_hang()
-
-  def on_device_hang(self):
-    def _str(st): return ' '.join(f'{k[0]}={getattr(st, k[0])}' for k in st._real_fields_)
-
-    # try to collect fault info if not already set from sleep().
-    if not self.queue_event_arr[1].memory_exception_data.gpu_id and not self.queue_event_arr[2].hw_exception_data.gpu_id:
-      with contextlib.suppress(RuntimeError): self.sleep(tm=1)
-
-    report = []
-    if self.queue_event_arr[1].memory_exception_data.gpu_id:
-      report += [f"MMU fault: 0x{self.queue_event_arr[1].memory_exception_data.va:X} | {_str(self.queue_event_arr[1].memory_exception_data.failure)}"]
-    if self.queue_event_arr[2].hw_exception_data.gpu_id: report += [f"HW fault: {_str(self.queue_event_arr[2].hw_exception_data)}"]
-
-    raise RuntimeError("\n".join(report))
-
-  def require_profile_mode(self, can_set_mode=True):
-    if self.dev.target[0] == 9: return
-    fn = f'{self.dev_sysfs_path}/power_dpm_force_performance_level'
-    if (perflevel:=FileIOInterface(fn).read().strip()) != 'profile_standard':
-      if can_set_mode:
-        atexit.register(lambda: os.system(f"echo '{perflevel}' | sudo tee {fn} > /dev/null"))
-        os.system(f"echo 'profile_standard' | sudo tee {fn} > /dev/null")
-        self.require_profile_mode(can_set_mode=False)
-      else:
-        raise RuntimeError("PMC/SQTT requires stable power state: run `amd-smi set -l stable_std` for KFD iface")
-
-  @functools.cached_property
-  def drm_dev_info(self) -> amdgpu_drm.struct_drm_amdgpu_info_device:
-    amdgpu_drm.DRM_IOCTL_AMDGPU_INFO(self.drm_fd, query=amdgpu_drm.AMDGPU_INFO_DEV_INFO,
-      return_pointer=ctypes.addressof(inf:=amdgpu_drm.struct_drm_amdgpu_info_device()), return_size=ctypes.sizeof(inf))
-    return inf
-  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return ((self.drm_dev_info.cu_bitmap[se % 4][sa + (se // 4) * 2] >> (2 * wgp)) & 0x3) == 0x3
-
-class PCIIface(PCIIfaceBase):
-  def __init__(self, dev, dev_id):
-    super().__init__(dev, dev_id, vendor=0x1002, devices=((0xffff, (0x74a1,0x744c,0x7480,0x7550,0x7551,0x7590,0x75a0)),), vram_bar=0,
-      va_start=AMMemoryManager.va_allocator.base, va_size=AMMemoryManager.va_allocator.size, dev_impl_t=AMDev)
-    self._compute_props()
-
-  def p2p_paddrs(self, paddrs:list[tuple[int,int]]) -> tuple[list[tuple[int,int]], AddrSpace]:
-    return ([(self.dev_impl.paddr2xgmi(p), sz) for p, sz in paddrs], AddrSpace.PEER) if self.dev_impl.is_hive() else super().p2p_paddrs(paddrs)
-
-  def require_profile_mode(self): return True
-  def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
-  def unmap(self, mem): self.free(mem)
-
-  def _compute_props(self):
-    self.ip_versions = self.dev_impl.ip_ver
-
-    gfxver = int(f"{self.dev_impl.ip_ver[am.GC_HWIP][0]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][1]:02d}{self.dev_impl.ip_ver[am.GC_HWIP][2]:02d}")
-    if self.dev_impl.gc_info.header.version_major == 2:
-      cu_per_sa = self.dev_impl.gc_info.gc_num_cu_per_sh
-      max_sh_per_se = self.dev_impl.gc_info.gc_num_sh_per_se
-    else:
-      cu_per_sa = 2 * (self.dev_impl.gc_info.gc_num_wgp0_per_sa + self.dev_impl.gc_info.gc_num_wgp1_per_sa)
-      max_sh_per_se = self.dev_impl.gc_info.gc_num_sa_per_se
-
-    array_count = max_sh_per_se * self.dev_impl.gc_info.gc_num_se * self.dev_impl.gfx.xccs
-    self.props = {'cu_per_simd_array': cu_per_sa, 'simd_count': 2 * cu_per_sa * array_count, 'simd_per_cu': 2, 'array_count': array_count,
-      'max_slots_scratch_cu': self.dev_impl.gc_info.gc_max_scratch_slots_per_cu, 'max_waves_per_simd': self.dev_impl.gc_info.gc_max_waves_per_simd,
-      'simd_arrays_per_engine': max_sh_per_se, 'lds_size_in_kb': self.dev_impl.gc_info.gc_lds_size, 'num_xcc': self.dev_impl.gfx.xccs,
-      'gfx_target_version': {90403: 90402}.get(gfxver, gfxver)}
-
-  def create_queue(self, queue_type, ring, gart, rptr, wptr, eop_buffer=None, cwsr_buffer=None, ctl_stack_size=0, ctx_save_restore_size=0,
-                   xcc_id=0, idx=0):
-    assert cwsr_buffer is None, "no cwsr buffer for am"
-
-    rcvr_params: tuple
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA:
-      doorbell_index = self.dev_impl.sdma.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
-        gart._buf.va_addr+wptr, idx)))
-    else:
-      doorbell_index = self.dev_impl.gfx.setup_ring(*(rcvr_params:=(ring._buf.va_addr, ring._buf.size, gart._buf.va_addr+rptr,
-        gart._buf.va_addr+wptr, eop_buffer._buf.va_addr, eop_buffer._buf.size, is_aql:=(queue_type==kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL), is_aql)))
-
-    (put_value := Buffer("CPU", 1, dtypes.uint64, preallocate=True))._buf.view.view(fmt='Q')[0] = 0
-    doorbell = Buffer("CPU", 1, dtypes.uint64, options=BufferSpec(external_ptr=self.dev_impl.doorbell64.addr + doorbell_index*8), preallocate=True)
-    return AMDQueueDesc(ring=ring, doorbell=doorbell, read_ptr=gart.view(1, dtypes.uint64, rptr).ensure_allocated(),
-      write_ptr=gart.view(1, dtypes.uint64, wptr).ensure_allocated(), put_value=put_value, eop_buffer=eop_buffer, params=rcvr_params)
-
-  def _collect_interrupts(self, reset=False, drain_only=False):
-    d = self.dev
-    if drain_only: d.iface.dev_impl.ih.drain()
-    else: d.iface.dev_impl.ih.interrupt_handler()
-
-    if reset and d.iface.dev_impl.recover(force=True):
-      cq = d.compute_queue
-      for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
-      d.iface.dev_impl.gfx.setup_ring(*cq.params)
-      (tl:=d.timeline._buf.cpu_view().view(fmt='Q'))[0] = tl[1]
-
-  def sleep(self, timeout):
-    if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
-      self.pci_dev.irq_fd.read(8 * events_cnt)
-    self._collect_interrupts()
-    if self.dev_impl.is_err_state: raise RuntimeError("Device is in error state")
-
-  def on_device_hang(self):
-    self._collect_interrupts(reset=True)
-    raise RuntimeError("Device hang detected")
-
-  def device_fini(self): self.dev_impl.fini()
-
-class USBIface(PCIIface):
-  def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
-    if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
-      raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
-    self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
-    self.dev_impl = AMDev(self.pci_dev)
-    self._compute_props()
-    self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
-    self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
-    self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
-
-  def _dma_region(self, ctrl_addr, sys_addr, size):
-    region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
-    return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
-
-  def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
-    # everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
-    return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
-
-  def sleep(self, timeout): pass
-
-  # we don't own the sram region, so the buffer never frees it
-  @functools.cached_property
-  def usb_sram(self) -> Buffer:
-    return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
-
-def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
-
-class AMDDevice(HCQ2Compiled):
-  timestamp_divider = 100.0  # AMD GPU clock: ticks/us
-  max_scratch_psize = 0
-  pm_encode = PatternMatcher([
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_compute", name="submit"), lambda ctx, submit: encode_submit(AMDComputeQueue(ctx, submit))),
-    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_amd_copy", name="submit"), lambda ctx, submit: encode_submit(AMDSDMAQueue(ctx, submit))),
-  ])
-
-  ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
-
-  def device_props(self): return self.iface.props
-
-  def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
-
-  def __init__(self, device:str=""):
-    self.iface = self._select_iface(device)
-    self.is_usb = isinstance(self.iface, USBIface)
-    if self.is_usb: self.rt_nbytes = 4 << 20
-
-    self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
-    self.arch = "gfx%d%x%x" % self.target
-    assert (self.target in ((9,4,2),(9,5,0))) or self.target[0] in (11, 12), f"Unsupported arch: {self.arch}"
-    if DEBUG >= 1: print(f"AMDDevice: opening {self.device_id} with target {self.target} arch {self.arch}")
-
-    self.xccs = self.iface.props.get('num_xcc', 1)
-    self.se_cnt = self.iface.props['array_count'] // self.iface.props['simd_arrays_per_engine'] // self.xccs
-    self.cu_cnt = self.iface.props['simd_count'] // self.iface.props['simd_per_cu'] // self.xccs
-    self.waves_per_cu = self.iface.props['max_waves_per_simd'] * self.iface.props['simd_per_cu']
-    self.wave_cnt = (self.cu_cnt * self.waves_per_cu) if self.target[0] != 9 else min(self.cu_cnt * 40, self.se_cnt * self.xccs * 512)
-
-    self.ip_off = importlib.import_module(f"tinygrad.runtime.autogen.am.{'vega' if self.target[0] == 9 else 'navi'}_offsets")
-    self.soc = import_soc(self.target)
-    self.pm4 = importlib.import_module(f"tinygrad.runtime.autogen.am.pm4_{'soc15' if self.target[0] == 9 else 'nv'}")
-    self.sdma = import_module('sdma', min(self.iface.ip_versions[am.SDMA0_HWIP], (6, 0, 0)))
-    self.gc = AMDIP('gc', self.iface.ip_versions[am.GC_HWIP],
-                    bases={i: tuple(getattr(self.ip_off, f'GC_BASE__INST{i}_SEG{s}', 0) for s in range(6)) for i in range(6)})
-
-    self.nbio = AMDIP('nbio' if self.target[0] < 12 else 'nbif', self.iface.ip_versions[am.NBIF_HWIP],
-                      bases={i: tuple(getattr(self.ip_off, f'NBIO_BASE__INST{i}_SEG{s}', 0) for s in range(9)) for i in range(6)})
-
-    self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
-    if self.is_aql:
-      self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
-      self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
-
-    self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
-    self.sdma_queues:dict = {}
-    self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
-
-    super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
-
-    # Scratch setup
-    self.max_private_segment_size = 0
-    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx.scratch_buffer(b.max_numel()))]) + self.pm_bufferize
-
-    if self.is_usb:
-      self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
-      raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
-
-    self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
-    if self.pmc_enabled:
-      self.iface.require_profile_mode()
-
-      self.pmc_sched:list[PMCSample] = []
-      self.pmc_counters = import_pmc(self.target)
-
-      # validate counters: SQ for SIMD busy/instruction counts, LDS stats, GRBM for GPU cycles, L2 cache hits/misses
-      l2, lds = ("TCC", "SQ") if self.target[0] == 9 else ("GL2C", "SQC")
-      pmc_default = f"SQ_BUSY_CYCLES,SQ_INSTS_VALU,SQ_INSTS_SALU,{lds}_LDS_IDX_ACTIVE,{lds}_LDS_BANK_CONFLICT,GRBM_GUI_ACTIVE,{l2}_HIT,{l2}_MISS"
-      for k in (PMC_COUNTERS:=getenv("PMC_COUNTERS", pmc_default).split(",")):
-        if k not in self.pmc_counters: raise RuntimeError(f"PMC counter {k} is not supported. Available: {','.join(self.pmc_counters.keys())}")
-
-      raise NotImplementedError("PMC start not migrated to hcq2 yet")
-
-    # SQTT is disabled by default because of runtime overhead and big file sizes (~200mb to Tensor.full() two 4096x4096 tensors and matmul them)
-    self.sqtt_enabled:bool = PROFILE > 0 and SQTT > 0
-    if self.sqtt_enabled:
-      self.iface.require_profile_mode()
-
-      SQTT_BUFFER_SIZE = getenv("SQTT_BUFFER_SIZE", 256) # in mb, per shader engine
-      self.sqtt_buffers = [self.allocator.alloc(SQTT_BUFFER_SIZE<<20, BufferSpec(nolru=True, uncached=True)) for _ in range(self.se_cnt * self.xccs)]
-      self.sqtt_wptrs = self.allocator.alloc(round_up(self.se_cnt * self.xccs * 4, 0x1000), BufferSpec(cpu_access=True, nolru=True))
-      self.sqtt_next_cmd_id = itertools.count(0)
-
-  def create_queue(self, queue_type, ring_size, ctx_save_restore_size=0, eop_buffer_size=0, ctl_stack_size=0, debug_memory_size=0, idx=0):
-    ring = Buffer(self.device, ring_size // 4, dtypes.uint32, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
-    gart = Buffer(self.device, 0x100, dtypes.uint8, options=BufferSpec(uncached=True, cpu_access=True), preallocate=True)
-
-    if queue_type == kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL:
-      self.aql_gart = gart
-      self.aql_desc = hsa.amd_queue_t(queue_properties=hsa.AMD_QUEUE_PROPERTIES_IS_PTR64 | hsa.AMD_QUEUE_PROPERTIES_ENABLE_PROFILING,
-        read_dispatch_id_field_base_byte_offset=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-        max_cu_id=(self.cu_cnt * self.xccs) - 1, max_wave_id=self.waves_per_cu - 1)
-      self.aql_gart._buf.cpu_view().view(fmt='B')[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    cwsr_buffer_size = round_up((ctx_save_restore_size + debug_memory_size) * self.xccs, mmap.PAGESIZE)
-    cwsr_buffer = Buffer(self.device, cwsr_buffer_size, dtypes.uint8, preallocate=True) if ctx_save_restore_size else None
-    eop_buffer = Buffer(self.device, eop_buffer_size, dtypes.uint8, preallocate=True) if eop_buffer_size else None
-
-    queue = (self.iface.create_queue(queue_type, ring, gart, rptr=getattr(hsa.amd_queue_t, 'read_dispatch_id').offset,
-             wptr=getattr(hsa.amd_queue_t, 'write_dispatch_id').offset, eop_buffer=eop_buffer, cwsr_buffer=cwsr_buffer,
-             ctx_save_restore_size=ctx_save_restore_size, ctl_stack_size=ctl_stack_size, idx=idx))
-
-    qname = f"{'COPY' if queue_type == kfd.KFD_IOC_QUEUE_TYPE_SDMA else 'COMPUTE'}:{idx}"
-    self.pm_bufferize = PatternMatcher([
-      (UPat(Ops.PARAM, tag=to_name(name, qname)), lambda ctx, b=getattr(queue, name): b) for name in ["ring", "write_ptr", "doorbell", "put_value"]
-    ]) + self.pm_bufferize
-
-    return queue
-
-  @functools.cached_property
-  def compute_queue(self) -> AMDQueueDesc:
-    # https://gitlab.freedesktop.org/agd5f/linux/-/blob/a1fc9f584c4aaf8bc1ebfa459fc57a3f26a290d8/drivers/gpu/drm/amd/amdkfd/kfd_queue.c#L391
-    sgrp_size_per_cu, hwreg_size_per_cu = 0x4000, 0x1000
-    lds_size_per_cu = self.iface.props["lds_size_in_kb"] << 10 if self.target[:2] == (9,5) else 0x10000
-    vgpr_size_per_cu = 0x60000 if self.target in {(11,0,0), (11,0,1), (11,5,1), (12,0,0), (12,0,1)} else 0x80000 if self.target[0] == 9 else 0x40000
-    wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
-    ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
-    return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
-      0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
-      ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
-      debug_memory_size=round_up(self.wave_cnt * 32, 64))
-
-  def sdma_queue(self, idx:int):
-    if getenv("AMD_DISABLE_SDMA"): return None
-    if idx in self.sdma_queues: return self.sdma_queues[idx]
-    with contextlib.suppress(OSError):
-      self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
-    return self.sdma_queues.get(idx, None)
-
-  def tmpring_size(self, private_segment_size):
-    private_segment_size = max(private_segment_size, 128)
-
-    lanes_per_wave = 64 # wave64
-    mem_alignment_size = 256 if self.target[0] != 9 else 1024
-    size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
-    size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-
-    # NOTE: xcc logic is correct only for GFX9.
-    max_scratch_waves = self.cu_cnt * self.iface.props['max_slots_scratch_cu'] * self.xccs
-    wave_scratch = ceildiv(lanes_per_wave * size_per_thread, mem_alignment_size)
-    num_waves = (size_per_xcc // (wave_scratch * mem_alignment_size)) // (self.se_cnt if self.target[0] != 9 else 1)
-
-    tmpring_t = getattr(hsa, f'union_COMPUTE_TMPRING_SIZE{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-    tmpring = int.from_bytes(tmpring_t(WAVES=min(num_waves, max_scratch_waves), WAVESIZE=wave_scratch), 'little')
-
-    if hasattr(self, 'aql_desc'):
-      gfx9_rsrc = {'NUM_FORMAT':hsa.BUF_NUM_FORMAT_UINT, 'DATA_FORMAT':hsa.BUF_DATA_FORMAT_32, 'ELEMENT_SIZE':1, 'INDEX_STRIDE':3}
-      rsrc = {'DST_SEL_X':hsa.SQ_SEL_X, 'DST_SEL_Y':hsa.SQ_SEL_Y, 'DST_SEL_Z':hsa.SQ_SEL_Z, 'DST_SEL_W':hsa.SQ_SEL_W, 'ADD_TID_ENABLE':1,
-              'TYPE':hsa.SQ_RSRC_BUF, **(gfx9_rsrc if self.target[0] == 9 else {'FORMAT':hsa.BUF_FORMAT_32_UINT, 'OOB_SELECT':2})}
-      rsrc1_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD1{"_GFX11" if self.target[0] != 9 else ""}_bitfields')
-      rsrc3_t = getattr(hsa, f'union_SQ_BUF_RSRC_WORD3{"_GFX"+str(self.target[0]) if self.target[0] != 9 else ""}_bitfields')
-
-      self.aql_desc.scratch_backing_memory_location = int(self.scratch.get_buf().va_addr)
-      self.aql_desc.scratch_wave64_lane_byte_size = self.max_private_segment_size * lanes_per_wave // 64
-      self.aql_desc.scratch_resource_descriptor[:] = [lo32(self.scratch.get_buf().va_addr),
-        int.from_bytes(rsrc1_t(BASE_ADDRESS_HI=hi32(self.scratch.get_buf().va_addr), SWIZZLE_ENABLE=1), 'little'),
-        lo32(size_per_xcc), int.from_bytes(bytes(rsrc3_t(**rsrc)), 'little')]
-      self.aql_desc.compute_tmpring_size = tmpring
-      self.aql_gart._buf.cpu_view()[:ctypes.sizeof(self.aql_desc)] = bytes(self.aql_desc)
-
-    return tmpring
-
-  def scratch_buffer(self, private_segment_size):
-    AMDDevice.max_scratch_psize = private_segment_size = max(private_segment_size, 128, AMDDevice.max_scratch_psize)
-    if self.max_private_segment_size < private_segment_size:
-      lanes_per_wave = 64 # wave64
-      mem_alignment_size = 256 if self.target[0] != 9 else 1024
-      size_per_thread = round_up(private_segment_size, mem_alignment_size // lanes_per_wave)
-      size_per_xcc = size_per_thread * lanes_per_wave * self.iface.props['max_slots_scratch_cu'] * self.cu_cnt
-      self.scratch = Buffer(self.device, size_per_xcc * self.xccs, dtypes.uint8, options=BufferSpec(nolru=True), preallocate=True)
-      self.max_private_segment_size = private_segment_size
-    return self.scratch
-
-  def on_device_hang(self): self.iface.on_device_hang()
-
-  def device_props(self): return self.iface.props

extra/models/llama.py

@@ -139,7 +139,7 @@ class TransformerBlock:
 
   def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]):
     h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask)
-    return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward()
+    return (h + self.feed_forward(self.ffn_norm(h))).clone().contiguous_backward()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):

extra/remote/bench.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import os, sys, time
-from tinygrad.runtime.support.system import RemotePCIDevice
+from extra.hcq1.remote import RemotePCIDevice
 
 LAT_N_RUNS = 500
 THROUGHPUT_N_RUNS = 8
@@ -18,7 +18,7 @@ if __name__ == "__main__":
   print(f"connected to {os.environ['REMOTE']}, device: {name}\n")
 
   # ping (minimal server round-trip, no device I/O)
-  from tinygrad.runtime.support.system import RemoteCmd
+  from extra.hcq1.remote import RemoteCmd
   sock = pci.sock
   for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
   st = time.perf_counter()

extra/remote/serve.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import socket, struct, sys
-from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
+from tinygrad.runtime.support.system import PCIDevice, System
+from extra.hcq1.remote import RemoteCmd
 from tinygrad.helpers import DEBUG, OSX
 
 def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)

mkdocs.yml

@@ -29,7 +29,6 @@ nav:
     - UOp: developer/uop.md
     - Runtime:
       - developer/runtime.md
-      - HCQ: developer/hcq.md
       - AM Driver: developer/am.md
   - tinybox: tinybox.md
 #- tinygrad: reference/

test/amd/hw/test_ds.py

@@ -7,6 +7,29 @@ Includes: ds_store_b32, ds_load_b32, ds_store_2addr_*, ds_load_2addr_*,
 import unittest
 from test.amd.hw.helpers import *
 
+class TestDSSwizzle(unittest.TestCase):
+  def test_modes_and_overlapping_registers(self):
+    for offset in (0x041f, 0x401f, 0x7c1f, 0x00a0, 0x801b, 0xc020, 0xc420, 0xc021, 0xe000, 0xe010, 0xe01f):
+      for dst in (0, 1):
+        with self.subTest(offset=hex(offset), dst=dst):
+          st = run_program([
+            v_add_nc_u32_e32(v[0], 1, v[255]),
+            ds_swizzle_b32(vdst=v[dst], addr=v[0], offset0=offset & 255, offset1=offset >> 8),
+            s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+          ], n_lanes=32)
+          self.assertEqual(sorted(st.vgpr[i][dst] for i in range(32)), [6]*32 if offset == 0x00a0 else list(range(1, 33)))
+
+  def test_inactive_sources_and_destinations(self):
+    st = run_program([
+      v_add_nc_u32_e32(v[0], 1, v[255]),
+      v_mov_b32_e32(v[1], 99),
+      s_mov_b32(EXEC_LO, 0x55555555),
+      ds_swizzle_b32(vdst=v[1], addr=v[0], offset0=0x1f, offset1=4),
+      s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+      s_mov_b32(EXEC_LO, 0xffffffff),
+    ], n_lanes=32)
+    self.assertEqual([st.vgpr[i][1] for i in range(32)], [0, 99]*16)
+
 class TestDS2Addr(unittest.TestCase):
   """Tests for DS_*_2ADDR instructions."""
 

test/amd/test_asm_kernel.py

@@ -3,7 +3,7 @@ import functools
 import numpy as np
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
+from tinygrad.engine.realize import run_linear, estimate_uop, lower_and_compile
 from tinygrad.renderer import Estimates
 from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import getenv
@@ -169,7 +169,7 @@ class TestAsmKernel(unittest.TestCase):
     if self.arch != "rdna3": self.skipTest("only rdna3")
     a = Tensor.full((16, 16), 1.).contiguous().realize()
     a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
-    linear = compile_linear(a.schedule_linear())
+    linear = lower_and_compile(a.schedule_linear())
     est = estimate_uop(linear.src[-1])
     self.assertEqual(est.ops, a.numel())
     self.assertEqual(est.mem, a.nbytes()*2)

test/amd/test_emu2_pcode.py

@@ -17,6 +17,34 @@ def _srcs():
 class TestBasicParsing(unittest.TestCase):
   """Test basic pcode parsing for common instruction patterns."""
 
+  def test_c_style_blocks_and_array_access(self):
+    code = """
+      for (i = 0; i < 4; i+=2) {
+        if (mode == 0) {
+          out[i+0] = input[i+1];
+          out[i+1] = input[i+0];
+        } elsif (mode == 1) {
+          out[i+0] = 7;
+          out[i+1] = 8;
+        } else { // identity
+          out[i+0] = input[i+0];
+          out[i+1] = input[i+1];
+        }
+      }
+    """
+    for mode, expected in enumerate(([11, 10, 13, 12], [7, 8, 7, 8], [10, 11, 12, 13])):
+      with self.subTest(mode=mode):
+        result, _ = parse_pcode(code, {'mode': UOp.const(mode, dtypes.uint32)}, {'input': lambda i: i + 10})
+        self.assertEqual([result[f'out@{i}'].simplify().val for i in range(4)], expected)
+
+  def test_colon_concatenation(self):
+    result, _ = parse_pcode('offset = hi:lo;', {'hi': UOp.const(0x12, dtypes.uint8), 'lo': UOp.const(0x34, dtypes.uint8)})
+    self.assertEqual(result['offset'].simplify().val, 0x1234)
+
+  def test_unclosed_c_block(self):
+    with self.assertRaisesRegex(AssertionError, 'unclosed pcode block'):
+      parse_pcode('if (1) {\nvalue = 2;')
+
   def test_v_add_f32(self):
     """Test parsing V_ADD_F32 pcode."""
     _, assigns = parse_pcode(PCODE[VOP2Op.V_ADD_F32_E32], _srcs())

test/amd/test_mockgpu_invalid.py

@@ -1,51 +1,29 @@
 #!/usr/bin/env python3
 """Test that invalid instructions raise exceptions through the mock GPU stack."""
-import unittest, subprocess, os, sys, time
+import unittest, subprocess, os, sys
 
 class TestMockGPUInvalidInstruction(unittest.TestCase):
   def test_unsupported_instruction_raises(self):
     """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
     test_code = '''
-import struct
-from dataclasses import replace
-from tinygrad import Device, Tensor
-from tinygrad.engine.realize import compile_linear
+import os, sys
+from tinygrad import Tensor
+from tinygrad.engine.realize import lower_and_compile, run_linear
 
-dev = Device["AMD"]
-a = Tensor([1.0]).realize()
-b = a + 1
-linear = compile_linear(b.schedule_linear())
-compiled_prg = linear.src[-1].src[0]
-lib = bytearray(compiled_prg.src[3].arg)
-
-# Find s_endpgm (0xBFB00000) and replace with V_MOVRELD_B32 (op=66) which has no pcode
-# VOP1 encoding: bits[31:25]=0x7E, op=bits[16:9], so op=66 -> 66<<9 = 0x8400
-found = False
-for i in range(0, len(lib) - 4, 4):
-  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
-    lib[i:i+4] = struct.pack("<I", 0x7E008400)
-    found = True
-    break
-assert found, "s_endpgm not found"
-
-patched_prg = dev.runtime(replace(compiled_prg.to_elf(), name="patched", lib=bytes(lib)))
-b.uop.buffer.allocate()
-patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
-dev.synchronize()
+linear = lower_and_compile((Tensor.empty(1) + 1).schedule_linear())
+binary = linear.src[-1].src[0].src[3]
+lib = binary.arg.replace(bytes.fromhex("0000b0bf"), bytes.fromhex("00fe017e"), 1)
+try:
+  run_linear(linear.substitute({binary: binary.replace(arg=lib)}, enter_calls=True))
+except ValueError as error:
+  print(error, file=sys.stderr, flush=True)
+  os._exit(1)
 '''
 
-    env = os.environ.copy()
-    env["DEV"] = "MOCKKFD+AMD"
-    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
-
-    st = time.perf_counter()
-    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
-    elapsed = time.perf_counter() - st
-
-    self.assertNotEqual(result.returncode, 0, "should have raised")
-    self.assertTrue("Error" in result.stderr, f"expected an error in stderr, got: {result.stderr[:500]}")
-    # Should exit immediately, not wait for the full timeout
-    self.assertLess(elapsed, 9.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
+    env = {**os.environ, "DEV": "MOCKKFD+AMD", "HCQ_RUNTIME_DEV": "PYTHON"}
+    result = subprocess.run([sys.executable, "-c", test_code], env=env, capture_output=True, text=True, timeout=9)
+    self.assertEqual(result.returncode, 1)
+    self.assertIn("unknown rdna3 format word=0x7e01fe00", result.stderr)
 
 if __name__ == "__main__":
   unittest.main()

test/amd/test_pdf.py

@@ -43,6 +43,16 @@ class TestPcodePDF(unittest.TestCase):
     self.assertEqual(pcode[('S_CMOVK_I32', 2)],
       "if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
 
+  def test_swizzle_spans_blocks_and_pages(self):
+    for arch in ('rdna3', 'rdna4'):
+      with self.subTest(arch=arch):
+        code = self.pcode[arch][('DS_SWIZZLE_B32', 53)]
+        self.assertIn('} elsif (offset >= 0xc000) {', code)
+        self.assertIn('thread_out[i+3]', code)
+        self.assertIn('xor_mask = offset[14:10];', code)
+        self.assertEqual(code.count('{'), code.count('}'))
+        self.assertTrue(code.endswith('\n}'))
+
   def test_pcode_no_examples(self):
     """Pseudocode should not contain example lines with '=>'."""
     for name in ARCHS:

test/amd/test_roundtrip.py

@@ -58,11 +58,11 @@ def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelSnapshot], dict[int, in
   """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
   from tinygrad import Tensor
   from tinygrad.uop.ops import Ops
-  from tinygrad.engine.realize import compile_linear, resolve_params, unwrap_multi
+  from tinygrad.engine.realize import lower_and_compile, resolve_params, unwrap_multi
   from tinygrad.runtime.support.elf import elf_loader
 
   out = op_fn(Tensor)
-  linear = compile_linear(out.schedule_linear())
+  linear = lower_and_compile(out.schedule_linear())
   kernels = []
   buf_pool: dict[int, int] = {}  # buffer id -> size
   buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY

test/amd/test_sqtt_profiler.py

@@ -1,5 +1,6 @@
 import unittest, contextlib
 from tinygrad import Device, Tensor, Context, TinyJit, dtypes
+from test.helpers import is_hcq2_device
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.device import Compiled, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
@@ -114,8 +115,7 @@ class TestSQTTProfiler(unittest.TestCase):
     kernel_name = sqtt[0]["name"]
     for i,e in enumerate(sqtt[1:], start=1): self.assertEqual(e["name"], f"{kernel_name} n{i+1}")
 
-  # TODO: can we trace SQTT for graphed kernels?
-  def test_jit_graph(self, kernel_count=3*1):
+  def test_jit_graph(self, kernel_count=3*(5 if is_hcq2_device() else 1)): # hcq2 traces the graphed kernels too
     @TinyJit
     def f(a): return ((a + 1).contiguous() + 2).contiguous().sum()
     t = Tensor.empty(32)

test/backend/test_arange.py

@@ -4,7 +4,7 @@ from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device, Variable
 from tinygrad.helpers import Context, getenv, DEV
 from tinygrad.engine.realize import run_linear, estimate_uop, compile_linear
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException
+from test.helpers import needs_second_gpu, check_schedule, assert_kernel_count, KernelCountException, is_hcq2_device
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, tensor, desired):
@@ -153,7 +153,7 @@ class TestIndexing(unittest.TestCase):
       GlobalCounters.reset()
       z = emb(x).realize()
       self.assertLessEqual(GlobalCounters.global_ops, op_limit)
-      assert_kernel_count(2)
+      assert_kernel_count(3 if is_hcq2_device() else 2)
     if getenv("CHECK", 1):
       import torch
       with torch.no_grad():

test/backend/test_assign.py

@@ -4,7 +4,7 @@ import numpy as np
 from tinygrad import Device, dtypes, Tensor, TinyJit, GlobalCounters, Variable
 from tinygrad.uop.ops import Ops, UOp
 from tinygrad.helpers import temp, DEV, Context
-from test.helpers import assert_kernel_count, needs_second_gpu
+from test.helpers import assert_kernel_count, needs_second_gpu, is_hcq2_device
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -43,7 +43,7 @@ class TestAssign(unittest.TestCase):
     # it should copy into the empty buffer
     GlobalCounters.reset()
     c.realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
 
   def test_assign_slice(self):
     X = Tensor([1,2,3,4]).realize()
@@ -619,7 +619,7 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(2)  # TODO: first copy is dead, could be 1
+    assert_kernel_count(5 if is_hcq2_device() else 3)  # TODO: first copy is dead, could be 2
     self.assertEqual(base.tolist(), [1,4,3])
 
   def test_nested_after_contiguous_store_no_init(self):
@@ -629,9 +629,17 @@ class TestAssign(unittest.TestCase):
     contig.assign(Tensor([1, 4, 3], dtype=dtypes.int64))
     GlobalCounters.reset()
     base.assign(contig).realize()
-    assert_kernel_count(1)
+    assert_kernel_count(2 if is_hcq2_device() else 1)
     self.assertEqual(base.tolist(), [1,4,3])
 
+  def test_assign_temporary_copy_reshape(self):
+    a = Tensor([[1., 2], [3, 4]], device="PYTHON")
+    c = Tensor.empty(2, 2).assign(a.to(None))
+    GlobalCounters.reset()
+    c.realize()
+    assert_kernel_count(2 if is_hcq2_device() else 1)
+    self.assertEqual(c.tolist(), [[1., 2], [3, 4]])
+
 class TestAssignOrdering(unittest.TestCase):
   """Tests for complex assign orderings that could differ between lazy and eager execution.
 
@@ -867,12 +875,8 @@ class TestAssignOrdering(unittest.TestCase):
     a.assign(b + 1)                                    # a == 11
     v1 = a * 3                                         # reads 11 -> 33
     a.assign(b + 100)                                  # a == 110
-    out = (a + v1).numpy()
-    try:
-      np.testing.assert_allclose(out, 143)
-    except AssertionError:
-      # TODO: broken now, v1 reads a after the second assign
-      np.testing.assert_allclose(out, 440)
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, ideally v1 is realized between the assigns
+      np.testing.assert_allclose((a + v1).numpy(), 143)
 
   def test_two_reads_between_three_assigns(self):
     a = Tensor.zeros(4).realize()
@@ -987,12 +991,9 @@ class TestAssignOrdering(unittest.TestCase):
       x.assign(x+1)
       return y+x
     a = Tensor([1.]).realize()
-    out = outer(a).item()
-    try:
+    with self.assertRaisesRegex(RuntimeError, "cycle"):  # TODO: broken now, ideally y is realized between the assigns
+      out = outer(a).item()
       self.assertEqual([out, a.item()], [7., 3.])
-    except AssertionError:
-      # TODO: broken now, the inner assign is run twice
-      self.assertEqual([out, a.item()], [6., 4.])
 
 class TestAssignToUnrealizedView(unittest.TestCase):
   def test_copy(self):
@@ -1013,6 +1014,24 @@ class TestAssignToUnrealizedView(unittest.TestCase):
     c[:, 1:2].assign(Tensor.ones(2,1, dtype=dtypes.int).contiguous().realize())
     self.assertEqual(c.tolist(), [[1,1],[2,1]])
 
+  def test_contiguous_partial_assign_realize(self):
+    x = Tensor([1., 2.]).realize()
+    y = (x + 1).contiguous()  # unrealized CONTIGUOUS
+    self.assertIs(y.uop.base.op, Ops.CONTIGUOUS)
+    # a partial write survives an explicit realize: the values are right, storage is an implementation detail
+    y[:1].assign(9.)
+    y.realize()
+    self.assertEqual(y.tolist(), [9., 3.])
+    # and it stays assigned across schedules
+    y[:1].assign(7.)
+    y.realize()
+    self.assertEqual(y.tolist(), [7., 3.])
+    # setitem syntax gives the same values, contiguous or not
+    for mk in (lambda xx: xx + 1, lambda xx: (xx + 1).contiguous()):
+      z = mk(Tensor([1., 2.]).realize())
+      z[:1] = 9.
+      self.assertEqual(z.tolist(), [9., 3.])
+
   def test_contiguous_backward(self):
     t = Tensor([[1,2],[3,4]]).contiguous().realize()
     cb = t.contiguous_backward()  # unrealized CONTIGUOUS_BACKWARD
@@ -1086,6 +1105,16 @@ class TestAssignToUnrealizedView(unittest.TestCase):
       # TODO: broken now, silently dropped
       self.assertEqual(c.tolist(), [[5,5],[5,5]])
 
+  def test_detach_assignment_preserves_earlier_update(self):
+    x = Tensor([1., 2.]).detach()
+    state = Tensor([0., 0.]).detach()
+    state.assign(state + x * 2)
+    result = state + 1
+    x.assign(x + 1).realize(state, result)
+    self.assertEqual(x.tolist(), [2., 3.])
+    self.assertEqual(state.tolist(), [2., 4.])
+    self.assertEqual(result.tolist(), [3., 5.])
+
 class TestPartialAssignToSharedBuffer(unittest.TestCase):
   def test_five_slices(self):
     big = Tensor.zeros(50).contiguous().realize()

test/backend/test_const_folding.py

@@ -84,10 +84,10 @@ class TestReduceOpsConstFolding(unittest.TestCase):
       np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
 
   def test_zero_size_realize_folded(self):
-    # non contiguous folded output doesn't realize
+    # folded output doesn't realize on its own
     _check_ast_count(0, Tensor.empty(1, 0).sum())
-    # contiguous folded const can still schedule
-    a = Tensor.empty(1, 0).sum().contiguous()
+    # explicit storage of the folded const still schedules, and the value is usable
+    a = Tensor.empty(1, 0).sum().clone()
     _check_ast_count(2, a+2)
     self.assertIs(a.uop.base.op, Ops.BUFFER)
     np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)

test/backend/test_jit.py

@@ -2,12 +2,11 @@
 import unittest
 import numpy as np
 
-from test.helpers import assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
+from test.helpers import is_hcq2_device, assert_jit_cache_len, call_is_graph, not_support_multi_device, needs_second_gpu, KernelCountException
 from test.unit.test_jit import _simple_test
 from tinygrad import Tensor, TinyJit, Device, dtypes
 from tinygrad.engine.jit import graph_class
 from tinygrad.helpers import JIT, DEV, GlobalCounters
-from tinygrad.runtime.support.hcq2 import HCQ_DEVS
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer.isa.x86 import X86Renderer
 
@@ -223,7 +222,7 @@ class TestJitPrune(unittest.TestCase):
     assert_jit_cache_len(w2_prune, 1)
 
 class TestJitFree(unittest.TestCase):
-  @unittest.skipIf(Device.DEFAULT.split(":")[0] in HCQ_DEVS - {"CPU"}, "hcq2 keeps refs to intermediate buffers")
+  @unittest.skipIf(is_hcq2_device(), "hcq2 keeps refs to intermediate buffers")
   def test_free_intermediates(self):
     ext_tensor = Tensor([1,24,23,45,1])
     @TinyJit
@@ -293,7 +292,7 @@ class TestJitGraphSplit(unittest.TestCase):
     if graph_t is None: return
 
     got = f.captured.linear.src
-    from tinygrad.runtime.graph.hcq import HCQGraph
+    from extra.hcq1.graph import HCQGraph
     from tinygrad.engine.jit import MultiGraphRunner
     if graph_t is HCQGraph:
       validate = hcqgraph

test/backend/test_multitensor.py

@@ -1,16 +1,12 @@
 import unittest, random
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, Variable
 from tinygrad.uop.ops import Ops, UOp, AxisType, graph_rewrite
-from tinygrad.helpers import getenv, prod, Context
+from tinygrad.helpers import prod, Context
 from tinygrad.nn.state import get_parameters
-from tinygrad.engine.realize import run_linear, compile_linear, lower_and_compile, pm_beam
+from tinygrad.engine.realize import run_linear, lower_and_compile, pm_beam
 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, check_schedule, assert_kernel_count, KernelCountException
 
-settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
-settings.load_profile("my_profile")
-
 d0 = f"{Device.DEFAULT}:0"
 d1 = f"{Device.DEFAULT}:1"
 d2 = f"{Device.DEFAULT}:2"
@@ -76,7 +72,7 @@ class TestMultiTensor(unittest.TestCase):
     X = Tensor.ones(256).contiguous().realize()
     X.shard_(devices_2, 0)
     out = (X + X)
-    linear = compile_linear(out.schedule_linear())
+    linear = lower_and_compile(out.schedule_linear())
     uops = [call.src[0].src[0] for call in linear.src if call.src[0].op is Ops.PROGRAM]
     run_linear(linear)
     self.assertEqual(len(set(uops)), 1, "function was relinearized")
@@ -129,17 +125,21 @@ class TestMultiTensor(unittest.TestCase):
       run_linear(linear, var_vals)
       np.testing.assert_equal(xt.numpy(), X_np[i*2:i*2+2])
 
-  @given(strat.sampled_from((devices_2, devices_3)),
-         strat.sampled_from((Ops.ADD, Ops.MUL, Ops.MAX)),
-         strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)))
-  def test_simple_reduce(self, devices, rop, shard_axis, reduce_axis):
-    N = 4 * len(devices)
-    X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis)
-    n = X.numpy()
-    f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop]
-    fX = f(X)
-    fn = f(n)
-    np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
+  def test_simple_reduce(self):
+    for devices, rop, shard_axis, reduce_axis in [
+      (devices_2, Ops.ADD, None, None), (devices_2, Ops.ADD, 0, 0), (devices_2, Ops.ADD, 0, 1),
+      (devices_2, Ops.ADD, 1, 0), (devices_2, Ops.ADD, 1, 1),
+      (devices_3, Ops.ADD, 0, 0), (devices_3, Ops.ADD, 1, 0),
+      (devices_2, Ops.MUL, 0, 1), (devices_2, Ops.MUL, 1, 1), (devices_3, Ops.MUL, 0, 0),
+      (devices_2, Ops.MAX, 0, 1), (devices_3, Ops.MAX, 1, 0)]:
+      with self.subTest(devices=len(devices), op=rop.name, shard_axis=shard_axis, reduce_axis=reduce_axis):
+        N = 4 * len(devices)
+        X = (Tensor.rand(N*N)-1).reshape(N, N).shard_(devices, shard_axis)
+        n = X.numpy()
+        f = {Ops.ADD: lambda x: x.sum(reduce_axis), Ops.MUL: lambda x: x.prod(reduce_axis), Ops.MAX: lambda x: x.max(reduce_axis)}[rop]
+        fX = f(X)
+        fn = f(n)
+        np.testing.assert_allclose(fX.numpy(), fn, rtol=1e-6, atol=1e-6)
 
   def test_stack(self):
     X = Tensor.rand(4, 4).shard_(devices_2, 0)
@@ -176,21 +176,21 @@ class TestMultiTensor(unittest.TestCase):
   def test_allreduce_naive_jit(self):
     with Context(RING=0):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
   def test_allreduce_ring_jit(self):
     with Context(RING=2):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
   def test_allreduce_all2all_jit(self):
     with Context(ALL2ALL=2):
       jit_allreduce = TinyJit(_test_allreduce)
-      for _ in range(5):
+      for _ in range(3):
         a,b = jit_allreduce(Tensor.rand(256, 256))
         np.testing.assert_almost_equal(a.numpy(), b.numpy(), decimal=5)
 
@@ -212,7 +212,7 @@ class TestMultiTensor(unittest.TestCase):
 
   def test_fuzz_allreduce(self):
     random.seed(41)
-    for it in range(2):
+    for it in range(1):
       for n in range(2, 4+1):
         shape = tuple([(n if i == 0 else 1) * random.randint(1, 10) for i in range(random.randint(1, 4))])
         t = Tensor.rand(shape).shard_(tuple([d0, d1, d2, d3][:n]), 0)
@@ -445,6 +445,7 @@ class TestMultiBufferView(unittest.TestCase):
 
 @unittest.skipIf(not_support_multi_device(), "need multi")
 class Test2DShard(unittest.TestCase):
+  @needs_second_gpu
   def setUp(self):
     self.devices_4 = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
     self.rng = UOp.range(4, -1, AxisType.DEVICE)
@@ -460,6 +461,15 @@ class Test2DShard(unittest.TestCase):
     out = t.contiguous().realize()
     np.testing.assert_equal(out.numpy(), ref.numpy())
 
+  def test_2d_shard_clone(self):
+    ref = Tensor.arange(16).reshape(4, 4).realize()
+    t = self._shard_2d(ref)
+    out = t.clone().realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy())
+    out.assign(out + 1).realize()
+    np.testing.assert_equal(out.numpy(), ref.numpy() + 1)
+    np.testing.assert_equal(t.numpy(), ref.numpy())
+
   def test_2d_shard_elementwise(self):
     ref = Tensor.arange(16).reshape(4, 4).contiguous().realize()
     t = self._shard_2d(ref)
@@ -513,7 +523,8 @@ class TestMultiTransformer(unittest.TestCase):
       else: v.shard_(device, axis=None)
 
     last_tok = 0
-    for i in range(5):
+    # i=0: bypasses jit, i=1: jit warmup, i=2: capture and run, i>=3: re-execute jit with new start_pos (catches stale bindings)
+    for i in range(4):
       real_tok = real_model(Tensor([[last_tok]], device=Device.DEFAULT), i).item()
       shard_tok = shard_model(Tensor([[last_tok]], device=device), i).item()
 

test/backend/test_ops.py

@@ -6,7 +6,6 @@ from tinygrad.helpers import getenv, DEBUG, DEV, IMAGE, Context
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.renderer.nir import NIRRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
 
 TINY_BACKEND = getenv("TINY_BACKEND")
 if TINY_BACKEND:
@@ -713,6 +712,9 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: 0**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: 0.7**x, vals=[[-2.,-1,0,1,2,3]])
     helper_test_op(None, lambda x: (-2)**x, vals=[[-2.,-1,0,1,2,3]])
+    # 2**52+2 - 0.5 rounds back to itself
+    helper_test_op(None, lambda x: x**(2.0**52), vals=[[0.5, 1., 2.]], forward_only=True)
+    helper_test_op(None, lambda x: x**(2.0**52+2), vals=[[0.5, 1., 2.]], forward_only=True)
     # float to power of int
     helper_test_op(None, lambda x: 0.7**x, lambda x: (0.7**x).clone(), vals=[[-2,-1,0,1,2,3]], forward_only=True)
 
@@ -816,8 +818,6 @@ class TestOps(unittest.TestCase):
     helper_test_op([], lambda: tor^0x1337, lambda: ten^0x1337, forward_only=True)
     helper_test_op([], lambda: 0x1337^tor, lambda: 0x1337^ten, forward_only=True)
 
-  # TODO: x86 PARAM dtype fails SPEC=2
-  @Context(SPEC=1 if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else 2)
   def test_and(self):
     data = [[1,-8,1],[32,1,6]]
     tor = torch.tensor(data, dtype=torch.int)
@@ -1132,9 +1132,12 @@ class TestOps(unittest.TestCase):
   def test_relu6(self):
     helper_test_op([(45,65)], torch.nn.functional.relu6, Tensor.relu6)
     helper_test_op([()], torch.nn.functional.relu6, Tensor.relu6)
+    helper_test_op(None, torch.nn.functional.relu6, Tensor.relu6, vals=[[6.71089e7, 2.68435e8, 1e9]])
+    helper_test_op(None, torch.nn.functional.relu6, Tensor.relu6, vals=[[0., 6.]])
   def test_hardswish(self):
     helper_test_op([(45,65)], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
     helper_test_op([()], torch.nn.functional.hardswish, Tensor.hardswish, grad_atol=1e-6)
+    helper_test_op(None, torch.nn.functional.hardswish, Tensor.hardswish, vals=[[-3., 3.]], grad_atol=1e-6)
   def test_mish(self):
     helper_test_op([(45,65)], torch.nn.functional.mish, Tensor.mish)
     helper_test_op([()], torch.nn.functional.mish, Tensor.mish)
@@ -3110,6 +3113,13 @@ class TestOps(unittest.TestCase):
                          lambda x: x.gather(dim=0, index=Tensor([2, 1, 0, 1, 2])),
                          vals=[[-float("inf"), 2., 3.]])
 
+  def test_gather_bool_index(self):
+    helper_test_op(None, lambda x,y: x.gather(dim=0, index=y.bool().long()),
+                         lambda x,y: x.gather(dim=0, index=y.cast(dtypes.bool).cast(dtypes.int)),
+                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
+    helper_test_op(None, lambda x,y: x[y.bool().long()], lambda x,y: x[y.cast(dtypes.bool).cast(dtypes.int)],
+                         vals=[[1., 2., 3.], [0.5, 0., 2.]], forward_only=True)
+
   def test_scatter(self):
     b = torch.randint(3, size=[3,4,5], dtype=torch.int64, requires_grad=False)
     a = Tensor(b.detach().cpu().numpy().astype(np.int32), dtype=dtypes.int32)

test/backend/test_profiler.py

@@ -2,7 +2,8 @@ import unittest, struct, contextlib, statistics, gc
 from tinygrad import Device, Tensor, dtypes, TinyJit
 from tinygrad.helpers import DEV, 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 extra.hcq1.hcq import HCQCompiled
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled
 from tinygrad.engine.realize import get_runtime
 from tinygrad.codegen import to_program
 
@@ -34,7 +35,18 @@ def helper_profile_filter_device(profile, device:str):
   assert len(dev_events) == 1, "only one device registration event is expected"
   return [x for x in profile if getattr(x, "device", None) == device], dev_events[0]
 
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT], (HCQCompiled, HCQ2Compiled)) or Device.DEFAULT == "METAL", "Dev not supported")
+class TestSimpleProfiler(unittest.TestCase):
+  @unittest.skipIf(Device.DEFAULT == "CPU", "fails in CPU")
+  def test_profiler(self):
+    start = len(Compiled.profile_events)
+    with Context(PROFILE=1):
+      Tensor.empty(32).add(1).realize()
+      Device[Device.DEFAULT].synchronize()
+    self.assertTrue(any(isinstance(e, (ProfileRangeEvent, ProfileGraphEvent)) for e in Compiled.profile_events[start:]))
+
 # TODO: support in HCQCompiled
+# TODO: support these tests in HCQ2
 is_cpu_hcq = Device.DEFAULT in {"CPU"}
 
 @unittest.skipUnless((issubclass(type(Device[Device.DEFAULT]), HCQCompiled) and not is_cpu_hcq) or Device.DEFAULT in {"METAL"}, "Dev not supported")

test/backend/test_setitem.py

@@ -1,6 +1,5 @@
-import unittest
+import unittest, operator
 from tinygrad import Tensor, TinyJit, Variable, dtypes, Device
-from tinygrad.helpers import Context
 import numpy as np
 
 class TestSetitem(unittest.TestCase):
@@ -163,21 +162,20 @@ class TestSetitem(unittest.TestCase):
       np.testing.assert_allclose(t.numpy(), n)
 
   def test_jit_setitem_variable_offset(self):
-    with Context(CHECK_OOB=0):
-      @TinyJit
-      def f(t:Tensor, a:Tensor, v:Variable):
-        t.shrink(((v,v+1), None)).assign(a).realize()
+    @TinyJit
+    def f(t:Tensor, a:Tensor, v:Variable):
+      t.shrink(((v,v+1), None)).assign(a).realize()
 
-      t = Tensor.zeros(6, 6).contiguous().realize()
-      n = np.zeros((6, 6))
+    t = Tensor.zeros(6, 6).contiguous().realize()
+    n = np.zeros((6, 6))
 
-      for i in range(6):
-        v = Variable("v", 0, 6).bind(i)
-        a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
-        n[i, :] = i+1
-        f(t, a, v)
-        np.testing.assert_allclose(t.numpy(), n)
-      np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
+    for i in range(6):
+      v = Variable("v", 0, 6).bind(i)
+      a = Tensor.full((1, 6), fill_value=i+1, dtype=dtypes.float).contiguous()
+      n[i, :] = i+1
+      f(t, a, v)
+      np.testing.assert_allclose(t.numpy(), n)
+    np.testing.assert_allclose(t.numpy(), [[1,1,1,1,1,1],[2,2,2,2,2,2],[3,3,3,3,3,3],[4,4,4,4,4,4],[5,5,5,5,5,5],[6,6,6,6,6,6]])
 
   def test_setitem_overlapping_inplace1(self):
     t = Tensor([[3.0], [2.0], [1.0]]).contiguous()
@@ -378,6 +376,32 @@ class TestWithGrad(unittest.TestCase):
     with self.assertRaises(RuntimeError):
       y[0] = 99.0
 
+  def test_unrealized_inplace_keeps_storage(self):
+    x = Tensor([1., 2.]).clone()
+    view = x[:1]
+    x += 3
+    x.realize()
+    self.assertEqual(x.tolist(), [4., 5.])
+    self.assertEqual(view.tolist(), [4.])
+
+  def test_unrealized_view_inplace_keeps_storage(self):
+    x = Tensor([1., 2.]).clone()
+    view = x[:1]
+    view += 3
+    view.realize()
+    self.assertEqual(x.tolist(), [4., 2.])
+    self.assertEqual(view.tolist(), [4.])
+
+  def test_set_augmented_backward(self):
+    for op, expected in ((operator.isub, [-1., -1.]), (operator.imul, [1., 2.]), (operator.itruediv, [-0.01, -0.005])):
+      with self.subTest(op=op.__name__):
+        z = Tensor([1.0, 2.0, 3.0, 4.0])
+        x = Tensor([10.0, 20.0])
+        z[:2] = op(z[:2], x)
+        z.sum().backward()
+        np.testing.assert_allclose(z.grad.numpy(), np.ones(4))
+        np.testing.assert_allclose(x.grad.numpy(), expected)
+
 class TestSetitemLoop(unittest.TestCase):
   def test_arange(self):
     N = 10

test/backend/test_uops.py

@@ -11,6 +11,7 @@ from tinygrad.engine.realize import run_linear
 from tinygrad.codegen import to_program
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.runtime.ops_python import PythonRenderer
 from test.helpers import to_uops_list
 
 def run_uops(uops_list:list[UOp], bufs:list[Buffer]):
@@ -56,8 +57,8 @@ def _test_uops_result(output_dtype, uops, res):
   run_uops([out], [buf])
   return np.frombuffer(buf.as_memoryview(), _to_np_dtype(output_dtype))[0]
 
-@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, CStyleLanguage) and
-                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires C-style pointer bitcast and 64-bit ints")
+@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, (CStyleLanguage, PythonRenderer)) and
+                     dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "requires buffer bitcast and 64-bit ints")
 class TestBitcastBufferView(unittest.TestCase):
   @Context(SPEC=2)
   def test_render(self):
@@ -85,6 +86,16 @@ class TestBitcastBufferView(unittest.TestCase):
     run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
     self.assertEqual(np.frombuffer(buf.as_memoryview(), dtype=np.uint64, count=2, offset=4).tolist(), [val ^ 0xff, val])
 
+  def test_vector_load_store(self):
+    for src_dt, dst_dt in [(dtypes.uint8, dtypes.uint32), (dtypes.uint32, dtypes.uint8)]:
+      with self.subTest(src=src_dt, dst=dst_dt):
+        src, dst = [UOp.param(i, dt, 16 // dt.itemsize) for i, dt in enumerate((src_dt, dst_dt))]
+        src, dst = [b.bitcast(dtypes.uint32).index(UOp.stack(*[UOp.const(i) for i in range(4)])) for b in (src, dst)]
+        bufs = [Buffer(Device.DEFAULT, 16 // dt.itemsize, dt, initial_value=bytes(range(16)) if i == 0 else bytes(16))
+                for i, dt in enumerate((src_dt, dst_dt))]
+        run_uops([dst.store(src.load())], bufs)
+        self.assertEqual(bytes(bufs[1].as_memoryview()), bytes(range(16)))
+
 class TestUOps(unittest.TestCase):
   def _equal(self, v1, v2):
     assert isinstance(v2, (float, int, bool))

test/backend/test_wait_loop.py

@@ -1,13 +1,12 @@
-import unittest, threading
-from tinygrad import Tensor, UOp
+import unittest, threading, functools
+from tinygrad import Tensor, UOp, Context
 from tinygrad.device import Device, Buffer, BufferSpec
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.engine.realize import run_linear
 from tinygrad.uop.ops import Ops, KernelInfo
+from tinygrad.renderer.isa.x86 import X86Renderer
 
-def wait_loop_kernel(C:UOp) -> UOp:
-  N = 10
-
+def wait_loop_kernel(C:UOp, N=10) -> UOp:
   # a RANGE with no src is a bound-less loop header: a jump target with no induction variable.
   # the compare and conditional backedge are expanded by the renderers from the loop RANGE/END
   l = UOp.loop(0)
@@ -42,6 +41,19 @@ def nested_loop_kernel(C:UOp) -> UOp:
 
   return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))
 
+def pressure_loop_kernel(C:UOp, n=13) -> UOp:
+  vs = [C[j+1].load() for j in range(n)]
+  l = UOp.loop(0)
+
+  i = UOp.placeholder((1,), dtypes.int, 0, addrspace=AddrSpace.REG)
+  i = i.after(i[0].store(0))
+
+  inc = i.after(l)[0].load() + 1
+  st = i[0].store(inc)
+  i = i.after(st.end(l, inc < sum(v & inc for v in vs)))
+
+  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="pressure_loop", opts_to_apply=()))
+
 def wait_ext_kernel() -> UOp:
   sig = UOp.param(0, dtypes.int, 1, volatile=True)
   l = UOp.loop(0)
@@ -100,6 +112,25 @@ class TestWaitLoop(unittest.TestCase):
     c.realize()
     self.assertEqual(c.item(), 25)
 
+  # TODO: x86's lower_loop builds an Ops.IF node after regalloc, which fails spec_full
+  @(unittest.expectedFailure if isinstance(Device[Device.DEFAULT].renderer, X86Renderer) else lambda f: f)
+  def test_wait_loop_spec(self):
+    c = Tensor.custom_kernel(Tensor.empty(1, dtype=dtypes.int), fxn=functools.partial(wait_loop_kernel, N=7))[0]
+    with Context(SPEC=2): c.realize()
+    self.assertEqual(c.item(), 7)
+
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, X86Renderer), "TODO: do-while loop under register pressure segfaults on x86")
+  def test_loop_carried_registers(self):
+    # more loads live across the backedge than any register file (x86 15 gprs, arm64 31, sass 255, rdna3 256 vgprs)
+    c = Tensor.custom_kernel(Tensor.ones(301, dtype=dtypes.int), fxn=functools.partial(pressure_loop_kernel, n=300))[0]
+    self.assertEqual(c[0].item(), 2)
+
+  def test_register_pressure_loop(self):
+    c = Tensor.zeros(16, dtype=dtypes.int).contiguous()
+    c = Tensor.custom_kernel(c, fxn=pressure_loop_kernel)[0]
+    c.realize()
+    self.assertEqual(c[0].item(), 1)
+
   def test_loop_in_loop(self):
     c = Tensor.empty(1, dtype=dtypes.int)
     c = Tensor.custom_kernel(c, fxn=loop_in_loop_kernel)[0]

test/device/test_hcq2.py

@@ -1,11 +1,12 @@
-import unittest, contextlib, ctypes, numpy as np
+import unittest, contextlib, ctypes, gc, numpy as np
 from unittest.mock import patch
-from tinygrad import Device, Tensor, TinyJit, Variable, dtypes
+from tinygrad import Device, Tensor, TinyJit, Variable, dtypes, GlobalCounters
 from tinygrad.device import Buffer
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import Context, dedup, partition
-from tinygrad.uop.ops import Ops, UOp, KernelInfo
-from tinygrad.engine.realize import lower_and_compile, run_linear
+from tinygrad.helpers import Context, dedup, partition, unwrap
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
+from tinygrad.engine.realize import compile_linear, link_linear, lower_and_compile, run_linear
+from tinygrad.codegen import do_to_program
 from tinygrad.renderer.cstyle import CStyleLanguage
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.c import init_c_struct_t
@@ -16,13 +17,26 @@ from test.helpers import call_is_hcq
 @contextlib.contextmanager
 def rt_views():
   calls, orig = [], HCQ2Compiled.rt_view
-  with patch.object(HCQ2Compiled, "rt_view", lambda s, *a, **kw: (calls.append(s), orig(s, *a, **kw))[1]): yield calls
+  def track(dev, *args, **kwargs):
+    calls.append(dev)
+    return orig(dev, *args, **kwargs)
+  with patch.object(HCQ2Compiled, "rt_view", track): yield calls
+
+def chain(x:Tensor, n:int) -> Tensor:
+  for _ in range(n): x = (x + 1).contiguous()
+  return x
 
 @contextlib.contextmanager
 def encoded_batches():
   batches, orig = [], hcq2.lower_and_compile
-  with patch.object(hcq2, "lower_and_compile", lambda l, *a, **kw: (batches.extend(c for c in l.src if call_is_hcq(c)), orig(l, *a, **kw))[1]):
-    yield batches
+  def track(l, *args, **kwargs):
+    batches.extend(c.without_after for c in l.src if call_is_hcq(c))
+    return orig(l, *args, **kwargs)
+  with patch.object(hcq2, "lower_and_compile", track): yield batches
+
+def eager_chain(x:Tensor, n:int=64) -> Tensor: # at hcq_compile's use_rt bound: an eager linear this big bakes its inputs and borrows ring slots
+  for _ in range(n): x = (x + 1).contiguous()
+  return x.realize()
 
 def patch_words(batch:UOp) -> list[UOp]:
   return [w for s in batch.src[0].toposort() if s.op is Ops.STORE and s.src[0].op is Ops.INDEX and s.src[0].src[1].op is Ops.STACK
@@ -31,41 +45,122 @@ def patch_words(batch:UOp) -> list[UOp]:
 def rt_params(batch:UOp) -> list[str]:
   return dedup([u.arg.name for w in patch_words(batch) for u in w.toposort() if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.GLOBAL])
 
-@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS - {"CPU"}), "non-CPU hcq2 device required")
-class TestHCQ2Core(unittest.TestCase):
+def cpu_buf(size:int=1, dtype=dtypes.uint8, **kwargs) -> UOp: return UOp.placeholder((size,), dtype, device="CPU", **kwargs)
+
+def lower_hcq(body:UOp) -> UOp:
+  return unwrap(hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test")).call(aux=hcq2.HCQInfo(("CPU",)))))
+
+class TestHCQ2Deps(unittest.TestCase):
+  def test_disjoint_write_preserves_dependencies(self):
+    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
+    for write in ([], [0]):
+      tracker = hcq2.HCQDepsTracker()
+      tracker.access_resources([b.shrink(((0, 4),))], write, 0)
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 8),))], [0], 1), [])
+      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
+
+  def test_partial_write_preserves_dependencies(self):
+    b = UOp.param(0, dtypes.uint8, 16, device="CPU")
+    for write in ([], [0]):
+      tracker = hcq2.HCQDepsTracker()
+      tracker.access_resources([b], write, 0)
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 1), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((0, 4),))], [0], 2), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((12, 16),))], [0], 3), [0])
+      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [], 4), [1])
+
+@unittest.skipUnless(all_devices_in(Device.DEFAULT, HCQ_DEVS), "hcq2 device required")
+class TestHCQ2Schedule(unittest.TestCase):
+  @staticmethod
+  def input(value:int=2) -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32).contiguous().realize()
+
+  def compiled(self, n:int, jit=False):
+    x, inputs = self.input(), []
+    if jit:
+      f = TinyJit(lambda a: chain(a, n).realize())
+      f(x)
+      return f(x), f.captured._linear, [x.uop.base]
+    out = chain(x, n)
+    return out, compile_linear(out.schedule_linear(), input_uops=inputs), inputs
+
   def test_jit_has_no_rt_buffers(self):
-    x = Tensor.ones(16).contiguous().realize()
-    @TinyJit
-    def f(a): return (a + 2).contiguous().realize()
-    f(x)
-
-    before = len(link_linear_cache)
-    with rt_views() as calls:
-      out = f(x)
-      self.assertGreater(len(link_linear_cache), before)
-      self.assertEqual(len(calls), 0)
-      (x + 1).contiguous().realize()
-      self.assertGreater(len(calls), 0)
-    self.assertEqual(out.tolist(), [3.0] * 16)
-
-  def test_jit_survives_ring_wrap(self):
-    # the ring recycles with no liveness tracking, so eager work that wraps it must not land on the jit's buffers
     dev = Device[Device.DEFAULT]
-    allocs = {host:dev.rt_allocator(True, host) for host in (False, True)}
-    for host in allocs: dev.rt_buffer(True, host) # cache the full-sized backing buffers before temporarily shrinking their allocators
-    with patch.object(allocs[False], "size", 1 << 13), patch.object(allocs[True], "size", 1 << 13):
-      x = Tensor.ones(24).contiguous().realize()
-      @TinyJit
-      def g(a): return (a * 3 - 1).contiguous().realize()
-      for _ in range(3): g(x)
+    rings = [dev.rt_buffer(True, host) for host in (False, True)]
+    ranges = [(b._buf.va_addr, b._buf.va_addr + b.nbytes) for b in rings]
+    for n in (1, 65):
+      with self.subTest(kernels=n):
+        x, f = self.input(), TinyJit(lambda a: chain(a, n).realize())
+        for _ in range(2): f(x)
+        for u in f.captured.linear.toposort():
+          if u.op is Ops.BUFFER and (buf:=u.buffer).device == dev.device:
+            addr = buf._buf.va_addr
+            self.assertFalse(any(addr < end and start < addr + buf.nbytes for start, end in ranges))
 
-      wrapped = 0
-      for i in range(48):
-        before = dev.rt_allocator(True, False).ptr
-        (x + i).contiguous().realize()
-        wrapped += dev.rt_allocator(True, False).ptr < before
-        self.assertEqual(g(x).tolist(), [2.0] * 24)
-      self.assertGreater(wrapped, 0)
+  def test_small_eager_cached(self):
+    _, compiled, inputs = self.compiled(1)
+    linked = link_linear(compiled, input_uops=inputs)
+    self.assertIs(link_linear(compiled, input_uops=inputs), linked)
+
+  def test_profile_slots_survive_indirect_access(self):
+    pm = PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat())),), allow_any_len=True),
+                          lambda buf: hcq2.rt_addr(buf, "CPU") if hcq2.unwrap_view(buf)[0].tag == "slots" else None)])
+    with patch.object(Device[Device.DEFAULT], "pm_lower", pm):
+      compiled = compile_linear(Tensor.ones(4).contiguous().schedule_linear(), profile=True)
+    self.assertFalse(any(param.op is Ops.PARAM and (param.arg.name or "").startswith("slots_")
+                         for param in compiled.src[0].without_after.src[0].toposort()))
+    call = link_linear(compiled).src[0].without_after
+    ((device, index),) = call.arg.aux.slots
+    self.assertEqual(device, Device.DEFAULT)
+    self.assertEqual(call.src[1 + index].buffer.dtype, dtypes.uint64)
+
+  def test_host_copies(self):
+    dev = Device[Device.DEFAULT]
+    if not dev.has_copy_queue: self.skipTest("copy queue required")
+    for host_device in ("CPU", "NPY", "DISK"):
+      for direct in (False, True):
+        for upload in (False, True):
+          with self.subTest(host_device=host_device, direct=direct, upload=upload):
+            host, gpu = UOp.new_buffer(host_device, 4, dtypes.uint8), UOp.new_buffer(dev.device, 4, dtypes.uint8)
+            src, dst = (host, gpu) if upload else (gpu, host)
+            linear = UOp(Ops.LINEAR, src=(src.copy_to_device(dst.device).call(dst, src),))
+            with patch.object(dev, "host_devs", frozenset({"CPU", host_device}) if direct else frozenset({"CPU"})):
+              compiled = compile_linear(linear, profile=False)
+            self.assertEqual(len(compiled.src), 1 if direct or host_device == "CPU" else 2)
+            self.assertEqual(sum(call_is_hcq(call) for call in compiled.src), 1)
+
+  def test_large_eager_not_cached(self):
+    _, compiled, inputs = self.compiled(65)
+    linked = link_linear(compiled, input_uops=inputs)
+    self.assertIsNot(link_linear(compiled, input_uops=inputs), linked)
+    self.assertNotIn(compiled, link_linear_cache)
+
+  def test_double_compile(self):
+    for n in (1, 65):
+      for jit in (False, True):
+        with self.subTest(kernels=n, jit=jit):
+          out, compiled, inputs = self.compiled(n, jit=jit)
+          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
+          before = tuple(inputs)
+          with rt_views() as borrowed:
+            for linear in (compiled, linked):
+              self.assertIs(compile_linear(linear, input_uops=None if jit else inputs), linear)
+          self.assertEqual(tuple(inputs), before)
+          self.assertFalse(borrowed)
+          run_linear(linked, input_uops=inputs, jit=True, wait=True)
+          self.assertEqual(out.tolist(), [2 + n] * 4)
+
+  def test_double_link(self):
+    for n in (1, 65):
+      for jit in (False, True):
+        with self.subTest(kernels=n, jit=jit):
+          out, compiled, inputs = self.compiled(n, jit=jit)
+          linked = link_linear(compiled, input_uops=inputs, allow_cache=not jit)
+          with rt_views() as borrowed:
+            again = link_linear(linked, input_uops=inputs, allow_cache=not jit)
+          self.assertIs(again, linked)
+          self.assertFalse(borrowed)
+          run_linear(again, input_uops=inputs, jit=True, wait=True)
+          self.assertEqual(out.tolist(), [2 + n] * 4)
 
   def test_jit_new_inputs_each_call(self):
     @TinyJit
@@ -85,6 +180,13 @@ class TestHCQ2Core(unittest.TestCase):
       vi = Variable("i", 1, 10).bind(i)
       np.testing.assert_allclose(f(a[:, :vi]).item(), (a[:, :i] + 1).sum().item(), atol=1e-5, rtol=1e-5)
 
+  def test_map_cpu_buffer_preserves_contents(self):
+    src = Buffer("CPU", 16, dtypes.uint8, preallocate=True)
+    data = bytes(range(16))
+    src.as_memoryview(force_zero_copy=True)[:] = data
+    src.get_buf(Device.DEFAULT)
+    self.assertEqual(bytes(src.as_memoryview(force_zero_copy=True)), data)
+
   def test_staged_copy_roundtrip(self):
     # a host buffer the device cannot read copies in chunks through a small ring of staging slots: every rotation must land bit-exact
     stage = Buffer("CPU", size:=1 << 16, dtypes.uint8, preallocate=True)
@@ -102,6 +204,7 @@ class TestHCQ2Core(unittest.TestCase):
       @TinyJit
       def f(a): return (a.sin() * 3).contiguous().realize()
       for _ in range(3): f(x)
+      eager_chain(x)
 
     jit, eager = partition(batches, lambda c: c.arg.aux.table >= 0)
     self.assertTrue(jit and eager, f"want both kinds of batch, got {len(jit)} jit and {len(eager)} eager")
@@ -124,9 +227,26 @@ class TestHCQ2Core(unittest.TestCase):
       return max(c.arg.aux.nargs for c in batches)
     self.assertEqual(nargs(2), nargs(12))
 
+  def test_caches_hold_no_buffers(self):
+    # an eager template caches without its buffers and the jit's linear compiles once uncached: freeing the tensors frees the device memory
+    def step(i):
+      x = Tensor(np.full(1024, i, np.float32)).to(Device.DEFAULT).realize()
+      @TinyJit
+      def f(a): return (a * 2 + 1).contiguous().realize()
+      for _ in range(3): out = f(x)
+      self.assertEqual(out.tolist(), [2.0 * i + 1] * 1024)
+    step(1) # warms the programs, templates and rings
+    gc.collect()
+    used = GlobalCounters.mem_used
+    for i in range(2, 5): step(i)
+    gc.collect()
+    self.assertEqual(GlobalCounters.mem_used, used)
+
   def test_device_state_survives_as_link_refs(self):
     # a buffer the commands only address, never a param of the body, is kept by the linked call as a ref of what its getaddr resolved into
-    dev, names = Device[Device.DEFAULT], {"AMD": ("scratch",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
+    dev = Device[Device.DEFAULT]
+    names = {"AMD": () if getattr(dev, "is_aql", False) else ("scratch",), # the aql descriptor holds the scratch, nothing addresses it
+             "NV": ("timeline",), "QCOM": ("_stack", "dummy")}[Device.DEFAULT.split(":")[0]]
     @TinyJit
     def f(a): return (a * 2 + 1).contiguous().realize()
     x = Tensor.ones(16).contiguous().realize()
@@ -136,32 +256,70 @@ class TestHCQ2Core(unittest.TestCase):
     refs = [u.buffer for u in call.src[1:] if u.op is Ops.BUFFER]
     for n in names: self.assertTrue(any(r is getattr(dev, n) for r in refs), f"{n} is not a ref of the call")
 
+  def test_usb_renumbering(self):
+    programs = []
+    with Context(HCQ_RUNTIME_DEV="CPU"), patch("tinygrad.codegen.do_to_program", wraps=do_to_program) as build:
+      for ids in ((0, 1, 2, 3), (2, 0, 3, 1), (1, 0, 2, 3), (0, 1, 3, 2), (100, 101, 102, 103)):
+        with self.subTest(ids=ids):
+          regs = [UOp.placeholder((1,), dtypes.uint32, slot=i, addrspace=AddrSpace.REG) for i in ids[:2]]
+          a, b = [r.after(r.index(0).store(v)) for r, v in zip(regs, (3, 5))]
+          i, j = [UOp.range(UOp(Ops.NOOP), n, dtype=dtypes.void, src=(a, b)) for n in ids[2:]]
+          out = cpu_buf(dtype=dtypes.uint32, tag="out")
+          body = out.index(0).store(a.after(i, j).index(0).load()*10 + b.index(0).load()).end(j, UOp.const(False)).end(i, UOp.const(False))
+          compiled = lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),)))
+          programs.append(compiled.src[0].without_after.src[0])
+          self.assertIs(programs[-1], programs[0])
+          linear = hcq2.hcq_link(compiled, allow_cache=False)
+          run_linear(linear, jit=True)
+          self.assertEqual(linear.src[0].without_after.src[1].buffer._buf.cpu_view().view(fmt='I')[0], 35)
+      self.assertLessEqual(build.call_count, 1)
+
+  def test_patched_view(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      ctx = hcq2.EncodeCtx(("CPU",))
+      inner = hcq2.patch(cpu_buf(8, tag="inner"), [(4, UOp.const(42, dtypes.uint32))], bytes(8))
+      inner = unwrap(hcq2.hoist_links(ctx, inner))
+      outer = hcq2.patch(cpu_buf(8, tag="outer"), [(0, inner[4:8].getaddr("CPU"))])
+      with patch.object(hcq2, "EncodeCtx", return_value=ctx): call = lower_hcq(outer.bitcast(dtypes.uint64).index(0).load())
+      self.assertEqual(call.without_after.arg.aux.nargs, 1)
+      self.assertTrue(all(s.op is Ops.STORE for s in call.src[1:]))
+      linked = hcq2.hcq_link(UOp(Ops.LINEAR, src=(call,)), allow_cache=False).src[0]
+      inner_buf, outer_buf = linked.src[1].buffer, linked.without_after.src[1].buffer
+      self.assertEqual(inner_buf._buf.cpu_view().view(fmt='I')[1], 42)
+      self.assertEqual(outer_buf._buf.cpu_view().view(fmt='Q')[0], inner_buf._buf.va_addr + 4)
+
 @unittest.skipUnless(isinstance(Device["CPU"].renderer, CStyleLanguage), "CALL is rendered in C style only")
 class TestHCQ2FFI(unittest.TestCase):
   @staticmethod
   def _run(body:UOp) -> list[Buffer]:
-    call = hcq2.lower_call(UOp.sink(body, arg=KernelInfo("test_ffi")).call(aux=hcq2.HCQInfo(("CPU",))))
-    assert call is not None
-    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(call,))), cache=False)
+    linear = hcq2.hcq_link(lower_and_compile(UOp(Ops.LINEAR, src=(lower_hcq(body),))), allow_cache=False)
     run_linear(linear, jit=True)
     return [u.buffer for u in linear.src[0].without_after.src[1:] if u.op is Ops.BUFFER]
 
   def test_ffi_ccall(self):
     with Context(HCQ_RUNTIME_DEV="CPU"):
-      out = UOp.placeholder((1,), dtypes.int32, slot=1, device="CPU", volatile=True, tag="ffi_result")
+      out = cpu_buf(dtype=dtypes.int32, slot=1, volatile=True, tag="ffi_result")
       bufs = self._run(out.index(0).store(hcq2.ccall(libc.dll.ffs, 0x10)))
     self.assertEqual(next(b for b in bufs if b.dtype is dtypes.int)._buf.cpu_view().view(fmt='i')[0], 5)
 
   def test_ffi_cstruct(self):
     struct_t = init_c_struct_t(16, (("u8", ctypes.c_uint8, 0), ("u16", ctypes.c_uint16, 2),
                                   ("u32", ctypes.c_uint32, 4), ("u64", ctypes.c_uint64, 8)))
-    UOp.placeholder((1,), dtypes.uint8, device="CPU") # reserve slot zero for device-owned placeholders
+    cpu_buf() # reserve slot zero for device-owned placeholders
     with Context(HCQ_RUNTIME_DEV="CPU"):
       s = hcq2.cstruct(struct_t, u8=0x12, u16=UOp.const(0x3456, dtypes.uint16), u32=0x789ABCDE, u64=0xFEDCBA9876543210)
       bufs = self._run(s.index(0).load())
     got = struct_t.from_buffer_copy(bytes(next(b for b in bufs if b.nbytes == ctypes.sizeof(struct_t))._buf.cpu_view()))
     self.assertEqual((got.u8, got.u16, got.u32, got.u64), (0x12, 0x3456, 0x789ABCDE, 0xFEDCBA9876543210))
 
+  def test_nested_cstruct_patches(self):
+    with Context(HCQ_RUNTIME_DEV="CPU"):
+      inner = hcq2.cstruct(init_c_struct_t(4, (("value", ctypes.c_uint32, 0),)), value=42)
+      outer = hcq2.cstruct(init_c_struct_t(8, (("ptr", ctypes.c_uint64, 0),)), ptr=inner.getaddr("CPU"))
+      out = cpu_buf(dtype=dtypes.uint32, tag="result")
+      copied = hcq2.ccall(libc.memcpy, out.index(0), outer.bitcast(dtypes.uint64).index(0).load(), 4)
+      bufs = self._run(out.after(copied).index(0).load())
+    self.assertEqual(next(b for b in bufs if b.dtype is dtypes.uint32)._buf.cpu_view().view(fmt='I')[0], 42)
 
 if __name__ == "__main__":
   unittest.main()

test/external/external_test_usb_asm24.py

@@ -1,15 +1,22 @@
 import unittest
 from tinygrad.helpers import Timing, getenv
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, TinyJit
+from tinygrad.runtime.support.usb import HALF, CHUNK, SLOT
 import numpy as np
 
-class TestDevCopySpeeds(unittest.TestCase):
+class USBTestCase(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.sz = getenv("SIZE", 2000000)
     cls.dev = Device["AMD"]
-    if not cls.dev.is_usb(): raise unittest.SkipTest("only test this on USB devices")
+    if not cls.dev.is_usb: raise unittest.SkipTest("only test this on USB devices")
+    cls.rng = np.random.default_rng(0)
 
+  def roundtrip(self, a:np.ndarray): # a copy in, a kernel, a copy out: the queue must order them
+    np.testing.assert_array_equal(a, Tensor(a, device="NPY").to(Device.DEFAULT).numpy())
+    np.testing.assert_array_equal(a + 1, (Tensor(a, device="NPY").to(Device.DEFAULT) + 1).numpy())
+
+class TestDevCopySpeeds(USBTestCase):
   def testCopyCPUtoDefault(self):
     for _ in range(10):
       t = Tensor.ones(self.sz, device="CPU", dtype='uchar').contiguous().realize()
@@ -24,15 +31,60 @@ class TestDevCopySpeeds(unittest.TestCase):
       with Timing(f"copyout of {t.nbytes()/1e6:.2f} MB:  ", on_exit=lambda ns: f" @ {t.nbytes()/ns * 1e3:.2f} MB/s"):
         t.to('CPU').realize()
 
+class TestUSBIntegrity(USBTestCase):
   def testValidateCopies(self):
     t = Tensor.randn(self.sz, device="CPU", dtype='uchar').contiguous().realize()
     x = t.to(Device.DEFAULT).realize()
     Device[Device.DEFAULT].synchronize()
-
     y = x.to('CPU').realize()
-
     np.testing.assert_equal(t.numpy(), y.numpy())
-    del x, y, t
+
+  def testBoundaries(self): # around the slot, the chunk and the read window
+    for size in (1, 3, 508, 509, SLOT - 513, SLOT - 512, SLOT - 511, CHUNK - 1, CHUNK, CHUNK + 1, 2 * CHUNK - 1, 2 * CHUNK, 2 * CHUNK + 31, HALF,
+                 2 * HALF, 1 << 20):
+      with self.subTest(size=size): self.roundtrip(self.rng.integers(0, 256, size, dtype=np.uint8))
+
+  def testManyCopiesInABatch(self):
+    for n in (2, 7, 64, 300): # 300 chunks: the fence byte wraps
+      with self.subTest(n=n):
+        arrs = [self.rng.integers(0, 256, int(s), dtype=np.uint8) for s in self.rng.integers(1, 5000, n)]
+        ts = [Tensor(a, device="NPY").to(Device.DEFAULT) for a in arrs]
+        Tensor.realize(*ts)
+        for t, a in zip(ts, arrs): np.testing.assert_array_equal(a, t.numpy())
+
+  def testMixedBatch(self): # copies out and in, in one batch: runs of both directions
+    arrs = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
+    ts = [Tensor(a, device="NPY").to(Device.DEFAULT).realize() for a in arrs]
+    more = [self.rng.integers(0, 256, s, dtype=np.uint8) for s in (5, CHUNK + 7, 9, 2 * CHUNK + 3, 11)]
+    outs = [t.to("NPY") for t in ts] + [Tensor(a, device="NPY").to(Device.DEFAULT) for a in more]
+    Tensor.realize(*outs)
+    for o, a in zip(outs, arrs + more): np.testing.assert_array_equal(a, o.numpy())
+
+  def testRepeatedBatches(self): # a batch numbers its chunks from 0: the same batch again must not see what the last one left behind
+    a = self.rng.integers(0, 256, 2 * CHUNK + 31, dtype=np.uint8)
+    for _ in range(5): self.roundtrip(a)
+    @TinyJit
+    def step(x:Tensor) -> Tensor: return (x + 1).realize()
+    src = Tensor(a, device="NPY")
+    for i in range(5):
+      x = src.to(Device.DEFAULT)
+      np.testing.assert_array_equal(a + 1, step(x).numpy())
+
+  def testStaleSentinel(self): # payloads full of the tags the queue waits for, in both directions, before and around the real chunks
+    tags = np.array([0x51000000 | k for k in range(8)], dtype=np.uint32)
+    for tag in tags: # every dword of every chunk is the tag of some chunk of the copy
+      with self.subTest(payload=hex(tag)):
+        a = np.full((2 * CHUNK + 31) // 4, tag, dtype=np.uint32).view(np.uint8)
+        self.roundtrip(a)
+    with self.subTest(case="copyout residue"): # a read fills the sram with tags, then small chunks land in both halves
+      a = np.tile(tags, 2 * CHUNK // 32).view(np.uint8)
+      np.testing.assert_array_equal(a, (Tensor(a, device="NPY").to(Device.DEFAULT) * 1).numpy())
+      for size in (31, CHUNK + 31, 2 * CHUNK + 31): self.roundtrip(np.tile(tags, size // 32 + 1).view(np.uint8)[:size])
+
+  def testRingWrap(self): # 64MB of chunks: the sdma ring (1MB on usb) wraps within the copy
+    a = self.rng.integers(0, 256, 64 << 20, dtype=np.uint8)
+    t = Tensor(a, device="NPY").to(Device.DEFAULT).realize()
+    np.testing.assert_array_equal(a, t.numpy())
 
 if __name__ == "__main__":
   unittest.main()

test/external/fuzz_fast_idiv.py

@@ -12,7 +12,8 @@ if __name__ == "__main__":
     if i % 1000 == 0:
       print(f"Progress: {i}")
     dt = random.choice(dtypes.ints)
-    u = UOp.variable('x', random.randint(dt.min, 0), random.randint(1, dt.max), dtype=dt)
+    vmax = random.randint(1, 2**random.randint(1, dt.max.bit_length()))
+    u = UOp.variable('x', random.randint(0, vmax-1) if vmax > 1 else 0, vmax, dtype=dt)
     d = random.randint(1, max(1, u.vmax)*2)
     if d in powers_of_two: continue
     expr = fast_idiv(Device[Device.DEFAULT].renderer, u, d)

test/external/speed_v_theoretical.py

@@ -1,6 +1,10 @@
 import unittest
-from tinygrad import Tensor, TinyJit, Device
-from tinygrad.helpers import Context, DEBUG, GlobalCounters
+from dataclasses import replace
+from itertools import islice
+from tinygrad import Tensor, Device
+from tinygrad.codegen import to_program
+from tinygrad.engine.realize import time_call
+from tinygrad.helpers import Context, DEBUG
 from tinygrad.nn import Conv2d
 from tinygrad.nn.state import get_parameters
 
@@ -10,6 +14,13 @@ class TestKernelSpeed(unittest.TestCase):
       # TODO: randn is 20% faster than rand for gemv
       return Tensor.randn(shape, dtype="half").realize()
 
+  def _time_kernel(self, out:Tensor, beam:int):
+    linear = out.schedule_linear()
+    self.assertEqual(len(linear.src), 1, "expected a single kernel")
+    call = linear.src[0]
+    prg = to_program(call.src[0].replace(arg=replace(call.src[0].arg, beam=beam)), Device[out.device].renderer)
+    return min(islice(time_call(call.replace(src=(prg, *call.src[1:])), clear_l2=True), 3, 10))
+
   def _compare(self, tm, tflops, gbs, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
       if DEBUG >= 1:
         print(f"{tm=:.6f}")
@@ -34,53 +45,30 @@ class TestKernelSpeed(unittest.TestCase):
 
   def _test_matmul(self, M, K=None, N=None, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
     # (MxK) @ (KxN)
-    @TinyJit
-    def f(a, b) -> Tensor: return (a @ b).realize()
-
     if N is None: N = M
     if K is None: K = M
-    tms = []
-    with Context(BEAM=3):
-      for i in range(10):
-        a = self._get_tensor(M, K)
-        b = self._get_tensor(K, N)
-        if i >= 3:
-          GlobalCounters.time_sum_s = 0
-          with Context(DEBUG=max(DEBUG.value, 2)): c = f(a, b)
-          tms.append(GlobalCounters.time_sum_s)
-        else:
-          c = f(a, b)
+    a = self._get_tensor(M, K)
+    b = self._get_tensor(K, N)
+    tm = self._time_kernel(c:=a @ b, beam=3)
 
     ops = 2 * M * N * K
     mems = a.dtype.itemsize * M * K + b.dtype.itemsize * K * N + c.dtype.itemsize * M * N
-    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
   def _test_conv_3x3(self, BS, CIN, COUT, H, W, nv_tflops=None, nv_gbs=None, amd_tflops=None, amd_gbs=None):
-    @TinyJit
-    def f(conv, x) -> Tensor: return conv(x).realize()
-    tms = []
     K = 3
     with Context(BEAM=0, DEBUG=0):
       conv = Conv2d(CIN, COUT, K, padding=1)
       Tensor.realize(*get_parameters(conv))
 
-    with Context(BEAM=2):
-      for i in range(10):
-        x = self._get_tensor(BS, CIN, H, W)
-        if i >= 3:
-          GlobalCounters.time_sum_s = 0
-          with Context(DEBUG=max(DEBUG.value, 2)): _c = f(conv, x)
-          tms.append(GlobalCounters.time_sum_s)
-        else:
-          _c = f(conv, x)
+    x = self._get_tensor(BS, CIN, H, W)
+    tm = self._time_kernel(_c:=conv(x), beam=2)
 
     # naive algo
     ops = 2 * BS * CIN * COUT * K * K * H * W
     mems = x.nbytes() + conv.weight.nbytes() + conv.bias.nbytes() + _c.nbytes()
-    tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)

test/helpers.py

@@ -65,6 +65,10 @@ def assert_kernel_count(expected:int):
   got = GlobalCounters.kernel_count
   if got != expected: raise KernelCountException(expected, got)
 
+def is_hcq2_device() -> bool: # an hcq2 device stages every copy from the host through a pinned buffer: such a copy is two calls, not one
+  from tinygrad.runtime.support.hcq2 import HCQ_DEVS
+  return Device.DEFAULT.split(":")[0] in HCQ_DEVS
+
 def call_is_graph(call:UOp) -> bool:
   ast = call.src[0]
   return ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "graph"

test/mockgpu/amd/amdgpu.py

@@ -327,8 +327,11 @@ class SDMAExecutor(AMDQueue):
 
   def _execute_copy(self):
     struct = sdma_pkts.copy_linear.from_address(self.base + self.rptr[0] % self.size)
-    count_cnt = to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF
-    ctypes.memmove(self.gpu.translate_addr(struct.dst_addr), self.gpu.translate_addr(struct.src_addr), count_cnt + 1)
+    count, off = (to_mv(self.base + self.rptr[0] % self.size + 4, 4).cast('I')[0] & 0x3FFFFFFF) + 1, 0
+    while off < count: # a page at a time: the physical pages of a range needn't be contiguous
+      n = min(count - off, 0x1000 - ((struct.src_addr + off) & 0xfff), 0x1000 - ((struct.dst_addr + off) & 0xfff))
+      ctypes.memmove(self.gpu.translate_addr(struct.dst_addr + off), self.gpu.translate_addr(struct.src_addr + off), n)
+      off += n
     self.rptr[0] += ctypes.sizeof(struct)
 
 class AMDGPURegisters:

test/mockgpu/amd/emu.py

@@ -69,7 +69,7 @@ from tinygrad.runtime.autogen.amd.cdna import ins as irc
 from tinygrad.renderer.amd.dsl import VCC_LO, EXEC_LO, SCC, ttmp, Inst
 from tinygrad.runtime.autogen.amd.common import Fmt, OpType
 from test.amd.helpers import decode_dpp16
-from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32
+from test.mockgpu.amd.pcode import parse_pcode, _FUNCS, _set_bits, _to_bool, _to_u32, _val_to_bits, _ftz_f32, _bitreverse, _countbits
 
 MASK32 = 0xFFFFFFFF
 
@@ -1566,9 +1566,19 @@ def _compile_mem_op(inst: ir3.DS|ir3.FLAT|ir3.GLOBAL|ir3.SCRATCH|ir4.DS|ir4.VFLA
   has_data1 = is_lds and hasattr(inst, 'data1') and inst.data1 is not None
   data1_reg = ctx.inst_field(type(inst).data1) if is_lds else _c(0)  # type: ignore[union-attr]
 
+  if is_lds and op_name == 'DS_SWIZZLE_B32':
+    # The manual's reverse_bits operates on five-bit lane indices; thread indices wrap within the wave.
+    funcs = {'reverse_bits': lambda x: _bitreverse(x, 32) >> _c(27), 'count_ones': _countbits,
+             'thread_in': lambda x: ctx.rvgpr_dyn(addr_reg, x & _c(ctx.wave_size - 1)),
+             'thread_valid': lambda x: _lane_active(exec_mask, x & _c(ctx.wave_size - 1))}
+    result, _ = parse_pcode(pcode, {'offset0': offset0.cast(dtypes.uint8), 'offset1': offset1.cast(dtypes.uint8)}, funcs)
+    values = [result[f'thread_out@{i}'] for i in range(ctx.wave_size)]
+    # Snapshot every source before writing: destination and source registers may be identical.
+    reads = UOp(Ops.STACK, src=tuple(values))
+    return UOp.sink(*(ctx.wvgpr_dyn(vdst_reg, _c(i), val, exec_mask, after=reads) for i, val in enumerate(values)), *ctx.inc_pc())
+
   # DS_PERMUTE/DS_BPERMUTE: cross-lane VGPR access via pcode
   if is_lds and 'PERMUTE' in op_name:
-    pcode = get_pcode(inst.op)
     srcs = {'ADDR': addr_reg, 'DATA0': vdata_reg, 'VDST': vdst_reg, 'OFFSET': offset,
             'EXEC': exec_mask.cast(dtypes.uint64), '_vgpr': ctx.vgpr, '_wave_size': ctx.wave_size}
     _, assigns = parse_pcode(pcode, srcs)

test/mockgpu/amd/pcode.py

@@ -630,6 +630,10 @@ class Parser:
         self.eat('DOT')
         dt_name = self.eat('IDENT').val
         return self._handle_mem_load(addr, DTYPES.get(dt_name, dtypes.uint32))
+      if name in self.funcs and self.try_eat('LBRACKET'):
+        index = self.parse()
+        self.eat('RBRACKET')
+        return self.funcs[name](index)
       if name == 'VGPR' and self.at('LBRACKET'):
         self.eat('LBRACKET')
         lane = self.parse()
@@ -1006,20 +1010,24 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 
     # for loop
     if first == 'for':
-      # Parse: for VAR in [SIZE']START : [SIZE']END do
-      p = Parser(toks, env, funcs)
-      p.eat_val('for', 'IDENT')
-      loop_var = p.eat('IDENT').val
-      p.eat_val('in', 'IDENT')
-      def parse_bound():
-        if p.at('NUM') and p.peek(1).type == 'QUOTE':
-          p.eat('NUM')
-          p.eat('QUOTE')
-        if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
-        return int(p.parse())
-      start_val = parse_bound()
-      p.eat('COLON')
-      end_val = parse_bound()
+      # C-style loops use an exclusive bound; for/in loops use an inclusive bound.
+      if m := re.fullmatch(r'for\s*\(\s*(\w+)\s*=\s*(\d+);\s*\1\s*<\s*(\d+);\s*\1\s*(\+\+|\+=\s*\d+)\s*\)', line):
+        loop_var, start_val, end_val = m[1], int(m[2]), int(m[3]) - 1
+        step = 1 if m[4] == '++' else int(m[4][2:])
+      else:
+        p = Parser(toks, env, funcs)
+        p.eat_val('for', 'IDENT')
+        loop_var = p.eat('IDENT').val
+        p.eat_val('in', 'IDENT')
+        def parse_bound():
+          if p.at('NUM') and p.peek(1).type == 'QUOTE':
+            p.eat('NUM')
+            p.eat('QUOTE')
+          if p.at('NUM'): return int(p.eat('NUM').val.rstrip('UuLl'))
+          return int(p.parse())
+        start_val = parse_bound()
+        p.eat('COLON')
+        end_val, step = parse_bound(), 1
       # Collect body
       i += 1
       body_lines: list[str] = []
@@ -1035,7 +1043,7 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       has_break = any('break' in bl.lower() for bl in body_lines)
       found_var = f'_found_{next(_break_var_ids)}' if has_break else None
       if found_var: env[found_var] = block_assigns[found_var] = _const(dtypes.bool, False)
-      for loop_i in range(start_val, end_val + 1):
+      for loop_i in range(start_val, end_val + 1, step):
         subst_lines = [_subst_loop_var(bl, loop_var, loop_i) for bl in body_lines if not (has_break and bl.strip().lower() == 'break')]
         _, iter_assigns, _ = parse_block(subst_lines, 0, {**env, **block_assigns}, funcs, assigns)
         if has_break:
@@ -1224,9 +1232,9 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
       var = toks[0].val
       j, idx_toks = _match_bracket(toks, 1)
       if j < len(toks) and toks[j].type == 'EQUALS':
+        idx_expr = parse_tokens(idx_toks, env, funcs)
         # Static index: var[NUM] = value
-        if len(idx_toks) == 1 and idx_toks[0].type == 'NUM':
-          idx = int(idx_toks[0].val.rstrip('UuLl'))
+        if isinstance(idx := _single_value(idx_expr), int):
           val = parse_tokens(toks[j+1:], env, funcs)
           existing = block_assigns.get(var, env.get(var))
           if existing is not None and isinstance(existing, UOp):
@@ -1238,7 +1246,6 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
         # Dynamic index: var[expr] = value where var has @-elements
         elems = [(k.split('@')[1], v) for k, v in {**env, **block_assigns}.items() if k.startswith(f'{var}@') and isinstance(v, UOp)]
         if elems:
-          idx_expr = parse_tokens(idx_toks, env, funcs)
           val = parse_tokens(toks[j+1:], env, funcs)
           for elem_idx_str, old_elem in elems:
             elem_idx = int(elem_idx_str)
@@ -1407,16 +1414,29 @@ def parse_block(lines: list[str], start: int, env: dict[str, VarVal], funcs: dic
 def parse_expr(expr: str, env: dict[str, VarVal], funcs: dict | None = None) -> UOp:
   return parse_tokens(tokenize(expr.strip().rstrip(';')), env, funcs)
 
-def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None) -> tuple[dict, list]:
+def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dict | None = None) -> tuple[dict, list]:
   env: dict = srcs.copy() if srcs else {}
   assigns: list[tuple[str, UOp]] = []
-  raw_lines = [l.strip().rstrip(';') for l in pcode.split('\n') if l.strip() and not l.strip().startswith('//')]
-  # TODO: pcode.py should tokenize full pcode string instead of line-by-line, then this hack can be removed
   lines: list[str] = []
-  for l in raw_lines:
-    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] = lines[-1] + ' ' + l
-    else: lines.append(l)
-  _, final, _ = parse_block(lines, 0, env, assigns=assigns)
+  blocks: list[str] = []
+  for raw in pcode.splitlines():
+    line = raw.split('//')[0].strip().rstrip(';')
+    if not line: continue
+    # Both block syntaxes share the same parser; braces supply the implicit end markers.
+    if line.startswith('}') and blocks:
+      end = blocks.pop()
+      line = line[1:].strip()
+      if not line.startswith(('elsif', 'else')): lines.append(end)
+    if m := re.match(r'(if|elsif|else|for)\b.*\{$', line):
+      blocks.append('endfor' if m[1] == 'for' else 'endif')
+      line = line[:-1].rstrip()
+      if m[1] in ('if', 'elsif'): line += ' then'
+    if not line: continue
+    line = re.sub(r'=\s*(\w+):(\w+)$', r'= {\1, \2}', line)
+    if lines and re.search(r'(&&|\|\||[&|+\-*/^])\s*$', lines[-1]): lines[-1] += ' ' + line
+    else: lines.append(line)
+  assert not blocks, "unclosed pcode block"
+  _, final, _ = parse_block(lines, 0, env, {**_FUNCS, **funcs} if funcs else None, assigns=assigns)
   sliced = set(d.split('[')[0] for d, _ in assigns if '[' in d)
   for var, val in final.items():
     if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA'] and isinstance(val, UOp):

test/mockgpu/nv/nvdriver.py

@@ -53,6 +53,7 @@ class NVDriver(VirtDriver):
                            VirtFile('/dev/nvidia-uvm', functools.partial(NVUVMFileDesc, driver=self))]
 
     self.root_handle = None
+    self.host_ranges: set[int] = set()
 
     self.gpus = {}
     self.next_fd = (1 << 29)
@@ -251,7 +252,9 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_ENABLE_PEER_ACCESS: pass # uvm and shared spaced are setup already, no emulation for now
     elif nr == nv_gpu.UVM_CREATE_EXTERNAL_RANGE:
       st = nv_gpu.UVM_CREATE_EXTERNAL_RANGE_PARAMS.from_address(argp)
-      libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
+      # Registered host memory already has a CPU mapping; MAP_FIXED would discard its contents.
+      if st.base not in self.host_ranges:
+        libc.mmap(st.base, st.length, mmap.PROT_READ|mmap.PROT_WRITE, libc.MAP_FIXED|mmap.MAP_SHARED|mmap.MAP_ANONYMOUS, -1, 0)
     elif nr == nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION:
       st = nv_gpu.UVM_MAP_EXTERNAL_ALLOCATION_PARAMS.from_address(argp)
       for gpu_attr_id in range(st.gpuAttributesCount):
@@ -265,6 +268,7 @@ class NVDriver(VirtDriver):
     elif nr == nv_gpu.UVM_REGISTER_CHANNEL: pass
     elif nr == nv_gpu.UVM_FREE:
       st = nv_gpu.UVM_FREE_PARAMS.from_address(argp)
+      self.host_ranges.discard(st.base)
       libc.munmap(st.base, st.length)
     else: raise RuntimeError(f"Unknown {nr} to nvidia-uvm")
     return 0
@@ -276,6 +280,7 @@ class NVDriver(VirtDriver):
       st:Any = nv_gpu.nv_ioctl_nvos02_parameters_with_fd.from_address(argp)
       # Track host memory (signal memory) - progress queues when written to
       if st.params.hClass == nv_gpu.NV01_MEMORY_SYSTEM_OS_DESCRIPTOR:
+        self.host_ranges.add(st.params.pMemory)
         self.track_address(st.params.pMemory, st.params.pMemory + st.params.limit + 1,
                            lambda mv,off: None, lambda mv, off: self._gpu_mmio_write(mv, off, None))
     return 0

test/mockgpu/nv/nvgpu.py

@@ -100,11 +100,11 @@ class GPFIFO:
     if qmd.release0_enable:
       rel0 = to_mv(qmd.release0_address_lower + (qmd.release0_address_upper << 32), 0x10).cast('Q')
       rel0[0] = qmd.release0_payload_lower + (qmd.release0_payload_upper << 32)
-      rel0[1] = int(time.perf_counter() * 1e9)
+      if qmd.release0_structure_size == 0: rel0[1] = int(time.perf_counter() * 1e9) # four words: the timestamp after the payload
     if qmd.release1_enable:
       rel1 = to_mv(qmd.release1_address_lower + (qmd.release1_address_upper << 32), 0x10).cast('Q')
       rel1[0] = qmd.release1_payload_lower + (qmd.release1_payload_upper << 32)
-      rel1[1] = int(time.perf_counter() * 1e9)
+      if qmd.release1_structure_size == 0: rel1[1] = int(time.perf_counter() * 1e9)
     if qmd.dependent_qmd0_enable:
       if qmd.dependent_qmd0_action == 1: self.execute_qmd(qmd.dependent_qmd0_pointer << 8)
       else: raise RuntimeError("unsupported dependent qmd action")
@@ -192,11 +192,10 @@ class GPFIFO:
       sz = self._state(nv_gpu.NVC6B5_LINE_LENGTH_IN)
       assert flags == 0x182, f"unsupported flags in _exec_nvc6b5_dma: {flags}"
       ctypes.memmove(dst, src, sz)
-    elif ((flags >> 3) & 0b11) != 0:
-      src = to_mv(self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 0x10).cast('Q')
-      val = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
-      src[0] = val
-      src[1] = int(time.perf_counter() * 1e9)
+    elif (semaphore_type:=((flags >> 3) & 0b11)) != 0:
+      to_mv(addr:=self._state64(nv_gpu.NVC6B5_SET_SEMAPHORE_A), 4).cast('I')[0] = self._state(nv_gpu.NVC6B5_SET_SEMAPHORE_PAYLOAD)
+      if semaphore_type == nv_gpu.NVC6B5_LAUNCH_DMA_SEMAPHORE_TYPE_RELEASE_FOUR_WORD_SEMAPHORE:
+        to_mv(addr + 8, 8).cast('Q')[0] = int(time.perf_counter() * 1e9)
     else: raise RuntimeError("unknown nvc6b5_dma flags")
 
   def _exec_pcas2(self):

test/null/test_device.py

@@ -89,7 +89,7 @@ class TestDevice(unittest.TestCase):
       except Exception as e: self.skipTest(f"skipping compiler test: not all compilers: {e}")
 
       imports = ("from tinygrad import Device; from tinygrad.runtime.support.compiler_amd import HIPCompiler; "
-                 "from tinygrad.runtime.support.compiler_amd import AMDLLVMCompiler")
+                 "from tinygrad.runtime.support.compiler_llvm import AMDLLVMCompiler")
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, AMDLLVMCompiler)"'],
                         shell=True, check=True, env={**os.environ, "DEV": "AMD:LLVM"})
       subprocess.run([f'python3 -c "{imports}; assert isinstance(Device[Device.DEFAULT].compiler, HIPCompiler)"'],

test/null/test_helpers.py

@@ -78,6 +78,13 @@ class TestContextVars(unittest.TestCase):
     test()
     self.assertEqual(VARIABLE.value, 0)
 
+  def test_decorator_recursive(self):
+    @Context(VARIABLE=1)
+    def test(n):
+      if n: test(n-1)
+    test(2)
+    self.assertEqual(VARIABLE.value, 0)
+
   def test_context_exit_reverts_updated_values(self):
     D = ContextVar("D", 1)
     D.value = 2

test/null/test_schedule.py

@@ -582,8 +582,8 @@ class TestSchedule(unittest.TestCase):
     p = P[0]
     p = p.pad(((1, 0), ))
     p = p.repeat([2])
-    # TODO: this should be 3 if fix store hazard worked correctly
-    check_schedule(p, 4)
+    # assign on a pending contiguous overwrites the whole value, no store hazard
+    check_schedule(p, 3)
 
   def test_conv2d(self, allowed=4, dtype=dtypes.float):
     self.enterContext(Context(DEFAULT_FLOAT=dtype))

test/null/test_tensor_uop_mixin.py

@@ -246,8 +246,10 @@ class TestTensorUOpRand(unittest.TestCase):
     self.assertIs(Tensor._threefry_random_bits(Tensor(key), Tensor(c0), Tensor(c1)).uop, UOp._threefry_random_bits(key, c0, c1))
   def test_rand(self):
     k, c = UOp.empty((2,), dtype=dtypes.uint32), UOp.zeros(2, dtype=dtypes.uint32)
-    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop, UOp._rand(k, c, (2, 2), dtypes.float32))
-    self.assertIs(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop, UOp._rand(k, c, (0, 3), dtypes.float32))
+    self.assertIs(_strip_unique(Tensor._rand(Tensor(k), Tensor(c), (2, 2), dtypes.float32).uop),
+                  _strip_unique(UOp._rand(k, c, (2, 2), dtypes.float32)))
+    self.assertIs(_strip_unique(Tensor._rand(Tensor(k), Tensor(c), (0, 3), dtypes.float32).uop),
+                  _strip_unique(UOp._rand(k, c, (0, 3), dtypes.float32)))
 
 class TestTensorUOpGather(unittest.TestCase):
   def _check(self, t, dim, idx):

test/null/test_uop_graph.py

@@ -406,6 +406,18 @@ class TestUOpGraph(unittest.TestCase):
     a = c.after(e)
     self.assertNotIn(r, a.ranges)
 
+  def test_external_call_preserves_ranges(self):
+    r = UOp.range(4, 0, dtype=dtypes.int)
+    fn = UOp.custom_function("external", UOp.const(0, dtypes.uint64))
+    call = fn.call(r + 1, ret_dtype=dtypes.int)
+    self.assertEqual(set(call.ranges), {r})
+
+  def test_conditional_end_preserves_outer_range(self):
+    outer, inner = UOp.range(4, 0), UOp.loop(1)
+    end = UOp.const(1).end(inner, outer < 2)
+    self.assertEqual(set(end.ranges), {outer})
+    self.assertEqual(set((outer + 1).after(end).ranges), {outer})
+
 class TestReduceCollapse(unittest.TestCase):
   def test_multi_range_reduce_add(self):
     """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""

test/null/test_uop_symbolic.py

@@ -919,6 +919,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a % -8) // 2, -4, 0, "(a%-8//2)")
     self.helper_test_variable((a % -8) % 2, 0, 1, "(a%2)")
 
+  def test_nested_div_mod_symbolic_inner_divisor(self):
+    a = Variable("a", 0, 100)
+    self.helper_test_variable((a % (Variable("n", 1, 10)*4)) // 2, 0, 19, "(a//2%(n*2))")
+    check_uop_against_string(self, (a % (Variable("n", 0, 10)*4) // 2).simplify(), "(a%(n*4)//2)")
+
   def test_floordiv_lt_negative_c(self):
     # x//d<c with negative c also reduces to x<c*d for d>0
     idx = Variable("idx", -20, 20)

test/null/test_uop_vmin_vmax.py

@@ -167,6 +167,17 @@ class TestVminVmaxProperties(unittest.TestCase):
     self.assertEqual(UOp.const(4.5).cast(dtypes.float).cast(dtypes.int)._min_max, (4, 4))
     x = UOp.const(4.5).cast(dtypes.float)
     self.assertIs(x.ne(x.cast(dtypes.int).cast(dtypes.float)).simplify().arg, True)
+    # a source reaching past the destination clamps to its edge
+    self.assertEqual(UOp.variable('x', 2e9, 3e9, dtypes.float).cast(dtypes.int)._min_max, (2000000000, dtypes.int.max))
+    # a source entirely past the destination has no value in it
+    self.assertEqual(UOp.variable('x', 3e9, 4e9, dtypes.float).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.variable('x', -4e9, -3e9, dtypes.float).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.variable('x', 200, 300, dtypes.int).cast(dtypes.char)._min_max, (dtypes.char.min, dtypes.char.max))
+    self.assertEqual(UOp.const(300, dtypes.char)._min_max, (dtypes.char.min, dtypes.char.max))
+    self.assertEqual(UOp.const(math.inf).cast(dtypes.int)._min_max, (dtypes.int.min, dtypes.int.max))
+    self.assertEqual(UOp.const(math.nan, dtypes.float)._min_max, (-math.inf, math.inf))
+    # a weak destination has no width to clamp to
+    self.assertEqual(UOp.variable('x', 5, 7, dtypes.int).cast(dtypes.weakfloat)._min_max, (5, 7))
 
   def test_vmin_vmax_cast_int_to_float_grid(self):
     # a cast to float only takes values on the float grid, so its bounds are the source bounds rounded at the destination

test/null/test_uops.py

@@ -269,7 +269,6 @@ class TestGatedStoreRewrite(unittest.TestCase):
     for x in gated_uops: self.assertIs(x.op, Ops.STORE)
     for x in gated_uops: self.assertEqual(len(x.src), 2)
 
-@unittest.skipIf(Device.DEFAULT == "METAL", "compiler bug")
 @unittest.skipUnless(Ops.SHR in Device[Device.DEFAULT].renderer.code_for_op, "fast_idiv requires SHR")
 class TestFastIdiv(unittest.TestCase):
   def test_division_power_of_two(self):
@@ -310,7 +309,7 @@ class TestFastIdiv(unittest.TestCase):
       self.assertNotIn(Ops.FLOORDIV, ops, f"For dtype={dt} FLOORDIV survived past late rewrite")
 
   @Context(DISABLE_FAST_IDIV=0)
-  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
+  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "fast_idiv widens uint32 to uint64")
   def test_fast_idiv_and_mod(self):
     g = UOp.param(0, dtypes.uint32, 4)
     c = UOp.const(3)
@@ -329,6 +328,25 @@ class TestFastIdiv(unittest.TestCase):
     self.assertIn(Ops.SHR, ops)
     self.assertNotIn(Ops.CMOD, ops)
 
+  @Context(DISABLE_FAST_IDIV=0)
+  def test_fast_idiv_nonpositive_divisor(self):
+    ridx = UOp.range(20, 0)
+    for d in (-3, 0):
+      for op in (Ops.CDIV, Ops.CMOD):
+        ops = [x.op for x in to_uops_list([ridx.alu(op, UOp.const(d))], ren=Device[Device.DEFAULT].renderer)]
+        self.assertNotIn(Ops.SHR, ops, f"fast_idiv fired on {op} by {d}")
+
+  @Context(DISABLE_FAST_IDIV=0)
+  @unittest.skipUnless(dtypes.uint64 in Device[Device.DEFAULT].renderer.supported_dtypes(), "needs a uint64 buffer")
+  def test_fast_idiv_cmod_kept_when_idiv_declines(self):
+    ren = Device[Device.DEFAULT].renderer
+    d = UOp.param(0, dtypes.int32, 4).index(UOp.const(0))
+    ops = [x.op for x in to_uops_list([UOp.range(30, 0).alu(Ops.CMOD, d)], ren=ren)]
+    self.assertIn(Ops.CMOD, ops, "CMOD by a non-const divisor should be left alone")
+    big = UOp.param(1, dtypes.uint64, 4).index(UOp.const(0))
+    ops = [x.op for x in to_uops_list([big.alu(Ops.CMOD, UOp.const(3, dtypes.uint64))], ren=ren)]
+    self.assertIn(Ops.CMOD, ops, "CMOD should be left alone when fast_idiv declines")
+
   @Context(DISABLE_FAST_IDIV=0)
   def test_fast_idiv_bounded_numerator_zero(self):
     x = UOp.variable("x", 0, 1, dtype=dtypes.int32)
@@ -342,6 +360,7 @@ class TestFastIdiv(unittest.TestCase):
     # this requires shifting out the powers of two before doing fast_idiv
     # (((ridx0>>6)*18725)>>17) instead of (int)((((long)(ridx0)*1198373)>>29))
     self.assertNotIn(dtypes.long, [x.dtype for x in uops])
+    self.assertNotIn(Ops.CDIV, [x.op for x in uops])
 
   @unittest.expectedFailure
   def test_fast_idiv_overflow(self):

test/null/test_validate_oob.py

@@ -122,6 +122,35 @@ class TestValidateOOB(unittest.TestCase):
       r = UOp.range(20, 0)
       i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
       to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])
+      # a float entirely out of the int range has no value, not an empty one
+      f = UOp.variable("f", 3e9, 4e9, dtypes.float32, param=True).cast(dtypes.int)
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(f).load()])
+
+  def test_float_cast_in_mask(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int, 1)
+      r = UOp.range(20, 0)
+      unknown = r.cast(dtypes.float).cast(dtypes.bool)  # a bool from a float is unconstrained
+      to_uops_list([buf.index(r.valid((r < 1) & unknown)).load()])
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.valid(unknown)).load()])
+
+  def test_bitcast_in_index(self):
+    with Context(CHECK_OOB=1, SPEC=2):
+      buf = UOp.param(0, dtypes.int, 16)
+      r = UOp.range(16, 0)
+      # the WEBGPU shift: int -> uint, shift, back to int
+      i = (r.cast(dtypes.int).bitcast(dtypes.uint) << UOp.const(1).cast(dtypes.uint)).bitcast(dtypes.int)
+      to_uops_list([buf.index(i.valid(i < 16)).load()])
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(i).load()])  # 0..30 oob
+      # a negative char reads as a large uchar
+      c = Variable("c", -128, -113).cast(dtypes.char)
+      to_uops_list([UOp.param(1, dtypes.int, 144).index(c.bitcast(dtypes.uchar).cast(dtypes.int)).load()])  # 128..143 valid
+      # the bits of a float are any int
+      with self.assertRaises(RuntimeError):
+        to_uops_list([buf.index(r.cast(dtypes.float).bitcast(dtypes.int)).load()])
 
   def test_bool_cast_in_mask(self):
     with Context(CHECK_OOB=1, SPEC=2):
@@ -157,40 +186,20 @@ class TestValidateOOB(unittest.TestCase):
       with self.assertRaises(RuntimeError):
         to_uops_list([buf_int.index(gidx.valid(ld_bool)).load()])  # gidx 0..15, buf_int size 8
 
-  # skipped tests (moved from test_uop_graph.py)
-  @unittest.skip("if not allowed in graph")
-  def test_in_bounds_access_gated_local(self):
-    with Context(CHECK_OOB=1):
-      # Define buffers
+  # local memory
+  def test_gated_local(self):
+    with Context(CHECK_OOB=1, SPEC=2):
       gbuf = UOp.param(0, dtypes.uint, 400)
       sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)
-
-      # Define indices, valids and barrier
       gidx = UOp(Ops.SPECIAL, src=(UOp.const(416),), arg="gidx0")
       lidx = UOp(Ops.SPECIAL, src=(UOp.const(10),), arg="lidx0")
-
-      gate = (gidx<400) & (lidx<8)
-
-      local_store = sbuf.index(lidx.valid(lidx<8)).store(UOp.const(1))
-
-      barrier = UOp(Ops.BARRIER, src=(local_store,))
-      if_barrier = UOp(Ops.IF, src=(gate, barrier))
-
-      # Load from local memory (after the IF/barrier)
-      local_load = UOp(Ops.LOAD, src=(sbuf.index(lidx), if_barrier))
-
-      # Store to global memory
-      global_store = UOp(Ops.STORE, src=(gbuf.index(gidx), local_load))
-      to_uops_list([global_store])
-
-  @unittest.skip("Bool load is not supported yet")
-  def test_load_mask(self):
-    with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, 16)
-      mask = UOp.param(0, dtypes.bool, 16)
-      ridx = UOp.range(20, 0)
-      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
-      to_uops_list([ld0])
+      store = sbuf.index(lidx.valid(lidx < 8)).store(UOp.const(1))
+      load = sbuf.after(store).index(lidx.valid(lidx < 8)).load()
+      to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(load)])  # valid: local store and load gated to 8, global store gated to 400
+      with self.assertRaises(RuntimeError):
+        to_uops_list([gbuf.index(gidx.valid(gidx < 400)).store(sbuf.after(store).index(lidx).load())])  # lidx 0..9 into 8
+      with self.assertRaises(RuntimeError):
+        to_uops_list([gbuf.index(gidx).store(load)])  # gidx 0..415 into 400
 
 if __name__ == "__main__":
   unittest.main()

test/null/test_viz.py

@@ -454,7 +454,7 @@ class TestVizIntegration(unittest.TestCase):
   def test_jit(self):
     with save_viz():
       @TinyJit
-      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(1).contiguous().assign(a.to(c.device)), b.assign(c.to(b.device))
+      def f(a, b, c): return (a+b).contiguous().mul(3), c.add(a.to(c.device)).contiguous(), b.assign(c.to(b.device))
       a, b, c = Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL"), Tensor.empty(16, device="NULL:1")
       for _ in range(3): Tensor.realize(*f(a, b, c))
     out = load_profile(cpu_events)
@@ -1073,10 +1073,11 @@ class TestCLI(unittest.TestCase):
       out = run_cli(*files, "-s", "NULL")
       aggregate = run_cli(*files, "-s", "NULL", "-t")
     self.assertEqual(len(out), 3*2)
-    # flops increases as N gets larger
+    # Operation count increases with N; FLOPS is a rate and also depends on the measured duration.
     gflops = [row["fmt"]["FLOPS"] for row in out]
-    self.assertGreater(gflops[4], gflops[2])
-    self.assertGreater(gflops[5], gflops[3])
+    flops = [rate * row["dur_ms"] * 1e-3 for rate, row in zip(gflops, out)]
+    self.assertGreater(flops[4], flops[2])
+    self.assertGreater(flops[5], flops[3])
     # aggregate flops
     self.assertEqual(len(aggregate), 2)
     agg_gflops = [row["fmt"]["FLOPS"] for row in aggregate]

test/testextra/test_hevc.py

@@ -98,7 +98,7 @@ class TestHevc(unittest.TestCase):
       Variable("pos", 0, max_hist + 1).bind(frame_pos), out_image_size, opaque[1], history)
 
     compiled = compile_linear(decoded.linear_with_vars()[0])
-    self.assertTrue(any(call.src[0].op is Ops.PROGRAM for call in compiled.src))
+    self.assertTrue(any(call.without_after.src[0].op is Ops.PROGRAM for call in compiled.src))
     encdec_calls = [call for call in compiled.src if call.src[0].op is Ops.CUSTOM_FUNCTION and call.src[0].arg == "encdec"]
     self.assertEqual(len(encdec_calls), 1)
 

test/unit/test_after.py

@@ -0,0 +1,91 @@
+import unittest
+from tinygrad import Tensor
+
+class TestAfterCounterexamples(unittest.TestCase):
+  def test_ordered_writes_allowed(self):
+    x = Tensor([0.]).realize().uop
+    a = x.after(x.store(1))
+    b = a.after(a.store(2))
+    self.assertEqual(Tensor(b).tolist(), [2.])
+
+  def test_disjoint_writes_allowed(self):
+    x = Tensor([0., 0.]).realize().uop
+    y = Tensor(x.after(x[:1].store(1), x[1:].store(2)))
+    self.assertEqual(y.tolist(), [1., 2.])
+
+  def test_read_modify_write_chain(self):
+    x = Tensor([2.]).clone()
+    x.assign(x + 1)
+    x.assign(x * 2)
+    self.assertEqual(x.tolist(), [6.])
+
+  def test_overwrite_cuts_gradient(self):
+    x = Tensor([2.])
+    y = x.clone()
+    y.assign(3)  # overwriting with a constant makes y independent of x
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0.])
+
+  def test_shared_state_readers(self):
+    x = Tensor([2.]).clone()
+    x.assign(x + 1)
+    a, b = x + 1, x * 2
+    Tensor.realize(a, b)
+    self.assertEqual(a.tolist(), [4.])
+    self.assertEqual(b.tolist(), [6.])
+
+  @unittest.expectedFailure
+  def test_chained_square_assign_gradient(self):
+    x = Tensor([2.0])
+    y = x.clone()
+    y.assign(y*y)
+    y.assign(y*y)
+    # y = x**4, so dy/dx = 4*x**3. Currently raises "cycle detected while indexing".
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [32.])
+
+  @unittest.expectedFailure
+  def test_partial_store_gradient(self):
+    x = Tensor([2., 3.]).realize()
+    y = Tensor(x.uop.after(x[:1].uop.store(4)))
+    # y = [4, x[1]]. Currently returns [0., 0.].
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
+
+  @unittest.expectedFailure
+  def test_partial_store_source_gradient(self):
+    x = Tensor([4.])
+    y = Tensor([2., 3.]).realize()
+    z = Tensor(y.uop.after(y[:1].uop.store(x.uop)))
+    # x contributes once, not twice. Currently returns [2.].
+    self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
+
+  def test_unrelated_store_gradient(self):
+    x = Tensor([2.]).realize()
+    y = x.clone()
+    z = Tensor(x.uop.after(y.uop.store(0)))
+    # Zeroing y does not change x.
+    self.assertEqual(z.sum().gradient(x)[0].tolist(), [1.])
+
+  @unittest.expectedFailure
+  def test_after_dependency_gradient(self):
+    x = Tensor([2., 3.])
+    y = x.clone()
+    y[:1].assign(0)
+    # View assign creates a nested AFTER; currently raises in backward.
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [0., 1.])
+
+  @unittest.expectedFailure
+  def test_unordered_overlapping_stores_rejected(self):
+    x = Tensor([0.]).realize().uop
+    # No ordering between the writes. Currently succeeds with [2.].
+    with self.assertRaises(RuntimeError):
+      Tensor(x.after(x.store(1), x.store(2))).realize()
+
+  @unittest.expectedFailure
+  def test_gradient_after_callify(self):
+    x = Tensor([2.]).realize()
+    y = x * 2
+    y.callify()
+    # Currently raises: "expected a CALL with unbound BUFFER outputs or a grad_fxn".
+    self.assertEqual(y.sum().gradient(x)[0].tolist(), [2.])
+
+if __name__ == "__main__":
+  unittest.main()

test/unit/test_callify.py

@@ -107,5 +107,20 @@ class TestCallify(unittest.TestCase):
     self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
     self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
 
+  def test_intermediate_clone_persists(self):
+    x = (Tensor([1, 2, 3]).realize() + 1).clone()
+    y = (x * 2).realize()
+    self.assertTrue(x.uop.has_buffer_identity())
+    self.assertEqual(x.tolist(), [2, 3, 4])
+    self.assertEqual(y.tolist(), [4, 6, 8])
+
+  def test_zero_size_cat_with_rng(self):
+    # Empty outputs must not replay a pending RNG counter update.
+    a = Tensor.rand(2, 2)
+    b = Tensor.rand(2, 0)
+    t = a.cat(b, dim=1).realize()
+    self.assertEqual(t.shape, (2, 2))
+    self.assertListEqual(t.tolist(), a.tolist())
+
 if __name__ == "__main__":
   unittest.main()

test/unit/test_dtype_weak.py

@@ -108,6 +108,13 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertIs(stacked.dtype, dtypes.weakfloat)
     self.assertEqual(stacked.tolist(), [2.0, -3.0])
 
+  def test_weakint_cast_truncates_for_every_consumer(self):
+    # a weakint cast of a float is a truncation whether a cast, a compare or an arithmetic op consumes it
+    x = Tensor([2.5, -3.5], dtype=dtypes.float32, device="CPU")
+    self.assertEqual(x.cast(dtypes.weakint).cast(dtypes.float32).tolist(), [2.0, -3.0])
+    self.assertEqual((x.cast(dtypes.weakint) * x).tolist(), [5.0, 10.5])
+    self.assertEqual(Tensor([0.5, -0.5], dtype=dtypes.float32, device="CPU").cast(dtypes.weakint).cast(dtypes.bool).tolist(), [False, False])
+
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),

test/unit/test_gradient.py

@@ -135,6 +135,10 @@ class TestTensorGradient(unittest.TestCase):
     (Tensor.rand(()) + w).backward()
     self.assertIsNone(w.grad)
 
+  def test_max_backward_many_ties(self):
+    t = Tensor.ones(70000, dtype=dtypes.half).contiguous()
+    np.testing.assert_allclose(t.max().gradient(t)[0].sum().numpy(), 1.0, atol=1e-3)
+
 class TestMultiOutputGradient(unittest.TestCase):
   @staticmethod
   def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:

test/unit/test_llm_amd.py

@@ -1,7 +1,8 @@
 import unittest
+from unittest.mock import patch
 import numpy as np
 from tinygrad import Tensor, UOp, dtypes, nn, function
-from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention
+from tinygrad.llm.kernels.amd import Linear, amd_custom_kernels_supported, q8_quantize, flash_attention, gated_delta_prefill
 from tinygrad.llm.gguf import ggml_data_to_tensor
 
 class TestQ8Quantize(unittest.TestCase):
@@ -28,6 +29,12 @@ class TestQ8Quantize(unittest.TestCase):
     # xsum holds the two per-16 sums per 32-wide group
     np.testing.assert_array_equal(gsum.numpy().reshape(2, 2), expected.reshape(2, 2, 16).sum(-1).astype(np.float32))
 
+  def test_quantize_rounding_ties(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    values = np.array([-127,127]+[i+0.5 for i in range(-15,15)],dtype=np.float32)
+    quant,_,_ = q8_quantize(Tensor(values),1,32)
+    np.testing.assert_array_equal(quant.bitcast(dtypes.int8).reshape(32).numpy(),np.rint(values).astype(np.int8))
+
   def test_q6_linear_compiles_in_function(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
@@ -44,22 +51,125 @@ class TestQ8Quantize(unittest.TestCase):
     self.assertEqual(linear.weight.uop.buf_uop.buffer.nbytes, 53*4)
     self.assertEqual(linear.weight.dtype, dtypes.uint32)
 
-  def test_q4_k_linear(self):
+  def test_q4_k_linear(self): self._test_quant_linear(12, 144)
+  def test_iq4_linear(self): self._test_quant_linear(23, 136)
+  def test_q5_linear(self): self._test_quant_linear(13, 176)
+
+  def test_quant_linear_partial_output_tile(self):
+    # Cover a sub-tile output, a trailing tile, and IQ4's larger-output tile selection.
+    for typ, size, outputs, tokens in ((12, 144, 16, 16), (12, 144, 48, 32), (13, 176, 48, 16), (23, 136, 4112, 32)):
+      with self.subTest(ggml_type=typ, out_features=outputs):
+        self._test_quant_linear(typ, size, in_features=256, out_features=outputs, token_counts=(tokens,))
+
+  def test_quant_linear_preserves_rope_permutation(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     rng = np.random.default_rng(42)
-    in_features, blocks = 2048, 16*2048//256
-    packed = rng.integers(0, 256, blocks*144, dtype=np.uint8)
-    for i in range(blocks): packed[i*144:i*144+4] = np.array([0.01, 0.002], dtype=np.float16).view(np.uint8)
-    raw = Tensor(np.pad(packed, (4, 0))).contiguous().realize()[4:]
-    decoded = ggml_data_to_tensor(raw, 16*in_features, 12).reshape(16, in_features)
+    for typ, size in ((12, 144), (13, 176), (14, 210), (23, 136)):
+      with self.subTest(ggml_type=typ):
+        packed = rng.integers(0, 256, (16, size), dtype=np.uint8)
+        packed[:, -2:] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        if typ != 14: packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        if typ in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
+        raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
+        decoded = ggml_data_to_tensor(raw, 16*256, typ).reshape(16, 256).half()
+        original = decoded.numpy()
+        x = rng.normal(size=(3, 256)).astype(np.float16)
+        for prefix in (None, 0, 4):
+          with self.subTest(prefix=prefix):
+            w = decoded.reshape(2, 8, 256)
+            if prefix is None:
+              weight = w.rearrange("n (h two) d -> n (two h) d", two=2)
+            else:
+              weight = w[:, :prefix].cat(w[:, prefix:].rearrange("n (h two) d -> n (two h) d", two=2), dim=1)
+            start = prefix or 0
+            rows = np.arange(16).reshape(2, 8)
+            order = np.concatenate((rows[:, :start], rows[:, start:].reshape(2, -1, 2).transpose(0, 2, 1).reshape(2, -1)), axis=1)
+            linear = Linear(256, 16, bias=False)
+            linear.weight = weight.reshape(16, 256)
+            np.testing.assert_allclose(linear(Tensor(x)).numpy(), x.astype(np.float32) @ original[order.flatten()].astype(np.float32).T,
+                                       rtol=3e-3, atol=2e-2)
+            self.assertIsNone(linear.ggml_type)
+
+  def test_quant_linear_rejects_unaligned_rows_and_integer_casts(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for width in (128, 256):
+      with self.subTest(width=width):
+        packed = np.zeros((2*width//256, 136), dtype=np.uint8)
+        packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+        packed[:, 8:] = np.arange(128, dtype=np.uint8)
+        raw = Tensor(np.pad(packed.flatten(), (4, 0))).realize()[4:]
+        weight = ggml_data_to_tensor(raw, 2*width, 23).reshape(2, width)
+        if width == 256: weight = weight.int().float()
+        expected = weight.numpy().sum(-1)[None]
+        linear = Linear(width, 2, bias=False)
+        linear.weight = weight
+        np.testing.assert_allclose(linear(Tensor.ones(1, width)).numpy(), expected, rtol=1e-3, atol=1e-3)
+        self.assertIsNone(linear.ggml_type)
+
+  def test_dense_gemv_preserves_integer_casts(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    linear = Linear(128, 1)
+    linear.weight = Tensor.full((1, 128), 0.75).contiguous().realize().int().float()
+    linear.bias = Tensor.full((1,), 0.75).contiguous().realize().int().float()
+    np.testing.assert_array_equal(linear(Tensor.ones(1, 128)).numpy(), 0)
+
+  def test_dense_gemv_float32_range(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    linear = Linear(128, 1, bias=False)
+    linear.weight = Tensor.full((1, 128), 1/128, dtype=dtypes.float32).realize()
+    np.testing.assert_array_equal(linear(Tensor.full((1, 128), 65536, dtype=dtypes.float32)).numpy(), 65536)
+
+  def test_gated_delta_state_and_precision(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for case in ("view", "reset", "half"):
+      with self.subTest(case=case):
+        q = Tensor.full((1, 1, 1, 32), 256 if case == "half" else 1, dtype=dtypes.half if case == "half" else dtypes.float32)
+        state = Tensor.full((1, 1, 32, 4), int(case == "reset"), dtype=dtypes.float32).contiguous().realize().transpose(-1, -2)
+        if case != "view": state = state.contiguous().realize()
+        start = Tensor(UOp.variable("start_pos", 0, 10).bind(0)) if case == "reset" else None
+        beta = Tensor.full((1, 1, 1), 1/2097152 if case == "half" else 1, dtype=dtypes.float32)
+        if case != "reset":
+          message = "recurrent state must be contiguous" if case == "view" else "recurrent Q/K must be float32"
+          with self.assertRaisesRegex(AssertionError, message):
+            gated_delta_prefill(q, q, Tensor.ones(1, 1, 1, 4), beta, Tensor.ones(1, 1, 1), state, start)
+          continue
+        out = gated_delta_prefill(q, q, Tensor.ones(1, 1, 1, 4), beta, Tensor.ones(1, 1, 1), state, start)
+        np.testing.assert_array_equal(out.numpy(), 32)
+        np.testing.assert_array_equal(state.numpy(), 1)
+
+  def test_dense_gemv_bias(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    w, bias = rng.normal(size=(32, 128)).astype(np.float16), rng.normal(size=32).astype(np.float16)
+    linear = Linear(128, 32)
+    linear.weight, linear.bias = Tensor(w), Tensor(bias)
+    for tokens in (1, 3):
+      with self.subTest(tokens=tokens):
+        x = rng.normal(size=(tokens, 128)).astype(np.float16)
+        np.testing.assert_allclose(linear(Tensor(x)).numpy(), x.astype(np.float32) @ w.astype(np.float32).T + bias, rtol=2e-3, atol=2e-3)
+
+  def _test_quant_linear(self, ggml_type, block_bytes, in_features=2048, out_features=64, token_counts=(1, 3, 32, 64, 128)):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    packed = rng.integers(0, 256, (out_features*in_features//256, block_bytes), dtype=np.uint8)
+    packed[:, :2] = np.array([0.001], dtype=np.float16).view(np.uint8)
+    if ggml_type in (12, 13): packed[:, 2:4] = np.array([0.0002], dtype=np.float16).view(np.uint8)
+    raw = Tensor(np.pad(packed.flatten(), (4, 0))).contiguous().realize()[4:]
+    decoded = ggml_data_to_tensor(raw, out_features*in_features, ggml_type).reshape(out_features, in_features)
     weight = decoded.numpy()
-    linear = Linear(in_features, 16, bias=False)
-    nn.state.load_state_dict(linear, {"weight":decoded}, verbose=False, realize=False)
-    x = rng.normal(size=(3, in_features)).astype(np.float32)
-    scale = np.maximum(np.abs(x).reshape(3, in_features//32, 32).max(-1, keepdims=True) / 127, 1e-8)
-    xq = np.clip(np.rint(x.reshape(3, in_features//32, 32) / scale), -127, 127) * scale
-    np.testing.assert_allclose(linear(Tensor(x)).numpy(), xq.reshape(3, in_features) @ weight.T, rtol=2e-3, atol=2e-2)
-    self.assertEqual(linear.ggml_type, 12)
+    linear = Linear(in_features, out_features, bias=False)
+    linear.weight = decoded
+    for tokens in token_counts:
+      with self.subTest(tokens=tokens):
+        x = rng.normal(size=(tokens, in_features)).astype(np.float32 if tokens == 3 else np.float16)
+        reference_x = x.astype(np.float32)
+        if tokens < 16:
+          grouped = reference_x.reshape(tokens, -1, 32)
+          scale = np.maximum(np.abs(grouped).max(-1, keepdims=True) / 127, 1e-8)
+          reference_x = (np.clip(np.rint(grouped/scale), -127, 127)*scale).reshape(tokens, in_features)
+        reference_w = weight if tokens < 16 else weight.astype(np.float16).astype(np.float32)
+        np.testing.assert_allclose(linear(Tensor(x)).numpy(), reference_x @ reference_w.T, rtol=3e-3, atol=2e-2)
+    self.assertEqual(linear.ggml_type, ggml_type)
 
   def test_q6_linear_multiple_tokens(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -86,6 +196,18 @@ class TestQ8Quantize(unittest.TestCase):
     self.assertTrue(generic.use_custom_quant)
     self.assertEqual(generic.ggml_type, 14)
 
+  def test_attention_fallback_shapes(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    for tokens, capacity, dim in ((1, 65, 64), (32, 64, 32), (32, 64, 384), (32, 64, 512)):
+      with self.subTest(tokens=tokens, capacity=capacity, dim=dim):
+        valid = 33
+        cache = np.full((2, 1, 1, capacity, dim), np.nan, dtype=np.float16)
+        cache[0, :, :, :valid] = 0
+        cache[1, :, :, :valid] = np.arange(valid)[:, None]
+        q = Tensor.zeros(1, 2, tokens, dim, dtype=dtypes.half)
+        expected = np.broadcast_to(np.arange(valid-tokens, valid)[None, None, :, None]/2, q.shape)
+        np.testing.assert_allclose(flash_attention(q, Tensor(cache), valid).numpy(), expected, rtol=1e-3, atol=1e-3)
+
   def test_attention_uses_physical_cache_length(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
     q, k, v = Tensor.zeros(1, 2, 1, 32), Tensor.randn(1, 1, 1, 32), Tensor.randn(1, 1, 1, 32)
@@ -94,14 +216,71 @@ class TestQ8Quantize(unittest.TestCase):
     out = flash_attention(q, assigned, 1).realize()
     np.testing.assert_allclose(out.numpy(), v.expand(1, 2, 1, 32).numpy(), rtol=2e-2, atol=2e-2)
 
-  def test_flash_attention_decode_gqa_output_layout(self):
+  def test_flash_attention_decode_symbolic_gqa(self):
+    with patch.object(Tensor, "scaled_dot_product_attention", side_effect=AssertionError("expected custom decode")):
+      self._test_flash_decode(8, 2, 256, 128, 37, symbolic=True)
+
+  def test_flash_attention_decode_gqa_tail(self): self._test_flash_decode(3, 1, 192, 64, 37)
+
+  def test_flash_attention_decode_gqa_output_layout(self): self._test_flash_decode(4, 1, 128, 256, 3)
+  def test_flash_attention_decode_large_gqa_group(self): self._test_flash_decode(8, 1, 256, 256, 73)
+
+  def _test_flash_decode(self, heads, kv_heads, dim, n, valid, symbolic=False):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
-    Tensor.manual_seed(42)
-    q = Tensor.randn(1, 4, 1, 128, dtype=dtypes.half).realize()
-    cache = Tensor.randn(2, 1, 1, 256, 128, dtype=dtypes.half).realize()
-    out = flash_attention(q, cache, 3).realize()
-    expected = q.scaled_dot_product_attention(cache[0, :, :, :3], cache[1, :, :, :3], enable_gqa=True)
-    np.testing.assert_allclose(out.numpy(), expected.numpy(), rtol=2e-3, atol=2e-3)
+    rng = np.random.default_rng(42)
+    q = rng.normal(size=(1, heads, 1, dim)).astype(np.float16)
+    cache = rng.normal(size=(2, 1, kv_heads, n, dim)).astype(np.float16)
+    k, v = (np.repeat(c[0, :, :valid].astype(np.float32), heads//kv_heads, axis=0) for c in cache)
+    scores = q[0].astype(np.float32) @ k.transpose(0, 2, 1) / np.sqrt(dim)
+    probs = np.exp(scores - scores.max(-1, keepdims=True))
+    expected = (probs / probs.sum(-1, keepdims=True)) @ v
+    cache_tensor = Tensor(cache)
+    if symbolic:
+      start_pos = UOp.variable("start_pos", 0, n-1).bind(valid-1)
+      valid = start_pos + 1
+      cache_tensor = Tensor(cache_tensor.realize().uop.after(Tensor(start_pos).uop))
+    np.testing.assert_allclose(flash_attention(Tensor(q), cache_tensor, valid).numpy(), expected[None], rtol=2e-3, atol=2e-3)
+
+  def test_prefill_attention_nonfinite_cache_tail(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    rng = np.random.default_rng(42)
+    q = Tensor.zeros(1, 2, 32, 128, dtype=dtypes.half)
+    values = rng.normal(size=(33, 128)).astype(np.float16)
+    expected = np.stack([values[:i+2].astype(np.float32).mean(0) for i in range(32)])[None, None].repeat(2, axis=1)
+    for tail in (np.nan, np.inf, -np.inf):
+      with self.subTest(tail=tail):
+        cache = np.full((2, 1, 1, 64, 128), tail, dtype=np.float16)
+        cache[0, :, :, :33] = 0
+        cache[1, :, :, :33] = values
+        valid = UOp.variable("valid_end", 32, 64).bind(33)
+        cache_tensor = Tensor(cache).realize()
+        assigned = Tensor(cache_tensor.uop.after(Tensor(valid).uop))
+        out = flash_attention(q, assigned, valid)
+        np.testing.assert_allclose(out.numpy(), expected, rtol=2e-3, atol=2e-3)
+
+  def test_flash_attention_decode_beyond_256_chunks(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    n = 257 * 64
+    q = Tensor.zeros(1, 1, 1, 32, dtype=dtypes.half).realize()
+    k = Tensor.zeros(1, 1, n, 32, dtype=dtypes.half)
+    v = Tensor.zeros(1, 1, n-64, 32, dtype=dtypes.half).cat(Tensor.ones(1, 1, 64, 32, dtype=dtypes.half), dim=2)
+    cache = Tensor.stack(k, v).contiguous().realize()
+    for valid, expected in ((1, 0), (n, 1/257)):
+      with self.subTest(valid=valid):
+        valid_kv_len = UOp.variable("valid_kv_len", 1, n).bind(valid)
+        assigned = Tensor(cache.uop.after(Tensor(valid_kv_len).uop))
+        np.testing.assert_allclose(flash_attention(q, assigned, valid_kv_len).numpy(), expected, rtol=2e-3, atol=2e-4)
+
+  def test_flash_attention_decode_long_context_random(self):
+    self._test_flash_decode(8, 2, 128, 257*64, 257*64-13)  # past 256 chunks, with a ragged tail
+
+  def test_flash_attention_decode_chunk_round_accumulator_range(self):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
+    valid_kv_len, max_kv_len = 6749, 6784  # three chunk rounds, with a ragged tail
+    q = Tensor.zeros(1, 8, 1, 32, dtype=dtypes.half).realize()
+    cache = Tensor.stack(Tensor.zeros(1, 1, max_kv_len, 32, dtype=dtypes.half),
+                         Tensor.full((1, 1, max_kv_len, 32), 5500, dtype=dtypes.half)).contiguous().realize()
+    np.testing.assert_allclose(flash_attention(q, cache, valid_kv_len).numpy(), 5500, rtol=2e-3, atol=2e-3)
 
   def test_prefill_attention_unaligned_start(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")

tinygrad/codegen/__init__.py

@@ -7,7 +7,7 @@ from tinygrad.uop.weak import pm_lower_weak, pm_commit_weak, pm_cast_const
 from tinygrad.uop.render import pyrender
 from tinygrad.uop.spec import type_verify, spec_tensor, spec_program
 from tinygrad.renderer import Renderer, Estimates
-from tinygrad.renderer.isa import ISARenderer, IselContext, PreRegAllocContext
+from tinygrad.renderer.isa import ISARenderer, IselContext
 from tinygrad.dtype import dtypes, AddrSpace
 
 # import all pattern matchers here
@@ -439,12 +439,13 @@ def do_linearize(ctx:Renderer, prg:UOp, sink:UOp) -> UOp:
   lst = line_rewrite(linearize(sink), pm_linearize_cleanups)
   # isa renderers need to allocate registers
   if isinstance(ctx, ISARenderer):
-    if ctx.pre_regalloc_matcher is not None: lst = line_rewrite(lst, ctx.pre_regalloc_matcher, PreRegAllocContext())
+    lin_ctx = ctx.linear_ctx_type(ctx)
+    lst = line_rewrite(lst, ctx.pre_regalloc_matcher, lin_ctx)
     # register definitions (INS without srcs) move to the top so regalloc sees their live ranges span the whole program (callee saved regs)
     lst = sorted(lst, key=lambda u: u.op is not Ops.INS or bool(u.src))
-    regalloc_ctx = LinearScanRegallocContext(lst, ctx)
+    regalloc_ctx = LinearScanRegallocContext(lin_ctx, lst, ctx)
     lst = line_rewrite(lst, pm_regalloc_rewrite, regalloc_ctx)
-    lst = line_rewrite(lst, ctx.post_regalloc_matcher, regalloc_ctx)
+    lst = line_rewrite(lst, ctx.post_regalloc_matcher, lin_ctx)
     if DEBUG >= 4: print(ctx.asm_str(lst, sink.arg.function_name))
   return prg.replace(src=prg.src + (UOp(Ops.LINEAR, src=tuple(lst)),))
 

tinygrad/codegen/decomp/op.py

@@ -18,29 +18,19 @@ def magicgu(vmax:int, d:int) -> tuple[int,int]:
   assert False
 
 def fast_idiv(ren: Renderer, x: UOp, d: int, dont_cast=False) -> UOp|None:
-  from tinygrad.renderer.cstyle import MetalRenderer
-  # NOTE: disable for METAL due to compiler bug. keccak with -O0 works but not with optimization
-  if isinstance(ren, MetalRenderer): return None
-  # If d is a power of two this is not valid for signed ints!
-  is_unsigned = x.vmin>=0 or x.dtype in dtypes.uints
-  assert d>0, "Sign should have been taken out of divisor"
-  vmin,vmax = max(x.vmin, x.dtype.min), min(x.vmax, x.dtype.max)
-  if vmin > -d and vmax < d: return x.const_like(0)
-  m,s = magicgu(max(vmax, abs(vmin)), d)
-  if m*vmin >= x.dtype.min and m*vmax <= x.dtype.max:
-    return ((x*m) >> s) if is_unsigned else ((x*m) >> s) + (x<0).where(x.ufix(1), 0)
+  if d <= 0 or x.vmin < 0: return None
+  if (vmax:=min(x.vmax, x.dtype.max)) < d: return x.const_like(0)
+  m,s = magicgu(vmax, d)
+  if m*vmax <= x.dtype.max: return (x*m) >> s
   # before we try casting to a larger dtype (slow), we see if there are powers of two in d we can shift to make x smaller
-  # use explicit Ops.CDIV (trunc) since the recursion assumes trunc semantics throughout
-  if (largest_factor_of_two_in_d := (d & -d)) > 1:
-    if (ret:=fast_idiv(ren, x.alu(Ops.CDIV, x.const_like(largest_factor_of_two_in_d)),
-                       d//largest_factor_of_two_in_d, dont_cast=True)) is not None: return ret
+  if (k := (d & -d).bit_length()-1) > 0:
+    if (ret:=fast_idiv(ren, x >> k, d >> k, dont_cast=True)) is not None: return ret
   if dont_cast: return None
   # the next integer width that holds x*m
   widen = {dtypes.int8:dtypes.int16, dtypes.int16:dtypes.int32, dtypes.int32:dtypes.int64, dtypes.int64:dtypes.uint64,
            dtypes.uint8:dtypes.uint16, dtypes.uint16:dtypes.uint32, dtypes.uint32:dtypes.uint64}
   if (next_dtype := widen.get(x.dtype)) is not None and next_dtype in ren.supported_dtypes():
-    if m*vmin >= next_dtype.min and m*vmax <= next_dtype.max:
-      return ((x.cast(next_dtype)*m) >> s).cast(x.dtype) if is_unsigned else ((x.cast(next_dtype)*m) >> s).cast(x.dtype) + (x<0).where(x.ufix(1), 0)
+    if m*vmax <= next_dtype.max: return ((x.cast(next_dtype)*m) >> s).cast(x.dtype)
   return None
 
 # ***** threefry *****
@@ -105,13 +95,12 @@ def get_late_rewrite_patterns(ops:tuple[Ops, ...], disable_fast_idiv:bool) -> Pa
       lambda x,c: (x+(l.const_like(l.vmin) if (l:=(x<0)).vmin==l.vmax else l).where(c-1, 0)) >> v
         if (v:=powers_of_two.get(c.val, 0)) else None)]
     if not disable_fast_idiv:
-      # fast_idiv handles non-pow2: only fire on non-negative inputs (signed magic-mul is unreliable for x<0)
-      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
-        lambda ctx, x, d: fast_idiv(ctx, x, d.val) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
-      # rewrite raw CMOD -> x - d*CDIV(x,d) so fast_idiv can pick up the CDIV. only on non-negative inputs;
+      # fast_idiv handles non-pow2 divisors on non-negative inputs
+      pat += [(UPat(Ops.CDIV, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))), lambda ctx, x, d: fast_idiv(ctx, x, d.val))]
+      # rewrite raw CMOD -> x - d*fast_idiv(x,d), only when fast_idiv can actually divide;
       # avoids disturbing floormod_to_mod's general-path output (which uses a trunc Ops.CMOD as an implementation detail)
-      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.var("d"))),
-        lambda x, d: x - d * x.alu(Ops.CDIV, d) if x.vmin >= 0 or x.dtype in dtypes.uints else None)]
+      pat += [(UPat(Ops.CMOD, src=(UPat.var("x", dtypes.ints), UPat.cvar("d"))),
+        lambda ctx, x, d: x - d * q if (q:=fast_idiv(ctx, x, d.val)) is not None else None)]
   if Ops.NEG in ops:
     pat += [(UPat.var('x')*-1, lambda ctx,x: x.alu(Ops.NEG))]
     if Ops.SUB in ops: pat += [(UPat.var('x')+UPat.var('y').alu(Ops.NEG), lambda ctx,x,y: x.alu(Ops.SUB, y))]

tinygrad/codegen/late/regalloc.py

@@ -1,38 +1,35 @@
 import itertools
 from tinygrad.helpers import dedup
 from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat
-from tinygrad.renderer.isa import ISARenderer, Register, greg
-from tinygrad.dtype import dtypes
+from tinygrad.renderer.isa import ISARenderer, Register, rdef, LinearContext
+from typing import Any
 
 PSEUDO_OPS = {Ops.CONST, Ops.CAST, Ops.BITCAST, Ops.NOOP, Ops.AFTER, Ops.BARRIER, Ops.GROUP, Ops.STACK}
 
 class LinearScanRegallocContext:
   # returns the uop that defines the virtual register
   def vdef(self, v:Register) -> UOp: return self.uops[self.live_range[v][0]]
-  def __init__(self, uops:list[UOp], ren:ISARenderer):
+  def __init__(self, ctx:LinearContext, uops:list[UOp], ren:ISARenderer):
     self.uops = uops
     self.ren = ren
     self.idx = itertools.count()
-    # the label associated with each loop NOTE: this is only used post regalloc and should be removed
-    self.loop_label: dict[UOp, str] = {}
 
     # compute live ranges
     self.live_range: dict[Register, list[int]] = {}
     lr = self.live_range
-    ranges: list[Register] = []
-    for i,u in enumerate(reversed(uops)):
+    loops: dict[int, int] = {} # the interval of each loop, from its RANGE to the last uop that reads that RANGE
+    for idx,u in reversed(list(enumerate(uops))):
       if u.op in PSEUDO_OPS: continue
       defs = u.tag if isinstance(u.tag, tuple) else ()
-      for v in defs + tuple(greg(s) for s in dedup(u.src)):
-        if isinstance(v, Register): lr.setdefault(v, []).insert(0, len(uops) - 1 - i)
+      for v in defs + tuple(rdef(s) for s in dedup(u.src)):
+        if isinstance(v, Register): lr.setdefault(v, []).insert(0, idx)
       for v in defs:
-        if v in lr and (n:=max((lr[rng][-1] for rng in ranges if lr[rng][0] <= lr[v][-1] < lr[rng][-1]), default=None)): lr[v].append(n)
-      if u.op is Ops.RANGE: ranges.append(greg(u))
+        if v in lr and (n:=max((e for s,e in loops.items() if s <= lr[v][-1] < e), default=None)): lr[v].append(n)
+      if u.op is Ops.RANGE: loops[idx] = max(j for j,x in enumerate(uops) if u in x.src)
 
     # allocate registers
-    self.stack_size: int = 0
     self.locals: dict[UOp, UOp] = {}
-    self.spills: dict[Register, UOp] = {} # mapping from virtual to stack slot
+    self.spills: dict[Register, Any] = {} # mapping from virtual to arbitrary spill slot
     self.reals: dict[int, dict[Register, Register]] = {} # mapping from virtual to real at each program point
     self.insert_before: dict[int, list[tuple[Register, Register]]] = {} # fills to be inserted at each program point
     live: dict[Register, Register] = {} # mapping from virtual to real that's currently assigned to it
@@ -49,11 +46,7 @@ class LinearScanRegallocContext:
     # assign register to spilled virtual and record load to be emitted before current uop, also assign it a stack slot
     def fill(v:Register, i:int, cons:tuple[Register, ...]|None=None) -> Register:
       if v not in self.spills:
-        # the value of a BUFFER is its 64bit address, XMM registers need 16 bytes
-        sz = 16 if v.cons[0].size == 16 else (8 if self.vdef(v).op is Ops.BUFFER else self.vdef(v).dtype.itemsize)
-        offset = self.stack_size + (sz - self.stack_size % sz) % sz
-        self.spills[v] = UOp.cconst(offset, dtypes.int32)
-        self.stack_size = offset + sz
+        self.spills[v] = ctx.assign_spill_slot(v, self.vdef(v))
       r = alloc(cons if cons is not None else v.cons, i)
       self.insert_before.setdefault(i, []).append((v, r))
       return r
@@ -64,7 +57,7 @@ class LinearScanRegallocContext:
       for s in u.src:
         # HACK: cause of later hacks to lower range
         if u.op is Ops.END: continue
-        if not isinstance(v:=greg(s), Register): continue
+        if not isinstance(v:=rdef(s), Register): continue
         if v not in live: live[v] = fill(v, i)
         self.reals.setdefault(i, {})[v] = live[v]
 
@@ -76,21 +69,16 @@ class LinearScanRegallocContext:
           cons = v.cons
           # two address instructions (src is reused by def) can only coalesce reused src. reused src goes first to get priority in case of a tiebreak
           if ren.is_two_address(u) and j == 0:
-            uses = tuple(live.get(greg(s)) for s in u.src)
+            uses = tuple(live.get(rdef(s)) for s in u.src)
             cons = ((uses[0],) if uses[0] in cons else ()) + tuple(r for r in cons if r not in uses)
           # HACK: cause the range is missing the comparison
           live[v] = alloc(cons, i+1 if u.op is not Ops.RANGE else i)
           self.reals.setdefault(i, {})[v] = live[v]
 
-      # allocate stack array
-      if u.op is Ops.BUFFER:
-        self.locals[u] = UOp.cconst(self.stack_size, dtypes.int32)
-        self.stack_size += u.max_numel() * u.dtype.itemsize
-
       # loop prologue, avoid loading inside the loop
       if u.op is Ops.RANGE:
         # we move to registers vars used in the loop sorted by next use, vars not used in the loop will not be reloaded in the epilogue
-        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < lr[greg(u)][-1] for l in lr[v])]
+        used_in_loop = [v for v in live.keys() | self.spills.keys() if any(i <= l < loops[i] for l in lr[v])]
         sorted_uses = sorted(used_in_loop, key=lambda k: (next(l-i for l in lr[k] if l >= i), lr[k][0], k.name, k.index))
         live_in: dict[Register, Register] = {}
         for v in sorted_uses:
@@ -113,22 +101,14 @@ def regalloc_rewrite(ctx:LinearScanRegallocContext, x:UOp):
   nsrc = []
   for j,s in enumerate(x.src):
     # v here is the virtual defined by the original s as s is the rewritten version
-    if i in ctx.reals and (v:=greg(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
+    if i in ctx.reals and (v:=rdef(ctx.uops[i].src[j])) in ctx.spills: nsrc.append(ctx.ren.fill(ctx.spills[v], ctx.vdef(v), ctx.reals[i][v]))
     else: nsrc.append(s)
   ndefs = tuple(ctx.reals[i][v] for v in x.tag) if isinstance(x.tag, tuple) else x.tag
-  if x.op is Ops.BUFFER: nx = ctx.ren.isel_matcher.rewrite(ctx.ren.stack_pointer().index(ctx.locals[x], tag=ndefs))
-  else: nx = x.replace(src=tuple(nsrc), tag=ndefs)
+  nx = x.replace(src=tuple(nsrc), tag=ndefs)
 
   before = [ctx.ren.fill(ctx.spills[v], ctx.vdef(v), r) for v,r in ctx.insert_before.get(i, [])]
   after = [ctx.ren.spill(ctx.spills[v], nx) for v in x.tag if v in ctx.spills] if isinstance(x.tag, tuple) else []
 
-  # alloc/dealloc stack
-  if ctx.stack_size > 0:
-    sp = ctx.ren.stack_pointer()
-    offset = UOp.cconst(ctx.stack_size, sp.dtype)
-    if i == 0: before = [ctx.ren.isel_matcher.rewrite(UOp(Ops.SUB, src=(sp, offset), tag=sp.tag))] + before
-    elif i == len(ctx.uops) - 2: before += [ctx.ren.isel_matcher.rewrite(UOp(Ops.ADD, src=(sp, offset), tag=sp.tag))]
-
   return nx, before + [nx] + after
 
 pm_regalloc_rewrite = PatternMatcher([

tinygrad/device.py

@@ -187,7 +187,7 @@ class Buffer:
     # zero copy with as_memoryview (disabled by default due to use after free)
     if (force_zero_copy or allow_zero_copy) and hasattr(self.allocator, '_as_buffer'):
       if not no_sync: self.allocator.dev.synchronize()
-      return self.allocator._as_buffer(self._buf)
+      if (mv:=self.allocator._as_buffer(self._buf)) is not None: return mv
     assert not force_zero_copy, "force zero copy was passed, but copy is required"
     Buffer("PYTHON", self.size, self.dtype, opaque=(mv:=memoryview(bytearray(self.nbytes)))).copy_from(self)
     return mv
@@ -280,9 +280,13 @@ class DepsTracker:
       if i in write:
         for dmap in [self.w_dependency_map, self.r_dependency_map]:
           kept = []
-          for st,en,dep in dmap[key]:
-            if st < min(s, en): kept.append((st, min(s, en), dep))
-            if max(e, st) < en: kept.append((max(e, st), en, dep))
+          for entry in dmap[key]:
+            st, en, dep = entry
+            if st == en: continue
+            if en <= s or e <= st: kept.append(entry)
+            else:
+              if st < s: kept.append((st, s, dep))
+              if e < en: kept.append((e, en, dep))
           dmap[key] = kept
         self.w_dependency_map[key].append((s, e, new_dependency))
       else: self.r_dependency_map[key].append((s, e, new_dependency))
@@ -337,8 +341,9 @@ class Compiled:
 
   has_copy_queue:bool = True
 
-  pm_encode:Any = None # per queue kind: queue ops -> flat command words
-  pm_lower:Any = None # per queue kind: custom_function(submit, cmdbuf) -> the queue push
+  pm_batch:Any = None
+  pm_encode:Any = None
+  pm_lower:Any = None
   pm_bufferize:Any = None
 
   def __init__(self, device:str, allocator:Allocator, renderers:list[type[Renderer]], runtime:type[Program[Self]]|None, graph=None, arch=None):

tinygrad/engine/jit.py

@@ -166,7 +166,7 @@ class CapturedJit(Generic[ReturnType]):
   expected_input_info: list[tuple[UOp, tuple[Variable, ...], DType, str]]  # (view, variables, dtype, device) per input
 
   @functools.cached_property
-  def linear(self) -> UOp: return link_linear(self._linear)
+  def linear(self) -> UOp: return link_linear(self._linear, allow_cache=False) # do not cache jit
 
   def __reduce__(self): return self.__class__, (self.ret, self._linear, self.expected_names, self.expected_input_info)
 

tinygrad/engine/realize.py

@@ -52,6 +52,8 @@ def get_call_name(call:UOp, bufs:Sequence[Buffer|UOp], var_vals:dict[str, int]|N
 # **************** Stat ****************
 
 def estimate_uop(call:UOp) -> Estimates:
+  call = call.without_after
+  if isinstance(call.arg.aux, HCQInfo): return call.arg.aux.estimates
   if (ast:=call.src[0]).op is Ops.PROGRAM: return ast.src[0].arg.estimates or Estimates()
   if ast.op is Ops.COPY or (ast.op is Ops.CUSTOM_FUNCTION and ast.arg == "encdec"):
     return Estimates(lds=(nbytes:=prod(call.src[1].shape) * call.src[1].dtype.itemsize), mem=nbytes)
@@ -130,7 +132,7 @@ class ExecContext:
   cache: bool = True
 
 def _resolve(b:UOp, inputs:tuple[UOp, ...]) -> UOp:
-  if b.op in (Ops.MSELECT, Ops.SHRINK): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
+  if b.op in (Ops.MSELECT, Ops.SHRINK, Ops.BITCAST): return b.replace(src=(_resolve(b.src[0], inputs), *b.src[1:]))
   if b.op is Ops.MSTACK: return b.replace(src=tuple(_resolve(x, inputs) for x in b.src))
   return inputs[b.arg.slot] if b.op is Ops.PARAM else b
 def resolve_params(call:UOp, inputs:tuple[UOp, ...]) -> list[UOp]: return [_resolve(b, inputs) for b in get_call_arg_uops(call)]
@@ -192,20 +194,21 @@ def exec_graph(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
 
 def exec_hcq(ctx:ExecContext, call:UOp, ast:UOp) -> list[float|None]:
   if (info:=call.arg.aux).inputs:
-    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr for u, dev in info.inputs]
+    addrs = [cast(Buffer, _resolve(u, ctx.input_uops).buffer).get_buf(dev).va_addr + off for u, dev, off in info.inputs]
     cast(Buffer, call.src[1 + info.table].buffer)._buf.cpu_view().view(fmt='Q')[:] = array.array('Q', addrs)
   ctx = replace(ctx, var_vals={**ctx.var_vals, **{k: v for d in info.device for k, v in cast(Any, Device[d]).var_vals.items()}})
   ets = exec_kernel(ctx, call, ast, devices=(HCQ_RUNTIME_DEV.value,))
   if not (ctx.wait or PROFILE): return ets
 
   slots = {d: cast(Buffer, call.src[1 + i].buffer) for d, i in info.slots}
-  def _prof_tm(device:str, name:str, prof:tuple[int, ...], profile_key:bytes) -> float|None:
-    (d:=cast(Any, Device[device])).prof_ents[(slots[device], prof[0])] = ProfileGraphEntry(device, name, prof[0], prof[1], profile_key)
-    if not ctx.wait: return None
-    d.synchronize(timeout=ctx.timeout)
+  for devs, name, _, prof, pkey in info.kernels:
+    for d in (devs if prof else ()): cast(Any, Device[d]).prof_ents[(slots[d], prof[0])] = ProfileGraphEntry(d, name, prof[0], prof[1], pkey)
+  if ctx.wait:
+    for device in info.device: cast(Any, Device[device]).synchronize(timeout=ctx.timeout)
+  def _prof_tm(device:str, prof:tuple[int, ...]) -> float:
     st, en = (slots[device]._buf.cpu_view().view(fmt='Q')[x] for x in prof)
-    return float(en-st) / d.timestamp_divider / 1e6
-  return ets + [_prof_tm(device, name, prof, profile_key) for devices,name,_,prof,profile_key in info.kernels if prof for device in devices]
+    return float(en-st) / cast(Any, Device[device]).timestamp_divider / 1e6
+  return ets + [_prof_tm(device, prof) if ctx.wait else None for devices, _, _, prof, _ in info.kernels if prof for device in devices]
 
 # flatten LINEAR-in-LINEAR: any nested LINEAR child gets inlined into its parent's src
 pm_flatten_linear = PatternMatcher([
@@ -283,17 +286,18 @@ def compile_linear(linear:UOp, beam:int|None=None, validate=False, input_uops:li
   linear = hcq_compile(linear, input_uops, bool(PROFILE or DEBUG >= 2) if profile is None else profile)
   return linear
 
-def link_linear(linear:UOp, cache=True) -> UOp: return hcq_link(linear, cache=cache)
+def link_linear(linear:UOp, input_uops:list[UOp]|None=None, allow_cache=True) -> UOp:
+  return hcq_link(linear, input_uops=input_uops, allow_cache=allow_cache)
 
 def run_linear(linear:UOp, var_vals:dict[str, int]|None=None, input_uops:Sequence[UOp]=(), update_stats=True, jit=False, wait=False):
   inputs = list(input_uops)
-  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), cache=False) # a one-shot link
+  if not jit: linear = link_linear(compile_linear(linear, validate=VALIDATE_WITH_CPU, input_uops=inputs), input_uops=inputs)
   ctx = ExecContext(var_vals or {}, tuple(inputs), update_stats, jit, wait or DEBUG>=2)
   for call in linear.src: track_stats(ctx, call.without_after, perf_counter_us(), pm_exec.rewrite(call.without_after, ctx))
 
 def time_call(call:UOp, var_vals:dict[str, int]|None=None, timeout:int|None=None, clear_l2:bool=False) -> Iterator[float]:
   ctx = ExecContext(var_vals or {}, update_stats=False, wait=True, timeout=timeout, cache=False)
-  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), cache=ctx.cache)
+  linear = link_linear(compile_linear(UOp(Ops.LINEAR, src=(call,)), beam=0, profile=True), allow_cache=ctx.cache)
   while True:
     if clear_l2:
       if hasattr(dev:=Device[call.src[1].device], 'invalidate_caches'): dev.invalidate_caches()

tinygrad/function.py

@@ -26,8 +26,9 @@ def invalid_outputs(uret:UOp) -> set[UOp]:
           if u.op is Ops.STORE and u.src[1].base.is_invalid and not u.src[0].buf_uop.is_realized}
 
 def renumber_invalid_outputs(uret:UOp) -> UOp:
+  invalid = invalid_outputs(uret)
   return uret.substitute({b:b.replace(arg=replace(b.arg, slot=i))
-                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid_outputs(uret))})
+                          for i,b in enumerate(x for x in uret.toposort(enter_calls=False) if x in invalid)})
 
 ReturnType = TypeVar('ReturnType')
 class _function(Generic[ReturnType]):

tinygrad/helpers.py

@@ -166,6 +166,8 @@ def stderr_log(msg:str): print(msg, end='', file=sys.stderr, flush=True)
 
 class Context(contextlib.ContextDecorator):
   def __init__(self, **kwargs): self.kwargs = kwargs
+  # ContextDecorator otherwise reuses self, so recursive calls overwrite old_context.
+  def _recreate_cm(self): return Context(**self.kwargs)
   def __enter__(self):
     self.old_context:dict[str, Any] = {k: ContextVar._cache[k].value for k in self.kwargs}
     for k,v in self.kwargs.items(): ContextVar._cache[k].value = v
@@ -239,7 +241,7 @@ TRANSCENDENTAL = ContextVar("TRANSCENDENTAL", 1)
 SPLIT_REDUCEOP, NO_MEMORY_PLANNER, LRU = ContextVar("SPLIT_REDUCEOP", 1), ContextVar("NO_MEMORY_PLANNER", 0), ContextVar("LRU", 1)
 RING, ALL2ALL, ALLREDUCE_CAST = ContextVar("RING", 1), ContextVar("ALL2ALL", 0), ContextVar("ALLREDUCE_CAST", 1)
 CACHELEVEL, IGNORE_BEAM_CACHE = ContextVar("CACHELEVEL", 2), ContextVar("IGNORE_BEAM_CACHE", 0)
-VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 0)
+VALIDATE_WITH_CPU, HCQ2 = ContextVar("VALIDATE_WITH_CPU", 0), ContextVar("HCQ2", 1)
 # TODO: this is broken for some indexing
 DISABLE_FAST_IDIV = ContextVar("DISABLE_FAST_IDIV", 1)
 FUSE_OPTIM = ContextVar("FUSE_OPTIM", 0)

tinygrad/llm/chat.html

@@ -1,24 +1,36 @@
-<!DOCTYPE html><html><head><title>tinygrad chat</title><style>
+<!DOCTYPE html><html><head><meta charset="utf-8"><title>tinygrad chat</title><style>
   * { margin: 0 }
   body { background: #212121; color: #e3e3e3; font-family: system-ui;
          height: 100vh; display: flex; flex-direction: column }
   #chat { flex: 1; overflow-y: auto; padding: 20px }
   .msg { padding: 10px 16px; margin: 8px 0; white-space: pre-wrap; border-radius: 18px }
+  table { border-collapse: collapse; table-layout: fixed; width: 100%; overflow-wrap: anywhere }
+  th, td { border: 1px solid #555; padding: 6px 10px; text-align: left }
+  a { color: #8ab4f8 } hr { border: 0; border-top: 1px solid #555 }
+  .answer { white-space: normal; line-height: 1.65 } .answer > * { margin: 12px 0 }
+  pre, blockquote { background: #2f2f2f; padding: 12px 16px; border-radius: 8px } pre { white-space: pre-wrap }
   .user { background: #2f2f2f; margin-left: auto; width: fit-content; max-width: 70% }
   #input { max-width: 768px; width: 100%; margin: 20px auto; padding: 14px 20px;
            background: #2f2f2f; color: inherit; font: inherit;
            border: none; outline: none; resize: none; border-radius: 24px; field-sizing: content }
 </style></head><body><div id="chat"></div>
 <textarea id="input" rows="1" placeholder="Ask anything" autofocus></textarea>
+<script src="/assets/cdn.jsdelivr.net/npm/[email protected]/dist/browser/markdown-it.umd.min.js"></script>
 <script>
-  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) { e.preventDefault(); send() } }
+  let generating = false;
+  input.onkeydown = (e) => { if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
+    e.preventDefault(); if (generating) return;
+    generating = true; send().finally(() => generating = false);
+  } };
   const msgs = [];
+  const md = markdownit();
   async function send() {
     if (!input.value.trim()) return;
     msgs.push({role: 'user', content: input.value.trim()});
     chat.innerHTML += '<div class="msg user">' + input.value.trim().replace(/</g, '&lt;') + '</div>';
     input.value = '';
     const d = document.createElement('div'); d.className = 'msg'; chat.appendChild(d);
+    d.innerHTML = '<span style="color:#888"></span><div class="answer"></div>'; const [thinking, answer] = d.children;
     const r = await fetch('/v1/chat/completions', {method: 'POST', headers: {'Content-Type': 'application/json'},
       body: JSON.stringify({model: 'llama', messages: msgs, stream: true, temperature: 0.7})});
     let buf = '', txt = '', rsn = '';
@@ -29,12 +41,11 @@
       const lines = buf.split('\n');
       buf = lines.pop();
       for (const ln of lines)
-        if (ln.startsWith('data: ') && !ln.includes('[DONE]'))
-          try { const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
-            if (dl?.reasoning_content) { const s = document.createElement('span'); s.style.color = '#888';
-              s.textContent = dl.reasoning_content; rsn += dl.reasoning_content; d.appendChild(s) }
-            if (dl?.content) { const s = document.createElement('span');
-              s.textContent = dl.content; txt += dl.content; d.appendChild(s) } } catch {}
+        if (ln.startsWith('data: ') && !ln.includes('[DONE]')) {
+          const dl = JSON.parse(ln.slice(6)).choices[0]?.delta;
+          if (dl?.reasoning_content) { rsn += dl.reasoning_content; thinking.textContent = rsn }
+          if (dl?.content) { txt += dl.content; answer.innerHTML = md.render(txt) }
+        }
       chat.scrollTop = chat.scrollHeight;
     }
     const m = {role:'assistant', content:txt}; if (rsn) m.reasoning_content = rsn; msgs.push(m);

tinygrad/llm/gguf.py

@@ -129,7 +129,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
       return (dl * (grid + delta)).flatten(-3)
     if ggml_type == 20:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32)
-      return d * Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)[q_to_uint8(blocks[:, 2:], 4)]
+      return d * Tensor.const(tuple(_ggml.kvalues_iq4nl), dtypes.float32)[q_to_uint8(blocks[:, 2:], 4)]
     if ggml_type == 21:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1))
       scales = (1 + 2 * q_to_uint8(blocks[:, 106:110].reshape((-1, 4, 1)), 4).reshape((-1, 8))).cast(dtypes.float32).reshape((-1, 8, 1, 1))
@@ -147,7 +147,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor:
     if ggml_type == 23:
       d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1))
       scale_shifts = Tensor.const((0, 2, 4, 6, 8, 10, 12, 14), dtypes.uint16)
-      iq4_xs_lut = Tensor(list(_ggml.kvalues_iq4nl), dtype=dtypes.float32, device=t.device)
+      iq4_xs_lut = Tensor.const(tuple(_ggml.kvalues_iq4nl), dtypes.float32)
       scales_l = Tensor.stack((sl:=blocks[:, 4:8]).bitwise_and(0xF), sl.rshift(4), dim=2).reshape((-1, 8))
       scales_h = blocks[:, 2:4].bitcast(dtypes.uint16).unsqueeze(-1).rshift(scale_shifts).bitwise_and(0x03).reshape((-1, 8)).cast(dtypes.uint8)
       scales = (scales_l.bitwise_or(scales_h.lshift(4)).bitcast(dtypes.int8) - 32).cast(dtypes.float32).reshape((-1, 8, 1))

tinygrad/llm/kernels/amd.py

@@ -3,9 +3,11 @@ import functools, math
 from typing import Callable, cast
 from tinygrad import Tensor, UOp, nn, Device, Context
 from tinygrad.device import Buffer
+from tinygrad.llm.gguf import ggml_data_to_tensor
 from tinygrad.dtype import AddrSpace, dtypes
 from tinygrad.helpers import prod
 from tinygrad.uop.ops import AxisType, KernelInfo, Ops, resolve
+from tinygrad.renderer.cstyle import HIPRenderer
 
 BLOCK_M, BLOCK_N, WARP_SIZE = 32, 32, 32
 WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
@@ -31,7 +33,7 @@ def amd_custom_kernels_supported(device:str|tuple[str, ...]|None) -> bool:
   if device is None or device.split(":")[0] != "AMD": return False
   # @function contexts set ALLOW_DEVICE_USAGE=0 (scheduling must not open devices); the device is always open here
   with Context(ALLOW_DEVICE_USAGE=1):
-    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11
+    return (t:=getattr(Device[device], "target", None)) is not None and t[0] == 11 and isinstance(Device[device].renderer, HIPRenderer)
 
 def warp_reduce(val:UOp, maximum:bool=False, full_wave:bool=False) -> UOp:
   for offset in ((16, 8, 4, 2, 1) if full_wave else (8, 4, 2, 1)):
@@ -54,15 +56,20 @@ class Linear(nn.Linear):
     super().__init__(in_features, out_features, bias)
     self.in_features, self.out_features = in_features, out_features
   def set_quantized(self, decoded:Tensor):
+    if self.in_features % GGML_BLOCK_SIZE: return
     packed_sizes = {decoded.numel() // 256 * type_size:typ for typ,type_size in QUANT_SIZES.items()}
     graph = decoded.uop.toposort()
     raw = next((u for u in graph if u.op is Ops.SHRINK and u.dtype == dtypes.uint8 and prod(u.shape) in packed_sizes), None)
     if raw is None: return
     ggml_type = packed_sizes[prod(raw.shape)]
-    # the packed byte rate alone can't distinguish same-rate formats (Q4_0 vs Q4_K, Q5_0 vs Q5_K, MXFP4 vs IQ4_XS).
-    # the supported formats are 256-wide superblocks: their decode views the packed bytes at the superblock width
-    # (ggml_data_to_tensor reshapes to (-1, QUANT_SIZES[type])), while same-rate 32-wide formats reshape to 17-22
-    if not any(u.op is Ops.RESHAPE and u.shape[-1:] == (QUANT_SIZES[ggml_type],) for u in graph): return
+    # Only unwrap storage/order-preserving views, then require the exact dequantization expression.
+    # This rejects subsequent arithmetic and permutations, including RoPE's concatenated query weights.
+    def unwrapped(u:UOp) -> UOp:
+      while u.op in (Ops.RESHAPE, Ops.CONTIGUOUS) or (u.op is Ops.CAST and dtypes.is_float(u.dtype) and dtypes.is_float(u.src[0].dtype)):
+        u = u.src[0]
+      return u
+    expected = ggml_data_to_tensor(Tensor(raw), self.in_features * self.out_features, ggml_type)
+    if unwrapped(decoded.uop).key != unwrapped(expected.uop).key: return
     raw_offset = raw.contiguous_view_offset()
     assert raw_offset is not None and raw_offset % 4 == 0 and raw.buf_uop.dtype == dtypes.uint8
     self.ggml_type = ggml_type
@@ -100,23 +107,18 @@ class Linear(nn.Linear):
     return super().__call__(x)
 
 def _amd_dp4a(a:UOp, b:UOp, c:UOp) -> UOp:
-  # int8 4-wide dot, widened to scalar multiply-adds (2% decode slower than the sudot4 builtin, but portable)
-  for i in range(4):
-    av = ((a >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    bv = ((b >> (8*i)) & 255).cast(dtypes.uint8).bitcast(dtypes.int8).int()
-    c = c + av*bv
-  return c
+  return UOp(Ops.CUSTOMI, src=(a, b, c), arg=("__builtin_amdgcn_sudot4(true, {}, true, {}, {}, false)", dtypes.int32))
 
 def _amd_byte_perm(a:UOp, b:UOp, selectors:UOp) -> UOp:
   return UOp(Ops.CUSTOMI, src=tuple(x.cast(dtypes.uint32) for x in (a, b, selectors)), arg=("__builtin_amdgcn_perm({}, {}, {})", dtypes.uint32))
 
-def _amd_load(ptr:UOp, lanes:int|None=None) -> UOp:
+def _amd_load(ptr:UOp, lanes:int|None=None, stream:bool=False) -> UOp:
   assert ptr.op is Ops.INDEX
   # nontemporal scalar load: streamed weights must not evict the activations/KV cache from L2
   if lanes is None: return ptr.load(arg="nontemporal")
   buf, coords = ptr.src[0], ptr.src[1:]
   idx = sum((coord*math.prod(buf.shape[i+1:]) for i,coord in enumerate(coords)), UOp.const(0))
-  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load()
+  return UOp(Ops.SHRINK, src=(buf.flatten(), idx, UOp.const(lanes))).load(arg="nontemporal" if stream else None)
 
 def _load_byte(raw:UOp, base:UOp, offset:UOp) -> UOp: return (raw[base + offset//4] >> ((offset&3)*8).cast(dtypes.uint32)) & 255
 def _half(value:UOp) -> UOp: return value.cast(dtypes.uint16).bitcast(dtypes.float16).float()
@@ -152,22 +154,19 @@ def iq4_half_lut(device:str) -> Tensor:
 @functools.cache
 def _q8_quantize_kernel(q:UOp, scale:UOp, xsum:UOp, x:UOp, tokens:int, in_features:int) -> UOp:
   groups = in_features//Q8_GROUP_SIZE
-  token_group, lane = UOp.range(tokens*groups, 0, axis_type=AxisType.GLOBAL), UOp.range(32, 1, axis_type=AxisType.LOCAL)
+  token_group, lane = UOp.range(tokens*groups, 0, AxisType.GLOBAL), UOp.range(32, -1, AxisType.WARP)
   token, group = token_group//groups, token_group%groups
-  x = x.reshape(tokens, groups, 32)
-  group_scale = (warp_reduce(x[token, group, lane].float().abs(), maximum=True, full_wave=True) / 127).maximum(1e-8)
-  word_lane = lane.minimum(7)
-  xs = tuple(x[token, group, word_lane*4+i].float() for i in range(4))
-  qs = tuple((v/group_scale).round().clip(-127, 127).cast(dtypes.int8) for v in xs)
-  word = sum((v.cast(dtypes.uint8).cast(dtypes.uint32) << (i*8) for i, v in enumerate(qs)), UOp.const(0, dtypes.uint32))
-  # per-16 sums of the quantized values (lanes 0-3 / 4-7): Q4_K/Q5_K need the 32-sum, Q6_K the 16-sums
-  part = (lane < 8).where(sum((v.cast(dtypes.int32) for v in qs), UOp.const(0, dtypes.int32)), UOp.const(0, dtypes.int32))
-  gsum = [warp_reduce(((lane & 4).eq(h*4)).where(part, UOp.const(0, dtypes.int32)), full_wave=True) for h in range(2)]
-  store_half = (lane & 4) >> 2
-  stores = (q[token, group, lane.valid(lane < 8)].store(word),
-            UOp.group(scale[token, group.valid(lane.eq(0))].store(group_scale),
-                      xsum[token, group, store_half.valid(lane.eq(0) | lane.eq(4))].store(
-                        store_half.eq(0).where(gsum[0].float(), gsum[1].float()))))
+  value = x.reshape(tokens, groups, 32)[token, group, lane].float()
+  # Quantize each input once, then pack four neighboring lanes into one word.
+  d = (warp_reduce(value.abs(), maximum=True, full_wave=True)/127).maximum(1e-8)
+  rounded = UOp(Ops.CUSTOM, src=(value/d,), arg=("__builtin_nearbyintf({0})", dtypes.float))
+  quant = rounded.clip(-127, 127).cast(dtypes.int8)
+  word = quant.cast(dtypes.uint8).cast(dtypes.uint32) << ((lane%4)*8).cast(dtypes.uint32)
+  for offset in (1, 2):
+    word |= UOp(Ops.CUSTOM, src=(word,), arg=(f"__builtin_amdgcn_ds_swizzle({{0}}, {0x1f | offset<<10})", dtypes.uint32))
+  stores = (q[token, group, (lane//4).valid((lane%4).eq(0))].store(word),
+            scale[token, group.valid(lane.eq(0))].store(d),
+            xsum[token, group, (lane//16).valid((lane%16).eq(0))].store(warp_reduce(quant.float())))
   return UOp.group(*stores).end(token_group, lane).sink(arg=KernelInfo(name="q8_quantize", opts_to_apply=()))
 
 def q8_quantize(x:Tensor, tokens:int, in_features:int) -> tuple[Tensor, Tensor, Tensor]:
@@ -221,8 +220,8 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
     # the packed rows were padded to 212 bytes (53 words) per 256-block in set_quantized: everything is word-aligned
     base = (output*in_features//GGML_BLOCK_SIZE+block)*Q6_WORDS
     # the subgroup's 8 ql words and 8 qh words are contiguous: two 16-byte vector loads each
-    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4) for half in range(2))
-    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4) for half in range(2))
+    lows = tuple(_amd_load(raw[base + (subgroup//4)*16 + (subgroup%2)*8 + half*4], 4, stream=True) for half in range(2))
+    highs = tuple(_amd_load(raw[base + 32 + (subgroup//4)*8 + half*4], 4, stream=True) for half in range(2))
     dots = [UOp.const(0, dtypes.int32)] * 2
     for word_idx in range(8):
       within = (subgroup*32 + word_idx*4)%128
@@ -240,6 +239,7 @@ def _quant_decode_kernel(out:UOp, raw:UOp, xq:UOp, xd:UOp, xs:UOp, out_features:
   return _decode_linear(out, out_features, group_count, group_dot, names[ggml_type])
 
 def _wmma_layout(out:UOp, out_features:int, token_tile:int, output_tiles:int):
+  if out_features % (16*output_tiles): output_tiles = 1
   output_waves = 2 if out_features % (32*output_tiles) == 0 else 1
   token_block, output_block = UOp.range(out.shape[0]//token_tile, 0), UOp.range(out_features//(16*output_tiles*output_waves), 1)
   # lane is a hardware WARP range (like the flash kernel): the fragment math stays visible without being
@@ -318,15 +318,9 @@ def _iq4_linear_f16_wmma_kernel(out:UOp, raw:UOp, x:UOp, lut:UOp, out_features:i
   def dequant(base:UOp, subgroup:UOp, half:int) -> tuple[UOp, ...]:
     d, scale = _iq4_scales(raw, base, subgroup)
     scale = scale * d
-    if out_features <= 6144:
-      pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
-                    for word in range(4) for byte in range(4))
-      return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
-    # a subgroup-half gathers the lo (half=0) or hi (half=1) nibbles of byte pairs of each packed word
-    lut_pairs = (lut[(((raw[base+2+subgroup*4+i] >> (8*j+4*half)) & 15) |
-                      (((raw[base+2+subgroup*4+i] >> (8*j+8+4*half)) & 15) << 4)).cast(dtypes.weakint)]
-                 for i in range(4) for j in (0, 2))
-    return tuple((_half((pair >> (i*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in lut_pairs for i in range(2))
+    pairs = tuple(lut[((raw[base + 2 + subgroup*4 + word] >> (byte*8)) & 255).cast(dtypes.weakint)]
+                  for word in range(4) for byte in range(4))
+    return tuple((_half((pair >> (half*16)) & 0xffff)*scale).cast(dtypes.float16) for pair in pairs)
   return _quant_linear_wmma(out, x, out_features, in_features, IQ4_WORDS, layout, dequant, "linear_iq4_xs_f16_wmma")
 
 def q8_linear(layer:Linear, x:Tensor) -> Tensor:
@@ -369,21 +363,20 @@ def _amd_f16_gemv_kernel(out:UOp, w:UOp, x:UOp, *rest:UOp, in_features:int, out_
     for j in range(val_chunk):
       acc = acc + w[out_row, i, lane*val_chunk + j].load().float() * x[token, i, lane*val_chunk + j].load().float()
   total = warp_reduce(acc, full_wave=True)
-  if bias is not None: total = total + bias[token, out_row].load().float()
+  if bias is not None: total = total + bias[out_row].load().float()
   return out[token, out_row.valid(lane.eq(0))].store(total).end(token, out_row, lane).sink(arg=KernelInfo(name="linear_f16_gemv", opts_to_apply=()))
 
 def _view_back(t:Tensor) -> Tensor:
-  """strip top-of-chain CAST(s) from a lazy weight: reading the raw file bytes in the kernel instead of
-  materializing the cast into a fresh buffer every step"""
+  # Widening half to float is exact; preserve casts that round or change the values.
   uop = t.uop
-  while uop.op is Ops.CAST: uop = uop.src[0]
+  while uop.op is Ops.CAST and uop.dtype == dtypes.float32 and uop.src[0].dtype in (dtypes.half, dtypes.bfloat16): uop = uop.src[0]
   return Tensor(uop).reshape(t.shape)
 
 def f16_gemv(layer:Linear, x:Tensor) -> Tensor:
   tokens = prod(x.shape[:-1])
   assert isinstance(tokens, int)
   weight = _view_back(layer.weight)
-  x = x.contiguous() if x.dtype == dtypes.half else x.cast(dtypes.half).contiguous()
+  x = x.contiguous()
   out = Tensor.empty(tokens, layer.out_features, dtype=dtypes.float32, device=x.device)
   fxn = functools.partial(_amd_f16_gemv_kernel, in_features=layer.in_features, out_features=layer.out_features, tokens=tokens)
   srcs = (out, weight.reshape(-1), x.reshape(tokens, layer.in_features)) + (() if layer.bias is None else (_view_back(layer.bias),))
@@ -402,22 +395,25 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   _, B, H_KV, N, D = cast(tuple[int, int, int, int, int], cache_kv.shape)
   _, H, M, _ = cast(tuple[int, int, int, int], q.shape)
   assert M == 1 and H % H_KV == 0 and D % WARP_SIZE == 0 and max_kv_len <= N and max_kv_len % block_n == 0
-  G, CHUNK, DPL, WAVES = H // H_KV, block_n, D // WARP_SIZE, waves
+  G, CHUNK, DPL, WAVES, PARTIALS = H // H_KV, block_n, D // WARP_SIZE, waves, out.shape[2]
   assert CHUNK % WAVES == 0
   SEC = CHUNK // WAVES  # keys each wave scans independently
-  live_chunks = (valid_kv_len+CHUNK-1)//CHUNK
-  live_chunks = min(live_chunks, out.shape[2]) if isinstance(live_chunks, int) else live_chunks.minimum(out.shape[2])
+  total_chunks = (valid_kv_len+CHUNK-1)//CHUNK
+  live_chunks = min(total_chunks, PARTIALS) if isinstance(total_chunks, int) else total_chunks.minimum(PARTIALS)
   block_bhkv, block_chunk = UOp.range(B*H_KV, 0, AxisType.GLOBAL), UOp.range(live_chunks, 1, AxisType.GLOBAL)
   lane, wave = UOp.range(WARP_SIZE, -1, axis_type=AxisType.WARP), UOp.range(WAVES, 3, axis_type=AxisType.LOCAL)
   b, kv_head = block_bhkv // H_KV, block_bhkv % H_KV
   # per-lane query fragments for every GQA head, kept packed in registers; unpacked at use
   qf = tuple(_vec_load(q[b, kv_head*G+h, 0, lane*DPL], DPL) for h in range(G))
   zerof = UOp.const(0, dtypes.float)
+  # Each block scans every PARTIALS-th chunk, keeping an online softmax across rounds.
+  chunk_round = UOp.range((total_chunks-1-block_chunk)//PARTIALS+1, 4, AxisType.REDUCE)
+  chunk_id = block_chunk + chunk_round*PARTIALS
   valids: list[UOp] = []
   scores: list[list[UOp]] = [[zerof]*G for _ in range(SEC)]
   vfrags: list[tuple[UOp, ...]] = [()]*SEC
   for j in range(SEC):
-    key = block_chunk*CHUNK + wave*SEC + j
+    key = chunk_id*CHUNK + wave*SEC + j
     valid = key < valid_kv_len
     valids.append(valid)
     kfrag = _vec_load(cache_kv[0, b, kv_head, key, lane*DPL], DPL)
@@ -425,23 +421,32 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
     vfrags[j] = tuple(valid.where(v, zerof) for v in _vec_load(cache_kv[1, b, kv_head, key, lane*DPL], DPL))
     for h in range(G):
       s = warp_reduce(sum((qf[h][i]*kfrag[i] for i in range(DPL)), UOp.const(0, dtypes.float)), full_wave=True) * (1/math.sqrt(D))
-      scores[j][h] = valid.where(s, UOp.const(-math.inf, dtypes.float))
-  ninf = UOp.const(-math.inf, dtypes.float)
-  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), ninf) for h in range(G)]
-  accs:list[list[UOp]] = [[UOp.const(0, dtypes.float)] * DPL for _ in range(G)]
-  row_sums:list[UOp] = [UOp.const(0, dtypes.float) for _ in range(G)]
+      scores[j][h] = valid.where(s, UOp.const(-1e30, dtypes.float))
+  # A finite initial max keeps fully masked waves from computing exp(-inf - -inf).
+  acc_reg, max_reg, sum_reg = _reg((G, DPL), 2, 0), _reg((G,), 3, -1e30), _reg((G,), 4, 0)
+  prev_acc, prev_max, prev_sum = acc_reg.after(chunk_round), max_reg.after(chunk_round), sum_reg.after(chunk_round)
+  row_max = [functools.reduce(UOp.maximum, (scores[j][h] for j in range(SEC)), prev_max[h].load()) for h in range(G)]
+  # Rescale the previous rounds to the new max, then accumulate this round's keys.
+  alpha = [((prev_max[h].load()-row_max[h])*LOG2E).exp2() for h in range(G)]
+  accs = [[alpha[h]*prev_acc[h, i].load() for i in range(DPL)] for h in range(G)]
+  row_sums = [alpha[h]*prev_sum[h].load() for h in range(G)]
   for j in range(SEC):
     for h in range(G):
-      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), UOp.const(0, dtypes.float))
+      beta = valids[j].where(((scores[j][h]-row_max[h])*LOG2E).exp2(), zerof)
       accs[h] = [a + beta*v for a, v in zip(accs[h], vfrags[j])]
       row_sums[h] = row_sums[h] + beta
+  update = UOp.group(acc_reg.store(UOp.stack(*(x for acc in accs for x in acc)).reshape(G, DPL)),
+                     max_reg.store(UOp.stack(*row_max)), sum_reg.store(UOp.stack(*row_sums))).end(chunk_round)
+  acc_reg, max_reg, sum_reg = acc_reg.after(update), max_reg.after(update), sum_reg.after(update)
   # exchange across the block's waves through LDS (fp16 halves LDS so more blocks fit per CU)
-  acc_lds = UOp.placeholder((WAVES, G, D), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
+  # Matching cache/LDS strides can reuse a loop-local cache index outside the loop. Pad that layout.
+  acc_lds = UOp.placeholder((WAVES, G, D + (LDS_PAD if G == SEC else 0)), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)[:, :, :D]
   ml_lds = UOp.placeholder((WAVES, G, 2), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
   lds_acc = acc_lds.reshape(WAVES, G, WARP_SIZE, DPL)
-  stores = [lds_acc[wave, h, lane].store(UOp.stack(*accs[h]).cast(dtypes.half)) for h in range(G)]
+  # Normalize before fp16 to avoid overflow. Nonempty waves have sum >= 1; empty waves keep their zero accumulator.
+  stores = [lds_acc[wave, h, lane].store((acc_reg[h].load() / sum_reg[h].load().maximum(1)).cast(dtypes.half)) for h in range(G)]
   # NOTE: duplicate stores of the same value from every lane are harmless here
-  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((row_max[h], row_sums[h]))]
+  stores += [ml_lds[wave, h, i].store(x) for h in range(G) for i, x in enumerate((max_reg[h].load(), sum_reg[h].load()))]
   barrier = UOp.barrier(UOp.group(*stores))
   acc_lds, ml_lds = acc_lds.after(barrier), ml_lds.after(barrier)
   tid = wave*WARP_SIZE + lane
@@ -449,14 +454,16 @@ def _amd_flash_attention_decode_partial(out, stats, q, cache_kv, valid_kv_len, m
   for i in range(-(-G*D//(WAVES*WARP_SIZE))):
     flat = tid + i*WAVES*WARP_SIZE
     h, d = flat // D, flat % D
-    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)), ninf)
-    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * acc_lds[w, h, d].load().float() for w in range(WAVES)), UOp.const(0, dtypes.float))
+    M = functools.reduce(UOp.maximum, (ml_lds[w, h, 0].load() for w in range(WAVES)))
+    # LDS holds normalized values; restore each wave's sum before combining.
+    val = sum((((ml_lds[w, h, 0].load()-M)*LOG2E).exp2() * ml_lds[w, h, 1].load() * acc_lds[w, h, d].load().float()
+               for w in range(WAVES)), zerof)
     oidx = out[b, kv_head*G + h, block_chunk, d]
     if G*D % (WAVES*WARP_SIZE): oidx = out[b, (kv_head*G + h).valid(flat < G*D), block_chunk, d]
     final_stores.append(oidx.store(val))
   hstat = tid
-  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)), ninf)
-  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), UOp.const(0, dtypes.float))
+  M = functools.reduce(UOp.maximum, (ml_lds[w, hstat, 0].load() for w in range(WAVES)))
+  L = sum((((ml_lds[w, hstat, 0].load()-M)*LOG2E).exp2() * ml_lds[w, hstat, 1].load() for w in range(WAVES)), zerof)
   q_head = (kv_head*G + hstat).valid(hstat < G) if WAVES*WARP_SIZE > G else kv_head*G + hstat
   final_stores += [stats[b, q_head, block_chunk, 0].store(M), stats[b, q_head, block_chunk, 1].store(L)]
   return UOp.group(*final_stores).end(lane, wave, block_chunk, block_bhkv).sink(arg=KernelInfo(name="flash_decode_partial", opts_to_apply=()))
@@ -493,10 +500,13 @@ def _amd_flash_decode_combine(o:UOp, partial:UOp, stats:UOp, live:int|UOp) -> UO
 
 def amd_flash_attention_decode(q:Tensor, cache_kv:Tensor, valid_kv_len:int|UOp, max_kv_len:int) -> Tensor:
   B, H, D = cache_kv.shape[1], q.shape[1], cache_kv.shape[4]
-  chunks = min(256, max_kv_len // 64)
+  chunks = min(48, max_kv_len // 64)
   partial = Tensor.empty(B, H, chunks, D, dtype="float32", device=q.device)
   stats = Tensor.empty(B, H, chunks, 2, dtype="float32", device=q.device)
-  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=16)
+  waves, group = 16, H // cache_kv.shape[2]
+  while waves * group * ((D+LDS_PAD)*2 + 8) > 65536: waves //= 2
+  assert waves > 0, "attention head group exceeds shared memory capacity"
+  fxn = functools.partial(_amd_flash_attention_decode_partial, valid_kv_len=valid_kv_len, max_kv_len=max_kv_len, block_n=64, waves=waves)
   partial, stats = Tensor.custom_kernel(partial, stats, q, cache_kv, fxn=fxn)[:2]
   live = (valid_kv_len+63)//64
   live = min(live, chunks) if isinstance(live, int) else live.minimum(chunks)
@@ -512,7 +522,7 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   k, v = cache[0].reshape(B*H_KV, physical_n, cache_dim), cache[1].reshape(B*H_KV, physical_n, cache_dim)
   assert k.shape == v.shape and BH % k.shape[0] == 0 and k.shape[2] == D
   gqa_group = BH // k.shape[0]
-  if isinstance(M, int) and isinstance(valid_kv_len, int): assert M % BLOCK_M == 0 and valid_kv_len % BLOCK_N == 0
+  if isinstance(M, int): assert M % BLOCK_M == 0
   assert isinstance(D, int) and D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0
   TM, TN, TD, SCALE = BLOCK_M//(WAVES_M*LANES_PER_WAVE_M), BLOCK_N//LANES_PER_WAVE_N, D//(WAVES_N*LANES_PER_WAVE_N), 1/math.sqrt(D)
   # query row 0 sits at sequence position q_base (the queries may be padded beyond valid_kv_len - q_base rows)
@@ -540,7 +550,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
   q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
   k_frag = KV_lds_k.reshape(TN, WMMA_N, D // WMMA_K, WMMA_K)[tn1, lane_n, k_qk]
-  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk)
+  # All waves must finish reading Q/K before their shared memory is reused for P/V.
+  qk_done = S_frag.store(UOp.wmma(q_frag, k_frag, S_frag.after(k_qk), *WMMA_ARG)).end(tm1, tn1).end(k_qk).barrier()
   S_reg = S_reg.after(qk_done, S_reg.store(S_reg * SCALE))
   rm, rn = UOp.range(TM, 250), UOp.range(TN, 251)
   q_idx = q_base + block_m * BLOCK_M + wave_m * WMMA_M + rm * LANES_PER_WAVE_M + lane_m
@@ -567,7 +578,8 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
   acc, l_i, m_i, beta_i = acc.after(correction), l_i.after(correction), m_i.after(correction), beta_i.after(correction)
   V_lds = UOp.placeholder((D, BLOCK_N + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :BLOCK_N]
   V_copy, load_v = V_lds.after(qk_done).permute(1, 0), UOp.range(KV_ELEMS_PER_THREAD, 390)
-  vval = v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float()
+  v_pos = n_tile*BLOCK_N + (tid*KV_ELEMS_PER_THREAD + load_v)//D
+  vval = (v_pos < valid_kv_len).where(v.reshape(physical_n*D)[n_tile*BLOCK_N*D + tid*KV_ELEMS_PER_THREAD + load_v].float(), 0)
   V_store = V_copy.reshape(THREADS_PER_BLOCK, KV_ELEMS_PER_THREAD)[tid, load_v].store(vval).end(load_v)
   pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
   P_lds, V_lds = P_lds.after(pv_barrier), V_lds.after(pv_barrier)
@@ -589,7 +601,16 @@ def _amd_flash_attention(o:UOp, q:UOp, cache:UOp, valid_kv_len:int|UOp, q_start:
 def flash_attention(q:Tensor, assigned_kv:Tensor, valid_end:int|UOp) -> Tensor:
   # cached flash attention on the half KV cache (already written through assigned_kv); valid_end stays bound at the graph level
   T_real, q_start = q.shape[2], None
-  if resolve(T_real == 1): return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, assigned_kv.shape[3]))
+  D, N, group = q.shape[3], assigned_kv.shape[3], q.shape[1] // assigned_kv.shape[2]
+  decode = resolve(T_real == 1, False)
+  # Non-power-of-two decode dimensions can lose tail-store masks. Q/P, K, and V use separate LDS allocations.
+  supported = D % 32 == 0 and (D & (D-1) == 0 and N % 64 == 0 and group*((D+LDS_PAD)*2+8) <= 65536 if decode else
+    D >= 64 and 2*(2*BLOCK_M*(D+LDS_PAD) + D*(BLOCK_N+LDS_PAD)) <= 65536 and N % BLOCK_N == 0 and q.max_shape[2] % BLOCK_M == 0)
+  if not supported:
+    k, v = (assigned_kv[i, :, :, :valid_end].float() for i in range(2))
+    mask = None if decode else Tensor.full((T_real, valid_end), -math.inf, dtype=dtypes.float32, device=q.device).triu(valid_end-T_real+1)
+    return q.float().scaled_dot_product_attention(k, v, attn_mask=mask, enable_gqa=True)
+  if decode: return amd_flash_attention_decode(q.half(), assigned_kv, valid_end, cast(int, N))
   if isinstance(T_real, UOp):
     # symbolic chunk: pad the queries to the static tile; garbage rows are sliced off
     T_pad = q.max_shape[2]
@@ -643,12 +664,14 @@ def gated_delta_prefill(q:Tensor, k:Tensor, v:Tensor, beta:Tensor, alpha:Tensor,
   assert q.shape == k.shape and v.shape[:3] == beta.shape == (batch, heads, tokens) and state.shape == (batch, heads, value_dim, key_dim)
   assert alpha.shape[:3] == (batch, heads, tokens) and (len(alpha.shape) == 3 or alpha.shape[-1] in (1, value_dim))
   assert key_dim % 32 == 0 and value_dim % 4 == 0
+  assert q.dtype == k.dtype == dtypes.float32, "recurrent Q/K must be float32"
+  assert state.uop.contiguous_view_offset() is not None, "recurrent state must be contiguous"
+  if start_pos is not None:
+    assert start_pos.uop.is_bound_var
+    state = Tensor(state.uop.after(start_pos.uop))
   core, kq = Tensor.empty_like(v), (q*k).sum(-1).contiguous()
   srcs = (core, q.contiguous(), k.contiguous(), v.contiguous(), beta.contiguous(), alpha.contiguous(), state, kq)
-  if start_pos is None: return Tensor.custom_kernel(*srcs, fxn=_gated_delta_prefill_kernel)[0]
   contig = tuple(x.uop if x.uop.op is Ops.AFTER else x.uop.contiguous() for x in srcs)
   params = tuple(UOp.placeholder_like(x, slot=i) for i,x in enumerate(contig))
-  assert start_pos.uop.is_bound_var
-  # the bound start_pos reaches the graph through the state AFTER chain, like the flash kernels' valid_end
-  call = _gated_delta_prefill_kernel(*params, kernel_var(start_pos.uop.src[0])).call(*contig)
+  call = _gated_delta_prefill_kernel(*params, None if start_pos is None else kernel_var(start_pos.uop.src[0])).call(*contig)
   return Tensor(contig[0].after(call))

tinygrad/llm/serve.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 import json, pathlib, re, time, typing, uuid
 from typing import TYPE_CHECKING
 from tinygrad.helpers import DEBUG, colored, stderr_log
-from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
+from tinygrad.viz.serve import TCPServerWithReuse, Handler as VizHandler
 if TYPE_CHECKING:
   from tinygrad.llm.cli import SimpleTokenizer
   from tinygrad.llm.model import Transformer
@@ -60,11 +60,12 @@ class StreamRouter:
     if emit: yield "content", emit
     if found: self.mode, self.buf = "tool", "<tool_call>" + self.buf
 
-class Handler(HTTPRequestHandler):
+class Handler(VizHandler):
   server: LLMServer
   def log_request(self, code='-', size='-'): pass
   def do_GET(self):
     if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
+    elif self.path.startswith("/assets/"): super().do_GET()
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
   def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
                 reasoning:bool=False):

tinygrad/mixin/elementwise.py

@@ -56,14 +56,14 @@ class ElementwiseMixin(CreationMixin):
     """
     return self.cast(dtypes.bool).ne(True)
 
-  def contiguous(self, **kwargs) -> Self:
+  def contiguous(self) -> Self:
     """
     Returns a contiguous tensor.
     """
     if self.dtype in dtypes.weaks: return self
     uop = self._uop
     if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
-    return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
+    return self._wrap_uop(uop.alu(Ops.CONTIGUOUS))
 
   def contiguous_backward(self) -> Self:
     """
@@ -705,7 +705,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([-9., -6., -3., 0., 3., 6., 9.]).relu6().numpy())
     ```
     """
-    return self.relu() - (self-6).relu()
+    return ((r:=self.relu()) < 6).where(r, 6)
 
   def hardswish(self) -> Self:
     """

tinygrad/mixin/gradient.py

@@ -7,7 +7,10 @@ from tinygrad.function import renumber_invalid_outputs
 
 def reduce_gradient(ctx:UOp, ret:UOp, op:Ops):
   if op == Ops.ADD: return (ctx._broadcast_to(ret.src[0].shape),)
-  if op == Ops.MAX: return (((mask:=ret.src[0].eq(ret).cast(ctx.dtype))/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))) * ctx,)
+  if op == Ops.MAX:
+    # count the ties in the acc dtype, the count can overflow the gradient dtype
+    mask = ret.src[0].eq(ret).cast(sum_acc_dtype(ctx.dtype))
+    return ((mask/mask._rop(Ops.ADD, tuple(range(ret.arg[1])))).cast(ctx.dtype) * ctx,)
   if op == Ops.MUL:
     # d(prod x)/dx_j = prod_{i!=j} x_i: ret/x_j whenever x_j != 0 (any zero makes ret 0), else the product of the others
     safe_x, axes = (is_zero:=(x:=ret.src[0]).eq(0)).where(1, x), tuple(range(ret.arg[1]))
@@ -96,8 +99,11 @@ pm_gradient = PatternMatcher([
   (UPat(Ops.SINK), lambda ctx: ctx.src),
   (UPat(Ops.AFTER, src=(UPat.var("d"), UPat(Ops.CALL, name="k"))), lambda ctx, d, k:
     (ctx, UOp.sink(*([ctx if i == k.src.index(d)-1 else UOp(Ops.NOOP) for i in range(len(k.src)-1)])))),
+  # ordering-only AFTER: store target is a different buffer, gradient flows straight through to dest
+  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="t"), UPat())))),
+   lambda ctx, dest, t: (ctx, None) if t.buf_uop is not dest.buf_uop else None),
   # clone/assign gradient passes through to val
-  (UPat(Ops.AFTER, src=(UPat(), UPat(Ops.STORE))), lambda ctx: (None, ctx)),
+  (UPat(Ops.AFTER, src=(UPat(name="dest"), UPat(Ops.STORE, src=(UPat(name="dest"), UPat())))), lambda ctx,dest: (None, ctx)),
   (UPat(Ops.STORE, src=(UPat(), UPat())), lambda ctx: (None, ctx)),
   # there's no gradient for bitcast
   (UPat(Ops.BITCAST), lambda: (None,)),

tinygrad/mixin/rand.py

@@ -40,7 +40,7 @@ class RandMixin(OpMixin):
   def _rand(cls, key:Self, counter:Self, shape:tuple[int, ...], dtype:DType, contiguous:bool=True) -> Self:
     bits = cls.random_bits(key, counter, ceildiv(prod(shape) * dtype.itemsize, 4))
     out = cls._bits_to_rand(bits, shape, dtype)
-    return out.contiguous() if contiguous else out
+    return out.clone() if contiguous else out
 
   @staticmethod
   def _next_counter(device:str, num:int):
@@ -293,7 +293,7 @@ class RandMixin(OpMixin):
     if not 0 <= p <= 1: raise ValueError(f"{p=} is out of range [0, 1]")
     if not TRAINING or p == 0: return self
     if p == 1: return self.const_like(0)
-    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).contiguous().where(self, 0) / (1.0 - p)
+    return (self.rand_like(dtype=dtypes.default_float, contiguous=False) >= p).clone().where(self, 0) / (1.0 - p)
 
   def scaled_dot_product_attention(self, key:Self, value:Self, attn_mask:Self|None=None, dropout_p:float=0.0,
                                    is_causal:bool=False, enable_gqa:bool=False) -> Self:

tinygrad/renderer/amd/generate.py

@@ -233,12 +233,15 @@ def extract_pcode(pages: list[list[tuple[float, float, str, str]]], name_to_op:
       sorted_lines = sorted(lines, key=lambda x: (x[0], -x[1]))
       # Stop at large Y gaps (>30) - indicates section break
       filtered = [sorted_lines[0]]
+      depth = sorted_lines[0][2].count("{") - sorted_lines[0][2].count("}")
       for j in range(1, len(sorted_lines)):
         prev_page, prev_y, _ = sorted_lines[j-1]
         curr_page, curr_y, _ = sorted_lines[j]
-        if curr_page == prev_page and prev_y - curr_y > 30: break
-        if curr_page != prev_page and prev_y > 60 and curr_y < 730: break
+        if depth == 0 and curr_page == prev_page and prev_y - curr_y > 30: break
+        if depth == 0 and curr_page != prev_page and prev_y > 60 and curr_y < 730: break
         filtered.append(sorted_lines[j])
+        code = sorted_lines[j][2].split("//")[0]
+        depth += code.count("{") - code.count("}")
       pcode_lines = [t.replace('Ê', '').strip() for _, _, t in filtered]
       if pcode_lines: pcode[(name, opcode)] = '\n'.join(pcode_lines)
   return pcode

tinygrad/renderer/amd/sqtt.py

@@ -670,7 +670,6 @@ def map_insts(data:bytes, lib:bytes, target:str) -> Iterator[tuple[PacketType, I
       for wave in range(10):
         if (p.inst >> (wave * 2)) & 3 == 3:
           inst = pc_map[pc:=wave_pc[(p.simd, wave)]]
-          wave_pc[(p.simd, wave)] += inst.size()
           yield (p, InstructionInfo(pc, wave, inst))
     # map INST events on this SIMD to the program counter, we know the waves
     elif isinstance(p, (VALUINST, INST, INST_RDNA4, IMMEDIATE)) and not (isinstance(p, (INST, INST_RDNA4)) and p.op.name.startswith("OTHER_")):
@@ -735,5 +734,6 @@ if __name__ == "__main__":
   sqtt_events = [e for e in data if type(e).__name__ == "ProfileSQTTEvent"]
   evt_num = getenv("SQTT_EVENT", -1)
   for i, event in enumerate(sqtt_events):
-    print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
-    print_packets(decode(event.blob))
+    if evt_num == -1 or i == evt_num:
+      print(f"\n=== event {i} {prg_names.get(event.kern, '')} ===")
+      print_packets(decode(event.blob))

tinygrad/renderer/isa/__init__.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 import itertools
 from dataclasses import dataclass, field
 from tinygrad.renderer import Renderer
-from tinygrad.uop.ops import PatternMatcher, UOp, Ops, consumer_map_from_toposort
+from tinygrad.uop.ops import PatternMatcher, UOp, Ops
+from typing import Any
 
 @dataclass(frozen=True)
 class Register:
@@ -16,35 +17,31 @@ class Register:
 
 class IselContext:
   def __init__(self, sink:UOp):
-    self.uses = consumer_map_from_toposort(sink.toposort())
     self.reg_n = itertools.count()
-    def arg_key(u:UOp):
-      if u.op is Ops.SPECIAL: return (2, u.arg)
-      return (0, u.arg.slot) if u.arg.addrspace is not None else (1, u.expr)
-    self.func_args = sorted([u for u in self.uses if u.op in {Ops.PARAM, Ops.SPECIAL}], key=arg_key)
+    def arg_key(u:UOp): return (1, u.arg) if u.op is Ops.SPECIAL else (0, u.arg.slot)
+    self.func_args = sorted([u for u in sink.toposort() if u.op in {Ops.PARAM, Ops.SPECIAL}], key=arg_key)
 
   def vreg(self, cons:tuple[Register, ...]|Register):
     return Register(f"v{next(self.reg_n)}", 0, _cons=cons if isinstance(cons, tuple) else (cons,))
 
-def greg(u:UOp):
-  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return greg(u.src[0])
-  if isinstance(u.tag, tuple): return u.tag[0]
-  return u.tag
+def rdef(u:UOp):
+  if u.op in {Ops.NOOP, Ops.AFTER, Ops.BITCAST} and u.src: return rdef(u.src[0])
+  return u.tag[0] if isinstance(u.tag, tuple) else u.tag
 
-@dataclass
-class PreRegAllocContext:
-  lock: UOp|None = None
-  clobbered: set[UOp] = field(default_factory=set)
+class LinearContext:
+  def __init__(self, ren:ISARenderer):
+    self.ren, self.stack_size = ren, 0
+    self.loop_label: dict[UOp, str] = {}
+  def assign_spill_slot(self, r:Register, u:UOp) -> Any: raise NotImplementedError("arch specific")
 
 class ISARenderer(Renderer):
   pre_isel_matcher: PatternMatcher
   isel_matcher: PatternMatcher
-  pre_regalloc_matcher: PatternMatcher|None = None
+  pre_regalloc_matcher: PatternMatcher
   post_regalloc_matcher: PatternMatcher
+  linear_ctx_type: type = LinearContext
 
   def is_two_address(self, x:UOp) -> bool: return False
-  def stack_pointer(self) -> UOp: raise NotImplementedError("arch specific")
-  def copy(self, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
-  def spill(self, disp:UOp, x:UOp) -> UOp: raise NotImplementedError("arch specific")
-  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
+  def spill(self, spill_slot:Any, x:UOp) -> UOp: raise NotImplementedError("arch specific")
+  def fill(self, spill_slot:Any, x:UOp, reg:Register) -> UOp: raise NotImplementedError("arch specific")
   def asm_str(self, uops:list[UOp], function_name:str) -> str: raise NotImplementedError("arch specific")

tinygrad/renderer/isa/x86.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 # flake8: noqa: E702
 # allow semicolons to put multiple ops on one line
 import sys, struct, functools
@@ -6,7 +7,7 @@ from dataclasses import replace
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
 from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
-from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
+from tinygrad.renderer.isa import ISARenderer, IselContext, Register, LinearContext, rdef
 from tinygrad.helpers import unwrap, Target
 
 # ***** X86 Ops *****
@@ -23,18 +24,13 @@ class X86Ops(FastEnum):
   VMOVSSm = auto(); VMOVSDm = auto(); VMOVUPSm = auto()
   # casts
   MOVZX = auto(); MOVSX = auto(); MOVSXD = auto()
-  VPMOVZXBW = auto(); VPMOVZXBD = auto(); VPMOVZXBQ = auto()
-  VPMOVZXWD = auto(); VPMOVZXWQ = auto(); VPMOVZXDQ = auto()
-  VPMOVSXBW = auto(); VPMOVSXBD = auto(); VPMOVSXBQ = auto()
-  VPMOVSXWD = auto(); VPMOVSXWQ = auto(); VPMOVSXDQ = auto()
-  VCVTDQ2PS = auto(); VCVTDQ2PD = auto(); VCVTTPS2DQ = auto(); VCVTTPD2DQ = auto()
-  VCVTPH2PS = auto(); VCVTPS2PH = auto(); VCVTPS2PD = auto(); VCVTPD2PS = auto()
+  VCVTPH2PS = auto(); VCVTPS2PH = auto()
   VCVTSS2SD = auto(); VCVTSD2SS = auto(); VCVTSI2SS = auto(); VCVTSI2SD = auto()
   VCVTTSS2SI = auto(); VCVTTSD2SI = auto()
   # bitcasts
   VMOVD = auto(); VMOVQ = auto(); VMOVDm = auto(); VMOVQm = auto()
   # comparisons
-  VCMPSS = auto(); VCMPSD = auto(); VCMPPS = auto(); VCMPPD = auto()
+  VCMPSS = auto(); VCMPSD = auto()
   SETNE = auto(); SETE = auto(); SETL = auto(); SETB = auto()
   # where
   CMOVNE = auto(); CMOVE = auto(); CMOVL = auto(); CMOVB = auto()
@@ -43,29 +39,17 @@ class X86Ops(FastEnum):
   JNE = auto(); JE = auto(); JL = auto(); JB = auto(); JGE = auto(); JMP = auto()
   # vectorize / gep
   VINSERTPS = auto(); VPSRLDQ = auto()
-  VPEXTRB = auto(); VPEXTRW = auto(); VPEXTRD = auto(); VPEXTRQ = auto()
-  VPINSRB = auto(); VPINSRW = auto(); VPINSRD = auto(); VPINSRQ = auto()
+  VPEXTRW = auto(); VPEXTRD = auto()
+  VPINSRW = auto(); VPINSRD = auto()
   # int binary
   IDIV = auto(); DIV = auto()
   ADD = auto(); ADDi = auto(); SUB = auto(); SUBi = auto(); IMUL = auto(); IMULi = auto()
   AND = auto(); ANDi = auto(); XOR = auto(); XORi = auto(); OR = auto(); ORi = auto()
   SHL = auto(); SHLi = auto(); SHR = auto(); SHRi = auto(); SAR = auto(); SARi = auto(); CMP = auto(); CMPi = auto()
   # float unary (sometimes not unary)
-  VROUNDSS = auto(); VROUNDSD = auto(); VROUNDPS = auto(); VROUNDPD = auto()
-  VSQRTSS = auto(); VSQRTSD = auto(); VSQRTPS = auto(); VSQRTPD = auto()
-  # float scalar / vector binary
-  VADDSS = auto(); VADDSD = auto(); VADDPS = auto(); VADDPD = auto()
-  VSUBSS = auto(); VSUBSD = auto(); VSUBPS = auto(); VSUBPD = auto()
-  VMULSS = auto(); VMULSD = auto(); VMULPS = auto(); VMULPD = auto()
-  VDIVSS = auto(); VDIVSD = auto(); VDIVPS = auto(); VDIVPD = auto()
-  # int vector binary
-  VPADDB = auto(); VPADDW = auto(); VPADDD = auto(); VPADDQ = auto()
-  VPSUBB = auto(); VPSUBW = auto(); VPSUBD = auto(); VPSUBQ = auto()
-  VPMULLW = auto(); VPMULLD = auto()
-  # packed bitwise
-  VPAND = auto(); VPOR = auto(); VPXOR = auto()
-  # packed variable shifts
-  VPSLLVD = auto(); VPSLLVQ = auto(); VPSRLVD = auto(); VPSRLVQ = auto(); VPSRAVD = auto()
+  VROUNDSS = auto(); VROUNDSD = auto(); VSQRTSS = auto(); VSQRTSD = auto()
+  # float binary
+  VADDSS = auto(); VADDSD = auto(); VSUBSS = auto(); VSUBSD = auto(); VMULSS = auto(); VMULSD = auto(); VDIVSS = auto(); VDIVSD = auto()
   # return
   RET = auto()
 
@@ -75,30 +59,18 @@ class X86GroupOp:
                 X86Ops.SUB, X86Ops.SUBi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi,
                 X86Ops.IDIV, X86Ops.DIV, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB}
 
-  # X86Ops whose first src can read from memory
-  ReadMem1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
-                X86Ops.VPMOVZXBW, X86Ops.VPMOVZXBD, X86Ops.VPMOVZXBQ, X86Ops.VPMOVZXWD, X86Ops.VPMOVZXWQ, X86Ops.VPMOVZXDQ,
-                X86Ops.VPMOVSXBW, X86Ops.VPMOVSXBD, X86Ops.VPMOVSXBQ, X86Ops.VPMOVSXWD, X86Ops.VPMOVSXWQ, X86Ops.VPMOVSXDQ,
-                X86Ops.VCVTDQ2PS, X86Ops.VCVTDQ2PD, X86Ops.VCVTTPS2DQ, X86Ops.VCVTTPD2DQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
-                X86Ops.VCVTPH2PS, X86Ops.VCVTPS2PD, X86Ops.VCVTPD2PS, X86Ops.VROUNDPS, X86Ops.VROUNDPD, X86Ops.VSQRTPS, X86Ops.VSQRTPD,
-                X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA}
-
-  # X86Ops whose second src can read from memory NOTE: some of these are TwoAddress so the second src is actually the first
-  ReadMem2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
-                X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VADDPS, X86Ops.VADDPD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VSUBPS, X86Ops.VSUBPD,
-                X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VMULPS, X86Ops.VMULPD, X86Ops.VDIVSS, X86Ops.VDIVSD, X86Ops.VDIVPS, X86Ops.VDIVPD,
-                X86Ops.VPADDB, X86Ops.VPADDW, X86Ops.VPADDD, X86Ops.VPADDQ, X86Ops.VPSUBB, X86Ops.VPSUBW, X86Ops.VPSUBD, X86Ops.VPSUBQ,
-                X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VCMPPS, X86Ops.VCMPPD,
-                X86Ops.VPMULLW, X86Ops.VPMULLD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD, X86Ops.VINSERTPS,
-                X86Ops.VPINSRB, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.VPINSRQ, X86Ops.VPAND, X86Ops.VPOR, X86Ops.VPXOR, X86Ops.VPSLLVD,
-                X86Ops.VPSLLVQ, X86Ops.VPSRLVD, X86Ops.VPSRLVQ, X86Ops.VPSRAVD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
-                X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
+  # X86Ops whose second src is the rm field, so that src is what can be a memory operand
+  Rm2nd = {X86Ops.ADD, X86Ops.SUB, X86Ops.AND, X86Ops.OR, X86Ops.XOR, X86Ops.IMUL, X86Ops.CMP,
+           X86Ops.VADDSS, X86Ops.VADDSD, X86Ops.VSUBSS, X86Ops.VSUBSD, X86Ops.VMULSS, X86Ops.VMULSD, X86Ops.VDIVSS, X86Ops.VDIVSD,
+           X86Ops.VBLENDVPS, X86Ops.VBLENDVPD, X86Ops.VCMPSS, X86Ops.VCMPSD, X86Ops.VROUNDSS, X86Ops.VROUNDSD, X86Ops.VSQRTSS, X86Ops.VSQRTSD,
+           X86Ops.VINSERTPS, X86Ops.VPINSRW, X86Ops.VPINSRD, X86Ops.CMOVNE, X86Ops.CMOVE, X86Ops.CMOVL, X86Ops.CMOVB,
+           X86Ops.VCVTSI2SS, X86Ops.VCVTSI2SD, X86Ops.VCVTSS2SD, X86Ops.VCVTSD2SS, X86Ops.IDIV, X86Ops.DIV}
 
   # X86Ops that can write to memory
   WriteMem = {X86Ops.MOVm, X86Ops.MOVi, X86Ops.VMOVSSm, X86Ops.VMOVSDm, X86Ops.VMOVUPSm, X86Ops.VMOVDm, X86Ops.VMOVQm,
               X86Ops.ADDi, X86Ops.SUBi, X86Ops.ANDi, X86Ops.ORi, X86Ops.XORi, X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR,
               X86Ops.SARi, X86Ops.SETNE, X86Ops.SETE, X86Ops.SETL, X86Ops.SETB,
-              X86Ops.VCVTPS2PH, X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ}
+              X86Ops.VCVTPS2PH, X86Ops.VPEXTRW, X86Ops.VPEXTRD}
 
   # X86Ops that read flags
   ReadFlags = {X86Ops.CMOVB, X86Ops.CMOVL, X86Ops.CMOVE, X86Ops.CMOVNE, X86Ops.SETB, X86Ops.SETL, X86Ops.SETE, X86Ops.SETNE, X86Ops.JB, X86Ops.JL,
@@ -109,11 +81,9 @@ class X86GroupOp:
                 X86Ops.SHL, X86Ops.SHLi, X86Ops.SHR, X86Ops.SHRi, X86Ops.SAR, X86Ops.SARi, X86Ops.AND, X86Ops.ANDi, X86Ops.XOR, X86Ops.XORi,
                 X86Ops.OR, X86Ops.ORi}
 
-  # X86Ops whose first src is the rm field
-  Rm1st = ReadMem1st | (ReadMem2nd & TwoAddress) | {X86Ops.VPSRLDQ}
-
-  # X86Ops whose second src is the rm field
-  Rm2nd = ReadMem2nd
+  # X86Ops whose first src is the rm field. a TwoAddress op drops its first src post regalloc, so its Rm2nd src ends up first
+  Rm1st = {X86Ops.MOV, X86Ops.VMOVSS, X86Ops.VMOVSD, X86Ops.VMOVUPS, X86Ops.MOVZX, X86Ops.MOVSX, X86Ops.MOVSXD, X86Ops.VMOVD, X86Ops.VMOVQ,
+           X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI, X86Ops.VCVTPH2PS, X86Ops.CMPi, X86Ops.IMULi, X86Ops.LEA, X86Ops.VPSRLDQ} | (Rm2nd & TwoAddress)
 
 # ***** X86 legalization *****
 
@@ -136,15 +106,12 @@ extra_matcher = PatternMatcher([
   # no int8 mul or cmove, cast to int16
   (UPat.var("a", dtypes.int8s) * UPat.var("b"), lambda a,b: (a.cast(dtypes.int16) * b.cast(dtypes.int16)).cast(a.dtype)),
   (UPat.var("m").where(UPat.var("a", (dtypes.bool,)+dtypes.int8s), UPat.var("b")),
-   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype) if a.max_numel() == 1 else None),
+   lambda m,a,b: m.where(a.cast(dtypes.int16), b.cast(dtypes.int16)).cast(a.dtype)),
   # float16 alus are done in float32
   (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
    src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
   (UPat(GroupOp.Comparison, src=[UPat(dtype=dtypes.float16), UPat()], name="x"),
    lambda x: UOp(x.op, src=tuple(s.cast(dtypes.float32) for s in x.src)).cast(x.dtype)),
-  # no cmpne for packed ints, y != x => !(y==x)
-  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
-   lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
   # a float WHERE blends at the width of its value, so it needs a comparison at that width to make the mask
   (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
    lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and promo_dtype(m.src) is not w.dtype else None),
@@ -179,8 +146,8 @@ def flag_gate(m:UOp) -> UOp|None:
 
 # legalize the new style graph for isel. NOTE: this runs after the spec is verified, some of these rewrites violate it
 pre_isel_matcher = PatternMatcher([
-  # widening a scalar uint32 is free, the 32bit write that produced it already zeroed the upper half
-  (UPat.var("y", dtypes.uint32).cast(dtypes.int64s, name="x"), lambda y,x: x.replace(op=Ops.BITCAST) if y.max_numel() == 1 else None),
+  # widening a uint32 is free, the 32bit write that produced it already zeroed the upper half
+  (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.replace(op=Ops.BITCAST)),
   (UPat.var("y", dtypes.ints+(dtypes.bool,)).cast(dtypes.ints, name="x"),
    lambda y,x: x.replace(op=Ops.BITCAST) if x.dtype.itemsize == y.dtype.itemsize else None),
   # gated load/store become a conditional move on the address, the load/store are unconditional
@@ -192,6 +159,9 @@ pre_isel_matcher = PatternMatcher([
 ])
 
 # ***** X86 registers *****
+def def_reg(dt:DType, reg:Register) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=(reg,))
+# undefined operand, used for VEX instructions
+def undef(): return UOp(Ops.NOOP)
 
 RAX = Register("rax", 0)
 RCX = Register("rcx", 1)
@@ -212,11 +182,12 @@ reg_strs = {"rax": {4:"eax", 2:"ax", 1:"al"}, "rcx": {4:"ecx", 2:"cx", 1:"cl"},
         "rsp": {4:"esp", 2:"sp", 1:"spl"}, "rbp": {4:"ebp", 2:"bp", 1:"bpl"}, "rsi": {4:"esi", 2:"si", 1:"sil"}, "rdi": {4:"edi", 2:"di", 1:"dil"},
         **{f"r{i}": {4:f"r{i}d", 2:f"r{i}w", 1:f"r{i}b"} for i in range(8, 16)}}
 
+stack_pointer = def_reg(dtypes.uint64, RSP)
+
 # ***** X86 instruction selection *****
 def base(x:UOp, i:int) -> UOp: return s.src[0] if (s:=x.src[i]).op is Ops.INDEX else s
 def lane(x:UOp, i:int) -> int: return s.src[1].src[0].val if (s:=x.src[i]).op is Ops.INDEX else 0
 def to_int(dt:DType): return {dtypes.float16: dtypes.int16, dtypes.float32: dtypes.int32, dtypes.float64: dtypes.int64}[dt]
-def def_reg(dt:DType, reg:Register|None=None) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=None if reg is None else (reg,))
 def imm(dt:DType, v:int) -> UOp: return UOp.cconst(truncate[dt](v), dt).rtag()
 def to_imm(c:UOp) -> UOp|None:
   if not (c.op is Ops.CAST and (v:=c.src[0]).op is Ops.CONST): return None
@@ -231,8 +202,7 @@ def cmp(x:UOp) -> UOp:
 # comparisons that produce masks, the mask has the width of the operands
 def mask(x:UOp) -> UOp:
   dt, v = x.src[0].dtype, imm(dtypes.uint8, {Ops.CMPLT: 1, Ops.CMPNE: 4, Ops.CMPEQ: 0}[x.op])
-  if dt is dtypes.float32: return x.ins(X86Ops.VCMPSS if x.max_numel() == 1 else X86Ops.VCMPPS, dtype=dt, src=x.src + (v,))
-  return x.ins(X86Ops.VCMPSD if x.max_numel() == 1 else X86Ops.VCMPPD, dtype=dt, src=x.src + (v,))
+  return x.ins(X86Ops.VCMPSS if dt is dtypes.float32 else X86Ops.VCMPSD, dtype=dt, src=x.src + (v,))
 
 # vinsertps xmm2, xmm0, xmm1, imm
 # inserts any 32 bit element in xmm1 into any position in xmm0 according to immm, result is written to xmm2
@@ -241,13 +211,13 @@ def vinsertps(x:UOp) -> UOp:
   def _insert(ret:UOp, i:int) -> UOp:
     s, v = base(x, i), lane(x, i)
     return x.ins(X86Ops.VINSERTPS, src=(ret, s, imm(dtypes.uint8, v << 6 | i << 4)))
-  return functools.reduce(_insert, range(len(x.src)), def_reg(x.dtype))
+  return functools.reduce(_insert, range(len(x.src)), undef())
 
-# vpinsq xmm2, xmm0, rax, imm
-# inserts element in rax into any position in xmm0, result is written to xmm2 according to imm
+# vpinsrd xmm2, xmm0, eax, imm
+# inserts the element in eax into any position in xmm0, result is written to xmm2 according to imm
 def vpins(x:UOp, srcs:tuple[UOp, ...]) -> UOp:
-  op = {1: X86Ops.VPINSRB, 2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD, 8: X86Ops.VPINSRQ}[x.dtype.itemsize]
-  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), def_reg(x.dtype))
+  op = {2: X86Ops.VPINSRW, 4: X86Ops.VPINSRD}[x.dtype.itemsize]
+  return functools.reduce(lambda ret,i: x.ins(op, src=(ret, srcs[i], imm(dtypes.uint8, i))), range(len(srcs)), undef())
 
 # we don't call ctx.vreg on the srcs to avoid duplicates, a rewrite will assign the tuple of valid registers to a vreg
 def idiv(ctx:IselContext, x:UOp) -> UOp:
@@ -300,14 +270,13 @@ def abi(ctx:IselContext, x:UOp) -> UOp|None:
   # the shape srcs of a PARAM are not values, tag them so they aren't materialized into registers
   def _reg_arg(r:Register) -> tuple[UOp, ...]: return (x.replace(arg=arg, src=tuple(s.rtag() for s in x.src), tag=(r,)),)
   def _stack_arg(disp:int):
-    return (def_reg(dtypes.uint64, RSP), UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), tag=disp), imm(dtypes.uint8, 8))
+    return (stack_pointer, UOp(Ops.NOOP), UOp(Ops.INS, arg=(X86Ops.FRAME_INDEX, dtypes.int32), src=(imm(dtypes.int32, disp),)), imm(dtypes.uint8, 8))
   if sys.platform == "win32": src = _reg_arg((RCX, RDX, GPR[8], GPR[9])[i]) if i < 4 else _stack_arg((i-3)*8+32)
   else: src = _reg_arg((RDI, RSI, RDX, RCX, GPR[8], GPR[9])[i]) if i < 6 else _stack_arg((i-5)*8)
   # this move "cleanses" the abi register constraint
   return x.ins(X86Ops.MOV, dtype=dt, src=src)
 
-GPR_DEST_OPS = {X86Ops.VPEXTRB, X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VPEXTRQ, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI,
-                 X86Ops.VMOVDm, X86Ops.VMOVQm}
+GPR_DEST_OPS = {X86Ops.VPEXTRW, X86Ops.VPEXTRD, X86Ops.VCVTTSS2SI, X86Ops.VCVTTSD2SI, X86Ops.VMOVDm, X86Ops.VMOVQm}
 XMM_OPS = {op for op in X86Ops if op.name.startswith('V')} - GPR_DEST_OPS
 
 def _is_vec_xmm(y: UOp) -> bool:
@@ -351,15 +320,14 @@ isel_matcher = PatternMatcher([
   # **** Op -> Op ****
   # range is lowered to acc, cmp, jmp after regalloc
   (UPat(Ops.RANGE, src=(UPat.cvar("c").cast(),), allow_any_len=True, name="x"), lambda c,x: x.replace(src=(imm(x.dtype, c.val),) + x.src[1:])),
-  (UPat(Ops.RANGE, name="x"), lambda ctx,x: x.replace(tag=(ctx.vreg(WGPR),)) if not isinstance(x.tag, tuple) else None),
   # really all a backedge END is is an IF with a tag referencing the RANGE start label
   (UPat(Ops.END, src=(UPat(), UPat(), UPat(GroupOp.Comparison, name="cond")), name="x"),
     lambda x,cond: cond.ins(X86Ops.LOOP_CMP, tag=cond.op, src=cond.src + x.src[:2])),
   # **** Op -> X86Op ****
   # add callee saved registers to the RET, these will be scheduled at the top of the kernel and will be saved/restored if they are used in regalloc
-  # so regalloc builds the prologue/epilogue naturally
+  # so regalloc builds the prologue/epilogue naturally. they all share the stack pointer define's dtype so the the stack pointer define is first
   (UPat(Ops.SINK, name="x"), lambda x:
-   x.replace(src=(x.ins(X86Ops.RET, src=x.src + tuple(def_reg(dtypes.uint64 if r in GPR else dtypes.float64, r) for r in CALLEE_SAVED)),)) \
+   x.replace(src=(x.ins(X86Ops.RET, src=x.src + (stack_pointer,) + tuple(def_reg(dtypes.uint64, r) for r in CALLEE_SAVED)),))
     if not x.src or x.src[0].op is not Ops.INS or x.src[0].arg[0] is not X86Ops.RET else None),
   # function abi constraints
   (UPat((Ops.PARAM, Ops.SPECIAL), name="x"), abi),
@@ -399,50 +367,22 @@ isel_matcher = PatternMatcher([
   (UPat(Ops.CMPEQ, name="x"), lambda x: x.ins(X86Ops.SETE, src=(cmp(x),))),
   (UPat(Ops.CMPNE, name="x"), lambda x: x.ins(X86Ops.SETNE, src=(cmp(x),))),
   # float unary
-  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPS)),
-  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y)) if x.max_numel() == 1 else x.ins(X86Ops.VSQRTPD)),
-  (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x:
-   x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPS, src=(y, imm(dtypes.uint8, 3)))),
-  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x:
-   x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3))) if x.max_numel() == 1 else x.ins(X86Ops.VROUNDPD, src=(y, imm(dtypes.uint8, 3)))),
+  (UPat.var("y", dtypes.float32).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSS, src=(y, y))),
+  (UPat.var("y", dtypes.float64).sqrt().named("x"), lambda y,x: x.ins(X86Ops.VSQRTSD, src=(y, y))),
+  (UPat.var("y", dtypes.float32).trunc().named("x"), lambda y,x: x.ins(X86Ops.VROUNDSS, src=(y, y, imm(dtypes.uint8, 3)))),
+  (UPat.var("y", dtypes.float64).trunc().named("x"), lambda y,x: x.ins(X86Ops.VROUNDSD, src=(y, y, imm(dtypes.uint8, 3)))),
   # for float16 we route the srcs through gprs, this is suboptimal for values in xmms, in that case we want vpunpcklwd
   (UPat(Ops.STACK, dtypes.float16, name="x"), lambda x: vpins(x, tuple(s.bitcast(dtypes.int16) for s in x.src))),
   (UPat(Ops.STACK, dtypes.float32, name="x"), vinsertps),
-  (UPat(Ops.STACK, dtypes.ints+(dtypes.bool,), name="x"), lambda x: vpins(x, x.src)),
+  (UPat(Ops.STACK, dtypes.int32s, name="x"), lambda x: vpins(x, x.src)),
   # INDEX on a vector register value extracts a single element
-  (UPat.var("y", dtypes.int8s+(dtypes.bool,)).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRB, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int16s).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.int32s).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPEXTRD, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
-  (UPat.var("y", dtypes.int64s).index(UPat.cvar("c").cast(), name="x"),
-   lambda y,c,x: x.ins(X86Ops.VPEXTRQ, src=(y, imm(dtypes.uint8, c.val))) if _is_vec_xmm(y) else None),
   (UPat.var("y", dtypes.floats).index(UPat.cvar("c").cast(), name="x"),
    lambda y,c,x: x.ins(X86Ops.VPSRLDQ, src=(y, imm(dtypes.uint8, c.val * x.dtype.itemsize))) if _is_vec_xmm(y) else None),
-  # packed bitwise
-  ((UPat() & UPat()).named("x"), lambda x: x.ins(X86Ops.VPAND) if x.max_numel() > 1 else None),
-  ((UPat() | UPat()).named("x"), lambda x: x.ins(X86Ops.VPOR) if x.max_numel() > 1 else None),
-  ((UPat() ^ UPat()).named("x"), lambda x: x.ins(X86Ops.VPXOR) if x.max_numel() > 1 else None),
-  # packed int binary
-  ((UPat(dtype=dtypes.int32s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int64s) << UPat()).named("x"), lambda x: x.ins(X86Ops.VPSLLVQ) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.uint32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.uint64) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRLVQ) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int32) >> UPat()).named("x"), lambda x: x.ins(X86Ops.VPSRAVD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int8s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDB) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int16s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDW) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int32s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDD) if x.max_numel() > 1 else None),
-  ((UPat(dtype=dtypes.int64s) + UPat()).named("x"), lambda x: x.ins(X86Ops.VPADDQ) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int8s, name="x"), lambda x: x.ins(X86Ops.VPSUBB) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPSUBW) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPSUBD) if x.max_numel() > 1 else None),
-  (UPat(Ops.SUB, dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPSUBQ) if x.max_numel() > 1 else None),
-  (UPat(Ops.MUL, dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMULLW) if x.max_numel() > 1 else None),
-  (UPat(Ops.MUL, dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMULLD) if x.max_numel() > 1 else None),
-  # scalar int binary
+  # int binary
   ((UPat(dtype=dtypes.ints).alu(Ops.CDIV, UPat())).named("x"), idiv),
-  # scalar int binary with immediate
+  # int binary with immediate
   (UPat.var("a", dtypes.ints) << UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHLi, src=(a, imm(dtypes.uint8, c.val)))),
   (UPat.var("a", dtypes.uints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SHRi, src=(a, imm(dtypes.uint8, c.val)))),
   (UPat.var("a", dtypes.sints) >> UPat.cvar("c").cast(), lambda a,c: a.ins(X86Ops.SARi, src=(a, imm(dtypes.uint8, c.val)))),
@@ -456,7 +396,7 @@ isel_matcher = PatternMatcher([
    lambda a,c: a.ins(X86Ops.XORi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.cvar().cast(name="c"))),
    lambda a,c: a.ins(X86Ops.SUBi, src=(a, i)) if (i:=to_imm(c)) is not None else None),
-  # scalar int binary with register
+  # int binary with register
   ((UPat(dtype=dtypes.ints) << UPat()).named("x"), lambda x: shift(x, X86Ops.SHL)),
   ((UPat(dtype=dtypes.uints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SHR)),
   ((UPat(dtype=dtypes.sints) >> UPat()).named("x"), lambda x: shift(x, X86Ops.SAR)),
@@ -467,47 +407,28 @@ isel_matcher = PatternMatcher([
   (UPat.var("a", dtypes.ints+(dtypes.bool,)) ^ UPat.var("b"), lambda a,b: a.ins(X86Ops.XOR, src=(a, b))),
   (UPat(Ops.SUB, dtypes.ints, (UPat.var("a"), UPat.var("b"))), lambda a,b: a.ins(X86Ops.SUB, src=(a, b))),
   # float binary
-  ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS if x.max_numel() == 1 else X86Ops.VADDPS)),
-  ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD if x.max_numel() == 1 else X86Ops.VADDPD)),
-  ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS if x.max_numel() == 1 else X86Ops.VMULPS)),
-  ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD if x.max_numel() == 1 else X86Ops.VMULPD)),
-  (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS if x.max_numel() == 1 else X86Ops.VSUBPS)),
-  (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD if x.max_numel() == 1 else X86Ops.VSUBPD)),
-  (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS if x.max_numel() == 1 else X86Ops.VDIVPS)),
-  (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD if x.max_numel() == 1 else X86Ops.VDIVPD)),
+  ((UPat(dtype=dtypes.float32) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSS)),
+  ((UPat(dtype=dtypes.float64) + UPat()).named("x"), lambda x: x.ins(X86Ops.VADDSD)),
+  ((UPat(dtype=dtypes.float32) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSS)),
+  ((UPat(dtype=dtypes.float64) * UPat()).named("x"), lambda x: x.ins(X86Ops.VMULSD)),
+  (UPat(Ops.SUB, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VSUBSS)),
+  (UPat(Ops.SUB, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VSUBSD)),
+  (UPat(Ops.FDIV, dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VDIVSS)),
+  (UPat(Ops.FDIV, dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VDIVSD)),
   # casts
-  (UPat(dtype=dtypes.int32).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PS) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.int32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTDQ2PD) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float32).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPS2DQ) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float64).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VCVTTPD2DQ) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float32).cast(dtypes.float64, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PD) if x.max_numel() > 1 else None),
-  (UPat(dtype=dtypes.float64).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPD2PS) if x.max_numel() > 1 else None),
   (UPat(dtype=dtypes.float32).cast(dtypes.float16, name="x"), lambda x: x.ins(X86Ops.VCVTPS2PH, src=x.src + (imm(dtypes.uint8, 4),))),
   (UPat(dtype=dtypes.float16).cast(dtypes.float32, name="x"), lambda x: x.ins(X86Ops.VCVTPH2PS)),
   (UPat(dtype=dtypes.float32).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSS2SI)),
   (UPat(dtype=dtypes.float64).cast(dtypes.int32s+dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VCVTTSD2SI)),
   (UPat.var("y", dtypes.float32).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSS2SD, src=(y, y))),
   (UPat.var("y", dtypes.float64).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSD2SS, src=(y, y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(def_reg(x.dtype), y))),
-  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(def_reg(x.dtype), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float32, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SS, src=(undef(), y))),
+  (UPat.var("y", (dtypes.int32, dtypes.int64)).cast(dtypes.float64, name="x"), lambda y,x: x.ins(X86Ops.VCVTSI2SD, src=(undef(), y))),
   (UPat(dtype=(dtypes.uint8, dtypes.uint16, dtypes.bool)).cast(dtypes.ints, name="x"), lambda x:
-   x.ins(X86Ops.MOVZX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
-  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD) if x.max_numel() == 1 else None),
-  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x:
-   x.ins(X86Ops.MOVSX) if x.max_numel() == 1 and x.src[0].dtype.itemsize < x.dtype.itemsize else None),
-  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else None),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBW)),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBD)),
-  (UPat(dtype=(dtypes.uint8, dtypes.bool)).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXBQ)),
-  (UPat(dtype=dtypes.uint16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWD)),
-  (UPat(dtype=dtypes.uint16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXWQ)),
-  (UPat(dtype=dtypes.uint32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVZXDQ)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int16s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBW)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBD)),
-  (UPat(dtype=dtypes.int8).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXBQ)),
-  (UPat(dtype=dtypes.int16).cast(dtypes.int32s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWD)),
-  (UPat(dtype=dtypes.int16).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXWQ)),
-  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.VPMOVSXDQ)),
+   x.ins(X86Ops.MOVZX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
+  (UPat(dtype=dtypes.int32).cast(dtypes.int64s, name="x"), lambda x: x.ins(X86Ops.MOVSXD)),
+  (UPat(dtype=dtypes.sints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOVSX) if x.src[0].dtype.itemsize < x.dtype.itemsize else None),
+  (UPat(dtype=dtypes.ints).cast(dtypes.ints, name="x"), lambda x: x.ins(X86Ops.MOV)),
   # bitcasts between scalar floats and ints
   (UPat.var("y", dtypes.float16).bitcast(dtypes.int16s).named("x"), lambda y,x: x.ins(X86Ops.VPEXTRW, src=(y, imm(dtypes.uint8, 0)))),
   (UPat(dtype=dtypes.int16s).bitcast(dtypes.float16).named("x"), lambda x: vpins(x, x.src)),
@@ -518,42 +439,41 @@ isel_matcher = PatternMatcher([
   # index on a buffer (or the stack pointer) computes an address, addresses are 64bit values
   (UPat((Ops.INDEX, Ops.SHRINK), name="x"), lambda x: lea(x) if not _is_vec_xmm(x.src[0]) else None),
   # TODO: fuse stores, very few cases -- store cmp becomes setcc, store gep int becomes vpextr, store bitcast to int becomes vmovd/q
-  # copy, load, store
-  # NOTE: copy here violates the spec, it only happens post register allocation when a reg to reg move needs to be inserted
-  (UPat(Ops.COPY, dtypes.floats, name="x"), lambda x: x.ins(_xmm_sz(x))),
-  (UPat(Ops.COPY, dtypes.ints+(dtypes.bool,), name="x"), lambda x: x.ins(X86Ops.MOV) if x.max_numel() == 1 else x.ins(_xmm_sz(x))),
+  # load, store
   (UPat(Ops.LOAD, dtypes.floats, src=(UPat(name="a"),), name="x"), lambda x,a:
-   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
+   x.ins(X86Ops.VPINSRW, src=(undef(),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
    x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat(Ops.LOAD, dtypes.ints+(dtypes.bool,), src=(UPat(name="a"),), name="x"), lambda x,a:
-   x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else
-   x.ins(X86Ops.VPINSRW, src=(def_reg(x.dtype, x.tag),) + fold_address(a) + (imm(dtypes.uint8, 0),)) if x.max_numel() * x.dtype.itemsize == 2 else
-   x.ins(_xmm_sz(x), src=fold_address(a))),
+   x.ins(X86Ops.MOV, src=fold_address(a)) if x.max_numel() == 1 else x.ins(_xmm_sz(x), src=fold_address(a))),
   (UPat.var("a").store(UPat.var("b", dtypes.floats), name="x"), lambda a,b,x:
    x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() * b.dtype.itemsize == 2 else
    x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,))),
   (UPat.var("a").store(UPat.var("b", dtypes.ints+(dtypes.bool,)), name="x"), lambda a,b,x:
-   x.ins(X86Ops.VPEXTRW, src=fold_address(a) + (b, imm(dtypes.uint8, 0))) if b.max_numel() > 1 and b.max_numel() * b.dtype.itemsize == 2 else
    x.ins(_xmm_sz_m(b), src=fold_address(a) + (b,)) if b.max_numel() > 1 else
    x.ins(X86Ops.MOVm, src=fold_address(a) + (b,)) if (i:=to_imm(b)) is None else x.ins(X86Ops.MOVi, src=fold_address(a) + (i,))),
   # allocate virtual registers
-  (UPat((Ops.INS, Ops.BUFFER), name="x"), alloc_vregs),
+  (UPat((Ops.INS, Ops.BUFFER, Ops.RANGE), name="x"), alloc_vregs),
 ])
 
 # ***** pre register allocation *****
-# this handles flag clobbers. Unfortunately x86 doesn't have a good way to store/restore the flag register (then regalloc would handle it)
-# so we rematerialize. This is different from rematerialization you might want to do in regalloc because it is not optional,
-# regalloc shouldn't rematerialize if a src of the instruction is dead, but here you need to as there's no fallback load from stack
-def flag_rematerialize(ctx:PreRegAllocContext, x:UOp):
-  flag_def = x if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags else x.src[-1] if x.arg[0] in X86GroupOp.ReadFlags else None
-  if flag_def is None: return None
-  if ctx.lock is not None and ctx.lock is not flag_def: ctx.clobbered.add(ctx.lock)
-  ctx.lock = flag_def
-  if flag_def not in ctx.clobbered: return None
-  ctx.clobbered.remove(flag_def)
-  return (x, [flag_def, x])
+# the flags belong to the last instruction that wrote them. x86 has no good way to store/restore them (then regalloc would
+# handle it), so a consumer that no longer owns its compare re-emits it. Unlike a regalloc rematerialization this is not
+# optional, there is no fallback load from stack
+def flag_rematerialize(ctx:X86LinearContext, x:UOp):
+  if x.op in (Ops.RANGE, Ops.END) or x.arg[0] in X86GroupOp.WriteFlags: ctx.lock = x
+  elif x.arg[0] in X86GroupOp.ReadFlags and ctx.lock is not (flag_def:=x.src[-1]):
+    ctx.lock = flag_def
+    return (x, [flag_def, x])
+  return None
+
+# TODO: dont use rewrite
+def alloc_buffer(ctx:X86LinearContext, x:UOp):
+  nx = isel_matcher.rewrite(stack_pointer.index(UOp.cconst(ctx.stack_size, dtypes.uint32), tag=x.tag))
+  ctx.stack_size += x.max_numel() * x.dtype.itemsize
+  return nx, [nx]
 
 pre_regalloc_matcher = PatternMatcher([
+  (UPat(Ops.BUFFER, name="x"), alloc_buffer),
   (UPat((Ops.INS, Ops.RANGE, Ops.END), name="x"), flag_rematerialize),
 ])
 
@@ -584,8 +504,14 @@ def lower_loop(ctx, x:UOp) -> tuple[UOp, list[UOp]]:
 
 # final rewrite to match the isa spec
 post_regalloc_matcher = PatternMatcher([
+  # the frame is allocated after the stack pointer define at the top of the program and freed before RET
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [x, x.ins(X86Ops.SUBi, src=(imm(dtypes.int32, ctx.stack_size),))])
+    if ctx.stack_size and x.arg[0] is X86Ops.DEFINE and rdef(x) == RSP else None),
+  (UPat(Ops.INS, name="x"), lambda ctx,x: (x, [stack_pointer.ins(X86Ops.ADDi, src=(imm(dtypes.int32, ctx.stack_size),)), x])
+    if ctx.stack_size and x.arg[0] is X86Ops.RET else None),
   # rewrite FRAME_INDEX to IMM now that the stack size is known
-  (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=UOp.cconst(ctx.stack_size + x.tag, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
+  (UPat(Ops.INS, src=(UPat.cvar("disp").cast(),), name="x"), lambda ctx,disp,x:
+    (nx:=UOp.cconst(ctx.stack_size + disp.val, x.dtype), [nx]) if x.arg[0] is X86Ops.FRAME_INDEX else None),
   # expand the cmp here so we can preserve rng src edge to get label from ctx
   (UPat(Ops.INS, name="x"), lambda ctx,x: lower_loop(ctx, x) if x.arg[0] is X86Ops.LOOP_CMP else None),
   # rewrite RANGE to ACC = 0 -> LABEL -> JUMP if ACC >= loop bound
@@ -594,7 +520,7 @@ post_regalloc_matcher = PatternMatcher([
   (UPat(Ops.END, name="x"), lower_end),
   # rewrite two address instructions to two address form, if reused src wasn't coalesced insert a move
   (UPat(Ops.INS, name="x"), lambda ctx,x: (nx:=x.replace(src=x.src[1:]),
-   [ctx.ren.copy(x.src[0], greg(x)), nx] if greg(x) != greg(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
+   [ctx.ren.copy(x.src[0], rdef(x)), nx] if rdef(x) != rdef(x.src[0]) else [nx]) if x.arg[0] in X86GroupOp.TwoAddress else None),
 ])
 
 # ***** X86 instruction encoding *****
@@ -604,9 +530,9 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
               vvvv_uop:UOp|None=None, imm_uop:UOp|None=None) -> bytes:
     nonlocal reg, opc
     # get the encoding values of the different fields
-    reg = cast(int, cast(Register, greg(reg_uop)).index if reg_uop is not None else reg)
-    rm = cast(Register, greg(rm_uop)).index
-    idx = cast(Register, greg(idx_uop)).index if idx_uop is not None and greg(idx_uop) is not None else 4
+    reg = cast(int, cast(Register, rdef(reg_uop)).index if reg_uop is not None else reg)
+    rm = cast(Register, rdef(rm_uop)).index
+    idx = cast(Register, rdef(idx_uop)).index if idx_uop is not None and rdef(idx_uop) is not None else 4
     # for a memory operand the rm size is the element size from the address, otherwise it's the size of the value in the register
     rm_sz = sz_uop.src[0].val if sz_uop is not None else rm_uop.dtype.itemsize
     reg_sz = reg_uop.dtype.itemsize if reg_uop is not None else 0
@@ -618,7 +544,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # r extends reg field, x extends index field, b extends rm or base field
     r, _x, b = reg >> 3, idx >> 3, rm >> 3
     if sel: # VEX bytes
-      vvvv = cast(Register, greg(vvvv_uop)).index if vvvv_uop is not None else 0
+      vvvv = (vd.index if isinstance(vd := rdef(vvvv_uop), Register) else reg) if vvvv_uop is not None else 0
       if sel == 1 and _x == b == we == 0: inst += bytes([0xC5, (~r & 0b1) << 7 | (~vvvv & 0b1111) << 3 | pp])
       else: inst += bytes([0xC4, (~r & 0b1) << 7 | (~_x & 0b1) << 6 | (~b & 0b1) << 5 | sel, we << 7 | (~vvvv & 0b1111) << 3 | pp])
     else: # optional PREFIX and REX bytes
@@ -663,7 +589,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
     # IMM byte
     if imm_uop is not None:
       if imm_uop.op is Ops.CAST: inst += struct.pack(unwrap(imm_uop.dtype.fmt), imm_uop.src[0].val)
-      elif isinstance(greg(imm_uop), Register): inst += bytes([(greg(imm_uop).index & 0b1111) << 4 | 0b0000])
+      elif isinstance(rdef(imm_uop), Register): inst += bytes([(rdef(imm_uop).index & 0b1111) << 4 | 0b0000])
     return inst
 
   # get the encoding structure of the uop
@@ -696,7 +622,7 @@ def encode(x:UOp, opc:int, reg:int|None=None, pp:int=0, sel:int=0, we:int=0) ->
 encodings = {
   # moves
   X86Ops.MOVABS: lambda x:
-   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | greg(x).index >> 3, 0xB8 + (greg(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
+   bytes([0b0100 << 4 | 0b1 << 3 | 0b00 << 2 | rdef(x).index >> 3, 0xB8 + (rdef(x).index & 0b111)]) + struct.pack(x.dtype.fmt, x.src[0].src[0].val),
   X86Ops.MOV: lambda x: encode(x, 0x8B), X86Ops.MOVi: lambda x: encode(x, 0xC7, reg=0),
   X86Ops.MOVm: lambda x: encode(x, 0x89), X86Ops.LEA: lambda x: encode(x, 0x8D),
   X86Ops.VMOVSS: lambda x: encode(x, 0x10, pp=2, sel=1), X86Ops.VMOVSSm: lambda x: encode(x, 0x11, pp=2, sel=1),
@@ -707,20 +633,11 @@ encodings = {
   # casts
   X86Ops.MOVZX: lambda x: encode(x, 0x0FB7),
   X86Ops.MOVSX: lambda x: encode(x, 0x0FBF), X86Ops.MOVSXD: lambda x: encode(x, 0x63),
-  X86Ops.VPMOVZXBW: lambda x: encode(x, 0x30, pp=1, sel=2), X86Ops.VPMOVZXBD: lambda x: encode(x, 0x31, pp=1, sel=2),
-  X86Ops.VPMOVZXBQ: lambda x: encode(x, 0x32, pp=1, sel=2), X86Ops.VPMOVZXWD: lambda x: encode(x, 0x33, pp=1, sel=2),
-  X86Ops.VPMOVZXWQ: lambda x: encode(x, 0x34, pp=1, sel=2), X86Ops.VPMOVZXDQ: lambda x: encode(x, 0x35, pp=1, sel=2),
-  X86Ops.VPMOVSXBW: lambda x: encode(x, 0x20, pp=1, sel=2), X86Ops.VPMOVSXBD: lambda x: encode(x, 0x21, pp=1, sel=2),
-  X86Ops.VPMOVSXBQ: lambda x: encode(x, 0x22, pp=1, sel=2), X86Ops.VPMOVSXWD: lambda x: encode(x, 0x23, pp=1, sel=2),
-  X86Ops.VPMOVSXWQ: lambda x: encode(x, 0x24, pp=1, sel=2), X86Ops.VPMOVSXDQ: lambda x: encode(x, 0x25, pp=1, sel=2),
   X86Ops.VCVTSS2SD: lambda x: encode(x, 0x5A, pp=2, sel=1), X86Ops.VCVTSD2SS: lambda x: encode(x, 0x5A, pp=3, sel=1),
   X86Ops.VCVTPH2PS: lambda x: encode(x, 0x13, pp=1, sel=2), X86Ops.VCVTPS2PH: lambda x: encode(x, 0x1D, pp=1, sel=3),
-  X86Ops.VCVTDQ2PS: lambda x: encode(x, 0x5B, pp=0, sel=1), X86Ops.VCVTDQ2PD: lambda x: encode(x, 0xE6, pp=2, sel=1),
-  X86Ops.VCVTPS2PD: lambda x: encode(x, 0x5A, pp=0, sel=1), X86Ops.VCVTPD2PS: lambda x: encode(x, 0x5A, pp=1, sel=1),
-  X86Ops.VCVTTPS2DQ: lambda x: encode(x, 0x5B, pp=2, sel=1), X86Ops.VCVTTPD2DQ: lambda x: encode(x, 0xE6, pp=1, sel=1),
-  # the int src is the 2nd src (the rm field), if it was folded into a memory operand its width is the element size of the address
-  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
-  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=(x.src[4].src[0].val if len(x.src) > 4 else x.src[1].dtype.itemsize) == 8),
+  # the int src is the 2nd src (the rm field), its width picks the 32 or 64 bit form
+  X86Ops.VCVTSI2SS: lambda x: encode(x, 0x2A, pp=2, sel=1, we=x.src[1].dtype.itemsize == 8),
+  X86Ops.VCVTSI2SD: lambda x: encode(x, 0x2A, pp=3, sel=1, we=x.src[1].dtype.itemsize == 8),
   X86Ops.VCVTTSS2SI: lambda x: encode(x, 0x2C, pp=2, sel=1, we=x.dtype.itemsize == 8),
   X86Ops.VCVTTSD2SI: lambda x: encode(x, 0x2C, pp=3, sel=1, we=x.dtype.itemsize == 8),
   # int division
@@ -738,46 +655,25 @@ encodings = {
   X86Ops.IMUL: lambda x: encode(x, 0x0FAF), X86Ops.IMULi: lambda x: encode(x, 0x69),
   X86Ops.SETB: lambda x: encode(x, 0x0F92, reg=0), X86Ops.SETL: lambda x: encode(x, 0x0F9C, reg=0),
   X86Ops.SETE: lambda x: encode(x, 0x0F94, reg=0), X86Ops.SETNE: lambda x: encode(x, 0x0F95, reg=0),
-  # packed bitwise NOTE: only bitwise and packed
-  X86Ops.VPAND: lambda x: encode(x, 0xDB, pp=1, sel=1), X86Ops.VPXOR: lambda x: encode(x, 0xEF, pp=1, sel=1),
-  X86Ops.VPOR: lambda x: encode(x, 0xEB, pp=1, sel=1),
   # unary
-  X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTPS: lambda x: encode(x, 0x51, pp=0, sel=1),
-  X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1), X86Ops.VSQRTPD: lambda x: encode(x, 0x51, pp=1, sel=1),
-  X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDPS: lambda x: encode(x, 0x08, pp=1, sel=3),
-  X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3), X86Ops.VROUNDPD: lambda x: encode(x, 0x09, pp=1, sel=3),
-  # packed int binary
-  X86Ops.VPSLLVD: lambda x: encode(x, 0x47, pp=1, sel=2), X86Ops.VPSLLVQ: lambda x: encode(x, 0x47, pp=1, sel=2, we=1),
-  X86Ops.VPSRLVD: lambda x: encode(x, 0x45, pp=1, sel=2), X86Ops.VPSRLVQ: lambda x: encode(x, 0x45, pp=1, sel=2, we=1),
-  X86Ops.VPMULLW: lambda x: encode(x, 0xD5, pp=1, sel=1), X86Ops.VPMULLD: lambda x: encode(x, 0x40, pp=1, sel=2),
-  X86Ops.VPADDB: lambda x: encode(x, 0xFC, pp=1, sel=1), X86Ops.VPADDW: lambda x: encode(x, 0xFD, pp=1, sel=1),
-  X86Ops.VPADDD: lambda x: encode(x, 0xFE, pp=1, sel=1), X86Ops.VPADDQ: lambda x: encode(x, 0xD4, pp=1, sel=1),
-  X86Ops.VPSUBB: lambda x: encode(x, 0xF8, pp=1, sel=1), X86Ops.VPSUBW: lambda x: encode(x, 0xF9, pp=1, sel=1),
-  X86Ops.VPSUBD: lambda x: encode(x, 0xFA, pp=1, sel=1), X86Ops.VPSUBQ: lambda x: encode(x, 0xFB, pp=1, sel=1),
-  X86Ops.VPSRAVD: lambda x: encode(x, 0x46, pp=1, sel=2),
-  # scalar / packed float binary
-  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDPS: lambda x: encode(x, 0x58, pp=0, sel=1),
-  X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1), X86Ops.VADDPD: lambda x: encode(x, 0x58, pp=1, sel=1),
-  X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBPS: lambda x: encode(x, 0x5C, pp=0, sel=1),
-  X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1), X86Ops.VSUBPD: lambda x: encode(x, 0x5C, pp=1, sel=1),
-  X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULPS: lambda x: encode(x, 0x59, pp=0, sel=1),
-  X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1), X86Ops.VMULPD: lambda x: encode(x, 0x59, pp=1, sel=1),
-  X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVPS: lambda x: encode(x, 0x5E, pp=0, sel=1),
-  X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1), X86Ops.VDIVPD: lambda x: encode(x, 0x5E, pp=1, sel=1),
-  X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPPS: lambda x: encode(x, 0xC2, pp=0, sel=1),
-  X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1), X86Ops.VCMPPD: lambda x: encode(x, 0xC2, pp=1, sel=1),
+  X86Ops.VSQRTSS: lambda x: encode(x, 0x51, pp=2, sel=1), X86Ops.VSQRTSD: lambda x: encode(x, 0x51, pp=3, sel=1),
+  X86Ops.VROUNDSS: lambda x: encode(x, 0x0A, pp=1, sel=3), X86Ops.VROUNDSD: lambda x: encode(x, 0x0B, pp=1, sel=3),
+  # float binary
+  X86Ops.VADDSS: lambda x: encode(x, 0x58, pp=2, sel=1), X86Ops.VADDSD: lambda x: encode(x, 0x58, pp=3, sel=1),
+  X86Ops.VSUBSS: lambda x: encode(x, 0x5C, pp=2, sel=1), X86Ops.VSUBSD: lambda x: encode(x, 0x5C, pp=3, sel=1),
+  X86Ops.VMULSS: lambda x: encode(x, 0x59, pp=2, sel=1), X86Ops.VMULSD: lambda x: encode(x, 0x59, pp=3, sel=1),
+  X86Ops.VDIVSS: lambda x: encode(x, 0x5E, pp=2, sel=1), X86Ops.VDIVSD: lambda x: encode(x, 0x5E, pp=3, sel=1),
+  X86Ops.VCMPSS: lambda x: encode(x, 0xC2, pp=2, sel=1), X86Ops.VCMPSD: lambda x: encode(x, 0xC2, pp=3, sel=1),
   # ternary
   X86Ops.CMOVB: lambda x: encode(x, 0x0F42), X86Ops.CMOVL: lambda x: encode(x, 0x0F4C),
   X86Ops.CMOVE: lambda x: encode(x, 0x0F44), X86Ops.CMOVNE: lambda x: encode(x, 0x0F45),
   X86Ops.VBLENDVPS: lambda x: encode(x, 0x4A, pp=1, sel=3), X86Ops.VBLENDVPD: lambda x: encode(x, 0x4B, pp=1, sel=3),
   # shuffles
   X86Ops.VPSRLDQ: lambda x: encode(x, 0x73, reg=3, pp=1, sel=1),
-  X86Ops.VPINSRB: lambda x: encode(x, 0x20, pp=1, sel=3), X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1),
-  X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3), X86Ops.VPINSRQ: lambda x: encode(x, 0x22, pp=1, sel=3, we=1),
+  X86Ops.VPINSRW: lambda x: encode(x, 0xC4, pp=1, sel=1), X86Ops.VPINSRD: lambda x: encode(x, 0x22, pp=1, sel=3),
   X86Ops.VINSERTPS: lambda x: encode(x, 0x21, pp=1, sel=3),
   # extract
-  X86Ops.VPEXTRB: lambda x: encode(x, 0x14, pp=1, sel=3), X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3),
-  X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3), X86Ops.VPEXTRQ: lambda x: encode(x, 0x16, pp=1, sel=3, we=1),
+  X86Ops.VPEXTRW: lambda x: encode(x, 0x15, pp=1, sel=3), X86Ops.VPEXTRD: lambda x: encode(x, 0x16, pp=1, sel=3),
   # jumps are encoded with a placeholder which gets patched later once the real offset is known
   X86Ops.JE: lambda x: bytes([0x0F, 0x84]) + int(0).to_bytes(4),
   X86Ops.JNE: lambda x: bytes([0x0F, 0x85]) + int(0).to_bytes(4),
@@ -788,6 +684,16 @@ encodings = {
   X86Ops.RET: lambda x: bytes([0xC3]),
 }
 
+class X86LinearContext(LinearContext):
+  def __init__(self, ren:X86Renderer):
+    super().__init__(ren)
+    self.lock: UOp|None = None
+  def assign_spill_slot(self, r:Register, u:UOp) -> int:
+    sz = r.cons[0].size
+    offset = self.stack_size + (sz - self.stack_size % sz) %sz
+    self.stack_size = offset + sz
+    return offset
+
 class X86Renderer(ISARenderer):
   device = "CPU"
   has_local = False
@@ -798,39 +704,36 @@ class X86Renderer(ISARenderer):
   pre_regalloc_matcher = pre_regalloc_matcher
   post_regalloc_matcher = post_regalloc_matcher
   code_for_op = {x: lambda: None for x in (Ops.SQRT, Ops.AND, Ops.OR, Ops.SHL, Ops.SHR, Ops.NEG, Ops.SUB, Ops.FDIV, Ops.CMPLT, Ops.CMPEQ)}
+  linear_ctx_type = X86LinearContext
   def __init__(self, target:Target):
     if target.arch.split(",")[0] != "x86_64": raise RuntimeError(f"X86Renderer only supports x86_64, got {target.arch}")
     super().__init__(target)
     from tinygrad.runtime.support.compiler_cpu import X86Compiler
     self.compiler = X86Compiler()
   def is_two_address(self, x:UOp) -> bool: return x.op is Ops.INS and x.arg[0] in X86GroupOp.TwoAddress
-  def stack_pointer(self) -> UOp: return def_reg(dtypes.uint64, RSP)
-  # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
-  def copy(self, x:UOp, reg:Register):
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
-    ret = isel_matcher.rewrite(UOp(Ops.COPY, src=(x,), tag=reg))
-    assert ret is not None, f"failed to copy {x}"
-    return ret
+  def copy(self, x:UOp, reg:Register) -> UOp: return x.ins(X86Ops.MOV, src=(x,), tag=reg)
 
-  def spill(self, disp:UOp, x:UOp) -> UOp:
-    if x.op is Ops.BUFFER: x = x.replace(arg=replace(x.arg, dtype=dtypes.uint64))
+  def spill(self, spill_slot:int, x:UOp) -> UOp:
     is_xmm = isinstance(x.tag, tuple) and x.tag[0].cons[0].size == 16
     op = X86Ops.VMOVUPSm if is_xmm else X86Ops.MOVm
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
+    disp = UOp.cconst(spill_slot, dtypes.int32)
+    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)) + (x,), arg=(op, dtypes.void), tag=x.tag)
 
-  def fill(self, disp:UOp, x:UOp, reg:Register) -> UOp:
+  # the value of a BUFFER is its address, it moves through registers and the stack as a 64bit int
+  def fill(self, spill_slot:int, x:UOp, reg:Register) -> UOp:
     is_xmm = reg.cons[0].size == 16
     dt = dtypes.uint64 if x.op is Ops.BUFFER else x.dtype
-    return UOp(Ops.INS, src=fold_address(self.stack_pointer().index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
+    disp = UOp.cconst(spill_slot, dtypes.int32)
+    return UOp(Ops.INS, src=fold_address(stack_pointer.index(disp)), arg=(X86Ops.VMOVUPS if is_xmm else X86Ops.MOV, dt), tag=(reg,))
 
   def asm_str(self, uops:list[UOp], function_name:str) -> str:
     def _format_op(x:UOp) -> str: return f"    {(o[7:-1] if (o:=str(x.arg[0]))[-1] in ('i', 'm') else o[7:]).lower():7s}"
     def _format_operands(x:UOp) -> str:
       def _format(src:tuple[UOp, ...]) -> list[str]:
         return [str(s.src[0].val) if s.op is Ops.CAST else reg_strs[o].get(s.dtype.itemsize, o) if \
-                (o:=str(greg(s))) in reg_strs else o for s in src if greg(s) is not None]
+                (o:=str(rdef(s))) in reg_strs else o for s in src if rdef(s) is not None]
       def _mem_adress(base:UOp, idx:UOp, disp:UOp, sz:UOp) -> list[str]:
-        return [f"[{greg(base)}" + (f" + {greg(idx)}*{sz.src[0].val}" if greg(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
+        return [f"[{rdef(base)}" + (f" + {rdef(idx)}*{sz.src[0].val}" if rdef(idx) else "") + (f" + {d}" if (d:=disp.src[0].val) else "") + "]"]
 
       if len(x.src) > 4 and x.arg[0] in X86GroupOp.WriteMem: ret = _mem_adress(*x.src[:4]) + _format(x.src[4:])
       elif len(x.src) > 3 and x.arg[0] in X86GroupOp.Rm1st: ret = _format((x,)) + _mem_adress(*x.src[:4]) + _format(x.src[4:])

tinygrad/runtime/autogen/am/__init__.py

@@ -22,7 +22,7 @@ reg_files = {
 reg_patterns = {
   "gc": ["GCVM", "GCMC_VM", "CP_(HQD|MQD|MEC|ME_CNTL|PERFMON|RB_WPTR_POLL_CNTL|INT_CNTL|STAT|PFP_PRGRM|ME_PRGRM|COHER_START)", "COMPUTE_",
          "(SQ|GL2C|TCC)_PERFCOUNTER", "SQ_THREAD_TRACE", "SPI_(CONFIG_CNTL|COMPUTE_QUEUE_RESET)", "GRBM", "SH_MEM", "RLC", "TCP", "GB_ADDR_CONFIG",
-         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[0-367]"],
+         "SDMA[01]_(WATCHDOG_CNTL|UTCL1_(CNTL|PAGE)|MCU_CNTL|F32_CNTL|CNTL|QUEUE0_|RLC_CGCG_CTRL)", "SCRATCH_REG[0-35-7]"],
   "mmhub": ["MMVM", "MMMC_VM", "MM_ATC_L2_MISC_CG"],
   "nbio": (nbio:=["BIF_BX_PF[01]_GPU_HDP_FLUSH", "BIF_BX_PF0_RSMU", "BIF_BX0_(REMAP_HDP_MEM_FLUSH_CNTL|BIF_DOORBELL_INT_CNTL|PCIE_INDEX2|PCIE_DATA2)",
                   "BIFC_(DOORBELL_ACCESS_EN_PF|GFX_INT_MONITOR_MASK)", "XCC_DOORBELL_FENCE", "DOORBELL0_CTRL_ENTRY", "GDC_S2A0_S2A_DOORBELL_ENTRY",

tinygrad/runtime/autogen/am/regs.py

@@ -514,6 +514,7 @@ gc_9_4_3 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regCP_COHER_START_DELAY': (8315, 1, {'start_delay_count': (0, 5)}),
@@ -1799,6 +1800,7 @@ gc_11_0_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -3378,6 +3380,7 @@ gc_11_0_3 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -4807,6 +4810,7 @@ gc_11_5_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),
@@ -6068,6 +6072,7 @@ gc_12_0_0 = {
   'regSCRATCH_REG1': (8257, 1, {'scratch_reg1': (0, 31)}),
   'regSCRATCH_REG2': (8258, 1, {'scratch_reg2': (0, 31)}),
   'regSCRATCH_REG3': (8259, 1, {'scratch_reg3': (0, 31)}),
+  'regSCRATCH_REG5': (8261, 1, {'scratch_reg5': (0, 31)}),
   'regSCRATCH_REG6': (8262, 1, {'scratch_reg6': (0, 31)}),
   'regSCRATCH_REG7': (8263, 1, {'scratch_reg7': (0, 31)}),
   'regRLC_GPM_PERF_COUNT_0': (8512, 1, {'feature_sel': (0, 3), 'se_index': (4, 7), 'sa_index': (8, 11), 'wgp_index': (12, 15), 'event_sel': (16, 17), 'unused': (18, 19), 'enable': (20, 20), 'reserved': (21, 31)}),

tinygrad/runtime/autogen/amd/rdna3/str_pcode.py

@@ -47,7 +47,7 @@ PCODE = {
   DSOp.DS_MIN_RTN_F32: 'tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nMEM[ADDR].f32 = src < tmp ? src : tmp;\nRETURN_DATA.f32 = tmp',
   DSOp.DS_MAX_RTN_F32: 'tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nMEM[ADDR].f32 = src > tmp ? src : tmp;\nRETURN_DATA.f32 = tmp',
   DSOp.DS_WRAP_RTN_B32: 'tmp = MEM[ADDR].u32;\nMEM[ADDR].u32 = tmp >= DATA.u32 ? tmp - DATA.u32 : tmp + DATA2.u32;\nRETURN_DATA = tmp',
-  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}',
+  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset >= 0xc000) {\n// rotate\nrotate = offset[9:5];\nmask = offset[4:0];\nif (offset[10]) {\nrotate = -rotate;\n}\nfor (i = 0; i < 64; i++) {\nj = (i & mask) | ((i + rotate) & ~mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset[15]) {\n// full data sharing within 4 consecutive threads\nfor (i = 0; i < 64; i+=4) {\nthread_out[i+0] = thread_valid[i+offset[1:0]]?thread_in[i+offset[1:0]]:0;\nthread_out[i+1] = thread_valid[i+offset[3:2]]?thread_in[i+offset[3:2]]:0;\nthread_out[i+2] = thread_valid[i+offset[5:4]]?thread_in[i+offset[5:4]]:0;\nthread_out[i+3] = thread_valid[i+offset[7:6]]?thread_in[i+offset[7:6]]:0;\n}\n} else { // offset[15] == 0\n// limited data sharing within 32 consecutive threads\nxor_mask = offset[14:10];\nor_mask = offset[9:5];\nand_mask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = (((i & 0x1f) & and_mask) | or_mask) ^ xor_mask;\nj |= (i & 0x20); // which group of 32\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n}',
   DSOp.DS_LOAD_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET.u32].b32',
   DSOp.DS_LOAD_2ADDR_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET0.u32 * 4U].b32;\nRETURN_DATA[63 : 32] = MEM[ADDR + OFFSET1.u32 * 4U].b32',
   DSOp.DS_LOAD_2ADDR_STRIDE64_B32: 'RETURN_DATA[31 : 0] = MEM[ADDR + OFFSET0.u32 * 256U].b32;\nRETURN_DATA[63 : 32] = MEM[ADDR + OFFSET1.u32 * 256U].b32',

tinygrad/runtime/autogen/amd/rdna4/str_pcode.py

@@ -44,7 +44,7 @@ PCODE = {
   DSOp.DS_CMPSTORE_RTN_B32: 'addr = CalcDsAddr(vgpr_a.b32, offset.b32);\ntmp = MEM[addr].b32;\nsrc = DATA.b32;\ncmp = DATA2.b32;\nMEM[addr].b32 = tmp == cmp ? src : tmp;\nRETURN_DATA.b32 = tmp',
   DSOp.DS_MIN_NUM_RTN_F32: "tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nif (isNAN(64'F(src.f32)) && isNAN(64'F(tmp.f32))) then\nMEM[ADDR].f32 = 32'F(cvtToQuietNAN(64'F(src.f32)))\nelsif isNAN(64'F(src.f32)) then\nMEM[ADDR].f32 = tmp.f32\nelsif isNAN(64'F(tmp.f32)) then\nMEM[ADDR].f32 = src.f32\nelsif ((src.f32 < tmp.f32) || ((abs(src.f32) == 0.0F) && (abs(tmp.f32) == 0.0F) && sign(src.f32) &&\n!sign(tmp.f32))) then\n// NOTE: -0<+0 is TRUE in this comparison\nMEM[ADDR].f32 = src.f32\nelse\nMEM[ADDR].f32 = tmp.f32\nendif;\nRETURN_DATA.f32 = tmp",
   DSOp.DS_MAX_NUM_RTN_F32: "tmp = MEM[ADDR].f32;\nsrc = DATA.f32;\nif (isNAN(64'F(src.f32)) && isNAN(64'F(tmp.f32))) then\nMEM[ADDR].f32 = 32'F(cvtToQuietNAN(64'F(src.f32)))\nelsif isNAN(64'F(src.f32)) then\nMEM[ADDR].f32 = tmp.f32\nelsif isNAN(64'F(tmp.f32)) then\nMEM[ADDR].f32 = src.f32\nelsif ((src.f32 > tmp.f32) || ((abs(src.f32) == 0.0F) && (abs(tmp.f32) == 0.0F) && !sign(src.f32) &&\nsign(tmp.f32))) then\n// NOTE: +0>-0 is TRUE in this comparison\nMEM[ADDR].f32 = src.f32\nelse\nMEM[ADDR].f32 = tmp.f32\nendif;\nRETURN_DATA.f32 = tmp",
-  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}',
+  DSOp.DS_SWIZZLE_B32: 'offset = offset1:offset0;\nif (offset >= 0xe000) {\n// FFT decomposition\nmask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = reverse_bits(i & 0x1f);\nj = (j >> count_ones(mask));\nj |= (i & mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset >= 0xc000) {\n// rotate\nrotate = offset[9:5];\nmask = offset[4:0];\nif (offset[10]) {\nrotate = -rotate;\n}\nfor (i = 0; i < 64; i++) {\nj = (i & mask) | ((i + rotate) & ~mask);\nj |= i & 0x20;\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n} elsif (offset[15]) {\n// full data sharing within 4 consecutive threads\nfor (i = 0; i < 64; i+=4) {\nthread_out[i+0] = thread_valid[i+offset[1:0]]?thread_in[i+offset[1:0]]:0;\nthread_out[i+1] = thread_valid[i+offset[3:2]]?thread_in[i+offset[3:2]]:0;\nthread_out[i+2] = thread_valid[i+offset[5:4]]?thread_in[i+offset[5:4]]:0;\nthread_out[i+3] = thread_valid[i+offset[7:6]]?thread_in[i+offset[7:6]]:0;\n}\n} else { // offset[15] == 0\n// limited data sharing within 32 consecutive threads\nxor_mask = offset[14:10];\nor_mask = offset[9:5];\nand_mask = offset[4:0];\nfor (i = 0; i < 64; i++) {\nj = (((i & 0x1f) & and_mask) | or_mask) ^ xor_mask;\nj |= (i & 0x20); // which group of 32\nthread_out[i] = thread_valid[j] ? thread_in[j] : 0;\n}\n}',
   DSOp.DS_LOAD_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET.u32].b32',
   DSOp.DS_LOAD_2ADDR_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET0.u32 * 4U].b32;\naddr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[63 : 32] = MEM[addr + OFFSET1.u32 * 4U].b32',
   DSOp.DS_LOAD_2ADDR_STRIDE64_B32: 'addr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[31 : 0] = MEM[addr + OFFSET0.u32 * 256U].b32;\naddr = CalcDsAddr(vgpr_a.b32, 0x0);\nRETURN_DATA[63 : 32] = MEM[addr + OFFSET1.u32 * 256U].b32',

tinygrad/runtime/ops_cpu.py

@@ -87,7 +87,7 @@ class CPUAllocator(HCQAllocator['CPUDevice']):
     ctypes.memmove(int(dest.va_addr), from_mv(src), len(src))
   def _copyout(self, dest:memoryview, src:HCQBuffer):
     self.dev.synchronize()
-    ctypes.memmove(from_mv(dest), int(src.va_addr), len(dest))
+    dest[:] = self._as_buffer(src)[:len(dest)]
   def _do_map(self, buf:HCQBuffer):
     if buf.view is None or not isinstance(buf.view, MMIOInterface): raise RuntimeError("Cannot map buffer without view to cpu")
     return HCQBuffer(buf.view.addr, buf.size, view=buf.view, owner=buf.owner)

tinygrad/runtime/ops_nv.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
-import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, weakref
+import os, ctypes, contextlib, re, functools, mmap, struct, array, sys, itertools
 assert sys.platform != 'win32'
-from typing import cast
+from typing import Any
 from dataclasses import dataclass
-from tinygrad.runtime.support.hcq import HCQCompiled, HCQAllocator, HCQBuffer, HWQueue, CLikeArgsState, HCQProgram, HCQSignal, BumpAllocator
-from tinygrad.runtime.support.hcq import MMIOInterface, FileIOInterface, hcq_filter_visible_devices, hcq_profile
-from tinygrad.uop.ops import sint
-from tinygrad.device import Compiled, BufferSpec, TinyELF
-from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, hi32, lo32, PROFILE, ContextVar, VIZ, ProfileEvent
+from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HWQueue, encode_submit, patch, to_name, unwrap_view
+from tinygrad.runtime.support.hcq import HCQBuffer, MMIOInterface, FileIOInterface, BumpAllocator, hcq_filter_visible_devices
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher
+from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
+from tinygrad.device import Buffer, BufferSpec, Compiled, Device, TinyELF
+from tinygrad.dtype import dtypes, DType
+from tinygrad.helpers import getenv, mv_address, round_up, data64, data64_le, prod, OSX, PROFILE, ContextVar, VIZ
+from tinygrad.helpers import ProfileEvent
 from tinygrad.renderer.ptx import PTXRenderer
 from tinygrad.renderer.cstyle import CUDARenderer, NVCCRenderer
 from tinygrad.runtime.autogen import nv_570, nv_580, nv_610, mesa
@@ -24,10 +27,7 @@ PMA = ContextVar("PMA", abs(VIZ.value)>=2)
 @dataclass(frozen=True)
 class ProfilePMAEvent(ProfileEvent): device:str; kern:str; blob:bytes; exec_tag:int; profile_key:bytes|None=None # noqa: E702
 
-class NVSignal(HCQSignal):
-  def _sleep(self, time_spent_since_last_sleep_ms:int):
-    # Reasonable to sleep for long workloads (which take more than 200ms) and only timeline signals.
-    if time_spent_since_last_sleep_ms > 200 and self.owner is not None: self.owner.iface.sleep(200)
+def hilo(addr:UOp) -> tuple[UOp, UOp]: return (addr >> 32).cast(dtypes.uint32), addr.cast(dtypes.uint32)
 
 def get_error_str(status): return f"{status}: {nv_gpu.nv_status_codes.get(status, 'Unknown error')}"
 
@@ -41,10 +41,13 @@ def nv_iowr(fd:FileIOInterface, nr, args, cmd=None):
   ret = fd.ioctl(cmd or ((3 << 30) | (ctypes.sizeof(args) & 0x1FFF) << 16 | (ord('F') & 0xFF) << 8 | (nr & 0xFF)), args)
   if ret != 0: raise RuntimeError(f"ioctl returned {ret}")
 
+def nvm(subc:int, mthd:int, *vals, typ=2) -> list:
+  return [(typ << 28) | (sum(v.dtype.itemsize // 4 if isinstance(v, UOp) else 1 for v in vals) << 16) | (subc << 13) | (mthd >> 2), *vals]
+
 class QMD:
   fields: dict[str, dict[str, tuple[int, int]]] = {}
 
-  def __init__(self, dev:NVDevice, view:MMIOInterface|None=None, **kwargs):
+  def __init__(self, dev:NVDevice, blob:bytearray|None=None):
     self.ver, self.sz = (5, 0x60) if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else (3, 0x40)
 
     # Init fields from module
@@ -52,272 +55,225 @@ class QMD:
       QMD.fields[pref] = {**{name[len(pref)+1:]: dt for name,dt in nv_gpu.__dict__.items() if name.startswith(pref) and isinstance(dt, tuple)},
         **{name[len(pref)+1:]+f"_{i}": dt(i) for name,dt in nv_gpu.__dict__.items() for i in range(8) if name.startswith(pref) and callable(dt)}}
 
-    self.mv, self.pref = (memoryview(bytearray(self.sz * 4)) if view is None else view), pref
-    if kwargs: self.write(**kwargs)
+    self.mv, self.pref = (bytearray(self.sz * 4) if blob is None else blob), pref
+    self.patches:dict[int, UOp] = {}
 
-  def _rw_bits(self, hi:int, lo:int, value:int|None=None):
-    mask = ((1 << (width:=hi - lo + 1)) - 1) << (lo % 8)
-    num = int.from_bytes(self.mv[lo//8:hi//8+1], "little")
+  def read(self, k:str) -> int:
+    hi, lo = QMD.fields[self.pref][k.upper()]
+    return (int.from_bytes(self.mv[lo//8:hi//8+1], "little") >> (lo % 8)) & ((1 << (hi - lo + 1)) - 1)
 
-    if value is None: return (num & mask) >> (lo % 8)
+  def write(self, **kwargs:int|UOp):
+    for k, v in kwargs.items():
+      hi, lo = QMD.fields[self.pref][k.upper()]
+      if isinstance(v, UOp):
+        assert lo % 8 == 0, f"{k} is not byte aligned"
+        self.patches[lo // 8] = v.ccast(next(t for t in (dtypes.uint64, dtypes.uint32, dtypes.uint16, dtypes.uint8) if t.itemsize * 8 <= hi - lo + 1))
+      else:
+        if v >> (hi - lo + 1): raise ValueError(f"{k}={v:#x} does not fit")
+        mask, num = ((1 << (hi - lo + 1)) - 1) << (lo % 8), int.from_bytes(self.mv[lo//8:hi//8+1], "little")
+        self.mv[lo//8:hi//8+1] = ((num & ~mask) | (v << (lo % 8))).to_bytes(hi//8 - lo//8 + 1, "little")
 
-    if value >= (1 << width): raise ValueError(f"{value:#x} does not fit.")
-    self.mv[lo//8:hi//8+1] = int((num & ~mask) | ((value << (lo % 8)) & mask)).to_bytes((hi//8 - lo//8 + 1), "little")
+  def set_addr(self, name:str, addr:UOp, sfx:str=""): self.write(**{f"{name}_lower{sfx}": addr, f"{name}_upper{sfx}": addr >> 32})
+  def set_constant_buf_addr(self, i:int, addr:UOp):
+    self.set_addr("constant_buffer_addr", addr >> (6 if self.ver >= 4 else 0), f"_shifted6_{i}" if self.ver >= 4 else f"_{i}")
+  def set_program_addr(self, addr:UOp):
+    self.set_addr("program_address", addr >> (4 if self.ver >= 4 else 0), "_shifted4" if self.ver >= 4 else "")
+    self.set_addr("program_prefetch_addr", addr >> 8, "_shifted")
+  def set_release(self, addr:UOp, payload:UOp, timestamp:bool=False) -> bool:
+    if (i:=next((i for i in range(2) if not self.read(f"release{i}_enable")), None)) is None: return False
+    self.set_addr(f"release_semaphore{i}_addr" if self.ver >= 4 else f"release{i}_address", addr)
+    self.set_addr(f"release_semaphore{i}_payload" if self.ver >= 4 else f"release{i}_payload", payload)
+    self.write(**{f"release{i}_enable": 1, f"release_structure_size_{i}" if self.ver >= 4 else f"release{i}_structure_size": 0 if timestamp else 2},
+               **({} if self.ver >= 4 else {f"release{i}_payload64b": 1}))
+    return True
+  @property
+  def grid(self) -> tuple[str, ...]:
+    return ("grid_width", "grid_height", "grid_depth") if self.ver >= 4 else ("cta_raster_width", "cta_raster_height", "cta_raster_depth")
 
-  def write(self, **kwargs):
-    for k,val in kwargs.items(): self._rw_bits(*QMD.fields[self.pref][k.upper()], value=val) # type: ignore [misc]
+# *****************
+# queues
 
-  def read(self, k, val=0): return self._rw_bits(*QMD.fields[self.pref][k.upper()])
+class NVQueue(HWQueue):
+  dev:NVDevice
 
-  def field_offset(self, k): return QMD.fields[self.pref][k.upper()][1] // 8
+  def nvm(self, subc:int, mthd:int, *vals, typ=2): self.q(*nvm(subc, mthd, *vals, typ=typ))
 
-  def set_constant_buf_addr(self, i, addr):
-    if self.ver < 4: self.write(**{f'constant_buffer_addr_upper_{i}':hi32(addr), f'constant_buffer_addr_lower_{i}':lo32(addr)})
-    else: self.write(**{f'constant_buffer_addr_upper_shifted6_{i}':hi32(addr >> 6), f'constant_buffer_addr_lower_shifted6_{i}':lo32(addr >> 6)})
+  def sem(self, addr:UOp, value:UOp, **flags:str):
+    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, addr, value.ccast(dtypes.uint64), nv_flags("NVC56F_SEM_EXECUTE", payload_size="64bit", **flags))
 
-class NVCommandQueue(HWQueue[HCQSignal, 'NVDevice', 'NVProgram', 'NVArgsState']):
-  def __init__(self):
-    self.active_qmd = None
-    super().__init__()
+  def wait(self, signal:UOp, value:UOp): self.sem(signal.getaddr(self.devs), value, operation="acq_circ_geq")
+  def signal(self, signal:UOp, value:UOp): self.release(signal, value)
+  def timestamp(self, signal:UOp): self.release(signal, UOp.const(0, dtypes.uint64), timestamp=True)
+  def release(self, signal:UOp, value:UOp, timestamp:bool=False):
+    self.sem(signal.getaddr(self.devs), value, operation="release", release_wfi="en", release_timestamp="en" if timestamp else "dis")
+    if not timestamp: self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
 
-  def __del__(self):
-    if self.binded_device is not None: self.binded_device.allocator.free(self.hw_page, self.hw_page.size, BufferSpec(cpu_access=True, nolru=True))
+  def submit(self, cmdbuf:UOp) -> UOp:
+    fifo, ib, off = self.dev.fifos[self.queue], *unwrap_view(cmdbuf)
 
-  def nvm(self, subchannel, mthd, *args, typ=2): self.q((typ << 28) | (len(args) << 16) | (subchannel << 13) | (mthd >> 2), *args)
+    ring, gpput, doorbell, put, gpentry = [UOp.placeholder((sz,), dt, device=self.devs, volatile=True, tag=to_name(nm, self.queue))
+      for nm, dt, sz in (("ring", dtypes.uint64, fifo.entries), ("gpput", dtypes.uint32, 1), ("doorbell", dtypes.uint32, 1),
+                         ("put_value", dtypes.uint64, 1), ("gpentry", dtypes.uint64, 1))]
+    gpentry = patch(gpentry, [(0, ib.getaddr(self.devs) + UOp.const(off | (cmdbuf.max_numel() // 4 << 42) | (1 << 41), dtypes.uint64))])
 
-  def setup(self, compute_class=None, copy_class=None, local_mem_window=None, shared_mem_window=None, local_mem=None, local_mem_tpc_bytes=None):
-    if compute_class: self.nvm(1, nv_gpu.NVC6C0_SET_OBJECT, compute_class)
-    if copy_class: self.nvm(4, nv_gpu.NVC6C0_SET_OBJECT, copy_class)
-    if local_mem_window: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_WINDOW_A, *data64(local_mem_window))
-    if shared_mem_window: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_SHARED_MEMORY_WINDOW_A, *data64(shared_mem_window))
-    if local_mem: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(local_mem))
-    if local_mem_tpc_bytes: self.nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(local_mem_tpc_bytes), 0xff)
-    return self
+    p = put.index(0).load()
+    written = UOp.barrier(ring.after(cmdbuf).index((p % fifo.entries).cast(dtypes.int)).store(gpentry.index(0).load()), put.index(0).store(p + 1))
+    queued = UOp.barrier(gpput.after(written).index(0).store(((p + 1) % fifo.entries).cast(dtypes.uint32)))
+    return doorbell.after(queued).index(0).store(UOp.const(fifo.token, dtypes.uint32))
 
-  def wait(self, signal:HCQSignal, value:sint=0):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit"))
-    self.active_qmd = None
-    return self
+class NVComputeQueue(NVQueue):
+  def __init__(self, ctx, submit):
+    super().__init__(ctx, submit)
 
-  def timestamp(self, signal:HCQSignal): return self.signal(signal, 0)
+    progs = [nv_build_program(self.dev, u.src[0], self.devs)[0] for u in self.lin.src if u.op is Ops.CALL]
+    self.qmd_sz = round_up(QMD(self.dev).sz * 4, 256)
+    self.stride = self.qmd_sz + max([p.kernargs_size for p in progs], default=0)
+    self.qmd_buf = UOp.placeholder((len(progs) * self.stride,), dtypes.uint8, device=self.devs, tag=to_name("qmd", self.queue))
+    self.qmds:list[QMD] = []
+    self.prev_qmd:QMD|None = None # the launch the next one chains onto
 
-  def bind(self, dev:NVDevice):
-    self.binded_device = dev
-    self.hw_page = dev.allocator.alloc(len(self._q) * 4, BufferSpec(cpu_access=True, nolru=True))
-    hw_view = self.hw_page.cpu_view().view(fmt='I')
-    for i, value in enumerate(self._q): hw_view[i] = value
+  def wait(self, signal:UOp, value:UOp):
+    self.prev_qmd = None
+    super().wait(signal, value)
 
-    # From now on, the queue is on the device for faster submission.
-    self._q = hw_view
+  def release(self, signal:UOp, value:UOp, timestamp:bool=False):
+    if self.prev_qmd is None or not self.prev_qmd.set_release(signal.getaddr(self.devs), value, timestamp):
+      self.prev_qmd = None
+      super().release(signal, value, timestamp)
 
-  def _submit_to_gpfifo(self, dev:NVDevice, gpfifo:GPFifo):
-    if dev == self.binded_device: cmdq_addr = self.hw_page.va_addr
-    else:
-      cmdq_addr = dev.cmdq_allocator.alloc(len(self._q) * 4, 16)
-      cmdq_wptr = (cmdq_addr - dev.cmdq_page.va_addr) // 4
-      dev.cmdq[cmdq_wptr : cmdq_wptr + len(self._q)] = array.array('I', self._q)
+  def submit(self, cmdbuf:UOp) -> UOp:
+    if self.qmds:
+      patches = [(i * self.stride + off, w) for i, q in enumerate(self.qmds) for off, w in q.patches.items()]
+      cmdbuf = cmdbuf.after(patch(self.qmd_buf, patches, b"".join(q.mv for q in self.qmds)))
+    return super().submit(cmdbuf)
 
-    gpfifo.ring[gpfifo.put_value % gpfifo.entries_count] = (cmdq_addr//4 << 2) | (len(self._q) << 42) | (1 << 41)
-    gpfifo.gpput[0] = (gpfifo.put_value + 1) % gpfifo.entries_count
-
-    System.memory_barrier()
-    dev.gpu_mmio[0x90 // 4] = gpfifo.token
-    gpfifo.put_value += 1
-
-class NVComputeQueue(NVCommandQueue):
   def memory_barrier(self):
+    self.prev_qmd = None
     self.nvm(1, nv_gpu.NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI,
              nv_flags("NVC6C0_INVALIDATE_SHADER_CACHES_NO_WFI", instruction="true", global_data="true", constant="true"))
-    self.active_qmd:QMD|None = None
-    return self
 
-  def exec(self, prg:NVProgram, args_state:NVArgsState, global_size:tuple[sint, ...], local_size:tuple[sint, ...]):
-    self.bind_args_state(args_state)
+  def exec(self, call:UOp, prg:UOp):
+    data, lib = nv_build_program(self.dev, prg, self.devs)
+    global_size, local_size = prg.arg.global_size, prg.arg.local_size
+    if prod(local_size) > 1024 or data.max_threads < prod(local_size):
+      raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {data.max_threads=}")
+    if any(g > mx for g,mx in zip(global_size, [2147483647, 65535, 65535]) if isinstance(g, int)) or \
+       any(l > mx for l,mx in zip(local_size, [1024, 1024, 64])):
+      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
 
-    qmd_buf = args_state.buf.offset(round_up(prg.constbufs[0][1], 1 << 8))
-    qmd_buf.cpu_view().view(size=prg.qmd.mv.nbytes, fmt='B')[:] = prg.qmd.mv
-    assert qmd_buf.va_addr < (1 << 40), f"large qmd addr {qmd_buf.va_addr:x}"
+    qmd_addr = self.qmd_buf.getaddr(self.devs) + UOp.const(len(self.qmds) * self.stride, dtypes.uint64)
+    qmd = QMD(self.dev, data.qmd.mv.ljust(self.stride, b"\0")) # the program's template, in a slot of its own
+    qmd.write(**dict(zip(qmd.grid, global_size)), **{f"cta_thread_dimension{j}": l for j, l in enumerate(local_size)})
+    qmd.set_program_addr(lib.getaddr(self.devs) + data.prog_off)
+    for j, (off, _) in data.constbufs.items():
+      qmd.set_constant_buf_addr(j, qmd_addr + UOp.const(self.qmd_sz, dtypes.uint64) if j == 0 else lib.getaddr(self.devs) + off)
+    bufs, vals = [get_call_arg_uops(call)[j] for j in prg.arg.globals], get_call_var_uops(call, prg)
+    qmd.mv[self.qmd_sz:(at:=self.qmd_sz + len(data.cbuf_0) * 4)] = array.array('I', data.cbuf_0).tobytes() # constant buffer 0: the driver params
+    qmd.patches |= {at + j * 8: b.getaddr(self.devs) for j, b in enumerate(bufs)} | {at + o: v.ccast(dt) for v, (o, dt) in zip(vals, data.vars)}
 
-    qmd = QMD(dev=prg.dev, view=qmd_buf.cpu_view()) # Save qmd for later update
+    if self.prev_qmd is None:
+      if self.dev.pma_enabled: self.nvm(1, nv_gpu.NVC6C0_PM_TRIGGER, 0)
+      self.nvm(1, nv_gpu.NVC6C0_SEND_PCAS_A, (qmd_addr >> 8).cast(dtypes.uint32))
+      self.nvm(1, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B_PCAS_ACTION_PREFETCH_SCHEDULE)
+    else: self.prev_qmd.write(dependent_qmd0_pointer=qmd_addr >> 8, dependent_qmd0_action=1, dependent_qmd0_prefetch=1, dependent_qmd0_enable=1)
+    self.qmds.append(qmd)
+    self.prev_qmd = qmd
 
-    self.bind_sints_to_mem(*global_size, mem=qmd_buf.cpu_view(), fmt='I', offset=qmd.field_offset('cta_raster_width' if qmd.ver<4 else 'grid_width'))
-    self.bind_sints_to_mem(*(local_size[:2]), mem=qmd_buf.cpu_view(), fmt='H', offset=qmd.field_offset('cta_thread_dimension0'))
-    self.bind_sints_to_mem(local_size[2], mem=qmd_buf.cpu_view(), fmt='B', offset=qmd.field_offset('cta_thread_dimension2'))
-    qmd.set_constant_buf_addr(0, args_state.buf.va_addr)
-
-    if self.active_qmd is None:
-      if prg.dev.pma_enabled: self.nvm(1, nv_gpu.NVC6C0_PM_TRIGGER, 0)
-      self.nvm(1, nv_gpu.NVC6C0_SEND_PCAS_A, qmd_buf.va_addr >> 8)
-      self.nvm(1, nv_gpu.NVC6C0_SEND_SIGNALING_PCAS2_B, 9)
-    else:
-      self.active_qmd.write(dependent_qmd0_pointer=qmd_buf.va_addr >> 8, dependent_qmd0_action=1, dependent_qmd0_prefetch=1, dependent_qmd0_enable=1)
-
-    self.active_qmd, self.active_qmd_buf = qmd, qmd_buf
-    return self
-
-  def signal(self, signal:HCQSignal, value:sint=0):
-    if self.active_qmd is not None:
-      for i in range(2):
-        if self.active_qmd.read(f'release{i}_enable') == 0:
-          self.active_qmd.write(**{f'release{i}_enable': 1})
-
-          addr_off = self.active_qmd.field_offset(f'release{i}_address_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_addr_lower')
-          self.bind_sints_to_mem(signal.value_addr & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=addr_off)
-          self.bind_sints_to_mem(signal.value_addr >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', mask=0xf, offset=addr_off+4)
-
-          val_off = self.active_qmd.field_offset(f'release{i}_payload_lower' if self.active_qmd.ver<4 else f'release_semaphore{i}_payload_lower')
-          self.bind_sints_to_mem(value & 0xffffffff, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off)
-          self.bind_sints_to_mem(value >> 32, mem=self.active_qmd_buf.cpu_view(), fmt='I', offset=val_off+4)
-          return self
-
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(signal.value_addr), *data64_le(value),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit", release_timestamp="en"))
-    self.nvm(0, nv_gpu.NVC56F_NON_STALL_INTERRUPT, 0x0)
-    self.active_qmd = None
-    return self
-
-  def write(self, b:HCQBuffer, val:sint, b64:bool=False):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(b.va_addr), *data64_le(val),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit" if b64 else "32bit"))
-    self.active_qmd = None
-    return self
-
-  def poll_bit(self, b:HCQBuffer, val:sint, mask:int):
-    self.nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(b.va_addr), *data64_le((~mask & 0xFFFFFFFF) if val == 0 else val),
-             nv_flags("NVC56F_SEM_EXECUTE", operation="acq_nor" if val == 0 else "acq_and", payload_size="32bit"))
-    self.active_qmd = None
-    return self
-
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.compute_gpfifo)
-
-class NVCopyQueue(NVCommandQueue):
-  def __init__(self, queue_idx=0):
-    self.queue_idx = queue_idx
-    super().__init__()
-
-  def copy(self, dest:HCQBuffer, src:HCQBuffer, copy_size:int):
-    for off in range(0, copy_size, step:=(1 << 31)):
-      self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *data64(src.va_addr+off), *data64(dest.va_addr+off))
-      self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(copy_size-off, step))
+class NVCopyQueue(NVQueue):
+  def copy(self, call:UOp):
+    dest, src = (a.getaddr(self.devs) for a in call.src[1:3])
+    for off in range(0, sz:=call.src[2].max_numel() * call.src[2].dtype.itemsize, step:=(1 << 31)):
+      self.nvm(4, nv_gpu.NVC6B5_OFFSET_IN_UPPER, *hilo(src + UOp.const(off, dtypes.uint64)), *hilo(dest + UOp.const(off, dtypes.uint64)))
+      self.nvm(4, nv_gpu.NVC6B5_LINE_LENGTH_IN, min(sz - off, step))
       self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA,
                nv_flags("NVC6B5_LAUNCH_DMA", data_transfer_type="non_pipelined", src_memory_layout="pitch", dst_memory_layout="pitch"))
-    return self
 
-  def signal(self, signal:HCQSignal, value:sint=0):
-    self.nvm(4, nv_gpu.NVC6B5_SET_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type="release_four_word_semaphore"))
-    return self
+  def semaphore(self, addr:UOp, value:UOp, typ:str): # a one word release writes just the payload, a four word one the timestamp after it
+    self.nvm(4, nv_gpu.NVC6B5_SET_SEMAPHORE_A, *hilo(addr), value.ccast(dtypes.uint32))
+    self.nvm(4, nv_gpu.NVC6B5_LAUNCH_DMA, nv_flags("NVC6B5_LAUNCH_DMA", flush_enable="true", semaphore_type=f"release_{typ}_word_semaphore"))
+  def timestamp(self, signal:UOp): self.semaphore(signal.getaddr(self.devs), UOp.const(0, dtypes.uint32), "four")
+  def signal(self, signal:UOp, value:UOp): self.semaphore(signal.getaddr(self.devs), value, "one")
 
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.dma_gpfifo)
+# *****************
+# programs
 
-class NVVideoQueue(NVCommandQueue):
-  def decode_hevc_chunk(self, pic_desc:HCQBuffer, in_buf:HCQBuffer, out_buf:HCQBuffer, out_buf_pos:int, hist_bufs:list[HCQBuffer], hist_pos:list[int],
-                        chroma_off:int, coloc_buf:HCQBuffer, filter_buf:HCQBuffer, intra_top_off:int, intra_unk_off:int|None, status_buf:HCQBuffer):
-    self.nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
-    self.nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
-             err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
-    self.nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, pic_desc.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, in_buf.va_addr >> 8)
-    for pos, buf in zip(hist_pos + [out_buf_pos], hist_bufs + [out_buf]):
-      self.nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
-      self.nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, coloc_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, status_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, pic_desc.offset(0x200).va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, filter_buf.va_addr >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_buf.va_addr + intra_top_off) >> 8)
-    if intra_unk_off is not None: self.nvm(4, 0x4dc, (filter_buf.va_addr + intra_unk_off) >> 8)
-    self.nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
-    return self
-
-  def signal(self, signal:HCQSignal, value:sint=0):
-    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_A, *data64(signal.value_addr), value)
-    self.nvm(4, nv_gpu.NVC9B0_SEMAPHORE_D, nv_flags("NVC9B0_SEMAPHORE_D", structure_size="four", payload_size="64bit"))
-    return self
-
-  def _submit(self, dev:NVDevice): self._submit_to_gpfifo(dev, dev.vid_gpfifo)
-
-class NVArgsState(CLikeArgsState):
-  def __init__(self, buf:HCQBuffer, prg:NVProgram, bufs:tuple[HCQBuffer, ...], vals:tuple[int, ...]=()):
-    if (is_mock:=isinstance(prg.dev.iface, MOCKIface)): prg.cbuf_0[80:82] = [len(bufs), len(vals)]
-    super().__init__(buf, prg, bufs, vals=() if is_mock else vals, prefix=prg.cbuf_0 or None)
-    # mock expects all vars to be 64 bit
-    if is_mock and vals: self.bind_sints_to_buf(*vals, buf=self.buf, fmt='q', offset=len(prg.cbuf_0)*4 + len(bufs)*8)
-
-class NVProgram(HCQProgram['NVDevice']):
+class NVProgramData:
   def __init__(self, dev:NVDevice, obj:TinyELF):
-    self.dev, self.name, self.lib = dev, obj.name, obj.lib
-    self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[va_addr, size]]
+    name, signature, mock = obj.name, obj.signature, isinstance(dev.iface, MOCKIface)
+    self.constbufs: dict[int, tuple[int, int]] = {0: (0, 0x160)} # dict[constbuf index, tuple[offset in the image, size]]
+    self.relocs: list[tuple[int, int, DType, int]] = [] # (byte offset in the image, symbol offset, width, shift) of the program's address
+    self.prog_off, self.cbuf_0, sections, relocs = 0, [], list[Any](), list[Any]()
+    image:bytes = obj.lib
 
     if (NAK:=isinstance(dev.renderer, NAKRenderer)):
-      image, self.cbuf_0 = memoryview(bytearray(obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):])), []
-      self.regs_usage, self.shmem_usage, self.lcmem_usage = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
-    elif isinstance(dev.iface, MOCKIface): image, sections, relocs = memoryview(bytearray(obj.lib) + b'\x00' * (4 - len(obj.lib)%4)).cast("I"), [], [] # type: ignore
-    else: image, sections, relocs = elf_loader(self.lib, force_section_align=128)
-    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
-    self.lib_gpu = self.dev.allocator.alloc(round_up((prog_sz:=image.nbytes), 0x1000) + 0x1000, buf_spec:=BufferSpec(nolru=True))
-    prog_addr = self.lib_gpu.va_addr
+      image = obj.lib[ctypes.sizeof(info:=mesa.struct_nak_shader_info.from_buffer_copy(obj.lib)):]
+      regs, shmem, lcmem = info.num_gprs, round_up(info.cs.smem_size, 128), round_up(info.slm_size, 16)
+    elif mock: image = obj.lib.ljust(round_up(len(obj.lib), 4), b'\x00') # for MOCKGPU the lib is PTX code, not an elf
+    else:
+      img, sections, relocs = elf_loader(obj.lib, force_section_align=128)
+      image = bytes(img)
+    prog_sz = len(image)
+
     if not NAK:
       # For MOCKGPU, the lib is PTX code, so some values are emulated.
-      self.regs_usage, self.shmem_usage, self.lcmem_usage, cbuf0_size = 0, 0x400, 0x240, 0x160 if isinstance(dev.iface, MOCKIface) else 0
-      for sh in sections: # pylint: disable=possibly-used-before-assignment
-        if sh.name == f".nv.shared.{self.name}": self.shmem_usage = round_up(0x400 + sh.header.sh_size, 128)
-        if sh.name == f".text.{self.name}": prog_addr, prog_sz = self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size
-        elif m:=re.match(r'\.nv\.constant(\d+)', sh.name):
-          self.constbufs[int(m.group(1))] = (self.lib_gpu.va_addr+sh.header.sh_addr, sh.header.sh_size)
+      regs, shmem, lcmem, cbuf0_size = 0, 0x400, 0x240, 0x160 if mock else 0
+      for sh in sections:
+        if sh.name == f".nv.shared.{name}": shmem = round_up(0x400 + sh.header.sh_size, 128)
+        if sh.name == f".text.{name}": self.prog_off, prog_sz = sh.header.sh_addr, sh.header.sh_size
+        elif m:=re.match(r'\.nv\.constant(\d+)', sh.name): self.constbufs[int(m.group(1))] = (sh.header.sh_addr, sh.header.sh_size)
         elif sh.name.startswith(".nv.info"):
           for typ, param, data in self._parse_elf_info(sh):
             if sh.name == f".nv.info.{obj.name}" and param == 0xa: cbuf0_size = struct.unpack_from("IH", data)[1] # EIATTR_PARAM_CBANK
-            elif sh.name == ".nv.info" and param == 0x12: self.lcmem_usage = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
-            elif sh.name == ".nv.info" and param == 0x2f: self.regs_usage = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
+            elif sh.name == ".nv.info" and param == 0x12: lcmem = struct.unpack_from("II", data)[1] + 0x240 # EIATTR_MIN_STACK_SIZE
+            elif sh.name == ".nv.info" and param == 0x2f: regs = struct.unpack_from("II", data)[1] # EIATTR_REGCOUNT
 
-      # Apply relocs
-      for apply_image_offset, rel_sym_offset, typ, _ in relocs: # pylint: disable=possibly-used-before-assignment
-        # These types are CUDA-specific, applying them here
-        if typ == 2: image[apply_image_offset:apply_image_offset+8] = struct.pack('<Q', self.lib_gpu.va_addr + rel_sym_offset) # R_CUDA_64
-        elif typ == 0x38: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) & 0xffffffff)
-        elif typ == 0x39: image[apply_image_offset+4:apply_image_offset+8] = struct.pack('<I', (self.lib_gpu.va_addr + rel_sym_offset) >> 32)
+      # These reloc types are CUDA-specific: they all want the program's own address, which is only known once the linear links.
+      for apply_image_offset, rel_sym_offset, typ, _ in relocs:
+        if typ == 2: self.relocs.append((apply_image_offset, rel_sym_offset, dtypes.uint64, 0)) # R_CUDA_64
+        elif typ == 0x38: self.relocs.append((apply_image_offset + 4, rel_sym_offset, dtypes.uint32, 0))
+        elif typ == 0x39: self.relocs.append((apply_image_offset + 4, rel_sym_offset, dtypes.uint32, 32))
         else: raise RuntimeError(f"unknown NV reloc {typ}")
 
       # Minimum cbuf_0 size for driver params: Blackwell needs index 223 (224 entries), older GPUs need index 11 (12 entries)
       min_cbuf0_entries = 224 if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A else 12
       self.cbuf_0 = [0] * max(cbuf0_size // 4, min_cbuf0_entries)
 
+    # the arguments follow the driver params in constant buffer 0: the buffers as 64 bit addresses, then the vars packed by their width
+    nbufs = sum(name is None for name, *_ in signature)
+    self.vars = list(TinyELF.iter_sig(signature[nbufs:], nbufs * 8))
+    if mock: # mockgpu reads the arg counts out of cbuf0 and wants every var 64 bit
+      self.cbuf_0[80:82], self.vars = [nbufs, len(self.vars)], [(nbufs * 8 + i * 8, dtypes.uint64) for i in range(len(self.vars))]
+
+    # NOTE: Ensure at least 4KB of space after the program to mitigate prefetch memory faults.
+    self.image = image.ljust(round_up(len(image), 0x1000) + 0x1000, b'\x00')
+    # constant buffer 0 holds the driver params and every argument after them, and starts 256 aligned like all constant buffers
+    self.kernargs_size = round_up(max(self.constbufs[0][1], len(self.cbuf_0) * 4 + len(signature) * 8), 256)
+
     # Ensure device has enough local memory to run the program
-    self.dev._ensure_has_local_memory(self.lcmem_usage)
-    self.dev.allocator._copyin(self.lib_gpu, image)
-    self.dev.synchronize()
+    dev._ensure_has_local_memory(lcmem)
 
     if dev.iface.compute_class >= nv_gpu.BLACKWELL_COMPUTE_A:
-      if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window)], 0xfffdc0
-      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'program_address_upper_shifted4':hi32(prog_addr>>4),
-        'program_address_lower_shifted4':lo32(prog_addr>>4), 'register_count':self.regs_usage, 'shared_memory_size_shifted7':self.shmem_usage>>7,
-        f'shader_local_memory_{"low" if NAK else "high"}_size_shifted4': self.dev.slm_per_thread>>4}
+      if not NAK: self.cbuf_0[188:192], self.cbuf_0[223] = [*data64_le(dev.shared_mem_window), *data64_le(dev.local_mem_window)], 0xfffdc0
+      qmd = {'qmd_major_version':5, 'qmd_type':nv_gpu.NVCEC0_QMDV05_00_QMD_TYPE_GRID_CTA, 'register_count':regs,
+        'shared_memory_size_shifted7':shmem>>7, f'shader_local_memory_{"low" if NAK else "high"}_size_shifted4':dev.slm_per_thread>>4}
     else:
-      if not NAK: self.cbuf_0[6:12] = [*data64_le(self.dev.shared_mem_window), *data64_le(self.dev.local_mem_window), *data64_le(0xfffdc0)]
-      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'program_address_upper':hi32(prog_addr), 'program_address_lower':lo32(prog_addr),
-        'shared_memory_size':self.shmem_usage, 'register_count_v':self.regs_usage,
-        f'shader_local_memory_{"low" if NAK else "high"}_size':self.dev.slm_per_thread}
+      if not NAK: self.cbuf_0[6:12] = [*data64_le(dev.shared_mem_window), *data64_le(dev.local_mem_window), *data64_le(0xfffdc0)]
+      qmd = {'qmd_major_version':3, 'sm_global_caching_enable':1, 'shared_memory_size':shmem, 'register_count_v':regs,
+        f'shader_local_memory_{"low" if NAK else "high"}_size':dev.slm_per_thread}
 
-    smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= self.shmem_usage) // 4096 + 1
+    smem_cfg = min(shmem_conf * 1024 for shmem_conf in [32, 64, 100] if shmem_conf * 1024 >= shmem) // 4096 + 1
 
-    self.qmd:QMD = QMD(dev, **qmd, qmd_group_id=0x3f, invalidate_texture_header_cache=1, invalidate_texture_sampler_cache=1,
+    # the program and constant buffer addresses are patched into a copy of this at exec, everything else is the same for every launch
+    self.qmd = QMD(dev)
+    self.qmd.write(**qmd, qmd_group_id=0x3f, invalidate_texture_header_cache=1, invalidate_texture_sampler_cache=1,
       invalidate_texture_data_cache=1, invalidate_shader_data_cache=1, api_visible_call_limit=1, sampler_index=1, barrier_count=1,
       cwd_membar_type=nv_gpu.NVC6C0_QMDV03_00_CWD_MEMBAR_TYPE_L1_SYSMEMBAR, constant_buffer_invalidate_0=1, min_sm_config_shared_mem_size=smem_cfg,
       target_sm_config_shared_mem_size=smem_cfg, max_sm_config_shared_mem_size=0x1a, program_prefetch_size=min(prog_sz>>8, 0x1ff),
-      sass_version=dev.sass_version, program_prefetch_addr_upper_shifted=prog_addr>>40, program_prefetch_addr_lower_shifted=prog_addr>>8)
-
-    for i,(addr,sz) in self.constbufs.items():
-      self.qmd.set_constant_buf_addr(i, addr)
-      self.qmd.write(**{f'constant_buffer_size_shifted4_{i}': sz, f'constant_buffer_valid_{i}': 1})
+      sass_version=dev.sass_version)
+    for i,(_,sz) in self.constbufs.items(): self.qmd.write(**{f'constant_buffer_size_shifted4_{i}': sz, f'constant_buffer_valid_{i}': 1})
 
     # Registers allocation granularity per warp is 256, warp allocation granularity is 4. Register file size is 65536.
-    self.max_threads = ((65536 // round_up(max(1, self.regs_usage) * 32, 256)) // 4) * 4 * 32
-
-    # NV's kernargs is constbuffer, then arguments to the kernel follows. Kernargs also appends QMD at the end of the kernel.
-    super().__init__(NVArgsState, self.dev, obj, kernargs_alloc_size=round_up(self.constbufs[0][1], 1 << 8) + (8 << 8))
-    weakref.finalize(self, self._fini, self.dev, self.lib_gpu, buf_spec)
+    self.max_threads = ((65536 // round_up(max(1, regs) * 32, 256)) // 4) * 4 * 32
 
   def _parse_elf_info(self, sh, start_off=0):
     while start_off < sh.header.sh_size:
@@ -325,18 +281,14 @@ class NVProgram(HCQProgram['NVDevice']):
       yield typ, param, sh.content[start_off+4:start_off+sz+4] if typ == 0x4 else sz
       start_off += (sz if typ == 0x4 else 0) + 4
 
-  def __call__(self, *bufs, global_size:tuple[int,int,int]=(1,1,1), local_size:tuple[int,int,int]=(1,1,1), vals:tuple[int|None, ...]=(),
-               wait=False, timeout:int|None=None):
-    if prod(local_size) > 1024 or self.max_threads < prod(local_size) or self.lcmem_usage > self.dev.slm_per_thread:
-      raise RuntimeError(f"Too many resources requested for launch, {prod(local_size)=}, {self.max_threads=}")
-    if any(cur > mx for cur,mx in zip(global_size, [2147483647, 65535, 65535])) or any(cur > mx for cur,mx in zip(local_size, [1024, 1024, 64])):
-      raise RuntimeError(f"Invalid global/local dims {global_size=}, {local_size=}")
-    res = super().__call__(*bufs, global_size=global_size, local_size=local_size, vals=vals, wait=wait, timeout=timeout)
-    if self.dev.pma_enabled:
-      self.dev.synchronize()
-      if pma_blob:=self.dev._prof_readback():
-        Compiled.profile_events += [ProfilePMAEvent(self.dev.device, self.name, pma_blob, self.dev.prof_exec_counter, self.profile_key)]
-    return res
+_nv_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[NVProgramData, UOp]] = {}
+def nv_build_program(dev:NVDevice, prg:UOp, devs:tuple[str, ...]) -> tuple[NVProgramData, UOp]:
+  if (cached:=_nv_program_cache.get(key:=(prg.src[3].arg, devs))) is None:
+    data = NVProgramData(dev, prg.to_elf())
+    buf = UOp.placeholder((len(data.image),), dtypes.uint8, next(UOp.unique_num), device=devs).rtag("program")
+    rows = [(off, ((buf.getaddr(devs) + sym) >> sh).ccast(dt)) for off, sym, dt, sh in data.relocs]
+    cached = _nv_program_cache[key] = (data, patch(buf, rows, data.image))
+  return cached
 
 class NVAllocator(HCQAllocator['NVDevice']):
   def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
@@ -350,22 +302,31 @@ class NVAllocator(HCQAllocator['NVDevice']):
     assert all(h.va_addr % 0x100 == 0 for h in hist + [bufin, bufout, desc_buf]), "all buffers must be 0x100 aligned"
 
     h, w = ((2 * shape[0]) // 3 if shape[0] % 3 == 0 else (2 * shape[0] - 1) // 3), shape[1]
-    self.dev._ensure_has_vid_hw(w, h)
+    dev, chroma_off = self.dev, round_up(w, 64) * round_up(h, 64)
+    dev._ensure_has_vid_hw(w, h)
 
-    q = NVVideoQueue().wait(self.dev.timeline_signal, self.dev.timeline_value - 1)
-    with hcq_profile(self.dev, queue=q, desc="HEVC Decode", enabled=PROFILE, dev_suff="NVDEC"):
-      q.decode_hevc_chunk(desc_buf, bufin, bufout, frame_pos, hist, [(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)],
-                          round_up(w, 64)*round_up(h, 64), self.dev.vid_coloc_buf, self.dev.vid_filter_buf, self.dev.intra_top_off,
-                          self.dev.intra_unk_off, self.dev.vid_stat_buf)
-    q.signal(self.dev.timeline_signal, self.dev.next_timeline()).submit(self.dev)
+    cmds = nvm(4, nv_gpu.NVC9B0_SET_APPLICATION_ID, nv_gpu.NVC9B0_SET_APPLICATION_ID_ID_HEVC)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_CONTROL_PARAMS, nv_flags("NVC9B0_SET_CONTROL_PARAMS", codec_type="hevc", testrun_env="prod_run", gptimer_on=1,
+                err_conceal_on=1, mbtimer_on=1, event_trace_logging_on=1))
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_DRV_PIC_SETUP_OFFSET, desc_buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_IN_BUF_BASE_OFFSET, bufin.va_addr >> 8)
+    for pos, buf in zip([(frame_pos-x) % (len(hist) + 1) for x in range(len(hist), 0, -1)] + [frame_pos], hist + [bufout]):
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_LUMA_OFFSET0 + pos*4, buf.va_addr >> 8)
+      cmds += nvm(4, nv_gpu.NVC9B0_SET_PICTURE_CHROMA_OFFSET0 + pos*4, buf.offset(chroma_off).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_COLOC_DATA_OFFSET, dev.vid_coloc_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_NVDEC_STATUS_OFFSET, dev.vid_stat_buf._buf.va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_TILE_SIZES_OFFSET, desc_buf.offset(0x200).va_addr >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_HEVC_SET_FILTER_BUFFER_OFFSET, (filter_addr:=dev.vid_filter_buf._buf.va_addr) >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_SET_INTRA_TOP_BUF_OFFSET, (filter_addr + dev.intra_top_off) >> 8)
+    if dev.intra_unk_off is not None: cmds += nvm(4, 0x4dc, (filter_addr + dev.intra_unk_off) >> 8)
+    cmds += nvm(4, nv_gpu.NVC9B0_EXECUTE, 0)
+    dev._submit_cmds(dev.fifos["NVDEC:0"], *cmds)
+
+# *****************
+# device
 
 @dataclass
-class GPFifo:
-  ring: MMIOInterface
-  gpput: MMIOInterface
-  entries_count: int
-  token: int
-  put_value: int = 0
+class GPFifo: ring: Buffer; gpput: Buffer; doorbell: Buffer; put_value: Buffer; entries: int; token: int # noqa: E702
 
 class NVKIface:
   root = None
@@ -454,7 +415,7 @@ class NVKIface:
     self.uvm(nv_gpu.UVM_REGISTER_GPU_VASPACE, nv_gpu.UVM_REGISTER_GPU_VASPACE_PARAMS(
       gpuUuid=self.gpu_uuid, rmCtrlFd=self.fd_ctl.fd, hClient=self.root, hVaSpace=vaspace))
 
-    for dev in [d for pg in HCQCompiled.peer_groups.values() for d in pg if isinstance(d, NVDevice) and not d.is_nvd()]:
+    for dev in [d for x in Device._opened_devices if isinstance(d:=Device[x], NVDevice) and not d.is_nvd()]:
       try: self.uvm(nv_gpu.UVM_ENABLE_PEER_ACCESS, nv_gpu.UVM_ENABLE_PEER_ACCESS_PARAMS(gpuUuidA=self.gpu_uuid, gpuUuidB=dev.iface.gpu_uuid))
       except RuntimeError as e: raise RuntimeError(f"{e}. Make sure GPUs #{self.gpu_minor} & #{dev.iface.gpu_minor} have P2P enabled.") from e
 
@@ -582,8 +543,13 @@ class PCIIface(PCIIfaceBase):
 
 class MOCKIface(NVKIface): count = 1
 
-class NVDevice(HCQCompiled[NVSignal]):
+class NVDevice(HCQ2Compiled):
   ifaces = [NVKIface, PCIIface, MOCKIface]
+  sleep_timeout_ms = 200
+  pm_encode = PatternMatcher([
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_compute", name="submit"), lambda ctx, submit: encode_submit(NVComputeQueue(ctx, submit))),
+    (UPat(Ops.CUSTOM_FUNCTION, arg="submit_nv_copy", name="submit"), lambda ctx, submit: encode_submit(NVCopyQueue(ctx, submit))),
+  ])
 
   def is_nvd(self) -> bool: return isinstance(self.iface, PCIIface)
 
@@ -610,19 +576,11 @@ class NVDevice(HCQCompiled[NVSignal]):
     channel_params = nv_gpu.NV_CHANNEL_GROUP_ALLOCATION_PARAMETERS(engineType=nv_gpu.NV2080_ENGINE_TYPE_GRAPHICS)
     self.channel_group = self.iface.rm_alloc(self.nvdevice, nv_gpu.KEPLER_CHANNEL_GROUP_A, channel_params)
 
-    self.gpfifo_area = self.iface.alloc(0x300000, contiguous=True, cpu_access=True, force_devmem=True,
+    self.gpfifo_mem = self.iface.alloc(0x300000, contiguous=True, cpu_access=True, force_devmem=True,
       map_flags=(nv_gpu.NVOS33_FLAGS_CACHING_TYPE_WRITECOMBINED<<23))
 
-    ctxshare_params = nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC)
-    ctxshare = self.iface.rm_alloc(self.channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A, ctxshare_params)
-
-    self.compute_gpfifo = self._new_gpu_fifo(self.gpfifo_area, ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
-    self.dma_gpfifo = self._new_gpu_fifo(self.gpfifo_area, ctxshare, self.channel_group, offset=0x100000, entries=0x10000, compute=False)
-    self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
-
-    self.cmdq_page:HCQBuffer = self.iface.alloc(0x200000, cpu_access=True)
-    self.cmdq_allocator = BumpAllocator(size=self.cmdq_page.size, base=int(self.cmdq_page.va_addr), wrap=True)
-    self.cmdq = self.cmdq_page.cpu_view().view(fmt='I')
+    self.ctxshare = self.iface.rm_alloc(self.channel_group, nv_gpu.FERMI_CONTEXT_SHARE_A,
+      nv_gpu.NV_CTXSHARE_ALLOCATION_PARAMETERS(hVASpace=vaspace, flags=nv_gpu.NV_CTXSHARE_ALLOCATION_FLAGS_SUBCONTEXT_ASYNC))
 
     self.num_gpcs, self.num_tpc_per_gpc, self.num_sm_per_tpc, self.max_warps_per_sm, self.sm_version = self._query_gpu_info('num_gpcs',
       'num_tpc_per_gpc', 'num_sm_per_tpc', 'max_warps_per_sm', 'sm_version')
@@ -631,19 +589,37 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.arch: str = "sm_120" if self.sm_version==0xa04 else f"sm_{(self.sm_version>>8)&0xff}{(val>>4) if (val:=self.sm_version&0xff) > 0xf else val}"
     self.sass_version = ((self.sm_version & 0xf00) >> 4) | (self.sm_version & 0xf)
 
-    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], NVProgram, NVSignal, NVComputeQueue,
-                     NVCopyQueue, arch=self.arch)
+    self.slm_per_thread = 0
+    self.shader_local_mem:Buffer|None = None
+    # Set windows addresses to not collide with other allocated buffers.
+    self.shared_mem_window, self.local_mem_window = 0x729400000000, 0x729300000000
 
-    self.pma_enabled = PMA.value > 0 and PROFILE >= 1
-    if self.pma_enabled: self._prof_init()
+    super().__init__(device, NVAllocator(self), [CUDARenderer, PTXRenderer, NVCCRenderer, NAKRenderer], None, arch=self.arch)
 
-    self._setup_gpfifos()
+    self.pma_enabled, self.pma_exec_counter = PMA.value > 0 and PROFILE >= 1, itertools.count(0)
 
-  def _new_gpu_fifo(self, gpfifo_area, ctxshare, channel_group, offset=0, entries=0x400, compute=False, video=False) -> GPFifo:
+  @functools.cached_property
+  def fifos(self) -> dict[str, GPFifo]:
+    self.gpfifo_buf = Buffer(self.device, self.gpfifo_mem.size, dtypes.uint8, options=BufferSpec(external_ptr=self.gpfifo_mem.va_addr, nolru=True)) \
+                        .allocate(opaque=self.gpfifo_mem)
+    compute = self._new_gpu_fifo("COMPUTE:0", self.ctxshare, self.channel_group, offset=0, entries=0x10000, compute=True)
+    copy = self._new_gpu_fifo("COPY:0", self.ctxshare, self.channel_group, offset=0x100000, entries=0x10000)
+    self.iface.rm_control(self.channel_group, nv_gpu.NVA06C_CTRL_CMD_GPFIFO_SCHEDULE, nv_gpu.NVA06C_CTRL_GPFIFO_SCHEDULE_PARAMS(bEnable=1))
+
+    self._submit_cmds(compute, *nvm(1, nv_gpu.NVC6C0_SET_OBJECT, self.iface.compute_class),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_WINDOW_A, *data64(self.local_mem_window)),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_SHARED_MEMORY_WINDOW_A, *data64(self.shared_mem_window)))
+    self._submit_cmds(copy, *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.dma_class))
+
+    if self.pma_enabled: self._prof_init() # the sampler binds to the channel group, so it only comes up once the channels do
+    return {"COMPUTE:0": compute, "COPY:0": copy}
+
+  def _new_gpu_fifo(self, name:str, ctxshare, channel_group, offset=0, entries=0x400, compute=False, video=False) -> GPFifo:
     notifier = self.iface.alloc(48 << 20, uncached=True)
-    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=gpfifo_area.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
-      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else gpfifo_area.meta.hMemory,
-      hUserdMemory=(ctypes.c_uint32*8)(gpfifo_area.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset), engineType=19 if video else 0,
+    params = nv_gpu.NV_CHANNELGPFIFO_ALLOCATION_PARAMETERS(gpFifoOffset=self.gpfifo_mem.va_addr+offset, gpFifoEntries=entries, hContextShare=ctxshare,
+      hObjectError=notifier.meta.hMemory, hObjectBuffer=self.virtmem if video else self.gpfifo_mem.meta.hMemory,
+      hUserdMemory=(ctypes.c_uint32*8)(self.gpfifo_mem.meta.hMemory), userdOffset=(ctypes.c_uint64*8)(entries*8+offset),
+      engineType=19 if video else 0,
       hVASpace=self.vaspace if video and self.is_nvd() else 0) # gsp has no default vaspace, rm maps the decoder ctx into its own
     gpfifo = self.iface.rm_alloc(channel_group, self.iface.gpfifo_class, params)
 
@@ -662,8 +638,14 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.NVC36F_CTRL_CMD_GPFIFO_GET_WORK_SUBMIT_TOKEN_PARAMS(workSubmitToken=-1))
     if ctxshare != 0: self.iface.setup_gpfifo_vm(gpfifo)
 
-    return GPFifo(ring=gpfifo_area.cpu_view().view(offset, entries*8, fmt='Q'), entries_count=entries, token=ws_token_params.workSubmitToken,
-                  gpput=gpfifo_area.cpu_view().view(offset + entries*8 + getattr(nv_gpu.AmpereAControlGPFifo, 'GPPut').offset, fmt='I'))
+    gpput_off = offset + entries*8 + getattr(nv_gpu.AmpereAControlGPFifo, 'GPPut').offset
+    fifo = GPFifo(ring=self.gpfifo_buf.view(entries, dtypes.uint64, offset).ensure_allocated(),
+      gpput=self.gpfifo_buf.view(1, dtypes.uint32, gpput_off).ensure_allocated(),
+      doorbell=Buffer("CPU", 1, dtypes.uint32, options=BufferSpec(external_ptr=self.gpu_mmio.addr + 0x90), preallocate=True),
+      put_value=Buffer("CPU", 1, dtypes.uint64, preallocate=True), entries=entries, token=ws_token_params.workSubmitToken)
+    self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag=to_name(n, name)), lambda ctx, b=getattr(fifo, n): b)
+                                        for n in ("ring", "gpput", "doorbell", "put_value")]) + self.pm_bufferize
+    return fifo
 
   def _query_gpu_info(self, *reqs):
     nvrs = [getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_'+r.upper(), getattr(nv_gpu,'NV2080_CTRL_GR_INFO_INDEX_LITTER_'+r.upper(), None)) for r in reqs]
@@ -678,34 +660,35 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.NV2080_CTRL_GR_GET_INFO_PARAMS(grInfoListSize=len(infos), grInfoList=ctypes.addressof(infos)))
     return [x.data for x in infos]
 
-  def _setup_gpfifos(self):
-    self.slm_per_thread, self.shader_local_mem = 0, None
+  def _push(self, fifo:GPFifo, cmds:list[int]): # a pushbuffer built in python: channel setup and video decode
+    (buf:=self.rt_view(len(cmds) * 4))._buf.cpu_view().view(fmt='I')[:] = array.array('I', cmds)
 
-    # Set windows addresses to not collide with other allocated buffers.
-    self.shared_mem_window, self.local_mem_window = 0x729400000000, 0x729300000000
+    put = fifo.put_value._buf.view.view(fmt='Q')
+    fifo.ring._buf.cpu_view().view(fmt='Q')[put[0] % fifo.entries] = buf._buf.va_addr | (len(cmds) << 42) | (1 << 41)
+    fifo.gpput._buf.cpu_view().view(fmt='I')[0] = (put[0] + 1) % fifo.entries
 
-    NVComputeQueue().setup(compute_class=self.iface.compute_class, local_mem_window=self.local_mem_window, shared_mem_window=self.shared_mem_window) \
-                    .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    System.memory_barrier()
+    self.gpu_mmio[0x90 // 4] = fifo.token
+    put[0] += 1
 
-    NVCopyQueue().wait(self.timeline_signal, self.timeline_value - 1) \
-                 .setup(copy_class=self.iface.dma_class) \
-                 .signal(self.timeline_signal, self.next_timeline()).submit(self)
-
-    self.synchronize()
+  def _submit_cmds(self, fifo:GPFifo, *cmds:int): # runs cmds once everything already submitted is done, then bumps the timeline
+    tl, addr = self.timeline._buf.cpu_view().view(fmt='Q'), self.timeline._buf.va_addr
+    self._push(fifo, nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1]),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="acq_circ_geq", payload_size="64bit")) + list(cmds) +
+                     nvm(0, nv_gpu.NVC56F_SEM_ADDR_LO, *data64_le(addr), *data64_le(tl[1] + 1),
+                         nv_flags("NVC56F_SEM_EXECUTE", operation="release", release_wfi="en", payload_size="64bit")))
+    tl[1] += 1
 
   def _ensure_has_local_memory(self, required):
     if self.slm_per_thread >= required: return
 
-    self.slm_per_thread, old_slm_per_thread = round_up(required, 32), self.slm_per_thread
+    self.slm_per_thread = round_up(required, 32)
     bytes_per_tpc = round_up(round_up(self.slm_per_thread * 32, 0x200) * self.max_warps_per_sm * self.num_sm_per_tpc, 0x8000)
-    self.shader_local_mem, ok = self._realloc(self.shader_local_mem, round_up(bytes_per_tpc*self.num_tpc_per_gpc*self.num_gpcs, 0x20000))
+    self.shader_local_mem = Buffer(self.device, round_up(bytes_per_tpc*self.num_tpc_per_gpc*self.num_gpcs, 0x20000), dtypes.uint8,
+                                   options=BufferSpec(nolru=True), preallocate=True)
 
-    # Realloc failed, restore the old value.
-    if not ok: self.slm_per_thread = old_slm_per_thread
-
-    cast(NVComputeQueue, NVComputeQueue().wait(self.timeline_signal, self.timeline_value - 1)) \
-                                         .setup(local_mem=self.shader_local_mem.va_addr, local_mem_tpc_bytes=bytes_per_tpc) \
-                                         .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    self._submit_cmds(self.fifos["COMPUTE:0"], *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_A, *data64(self.shader_local_mem._buf.va_addr)),
+                       *nvm(1, nv_gpu.NVC6C0_SET_SHADER_LOCAL_MEMORY_NON_THROTTLED_A, *data64(bytes_per_tpc), 0xff))
 
   def _ensure_has_vid_hw(self, w, h):
     if self.iface.viddec_class is None: raise RuntimeError(f"{self.device} Video decoder class not available.")
@@ -716,18 +699,20 @@ class NVDevice(HCQCompiled[NVSignal]):
     self.intra_unk_off = (round_up(self.intra_top_off, 0x10000) + (64 << 10)) if intra_unk_size > 0 else None
     filter_sz = round_up(round_up(self.intra_top_off, 0x10000) + (64 << 10) + intra_unk_size, 2 << 20)
 
-    if not hasattr(self, 'vid_gpfifo'):
-      self.vid_gpfifo = self._new_gpu_fifo(self.gpfifo_area, 0, self.nvdevice, offset=0x200000, entries=2048, compute=False, video=True)
-      self.vid_coloc_buf, self.vid_filter_buf = (self.allocator.alloc(sz, BufferSpec(zero=True)) for sz in [coloc_sz, filter_sz])
-      self.vid_stat_buf = self.allocator.alloc(0x1000, BufferSpec(zero=True))
-      NVVideoQueue().wait(self.timeline_signal, self.timeline_value - 1) \
-                    .setup(copy_class=self.iface.viddec_class) \
-                    .signal(self.timeline_signal, self.next_timeline()).submit(self)
+    def _vid_buf(sz): return Buffer(self.device, sz, dtypes.uint8, options=BufferSpec(zero=True, nolru=True), preallocate=True)
+    if "NVDEC:0" not in self.fifos:
+      self.fifos["NVDEC:0"] = self._new_gpu_fifo("NVDEC:0", 0, self.nvdevice, offset=0x200000, entries=2048, video=True)
+      self.vid_coloc_buf, self.vid_filter_buf, self.vid_stat_buf = _vid_buf(coloc_sz), _vid_buf(filter_sz), _vid_buf(0x1000)
+      self._submit_cmds(self.fifos["NVDEC:0"], *nvm(4, nv_gpu.NVC6C0_SET_OBJECT, self.iface.viddec_class))
     else:
-      if coloc_sz > self.vid_coloc_buf.size: self.vid_coloc_buf,_= self._realloc(self.vid_coloc_buf, coloc_sz, BufferSpec(zero=True), force=True)
-      if filter_sz > self.vid_filter_buf.size: self.vid_filter_buf,_= self._realloc(self.vid_filter_buf, filter_sz, BufferSpec(zero=True), force=True)
+      if coloc_sz > self.vid_coloc_buf.nbytes: self.vid_coloc_buf = _vid_buf(coloc_sz)
+      if filter_sz > self.vid_filter_buf.nbytes: self.vid_filter_buf = _vid_buf(filter_sz)
 
-  def hw_copy_queues(self): return super().hw_copy_queues() + ([("NVDEC:0", NVVideoQueue)] if hasattr(self, 'vid_gpfifo') else [])
+  def collect_prof(self):
+    # the pc samples of a whole batch come back as one stream, so they are reported against the first kernel of it
+    if self.pma_enabled and (ents:=list(self.prof_ents.values())) and (blob:=self._prof_readback()) is not None:
+      Compiled.profile_events.append(ProfilePMAEvent(self.device, str(ents[0].name), blob, next(self.pma_exec_counter), ents[0].profile_key))
+    super().collect_prof()
 
   def invalidate_caches(self):
     if self.is_nvd(): self.iface.rm_control(self.subdevice, nv_gpu.NV2080_CTRL_CMD_INTERNAL_BUS_FLUSH_WITH_SYSMEMBAR, None)
@@ -845,4 +830,4 @@ class NVDevice(HCQCompiled[NVSignal]):
       nv_gpu.struct_NVB0CC_CTRL_PMA_STREAM_UPDATE_GET_PUT_PARAMS(bytesConsumed=params.bytesAvailable))
     return pma_data
 
-  def device_props(self): return {'arch': self.arch, 'sm_version': self.sm_version}
+  def device_props(self) -> dict[str, Any]: return {'arch': self.arch, 'sm_version': self.sm_version}

tinygrad/runtime/ops_python.py

@@ -18,9 +18,11 @@ def _load(m, i, dtype: DType):
   if (w:=m.nbytes // len(m)) >= dtype.itemsize: return from_storage_scalar(m[i], dtype)
   return sum(m[i+k] << (8*w*k) for k in range(dtype.itemsize // w)) # a bitcast can read wider than the buffer, _store splits it the same way
 
+def _step(m, dtype: DType): return max(1, dtype.itemsize // (m.nbytes // len(m))) # storage elements per lane
+
 def load(inp, j, dtype: DType):
-  if len(inp) >= 3: return [_load(m, x+j if x is not None else None, dtype) if gate else default for (m,x),default,gate in zip(*inp[:3])]
-  return [_load(m, x+j if x is not None else None, dtype) for m,x in inp[0]]
+  if len(inp) >= 3: return [_load(m, x+j*_step(m, dtype) if x is not None else None, dtype) if gate else alt for (m,x),alt,gate in zip(*inp[:3])]
+  return [_load(m, x+j*_step(m, dtype) if x is not None else None, dtype) for m,x in inp[0]]
 
 def _store(m, i, v, dtype: DType):
   if i < 0 or i >= len(m): raise IndexError(f"store out of bounds, size is {len(m)}, access is {i}, value is {v}")
@@ -86,7 +88,7 @@ class PythonProgram(Program['PythonDevice']):
           store_gate = exec_masks[-1]
           for j,val in enumerate(src_values[1] if u.max_numel() > 1 else [src_values[1]]):
             for (m,o),v,g in zip(src_values[0], val, store_gate):
-              if g: _store(m, o+j, v, src_dtypes[1])
+              if g: _store(m, o+j*_step(m, src_dtypes[1]), v, src_dtypes[1])
           i += 1
           continue
         if u.op is Ops.AFTER or (u.op is Ops.BITCAST and u.addrspace in (AddrSpace.GLOBAL, AddrSpace.LOCAL)): values[u] = src_values[0]

tinygrad/runtime/ops_qcom.py

@@ -53,14 +53,6 @@ def _read_lib(lib, off) -> int: return struct.unpack("I", lib[off:off+4])[0]
 
 class QCOMComputeQueue(HWQueue):
   dev:QCOMDevice
-  q_rewrite = PatternMatcher([
-    (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), lambda ctx, call, prg: ctx.exec(call, prg)),
-    (UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda ctx: ctx.memory_barrier()),
-    (UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.wait(dst, val)),
-    (UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), lambda ctx, dst: ctx.timestamp(dst)),
-    (UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), lambda ctx, dst, val: ctx.signal(dst, val)),
-  ])
-
   def cmd(self, opcode:int, *vals): self.q(pkt7_hdr(opcode, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)
 
   def reg(self, reg:int, *vals): self.q(pkt4_hdr(reg, sum(x.dtype.itemsize // 4 if isinstance(x, UOp) else 1 for x in vals)), *vals)

tinygrad/runtime/support/am/amdev.py

@@ -144,7 +144,7 @@ class AMMemoryManager(MemoryManager):
     self.dev.gmc.flush_tlb(ip='MM', vmid=0)
 
 class AMDev:
-  Version = 0xA0000008
+  Version = 0xA000000D
 
   def _disable_aspm(self):
     # L1 across retimers makes reads oscillate to 0xffffffff; power on defaults it enabled. Clearing the GPU endpoint
@@ -200,6 +200,7 @@ class AMDev:
         self.smu.mode1_reset()
       self.pci_dev.write_config_flush(pci.PCI_COMMAND, self.pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
       self.init_hw(self.soc, self.gmc, self.ih, *(() if self.is_vf else (self.psp, self.smu)))
+    elif not self.is_vf: self.psp._tmr_init()
 
     # Booting done
     self.is_booting = False
@@ -213,6 +214,7 @@ class AMDev:
         self.smu.set_clocks(level=None)
       else: self.smu.set_clocks(level=-1) # last level, max perf.
       for ip in [self.soc, self.gfx]: ip.set_clockgating_state()
+      self.reg("regSCRATCH_REG5").write(self.psp.tmr_size) # scratch registers are writable after GFX initialization
       self.reg("regSCRATCH_REG7").write(AMDev.Version)
       self.reg("regSCRATCH_REG6").write(1) # set initialized state.
 
@@ -222,7 +224,7 @@ class AMDev:
     self.smi_dev, self.is_err_state = smi_dev, False
 
     # Memory manager & firmware
-    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(32 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
+    self.mm = AMMemoryManager(self, self.vram_size - self.reserved_vram_size, boot_size=(3 << 20), pt_t=AMPageTableEntry, va_shifts=[12, 21, 30, 39],
       va_bits=48, first_lv=am.AMDGPU_VM_PDB2, va_base=AMMemoryManager.va_allocator.base, reserve_ptable=not self.large_bar,
       palloc_ranges=[(1 << (i + 12), (2 << 20) if i >= 9 else 0x1000) for i in range(9 * (3 - am.AMDGPU_VM_PDB2), -1, -1)])
     self.fw = AMFirmware(self)

tinygrad/runtime/support/am/ip.py

@@ -603,10 +603,9 @@ class AM_PSP(AM_IP):
     self.ring_size = 0x10000
     self.ring_paddr = self.adev.mm.palloc(self.ring_size, zero=False, boot=True)
 
-    self.max_tmr_size, self.tmr_size = 0x1300000, 0
+    self.tmr_size, self.tmr_paddr = 0, 0
     self.boot_time_tmr = self.adev.ip_ver[am.MP0_HWIP] in {(13,0,6), (13,0,14), (14,0,2), (14,0,3)}
     self.autoload_tmr = self.adev.ip_ver[am.MP0_HWIP] not in {(13,0,6), (13,0,14)}
-    self.tmr_paddr = self.adev.mm.palloc(self.max_tmr_size, align=am.PSP_TMR_ALIGNMENT, zero=False, boot=True) if not self.boot_time_tmr else 0
 
   def init_hw(self):
     spl_key = am.PSP_FW_TYPE_PSP_SPL if self.adev.ip_ver[am.MP0_HWIP] >= (14,0,0) else am.PSP_FW_TYPE_PSP_KDB
@@ -655,10 +654,13 @@ class AM_PSP(AM_IP):
     return self._wait_for_bootloader() if compid != am.PSP_BL__LOAD_SOSDRV else 0
 
   def _tmr_init(self):
-    # Load TOC and calculate TMR size
-    self._prep_msg1(fwm:=self.adev.fw.sos_fw[am.PSP_FW_TYPE_PSP_TOC])
-    self.tmr_size = self._load_toc_cmd(len(fwm)).resp.tmr_size
-    assert self.tmr_size <= self.max_tmr_size
+    if self.adev.partial_boot: self.tmr_size = self.adev.reg("regSCRATCH_REG5").read()
+    else:
+      # Load TOC and calculate TMR size
+      self._prep_msg1(fwm:=self.adev.fw.sos_fw[am.PSP_FW_TYPE_PSP_TOC])
+      self.tmr_size = self._load_toc_cmd(len(fwm)).resp.tmr_size
+    # First runtime allocation on both full and partial boots, so the resident TMR keeps the same address.
+    if not self.boot_time_tmr: self.tmr_paddr = self.adev.mm.pa_allocator.alloc(self.tmr_size, am.PSP_TMR_ALIGNMENT)
 
   def _ring_create(self):
     # If the ring is already created, destroy it