modeld: tinygrad jit copies and update tinygrad (#39187)

openpilot/selfdrive/modeld/modeld.py

@@ -167,6 +167,13 @@ def input_view(buffer: Buffer, shape: tuple[int, ...], dtype: DType, offset: int
   return Tensor(UOp.from_buffer(view, dtype)).reshape(shape)
 
 
+def compile_copy(tensor: Tensor) -> Callable[[], Tensor]:
+  from tinygrad.engine.jit import CapturedJit
+  from tinygrad.engine.realize import compile_linear
+  jit = CapturedJit(tensor, compile_linear(tensor.schedule_linear()), [], [])
+  return lambda: jit([], {})
+
+
 class ModelState:
   lat_delay: float
   frame_drop_ratio: float
@@ -200,6 +207,7 @@ class ModelState:
       state = self.input_queues[name]
       assert all_int(state.shape)
       self.outputs[next_name] = input_view(state._buffer(), state.shape, state.dtype, 0)
+    self.copy_outputs = None if selected_legacy() else compile_copy(self.outputs['outputs'].to('CPU'))
     self.parser = Parser()
 
   def pack_inputs(self) -> None:
@@ -210,8 +218,15 @@ class ModelState:
                                    if name not in self.state_pairs and name != 'new_img'}
     npy_size = sum(round_up(math.prod(shape) * 4, 128) for shape in shapes.values())
     self.packed_input = np.zeros(npy_size + 2 * self.frame_copy_size, dtype=np.uint8)
-    self.input_host = Tensor(self.packed_input, device='NPY')._buffer()
-    self.input_device = Tensor(self.packed_input, device=self.model_device)._buffer()
+    if selected_legacy():
+      self.input_host = Tensor(self.packed_input, device='NPY')._buffer()
+      self.input_device = Tensor(self.packed_input, device=self.model_device)._buffer()
+    else:
+      host = Tensor(self.packed_input, device='NPY')
+      device = host.to(self.model_device)
+      self.copy_inputs = compile_copy(device)
+      self.copy_inputs()
+      self.input_host, self.input_device = host._buffer(), device._buffer()
     self.npy = {}
     offset = 0
     for name, shape in shapes.items():
@@ -237,7 +252,10 @@ class ModelState:
     self.npy['traffic_convention'][:] = inputs['traffic_convention']
     self.npy['action_t'][:] = inputs['action_t']
 
-    self.input_device.copy_from(self.input_host)
+    if selected_legacy():
+      self.input_device.copy_from(self.input_host)
+    else:
+      self.copy_inputs()
     self.input_queues['new_img'] = self.run_warp(**self.warp_inputs)
     self.run_model(output_buffers=self.outputs, **self.input_queues)
     if (self.chestnut and not selected_legacy() and
@@ -245,7 +263,11 @@ class ModelState:
       raise RuntimeError(f"chestnut failed with {err}")
     if after_enqueue is not None:
       after_enqueue()
-    model_output = self.outputs['outputs'].numpy()[0]
+    if selected_legacy():
+      model_output = self.outputs['outputs'].numpy()[0]
+    else:
+      assert self.copy_outputs is not None
+      model_output = self.copy_outputs().numpy()[0]
     if self.chestnut and not np.all(np.isfinite(model_output)):
       raise RuntimeError("model output not finite")
     outputs_dict = self.parser.parse_outputs(self.slice_outputs(model_output, self.output_slices))

openpilot/starpilot/flm/producer_clock_equivalence.py

@@ -18,7 +18,7 @@ PRODUCER_SOURCE_HASHES = (
   ('openpilot', 'system/camerad/cameras/spectra.cc', '4fa26c17f0fb792a4ab6815657beccddee1d641faa2fd02e5aa7393f717a2f3a'),
   ('openpilot', 'system/camerad/cameras/camera_qcom2.cc', 'b34c24c9984a0ac086f7b961cb7ff09d0501231da50aaad0e676bb3dc3f2edd9'),
   ('openpilot', 'system/camerad/cameras/camera_common.cc', 'bd9a416a6ee21047b4deec6df04276d8717cd92fe14de43f3e9daff86b62bb8f'),
-  ('openpilot', 'selfdrive/modeld/modeld.py', '8d0fec445484b325e55cf4ee3bee95b7061d99b7e125611358eef407aa10e871'),
+  ('openpilot', 'selfdrive/modeld/modeld.py', '4ffe0193c64acab99f47b594ed98164e7ab6285f04f18205525d1b9931b9a086'),
   ('openpilot', 'selfdrive/modeld/fill_model_msg.py', 'b3916c34d1073d0c26a7a8adc51e7b148f63a1ac68b388226170fd0f8b8b56b4'),
   ('openpilot', 'system/sensord/sensord.py', '7c79ab07d775622a43001e5a9b9ccc9a7bafb065940185849b33eefc2fd26420'),
   ('openpilot', 'system/sensord/sensors/lsm6ds3_gyro.py', '3e561f3d8572e1477c9b4b810c7570837d51e81c5672fe21ec67a68cae00dd92'),

openpilot/starpilot/flm/tests/test_upstream_producer_transport.py

@@ -9,7 +9,7 @@ from openpilot.starpilot.flm import producer_clock_equivalence as clock
 
 
 APPROVED = {
-  'selfdrive/modeld/modeld.py': '8d0fec445484b325e55cf4ee3bee95b7061d99b7e125611358eef407aa10e871',
+  'selfdrive/modeld/modeld.py': '4ffe0193c64acab99f47b594ed98164e7ab6285f04f18205525d1b9931b9a086',
   'visionipc/visionipc_py.cc': '6c131c4f5eccb9f822f1aaae63733f26b1427c52340b61c9cdaf6cd4aee2f89d',
   'visionipc/visionipc_server.cc': 'a4325a5171eacb931eb271d06a992d2f27216040df113d92099f62cdfa67d81a',
   'visionipc/visionipc_client.cc': 'd142cd60eedabfc581e9770fcd6f0ca81b762ba9dbc608b54dd20a27a317a52d',

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

@@ -63,15 +63,11 @@ runs:
         echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
         # no buffers should be over 300MB in CI
         echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
-        if [[ "$RUNNER_OS" != "Windows" ]]; then
-          if [[ "$RUNNER_OS" == "macOS" ]]; then
-            sudo install -d -o "$USER" /opt/venv /opt/python
-          fi
-          echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
-          echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
-        else
-          echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
+        if [[ "$RUNNER_OS" == "macOS" ]]; then
+          sudo install -d -o "$USER" /opt/venv /opt/python
         fi
+        echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
+        echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
 
     - name: Set up uv
       uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
@@ -131,12 +127,7 @@ runs:
       run: uv cache prune --ci
     - name: Configure venv
       shell: bash
-      run: |
-        if [[ "$RUNNER_OS" == "Windows" ]]; then
-          echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
-        else
-          echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
-        fi
+      run: echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
 
     # ******************* apt *******************
     - name: Setup apt

tinygrad_repo/.github/workflows/autogen.yml

@@ -41,7 +41,7 @@ jobs:
       run: |
         find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
         python3 -c "from tinygrad.runtime.autogen import opencl"
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
+        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv_615, nv"
         python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
         python3 -c "from tinygrad.runtime.autogen.am import *"
         python3 -c "from tinygrad.runtime.autogen.nv_regs import *"

tinygrad_repo/.github/workflows/benchmark.yml

@@ -6,6 +6,8 @@ env:
   ASSERT_PROCESS_REPLAY: "0"
   PYTHONPATH: .
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  INFLUXDB_HOST: ${{ secrets.INFLUXDB_HOST }}
+  INFLUXDB_TOKEN: ${{ secrets.INFLUXDB_TOKEN }}
 
 on:
   push:
@@ -81,7 +83,7 @@ jobs:
 
   llmbenchmark:
     name: Benchmark ${{ matrix.model }} (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && '7900xtx' || '4090' }}"
     strategy:
       fail-fast: false
       matrix:
@@ -89,7 +91,7 @@ jobs:
         model: ['llama3.2:3b-f16', 'qwen3.8:27b', 'olmoe']
         # qwen3.8:27b doesn't fit on mac
         exclude: [{ dev: 'METAL', model: 'qwen3.8:27b' }, { dev: 'AMD', model: 'olmoe' }, { dev: 'NV', model: 'olmoe' }]
-    timeout-minutes: 15
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -99,14 +101,12 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+      run: sudo modprobe -a nvidia nvidia_uvm
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: AM_RESET=1 python3 -c 'from tinygrad import Device; Device[Device.DEFAULT]'
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -123,7 +123,7 @@ jobs:
 
   cifarbenchmark:
     name: HLB-CIFAR10 (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && '7900xtx' || '4090' }}"
     strategy:
       fail-fast: false
       matrix:
@@ -138,14 +138,12 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+      run: sudo modprobe -a nvidia nvidia_uvm
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: AM_RESET=1 python3 -c 'from tinygrad import Device; Device[Device.DEFAULT]'
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -170,7 +168,7 @@ jobs:
 
   mlperfbenchmark:
     name: MLPerf (${{ matrix.dev }})
-    runs-on: [self-hosted, Linux, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'AMD' && '7900xtx' || '4090' }}"
     strategy:
       fail-fast: false
       matrix:
@@ -185,18 +183,16 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+      run: sudo modprobe -a nvidia nvidia_uvm
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: AM_RESET=1 python3 -c 'from tinygrad import Device; Device[Device.DEFAULT]'
     - name: Symlink models and datasets
       run: |
         mkdir -p extra/datasets
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+        ln -s /datasets/imagenet extra/datasets/imagenet
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -211,7 +207,7 @@ jobs:
 
   sdbenchmark:
     name: Stable Diffusion (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && '7900xtx' || '4090' }}"
     strategy:
       fail-fast: false
       matrix:
@@ -226,14 +222,12 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+      run: sudo modprobe -a nvidia nvidia_uvm
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: AM_RESET=1 python3 -c 'from tinygrad import Device; Device[Device.DEFAULT]'
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -255,7 +249,7 @@ jobs:
 
   multigpubenchmark:
     name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'AMD' && '4x7900xtx' || '4x4090' }}"
     strategy:
       fail-fast: false
       matrix:
@@ -273,17 +267,13 @@ jobs:
     - name: Setup (AMD)
       if: ${{ matrix.dev == 'AMD' }}
       run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
-    - name: Setup (NV)
-      if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+        AM_RESET=1 python3 -c 'from tinygrad import Device; [Device[f"AMD:{i}"] for i in range(4)]'
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
         mkdir -p extra/datasets
-        ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
-        ln -s /raid/datasets/imagenet extra/datasets/imagenet
+        ln -s /weights/LLaMA-3 weights/LLaMA-3
+        ln -s /datasets/imagenet extra/datasets/imagenet
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -293,26 +283,29 @@ jobs:
       run: python3 test/external/process_replay/reset.py
     - name: Run LLaMA-3 8B on 4 GPUs with BEAM
       run: BENCHMARK_LOG=llama3_beam_4gpu JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
+    - name: Run full CIFAR training steps w 4 GPUS
+      run: time BENCHMARK_LOG=cifar_multigpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=525 BS=1024 GPUS=4 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
     - name: Run MLPerf resnet eval on training data
-      run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
-    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
+      run: time BENCHMARK_LOG=resnet_eval BS=512 GPUS=4 MODEL=resnet python3 examples/mlperf/model_eval.py
+    - name: Run 10 MLPerf ResNet50 training steps (4 gpu)
+      run: BENCHMARK_LOG=resnet_10steps_multigpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1024 GPUS=4 MODEL=resnet python3 examples/mlperf/model_train.py
+    - name: Run 10 MLPerf Bert training steps (4 gpu)
       # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+      run: BENCHMARK_LOG=bert_10steps_multigpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=48 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
+    - name: Test benchmark allreduce
+      if: ${{ matrix.dev == 'NV' }}
+      run: GPUS=4 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
   tests:
     name: Tests (DEV=${{ matrix.dev }})
-    runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
+    runs-on: "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && '7900xtx' || '4090' }}"
     strategy:
       fail-fast: false
       matrix:
         dev: ['METAL', 'AMD', 'NV']
-    timeout-minutes: 11
+    timeout-minutes: 14
     defaults:
       run:
         shell: bash -e -o pipefail {0}
@@ -322,14 +315,12 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v6
-    - name: Setup (AMD)
-      if: ${{ matrix.dev == 'AMD' }}
-      run: |
-        ./extra/hcq/hcq_smi.py amd rmmod --expect
-        ./extra/hcq/hcq_smi.py amd kill_pids --sudoless
     - name: Setup (NV)
       if: ${{ matrix.dev == 'NV' }}
-      run: ./extra/hcq/hcq_smi.py nv kill_pids --sudoless
+      run: sudo modprobe -a nvidia nvidia_uvm
+    - name: Setup (AMD)
+      if: ${{ matrix.dev == 'AMD' }}
+      run: AM_RESET=1 python3 -c 'from tinygrad import Device; Device[Device.DEFAULT]'
     - name: setup staging db
       if: github.ref == 'refs/heads/update_benchmark_staging'
       run: |
@@ -388,9 +379,6 @@ jobs:
       run: IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Train MNIST
       run: time TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py
-    - name: Test benchmark allreduce
-      if: ${{ matrix.dev == 'NV' }}
-      run: python test/external/external_benchmark_multitensor_allreduce.py
     # TODO: HEVC decode timing test
     # - name: HEVC Decode Benchmark
     #   if: ${{ matrix.dev == 'NV' }}
@@ -513,10 +501,10 @@ jobs:
       uses: ./.github/actions/process-replay
 
   testcommausbgpubenchmark:
-    name: UsbGPU Benchmark (comma)
+    name: openpilot UsbGPU Benchmark
     env:
       LFS_SHA: "51b12d20a25b76fa8a5717c2e5af01f148e469e2"
-    runs-on: [self-hosted, Linux, comma4]
+    runs-on: usb3-9060
     timeout-minutes: 14
     defaults:
       run:
@@ -530,18 +518,20 @@ jobs:
       run: |
         echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: reset chestnut
+      run: python3 extra/usbgpu/debug.py -rnw
     - name: openpilot compile big_driving_supercombo
       run: |
-        GMMU=0 DEV=USB+AMD:LLVM taskset -c 4-7 python3 examples/openpilot/compile_onnx.py "https://huggingface.co/commaai/openpilot-lfs/resolve/$LFS_SHA/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx" openpilot.pkl --out-of-band --retargetable
-        BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=25 taskset -c 4-7 python3 examples/openpilot/load_pickle.py --run openpilot.pkl --out-of-band --retarget
+        GMMU=0 DEV=USB+AMD:LLVM python3 examples/openpilot/compile_onnx.py "https://huggingface.co/commaai/openpilot-lfs/resolve/$LFS_SHA/openpilot/selfdrive/modeld/models/big_driving_supercombo.onnx" openpilot.pkl --out-of-band --retargetable
+        BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo GMMU=0 DEV=USB+AMD:LLVM ASSERT_MIN_STEP_TIME=25 python3 examples/openpilot/load_pickle.py --run openpilot.pkl --out-of-band --retarget
     - name: openpilot load_pickle big_driving_supercombo
-      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 taskset -c 4-7 python3 examples/openpilot/load_pickle.py openpilot.pkl --out-of-band
+      run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_load_pickle GMMU=0 DEV=USB+AMD ASSERT_MIN_LOAD_TIME=25 python3 examples/openpilot/load_pickle.py openpilot.pkl --out-of-band
     - name: compile driving warp
-      run: GMMU=0 DEV=USB+AMD:LLVM taskset -c 4-7 python3 examples/openpilot/compile_warp.py --frame 1928,1208,2048,1216,608,3735552 --warp-to 512x256 --layout yuv420 --frames 2 --output warp.pkl --retargetable
+      run: GMMU=0 DEV=USB+AMD:LLVM python3 examples/openpilot/compile_warp.py --frame 1928,1208,2048,1216,608,3735552 --warp-to 512x256 --layout yuv420 --frames 2 --output warp.pkl --retargetable
     - name: run warp pickle
-      run: BENCHMARK_LOG=usbgpu_openpilot_driving_warp GMMU=0 DEV=USB+AMD taskset -c 4-7 python3 examples/openpilot/load_pickle.py --run warp.pkl --retarget
+      run: BENCHMARK_LOG=usbgpu_openpilot_driving_warp GMMU=0 DEV=USB+AMD python3 examples/openpilot/load_pickle.py --run warp.pkl --retarget
     - name: Test copy speeds
-      run: SIZE=64000000 GMMU=0 DEV=USB+AMD taskset -c 4-7 python3 test/external/external_test_usb_asm24.py
+      run: SIZE=64000000 GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py
 
   drivertests:
     name: PCI Driver Tests (DEV=${{ matrix.dev }})
@@ -613,16 +603,3 @@ jobs:
         trap - EXIT
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-
-  llvmspeed:
-    name: LLVM Speed
-    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 10
-    if: github.repository_owner == 'tinygrad'
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v6
-    - name: Speed Test
-      run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
-    - name: Speed Test (BEAM=2)
-      run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py

tinygrad_repo/.github/workflows/platform.yml

@@ -144,39 +144,3 @@ jobs:
       run: python -m pytest -n=auto test/test_tiny.py test/device/$(python -c 'from tinygrad import Device; print(Device.DEFAULT.lower())') --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-
-# ****** Windows Tests ******
-
-  testwindows:
-    strategy:
-      fail-fast: false
-      matrix:
-        dev:
-          - 'CPU:CLANG'
-          - 'CPU:LLVM'
-          - 'CPU:X86'
-          - 'WEBGPU'
-
-    name: Windows (DEV=${{ matrix.dev }})
-    runs-on: windows-2025
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v6
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: windows-${{ matrix.dev }}-minimal
-          deps: testing_minimal
-          pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
-      - name: Set env
-        shell: bash
-        run: printf "DEV=${{ matrix.dev }}${{ matrix.dev == 'CPU:CLANG' && '\nCPU_COUNT=2' || '' }}" >> $GITHUB_ENV
-      - name: Check Device.DEFAULT and print some source
-        shell: bash
-        run: |
-          python -c "from tinygrad import Device; from tinygrad.helpers import Target; assert Device.DEFAULT == Target.parse('${{ matrix.dev }}').device"
-          DEBUG=4 python test/test_tiny.py TestTiny.test_plus
-      - name: Run test_tiny and device tests
-        shell: bash
-        run: python -m pytest -n=auto test/test_tiny.py test/device/$(python -c 'from tinygrad import Device; print(Device.DEFAULT.lower())') --durations=20

tinygrad_repo/.github/workflows/test.yml

@@ -228,7 +228,7 @@ jobs:
     - name: Run test as module
       run: DEV=CPU python test/null/test_device.py TestRunAsModule.test_module_runs
     - name: Run GC tests
-      run: python test/external/external_uop_gc.py
+      run: DEV=CPU python test/external/external_uop_gc.py
     - name: External Benchmark Schedule
       run: python3 test/external/external_benchmark_schedule.py
     - name: Run process replay tests
@@ -289,7 +289,7 @@ jobs:
           opencl: 'true'
       - name: Test CL IMAGE=1 ops
         run: |
-          DEV=CL IMAGE=1 python -m pytest -n=auto test/runtime/test_ops.py --durations=20
+          DEV=CL IMAGE=1 python -m pytest -n=auto test/runtime/test_ops.py test/runtime/test_assign.py --durations=20
           DEV=CL IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
@@ -562,6 +562,14 @@ jobs:
         run: python -m pytest test/null/test_hcq2.py test/runtime/test_hcq2.py
       - name: Run disk copy tests on MOCKPCI
         run: python -m pytest test/runtime/test_disk_tensor.py -k test_copy_from_disk
+      - name: Run test_tiny on MOCKPCI over REMOTE
+        run: |
+          python3 extra/remote/serve.py 6482 &
+          trap 'kill $!' EXIT
+          sleep 1
+          REMOTE=127.0.0.1:6482 python -m pytest test/test_tiny.py
+      - name: Run tests on MOCK+HIP
+        run: DEV=MOCK+HIP python -m pytest -n=auto test/test_tiny.py test/runtime/test_jit.py test/runtime/test_hcq2.py test/runtime/test_dtype.py
 
   testamd:
     strategy:

tinygrad_repo/.pre-commit-config.yaml

@@ -1,4 +1,3 @@
-# on Windows -- $env:SKIP="tests,example"
 repos:
   - repo: local
     hooks:

tinygrad_repo/docs/developer/am.md

@@ -5,10 +5,7 @@ AM driver is a userspace driver targeting AMD's RDNA3/RDNA4. You only need tinyg
 ## How to run?
 Make sure that amdgpu module is unloaded and just run tinygrad with `DEV=AMD`!
 
-Optional requirements:
-
-* System without IOMMU for P2P / SDMA support
-* vfio-pci module for IRQ handling
+If the IOMMU translates (not `iommu=pt`), tinygrad drives the GPU through vfio-pci (`sudo modprobe vfio-pci disable_idle_d3=1`).
 
 ## Environment Variables
 

tinygrad_repo/examples/hlb_cifar10.py

@@ -261,7 +261,8 @@ def train_cifar():
         net_ema_param.assign(net_param.numpy())
 
     @TinyJit
-    def update(self, net, decay):
+    def update(self, net, step):
+      decay = projected_ema_decay_val*(Tensor(step)/STEPS)**hyp['ema']['decay_pow']
       for net_ema_param, (param_name, net_param) in zip(get_state_dict(self.net_ema).values(), get_state_dict(net).items()):
         # batchnorm currently is not being tracked
         if not ("num_batches_tracked" in param_name) and not ("running" in param_name):
@@ -320,6 +321,7 @@ def train_cifar():
   lr_sched_non_bias = OneCycleLR(opt_non_bias, max_lr=hyp['opt']['non_bias_lr'], pct_start=pct_start, div_factor=initial_div_factor, final_div_factor=1./(initial_div_factor*final_lr_ratio), total_steps=STEPS)
 
   def train_step(model, optimizer, lr_scheduler, X, Y):
+    X, Y = X.shard(GPUS, axis=0), Y.shard(GPUS, axis=0)  # copies to GPUS (and back to CPU on return) happen inside the jit
     out = model(X)
     loss_batchsize_scaler = 512/BS
     loss = cross_entropy(out, Y, reduction='none', label_smoothing=hyp['opt']['label_smoothing']).mul(hyp['opt']['loss_scale_scaler']*loss_batchsize_scaler).sum().div(hyp['opt']['loss_scale_scaler'])
@@ -328,16 +330,17 @@ def train_cifar():
       # index 0 for bias and 1 for non-bias
       optimizer.zero_grad()
       loss.backward()
-      return loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
-    return loss.realize()
+      loss.realize(*optimizer.schedule_step(), *lr_scheduler[0].schedule_step(), *lr_scheduler[1].schedule_step())
+    return loss.to("CPU"), optimizer[1].lr.to("CPU")
 
   train_step_jitted = TinyJit(train_step)
 
   def eval_step(model, X, Y):
+    X, Y = X.shard(GPUS, axis=0), Y.shard(GPUS, axis=0)
     out = model(X, training=False)
     loss = cross_entropy(out, Y, reduction='mean')
     correct = out.argmax(axis=1) == Y.argmax(axis=1)
-    return correct.realize(), loss.realize()
+    return correct.to("CPU"), loss.to("CPU")
   eval_step_jitted     = TinyJit(eval_step)
   eval_step_ema_jitted = TinyJit(eval_step)
 
@@ -365,10 +368,6 @@ def train_cifar():
         losses = []
         losses_ema = []
         for Xt, Yt in fetch_batches(X_test, Y_test, BS=EVAL_BS, is_train=False):
-          if len(GPUS) > 1:
-            Xt.shard_(GPUS, axis=0)
-            Yt.shard_(GPUS, axis=0)
-
           correct, loss = eval_step_jitted(model, Xt, Yt)
           losses.append(loss.numpy().tolist())
           corrects.extend(correct.numpy().tolist())
@@ -392,25 +391,21 @@ def train_cifar():
 
       with WallTimeEvent(BenchEvent.STEP):
         X, Y = next(batcher)
-        if len(GPUS) > 1:
-          X.shard_(GPUS, axis=0)
-          Y.shard_(GPUS, axis=0)
-
         with Context(BEAM=getenv("LATEBEAM", BEAM.value), WINO=getenv("LATEWINO", WINO.value)):
-          loss = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
+          loss, lr = train_step_jitted(model, optim.OptimizerGroup(opt_bias, opt_non_bias), [lr_sched_bias, lr_sched_non_bias], X, Y)
           et = time.monotonic()
           loss_cpu = loss.numpy()
         # EMA for network weights
         if getenv("EMA") and i > hyp['ema']['steps'] and (i+1) % hyp['ema']['every_n_steps'] == 0:
           if model_ema is None:
             model_ema = modelEMA(W, model)
-          model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))
+          model_ema.update(model, Variable("step", 0, STEPS).bind(i))
 
       cl = time.monotonic()
       step_times.append((cl-st)*1000.0)
-      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+      device_str = Device.DEFAULT if len(GPUS) == 1 else f"{Device.DEFAULT} * {len(GPUS)}"
       #  53  221.74 ms run,    2.22 ms python,  219.52 ms CL,  803.39 loss, 0.000807 LR, 4.66 GB used,   3042.49 GFLOPS,    674.65 GOPS
-      print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
+      print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
       st = cl
       i += 1
 

tinygrad_repo/examples/mlperf/model_train.py

@@ -257,6 +257,7 @@ def train_resnet():
           min_time = min(step_times)
           assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
 
+        if not INITMLPERF: batch_loader.close()
         if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
         return
     if MLLOGGER and RUNMLPERF:

tinygrad_repo/extra/amdpci/setup_vfio.sh

@@ -1,2 +0,0 @@
-#!/bin/bash
-sudo modprobe vfio-pci disable_idle_d3=1

tinygrad_repo/extra/export_model.py

@@ -39,7 +39,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
     prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
     info = prg.arg
     functions[prg.src[0].arg.function_name] = prg.src[2].arg
-    cargs = [name_of(bu, i == 0) for i, bu in enumerate(arg_uops)] + list(info.vars)
+    cargs = [p if p.addrspace is AddrSpace.ALU else name_of(call.src[1+p.arg.slot], p.arg.slot == 0) for p in prg.kernel_params]
     statements.append((prg.src[0].arg.function_name, cargs, info.global_size, info.local_size))
 
   return functions, statements, {name:(size, dtype, key) for name, size, dtype, key in bufs.values()}, bufs_to_save
@@ -56,7 +56,7 @@ def jit_model(model, *args) -> Tuple[UOp, List[Buffer]]:
   # run twice to trigger JIT capture
   for _ in range(2): the_output = run(*args)
   assert run.captured is not None
-  return run.captured.linear, [o.uop.base.realized for o in the_output]
+  return run.captured._linears[()][0], [o.uop.base.realized for o in the_output]
 
 def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,int,int]], bufs:Dict[str,Tuple[str,int,int]],
   bufs_to_save:Dict[str,Tensor], input_names:List[str], output_names:List[str], weight_names={}, model_name="model", symbolic_vars={}, wasm=False) -> str:
@@ -72,7 +72,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
       weight = ''.join(["\\x%02X"%x for x in cl.as_memoryview()])
       cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}(&(struct {name}_args){{{', '.join(args)}}});" for (name, args, _global_size, _local_size) in statements] + ["}"]
     return '\n'.join(headers + cprog)
   else:
     if bufs_to_save:
@@ -88,7 +88,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
 
     cprog += [f"void net({forward_args})"] + ["{"]
     get_weight_ptr = lambda x: f"({dtype_map[bufs_to_save[x][1]]} *)bufs[{buf_to_name[x]['idx']}]" if x in bufs_to_save else x
-    cprog += [f"  {name}({', '.join(map(get_weight_ptr, args))});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"  {name}(&(struct {name}_args){{{', '.join(map(get_weight_ptr, args))}}});" for (name, args, _global_size, _local_size) in statements] + ["}"]
     weightMapping = "" if not bufs_to_save else f"""\nconst weightNames = [{", ".join([f'"{weight_name}"' for weight_name in [v["name"] for v in buf_to_name.values()]])}];
 const {model_name}_name_to_id = Object.fromEntries(weightNames.map((name, index) => [name, index]));\n"""
     top = f"""import {model_name}Module from './{model_name}.js'{weightMapping}"""

tinygrad_repo/extra/gemm/moe_gemm.py

@@ -101,11 +101,12 @@ def custom_grouped_mx_gemm_bw(gradient:UOp, kernel:UOp, w_stored:bool=False) ->
 
 def grouped_mx_gemm_backward(g:Tensor, aq:Tensor, ae8:Tensor, bq:Tensor, be8:Tensor, expert_off:Tensor, *,
                              w_stored:bool=False, quantized_gradient:tuple[Tensor, Tensor]|None=None) -> tuple[Tensor, Tensor]:
+  from extra.gptoss_kernels.transpose_requantize import transpose_requantize_mxfp8
   E, N = bq.shape[:2]
   g = g.reshape(aq.shape[0], N).cast(dtypes.bfloat16)
   x_phys = (aq.cast(dtypes.bfloat16) * _mx_block_scale(ae8).cast(dtypes.bfloat16))
-  w_phys = (bq.cast(dtypes.bfloat16) * _mx_block_scale_3d(be8).cast(dtypes.bfloat16))
-  grad_x = grouped_mx_gemm(g if quantized_gradient is None else quantized_gradient, w_phys.transpose(1, 2), expert_off)
+  w_transposed = transpose_requantize_mxfp8(bq, be8)
+  grad_x = grouped_mx_gemm(g if quantized_gradient is None else quantized_gradient, w_transposed, expert_off)
   grad_w = grouped_mx_wgrad(g, x_phys, expert_off, E)
   grad_xq = grad_x * _mx_block_scale(ae8).cast(dtypes.bfloat16)
   grad_wq = grad_w.contiguous() if w_stored else (grad_w * _mx_block_scale_3d(be8).cast(dtypes.bfloat16)).contiguous()
@@ -113,14 +114,14 @@ def grouped_mx_gemm_backward(g:Tensor, aq:Tensor, ae8:Tensor, bq:Tensor, be8:Ten
 
 _grouped_bw_stored = functools.partial(custom_grouped_mx_gemm_bw, w_stored=True)
 
-def grouped_mx_gemm(x:Tensor|tuple[Tensor, Tensor], w:Tensor|tuple[Tensor, Tensor], expert_off:Tensor) -> Tensor:
+def grouped_mx_gemm(x:Tensor|tuple[Tensor, Tensor], w:Tensor|tuple[Tensor, Tensor]|tuple[Tensor, Tensor, Tensor], expert_off:Tensor) -> Tensor:
   if isinstance(x, tuple):
     x_q, x_e8 = x
     x_si, x = mx_pack(x_e8), x_q
   else:
     x_q, x_e8, x_si = quantize_mxfp8(x)
   if (pre_quantized := isinstance(w, tuple)):
-    w_q, w_e8 = w
+    w_q, w_e8 = w[:2]
     E, N, K2 = w_q.shape
   else:
     E, N, K2 = w.shape
@@ -130,7 +131,7 @@ def grouped_mx_gemm(x:Tensor|tuple[Tensor, Tensor], w:Tensor|tuple[Tensor, Tenso
   assert M % 256 == 0 and N % 256 == 0 and K % 128 == 0, f"grouped mxfp8 needs M%256,N%256,K%128, got {M,K,N}"
   dname = (x.device[0] if isinstance(x.device, tuple) else x.device).split(":")[0]
   if not pre_quantized: w_q, w_e8, _ = quantize_mxfp8(w)
-  w_si = mx_pack_3d(w_e8)
+  w_si = w[2] if pre_quantized and len(w) == 3 else mx_pack_3d(w_e8)
   xe_in, out_shape = x_e8.reshape(M, K // 32), (M, N)
   if isinstance(x.device, tuple) and (row_axis := x.uop.axis) is not None:
     ndev = len(x.device)

tinygrad_repo/extra/gemm/triton_nv_matmul.py

@@ -100,7 +100,7 @@ if __name__ == "__main__":
   gsize, lsize = info.launch_dims({})
   tflops = []
   for i in range(5):
-    tm = rt(*[b._buf for b in prg_bufs], global_size=gsize, local_size=lsize, vals=info.vals({}), wait=True)
+    tm = rt(*[b._buf for b in prg_bufs], global_size=gsize, local_size=lsize, wait=True)
     tflops.append((2*M*K*N/tm)*1e-12)
   print(f"TFLOPS: {max(tflops):.2f}")
 

tinygrad_repo/extra/gptoss_kernels/qkv_rope/__init__.py

@@ -15,14 +15,19 @@ def _forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs:UOp) -> UOp:
 
 @functools.cache
 def _backward_kernel(out:UOp, dq:UOp, dk:UOp, dv:UOp, freqs:UOp) -> UOp:
-  assert out.shape == (2, 8192, 5120) and dq.shape == (2, 64, 8192, 64) and dk.shape == dv.shape == (16, 8192, 8, 64)
+  assert out.shape == (2, 8192, 5120) and dq.shape == (2, 64, 8192, 64) and dk.shape == dv.shape
+  assert dk.shape[1:] == (8192, 8, 64) and dk.shape[0] in (2, 8)
+  partials = dk.shape[0] // 2
+  q_heads_per_wg = 8 if partials == 4 else 2
   sink = UOp.sink(out.base, dq.base, dk.base, dv.base, freqs.base,
-                  UOp.special(256, "lidx0"), UOp.special(2, "gidx0"), UOp.special(128, "gidx1"), UOp.special(80, "gidx2"),
+                  UOp.special(256, "lidx0"), UOp.special(2, "gidx0"), UOp.special(128, "gidx1"),
+                  UOp.special(64 // q_heads_per_wg + 16, "gidx2"),
                   arg=KernelInfo("fused_qkv_rope_backward_gptoss"))
   src = (pathlib.Path(__file__).parent/"backward.cpp").read_text()
   include = pathlib.Path(__file__).parents[2]/"thunder"/"amd"/"include"
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
-                               UOp(Ops.BINARY, arg=compile_hip(src, [f"-I{include}", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS"]))))
+                               UOp(Ops.BINARY, arg=compile_hip(src, [f"-I{include}", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS",
+                                                                  f"-DKV_PARTIALS={partials}"]))))
 
 def _backward(dq:UOp, dk:UOp, dv:UOp, call:UOp) -> tuple:
   from extra.thunder.amd.fa import _fa_native_grads

tinygrad_repo/extra/gptoss_kernels/qkv_rope/backward.cpp

@@ -23,9 +23,11 @@ using namespace kittens;
 constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
 constexpr int HALF_D = ATTN_D / 2;
 constexpr int PACKED_H = ATTN_H_KV * (GROUP_SIZE + 2);
-// The current D=64 FA emits one dK/dV partial per query head in each GQA group.
-constexpr int KV_PARTIALS = GROUP_SIZE;
-constexpr int Q_HEADS_PER_WG = 1;
+// Full-causal FA groups two query heads; sliding-window FA groups all eight.
+#ifndef KV_PARTIALS
+#define KV_PARTIALS 4
+#endif
+constexpr int Q_HEADS_PER_WG = KV_PARTIALS == 4 ? 8 : 2;
 constexpr int Q_FIELDS = ATTN_H / Q_HEADS_PER_WG;
 constexpr int NUM_WARPS = 4;
 constexpr int TILE_N = 16;
@@ -49,7 +51,7 @@ __device__ __forceinline__ void load_fa_shuffled(RT &dst, const GL &src, const C
     #pragma unroll
     for (int j = 0; j < dst.width; j++) {
       const float4 loaded = std::bit_cast<float4>(llvm_amdgcn_raw_buffer_load_b128(
-        std::bit_cast<i32x4>(br), (i * tile_row_stride + j * tile_stride + lane * 8) * sizeof(U), 0, 0));
+        std::bit_cast<i32x4>(br), (i * tile_row_stride + j * tile_stride + lane * 8) * sizeof(U), 0, 2)); // non-temporal dQ loads
       const U2 *packed = reinterpret_cast<const U2*>(&loaded);
       #pragma unroll
       for (int k = 0; k < dst.packed_per_base_tile; k++) dst.tiles[i][j].data[k] = packed[k];

tinygrad_repo/extra/gptoss_kernels/transpose_requantize/__init__.py

@@ -0,0 +1,27 @@
+import functools, pathlib
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from extra.llama_kernels import alloc_like, compile_hip
+
+@functools.cache
+def _kernel(q:UOp, e8:UOp, si:UOp, src_q:UOp, src_e8:UOp) -> UOp:
+  E, M, N = src_q.shape
+  assert M % 128 == 0 and N % 128 == 0
+  assert src_e8.shape == (E, M, N // 32)
+  assert q.shape == (E, N, M) and e8.shape == (E, N, M // 32) and si.shape == (E, M // 128, N)
+  sink = UOp.sink(q.base, e8.base, si.base, src_q.base, src_e8.base,
+                  UOp.special(128, "lidx0"), UOp.special(E * (M // 32) * (N // 128), "gidx0"),
+                  arg=KernelInfo(f"transpose_requantize_mxfp8_{E}_{M}_{N}"))
+  src = (pathlib.Path(__file__).parent/"transpose_requantize.cpp").read_text()
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src),
+                               UOp(Ops.BINARY, arg=compile_hip(src, [f"-DE_DIM={E}", f"-DM_DIM={M}", f"-DN_DIM={N}"]))))
+
+def transpose_requantize_mxfp8(q:Tensor, e8:Tensor) -> tuple[Tensor, Tensor, Tensor]:
+  E, M, N = q.shape
+  assert q.dtype == dtypes.fp8e4m3 and e8.dtype == dtypes.uint8 and e8.shape == (E, M, N // 32)
+  axis = q.uop.axis
+  out_axis = None if axis is None else (0 if axis == 0 else 3 - axis)
+  qt = alloc_like((E, N, M), q.dtype, q.device, out_axis)
+  et = alloc_like((E, N, M // 32), dtypes.uint8, q.device, out_axis)
+  si = alloc_like((E, M // 128, N), dtypes.uint32, q.device, axis)
+  return tuple(Tensor.custom_kernel(qt, et, si, q, e8, fxn=_kernel)[:3])

tinygrad_repo/extra/gptoss_kernels/transpose_requantize/transpose_requantize.cpp

@@ -0,0 +1,58 @@
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp8.h>
+#include <hip/hip_bf16.h>
+
+constexpr int BLK = 32;
+constexpr int TILE_M = BLK;
+constexpr int TILE_N = 128;
+constexpr int N_TILES_N = N_DIM / TILE_N;
+constexpr int REQUANT_M_GROUPS = M_DIM / 128;
+static_assert(M_DIM % 128 == 0 && N_DIM % TILE_N == 0);
+typedef short short2v __attribute__((ext_vector_type(2)));
+
+// Source15's M32 schedule: dequantize with its BF16 rounding boundary, then transpose and requantize.
+extern "C" __global__ __launch_bounds__(128) void
+transpose_requantize_mxfp8(__hip_fp8_storage_t* __restrict__ q_out,
+                          uint8_t* __restrict__ e8_out,
+                          uint32_t* __restrict__ si_out,
+                          const __hip_fp8_storage_t* __restrict__ q_in,
+                          const uint8_t* __restrict__ e8_in) {
+  const int tid = threadIdx.x;
+  int tile = blockIdx.x;
+  const int tile_n = tile % N_TILES_N;
+  tile /= N_TILES_N;
+  const int tile_m = tile % (M_DIM / BLK);
+  const int expert = tile / (M_DIM / BLK);
+  const int m_group = tile_m / 4, sub = tile_m % 4;
+  const long long n = tile_n * TILE_N + tid;
+  float vals[TILE_M];
+  float amax = 0.0f;
+  #pragma unroll
+  for (int mm = 0; mm < TILE_M; mm++) {
+    const long long m = tile_m * TILE_M + mm;
+    const long long src = ((long long)expert * M_DIM + m) * N_DIM + n;
+    const __hip_fp8_e4m3 fq = __builtin_bit_cast(__hip_fp8_e4m3, q_in[src]);
+    const long long src_e8 = ((long long)expert * M_DIM + m) * (N_DIM / BLK) + n / BLK;
+    const unsigned src_scale = e8_in[src_e8];
+    const float dqscale = __builtin_bit_cast(float, src_scale ? src_scale << 23 : 0x00400000u);
+    const float v = (float)(__hip_bfloat16)((float)fq * dqscale);
+    vals[mm] = v;
+    amax = fmaxf(amax, fabsf(v));
+  }
+  int e8 = (int)((__builtin_bit_cast(unsigned, amax) >> 23) & 0xFFu);
+  e8 = max(0, min(254, e8));
+  const float scale_f = __builtin_bit_cast(float, (unsigned)e8 << 23);
+  __hip_fp8_storage_t out[TILE_M];
+  #pragma unroll
+  for (int mm = 0; mm < TILE_M; mm += 2) {
+    short2v acc = {0, 0};
+    acc = __builtin_amdgcn_cvt_scalef32_pk_fp8_f32(acc, vals[mm], vals[mm + 1], scale_f, false);
+    const unsigned char *rb = reinterpret_cast<const unsigned char *>(&acc);
+    out[mm] = rb[0]; out[mm + 1] = rb[1];
+  }
+  const long long obase = ((long long)expert * N_DIM + n) * M_DIM + (long long)tile_m * TILE_M;
+  *reinterpret_cast<uint4*>(&q_out[obase]) = *reinterpret_cast<uint4*>(&out[0]);
+  *reinterpret_cast<uint4*>(&q_out[obase + 16]) = *reinterpret_cast<uint4*>(&out[16]);
+  e8_out[((long long)expert * N_DIM + n) * (M_DIM / BLK) + tile_m] = (uint8_t)e8;
+  reinterpret_cast<uint8_t *>(si_out)[(((long long)expert * REQUANT_M_GROUPS + m_group) * N_DIM + n) * 4 + sub] = (uint8_t)e8;
+}

tinygrad_repo/extra/hcq1/graph.py

@@ -135,7 +135,6 @@ class HCQGraph(GraphRunner):
     # When profiling allocate 2 signals for each jit item to measure speed. The jth jit item have signals at 2*j and 2*j+1.
     # TODO: This logic might allocate a few extra signals...
     self.prof_signals: list[HCQSignal] = []
-    self.prof_graph_deps: list[list[int]] = []
     self.prof_graph_entries: list[ProfileGraphEntry] = []
 
     self.last_j: dict[HWQueue, int|None] = collections.defaultdict(lambda: None)
@@ -205,7 +204,6 @@ class HCQGraph(GraphRunner):
         prof_name = enqueue_dev.device if runtime is not None else f"{enqueue_dev.device}:SDMA:{queue_idx}"
         self.prof_graph_entries.append(ProfileGraphEntry(prof_name, prof_ji_desc, sig_st, j * 2 + 1,
                                                          runtime.profile_key if runtime is not None else None))
-        self.prof_graph_deps.append([d - 1 for _, d in rdeps])
 
       self.last_j[enqueue_queue] = j
 
@@ -368,7 +366,7 @@ class HCQGraph(GraphRunner):
 
   def collect_timestamps(self):
     # NOTE: Append to any device is fine...
-    self.devices[0].profile_events += [ProfileGraphEvent(self.prof_graph_entries, self.prof_graph_deps, [s.timestamp for s in self.prof_signals])]
+    self.devices[0].profile_events += [ProfileGraphEvent(self.prof_graph_entries, [s.timestamp for s in self.prof_signals])]
 
   def dev_name(self, dev) -> str: return dev.device.replace(":", "_")
 

tinygrad_repo/extra/llama_kernels/fused_ce/__init__.py

@@ -57,7 +57,7 @@ def _fused_ce_loss_bwd(gradient:UOp, kernel:UOp, label_smoothing:float):
     seq_per_dev = SEQ
   # NOTE: .mean() backward gives same grad per row (1/N), so broadcast is safe; take scalar
   scale = Tensor(gradient, device=device).float().reshape(-1)[0:1].contiguous()
-  logits_t = Tensor(logits_u.after(kernel), device=device)
+  logits_t = Tensor(logits_u, device=device)
   lse_t = Tensor(lse_u.after(kernel), device=device)
   targets_t = Tensor(targets_u, device=device)
   fxn = functools.partial(_custom_fused_ce_loss_bwd, vocab=VOCAB, rows=rows_per_dev, seq=seq_per_dev, label_smoothing=label_smoothing)

tinygrad_repo/extra/perfetto/to_perfetto.py

@@ -12,16 +12,12 @@ def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
 def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
   name = ev.name.display_name if isinstance(ev.name, TracingKey) else ev.name
   return [{"name": name, "ph": "X", "ts": prep_ts(ev.device, ev.st), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device)}]
-def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
+def graph_ev_to_perfetto_json(ev:ProfileGraphEvent):
   ret = []
-  for i,e in enumerate(ev.ents):
+  for e in ev.ents:
     st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
     name = e.name.display_name if isinstance(e.name, TracingKey) else e.name
     ret += [{"name": name, "ph": "X", "ts": prep_ts(e.device, st), "dur": float(en-st), **dev_to_pid(e.device)}]
-    for dep in ev.deps[i]:
-      d = ev.ents[dep]
-      ret += [{"ph": "s", **dev_to_pid(d.device), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id]), "bp": "e"}]
-      ret += [{"ph": "f", **dev_to_pid(e.device), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st), "bp": "e"}]
   return ret
 def to_perfetto(profile:list[ProfileEvent]):
   # Start json with devices.
@@ -30,7 +26,7 @@ def to_perfetto(profile:list[ProfileEvent]):
   prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
   for ev in tqdm(profile, desc="preparing profile"):
     if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)
-    elif isinstance(ev, ProfileGraphEvent): prof_json += graph_ev_to_perfetto_json(ev, reccnt=len(prof_json))
+    elif isinstance(ev, ProfileGraphEvent): prof_json += graph_ev_to_perfetto_json(ev)
   return {"traceEvents": prof_json}
 
 if __name__ == "__main__":

tinygrad_repo/extra/remote/serve.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-import socket, struct, sys, signal
+import socket, struct, sys, signal, ctypes, time, json
 from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System, REMOTE_REQ, REMOTE_RESP
 from tinygrad.runtime.support.am.amdev import AMMemoryManager
 from tinygrad.runtime.support.system import FileIOInterface
@@ -12,6 +12,7 @@ def resp_err(msg): return resp(len(err:=msg.encode()), status=1) + err
 discovered_devices: list[tuple[type, str]] = []
 opened_devices: dict[int, PCIDevice] = {}
 mapped_bars: dict[tuple[int, int], object] = {}
+sysmem_iovas: dict[int, int] = {}
 programs: list = []
 
 def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
@@ -66,6 +67,7 @@ def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
     else: bar_view[arg0:arg0+arg1] = data
   elif cmd == RemoteCmd.MAP_SYSMEM:
     memview, paddrs = System.alloc_sysmem(arg0, vaddr=arg2, contiguous=bool(arg1))
+    if System.vfio is not None: sysmem_iovas[memview.addr] = paddrs[0]
     paddrs_bytes = struct.pack(f'<{len(paddrs) + 1}Q', memview.addr, *paddrs)
     conn.sendall(resp(len(paddrs_bytes)) + paddrs_bytes)
   elif cmd == RemoteCmd.SYSMEM_READ:
@@ -73,13 +75,19 @@ def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
   elif cmd == RemoteCmd.SYSMEM_WRITE:
     to_mv(arg0, arg1)[:] = conn.recv(arg1, socket.MSG_WAITALL)
   elif cmd == RemoteCmd.UNMAP_SYSMEM:
+    if (iova:=sysmem_iovas.pop(arg0, None)) is not None: System.vfio_unmap(iova, arg1)
     FileIOInterface.munmap(arg0, arg1)
     conn.sendall(resp())
   elif cmd == RemoteCmd.LOAD_PROG:
-    programs.append(Device["CPU"].runtime(TinyELF(conn.recv(arg0, socket.MSG_WAITALL), "hcq_submit", Target("CPU"), ())))
+    lib = conn.recv(arg0, socket.MSG_WAITALL)
+    links = {n:programs[i] for n,i in json.loads(conn.recv(arg1, socket.MSG_WAITALL)).items()}
+    programs.append(Device["CPU"].runtime(TinyELF(lib, "hcq_submit", Target("CPU"), ()), links=links))
     conn.sendall(resp(len(programs) - 1))
   elif cmd == RemoteCmd.EXEC_PROG:
-    et = programs[arg0](*struct.unpack(f'<{arg1}Q', conn.recv(arg1 * 8, socket.MSG_WAITALL)), wait=bool(arg2))
+    kernargs = ctypes.create_string_buffer(conn.recv(arg1, socket.MSG_WAITALL))
+    st = time.perf_counter()
+    programs[arg0].fxn(kernargs)
+    et = time.perf_counter() - st
     if (mock:=sys.modules.get("test.mockgpu.mockgpu")) is not None: # native programs bypass the mock's memoryview hooks
       for d in mock.drivers: d._emulate_execute()
     if arg2: conn.sendall(resp(int(et * 1e9)))

tinygrad_repo/extra/thunder/amd/fa.py

@@ -65,16 +65,20 @@ def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:
 
 def _fa_native_grads(dq:UOp, dk:UOp, dv:UOp) -> tuple[UOp, UOp, UOp]|None:
   def unwrap_partial(x:UOp) -> UOp|None:
-    expected = (Ops.CAST, Ops.REDUCE, Ops.PERMUTE, Ops.CAST, Ops.RESHAPE, Ops.AFTER)
-    for op in expected:
-      if x.op is not op: return None
-      if op is not Ops.AFTER: x = x.src[0]
-    return x
+    for expected in ((Ops.CAST, Ops.REDUCE, Ops.PERMUTE, Ops.CAST, Ops.RESHAPE, Ops.AFTER),
+                     (Ops.CAST, Ops.RESHAPE, Ops.CAST, Ops.RESHAPE, Ops.AFTER)):
+      cur = x
+      for op in expected:
+        if cur.op is not op: break
+        if op is not Ops.AFTER: cur = cur.src[0]
+      else: return cur
+    return None
   dq_native, dk_partial, dv_partial = dq.base, unwrap_partial(dk), unwrap_partial(dv)
   if dq_native.op is not Ops.AFTER or dk_partial is None or dv_partial is None: return None
   B, N, H, D, H_KV = dq.shape[0], dq.shape[1], dq.shape[2], dq.shape[3], dk.shape[2]
-  heads_per_wg = 2 if D == 128 and (H // H_KV) % 2 == 0 else 1
-  partials = (H // H_KV) // heads_per_wg
+  if dk_partial.shape[0] % B: return None
+  partials = dk_partial.shape[0] // B
+  if partials == 0 or (H // H_KV) % partials: return None
   if dq_native.shape != (B, H, N, D) or dk_partial.shape != (B * partials, N, H_KV, D) or dv_partial.shape != dk_partial.shape: return None
   return dq_native, dk_partial, dv_partial
 
@@ -126,6 +130,12 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
   return _sharded_empty(ref.shape, ref, axis)
 
+def _bwd_heads_per_wg(D:int, group_size:int, window:int) -> int:
+  if D == 64:
+    assert group_size == 8, "D=64 backward is specialized for GPT-OSS GQA8"
+    return 8 if window else 2
+  return 2 if D == 128 and group_size % 2 == 0 else 1
+
 def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
   def grad(dou:UOp, ker:UOp) -> tuple:
     do = Tensor(dou, device=dou.device)
@@ -137,7 +147,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
 
     dq = _sharded_empty((B, H, N, D), xq, axis=shard_axis_t)
     GROUP_SIZE = H_local // H_KV_local
-    HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
+    HEADS_PER_WG = _bwd_heads_per_wg(D, GROUP_SIZE, window)
     dk_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xk, axis=shard_axis)
     dv_partial = _sharded_empty((B * GROUP_SIZE // HEADS_PER_WG, N, H_KV, D), xv, axis=shard_axis)
 
@@ -260,7 +270,7 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
   if not getenv("NO_HIPCC"):
     lib = bytearray(lib)
     rodata_off = next(sh.header.sh_offset for sh in elf_loader(bytes(lib))[1] if sh.name == ".rodata")
-    struct.pack_into('<I', lib, rodata_off, 160000)
+    struct.pack_into('<I', lib, rodata_off, 0 if D == 64 else 160000)
     lib = bytes(lib)
 
   return UOp(Ops.PROGRAM,
@@ -268,14 +278,19 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
 
 @functools.cache
 def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
-  code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
+  source = ("fa_bwd_causal_d64_window.cpp" if window else "fa_bwd_causal_d64.cpp") if D == 64 else "fa_bwd_causal.cpp"
+  code = (pathlib.Path(__file__).parent / source).read_text()
   compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
                   f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
 
   BLOCK_SIZE_KV = 256
   GROUP_SIZE = H // H_KV
-  HEADS_PER_WG = 2 if D == 128 and GROUP_SIZE % 2 == 0 else 1
+  HEADS_PER_WG = _bwd_heads_per_wg(D, GROUP_SIZE, window)
   NUM_WARPS = 4
+  if D == 64: compile_args += [f"-DBWD_HEADS_PER_WG={HEADS_PER_WG}", f"-DBWD_BLOCK_SIZE_KV={BLOCK_SIZE_KV}", f"-DNUM_WARPS={NUM_WARPS}"]
+  relax_dq_sync = bool(getenv("GPTOSS_FA_BWD_RELAX_DQ_SYNC", 0) and arch == "gfx950" and
+                       (B, N, H, H_KV, D, window) == (2, 8192, 64, 8, 64, 0))
+  if relax_dq_sync: compile_args.append("-DGPTOSS_FA_BWD_RELAX_DQ_SYNC=1")
   NUM_THREADS = 64 * NUM_WARPS
   gsz = (H // HEADS_PER_WG, N // BLOCK_SIZE_KV, B)
   lsz = (NUM_THREADS, 1, 1)
@@ -287,7 +302,7 @@ def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_ve
   estimates = Estimates(ops=5*B*H*N*N*D, lds=mem, mem=mem)
   sink = UOp.sink(dq.base, dk.base, dv.base, do.base, q.base, k.base, v.base, l_vec.base, delta_vec.base,
                   threadIdx_x, blockIdx_x, blockIdx_y, blockIdx_z,
-                  arg=KernelInfo(name="custom_fa_backward", estimates=estimates))
+                  arg=KernelInfo(name="custom_fa_backward_d64_full_rs1" if relax_dq_sync else "custom_fa_backward", estimates=estimates))
 
   lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
   if not getenv("NO_HIPCC"):

tinygrad_repo/pyproject.toml

@@ -30,7 +30,6 @@ packages = [
   'tinygrad.nn',
   'tinygrad.renderer',
   'tinygrad.renderer.amd',
-  'tinygrad.renderer.isa',
   'tinygrad.runtime',
   'tinygrad.runtime.autogen',
   'tinygrad.runtime.autogen.am',
@@ -59,10 +58,10 @@ tinygrad = ["py.typed"]
 # triton = ["triton-nightly>=2.1.0.dev20231014192330"]
 linting = [
   "pylint",
-  "mypy==1.19.1",
+  "mypy==2.3.1",
   "typing-extensions",
   "pre-commit",
-  "ruff==0.14.10",
+  "ruff==0.16.4",
   "numpy",
   "typeguard",
 ]
@@ -76,7 +75,7 @@ testing_minimal = [
   "pytest-xdist",
   "pytest-split",
   "hypothesis>=6.148.9",
-  "z3-solver<4.15.4",  # 4.15.4 has a segfault when creating many z3.Context()
+  "z3-solver==5.1.0.0",
 ]
 testing_unit = ["tinygrad[testing_minimal]", "tqdm", "safetensors", "tabulate", "openai", "gguf>=0.18", "capstone"]
 testing = [

tinygrad_repo/spec/tinyspec.tex

@@ -59,7 +59,7 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \end{tabular}
 
 \smallskip
-\texttt{ParamArg} contains slot, dtype, value bounds, alignment, name, addrspace, device, volatility, optional image shape, and an optional bound device buffer (absent for \op{Alloc}s). \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}. The flat size $n$ is stated as a concrete weakint \op{Const} in src (an empty \op{Stack} for a scalar), followed by a \texttt{DEVICE} \op{Range} for multi-device storage.
+\texttt{ParamArg} contains slot, dtype, value bounds, name, addrspace, device, volatility, optional image shape, and an optional bound device buffer (absent for \op{Alloc}s). \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}. The flat size $n$ is stated as a concrete weakint \op{Const} in src (an empty \op{Stack} for a scalar), followed by a \texttt{DEVICE} \op{Range} for multi-device storage.
 
 %% ============================================================
 \subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
@@ -158,22 +158,24 @@ Assign is \op{Store} followed by \op{After}: write the value, then return the bu
 \textbf{Arity} & \textbf{src} & \textbf{Op} & \textbf{Semantics} \\
 \midrule
 Unary & $(T,)$
-  & \op{Recip}
+  & \op{Cast}
+  & Convert to target dtype (specified in arg). \\
+& & \op{Recip}
   & $1/x$ \\
 & & \op{Trunc}
   & $\mathrm{trunc}(x)$: round toward zero. \\
-& & \op{Cast}
-  & Convert to target dtype (specified in arg). \\
   \\[4pt]
 Binary & $(A, B)$
   & \op{Add}, \op{Mul}, \op{Max}
   & $a+b$, $a \cdot b$, $\max(a,b)$ \\
-& & \op{Floordiv}, \op{Floormod}
+& & \op{Div}, \op{Mod}
   & $\lfloor a/b \rfloor$, $a-b\lfloor a/b \rfloor$ \\
-& & \op{CmpLt}, \op{CmpNe}
-  & $[a < b]$, $[a \ne b]$ \\
-& & \op{Xor}, \op{Or}, \op{And}, \op{Shr}, \op{Shl}
-  & $a \oplus b$, $a \mid b$, $a \mathbin{\&} b$, $a \gg b$, $a \ll b$ \\[4pt]
+& & \op{CmpLt}, \op{CmpEq}
+  & $[a < b]$, $[a = b]$ \\
+& & \op{Xor}, \op{Or}, \op{And}
+  & $a \oplus b$, $a \mid b$, $a \mathbin{\&} b$ \\
+& & \op{Shl}, \op{Shr}
+  & $a \ll b$, $a \gg b$ \\[4pt]
 Ternary & $(P, A, B)$
   & \op{Where}
   & $A[\mathbf{i}]$ if $P[\mathbf{i}] \ne 0$, else $B[\mathbf{i}]$ \\
@@ -182,6 +184,10 @@ Ternary & $(P, A, B)$
 \end{tabular}
 
 \medskip
+\op{CmpEq} is primitive elementwise equality, returning bool (as do \op{CmpLt} and \op{CmpNe}). For floating-point inputs, equality is false if either operand is NaN, and $+0 = -0$. \op{CmpNe} is its logical complement, including for NaNs.
+
+\medskip
+\noindent\begin{minipage}{\linewidth}
 \textbf{Decomposed elementwise ops} --- defined in terms of the primitives above.
 
 \smallskip
@@ -191,14 +197,14 @@ Ternary & $(P, A, B)$
 \midrule
 \op{Neg} & \op{Mul}($A$, $-1$) & $-x$ \\
 \op{Sub} & \op{Add}($A$, \op{Neg}($B$)) & $a - b$ \\
-\op{Div} & \op{Mul}($A$, \op{Recip}($B$)) & $a / b$ \\
-\op{Cdiv} & $q+c$ & Integer division toward zero. \\
-\op{Cmod} & $r-bc$ & Remainder with dividend's sign (or zero). \\
+\op{Fdiv} & \op{Mul}($A$, \op{Recip}($B$)) & Floating-point division $a / b$. \\
+\op{Truncdiv} & $\operatorname{sgn}(a)\operatorname{sgn}(b)\lfloor |a|/|b| \rfloor$ & Integer division toward zero. \\
+\op{Truncmod} & $a - b \cdot \op{Truncdiv}(a,b)$ & Remainder with dividend's sign (or zero). \\
 \op{CmpGt} & \op{CmpLt}($B$, $A$) & $[a > b]$ \\
-\op{CmpGe} & \op{CmpNe}(\op{CmpLt}($A$, $B$),\, $1$) & $[a \ge b]$ \\
-\op{CmpLe} & \op{CmpNe}(\op{CmpLt}($B$, $A$),\, $1$) & $[a \le b]$ \\
-\op{CmpEq} & \op{CmpNe}(\op{CmpNe}($A$, $B$),\, $1$) & $[a = b]$ \\
-\op{Not} & \op{CmpNe}($A$, $1$) & $\lnot a$ \\[4pt]
+\op{CmpGe} & \op{CmpEq}(\op{CmpLt}($A$, $B$),\, $0$) & $[a \ge b]$ \\
+\op{CmpLe} & \op{CmpEq}(\op{CmpLt}($B$, $A$),\, $0$) & $[a \le b]$ \\
+\op{CmpNe} & \op{CmpEq}(\op{CmpEq}($A$, $B$),\, $0$) & $[a \ne b]$ \\
+\op{Not} & \op{CmpEq}($A$, $0$) & $\lnot a$ \\[4pt]
 \op{Exp2} & polynomial approx + \op{Mul}, \op{Add} & $2^x$ \\
 \op{Log2} & exponent extract + polynomial approx & $\log_2 x$ \\
 \op{Sin} & argument reduction + polynomial approx & $\sin x$ \\
@@ -208,6 +214,7 @@ Ternary & $(P, A, B)$
 \op{Threefry} & 5 rounds of add-rotate-xor (ARX) & Threefry 2x32 PRNG \\
 \bottomrule
 \end{tabular}
+\end{minipage}
 
 %% ============================================================
 \subsection*{{\color{markerorange}Marker Ops} \normalfont\small--- identity on data}
@@ -271,7 +278,7 @@ ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$
 \op{Mul} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min,\max]$ of products \\
 \op{Max} & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\max(a,b),\, \max(A,B)]$ \\
 Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
-\op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
+\shortstack[l]{\op{CmpLt}, \op{CmpEq},\\\op{CmpNe}} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
 \op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
 \op{Call} & void & --- & first non-null src device & --- \\
 \op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\

tinygrad_repo/test/amd/hw/helpers.py

@@ -210,64 +210,22 @@ def run_program_emu(instructions: list, n_lanes: int = 1) -> WaveState:
 
 def run_program_hw(instructions: list, n_lanes: int = 1) -> WaveState:
   """Run instructions on real AMD hardware via HIPCompiler and the AMD runtime."""
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
-  from tinygrad.helpers import Target
+  from tinygrad.device import Device, Buffer
+  from tinygrad.dtype import dtypes, AddrSpace
+  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.engine.realize import run_linear
 
   dev = Device["AMD"]
-  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
-
   prologue, epilogue = get_prologue_epilogue(n_lanes)
-  code = assemble(prologue + instructions + epilogue)
-
-  byte_str = ', '.join(f'0x{b:02x}' for b in code)
-  asm_src = f""".text
-.globl test
-.p2align 8
-.type test,@function
-test:
-.byte {byte_str}
-
-.rodata
-.p2align 6
-.amdhsa_kernel test
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 96
-  .amdhsa_wavefront_size32 1
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_kernarg_size 8
-  .amdhsa_group_segment_fixed_size 65536
-  .amdhsa_private_segment_fixed_size 65536
-  .amdhsa_enable_private_segment 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: test
-    .symbol: test.kd
-    .kernarg_segment_size: 8
-    .group_segment_fixed_size: 65536
-    .private_segment_fixed_size: 65536
-    .kernarg_segment_align: 8
-    .wavefront_size: 32
-    .sgpr_count: 96
-    .vgpr_count: 256
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata
-"""
-
-  lib = compiler.compile(asm_src)
-  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
-
   buf_sz = _out_bytes(n_lanes)
   out_gpu = Buffer(dev.device, buf_sz, preallocate=True)
   assert out_gpu._buf % 16 == 0, f"buffer not 16-byte aligned: 0x{out_gpu._buf:x}"
-  prg(out_gpu._buf, global_size=(1, 1, 1), local_size=(n_lanes, 1, 1), wait=True)
+  lds = UOp.placeholder((65536,), dtypes.uint8, 0, AddrSpace.LOCAL)
+  scratch = UOp.alloc((512,), dtypes.uint8, name="scratch", addrspace=AddrSpace.GLOBAL)
+  sink = UOp.sink(UOp.param(0, dtypes.uint32, (buf_sz // 4,)), UOp.special(n_lanes, "lidx0"), lds, scratch, arg=KernelInfo("test"))
+  body = UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in prologue + instructions + epilogue))
+  prg = UOp(Ops.PROGRAM, src=(sink, body))
+  run_linear(UOp(Ops.LINEAR, src=(prg.call(UOp.from_buffer(out_gpu, dtypes.uint32)),)), wait=True)
 
   out_buf = out_gpu.as_memoryview()
 

tinygrad_repo/test/amd/hw/test_cdna_ds.py

@@ -0,0 +1,18 @@
+import unittest
+import tinygrad.runtime.autogen.amd.cdna.ins as cdna
+from test.amd.hw.test_cdna_vop3 import run_cdna_lanes
+
+class TestCDNADSSwizzle(unittest.TestCase):
+  def test_modes(self):
+    # the source lane of lanes 0-7 for each mode of the ISA pseudocode
+    for offset, src in {0x041f: [1, 0, 3, 2, 5, 4, 7, 6],                  # bitmask: xor 1
+                        0x00a0: [5]*8,                                     # bitmask: or 5, broadcast
+                        0x801b: [3, 2, 1, 0, 7, 6, 5, 4],                  # quad permute: reversed
+                        0xc020: [1, 2, 3, 4, 5, 6, 7, 8],                  # rotate by 1
+                        0xe000: [0, 16, 8, 24, 4, 20, 12, 28]}.items():    # fft: bit reversed
+      with self.subTest(offset=hex(offset)):
+        out = run_cdna_lanes([cdna.ds_swizzle_b32(vdst=cdna.v[2], addr=cdna.v[255], offset0=offset & 255, offset1=offset >> 8), cdna.s_waitcnt(0)])
+        self.assertEqual(out[:8], src)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/amd/hw/test_cdna_sdwa.py

@@ -10,6 +10,7 @@ class TestCDNASDWA(unittest.TestCase):
       cdna.v_mov_b32_e32(cdna.v[0], cdna.s[0]),
       cdna.v_mov_b32_e32(cdna.v[13], 1),
       cdna.v_add_co_u32_e32(cdna.v[0], cdna.SDWA, cdna.v[13], vsrc0=cdna.v[0], dst_sel=6, src0_sel=6),
+      cdna.s_nop(1),
       cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO),
       cdna.v_lshlrev_b32_e32(cdna.v[2], 31, cdna.v[2]),
       cdna.v_or_b32_e32(cdna.v[2], cdna.v[2], cdna.v[0]),

tinygrad_repo/test/amd/hw/test_cdna_vop3.py

@@ -5,98 +5,59 @@ gfx950 hardware when USE_HW=1.
 """
 import ctypes, struct, unittest
 import tinygrad.runtime.autogen.amd.cdna.ins as cdna
-from tinygrad.helpers import Target
 from tinygrad.renderer.amd.dsl import NULL
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
-LANES = 1
-
-def _code(instructions: list, out_reg: int = 2, out_addr: int | None = None) -> bytes:
-  load_out_addr = [
-    cdna.s_mov_b32(cdna.s[92], out_addr & 0xffffffff),
-    cdna.s_mov_b32(cdna.s[93], out_addr >> 32),
-  ] if out_addr is not None else [
-    cdna.s_load_dwordx2(cdna.s[92:93], cdna.s[80:81], 0, soffset=NULL),
-    cdna.s_waitcnt(0),
-  ]
-  return assemble([
+def _code(instructions: list, out_reg: int = 2) -> list:
+  return [
     cdna.s_mov_b32(cdna.s[80], cdna.s[0]),
     cdna.s_mov_b32(cdna.s[81], cdna.s[1]),
     cdna.v_mov_b32_e32(cdna.v[255], cdna.v[0]),
     *instructions,
-    *load_out_addr,
+    cdna.s_load_dwordx2(cdna.s[92:93], cdna.s[80:81], 0, soffset=NULL, imm=1),
+    cdna.s_waitcnt(0),
     cdna.v_lshlrev_b32_e32(cdna.v[240], 2, cdna.v[255]),
     cdna.global_store_dword(addr=cdna.v[240], data=cdna.v[out_reg], saddr=cdna.s[92:93], offset=0),
     cdna.s_endpgm(),
-  ])
+  ]
 
-def _run_emu(instructions: list, out_reg: int = 2) -> int:
-  out_buf = (ctypes.c_uint32 * LANES)(*([0] * LANES))
+def _run_emu(instructions: list, out_reg: int = 2, lanes: int = 1) -> list[int]:
+  out_buf = (ctypes.c_uint32 * lanes)(*([0] * lanes))
   args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
-  code = _code(instructions, out_reg)
+  code = assemble(_code(instructions, out_reg))
   kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
-  result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, LANES, 1, 1, ctypes.addressof(args),
+  result = run_asm(ctypes.addressof(kernel_buf), len(code), 1, 1, 1, lanes, 1, 1, ctypes.addressof(args),
                    0x19c | (128 << 15), 0x10000, arch="cdna")
   assert result == 0, f"run_asm failed with {result}"
-  return out_buf[0]
+  return list(out_buf)
 
-def _run_hw(instructions: list, out_reg: int = 2) -> int:
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+def _run_hw(instructions: list, out_reg: int = 2, lanes: int = 1) -> list[int]:
+  from tinygrad.device import Device, Buffer
+  from tinygrad.dtype import dtypes, AddrSpace
+  from tinygrad.uop.ops import UOp, Ops, KernelInfo
+  from tinygrad.engine.realize import run_linear
 
   dev = Device["AMD"]
   if dev.arch != "gfx950": raise unittest.SkipTest("requires gfx950 hardware")
-  out_gpu = Buffer(dev.device, LANES * 4, preallocate=True)
-  code = _code(instructions, out_reg, out_gpu._buf)
-  byte_str = ", ".join(f"0x{b:02x}" for b in code)
-  asm_src = f""".text
-.globl test
-.p2align 8
-.type test,@function
-test:
-.byte {byte_str}
-
-.rodata
-.p2align 6
-.amdhsa_kernel test
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 96
-  .amdhsa_accum_offset 256
-  .amdhsa_kernarg_size 0
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: test
-    .symbol: test.kd
-    .kernarg_segment_size: 0
-    .group_segment_fixed_size: 0
-    .private_segment_fixed_size: 0
-    .kernarg_segment_align: 8
-    .wavefront_size: 64
-    .sgpr_count: 96
-    .vgpr_count: 256
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata
-"""
-  prg = dev.runtime(TinyELF(HIPCompiler(dev.arch).compile(asm_src), "test", Target("AMD", arch=dev.arch), ()))
-  prg(global_size=(1, 1, 1), local_size=(LANES, 1, 1), wait=True)
+  out_gpu = Buffer(dev.device, lanes * 4, preallocate=True)
+  lds = UOp.placeholder((65536,), dtypes.uint8, 0, AddrSpace.LOCAL)
+  sink = UOp.sink(UOp.param(0, dtypes.uint32, (lanes,)), UOp.special(lanes, "lidx0"), lds, arg=KernelInfo("test"))
+  body = UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in _code(instructions, out_reg)))
+  prg = UOp(Ops.PROGRAM, src=(sink, body))
+  run_linear(UOp(Ops.LINEAR, src=(prg.call(UOp.from_buffer(out_gpu, dtypes.uint32)),)), wait=True)
   out = out_gpu.as_memoryview()
-  return struct.unpack("<I", out)[0]
+  return list(struct.unpack(f"<{lanes}I", out))
 
-def run_cdna(instructions: list, out_reg: int = 2) -> int:
-  emu = _run_emu(instructions, out_reg)
+def run_cdna_lanes(instructions: list, out_reg: int = 2, lanes: int = 64) -> list[int]:
+  emu = _run_emu(instructions, out_reg, lanes)
   if not USE_HW: return emu
-  hw = _run_hw(instructions, out_reg)
-  if emu != hw: raise AssertionError(f"Emulator vs Hardware mismatch: emu=0x{emu:08x} hw=0x{hw:08x}")
+  hw = _run_hw(instructions, out_reg, lanes)
+  if emu != hw: raise AssertionError(f"Emulator vs Hardware mismatch: emu={[hex(x) for x in emu]} hw={[hex(x) for x in hw]}")
   return hw
 
+def run_cdna(instructions: list, out_reg: int = 2) -> int: return run_cdna_lanes(instructions, out_reg, 1)[0]
+
 class TestCDNAVOP3(unittest.TestCase):
   def test_cvt_pk_fp8_f32_preserves_upper_half(self):
     """V_CVT_PK_FP8_F32 with OPSEL[3]=0 writes only D[15:0]."""
@@ -126,3 +87,27 @@ class TestCDNAVOP3(unittest.TestCase):
           cdna.v_cvt_pk_bf8_f32(cdna.v[2], cdna.v[0], cdna.v[1]),
         ])
         self.assertEqual(out, 0xdead3c00 | expected)
+
+# CDNA has no RDNA true16: v128+ in a 16-bit e32 operand is that register, not the hi half of v[n-128]
+class TestCDNATrue16(unittest.TestCase):
+  def test_v_add_f16_e32(self):
+    out = run_cdna([
+      cdna.v_mov_b32_e32(cdna.v[130], 0x4000),  # 2.0
+      cdna.v_mov_b32_e32(cdna.v[131], 0x4200),  # 3.0
+      cdna.v_add_f16_e32(cdna.v[132], cdna.v[130], cdna.v[131]),
+      cdna.v_mov_b32_e32(cdna.v[2], cdna.v[132]),
+    ])
+    self.assertEqual(out & 0xffff, 0x4500)  # 5.0
+
+  def test_v_cmp_lt_f16_e32(self):
+    out = run_cdna([
+      cdna.v_mov_b32_e32(cdna.v[130], 0x4000),  # 2.0
+      cdna.v_mov_b32_e32(cdna.v[131], 0x4200),  # 3.0
+      cdna.v_cmp_lt_f16_e32(cdna.v[130], cdna.v[131]),
+      cdna.s_nop(1),
+      cdna.v_mov_b32_e32(cdna.v[2], cdna.VCC_LO),
+    ])
+    self.assertEqual(out & 1, 1)
+
+if __name__ == "__main__":
+  unittest.main()

tinygrad_repo/test/amd/hw/test_dpp.py

@@ -6,14 +6,13 @@ when USE_HW=1.
 """
 import ctypes, unittest
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-from tinygrad.helpers import Target
 from test.amd.hw.helpers import USE_HW, assemble
 from test.mockgpu.amd.emu import run_asm
 
 WAVE64 = 64
 
-def _wave64_code(instructions: list, out_reg: int = 1) -> bytes:
-  return assemble([
+def _wave64_code(instructions: list, out_reg: int = 1) -> list:
+  return [
     s_mov_b32(s[80], s[0]),
     s_mov_b32(s[81], s[1]),
     v_mov_b32_e32(v[255], v[0]),
@@ -23,12 +22,12 @@ def _wave64_code(instructions: list, out_reg: int = 1) -> bytes:
     v_lshlrev_b32_e32(v[240], 2, v[255]),
     global_store_b32(addr=v[240], data=v[out_reg], saddr=s[92:93], offset=0),
     s_endpgm(),
-  ])
+  ]
 
 def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
   out_buf = (ctypes.c_uint32 * WAVE64)(*([0] * WAVE64))
   args = (ctypes.c_uint64 * 1)(ctypes.addressof(out_buf))
-  code = _wave64_code(instructions, out_reg)
+  code = assemble(_wave64_code(instructions, out_reg))
   kernel_buf = (ctypes.c_char * len(code)).from_buffer_copy(code)
   rsrc2 = 0x19c | (128 << 15)
   scratch_size = 0x10000
@@ -37,55 +36,19 @@ def _run_wave64_emu(instructions: list, out_reg: int = 1) -> list[int]:
   return list(out_buf)
 
 def _run_wave64_hw(instructions: list, out_reg: int = 1) -> list[int]:
-  from tinygrad.device import Device, TinyELF, Buffer
-  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+  from tinygrad.device import Device, Buffer
+  from tinygrad.dtype import dtypes, AddrSpace
+  from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+  from tinygrad.engine.realize import run_linear
 
   dev = Device["AMD"]
-  compiler = HIPCompiler(dev.arch)  # type: ignore[attr-defined]
-  code = _wave64_code(instructions, out_reg)
-  byte_str = ', '.join(f'0x{b:02x}' for b in code)
-  asm_src = f""".text
-.globl test
-.p2align 8
-.type test,@function
-test:
-.byte {byte_str}
-
-.rodata
-.p2align 6
-.amdhsa_kernel test
-  .amdhsa_next_free_vgpr 256
-  .amdhsa_next_free_sgpr 96
-  .amdhsa_user_sgpr_kernarg_segment_ptr 1
-  .amdhsa_kernarg_size 8
-  .amdhsa_group_segment_fixed_size 65536
-  .amdhsa_private_segment_fixed_size 65536
-  .amdhsa_enable_private_segment 1
-.end_amdhsa_kernel
-
-.amdgpu_metadata
----
-amdhsa.version:
-  - 1
-  - 0
-amdhsa.kernels:
-  - .name: test
-    .symbol: test.kd
-    .kernarg_segment_size: 8
-    .group_segment_fixed_size: 65536
-    .private_segment_fixed_size: 65536
-    .kernarg_segment_align: 8
-    .wavefront_size: 64
-    .sgpr_count: 96
-    .vgpr_count: 256
-    .max_flat_workgroup_size: 1024
-...
-.end_amdgpu_metadata
-"""
-  lib = compiler.compile(asm_src)
-  prg = dev.runtime(TinyELF(lib, "test", Target("AMD", arch=dev.arch), ()))
   out_gpu = Buffer(dev.device, WAVE64 * 4, preallocate=True)
-  prg(out_gpu._buf, global_size=(1, 1, 1), local_size=(WAVE64, 1, 1), wait=True)
+  lds = UOp.placeholder((65536,), dtypes.uint8, 0, AddrSpace.LOCAL)
+  sink = UOp.sink(UOp.param(0, dtypes.uint32, (WAVE64,)), UOp.special(WAVE64, "lidx0"),
+                  lds.end(UOp.range(WAVE64, 0, AxisType.WARP)), arg=KernelInfo("test"))
+  body = UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in _wave64_code(instructions, out_reg)))
+  prg = UOp(Ops.PROGRAM, src=(sink, body))
+  run_linear(UOp(Ops.LINEAR, src=(prg.call(UOp.from_buffer(out_gpu, dtypes.uint32)),)), wait=True)
   out = out_gpu.as_memoryview()
   return [int.from_bytes(out[i*4:(i+1)*4], 'little') for i in range(WAVE64)]
 

tinygrad_repo/test/amd/hw/test_ds.py

@@ -607,6 +607,16 @@ class TestDSRegisterWidth(unittest.TestCase):
     self.assertEqual(st.vgpr[0][1], 0xDEADBEEF)
     self.assertEqual(st.vgpr[0][2], 0x11111111, "v2 should be untouched")
 
+  def test_ds_load_inactive_lane(self):
+    st = run_program([
+      v_mov_b32_e32(v[0], 0x40000000), # set lds base addr to 1GB
+      v_mov_b32_e32(v[1], 99),
+      s_mov_b32(EXEC_LO, 0), # all lanes are exec 0
+      ds_load_b64(addr=v[0], vdst=v[1:2]),
+      s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+      s_mov_b32(EXEC_LO, 1),
+    ])
+    self.assertEqual(st.vgpr[0][1], 99)
 
 class TestDS2AddrStride64(unittest.TestCase):
   """Tests for DS_*_2ADDR_STRIDE64 (offset * 256 for B32, offset * 512 for B64)."""
@@ -812,7 +822,7 @@ class TestDsPermute(unittest.TestCase):
   def test_ds_bpermute_b32_identity(self):
     """DS_BPERMUTE_B32 with identity permutation (each lane reads from itself)."""
     instructions = [
-      v_mov_b32_e32(v[0], 0),  # addr = 0 (read from lane 0)
+      v_mov_b32_e32(v[0], 128),  # addr = 128 / 4 % 32 = 0 (read from lane 0)
       v_mov_b32_e32(v[1], 0xCAFEBABE),  # data in lane 0
       ds_bpermute_b32(v[2], v[0], v[1]),
       s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),

tinygrad_repo/test/amd/hw/test_sop.py

@@ -234,6 +234,7 @@ class TestBranch(unittest.TestCase):
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.vgpr[0][0], 42, "Branch should NOT be taken (VCC_LO is 1)")
 
+  @unittest.skip("TODO: fix forward branch lifting")
   def test_cbranch_vccnz_branches_on_vcc_lo(self):
     """S_CBRANCH_VCCNZ branches when VCC_LO is non-zero."""
     instructions = [

tinygrad_repo/test/amd/hw/test_vop1.py

@@ -15,6 +15,7 @@ class TestMov(unittest.TestCase):
     instructions = [
       s_mov_b32(s[0], 42),
       v_mov_b32_e32(v[0], s[0]),
+      v_nop_e32()
     ]
     st = run_program(instructions, n_lanes=1)
     self.assertEqual(st.vgpr[0][0], 42)

tinygrad_repo/test/amd/hw/test_vop3.py

@@ -1184,7 +1184,7 @@ class TestF64Ops(unittest.TestCase):
     # run_program with USE_HW=1 will verify exact bit match with hardware
     st = run_program(instructions, n_lanes=1)
     result_bits = st.vgpr[0][6] | (st.vgpr[0][7] << 32)
-    self.assertNotEqual(result_bits, 0, "Result should not be zero")
+    self.assertAlmostEqual(i642f(result_bits), a*b+c, delta=math.ulp(a*b))
 
 
 class TestMad64More(unittest.TestCase):

tinygrad_repo/test/amd/test_asm_kernel.py

@@ -14,6 +14,7 @@ import tinygrad.runtime.autogen.amd.rdna3.ins as r3
 import tinygrad.runtime.autogen.amd.rdna4.ins as r4
 from tinygrad.renderer.amd.dsl import s, v, NULL, Reg, Inst
 from test.amd.helpers import TARGET_TO_ARCH
+from test.mockgpu.amd.emu import RELIFT
 from extra.gemm.amd_asm_matmul import Kernel
 
 # small pattern matcher converting CALL to INS
@@ -204,14 +205,10 @@ def custom_data_deps(A:UOp) -> UOp:
   sink = UOp.sink(A.base, threads, arg=KernelInfo("custom_data_deps"))
   return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
 
-# import contextvar to use it
-from test.mockgpu.amd.emu import ASM_CALL # noqa: F401
-
 @unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestAsmKernel(unittest.TestCase):
   def setUp(self):
     self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
-    self.enterContext(Context(ASM_CALL=1))
 
   def test_simple(self):
     if self.arch != "rdna3": self.skipTest("only rdna3")
@@ -256,8 +253,9 @@ class TestAsmKernel(unittest.TestCase):
     a.realize()
     self.assertTrue((a.numpy() == 6.0).all())
 
-  @unittest.expectedFailure
+  @Context(RELIFT=1)
   def test_cfg_branch_diamond(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
     def cfg_kernel(out:UOp):
       k = Kernel()
       k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
@@ -282,6 +280,7 @@ class TestAsmKernel(unittest.TestCase):
     self.assertListEqual(out.tolist(), [2.0, 1.0, 3.0])
 
   def test_cfg_loop(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
     def cfg_kernel(out:UOp):
       k = Kernel()
       k.emit(s_load_b64(s[0:1], s[0:1], soffset=NULL))
@@ -317,7 +316,7 @@ class TestAsmKernel(unittest.TestCase):
 
   def test_readlane(self):
     if self.arch != "rdna3": self.skipTest("only rdna3")
-    def readlane_kernel(out:UOp):
+    def readlane_kernel(out:UOp, threads:int):
       insts = [
         s_load_b64(s[0:1], s[0:1], soffset=NULL),
         s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
@@ -328,10 +327,11 @@ class TestAsmKernel(unittest.TestCase):
         global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]),
         s_endpgm(),
       ]
-      sink = UOp.sink(out.base, UOp.special(64, "lidx0"), arg=KernelInfo("wave_sgprs"))
+      sink = UOp.sink(out.base, UOp.special(threads, "lidx0"), arg=KernelInfo(f"wave_sgprs_{threads}"))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
-    out = Tensor(np.full(64, 0xFFFFFFFF, dtype=np.uint32)).realize().custom_kernel(fxn=readlane_kernel)[0]
-    self.assertListEqual(out.tolist(), [0]*32 + [32]*32)
+    for t in [31, 55, 64]:
+      out = Tensor(np.full(t, 0xFFFFFFFF, dtype=np.uint32)).realize().custom_kernel(fxn=functools.partial(readlane_kernel, threads=t))[0]
+      self.assertListEqual(out.tolist(), [i//32*32 for i in range(t)])
 
   def test_plus_tensor(self):
     out = Tensor.arange(1, 4).clone() + Tensor.arange(4, 7).clone()
@@ -344,5 +344,21 @@ class TestAsmKernel(unittest.TestCase):
     out = a@b
     self.assertEqual(out.tolist(), (a.numpy()@b.numpy()).tolist())
 
+  def test_relift_simple(self):
+    if self.arch != "rdna3": self.skipTest("only rdna3")
+    def k(out:UOp):
+      insts = [
+        s_load_b64(s[0:1], s[0:1], soffset=NULL), s_waitcnt_lgkmcnt(sdst=NULL, simm16=0),
+        s_mov_b32(s[2], 2), s_add_u32(s[2], s[2], 2),
+        v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], s[2]),
+        global_store_b32(addr=v[0], data=v[1], saddr=s[0:1]), s_endpgm(),
+      ]
+      sink = UOp.sink(out.base, arg=KernelInfo("k"))
+      return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
+    out = Tensor.empty(1, dtype=dtypes.uint32).custom_kernel(fxn=k)[0]
+    with Context(RELIFT=max(RELIFT.value, 1)):
+      out.realize()
+    self.assertListEqual(out.tolist(), [4])
+
 if __name__ == "__main__":
   unittest.main()

tinygrad_repo/test/amd/test_pdf.py

@@ -44,9 +44,9 @@ class TestPcodePDF(unittest.TestCase):
       "if SCC then\nD0.i32 = 32'I(signext(SIMM16.i16))\nendif")
 
   def test_swizzle_spans_blocks_and_pages(self):
-    for arch in ('rdna3', 'rdna4'):
+    for arch, opcode in (('rdna3', 53), ('rdna4', 53), ('cdna', 61)):
       with self.subTest(arch=arch):
-        code = self.pcode[arch][('DS_SWIZZLE_B32', 53)]
+        code = self.pcode[arch][('DS_SWIZZLE_B32', opcode)]
         self.assertIn('} elsif (offset >= 0xc000) {', code)
         self.assertIn('thread_out[i+3]', code)
         self.assertIn('xor_mask = offset[14:10];', code)

tinygrad_repo/test/amd/test_sqtt_encoder.py

@@ -1,108 +0,0 @@
-#!/usr/bin/env python3
-"""Tests for SQTT encoder: verifies the emulator produces correct SQTT traces for known kernels.
-
-Run with: DEV=MOCK+AMD python -m pytest test/amd/test_sqtt_encoder.py -v
-"""
-import ctypes, unittest
-from tinygrad.helpers import Context
-from tinygrad.renderer.amd.sqtt import decode, LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST, InstOp
-from tinygrad.runtime.autogen.amd.rdna3.ins import *
-
-def _run_kernel(instructions: list, lx=1, ly=1, lz=1, gx=1, gy=1, gz=1, args_ptr=0) -> bytes:
-  """Assemble instructions, run on emulator with PROFILE=1, return the SQTT blob."""
-  from test.mockgpu.amd.emu import run_asm, sqtt_traces
-  code = b''.join(inst.to_bytes() for inst in instructions)
-  buf = (ctypes.c_char * len(code))(*code)
-  lib = ctypes.addressof(buf)
-  sqtt_traces.clear()
-  with Context(PROFILE=1):
-    run_asm(lib, len(code), gx, gy, gz, lx, ly, lz, args_ptr)
-  assert len(sqtt_traces) == 1, f"expected 1 trace, got {len(sqtt_traces)}"
-  return sqtt_traces.pop()
-
-class TestSQTTEncoder(unittest.TestCase):
-
-  def test_simple_salu(self):
-    """A simple s_mov + s_endpgm kernel emits SALU INST packet."""
-    blob = _run_kernel([s_mov_b32(s[0], 42), s_endpgm()])
-    packets = list(decode(blob))
-    inst_pkts = [p for p in packets if isinstance(p, INST)]
-    self.assertEqual(len(inst_pkts), 1)
-    self.assertEqual(inst_pkts[0].op, InstOp.SALU)
-
-  def test_valu_emits_valuinst(self):
-    """Regular VALU ops emit VALUINST packets."""
-    blob = _run_kernel([v_mov_b32_e32(v[0], 0), v_add_f32_e32(v[1], v[0], v[0]), s_endpgm()])
-    packets = list(decode(blob))
-    valu_pkts = [p for p in packets if isinstance(p, VALUINST)]
-    self.assertEqual(len(valu_pkts), 2)
-    # no INST packets for regular VALU
-    self.assertEqual(len([p for p in packets if isinstance(p, INST)]), 0)
-
-  def test_waitcnt_emits_immediate(self):
-    """s_waitcnt and s_nop emit IMMEDIATE packets."""
-    blob = _run_kernel([s_nop(simm16=0), s_waitcnt(simm16=0), s_endpgm()])
-    imm_pkts = [p for p in decode(blob) if isinstance(p, IMMEDIATE)]
-    self.assertEqual(len(imm_pkts), 2)  # s_nop + s_waitcnt
-
-  def test_endpgm_skipped(self):
-    """s_endpgm does not emit any packet."""
-    blob = _run_kernel([s_endpgm()])
-    packets = list(decode(blob))
-    self.assertEqual(len([p for p in packets if isinstance(p, INST)]), 0)
-    self.assertEqual(len([p for p in packets if isinstance(p, IMMEDIATE)]), 0)
-
-  def test_wave_lifecycle(self):
-    """Every WAVESTART has a matching WAVEEND."""
-    blob = _run_kernel([s_mov_b32(s[0], 0), s_endpgm()])
-    packets = list(decode(blob))
-    self.assertEqual(sum(1 for p in packets if isinstance(p, WAVESTART)), sum(1 for p in packets if isinstance(p, WAVEEND)))
-
-  def test_layout_header(self):
-    """First packet is LAYOUT_HEADER with layout=3."""
-    blob = _run_kernel([s_endpgm()])
-    packets = list(decode(blob))
-    self.assertIsInstance(packets[0], LAYOUT_HEADER)
-    self.assertEqual(packets[0].layout, 3)
-
-  def test_blob_32byte_aligned(self):
-    """SQTT blob is 32-byte aligned."""
-    blob = _run_kernel([s_mov_b32(s[0], 0), s_mov_b32(s[1], 1), s_endpgm()])
-    self.assertEqual(len(blob) % 32, 0)
-
-  def test_multiple_waves(self):
-    """Multiple wavefronts each get their own WAVESTART/WAVEEND."""
-    blob = _run_kernel([s_mov_b32(s[0], 0), s_endpgm()], lx=64)  # 64 threads = 2 waves (WAVE_SIZE=32)
-    packets = list(decode(blob))
-    self.assertEqual(sum(1 for p in packets if isinstance(p, WAVESTART)), 2)
-    self.assertEqual(sum(1 for p in packets if isinstance(p, WAVEEND)), 2)
-
-  def test_branch_taken_and_not_taken(self):
-    """A loop with s_cbranch_scc1 emits JUMP when taken, JUMP_NO on final iteration."""
-    # s[0] = 2; loop: s[0] -= 1; cmp s[0] != 0 (SCC=1 if true); cbranch_scc1 loop; endpgm
-    # iteration 1: s[0]=2→1, SCC=1 (1!=0), branch taken (JUMP)
-    # iteration 2: s[0]=1→0, SCC=0 (0==0), branch not taken (JUMP_NO)
-    blob = _run_kernel([s_mov_b32(s[0], 2), s_sub_u32(s[0], s[0], 1), s_cmp_lg_u32(s[0], 0), s_cbranch_scc1(simm16=-3), s_endpgm()])
-    inst_pkts = [p for p in decode(blob) if isinstance(p, INST)]
-    ops = [p.op for p in inst_pkts]
-    self.assertIn(InstOp.JUMP, ops)
-    self.assertIn(InstOp.JUMP_NO, ops)
-
-  def test_timestamps_monotonic(self):
-    """Timestamps are monotonically non-decreasing."""
-    blob = _run_kernel([s_mov_b32(s[0], 0), s_mov_b32(s[1], 1), s_mov_b32(s[2], 2), s_endpgm()])
-    times = [p._time for p in decode(blob)]
-    self.assertEqual(times, sorted(times))
-
-  def test_no_trace_without_profile(self):
-    """No SQTT trace is emitted when PROFILE=0."""
-    from test.mockgpu.amd.emu import run_asm, sqtt_traces
-    code = s_endpgm().to_bytes()
-    buf = (ctypes.c_char * len(code))(*code)
-    sqtt_traces.clear()
-    with Context(PROFILE=0):
-      run_asm(ctypes.addressof(buf), len(code), 1, 1, 1, 1, 1, 1, 0)
-    self.assertEqual(len(sqtt_traces), 0)
-
-if __name__ == "__main__":
-  unittest.main()

tinygrad_repo/test/device/amd/test_llm.py

@@ -2,15 +2,16 @@ import unittest
 from unittest.mock import patch
 import numpy as np
 from tinygrad import Tensor, UOp, dtypes, nn, function, Device, TinyJit
+from tinygrad.helpers import DEV
 from tinygrad.llm.kernels.amd import (Linear, amd_custom_kernels_supported, q8_quantize, flash_attention, gated_delta_prefill,
-                                      QUANT_SIZES, HALFWORD_QUANTS, _wmma_rdna4)
+                                      QUANT_SIZES, HALFWORD_QUANTS, SMALL_BLOCK_QUANTS, _is_rdna4, _is_cdna, Q4_K, IQ3_S, ExpertWeights)
 from tinygrad.llm.gguf import ggml_data_to_tensor
 from test.runtime.test_llm_quantized import QuantLinearMixin
 
 class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
   def test_wmma_architecture_selection(self):
     if not amd_custom_kernels_supported(device:=Tensor.empty(1).device): self.skipTest("RDNA3/4 required")
-    self.assertEqual(_wmma_rdna4(device), Device[device].target[0] == 12)
+    self.assertEqual(_is_rdna4(device), Device[device].target[0] == 12)
 
   def test_values_and_scales(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3 required")
@@ -58,6 +59,7 @@ class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
   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_q8_0_linear_small_blocks(self): self._test_quant_linear(8, 272, in_features=288, out_features=16, token_counts=(1, 16))
   def test_q6_linear_odd_blocks(self): self._test_quant_linear(14, 210, in_features=768, out_features=3, token_counts=(1, 3, 16))
   def test_q2_k_linear(self): self._test_quant_linear(10, 84, in_features=768, token_counts=(1, 3, 16))
   def test_q3_k_linear(self): self._test_quant_linear(11, 110, in_features=768, token_counts=(1, 3, 16))
@@ -86,7 +88,7 @@ class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
     for typ, size in QUANT_SIZES.items():
       with self.subTest(ggml_type=typ):
         packed = rng.integers(0, 256, (48*3, size), dtype=np.uint8)
-        blocks = packed.reshape(-1, 18) if typ == 20 else packed
+        blocks = packed.reshape(-1, SMALL_BLOCK_QUANTS[typ]*2) if typ in SMALL_BLOCK_QUANTS else packed
         offset = size-2 if typ in (11, 14) else 80 if typ == 10 else 0
         blocks[:, offset:offset+2] = np.resize(scales, len(blocks)).view(np.uint8).reshape(-1, 2)
         if typ in (10, 12, 13):
@@ -220,6 +222,30 @@ class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
         np.testing.assert_array_equal(out.numpy(), 32)
         np.testing.assert_array_equal(state.numpy(), 1)
 
+  def _test_moe_decode(self, typ:int):
+    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3/4 required")
+    rng = np.random.default_rng(5)
+    E, T, K, IN, OUT, size = 4, 1, 2, 256, 16, QUANT_SIZES[typ]
+    blocks = rng.integers(0, 256, (E*OUT*IN//256, size), dtype=np.uint8)
+    blocks[:, :2] = np.array([0.001], np.float16).view(np.uint8)
+    blocks[:, 2:4] = np.array([0.0001], np.float16).view(np.uint8)
+    decoded = ggml_data_to_tensor(Tensor(blocks.flatten()).contiguous().realize(), E*OUT*IN, typ).reshape(E, OUT, IN)
+    w = decoded.numpy()
+    act = rng.standard_normal((1, T, K, IN)).astype(np.float32)
+    sel = np.stack([rng.choice(E, K, replace=False) for _ in range(T)])[None].astype(np.int32)
+    ew = ExpertWeights(E, IN, OUT)
+    ew.weight = decoded.half()
+    ew.set_quantized()
+    self.assertEqual(ew.ggml_type, typ)
+    got = ew(Tensor(sel), Tensor(act)).numpy()[0]
+    for t in range(T):
+      for j in range(K):
+        ref = w[sel[0, t, j]] @ act[0, t, j]
+        np.testing.assert_allclose(got[t, j], ref, rtol=2e-2, atol=2e-2*np.abs(ref).max())
+
+  def test_moe_decode(self): self._test_moe_decode(Q4_K)
+  def test_moe_decode_iq3_s(self): self._test_moe_decode(IQ3_S)
+
   def test_gated_delta_random_scan(self):
     if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3/4 required")
     rng = np.random.default_rng(42)
@@ -342,7 +368,7 @@ class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
           np.testing.assert_allclose(run(q, kv, end).numpy(), expected, rtol=1e-3, atol=1e-3)
 
   def test_prefill_attention_random(self):
-    if not amd_custom_kernels_supported(Tensor.empty(1).device): self.skipTest("RDNA3/4 required")
+    if not amd_custom_kernels_supported(d:=Tensor.empty(1).device) or _is_cdna(d): self.skipTest("RDNA3/4 required")
     rng = np.random.default_rng(42)
     cases = ((2, 4, 2, 32, 64, 45, 64, False), (1, 3, 1, 64, 128, 83, 128, False),
              (1, 2, 1, 96, 256, 149, 192, False), (1, 4, 2, 17, 128, 128, 128, True), (1, 2, 1, 3, 64, 33, 64, True),
@@ -410,7 +436,8 @@ class TestQ8Quantize(QuantLinearMixin, unittest.TestCase):
     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")
+    if not amd_custom_kernels_supported(d:=Tensor.empty(1).device): self.skipTest("RDNA3/4 or CDNA3/4 required")
+    if _is_cdna(d) and DEV.interface.startswith("MOCK"): self.skipTest("too slow on the CDNA emulator")
     rng = np.random.default_rng(42)
     start_pos = 1718
     q = Tensor.zeros(1, 8, 32, 128)

tinygrad_repo/test/device/cl/test_ocl.py

@@ -1,19 +1,7 @@
 import unittest
-from unittest.mock import patch
 from tinygrad import Device
-from tinygrad.device import Buffer, TinyELF
-from tinygrad.helpers import Target
-from tinygrad.runtime.ops_cl import CLDevice, CLAllocator, CLCompiler
-
-@unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
-class TestCLCompileCache(unittest.TestCase):
-  def test_compile_cached(self):
-    device = Device[Device.DEFAULT]
-    src = "__kernel void cached_test(__global int* a) { a[0] = 1; }"
-    obj = TinyELF(src.encode(), "cached_test", Target("CL"), ())
-    device.runtime(obj)
-    with patch.object(CLCompiler, 'compile', side_effect=RuntimeError("compile should not be called on cache hit")):
-      device.runtime(obj)
+from tinygrad.device import Buffer
+from tinygrad.runtime.ops_cl import CLDevice, CLAllocator
 
 @unittest.skipUnless(Device.DEFAULT == "CL", "Runs only on OpenCL")
 class TestCLError(unittest.TestCase):
@@ -28,7 +16,7 @@ class TestCLError(unittest.TestCase):
   def test_invalid_kernel_name(self):
     device = Device[Device.DEFAULT]
     with self.assertRaises(RuntimeError) as err:
-      device.runtime(TinyELF(b"__kernel void test(__global int* a) { a[0] = 1; }", "", Target("CL"), ()))
+      device.function(b"__kernel void test(__global int* a) { a[0] = 1; }", "")
     assert str(err.exception) == "OpenCL Error -46: CL_INVALID_KERNEL_NAME"
 
   def test_unaligned_copy(self):

tinygrad_repo/test/device/cpu/test_cpu.py

@@ -3,7 +3,7 @@ from contextlib import redirect_stdout
 from tinygrad import Tensor, Device
 from tinygrad.helpers import Target
 from tinygrad.renderer.nir import LVPRenderer
-from tinygrad.renderer.isa.x86 import X86Renderer
+from tinygrad.renderer.x86 import X86Renderer
 from tinygrad.codegen import to_program
 
 @unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")

tinygrad_repo/test/device/cpu/test_custom_kernel.py

@@ -1,14 +1,13 @@
 import unittest
 from tinygrad import Tensor, UOp, Context, GlobalCounters
 from tinygrad.uop.ops import Ops, KernelInfo
-from tinygrad.renderer.cstyle import ClangRenderer
 from test.helpers import assert_kernel_count
 
 class TestCustomKernel(unittest.TestCase):
   @Context(DEV="CPU")
   def test_simple_from_source(self):
     a = Tensor.arange(4).clone().realize()
-    src = ClangRenderer.kernel_typedef + " test_src(int* restrict a) { a[0] = 1; }"
+    src = "struct args_t { int* a; }; void test_src(const struct args_t* args) { args->a[0] = 1; }"
     def custom_src_kernel(A:UOp, B:UOp) -> UOp:
       sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src),))
@@ -18,7 +17,7 @@ class TestCustomKernel(unittest.TestCase):
   @Context(DEV="CPU")
   def test_simple_from_source_alt(self):
     a = Tensor.arange(4).clone().realize()
-    src = ClangRenderer.kernel_typedef + " copy(int* restrict out, int* restrict in) { for (int i = 0; i < 4; i++) out[i] = in[i]; }"
+    src = "struct args_t { int* out; int* in; }; void copy(const struct args_t* args) { for (int i = 0; i < 4; i++) args->out[i] = args->in[i]; }"
     def custom_src_kernel(out:UOp, inp:UOp) -> UOp:
       sink = UOp.sink(out, inp, arg=KernelInfo(name="copy"))
       return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src),))

tinygrad_repo/test/device/dsp/test_quantize_onnx.py

@@ -115,9 +115,11 @@ class TestQuantizeOnnx(unittest.TestCase):
         unsigned_char128 hi128;
       };
     };
-    __attribute__((noinline)) void fxn(unsigned char* restrict __attribute__((align_value(128))) data0,
-                                       unsigned char* restrict __attribute__((align_value(128))) data1,
-                                       signed char* restrict __attribute__((align_value(128))) data2) {
+    struct fxn_args { unsigned char* data0; unsigned char* data1; signed char* data2; };
+    __attribute__((noinline)) void fxn(const struct fxn_args* args) {
+      unsigned char* restrict __attribute__((align_value(128))) data0 = args->data0;
+      unsigned char* restrict __attribute__((align_value(128))) data1 = args->data1;
+      signed char* restrict __attribute__((align_value(128))) data2 = args->data2;
       for (int ridx0 = 0; ridx0 < 512; ridx0++) {
         int alu0 = (ridx0<<9);
         for (int ridx1 = 0; ridx1 < 4; ridx1++) {

tinygrad_repo/test/device/metal/test_metal.py

@@ -1,11 +1,11 @@
 import unittest
 from tinygrad.device import CompileError, Device, BufferSpec
 if Device.DEFAULT=="METAL":
-  from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler
+  from tinygrad.runtime.ops_metal import MetalCompiler
 @unittest.skipIf(Device.DEFAULT!="METAL", "Metal support required")
 class TestMetal(unittest.TestCase):
   def test_alloc_oom(self):
-    device = MetalDevice("metal")
+    device = Device["METAL"]
     with self.assertRaises(MemoryError):
       device.allocator.alloc(10000000000000000000)
 
@@ -38,7 +38,7 @@ class TestMetal(unittest.TestCase):
     assert ret is not None
 
   def test_failed_newLibraryWithData(self):
-    device = MetalDevice("metal")
+    device = Device["METAL"]
     compiler = MetalCompiler()
     compiled = compiler.compile("""
 #include <metal_stdlib>

tinygrad_repo/test/device/qcom/test_qcom.py

@@ -40,7 +40,7 @@ class TestQCOM(unittest.TestCase):
   def test_dcache_flush(self):
     from tinygrad.runtime.ops_qcom import dcache_flush
     buf = (ctypes.c_uint8 * 64)()
-    dcache_flush().fxn(buf, 0)
+    dcache_flush()(ctypes.addressof(buf), 0)
 
 if __name__ == '__main__':
   unittest.main()

tinygrad_repo/test/device/webgpu/test_renderer_failures.py

@@ -13,6 +13,10 @@ class TestWGSLFailures(unittest.TestCase):
     ret = _setup_and_test_alu(Ops.MUL, 5.0, UOp.const(float("inf")).cast(dtypes.float32))
     self.assertEqual(ret[0], float("inf"))
 
+  def test_uint_sub_large_const(self):
+    x = Tensor([0xffffffff], dtype=dtypes.uint32).contiguous().realize()
+    self.assertEqual((x - 0x80000001).tolist(), [0x7ffffffe])
+
   # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
   def test_gated_load(self):
     a = UOp.param(0, dtypes.int, 4)

tinygrad_repo/test/external/external_benchmark_hip_compile.py

@@ -1,31 +0,0 @@
-import random, os
-from tinygrad.helpers import Timing
-from tinygrad.runtime.ops_hip import compile_hip, HIPDevice
-from tinygrad.runtime.ops_cl import compile_cl, CLDevice
-
-# OMP_NUM_THREADS=1 strace -tt -f -e trace=file python3 test/external/external_benchmark_hip_compile.py
-# AMD_COMGR_REDIRECT_LOGS=stdout AMD_COMGR_EMIT_VERBOSE_LOGS=1 python3 test/external/external_benchmark_hip_compile.py
-
-# issue is in https://github.com/ROCm-Developer-Tools/clr/
-
-if __name__ == "__main__":
-  HIPDevice()
-  CLDevice()
-
-  # warmup
-  name = "none"+str(random.randint(0, 1000000))
-  compile_cl.__wrapped__(f"void {name}() {{}}")
-  print("compile cl warmed up")
-  compile_hip.__wrapped__(f"void {name}() {{}}")
-  print("compile hip warmed up")
-
-  print("**** benchmark ****")
-  name = "none"+str(random.randint(0, 1000000))
-  # this uses AMD_COMGR_ACTION_COMPILE_SOURCE_TO_BC, then it links the lib on the next step
-  with Timing("compile cl:  "): compile_cl.__wrapped__(f"void {name}() {{}}")
-  # this uses AMD_COMGR_ACTION_COMPILE_SOURCE_WITH_DEVICE_LIBS_TO_BC, much slower
-  with Timing("compile hip: "): compile_hip.__wrapped__(f"void {name}() {{}}")
-  os._exit(0)
-
-
-

tinygrad_repo/test/external/external_benchmark_op_conv.py

@@ -95,5 +95,5 @@ gs = sorted(dedup([u for u in ast.toposort() if u.op is Ops.PARAM]), key=lambda
 bufs = [Buffer(ps.arg.device, g.max_numel() * g.dtype.itemsize).ensure_allocated() for g in gs]
 
 gsize, lsize = ps.arg.launch_dims({})
-t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, vals=ps.arg.vals({}), wait=True)
+t = rt(*[b._buf for b in bufs], global_size=gsize, local_size=lsize, wait=True)
 print(f"{t*1e6:.2f} us")

tinygrad_repo/test/external/external_cl_half_max.py

@@ -1,13 +0,0 @@
-from tinygrad.runtime.ops_cl import CLDevice, CLProgram, compile_cl
-
-if __name__ == "__main__":
-  dev = CLDevice()
-  lib = compile_cl("""
-#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void test(__global half *out, __global half *a, __global half *b) {
-  int gid = get_global_id(0);
-  out[gid] = max(a[gid], b[gid]);
-}
-""")
-  prg = CLProgram(dev, "test", lib)
-

tinygrad_repo/test/external/external_gpu_fail_osx.py

@@ -1,16 +0,0 @@
-# ugh, OS X OpenCL doesn't support half
-from tinygrad.runtime.ops_cl import CLDevice, CLProgram, CLCompiler
-
-src = """#pragma OPENCL EXTENSION cl_khr_fp16 : enable
-__kernel void max_half(__global half* data0, const __global half* data1) {
-  int gidx0 = get_group_id(0);
-  data0[gidx0] = max(data1[gidx0], (half)0.0);
-}"""
-
-if __name__ == "__main__":
-  dev = CLDevice()
-  print("created device")
-  lib = CLCompiler(dev, "test").compile(src)
-  print("created lib", len(lib))
-  prg = CLProgram(dev, "max_half", lib)
-  print("created prg")

tinygrad_repo/test/external/external_osx_profiling.py

@@ -1,41 +0,0 @@
-from tinygrad.runtime.ops_cl import CLProgram, CL, CLBuffer
-from tinygrad import dtypes
-import time
-
-N = 1000000
-a = CLBuffer(N, dtypes.float32)
-b = CLBuffer(N, dtypes.float32)
-c = CLBuffer(N, dtypes.float32)
-
-prg = CLProgram("test", """__kernel void test(__global float *a, __global float *b, __global float *c) {
-  int idx = get_global_id(0);
-  a[idx] = b[idx] + c[idx];
-}""")
-prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
-t1 = time.monotonic_ns()
-e1 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
-CL.synchronize()
-t2 = time.monotonic_ns()
-time.sleep(3)
-t3 = time.monotonic_ns()
-e2 = prg.clprgs[0](CL.cl_queue[0], [N,], None, a._buf, b._buf, c._buf)
-CL.synchronize()
-t4 = time.monotonic_ns()
-
-print(e1.profile.queued)
-print(e1.profile.submit)
-print(e1.profile.start)
-print(e1.profile.end)
-
-print(e1, e2)
-print(t2-t1, e1.profile.end - e1.profile.start)
-print(t4-t3, e2.profile.end - e2.profile.start)
-print(t3-t2, e2.profile.queued-e1.profile.end)
-print((t3-t2) / (e2.profile.start-e1.profile.end), "ratio")
-
-print("ratio since boot", t1/e1.profile.start)
-
-print(e1.profile.start)
-print(e1.profile.end)
-print(e2.profile.start)
-print(e2.profile.end)

tinygrad_repo/test/external/external_test_speed_llama.py

@@ -10,7 +10,7 @@ from tinygrad.helpers import Profiling
 
 class FakeProgram:
   def __init__(self, name:str, lib:bytes, *args, **kwargs): pass
-  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False, **kw): pass
+  def __call__(self, *args, global_size, local_size, wait=False, **kw): pass
 
 class FakeAllocator(Allocator[Compiled]):
   def _alloc(self, sz, options): return BufferStorage(None)

tinygrad_repo/test/external/fuzz_symbolic_div.py

@@ -16,7 +16,7 @@ def get_random_term(ranges, factors):
 def get_random_expr(ranges, factors):
   num_terms = random.randint(2,4)
   x = UOp.usum(*[get_random_term(ranges, factors) for _ in range(num_terms)])
-  return x.alu(random.choice([Ops.CDIV, Ops.CMOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
+  return x.alu(random.choice([Ops.TRUNCDIV, Ops.TRUNCMOD]), x.ufix(random.choice(factors)*random.choice([1, 1, 1, -1])))
 
 if __name__ == "__main__":
   skipped = 0

tinygrad_repo/test/external/process_replay/process_replay.py

@@ -38,6 +38,7 @@ SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE"
 # uncomment this to disable by default
 #SKIP_PROCESS_REPLAY = not ASSERT_DIFF and not ((k:="[p]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", ""))
 if REF == "master": SKIP_PROCESS_REPLAY = True
+if REF == "update_benchmark": SKIP_PROCESS_REPLAY = "[pr]" not in os.getenv("COMMIT_MESSAGE", "").lower()
 class ProcessReplayWarning(Warning): pass
 
 # *** replay the function and convert return values to string
@@ -53,7 +54,11 @@ def replay_to_program(p:UOp, ast:UOp, renderer:Renderer) -> tuple[str, str, tupl
   def to_str(ret:UOp) -> str:
     src = next(x.arg for x in ret.src if x.op is Ops.SOURCE)
     # PYTHON renderer pickles UOps, first unpickle and decode here
-    if device.startswith("PYTHON"): return "\n".join([str(x) for x in pickle.loads(base64.b64decode(src))])
+    # a nested PROGRAM prints its LINEAR, its SOURCE/BINARY are pickles of that LINEAR whose bytes depend on object identity
+    if device.startswith("PYTHON"):
+      uops = pickle.loads(base64.b64decode(src))
+      pickles = {u:u.replace(arg=None) for u in UOp.sink(*uops).toposort(enter_calls=True) if u.op in (Ops.SOURCE, Ops.BINARY)}
+      return "\n".join([str(x.substitute(pickles, enter_calls=True)) for x in uops])
     return src
   # properly color the name arg
   ast_repr = codecs.decode(str(input_ast), "unicode_escape")
@@ -64,12 +69,12 @@ replayers["do_to_program"] = replay_to_program
 
 # *** run replayers on captured rows and print diffs
 
-def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
+def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]], table_name=TABLE_NAME) -> None:
   if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
   if early_stop.is_set(): return None
   conn = db_connection()
   cur = conn.cursor()
-  cur.execute(f"SELECT val FROM '{TABLE_NAME}' LIMIT ? OFFSET ?", (PAGE_SIZE, offset))
+  cur.execute(f"SELECT val FROM '{table_name}' LIMIT ? OFFSET ?", (PAGE_SIZE, offset))
   changed = 0
   for row in cur.fetchall():
     if changed > MAX_DIFF_PCT:
@@ -100,28 +105,32 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
 
 # *** generic runner to map rows of a table to a function in parallel
 
-def _pmap(fxns:dict[str, Callable]) -> None:
+def _pmap(fxns:dict[str, Callable], table_name:str=TABLE_NAME) -> None:
   conn = db_connection()
   cur = conn.cursor()
-  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
+  try: row_count = cur.execute(f"select count(*) from '{table_name}'").fetchone()[0]
   except sqlite3.OperationalError:
-    raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
+    raise RuntimeError(f"{table_name} isn't accessible in master, did DB_VERSION change?")
   finally:
     cur.close()
 
   with ProcessPoolExecutor(multiprocessing.cpu_count(), mp_context=multiprocessing.get_context("spawn")) as pool:
     with tqdm(total=row_count) as bar:
-      for _ in pool.map(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(min(s, row_count-bar.n))
+      for _ in pool.map(functools.partial(diff, fxns=fxns, table_name=table_name), range(0, row_count, s:=min(PAGE_SIZE, row_count))):
+        bar.update(min(s, row_count-bar.n))
 
 # *** main loop
 
-if __name__ == "__main__":
+def main(fxns:dict[str, Callable], table_name:str=TABLE_NAME) -> None:
   if SKIP_PROCESS_REPLAY:
     logging.info("skipping process replay.")
     exit(0)
 
   logging.info(f"running process replay with {ASSERT_DIFF=}")
-  try: _pmap(replayers)
+  try: _pmap(fxns, table_name)
   except Exception as e:
     logging.info(f"process replay err: {e}")
     exit(int(ASSERT_DIFF))
+
+if __name__ == "__main__":
+  main(replayers)

tinygrad_repo/test/external/process_replay/reset.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
 from tinygrad.helpers import db_connection, VERSION
 cur = db_connection()
-cur.execute(f"drop table if exists process_replay_{VERSION}")
+for table in ("process_replay", "emu_process_replay"):
+  cur.execute(f"drop table if exists {table}_{VERSION}")

tinygrad_repo/test/helpers.py

@@ -1,6 +1,6 @@
-import os, sys, time, struct, functools, unittest
+import os, sys, gc, time, struct, functools, unittest, contextlib
 from dataclasses import replace
-from typing import Any, Callable
+from typing import Any, Callable, Iterator
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.device import Buffer
@@ -22,7 +22,7 @@ from tinygrad.runtime.ops_python import PythonRenderer
 def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp:
   if ren is None: ren = Renderer(Target())
   if sink.arg is None: sink = sink.replace(arg=KernelInfo())
-  return full_rewrite_to_sink(sink, ren, optimize=sink.tag is None)
+  return full_rewrite_to_sink(sink, ren, optimize=sink.arg is None or not sink.arg.optimized)
 
 def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]:
   """Extract linearized UOps from a sink. Test helper that only does linearization (no render)."""
@@ -66,6 +66,13 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
     raise KernelCountException(allowed, kernel_cnt)
   return linear, var_vals
 
+@contextlib.contextmanager
+def measure_mem() -> Iterator[Callable[[], int]]:
+  Device.default.synchronize()
+  gc.collect()
+  base = GlobalCounters.mem_used
+  yield lambda: (gc.collect(), GlobalCounters.mem_used - base)[1]
+
 def assert_kernel_count(expected:int):
   got = GlobalCounters.kernel_count
   if got != expected: raise KernelCountException(expected, got)
@@ -82,7 +89,7 @@ def jit_cache_count(linear:UOp) -> int:
   return sum(len(call.without_after.arg.aux.kernels) if call_is_hcq(call) else 1 for call in linear.src)
 
 def assert_jit_cache_len(fxn, expected_len):
-  linear = fxn.captured.linear if fxn.captured is not None else None
+  linear = fxn.captured._linear if fxn.captured is not None else None
   if linear is None or not linear.src:
     if expected_len != 0: raise KernelCountException(expected_len, 0)
     return
@@ -149,7 +156,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
   prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
   prog = dev.runtime(prg.to_elf())
   out_buf = Buffer("PYTHON", uop.dtype.itemsize, preallocate=True)
-  prog(out_buf._buf, *[b.buf for b in bufs], vals=vals)
+  prog(out_buf._buf, *[b.buf for b in bufs], *vals)
   return out_buf.as_memoryview().cast(uop.dtype.fmt or "").tolist()[0]
 
 def to_uops_list(u:list[UOp], ren=None) -> list[UOp]:

tinygrad_repo/test/mockgpu/am/amdriver.py

@@ -97,6 +97,7 @@ class AMDriver(VirtDriver):
 
   def open(self, name, flags, mode, virtfile): return virtfile.fdcls(self._alloc_fd())
 
+  def progress(self): self._emulate_execute()
   def _emulate_execute(self):
     if self._executing: return
     self._executing = True

tinygrad_repo/test/mockgpu/am/amgpu.py

@@ -1,6 +1,6 @@
 # mypy: ignore-errors
 from __future__ import annotations
-import ctypes, struct, functools, os, mmap
+import ctypes, struct, functools, os, mmap, tempfile
 from tinygrad.runtime.autogen.am import am, smu_14_0_2
 from tinygrad.runtime.autogen import libc
 from tinygrad.runtime.support.amd import AMDReg, import_asic_regs
@@ -8,6 +8,13 @@ from test.mockgpu.amd.amdgpu import AMDGPU
 
 VRAM_SIZE = 512 << 20
 
+def memfd(name:str) -> int:
+  if hasattr(os, "memfd_create"): return os.memfd_create(name)
+  # macos has no memfd_create, use an unlinked temp file
+  fd, path = tempfile.mkstemp(prefix=name)
+  os.unlink(path)
+  return fd
+
 IP_VERSIONS = {
   am.GC_HWIP: (12, 0, 0), am.SDMA0_HWIP: (7, 0, 0), am.MMHUB_HWIP: (4, 1, 0), am.NBIO_HWIP: (6, 3, 1),
   am.MP0_HWIP: (14, 0, 2), am.MP1_HWIP: (14, 0, 2), am.HDP_HWIP: (7, 0, 0), am.OSSSYS_HWIP: (7, 0, 0),
@@ -263,11 +270,11 @@ class MockMMIOInterface:
 class MockAMGPU(AMDGPU):
   def __init__(self, gpuid:int=0):
     super().__init__(gpuid)
-    self.vram_fd = libc.memfd_create(b"vram", libc.MFD_CLOEXEC)
+    self.vram_fd = memfd("vram")
     os.ftruncate(self.vram_fd, VRAM_SIZE)
     self.vram_addr = libc.mmap(0, VRAM_SIZE, mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, self.vram_fd, 0)
     self.vram = (ctypes.c_ubyte * VRAM_SIZE).from_address(self.vram_addr)
-    self.doorbell_fd = libc.memfd_create(b"doorbell", libc.MFD_CLOEXEC)
+    self.doorbell_fd = memfd("doorbell")
     os.ftruncate(self.doorbell_fd, 0x2000)
     self.arch = "rdna4"
     self._sysmem_map:dict[int,int] = {}

tinygrad_repo/test/mockgpu/amd/amddriver.py

@@ -171,6 +171,7 @@ class AMDDriver(VirtDriver):
       raise RuntimeError(f"unsupported kfd ioctl, {nr} {name}")
     return 0
 
+  def progress(self): self._emulate_execute()
   def _emulate_execute(self):
     if self._executing: return  # prevent re-entrancy
     self._executing = True

tinygrad_repo/test/mockgpu/amd/amdgpu.py

@@ -1,7 +1,6 @@
 import ctypes, time
 from dataclasses import replace
 from test.mockgpu.gpu import VirtGPU
-from test.mockgpu.helpers import PythonRemu
 from tinygrad.helpers import getbits, to_mv, getenv, DEV
 from tinygrad.runtime.support import c
 
@@ -41,8 +40,6 @@ WAIT_REG_MEM_FUNCTION_EQ  = 3 # ==
 WAIT_REG_MEM_FUNCTION_NEQ = 4 # !=
 WAIT_REG_MEM_FUNCTION_GEQ = 5 # >=
 
-remu = PythonRemu()
-
 def create_sdma_packets():
   # TODO: clean up this, if we want to keep it
   structs = {}
@@ -77,6 +74,7 @@ class PM4Executor(AMDQueue):
   def __init__(self, gpu, base, size, rptr, wptr):
     self.gpu = gpu
     self.ib_executor: PM4Executor|None = None
+    self.kernel_name, self.kernel_text_size  = "", 0
     super().__init__(base, size, rptr, wptr)
 
   def _next_dword(self):
@@ -109,6 +107,9 @@ class PM4Executor(AMDQueue):
       elif op == amd_gpu.PACKET3_DISPATCH_DIRECT: self._exec_dispatch_direct(n)
       elif op == amd_gpu.PACKET3_INDIRECT_BUFFER: self._exec_indirect_buffer(n)
       elif op == amd_gpu.PACKET3_EVENT_WRITE: self._exec_event_write(n)
+      elif op == amd_gpu.PACKET3_NOP:
+        data = b"".join(self._next_dword().to_bytes(4,"little") for _ in range(n+1))
+        self.kernel_name, self.kernel_text_size = data[:-4].rstrip(b"\x00").decode(), int.from_bytes(data[-4:], "little")
       else: raise RuntimeError(f"PM4: Unknown opcode: {op}")
     return (self.rptr[0] - prev_rptr) + executed_in_ib
 
@@ -200,8 +201,8 @@ class PM4Executor(AMDQueue):
       except KeyError: user_data.append(0)
 
     prg_sz = 0
-    for st,sz in self.gpu.mapped_ranges:
-      if st <= prg_addr < st+sz: prg_sz = sz - (prg_addr - st)
+    for st,sz in sorted(self.gpu.mapped_ranges):
+      if st <= prg_addr+prg_sz < st+sz: prg_sz = st+sz-prg_addr
 
     # Get scratch size from COMPUTE_TMPRING_SIZE register
     # WAVESIZE = ceildiv(lanes * size_per_thread, mem_alignment_size)
@@ -212,14 +213,10 @@ class PM4Executor(AMDQueue):
     wavesize = (tmpring_size >> 12) & 0x3FFF  # WAVESIZE field is bits 12:25
     scratch_size = wavesize * (16 if self.gpu.arch == "cdna" else 4)  # per-thread scratch size in bytes
 
-    assert prg_sz > 0, "Invalid prg ptr (not found in mapped ranges)"
-    # Pass valid memory ranges, rsrc2, scratch_size, arch, and user data registers to the emulator
-    remu.valid_mem_ranges = self.gpu.mapped_ranges
-    remu.rsrc2 = rsrc2
-    remu.scratch_size = scratch_size
-    remu.arch = self.gpu.arch
-    remu.user_data = user_data
-    err = remu.run_asm(prg_addr, prg_sz, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr)
+    assert 0 < self.kernel_text_size <= prg_sz, "Invalid prg ptr or text size"
+    from test.mockgpu.amd.emu import run_asm
+    err = run_asm(prg_addr, self.kernel_text_size, gl[0], gl[1], gl[2], lc[0], lc[1], lc[2], args_addr,
+                  rsrc2, scratch_size, self.gpu.arch, user_data, self.kernel_name)
     if err != 0: raise RuntimeError("remu does not support the new instruction introduced in this kernel")
 
   def _exec_indirect_buffer(self, n):
@@ -238,9 +235,6 @@ class PM4Executor(AMDQueue):
     event_dw = self._next_dword()
     match (event_dw & 0xFF): # event type
       case SQTT_EVENTS.THREAD_TRACE_FINISH:
-        # Get the most recent trace from the emulator (if available)
-        from test.mockgpu.amd.emu import sqtt_traces
-        blob = sqtt_traces.pop(0) if sqtt_traces else b''
         old_idx = self.gpu.regs.grbm_index
         for se in range(self.gpu.regs.n_se):
           self.gpu.regs.grbm_index = 0b011 << 29 | se << 16 # select se, broadcast sa and instance
@@ -249,13 +243,7 @@ class PM4Executor(AMDQueue):
             buf_addr = ((self.gpu.regs[regSQ_THREAD_TRACE_BUF0_SIZE]&0xf)<<32|self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE])<<12
           else:
             buf_addr = ((self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE_HI])<<32|self.gpu.regs[regSQ_THREAD_TRACE_BUF0_BASE_LO])<<12
-          # Use real trace blob for SE 0 (which has itrace enabled), empty blob for other SEs
-          se_blob = blob if se == 0 else b''
-
-          # Write blob to trace buffer
-          if se_blob: ctypes.memmove(buf_addr, se_blob, len(se_blob))
-          # RDNA3 has absolute address for wptr, RDNA4 has relative
-          wptr_val = (((buf_addr if MOCKGPU_ARCH == "rdna3" else 0) + len(se_blob)) // 32) & 0x1FFFFFFF
+          wptr_val = ((buf_addr if MOCKGPU_ARCH == "rdna3" else 0) // 32) & 0x1FFFFFFF
           self.gpu.regs[regSQ_THREAD_TRACE_WPTR] = wptr_val
         self.gpu.regs.grbm_index = old_idx
       case _: pass # NOTE: for now most events aren't emulated

tinygrad_repo/test/mockgpu/amd/call.py

@@ -1,136 +0,0 @@
-import ctypes, itertools
-from tinygrad.viz.serve import amd_decode, get_cfg, COND_TAKEN, COND_NOT_TAKEN
-from tinygrad.uop.ops import UOp, sint, Ops, KernelInfo, PatternMatcher, UPat, graph_rewrite, rewrite_group, uopfunc
-from tinygrad.codegen import to_program
-from tinygrad.device import Device
-from tinygrad.dtype import Invalid, AddrSpace, dtypes
-from tinygrad.helpers import Context, getenv, TracingKey, dedup, unwrap
-from tinygrad.runtime.autogen import hsa
-from tinygrad.renderer.amd.dsl import EXEC_LO, ttmp
-from test.mockgpu.amd.emu import _Ctx, _get_handler, _wave_size, _canonical_info, PC_LO_IDX, PC_HI_IDX, SGPR_COUNT, SCRATCH_STRIDE_IDX, F32_INLINE
-
-asm_call_counter = itertools.count(1)
-
-# this is meant to replace the old emulator
-@uopfunc
-def init_wave(wg:UOp, wave:UOp, sgpr:UOp, vgpr:UOp, lds:UOp, args_ptr:UOp, gx:int, gy:int, lx:int, ly:int, total_threads:int, wave_size:int,
-              lds_size:int, scratch_size:int, rsrc2:int, arch:str="rdna3", user_data_count:int=0, accvgpr:UOp|None=None):
-  # define ranges inside a wave
-  li = UOp.range((wave.eq(0)).where(max(lds_size//4, 1), 0), 2, dtype=dtypes.int)
-  si = UOp.range(SGPR_COUNT, 3, dtype=dtypes.int)
-  vi = UOp.range(256*wave_size, 4, dtype=dtypes.int)
-  # zero ALLOCs
-  zero_lds = lds.index(li).store(0).end(li)
-  zero_sgpr = sgpr.after(zero_lds).index(si).store(0).end(si)
-  zero_vgpr = vgpr.after(zero_sgpr).index(vi).store(0)
-  zero_agpr = unwrap(accvgpr).after(zero_sgpr).index(vi).store(0) if wave_size == 64 else UOp(Ops.NOOP)
-  clear_vgpr = UOp.group(zero_vgpr, zero_agpr).end(vi)
-  # set RANGE registers
-  gidx, gidy, gidz = wg%gx, (wg//gx)%gy, wg//(gx*gy)
-  n_lanes = (total_threads-wave*wave_size).minimum(wave_size)
-  initial:list[tuple[int, sint]] = [*((128+i, i) for i in range(65)), *((193+i, (-i-1)&0xFFFFFFFF) for i in range(16)), *F32_INLINE.items()]
-  if user_data_count: initial += [(i, UOp.variable(f"user_data_{i}", 0, dtypes.uint32.max, dtypes.uint32)) for i in range(user_data_count)]
-  else: initial += [(0, args_ptr.cast(dtypes.uint32)), (1, (args_ptr>>32).cast(dtypes.uint32))]
-  if arch == "rdna4": initial += [(ttmp[7].offset, (gidy&0xFFFF)|((gidz&0xFFFF)<<16)), (ttmp[9].offset, gidx)]
-  else:
-    sgpr_id = (rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_USER_SGPR_COUNT_SHIFT
-    for enabled, gid in [(hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_X, gidx),
-                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Y, gidy),
-                         (hsa.AMD_COMPUTE_PGM_RSRC_TWO_ENABLE_SGPR_WORKGROUP_ID_Z, gidz)]:
-      if rsrc2 & enabled:
-        initial.append((sgpr_id, gid))
-        sgpr_id += 1
-  initial += [(EXEC_LO.offset, ((UOp.const(1, dtypes.uint64)<<n_lanes.minimum(32).cast(dtypes.uint64))-1).cast(dtypes.uint32)),
-              (SCRATCH_STRIDE_IDX, scratch_size), (SGPR_COUNT-16+4, (wave&15)|((wave&3)<<4))]
-  if wave_size == 64: initial.append((EXEC_LO.offset+1,
-                                     ((UOp.const(1, dtypes.uint64)<<(n_lanes-32).maximum(0).cast(dtypes.uint64))-1).cast(dtypes.uint32)))
-  lane = UOp.range(wave_size, 5, dtype=dtypes.int)
-  tid = wave*wave_size+lane
-  init_vgpr = vgpr.after(clear_vgpr).index(lane.valid(tid<total_threads)).store(
-      (((tid//(lx*ly))<<20)|(((tid//lx)%ly)<<10)|(tid%lx)).cast(dtypes.uint32)).end(lane)
-  return UOp.sink(init_vgpr, *(sgpr.after(clear_vgpr).index(i).store(UOp.const(v, dtypes.uint32)) for i,v in dict(initial).items()))
-
-def pc_index(idx:int) -> UPat:
-  reg, null = UPat.const(idx).cast(), UPat.const(124).cast()
-  return UPat.any(reg, reg.ne(null).where(reg, UPat.const(Invalid)))
-
-def move_const_idxs(call:UOp) -> UOp|None:
-  idxs = dedup(u.src[1] for u in call.body.toposort() if u.op is Ops.INDEX and u.src[0].op is Ops.PARAM and u.src[0].arg.name == "vmem"
-               and u.src[1].op is Ops.CONST)
-  if not idxs: return None
-  rep = {idx:UOp.param(len(call.src)-1+i, idx.commit_dtype(), name=f"inst_{i}", addrspace=AddrSpace.ALU) for i,idx in enumerate(idxs)}
-  return call.replace(src=(call.body.substitute(rep, walk=True), *call.src[1:], *idxs))
-
-pm_asm_call = PatternMatcher([
-  # remove PC from CALL body
-  (UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.PARAM, name="buf").index(UPat.any(pc_index(PC_LO_IDX), pc_index(PC_HI_IDX))),), allow_any_len=True),
-   lambda buf: UOp(Ops.NOOP) if buf.arg.name == "sgpr" else None),
-  # move CONST outside CALL body
-  (UPat(Ops.CALL, src=(UPat(Ops.SINK),), allow_any_len=True, name="call"), move_const_idxs),
-])
-
-@rewrite_group(name=lambda *args,ret,**_: TracingKey(f"Lift {(k:=ret.src[0].arg).name}", (("lift", k.function_name),)))
-def lift(lib:int, lib_sz:int, gx:int, gy:int, gz:int, lx:int, ly:int, lz:int, rsrc2:int, scratch_size:int, arch:str="rdna3",
-         user_data_count:int=0, backend:str|None=None) -> UOp:
-  backend = getenv("ASM_CALL_BACKEND", "CPU") if backend is None else backend
-  # decode
-  lib_bytes = ctypes.string_at(lib, lib_sz)
-  insts = amd_decode(lib_bytes, arch)
-  cfg = get_cfg(insts)["data"]
-  # construct CALL graph
-  wave_size, total_threads = _wave_size(arch), lx*ly*lz
-  n_waves = (total_threads+wave_size-1)//wave_size
-  wg = UOp.range(gx*gy*gz, 0)
-  wave = UOp.range(n_waves, 1)
-  # alloc register and LDS buffers
-  sgpr = UOp.alloc((SGPR_COUNT*n_waves,), dtypes.uint32, 0, AddrSpace.REG)
-  vgpr = UOp.alloc((256*wave_size*n_waves,), dtypes.uint32, 1, AddrSpace.REG)
-  lds_size = ((rsrc2 & hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE) >> hsa.AMD_COMPUTE_PGM_RSRC_TWO_GRANULATED_LDS_SIZE_SHIFT)*512
-  lds = UOp.alloc((max(lds_size//4, 1),), dtypes.uint32, 3, AddrSpace.REG)
-  scratch = UOp.alloc((max(scratch_size*wave_size*n_waves, 1),), dtypes.uint8, 4, AddrSpace.REG)
-  accvgpr = UOp.alloc((256*wave_size*n_waves,), dtypes.uint32, 5, AddrSpace.REG) if wave_size == 64 else vgpr
-  args_ptr = UOp.variable("args_ptr", 0, dtypes.uint64.max, dtypes.uint64)
-  lib_addr = UOp.variable("lib", 0, dtypes.uint64.max, dtypes.uint64)
-  inst_addr = UOp.param(-1, dtypes.uint64, name="inst", addrspace=AddrSpace.ALU)
-  init = init_wave(wg, wave, sgpr.index(wave*SGPR_COUNT), vgpr.index(wave*256*wave_size), lds, args_ptr, gx, gy, lx, ly, total_threads, wave_size,
-                   lds_size, scratch_size, rsrc2, arch, user_data_count, *([accvgpr.index(wave*256*wave_size)] if wave_size == 64 else []))
-  user_vars = [u for u in init.body.toposort() if u.is_variable]
-  init = init.replace(src=(init.body.substitute({u:UOp.param(len(init.src)-1+i, u.dtype, name=u.arg.name, addrspace=AddrSpace.ALU)
-                                                 for i,u in enumerate(user_vars)}), *init.src[1:], *user_vars))
-  ctx = _Ctx(4, wave_size)
-  end = init.end(wave)
-  mem = {ctx.sgpr:(sgpr,SGPR_COUNT),ctx.vgpr:(vgpr,256*wave_size),ctx.vmem:(ctx.vmem,0),ctx.lds:(lds,0),ctx.scratch:(scratch,scratch_size*wave_size)}
-  if wave_size == 64: mem[ctx.accvgpr] = (accvgpr, 256*wave_size)
-  for block_idx, (block_pc, block) in enumerate(cfg["blocks"].items()):
-    wave = UOp.range(n_waves, lib_sz+block_idx)
-    afters = {b:(buf.shrink(((wave*stride, (wave+1)*stride),)) if stride else buf).after(end, wave) for b,(buf,stride) in mem.items()}
-    loop_path = cfg["paths"][block_pc].get(block_pc)
-    loop = UOp.loop(block_pc) if loop_path is not None else None
-    if loop is not None: afters = {b:arg.after(loop) for b,arg in afters.items()}
-    branch_cond:UOp|None = None
-    for off in block:
-      inst = insts[off]
-      inst_st = str(inst)
-      if inst_st.startswith("s_code_end"): continue
-      if inst_st.startswith(("s_getpc", "s_setpc")): raise AssertionError("getpc and setpc are not allowed in ASM_CALL")
-      ctx = _Ctx(inst.size(), _wave_size(arch), inst_addr=inst_addr)
-      sink = _get_handler(inst)(inst, ctx)
-      *_, canonical_name = _canonical_info(inst, ctx, lib_bytes[off:])
-      bufs = sorted((u for u in sink.toposort() if u.op is Ops.PARAM), key=lambda u: u.arg.slot)
-      args = [lib_addr+off if b is inst_addr else afters.get(b, b.after(loop) if loop is not None else b) for b in bufs]
-      if ctx.branch_cond is not None and loop_path is not None:
-        branch_cond = ctx.branch_cond.substitute(dict(zip(bufs, args)), walk=True)
-        continue
-      body = sink.substitute({b:b.param_like(i, name=b.arg.name) for i,b in enumerate(bufs)})
-      call = body.call(*args, name=canonical_name)
-      afters.update((b, arg.after(call)) for b, arg in zip(bufs, args) if b is not inst_addr)
-    if loop is not None:
-      assert branch_cond is not None and loop_path in (COND_TAKEN, COND_NOT_TAKEN)
-      if loop_path is COND_NOT_TAKEN: branch_cond = branch_cond.logical_not()
-      backedge = UOp.sink(*afters.values()).backedge(loop, branch_cond)
-      afters = {b:arg.after(backedge) for b,arg in afters.items()}
-    end = UOp.group(*afters.values()).end(wave)
-  sink = UOp.sink(end.end(wg), arg=KernelInfo(name=f"asm_call n{next(asm_call_counter)}", opts_to_apply=()))
-  sink = graph_rewrite(sink, pm_asm_call, name="pm_asm_call", bottom_up=True, enter_calls=True)
-  with Context(NOOPT=1, CHECK_OOB=0, TUPLE_ORDER=0, EMULATED_DTYPES="", CAPTURE_PROCESS_REPLAY=0):
-    return to_program(sink, Device[backend].renderer)

tinygrad_repo/test/mockgpu/amd/pcode.py

@@ -568,8 +568,9 @@ class Parser:
         if op == '*': return left * right
         return (left // right) if dtypes.is_int(left.dtype) else (left / right)
       case '**': return UOp(Ops.EXP2, src=(right.cast(left.dtype),)) if _single_value(left) == 2.0 else left
+      case '%': return left % right
 
-  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/'), ('**',)]
+  _PREC = [('||',), ('&&',), ('|',), ('^',), ('&',), ('==', '!=', '<>'), ('>=', '<=', '>', '<'), ('>>', '<<'), ('+', '-'), ('*', '/', '%'), ('**',)]
 
   def binop(self, prec: int) -> UOp:
     if prec >= len(self._PREC): return self.unary()
@@ -1449,7 +1450,7 @@ def parse_pcode(pcode: str, srcs: dict[str, UOp | int] | None = None, funcs: dic
   _, 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):
+    if var in ['D0', 'S0', 'SCC', 'VCC', 'EXEC', 'PC', 'RETURN_DATA', 'VDATA', 'SDATA'] and isinstance(val, UOp):
       if var in sliced and not any(re.match(rf'{var}\.\w+\s*=', l) for l in lines): continue
       for l in lines:
         if (m := re.match(rf'{var}\.(\w+(?:\[\w+\])?)', l)):

tinygrad_repo/test/mockgpu/amd/replay.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+# usage: replay.py [clang,llvm,amd] [asm]
+import functools, re, sys
+from dataclasses import replace
+from test.mockgpu.amd.emu import lift
+from tinygrad.codegen import to_program
+from tinygrad.device import Device
+from tinygrad.helpers import VERSION, Target
+from tinygrad.renderer import Renderer
+from tinygrad.renderer.cstyle import ClangRenderer, HIPRenderer
+from tinygrad.renderer.llvmir import CPULLVMRenderer
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.viz.serve import amd_decode, get_arch, get_cfg, get_elf_section, get_stdout, parse_branch
+
+def get_output(p:UOp, renderer:Renderer, asm:bool=False) -> str:
+  nodes = p.toposort(enter_calls=True)
+  parts = [x.arg for x in nodes if x.op is Ops.SOURCE]
+  for prg in (x for x in nodes if asm and x.op is Ops.PROGRAM):
+    binary = next(x.arg for x in prg.src if x.op is Ops.BINARY)
+    if renderer.target.device != "AMD":
+      disasm = get_stdout(lambda: renderer.compiler.disassemble(binary))
+      parts.append(re.sub(r"^.*:\s+file format", "<binary>: file format", disasm, flags=re.MULTILINE))
+      continue
+    text = get_elf_section(binary, ".text")
+    insts = amd_decode(text.content, get_arch(renderer.target.arch), text.header.sh_addr)
+    cfg = get_cfg(insts)
+    labels = {pc:f"%bb.{i}" for i,pc in enumerate(cfg["blocks"])}
+    blocks = []
+    for start, pcs in cfg["blocks"].items():
+      lines = [f"{labels[start][1:]}:"]
+      if successors:=cfg["paths"][start]: lines.append("  successors: " + ", ".join(labels[pc] for pc in successors))
+      for pc in pcs:
+        op = getattr(inst:=insts[pc], "op_name", "").lower()
+        if (offset:=parse_branch(inst, op)) is not None: lines.append(f"  {op} {labels[pc+inst.size()+offset]}")
+        else: lines.append(f"  {inst}")
+      blocks.append("\n".join(lines))
+    parts.append("\n\n".join(blocks))
+  return "\n".join(parts)
+
+def replay_asm(ast:UOp, code:bytes, arch:str, renderer:Renderer, dispatch:tuple, name:str,
+               backends:tuple[str, ...], asm:bool=False) -> tuple[str, str, tuple]:
+  compare = lift(code, arch, renderer, dispatch, name)
+  good_outs, compare_outs = [], []
+  for backend in backends:
+    target = Target("AMD", "HIP", {"rdna3":"gfx1100", "rdna4":"gfx1200", "cdna":"gfx950"}[arch]) if backend == "amd" else \
+      replace(Device["CPU"].renderer.target, renderer=backend.upper())
+    ren = {"clang":ClangRenderer, "llvm":CPULLVMRenderer, "amd":HIPRenderer}[backend](target)
+    good_outs.append(get_output(to_program(ast, ren), ren, asm))
+    compare_outs.append(get_output(to_program(compare, ren), ren, asm))
+  return "\n".join(good_outs), "\n".join(compare_outs), (name, backends, dispatch)
+
+if __name__ == "__main__":
+  from test.external.process_replay.process_replay import main
+  backends = tuple((sys.argv[1] if len(sys.argv) > 1 else "clang").split(","))
+  main({"lift":functools.partial(replay_asm, backends=backends, asm=bool(sys.argv[2:]))}, f"emu_process_replay_{VERSION}")

tinygrad_repo/test/mockgpu/amd/sqtt_enc.py

@@ -1,100 +0,0 @@
-# SQTT trace encoder for the emulator (the decoder lives in tinygrad/renderer/amd/sqtt.py).
-# run_asm emits packets inline as instructions execute; finished traces end up in emu.sqtt_traces.
-from __future__ import annotations
-from tinygrad.renderer.amd.dsl import Inst
-from tinygrad.renderer.amd.sqtt import (_build_decode_tables, PACKET_TYPES_RDNA3, PacketType, InstOp,
-                                        LAYOUT_HEADER, WAVESTART, WAVEEND, INST, IMMEDIATE, VALUINST)
-
-_NIB_COUNTS = {cls: nc for _, (cls, nc, *_) in _build_decode_tables(PACKET_TYPES_RDNA3)[0].items()}
-
-def _emit_nibbles(nibbles: list[int], pkt_cls: type[PacketType], **kwargs):
-  raw = pkt_cls.encoding.default
-  for k, v in kwargs.items(): raw = pkt_cls.__dict__[k].set(raw, v)
-  nibbles.extend((raw >> (i * 4)) & 0xF for i in range(_NIB_COUNTS[pkt_cls]))
-
-def make_encoder():
-  """Build an SQTT trace encoder for the emulator. Returns (emit, finish, finalize)."""
-  from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp as SOPPOp3
-  from tinygrad.runtime.autogen.amd.rdna4.enum import SOPPOp as SOPPOp4
-  from tinygrad.runtime.autogen.amd.rdna3 import ins as ir3
-  from tinygrad.runtime.autogen.amd.rdna4 import ins as ir4
-  from tinygrad.runtime.autogen.amd.cdna import ins as irc
-  import re
-
-  def _kinds(*names: str) -> tuple[type[Inst], ...]:
-    return tuple(getattr(m, n) for m in (ir3, ir4, irc) for n in names if hasattr(m, n))
-  _SOPP, _SMEM, _DS = _kinds('SOPP'), _kinds('SMEM'), _kinds('DS')
-  _GLOBAL, _FLAT, _SCRATCH = _kinds('GLOBAL', 'VGLOBAL'), _kinds('FLAT', 'VFLAT'), _kinds('SCRATCH', 'VSCRATCH')
-  _VALU = _kinds('VOP1', 'VOP2', 'VOP3', 'VOP3P', 'VOP3PX2', 'VOPC', 'VOPD', 'VOP3SD', 'VOP3_SDST', 'VOP1_SDST')
-
-  # SOPP classification sets
-  _SOPP_SKIP = {SOPPOp3.S_ENDPGM.value, SOPPOp3.S_ENDPGM_SAVED.value, SOPPOp3.S_ENDPGM_ORDERED_PS_DONE.value, SOPPOp3.S_DELAY_ALU.value}
-  _SOPP_IMMEDIATE = {SOPPOp3.S_NOP.value, SOPPOp3.S_CLAUSE.value, SOPPOp3.S_WAITCNT.value, SOPPOp3.S_WAITCNT_DEPCTR.value,
-                     SOPPOp3.S_WAIT_IDLE.value, SOPPOp3.S_WAIT_EVENT.value, SOPPOp3.S_SLEEP.value, SOPPOp3.S_SET_INST_PREFETCH_DISTANCE.value}
-  for _op in (SOPPOp4.S_WAIT_ALU, SOPPOp4.S_WAIT_LOADCNT, SOPPOp4.S_WAIT_STORECNT, SOPPOp4.S_WAIT_SAMPLECNT,
-              SOPPOp4.S_WAIT_BVHCNT, SOPPOp4.S_WAIT_EXPCNT, SOPPOp4.S_WAIT_DSCNT, SOPPOp4.S_WAIT_KMCNT,
-              SOPPOp4.S_WAIT_LOADCNT_DSCNT, SOPPOp4.S_WAIT_STORECNT_DSCNT):
-    _SOPP_IMMEDIATE.add(_op.value)
-  _SOPP_BARRIER = {SOPPOp3.S_BARRIER.value}
-  if hasattr(SOPPOp4, 'S_BARRIER_WAIT'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_WAIT.value)
-  if hasattr(SOPPOp4, 'S_BARRIER_LEAVE'): _SOPP_BARRIER.add(SOPPOp4.S_BARRIER_LEAVE.value)
-  _SOPP_BRANCH = {SOPPOp3.S_BRANCH.value, SOPPOp3.S_CBRANCH_SCC0.value, SOPPOp3.S_CBRANCH_SCC1.value,
-                  SOPPOp3.S_CBRANCH_VCCZ.value, SOPPOp3.S_CBRANCH_VCCNZ.value,
-                  SOPPOp3.S_CBRANCH_EXECZ.value, SOPPOp3.S_CBRANCH_EXECNZ.value}
-
-  # VALU sub-classification patterns
-  _VALUT_4_RE = re.compile(r'V_(EXP|LOG|RCP|RSQ|SQRT|SIN|COS|CEIL|FLOOR|TRUNC|RNDNE|FRACT|FREXP)_')
-  _VALUB_2_RE = re.compile(r'V_(LSHLREV|LSHRREV|ASHRREV)_(B|I)64')
-  _VALUB_4_RE = re.compile(r'V_MAD_(U|I)64')
-  _VALUB_16_RE = re.compile(r'V_\w+_F64')
-
-  def _valu_op(op_name: str) -> InstOp|None:
-    if 'CMPX' in op_name: return InstOp.VALU1_WR_EXEC
-    if _VALUB_2_RE.search(op_name): return InstOp.VALUB_2
-    if _VALUB_4_RE.search(op_name): return InstOp.VALUB_4
-    if _VALUB_16_RE.search(op_name): return InstOp.VALUB_16
-    if _VALUT_4_RE.search(op_name): return InstOp.VALUT_4
-    return None
-
-  def _mem_op(t: type[Inst], op_name: str) -> InstOp:
-    is_store = "STORE" in op_name
-    if issubclass(t, _DS): return InstOp.LDS_WR_2 if is_store else InstOp.LDS_RD
-    if issubclass(t, _GLOBAL): return InstOp.SGMEM_WR_2 if is_store else InstOp.SGMEM_RD_1
-    if issubclass(t, _FLAT) or issubclass(t, _SCRATCH): return InstOp.FLAT_WR_3 if is_store else InstOp.FLAT_RD_2
-    return InstOp.SALU
-
-  nibbles: list[int] = []
-  started: set[int] = set()
-  _emit_nibbles(nibbles, LAYOUT_HEADER, layout=3, sel_a=6)
-
-  def emit(wave_id: int, inst: Inst, branch_taken: bool|None):
-    """Emit an SQTT packet for one executed instruction."""
-    w = wave_id & 0x1F
-    if wave_id not in started:
-      _emit_nibbles(nibbles, WAVESTART, delta=1, simd=0, wgp=0, wave=w, id7=wave_id)
-      started.add(wave_id)
-    inst_type, inst_op, op_name = type(inst), inst.op.value if hasattr(inst, 'op') else 0, inst.op.name if hasattr(inst, 'op') else ""
-    if issubclass(inst_type, _SOPP):
-      if inst_op in _SOPP_SKIP: return
-      if inst_op in _SOPP_IMMEDIATE: _emit_nibbles(nibbles, IMMEDIATE, delta=1, wave=w)
-      elif inst_op in _SOPP_BARRIER: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.BARRIER)
-      elif inst_op in _SOPP_BRANCH: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.JUMP if branch_taken else InstOp.JUMP_NO)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SALU)
-    elif issubclass(inst_type, _VALU):
-      if (op := _valu_op(op_name)) is None: _emit_nibbles(nibbles, VALUINST, delta=1, wave=w)
-      else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=op)
-    elif issubclass(inst_type, _SMEM): _emit_nibbles(nibbles, INST, delta=1, wave=w, op=InstOp.SMEM_RD)
-    else: _emit_nibbles(nibbles, INST, delta=1, wave=w, op=_mem_op(inst_type, op_name))
-
-  def finish(wave_id: int):
-    """Emit WAVEEND for a completed wave."""
-    if wave_id in started: _emit_nibbles(nibbles, WAVEEND, delta=1, simd=0, wgp=0, wave=wave_id & 0x1F)
-
-  def finalize() -> bytes:
-    """Pad and return the encoded SQTT blob."""
-    while len(nibbles) % 2 != 0: nibbles.append(0)
-    nibbles.extend([0] * 32)
-    while len(nibbles) % 64 != 0: nibbles.append(0)
-    return bytes(nibbles[i] | ((nibbles[i + 1] if i + 1 < len(nibbles) else 0) << 4) for i in range(0, len(nibbles), 2))
-
-  return emit, finish, finalize

tinygrad_repo/test/mockgpu/cuda/cuda.py

@@ -62,6 +62,7 @@ def cuMemFree_v2(dptr): return free(dptr)
 def cuMemFreeHost(p): return free(p)
 def cuMemHostRegister_v2(p, size, flags): return cuda.CUDA_SUCCESS
 def cuMemHostUnregister(p): return cuda.CUDA_SUCCESS
+def cuMemHostGetDevicePointer_v2(pdptr, p, flags): return out(pdptr, p)
 
 def cuModuleLoadData(module, image):
   modules.append(ctypes.create_string_buffer(ctypes.string_at(image)))

tinygrad_repo/test/mockgpu/driver.py

@@ -37,3 +37,4 @@ class VirtDriver:
     self.tracked_addresses = []
   def track_address(self, staddr, enaddr, rcb, wcb): self.tracked_addresses.append((staddr, enaddr, rcb, wcb))
   def open(self, name, flags, mode, fdcls): raise NotImplementedError()
+  def progress(self): pass # runs the queues

tinygrad_repo/test/mockgpu/helpers.py

@@ -6,15 +6,3 @@ gpuocelot_lib = c.DLL("ocelot", "gpuocelot")
                     ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int)
 def ptx_run(source:bytes, n_args:int, args:c.POINTER[ctypes.c_void_p], blck_x:int, blck_y:int, blck_z:int,
             grid_x:int, grid_y:int, grid_z:int, shared_mem_size:int): pass
-
-class PythonRemu:
-  """Python RDNA3/RDNA4 emulator wrapper used by mockgpu."""
-  valid_mem_ranges: set[tuple[int, int]] = set()
-  rsrc2: int = 0x19c  # Default: USER_SGPR_COUNT=14, enable X and Y workgroup IDs
-  scratch_size: int = 0  # private_segment_fixed_size from kernel descriptor
-  arch: str = "rdna3"  # Architecture: rdna3 or rdna4
-  user_data: list[int] = []  # All COMPUTE_USER_DATA registers (loaded into s[0:N])
-
-  def run_asm(self, lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
-    from test.mockgpu.amd.emu import run_asm
-    return run_asm(lib, lib_sz, gx, gy, gz, lx, ly, lz, args_ptr, self.rsrc2, self.scratch_size, self.arch, self.user_data)

tinygrad_repo/test/mockgpu/hip/hip.py

@@ -0,0 +1,38 @@
+# the hip api on top of mock cuda, kernels run on the amd emulator
+import ctypes
+from types import SimpleNamespace
+from tinygrad.runtime.autogen import hip, libc
+from tinygrad.runtime.ops_amd import _amd_program_image
+from tinygrad.runtime.ops_hip import HIPIFY
+from tinygrad.runtime.support.elf import elf_loader
+from test.mockgpu.cuda import cuda as mc
+from test.mockgpu.amd.emu import run_asm
+
+programs:dict[int, tuple] = {} # a function handle is the address of its module's loaded image
+
+def run_kernel(func:int, args:int, size:int, gx:int, gy:int, gz:int, lx:int, ly:int, lz:int, smem:int) -> bool:
+  image, data, text_end = programs[func]
+  entry, entry_size = ctypes.addressof(image) + data.entry_point_offset, text_end - data.entry_point_offset
+  return run_asm(entry, entry_size, gx, gy, gz, lx, ly, lz, args, data.rsrc2, data.private_segment_size) == 0
+mc.run_kernel = run_kernel
+
+def hipModuleLoadData(module, image):
+  elf = libc.Elf64_Ehdr.from_address(image) # the section headers end the elf
+  gfx1100 = SimpleNamespace(iface=SimpleNamespace(props={"lds_size_in_kb": 64}), target=(11, 0, 0))
+  lib = ctypes.string_at(image, elf.e_shoff + elf.e_shnum * elf.e_shentsize)
+  data, img = _amd_program_image(gfx1100, lib)
+  text = next(sh.header for sh in elf_loader(lib)[1] if sh.name == ".text")
+  buf = ctypes.create_string_buffer(img, len(img))
+  programs[ctypes.addressof(buf)] = (buf, data, text.sh_addr + text.sh_size)
+  return mc.out(module, ctypes.addressof(buf))
+
+def hipGetDevicePropertiesR0600(prop, device_id):
+  prop[0].gcnArchName = b"gfx1100"
+  return hip.hipSuccess
+
+fns = {fn.__name__: getattr(mc, cu) for cu, fn in HIPIFY.items() if fn.__name__.startswith("hip")}
+fns |= dict(hipHostMalloc=mc.cuMemHostAlloc, hipModuleLoadData=hipModuleLoadData, hipGetDevicePropertiesR0600=hipGetDevicePropertiesR0600,
+  hipMemcpyAsync=lambda dst, src, size, kind, stream: mc.cuMemcpyAsync(dst, src, size, stream),
+  hipStreamWaitValue64=lambda stream, ptr, value, flags, mask: mc.cuStreamWaitValue64_v2(stream, ptr, value, flags))
+for name, fn in fns.items(): setattr(hip.dll, name, ctypes.CFUNCTYPE(getattr(hip, name).restype, *getattr(hip, name).argtypes)(fn))
+hip.dll._loaded_[hip.dll.nm] = hip.dll

tinygrad_repo/test/mockgpu/mockgpu.py

@@ -1,8 +1,10 @@
 import ctypes, time, os, builtins, fcntl, typing
-from tinygrad.helpers import DEV, dedup, to_tuple
-from tinygrad.device import Compiled
+from tinygrad.helpers import DEV, dedup, to_tuple, round_up
+from tinygrad.device import Compiled, Device
+from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, UPat, Ops, PatternMatcher
 from tinygrad.runtime.support.system import FileIOInterface, System
+from tinygrad.runtime.support.hcq2 import HWQueue
 from tinygrad.runtime.autogen import libc
 from test.mockgpu.nv.nvdriver import NVDriver
 from test.mockgpu.amd.amddriver import AMDDriver
@@ -21,10 +23,34 @@ setattr(libc.dll, "mockgpu_doorbell", mockgpu_doorbell)
 
 mocked = {t.device for t in DEV.value if t.interface.startswith("MOCK")}
 def hook_doorbells(s:UOp) -> UOp|None:
-  hooks = tuple(UOp.custom_function("mockgpu_doorbell").call(b.after(c).getaddr("CPU")) for c in s.toposort() if c.op is Ops.CALL
-                for b in c.src[1:] if b.op is Ops.ALLOC and "doorbell" in str(b.tag) and to_tuple(b.device)[0].split(":")[0] in mocked)
+  hooks = tuple(UOp.custom_function("mockgpu_doorbell").call(b.after(c).getaddr(Device[d].host)) for c in s.toposort() if c.op is Ops.CALL
+                for b in c.src[1:] if b.op is Ops.ALLOC and "doorbell" in str(b.arg.name) and (d:=to_tuple(b.device)[0]).split(":")[0] in mocked)
   return None if all(h in s.src for h in hooks) else s.replace(src=tuple(dedup(s.src + hooks)))
-Compiled.pm_lower = PatternMatcher([(UPat(Ops.SINK, name="s"), hook_doorbells)])
+Compiled.pm_lower += PatternMatcher([(UPat(Ops.SINK, name="s"), hook_doorbells)])
+
+@ctypes.CFUNCTYPE(None, ctypes.c_uint64)
+def mockgpu_progress(_:int): # compiled spin loops wait on the gpus, which run only when called
+  for d in drivers: d.progress()
+setattr(libc.dll, "mockgpu_progress", mockgpu_progress)
+
+def kick_gpus(body:UOp, loop:UOp, cond:UOp) -> UOp|None: # each iteration of a spin loop runs the gpus
+  kick = UOp.custom_function("mockgpu_progress").call(UOp.const(0, dtypes.uint64).after(body))
+  return None if kick in cond.src else body.backedge(loop, cond.after(kick))
+Compiled.pm_lower += PatternMatcher([(UPat(Ops.BACKEDGE, src=(UPat(name="body"), UPat(name="loop"), UPat(name="cond"))), kick_gpus)])
+
+# this links AMD emulator's kernel back to the original kernel
+def emit_kernel_info(ctx:HWQueue, prg:UOp) -> None:
+  from tinygrad.runtime.ops_amd import AMDComputeQueue, _amd_program_image
+  if isinstance(ctx, AMDComputeQueue):
+    from tinygrad.viz.serve import get_elf_section
+    name = prg.src[0].arg.name.encode() + b"\x00"
+    text = get_elf_section(prg.src[3].arg, ".text").header
+    text_size = text.sh_addr + text.sh_size - _amd_program_image(ctx.dev, prg.src[3].arg)[0].entry_point_offset
+    ctx.pkt3(ctx.pm4.PACKET3_NOP, *memoryview(name.ljust(round_up(len(name), 4), b"\x00")).cast("I"), text_size)
+
+HWQueue.q_rewrite = PatternMatcher([
+  (UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), allow_any_len=True), emit_kernel_info),
+]) + HWQueue.q_rewrite
 
 original_memoryview = builtins.memoryview
 class TrackedMemoryView:
@@ -44,6 +70,7 @@ class TrackedMemoryView:
     self.mv = self.mv.cast('B').cast(new_type, **kwargs)
     return self
 
+  def __buffer__(self, flags): return self.mv.__buffer__(flags)
   @property
   def nbytes(self): return self.mv.nbytes
   def __len__(self): return len(self.mv)

tinygrad_repo/test/mockgpu/nv/nvdriver.py

@@ -291,6 +291,7 @@ class NVDriver(VirtDriver):
         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
+  def progress(self): self._gpu_mmio_write(None, 0, None)
   def _gpu_mmio_write(self, mv, off, gpu):
     if self._executing: return  # prevent re-entrancy
     self._executing = True

tinygrad_repo/test/mockgpu/usb.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
-import ctypes, mmap, struct, sys
+import ctypes, mmap, struct
 from typing import Any, Callable
-from tinygrad.runtime.autogen import libusb
-if sys.platform != "win32": from tinygrad.runtime.autogen import libc
+from tinygrad.runtime.autogen import libusb, libc
 
 class MockUSB:
   def __init__(self, mem):

tinygrad_repo/test/null/test_autogen.py

@@ -1,10 +1,9 @@
 import ctypes, struct, subprocess, tempfile, unittest
-from tinygrad.helpers import OSX, WIN
+from tinygrad.helpers import OSX
 from tinygrad.runtime.support.c import DLL, record, Field
 from tinygrad.runtime.support import c
 from tinygrad.runtime.support.autogen import gen
 
-@unittest.skipIf(WIN, "doesn't compile on windows")
 class TestC(unittest.TestCase):
   def compile(self, src):
     with tempfile.NamedTemporaryFile(suffix=".so") as f:
@@ -322,7 +321,6 @@ class TestC(unittest.TestCase):
     self.assertEqual(result, 43)  # 42 + 1
 
 @unittest.skipIf(OSX and ('MTLCompiler' in DLL._loaded_ or 'llvm' in DLL._loaded_), "libclang can't be loaded after MTLCompiler or llvm on OSX")
-@unittest.skipIf(WIN, "doesn't compile on windows")
 class TestAutogen(unittest.TestCase):
   def run_gen(self, contents):
     with tempfile.NamedTemporaryFile(mode='w', suffix='.h') as f:

tinygrad_repo/test/null/test_call.py

@@ -122,7 +122,7 @@ class TestCallCodegen(unittest.TestCase):
     v = UOp.variable("external", 1, 8, dtype=dtypes.int)
     prg = to_program(p.index(0).store(v).sink(arg=KernelInfo("scalar")),
                      ClangRenderer(Target("CPU", arch="x86_64,x86-64")))
-    self.assertEqual(prg.arg.vars[0].arg.slot, 1)
+    self.assertEqual(prg.arg.params[1].arg.slot, 1)
     linear = UOp(Ops.LINEAR, src=(prg.call(p, v.bind(2)),))
     # Slot 1 in the enclosing call must not rewrite the already-compiled program's ABI slot 1.
     resolved = resolve_linear_call(linear.call(out, UOp.variable("other", 1, 8, dtype=dtypes.int).bind(3)))

tinygrad_repo/test/null/test_const_folding.py

@@ -19,7 +19,7 @@ class TestWeakConstFolding(unittest.TestCase):
     self.assertEqual((out.op, out.dtype, out.val), (Ops.CONST, dtypes.weakint, 2**41))
 
   def test_float_unaries(self):
-    for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIPROCAL):
+    for op in (Ops.SIN, Ops.LOG2, Ops.EXP2, Ops.SQRT, Ops.RECIP):
       out = UOp.const(4.0).alu(op).simplify()
       self.assertEqual((out.op, out.dtype), (Ops.CONST, dtypes.weakfloat))
 
@@ -32,7 +32,7 @@ class TestWeakConstFolding(unittest.TestCase):
     self.assertTrue(nan.ne(nan).simplify().val)
 
   def test_invalid_poison(self):
-    self.assertTrue(UOp.invalid().alu(Ops.CDIV, UOp.const(0)).simplify().is_invalid)
+    self.assertTrue(UOp.invalid().alu(Ops.TRUNCDIV, UOp.const(0)).simplify().is_invalid)
 
 class TestBitcastConstFolding(unittest.TestCase):
   def test_out_of_range_source_value(self):

tinygrad_repo/test/null/test_device.py

@@ -3,7 +3,7 @@ import unittest, os, subprocess
 from unittest.mock import patch
 from tinygrad import Tensor, dtypes
 from tinygrad.device import Device, Compiler, enumerate_devices_str
-from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, WIN, OSX, DEV
+from tinygrad.helpers import diskcache_get, diskcache_put, getenv, Context, Target, OSX, DEV
 from tinygrad.runtime.support.c import DLL
 
 class TestDevice(unittest.TestCase):
@@ -71,7 +71,6 @@ class TestDevice(unittest.TestCase):
     result = subprocess.run(['python3', '-c', 'import runpy, tinygrad; runpy.run_module("tinygrad.device")'], capture_output=True)
     self.assertEqual(result.returncode, 0, result.stderr.decode())
 
-  @unittest.skipIf(WIN, "skipping windows test") # TODO: subprocess causes memory violation?
   def test_env_overwrite_default_compiler(self):
     if Device.DEFAULT == "CPU":
       from tinygrad.runtime.support.compiler_cpu import ClangCompiler
@@ -103,7 +102,6 @@ class TestDevice(unittest.TestCase):
                         shell=True, check=True, env={**os.environ, "DEV": "AMD:HIP"})
     else: self.skipTest("only run on CPU/AMD")
 
-  @unittest.skipIf(WIN, "skipping windows test")
   def test_env_online(self):
     from tinygrad.runtime.support.compiler_cpu import ClangCompiler
     from tinygrad.runtime.support.compiler_llvm import CPULLVMCompiler

tinygrad_repo/test/null/test_dtype_weak.py

@@ -40,7 +40,7 @@ class TestWeakPromotion(unittest.TestCase):
     self.assertIs(denom.dtype, dtypes.weakfloat)  # the setup: the denominator expression itself is weak
     x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom)
     self.assertIs(y.dtype, dtypes.float32)
-    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
+    recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIP]
     self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
     with Context(DEFAULT_FLOAT=dtypes.float16):
       committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_commit_weak)

tinygrad_repo/test/null/test_encodings.py

@@ -3,14 +3,13 @@ from dataclasses import replace
 from tinygrad.helpers import Target
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.dtype import dtypes, DType
-from tinygrad.renderer.isa import Register
-from tinygrad.renderer.isa.x86 import X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm, def_reg
+from tinygrad.renderer.x86 import Register, X86Ops, X86Renderer, RBP, RDI, RSP, RSI, RAX, RDX, XMM, GPR, imm
 
-def ins(op, dt, src, tag=None): return UOp(Ops.INS, arg=(op, dt), src=src, tag=tag)
-# the operand width is carried by the register in the tag, not by the dtype
+# the operand width is carried by the register, not by the dtype
 def reg(r:Register, size:int) -> Register: return replace(r, size=size)
-# the element size of a memory operand comes from the dtype of its base
-def ptr(r:Register, dt:DType) -> UOp: return UOp(Ops.INS, arg=(X86Ops.DEFINE, dt), tag=(r,))
+def ins(op, dt, src, r:Register|None=None): return UOp(Ops.INS, arg=(op, dt) + ((r,) if r is not None else ()), src=src)
+def reg_placeholder(r:Register): return UOp(Ops.INS, arg=(X86Ops.REG, dtypes.uint64, r))
+def alloc(dt:DType, r:Register): return reg_placeholder(r)
 
 class TestEncodingsX86(unittest.TestCase):
   # NOTE: x86 supports a single displacement as memory address and index without base memory address
@@ -19,133 +18,136 @@ class TestEncodingsX86(unittest.TestCase):
 
   # displacement of 0 isn't emitted
   def test_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RDI, dtypes.int32), UOp(Ops.NOOP), imm(dtypes.int8, 0)), reg(RDI, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, (alloc(dtypes.int32, RDI), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), reg(RDI, 4))
     # mov edi, dword ptr [rdi]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 3F"))
 
   # rsp/r12 require a sib byte when used as base memory address
   def test_rsp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RSP, dtypes.int32), UOp(Ops.NOOP), imm(dtypes.int8, 0)), reg(RSP, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, (alloc(dtypes.int32, RSP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), reg(RSP, 4))
     # mov esp, dword ptr [rsp]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 24 24"))
 
   # rbp/r13 require a displacement when used as base memory address
   def test_rbp_base_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RBP, dtypes.int32), UOp(Ops.NOOP), imm(dtypes.int8, 0)), reg(RBP, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, (alloc(dtypes.int32, RBP), UOp(Ops.NOOP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), reg(RBP, 4))
     # mov ebp, dword ptr [rbp + 0]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 6D 00"))
 
   # test [base + index*scale]
   def test_base_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RAX, dtypes.int32), def_reg(RDX), imm(dtypes.int8, 0)), reg(RAX, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, (alloc(dtypes.int32, RAX), reg_placeholder(RDX), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), reg(RAX, 4))
     # mov eax, dword ptr [rax + rdx*4]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 04 90"))
 
   # rsp as index means no index
   def test_rsp_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RAX, dtypes.int32), def_reg(RSP), imm(dtypes.int8, 0)), reg(RAX, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, (alloc(dtypes.int32, RAX), reg_placeholder(RSP), imm(dtypes.int8, 0), imm(dtypes.uint8, 4)), reg(RAX, 4))
     # mov eax, dword ptr [rax]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 00"))
 
   # however r12 is a valid index
   def test_r12_index_address(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RAX, dtypes.int32), def_reg(GPR[12]), imm(dtypes.int8, 0)), reg(RAX, 4))
+    address = (alloc(dtypes.int32, RAX), reg_placeholder(GPR[12]), imm(dtypes.int8, 0), imm(dtypes.uint8, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, address, reg(RAX, 4))
     # mov eax, dword ptr [rax + r12*4]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("42 8B 04 A0"))
 
   # test [base + index*scale + 8bit disp]
   def test_complex_address_8bit_disp(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RDI, dtypes.int32), def_reg(RSI), imm(dtypes.int8, 10)), reg(RDI, 4))
+    address = (alloc(dtypes.int32, RDI), reg_placeholder(RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, address, reg(RDI, 4))
     # mov edi, dword ptr [rdi + rsi*4 + 0xa]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B 7C B7 0A"))
 
   # test [base + index*scale + 32bit disp]
   def test_complex_address_32bit_disp(self):
-    load = ins(X86Ops.MOV, dtypes.int32, (ptr(RDI, dtypes.int32), def_reg(RSI), imm(dtypes.int32, 10000)), reg(RDI, 4))
+    address = (alloc(dtypes.int32, RDI), reg_placeholder(RSI), imm(dtypes.int32, 10000), imm(dtypes.uint8, 4))
+    load = ins(X86Ops.MOV, dtypes.int32, address, reg(RDI, 4))
     # mov edi, dword ptr [rdi + rsi*4 + 0x2710]
     self.assertEqual(bytes.fromhex(self.encode(load)), bytes.fromhex("8B BC B7 10 27 00 00"))
 
   # 8bit variants of legacy instructions subtract 1 from opcode
   def test_8bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(reg(RDX, 1)),), reg(RAX, 4))
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (reg_placeholder(reg(RDX,1)),), reg(RAX, 4))
     # movsx eax, dl
     self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("0F BE C2"))
 
   # accessing lower 8 bits of rsp, rbp, rsi, rdi requires rex prefix
   def test_lower_8bits_reg(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int32, (def_reg(reg(RDI, 1)),), reg(RAX, 4))
+    cast = ins(X86Ops.MOVSX, dtypes.int32, (reg_placeholder(reg(RDI,1)),), reg(RAX, 4))
     # movsx eax, dil
     self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("40 0F BE C7"))
 
   # test 16 bit variant of legacy instruction
   def test_16bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int16, (def_reg(reg(RDX, 1)),), reg(RAX, 2))
+    cast = ins(X86Ops.MOVSX, dtypes.int16, (reg_placeholder(reg(RDX,1)),), reg(RAX, 2))
     # movsx ax, dl
     self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("66 0F BE C2"))
 
   # test 64 bit variant of legacy instruction
   def test_64bit_legacy_encoding(self):
-    cast = ins(X86Ops.MOVSX, dtypes.int64, (def_reg(reg(RDX, 1)),), reg(RAX, 8))
+    cast = ins(X86Ops.MOVSX, dtypes.int64, (reg_placeholder(reg(RDX, 1)),), reg(RAX, 8))
     # movsx rax, dl
     self.assertEqual(bytes.fromhex(self.encode(cast)), bytes.fromhex("48 0F BE C2"))
 
-  # the width comes from the register in the tag, the dtype doesn't take part in the encoding
+  # the width comes from the register, the dtype doesn't take part in the encoding
   def test_width_from_register(self):
     for size, expected in ((1, "8A C2"), (2, "66 8B C2"), (4, "8B C2"), (8, "48 8B C2")):
-      mov = ins(X86Ops.MOV, dtypes.float32, (def_reg(reg(RDX, size)),), reg(RAX, size))
+      mov = ins(X86Ops.MOV, dtypes.float32, (reg_placeholder(reg(RDX, size)),), reg(RAX, size))
       self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex(expected))
 
   # test compact vex encoding
   def test_compact_vex_encoding(self):
-    xmm0, xmm1 = def_reg(XMM[0]), def_reg(XMM[1])
+    xmm0, xmm1 = reg_placeholder(XMM[0]), reg_placeholder(XMM[1])
     add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm1), XMM[0])
     # vaddss xmm0, xmm0, xmm1
     self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C5 FA 58 C1"))
 
   # test long vex encoding
   def test_long_vex_encoding(self):
-    xmm0, xmm8 = def_reg(XMM[0]), def_reg(XMM[8])
+    xmm0, xmm8 = reg_placeholder(XMM[0]), reg_placeholder(XMM[8])
     add = ins(X86Ops.VADDSS, dtypes.float32, (xmm0, xmm8), XMM[0])
     # vaddss xmm0, xmm0, xmm8
     self.assertEqual(bytes.fromhex(self.encode(add)), bytes.fromhex("C4 C1 7A 58 C0"))
 
   # test encoding where register is in the immediate field
   def test_reg_in_imm_field(self):
-    xmm0, xmm1, xmm2 = def_reg(XMM[0]), def_reg(XMM[1]), def_reg(XMM[2])
+    xmm0, xmm1, xmm2 = reg_placeholder(XMM[0]), reg_placeholder(XMM[1]), reg_placeholder(XMM[2])
     blend = ins(X86Ops.VBLENDVPS, dtypes.float32, (xmm0, xmm1, xmm2), XMM[0])
     # vblendvps xmm0, xmm0, xmm1, xmm2
     self.assertEqual(bytes.fromhex(self.encode(blend)), bytes.fromhex("C4 E3 79 4A C1 20"))
 
   # when writting to mem the uop takes the store form where dtype is void and there's no definition
   def test_write_mem(self):
-    address = (ptr(RDI, dtypes.int32), def_reg(RSI), imm(dtypes.int8, 10))
-    xmm0 = def_reg(XMM[0])
+    address = (alloc(dtypes.int32, RDI), reg_placeholder(RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
+    xmm0 = reg_placeholder(XMM[0])
     extr = ins(X86Ops.VPEXTRD, dtypes.void, address + (xmm0, imm(dtypes.uint8, 0)))
     # vpextrd dword ptr [rdi + rsi*4 + 0xa], xmm0, 0
     self.assertEqual(bytes.fromhex(self.encode(extr)), bytes.fromhex("C4 E3 79 16 44 B7 0A 00"))
 
   # test two address instruction with fused load works
   def test_two_address_load(self):
-    address = (ptr(RDI, dtypes.int32), def_reg(RSI), imm(dtypes.int8, 10))
+    address = (alloc(dtypes.int32, RDI), reg_placeholder(RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 4))
     cmove = ins(X86Ops.CMOVE, dtypes.int32, address, reg(RAX, 4))
     # cmove eax, dword ptr [rdi + rsi*4 + 0xa]
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 44 B7 0A"))
 
   # test instruction where displacement and imm have the same value
   def test_disp_imm_same_value(self):
-    address = (ptr(RDI, dtypes.int8), def_reg(RSI), imm(dtypes.int8, 10))
+    address = (alloc(dtypes.int8, RDI), reg_placeholder(RSI), imm(dtypes.int8, 10), imm(dtypes.uint8, 1))
     mov = ins(X86Ops.MOVi, dtypes.void, address + (imm(dtypes.int8, 10),))
     # mov byte ptr [rdi + rsi + 0xa], 0xa
     self.assertEqual(bytes.fromhex(self.encode(mov)), bytes.fromhex("40 C6 44 37 0A 0A"))
 
-    address = (ptr(RDI, dtypes.int32), def_reg(RSI), imm(dtypes.int32, 10))
+    address = (alloc(dtypes.int32, RDI), reg_placeholder(RSI), imm(dtypes.int32, 10), imm(dtypes.uint8, 4))
     imul = ins(X86Ops.IMULi, dtypes.int32, address + (imm(dtypes.int32, 10),), reg(RDI, 4))
     # imul edi, dword ptr [rdi + rsi*4 + 0xa], 0xa
     self.assertEqual(bytes.fromhex(self.encode(imul)), bytes.fromhex("69 BC B7 0A 00 00 00 0A 00 00 00"))
 
   # cmoves have the cmp as the last src even though it is not explicitly used, the cmp doesn't define a reg and is ignored in the encoding
   def test_cmove_ignore_cmp(self):
-    cmove = ins(X86Ops.CMOVE, dtypes.int32, (def_reg(reg(RAX, 4)), UOp(Ops.INS, arg=(X86Ops.CMP, dtypes.void))), reg(RDX, 4))
+    cmove = ins(X86Ops.CMOVE, dtypes.int32, (reg_placeholder(reg(RAX, 4)), UOp(Ops.INS, arg=(X86Ops.CMP, dtypes.void))), reg(RDX, 4))
     # cmove edx, eax
     self.assertEqual(bytes.fromhex(self.encode(cmove)), bytes.fromhex("0F 44 D0"))
 

tinygrad_repo/test/null/test_gen_float4.py

@@ -117,7 +117,7 @@ class TestFloat4(unittest.TestCase):
     assert TestFloat4.count_float4(uops) == (1, 1)
 
   def test_float4_aligned_variable(self):
-    x = Variable('x', 0, 4, multiple_of=4).bind(4)
+    x = Variable('k', 0, 1).bind(1)*4
     a = Tensor.empty(4).realize()
     b = Tensor.empty(12).realize().shrink(((x, x+4),))
     c = a + b
@@ -131,7 +131,7 @@ class TestFloat4(unittest.TestCase):
     assert TestFloat4.count_float4(uops) == (2, 1)
 
   def test_float4_unaligned_variable(self):
-    x = Variable('x', 0, 4, multiple_of=2).bind(4)
+    x = Variable('k', 0, 2).bind(2)*2
     a = Tensor.empty(4).realize()
     b = Tensor.empty(12).realize().shrink(((x, x+4),))
     c = a + b

tinygrad_repo/test/null/test_graph_rewrite.py

@@ -116,14 +116,6 @@ class TestSubstitute(unittest.TestCase):
     ret = substitute(ret, {a.sin():a.sqrt(), n1.sin():n1.sqrt()})
     self.assertIs(ret, a.sqrt().sqrt())
 
-  def test_tagged_replace(self):
-    a = UOp.variable('a', 0, 10)
-    b = UOp.variable('b', 0, 10)
-    ret = (a+4).replace(tag=1)
-    ret = substitute(ret, {a:b})
-    # the srcs are rewritten but we keep tag
-    self.assertIs(ret, (b+4).replace(tag=1))
-
 matchers = strat.sampled_from([PatternMatcher, TrackedPatternMatcher])
 
 class TestRecurse(unittest.TestCase):

tinygrad_repo/test/null/test_hcq2.py

@@ -1,24 +1,26 @@
 import unittest
-from typing import cast
 from types import SimpleNamespace
 from collections import defaultdict
 from unittest.mock import patch
-from tinygrad import Device, Tensor, TinyJit, dtypes
-from tinygrad.device import Buffer, Compiled, ProfileGraphEvent
+from tinygrad import Device, Tensor, TinyJit, Variable, dtypes
+from tinygrad.device import Compiled, ProfileGraphEvent
 from tinygrad.helpers import Context, unwrap, to_tuple
-from tinygrad.dtype import AddrSpace
 from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, KernelInfo
-from tinygrad.engine.realize import compile_linear, link_linear, get_call_arg_uops, lower_and_compile, run_linear
+from tinygrad.engine.realize import compile_linear, get_call_arg_uops, lower_and_compile, run_linear
 import tinygrad.runtime.support.hcq2 as hcq2
 from tinygrad.runtime.support.hcq2 import HCQInfo
 
 def lower_hcq(*body:UOp) -> UOp: # body through lower_call, as a linear
-  return UOp(Ops.LINEAR, src=(unwrap(hcq2.lower_call(UOp.sink(*body, arg=KernelInfo("test"), tag=1).call(aux=HCQInfo(("CPU",))))),))
+  return UOp(Ops.LINEAR, src=(unwrap(hcq2.lower_call(UOp.sink(*body, arg=KernelInfo("test", optimized=True)).call(aux=HCQInfo(("CPU",))))),))
 
 def chain(x:Tensor, n:int) -> Tensor:
   for _ in range(n): x = (x + 1).contiguous()
   return x
 
+def add_n(n:UOp, C:UOp, B:UOp) -> UOp:
+  i = UOp.range(C.numel(), 0)
+  return C[i].store(B[i] + n).end(i).sink(arg=KernelInfo(name="add_n"))
+
 def chain_input(value:int=2, device="NULL") -> Tensor: return Tensor.full((4,), value, dtype=dtypes.int32, device=device).contiguous().realize()
 
 def compiled_chain(n:int, jit=False, device="NULL") -> tuple[Tensor, UOp, list[UOp]]:
@@ -39,14 +41,21 @@ def scheduled(*ts:Tensor, **kwargs) -> list[UOp]:
     lin = orig(l, profile)
     batches.extend(c for c in lin.src if c.op is Ops.CALL and isinstance(c.arg.aux, HCQInfo))
     return lin
-  with patch.object(hcq2, "sched_batches", track): compile_linear(ts[0].schedule_linear(*ts[1:]), **kwargs)
+  with patch.object(hcq2, "sched_batches", track): compile_linear(ts[0].linear_with_vars(*ts[1:])[0], **kwargs)
   return batches
 
+def submits(batch:UOp) -> list[UOp]: # the submit calls of a batch sink in order
+  calls, s = list[UOp](), batch.src[0]
+  while s.op is Ops.AFTER: calls, s = [s.src[0]] + calls, s.src[-1]
+  return calls
 def queues(batch:UOp) -> dict[tuple[str, str], list[UOp]]:
-  return {(lin.arg[0][0], lin.arg[1]): list(lin.src) for lin in (s.src[1].without_after for s in batch.body.src)}
-def calls(batch:UOp) -> list[UOp]: return [c for cmds in queues(batch).values() for c in cmds if c.op is Ops.CALL]
+  return {(q.device, q.arg.name): list(c.body.src) for c in submits(batch.body) for q in [c.src[1]]}
+def ins_name(c:UOp) -> str|None: return c.body.arg.name if c.op is Ops.CALL and c.body.op is Ops.CUSTOM_FUNCTION else None
+def calls(batch:UOp) -> list[UOp]: return [c for cmds in queues(batch).values() for c in cmds if c.op is Ops.CALL and ins_name(c) is None]
 def devices_of(call:UOp) -> set[str]: return {to_tuple(a.device)[0] for a in get_call_arg_uops(call)}
 
+def is_slots(buf:UOp) -> bool: return (b:=hcq2.unwrap_view(buf)[0]).op is Ops.ALLOC and b.arg.name == "slots"
+
 def word(u:UOp) -> tuple[UOp, int]:
   if u.op is Ops.INDEX: return (v:=hcq2.unwrap_view(u.src[0]))[0], v[1] + u.src[1].val * u.dtype.itemsize
   return hcq2.unwrap_view(u)
@@ -55,8 +64,8 @@ T = 5 # the timeline value the last submitted batch of every device signals
 def run(batch:UOp, done:dict[str, int]|None=None, prio:list|None=None) -> tuple[list, dict, dict]:
   qs, cs, mem = queues(batch), calls(batch), defaultdict(int)
   for d in batch.arg.aux.device:
-    mem[word(hcq2.timeline((d,)))] = (done or {}).get(d, T)
-    mem[word(hcq2.timeline((d,)).index(1))] = T
+    mem[word(hcq2.timeline(d))] = (done or {}).get(d, T)
+    mem[word(hcq2.timeline(d).index(1))] = T
   def val(u:UOp) -> int:
     if u.op is Ops.CONST: return u.val
     if u.op is Ops.LOAD: return mem[word(u.src[0])]
@@ -66,16 +75,16 @@ def run(batch:UOp, done:dict[str, int]|None=None, prio:list|None=None) -> tuple[
     for q in (list(qs) if prio is None else prio):
       if not (cmds:=qs[q]): continue
       c = cmds[0]
-      if c.op is Ops.INS and c.arg[0].startswith("wait"):
-        sig, target = mem[word(c.src[0])], val(c.src[1])
-        if sig != target if c.arg[0] == "wait_eq" else sig < target: continue
-      elif c.op is Ops.INS and c.arg[0] == "store":
-        mem[w:=word(c.src[0])] = val(c.src[1])
-        if w[0].tag == "timeline": log.append(w[0].device)
-      elif c.op is Ops.CALL: log.append(cs.index(c))
+      if (name:=ins_name(c)) is not None and name.startswith("wait"):
+        sig, target = mem[word(c.src[1])], val(c.src[2])
+        if sig != target if name == "wait_eq" else sig < target: continue
+      elif name == "store":
+        mem[w:=word(c.src[1])] = val(c.src[2])
+        if w[0].arg.name == "timeline": log.append(w[0].device)
+      elif c in cs: log.append(cs.index(c))
       cmds.pop(0)
       break
-    else: return log, {q: len(cmds) for q, cmds in qs.items()}, {d: mem[word(hcq2.timeline((d,)))] for d in batch.arg.aux.device}
+    else: return log, {q: len(cmds) for q, cmds in qs.items()}, {d: mem[word(hcq2.timeline(d))] for d in batch.arg.aux.device}
 
 def rotations(batch:UOp) -> list[list]: return [(qs:=list(queues(batch)))[i:] + qs[:i] for i in range(len(queues(batch)))]
 def orders(batch:UOp) -> set[tuple[int, ...]]: return {tuple(x for x in run(batch, prio=p)[0] if isinstance(x, int)) for p in rotations(batch)}
@@ -85,44 +94,44 @@ class TestHCQ2Deps(unittest.TestCase):
     b = UOp.param(0, dtypes.float32, 64, device=("NULL", "NULL:1"))
     for view in [b.mselect(0).shrink(((8, 16),)), b.shrink(((8, 16),)).mselect(0), b.shrink(((4, 32),)).mselect(0).shrink(((4, 12),))]:
       tracker = hcq2.DepsTracker()
-      tracker.access_resources([view], [0], 0)
-      self.assertEqual(tracker.access_resources([b.mselect(1)], [], 1), [])
-      self.assertEqual(tracker.access_resources([b.mselect(0).shrink(((16, 24),))], [], 2), [])
-      self.assertEqual(tracker.access_resources([b.mselect(0).shrink(((12, 20),))], [], 3), [0])
+      tracker.access_resources([(view, True)], 0)
+      self.assertEqual(tracker.access_resources([(b.mselect(1), False)], 1), [])
+      self.assertEqual(tracker.access_resources([(b.mselect(0).shrink(((16, 24),)), False)], 2), [])
+      self.assertEqual(tracker.access_resources([(b.mselect(0).shrink(((12, 20),)), False)], 3), [0])
 
   def test_disjoint_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
-    for write in ([], [0]):
+    for write in (False, True):
       tracker = hcq2.DepsTracker()
-      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])
+      tracker.access_resources([(b.shrink(((0, 4),)), write)], 0)
+      self.assertEqual(tracker.access_resources([(b.shrink(((4, 8),)), True)], 1), [])
+      self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0])
 
   def test_partial_write_preserves_dependencies(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
-    for write in ([], [0]):
+    for write in (False, True):
       tracker = hcq2.DepsTracker()
-      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])
+      tracker.access_resources([(b, write)], 0)
+      self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 1), [0])
+      self.assertEqual(tracker.access_resources([(b.shrink(((0, 4),)), True)], 2), [0])
+      self.assertEqual(tracker.access_resources([(b.shrink(((12, 16),)), True)], 3), [0])
+      self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), False)], 4), [1])
 
   def test_write_waits_for_all_readers(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
     tracker = hcq2.DepsTracker()
-    self.assertEqual(tracker.access_resources([b], [], 0), [])
-    self.assertEqual(tracker.access_resources([b], [], 1), [])
-    self.assertEqual(tracker.access_resources([b], [0], 2), [0, 1])
-    self.assertEqual(tracker.access_resources([b], [], 3), [2])
+    self.assertEqual(tracker.access_resources([(b, False)], 0), [])
+    self.assertEqual(tracker.access_resources([(b, False)], 1), [])
+    self.assertEqual(tracker.access_resources([(b, True)], 2), [0, 1])
+    self.assertEqual(tracker.access_resources([(b, False)], 3), [2])
 
   def test_aliases_do_not_wait_on_themselves(self):
     b = UOp.param(0, dtypes.uint8, 16, device="NULL")
-    for write in ([0], [1], [0, 1]):
+    for whole, part in ((True, False), (False, True), (True, True)):
       tracker = hcq2.DepsTracker()
-      tracker.access_resources([b], [0], 0)
-      self.assertEqual(tracker.access_resources([b, b.shrink(((4, 12),))], write, 1), [0])
-      self.assertEqual(tracker.access_resources([b.shrink(((4, 12),))], [0], 2), [1])
+      tracker.access_resources([(b, True)], 0)
+      self.assertEqual(tracker.access_resources([(b, whole), (b.shrink(((4, 12),)), part)], 1), [0])
+      self.assertEqual(tracker.access_resources([(b.shrink(((4, 12),)), True)], 2), [1])
 
 class TestHCQ2Schedule(unittest.TestCase):
   def setUp(self):
@@ -145,11 +154,21 @@ class TestHCQ2Schedule(unittest.TestCase):
 
   def scheduled(self, *ts:Tensor) -> list[UOp]: return [self.check(b) for b in scheduled(*ts)]
   def batch(self, *ts:Tensor) -> UOp: return unwrap(self.scheduled(*ts)[0])
+  def kernels(self, batch:UOp) -> list[str]:
+    return [ins_name(c) or "kernel" for c in queues(batch)[("NULL", "COMPUTE:0")] if ins_name(c) in ("exec_barrier", None)]
 
   def test_kernels_run_in_order(self): self.assertEqual(orders(self.batch(chain(self.x, 3))), {(0, 1, 2)})
 
-  def test_a_peer_kernel_runs_after_the_copy_that_feeds_it(self):
-    self.assertEqual(orders(self.batch((self.x.to("NULL:1") + 1).contiguous())), {(0, 1)})
+  def test_exec_barrier(self):
+    self.assertEqual(self.kernels(self.batch(chain(self.x, 3))), ["kernel", "exec_barrier", "kernel", "exec_barrier", "kernel"])
+    a = (self.x + 1).contiguous()
+    self.assertEqual(self.kernels(self.batch((a + 1).contiguous(), (a + 2).contiguous())), ["kernel", "exec_barrier", "kernel", "kernel"])
+
+  def test_a_copy_waits_for_the_kernel_when_the_scalar_is_first(self):
+    n, B = Tensor(Variable("n", 0, 10, dtypes.int).bind(5)), Tensor.ones(4, dtype=dtypes.int).contiguous().realize()
+    b = self.batch(Tensor.custom_kernel(n, Tensor.empty(4, dtype=dtypes.int), B, fxn=add_n)[1].to("NULL:1"))
+    kernel = next(i for i, c in enumerate(calls(b)) if c.body.op is Ops.PROGRAM)
+    self.assertEqual({o[0] for o in orders(b)}, {kernel})
 
   def test_lanes_of_a_sharded_kernel_do_not_wait_for_each_other(self):
     s = Tensor.ones(8).contiguous().realize().shard(("NULL", "NULL:1"), axis=0).contiguous().realize()
@@ -188,33 +207,23 @@ class TestHCQ2Profile(unittest.TestCase):
 
   def test_slots_addressed_by_the_device(self):
     pm = PatternMatcher([(UPat((Ops.LOAD, Ops.STORE), src=(UPat(Ops.INDEX, src=(UPat.var("buf"), UPat())),), allow_any_len=True),
-                          lambda buf: buf.getaddr(Device["NULL"].host) if hcq2.unwrap_view(buf)[0].tag == "slots" else None)])
-    with patch.object(Device["NULL"], "pm_lower", pm): self.test_profiling_reports_a_range_per_kernel(n=3)
+                          lambda buf: buf.getaddr(Device["NULL"].host) if is_slots(buf) else None)])
+    with patch.object(Compiled, "pm_lower", pm + Compiled.pm_lower): self.test_profiling_reports_a_range_per_kernel(n=3)
 
 class TestHCQ2Link(unittest.TestCase):
   def setUp(self): self.enterContext(Context(DEV="NULL"))
 
-  def test_links_serve_any_input(self):
-    a, inputs = chain_input(), list[UOp]()
-    linear = compile_linear(chain(a, 2).schedule_linear(), input_uops=inputs, cache=True)
-    linked = link_linear(linear, input_uops=inputs)
-    self.assertIs(link_linear(linear, input_uops=[chain_input(3).uop.base, *inputs[1:]]), linked)
-    bufs = [cast(Buffer, u.buffer) for u in linked.toposort() if u.op is Ops.BUFFER and u.addrspace is AddrSpace.GLOBAL]
-    self.assertNotIn(a.uop.base.buffer, bufs)
-    words = [w for b in bufs if b.options.external_ptr and b.nbytes % 8 == 0 for w in b.host.view(fmt='Q')[:]]
-    self.assertNotIn(cast(Buffer, a.uop.base.buffer)._buf, words)
-
   def test_eager_templates_compile_once(self): self.assertIs(compiled_chain(3)[1], compiled_chain(3)[1])
 
   def test_repeated_word_loops(self): # for (i..10) cmdbuf[off[i]] = var
-    var, offs = UOp.alloc((1,), dtypes.uint32, device="CPU").rtag("var"), [4 * i * i for i in range(10)]
+    var, offs = UOp.alloc((1,), dtypes.uint32, device="CPU", name="var"), [4 * i * i for i in range(10)]
     var = hcq2.patch(var, [], bytes(4)) # initialize at link: the allocator may return a reused buffer
     hq = SimpleNamespace(blob=bytearray(offs[-1] + 4), patches=[(o, var.index(0).load()) for o in offs], devs=("CPU",), queue="COPY:0", deps=())
     lowered = lower_hcq(var.after(hcq2.encode_cmdbuf(hq, device="CPU")).index(0).store(var.index(0).load() + 1))
     self.assertEqual(len([u for u in lowered.src[0].without_after.src[0].toposort() if u.op is Ops.RANGE]), 1)
-    linked = hcq2.hcq_link(lower_and_compile(lowered), allow_cache=False)
+    linked = hcq2.hcq_link(lower_and_compile(lowered))
     args = zip(lowered.src[0].without_after.src[1:], linked.src[0].without_after.src[1:])
-    cmdbuf = next(b.buffer for p, b in args if p.without_after.tag == "cmdbuf_copy_0")
+    cmdbuf = next(b.buffer for p, b in args if p.without_after.arg.name == "cmdbuf_copy_0")
     for step in range(3):
       run_linear(linked, jit=True)
       self.assertEqual([cmdbuf.host.view(fmt="I")[o // 4] for o in offs], [step] * 10)

tinygrad_repo/test/null/test_isel.py

@@ -1,9 +1,10 @@
 import unittest
-from tinygrad.helpers import Target
+from tinygrad import Tensor
+from tinygrad.helpers import Target, Context
+from tinygrad.codegen import to_program
 from tinygrad.uop import Ops
 from tinygrad.uop.ops import UOp, dtypes, graph_rewrite
-from tinygrad.renderer.isa.x86 import X86Renderer, X86Ops
-from tinygrad.renderer.isa import IselContext
+from tinygrad.renderer.x86 import X86Renderer, X86Ops, IselContext
 
 # INDEX on a register value with a constant index extracts a single element (the old GEP)
 def lane(y:UOp, i:int) -> UOp: return y.index(UOp.cconst(i, dtypes.int))
@@ -24,7 +25,7 @@ class TestIselX86(unittest.TestCase):
     a = UOp.variable("a", 0, 0, dtypes.int32)
     b = UOp.variable("b", 0, 0, dtypes.int32)
     c = (a < b).where(a, b)
-    d = (a != b).where(a, b)
+    d = a.alu(Ops.CMPNE, b).where(a, b)
     f = c + d
     n = self.isel_rewrite(f)
     self.assertTrue(n.src[0].arg[0] is X86Ops.CMOVL and n.src[1].arg[0] is X86Ops.CMOVNE)
@@ -49,5 +50,9 @@ class TestIselX86(unittest.TestCase):
     # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
     self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
 
+  @Context(SPEC=2, NOOPT=1)
+  def test_register_store_is_effect(self):
+    to_program(Tensor.empty(16, 16).sum(1).schedule_linear().src[-1].src[0], X86Renderer(Target(device="CPU", arch="x86_64")))
+
 if __name__ == "__main__":
   unittest.main()

tinygrad_repo/test/null/test_linearizer.py

@@ -1,11 +1,11 @@
 import unittest
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import Ops, GroupOp, AxisType
+from tinygrad.uop.ops import UOp, Ops, GroupOp, AxisType
 from tinygrad.device import Device
 from tinygrad.tensor import Tensor
 from tinygrad.codegen import to_program
 from tinygrad.dtype import DType, dtypes, AddrSpace
-from tinygrad.renderer.isa import ISARenderer
+from tinygrad.renderer.x86 import ISARenderer
 from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.renderer.llvmir import AMDLLVMRenderer
 from tinygrad.helpers import Target
@@ -59,6 +59,20 @@ class TestLinearizer(unittest.TestCase):
     num_ops = len([uop for uop in uops if uop.op in GroupOp.ALU])
     assert num_ops == 0, "more alu uops than needed"
 
+  def test_stack_select_folds_into_index(self):
+    x = Tensor.empty(7, 4).realize()
+    r = Tensor.stack(x.const_like(0), x, x.const_like(0)).reshape(21, 4).shrink(((4, 19), None)).contiguous()
+    uops = tuple(to_program(replace_opts(r.schedule_linear().src[-1].src[0], []), renderer=Device[Device.DEFAULT].renderer).src[1].src)
+    self.assertEqual([u.op for u in uops if u.op in {Ops.DIV, Ops.MOD, Ops.TRUNCDIV, Ops.TRUNCMOD}], [])
+
+  def test_const_table_even_index(self):
+    # idx*2 is even, so the gather never selects the odd entry 5 and the store folds to 0
+    idx = UOp.variable("idx", 0, 3).bind(0)
+    out = Tensor.const((0, 0, 0, 5, 0, 0, 0, 0), dtypes.int)[idx*2].to("NULL", force=True)
+    uops = to_program(out.linear_with_vars()[0].src[-1].src[0], Device["NULL"].renderer).src[1].src
+    self.assertNotIn(Ops.WHERE, [u.op for u in uops])
+    self.assertEqual(next(u for u in uops if u.op is Ops.STORE).src[1].vmax, 0)
+
   def test_sum_acc_dtype(self):
     for tensor_dtype, acc_dtype in (
       (dtypes.bool, dtypes.int), (dtypes.int16, dtypes.int), (dtypes.float16, dtypes.float), (dtypes.bfloat16, dtypes.float)):

tinygrad_repo/test/null/test_llm_tokenizer.py

@@ -1,8 +1,7 @@
-import unittest, base64, functools, re, sys, time, unicodedata
+import unittest, base64, functools, re, time, unicodedata, json, pathlib, tempfile
 from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
 from tinygrad.helpers import fetch
 
-@unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
 class TestLLMTokenizer(unittest.TestCase):
   @functools.cached_property
   def llama_tok(self):
@@ -105,8 +104,43 @@ class TestLLMTokenizer(unittest.TestCase):
     self.assertTrue(tok.is_end(0))  # eom ends generation
     self.assertFalse(tok.is_end(2))
 
+  def test_hf_tokens(self):
+    with tempfile.TemporaryDirectory() as folder:
+      path = pathlib.Path(folder)
+      vocab = {"h": 0, "i": 1, "hi": 2, "Ġ": 3, "!": 4, "[gMASK]": 5, "hello world": 9, "1": 10, "2": 11, "3": 12, "123": 13}
+      added = [dict(content=t, id=i, special=i != 9) for t, i in (("[gMASK]", 5), ("<eos>", 6), ("<eot>", 7), ("<eom>", 8), ("hello world", 9))]
+      (path / "tokenizer.json").write_text(json.dumps(dict(model=dict(vocab=vocab), added_tokens=added)))
+      (path / "config.json").write_text(json.dumps(dict(eos_token_id=[6, 7, 8])))
+      for preset in ("glm4", "llama3", "tekken"):
+        with self.subTest(preset=preset):
+          tok = SimpleTokenizer.from_hf(str(path), preset=preset)
+          self.assertEqual(tok.preset, preset)
+          self._test_coding(tok, "hi hi!", [2, 3, 2, 4])
+          self._test_coding(tok, "h!i", [0, 4, 1])
+          self._test_coding(tok, "hi [gMASK]hello world!<eos>", [2, 3, 5, 9, 4, 6])
+          self._test_coding(tok, "123", [10, 11, 12] if preset == "tekken" else [13])
+          self.assertEqual((tok.eos_id, tok.eot_id, tok.eom_id), (6, 7, 8))
+          self.assertEqual([tok.is_end(i) for i in range(5, 10)], [False, True, True, True, False])
+
+  def test_hf_metadata(self):
+    with tempfile.TemporaryDirectory() as folder:
+      path = pathlib.Path(folder)
+      (path / "tokenizer.json").write_text(json.dumps(dict(model=dict(vocab={"<s>": 0}))))
+      (path / "config.json").write_text(json.dumps(dict(eos_token_id=1)))
+      tok = SimpleTokenizer.from_hf(path, preset="llama3")
+      self.assertIsNone(tok.bos_id)
+      self.assertIsNone(tok.chat_template)
+      self.assertEqual((tok.eos_id, tok.eot_id, tok.eom_id), (1, None, None))
+      self.assertTrue(tok.is_end(1))
+      (path / "tokenizer_config.json").write_text(json.dumps(dict(bos_token="<s>", chat_template="config template")))
+      tok = SimpleTokenizer.from_hf(path, preset="llama3")
+      self.assertEqual(tok.bos_id, 0)
+      self.assertEqual(tok.chat_template, "config template")
+      (path / "chat_template.jinja").write_text("{{ messages[0].content }}")
+      self.assertEqual(SimpleTokenizer.from_hf(path, preset="llama3").chat_template, "{{ messages[0].content }}")
+
   def test_tekken_gpt4o_split(self):
-    split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o")}
+    split = {p: SimpleTokenizer({}, {}, p)._split_to_word.findall for p in ("tekken", "gpt-4o", "kimi-k2")}
     shared = {
       "HelloWorld": ["Hello", "World"],
       " ÜNICODE": [" ÜNICODE"],  # Ü: non-ascii upper joins the run
@@ -117,6 +151,7 @@ class TestLLMTokenizer(unittest.TestCase):
     for s, want in shared.items():
       self.assertEqual(split["tekken"](s), want, f"tekken {s!r}")
       self.assertEqual(split["gpt-4o"](s), want, f"gpt-4o {s!r}")
+      self.assertEqual(split["kimi-k2"](s), want, f"kimi-k2 {s!r}")
     differ = [
       ("12345", list("12345"), ["123", "45"]),  # digits: tekken single, o200k groups {1,3}
       ("it's I'M don'T", ["it", "'s", " I", "'M", " don", "'T"], ["it's", " I'M", " don'T"]),  # contraction: o200k inline suffix
@@ -124,6 +159,8 @@ class TestLLMTokenizer(unittest.TestCase):
     for s, tk, go in differ:
       self.assertEqual(split["tekken"](s), tk, f"tekken {s!r}")
       self.assertEqual(split["gpt-4o"](s), go, f"gpt-4o {s!r}")
+      self.assertEqual(split["kimi-k2"](s), go, f"kimi-k2 {s!r}")
+    self.assertEqual(split["kimi-k2"]("a.\n/b"), ["a", ".\n", "/b"])
 
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""

tinygrad_repo/test/null/test_multitensor.py

@@ -1,6 +1,7 @@
 import gc, unittest
 from tinygrad import Tensor, Device, nn, GlobalCounters, TinyJit, dtypes, UOp
 from tinygrad.uop.ops import Ops
+from tinygrad.dtype import AddrSpace
 from tinygrad.helpers import Context
 from tinygrad.nn.state import get_parameters, get_state_dict
 from test.helpers import not_support_multi_device, needs_second_gpu
@@ -131,8 +132,6 @@ class TestMultiRamUsage(unittest.TestCase):
       return x, w1, w2
 
     def run_layers(n_layers):
-      GlobalCounters.reset()
-
       @TinyJit
       def f(x, w1, w2):
         for _ in range(n_layers):
@@ -144,8 +143,10 @@ class TestMultiRamUsage(unittest.TestCase):
         r = f(*a)
         del a, r
 
-      gc.collect()
-      return GlobalCounters.mem_used
+      # get the data itself. cmdbuf memory might grow a bit
+      linked = next(iter(f.captured._linears.values()))[0]
+      bufs = [u.buffer for u in linked.toposort(enter_calls=True) if u.op is Ops.BUFFER and u.addrspace is AddrSpace.GLOBAL]
+      return sum(b.nbytes for b in {b for mb in bufs for b in getattr(mb, "bufs", [mb])} if not b.options.cpu_access)
 
     mem_2 = run_layers(2)
     mem_4 = run_layers(4)

tinygrad_repo/test/null/test_pattern_matcher.py

@@ -5,23 +5,23 @@ from tinygrad.uop.ops import PatternMatcher, UPat
 
 class TestPatternMatcher(unittest.TestCase):
   def test_simple_match(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x", dtype=dtypes.weakfloat), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp.const(1)
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_upat_any(self):
     def test(a, x=None, y=None, z=None):
       #print(x,y,z)
-      if y is not None: return (a+y).rtag()
+      if y is not None: return (a+y).sink()
     matcher = PatternMatcher([
       (UPat.var("a")+UPat.any(UPat.var("x"), UPat.var("y"), UPat.var("z")), test),
     ])
     v1 = UOp.variable("a", 0, 10)
     v2 = UOp.variable("b", 0, 10)
     c1 = v1+v2
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
 
   def test_minimum_len(self):
     matcher = PatternMatcher([
@@ -60,43 +60,43 @@ class TestPatternMatcher(unittest.TestCase):
     self.assertEqual(len(ctx), 1)
 
   def test_uop(self):
-    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp(Ops.ADD, src=(c1, c1))
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_uop_set(self):
-    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat((Ops.CONST, Ops.CAST), name="x"), lambda x: x.sink())])
     c1 = UOp.const(False)
     c2 = UOp(Ops.CAST, arg=dtypes.int, src=(c1,))
     c3 = UOp.const(1.0)
     c4 = UOp(Ops.ADD, src=(c3, c3))
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
-    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
+    self.assertEqual(matcher.rewrite(c2), c2.sink())
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_arg(self):
     matcher = PatternMatcher([
-      (UPat(Ops.CONST, arg=0, name="x"), lambda x: x.rtag()),
-      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.rtag()),
-      (UPat(Ops.MAX, name="x"), lambda x: x.rtag()),
+      (UPat(Ops.CONST, arg=0, name="x"), lambda x: x.sink()),
+      (UPat(Ops.CONST, arg=False, name="x"), lambda x: x.sink()),
+      (UPat(Ops.MAX, name="x"), lambda x: x.sink()),
     ])
     c1 = UOp.const(0.0)
     c2 = UOp.const(False)
     c3 = UOp(Ops.MAX, src=(c1, c1))
     c4 = UOp(Ops.MUL, src=(c1, c1))
     c5 = UOp.const(-1)
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
-    self.assertEqual(matcher.rewrite(c2), c2.rtag())
-    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
+    self.assertEqual(matcher.rewrite(c2), c2.sink())
+    self.assertEqual(matcher.rewrite(c3), c3.sink())
     self.assertEqual(matcher.rewrite(c4), None)
     self.assertEqual(matcher.rewrite(c5), None)
 
   def test_filter_arg(self):
     matcher = PatternMatcher([
       (UPat(Ops.MUL, src=[UPat(Ops.CONST, name="c"), UPat(Ops.CONST, arg=2)], name="x"),
-       lambda x,c: x.rtag() if c.val in {1, -1} else None)
+       lambda x,c: x.sink() if c.val in {1, -1} else None)
     ])
     y1 = UOp.const(1)
     y2 = UOp.const(2)
@@ -106,46 +106,46 @@ class TestPatternMatcher(unittest.TestCase):
     c3 = UOp(Ops.MUL, src=(y3, y2))
     c4 = UOp(Ops.MUL, src=(y2, y1))
     c5 = UOp(Ops.MUL, src=(y2, y3))
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
     self.assertEqual(matcher.rewrite(c2), None)
-    self.assertEqual(matcher.rewrite(c3), c3.rtag())
-    self.assertEqual(matcher.rewrite(c4), c4.rtag())
-    self.assertEqual(matcher.rewrite(c5), c5.rtag())
+    self.assertEqual(matcher.rewrite(c3), c3.sink())
+    self.assertEqual(matcher.rewrite(c4), c4.sink())
+    self.assertEqual(matcher.rewrite(c5), c5.sink())
 
   def test_dup_name(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.rtag())])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST, name="y"), UPat(Ops.CONST, name="y"))), lambda x, y: x.sink())])
     y1 = UOp.const(1.0)
     y2 = UOp.const(1.0)
     c1 = UOp(Ops.ADD, src=(y1, y1))
     c2 = UOp(Ops.ADD, src=(y1, y2))
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
-    self.assertEqual(matcher.rewrite(c2), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
+    self.assertEqual(matcher.rewrite(c2), c1.sink())
 
   def test_dtype(self):
     # a concrete const dtype lives on the pair's CAST
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype=dtypes.float32), lambda x: x.sink())])
     c1 = UOp.const(1.0).cast(dtypes.float32)
     c2 = UOp.const(1.0).cast(dtypes.float64)
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
     self.assertEqual(matcher.rewrite(c2), None)
 
   def test_dtype_set(self):
-    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(Ops.CAST, name="x", dtype={dtypes.float32, dtypes.float64}), lambda x: x.sink())])
     c1 = UOp.const(1.0).cast(dtypes.float32)
     c2 = UOp.const(1.0).cast(dtypes.float64)
     c3 = UOp.const(1.0).cast(dtypes.float16)
     c4 = UOp.const(1).cast(dtypes.int)
-    self.assertEqual(matcher.rewrite(c1), c1.rtag())
-    self.assertEqual(matcher.rewrite(c2), c2.rtag())
+    self.assertEqual(matcher.rewrite(c1), c1.sink())
+    self.assertEqual(matcher.rewrite(c2), c2.sink())
     self.assertEqual(matcher.rewrite(c3), None)
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_src_one(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=(UPat(Ops.CONST), UPat(Ops.CONST))), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
-    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c3), c3.sink())
     self.assertEqual(matcher.rewrite(c2), None)
     # that CONST/ALU -> ALU/CONST rewrite is now instant
     """
@@ -158,7 +158,7 @@ class TestPatternMatcher(unittest.TestCase):
     """
 
   def test_src_permutations(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=[UPat(Ops.CONST), UPat(GroupOp.ALU)]), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
@@ -166,21 +166,21 @@ class TestPatternMatcher(unittest.TestCase):
     c5 = UOp(Ops.ADD, src=(c2,c3))
     c6 = UOp(Ops.ADD, src=(c3,c4))
     self.assertEqual(matcher.rewrite(c3), None)
-    self.assertEqual(matcher.rewrite(c4), c4.rtag())
-    self.assertEqual(matcher.rewrite(c5), c5.rtag())
+    self.assertEqual(matcher.rewrite(c4), c4.sink())
+    self.assertEqual(matcher.rewrite(c5), c5.sink())
     self.assertEqual(matcher.rewrite(c6), None)
 
   def test_src_repeat(self):
-    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(GroupOp.ALU, name="x", src=UPat(Ops.CONST)), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp.const(2.0)
     c3 = UOp(Ops.ADD, src=(c1,c2))
     c4 = UOp(Ops.ADD, src=(c2,c3))
-    self.assertEqual(matcher.rewrite(c3), c3.rtag())
+    self.assertEqual(matcher.rewrite(c3), c3.sink())
     self.assertEqual(matcher.rewrite(c4), None)
 
   def test_allow_len(self):
-    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.rtag())])
+    matcher = PatternMatcher([(UPat(Ops.MULACC, name="x", src=(UPat(Ops.CONST),), allow_any_len=True), lambda x: x.sink())])
     c1 = UOp.const(1.0)
     c2 = UOp.const(2.0)
     c3 = UOp.const(3.0)
@@ -189,7 +189,7 @@ class TestPatternMatcher(unittest.TestCase):
     c6 = UOp(Ops.MULACC, src=(c1,c2,c3))
     self.assertEqual(matcher.rewrite(c4), None)
     self.assertEqual(matcher.rewrite(c5), None)
-    self.assertEqual(matcher.rewrite(c6), c6.rtag())
+    self.assertEqual(matcher.rewrite(c6), c6.sink())
 
   def test_deep_src_permutations(self):
     c1 = UOp.const(1.0)

tinygrad_repo/test/null/test_randomness.py

@@ -1,22 +1,7 @@
 import unittest
 from tinygrad import dtypes, Tensor
-from tinygrad.codegen import to_program
-from tinygrad.helpers import Target
-from tinygrad.renderer.cstyle import ClangRenderer
-from tinygrad.uop.ops import Ops
 
 class TestRandomness(unittest.TestCase):
-  def test_threefry_doesnt_use_long(self):
-    # Use a renderer without native THREEFRY or 64-bit pointer arithmetic in the UOps.
-    renderer = ClangRenderer(Target("CPU", arch="x86_64,x86-64"))
-    linear = Tensor.rand(20).schedule_linear()
-    for call in linear.src:
-      ast = call.src[0]
-      if ast.op is Ops.SINK:
-        prg = to_program(ast, renderer=renderer)
-        for u in tuple(prg.src[1].src):
-          self.assertNotIn(u.dtype, {dtypes.long, dtypes.ulong}, msg=f"long found in {prg.src[0].arg.name}")
-
   def test_rand_like(self):
     empty = Tensor.empty((80, 44))
     rand = Tensor.rand_like(empty)

tinygrad_repo/test/null/test_real_world.py

@@ -30,7 +30,7 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
       tms.append(time.perf_counter_ns() - st)
     mem_used = (GlobalCounters.mem_used - global_mem_used) / 1e9
 
-    kernels_used = jit_cache_count(model.captured.linear) if getattr(model, "captured", None) is not None else None
+    kernels_used = jit_cache_count(model.captured._linear) if getattr(model, "captured", None) is not None else None
     print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
     assert mem_used < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.3f} GB - {mem_used:.3} GB used"
     assert (max_memory_allowed - mem_used) / max_memory_allowed < 0.2, f"{max_memory_allowed:.3f} GB is too far from {mem_used:.3} GB used"

tinygrad_repo/test/null/test_renderer_failures.py

@@ -38,7 +38,7 @@ class TestWGSLFailures(unittest.TestCase):
     b = UOp.param(0, dtypes.char, 4)
     word = b.index(UOp.const(0).cast(dtypes.int)).cast(dtypes.uint32)
     store = word.store(word.load() & UOp.const(0xffffff00).cast(dtypes.uint32))
-    src = WGSLRenderer(Target("WEBGPU")).render(UOp.sink(store, arg=KernelInfo()).toposort())
+    src = WGSLRenderer(Target("WEBGPU")).render(list(UOp.sink(store, arg=KernelInfo()).toposort()))
     self.assertIn("atomicAnd(&data0_4[0],4294967040u);", src)
     self.assertNotIn("atomicAdd", src)
 

tinygrad_repo/test/null/test_schedule.py

@@ -1,13 +1,13 @@
 # schedule tests that pass on NULL backend (no copyout needed)
-import unittest, time, gc
+import unittest, time
 from typing import cast
 from tinygrad import nn, dtypes, Device, Tensor, getenv
-from tinygrad.helpers import GlobalCounters, Context, all_same
+from tinygrad.helpers import Context, all_same
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat, KernelInfo
 from tinygrad.dtype import AddrSpace
 from tinygrad.engine.realize import run_linear, compile_linear
 from tinygrad.codegen import to_program, full_rewrite_to_sink
-from test.helpers import check_schedule, assert_kernel_count, KernelCountException, jit_cache_count
+from test.helpers import check_schedule, KernelCountException, jit_cache_count, measure_mem
 
 def _realize_weights(m):
   for p in nn.state.get_parameters(m): p.realize()
@@ -209,16 +209,6 @@ class TestSchedule(unittest.TestCase):
     expr = a/a
     check_schedule(expr, 1)
 
-  def test_div_collapse(self):
-    a = Tensor.full((4,), 1.0).contiguous().realize()
-    b = Tensor.full((4,), 2.0).contiguous().realize()
-    c = Tensor.full((4,), 3.0).contiguous().realize()
-    GlobalCounters.reset()
-    expr = (a/b)/c
-    expr.realize()
-    assert_kernel_count(1)
-    self.assertLessEqual(GlobalCounters.global_ops, 4*3)
-
   # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
   # should contiguous dedup?
   @unittest.skip("we do the exact opposite now")
@@ -611,13 +601,10 @@ class TestSchedule(unittest.TestCase):
   def test_conv2d_half(self): self.test_conv2d(4, dtype=dtypes.half)
 
   def test_schedule_mem_used_with_inputs(self):
-    Tensor.ones(256).contiguous().realize() # hcq2 caches the linked schedule with its buffers
-    gc.collect()
-    base = GlobalCounters.mem_used
-    x = Tensor.ones(256).contiguous().realize()
-    (x+Tensor.ones(256).contiguous()).schedule_linear()
-    gc.collect()
-    self.assertEqual(GlobalCounters.mem_used-base, 1024)
+    with measure_mem() as used:
+      x = Tensor.ones(256).contiguous().realize()
+      (x+Tensor.ones(256).contiguous()).schedule_linear()
+    self.assertEqual(used(), 1024)
 
   def test_image_dot_f16_fusion(self):
     with Context(FLOAT16=1, OPENPILOT_HACKS=1):
@@ -722,7 +709,8 @@ class TestSchedule(unittest.TestCase):
       for call in linear.src:
         if call.src[0].op is not Ops.SINK: continue
         sink = full_rewrite_to_sink(call.src[0], Device[call.device].renderer)
-        reduce_kernels += any(u.op is Ops.STORE and u.addrspace is AddrSpace.REG and u.ranges for u in sink.toposort())
+        # a reduce kernel is identified by a STORE into a REG buffer (STORE has void dtype, check the addrspace of its target)
+        reduce_kernels += any(u.op is Ops.STORE and u.src[0].addrspace is AddrSpace.REG and u.ranges for u in sink.toposort())
       self.assertEqual(reduce_kernels, 1)
 
   def test_push_through_reshape(self):
@@ -1413,7 +1401,7 @@ class TestSchedule(unittest.TestCase):
 
   def test_pow_neg_05_is_rsqrt(self):
     t = Tensor([1.0, 2.0, 3.0]) ** -0.5
-    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIPROCAL, Ops.SQRT])
+    self.assertEqual(self._alu_from_tensor(t), [Ops.RECIP, Ops.SQRT])
 
   def test_pow_2_has_1_mul(self):
     t = Tensor([1.0, 2.0, 3.0]) ** Tensor(2.0)
@@ -1698,13 +1686,10 @@ class TestSchedule(unittest.TestCase):
     check_schedule(out, 2)
 
   def test_schedule_mem_used(self):
-    Tensor.ones(256).contiguous().realize() # hcq2 caches the linked schedule with its buffers
-    gc.collect()
-    base = GlobalCounters.mem_used
-    Tensor.ones(256).contiguous().realize()
-    Tensor.ones(5, 5).contiguous().schedule_linear()
-    gc.collect()
-    self.assertEqual(GlobalCounters.mem_used-base, 0)
+    with measure_mem() as used:
+      Tensor.ones(256).contiguous().realize()
+      Tensor.ones(5, 5).contiguous().schedule_linear()
+    self.assertEqual(used(), 0)
 
   def test_const_schedule(self):
     constv = Tensor.empty(2, 2).const_like(10).uop

tinygrad_repo/test/null/test_simplify_valid_idx.py

@@ -2,15 +2,16 @@ import unittest, itertools
 from dataclasses import replace
 
 from tinygrad.codegen.late.coalesce import indexing_simplify, pm_simplify_add_image
+from tinygrad import Tensor
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import UOp, Ops, graph_rewrite
 from tinygrad.uop.weak import pm_commit_weak
-from tinygrad.uop.symbolic import simplify_valid, sym, pm_move_where_on_load
+from tinygrad.uop.symbolic import simplify_valid, uop_given_valid, sym, pm_move_where_on_load
 from tinygrad.helpers import Context
 from test.helpers import full_rewrite
 from test.null.test_uop_symbolic import check_uop_against_string
 
-# symbolic-only idx + valid simplification (no late lowering of FLOORDIV/FLOORMOD)
+# symbolic-only idx + valid simplification (no late lowering of DIV/MOD)
 def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load, name="simplify_valid_idx")
 # image-aware idx + valid simplification: adds the codegen-layer matcher that drops provably in-bounds gates
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")
@@ -40,7 +41,7 @@ class TestValidIdxSimplification(unittest.TestCase):
     load = get_gated_load_uop(gate, idx)
     self.check(load,
       "0",
-      "(((lidx0+(gidx0*4))<19)!=True)")
+      "((lidx0+(gidx0*4))<19).eq(False)")
 
   def test_simplify_within_valid1(self):
     ridx0 = Range(0, 4)
@@ -61,6 +62,13 @@ class TestValidIdxSimplification(unittest.TestCase):
     valid = (alu0 < 57) & (alu0 >= 1)
     self.assertIsNone(simplify_valid(valid))
 
+  def test_simplex_const_term_keeps_lanes_aligned(self):
+    # the 1 in the valid is not a simplex variable, lane +1 must simplify like the other lanes
+    a, b, c, d = Variable("a", 0, 7), Variable("b", 0, 7), Variable("c", 0, 3), Variable("d", 0, 1)
+    valid = ((a+b+1)<1).ne(True) & ((c+d)<1).ne(True)
+    base = ((a*8+b+c*128+d*64+1000)%1024)*4
+    self.assertTrue(all(uop_given_valid(valid, base+k).src[0] is uop_given_valid(valid, base+0).src[0] for k in range(1, 4)))
+
   def test_bitwise_and_is_not_a_valid(self):
     ridx0 = Range(0, 16)
     self.assertEqual(simplify_valid_idx(UOp.sink((ridx0 & UOp.const(12, dtypes.int)) & ridx0)).src[0].render(), "((i32)(r0)&12&(i32)(r0))")
@@ -191,7 +199,7 @@ class TestValidIdxSimplification(unittest.TestCase):
     load = get_gated_load_uop(valid, idx)
     self.check(load,
       "1",
-      "((((r0+r1)<1)!=True)&(((r2+r3)<1)!=True))")
+      "(((r0+r1)<1).eq(False)&((r2+r3)<1).eq(False))")
 
   def test_valid_with_non_const_rhs(self):
     ridx0 = Range(0, 1024)
@@ -232,11 +240,11 @@ class TestImageSimplification(unittest.TestCase):
     # should match any one of the AND clause and drop the matched statement from valid
     valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
     load = get_load_image_uop(shape, valid, (gidx0+1, gidx1-1))
-    self.check(load, "((gidx0<1)!=True)", "(gidx0+1)", "(gidx1+-1)")
+    self.check(load, "(gidx0<1).eq(False)", "(gidx0+1)", "(gidx1+-1)")
 
     valid = (gidx0<1).ne(True) & (gidx1<1).ne(True)
     load = get_load_image_uop(shape, valid, (gidx0, gidx1-1))
-    self.check(load, "((gidx0<1)!=True)", "gidx0", "(gidx1+-1)")
+    self.check(load, "(gidx0<1).eq(False)", "gidx0", "(gidx1+-1)")
 
   def test_idx_lt_bound(self):
     # (idx1 < image_bound) ? (..., idx1) : 0 can drop the valid
@@ -471,7 +479,7 @@ class TestImageSimplification(unittest.TestCase):
       self.check(load, None, "(gidx0+lidx0*1024+r0*1024+lidx1*128+-3168)", "0")
     except AssertionError:
       # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<19))",
+      self.check(load, "((lidx1<1).eq(False)&((lidx0+r0)<19))",
                        "(gidx0+lidx1*128+(lidx0*1024+r0*1024)+-3168)", "0")
 
   def test_simplify10(self):
@@ -490,7 +498,7 @@ class TestImageSimplification(unittest.TestCase):
       self.check(load, None, "(lidx2+gidx0*4+lidx0*1024+r0*1024+lidx1*256+-3264)", "0")
     except AssertionError:
       # TODO: fold valid
-      self.check(load, "(((lidx1<1)!=True)&((lidx0+r0)<11))",
+      self.check(load, "((lidx1<1).eq(False)&((lidx0+r0)<11))",
                        "(lidx2+gidx0*4+lidx1*256+(lidx0*1024+r0*1024)+-3264)", "0")
 
   def test_drop_non_monotonic_window(self):
@@ -513,6 +521,10 @@ class TestImageSimplification(unittest.TestCase):
     off = graph_rewrite(load.sink(), pm_commit_weak+indexing_simplify).src[0].src[0]
     self.assertEqual(off.src[1].get_valid(), UOp.const(True))
 
+  def test_int64_index_into_image_with_emulated_long(self):
+    with Context(DEV="NULL:OPENCL:,IMAGE_PITCH_ALIGNMENT=64", IMAGE=1, EMULATED_DTYPES="long"):
+      Tensor.empty(96, 128)[Tensor([3, 1, 5], dtype=dtypes.int64)].realize()
+
 class TestImageStore(unittest.TestCase):
   def test_half_store_converts_lane_by_lane(self):
     # a half4 stored to a half image converts to float per lane: a half4->float4 CAST is not valid OpenCL

tinygrad_repo/test/null/test_symbolic_tensor.py

@@ -67,7 +67,7 @@ class TestSymbolicReshape(unittest.TestCase):
 
   def test_reshape_infer_multiple_of(self):
     for multiple in (1, 2, 3, 4, 8):
-      var = Variable("n", multiple, 8*multiple, multiple_of=multiple)
+      var = Variable("n", multiple, 8*multiple)
       for n in (var, var.bind(2*multiple)):
         with self.subTest(multiple=multiple, bound=n.is_bound_var):
           self.assertEqual(Tensor.empty(8*multiple, 4)[:n].reshape(n, -1).shape, (n, 4))

tinygrad_repo/test/null/test_tensor_metadata.py

@@ -1,127 +0,0 @@
-import unittest
-from tinygrad import Tensor, dtypes
-from tinygrad.tensor import _METADATA
-from tinygrad.engine.realize import capturing
-from tinygrad.helpers import Context
-
-@unittest.skip("tensor metadata is no longer supported")
-class TestTensorMetadata(unittest.TestCase):
-  def setUp(self) -> None:
-    _METADATA.set(None)
-    self._ctx = Context(SCACHE=0)
-    self._ctx.__enter__()
-  def tearDown(self) -> None:
-    self._ctx.__exit__(None, None, None)
-
-  @unittest.skip("why would this be true?")
-  def test_exclude_noop_metadata(self):
-    a = Tensor.rand(4, 4)*1
-    self.assertEqual(a.uop.metadata[0].name, "__mul__")
-    k = a.schedule_linear().src[-1]
-    self.assertEqual([m.name for m in k.arg.metadata], ["rand"])
-
-  @unittest.skip("metadata not reaching kernel schedule")
-  def test_exclude_const_metadata(self):
-    a = Tensor.arange(4)
-    b = Tensor.full((4,), -1, dtype=dtypes.int).contiguous()
-    sched = a.schedule_linear(b)
-    self.assertEqual([m.name for m in sched.src[0].arg.metadata], ["arange"])
-    self.assertEqual([m.name for m in sched.src[1].arg.metadata], ["contiguous"])
-
-  def test_matmul(self):
-    x = Tensor.rand(3)
-    W = Tensor.rand(3, 3)
-    out = x.matmul(W)
-    self.assertEqual(out.uop.metadata[0].name, "matmul")
-    si = out.schedule_linear().src[-1]
-    self.assertEqual(len(si.arg.metadata), 1)
-    self.assertEqual(si.arg.metadata[0].name, "matmul")
-
-  def test_relu(self):
-    x = Tensor.rand(3)
-    out = x.relu()
-    self.assertEqual(out.uop.metadata[0].name, "relu")
-    si = out.schedule_linear().src[-1]
-    self.assertEqual(len(si.arg.metadata), 1)
-    self.assertEqual(si.arg.metadata[0].name, "relu")
-
-  @unittest.skip("assign metadata no longer captured")
-  def test_assign(self):
-    x = Tensor.empty(10, 10).realize()
-    x.assign(Tensor.ones(10, 10).contiguous())
-    si = x.schedule_linear().src[-1]
-    self.assertEqual(len(si.arg.metadata), 1)
-    self.assertEqual(si.arg.metadata[0].name, "assign")
-
-  def test_complex(self):
-    x = Tensor.rand(3)
-    y = Tensor.rand(3)
-    out = x.relu() * y.sigmoid()
-    self.assertEqual(out.uop.metadata[0].name, "__mul__")
-    self.assertEqual(out.uop.src[0].metadata[0].name, "relu")
-    self.assertEqual(out.uop.src[1].metadata[0].name, "sigmoid")
-    si = out.schedule_linear().src[-1]
-    self.assertEqual(len(si.arg.metadata), 3)
-    self.assertEqual(set(m.name for m in si.arg.metadata), {"relu", "sigmoid", "__mul__"})
-
-  @unittest.skip("flaky")
-  def test_complex_backward(self):
-    x = Tensor.rand(3).realize()
-    y = Tensor.rand(3).realize()
-    out = (x.relu() * y.sigmoid()).sum()
-    self.assertEqual(out.uop.metadata[0].name, "sum")
-    out.backward()
-    self.assertEqual(x.grad.uop.metadata[0].name, "relu")
-    #self.assertTrue(x.grad.uop.metadata[0].backward)  # TODO: backward flag is False
-    self.assertEqual(y.grad.uop.metadata[0].name, "sigmoid")
-    #self.assertTrue(y.grad.uop.metadata[0].backward)  # TODO: backward flag is False
-    si = out.schedule_linear(x.grad, y.grad).src[-1]
-    #self.assertEqual(len(si.arg.metadata), 3, f"failed with {si.arg.metadata}")
-    # skip numpy, this is schedule cache
-    self.assertSetEqual(set(m.name for m in si.arg.metadata if m.name != "numpy"), {"sigmoid", "relu"})
-    #bw = [m for m in si.metadata if m.backward]
-    #self.assertEqual(len(bw), 1)
-    #self.assertEqual(bw[0].name, "sigmoid")
-
-  def test_tracemeta_0(self):
-    with Context(TRACEMETA=0):
-      x = Tensor.rand(3)
-      y = Tensor.rand(3)
-      out = (x.relu() * y.sigmoid()).sum()
-      self.assertIsNone(out.uop.metadata)
-      self.assertIsNone(out.uop.src[0].metadata)
-      si = out.schedule_linear().src[-1]
-      self.assertEqual(si.arg.metadata, ())
-
-  def _has_metadata(self, h, name):
-    linears = []
-    capturing.append(type("", (), {"add_linear": lambda _, linear: linears.append(linear)})())
-    try: h.realize()
-    finally: capturing.clear()
-    calls = [call for linear in linears for call in linear.src]
-    return any(m.name == name for call in calls for m in call.arg.metadata)
-
-  def test_metadata_survives_realize_pending_assign(self):
-    shared = Tensor.rand(4)
-    c = Tensor.zeros(8).contiguous().realize()
-    c[:4].assign(shared)
-    self.assertTrue(self._has_metadata(c[:4].relu(), "relu"))
-
-  @unittest.expectedFailure
-  def test_metadata_lost_realize_pending_assign(self):
-    shared = Tensor.rand(4)
-    c = Tensor.zeros(8).contiguous().realize()
-    c[:4].assign(shared)
-    self.assertTrue(self._has_metadata((c[:4] + shared).relu(), "relu"))
-
-class TestTraceMetaShutdown(unittest.TestCase):
-  def test_tracemeta_del_no_shutdown_error(self):
-    import subprocess, os
-    result = subprocess.run(['python3', '-c', 'from tinygrad import Tensor\n'
-                             'x=Tensor.eye(3); (x@x).sum().backward()'],
-                            env={**os.environ, "TRACEMETA": "2"}, capture_output=True)
-    self.assertEqual(result.returncode, 0)
-    self.assertNotIn(b"Exception", result.stderr)
-
-if __name__ == '__main__':
-  unittest.main()

tinygrad_repo/test/null/test_uir_render.py

@@ -2,7 +2,7 @@ import unittest, re, ast as pyast
 from typing import Any
 from tinygrad import Tensor
 from tinygrad.dtype import dtypes, DType, AddrSpace, Invalid
-from tinygrad.uop import Ops
+from tinygrad.uop import Ops, GroupOp
 from tinygrad.uop.ops import AxisType, UOp, graph_rewrite, ParamArg, KernelInfo
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.uop.render import render_uir
@@ -36,11 +36,10 @@ def _parse_const(tok:str) -> UOp:
 
 def _parse_paramarg(rest:str, name:str|None) -> ParamArg:
   kv = _kv(rest)
-  def i(k:str) -> int|None: return int(kv[k]) if k in kv else None
   dev = kv.get("device")
   return ParamArg(int(kv["slot"]), _DTYPES_BY_NAME[kv["dtype"]],
                   tuple(map(int, kv["bounds"].strip("[]").split(","))) if "bounds" in kv else None,
-                  i("multiple_of"), name, AddrSpace[kv["addrspace"]] if "addrspace" in kv else AddrSpace.GLOBAL,
+                  name, AddrSpace[kv["addrspace"]] if "addrspace" in kv else AddrSpace.GLOBAL,
                   pyast.literal_eval(dev) if dev and dev[0] in "'(" else dev, kv.get("volatile") == "true")
 
 def parse_ssa(text:str) -> UOp:
@@ -66,7 +65,7 @@ def parse_ssa(text:str) -> UOp:
       name, argstr = argstr[1:end], argstr[end+1:].strip()
     srcs = [parse_tok(t) for t in _src_tok.findall(srcstr)]
     arg: Any = None
-    if op in {Ops.PARAM, Ops.BUFFER, Ops.ALLOC}: arg = _parse_paramarg(argstr, name)
+    if op in GroupOp.Defines: arg = _parse_paramarg(argstr, name)
     elif op is Ops.RANGE:
       t, *ids = argstr.split()
       arg = (AxisType[t], *map(int, ids))
@@ -86,7 +85,7 @@ def parse_ssa(text:str) -> UOp:
 
 def _strip_buffers(root:UOp) -> UOp:
   # realized BUFFERs carry runtime state the wire can't; a BUFFER with no binding is exactly an ALLOC
-  subs = {b: b.replace(op=Ops.ALLOC, arg=ParamArg(b.arg.slot, b.arg.dtype, b.arg.vmin_vmax, b.arg.multiple_of,
+  subs = {b: b.replace(op=Ops.ALLOC, arg=ParamArg(b.arg.slot, b.arg.dtype, b.arg.vmin_vmax,
                                                  b.arg.name, b.arg.addrspace, b.arg.device, b.arg.volatile))
           for b in root.toposort() if b.op is Ops.BUFFER and isinstance(b.arg, ParamArg)}
   return root.substitute(subs, walk=True, name="strip buffers for wire format test") if subs else root

tinygrad_repo/test/null/test_uop_graph.py

@@ -1,9 +1,9 @@
 import unittest, pytest, weakref
 from tinygrad import dtypes, Variable
 from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, GroupOp, AxisType, broadcast_axes, KernelInfo
+from tinygrad.uop.ops import Ops, UOp, UPat, PatternMatcher, graph_rewrite, AxisType, broadcast_axes, KernelInfo
 from tinygrad.uop.symbolic import sym
-from tinygrad.codegen.simplify import pm_reduce_unparented
+from tinygrad.codegen.simplify import pm_reduce_collapse
 from test.helpers import full_rewrite, to_uops_list
 
 simple_pm = PatternMatcher([
@@ -19,9 +19,6 @@ class TestTuplize(unittest.TestCase):
     a, b = UOp.const(1), UOp.const(2)
     pairs = [(a, UOp.const(1)), (a, b), (a, a+b), (a+b, b+a), (UOp.sink(a, b), UOp.sink(a, b))]
     for x, y in pairs: self.assertEqual(x.tuplize == y.tuplize, x is y)
-    # tags are not part of the order: tag-only differences order neither way, like value-equal plain tuples
-    self.assertFalse(a.tuplize < a.rtag("tagged").tuplize)
-    self.assertFalse(a.rtag("tagged").tuplize < a.tuplize)
 
   def test_deep_shared_subgraphs(self):
     # long equal prefixes with a difference at the bottom: correct order, and the pathological case for value equality
@@ -31,7 +28,7 @@ class TestTuplize(unittest.TestCase):
     self.assertNotEqual(a.tuplize < b.tuplize, b.tuplize < a.tuplize)
 
   def test_does_not_retain_uops(self):
-    a = UOp.const(1).rtag(object())
+    a = UOp(Ops.NOOP, arg=object())
     ref = weakref.ref(a)
     _ = a.tuplize
     del a
@@ -300,19 +297,6 @@ class TestUOpGraph(unittest.TestCase):
       assert u.op is not Ops.WHERE
       if u.op is Ops.LOAD: self.assertIs(u.src[1], expected)
 
-  def test_where_on_gated_load_with_cast(self):
-    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, 100)
-    gate_idx = ridx0.valid((ridx0<50))
-    ld = d0.index(gate_idx).cast(dtypes.float)
-    w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, 100)
-    uops = to_uops_list([out.index(ridx0).store(w)])
-    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
-    for u in uops:
-      assert u.op is not Ops.WHERE
-      if u.op is Ops.LOAD and u.src[0].src[0].op is Ops.PARAM: self.assertIs(u.src[1], expected)
-
   def test_where_on_casted_gated_load_extra_cond(self):
     ridx0 = UOp.range(100, 0)
     d0 = UOp.param(0, dtypes.float, 100)
@@ -448,7 +432,6 @@ class TestUOpGraph(unittest.TestCase):
 class TestReduceCollapse(unittest.TestCase):
   def test_multi_range_reduce_add(self):
     """Test that (x + y).reduce(r1, r2) distributes over multiple ranges"""
-    from tinygrad.codegen.simplify import pm_reduce_collapse
     # Create two ranges
     r1 = UOp.range(3, 0)
     r2 = UOp.range(4, 1)
@@ -462,14 +445,12 @@ class TestReduceCollapse(unittest.TestCase):
     # Should become add of two separate reduces
     self.assertEqual(result.op, Ops.ADD)
 
-  def test_bool_reduce_unparented_keeps_dtype(self):
-    for op in (Ops.ADD, Ops.MUL):
-      self.assertIs(graph_rewrite(UOp.const(True).reduce(UOp.range(4, 0), arg=op), pm_reduce_unparented).dtype, dtypes.bool)
-
-  def test_max_reduce_unparented_empty_range(self):
-    n = Variable("n", 0, 4)
-    ret = graph_rewrite(UOp.const(1.0).reduce(UOp.range(n, 0), arg=Ops.MAX), pm_reduce_unparented)
-    self.assertEqual(ret.substitute({n: n.const_like(0)}).simplify().arg, float("-inf"))
+  def test_lt_lift_skips_unsigned(self):
+    # x+y<c -> x<c-y and x*y<c -> x<ceil(c/y) only hold when x+y and x*y cannot wrap
+    for dt in (dtypes.int32, dtypes.uint32):
+      x, y, c = [Variable(n, 1, 100, dt) for n in "xyc"]
+      for lt in ((x+y) < c, (x*y) < c):
+        self.assertEqual(graph_rewrite(lt, pm_reduce_collapse) is lt, dt is dtypes.uint32)
 
   def test_reduce_shapeless_const_unroll(self):
     """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
@@ -534,21 +515,6 @@ class TestConstBufferize(unittest.TestCase):
     const_vals = [u.val for u in result.toposort() if u.op is Ops.CONST and u.dtype is dtypes.weakfloat]
     self.assertIn(3.14, const_vals)
 
-class TestUOpTags(unittest.TestCase):
-  def test_inc_by_one(self):
-    g = UOp.const(1) + UOp.const(1)
-    assert g.ssimplify() == 2
-    pm_plus_1 = PatternMatcher([(UPat(Ops.CONST, name="x"), lambda x: UOp.const(x.val+1, x.dtype).rtag(1) if x.tag is None else None)])
-    pm_strip_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
-    g = graph_rewrite(g, pm_plus_1)
-    assert g.ssimplify() == 4
-    g = graph_rewrite(g, pm_plus_1)
-    assert g.ssimplify() == 4
-    g = graph_rewrite(g, pm_strip_tags)
-    assert g.ssimplify() == 4
-    g = graph_rewrite(g, pm_plus_1)
-    assert g.ssimplify() == 6
-
 class TestUOpGetItem(unittest.TestCase):
   def _placeholder(self, shape, dtype=dtypes.half):
     return UOp.placeholder(shape, dtype, slot=0, addrspace=AddrSpace.LOCAL)

tinygrad_repo/test/null/test_uop_symbolic.py

@@ -45,13 +45,13 @@ class TestSymbolic(unittest.TestCase):
 
   def test_cmp_simple(self):
     self.helper_test_variable(Variable("a", 3, 8) < 4, 0, 1, "(a<4)")
-    self.helper_test_variable(Variable("a", 3, 8) >= 8, 0, 1, "((a<8)!=True)")
+    self.helper_test_variable(Variable("a", 3, 8) >= 8, 0, 1, "(a<8).eq(False)")
 
   def test_ge(self):
     self.helper_test_variable(Variable("a", 3, 8) >= 77, 0, 0, "False")
     self.helper_test_variable(Variable("a", 3, 8) >= 9, 0, 0, "False")
-    self.helper_test_variable(Variable("a", 3, 8) >= 8, 0, 1, "((a<8)!=True)")
-    self.helper_test_variable(Variable("a", 3, 8) >= 4, 0, 1, "((a<4)!=True)")
+    self.helper_test_variable(Variable("a", 3, 8) >= 8, 0, 1, "(a<8).eq(False)")
+    self.helper_test_variable(Variable("a", 3, 8) >= 4, 0, 1, "(a<4).eq(False)")
     self.helper_test_variable(Variable("a", 3, 8) >= 3, 1, 1, "True")
     self.helper_test_variable(Variable("a", 3, 8) >= 2, 1, 1, "True")
 
@@ -108,7 +108,7 @@ class TestSymbolic(unittest.TestCase):
 
   def test_multiple_of_cancellation(self):
     for multiple in (1, 2, 3, 4, 8):
-      var = UOp.variable("n", multiple, 8*multiple, multiple_of=multiple)
+      var = UOp.variable("n", multiple, 8*multiple)
       for n in (var, var.bind(2*multiple)):
         for factor in (-multiple, -1, 2, multiple, multiple+1):
           with self.subTest(multiple=multiple, bound=n.is_bound_var, factor=factor):
@@ -117,14 +117,6 @@ class TestSymbolic(unittest.TestCase):
             self.assertEqual(((-n*factor)//(-n)).ssimplify(), factor)
             self.assertIs(UOp.gcd(n*factor, n).simplify(), n)
 
-  def test_gcd_multiple_of_without_shared_factors(self):
-    n = UOp.variable("n", 6, 60, multiple_of=6)
-    m = UOp.variable("m", 4, 40, multiple_of=4)
-    self.assertEqual(UOp.gcd(n, m).ssimplify(), 2)
-    self.assertEqual((n%3).ssimplify(), 0)
-    self.assertEqual((m%2).ssimplify(), 0)
-    self.assertIs((n<12).simplify(), n<12)
-
   def test_divides_exact(self):
     a = Variable("a", 1, 8)
     b = Variable("b", 1, 8)
@@ -373,14 +365,14 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(Variable("a", 0, 5)*4 < 13, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", 0, 5)*4 < 16, 0, 1, "(a<4)")
     self.helper_test_variable(Variable("a", -5, 5)*4 < -13, 0, 1, "(a<-3)")
-    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "((a*-1)<4)")
+    self.helper_test_variable(Variable("a", -5, 5)*-4 < 13, 0, 1, "(-4<a)")
     c0, c1 = 2, 2**54+1
     self.helper_test_variable(Variable("a", 0, c1)*c0 < c1, 0, 1, f"(a<{2**53+1})")
     c0, c1 = -2, -(2**54-1)
-    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"((a*-1)<{-(2**53-1)})")
-    self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "((a*-1)<0)")
-    self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "((a<3)!=True)")
-    self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "((a<4)!=True)")
+    self.helper_test_variable(Variable("a", 0, -c1)*c0 < c1, 0, 1, f"({2**53-1}<a)")
+    self.helper_test_variable(Variable("a", 0, 5)*(-2) < 0, 0, 1, "(0<a)")
+    self.helper_test_variable(Variable("a", 0, 5)*4 >= 12, 0, 1, "(a<3).eq(False)")
+    self.helper_test_variable(Variable("a", 0, 5)*4 >= 13, 0, 1, "(a<4).eq(False)")
 
   def test_div_div(self):
     self.helper_test_variable((Variable("a", 0, 1800)//10)//9, 0, 20, "(a//90)")
@@ -476,13 +468,26 @@ class TestSymbolic(unittest.TestCase):
     z = UOp.const(0.0)
     self.assertTrue(math.isnan(graph_rewrite(z/z, symbolic_simple, bottom_up=True).arg))
 
+  def test_two_stage_const_folding(self):
+    x = UOp.variable("x", 0, 10)
+    self.assertIs(graph_rewrite((x+1)+2, symbolic_simple), x+3)
+
+  def test_unpack_uint64_high(self):
+    x, y = UOp.variable("x", 0, 2**32-1, dtypes.uint64), UOp.variable("y", 0, 2**32-1, dtypes.uint32)
+    self.assertIs(graph_rewrite((((x<<32) | y.cast(dtypes.uint64)) >> 32).cast(dtypes.uint32), symbolic_simple), x.cast(dtypes.uint32))
+
   def test_masked_shr_fold(self):
     x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
     self.helper_test_variable((x & -4) >> 2, 0, 63, "(x>>2)")
 
-  def test_masked_idiv_fold(self):
+  def test_masked_shr_fold_committed_mask(self):
     x = UOp.variable('x', 0, 255, dtype=dtypes.uint32)
-    self.helper_test_variable((x & -4) // 4, 0, 63, "(x//4)")
+    self.helper_test_variable((x & UOp.const(-4, dtypes.uint32)) >> 2, 0, 63, "(x>>2)")
+
+  def test_signed_const_wraps(self):
+    c = UOp.const(2**31, dtypes.int32)
+    self.assertEqual(c.val, -2**31)
+    self.assertIs((UOp.const(1, dtypes.int32) < c).simplify().arg, False)
 
   def test_bool_or_not_tautology(self):
     a = Variable("a", 0, 10)
@@ -772,14 +777,14 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((30 * b + 1) % 18 + ((30 * b + 1) // 18) * 18, 1, 3001, "((b*30)+1)")
 
   def test_div_partial_quotient(self):
-    # FLOORDIV should extract partial quotients when const_factor > divisor, matching what FLOORMOD already does
+    # DIV should extract partial quotients when const_factor > divisor, matching what MOD already does
     # (f*x+c)//d -> (f%d*x+c)//d + (f//d)*x  when f >= d
     b = Variable("b", 0, 100)
     self.helper_test_variable((31*b+1)//18, 0, 172, "(((b*13)+1)//18+b)")
     self.helper_test_variable((19*b+3)//7, 0, 271, "(((b*5)+3)//7+(b*2))")
 
   def test_gcd_with_remainder(self):
-    # gcd_with_remainder: factor GCD out of non-constant terms and denominator
+    # nest_by_factor on the GCD of non-constant terms and denominator
     a = Variable("a", 0, 2)
     self.helper_test_variable((a*4)//6, 0, 1, "(a*2//3)")
     self.helper_test_variable((a*4+1)//6, 0, 1, "(a*2//3)")
@@ -790,6 +795,11 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a*4+2)%6, 0, 4, "((a*2+1)%3*2)")
     self.helper_test_variable((a*4+3)%6, 1, 5, "((a*2+1)%3*2+1)")
 
+  def test_gcd_with_remainder_negative(self):
+    a, b = Variable("a", -3, 5), Variable("b", 0, 7)
+    self.helper_test_variable((a*18+1)//4, -14, 22, "(a*9//2)")
+    self.helper_test_variable((a*12+b*18+5)//12, -3, 15, "((a*2+b*3)//2)")
+
   def test_div_by_factor_tie_break(self):
     a = Variable("a", 0, 1)
     b = Variable("b", 0, 1)
@@ -817,7 +827,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((x%4)*-2 + ((x//4)%2)*-8, -14, 0, "(x%8*-2)")
 
   def test_mod_nest_by_factor(self):
-    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for FLOORMOD
+    # (a*f+b) % (f*k) = (a%k)*f + b when 0<=b<f — mirrors nest_div_by_factor for MOD
     gidx0 = Variable("gidx0", 0, 15)
     lidx0 = Variable("lidx0", 0, 3)
     # f=4, k=2, c=8: (gidx0*4+lidx0)%8 = (gidx0%2)*4 + lidx0
@@ -831,7 +841,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a*3+b)%9, 0, 8, "(b+a%3*3)")
 
   def test_mod_nest_by_factor_with_const(self):
-    # nest_by_factor FLOORMOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
+    # nest_by_factor MOD with non-zero constant offset: (a*f+b+const) % (f*k) = (a%k)*f + b + const when 0<=b+const<f
     a = Variable("a", 0, 7)
     b = Variable("b", 0, 1)
     # f=4, k=2, const=2: (a*4+b+2)%8 = (a%2)*4 + b + 2
@@ -843,7 +853,7 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((a*3+b+1)%6, 1, 5, "(b+a%2*3+1)")
 
   def test_div_nest_by_factor_with_const(self):
-    # nest_by_factor FLOORDIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
+    # nest_by_factor DIV: (160*a + 5*b + 4*c + K) // 60 should pick div=5 (clean) over div=4 (dirty)
     a = Variable("a", 0, 2)
     b = Variable("b", 0, 31)
     c = Variable("c", 0, 1)
@@ -924,16 +934,12 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable((idx<4).where(idx//4, uconst(-1)), -1, 6, "(idx<4).where((idx//4), -1)")
 
   def test_floordiv_lt(self):
-    # x//d<c <=> x<c*d for d>0, and <=> c*d<x for d<0
+    # x//d<c <=> x<c*d for d>0
     idx = Variable("idx", 0, 24)
     self.helper_test_variable((idx//4<3), 0, 1, "(idx<12)")
     self.helper_test_variable(((idx-20)//4<-3), 0, 1, "(idx<8)")
     self.helper_test_variable(((idx-10)//4<0), 0, 1, "(idx<10)")
-    self.helper_test_variable((idx//-4<-3), 0, 1, "(12<idx)")
-    self.helper_test_variable((idx//-4<-5), 0, 1, "(20<idx)")
     self.helper_test_variable((idx//-4<-6), 0, 0, "False")
-    self.helper_test_variable(((idx-10)//-4<0), 0, 1, "(8<(idx+-2))")
-    self.helper_test_variable(((idx-20)//-4<2), 0, 1, "(12<idx)")
 
   def test_nested_div_mod_negative_inner_divisor(self):
     # (x % (k*c)) // c -> (x // c) % k requires k>0; (x % (k*c)) % c -> x % c is unconditional for c>0
@@ -959,16 +965,16 @@ class TestSymbolic(unittest.TestCase):
     c = Variable("c", 0, 3)
     d = Variable("d", -3, 3)
     self.helper_test_variable((a<2), 0, 1, "(a<2)")
-    self.helper_test_variable((a<=2), 0, 1, "((2<a)!=True)")
+    self.helper_test_variable((a<=2), 0, 1, "(2<a).eq(False)")
     self.helper_test_variable((a>1), 0, 1, "(1<a)")
-    self.helper_test_variable((a>=1), 0, 1, "((a<1)!=True)")
-    self.helper_test_variable((a<1).ne(True), 0, 1, "((a<1)!=True)")
-    self.helper_test_variable((a+b<1).ne(True), 0, 1, "(((a+b)<1)!=True)")
-    self.helper_test_variable((a*3+b*4<1).ne(True), 0, 1, "(((a+b)<1)!=True)")
-    self.helper_test_variable((a*(-3)+b*4<1).ne(True), 0, 1, "((((a*-3)+(b*4))<1)!=True)")  # negative coeff, should not be simplified
-    self.helper_test_variable((a*3+d*4<1).ne(True), 0, 1, "((((a*3)+(d*4))<1)!=True)")  # var can be negative, should not be simplified
-    self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
-    self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "((((a+b)+c)<1)!=True)")
+    self.helper_test_variable((a>=1), 0, 1, "(a<1).eq(False)")
+    self.helper_test_variable((a<1).ne(True), 0, 1, "(a<1).eq(False)")
+    self.helper_test_variable((a+b<1).ne(True), 0, 1, "((a+b)<1).eq(False)")
+    self.helper_test_variable((a*3+b*4<1).ne(True), 0, 1, "((a+b)<1).eq(False)")
+    self.helper_test_variable((a*(-3)+b*4<1).ne(True), 0, 1, "(((a*-3)+(b*4))<1).eq(False)")  # negative coeff, should not be simplified
+    self.helper_test_variable((a*3+d*4<1).ne(True), 0, 1, "(((a*3)+(d*4))<1).eq(False)")  # var can be negative, should not be simplified
+    self.helper_test_variable((a+b+c*2<1).ne(True), 0, 1, "(((a+b)+c)<1).eq(False)")
+    self.helper_test_variable((a+b*2+c*4<1).ne(True), 0, 1, "(((a+b)+c)<1).eq(False)")
 
   def test_mul_by_zero_casted_to_emulated_dtype(self):
     # a float zero cast to bfloat16 is one committed const, so the mul-by-zero fold still sees it
@@ -980,13 +986,13 @@ class TestSymbolic(unittest.TestCase):
     # CAST(bool -> int) != 0  ->  cond
     self.helper_test_variable(cond.cast(dtypes.int).ne(0), 0, 1, "(a<2)")
     # CAST(bool -> int) != 1  ->  !cond
-    self.helper_test_variable(cond.cast(dtypes.int).ne(1), 0, 1, "((a<2)!=True)")
+    self.helper_test_variable(cond.cast(dtypes.int).ne(1), 0, 1, "(a<2).eq(False)")
     # CAST(bool -> int) != c (c not in {0,1})  ->  always True (CAST is 0 or 1)
     self.helper_test_variable(cond.cast(dtypes.int).ne(2), 1, 1, "True")
     self.helper_test_variable(cond.cast(dtypes.int).ne(-1), 1, 1, "True")
     # CAST(bool -> index) folds too
     self.helper_test_variable(cond.cast(dtypes.weakint).ne(0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.cast(dtypes.weakint).ne(1), 0, 1, "((a<2)!=True)")
+    self.helper_test_variable(cond.cast(dtypes.weakint).ne(1), 0, 1, "(a<2).eq(False)")
 
   def test_where_removal(self):
     cond = Variable("a", 0, 3) < 2
@@ -994,13 +1000,15 @@ class TestSymbolic(unittest.TestCase):
     self.helper_test_variable(cond, 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0), 0, 1, "(a<2)")
     self.helper_test_variable(cond.where(u1, u0).where(u1, u0), 0, 1, "(a<2)")
-    self.helper_test_variable(cond.where(u0, u1), 0, 1, "((a<2)!=True)")
+    self.helper_test_variable(cond.where(u0, u1), 0, 1, "(a<2).eq(False)")
     self.helper_test_variable(cond.where(u0, u1).where(u0, u1), 0, 1, "(a<2)")
 
   def test_equivalent_const_max(self):
     x = Variable("x", -10, 10)
     self.helper_test_variable((x < 0).where(0, x), 0, 10, "x.maximum(0)")
     self.helper_test_variable((0 < x).where(x, 0), 0, 10, "x.maximum(0)")
+    self.helper_test_variable(((x+3) < 0).where(0, x+3), 0, 13, "(x+3).maximum(0)")
+    self.helper_test_variable((5 < x).where(5, x), -10, 5, "(5<x).where(5, x)")
 
   def test_where_combine(self):
     cond = Variable("x", 0, 3) < 2
@@ -1161,12 +1169,6 @@ class TestSymbolic(unittest.TestCase):
     r0 = UOp.range((Variable("a", 1, 10)<5).cast(dtypes.weakint), 0)
     self.helper_test_variable(r0, 0, 0, "r0")
 
-  def test_const_reciprocal(self):
-    a = Variable("a", 1, 10, dtypes.float)
-    # TODO: bounds for reciprocal
-    # TODO: should z3 work?
-    self.helper_test_variable(2*(2*a).reciprocal(), -math.inf, math.inf, "a.reciprocal()", test_z3=False)
-
   def test_trunc_noop(self):
     a = Variable("a", 1, 10, dtypes.int)
     self.helper_test_variable(a.trunc(), 1, 10, "a", test_z3=False)
@@ -1266,6 +1268,13 @@ class TestSymInfer(unittest.TestCase):
     # floor: 1 % -1000 = -999, 1 // -1000 = -1
     assert sym_infer(a%b, var_vals) == -999
     assert sym_infer(a//b, var_vals) == -1
+  def test_sym_infer_truncdiv_truncmod(self):
+    a = Variable("a", -10, 10)
+    b = Variable("b", -10, 10)
+    for av, bv, div, mod in ((7, -3, -2, 1), (-7, 3, -2, -1), (7, 0, 0, 7)):
+      var_vals = {a.expr: av, b.expr: bv}
+      assert sym_infer(a.alu(Ops.TRUNCDIV, b), var_vals) == div
+      assert sym_infer(a.alu(Ops.TRUNCMOD, b), var_vals) == mod
   def test_sym_infer_with_cast(self):
     a = Variable("a", 0, 100, dtypes.int)
     assert sym_infer(a.cast(dtypes.long) + 1, {a.expr: 5}) == 6
@@ -1284,10 +1293,6 @@ class TestSymInfer(unittest.TestCase):
     assert sym_infer(c.bitcast(dtypes.int), {c.expr: 0xFFFFFFFF}) == -1
 
     assert sym_infer(UOp.const(1.5).cast(dtypes.float).bitcast(dtypes.uint), {}) == 1069547520
-  def test_sym_infer_tagged(self):
-    a = Variable("a", 0, 10)
-    assert sym_infer((a*3+1).rtag("t"), {a.expr: 2}) == 7
-
   def test_sym_infer_deeply_nested(self):
     # build an expression that exceeds Python's nested parentheses limit for eval
     # max(x, negative_const) can't be simplified when x can be negative, so nesting compounds