hwtest fixes for rdna3 dsl

* improve asm dsl syntax

* improve asm dsl syntax

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

@@ -70,13 +70,13 @@ runs:
       uses: actions/cache@v4
       with:
         path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
     - name: Cache downloads (macOS)
       if: inputs.key != '' && runner.os == 'macOS'
       uses: actions/cache@v4
       with:
         path: ~/Library/Caches/tinygrad/downloads/
-        key: osx-downloads-cache-${{ inputs.key }}-${{ env.CACHE_VERSION }}
+        key: downloads-${{ github.job }}-${{ inputs.key }}-${{ env.CACHE_VERSION }}
 
     # **** Python deps ****
 

.github/workflows/benchmark.yml

@@ -14,12 +14,6 @@ on:
       - update_benchmark
       - update_benchmark_staging
   workflow_dispatch:
-    inputs:
-      run_process_replay:
-        description: "Run process replay tests"
-        required: false
-        default: false
-        type: boolean
 
 jobs:
   testmacbenchmark:
@@ -39,6 +33,7 @@ jobs:
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
+        mkdir -p extra/disassemblers
         ln -s ~/tinygrad/extra/disassemblers/applegpu extra/disassemblers/applegpu
         ln -s ~/tinygrad/weights/sd-v1-4.ckpt weights/sd-v1-4.ckpt
         ln -s ~/tinygrad/weights/bpe_simple_vocab_16e6.txt.gz weights/bpe_simple_vocab_16e6.txt.gz
@@ -124,18 +119,6 @@ jobs:
     # TODO: too slow
     # - name: Run 10 CIFAR training steps w winograd
     #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
-    - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
-    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-    - name: UsbGPU (USB4/TB) boot time
-      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
-    - name: UsbGPU (USB4/TB) tiny tests
-      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -169,6 +152,37 @@ jobs:
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3.11 process_replay.py
 
+  testusbgpu:
+    name: UsbGPU Benchmark
+    env:
+      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
+    runs-on: [self-hosted, macOS]
+    timeout-minutes: 10
+    defaults:
+      run:
+        shell: bash -e -o pipefail {0}
+    if: github.repository_owner == 'tinygrad'
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: setup staging db
+      if: github.ref == 'refs/heads/update_benchmark_staging'
+      run: |
+        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
+        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
+    - name: UsbGPU boot time
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU tiny tests
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+    - name: UsbGPU copy speeds
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    - name: UsbGPU (USB4/TB) boot time
+      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
+    - name: UsbGPU (USB4/TB) tiny tests
+      run: PYTHONPATH=. NV=1 NV_IFACE=PCI NV_NAK=1 python3.11 test/test_tiny.py
+
   testnvidiabenchmark:
     name: tinybox green Benchmark
     runs-on: [self-hosted, Linux, tinyboxgreen]
@@ -415,13 +429,15 @@ jobs:
     #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
       run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test tensor cores
-      run: |
-        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
-        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Test tensor cores AMD_LLVM=0
+      run: AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+    # TODO: this is flaky
+    # - name: Test tensor cores AMD_LLVM=1
+    #   run: AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
     - name: Run Tensor Core GEMM (AMD)
-      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+      run: |
+        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+        AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
     - name: Test AMD=1
       run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
     #- name: Test HIP=1
@@ -541,8 +557,6 @@ jobs:
       run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
     - name: Run full CIFAR training steps w 6 GPUS
       run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD Training)
@@ -554,7 +568,6 @@ jobs:
           train_cifar_wino.txt
           train_cifar_one_gpu.txt
           train_cifar_six_gpu.txt
-          train_cifar_six_gpu_remote.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 

.github/workflows/test.yml

@@ -236,13 +236,13 @@ jobs:
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
         python3 -m ruff check examples/mlperf/ --ignore E501
+        python3 -m ruff check extra/thunder/tiny/ --ignore E501 --ignore F841 --ignore E722
     - name: Run mypy
       run: |
         python -m mypy --strict-equality --lineprecision-report .
         cat lineprecision.txt
-    # broken because of UPatAny
-    #- name: Run TYPED=1
-    #  run: TYPED=1 python -c "import tinygrad"
+    - name: Run TYPED=1
+      run: TYPED=1 python -c "import tinygrad"
 
   unittest:
     name: Unit Tests
@@ -310,7 +310,7 @@ jobs:
         deps: testing_unit
         python-version: '3.14'
     - name: Test SPEC=2
-      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
+      run: IGNORE_OOB=0 SPEC=2 PYTHONPATH="." pytest --maxfail=10 -n auto --durations=30 --ignore=test/models --ignore test/test_custom_kernel.py --ignore test/unit/test_hashing.py --timeout 60 -k "not test_setitem_big" --splits 2 --group ${{ matrix.group }}
 
   fuzzing:
     name: Fuzzing
@@ -641,7 +641,7 @@ jobs:
         if: matrix.backend=='amdllvm'
         run: python test/device/test_amd_llvm.py
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py --durations=20
       - name: Run pytest (amd)
         run: python -m pytest test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
@@ -721,71 +721,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  amdremote:
-    name: Linux (remote)
-    runs-on: ubuntu-22.04
-    timeout-minutes: 20
-    env:
-      REMOTE: 1
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: linux-remote
-          deps: testing_minimal
-          amd: 'true'
-          llvm: 'true'
-          opencl: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "CL" 7667
-          start_server "remote-server-cpu" "CPU" 8667
-      - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
-      - name: Run REMOTE=1 Test (CL)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
-
 # ****** OSX Tests ******
 
   testmetal:
@@ -883,30 +818,6 @@ jobs:
     - name: Test ONNX Runner (WEBGPU)
       run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py
 
-  osxremote:
-   name: MacOS (remote metal)
-   runs-on: macos-15
-   timeout-minutes: 10
-   env:
-     REMOTE: 1
-     REMOTEDEV: METAL
-   steps:
-     - name: Checkout Code
-       uses: actions/checkout@v4
-     - name: Setup Environment
-       uses: ./.github/actions/setup-tinygrad
-       with:
-         key: macos-remote
-         deps: testing_minimal
-     - name: Check Device.DEFAULT and print some source
-       run: |
-         python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-         python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'METAL', Device.default.properties.real_device"
-         DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-     - name: Run REMOTE=1 Test
-       run: |
-         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
-
   osxtests:
     strategy:
       fail-fast: false

.pre-commit-config.yaml

@@ -27,8 +27,8 @@ repos:
         always_run: true
         pass_filenames: false
       - id: tests
-        name: subset of tests
-        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        name: comprehensive test suite
+        entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_schedule.py test/test_assign.py test/test_tensor.py test/test_jit.py test/unit/test_schedule_cache.py test/unit/test_pattern_matcher.py test/unit/test_uop_symbolic.py test/unit/test_helpers.py
         language: system
         always_run: true
         pass_filenames: false

CLAUDE.md

@@ -88,13 +88,15 @@ VIZ=1 python -c "from tinygrad import Tensor; Tensor.ones(10).sum().realize()"
 ## Debugging Tips
 
 1. **Print UOp graphs**: `print(tensor.uop)` or `print(tensor.uop.sink())`
-2. **Check schedule**: `tensor.schedule()` returns list of ScheduleItems
+2. **Check schedule**: `tensor.schedule()` returns list of ExecItems
 3. **Trace graph rewrites**: Use `VIZ=1` or add print in PatternMatcher callbacks
 4. **Find UOps by type**: `[u for u in uop.toposort() if u.op is Ops.SOMETHING]`
 
 ## Workflow Rules
 
 - **NEVER commit without explicit user approval** - always show the diff and wait for approval
+- **NEVER amend commits** - always create a new commit instead
+- Run `pre-commit run --all-files` before committing to catch linting/type errors
 - Run tests before proposing commits
 - Test with `SPEC=2` when modifying UOp-related code
 
@@ -132,6 +134,18 @@ The schedule cache strips values from BIND nodes so different bound values (e.g.
 - Only extract var_vals when schedule is non-empty (no kernels = no vars needed)
 - PatternMatchers are slow to construct - define at module level, not in functions
 
+### Readability Over Speed
+Don't add complexity for marginal performance gains. Simpler code that's slightly slower is often better:
+```python
+# BAD: "optimized" with extra complexity
+if has_afters:  # skip toposort if no AFTERs
+  after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
+
+# GOOD: simple, always works
+after_map = [(u, u.buf_uop) for u in big_sink.toposort() if u.op is Ops.AFTER]
+```
+The conditional check adds complexity, potential bugs, and often negligible speedup. Only optimize when profiling shows a real bottleneck.
+
 ### Testing LLM Changes
 ```bash
 # Quick smoke test

docs/abstractions3.py

@@ -38,25 +38,19 @@ optim.schedule_step()   # this will step the optimizer without running realize
 # The weight Tensors have been assigned to, but not yet realized. Everything is still lazy at this point
 # l1.uop and l2.uop define a computation graph
 
-from tinygrad.engine.schedule import ScheduleItem
-schedule: List[ScheduleItem] = Tensor.schedule(l1, l2)
+from tinygrad.engine.schedule import ExecItem
+schedule: List[ExecItem] = Tensor.schedule(l1, l2)
 
 print(f"The schedule contains {len(schedule)} items.")
 for si in schedule: print(str(si)[:80])
 
 # *****
-# 4. Lower a schedule.
+# 4. Lower and run the schedule.
 
-from tinygrad.engine.realize import lower_schedule_item, ExecItem
-lowered: List[ExecItem] = [lower_schedule_item(si) for si in tqdm(schedule)]
+for si in tqdm(schedule): si.run()
 
 # *****
-# 5. Run the schedule
-
-for ei in tqdm(lowered): ei.run()
-
-# *****
-# 6. Print the weight change
+# 5. Print the weight change
 
 print("first weight change\n", l1.numpy()-l1n)
 print("second weight change\n", l2.numpy()-l2n)

docs/developer/developer.md

@@ -13,19 +13,19 @@ There's also a [doc describing speed](../developer/speed.md)
 
 Everything in [Tensor](../tensor/index.md) is syntactic sugar around constructing a graph of [UOps](../developer/uop.md).
 
-The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
+The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view. Inputs to a base can be either base or view, inputs to a view can only be a single base.
 
 ## Scheduling
 
-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ExecItem`. One `ExecItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
 
-::: tinygrad.engine.schedule.ScheduleItem
+::: tinygrad.engine.schedule.ExecItem
 
 ## Lowering
 
-The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ScheduleItem` to `ExecItem` with
+The code in [realize](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/realize.py) lowers `ExecItem` by populating its `prg` field with
 
-::: tinygrad.engine.realize.lower_schedule
+::: tinygrad.engine.realize.run_schedule
 
 There's a ton of complexity hidden behind this, see the `codegen/` directory.
 

docs/developer/layout.md

@@ -26,9 +26,9 @@ Transforms the ast into an optimized ast. This is where BEAM search and heuristi
 
 ## tinygrad/codegen
 
-Transform the optimized ast into a linearized list of UOps.
+Transform the optimized ast into a linearized and rendered program.
 
-::: tinygrad.codegen.full_rewrite
+::: tinygrad.codegen.get_program
     options:
         members: false
         show_labels: false

eslint.config.mjs

@@ -1,9 +0,0 @@
-import globals from "globals";
-import pluginJs from "@eslint/js";
-import pluginHtml from "eslint-plugin-html";
-
-export default [
-  {files: ["**/*.html"], plugins: {html: pluginHtml}, rules:{"max-len": ["error", {"code": 150}]}},
-  {languageOptions: {globals: globals.browser}},
-  pluginJs.configs.recommended,
-];

examples/coder.py

@@ -1,93 +0,0 @@
-#!/usr/bin/env python3
-import os, sys, traceback
-sys.path.append(os.getcwd())
-
-from io import StringIO
-from contextlib import redirect_stdout
-from tinygrad import Tensor, nn
-from tinygrad.helpers import Timing, colored, getenv, fetch
-from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
-from sentencepiece import SentencePieceProcessor
-
-def create_fixed_tokenizer(output_file):
-  print("creating fixed tokenizer")
-  import extra.junk.sentencepiece_model_pb2 as spb2
-  mp = spb2.ModelProto()
-  mp.ParseFromString(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/tokenizer.model?download=true").read_bytes())
-  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
-  mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
-  with open(output_file, "wb") as f:
-    f.write(mp.SerializeToString())
-
-# example:
-# echo -en "write 2+2\nwrite hello world\ny\n" | TEMP=0 python3 examples/coder.py
-
-if __name__ == "__main__":
-  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/config.json
-  with Timing("create model: "):
-    model = Transformer(4096, 14336, n_heads=32, n_layers=32, norm_eps=1e-5, vocab_size=32002, n_kv_heads=8, max_context=4096, jit=getenv("JIT", 1))
-
-  with Timing("download weights: "):
-    part1 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00001-of-00002.bin?download=true"))
-    part2 = nn.state.torch_load(fetch("https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/resolve/main/pytorch_model-00002-of-00002.bin?download=true"))
-
-  with Timing("weights -> model: "):
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part1, 32, 32, 8)), strict=False)
-    nn.state.load_state_dict(model, fix_bf16(convert_from_huggingface(part2, 32, 32, 8)), strict=False)
-
-  if not os.path.isfile("/tmp/tokenizer.model"): create_fixed_tokenizer("/tmp/tokenizer.model")
-  spp = SentencePieceProcessor(model_file="/tmp/tokenizer.model")
-
-  # https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json
-  #   "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
-  IM_END = 32000
-  IM_START = 32001
-  def encode_prompt(k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
-  def start_prompt(k): return [IM_START]+spp.encode(f"{k}\n")
-  def output(outputted, toks, color):
-    cur = spp.decode(toks)[len(outputted):]
-    sys.stdout.write(colored(cur, color))
-    sys.stdout.flush()
-    outputted += cur
-    return outputted
-
-  # *** app below this line ***
-
-  toks = [spp.bos_id()] + encode_prompt("system", "You are Quentin. Quentin is a useful assistant who writes Python code to answer questions. He keeps the code as short as possible and doesn't read from user input")
-
-  PROMPT = getenv("PROMPT", 1)
-  temperature = getenv("TEMP", 0.7)
-
-  start_pos = 0
-  outputted = output("", toks, "green")
-  turn = True
-  while 1:
-    if PROMPT:
-      toks += encode_prompt("user", input("Q: ")) + start_prompt("assistant")
-    else:
-      toks += start_prompt("user" if turn else "assistant")
-      turn = not turn
-    old_output_len = len(outputted)
-    while 1:
-      tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
-      start_pos = len(toks)
-      toks.append(tok)
-      outputted = output(outputted, toks, "blue" if not turn else "cyan")
-      if tok == IM_END: break
-      if tok == spp.eos_id(): break
-      new_output = outputted[old_output_len:]
-
-      if new_output.endswith("```") and '```python\n' in new_output:
-        python_code = new_output.split('```python\n')[1].split("```")[0]
-        # AI safety. Warning to user. Do not press y if the AI is trying to do unsafe things.
-        if input(colored(f" <-- PYTHON DETECTED, RUN IT? ", "red")).lower() == 'y':
-          my_stdout = StringIO()
-          try:
-            with redirect_stdout(my_stdout): exec(python_code)
-            result = my_stdout.getvalue()
-          except Exception as e:
-            result = ''.join(traceback.format_exception_only(e))
-          toks += spp.encode(f"\nOutput:\n```\n{result}```")
-          outputted = output(outputted, toks, "yellow")
-          old_output_len = len(outputted)
-    print("")

examples/conversation.py

@@ -1,341 +0,0 @@
-import argparse
-import multiprocessing as mp
-import os
-import re
-import sys
-import time
-from contextlib import contextmanager
-from pathlib import Path
-
-import numpy as np
-import pyaudio
-import yaml
-from llama import LLaMa
-from vits import MODELS as VITS_MODELS
-from vits import Y_LENGTH_ESTIMATE_SCALARS, HParams, Synthesizer, TextMapper, get_hparams_from_file, load_model
-from whisper import init_whisper, transcribe_waveform
-from sentencepiece import SentencePieceProcessor
-
-from tinygrad.helpers import Timing, fetch
-from tinygrad import Tensor, dtypes
-
-# Whisper constants
-RATE = 16000
-CHUNK = 1600
-
-# LLaMa constants
-IM_START = 32001
-IM_END = 32002
-
-
-# Functions for encoding prompts to chatml md
-def encode_prompt(spp, k, v): return [IM_START]+spp.encode(f"{k}\n{v}")+[IM_END]+spp.encode("\n")
-def start_prompt(spp, k): return [IM_START]+spp.encode(f"{k}\n")
-
-def chunks(lst, n):
-  for i in range(0, len(lst), n): yield lst[i:i + n]
-
-def create_fixed_tokenizer():
-  """Function needed for extending tokenizer with additional chat tokens"""
-  import extra.junk.sentencepiece_model_pb2 as spb2
-  tokenizer_path = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/tokenizer.model")
-  if SentencePieceProcessor(model_file=str(tokenizer_path)).vocab_size() != 32003:
-    print("creating fixed tokenizer")
-    mp = spb2.ModelProto()
-    mp.ParseFromString(tokenizer_path.read_bytes())
-    # https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/blob/main/added_tokens.json
-    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="[PAD]", score=0))
-    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_start|>", score=0))
-    mp.pieces.append(spb2.ModelProto.SentencePiece(piece="<|im_end|>", score=0))
-    tokenizer_path.write_bytes(mp.SerializeToString())
-  return tokenizer_path
-
-def llama_prepare(llama: LLaMa, temperature: float, pre_prompt_path: Path) -> tuple[list[int], str, str, str]:
-  """Prepares a llama model from a specified pre-prompt file"""
-  with open(str(pre_prompt_path)) as f:
-    config = yaml.safe_load(f.read())
-  toks = [llama.tokenizer.bos_id()] + encode_prompt(llama.tokenizer, "system", config["pre_prompt"].replace("\n", " "))
-  for i in config["examples"]:
-    toks += encode_prompt(llama.tokenizer, config["user_delim"], i["user_prompt"])
-    toks += encode_prompt(llama.tokenizer, config["resp_delim"], i["resp_prompt"])
-  llama.model(Tensor([toks]), 0, temperature).realize()  # NOTE: outputs are not used
-  return toks, config["user_delim"], config["resp_delim"], len(toks), llama.tokenizer.decode(toks)
-
-def llama_generate(
-  llama: LLaMa,
-  toks: list[int],
-  outputted: str,
-  prompt: str,
-  start_pos: int,
-  user_delim: str,
-  resp_delim: str,
-  temperature=0.7,
-  max_tokens=1000
-):
-  """Generates an output for the specified prompt"""
-  toks += encode_prompt(llama.tokenizer, user_delim, prompt)
-  toks += start_prompt(llama.tokenizer, resp_delim)
-
-  outputted = llama.tokenizer.decode(toks)
-  init_length = len(outputted)
-  for _ in range(max_tokens):
-    token = llama.model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
-    start_pos = len(toks)
-    toks.append(token)
-
-    cur = llama.tokenizer.decode(toks)
-
-    # Print is just for debugging
-    sys.stdout.write(cur[len(outputted):])
-    sys.stdout.flush()
-    outputted = cur
-    if toks[-1] == IM_END: break
-  else:
-    toks.append(IM_END)
-  print() # because the output is flushed
-  return outputted, start_pos, outputted[init_length:].replace("<|im_end|>", "")
-
-def tts(
-  text_to_synthesize: str,
-  synth: Synthesizer,
-  hps: HParams,
-  emotion_embedding: Path,
-  speaker_id: int,
-  model_to_use: str,
-  noise_scale: float,
-  noise_scale_w: float,
-  length_scale: float,
-  estimate_max_y_length: bool,
-  text_mapper: TextMapper,
-  model_has_multiple_speakers: bool,
-  pad_length=600,
-  vits_pad_length=1000
-):
-  if model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
-
-  # Convert the input text to a tensor.
-  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
-  init_shape = stn_tst.shape
-  assert init_shape[0] < pad_length, "text is too long"
-  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), value=1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
-  sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
-
-  # Perform inference.
-  audio_tensor = synth.infer(x_tst, x_tst_lengths, sid, noise_scale, length_scale, noise_scale_w, emotion_embedding=emotion_embedding,
-                             max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[model_to_use] if estimate_max_y_length else None, pad_length=vits_pad_length)[0, 0]
-  # Save the audio output.
-  audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
-  return audio_data
-
-def init_vits(
-  model_to_use: str,
-  emotion_path: Path,
-  speaker_id: int,
-  seed: int,
-):
-  model_config = VITS_MODELS[model_to_use]
-
-  # Load the hyperparameters from the config file.
-  hps = get_hparams_from_file(fetch(model_config[0]))
-
-  # If model has multiple speakers, validate speaker id and retrieve name if available.
-  model_has_multiple_speakers = hps.data.n_speakers > 0
-  if model_has_multiple_speakers:
-    if speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {speaker_id} is invalid for this model.")
-    if hps.__contains__("speakers"): # maps speaker ids to names
-      speakers = hps.speakers
-      if isinstance(speakers, list): speakers = {speaker: i for i, speaker in enumerate(speakers)}
-
-  # Load emotions if any. TODO: find an english model with emotions, this is untested atm.
-  emotion_embedding = None
-  if emotion_path is not None:
-    if emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(emotion_path), dtype=dtypes.int64).unsqueeze(0)
-    else: raise ValueError("Emotion path must be a .npy file.")
-
-  # Load symbols, instantiate TextMapper and clean the text.
-  if hps.__contains__("symbols"): symbols = hps.symbols
-  elif model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
-  else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
-  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
-
-  # Load the model.
-  if seed is not None:
-    Tensor.manual_seed(seed)
-    np.random.seed(seed)
-  net_g = load_model(text_mapper.symbols, hps, model_config)
-
-  return net_g, emotion_embedding, text_mapper, hps, model_has_multiple_speakers
-
-@contextmanager
-def output_stream(num_channels: int, sample_rate: int):
-  try:
-    p = pyaudio.PyAudio()
-    stream = p.open(format=pyaudio.paInt16, channels=num_channels, rate=sample_rate, output=True)
-    yield stream
-  except KeyboardInterrupt: pass
-  finally:
-    stream.stop_stream()
-    stream.close()
-    p.terminate()
-
-@contextmanager
-def log_writer():
-  try:
-    logs = []
-    yield logs
-  finally:
-    sep = "="*os.get_terminal_size()[1]
-    print(f"{sep[:-1]}\nCHAT LOG")
-    print(*logs, sep="\n")
-    print(sep)
-
-def listener(q: mp.Queue, event: mp.Event):
-  try:
-    p = pyaudio.PyAudio()
-    stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
-    did_print = False
-    while True:
-      data = stream.read(CHUNK) # read data to avoid overflow
-      if event.is_set():
-        if not did_print:
-          print("listening")
-          did_print = True
-        q.put(((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3))
-      else:
-        did_print = False
-  finally:
-    stream.stop_stream()
-    stream.close()
-    p.terminate()
-
-def mp_output_stream(q: mp.Queue, counter: mp.Value, num_channels: int, sample_rate: int):
-  with output_stream(num_channels, sample_rate) as stream:
-    while True:
-      try:
-        stream.write(q.get())
-        counter.value += 1
-      except KeyboardInterrupt:
-        break
-
-if __name__ == "__main__":
-  import nltk
-  nltk.download("punkt")
-  # Parse CLI arguments
-  parser = argparse.ArgumentParser("Have a tiny conversation with tinygrad")
-
-  # Whisper args
-  parser.add_argument("--whisper_model_name", type=str, default="tiny.en")
-
-  # LLAMA args
-  parser.add_argument("--llama_pre_prompt_path", type=Path, default=Path(__file__).parent / "conversation_data" / "pre_prompt_stacy.yaml", help="Path to yaml file which contains all pre-prompt data needed. ")
-  parser.add_argument("--llama_count", type=int, default=1000, help="Max number of tokens to generate")
-  parser.add_argument("--llama_temperature", type=float, default=0.7, help="Temperature in the softmax")
-  parser.add_argument("--llama_quantize", type=str, default=None, help="Quantize the weights to int8 or nf4 in memory")
-  parser.add_argument("--llama_model", type=Path, default=None, help="Folder with the original weights to load, or single .index.json, .safetensors or .bin file")
-  parser.add_argument("--llama_gen", type=str, default="tiny", required=False, help="Generation of the model to use")
-  parser.add_argument("--llama_size", type=str, default="1B-Chat", required=False, help="Size of model to use")
-  parser.add_argument("--llama_tokenizer", type=Path, default=None, required=False, help="Path to llama tokenizer.model")
-
-  # vits args
-  parser.add_argument("--vits_model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
-  parser.add_argument("--vits_speaker_id", type=int, default=12, help="Specify the speaker ID. Default is 6.")
-  parser.add_argument("--vits_noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
-  parser.add_argument("--vits_noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
-  parser.add_argument("--vits_length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
-  parser.add_argument("--vits_seed", type=int, default=None, help="Specify the seed (set to None if no seed). Default is 1337.")
-  parser.add_argument("--vits_num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
-  parser.add_argument("--vits_sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
-  parser.add_argument("--vits_emotion_path", type=Path, default=None, help="Specify the path to emotion reference.")
-  parser.add_argument("--vits_estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
-  parser.add_argument("--vits_vocab_path", type=Path, default=None, help="Path to the TTS vocabulary.")
-
-  # conversation args
-  parser.add_argument("--max_sentence_length", type=int, default=20, help="Max words in one sentence to pass to vits")
-
-  args = parser.parse_args()
-
-  # Init models
-  model, enc = init_whisper(args.whisper_model_name)
-  synth, emotion_embedding, text_mapper, hps, model_has_multiple_speakers = init_vits(args.vits_model_to_use, args.vits_emotion_path, args.vits_speaker_id, args.vits_seed)
-
-  # Download tinyllama chat as a default model
-  if args.llama_model is None:
-    args.llama_model = fetch("https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v0.4/resolve/main/model.safetensors", "tinyllamachat.safetensors")
-    args.llama_gen = "tiny"
-    args.llama_size = "1B-Chat"
-  # Add 3 more tokens to the tokenizer
-  if args.llama_gen == "tiny" and args.llama_size.endswith("Chat"): args.llama_tokenizer = create_fixed_tokenizer()
-  tokenizer_path = args.llama_tokenizer or args.llama_model.parent / "tokenizer.model"
-  llama = LLaMa.build(args.llama_model, tokenizer_path, args.llama_gen, args.llama_size, args.llama_quantize)
-  toks, user_delim, resp_delim, start_pos, outputted = llama_prepare(llama, args.llama_temperature, args.llama_pre_prompt_path)
-
-  # Start child process for mic input
-  q = mp.Queue()
-  is_listening_event = mp.Event()
-  p = mp.Process(target=listener, args=(q, is_listening_event,))
-  p.daemon = True
-  p.start()
-
-  # Start child process for speaker output
-  out_q = mp.Queue()
-  out_counter = mp.Value("i", 0)
-  out_p = mp.Process(target=mp_output_stream, args=(out_q, out_counter, args.vits_num_channels, hps.data.sampling_rate,))
-  out_p.daemon = True
-  out_p.start()
-
-  # JIT tts
-  for i in ["Hello, I'm a chat bot", "I am capable of doing a lot of things"]:
-    tts(
-      i, synth, hps, emotion_embedding,
-      args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
-      args.vits_noise_scale_w, args.vits_length_scale,
-      args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
-    )
-
-  # Start the pipeline
-  with log_writer() as log:
-    while True:
-      tokens = [enc._special_tokens["<|startoftranscript|>"], enc._special_tokens["<|notimestamps|>"]]
-      total = np.array([])
-      out_counter.value = 0
-
-      s = time.perf_counter()
-      is_listening_event.set()
-      prev_text = None
-      while True:
-        for _ in range(RATE // CHUNK): total = np.concatenate([total, q.get()])
-        txt = transcribe_waveform(model, enc, [total], truncate=True)
-        print(txt, end="\r")
-        if txt == "[BLANK_AUDIO]" or re.match(r"^\([\w+ ]+\)$", txt.strip()): continue
-        if prev_text is not None and prev_text == txt:
-          is_listening_event.clear()
-          break
-        prev_text = txt
-      print() # to avoid llama printing on the same line
-      log.append(f"{user_delim.capitalize()}: {txt}")
-
-      # Generate with llama
-      with Timing("llama generation: "):
-        outputted, start_pos, response = llama_generate(
-          llama, toks, outputted, txt, start_pos,
-          user_delim=user_delim, resp_delim=resp_delim, temperature=args.llama_temperature,
-          max_tokens=args.llama_count
-        )
-        log.append(f"{resp_delim.capitalize()}: {response}")
-
-      # Convert to voice
-      with Timing("tts: "):
-        sentences = nltk.sent_tokenize(response.replace('"', ""))
-        for i in sentences:
-          total = np.array([], dtype=np.int16)
-          for j in chunks(i.split(), args.max_sentence_length):
-            audio_data = tts(
-              " ".join(j), synth, hps, emotion_embedding,
-              args.vits_speaker_id, args.vits_model_to_use, args.vits_noise_scale,
-              args.vits_noise_scale_w, args.vits_length_scale,
-              args.vits_estimate_max_y_length, text_mapper, model_has_multiple_speakers
-            )
-            total = np.concatenate([total, audio_data])
-          out_q.put(total.tobytes())
-      while out_counter.value < len(sentences): continue
-      log.append(f"Total: {time.perf_counter() - s}")

examples/efficientnet.py

@@ -1,89 +0,0 @@
-# load weights from
-# https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth
-# a rough copy of
-# https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
-import sys
-import ast
-import time
-import numpy as np
-from PIL import Image
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, fetch, Timing
-from tinygrad.engine.jit import TinyJit
-from extra.models.efficientnet import EfficientNet
-np.set_printoptions(suppress=True)
-
-# TODO: you should be able to put these in the jitted function
-bias = Tensor([0.485, 0.456, 0.406])
-scale = Tensor([0.229, 0.224, 0.225])
-
-@TinyJit
-def _infer(model, img):
-  img = img.permute((2,0,1))
-  img = img / 255.0
-  img = img - bias.reshape((1,-1,1,1))
-  img = img / scale.reshape((1,-1,1,1))
-  return model.forward(img).realize()
-
-def infer(model, img):
-  # preprocess image
-  aspect_ratio = img.size[0] / img.size[1]
-  img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
-
-  img = np.array(img)
-  y0,x0=(np.asarray(img.shape)[:2]-224)//2
-  retimg = img = img[y0:y0+224, x0:x0+224]
-
-  # if you want to look at the image
-  """
-  import matplotlib.pyplot as plt
-  plt.imshow(img)
-  plt.show()
-  """
-
-  # run the net
-  out = _infer(model, Tensor(img.astype("float32"))).numpy()
-
-  # if you want to look at the outputs
-  """
-  import matplotlib.pyplot as plt
-  plt.plot(out[0])
-  plt.show()
-  """
-  return out, retimg
-
-if __name__ == "__main__":
-  # instantiate my net
-  model = EfficientNet(getenv("NUM", 0))
-  model.load_from_pretrained()
-
-  # category labels
-  lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
-
-  # load image and preprocess
-  url = sys.argv[1] if len(sys.argv) >= 2 else "https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/showcase/stable_diffusion_by_tinygrad.jpg"
-  if url == 'webcam':
-    import cv2
-    cap = cv2.VideoCapture(0)
-    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
-    while 1:
-      _ = cap.grab() # discard one frame to circumvent capture buffering
-      ret, frame = cap.read()
-      img = Image.fromarray(frame[:, :, [2,1,0]])
-      lt = time.monotonic_ns()
-      out, retimg = infer(model, img)
-      print(f"{(time.monotonic_ns()-lt)*1e-6:7.2f} ms", np.argmax(out), np.max(out), lbls[np.argmax(out)])
-      SCALE = 3
-      simg = cv2.resize(retimg, (224*SCALE, 224*SCALE))
-      retimg = cv2.cvtColor(simg, cv2.COLOR_RGB2BGR)
-      cv2.imshow('capture', retimg)
-      if cv2.waitKey(1) & 0xFF == ord('q'):
-        break
-    cap.release()
-    cv2.destroyAllWindows()
-  else:
-    img = Image.open(fetch(url))
-    for i in range(getenv("CNT", 1)):
-      with Timing("did inference in "):
-        out, _ = infer(model, img)
-        print(np.argmax(out), np.max(out), lbls[np.argmax(out)])

examples/flux1.py

@@ -1,498 +0,0 @@
-# pip3 install sentencepiece
-
-# This file incorporates code from the following:
-# Github Name                    | License | Link
-# black-forest-labs/flux         | Apache  | https://github.com/black-forest-labs/flux/tree/main/model_licenses
-
-from tinygrad import Tensor, nn, dtypes, TinyJit
-from tinygrad.nn.state import safe_load, load_state_dict
-from tinygrad.helpers import fetch, tqdm, colored
-from sdxl import FirstStage
-from extra.models.clip import FrozenClosedClipEmbedder
-from extra.models.t5 import T5Embedder
-import numpy as np
-
-import math, time, argparse, tempfile
-from typing import List, Dict, Optional, Union, Tuple, Callable
-from dataclasses import dataclass
-from pathlib import Path
-from PIL import Image
-
-urls:dict = {
-  "flux-schnell": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/flux1-schnell.safetensors",
-  "flux-dev": "https://huggingface.co/camenduru/FLUX.1-dev/resolve/main/flux1-dev.sft",
-  "ae": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/ae.safetensors",
-  "T5_1_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00001-of-00002.safetensors",
-  "T5_2_of_2": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder_2/model-00002-of-00002.safetensors",
-  "T5_tokenizer": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/tokenizer_2/spiece.model",
-  "clip": "https://huggingface.co/black-forest-labs/FLUX.1-schnell/resolve/main/text_encoder/model.safetensors"
-}
-
-def tensor_identity(x:Tensor) -> Tensor: return x
-
-class AutoEncoder:
-  def __init__(self, scale_factor:float, shift_factor:float):
-    self.decoder = FirstStage.Decoder(128, 3, 3, 16, [1, 2, 4, 4], 2, 256)
-    self.scale_factor = scale_factor
-    self.shift_factor = shift_factor
-
-  def decode(self, z:Tensor) -> Tensor:
-    z = z / self.scale_factor + self.shift_factor
-    return self.decoder(z)
-
-# Conditioner
-class ClipEmbedder(FrozenClosedClipEmbedder):
-  def __call__(self, texts:Union[str, List[str], Tensor]) -> Tensor:
-    if isinstance(texts, str): texts = [texts]
-    assert isinstance(texts, (list,tuple)), f"expected list of strings, got {type(texts).__name__}"
-    tokens = Tensor.cat(*[Tensor(self.tokenizer.encode(text)) for text in texts], dim=0)
-    return self.transformer.text_model(tokens.reshape(len(texts),-1))[:, tokens.argmax(-1)]
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
-def attention(q:Tensor, k:Tensor, v:Tensor, pe:Tensor) -> Tensor:
-  q, k = apply_rope(q, k, pe)
-  x = Tensor.scaled_dot_product_attention(q, k, v)
-  return x.rearrange("B H L D -> B L (H D)")
-
-def rope(pos:Tensor, dim:int, theta:int) -> Tensor:
-  assert dim % 2 == 0
-  scale = Tensor.arange(0, dim, 2, dtype=dtypes.float32, device=pos.device) / dim # NOTE: this is torch.float64 in reference implementation
-  omega = 1.0 / (theta**scale)
-  out = Tensor.einsum("...n,d->...nd", pos, omega)
-  out = Tensor.stack(Tensor.cos(out), -Tensor.sin(out), Tensor.sin(out), Tensor.cos(out), dim=-1)
-  out = out.rearrange("b n d (i j) -> b n d i j", i=2, j=2)
-  return out.float()
-
-def apply_rope(xq:Tensor, xk:Tensor, freqs_cis:Tensor) -> Tuple[Tensor, Tensor]:
-  xq_ = xq.float().reshape(*xq.shape[:-1], -1, 1, 2)
-  xk_ = xk.float().reshape(*xk.shape[:-1], -1, 1, 2)
-  xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
-  xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
-  return xq_out.reshape(*xq.shape).cast(xq.dtype), xk_out.reshape(*xk.shape).cast(xk.dtype)
-
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
-class EmbedND:
-  def __init__(self, dim:int, theta:int, axes_dim:List[int]):
-    self.dim = dim
-    self.theta = theta
-    self.axes_dim = axes_dim
-
-  def __call__(self, ids:Tensor) -> Tensor:
-    n_axes = ids.shape[-1]
-    emb = Tensor.cat(*[rope(ids[..., i], self.axes_dim[i], self.theta) for i in range(n_axes)], dim=-3)
-    return emb.unsqueeze(1)
-
-class MLPEmbedder:
-  def __init__(self, in_dim:int, hidden_dim:int):
-    self.in_layer = nn.Linear(in_dim, hidden_dim, bias=True)
-    self.out_layer = nn.Linear(hidden_dim, hidden_dim, bias=True)
-
-  def __call__(self, x:Tensor) -> Tensor:
-    return self.out_layer(self.in_layer(x).silu())
-
-class QKNorm:
-  def __init__(self, dim:int):
-    self.query_norm = nn.RMSNorm(dim)
-    self.key_norm = nn.RMSNorm(dim)
-
-  def __call__(self, q:Tensor, k:Tensor) -> Tuple[Tensor, Tensor]:
-    return self.query_norm(q), self.key_norm(k)
-
-class SelfAttention:
-  def __init__(self, dim:int, num_heads:int = 8, qkv_bias:bool = False):
-    self.num_heads = num_heads
-    head_dim = dim // num_heads
-
-    self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
-    self.norm = QKNorm(head_dim)
-    self.proj = nn.Linear(dim, dim)
-
-  def __call__(self, x:Tensor, pe:Tensor) -> Tensor:
-    qkv = self.qkv(x)
-    q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-    q, k = self.norm(q, k)
-    x = attention(q, k, v, pe=pe)
-    return self.proj(x)
-
-@dataclass
-class ModulationOut:
-  shift:Tensor
-  scale:Tensor
-  gate:Tensor
-
-class Modulation:
-  def __init__(self, dim:int, double:bool):
-    self.is_double = double
-    self.multiplier = 6 if double else 3
-    self.lin = nn.Linear(dim, self.multiplier * dim, bias=True)
-
-  def __call__(self, vec:Tensor) -> Tuple[ModulationOut, Optional[ModulationOut]]:
-    out = self.lin(vec.silu())[:, None, :].chunk(self.multiplier, dim=-1)
-    return ModulationOut(*out[:3]), ModulationOut(*out[3:]) if self.is_double else None
-
-class DoubleStreamBlock:
-  def __init__(self, hidden_size:int, num_heads:int, mlp_ratio:float, qkv_bias:bool = False):
-    mlp_hidden_dim = int(hidden_size * mlp_ratio)
-    self.num_heads = num_heads
-    self.hidden_size = hidden_size
-    self.img_mod = Modulation(hidden_size, double=True)
-    self.img_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-    self.img_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-    self.img_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-    self.img_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
-
-    self.txt_mod = Modulation(hidden_size, double=True)
-    self.txt_norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-    self.txt_attn = SelfAttention(dim=hidden_size, num_heads=num_heads, qkv_bias=qkv_bias)
-
-    self.txt_norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-    self.txt_mlp = [nn.Linear(hidden_size, mlp_hidden_dim, bias=True), Tensor.gelu, nn.Linear(mlp_hidden_dim, hidden_size, bias=True)]
-
-  def __call__(self, img:Tensor, txt:Tensor, vec:Tensor, pe:Tensor) -> tuple[Tensor, Tensor]:
-    img_mod1, img_mod2 = self.img_mod(vec)
-    txt_mod1, txt_mod2 = self.txt_mod(vec)
-    assert img_mod2 is not None and txt_mod2 is not None
-    # prepare image for attention
-    img_modulated = self.img_norm1(img)
-    img_modulated = (1 + img_mod1.scale) * img_modulated + img_mod1.shift
-    img_qkv = self.img_attn.qkv(img_modulated)
-    img_q, img_k, img_v = img_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-    img_q, img_k = self.img_attn.norm(img_q, img_k)
-
-    # prepare txt for attention
-    txt_modulated = self.txt_norm1(txt)
-    txt_modulated = (1 + txt_mod1.scale) * txt_modulated + txt_mod1.shift
-    txt_qkv = self.txt_attn.qkv(txt_modulated)
-    txt_q, txt_k, txt_v = txt_qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-    txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k)
-
-    # run actual attention
-    q = Tensor.cat(txt_q, img_q, dim=2)
-    k = Tensor.cat(txt_k, img_k, dim=2)
-    v = Tensor.cat(txt_v, img_v, dim=2)
-
-    attn = attention(q, k, v, pe=pe)
-    txt_attn, img_attn = attn[:, : txt.shape[1]], attn[:, txt.shape[1] :]
-
-    # calculate the img bloks
-    img = img + img_mod1.gate * self.img_attn.proj(img_attn)
-    img = img + img_mod2.gate * ((1 + img_mod2.scale) * self.img_norm2(img) + img_mod2.shift).sequential(self.img_mlp)
-
-    # calculate the txt bloks
-    txt = txt + txt_mod1.gate * self.txt_attn.proj(txt_attn)
-    txt = txt + txt_mod2.gate * ((1 + txt_mod2.scale) * self.txt_norm2(txt) + txt_mod2.shift).sequential(self.txt_mlp)
-    return img, txt
-
-
-class SingleStreamBlock:
-  """
-  A DiT block with parallel linear layers as described in
-  https://arxiv.org/abs/2302.05442 and adapted modulation interface.
-  """
-
-  def __init__(self,hidden_size:int, num_heads:int, mlp_ratio:float=4.0, qk_scale:Optional[float]=None):
-    self.hidden_dim = hidden_size
-    self.num_heads = num_heads
-    head_dim = hidden_size // num_heads
-    self.scale = qk_scale or head_dim**-0.5
-
-    self.mlp_hidden_dim = int(hidden_size * mlp_ratio)
-    # qkv and mlp_in
-    self.linear1 = nn.Linear(hidden_size, hidden_size * 3 + self.mlp_hidden_dim)
-    # proj and mlp_out
-    self.linear2 = nn.Linear(hidden_size + self.mlp_hidden_dim, hidden_size)
-
-    self.norm = QKNorm(head_dim)
-
-    self.hidden_size = hidden_size
-    self.pre_norm = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-
-    self.mlp_act = Tensor.gelu
-    self.modulation = Modulation(hidden_size, double=False)
-
-  def __call__(self, x:Tensor, vec:Tensor, pe:Tensor) -> Tensor:
-    mod, _ = self.modulation(vec)
-    x_mod = (1 + mod.scale) * self.pre_norm(x) + mod.shift
-    qkv, mlp = Tensor.split(self.linear1(x_mod), [3 * self.hidden_size, self.mlp_hidden_dim], dim=-1)
-    q, k, v = qkv.rearrange("B L (K H D) -> K B H L D", K=3, H=self.num_heads)
-    q, k = self.norm(q, k)
-
-    # compute attention
-    attn = attention(q, k, v, pe=pe)
-    # compute activation in mlp stream, cat again and run second linear layer
-    output = self.linear2(Tensor.cat(attn, self.mlp_act(mlp), dim=2))
-    return x + mod.gate * output
-
-
-class LastLayer:
-  def __init__(self, hidden_size:int, patch_size:int, out_channels:int):
-    self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-    self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-    self.adaLN_modulation:List[Callable[[Tensor], Tensor]] = [Tensor.silu, nn.Linear(hidden_size, 2 * hidden_size, bias=True)]
-
-  def __call__(self, x:Tensor, vec:Tensor) -> Tensor:
-    shift, scale = vec.sequential(self.adaLN_modulation).chunk(2, dim=1)
-    x = (1 + scale[:, None, :]) * self.norm_final(x) + shift[:, None, :]
-    return self.linear(x)
-
-def timestep_embedding(t:Tensor, dim:int, max_period:int=10000, time_factor:float=1000.0) -> Tensor:
-  """
-  Create sinusoidal timestep embeddings.
-  :param t: a 1-D Tensor of N indices, one per batch element.
-                    These may be fractional.
-  :param dim: the dimension of the output.
-  :param max_period: controls the minimum frequency of the embeddings.
-  :return: an (N, D) Tensor of positional embeddings.
-  """
-  t = time_factor * t
-  half = dim // 2
-  freqs = Tensor.exp(-math.log(max_period) * Tensor.arange(0, stop=half, dtype=dtypes.float32) / half).to(t.device)
-
-  args = t[:, None].float() * freqs[None]
-  embedding = Tensor.cat(Tensor.cos(args), Tensor.sin(args), dim=-1)
-  if dim % 2:  embedding = Tensor.cat(*[embedding, Tensor.zeros_like(embedding[:, :1])], dim=-1)
-  if Tensor.is_floating_point(t):  embedding = embedding.cast(t.dtype)
-  return embedding
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/model.py
-class Flux:
-  """
-  Transformer model for flow matching on sequences.
-  """
-
-  def __init__(
-      self,
-      guidance_embed:bool,
-      in_channels:int = 64,
-      vec_in_dim:int = 768,
-      context_in_dim:int = 4096,
-      hidden_size:int = 3072,
-      mlp_ratio:float = 4.0,
-      num_heads:int = 24,
-      depth:int = 19,
-      depth_single_blocks:int = 38,
-      axes_dim:Optional[List[int]] = None,
-      theta:int = 10_000,
-      qkv_bias:bool = True,
-      ):
-
-    axes_dim = axes_dim or [16, 56, 56]
-    self.guidance_embed = guidance_embed
-    self.in_channels = in_channels
-    self.out_channels = self.in_channels
-    if hidden_size % num_heads != 0:
-      raise ValueError(f"Hidden size {hidden_size} must be divisible by num_heads {num_heads}")
-    pe_dim = hidden_size // num_heads
-    if sum(axes_dim) != pe_dim:
-      raise ValueError(f"Got {axes_dim} but expected positional dim {pe_dim}")
-    self.hidden_size = hidden_size
-    self.num_heads = num_heads
-    self.pe_embedder = EmbedND(dim=pe_dim, theta=theta, axes_dim=axes_dim)
-    self.img_in = nn.Linear(self.in_channels, self.hidden_size, bias=True)
-    self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size)
-    self.vector_in = MLPEmbedder(vec_in_dim, self.hidden_size)
-    self.guidance_in:Callable[[Tensor], Tensor] = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size) if guidance_embed else tensor_identity
-    self.txt_in = nn.Linear(context_in_dim, self.hidden_size)
-
-    self.double_blocks = [DoubleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias) for _ in range(depth)]
-    self.single_blocks = [SingleStreamBlock(self.hidden_size, self.num_heads, mlp_ratio=mlp_ratio) for _ in range(depth_single_blocks)]
-    self.final_layer = LastLayer(self.hidden_size, 1, self.out_channels)
-
-  def __call__(self, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, timesteps:Tensor, y:Tensor, guidance:Optional[Tensor] = None) -> Tensor:
-    if img.ndim != 3 or txt.ndim != 3:
-      raise ValueError("Input img and txt tensors must have 3 dimensions.")
-    # running on sequences img
-    img = self.img_in(img)
-    vec = self.time_in(timestep_embedding(timesteps, 256))
-    if self.guidance_embed:
-      if guidance is None:
-        raise ValueError("Didn't get guidance strength for guidance distilled model.")
-      vec = vec + self.guidance_in(timestep_embedding(guidance, 256))
-    vec = vec + self.vector_in(y)
-    txt = self.txt_in(txt)
-    ids = Tensor.cat(txt_ids, img_ids, dim=1)
-    pe = self.pe_embedder(ids)
-    for double_block in self.double_blocks:
-      img, txt = double_block(img=img, txt=txt, vec=vec, pe=pe)
-
-    img = Tensor.cat(txt, img, dim=1)
-    for single_block in self.single_blocks:
-      img = single_block(img, vec=vec, pe=pe)
-
-    img = img[:, txt.shape[1] :, ...]
-
-    return self.final_layer(img, vec)  # (N, T, patch_size ** 2 * out_channels)
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/util.py
-def load_flow_model(name:str, model_path:str):
-  # Loading Flux
-  print("Init model")
-  model = Flux(guidance_embed=(name != "flux-schnell"))
-  if not model_path: model_path = fetch(urls[name])
-  state_dict = {k.replace("scale", "weight"): v for k, v in safe_load(model_path).items()}
-  load_state_dict(model, state_dict)
-  return model
-
-def load_T5(max_length:int=512):
-  # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
-  print("Init T5")
-  T5 = T5Embedder(max_length, fetch(urls["T5_tokenizer"]))
-  pt_1 = fetch(urls["T5_1_of_2"])
-  pt_2 = fetch(urls["T5_2_of_2"])
-  load_state_dict(T5.encoder, safe_load(pt_1) | safe_load(pt_2), strict=False)
-  return T5
-
-def load_clip():
-  print("Init Clip")
-  clip = ClipEmbedder()
-  load_state_dict(clip.transformer, safe_load(fetch(urls["clip"])))
-  return clip
-
-def load_ae() -> AutoEncoder:
-  # Loading the autoencoder
-  print("Init AE")
-  ae = AutoEncoder(0.3611, 0.1159)
-  load_state_dict(ae, safe_load(fetch(urls["ae"])))
-  return ae
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/sampling.py
-def prepare(T5:T5Embedder, clip:ClipEmbedder, img:Tensor, prompt:Union[str, List[str]]) -> Dict[str, Tensor]:
-  bs, _, h, w = img.shape
-  if bs == 1 and not isinstance(prompt, str):
-    bs = len(prompt)
-
-  img = img.rearrange("b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
-  if img.shape[0] == 1 and bs > 1:
-    img = img.expand((bs, *img.shape[1:]))
-
-  img_ids = Tensor.zeros(h // 2, w // 2, 3).contiguous()
-  img_ids[..., 1] = img_ids[..., 1] + Tensor.arange(h // 2)[:, None]
-  img_ids[..., 2] = img_ids[..., 2] + Tensor.arange(w // 2)[None, :]
-  img_ids = img_ids.rearrange("h w c -> 1 (h w) c")
-  img_ids = img_ids.expand((bs, *img_ids.shape[1:]))
-
-  if isinstance(prompt, str):
-    prompt = [prompt]
-  txt = T5(prompt).realize()
-  if txt.shape[0] == 1 and bs > 1:
-    txt = txt.expand((bs, *txt.shape[1:]))
-  txt_ids = Tensor.zeros(bs, txt.shape[1], 3)
-
-  vec = clip(prompt).realize()
-  if vec.shape[0] == 1 and bs > 1:
-    vec = vec.expand((bs, *vec.shape[1:]))
-
-  return {"img": img, "img_ids": img_ids.to(img.device), "txt": txt.to(img.device), "txt_ids": txt_ids.to(img.device), "vec": vec.to(img.device)}
-
-
-def get_schedule(num_steps:int, image_seq_len:int, base_shift:float=0.5, max_shift:float=1.15, shift:bool=True) -> List[float]:
-  # extra step for zero
-  step_size = -1.0 / num_steps
-  timesteps = Tensor.arange(1, 0 + step_size, step_size)
-
-  # shifting the schedule to favor high timesteps for higher signal images
-  if shift:
-    # estimate mu based on linear estimation between two points
-    mu = 0.5 + (max_shift - base_shift) * (image_seq_len - 256) / (4096 - 256)
-    timesteps = math.exp(mu) / (math.exp(mu) + (1 / timesteps - 1))
-  return timesteps.tolist()
-
-@TinyJit
-def run(model, *args): return model(*args).realize()
-
-def denoise(model, img:Tensor, img_ids:Tensor, txt:Tensor, txt_ids:Tensor, vec:Tensor, timesteps:List[float], guidance:float=4.0) -> Tensor:
-  # this is ignored for schnell
-  guidance_vec = Tensor((guidance,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
-  for t_curr, t_prev in tqdm(list(zip(timesteps[:-1], timesteps[1:])), "Denoising"):
-    t_vec = Tensor((t_curr,), device=img.device, dtype=img.dtype).expand((img.shape[0],))
-    pred = run(model, img, img_ids, txt, txt_ids, t_vec, vec, guidance_vec)
-    img = img + (t_prev - t_curr) * pred
-
-  return img
-
-def unpack(x:Tensor, height:int, width:int) -> Tensor:
-  return x.rearrange("b (h w) (c ph pw) -> b c (h ph) (w pw)", h=math.ceil(height / 16), w=math.ceil(width / 16), ph=2, pw=2)
-
-# https://github.com/black-forest-labs/flux/blob/main/src/flux/cli.py
-if __name__ == "__main__":
-  default_prompt = "bananas and a can of coke"
-  parser = argparse.ArgumentParser(description="Run Flux.1", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-
-  parser.add_argument("--name",       type=str,   default="flux-schnell", help="Name of the model to load")
-  parser.add_argument("--model_path", type=str,   default="",             help="path of the model file")
-  parser.add_argument("--width",      type=int,   default=512,            help="width of the sample in pixels (should be a multiple of 16)")
-  parser.add_argument("--height",     type=int,   default=512,            help="height of the sample in pixels (should be a multiple of 16)")
-  parser.add_argument("--seed",       type=int,   default=None,           help="Set a seed for sampling")
-  parser.add_argument("--prompt",     type=str,   default=default_prompt, help="Prompt used for sampling")
-  parser.add_argument('--out',        type=str,   default=Path(tempfile.gettempdir()) / "rendered.png", help="Output filename")
-  parser.add_argument("--num_steps",  type=int,   default=None,           help="number of sampling steps (default 4 for schnell, 50 for guidance distilled)") #noqa:E501
-  parser.add_argument("--guidance",   type=float, default=3.5,            help="guidance value used for guidance distillation")
-  parser.add_argument("--output_dir", type=str,   default="output",       help="output directory")
-  args = parser.parse_args()
-
-  if args.name not in ["flux-schnell", "flux-dev"]:
-    raise ValueError(f"Got unknown model name: {args.name}, chose from flux-schnell and flux-dev")
-
-  if args.num_steps is None:
-    args.num_steps = 4 if args.name == "flux-schnell" else 50
-
-  # allow for packing and conversion to latent space
-  height = 16 * (args.height // 16)
-  width = 16 * (args.width // 16)
-
-  if args.seed is None: args.seed = Tensor._seed
-  else: Tensor.manual_seed(args.seed)
-
-  print(f"Generating with seed {args.seed}:\n{args.prompt}")
-  t0 = time.perf_counter()
-
-  # prepare input noise
-  x = Tensor.randn(1, 16, 2 * math.ceil(height / 16), 2 * math.ceil(width / 16), dtype="bfloat16")
-
-  # load text embedders
-  T5 = load_T5(max_length=256 if args.name == "flux-schnell" else 512)
-  clip = load_clip()
-
-  # embed text to get inputs for model
-  inp = prepare(T5, clip, x, prompt=args.prompt)
-  timesteps = get_schedule(args.num_steps, inp["img"].shape[1], shift=(args.name != "flux-schnell"))
-
-  # done with text embedders
-  del T5, clip
-
-  # load model
-  model = load_flow_model(args.name, args.model_path)
-
-  # denoise initial noise
-  x = denoise(model, **inp, timesteps=timesteps, guidance=args.guidance)
-
-  # done with model
-  del model, run
-
-  # load autoencoder
-  ae = load_ae()
-
-  # decode latents to pixel space
-  x = unpack(x.float(), height, width)
-  x = ae.decode(x).realize()
-
-  t1 = time.perf_counter()
-  print(f"Done in {t1 - t0:.1f}s. Saving {args.out}")
-
-  # bring into PIL format and save
-  x = x.clamp(-1, 1)
-  x = x[0].rearrange("c h w -> h w c")
-  x = (127.5 * (x + 1.0)).cast("uint8")
-
-  img = Image.fromarray(x.numpy())
-
-  img.save(args.out)
-
-  # validation!
-  if args.prompt == default_prompt and args.name=="flux-schnell" and args.seed == 0 and args.width == args.height == 512:
-    ref_image = Tensor(np.array(Image.open("examples/flux1_seed0.png")))
-    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
-    assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
-    print(colored(f"output validated with {distance=}", "green"))

examples/mask_rcnn.py

@@ -1,299 +0,0 @@
-from extra.models.mask_rcnn import MaskRCNN
-from extra.models.resnet import ResNet
-from extra.models.mask_rcnn import BoxList
-from torch.nn import functional as F
-from torchvision import transforms as T
-from torchvision.transforms import functional as Ft
-import random
-from tinygrad.tensor import Tensor
-from PIL import Image
-import numpy as np
-import torch
-import argparse
-import cv2
-
-
-class Resize:
-  def __init__(self, min_size, max_size):
-    if not isinstance(min_size, (list, tuple)):
-      min_size = (min_size,)
-    self.min_size = min_size
-    self.max_size = max_size
-
-  # modified from torchvision to add support for max size
-  def get_size(self, image_size):
-    w, h = image_size
-    size = random.choice(self.min_size)
-    max_size = self.max_size
-    if max_size is not None:
-      min_original_size = float(min((w, h)))
-      max_original_size = float(max((w, h)))
-      if max_original_size / min_original_size * size > max_size:
-        size = int(round(max_size * min_original_size / max_original_size))
-
-      if (w <= h and w == size) or (h <= w and h == size):
-        return (h, w)
-
-      if w < h:
-        ow = size
-        oh = int(size * h / w)
-      else:
-        oh = size
-        ow = int(size * w / h)
-
-      return (oh, ow)
-
-  def __call__(self, image):
-    size = self.get_size(image.size)
-    image = Ft.resize(image, size)
-    return image
-
-
-class Normalize:
-  def __init__(self, mean, std, to_bgr255=True):
-    self.mean = mean
-    self.std = std
-    self.to_bgr255 = to_bgr255
-
-  def __call__(self, image):
-    if self.to_bgr255:
-      image = image[[2, 1, 0]] * 255
-    else:
-      image = image[[0, 1, 2]] * 255
-    image = Ft.normalize(image, mean=self.mean, std=self.std)
-    return image
-
-transforms = lambda size_scale: T.Compose(
-  [
-    Resize(int(800*size_scale), int(1333*size_scale)),
-    T.ToTensor(),
-    Normalize(
-      mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], to_bgr255=True
-    ),
-  ]
-)
-
-def expand_boxes(boxes, scale):
-  w_half = (boxes[:, 2] - boxes[:, 0]) * .5
-  h_half = (boxes[:, 3] - boxes[:, 1]) * .5
-  x_c = (boxes[:, 2] + boxes[:, 0]) * .5
-  y_c = (boxes[:, 3] + boxes[:, 1]) * .5
-
-  w_half *= scale
-  h_half *= scale
-
-  boxes_exp = torch.zeros_like(boxes)
-  boxes_exp[:, 0] = x_c - w_half
-  boxes_exp[:, 2] = x_c + w_half
-  boxes_exp[:, 1] = y_c - h_half
-  boxes_exp[:, 3] = y_c + h_half
-  return boxes_exp
-
-
-def expand_masks(mask, padding):
-  N = mask.shape[0]
-  M = mask.shape[-1]
-  pad2 = 2 * padding
-  scale = float(M + pad2) / M
-  padded_mask = mask.new_zeros((N, 1, M + pad2, M + pad2))
-  padded_mask[:, :, padding:-padding, padding:-padding] = mask
-  return padded_mask, scale
-
-
-def paste_mask_in_image(mask, box, im_h, im_w, thresh=0.5, padding=1):
-  # TODO: remove torch
-  mask = torch.tensor(mask.numpy())
-  box = torch.tensor(box.numpy())
-  padded_mask, scale = expand_masks(mask[None], padding=padding)
-  mask = padded_mask[0, 0]
-  box = expand_boxes(box[None], scale)[0]
-  box = box.to(dtype=torch.int32)
-
-  TO_REMOVE = 1
-  w = int(box[2] - box[0] + TO_REMOVE)
-  h = int(box[3] - box[1] + TO_REMOVE)
-  w = max(w, 1)
-  h = max(h, 1)
-
-  mask = mask.expand((1, 1, -1, -1))
-
-  mask = mask.to(torch.float32)
-  mask = F.interpolate(mask, size=(h, w), mode='bilinear', align_corners=False)
-  mask = mask[0][0]
-
-  if thresh >= 0:
-    mask = mask > thresh
-  else:
-    mask = (mask * 255).to(torch.uint8)
-
-  im_mask = torch.zeros((im_h, im_w), dtype=torch.uint8)
-  x_0 = max(box[0], 0)
-  x_1 = min(box[2] + 1, im_w)
-  y_0 = max(box[1], 0)
-  y_1 = min(box[3] + 1, im_h)
-
-  im_mask[y_0:y_1, x_0:x_1] = mask[
-                              (y_0 - box[1]): (y_1 - box[1]), (x_0 - box[0]): (x_1 - box[0])
-                              ]
-  return im_mask
-
-
-class Masker:
-  def __init__(self, threshold=0.5, padding=1):
-    self.threshold = threshold
-    self.padding = padding
-
-  def forward_single_image(self, masks, boxes):
-    boxes = boxes.convert("xyxy")
-    im_w, im_h = boxes.size
-    res = [
-      paste_mask_in_image(mask[0], box, im_h, im_w, self.threshold, self.padding)
-      for mask, box in zip(masks, boxes.bbox)
-    ]
-    if len(res) > 0:
-      res = torch.stack(*res, dim=0)[:, None]
-    else:
-      res = masks.new_empty((0, 1, masks.shape[-2], masks.shape[-1]))
-    return Tensor(res.numpy())
-
-  def __call__(self, masks, boxes):
-    if isinstance(boxes, BoxList):
-      boxes = [boxes]
-
-    results = []
-    for mask, box in zip(masks, boxes):
-      result = self.forward_single_image(mask, box)
-      results.append(result)
-    return results
-
-
-masker = Masker(threshold=0.5, padding=1)
-
-def select_top_predictions(predictions, confidence_threshold=0.9):
-  scores = predictions.get_field("scores").numpy()
-  keep = [idx for idx, score in enumerate(scores) if score > confidence_threshold]
-  return predictions[keep]
-
-def compute_prediction(original_image, model, confidence_threshold, size_scale=1.0):
-  image = transforms(size_scale)(original_image).numpy()
-  image = Tensor(image, requires_grad=False)
-  predictions = model(image)
-  prediction = predictions[0]
-  prediction = select_top_predictions(prediction, confidence_threshold)
-  width, height = original_image.size
-  prediction = prediction.resize((width, height))
-
-  if prediction.has_field("mask"):
-    masks = prediction.get_field("mask")
-    masks = masker([masks], [prediction])[0]
-    prediction.add_field("mask", masks)
-  return prediction
-
-def compute_prediction_batched(batch, model, size_scale=1.0):
-  imgs = []
-  for img in batch:
-    imgs.append(transforms(size_scale)(img).numpy())
-  image = [Tensor(image, requires_grad=False) for image in imgs]
-  predictions = model(image)
-  del image
-  return predictions
-
-palette = np.array([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
-
-def findContours(*args, **kwargs):
-  if cv2.__version__.startswith('4'):
-    contours, hierarchy = cv2.findContours(*args, **kwargs)
-  elif cv2.__version__.startswith('3'):
-    _, contours, hierarchy = cv2.findContours(*args, **kwargs)
-  return contours, hierarchy
-
-def compute_colors_for_labels(labels):
-  l = labels[:, None]
-  colors = l * palette
-  colors = (colors % 255).astype("uint8")
-  return colors
-
-def overlay_mask(image, predictions):
-  image = np.asarray(image)
-  masks = predictions.get_field("mask").numpy()
-  labels = predictions.get_field("labels").numpy()
-
-  colors = compute_colors_for_labels(labels).tolist()
-
-  for mask, color in zip(masks, colors):
-    thresh = mask[0, :, :, None]
-    contours, hierarchy = findContours(
-        thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
-    )
-    image = cv2.drawContours(image, contours, -1, color, 3)
-
-  composite = image
-
-  return composite
-
-CATEGORIES = [
-    "__background", "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
-    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant",
-    "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard",
-    "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle",
-    "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
-    "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table",
-    "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster",
-    "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
-]
-
-def overlay_boxes(image, predictions):
-  labels = predictions.get_field("labels").numpy()
-  boxes = predictions.bbox
-  image = np.asarray(image)
-  colors = compute_colors_for_labels(labels).tolist()
-
-  for box, color in zip(boxes, colors):
-    box = torch.tensor(box.numpy())
-    box = box.to(torch.int64)
-    top_left, bottom_right = box[:2].tolist(), box[2:].tolist()
-    image = cv2.rectangle(
-        image, tuple(top_left), tuple(bottom_right), tuple(color), 1
-    )
-
-  return image
-
-def overlay_class_names(image, predictions):
-  scores = predictions.get_field("scores").numpy().tolist()
-  labels = predictions.get_field("labels").numpy().tolist()
-  labels = [CATEGORIES[int(i)] for i in labels]
-  boxes = predictions.bbox.numpy()
-  image = np.asarray(image)
-  template = "{}: {:.2f}"
-  for box, score, label in zip(boxes, scores, labels):
-    x, y = box[:2]
-    s = template.format(label, score)
-    x, y = int(x), int(y)
-    cv2.putText(
-        image, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .5, (255, 255, 255), 1
-    )
-
-  return image
-
-
-if __name__ == '__main__':
-  parser = argparse.ArgumentParser(description='Run MaskRCNN', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument('--image', type=str, help="Path of the image to run")
-  parser.add_argument('--threshold', type=float, default=0.7, help="Detector threshold")
-  parser.add_argument('--size_scale', type=float, default=1.0, help="Image resize multiplier")
-  parser.add_argument('--out', type=str, default="/tmp/rendered.png", help="Output filename")
-  args = parser.parse_args()
-
-  resnet = ResNet(50, num_classes=None, stride_in_1x1=True)
-  model_tiny = MaskRCNN(resnet)
-  model_tiny.load_from_pretrained()
-  img = Image.open(args.image)
-  top_result_tiny = compute_prediction(img, model_tiny, confidence_threshold=args.threshold, size_scale=args.size_scale)
-  bbox_image = overlay_boxes(img, top_result_tiny)
-  mask_image = overlay_mask(bbox_image, top_result_tiny)
-  final_image = overlay_class_names(mask_image, top_result_tiny)
-
-  im = Image.fromarray(final_image)
-  print(f"saving {args.out}")
-  im.save(args.out)
-  im.show()

examples/mlperf/dataloader.py

@@ -763,48 +763,26 @@ class BlendedGPTDataset:
 
     return dataset_idx, dataset_sample_idx
 
-def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
+def get_llama3_dataset(samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False) -> BlendedGPTDataset:
+  if small:
+    if val:
+      return BlendedGPTDataset(
+        [base_dir / "c4-validation-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
+    return BlendedGPTDataset(
+      [base_dir / "c4-train.en_6_text_document"], [1.0], samples, seqlen, seed, shuffle=True)
   if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "validation" / "c4-validationn-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-      base_dir / "c4-train.en_7_text_document",
-    ], [
-      1.0, 1.0
-    ], samples, seqlen, seed, True)
+    return BlendedGPTDataset(
+      [base_dir / "validation" / "c4-validationn-91205-samples.en_text_document"], [1.0], samples, seqlen, seed, shuffle=False)
+  return BlendedGPTDataset(
+    [base_dir / "c4-train.en_6_text_document", base_dir / "c4-train.en_7_text_document"], [1.0, 1.0], samples, seqlen, seed, shuffle=True)
 
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
+def iterate_llama3_dataset(dataset:BlendedGPTDataset, bs:int):
+  for b in range(math.ceil(dataset.samples / bs)):
+    batch = [dataset.get(b * bs + i) for i in range(bs)]
     yield Tensor.stack(batch, dim=0)
 
-def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-validation-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, True)
-
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
-    yield Tensor.stack(batch, dim=0)
+def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True, small:bool=False):
+  return iterate_llama3_dataset(get_llama3_dataset(samples, seqlen, base_dir, seed, val, small), bs)
 
 if __name__ == "__main__":
   def load_unet3d(val):

examples/mlperf/model_eval.py

@@ -204,43 +204,6 @@ def eval_bert():
 
     st = time.perf_counter()
 
-def eval_mrcnn():
-  from tqdm import tqdm
-  from extra.models.mask_rcnn import MaskRCNN
-  from extra.models.resnet import ResNet
-  from extra.datasets.coco import BASEDIR, images, convert_prediction_to_coco_bbox, convert_prediction_to_coco_mask, accumulate_predictions_for_coco, evaluate_predictions_on_coco, iterate
-  from examples.mask_rcnn import compute_prediction_batched, Image
-  mdl = MaskRCNN(ResNet(50, num_classes=None, stride_in_1x1=True))
-  mdl.load_from_pretrained()
-
-  bbox_output = '/tmp/results_bbox.json'
-  mask_output = '/tmp/results_mask.json'
-
-  accumulate_predictions_for_coco([], bbox_output, rm=True)
-  accumulate_predictions_for_coco([], mask_output, rm=True)
-
-  #TODO: bs > 1 not as accurate
-  bs = 1
-
-  for batch in tqdm(iterate(images, bs=bs), total=len(images)//bs):
-    batch_imgs = []
-    for image_row in batch:
-      image_name = image_row['file_name']
-      img = Image.open(BASEDIR/f'val2017/{image_name}').convert("RGB")
-      batch_imgs.append(img)
-    batch_result = compute_prediction_batched(batch_imgs, mdl)
-    for image_row, result in zip(batch, batch_result):
-      image_name = image_row['file_name']
-      box_pred = convert_prediction_to_coco_bbox(image_name, result)
-      mask_pred = convert_prediction_to_coco_mask(image_name, result)
-      accumulate_predictions_for_coco(box_pred, bbox_output)
-      accumulate_predictions_for_coco(mask_pred, mask_output)
-    del batch_imgs
-    del batch_result
-
-  evaluate_predictions_on_coco(bbox_output, iou_type='bbox')
-  evaluate_predictions_on_coco(mask_output, iou_type='segm')
-
 def eval_llama3():
   from extra.models.llama import Transformer
   from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
@@ -271,12 +234,9 @@ def eval_llama3():
     loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
     return loss.flatten().float()
 
-  if SMALL:
-    from examples.mlperf.dataloader import batch_load_llama3_small
-    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
+  from examples.mlperf.dataloader import get_llama3_dataset, iterate_llama3_dataset
+  eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
+  iter = iterate_llama3_dataset(eval_dataset, BS)
 
   losses = []
   for tokens in tqdm(iter, total=5760//BS):
@@ -541,7 +501,7 @@ if __name__ == "__main__":
   # inference only
   Tensor.training = False
 
-  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,mrcnn").split(",")
+  models = getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert").split(",")
   for m in models:
     nm = f"eval_{m}"
     if nm in globals():

examples/mlperf/model_train.py

@@ -1314,12 +1314,21 @@ def train_llama3():
   opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
   opt_end_learning_rate = getenv("END_LR", 8e-7)
 
-  # TODO: confirm weights are in bf16
+  # ** init wandb **
+  WANDB = getenv("WANDB")
+  if WANDB:
+    import wandb
+    wandb_args = {"id": wandb_id, "resume": "must"} if (wandb_id := getenv("WANDB_RESUME", "")) else {}
+    wandb.init(config=config, **wandb_args, project="MLPerf-LLaMA3")
+
+  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
   # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  if not SMALL: model_params |= {"vocab_size": 32000}
+  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
+  params = get_parameters(model)
+  # weights are all bfloat16 for now
+  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
 
   if getenv("FAKEDATA"):
     for v in get_parameters(model):
@@ -1409,55 +1418,62 @@ def train_llama3():
   # ** data iters **
   def fake_data(bs, samples):
     for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=model_params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
 
   def get_train_iter():
     if getenv("FAKEDATA", 0):
       return fake_data(BS, SAMPLES)
     else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(BS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL), small=bool(SMALL))
+
+  if getenv("FAKEDATA", 0):
+    eval_dataset = None
+  else:
+    from examples.mlperf.dataloader import get_llama3_dataset
+    eval_dataset = get_llama3_dataset(5760, SEQLEN, BASEDIR, val=True, small=bool(SMALL))
 
   def get_eval_iter():
-    if getenv("FAKEDATA", 0):
+    if eval_dataset is None:
       return fake_data(EVAL_BS, 5760)
-    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+    from examples.mlperf.dataloader import iterate_llama3_dataset
+    return iterate_llama3_dataset(eval_dataset, EVAL_BS)
 
   iter = get_train_iter()
   i, sequences_seen = resume_ckpt, 0
   for tokens in tqdm(iter, total=SAMPLES//GBS):
-    t = time.perf_counter()
     GlobalCounters.reset()
-    loss, lr = train_step(model, tokens)
-    loss = loss.float().item()
+    if getenv("TRAIN", 1):
+      t = time.perf_counter()
+      loss, lr = train_step(model, tokens)
+      loss = loss.float().item()
+      lr = lr.item()
 
-    i += 1
-    sequences_seen += tokens.shape[0]
+      i += 1
+      sequences_seen += tokens.shape[0]
 
-    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
-    if (fname:=getenv("LOSS_FILE", "")):
-      with open(fname, "a") as f:
-        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+      sec = time.perf_counter()-t
+      mem_gb = GlobalCounters.mem_used / 1e9
+      gflops = GlobalCounters.global_ops / 1e9 / sec
+      tqdm.write(
+        f"{i:5} {sec:.2f} s run, {loss:.4f} loss, {lr:.12f} LR, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS")
 
-    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
-      tqdm.write("saving checkpoint")
-      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
-      fn = f"{ckpt_dir}/llama3_{i}.safe"
-      safe_save(get_state_dict(model), fn)
+      if (fname:=getenv("LOSS_FILE", "")):
+        with open(fname, "a") as f:
+          f.write(f"{i} {loss:.4f} {lr:.12f} {mem_gb:.2f}\n")
 
-      tqdm.write("saving optim checkpoint")
-      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-      safe_save(get_state_dict(scheduler), fn)
+      if WANDB:
+        wandb.log({"lr": lr, "train/loss": loss, "train/step_time": sec, "train/GFLOPS": gflops, "train/sequences_seen": sequences_seen})
+
+      if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+        tqdm.write("saving checkpoint")
+        if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
+        fn = f"{ckpt_dir}/llama3_{i}.safe"
+        safe_save(get_state_dict(model), fn)
+
+        tqdm.write("saving optim checkpoint")
+        fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
+        safe_save(get_state_dict(scheduler), fn)
 
     if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
       tqdm.write(f"evaluating after {sequences_seen} sequences")
@@ -1473,6 +1489,9 @@ def train_llama3():
 
       tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
 
+      if WANDB:
+        wandb.log({"eval/log_perplexity": log_perplexity, "eval/sequences_seen": sequences_seen})
+
       if log_perplexity < EVAL_TARGET:
         tqdm.write(f"target achieved after {sequences_seen} sequences")
         if getenv("CKPT"):

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

@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI350X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=5000000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI350x_${DATETIME}_${SEED}.log"
+
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

examples/openelm.py

@@ -1,118 +0,0 @@
-import json, pprint
-from tinygrad import fetch, nn, Tensor
-from tinygrad.helpers import DEBUG
-
-class FeedForward:
-  def __init__(self, model_dim, intermediate_dim):
-    self.proj_1 = nn.Linear(model_dim, 2*intermediate_dim, bias=False)
-    self.proj_2 = nn.Linear(intermediate_dim, model_dim, bias=False)
-
-  def __call__(self, x):
-    y_12 = self.proj_1(x)
-    y_1, y_2 = y_12.chunk(2, dim=-1)
-    return self.proj_2(y_1.silu() * y_2)
-
-# NOTE: this RoPE doesn't match LLaMA's?
-def _rotate_half(x: Tensor) -> Tensor:
-  x1, x2 = x.chunk(2, dim=-1)
-  return Tensor.cat(-x2, x1, dim=-1)
-
-def _apply_rotary_pos_emb(x: Tensor, pos_sin: Tensor, pos_cos: Tensor) -> Tensor:
-  return (x * pos_cos) + (_rotate_half(x) * pos_sin)
-
-class Attention:
-  def __init__(self, model_dim, num_query_heads, num_kv_heads, head_dim):
-    self.qkv_proj = nn.Linear(model_dim, (num_query_heads + num_kv_heads*2) * head_dim, bias=False)
-    self.num_query_heads, self.num_kv_heads = num_query_heads, num_kv_heads
-    self.head_dim = head_dim
-    self.q_norm = nn.RMSNorm(head_dim)
-    self.k_norm = nn.RMSNorm(head_dim)
-    self.out_proj = nn.Linear(num_query_heads * head_dim, model_dim, bias=False)
-
-  def __call__(self, x:Tensor) -> Tensor:
-    batch_size, seq_len, embed_dim = x.shape
-    qkv = self.qkv_proj(x)
-    qkv = qkv.reshape(batch_size, seq_len, self.num_query_heads+self.num_kv_heads*2, self.head_dim).transpose(1, 2)
-    xq,xk,xv = qkv.split([self.num_query_heads, self.num_kv_heads, self.num_kv_heads], dim=1)
-    xq = self.q_norm(xq)
-    xk = self.k_norm(xk)
-
-    # add positional embedding (how many kernels is this?)
-    freq_constant = 10000
-    inv_freq = 1.0 / (freq_constant ** (Tensor.arange(0, self.head_dim, 2) / self.head_dim))
-    pos_index_theta = Tensor.einsum("i,j->ij", Tensor.arange(seq_len), inv_freq)
-    emb = Tensor.cat(pos_index_theta, pos_index_theta, dim=-1)
-    cos_emb, sin_emb = emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]
-    xq = _apply_rotary_pos_emb(xq, sin_emb, cos_emb)
-    xk = _apply_rotary_pos_emb(xk, sin_emb, cos_emb)
-
-    # grouped-query attention
-    num_groups = self.num_query_heads // self.num_kv_heads
-    xk = xk.repeat_interleave(num_groups, dim=1)
-    xv = xv.repeat_interleave(num_groups, dim=1)
-
-    # masked attention
-    #start_pos = 0
-    #mask = Tensor.full((1, 1, seq_len, start_pos+seq_len), float("-inf"), dtype=xq.dtype, device=xq.device).triu(start_pos+1)
-    #attn_output = xq.scaled_dot_product_attention(xk, xv, mask).transpose(1, 2)
-
-    # causal is fine, no mask needed
-    attn_output = xq.scaled_dot_product_attention(xk, xv, is_causal=True).transpose(1, 2)
-    return self.out_proj(attn_output.reshape(batch_size, seq_len, self.num_query_heads * self.head_dim))
-
-class Layer:
-  def __init__(self, model_dim, intermediate_dim, num_query_heads, num_kv_heads, head_dim):
-    self.ffn = FeedForward(model_dim, intermediate_dim)
-    self.attn = Attention(model_dim, num_query_heads, num_kv_heads, head_dim)
-    self.ffn_norm = nn.RMSNorm(model_dim)
-    self.attn_norm = nn.RMSNorm(model_dim)
-
-  def __call__(self, x:Tensor) -> Tensor: # (batch, seq_len, embed_dim)
-    x = x + self.attn(self.attn_norm(x))
-    x = x + self.ffn(self.ffn_norm(x))
-    return x
-
-# stupidly complex
-def make_divisible(v, divisor):
-  new_v = max(divisor, int(v + divisor / 2) // divisor * divisor)
-  if new_v < 0.9 * v: new_v += divisor
-  return new_v
-
-class Transformer:
-  def __init__(self, cfg):
-    if DEBUG >= 3: pprint.pp(cfg)
-    self.layers = [Layer(cfg['model_dim'], make_divisible(int(cfg["model_dim"] * cfg['ffn_multipliers'][i]), cfg['ffn_dim_divisor']),
-                         cfg['num_query_heads'][i], cfg['num_kv_heads'][i], cfg['head_dim']) for i in range(cfg['num_transformer_layers'])]
-    self.norm = nn.RMSNorm(cfg['model_dim'])
-    self.token_embeddings = nn.Embedding(cfg['vocab_size'], cfg['model_dim'])
-
-  def __call__(self, tokens:Tensor):
-    # _bsz, seqlen = tokens.shape
-    x = self.token_embeddings(tokens)
-    for l in self.layers: x = l(x)
-    return self.norm(x) @ self.token_embeddings.weight.T
-
-if __name__ == "__main__":
-  #model_name = "OpenELM-270M-Instruct"
-  model_name = "OpenELM-270M"  # this is fp32
-  model = Transformer(json.loads(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/config.json?download=true").read_bytes()))
-  weights = nn.state.safe_load(fetch(f"https://huggingface.co/apple/{model_name}/resolve/main/model.safetensors?download=true"))
-  if DEBUG >= 3:
-    for k, v in weights.items(): print(k, v.shape)
-  nn.state.load_state_dict(model, {k.removeprefix("transformer."):v for k,v in weights.items()})
-
-  from sentencepiece import SentencePieceProcessor
-  tokenizer = SentencePieceProcessor(fetch("https://github.com/karpathy/llama2.c/raw/master/tokenizer.model").as_posix())
-  toks = [tokenizer.bos_id()] + tokenizer.encode("Some car brands include")
-  for i in range(100):
-    ttoks = Tensor([toks])
-    out = model(ttoks).realize()
-    t0 = out[0].argmax(axis=-1).tolist()
-    toks.append(t0[-1])
-    # hmmm...passthrough still doesn't match (it shouldn't, it outputs the most likely)
-    print(tokenizer.decode(toks))
-    #print(toks)
-    #print(tokenizer.decode(t0))
-    #print(t0)
-
-

examples/other_mnist/beautiful_mnist_mlx.py

@@ -1,55 +0,0 @@
-from tinygrad.helpers import trange
-from tinygrad.nn.datasets import mnist
-import mlx.core as mx
-import mlx.nn as nn
-import mlx.optimizers as optim
-from functools import partial
-
-class Model(nn.Module):
-  def __init__(self):
-    super().__init__()
-    self.c1 = nn.Conv2d(1, 32, 5)
-    self.c2 = nn.Conv2d(32, 32, 5)
-    self.bn1 = nn.BatchNorm(32)
-    self.m1 = nn.MaxPool2d(2)
-    self.c3 = nn.Conv2d(32, 64, 3)
-    self.c4 = nn.Conv2d(64, 64, 3)
-    self.bn2 = nn.BatchNorm(64)
-    self.m2 = nn.MaxPool2d(2)
-    self.lin = nn.Linear(576, 10)
-  def __call__(self, x):
-    x = mx.maximum(self.c1(x), 0)
-    x = mx.maximum(self.c2(x), 0)
-    x = self.m1(self.bn1(x))
-    x = mx.maximum(self.c3(x), 0)
-    x = mx.maximum(self.c4(x), 0)
-    x = self.m2(self.bn2(x))
-    return self.lin(mx.flatten(x, 1))
-
-if __name__ == "__main__":
-  X_train, Y_train, X_test, Y_test = mnist()
-  X_train = mx.array(X_train.float().permute((0,2,3,1)).numpy())
-  Y_train = mx.array(Y_train.numpy())
-  X_test = mx.array(X_test.float().permute((0,2,3,1)).numpy())
-  Y_test = mx.array(Y_test.numpy())
-
-  model = Model()
-  optimizer = optim.Adam(1e-3)
-  def loss_fn(model, x, y): return nn.losses.cross_entropy(model(x), y).mean()
-
-  state = [model.state, optimizer.state]
-  @partial(mx.compile, inputs=state, outputs=state)
-  def step(samples):
-    # Compiled functions will also treat any inputs not in the parameter list as constants.
-    X,Y = X_train[samples], Y_train[samples]
-    loss_and_grad_fn = nn.value_and_grad(model, loss_fn)
-    loss, grads = loss_and_grad_fn(model, X, Y)
-    optimizer.update(model, grads)
-    return loss
-
-  test_acc = float('nan')
-  for i in (t:=trange(70)):
-    samples = mx.random.randint(0, X_train.shape[0], (512,))  # putting this in JIT didn't work well
-    loss = step(samples)
-    if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
-    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")

examples/rl/lightupbutton.py

@@ -1,45 +0,0 @@
-import gymnasium as gym
-import numpy as np
-from gymnasium.envs.registration import register
-
-# a very simple game
-# one of <size> lights will light up
-# take the action of the lit up light
-# in <hard_mode>, you act differently based on the step number and need to track this
-
-class PressTheLightUpButton(gym.Env):
-  metadata = {"render_modes": []}
-  def __init__(self, render_mode=None, size=2, game_length=10, hard_mode=False):
-    self.size, self.game_length = size, game_length
-    self.observation_space = gym.spaces.Box(0, 1, shape=(self.size,), dtype=np.float32)
-    self.action_space = gym.spaces.Discrete(self.size)
-    self.step_num = 0
-    self.done = True
-    self.hard_mode = hard_mode
-
-  def _get_obs(self):
-    obs = [0]*self.size
-    if self.step_num < len(self.state):
-      obs[self.state[self.step_num]] = 1
-    return np.array(obs, dtype=np.float32)
-
-  def reset(self, seed=None, options=None):
-    super().reset(seed=seed)
-    self.state = np.random.randint(0, self.size, size=self.game_length)
-    self.step_num = 0
-    self.done = False
-    return self._get_obs(), {}
-
-  def step(self, action):
-    target = ((action + self.step_num) % self.size) if self.hard_mode else action
-    reward = int(target == self.state[self.step_num])
-    self.step_num += 1
-    if not reward:
-      self.done = True
-    return self._get_obs(), reward, self.done, self.step_num >= self.game_length, {}
-
-register(
-  id="PressTheLightUpButton-v0",
-  entry_point="examples.rl.lightupbutton:PressTheLightUpButton",
-  max_episode_steps=None,
-)

examples/serious_mnist.py

@@ -1,136 +0,0 @@
-#!/usr/bin/env python
-#inspired by https://github.com/Matuzas77/MNIST-0.17/blob/master/MNIST_final_solution.ipynb
-import sys
-import numpy as np
-from tinygrad.nn.state import get_parameters
-from tinygrad.tensor import Tensor
-from tinygrad.nn import BatchNorm2d, optim
-from tinygrad.helpers import getenv
-from extra.datasets import fetch_mnist
-from extra.augment import augment_img
-from extra.training import train, evaluate
-GPU = getenv("GPU")
-QUICK = getenv("QUICK")
-DEBUG = getenv("DEBUG")
-
-class SqueezeExciteBlock2D:
-  def __init__(self, filters):
-    self.filters = filters
-    self.weight1 = Tensor.scaled_uniform(self.filters, self.filters//32)
-    self.bias1 = Tensor.scaled_uniform(1,self.filters//32)
-    self.weight2 = Tensor.scaled_uniform(self.filters//32, self.filters)
-    self.bias2 = Tensor.scaled_uniform(1, self.filters)
-
-  def __call__(self, input):
-    se = input.avg_pool2d(kernel_size=(input.shape[2], input.shape[3])) #GlobalAveragePool2D
-    se = se.reshape(shape=(-1, self.filters))
-    se = se.dot(self.weight1) + self.bias1
-    se = se.relu()
-    se = se.dot(self.weight2) + self.bias2
-    se = se.sigmoid().reshape(shape=(-1,self.filters,1,1)) #for broadcasting
-    se = input.mul(se)
-    return se
-
-class ConvBlock:
-  def __init__(self, h, w, inp, filters=128, conv=3):
-    self.h, self.w = h, w
-    self.inp = inp
-    #init weights
-    self.cweights = [Tensor.scaled_uniform(filters, inp if i==0 else filters, conv, conv) for i in range(3)]
-    self.cbiases = [Tensor.scaled_uniform(1, filters, 1, 1) for i in range(3)]
-    #init layers
-    self._bn = BatchNorm2d(128)
-    self._seb = SqueezeExciteBlock2D(filters)
-
-  def __call__(self, input):
-    x = input.reshape(shape=(-1, self.inp, self.w, self.h))
-    for cweight, cbias in zip(self.cweights, self.cbiases):
-      x = x.pad(padding=[1,1,1,1]).conv2d(cweight).add(cbias).relu()
-    x = self._bn(x)
-    x = self._seb(x)
-    return x
-
-class BigConvNet:
-  def __init__(self):
-    self.conv = [ConvBlock(28,28,1), ConvBlock(28,28,128), ConvBlock(14,14,128)]
-    self.weight1 = Tensor.scaled_uniform(128,10)
-    self.weight2 = Tensor.scaled_uniform(128,10)
-
-  def parameters(self):
-    if DEBUG: #keeping this for a moment
-      pars = [par for par in get_parameters(self) if par.requires_grad]
-      no_pars = 0
-      for par in pars:
-        print(par.shape)
-        no_pars += np.prod(par.shape)
-      print('no of parameters', no_pars)
-      return pars
-    else:
-      return get_parameters(self)
-
-  def save(self, filename):
-    with open(filename+'.npy', 'wb') as f:
-      for par in get_parameters(self):
-        #if par.requires_grad:
-        np.save(f, par.numpy())
-
-  def load(self, filename):
-    with open(filename+'.npy', 'rb') as f:
-      for par in get_parameters(self):
-        #if par.requires_grad:
-        try:
-          par.numpy()[:] = np.load(f)
-          if GPU:
-            par.gpu()
-        except:
-          print('Could not load parameter')
-
-  def forward(self, x):
-    x = self.conv[0](x)
-    x = self.conv[1](x)
-    x = x.avg_pool2d(kernel_size=(2,2))
-    x = self.conv[2](x)
-    x1 = x.avg_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
-    x2 = x.max_pool2d(kernel_size=(14,14)).reshape(shape=(-1,128)) #global
-    xo = x1.dot(self.weight1) + x2.dot(self.weight2)
-    return xo
-
-
-if __name__ == "__main__":
-  lrs = [1e-4, 1e-5] if QUICK else [1e-3, 1e-4, 1e-5, 1e-5]
-  epochss = [2, 1] if QUICK else [13, 3, 3, 1]
-  BS = 32
-
-  lmbd = 0.00025
-  lossfn = lambda out,y: out.sparse_categorical_crossentropy(y) + lmbd*(model.weight1.abs() + model.weight2.abs()).sum()
-  X_train, Y_train, X_test, Y_test = fetch_mnist()
-  X_train = X_train.reshape(-1, 28, 28).astype(np.uint8)
-  X_test = X_test.reshape(-1, 28, 28).astype(np.uint8)
-  steps = len(X_train)//BS
-  np.random.seed(1337)
-  if QUICK:
-    steps = 1
-    X_test, Y_test = X_test[:BS], Y_test[:BS]
-
-  model = BigConvNet()
-
-  if len(sys.argv) > 1:
-    try:
-      model.load(sys.argv[1])
-      print('Loaded weights "'+sys.argv[1]+'", evaluating...')
-      evaluate(model, X_test, Y_test, BS=BS)
-    except:
-      print('could not load weights "'+sys.argv[1]+'".')
-
-  if GPU:
-    params = get_parameters(model)
-    [x.gpu_() for x in params]
-
-  for lr, epochs in zip(lrs, epochss):
-    optimizer = optim.Adam(model.parameters(), lr=lr)
-    for epoch in range(1,epochs+1):
-      #first epoch without augmentation
-      X_aug = X_train if epoch == 1 else augment_img(X_train)
-      train(model, X_aug, Y_train, optimizer, steps=steps, lossfn=lossfn, BS=BS)
-      accuracy = evaluate(model, X_test, Y_test, BS=BS)
-      model.save(f'examples/checkpoint{accuracy * 1e6:.0f}')

examples/simple_conv_bn.py

@@ -1,17 +0,0 @@
-from tinygrad.tensor import Tensor
-from tinygrad.nn import Conv2d, BatchNorm2d
-from tinygrad.nn.state import get_parameters
-
-if __name__ == "__main__":
-  with Tensor.train():
-
-    BS, C1, H, W = 4, 16, 224, 224
-    C2, K, S, P = 64, 7, 2, 1
-
-    x = Tensor.uniform(BS, C1, H, W)
-    conv = Conv2d(C1, C2, kernel_size=K, stride=S, padding=P)
-    bn = BatchNorm2d(C2, track_running_stats=False)
-    for t in get_parameters([x, conv, bn]): t.realize()
-
-    print("running network")
-    x.sequential([conv, bn]).numpy()

examples/so_vits_svc.py

@@ -1,669 +0,0 @@
-# original implementation: https://github.com/svc-develop-team/so-vits-svc
-from __future__ import annotations
-import sys, logging, time, io, math, argparse, operator, numpy as np
-from functools import partial, reduce
-from pathlib import Path
-from typing import Tuple, Optional, Type
-from tinygrad import nn, dtypes, Tensor
-from tinygrad.helpers import getenv, fetch
-from tinygrad.nn.state import torch_load
-from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, get_hparams_from_file, load_checkpoint, weight_norm, HParams
-from examples.sovits_helpers import preprocess
-import soundfile
-
-DEBUG = getenv("DEBUG")
-
-F0_BIN = 256
-F0_MAX = 1100.0
-F0_MIN = 50.0
-F0_MEL_MIN = 1127 * np.log(1 + F0_MIN / 700)
-F0_MEL_MAX = 1127 * np.log(1 + F0_MAX / 700)
-
-class SpeechEncoder:
-  def __init__(self, hidden_dim, model:ContentVec): self.hidden_dim, self.model = hidden_dim, model
-  def encode(self, ): raise NotImplementedError("implement me")
-  @classmethod
-  def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
-    contentvec = ContentVec.load_from_pretrained(checkpoint_path, checkpoint_url)
-    return cls(contentvec)
-
-class ContentVec256L9(SpeechEncoder):
-  def __init__(self, model:ContentVec): super().__init__(hidden_dim=256, model=model)
-  def encode(self, wav: Tensor):
-    feats = wav
-    if len(feats.shape) == 2:  # double channels
-      feats = feats.mean(-1)
-    assert len(feats.shape) == 1, feats.dim()
-    feats = feats.reshape(1, -1)
-    padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
-    logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=9)
-    feats = self.model.final_proj(logits[0])
-    return feats.transpose(1,2)
-
-class ContentVec768L12(SpeechEncoder):
-  def __init__(self, model:ContentVec): super().__init__(hidden_dim=768, model=model)
-  def encode(self, wav: Tensor):
-    feats = wav
-    if len(feats.shape) == 2:  # double channels
-      feats = feats.mean(-1)
-    assert len(feats.shape) == 1, feats.dim()
-    feats = feats.reshape(1, -1)
-    padding_mask = Tensor.zeros_like(feats).cast(dtypes.bool)
-    logits = self.model.extract_features(feats.to(wav.device), padding_mask=padding_mask.to(wav.device), output_layer=12)
-    return logits[0].transpose(1,2)
-
-# original code for contentvec: https://github.com/auspicious3000/contentvec/
-class ContentVec:
-  # self.final_proj dims are hardcoded and depend on fairseq.data.dictionary Dictionary in the checkpoint. This param can't yet be loaded since there is no pickle for it. See with DEBUG=2.
-  # This means that the ContentVec only works with the hubert weights used in all SVC models
-  def __init__(self, cfg: HParams):
-    self.feature_grad_mult, self.untie_final_proj = cfg.feature_grad_mult, cfg.untie_final_proj
-    feature_enc_layers = eval(cfg.conv_feature_layers)
-    self.embed = feature_enc_layers[-1][0]
-    final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim
-    self.feature_extractor = ConvFeatureExtractionModel(conv_layers=feature_enc_layers, dropout=0.0, mode=cfg.extractor_mode, conv_bias=cfg.conv_bias)
-    self.post_extract_proj = nn.Linear(self.embed, cfg.encoder_embed_dim) if self.embed != cfg.encoder_embed_dim else None
-    self.encoder = TransformerEncoder(cfg)
-    self.layer_norm = nn.LayerNorm(self.embed)
-    self.final_proj = nn.Linear(cfg.encoder_embed_dim, final_dim * 1) if self.untie_final_proj else nn.Linear(cfg.encoder_embed_dim, final_dim)
-    self.mask_emb = Tensor.uniform(cfg.encoder_embed_dim, dtype=dtypes.float32)
-    self.label_embs_concat = Tensor.uniform(504, final_dim, dtype=dtypes.float32)
-  def forward_features(self, source, padding_mask):
-    if self.feature_grad_mult > 0:
-      features = self.feature_extractor(source, padding_mask)
-      if self.feature_grad_mult != 1.0: pass  # training: GradMultiply.forward(features, self.feature_grad_mult)
-    else:
-      features = self.feature_extractor(source, padding_mask)
-    return features
-  def forward_padding_mask(self, features, padding_mask):  # replaces original forward_padding_mask for batch inference
-    lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1)  # ensure its bool for tilde
-    lengths = (lengths_org - 400).float().div(320).floor().cast(dtypes.int64) + 1  # intermediate float to divide
-    padding_mask = lengths_to_padding_mask(lengths)
-    return padding_mask
-  def extract_features(self, source: Tensor, spk_emb:Tensor=None, padding_mask=None, ret_conv=False, output_layer=None, tap=False):
-    features = self.forward_features(source, padding_mask)
-    if padding_mask is not None:
-      padding_mask = self.forward_padding_mask(features, padding_mask)
-    features = features.transpose(1, 2)
-    features = self.layer_norm(features)
-    if self.post_extract_proj is not None:
-      features = self.post_extract_proj(features)
-    x, _ = self.encoder(features, spk_emb, padding_mask=padding_mask, layer=(None if output_layer is None else output_layer - 1), tap=tap)
-    res = features if ret_conv else x
-    return res, padding_mask
-  @classmethod
-  def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
-    fetch(checkpoint_url, checkpoint_path)
-    cfg = load_fairseq_cfg(checkpoint_path)
-    enc = cls(cfg.model)
-    _ = load_checkpoint_enc(checkpoint_path, enc, None)
-    logging.debug(f"{cls.__name__}: Loaded model with cfg={cfg}")
-    return enc
-
-class TransformerEncoder:
-  def __init__(self, cfg: HParams):
-    def make_conv() -> nn.Conv1d:
-      layer = nn.Conv1d(self.embedding_dim, self.embedding_dim, kernel_size=cfg.conv_pos, padding=cfg.conv_pos // 2, groups=cfg.conv_pos_groups)
-      std = std = math.sqrt(4 / (cfg.conv_pos * self.embedding_dim))
-      layer.weight, layer.bias = (Tensor.normal(*layer.weight.shape, std=std)), (Tensor.zeros(*layer.bias.shape))
-      # for training: layer.weights need to be weight_normed
-      return layer
-    self.dropout, self.embedding_dim, self.layer_norm_first, self.layerdrop, self.num_layers, self.num_layers_1 = cfg.dropout, cfg.encoder_embed_dim, cfg.layer_norm_first, cfg.encoder_layerdrop, cfg.encoder_layers, cfg.encoder_layers_1
-    self.pos_conv, self.pos_conv_remove = [make_conv()], (1 if cfg.conv_pos % 2 == 0 else 0)
-    self.layers = [
-      TransformerEncoderLayer(self.embedding_dim, cfg.encoder_ffn_embed_dim, cfg.encoder_attention_heads, self.dropout, cfg.attention_dropout, cfg.activation_dropout, cfg.activation_fn, self.layer_norm_first, cond_layer_norm=(i >= cfg.encoder_layers))
-      for i in range(cfg.encoder_layers + cfg.encoder_layers_1)
-      ]
-    self.layer_norm = nn.LayerNorm(self.embedding_dim)
-    self.cond_layer_norm = CondLayerNorm(self.embedding_dim) if cfg.encoder_layers_1 > 0 else None
-    # training: apply init_bert_params
-  def __call__(self, x, spk_emb, padding_mask=None, layer=None, tap=False):
-    x, layer_results = self.extract_features(x, spk_emb, padding_mask, layer, tap)
-    if self.layer_norm_first and layer is None:
-      x = self.cond_layer_norm(x, spk_emb) if (self.num_layers_1 > 0) else self.layer_norm(x)
-    return x, layer_results
-  def extract_features(self, x: Tensor, spk_emb: Tensor, padding_mask=None, tgt_layer=None, tap=False):
-    if tgt_layer is not None:  # and not self.training
-      assert tgt_layer >= 0 and tgt_layer < len(self.layers)
-    if padding_mask is not None:
-      # x[padding_mask] = 0
-      assert padding_mask.shape == x.shape[:len(padding_mask.shape)]  # first few dims of x must match padding_mask
-      tmp_mask = padding_mask.unsqueeze(-1).repeat((1, 1, x.shape[-1]))
-      tmp_mask = tilde(tmp_mask.cast(dtypes.bool))
-      x = tmp_mask.where(x, 0)
-    x_conv = self.pos_conv[0](x.transpose(1,2))
-    if self.pos_conv_remove > 0: x_conv = x_conv[:, :, : -self.pos_conv_remove]
-    x_conv = x_conv.gelu().transpose(1, 2)
-    x = (x + x_conv).transpose(0, 1)  # B x T x C -> T x B x C
-    if not self.layer_norm_first: x = self.layer_norm(x)
-    x = x.dropout(p=self.dropout)
-    layer_results = []
-    r = None
-    for i, layer in enumerate(self.layers):
-      if i < self.num_layers:  # if (not self.training or (dropout_probability > self.layerdrop)) and (i < self.num_layers):
-        assert layer.cond_layer_norm == False
-        x = layer(x, self_attn_padding_mask=padding_mask, need_weights=False)
-        if tgt_layer is not None or tap:
-          layer_results.append(x.transpose(0, 1))
-      if i>= self.num_layers:
-        assert layer.cond_layer_norm == True
-        x = layer(x, emb=spk_emb, self_attn_padding_mask=padding_mask, need_weights=False)
-      if i == tgt_layer:
-        r = x
-        break
-    if r is not None:
-      x = r
-    x = x.transpose(0, 1)  # T x B x C -> B x T x C
-    return x, layer_results
-
-class TransformerEncoderLayer:
-  def __init__(self, embedding_dim=768.0, ffn_embedding_dim=3072.0, num_attention_heads=8.0, dropout=0.1, attention_dropout=0.1, activation_dropout=0.1, activation_fn="relu", layer_norm_first=False, cond_layer_norm=False):
-    def get_activation_fn(activation):
-      if activation == "relu": return Tensor.relu
-      if activation == "gelu": return Tensor.gelu
-      else: raise RuntimeError(f"activation function={activation} is not forseen")
-    self.embedding_dim, self.dropout, self.activation_dropout, self.layer_norm_first, self.num_attention_heads, self.cond_layer_norm, self.activation_fn = embedding_dim, dropout, activation_dropout, layer_norm_first, num_attention_heads, cond_layer_norm, get_activation_fn(activation_fn)
-    self.self_attn = MultiHeadAttention(self.embedding_dim, self.num_attention_heads)
-    self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
-    self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
-    self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
-    self.final_layer_norm = nn.LayerNorm(self.embedding_dim) if not cond_layer_norm else CondLayerNorm(self.embedding_dim)
-  def __call__(self, x:Tensor, self_attn_mask:Tensor=None, self_attn_padding_mask:Tensor=None, emb:Tensor=None, need_weights=False):
-    #self_attn_padding_mask = self_attn_padding_mask.reshape(x.shape[0], 1, 1, self_attn_padding_mask.shape[1]).expand(-1, self.num_attention_heads, -1, -1).reshape(x.shape[0] * self.num_attention_heads, 1, self_attn_padding_mask.shape[1]) if self_attn_padding_mask is not None else None
-    assert self_attn_mask is None and self_attn_padding_mask is not None
-    residual = x
-    if self.layer_norm_first:
-      x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
-      x = self.self_attn(x=x, mask=self_attn_padding_mask)
-      x = x.dropout(self.dropout)
-      x = residual + x
-      x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
-      x = self.activation_fn(self.fc1(x))
-      x = x.dropout(self.activation_dropout)
-      x = self.fc2(x)
-      x = x.dropout(self.dropout)
-      x = residual + x
-    else:
-      x = self.self_attn(x=x, mask=self_attn_padding_mask)
-      x = x.dropout(self.dropout)
-      x = residual + x
-      x = self.self_attn_layer_norm(x) if not self.cond_layer_norm else self.self_attn_layer_norm(x, emb)
-      residual = x
-      x = self.activation_fn(self.fc1(x))
-      x = x.dropout(self.activation_dropout)
-      x = self.fc2(x)
-      x = x.dropout(self.dropout)
-      x = residual + x
-      x = self.final_layer_norm(x) if not self.cond_layer_norm else self.final_layer_norm(x, emb)
-    return x
-
-class MultiHeadAttention:
-  def __init__(self, n_state, n_head):
-    self.n_state, self.n_head = n_state, n_head
-    self.q_proj, self.k_proj, self.v_proj, self.out_proj = [nn.Linear(n_state, n_state) for _ in range(4)]
-  def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None):
-    x = x.transpose(0,1)  # TxBxC -> BxTxC
-    q, k, v = self.q_proj(x), self.k_proj(xa or x), self.v_proj(xa or x)
-    q, k, v = [x.reshape(*q.shape[:2], self.n_head, -1) for x in (q, k, v)]
-    wv = Tensor.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), None).transpose(1, 2).reshape(*x.shape[:2], -1)
-    ret =  self.out_proj(wv).transpose(0,1)  # BxTxC -> TxBxC
-    return ret
-
-class ConvFeatureExtractionModel:
-  def __init__(self, conv_layers, dropout=.0, mode="default", conv_bias=False):
-    assert mode in {"default", "group_norm_masked", "layer_norm"}
-    def block(n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False):
-      def make_conv():
-        conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias)
-        conv.weight = Tensor.kaiming_normal(*conv.weight.shape)
-        return conv
-      assert (is_layer_norm and is_group_norm) == False, "layer norm and group norm are exclusive"
-      if is_layer_norm:
-        return [make_conv(), partial(Tensor.dropout, p=dropout),[partial(Tensor.transpose, dim0=-2, dim1=-1), nn.LayerNorm(dim, elementwise_affine=True), partial(Tensor.transpose, dim0=-2, dim1=-1)], Tensor.gelu]
-      elif is_group_norm and mode == "default":
-        return [make_conv(), partial(Tensor.dropout, p=dropout), nn.GroupNorm(dim, dim, affine=True), Tensor.gelu]
-      elif is_group_norm and mode == "group_norm_masked":
-        return [make_conv(), partial(Tensor.dropout, p=dropout), GroupNormMasked(dim, dim, affine=True), Tensor.gelu]
-      else:
-        return [make_conv(), partial(Tensor.dropout, p=dropout), Tensor.gelu]
-    in_d, self.conv_layers, self.mode = 1, [], mode
-    for i, cl in enumerate(conv_layers):
-      assert len(cl) == 3, "invalid conv definition: " + str(cl)
-      (dim, k, stride) = cl
-      if i == 0: self.cl = cl
-      self.conv_layers.append(block(in_d, dim, k, stride, is_layer_norm=(mode == "layer_norm"), is_group_norm=((mode == "default" or mode == "group_norm_masked") and i == 0), conv_bias=conv_bias))
-      in_d = dim
-  def __call__(self, x:Tensor, padding_mask:Tensor):
-    x = x.unsqueeze(1)  # BxT -> BxCxT
-    if self.mode == "group_norm_masked":
-      if padding_mask is not None:
-        _, k, stride = self.cl
-        lengths_org = tilde(padding_mask.cast(dtypes.bool)).cast(dtypes.int64).sum(1)  # ensure padding_mask is bool for tilde
-        lengths = (((lengths_org - k) / stride) + 1).floor().cast(dtypes.int64)
-        padding_mask = tilde(lengths_to_padding_mask(lengths)).cast(dtypes.int64)  # lengths_to_padding_mask returns bool tensor
-      x = self.conv_layers[0][0](x)  # padding_mask is numeric
-      x = self.conv_layers[0][1](x)
-      x = self.conv_layers[0][2](x, padding_mask)
-      x = self.conv_layers[0][3](x)
-    else:
-      x = x.sequential(self.conv_layers[0])  # default
-    for _, conv in enumerate(self.conv_layers[1:], start=1):
-      conv = reduce(lambda a,b: operator.iconcat(a,b if isinstance(b, list) else [b]), conv, [])  # flatten
-      x = x.sequential(conv)
-    return x
-
-class CondLayerNorm:  # https://github.com/auspicious3000/contentvec/blob/main/contentvec/modules/cond_layer_norm.py#L10
-  def __init__(self, dim_last, eps=1e-5, dim_spk=256, elementwise_affine=True):
-    self.dim_last, self.eps, self.dim_spk, self.elementwise_affine = dim_last, eps, dim_spk, elementwise_affine
-    if self.elementwise_affine:
-      self.weight_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
-      self.bias_ln = nn.Linear(self.dim_spk, self.dim_last, bias=False)
-      self.weight_ln.weight, self.bias_ln.weight = (Tensor.ones(*self.weight_ln.weight.shape)), (Tensor.zeros(*self.bias_ln.weight.shape))
-  def __call__(self, x: Tensor, spk_emb: Tensor):
-    axis = tuple(-1-i for i in range(len(x.shape[1:])))
-    x = x.layernorm(axis=axis, eps=self.eps)
-    if not self.elementwise_affine: return x
-    weights, bias = self.weight_ln(spk_emb), self.bias_ln(spk_emb)
-    return weights * x + bias
-
-class GroupNormMasked:  # https://github.com/auspicious3000/contentvec/blob/d746688a32940f4bee410ed7c87ec9cf8ff04f74/contentvec/modules/fp32_group_norm.py#L16
-  def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
-    self.num_groups, self.num_channels, self.eps, self.affine = num_groups, num_channels, eps, affine
-    self.weight, self.bias = (Tensor.ones(num_channels)), (Tensor.zeros(num_channels)) if self.affine else (None, None)
-  def __call__(self, x:Tensor, mask:Tensor):
-    bsz, n_c, length = x.shape
-    assert n_c % self.num_groups == 0
-    x = x.reshape(bsz, self.num_groups, n_c // self.num_groups, length)
-    if mask is None: mask = Tensor.ones_like(x)
-    else: mask = mask.reshape(bsz, 1, 1, length)
-    x = x * mask
-    lengths = mask.sum(axis=3, keepdim=True)
-    assert x.shape[2] == 1
-    mean_ = x.mean(dim=3, keepdim=True)
-    mean = mean_ * length / lengths
-    var = (((x.std(axis=3, keepdim=True) ** 2) + mean_**2) * length / lengths - mean**2) + self.eps
-    return x.add(-mean).div(var.sqrt()).reshape(bsz, n_c, length).mul(self.weight.reshape(1,-1,1)).add(self.bias.reshape(1,-1,1))
-
-class Synthesizer:
-  def __init__(self, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels, ssl_dim, n_speakers, sampling_rate=44100, vol_embedding=False, n_flow_layer=4, **kwargs):
-    self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.vol_embedding = spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, vol_embedding
-    self.emb_g = nn.Embedding(n_speakers, gin_channels)
-    if vol_embedding: self.emb_vol = nn.Linear(1, hidden_channels)
-    self.pre = nn.Conv1d(ssl_dim, hidden_channels, kernel_size=5, padding=2)
-    self.enc_p = TextEncoder(inter_channels, hidden_channels, kernel_size, n_layers, filter_channels=filter_channels, n_heads=n_heads, p_dropout=p_dropout)
-    self.dec = Generator(sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels)
-    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
-    self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, n_flow_layer, gin_channels=gin_channels)
-    self.emb_uv = nn.Embedding(vocab_size=2, embed_size=hidden_channels)
-  def infer(self, c:Tensor, f0:Tensor, uv:Tensor, g:Tensor=None, noise_scale=0.35, seed=52468, vol=None) -> Tuple[Tensor, Tensor]:
-    Tensor.manual_seed(getenv('SEED', seed))
-    c_lengths = (Tensor.ones([c.shape[0]]) * c.shape[-1]).to(c.device)
-    if len(g.shape) == 1: g = g.unsqueeze(0)
-    g = self.emb_g(g).transpose(1, 2)
-    x_mask = sequence_mask(c_lengths, c.shape[2]).unsqueeze(1).cast(c.dtype)
-    vol = self.emb_vol(vol[:,:,None]).transpose(1,2) if vol is not None and self.vol_embedding else 0
-    x = self.pre(c) * x_mask + self.emb_uv(uv.cast(dtypes.int64)).transpose(1, 2) + vol
-    z_p, _, _, c_mask = self.enc_p.forward(x, x_mask, f0=self._f0_to_coarse(f0), noise_scale=noise_scale)
-    z = self.flow.forward(z_p, c_mask, g=g, reverse=True)
-    o = self.dec.forward(z * c_mask, g=g, f0=f0)
-    return o,f0
-  def _f0_to_coarse(self, f0 : Tensor):
-    f0_mel = 1127 * (1 + f0 / 700).log()
-    a = (F0_BIN - 2) / (F0_MEL_MAX - F0_MEL_MIN)
-    b = F0_MEL_MIN * a - 1.
-    f0_mel = (f0_mel > 0).where(f0_mel * a - b, f0_mel)
-    f0_coarse = f0_mel.ceil().cast(dtype=dtypes.int64)
-    f0_coarse = f0_coarse * (f0_coarse > 0)
-    f0_coarse = f0_coarse + ((f0_coarse < 1) * 1)
-    f0_coarse = f0_coarse * (f0_coarse < F0_BIN)
-    f0_coarse = f0_coarse + ((f0_coarse >= F0_BIN) * (F0_BIN - 1))
-    return f0_coarse
-  @classmethod
-  def load_from_pretrained(cls, config_path:str, config_url:str, weights_path:str, weights_url:str) -> Synthesizer:
-    fetch(config_url, config_path)
-    hps = get_hparams_from_file(config_path)
-    fetch(weights_url, weights_path)
-    net_g = cls(hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, **hps.model)
-    _ = load_checkpoint(weights_path, net_g, None, skip_list=["f0_decoder"])
-    logging.debug(f"{cls.__name__}:Loaded model with hps: {hps}")
-    return net_g, hps
-
-class TextEncoder:
-  def __init__(self, out_channels, hidden_channels, kernel_size, n_layers, gin_channels=0, filter_channels=None, n_heads=None, p_dropout=None):
-    self.out_channels, self.hidden_channels, self.kernel_size, self.n_layers, self.gin_channels = out_channels, hidden_channels, kernel_size, n_layers, gin_channels
-    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-    self.f0_emb = nn.Embedding(256, hidden_channels)  # n_vocab = 256
-    self.enc_ = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
-  def forward(self, x, x_mask, f0=None, noise_scale=1):
-    x = x + self.f0_emb(f0).transpose(1, 2)
-    x = self.enc_.forward(x * x_mask, x_mask)
-    stats = self.proj(x) * x_mask
-    m, logs = split(stats, self.out_channels, dim=1)
-    z = (m + randn_like(m) * logs.exp() * noise_scale) * x_mask
-    return z, m, logs, x_mask
-
-class Upsample:
-  def __init__(self, scale_factor):
-    assert scale_factor % 1 == 0, "Only integer scale factor allowed."
-    self.scale = int(scale_factor)
-  def forward(self, x:Tensor):
-    repeats = tuple([1] * len(x.shape) + [self.scale])
-    new_shape = (*x.shape[:-1], x.shape[-1] * self.scale)
-    return x.unsqueeze(-1).repeat(repeats).reshape(new_shape)
-
-class SineGen:
-  def __init__(self, samp_rate, harmonic_num=0, sine_amp=0.1, noise_std=0.003, voice_threshold=0, flag_for_pulse=False):
-    self.sine_amp, self.noise_std, self.harmonic_num, self.sampling_rate, self.voiced_threshold, self.flag_for_pulse = sine_amp, noise_std, harmonic_num, samp_rate, voice_threshold, flag_for_pulse
-    self.dim = self.harmonic_num + 1
-  def _f02uv(self, f0): return (f0 > self.voiced_threshold).float()  #generate uv signal
-  def _f02sine(self, f0_values):
-    def padDiff(x : Tensor): return (x.pad((0,0,-1,1)) - x).pad((0,0,0,-1))
-    def mod(x: Tensor, n: int) -> Tensor: return x - n * x.div(n).floor()  # this is what the % operator does in pytorch.
-    rad_values = mod((f0_values / self.sampling_rate) , 1)  # convert to F0 in rad
-    rand_ini = Tensor.rand(f0_values.shape[0], f0_values.shape[2], device=f0_values.device)  # initial phase noise
-
-    #rand_ini[:, 0] = 0
-    m = Tensor.ones(f0_values.shape[0]).unsqueeze(1).pad((0,f0_values.shape[2]-1,0,0)).cast(dtypes.bool)
-    m = tilde(m)
-    rand_ini = m.where(rand_ini, 0)
-
-    #rad_values[:, 0, :] = rad_values[:, 0, :] + rand_ini
-    tmp = rad_values[:, 0, :] + rand_ini
-    m = Tensor.ones(tmp.shape).pad((0,0,0,rad_values.shape[1]-1,0)).cast(dtypes.bool)
-    m = tilde(m)
-    tmp = tmp.unsqueeze(1).pad((0,0,0,rad_values.shape[1]-1,0))
-    rad_values = m.where(rad_values, tmp)
-
-    tmp_over_one = mod(rad_values.cumsum(1), 1)
-    tmp_over_one_idx = padDiff(tmp_over_one) < 0
-    cumsum_shift = Tensor.zeros_like(rad_values)
-
-    #cumsum_shift[:, 1:, :] = tmp_over_one_idx * -1.0
-    tmp_over_one_idx = (tmp_over_one_idx * -1.0).pad((0,0,1,0))
-    cumsum_shift = tmp_over_one_idx
-
-    sines = ((rad_values + cumsum_shift).cumsum(1) * 2 * np.pi).sin()
-    return sines
-  def forward(self, f0, upp=None):
-    fn = f0.mul(Tensor([[range(1, self.harmonic_num + 2)]], dtype=dtypes.float32).to(f0.device))
-    sine_waves = self._f02sine(fn) * self.sine_amp  #generate sine waveforms
-    uv = self._f02uv(f0)  # generate uv signal
-    noise_amp = uv * self.noise_std + (1 - uv) * self.sine_amp / 3
-    noise = noise_amp * randn_like(sine_waves)
-    sine_waves = sine_waves * uv + noise
-    return sine_waves, uv, noise
-
-class SourceHnNSF:
-  def __init__(self, sampling_rate, harmonic_num=0, sine_amp=0.1, add_noise_std=0.003, voiced_threshold=0):
-    self.sine_amp, self.noise_std = sine_amp, add_noise_std
-    self.l_sin_gen = SineGen(sampling_rate, harmonic_num, sine_amp, add_noise_std, voiced_threshold)
-    self.l_linear = nn.Linear(harmonic_num + 1, 1)
-  def forward(self, x, upp=None):
-    sine_waves, uv, _ = self.l_sin_gen.forward(x, upp)
-    sine_merge = self.l_linear(sine_waves.cast(self.l_linear.weight.dtype)).tanh()
-    noise = randn_like(uv) * self.sine_amp / 3
-    return sine_merge, noise, uv
-
-# most of the hifigan in standard vits is reused here, but need to upsample and construct harmonic source from f0
-class Generator:
-  def __init__(self, sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels):
-    self.sampling_rate, self.inter_channels, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.gin_channels = sampling_rate, inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels
-    self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
-    self.conv_pre = nn.Conv1d(inter_channels, upsample_initial_channel, 7, 1, padding=3)
-    self.f0_upsamp = Upsample(scale_factor=np.prod(upsample_rates))
-    self.m_source = SourceHnNSF(sampling_rate, harmonic_num=8)
-    resblock = ResBlock1 if resblock == '1' else ResBlock2
-    self.ups, self.noise_convs, self.resblocks = [], [], []
-    for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
-      c_cur = upsample_initial_channel//(2**(i+1))
-      self.ups.append(nn.ConvTranspose1d(upsample_initial_channel//(2**i), c_cur, k, u, padding=(k-u)//2))
-      stride_f0 = int(np.prod(upsample_rates[i + 1:]))
-      self.noise_convs.append(nn.Conv1d(1, c_cur, kernel_size=stride_f0 * 2, stride=stride_f0, padding=(stride_f0+1) // 2) if (i + 1 < len(upsample_rates)) else nn.Conv1d(1, c_cur, kernel_size=1))
-    for i in range(len(self.ups)):
-      ch = upsample_initial_channel // (2 ** (i + 1))
-      for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
-        self.resblocks.append(resblock(ch, k, d))
-    self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3)
-    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
-    self.upp = np.prod(upsample_rates)
-  def forward(self, x, f0, g=None):
-    f0 = self.f0_upsamp.forward(f0[:, None]).transpose(1, 2)  # bs,n,t
-    har_source, _, _ = self.m_source.forward(f0, self.upp)
-    har_source = har_source.transpose(1, 2)
-    x = self.conv_pre(x)
-    if g is not None:  x = x + self.cond(g)
-    for i in range(self.num_upsamples):
-      x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
-      x_source = self.noise_convs[i](har_source)
-      x = x + x_source
-      for j in range(self.num_kernels):
-        if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
-        else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
-      x = xs / self.num_kernels
-    return self.conv_post(x.leaky_relu()).tanh()
-
-# **** helpers ****
-
-def randn_like(x:Tensor) -> Tensor: return Tensor.randn(*x.shape, dtype=x.dtype).to(device=x.device)
-
-def tilde(x: Tensor) -> Tensor:
-  if x.dtype == dtypes.bool: return (1 - x).cast(dtypes.bool)
-  return (x + 1) * -1  # this seems to be what the ~ operator does in pytorch for non bool
-
-def lengths_to_padding_mask(lens:Tensor) -> Tensor:
-  bsz, max_lens = lens.shape[0], lens.max().numpy().item()
-  mask = Tensor.arange(max_lens).to(lens.device).reshape(1, max_lens)
-  mask = mask.expand(bsz, -1) >= lens.reshape(bsz, 1).expand(-1, max_lens)
-  return mask.cast(dtypes.bool)
-
-def repeat_expand_2d_left(content, target_len): # content : [h, t]
-  src_len = content.shape[-1]
-  temp = np.arange(src_len+1) * target_len / src_len
-  current_pos, cols = 0, []
-  for i in range(target_len):
-    if i >= temp[current_pos+1]:
-      current_pos += 1
-    cols.append(content[:, current_pos])
-  return Tensor.stack(*cols).transpose(0, 1)
-
-def load_fairseq_cfg(checkpoint_path):
-  assert Path(checkpoint_path).is_file()
-  state = torch_load(checkpoint_path)
-  cfg = state["cfg"] if ("cfg" in state and state["cfg"] is not None) else None
-  if cfg is None: raise RuntimeError(f"No cfg exist in state keys = {state.keys()}")
-  return HParams(**cfg)
-
-def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip_list=[]):
-  assert Path(checkpoint_path).is_file()
-  start_time = time.time()
-  checkpoint_dict = torch_load(checkpoint_path)
-  saved_state_dict = checkpoint_dict['model']
-  weight_g, weight_v, parent = None, None, None
-  for key, v in saved_state_dict.items():
-    if any(layer in key for layer in skip_list): continue
-    try:
-      obj, skip = model, False
-      for k in key.split('.'):
-        if k.isnumeric(): obj = obj[int(k)]
-        elif isinstance(obj, dict): obj = obj[k]
-        else:
-          if k in ["weight_g", "weight_v"]:
-            parent, skip = obj, True
-            if k == "weight_g": weight_g = v
-            else: weight_v = v
-          if not skip:
-            parent = obj
-            obj = getattr(obj, k)
-      if weight_g and weight_v:
-        setattr(obj, "weight_g", weight_g.numpy())
-        setattr(obj, "weight_v", weight_v.numpy())
-        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
-        weight_g, weight_v, parent, skip = None, None, None, False
-      if not skip and obj.shape == v.shape:
-        if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))):  # cast
-          obj.assign(v.to(obj.device).float())
-        else:
-          obj.assign(v.to(obj.device))
-      elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
-    except Exception as e: raise e
-  logging.info(f"Loaded checkpoint '{checkpoint_path}' in {time.time() - start_time:.4f}s")
-  return model, optimizer
-
-def pad_array(arr, target_length):
-  current_length = arr.shape[0]
-  if current_length >= target_length: return arr
-  pad_width = target_length - current_length
-  pad_left = pad_width // 2
-  pad_right = pad_width - pad_left
-  padded_arr = np.pad(arr, (pad_left, pad_right), 'constant', constant_values=(0, 0))
-  return padded_arr
-
-def split_list_by_n(list_collection, n, pre=0):
-  for i in range(0, len(list_collection), n):
-    yield list_collection[i-pre if i-pre>=0 else i: i + n]
-
-def get_sid(spk2id:HParams, speaker:str) -> Tensor:
-  speaker_id = spk2id[speaker]
-  if not speaker_id and type(speaker) is int:
-    if len(spk2id.__dict__) >= speaker: speaker_id = speaker
-  if speaker_id is None: raise RuntimeError(f"speaker={speaker} not in the speaker list")
-  return Tensor([int(speaker_id)], dtype=dtypes.int64).unsqueeze(0)
-
-def get_encoder(ssl_dim) -> Type[SpeechEncoder]:
-  if ssl_dim == 256: return ContentVec256L9
-  if ssl_dim == 768: return ContentVec768L12
-
-#########################################################################################
-# CODE: https://github.com/svc-develop-team/so-vits-svc
-#########################################################################################
-# CONTENTVEC:
-#   CODE: https://github.com/auspicious3000/contentvec
-#   PAPER: https://arxiv.org/abs/2204.09224
-#########################################################################################
-# INSTALLATION: dependencies are for preprocessing and loading/saving audio.
-# pip3 install soundfile librosa praat-parselmouth
-#########################################################################################
-# EXAMPLE USAGE:
-# python3 examples/so_vits_svc.py --model tf2spy --file ~/recording.wav
-#########################################################################################
-# DEMO USAGE (uses audio sample from LJ-Speech):
-# python3 examples/so_vits_svc.py --model saul_goodman
-#########################################################################################
-SO_VITS_SVC_PATH = Path(__file__).parents[1] / "weights/So-VITS-SVC"
-VITS_MODELS = { # config_path, weights_path, config_url, weights_url
-  "saul_goodman" : (SO_VITS_SVC_PATH / "config_saul_gman.json", SO_VITS_SVC_PATH / "pretrained_saul_gman.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/Saul_Goodman_80000/G_80000.pth"),
-  "drake" : (SO_VITS_SVC_PATH / "config_drake.json", SO_VITS_SVC_PATH / "pretrained_drake.pth", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/config_aubrey.json", "https://huggingface.co/jaspa/so-vits-svc/resolve/main/aubrey/pretrained_aubrey.pth"),
-  "cartman" : (SO_VITS_SVC_PATH / "config_cartman.json", SO_VITS_SVC_PATH / "pretrained_cartman.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/EricCartman/G_10200.pth"),
-  "tf2spy" : (SO_VITS_SVC_PATH / "config_tf2spy.json", SO_VITS_SVC_PATH / "pretrained_tf2spy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_spy_60k/G_60000.pth"),
-  "tf2heavy" : (SO_VITS_SVC_PATH / "config_tf2heavy.json", SO_VITS_SVC_PATH / "pretrained_tf2heavy.pth", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/config.json", "https://huggingface.co/Amo/so-vits-svc-4.0_GA/resolve/main/ModelsFolder/TF2_heavy_100k/G_100000.pth"),
-  "lady_gaga" : (SO_VITS_SVC_PATH / "config_gaga.json", SO_VITS_SVC_PATH / "pretrained_gaga.pth", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/config.json", "https://huggingface.co/marcoc2/so-vits-svc-4.0-models/resolve/main/LadyGaga/G_14400.pth")
-}
-ENCODER_MODELS = { # weights_path, weights_url
-  "contentvec": (SO_VITS_SVC_PATH / "contentvec_checkpoint.pt", "https://huggingface.co/lj1995/VoiceConversionWebUI/resolve/main/hubert_base.pt")
-}
-ENCODER_MODEL = "contentvec"
-DEMO_PATH, DEMO_URL = Path(__file__).parents[1] / "temp/LJ037-0171.wav", "https://keithito.com/LJ-Speech-Dataset/LJ037-0171.wav"
-if __name__=="__main__":
-  logging.basicConfig(stream=sys.stdout, level=(logging.INFO if DEBUG < 1 else logging.DEBUG))
-  parser = argparse.ArgumentParser()
-  parser.add_argument("-m", "--model", default=None, help=f"Specify the model to use. All supported models: {VITS_MODELS.keys()}", required=True)
-  parser.add_argument("-f", "--file", default=DEMO_PATH, help=f"Specify the path of the input file")
-  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
-  parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
-  parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
-  parser.add_argument("--speaker", default=None, help="If not specified, the first available speaker is chosen. Usually there is only one speaker per model.")
-  parser.add_argument("--noise_scale", default=0.4)
-  parser.add_argument("--tran", default=0.0, help="Pitch shift, supports positive and negative (semitone) values. Default 0.0")
-  parser.add_argument("--pad_seconds", default=0.5)
-  parser.add_argument("--lg_num", default=0.0)
-  parser.add_argument("--clip_seconds", default=0.0)
-  parser.add_argument("--slice_db", default=-40)
-  args = parser.parse_args()
-
-  vits_model = args.model
-  encoder_location, vits_location = ENCODER_MODELS[ENCODER_MODEL], VITS_MODELS[vits_model]
-
-  Tensor.training = False
-  # Get Synthesizer and ContentVec
-  net_g, hps = Synthesizer.load_from_pretrained(vits_location[0], vits_location[2], vits_location[1], vits_location[3])
-  Encoder = get_encoder(hps.model.ssl_dim)
-  encoder = Encoder.load_from_pretrained(encoder_location[0], encoder_location[1])
-
-  # model config args
-  target_sample, spk2id, hop_length, target_sample = hps.data.sampling_rate, hps.spk, hps.data.hop_length, hps.data.sampling_rate
-  vol_embedding = hps.model.vol_embedding if hasattr(hps.data, "vol_embedding") and hps.model.vol_embedding is not None else False
-
-  # args
-  slice_db, clip_seconds, lg_num, pad_seconds, tran, noise_scale, audio_path = args.slice_db, args.clip_seconds, args.lg_num, args.pad_seconds, args.tran, args.noise_scale, args.file
-  speaker = args.speaker if args.speaker is not None else list(hps.spk.__dict__.keys())[0]
-
-  ### Loading audio and slicing ###
-  if audio_path == DEMO_PATH: fetch(DEMO_URL, DEMO_PATH)
-  assert Path(audio_path).is_file() and Path(audio_path).suffix == ".wav"
-  chunks = preprocess.cut(audio_path, db_thresh=slice_db)
-  audio_data, audio_sr = preprocess.chunks2audio(audio_path, chunks)
-
-  per_size = int(clip_seconds * audio_sr)
-  lg_size = int(lg_num * audio_sr)
-
-  ### Infer per slice ###
-  global_frame = 0
-  audio = []
-  for (slice_tag, data) in audio_data:
-    print(f"\n====segment start, {round(len(data) / audio_sr, 3)}s====")
-    length = int(np.ceil(len(data) / audio_sr * target_sample))
-
-    if slice_tag:
-      print("empty segment")
-      _audio = np.zeros(length)
-      audio.extend(list(pad_array(_audio, length)))
-      global_frame += length // hop_length
-      continue
-
-    datas = [data] if per_size == 0 else split_list_by_n(data, per_size, lg_size)
-
-    for k, dat in enumerate(datas):
-      per_length = int(np.ceil(len(dat) / audio_sr * target_sample)) if clip_seconds!=0 else length
-      pad_len = int(audio_sr * pad_seconds)
-      dat = np.concatenate([np.zeros([pad_len]), dat, np.zeros([pad_len])])
-      raw_path = io.BytesIO()
-      soundfile.write(raw_path, dat, audio_sr, format="wav")
-      raw_path.seek(0)
-
-      ### Infer START ###
-      wav, sr = preprocess.load_audiofile(raw_path)
-      wav = preprocess.sinc_interp_resample(wav, sr, target_sample)[0]
-      wav16k, f0, uv = preprocess.get_unit_f0(wav, tran, hop_length, target_sample)
-      sid = get_sid(spk2id, speaker)
-      n_frames = f0.shape[1]
-
-      # ContentVec infer
-      start = time.time()
-      c = encoder.encode(wav16k)
-      c = repeat_expand_2d_left(c.squeeze(0).realize(), f0.shape[1])  # interpolate speech encoding to match f0
-      c = c.unsqueeze(0).realize()
-      enc_time = time.time() - start
-
-      # VITS infer
-      vits_start = time.time()
-      out_audio, f0 = net_g.infer(c, f0=f0, uv=uv, g=sid, noise_scale=noise_scale, vol=None)
-      out_audio = out_audio[0,0].float().realize()
-      vits_time = time.time() - vits_start
-
-      infer_time = time.time() - start
-      logging.info("total infer time:{:.2f}s, speech_enc time:{:.2f}s, vits time:{:.2f}s".format(infer_time, enc_time, vits_time))
-      ### Infer END ###
-
-      out_sr, out_frame = out_audio.shape[-1], n_frames
-      global_frame += out_frame
-      _audio = out_audio.numpy()
-      pad_len = int(target_sample * pad_seconds)
-      _audio = _audio[pad_len:-pad_len]
-      _audio = pad_array(_audio, per_length)
-      audio.extend(list(_audio))
-
-  audio = np.array(audio)
-  out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model}{f'_spk_{speaker}'}_{args.base_name}.wav")
-  out_path.parent.mkdir(parents=True, exist_ok=True)
-  soundfile.write(out_path, audio, target_sample, format="flac")
-  logging.info(f"Saved audio output to {out_path}")

examples/sovits_helpers/preprocess.py

@@ -1,204 +0,0 @@
-import math
-from typing import Optional, Tuple
-from tinygrad import Tensor, dtypes
-import librosa
-import soundfile
-import numpy as np
-import parselmouth
-
-class PMF0Predictor:  # from https://github.com/svc-develop-team/so-vits-svc/
-  def __init__(self,hop_length=512,f0_min=50,f0_max=1100,sampling_rate=44100):
-    self.hop_length, self.f0_min, self.f0_max, self.sampling_rate, self.name = hop_length, f0_min, f0_max, sampling_rate, "pm"
-  def interpolate_f0(self,f0):
-    vuv_vector = np.zeros_like(f0, dtype=np.float32)
-    vuv_vector[f0 > 0.0] = 1.0
-    vuv_vector[f0 <= 0.0] = 0.0
-    nzindex = np.nonzero(f0)[0]
-    data = f0[nzindex]
-    nzindex = nzindex.astype(np.float32)
-    time_org = self.hop_length / self.sampling_rate * nzindex
-    time_frame = np.arange(f0.shape[0]) * self.hop_length / self.sampling_rate
-    if data.shape[0] <= 0: return np.zeros(f0.shape[0], dtype=np.float32),vuv_vector
-    if data.shape[0] == 1: return np.ones(f0.shape[0], dtype=np.float32) * f0[0],vuv_vector
-    f0 = np.interp(time_frame, time_org, data, left=data[0], right=data[-1])
-    return f0,vuv_vector
-  def compute_f0(self,wav,p_len=None):
-    x = wav
-    if p_len is None: p_len = x.shape[0]//self.hop_length
-    else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
-    time_step = self.hop_length / self.sampling_rate * 1000
-    f0 = parselmouth.Sound(x, self.sampling_rate) \
-                    .to_pitch_ac(time_step=time_step / 1000, voicing_threshold=0.6,pitch_floor=self.f0_min, pitch_ceiling=self.f0_max) \
-                    .selected_array['frequency']
-    pad_size=(p_len - len(f0) + 1) // 2
-    if(pad_size>0 or p_len - len(f0) - pad_size>0):
-      f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
-    f0,uv = self.interpolate_f0(f0)
-    return f0
-  def compute_f0_uv(self,wav,p_len=None):
-    x = wav
-    if p_len is None: p_len = x.shape[0]//self.hop_length
-    else: assert abs(p_len-x.shape[0]//self.hop_length) < 4, "pad length error"
-    time_step = self.hop_length / self.sampling_rate * 1000
-    f0 = parselmouth.Sound(x, self.sampling_rate).to_pitch_ac(
-      time_step=time_step / 1000, voicing_threshold=0.6,
-      pitch_floor=self.f0_min, pitch_ceiling=self.f0_max).selected_array['frequency']
-    pad_size=(p_len - len(f0) + 1) // 2
-    if(pad_size>0 or p_len - len(f0) - pad_size>0):
-      f0 = np.pad(f0,[[pad_size,p_len - len(f0) - pad_size]], mode='constant')
-    f0,uv = self.interpolate_f0(f0)
-    return f0,uv
-
-class Slicer:  # from https://github.com/svc-develop-team/so-vits-svc/
-  def __init__(self, sr: int, threshold: float = -40., min_length: int = 5000, min_interval: int = 300, hop_size: int = 20, max_sil_kept: int = 5000):
-    if not min_length >= min_interval >= hop_size:
-      raise ValueError('The following condition must be satisfied: min_length >= min_interval >= hop_size')
-    if not max_sil_kept >= hop_size:
-      raise ValueError('The following condition must be satisfied: max_sil_kept >= hop_size')
-    min_interval = sr * min_interval / 1000
-    self.threshold = 10 ** (threshold / 20.)
-    self.hop_size = round(sr * hop_size / 1000)
-    self.win_size = min(round(min_interval), 4 * self.hop_size)
-    self.min_length = round(sr * min_length / 1000 / self.hop_size)
-    self.min_interval = round(min_interval / self.hop_size)
-    self.max_sil_kept = round(sr * max_sil_kept / 1000 / self.hop_size)
-  def _apply_slice(self, waveform, begin, end):
-    if len(waveform.shape) > 1: return waveform[:, begin * self.hop_size: min(waveform.shape[1], end * self.hop_size)]
-    else: return waveform[begin * self.hop_size: min(waveform.shape[0], end * self.hop_size)]
-  def slice(self, waveform):
-    samples = librosa.to_mono(waveform) if len(waveform.shape) > 1 else waveform
-    if samples.shape[0] <= self.min_length: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}
-    rms_list = librosa.feature.rms(y=samples, frame_length=self.win_size, hop_length=self.hop_size).squeeze(0)
-    sil_tags, silence_start, clip_start = [], None, 0
-    for i, rms in enumerate(rms_list):
-      if rms < self.threshold:  # Keep looping while frame is silent.
-        if silence_start is None:  # Record start of silent frames.
-          silence_start = i
-        continue
-      if silence_start is None: continue  # Keep looping while frame is not silent and silence start has not been recorded.
-      # Clear recorded silence start if interval is not enough or clip is too short
-      is_leading_silence = silence_start == 0 and i > self.max_sil_kept
-      need_slice_middle = i - silence_start >= self.min_interval and i - clip_start >= self.min_length
-      if not is_leading_silence and not need_slice_middle:
-        silence_start = None
-        continue
-      if i - silence_start <= self.max_sil_kept:  # Need slicing. Record the range of silent frames to be removed.
-        pos = rms_list[silence_start: i + 1].argmin() + silence_start
-        sil_tags.append((0, pos) if silence_start == 0 else (pos, pos))
-        clip_start = pos
-      elif i - silence_start <= self.max_sil_kept * 2:
-        pos = rms_list[i - self.max_sil_kept: silence_start + self.max_sil_kept + 1].argmin()
-        pos += i - self.max_sil_kept
-        pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
-        pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
-        if silence_start == 0:
-          sil_tags.append((0, pos_r))
-          clip_start = pos_r
-        else:
-          sil_tags.append((min(pos_l, pos), max(pos_r, pos)))
-          clip_start = max(pos_r, pos)
-      else:
-        pos_l = rms_list[silence_start: silence_start + self.max_sil_kept + 1].argmin() + silence_start
-        pos_r = rms_list[i - self.max_sil_kept: i + 1].argmin() + i - self.max_sil_kept
-        sil_tags.append((0, pos_r) if silence_start == 0 else (pos_l, pos_r))
-        clip_start = pos_r
-      silence_start = None
-    total_frames = rms_list.shape[0]
-    if silence_start is not None and total_frames - silence_start >= self.min_interval:  # Deal with trailing silence.
-      silence_end = min(total_frames, silence_start + self.max_sil_kept)
-      pos = rms_list[silence_start: silence_end + 1].argmin() + silence_start
-      sil_tags.append((pos, total_frames + 1))
-    if len(sil_tags) == 0: return {"0": {"slice": False, "split_time": f"0,{len(waveform)}"}}  # Apply and return slices.
-    chunks = []
-    if sil_tags[0][0]:
-      chunks.append({"slice": False, "split_time": f"0,{min(waveform.shape[0], sil_tags[0][0] * self.hop_size)}"})
-    for i in range(0, len(sil_tags)):
-      if i: chunks.append({"slice": False, "split_time": f"{sil_tags[i - 1][1] * self.hop_size},{min(waveform.shape[0], sil_tags[i][0] * self.hop_size)}"})
-      chunks.append({"slice": True, "split_time": f"{sil_tags[i][0] * self.hop_size},{min(waveform.shape[0], sil_tags[i][1] * self.hop_size)}"})
-    if sil_tags[-1][1] * self.hop_size < len(waveform):
-      chunks.append({"slice": False, "split_time": f"{sil_tags[-1][1] * self.hop_size},{len(waveform)}"})
-    chunk_dict = {}
-    for i in range(len(chunks)): chunk_dict[str(i)] = chunks[i]
-    return chunk_dict
-
-# sinc_interp_hann audio resampling
-class Resample:
-  def __init__(self, orig_freq:int=16000, new_freq:int=16000, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None, dtype:Optional[dtypes]=None):
-    self.orig_freq, self.new_freq, self.lowpass_filter_width, self.rolloff, self.beta = orig_freq, new_freq, lowpass_filter_width, rolloff, beta
-    self.gcd = math.gcd(int(self.orig_freq), int(self.new_freq))
-    self.kernel, self.width = self._get_sinc_resample_kernel(dtype) if self.orig_freq != self.new_freq else (None, None)
-  def __call__(self, waveform:Tensor) -> Tensor:
-    if self.orig_freq == self.new_freq: return waveform
-    return self._apply_sinc_resample_kernel(waveform)
-  def _apply_sinc_resample_kernel(self, waveform:Tensor):
-    if not waveform.is_floating_point(): raise TypeError(f"Waveform tensor expected to be of type float, but received {waveform.dtype}.")
-    orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
-    shape = waveform.shape
-    waveform = waveform.reshape(-1, shape[-1])  # pack batch
-    num_wavs, length = waveform.shape
-    target_length = int(math.ceil(new_freq * length / orig_freq))
-    waveform = waveform.pad((self.width, self.width + orig_freq))
-    resampled = waveform[:, None].conv2d(self.kernel, stride=orig_freq)
-    resampled = resampled.transpose(1, 2).reshape(num_wavs, -1)
-    resampled = resampled[..., :target_length]
-    resampled = resampled.reshape(shape[:-1] + resampled.shape[-1:])  # unpack batch
-    return resampled
-  def _get_sinc_resample_kernel(self, dtype=None):
-    orig_freq, new_freq = (int(self.orig_freq) // self.gcd), (int(self.new_freq) // self.gcd)
-    if self.lowpass_filter_width <= 0: raise ValueError("Low pass filter width should be positive.")
-    base_freq = min(orig_freq, new_freq)
-    base_freq *= self.rolloff
-    width = math.ceil(self.lowpass_filter_width * orig_freq / base_freq)
-    idx = Tensor.arange(-width, width + orig_freq, dtype=(dtype if dtype is not None else dtypes.float32))[None, None] / orig_freq
-    t = Tensor.arange(0, -new_freq, -1, dtype=dtype)[:, None, None] / new_freq + idx
-    t *= base_freq
-    t = t.clip(-self.lowpass_filter_width, self.lowpass_filter_width)
-    window = (t * math.pi / self.lowpass_filter_width / 2).cos() ** 2
-    t *= math.pi
-    scale = base_freq / orig_freq
-    kernels = Tensor.where(t == 0, Tensor(1.0, dtype=t.dtype).to(t.device), t.sin() / t)
-    kernels *= window * scale
-    if dtype is None: kernels = kernels.cast(dtype=dtypes.float32)
-    return kernels, width
-
-def sinc_interp_resample(x:Tensor, orig_freq:int=16000, new_freq:int=1600, lowpass_filter_width:int=6, rolloff:float=0.99, beta:Optional[float]=None):
-  resamp = Resample(orig_freq, new_freq, lowpass_filter_width, rolloff, beta, x.dtype)
-  return resamp(x)
-
-def cut(audio_path, db_thresh=-30, min_len=5000):
-  audio, sr = librosa.load(audio_path, sr=None)
-  slicer = Slicer(sr=sr, threshold=db_thresh, min_length=min_len)
-  chunks = slicer.slice(audio)
-  return chunks
-
-def chunks2audio(audio_path, chunks):
-  chunks = dict(chunks)
-  audio, sr = load_audiofile(audio_path)
-  if len(audio.shape) == 2 and audio.shape[1] >= 2:
-    audio = audio.mean(0).unsqueeze(0)
-  audio = audio.numpy()[0]
-  result = []
-  for k, v in chunks.items():
-    tag = v["split_time"].split(",")
-    if tag[0] != tag[1]:
-      result.append((v["slice"], audio[int(tag[0]):int(tag[1])]))
-  return result, sr
-
-def load_audiofile(filepath:str, frame_offset:int=0, num_frames:int=-1, channels_first:bool=True):
-  with soundfile.SoundFile(filepath, "r") as file_:
-    frames = file_._prepare_read(frame_offset, None, num_frames)
-    waveform = file_.read(frames, "float32", always_2d=True)
-    sample_rate = file_.samplerate
-  waveform = Tensor(waveform)
-  if channels_first: waveform = waveform.transpose(0, 1)
-  return waveform, sample_rate
-
-def get_unit_f0(wav:Tensor, tran, hop_length, target_sample, f0_filter=False) -> Tuple[Tensor,Tensor,Tensor]:
-  f0_predictor = PMF0Predictor(hop_length, sampling_rate=target_sample)
-  f0, uv = f0_predictor.compute_f0_uv(wav.numpy())
-  if f0_filter and sum(f0) == 0: raise RuntimeError("No voice detected")
-  f0 = Tensor(f0.astype(np.float32)).float()
-  f0 = (f0 * 2 ** (tran / 12)).unsqueeze(0)
-  uv = Tensor(uv.astype(np.float32)).float().unsqueeze(0)
-  wav16k = sinc_interp_resample(wav[None,:], target_sample, 16000)[0]
-  return wav16k.realize(), f0.realize(), uv.realize()

examples/train_efficientnet.py

@@ -1,104 +0,0 @@
-import traceback
-import time
-from multiprocessing import Process, Queue
-import numpy as np
-from tinygrad.nn.state import get_parameters
-from tinygrad.nn import optim
-from tinygrad.helpers import getenv, trange
-from tinygrad.tensor import Tensor
-from extra.datasets import fetch_cifar
-from extra.models.efficientnet import EfficientNet
-
-class TinyConvNet:
-  def __init__(self, classes=10):
-    conv = 3
-    inter_chan, out_chan = 8, 16   # for speed
-    self.c1 = Tensor.uniform(inter_chan,3,conv,conv)
-    self.c2 = Tensor.uniform(out_chan,inter_chan,conv,conv)
-    self.l1 = Tensor.uniform(out_chan*6*6, classes)
-
-  def forward(self, x):
-    x = x.conv2d(self.c1).relu().max_pool2d()
-    x = x.conv2d(self.c2).relu().max_pool2d()
-    x = x.reshape(shape=[x.shape[0], -1])
-    return x.dot(self.l1)
-
-if __name__ == "__main__":
-  IMAGENET = getenv("IMAGENET")
-  classes = 1000 if IMAGENET else 10
-
-  TINY = getenv("TINY")
-  TRANSFER = getenv("TRANSFER")
-  if TINY:
-    model = TinyConvNet(classes)
-  elif TRANSFER:
-    model = EfficientNet(getenv("NUM", 0), classes, has_se=True)
-    model.load_from_pretrained()
-  else:
-    model = EfficientNet(getenv("NUM", 0), classes, has_se=False)
-
-  parameters = get_parameters(model)
-  print("parameter count", len(parameters))
-  optimizer = optim.Adam(parameters, lr=0.001)
-
-  BS, steps = getenv("BS", 64 if TINY else 16), getenv("STEPS", 2048)
-  print(f"training with batch size {BS} for {steps} steps")
-
-  if IMAGENET:
-    from extra.datasets.imagenet import fetch_batch
-    def loader(q):
-      while 1:
-        try:
-          q.put(fetch_batch(BS))
-        except Exception:
-          traceback.print_exc()
-    q = Queue(16)
-    for i in range(2):
-      p = Process(target=loader, args=(q,))
-      p.daemon = True
-      p.start()
-  else:
-    X_train, Y_train, _, _ = fetch_cifar()
-    X_train = X_train.reshape((-1, 3, 32, 32))
-    Y_train = Y_train.reshape((-1,))
-
-  with Tensor.train():
-    for i in (t := trange(steps)):
-      if IMAGENET:
-        X, Y = q.get(True)
-      else:
-        samp = np.random.randint(0, X_train.shape[0], size=(BS))
-        X, Y = X_train.numpy()[samp], Y_train.numpy()[samp]
-
-      st = time.time()
-      out = model.forward(Tensor(X.astype(np.float32), requires_grad=False))
-      fp_time = (time.time()-st)*1000.0
-
-      y = np.zeros((BS,classes), np.float32)
-      y[range(y.shape[0]),Y] = -classes
-      y = Tensor(y, requires_grad=False)
-      loss = out.log_softmax().mul(y).mean()
-
-      optimizer.zero_grad()
-
-      st = time.time()
-      loss.backward()
-      bp_time = (time.time()-st)*1000.0
-
-      st = time.time()
-      optimizer.step()
-      opt_time = (time.time()-st)*1000.0
-
-      st = time.time()
-      loss = loss.numpy()
-      cat = out.argmax(axis=1).numpy()
-      accuracy = (cat == Y).mean()
-      finish_time = (time.time()-st)*1000.0
-
-      # printing
-      t.set_description("loss %.2f accuracy %.2f -- %.2f + %.2f + %.2f + %.2f = %.2f" %
-        (loss, accuracy,
-        fp_time, bp_time, opt_time, finish_time,
-        fp_time + bp_time + opt_time + finish_time))
-
-      del out, y, loss

examples/vit.py

@@ -1,46 +0,0 @@
-import ast
-import numpy as np
-from PIL import Image
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv, fetch
-from extra.models.vit import ViT
-"""
-fn = "gs://vit_models/augreg/Ti_16-i21k-300ep-lr_0.001-aug_none-wd_0.03-do_0.0-sd_0.0.npz"
-import tensorflow as tf
-with tf.io.gfile.GFile(fn, "rb") as f:
-  dat = f.read()
-  with open("cache/"+ fn.rsplit("/", 1)[1], "wb") as g:
-    g.write(dat)
-"""
-
-Tensor.training = False
-if getenv("LARGE", 0) == 1:
-  m = ViT(embed_dim=768, num_heads=12)
-else:
-  # tiny
-  m = ViT(embed_dim=192, num_heads=3)
-m.load_from_pretrained()
-
-# category labels
-lbls = ast.literal_eval(fetch("https://gist.githubusercontent.com/yrevar/942d3a0ac09ec9e5eb3a/raw/238f720ff059c1f82f368259d1ca4ffa5dd8f9f5/imagenet1000_clsidx_to_labels.txt").read_text())
-
-#url = "https://upload.wikimedia.org/wikipedia/commons/4/41/Chicken.jpg"
-url = "https://repository-images.githubusercontent.com/296744635/39ba6700-082d-11eb-98b8-cb29fb7369c0"
-
-# junk
-img = Image.open(fetch(url))
-aspect_ratio = img.size[0] / img.size[1]
-img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
-img = np.array(img)
-y0,x0=(np.asarray(img.shape)[:2]-224)//2
-img = img[y0:y0+224, x0:x0+224]
-img = np.moveaxis(img, [2,0,1], [0,1,2])
-img = img.astype(np.float32)[:3].reshape(1,3,224,224)
-img /= 255.0
-img -= 0.5
-img /= 0.5
-
-out = m.forward(Tensor(img))
-outnp = out.numpy().ravel()
-choice = outnp.argmax()
-print(out.shape, choice, outnp[choice], lbls[choice])

examples/vits.py

@@ -1,740 +0,0 @@
-import json, logging, math, re, sys, time, wave, argparse, numpy as np
-from phonemizer.phonemize import default_separator, _phonemize
-from phonemizer.backend import EspeakBackend
-from phonemizer.punctuation import Punctuation
-from functools import reduce
-from pathlib import Path
-from typing import List
-from tinygrad import nn, dtypes
-from tinygrad.helpers import fetch
-from tinygrad.nn.state import torch_load
-from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit
-from unidecode import unidecode
-
-LRELU_SLOPE = 0.1
-
-class Synthesizer:
-  def __init__(self, n_vocab, spec_channels, segment_size, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, n_speakers=0, gin_channels=0, use_sdp=True, emotion_embedding=False, **kwargs):
-    self.n_vocab, self.spec_channels, self.inter_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.resblock, self.resblock_kernel_sizes, self.resblock_dilation_sizes, self.upsample_rates, self.upsample_initial_channel, self.upsample_kernel_sizes, self.segment_size, self.n_speakers, self.gin_channels, self.use_sdp = n_vocab, spec_channels, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, segment_size, n_speakers, gin_channels, use_sdp
-    self.enc_p = TextEncoder(n_vocab, inter_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding)
-    self.dec = Generator(inter_channels, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=gin_channels)
-    self.enc_q = PosteriorEncoder(spec_channels, inter_channels, hidden_channels, 5, 1, 16, gin_channels=gin_channels)
-    self.flow = ResidualCouplingBlock(inter_channels, hidden_channels, 5, 1, 4, gin_channels=gin_channels)
-    self.dp = StochasticDurationPredictor(hidden_channels, 192, 3, 0.5, 4, gin_channels=gin_channels) if use_sdp else DurationPredictor(hidden_channels, 256, 3, 0.5, gin_channels=gin_channels)
-    if n_speakers > 1: self.emb_g = nn.Embedding(n_speakers, gin_channels)
-  def infer(self, x, x_lengths, sid=None, noise_scale=1.0, length_scale=1, noise_scale_w=1., max_len=None, emotion_embedding=None, max_y_length_estimate_scale=None, pad_length=-1):
-    x, m_p, logs_p, x_mask = self.enc_p.forward(x.realize(), x_lengths.realize(), emotion_embedding.realize() if emotion_embedding is not None else emotion_embedding)
-    g = self.emb_g(sid.reshape(1, 1)).squeeze(1).unsqueeze(-1) if self.n_speakers > 0 else None
-    logw = self.dp.forward(x, x_mask.realize(), g=g.realize(), reverse=self.use_sdp, noise_scale=noise_scale_w if self.use_sdp else 1.0)
-    w_ceil = Tensor.ceil(logw.exp() * x_mask * length_scale)
-    y_lengths = Tensor.maximum(w_ceil.sum([1, 2]), 1).cast(dtypes.int64)
-    return self.generate(g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length)
-  def generate(self, g, logs_p, m_p, max_len, max_y_length_estimate_scale, noise_scale, w_ceil, x, x_mask, y_lengths, pad_length):
-    max_y_length = y_lengths.max().item() if max_y_length_estimate_scale is None else max(15, x.shape[-1]) * max_y_length_estimate_scale
-    y_mask = sequence_mask(y_lengths, max_y_length).unsqueeze(1).cast(x_mask.dtype)
-    attn_mask = x_mask.unsqueeze(2) * y_mask.unsqueeze(-1)
-    attn = generate_path(w_ceil, attn_mask)
-    m_p_2 = attn.squeeze(1).matmul(m_p.transpose(1, 2)).transpose(1, 2)        # [b, t', t], [b, t, d] -> [b, d, t']
-    logs_p_2 = attn.squeeze(1).matmul(logs_p.transpose(1, 2)).transpose(1, 2)  # [b, t', t], [b, t, d] -> [b, d, t']
-    z_p = m_p_2 + Tensor.randn(*m_p_2.shape, dtype=m_p_2.dtype) * logs_p_2.exp() * noise_scale
-    row_len = y_mask.shape[2]
-    if pad_length > -1:
-      # Pad flow forward inputs to enable JIT
-      assert pad_length > row_len, "pad length is too small"
-      y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len))).cast(z_p.dtype)
-      # New y_mask tensor to remove sts mask
-      y_mask = Tensor(y_mask.numpy(), device=y_mask.device, dtype=y_mask.dtype, requires_grad=y_mask.requires_grad)
-      z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), value=1).unsqueeze(0)
-    z = self.flow.forward(z_p.realize(), y_mask.realize(), g=g.realize(), reverse=True)
-    result_length = reduce(lambda x, y: x * y, self.dec.upsample_rates, row_len)
-    o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)[:, :, :result_length]
-    if max_y_length_estimate_scale is not None:
-      length_scaler = o.shape[-1] / max_y_length
-      o.realize()
-      real_max_y_length = y_lengths.max().numpy()
-      if real_max_y_length > max_y_length:
-        logging.warning(f"Underestimated max length by {(((real_max_y_length / max_y_length) * 100) - 100):.2f}%, recomputing inference without estimate...")
-        return self.generate(g, logs_p, m_p, max_len, None, noise_scale, w_ceil, x, x_mask, y_lengths)
-      if real_max_y_length < max_y_length:
-        overestimation = ((max_y_length / real_max_y_length) * 100) - 100
-        logging.info(f"Overestimated max length by {overestimation:.2f}%")
-        if overestimation > 10: logging.warning("Warning: max length overestimated by more than 10%")
-      o = o[:, :, :(real_max_y_length * length_scaler).astype(np.int32)]
-    return o
-
-class StochasticDurationPredictor:
-  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, n_flows=4, gin_channels=0):
-    filter_channels = in_channels # it needs to be removed from future version.
-    self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.n_flows, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, n_flows, gin_channels
-    self.log_flow, self.flows = Log(), [ElementwiseAffine(2)]
-    for _ in range(n_flows):
-      self.flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
-      self.flows.append(Flip())
-    self.post_pre, self.post_proj = nn.Conv1d(1, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
-    self.post_convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
-    self.post_flows = [ElementwiseAffine(2)]
-    for _ in range(4):
-      self.post_flows.append(ConvFlow(2, filter_channels, kernel_size, n_layers=3))
-      self.post_flows.append(Flip())
-    self.pre, self.proj = nn.Conv1d(in_channels, filter_channels, 1), nn.Conv1d(filter_channels, filter_channels, 1)
-    self.convs = DDSConv(filter_channels, kernel_size, n_layers=3, p_dropout=p_dropout)
-    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, filter_channels, 1)
-  @TinyJit
-  def forward(self, x: Tensor, x_mask, w=None, g=None, reverse=False, noise_scale=1.0):
-    x = self.pre(x.detach())
-    if g is not None: x = x + self.cond(g.detach())
-    x = self.convs.forward(x, x_mask)
-    x = self.proj(x) * x_mask
-    if not reverse:
-      flows = self.flows
-      assert w is not None
-      log_det_tot_q = 0
-      h_w = self.post_proj(self.post_convs.forward(self.post_pre(w), x_mask)) * x_mask
-      e_q = Tensor.randn(w.size(0), 2, w.size(2), dtype=x.dtype).to(device=x.device) * x_mask
-      z_q = e_q
-      for flow in self.post_flows:
-        z_q, log_det_q = flow.forward(z_q, x_mask, g=(x + h_w))
-        log_det_tot_q += log_det_q
-      z_u, z1 = z_q.split([1, 1], 1)
-      u = z_u.sigmoid() * x_mask
-      z0 = (w - u) * x_mask
-      log_det_tot_q += Tensor.sum((z_u.logsigmoid() + (-z_u).logsigmoid()) * x_mask, [1,2])
-      log_q = Tensor.sum(-0.5 * (math.log(2*math.pi) + (e_q**2)) * x_mask, [1,2]) - log_det_tot_q
-      log_det_tot = 0
-      z0, log_det = self.log_flow.forward(z0, x_mask)
-      log_det_tot += log_det
-      z = z0.cat(z1, 1)
-      for flow in flows:
-        z, log_det = flow.forward(z, x_mask, g=x, reverse=reverse)
-        log_det_tot = log_det_tot + log_det
-      nll = Tensor.sum(0.5 * (math.log(2*math.pi) + (z**2)) * x_mask, [1,2]) - log_det_tot
-      return (nll + log_q).realize() # [b]
-    flows = list(reversed(self.flows))
-    flows = flows[:-2] + [flows[-1]] # remove a useless vflow
-    z = Tensor.randn(x.shape[0], 2, x.shape[2], dtype=x.dtype).to(device=x.device) * noise_scale
-    for flow in flows: z = flow.forward(z, x_mask, g=x, reverse=reverse)
-    z0, z1 = z.split([1, 1], 1)
-    return z0.realize()
-
-class DurationPredictor:
-  def __init__(self, in_channels, filter_channels, kernel_size, p_dropout, gin_channels=0):
-    self.in_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.gin_channels = in_channels, filter_channels, kernel_size, p_dropout, gin_channels
-    self.conv_1, self.norm_1 = nn.Conv1d(in_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
-    self.conv_2, self.norm_2 = nn.Conv1d(filter_channels, filter_channels, kernel_size, padding=kernel_size//2), LayerNorm(filter_channels)
-    self.proj = nn.Conv1d(filter_channels, 1, 1)
-    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, in_channels, 1)
-  def forward(self, x: Tensor, x_mask, g=None):
-    x = x.detach()
-    if g is not None: x = x + self.cond(g.detach())
-    x = self.conv_1(x * x_mask).relu()
-    x = self.norm_1(x).dropout(self.p_dropout)
-    x = self.conv_2(x * x_mask).relu(x)
-    x = self.norm_2(x).dropout(self.p_dropout)
-    return self.proj(x * x_mask) * x_mask
-
-class TextEncoder:
-  def __init__(self, n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, emotion_embedding):
-    self.n_vocab, self.out_channels, self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout = n_vocab, out_channels, hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
-    if n_vocab!=0:self.emb = nn.Embedding(n_vocab, hidden_channels)
-    if emotion_embedding: self.emo_proj = nn.Linear(1024, hidden_channels)
-    self.encoder = Encoder(hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout)
-    self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
-  @TinyJit
-  def forward(self, x: Tensor, x_lengths: Tensor, emotion_embedding=None):
-    if self.n_vocab!=0: x = (self.emb(x) * math.sqrt(self.hidden_channels))
-    if emotion_embedding: x = x + self.emo_proj(emotion_embedding).unsqueeze(1)
-    x = x.transpose(1, -1)  # [b, t, h] -transpose-> [b, h, t]
-    x_mask = sequence_mask(x_lengths, x.shape[2]).unsqueeze(1).cast(x.dtype)
-    x = self.encoder.forward(x * x_mask, x_mask)
-    m, logs = (self.proj(x) * x_mask).split(self.out_channels, dim=1)
-    return x.realize(), m.realize(), logs.realize(), x_mask.realize()
-
-class ResidualCouplingBlock:
-  def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows=4, gin_channels=0):
-    self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.n_flows, self.gin_channels = channels, hidden_channels, kernel_size, dilation_rate, n_layers, n_flows, gin_channels
-    self.flows = []
-    for _ in range(n_flows):
-      self.flows.append(ResidualCouplingLayer(channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels, mean_only=True))
-      self.flows.append(Flip())
-  @TinyJit
-  def forward(self, x, x_mask, g=None, reverse=False):
-    for flow in reversed(self.flows) if reverse else self.flows: x = flow.forward(x, x_mask, g=g, reverse=reverse)
-    return x.realize()
-
-class PosteriorEncoder:
-  def __init__(self, in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0):
-    self.in_channels, self.out_channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels = in_channels, out_channels, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels
-    self.pre, self.proj = nn.Conv1d(in_channels, hidden_channels, 1), nn.Conv1d(hidden_channels, out_channels * 2, 1)
-    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=gin_channels)
-  def forward(self, x, x_lengths, g=None):
-    x_mask = sequence_mask(x_lengths, x.size(2)).unsqueeze(1).cast(x.dtype)
-    stats = self.proj(self.enc.forward(self.pre(x) * x_mask, x_mask, g=g)) * x_mask
-    m, logs = stats.split(self.out_channels, dim=1)
-    z = (m + Tensor.randn(m.shape, m.dtype) * logs.exp()) * x_mask
-    return z, m, logs, x_mask
-
-class Generator:
-  def __init__(self, initial_channel, resblock, resblock_kernel_sizes, resblock_dilation_sizes, upsample_rates, upsample_initial_channel, upsample_kernel_sizes, gin_channels=0):
-    self.num_kernels, self.num_upsamples = len(resblock_kernel_sizes), len(upsample_rates)
-    self.conv_pre = nn.Conv1d(initial_channel, upsample_initial_channel, 7, 1, padding=3)
-    resblock = ResBlock1 if resblock == '1' else ResBlock2
-    self.ups = [nn.ConvTranspose1d(upsample_initial_channel//(2**i), upsample_initial_channel//(2**(i+1)), k, u, padding=(k-u)//2) for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes))]
-    self.resblocks = []
-    self.upsample_rates = upsample_rates
-    for i in range(len(self.ups)):
-      ch = upsample_initial_channel // (2 ** (i + 1))
-      for _, (k, d) in enumerate(zip(resblock_kernel_sizes, resblock_dilation_sizes)):
-        self.resblocks.append(resblock(ch, k, d))
-    self.conv_post = nn.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
-    if gin_channels != 0: self.cond = nn.Conv1d(gin_channels, upsample_initial_channel, 1)
-  @TinyJit
-  def forward(self, x: Tensor, g=None):
-    x = self.conv_pre(x)
-    if g is not None:  x = x + self.cond(g)
-    for i in range(self.num_upsamples):
-      x = self.ups[i](x.leaky_relu(LRELU_SLOPE))
-      xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
-      x = (xs / self.num_kernels).realize()
-    res = self.conv_post(x.leaky_relu()).tanh().realize()
-    return res
-
-class LayerNorm(nn.LayerNorm):
-  def __init__(self, channels, eps=1e-5): super().__init__(channels, eps, elementwise_affine=True)
-  def forward(self, x: Tensor): return self.__call__(x.transpose(1, -1)).transpose(1, -1)
-
-class WN:
-  def __init__(self, hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels=0, p_dropout=0):
-    assert (kernel_size % 2 == 1)
-    self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.gin_channels, self.p_dropout = hidden_channels, kernel_size, dilation_rate, n_layers, gin_channels, p_dropout
-    self.in_layers, self.res_skip_layers = [], []
-    if gin_channels != 0: self.cond_layer = nn.Conv1d(gin_channels, 2 * hidden_channels * n_layers, 1)
-    for i in range(n_layers):
-      dilation = dilation_rate ** i
-      self.in_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels, kernel_size, dilation=dilation, padding=int((kernel_size * dilation - dilation) / 2)))
-      self.res_skip_layers.append(nn.Conv1d(hidden_channels, 2 * hidden_channels if i < n_layers - 1 else hidden_channels, 1))
-  def forward(self, x, x_mask, g=None, **kwargs):
-    output = Tensor.zeros_like(x)
-    if g is not None: g = self.cond_layer(g)
-    for i in range(self.n_layers):
-      x_in = self.in_layers[i](x)
-      if g is not None:
-        cond_offset = i * 2 * self.hidden_channels
-        g_l = g[:, cond_offset:cond_offset + 2 * self.hidden_channels, :]
-      else:
-        g_l = Tensor.zeros_like(x_in)
-      acts = fused_add_tanh_sigmoid_multiply(x_in, g_l, self.hidden_channels)
-      res_skip_acts = self.res_skip_layers[i](acts)
-      if i < self.n_layers - 1:
-        x = (x + res_skip_acts[:, :self.hidden_channels, :]) * x_mask
-        output = output + res_skip_acts[:, self.hidden_channels:, :]
-      else:
-        output = output + res_skip_acts
-    return output * x_mask
-
-class ResBlock1:
-  def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
-    self.convs1 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(3)]
-    self.convs2 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) for _ in range(3)]
-  def forward(self, x: Tensor, x_mask=None):
-    for c1, c2 in zip(self.convs1, self.convs2):
-      xt = x.leaky_relu(LRELU_SLOPE)
-      xt = c1(xt if x_mask is None else xt * x_mask).leaky_relu(LRELU_SLOPE)
-      x = c2(xt if x_mask is None else xt * x_mask) + x
-    return x if x_mask is None else x * x_mask
-
-class ResBlock2:
-  def __init__(self, channels, kernel_size=3, dilation=(1, 3)):
-    self.convs = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[i], padding=get_padding(kernel_size, dilation[i])) for i in range(2)]
-  def forward(self, x, x_mask=None):
-    for c in self.convs:
-      xt = x.leaky_relu(LRELU_SLOPE)
-      xt = c(xt if x_mask is None else xt * x_mask)
-      x = xt + x
-    return x if x_mask is None else x * x_mask
-
-class DDSConv: # Dilated and Depth-Separable Convolution
-  def __init__(self, channels, kernel_size, n_layers, p_dropout=0.):
-    self.channels, self.kernel_size, self.n_layers, self.p_dropout = channels, kernel_size, n_layers, p_dropout
-    self.convs_sep, self.convs_1x1, self.norms_1, self.norms_2 = [], [], [], []
-    for i in range(n_layers):
-      dilation = kernel_size ** i
-      padding = (kernel_size * dilation - dilation) // 2
-      self.convs_sep.append(nn.Conv1d(channels, channels, kernel_size, groups=channels, dilation=dilation, padding=padding))
-      self.convs_1x1.append(nn.Conv1d(channels, channels, 1))
-      self.norms_1.append(LayerNorm(channels))
-      self.norms_2.append(LayerNorm(channels))
-  def forward(self, x, x_mask, g=None):
-    if g is not None: x = x + g
-    for i in range(self.n_layers):
-      y = self.convs_sep[i](x * x_mask)
-      y = self.norms_1[i].forward(y).gelu()
-      y = self.convs_1x1[i](y)
-      y = self.norms_2[i].forward(y).gelu()
-      x = x + y.dropout(self.p_dropout)
-    return x * x_mask
-
-class ConvFlow:
-  def __init__(self, in_channels, filter_channels, kernel_size, n_layers, num_bins=10, tail_bound=5.0):
-    self.in_channels, self.filter_channels, self.kernel_size, self.n_layers, self.num_bins, self.tail_bound = in_channels, filter_channels, kernel_size, n_layers, num_bins, tail_bound
-    self.half_channels = in_channels // 2
-    self.pre = nn.Conv1d(self.half_channels, filter_channels, 1)
-    self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
-    self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
-  def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = x.split([self.half_channels] * 2, 1)
-    h = self.proj(self.convs.forward(self.pre(x0), x_mask, g=g)) * x_mask
-    b, c, t = x0.shape
-    h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
-    un_normalized_widths = h[..., :self.num_bins] / math.sqrt(self.filter_channels)
-    un_normalized_heights = h[..., self.num_bins:2*self.num_bins] / math.sqrt(self.filter_channels)
-    un_normalized_derivatives = h[..., 2 * self.num_bins:]
-    x1, log_abs_det = piecewise_rational_quadratic_transform(x1, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=reverse, tails='linear', tail_bound=self.tail_bound)
-    x = x0.cat(x1, dim=1) * x_mask
-    return x if reverse else (x, Tensor.sum(log_abs_det * x_mask, [1,2]))
-
-class ResidualCouplingLayer:
-  def __init__(self, channels, hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=0, gin_channels=0, mean_only=False):
-    assert channels % 2 == 0, "channels should be divisible by 2"
-    self.channels, self.hidden_channels, self.kernel_size, self.dilation_rate, self.n_layers, self.mean_only = channels, hidden_channels, kernel_size, dilation_rate, n_layers, mean_only
-    self.half_channels = channels // 2
-    self.pre = nn.Conv1d(self.half_channels, hidden_channels, 1)
-    self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
-    self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
-  def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = x.split([self.half_channels] * 2, 1)
-    stats = self.post(self.enc.forward(self.pre(x0) * x_mask, x_mask, g=g)) * x_mask
-    if not self.mean_only:
-      m, logs = stats.split([self.half_channels] * 2, 1)
-    else:
-      m = stats
-      logs = Tensor.zeros_like(m)
-    if not reverse: return x0.cat((m + x1 * logs.exp() * x_mask), dim=1)
-    return x0.cat(((x1 - m) * (-logs).exp() * x_mask), dim=1)
-
-class Log:
-  def forward(self, x : Tensor, x_mask, reverse=False):
-    if not reverse:
-      y = x.maximum(1e-5).log() * x_mask
-      return y, (-y).sum([1, 2])
-    return x.exp() * x_mask
-
-class Flip:
-  def forward(self, x: Tensor, *args, reverse=False, **kwargs):
-    return x.flip([1]) if reverse else (x.flip([1]), Tensor.zeros(x.shape[0], dtype=x.dtype).to(device=x.device))
-
-class ElementwiseAffine:
-  def __init__(self, channels): self.m, self.logs = Tensor.zeros(channels, 1), Tensor.zeros(channels, 1)
-  def forward(self, x, x_mask, reverse=False, **kwargs): # x if reverse else y, logdet
-    return (x - self.m) * Tensor.exp(-self.logs) * x_mask if reverse \
-      else ((self.m + Tensor.exp(self.logs) * x) * x_mask, Tensor.sum(self.logs * x_mask, [1, 2]))
-
-class MultiHeadAttention:
-  def __init__(self, channels, out_channels, n_heads, p_dropout=0., window_size=None, heads_share=True, block_length=None, proximal_bias=False, proximal_init=False):
-    assert channels % n_heads == 0
-    self.channels, self.out_channels, self.n_heads, self.p_dropout, self.window_size, self.heads_share, self.block_length, self.proximal_bias, self.proximal_init = channels, out_channels, n_heads, p_dropout, window_size, heads_share, block_length, proximal_bias, proximal_init
-    self.attn, self.k_channels  = None, channels // n_heads
-    self.conv_q, self.conv_k, self.conv_v = [nn.Conv1d(channels, channels, 1) for _ in range(3)]
-    self.conv_o = nn.Conv1d(channels, out_channels, 1)
-    if window_size is not None: self.emb_rel_k, self.emb_rel_v = [Tensor.randn(1 if heads_share else n_heads, window_size * 2 + 1, self.k_channels) * (self.k_channels ** -0.5) for _ in range(2)]
-  def forward(self, x, c, attn_mask=None):
-    q, k, v = self.conv_q(x), self.conv_k(c), self.conv_v(c)
-    x, self.attn = self.attention(q, k, v, mask=attn_mask)
-    return self.conv_o(x)
-  def attention(self, query: Tensor, key: Tensor, value: Tensor, mask=None):# reshape [b, d, t] -> [b, n_h, t, d_k]
-    b, d, t_s, t_t = key.shape[0], key.shape[1], key.shape[2], query.shape[2]
-    query = query.reshape(b, self.n_heads, self.k_channels, t_t).transpose(2, 3)
-    key = key.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
-    value = value.reshape(b, self.n_heads, self.k_channels, t_s).transpose(2, 3)
-    scores = (query / math.sqrt(self.k_channels)) @ key.transpose(-2, -1)
-    if self.window_size is not None:
-      assert t_s == t_t, "Relative attention is only available for self-attention."
-      key_relative_embeddings = self._get_relative_embeddings(self.emb_rel_k, t_s)
-      rel_logits = self._matmul_with_relative_keys(query / math.sqrt(self.k_channels), key_relative_embeddings)
-      scores = scores + self._relative_position_to_absolute_position(rel_logits)
-    if mask is not None:
-      scores = Tensor.where(mask, scores, -1e4)
-      if self.block_length is not None:
-        assert t_s == t_t, "Local attention is only available for self-attention."
-        scores = Tensor.where(Tensor.ones_like(scores).triu(-self.block_length).tril(self.block_length), scores, -1e4)
-    p_attn = scores.softmax(axis=-1)  # [b, n_h, t_t, t_s]
-    output = p_attn.matmul(value)
-    if self.window_size is not None:
-      relative_weights = self._absolute_position_to_relative_position(p_attn)
-      value_relative_embeddings = self._get_relative_embeddings(self.emb_rel_v, t_s)
-      output = output + self._matmul_with_relative_values(relative_weights, value_relative_embeddings)
-    output = output.transpose(2, 3).contiguous().reshape(b, d, t_t)  # [b, n_h, t_t, d_k] -> [b, d, t_t]
-    return output, p_attn
-  def _matmul_with_relative_values(self, x, y): return x.matmul(y.unsqueeze(0))                 # x: [b, h, l, m], y: [h or 1, m, d], ret: [b, h, l, d]
-  def _matmul_with_relative_keys(self, x, y): return x.matmul(y.unsqueeze(0).transpose(-2, -1)) # x: [b, h, l, d], y: [h or 1, m, d], re, : [b, h, l, m]
-  def _get_relative_embeddings(self, relative_embeddings, length):
-    pad_length, slice_start_position = max(length - (self.window_size + 1), 0), max((self.window_size + 1) - length, 0)
-    padded_relative_embeddings = relative_embeddings if pad_length <= 0\
-      else relative_embeddings.pad(convert_pad_shape([[0, 0], [pad_length, pad_length], [0, 0]]))
-    return padded_relative_embeddings[:, slice_start_position:(slice_start_position + 2 * length - 1)] #used_relative_embeddings
-  def _relative_position_to_absolute_position(self, x: Tensor): # x: [b, h, l, 2*l-1] -> [b, h, l, l]
-    batch, heads, length, _ = x.shape
-    x = x.pad(convert_pad_shape([[0,0],[0,0],[0,0],[0,1]]))
-    x_flat = x.reshape([batch, heads, length * 2 * length]).pad(convert_pad_shape([[0,0],[0,0],[0,length-1]]))
-    return x_flat.reshape([batch, heads, length+1, 2*length-1])[:, :, :length, length-1:]
-  def _absolute_position_to_relative_position(self, x: Tensor): # x: [b, h, l, l] -> [b, h, l, 2*l-1]
-    batch, heads, length, _ = x.shape
-    x = x.pad(convert_pad_shape([[0, 0], [0, 0], [0, 0], [0, length-1]]))
-    x_flat = x.reshape([batch, heads, length**2 + length*(length -1)]).pad(convert_pad_shape([[0, 0], [0, 0], [length, 0]]))
-    return x_flat.reshape([batch, heads, length, 2*length])[:,:,:,1:]
-
-class FFN:
-  def __init__(self, in_channels, out_channels, filter_channels, kernel_size, p_dropout=0., activation=None, causal=False):
-    self.in_channels, self.out_channels, self.filter_channels, self.kernel_size, self.p_dropout, self.activation, self.causal = in_channels, out_channels, filter_channels, kernel_size, p_dropout, activation, causal
-    self.padding = self._causal_padding if causal else self._same_padding
-    self.conv_1, self.conv_2 = nn.Conv1d(in_channels, filter_channels, kernel_size), nn.Conv1d(filter_channels, out_channels, kernel_size)
-  def forward(self, x, x_mask):
-    x = self.conv_1(self.padding(x * x_mask))
-    x = x * (1.702 * x).sigmoid() if self.activation == "gelu" else x.relu()
-    return self.conv_2(self.padding(x.dropout(self.p_dropout) * x_mask)) * x_mask
-  def _causal_padding(self, x):return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [self.kernel_size - 1, 0]]))
-  def _same_padding(self, x): return x if self.kernel_size == 1 else x.pad(convert_pad_shape([[0, 0], [0, 0], [(self.kernel_size - 1) // 2, self.kernel_size // 2]]))
-
-class Encoder:
-  def __init__(self, hidden_channels, filter_channels, n_heads, n_layers, kernel_size=1, p_dropout=0., window_size=4, **kwargs):
-    self.hidden_channels, self.filter_channels, self.n_heads, self.n_layers, self.kernel_size, self.p_dropout, self.window_size = hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout, window_size
-    self.attn_layers, self.norm_layers_1, self.ffn_layers, self.norm_layers_2 = [], [], [], []
-    for _ in range(n_layers):
-      self.attn_layers.append(MultiHeadAttention(hidden_channels, hidden_channels, n_heads, p_dropout=p_dropout, window_size=window_size))
-      self.norm_layers_1.append(LayerNorm(hidden_channels))
-      self.ffn_layers.append(FFN(hidden_channels, hidden_channels, filter_channels, kernel_size, p_dropout=p_dropout))
-      self.norm_layers_2.append(LayerNorm(hidden_channels))
-  def forward(self, x, x_mask):
-    attn_mask, x = x_mask.unsqueeze(2) * x_mask.unsqueeze(-1), x * x_mask
-    for i in range(self.n_layers):
-      y = self.attn_layers[i].forward(x, x, attn_mask).dropout(self.p_dropout)
-      x = self.norm_layers_1[i].forward(x + y)
-      y = self.ffn_layers[i].forward(x, x_mask).dropout(self.p_dropout)
-      x = self.norm_layers_2[i].forward(x + y)
-    return x * x_mask
-
-DEFAULT_MIN_BIN_WIDTH, DEFAULT_MIN_BIN_HEIGHT, DEFAULT_MIN_DERIVATIVE = 1e-3, 1e-3, 1e-3
-def piecewise_rational_quadratic_transform(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails=None, tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
-  if tails is None: spline_fn, spline_kwargs = rational_quadratic_spline, {}
-  else: spline_fn, spline_kwargs = unconstrained_rational_quadratic_spline, {'tails': tails, 'tail_bound': tail_bound}
-  return spline_fn(inputs=inputs, un_normalized_widths=un_normalized_widths, un_normalized_heights=un_normalized_heights, un_normalized_derivatives=un_normalized_derivatives, inverse=inverse, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative, **spline_kwargs)
-def unconstrained_rational_quadratic_spline(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails='linear', tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
-  if not tails == 'linear': raise RuntimeError('{} tails are not implemented.'.format(tails))
-  constant = np.log(np.exp(1 - min_derivative) - 1).item()
-  un_normalized_derivatives = cat_lr(un_normalized_derivatives, constant, constant)
-  output, log_abs_det = rational_quadratic_spline(inputs=inputs.squeeze(dim=0).squeeze(dim=0), unnormalized_widths=un_normalized_widths.squeeze(dim=0).squeeze(dim=0), unnormalized_heights=un_normalized_heights.squeeze(dim=0).squeeze(dim=0), unnormalized_derivatives=un_normalized_derivatives.squeeze(dim=0).squeeze(dim=0), inverse=inverse, left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative)
-  return output.unsqueeze(dim=0).unsqueeze(dim=0), log_abs_det.unsqueeze(dim=0).unsqueeze(dim=0)
-def rational_quadratic_spline(inputs: Tensor, unnormalized_widths: Tensor, unnormalized_heights: Tensor, unnormalized_derivatives: Tensor, inverse=False, left=0., right=1., bottom=0., top=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
-  num_bins = unnormalized_widths.shape[-1]
-  if min_bin_width * num_bins > 1.0: raise ValueError('Minimal bin width too large for the number of bins')
-  if min_bin_height * num_bins > 1.0: raise ValueError('Minimal bin height too large for the number of bins')
-  widths = min_bin_width + (1 - min_bin_width * num_bins) * unnormalized_widths.softmax(axis=-1)
-  cum_widths = cat_lr(((right - left) * widths[..., :-1].cumsum(axis=1) + left), left, right + 1e-6 if not inverse else right)
-  widths = cum_widths[..., 1:] - cum_widths[..., :-1]
-  derivatives = min_derivative + (unnormalized_derivatives.exp()+1).log()
-  heights = min_bin_height + (1 - min_bin_height * num_bins) * unnormalized_heights.softmax(axis=-1)
-  cum_heights = cat_lr(((top - bottom) * heights[..., :-1].cumsum(axis=1) + bottom), bottom, top + 1e-6 if inverse else top)
-  heights = cum_heights[..., 1:] - cum_heights[..., :-1]
-  bin_idx = ((inputs[..., None] >= (cum_heights if inverse else cum_widths)).sum(axis=-1) - 1)[..., None]
-  input_cum_widths = gather(cum_widths, bin_idx, axis=-1)[..., 0]
-  input_bin_widths = gather(widths, bin_idx, axis=-1)[..., 0]
-  input_cum_heights = gather(cum_heights, bin_idx, axis=-1)[..., 0]
-  input_delta = gather(heights / widths, bin_idx, axis=-1)[..., 0]
-  input_derivatives = gather(derivatives, bin_idx, axis=-1)[..., 0]
-  input_derivatives_plus_one = gather(derivatives[..., 1:], bin_idx, axis=-1)[..., 0]
-  input_heights = gather(heights, bin_idx, axis=-1)[..., 0]
-  if inverse:
-    a = ((inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta) + input_heights * (input_delta - input_derivatives))
-    b = (input_heights * input_derivatives - (inputs - input_cum_heights) * (input_derivatives + input_derivatives_plus_one - 2 * input_delta))
-    c = - input_delta * (inputs - input_cum_heights)
-    discriminant = b.square() - 4 * a * c
-    # assert (discriminant.numpy() >= 0).all()
-    root = (2 * c) / (-b - discriminant.sqrt())
-    theta_one_minus_theta = root * (1 - root)
-    denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
-    derivative_numerator = input_delta.square() * (input_derivatives_plus_one * root.square() + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - root).square())
-    return root * input_bin_widths + input_cum_widths, -(derivative_numerator.log() - 2 * denominator.log())
-  theta = (inputs - input_cum_widths) / input_bin_widths
-  theta_one_minus_theta = theta * (1 - theta)
-  numerator = input_heights * (input_delta * theta.pow(2) + input_derivatives * theta_one_minus_theta)
-  denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta) * theta_one_minus_theta)
-  derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2) + 2 * input_delta * theta_one_minus_theta + input_derivatives * (1 - theta).pow(2))
-  return input_cum_heights + numerator / denominator, derivative_numerator.log() - 2 * denominator.log()
-
-def sequence_mask(length: Tensor, max_length): return Tensor.arange(max_length, dtype=length.dtype, device=length.device).unsqueeze(0) < length.unsqueeze(1)
-def generate_path(duration: Tensor, mask: Tensor):  # duration: [b, 1, t_x], mask: [b, 1, t_y, t_x]
-  b, _, t_y, t_x = mask.shape
-  path = sequence_mask(duration.cumsum(axis=2).reshape(b * t_x), t_y).cast(mask.dtype).reshape(b, t_x, t_y)
-  path = path - path.pad(convert_pad_shape([[0, 0], [1, 0], [0, 0]]))[:, :-1]
-  return path.unsqueeze(1).transpose(2, 3) * mask
-def fused_add_tanh_sigmoid_multiply(input_a: Tensor, input_b: Tensor, n_channels: int):
-  n_channels_int, in_act = n_channels, input_a + input_b
-  t_act, s_act = in_act[:, :n_channels_int, :].tanh(), in_act[:, n_channels_int:, :].sigmoid()
-  return t_act * s_act
-
-def cat_lr(t, left, right): return Tensor.full(get_shape(t), left).cat(t, dim=-1).cat(Tensor.full(get_shape(t), right), dim=-1)
-def get_shape(tensor):
-  (shape := list(tensor.shape))[-1] = 1
-  return tuple(shape)
-def convert_pad_shape(pad_shape): return tuple(tuple(x) for x in pad_shape)
-def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2)
-
-def gather(x, indices, axis):
-  indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(0, axis)
-  permute_args = list(range(x.ndim))
-  permute_args[0], permute_args[axis] = permute_args[axis], permute_args[0]
-  permute_args.append(permute_args.pop(0))
-  x = x.permute(*permute_args)
-  reshape_arg = [1] * x.ndim + [x.shape[-1]]
-  return ((indices.unsqueeze(indices.ndim).expand(*indices.shape, x.shape[-1]) ==
-           Tensor.arange(x.shape[-1]).reshape(*reshape_arg).expand(*indices.shape, x.shape[-1])) * x).sum(indices.ndim).transpose(0, axis)
-
-def norm_except_dim(v, dim):
-  if dim == -1: return np.linalg.norm(v)
-  if dim == 0:
-    (output_shape := [1] * v.ndim)[0] = v.shape[0]
-    return np.linalg.norm(v.reshape(v.shape[0], -1), axis=1).reshape(output_shape)
-  if dim == v.ndim - 1:
-    (output_shape := [1] * v.ndim)[-1] = v.shape[-1]
-    return np.linalg.norm(v.reshape(-1, v.shape[-1]), axis=0).reshape(output_shape)
-  transposed_v = np.transpose(v, (dim,) + tuple(i for i in range(v.ndim) if i != dim))
-  return np.transpose(norm_except_dim(transposed_v, 0), (dim,) + tuple(i for i in range(v.ndim) if i != dim))
-def weight_norm(v: Tensor, g: Tensor, dim):
-  v, g = v.numpy(), g.numpy()
-  return Tensor(v * (g / norm_except_dim(v, dim)))
-
-# HPARAMS LOADING
-def get_hparams_from_file(path):
-  with open(path, "r") as f:
-    data = f.read()
-  return HParams(**json.loads(data))
-class HParams:
-  def __init__(self, **kwargs):
-    for k, v in kwargs.items(): self[k] = v if type(v) != dict else HParams(**v)
-  def keys(self): return self.__dict__.keys()
-  def items(self): return self.__dict__.items()
-  def values(self): return self.__dict__.values()
-  def __len__(self): return len(self.__dict__)
-  def __getitem__(self, key): return getattr(self, key)
-  def __setitem__(self, key, value): return setattr(self, key, value)
-  def __contains__(self, key): return key in self.__dict__
-  def __repr__(self): return self.__dict__.__repr__()
-
-# MODEL LOADING
-def load_model(symbols, hps, model) -> Synthesizer:
-  net_g = Synthesizer(len(symbols), hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, n_speakers = hps.data.n_speakers, **hps.model)
-  _ = load_checkpoint(fetch(model[1]), net_g, None)
-  return net_g
-def load_checkpoint(checkpoint_path, model: Synthesizer, optimizer=None, skip_list=[]):
-  assert Path(checkpoint_path).is_file()
-  start_time = time.time()
-  checkpoint_dict = torch_load(checkpoint_path)
-  iteration, learning_rate = checkpoint_dict['iteration'], checkpoint_dict['learning_rate']
-  if optimizer: optimizer.load_state_dict(checkpoint_dict['optimizer'])
-  saved_state_dict = checkpoint_dict['model']
-  weight_g, weight_v, parent = None, None, None
-  for key, v in saved_state_dict.items():
-    if any(layer in key for layer in skip_list): continue
-    try:
-      obj, skip = model, False
-      for k in key.split('.'):
-        if k.isnumeric(): obj = obj[int(k)]
-        elif isinstance(obj, dict): obj = obj[k]
-        else:
-          if isinstance(obj, (LayerNorm, nn.LayerNorm)) and k in ["gamma", "beta"]:
-            k = "weight" if k == "gamma" else "bias"
-          elif k in ["weight_g", "weight_v"]:
-            parent, skip = obj, True
-            if k == "weight_g": weight_g = v
-            else: weight_v = v
-          if not skip: obj = getattr(obj, k)
-      if weight_g is not None and weight_v is not None:
-        setattr(obj, "weight_g", weight_g.numpy())
-        setattr(obj, "weight_v", weight_v.numpy())
-        obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
-        weight_g, weight_v, parent, skip = None, None, None, False
-      if not skip and obj.shape == v.shape: obj.assign(v.to(obj.device))
-      elif not skip: logging.error(f"MISMATCH SHAPE IN {key}, {obj.shape} {v.shape}")
-    except Exception as e: raise e
-  logging.info(f"Loaded checkpoint '{checkpoint_path}' (iteration {iteration}) in {time.time() - start_time:.4f}s")
-  return model, optimizer, learning_rate, iteration
-
-# Used for cleaning input text and mapping to symbols
-class TextMapper: # Based on https://github.com/keithito/tacotron
-  def __init__(self, symbols, apply_cleaners=True):
-    self.apply_cleaners, self.symbols, self._inflect = apply_cleaners, symbols, None
-    self._symbol_to_id, _id_to_symbol = {s: i for i, s in enumerate(symbols)}, {i: s for i, s in enumerate(symbols)}
-    self._whitespace_re, self._abbreviations = re.compile(r'\s+'), [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [('mrs', 'misess'), ('mr', 'mister'), ('dr', 'doctor'), ('st', 'saint'), ('co', 'company'), ('jr', 'junior'), ('maj', 'major'), ('gen', 'general'), ('drs', 'doctors'), ('rev', 'reverend'), ('lt', 'lieutenant'), ('hon', 'honorable'), ('sgt', 'sergeant'), ('capt', 'captain'), ('esq', 'esquire'), ('ltd', 'limited'), ('col', 'colonel'), ('ft', 'fort'), ]]
-    self.phonemizer = EspeakBackend(
-        language="en-us", punctuation_marks=Punctuation.default_marks(), preserve_punctuation=True, with_stress=True,
-    )
-  def text_to_sequence(self, text, cleaner_names):
-    if self.apply_cleaners:
-      for name in cleaner_names:
-        cleaner = getattr(self, name)
-        if not cleaner: raise ModuleNotFoundError('Unknown cleaner: %s' % name)
-        text = cleaner(text)
-    else: text = text.strip()
-    return [self._symbol_to_id[symbol] for symbol in text]
-  def get_text(self, text, add_blank=False, cleaners=('english_cleaners2',)):
-    text_norm = self.text_to_sequence(text, cleaners)
-    return Tensor(self.intersperse(text_norm, 0) if add_blank else text_norm, dtype=dtypes.int64)
-  def intersperse(self, lst, item):
-    (result := [item] * (len(lst) * 2 + 1))[1::2] = lst
-    return result
-  def phonemize(self, text, strip=True): return _phonemize(self.phonemizer, text, default_separator, strip, 1, False, False)
-  def filter_oov(self, text): return "".join(list(filter(lambda x: x in self._symbol_to_id, text)))
-  def base_english_cleaners(self, text): return self.collapse_whitespace(self.phonemize(self.expand_abbreviations(unidecode(text.lower()))))
-  def english_cleaners2(self, text): return self.base_english_cleaners(text)
-  def transliteration_cleaners(self, text): return self.collapse_whitespace(unidecode(text.lower()))
-  def cjke_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text).replace('ɑ', 'a').replace('ɔ', 'o').replace('ɛ', 'e').replace('ɪ', 'i').replace('ʊ', 'u')))
-  def cjke_cleaners2(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_ipa2(text)))
-  def cjks_cleaners(self, text): return re.sub(r'([^\.,!\?\-…~])$', r'\1.', re.sub(r'\s+$', '', self.english_to_lazy_ipa(text)))
-  def english_to_ipa2(self, text):
-    _ipa_to_ipa2 = [(re.compile('%s' % x[0]), x[1]) for x in [ ('r', 'ɹ'), ('ʤ', 'dʒ'), ('ʧ', 'tʃ')]]
-    return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _ipa_to_ipa2, self.mark_dark_l(self.english_to_ipa(text))).replace('...', '…')
-  def mark_dark_l(self, text): return re.sub(r'l([^aeiouæɑɔəɛɪʊ ]*(?: |$))', lambda x: 'ɫ' + x.group(1), text)
-  def english_to_ipa(self, text):
-    import eng_to_ipa as ipa
-    return self.collapse_whitespace(ipa.convert(self.normalize_numbers(self.expand_abbreviations(unidecode(text).lower()))))
-  def english_to_lazy_ipa(self, text):
-    _lazy_ipa = [(re.compile('%s' % x[0]), x[1]) for x in [('r', 'ɹ'), ('æ', 'e'), ('ɑ', 'a'), ('ɔ', 'o'), ('ð', 'z'), ('θ', 's'), ('ɛ', 'e'), ('ɪ', 'i'), ('ʊ', 'u'), ('ʒ', 'ʥ'), ('ʤ', 'ʥ'), ('ˈ', '↓')]]
-    return reduce(lambda t, rx: re.sub(rx[0], rx[1], t), _lazy_ipa, self.english_to_ipa(text))
-  def expand_abbreviations(self, text): return reduce(lambda t, abbr: re.sub(abbr[0], abbr[1], t), self._abbreviations, text)
-  def collapse_whitespace(self, text): return re.sub(self._whitespace_re, ' ', text)
-  def normalize_numbers(self, text):
-    import inflect
-    self._inflect = inflect.engine()
-    text = re.sub(re.compile(r'([0-9][0-9\,]+[0-9])'), self._remove_commas, text)
-    text = re.sub(re.compile(r'£([0-9\,]*[0-9]+)'), r'\1 pounds', text)
-    text = re.sub(re.compile(r'\$([0-9\.\,]*[0-9]+)'), self._expand_dollars, text)
-    text = re.sub(re.compile(r'([0-9]+\.[0-9]+)'), self._expand_decimal_point, text)
-    text = re.sub(re.compile(r'[0-9]+(st|nd|rd|th)'), self._expand_ordinal, text)
-    text = re.sub(re.compile(r'[0-9]+'), self._expand_number, text)
-    return text
-  def _remove_commas(self, m): return m.group(1).replace(',', '') # george won't like this
-  def _expand_dollars(self, m):
-    match = m.group(1)
-    parts = match.split('.')
-    if len(parts) > 2: return match + ' dollars'  # Unexpected format
-    dollars, cents = int(parts[0]) if parts[0] else 0, int(parts[1]) if len(parts) > 1 and parts[1] else 0
-    if dollars and cents: return '%s %s, %s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars', cents, 'cent' if cents == 1 else 'cents')
-    if dollars: return '%s %s' % (dollars, 'dollar' if dollars == 1 else 'dollars')
-    if cents: return '%s %s' % (cents, 'cent' if cents == 1 else 'cents')
-    return 'zero dollars'
-  def _expand_decimal_point(self, m): return m.group(1).replace('.', ' point ')
-  def _expand_ordinal(self, m): return self._inflect.number_to_words(m.group(0))
-  def _expand_number(self, _inflect, m):
-    num = int(m.group(0))
-    if 1000 < num < 3000:
-      if num == 2000: return 'two thousand'
-      if 2000 < num < 2010: return 'two thousand ' + self._inflect.number_to_words(num % 100)
-      if num % 100 == 0: return self._inflect.number_to_words(num // 100) + ' hundred'
-      return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
-    return self._inflect.number_to_words(num, andword='')
-
-#########################################################################################
-# PAPER: https://arxiv.org/abs/2106.06103
-# CODE: https://github.com/jaywalnut310/vits/tree/main
-#########################################################################################
-# INSTALLATION: this is based on default config, dependencies are for preprocessing.
-# vctk, ljs                      | pip3 install unidecode phonemizer          | phonemizer requires [eSpeak](https://espeak.sourceforge.net) backend to be installed on your system
-# mmts-tts                       | pip3 install unidecode                     |
-# uma_trilingual, cjks, voistock | pip3 install unidecode inflect eng_to_ipa  |
-#########################################################################################
-# Some good speakers to try out, there may be much better ones, I only tried out a few:
-# male vctk 1  | --model_to_use vctk --speaker_id 2
-# male vctk 2  | --model_to_use vctk --speaker_id 6
-# anime lady 1 | --model_to_use uma_trilingual --speaker_id 36
-# anime lady 2 | --model_to_use uma_trilingual --speaker_id 121
-#########################################################################################
-VITS_PATH = Path(__file__).parents[1] / "weights/VITS/"
-MODELS = { # config_url, weights_url
-  "ljs": ("https://raw.githubusercontent.com/jaywalnut310/vits/main/configs/ljs_base.json", "https://drive.google.com/uc?export=download&id=1q86w74Ygw2hNzYP9cWkeClGT5X25PvBT&confirm=t"),
-  "vctk": ("https://huggingface.co/csukuangfj/vits-vctk/resolve/main/vctk_base.json", "https://huggingface.co/csukuangfj/vits-vctk/resolve/main/pretrained_vctk.pth"),
-  "mmts-tts": ("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/config.json", "https://huggingface.co/facebook/mms-tts/resolve/main/full_models/eng/G_100000.pth"),
-  "uma_trilingual": ("https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/raw/main/configs/uma_trilingual.json", "https://huggingface.co/spaces/Plachta/VITS-Umamusume-voice-synthesizer/resolve/main/pretrained_models/G_trilingual.pth"),
-  "cjks": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/14/model.pth"),
-  "voistock": ("https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/config.json", "https://huggingface.co/spaces/skytnt/moe-tts/resolve/main/saved_model/15/model.pth"),
-}
-Y_LENGTH_ESTIMATE_SCALARS = {"ljs": 2.8, "vctk": 1.74, "mmts-tts": 1.9, "uma_trilingual": 2.3, "cjks": 3.3, "voistock": 3.1}
-if __name__ == '__main__':
-  logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--model_to_use", default="vctk", help="Specify the model to use. Default is 'vctk'.")
-  parser.add_argument("--speaker_id", type=int, default=6, help="Specify the speaker ID. Default is 6.")
-  parser.add_argument("--out_path", default=None, help="Specify the full output path. Overrides the --out_dir and --name parameter.")
-  parser.add_argument("--out_dir", default=str(Path(__file__).parents[1] / "temp"), help="Specify the output path.")
-  parser.add_argument("--base_name", default="test", help="Specify the base of the output file name. Default is 'test'.")
-  parser.add_argument("--text_to_synthesize", default="""Hello person. If the code you are contributing isn't some of the highest quality code you've written in your life, either put in the effort to make it great, or don't bother.""", help="Specify the text to synthesize. Default is a greeting message.")
-  parser.add_argument("--noise_scale", type=float, default=0.667, help="Specify the noise scale. Default is 0.667.")
-  parser.add_argument("--noise_scale_w", type=float, default=0.8, help="Specify the noise scale w. Default is 0.8.")
-  parser.add_argument("--length_scale", type=float, default=1, help="Specify the length scale. Default is 1.")
-  parser.add_argument("--seed", type=int, default=1337, help="Specify the seed (set to None if no seed). Default is 1337.")
-  parser.add_argument("--num_channels", type=int, default=1, help="Specify the number of audio output channels. Default is 1.")
-  parser.add_argument("--sample_width", type=int, default=2, help="Specify the number of bytes per sample, adjust if necessary. Default is 2.")
-  parser.add_argument("--emotion_path", type=str, default=None, help="Specify the path to emotion reference.")
-  parser.add_argument("--estimate_max_y_length", type=str, default=False, help="If true, overestimate the output length and then trim it to the correct length, to prevent premature realization, much more performant for larger inputs, for smaller inputs not so much. Default is False.")
-  args = parser.parse_args()
-
-  model_config = MODELS[args.model_to_use]
-
-  # Load the hyperparameters from the config file.
-  hps = get_hparams_from_file(fetch(model_config[0]))
-
-  # If model has multiple speakers, validate speaker id and retrieve name if available.
-  model_has_multiple_speakers = hps.data.n_speakers > 0
-  if model_has_multiple_speakers:
-    logging.info(f"Model has {hps.data.n_speakers} speakers")
-    if args.speaker_id >= hps.data.n_speakers: raise ValueError(f"Speaker ID {args.speaker_id} is invalid for this model.")
-    speaker_name = "?"
-    if hps.__contains__("speakers"): # maps speaker ids to names
-      speakers = hps.speakers
-      if isinstance(speakers, List): speakers = {speaker: i for i, speaker in enumerate(speakers)}
-      speaker_name = next((key for key, value in speakers.items() if value == args.speaker_id), None)
-    logging.info(f"You selected speaker {args.speaker_id} (name: {speaker_name})")
-
-  # Load emotions if any. TODO: find an english model with emotions, this is untested atm.
-  emotion_embedding = None
-  if args.emotion_path is not None:
-    if args.emotion_path.endswith(".npy"): emotion_embedding = Tensor(np.load(args.emotion_path), dtype=dtypes.int64).unsqueeze(0)
-    else: raise ValueError("Emotion path must be a .npy file.")
-
-  # Load symbols, instantiate TextMapper and clean the text.
-  if hps.__contains__("symbols"): symbols = hps.symbols
-  elif args.model_to_use == "mmts-tts": symbols = [x.replace("\n", "") for x in fetch("https://huggingface.co/facebook/mms-tts/raw/main/full_models/eng/vocab.txt").open(encoding="utf-8").readlines()]
-  else: symbols = ['_'] + list(';:,.!?¡¿—…"«»“” ') + list('ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz') + list("ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ")
-  text_mapper = TextMapper(apply_cleaners=True, symbols=symbols)
-
-  # Load the model.
-  if args.seed is not None:
-    Tensor.manual_seed(args.seed)
-    np.random.seed(args.seed)
-  net_g = load_model(text_mapper.symbols, hps, model_config)
-  logging.debug(f"Loaded model with hps: {hps}")
-
-  # Convert the input text to a tensor.
-  text_to_synthesize = args.text_to_synthesize
-  if args.model_to_use == "mmts-tts": text_to_synthesize = text_mapper.filter_oov(text_to_synthesize.lower())
-  stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
-  logging.debug(f"Converted input text to tensor \"{text_to_synthesize}\" -> Tensor({stn_tst.shape}): {stn_tst.numpy()}")
-  x_tst, x_tst_lengths = stn_tst.unsqueeze(0), Tensor([stn_tst.shape[0]], dtype=dtypes.int64)
-  sid = Tensor([args.speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
-
-  # Perform inference.
-  start_time = time.time()
-  audio_tensor = net_g.infer(x_tst, x_tst_lengths, sid, args.noise_scale, args.length_scale, args.noise_scale_w, emotion_embedding=emotion_embedding,
-                             max_y_length_estimate_scale=Y_LENGTH_ESTIMATE_SCALARS[args.model_to_use] if args.estimate_max_y_length else None)[0, 0].realize()
-  logging.info(f"Inference took {(time.time() - start_time):.2f}s")
-
-  # Save the audio output.
-  audio_data = (np.clip(audio_tensor.numpy(), -1.0, 1.0) * 32767).astype(np.int16)
-  out_path = Path(args.out_path or Path(args.out_dir)/f"{args.model_to_use}{f'_sid_{args.speaker_id}' if model_has_multiple_speakers else ''}_{args.base_name}.wav")
-  out_path.parent.mkdir(parents=True, exist_ok=True)
-  with wave.open(str(out_path), 'wb') as wav_file:
-    wav_file.setnchannels(args.num_channels)
-    wav_file.setsampwidth(args.sample_width)
-    wav_file.setframerate(hps.data.sampling_rate)
-    wav_file.setnframes(len(audio_data))
-    wav_file.writeframes(audio_data.tobytes())
-  logging.info(f"Saved audio output to {out_path}")

extra/amdpci/am_smi.py

@@ -26,11 +26,13 @@ def color_temp(temp):
 def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")
 
 def draw_bar(percentage, width=40, fill='|', empty=' ', opt_text='', color='cyan'):
+  percentage = 0.0 if percentage != percentage else percentage  # NaN guard
+  percentage = max(0.0, min(1.0, float(percentage)))
   filled_width = int(width * percentage)
   if not opt_text: opt_text = f'{percentage*100:.1f}%'
 
   bar = fill * filled_width + empty * (width - filled_width)
-  bar = (bar[:-len(opt_text)] + opt_text) if opt_text else bar
+  if opt_text and len(opt_text) <= len(bar): bar = (bar[:-len(opt_text)] + opt_text)
   bar = colored(bar[:filled_width], color) + bar[filled_width:]
   return f'[{bar}]'
 
@@ -88,6 +90,7 @@ class SMICtx:
     self.opened_pci_resources = {}
     self.prev_lines_cnt = 0
     self.prev_terminal_width = 0
+    self.prev_terminal_height = 0
 
     remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
     lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
@@ -95,6 +98,20 @@ class SMICtx:
     for k,v in self.lspci.items():
       for part in remove_parts: self.lspci[k] = self.lspci[k].replace(part, "").strip().rstrip()
 
+  def _smuq10_round(self, v:int) -> int:
+    v = int(v)
+    return (v + 512) >> 10  # SMUQ10_ROUND
+
+  def _fmt_kb(self, kb:int) -> str:
+    kb = int(kb)
+    if kb < 1024: return f"{kb}KB"
+    mb = kb / 1024.0
+    if mb < 1024: return f"{mb:.1f}MB"
+    gb = mb / 1024.0
+    if gb < 1024: return f"{gb:.2f}GB"
+    tb = gb / 1024.0
+    return f"{tb:.2f}TB"
+
   def _open_am_device(self, pcibus):
     if pcibus not in self.opened_pci_resources:
       bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
@@ -116,6 +133,7 @@ class SMICtx:
   def rescan_devs(self):
     pattern = os.path.join('/tmp', 'am_*.lock')
     for d in [f[8:-5] for f in glob.glob(pattern)]:
+      if d.startswith("usb"): continue
       if d not in self.opened_pcidevs:
         self._open_am_device(d)
 
@@ -131,21 +149,53 @@ class SMICtx:
         os.system('clear')
         if DEBUG >= 2: print(f"Removed AM device {d.pcibus}")
 
-  def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}
+  def collect(self):
+    tables = {}
+    for dev in self.devs:
+      match dev.ip_ver[am.MP1_HWIP]:
+        case (13,0,6): table_t = dev.smu.smu_mod.MetricsTableX_t
+        case (13,0,12): table_t = dev.smu.smu_mod.MetricsTableV2_t
+        case _: table_t = dev.smu.smu_mod.SmuMetricsExternal_t
+      tables[dev] = dev.smu.read_table(table_t, dev.smu.smu_mod.SMU_TABLE_SMU_METRICS) if dev.pci_state == "D0" else None
+    return tables
 
-  def get_gfx_activity(self, dev, metrics): return metrics.SmuMetrics.AverageGfxActivity
-  def get_mem_activity(self, dev, metrics): return metrics.SmuMetrics.AverageUclkActivity
+  def _pick_nonzero_avg(self, vals) -> int:
+    xs = [x for x in vals if x > 0]
+    return int(sum(xs) / len(xs)) if xs else 0
+
+  def get_gfx_activity(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.SocketGfxBusy)))
+      case _: return metrics.SmuMetrics.AverageGfxActivity
+
+  def get_mem_activity(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return max(0, min(100, self._smuq10_round(metrics.DramBandwidthUtilization)))
+      case _: return metrics.SmuMetrics.AverageUclkActivity
 
   def get_temps(self, dev, metrics, compact=False):
-    temps_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_TEMP_e__enumvalues.items()
-                  if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
-    if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
-    return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6):
+        temps = {
+          "Hotspot": self._smuq10_round(metrics.MaxSocketTemperature),
+          "HBM": self._smuq10_round(metrics.MaxHbmTemperature),
+          "VR": self._smuq10_round(metrics.MaxVrTemperature),
+        }
+        if compact: return {k: temps[k] for k in ("Hotspot", "HBM") if temps.get(k, 0) != 0}
+        return {k: v for k, v in temps.items() if v != 0}
+      case _:
+        temps_keys = [(k, name) for k, name in dev.smu.smu_mod.TEMP_e.items()
+                      if k < dev.smu.smu_mod.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
+        if compact: temps_keys = [(k, name) for k, name in temps_keys if k in (dev.smu.smu_mod.TEMP_HOTSPOT, dev.smu.smu_mod.TEMP_MEM)]
+        return {name: metrics.SmuMetrics.AvgTemperature[k] for k, name in temps_keys}
 
   def get_voltage(self, dev, metrics, compact=False):
-    voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.c__EA_SVI_PLANE_e__enumvalues.items()
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return {}
+      case _:
+        voltage_keys = [(k, name) for k, name in dev.smu.smu_mod.SVI_PLANE_e.items()
                         if k < dev.smu.smu_mod.SVI_PLANE_COUNT and metrics.SmuMetrics.AvgVoltage[k] != 0]
-    return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
+        return {name: metrics.SmuMetrics.AvgVoltage[k] for k, name in voltage_keys}
 
   def get_busy_threshold(self, dev):
     match dev.ip_ver[am.MP1_HWIP]:
@@ -153,22 +203,40 @@ class SMICtx:
       case _: return 15
 
   def get_gfx_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-          metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
+    if metrics is None: return 0
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return self._smuq10_round(metrics.GfxclkFrequency[0])
+      case _:
+        return metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if self.get_gfx_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageGfxclkFrequencyPreDs
 
   def get_mem_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-           metrics.SmuMetrics.AverageMemclkFrequencyPreDs
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return self._smuq10_round(metrics.UclkFrequency)
+      case _:
+        return metrics.SmuMetrics.AverageMemclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageMemclkFrequencyPreDs
 
   def get_fckl_freq(self, dev, metrics):
-    return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
-           metrics.SmuMetrics.AverageFclkFrequencyPreDs
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return self._smuq10_round(metrics.FclkFrequency)
+      case _:
+        return metrics.SmuMetrics.AverageFclkFrequencyPostDs if self.get_mem_activity(dev, metrics) <= self.get_busy_threshold(dev) else \
+               metrics.SmuMetrics.AverageFclkFrequencyPreDs
 
-  def get_fan_rpm_pwm(self, dev, metrics): return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
+  def get_fan_rpm_pwm(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return None, None
+      case _: return metrics.SmuMetrics.AvgFanRpm, metrics.SmuMetrics.AvgFanPwm
 
-  def get_power(self, dev, metrics): return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
+  def get_power(self, dev, metrics):
+    match dev.ip_ver[am.MP1_HWIP]:
+      case (13,0,6): return self._smuq10_round(metrics.SocketPower), self._smuq10_round(metrics.MaxSocketPowerLimit)
+      case _: return metrics.SmuMetrics.AverageSocketPower, metrics.SmuMetrics.dGPU_W_MAX
 
   def get_mem_usage(self, dev):
+    return 0
+
     usage = 0
     pt_stack = [dev.mm.root_page_table]
     while len(pt_stack) > 0:
@@ -177,7 +245,7 @@ class SMICtx:
         entry = pt.entries[i]
 
         if (entry & am.AMDGPU_PTE_VALID) == 0: continue
-        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(entry):
+        if pt.lv!=am.AMDGPU_VM_PTB and not dev.gmc.is_pte_huge_page(pt.lv, entry):
           pt_stack.append(AMPageTableEntry(dev, entry & 0x0000FFFFFFFFF000, lv=pt.lv+1))
           continue
         if (entry & am.AMDGPU_PTE_SYSTEM) != 0: continue
@@ -219,23 +287,28 @@ class SMICtx:
       temps_table_compact = ["Temps (°C):" + '/'.join([f"{color_temp(val)} {name}" for name, val in temps_data_compact.items()])]
 
       fan_rpm, fan_pwm = self.get_fan_rpm_pwm(dev, metrics)
-      power_table = ["=== Power ==="] + [f"Fan Speed: {fan_rpm} RPM"] + [f"Fan Power: {fan_pwm}%"]
+      power_table = ["=== Power ==="]
+      power_table += ["Fan: N/A"] if fan_rpm is None or fan_pwm is None else [f"Fan Speed: {fan_rpm} RPM", f"Fan Power: {fan_pwm}%"]
 
       total_power, max_power = self.get_power(dev, metrics)
-      power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
-      power_line_compact = [f"Power:       " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
+      if max_power > 0:
+        power_line = [f"Power: " + draw_bar(total_power / max_power, 16, opt_text=f"{total_power}/{max_power}W")]
+        power_line_compact = [f"Power:       " + draw_bar(total_power / max_power, activity_line_width, opt_text=f"{total_power}/{max_power}W")]
+      else:
+        power_line = ["Power: N/A"]
+        power_line_compact = ["Power: N/A"]
 
       voltage_data = self.get_voltage(dev, metrics)
-      voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()]
+      voltage_table = None if not voltage_data else (["=== Voltages ==="] + [f"{name:<20}: {color_voltage(voltage)}" for name, voltage in voltage_data.items()])
 
       gfx_freq = self.get_gfx_freq(dev, metrics)
       mclk_freq = self.get_mem_freq(dev, metrics)
       fclk_freq = self.get_fckl_freq(dev, metrics)
-
       frequency_table = ["=== Frequencies ===", f"GFXCLK: {gfx_freq:>4} MHz", f"FCLK  : {fclk_freq:>4} MHz", f"MCLK  : {mclk_freq:>4} MHz"]
 
       if self.prev_terminal_width >= 231:
-        power_table += power_line + [""] + voltage_table
+        power_table += power_line
+        if voltage_table is not None: power_table += [""] + voltage_table
         activity_line += [""]
       elif self.prev_terminal_width >= 171:
         power_table += power_line + [""] + frequency_table
@@ -307,4 +380,5 @@ if __name__ == "__main__":
       smi_ctx.draw(args.list)
       if args.list: break
       time.sleep(1)
-  except KeyboardInterrupt: print("Exiting...")
+  except KeyboardInterrupt:
+    print("Exiting...")

extra/amdpci/hive_reset.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python3
+
+from tinygrad.helpers import Context
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
+from tinygrad.runtime.support.am.amdev import AMDev
+
+if __name__ == "__main__":
+  gpus = System.pci_scan_bus(0x1002, [(0xffff, [0x74a1, 0x75a0])])
+  pcidevs = [PCIDevice(f"reset:{gpu}", gpu, bars=[0, 2, 5]) for gpu in gpus]
+  amdevs = []
+  with Context(DEBUG=2):
+    for pcidev in pcidevs:
+      amdevs.append(AMDev(pcidev, reset_mode=True))
+    for amdev in amdevs: amdev.smu.mode1_reset()

extra/amdpci/proclogs.py

@@ -2,6 +2,10 @@ import re, ctypes, sys, importlib
 from tinygrad.helpers import getenv
 
 from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+
+class GFXFake:
+  def __init__(self): self.xccs = 8
+
 class AMDFake(AMDev):
   def __init__(self, pci_dev, dma_regions=None):
     self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
@@ -9,6 +13,8 @@ class AMDFake(AMDev):
     self._run_discovery()
     self._build_regs()
 
+    self.gfx = GFXFake()
+
 amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface

extra/assembly/assembly.py

@@ -1,189 +0,0 @@
-from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
-from tinygrad.uop.ops import BinaryOps, UnaryOps
-from tinygrad.dtype import DType, dtypes
-from tinygrad.helpers import DEBUG
-from tinygrad.uop.ops import Variable, NumNode, MulNode, DivNode, ModNode, LtNode, SumNode, AndNode
-import functools
-import math
-from collections import defaultdict
-
-_type_to_letter = {dtypes.float32: 'f', dtypes.bool: 'p', dtypes.int32: 'i', dtypes.int64: 'a', dtypes.uint32: 'u', dtypes.uint64: 'b', dtypes.float.vec(4): 'x', dtypes.uint8: 'uc', dtypes.float16: 'h',
-                   dtypes.int8: 'c', dtypes.uint16: 'us', dtypes.float64: 'd'}
-
-class Register(NamedTuple):
-  nm:str
-  dtype:DType
-  scalar:bool
-  off:Optional[int] = None
-  def __repr__(self): return self.nm if self.off is None else f"{self.nm}:{self.off}"
-  def subregs(self):
-    if self.dtype == dtypes.float.vec(4):
-      return [Register(self.nm, dtypes.float, False, off=off) for off in range(4)]
-    return []
-
-class AssemblyInstruction(NamedTuple):
-  op: Ops
-  out: Optional[Register]
-  vin: List[Union[Register, int, float]]
-  arg: Any = None
-
-# warp size of 32, s registers are shared across the warp, v are 32-wide vectors
-class AssemblyLanguage:
-  supports_load3: bool = False
-  sin_is_sin2pi: bool = False
-  no_div: bool = False
-  #TODO: these should be global vars
-  cnts:DefaultDict[Tuple[DType, bool], int] = defaultdict(int)
-  tor: Dict[Any, Register] = {}
-  ins: List[AssemblyInstruction] = []
-
-  def type_to_letter(self,x): return _type_to_letter[x[0]].upper() if x[1] else _type_to_letter[x[0]]
-  def newreg(self, tok, dtype=dtypes.float32, scalar=False) -> Register:
-    self.tor[tok] = ret = Register(f"%{self.type_to_letter((dtype, scalar))}{self.cnts[(dtype, scalar)]}", dtype, scalar)
-    if dtype == dtypes.float.vec(4):
-      for off in range(4):
-        self.tor[tok] = Register(ret.nm, dtypes.float, ret.scalar, off)
-    self.cnts[(dtype, scalar)] += 1
-    return ret
-
-  def render_numnode(self, b) -> Register:
-    key = ("num", b)
-    if key not in self.tor: self.ins.append(AssemblyInstruction(Ops.LOAD, self.newreg(key, scalar=True, dtype=dtypes.int32), [], b))
-    return self.tor[key]
-
-  def render_alu(self, op, a:Register, b:Union[Register, int, float], dtype=dtypes.int32) -> Register:
-    key = (op, a, b)
-    if key not in self.tor:
-      #if not isinstance(b, Register): b = render_numnode(b)
-      self.ins.append(AssemblyInstruction(Ops.ALU, self.newreg(key, dtype=dtype, scalar=a.scalar and (not isinstance(b, Register) or b.scalar)), [a, b], op))
-    return self.tor[key]
-
-  def render_cast(self, a:Register, new_dtype:DType) -> Register:
-    if a.dtype == new_dtype: return a
-    key = (a, new_dtype)
-    if key not in self.tor:
-      self.ins.append(AssemblyInstruction(Ops.CAST, self.newreg(key, dtype=new_dtype), [a]))
-    return self.tor[key]
-
-  render_ops: Any = { Variable: lambda self, ops, ctx: ctx.tor[self], NumNode: lambda self, ops, ctx: ctx.render_numnode(self.b),
-                 MulNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MUL, self.a.render(ops, ctx), self.b),
-                 DivNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.DIV, self.a.render(ops, ctx), self.b),
-                 ModNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.MOD, self.a.render(ops, ctx), self.b),
-                 LtNode: lambda self, ops, ctx: ctx.render_alu(BinaryOps.CMPLT, self.a.render(ops, ctx), self.b, dtype=dtypes.bool),
-    SumNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.ADD, a, b.render(ops,ctx)), self.nodes[1:], self.nodes[0].render(ops,ctx)),
-    AndNode: lambda self,ops,ctx: functools.reduce(lambda a,b: ctx.render_alu(BinaryOps.MUL, a, b.render(ops,ctx), dtype=dtypes.bool), self.nodes[1:], self.nodes[0].render(ops,ctx)) }
-
-  def addr_w_offset(self, args):
-    assert isinstance(args, MemOp)
-    idx = args.idx*args.memory_dtype.itemsize
-    off = 0  # TODO: should this be None?
-    if isinstance(idx, SumNode):
-      nums = [n.b for n in idx.nodes if isinstance(n, NumNode)]
-      if nums and nums[0] < 4096 and (idx-nums[0]).min >= 0:  # TODO: different for each GPU?
-        idx -= nums[0]
-        off = cast(int, nums[0])
-    reg = idx.render(self.render_ops, self)
-    if self.supports_load3:
-      if reg.scalar:
-        new_reg = self.newreg((reg.nm, 'vec'), dtype=reg.dtype)
-        self.ins.append(AssemblyInstruction(Ops.ALU, new_reg, [reg], UnaryOps.NOOP))
-        reg = new_reg
-      return self.tor[args.name], reg, off
-    reg = self.render_alu(BinaryOps.ADD, self.render_cast(reg, dtypes.uint64), self.tor[args.name], dtype=dtypes.uint64)
-    return reg, None, off
-
-def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
-  #TODO: Do not use clear()
-  lang.ins.clear()
-  lang.tor.clear()
-  lang.cnts.clear()
-  buf_to_dtype = {args:dtype for uop,dtype,_,args,_ in uops if uop == Ops.DEFINE_GLOBAL}
-  global_size, local_size = [], []
-  skipload_branch = 0
-  lang.ins += [AssemblyInstruction(Ops.SPECIAL, lang.newreg(buf, dtype=dtypes.uint64, scalar=True), [], buf) for buf in buf_to_dtype]
-  for u in uops:
-    uop,dtype,vin,args,_ = u
-    if uop == Ops.DEFINE_LOCAL:
-      lang.ins.append(AssemblyInstruction(Ops.DEFINE_LOCAL, None, [], args))
-      lang.ins.append(AssemblyInstruction(Ops.ALU, lang.newreg(args[0], dtype=dtypes.uint64), [args[0]], UnaryOps.NOOP))
-    elif uop == Ops.LOOP:
-      if args[1] == "global":
-        for i,var in enumerate(args[0]):
-          global_size.append(var.max+1)
-          lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"gid{len(args[0])-1-i}"))
-      elif args[1] == "local":
-        for i,var in enumerate(args[0]):
-          local_size.append(var.max+1)
-          lang.ins.append(AssemblyInstruction(Ops.SPECIAL, lang.newreg(var, dtype=dtypes.int32), [], f"lid{len(args[0])-1-i}"))
-      else:
-        for var in args[0]:
-          if not isinstance(var, NumNode):  # TODO: why is this coming through?
-            lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(var, dtype=dtypes.int32, scalar=True), [], 0))
-            lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], "$loop_"+var.expr))
-    elif uop == Ops.ENDLOOP:
-      if args[1] not in ["global", "local", "global+local"]:
-        for var in reversed(args[0]):
-          if not isinstance(var, NumNode):  # TODO: why is this coming through?
-            lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[var], [lang.tor[var], 1], BinaryOps.ADD))
-            pred = lang.render_alu(BinaryOps.CMPLT, lang.tor[var], var.max+1, dtypes.bool)
-            lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], ("$loop_"+var.expr, True)))
-      elif args[1] == "global+local":
-        for i, var in enumerate(reversed(args[0])):
-          lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"gid{i}")))
-      elif args[1] == 'local':
-        for i, var in enumerate(reversed(args[0])):
-          lang.ins.append(AssemblyInstruction(Ops.ENDLOOP, None, [lang.tor[var]], (var.max+1, f"lid{i}")))
-    elif uop == Ops.CAST:
-      # TODO: we should reconsider outputting CAST in the linearizer. these are needless copies
-      out = lang.newreg(u, dtype)
-      for i,sr in enumerate(out.subregs()):
-        lang.ins.append(AssemblyInstruction(Ops.ALU, sr, [lang.tor[vin[i]]], UnaryOps.NOOP))
-    elif uop == Ops.ALU:
-      out = lang.newreg(u, dtype) if u not in lang.tor else lang.tor[u]
-      # this is the only thing that can violate SSA
-      if args in [BinaryOps.CMPLT]:
-        pred_reg = lang.newreg((u, 'pred'), dtype=dtypes.bool)
-        lang.ins.append(AssemblyInstruction(Ops.ALU, pred_reg, [lang.tor[x] for x in vin], args))
-        lang.ins.append(AssemblyInstruction(Ops.CAST, out, [pred_reg], args))
-      elif args == BinaryOps.DIV and lang.no_div:
-        tmp = lang.newreg((u, "rcp"))
-        lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[1]]], UnaryOps.RECIP))
-        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[vin[0]], tmp], BinaryOps.MUL))
-      elif args == UnaryOps.SIN and lang.sin_is_sin2pi:
-        tmp = lang.newreg((u, "2pi"))
-        lang.ins.append(AssemblyInstruction(Ops.ALU, tmp, [lang.tor[vin[0]], 1/(math.pi*2)], BinaryOps.MUL))
-        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
-      else:
-        lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
-    elif uop == Ops.DEFINE_REG:
-      reg = lang.newreg(u, dtype=dtype)
-      lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
-    elif uop == Ops.SPECIAL:
-      lang.tor[u] = lang.tor[args]
-    elif uop == Ops.CONST:
-      lang.ins.append(AssemblyInstruction(Ops.LOAD, lang.newreg(u, dtype=dtype), [], args))
-    elif uop == Ops.LOAD:
-      idx, treg, off = lang.addr_w_offset(args)
-      reg = lang.newreg(u, dtype=dtype, scalar=(idx.scalar and (not isinstance(treg, Register) or treg.scalar)))
-      if args.valid.min == 0:
-        lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], 0))
-        if args.valid.max == 1:
-          pred = args.valid.render(lang.render_ops, lang)
-          lang.ins.append(AssemblyInstruction(Ops.COND_BRANCH, None, [pred], (f"$skipload_{skipload_branch}", False)))
-      if args.valid.max == 1:
-          # NOTE: you can't compute the index in here, because it assumes it's all available later
-        lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [idx] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
-      if args.valid.min == 0 and args.valid.max == 1:
-        lang.ins.append(AssemblyInstruction(Ops.LABEL, None, [], f"$skipload_{skipload_branch}"))
-        skipload_branch += 1
-    elif uop == Ops.STORE:
-      if args is None:
-        lang.ins.append(AssemblyInstruction(Ops.ALU, lang.tor[vin[0]], [lang.tor[vin[1]]], UnaryOps.NOOP))
-      else:
-        idx, treg, off = lang.addr_w_offset(args)
-        lang.ins.append(AssemblyInstruction(Ops.STORE, None, [idx, lang.tor[vin[0]]] + ([treg] if treg is not None else []), (off, 'global' if not args.local else 'shared', args.memory_dtype if args.memory_dtype != dtypes.float else None)))
-
-  if DEBUG >= 4:
-    for tins in lang.ins: print(tins)
-  return global_size, local_size

extra/assembly/assembly_arm64.py

@@ -1,177 +0,0 @@
-import struct
-from platform import system
-from typing import Tuple, Dict, List, Optional
-from tinygrad import dtypes
-from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.opt.kernel import Ops, UOp
-from tinygrad.helpers import CI
-from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-
-def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
-def compute_offsets(total):
-  quotient, remainder = divmod(total, 4096)
-  return [4096]*quotient + [remainder] if remainder else [4096]*quotient
-
-#NOTE: Darwin needs names to start with a "_"
-def get_name(name): return ('_' if system() == 'Darwin' else '') + name
-
-class ARM64Language(AssemblyLanguage): pass
-
-def specialize_to_arm64(fn_nm, asm):
-  var_size = 16
-  prev_uop:Optional[Ops] = None
-  ins = []
-  x_regs = ['x' + str(i) for i in reversed(range(12))]
-  s_regs = ['s' + str(i) for i in reversed(range(3,32)) if i <= 7 or i >= 16]
-  type_to_reg = {dtypes.double: "d", dtypes.half: 'h', dtypes.float32: 's', dtypes.bool: 'w', dtypes.int8:'w', dtypes.int32: 'w', dtypes.int64: 'x', dtypes.uint8:'w', dtypes.uint32: 'w', dtypes.uint64: 'x'}
-  alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
-          BinaryOps.MOD: "", BinaryOps.CMPLT: "subs",
-          UnaryOps.NOOP: "mov", UnaryOps.NEG: "neg",
-          UnaryOps.SIN:'bl ' + get_name('sinf'), UnaryOps.LOG2: 'bl ' + get_name("log2f"), UnaryOps.EXP2: 'bl ' + get_name("exp2f"), UnaryOps.SQRT: 'bl ' + get_name("sqrtf"),
-          TernaryOps.MULACC: "madd", TernaryOps.WHERE: "fcsel"}
-
-  def mov_imm(value, reg):
-    # Manually move value into reg if value can't fit
-    if value.__class__ is not float and abs(value) > abs(65535):
-      ins.append(f"movz w15, #{value & 0xffff}")
-      ins.append(f"movk w15, #{(value >> 16) & 0xffff}, lsl #16")
-      ins.append(f"sxtw {reg}, w15")
-    elif reg[0] == 's':
-      ins.append(f"movz x15, 0x{float_to_hex(value)[4:]}")
-      ins.append(f"movk x15, 0x{float_to_hex(value)[:4]}, lsl #16")
-      ins.append("str x15, [sp, 16]")
-      ins.append(f"ldr {reg}, [sp, 16]")
-    else:
-      ins.append(f"mov {reg}, #{value}")
-
-  # Get variables intervals
-  live_range:Dict[str, List[int]] = {}
-  for i, (uop, out, vin, arg) in enumerate(asm):
-    for var in ([v for v in [out] + vin if v is not None and v.__class__ is not int]):
-      live_range[var.nm] = [i,i] if var.nm not in live_range else [live_range[var.nm][0], i]
-
-  mem_vars:Dict[str, int] = {}
-  rtor:Dict[str, str] = {}
-  def allocate_regs(mvars):
-    nonlocal var_size
-    for v in [v for v in mvars if v is not None and v.__class__ is not int and v.nm not in rtor]:
-      available_regs = s_regs if dtypes.is_float(v[1]) else x_regs
-      #NOTE: Very simple spill, everything that don't fit in regs goes to mem
-      if not available_regs:
-        # ARM needs the stack 16-byte aligned
-        var_size += 16
-        available_regs.append('s0' if dtypes.is_float(out[1]) else 'x12')
-        mem_vars[v.nm] = var_size
-      rtor[v.nm] = available_regs.pop()
-
-  temp_floats = ['s0', 's1', 's2']
-  temp_ints = ['x12', 'x13', 'x16']
-  for i, (uop, out, vin, arg) in enumerate(asm):
-    # Clear regs out of interval
-    for var, reg in list(rtor.items()):
-      available_regs = s_regs if reg[0] == 's' else x_regs
-      if var[1] not in 'B' and var not in mem_vars and i > live_range[var][1]:
-        available_regs.append(rtor.pop(var))
-    # Assign a registers to the variables using live ranges.
-    allocate_regs([out] + vin)
-    # Assign temp regs to vin and load them before direct use
-    for i, v in enumerate([v for v in vin if v.__class__ is not int and v.nm in mem_vars]):
-      rtor[v.nm] = temp_floats[i] if dtypes.is_float(v[1]) else temp_ints[i]
-      # ARM64 addressing constraints https://devblogs.microsoft.com/oldnewthing/20220728-00/?p=106912
-      ins.append(f"mov x15, {mem_vars[v.nm]}")
-      ins.append(f"ldr {rtor[v.nm]}, [sp, x15]")
-
-    if uop == Ops.SPECIAL:
-      if arg.startswith('data'):
-        # data 8 to n into the stack
-        if int(arg[4:]) >= 8:
-          ins.append(f"ldr x15, [x17, #{(int(arg[4:]) - 8) * 8}]")
-          ins.append(f"mov {rtor[out.nm]}, x15")
-      else:
-        ins.append(f"mov {rtor[out.nm]}, #0")
-        ins.append(f"loop_{arg}:")
-    elif uop == Ops.CAST:
-      if arg == BinaryOps.CMPLT:
-        if rtor[out.nm][0] == 's':
-          mov_imm(0.0, 's0')
-          mov_imm(1.0, 's1')
-          ins.append(f"fcsel {rtor[out.nm]}, s1, s0, lt")
-        if rtor[out.nm][0] == 'x':
-          mov_imm(0, 'x14')
-          mov_imm(1, 'x15')
-          ins.append(f"csel {rtor[out.nm]}, x15, x14, lt")
-      else:
-        ins.append(f"sxtw {rtor[out.nm]}, w{rtor[vin[0].nm][1:]}")
-    elif uop == Ops.ALU:
-      if len(vin)==2 and vin[1].__class__ is int: mov_imm(vin[1], 'x15')
-      if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
-        ins.append(f"ands {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
-      elif arg == TernaryOps.WHERE:
-        ins.append(f"fcmp {rtor[vin[0].nm]}, #0.0" if rtor[vin[0].nm][0] == 's' else f"cmp {rtor[vin[0].nm]}, #0")
-        ins.append(f"{alu[arg]} {rtor[out.nm]}, {rtor[vin[1].nm]}, {rtor[vin[2].nm]}, ne")
-      elif arg in [UnaryOps.LOG2, UnaryOps.SIN, UnaryOps.EXP2, UnaryOps.SQRT]:
-        #NOTE: Not a real instruction, use to emulate a ext call in unicorn
-        if CI: ins.append(f"{alu[arg]} {rtor[out.nm]} {rtor[vin[0].nm]}")
-        else:
-          save_regs = [k for k in rtor.keys() if k != out.nm and k not in mem_vars]
-          ins.append(f"sub sp, sp, #{(len(save_regs))*16}")
-          # Save the registers before they are cleared by func call
-          for i,k in enumerate(save_regs,1):
-            ins.append(f"str {rtor[k]}, [sp, #{16*i}]")
-          ins.append("stp x29, x30, [sp, #0]!")
-          ins.append("mov x29, sp")
-          ins.append(f"fmov s0, {rtor[vin[0].nm]}")
-          ins.append(alu[arg])
-          ins.append(f"fmov {rtor[out.nm]}, s0")
-          ins.append("mov sp, x29")
-          ins.append("ldp x29, x30, [sp], #0")
-          for i,k in enumerate(save_regs,1):
-            ins.append(f"ldr {rtor[k]}, [sp, #{16*i}]")
-          ins.append(f"add sp, sp, #{len(save_regs)*16}")
-      elif arg == BinaryOps.CMPLT:
-        ins.append(f"{alu[arg]} {','.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}" if not dtypes.is_float(vin[0][1]) else f"fcmp {rtor[vin[0].nm]}, {rtor[vin[1].nm]}")
-      elif arg == BinaryOps.MOD:
-        rhs = 'x15' if vin[1].__class__ is int else rtor[vin[1].nm]
-        ins.append(f"udiv x14, {rtor[vin[0].nm]}, {rhs}")
-        ins.append(f"msub {rtor[out.nm]}, x14, {rhs}, {rtor[vin[0].nm]}")
-      else:
-        ins.append(f"{'f' if dtypes.is_float(vin[0][1]) else 's' if arg == BinaryOps.DIV else ''}{alu[arg]} {', '.join('x15' if v.__class__ is int else rtor[v.nm] for v in [out] + vin)}")
-    elif uop == Ops.LOAD:
-      if arg.__class__ in (int, float):
-        mov_imm(arg, rtor[out.nm])
-      else:
-        #NOTE: if need casting load var in s/h0 or x/w12 temp regs
-        reg_in = type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[out.nm]
-        mov_imm(arg[0], "x15")
-        ins.append(f"add x15, {rtor[vin[0].nm]}, x15")
-        ins.append(f"ldr{'sb' if arg[2] is not None and arg[2] in (dtypes.int8, dtypes.uint8, dtypes.bool) else ''} {reg_in}, [x15]")
-        if arg[2] is not None: ins.append(f"{'fcvt' if arg[2] in [dtypes.half, dtypes.double] else 'scvtf'} {rtor[out.nm]}, {reg_in}")
-    elif uop == Ops.STORE:
-      #NOTE: if need casting load var in s/h0 or x/w12 temp regs
-      reg_out = (type_to_reg[arg[2]] + ('0' if dtypes.is_float(arg[2]) else '12') if arg[2] is not None else rtor[vin[1].nm])
-      if arg[2] is not None: ins.append(f"fcvt{'zs' if arg[2] not in [dtypes.half, dtypes.double] else '' } {reg_out}, {rtor[vin[1].nm]}")
-      ins.append(f"mov x15, #{arg[0]}")
-      ins.append(f"str {reg_out}, [{rtor[vin[0].nm]}, x15, lsl #0]")
-    elif uop == Ops.COND_BRANCH:
-      #TODO: this is a hack it shouldn't always be a cmp before a cond branch?
-      if prev_uop == Ops.LOAD:
-        ins.append(f"cmp {rtor[vin[0].nm]}, #0")
-      ins.append(f"b.{'lt' if arg[1] else 'ge'} {arg[0][1:]}")
-    elif uop == Ops.LABEL:
-      ins.append(f"{arg[1:]}:")
-    elif uop == Ops.ENDLOOP:
-      mov_imm(arg[0], "x15")
-      ins.append(f"add {rtor[vin[0].nm]}, {rtor[vin[0].nm]}, #1")
-      ins.append(f"cmp {rtor[vin[0].nm]}, x15")
-      ins.append(f"b.lt loop_{arg[1]}")
-    prev_uop = uop
-    # store regs into memory if needed
-    if out is not None and out.nm in mem_vars:
-      ins.append(f"mov x15, {mem_vars[out.nm]}")
-      ins.append(f"str {rtor[out.nm]}, [sp, x15]")
-  return "\n".join([f"//varsize {var_size}",".arch armv8-a",".text", f".global {get_name(fn_nm)}",".p2align 2", f"{get_name(fn_nm)}:", "mov x17, sp"] + [f"sub sp, sp, #{offset}" for offset in compute_offsets(var_size)]+ ins + [f"add sp, sp, #{offset}" for offset in compute_offsets(var_size)] +["ret", "\n"])
-
-def uops_to_arm64_asm(fn_nm:str, uops:List[UOp]) -> Tuple[str, List[int], List[int], bool]:
-  lang = ARM64Language()
-  global_size, local_size = uops_to_asmstyle(lang, fn_nm, uops)
-  return specialize_to_arm64(fn_nm, lang.ins), global_size[::-1], local_size[::-1], True

extra/assembly/assembly_ptx.py

@@ -1,105 +0,0 @@
-from typing import List
-import struct
-from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.opt.kernel import Ops, UOp
-from tinygrad import dtypes
-from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_cuda import arch
-
-dtype_to_nvtype = {dtypes.float32: "f32", dtypes.float16: "f16", dtypes.int64: "s64", dtypes.int32: "s32", dtypes.int8: "s8", dtypes.bool: "pred", dtypes.uint64: "u64", dtypes.uint32: "u32", dtypes.uint16: "u16", dtypes.uint8: "u8", "bits16": "b16", dtypes.float64: "f64"}
-def float_to_hex(x): return "%02X%02X%02X%02X" % tuple(struct.pack("f",x)[::-1])
-
-def ptx_needs_cast(dest_dtype, src_dtype): return dtypes.is_float(dest_dtype) and dtypes.is_int(src_dtype) or dtypes.is_int(dest_dtype) and dtypes.is_float(src_dtype) or (dtypes.is_float(src_dtype) and dtypes.is_float(dest_dtype) and dest_dtype.itemsize != src_dtype.itemsize)
-
-def render_cast(ins, inp, out):
-  if inp.dtype == dtypes.bool and (dtypes.is_float(out.dtype) or dtypes.is_int(out.dtype)):
-    ins.append(f"selp.{dtype_to_nvtype[out.dtype]} {out}, {'0f3F800000, 0f00000000' if dtypes.is_float(out.dtype) else '1, 0'}, {inp};")
-  elif out.dtype == dtypes.bool:
-    if inp.dtype == dtypes.bool:
-      ins.append(f"mov.pred {out}, {inp};")
-    else:
-      ins.append(f"setp.ne.{dtype_to_nvtype[inp.dtype]} {out}, {'0f00000000' if dtypes.is_float(inp.dtype) else '0'}, {inp};")
-  else:
-    round_mod = ".rzi" if dtypes.is_int(out.dtype) and dtypes.is_float(inp.dtype) else '.rz' if dtypes.is_float(out.dtype) and (dtypes.is_int(inp.dtype) or dtypes.is_float(inp.dtype) and inp.dtype.itemsize > out.dtype.itemsize) else ''
-    ins.append(f"cvt{round_mod}.{dtype_to_nvtype[out.dtype]}.{dtype_to_nvtype[inp.dtype]} {out}, {inp};")
-
-# https://docs.nvidia.com/cuda/parallel-thread-execution/#
-
-class PTXLanguage(AssemblyLanguage):
-  supports_constant_folding: bool = True
-
-def specialize_to_ptx(lang, function_name):
-  param_cnt = 0
-  ins = []
-  alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", BinaryOps.DIV: "div", BinaryOps.MAX: "max",
-         BinaryOps.MOD: "rem", BinaryOps.CMPLT: "setp.lt", UnaryOps.SQRT: "sqrt.approx",
-         UnaryOps.NOOP: "mov", UnaryOps.NEG: "neg",
-         UnaryOps.SIN: "sin.approx", UnaryOps.LOG2: "lg2.approx", UnaryOps.EXP2: "ex2.approx.ftz",
-         TernaryOps.MULACC: "fma.rn", TernaryOps.WHERE: "selp"}
-  for uop, out, vin, arg in lang.ins:
-    if uop == Ops.ENDLOOP:
-      ins.append("bar.sync 0;")
-    elif uop == Ops.DEFINE_LOCAL:
-      ins.append(f".shared .align 4 .b8 {arg[0]}[{arg[1]*4}];")
-    elif uop == Ops.SPECIAL:
-      if arg.startswith('data'):
-        param_cnt += 1
-        ins.append(f"ld.param.u64 {out}, [{arg}];")
-        # TODO: we sometimes want this to be local, nvcc converts to global most of the time, not sure when we would need to?
-        # ins.append(f"cvta.to.global.u64 {out}, {out};")
-      elif arg.startswith('gid'):
-        ins.append(f"mov.u32 {out}, %ctaid.{'xyz'[int(arg[3:])]};")
-      elif arg.startswith('lid'):
-        ins.append(f"mov.u32 {out}, %tid.{'xyz'[int(arg[3:])]};")
-    elif uop == Ops.ALU:
-      if arg == BinaryOps.MUL and out.dtype == dtypes.bool:
-        ins.append(f"and.pred {out}, {', '.join(str(x) for x in vin)};")
-      else:
-        otype = vin[0].dtype if arg in [BinaryOps.CMPLT] else out.dtype
-        if arg == TernaryOps.WHERE:
-          if vin[0].dtype == dtypes.bool:
-            reg = vin[0]
-          else:
-            reg = lang.newreg((vin[0], 'bool'), dtypes.bool)
-            ins.append(f"setp.ne.{dtype_to_nvtype[vin[0].dtype]} {reg}, {'0f00000000' if dtypes.is_float(vin[0].dtype) else '0'}, {vin[0]};")
-          vin = vin[1:] + [reg]
-        ins.append(f"{alu[arg]}{'.lo' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 else ''}{'.rn' if arg == BinaryOps.DIV and out.dtype == dtypes.float32 else ''}.{dtype_to_nvtype[otype]} {out}, {', '.join(str(x) for x in vin)};")
-    elif uop == Ops.LOAD:
-      if arg.__class__ in (int, float):
-        ins.append(f"mov.{dtype_to_nvtype[out.dtype]} {out}, {'0f'+float_to_hex(arg) if dtypes.is_float(out.dtype) else int(arg)};")
-      elif arg[2] is not None and (arg[2] == dtypes.bool or arg[2] != out.dtype):
-        dt = ('u16', dtypes.uint16) if arg[2] == dtypes.bool == out.dtype else ('u8', dtypes.uint8) if arg[2] == dtypes.bool else ('b16', dtypes.float16) if arg[2] == dtypes.half else (dtype_to_nvtype[arg[2]], arg[2])
-        reg = lang.newreg((out, dt[0]), dtype=dt[1])
-        ins.append(f"ld.{arg[1]}.{dt[0]} {reg}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];")
-        render_cast(ins, reg, out)
-      else:
-        ins.append(f"ld.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} {out}, [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}];")
-    elif uop == Ops.STORE:
-      if ptx_needs_cast(dtypes.float if arg[2] is None else arg[2], vin[1].dtype) or arg[2] == dtypes.bool:
-        if arg[2] == dtypes.bool != vin[1].dtype:
-          prereg = lang.newreg((vin[1],'bool'), dtype=dtypes.bool)
-          render_cast(ins, vin[1], prereg)
-        else: prereg = vin[1]
-        reg = lang.newreg((prereg, dtypes.uint16 if arg[2] == dtypes.bool else arg[2]), dtype=dtypes.uint16 if arg[2] == dtypes.bool else dtypes.float if arg[2] is None else arg[2])
-        render_cast(ins, prereg, reg)
-        ins.append(f"st.{arg[1]}.{dtype_to_nvtype['bits16' if arg[2] == dtypes.float16 else dtypes.uint8 if arg[2] == dtypes.bool else dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {reg};")
-      else:
-        ins.append(f"st.{arg[1]}.{dtype_to_nvtype[dtypes.float if arg[2] is None else arg[2]]} [{vin[0]}{f'+{arg[0]}' if arg[0] is not None else ''}], {vin[1]};")
-    elif uop == Ops.CAST:
-      render_cast(ins, vin[0], out)
-    elif uop == Ops.LABEL:
-      ins.append(f"{arg}:")
-    elif uop == Ops.COND_BRANCH:
-      ins.append(f"@{'!' if not arg[1] else ''}{vin[0]} bra {arg[0]};")
-
-  ins_prefix = [".version 7.8", ".target " + arch(), ".address_size 64",
-                f".visible .entry {function_name}({', '.join(f'.param .u64 data{i}' for i in range(param_cnt))}) {{"]
-  for arg in [(dtype, lang.type_to_letter(dtype), c) for dtype,c in lang.cnts.items()]: ins_prefix.append(f".reg .{dtype_to_nvtype[arg[0][0]]} %{arg[1]}<{arg[2]}>;",)
-  ins = ins_prefix + ins
-  ins += ["ret;", "}"]
-  return '\n'.join(ins)
-
-def uops_to_ptx_asm(function_name:str, uops:List[UOp]):
-  lang = PTXLanguage()
-  global_size, local_size = uops_to_asmstyle(lang, function_name, uops)
-  return specialize_to_ptx(lang, function_name), global_size[::-1], local_size[::-1], True

extra/assembly/assembly_rdna.py

@@ -1,203 +0,0 @@
-import yaml
-from typing import Tuple, Set, Dict
-from tinygrad import dtypes
-from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.codegen.opt.kernel import Ops
-from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_cl import ROCM_LLVM_PATH
-
-# ugh, is this really needed?
-from extra.helpers import enable_early_exec
-early_exec = enable_early_exec()
-
-boilerplate_start = """
-.global _start
-_start:
-.rodata
-.align 0x10
-.global code.kd
-.type code.kd,STT_OBJECT
-.amdhsa_kernel code"""
-
-code_start = """.end_amdhsa_kernel
-.text
-code:
-"""
-
-# https://github.com/RadeonOpenCompute/ROCm_Documentation/blob/master/ROCm_Compiler_SDK/ROCm-Codeobj-format.rst
-# https://github.com/ROCm-Developer-Tools/ROCm-ComputeABI-Doc/blob/master/AMDGPU-ABI.md#initial-kernel-register-state
-# RDNA3 is actually a SIMD machine!
-class RDNACodegen(AssemblyCodegen):
-  supports_float4: bool = True
-  supports_float4_alu: bool = True
-  supports_load3: bool = True
-  sin_is_sin2pi: bool = True
-  no_div: bool = True
-
-  def specialize(self, asm) -> Tuple[str, str]:
-    args = []
-    for i,b in enumerate(self.bufs): args.append({'.address_space': 'global', '.name': f'buf_{i}', '.offset': i*8, '.size': 8, '.type_name': b.dtype.name+"*", '.value_kind': 'global_buffer'})
-    ins = []
-
-    v_cnt = 3  # v[0:2] is local_xyz
-    s_cnt = 5  # s[0:1] is the address, s[2:4] is global_xyz
-
-    dtype_to_rdnatype = {dtypes.float32: "f32", dtypes.int64: "i64", dtypes.int32: "i32", dtypes.uint64: "u64", dtypes.bool: "i32"}
-    alu = {BinaryOps.ADD: "add", BinaryOps.SUB: "sub", BinaryOps.MUL: "mul", TernaryOps.MULACC: "fma",
-           BinaryOps.MAX: "max", UnaryOps.RECIP: "rcp",
-           UnaryOps.NOOP: "mov", UnaryOps.SIN: "sin", UnaryOps.LOG2: "log", UnaryOps.EXP2: "exp",
-           BinaryOps.CMPLT: "cmp_lt"}
-
-    pend_regs:Set[Register] = set()
-    rtor:Dict[Register, str] = {}
-    def reg_in(x):
-      nonlocal pend_regs
-      #print("reg_in", x, rtor[x], pend_regs)
-      if x in pend_regs:
-        #print("clear")
-        ins.append('s_waitcnt lgkmcnt(0), vmcnt(0)')
-        pend_regs.clear()
-      return rtor[x]
-    def reg_out(x):
-      return rtor[x]
-    for uop, out, vin, arg in asm:
-      if uop == Ops.DEFINE_REGISTER:
-        if arg[0][0] in [dtypes.uint32, dtypes.uint64, dtypes.int64, dtypes.int32, dtypes.float32, dtypes.float.vec(4)]:
-          for i in range(arg[2]):
-            # TODO: Re-use gaps created by this to avoid wasting registers
-            align = int(arg[0][0].itemsize / 4)
-            if arg[0][1]:
-              s_cnt += s_cnt % align
-              reg_name = f"s[{s_cnt}:{s_cnt + align - 1}]" if align > 1 else f"s{s_cnt}"
-              s_cnt += align
-            else:
-              v_cnt += v_cnt % align
-              reg_name = f"v[{v_cnt}:{v_cnt + align - 1}]" if align > 1 else f"v{v_cnt}"
-              v_cnt += align
-            rtor[Register(f"%{arg[1]}{i}", *arg[0])] = reg_name
-
-            if arg[0][0] == dtypes.float.vec(4):
-              for off in range(4):
-                reg_name = f"s{s_cnt-align+off}" if arg[0][1] else f"v{v_cnt-align+off}"
-                rtor[Register(f"%{arg[1]}{i}", dtypes.float, False, off=off)] = reg_name
-        elif arg[0][0] == dtypes.bool:
-          for i in range(arg[2]):
-            reg_name = "scc" if arg[0][1] else "vcc_lo" # `_lo` suffix since we're running wavefront_size=32
-            rtor[Register(f"%{arg[1]}{i}", *arg[0])] = reg_name
-        else:
-          raise NotImplementedError("DEFINE_REGISTER not implemented for arg: ", arg)
-      elif uop == Ops.SPECIAL:
-        if arg.startswith('buf'):
-          i = int(arg[3:])
-          ins.append(f's_load_b64 {reg_out(out)}, s[0:1], {i*8}')
-          pend_regs.add(out)
-          for r in out.subregs(): pend_regs.add(r)
-        elif arg.startswith('gid'):
-          ins.append(f'v_mov_b32 {reg_out(out)}, s{2+int(arg[3])}')
-          # the docs lied, this is actually y
-          if int(arg[3]) == 2: ins.append("v_bfe_u32 v2, v0, 20, 10")  # untested
-          if int(arg[3]) == 1: ins.append("v_bfe_u32 v1, v0, 10, 10")
-          elif int(arg[3]) == 0: ins.append("v_and_b32_e32 v0, 0x3ff, v0")
-          # get local size
-          offset = len(args)*8
-          args.append({".offset": offset, ".value_kind": f"hidden_group_size_{'xyz'[int(arg[3])]}", ".size": 8})
-          ins.append(f's_load_b32 s{2+int(arg[3])}, s[0:1], {offset}')
-          ins.append('s_waitcnt vmcnt(0) lgkmcnt(0)')
-          pend_regs.clear()
-          ins.append(f'v_mul_i32_i24 {reg_out(out)}, {reg_out(out)}, s{2+int(arg[3])}')
-          ins.append(f'v_add_nc_u32 {reg_out(out)}, v{int(arg[3])}, {reg_out(out)}')
-      elif uop == Ops.CONST:
-        if arg == float('inf'): arg = "0x7f800000"
-        elif arg == float('-inf'): arg = "0xff800000"
-        if out.dtype == dtypes.float.vec(4):
-          for off in range(4):
-            ins.append(f"{'s_' if out.scalar else 'v_'}mov_b32 {reg_out(Register(out.nm, dtypes.float, False, off=off))}, {arg}")
-        else:
-          ins.append(f"{'s_' if out.scalar else 'v_'}mov_b32 {reg_out(out)}, {arg}")
-      elif uop == Ops.ALU:
-        if arg in [BinaryOps.CMPLT]:
-          ins.append(f"{'s' if out.scalar else 'v'}_{alu[arg]}_{dtype_to_rdnatype[out.dtype]} {', '.join(reg_in(x) if x.__class__ is Register else str(x) for x in vin)}")
-        else:
-          alu_arg = alu[arg]
-          if arg == TernaryOps.MULACC and out == vin[2]:
-            alu_arg = "fmac"
-            vin = vin[0:2]
-          if out.dtype == dtypes.float.vec(4):
-            for rr in zip(*[x.subregs() if x.dtype == dtypes.float.vec(4) else [x,x,x,x] for x in [out]+vin]):
-              ins.append(f"{'s_' if rr[0].scalar else 'v_'}{alu_arg}_{dtype_to_rdnatype[rr[0].dtype]} {reg_out(rr[0])}, {', '.join(reg_in(x) if x.__class__ is Register else str(x) for x in rr[1:])}")
-          else:
-            ins.append(f"{'s_' if out.scalar else 'v_'}{alu_arg}_{dtype_to_rdnatype[out.dtype] if arg != UnaryOps.NOOP else 'b32'}{'_i24' if arg == BinaryOps.MUL and out.dtype != dtypes.float32 and not out.scalar else ''} {reg_out(out)}, {', '.join(reg_in(x) if x.__class__ is Register else str(x) for x in vin)}")
-      elif uop == Ops.LOAD:
-        if out.scalar:
-          # swap arg order
-          ins.append(f's_load_b32 {reg_out(out)}, {reg_in(vin[0])}, {reg_in(vin[1])} offset:{arg[0]}')
-        else:
-          ins.append(f'global_load_{"b128" if out.dtype == dtypes.float.vec(4) else "b32"} {reg_out(out)}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
-        pend_regs.add(out)
-        for r in out.subregs(): pend_regs.add(r)
-      elif uop == Ops.STORE:
-        ins.append(f'global_store_{"b128" if vin[1].dtype == dtypes.float.vec(4) else "b32"} {reg_in(vin[2])}, {reg_in(vin[1])}, {reg_in(vin[0])} offset:{arg[0]}')
-      elif uop == Ops.LABEL:
-        ins.append(f"{arg}:")
-      elif uop == Ops.COND_BRANCH:
-        ins.append(f"s_cbranch_scc{'1' if arg[1] else '0'} {arg[0]}")
-      elif uop == Ops.CAST:
-        if vin[0].dtype == dtypes.bool:
-          if out.dtype == dtypes.float32:
-            ins.append(f"v_cndmask_b32 {reg_out(out)}, 0.0, 1.0, {reg_in(vin[0])}")
-        else:
-          raise NotImplementedError(f"cast {vin[0].dtype} -> {out.dtype}")
-      else:
-        raise NotImplementedError(uop)
-
-    ins += ['s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)', 's_endpgm', 's_code_end']
-
-    # dual alu group
-    seen = set()
-    new_ins = []
-    for i,tins in enumerate(ins):
-      if tins in seen: continue
-      if tins.startswith("v_fmac_f32"):
-        for gins in reversed(ins[i+1:]):
-          if gins in seen: continue
-          if gins.startswith("v_fmac_f32"):
-            r0 = [int(x[1:].strip(',')) for x in tins.split(" ")[1:]]
-            r1 = [int(x[1:].strip(',')) for x in gins.split(" ")[1:]]
-            if r0[0]%2 == r1[0]%2: continue
-            if r0[1]%2 == r1[1]%2: continue
-            if r0[2]%2 == r1[2]%2: continue
-            new_ins.append(tins.replace("v_", "v_dual_")+" :: " + gins.replace("v_", "v_dual_"))
-            seen.add(tins)
-            seen.add(gins)
-            break
-      if tins not in seen:
-        new_ins.append(tins)
-    ins = new_ins
-
-    return 'code', self.assemble(args, ins, v_cnt, s_cnt)
-
-  def assemble(self, args, ins, v_cnt, s_cnt):
-    kernel_desc = {'.amdhsa_group_segment_fixed_size': 0, '.amdhsa_private_segment_fixed_size': 0, '.amdhsa_kernarg_size': 0,
-                   '.amdhsa_next_free_vgpr': v_cnt,   # this matters!
-                   '.amdhsa_reserve_vcc': 0, '.amdhsa_reserve_xnack_mask': 0,
-                   '.amdhsa_next_free_sgpr': s_cnt,
-                   '.amdhsa_float_round_mode_32': 0, '.amdhsa_float_round_mode_16_64': 0, '.amdhsa_float_denorm_mode_32': 3, '.amdhsa_float_denorm_mode_16_64': 3, '.amdhsa_dx10_clamp': 1, '.amdhsa_ieee_mode': 1,
-                   '.amdhsa_fp16_overflow': 0, '.amdhsa_workgroup_processor_mode': 1, '.amdhsa_memory_ordered': 1, '.amdhsa_forward_progress': 0, '.amdhsa_enable_private_segment': 0,
-                   '.amdhsa_system_sgpr_workgroup_id_x': 1, '.amdhsa_system_sgpr_workgroup_id_y': 1, '.amdhsa_system_sgpr_workgroup_id_z': 1,
-                   '.amdhsa_system_sgpr_workgroup_info': 0, '.amdhsa_system_vgpr_workitem_id': 2, # is amdhsa_system_vgpr_workitem_id real?
-                   '.amdhsa_exception_fp_ieee_invalid_op': 0, '.amdhsa_exception_fp_denorm_src': 0, '.amdhsa_exception_fp_ieee_div_zero': 0, '.amdhsa_exception_fp_ieee_overflow': 0, '.amdhsa_exception_fp_ieee_underflow': 0,
-                   '.amdhsa_exception_fp_ieee_inexact': 0, '.amdhsa_exception_int_div_zero': 0, '.amdhsa_user_sgpr_dispatch_ptr': 0, '.amdhsa_user_sgpr_queue_ptr': 0, '.amdhsa_user_sgpr_kernarg_segment_ptr': 1,
-                   '.amdhsa_user_sgpr_dispatch_id': 0, '.amdhsa_user_sgpr_private_segment_size': 0, '.amdhsa_wavefront_size32': 1, '.amdhsa_uses_dynamic_stack': 0}
-
-    metadata = {'amdhsa.kernels': [{'.args': args,
-                  '.group_segment_fixed_size': 0, '.kernarg_segment_align': 8, '.kernarg_segment_size': args[-1][".offset"] + args[-1][".size"],
-                  '.language': 'OpenCL C', '.language_version': [1, 2], '.max_flat_workgroup_size': 256,
-                  '.name': 'code', '.private_segment_fixed_size': 0, '.sgpr_count': s_cnt, '.sgpr_spill_count': 0,
-                  '.symbol': 'code.kd', '.uses_dynamic_stack': False, '.vgpr_count': v_cnt, '.vgpr_spill_count': 0,
-                  '.wavefront_size': 32}],
-                'amdhsa.target': 'amdgcn-amd-amdhsa--gfx1100', 'amdhsa.version': [1, 2]}
-
-    code = boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) + "\n" +  code_start + '\n'.join(ins) + "\n.amdgpu_metadata\n" + yaml.dump(metadata) + ".end_amdgpu_metadata"
-    obj = early_exec(([ROCM_LLVM_PATH / "llvm-mc", '--arch=amdgcn', '--mcpu=gfx1100', '--triple=amdgcn-amd-amdhsa', '--filetype=obj', '-'], code.encode("utf-8")))
-    asm = early_exec(([ROCM_LLVM_PATH / "ld.lld", "/dev/stdin", "-o", "/dev/stdout", "--pie"], obj))
-    return asm

extra/assembly/ptx/test.py

@@ -1,23 +0,0 @@
-#!/usr/bin/env python3
-import numpy as np
-from tinygrad.runtime.ops_cuda import CUDAProgram, RawCUDABuffer
-
-if __name__ == "__main__":
-  test = RawCUDABuffer.fromCPU(np.zeros(10, np.float32))
-  prg = CUDAProgram("test", """
-  .version 7.8
-  .target sm_86
-  .address_size 64
-  .visible .entry test(.param .u64 x) {
-    .reg .b32       %r<2>;
-    .reg .b64       %rd<3>;
-
-    ld.param.u64    %rd1, [x];
-    cvta.to.global.u64      %rd2, %rd1;
-    mov.u32         %r1, 0x40000000; // 2.0 in float
-    st.global.u32   [%rd2], %r1;
-    ret;
-  }""", binary=True)
-  prg([1], [1], test)
-  print(test.toCPU())
-

extra/assembly/rdna3/alu.py

@@ -0,0 +1,254 @@
+# Pure combinational ALU functions for RDNA3 emulation
+from __future__ import annotations
+import struct, math
+from typing import Callable
+from extra.assembly.rdna3.autogen import SOP1Op, SOP2Op, SOPCOp, SOPKOp, VOP1Op, VOP2Op, VOP3Op
+
+# Format base offsets for unified opcode space
+SOP2_BASE, SOP1_BASE, SOPC_BASE, SOPK_BASE = 0x000, 0x100, 0x200, 0x300
+VOP2_BASE, VOP1_BASE = 0x100, 0x180
+
+# Float conversion helpers
+_I, _f, _H, _e = struct.Struct('<I'), struct.Struct('<f'), struct.Struct('<H'), struct.Struct('<e')
+def f32(i: int) -> float: return _f.unpack(_I.pack(i & 0xffffffff))[0]
+def i32(f: float) -> int:
+  if math.isinf(f): return 0x7f800000 if f > 0 else 0xff800000
+  try: return _I.unpack(_f.pack(f))[0]
+  except (OverflowError, struct.error): return 0x7f800000 if f > 0 else 0xff800000
+def f16(i: int) -> float: return _e.unpack(_H.pack(i & 0xffff))[0]
+def i16(f: float) -> int:
+  if math.isinf(f): return 0x7c00 if f > 0 else 0xfc00
+  try: return _H.unpack(_e.pack(f))[0]
+  except (OverflowError, struct.error): return 0x7c00 if f > 0 else 0xfc00
+def sext(v: int, b: int) -> int: return v - (1 << b) if v & (1 << (b-1)) else v
+def clz(x: int) -> int: return 32 - x.bit_length() if x else 32
+def cls(x: int) -> int: x &= 0xffffffff; return 31 if x in (0, 0xffffffff) else clz(~x & 0xffffffff if x >> 31 else x) - 1
+def _cvt_i32_f32(v): return (0x7fffffff if v > 0 else 0x80000000) if math.isinf(v) else (0 if math.isnan(v) else max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff)
+def _cvt_u32_f32(v): return (0xffffffff if v > 0 else 0) if math.isinf(v) else (0 if math.isnan(v) or v < 0 else min(0xffffffff, int(v)))
+
+# SALU: op -> fn(s0, s1, scc_in) -> (result, scc_out)
+SALU: dict[int, Callable] = {
+  # SOP2
+  SOP2_BASE + SOP2Op.S_ADD_U32: lambda a, b, scc: ((a + b) & 0xffffffff, int((a + b) >= 0x100000000)),
+  SOP2_BASE + SOP2Op.S_SUB_U32: lambda a, b, scc: ((a - b) & 0xffffffff, int(b > a)),
+  SOP2_BASE + SOP2Op.S_ADDC_U32: lambda a, b, scc: ((r := a + b + scc) & 0xffffffff, int(r >= 0x100000000)),
+  SOP2_BASE + SOP2Op.S_SUBB_U32: lambda a, b, scc: ((a - b - scc) & 0xffffffff, int((b + scc) > a)),
+  SOP2_BASE + SOP2Op.S_ADD_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 32)) & 0xffffffff, int(((a >> 31) == (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
+  SOP2_BASE + SOP2Op.S_SUB_I32: lambda a, b, scc: ((r := sext(a, 32) - sext(b, 32)) & 0xffffffff, int(((a >> 31) != (b >> 31)) and ((a >> 31) != ((r >> 31) & 1)))),
+  SOP2_BASE + SOP2Op.S_AND_B32: lambda a, b, scc: ((r := a & b), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_OR_B32: lambda a, b, scc: ((r := a | b), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_XOR_B32: lambda a, b, scc: ((r := a ^ b), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_AND_NOT1_B32: lambda a, b, scc: ((r := a & (~b & 0xffffffff)), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_OR_NOT1_B32: lambda a, b, scc: ((r := a | (~b & 0xffffffff)), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_LSHL_B32: lambda a, b, scc: ((r := (a << (b & 0x1f)) & 0xffffffff), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_LSHR_B32: lambda a, b, scc: ((r := a >> (b & 0x1f)), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_ASHR_I32: lambda a, b, scc: ((r := sext(a, 32) >> (b & 0x1f)) & 0xffffffff, int(r != 0)),
+  SOP2_BASE + SOP2Op.S_MUL_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 32)) & 0xffffffff, scc),
+  SOP2_BASE + SOP2Op.S_MUL_HI_U32: lambda a, b, scc: (((a * b) >> 32) & 0xffffffff, scc),
+  SOP2_BASE + SOP2Op.S_MUL_HI_I32: lambda a, b, scc: (((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff, scc),
+  SOP2_BASE + SOP2Op.S_MIN_I32: lambda a, b, scc: (a, 1) if sext(a, 32) < sext(b, 32) else (b, 0),
+  SOP2_BASE + SOP2Op.S_MIN_U32: lambda a, b, scc: (a, 1) if a < b else (b, 0),
+  SOP2_BASE + SOP2Op.S_MAX_I32: lambda a, b, scc: (a, 1) if sext(a, 32) > sext(b, 32) else (b, 0),
+  SOP2_BASE + SOP2Op.S_MAX_U32: lambda a, b, scc: (a, 1) if a > b else (b, 0),
+  SOP2_BASE + SOP2Op.S_CSELECT_B32: lambda a, b, scc: (a if scc else b, scc),
+  SOP2_BASE + SOP2Op.S_BFE_U32: lambda a, b, scc: ((r := ((a >> (b & 0x1f)) & ((1 << ((b >> 16) & 0x7f)) - 1)) if (b >> 16) & 0x7f else 0), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_BFE_I32: lambda a, b, scc: ((r := sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := (b >> 16) & 0x7f) else 0), int(r != 0)),
+  SOP2_BASE + SOP2Op.S_PACK_LL_B32_B16: lambda a, b, scc: ((a & 0xffff) | ((b & 0xffff) << 16), scc),
+  SOP2_BASE + SOP2Op.S_PACK_LH_B32_B16: lambda a, b, scc: ((a & 0xffff) | (b & 0xffff0000), scc),
+  SOP2_BASE + SOP2Op.S_PACK_HH_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | (b & 0xffff0000), scc),
+  SOP2_BASE + SOP2Op.S_PACK_HL_B32_B16: lambda a, b, scc: (((a >> 16) & 0xffff) | ((b & 0xffff) << 16), scc),
+  SOP2_BASE + SOP2Op.S_ADD_F32: lambda a, b, scc: (i32(f32(a) + f32(b)), scc),
+  SOP2_BASE + SOP2Op.S_SUB_F32: lambda a, b, scc: (i32(f32(a) - f32(b)), scc),
+  SOP2_BASE + SOP2Op.S_MUL_F32: lambda a, b, scc: (i32(f32(a) * f32(b)), scc),
+  # SOP1
+  SOP1_BASE + SOP1Op.S_MOV_B32: lambda a, b, scc: (a, scc),
+  SOP1_BASE + SOP1Op.S_NOT_B32: lambda a, b, scc: ((r := (~a) & 0xffffffff), int(r != 0)),
+  SOP1_BASE + SOP1Op.S_BREV_B32: lambda a, b, scc: (int(f'{a & 0xffffffff:032b}'[::-1], 2), scc),
+  SOP1_BASE + SOP1Op.S_CLZ_I32_U32: lambda a, b, scc: (clz(a), scc),
+  SOP1_BASE + SOP1Op.S_CLS_I32: lambda a, b, scc: (cls(a), scc),
+  SOP1_BASE + SOP1Op.S_SEXT_I32_I8: lambda a, b, scc: (sext(a & 0xff, 8) & 0xffffffff, scc),
+  SOP1_BASE + SOP1Op.S_SEXT_I32_I16: lambda a, b, scc: (sext(a & 0xffff, 16) & 0xffffffff, scc),
+  SOP1_BASE + SOP1Op.S_ABS_I32: lambda a, b, scc: ((r := abs(sext(a, 32)) & 0xffffffff), int(r != 0)),
+  SOP1_BASE + SOP1Op.S_CVT_F32_I32: lambda a, b, scc: (i32(float(sext(a, 32))), scc),
+  SOP1_BASE + SOP1Op.S_CVT_F32_U32: lambda a, b, scc: (i32(float(a)), scc),
+  SOP1_BASE + SOP1Op.S_CVT_I32_F32: lambda a, b, scc: (_cvt_i32_f32(f32(a)), scc),
+  SOP1_BASE + SOP1Op.S_CVT_U32_F32: lambda a, b, scc: (_cvt_u32_f32(f32(a)), scc),
+  SOP1_BASE + SOP1Op.S_CEIL_F32: lambda a, b, scc: (i32(math.ceil(f32(a))), scc),
+  SOP1_BASE + SOP1Op.S_FLOOR_F32: lambda a, b, scc: (i32(math.floor(f32(a))), scc),
+  SOP1_BASE + SOP1Op.S_TRUNC_F32: lambda a, b, scc: (i32(math.trunc(f32(a))), scc),
+  SOP1_BASE + SOP1Op.S_RNDNE_F32: lambda a, b, scc: (i32(round(f32(a))), scc),
+  SOP1_BASE + SOP1Op.S_CVT_F16_F32: lambda a, b, scc: (i16(f32(a)), scc),
+  SOP1_BASE + SOP1Op.S_CVT_F32_F16: lambda a, b, scc: (i32(f16(a)), scc),
+  # SOPC
+  SOPC_BASE + SOPCOp.S_CMP_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 32))),
+  SOPC_BASE + SOPCOp.S_CMP_EQ_U32: lambda a, b, scc: (0, int(a == b)),
+  SOPC_BASE + SOPCOp.S_CMP_LG_U32: lambda a, b, scc: (0, int(a != b)),
+  SOPC_BASE + SOPCOp.S_CMP_GT_U32: lambda a, b, scc: (0, int(a > b)),
+  SOPC_BASE + SOPCOp.S_CMP_GE_U32: lambda a, b, scc: (0, int(a >= b)),
+  SOPC_BASE + SOPCOp.S_CMP_LT_U32: lambda a, b, scc: (0, int(a < b)),
+  SOPC_BASE + SOPCOp.S_CMP_LE_U32: lambda a, b, scc: (0, int(a <= b)),
+  SOPC_BASE + SOPCOp.S_BITCMP0_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) == 0)),
+  SOPC_BASE + SOPCOp.S_BITCMP1_B32: lambda a, b, scc: (0, int((a & (1 << (b & 0x1f))) != 0)),
+  # SOPK
+  SOPK_BASE + SOPKOp.S_MOVK_I32: lambda a, b, scc: (sext(b, 16) & 0xffffffff, scc),
+  SOPK_BASE + SOPKOp.S_CMOVK_I32: lambda a, b, scc: ((sext(b, 16) & 0xffffffff) if scc else a, scc),
+  SOPK_BASE + SOPKOp.S_ADDK_I32: lambda a, b, scc: ((r := sext(a, 32) + sext(b, 16)) & 0xffffffff, int(((a >> 31) == ((b >> 15) & 1)) and ((a >> 31) != ((r >> 31) & 1)))),
+  SOPK_BASE + SOPKOp.S_MULK_I32: lambda a, b, scc: ((sext(a, 32) * sext(b, 16)) & 0xffffffff, scc),
+  SOPK_BASE + SOPKOp.S_CMPK_EQ_I32: lambda a, b, scc: (0, int(sext(a, 32) == sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_LG_I32: lambda a, b, scc: (0, int(sext(a, 32) != sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_GT_I32: lambda a, b, scc: (0, int(sext(a, 32) > sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_GE_I32: lambda a, b, scc: (0, int(sext(a, 32) >= sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_LT_I32: lambda a, b, scc: (0, int(sext(a, 32) < sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_LE_I32: lambda a, b, scc: (0, int(sext(a, 32) <= sext(b, 16))),
+  SOPK_BASE + SOPKOp.S_CMPK_EQ_U32: lambda a, b, scc: (0, int(a == (b & 0xffff))),
+  SOPK_BASE + SOPKOp.S_CMPK_LG_U32: lambda a, b, scc: (0, int(a != (b & 0xffff))),
+  SOPK_BASE + SOPKOp.S_CMPK_GT_U32: lambda a, b, scc: (0, int(a > (b & 0xffff))),
+  SOPK_BASE + SOPKOp.S_CMPK_GE_U32: lambda a, b, scc: (0, int(a >= (b & 0xffff))),
+  SOPK_BASE + SOPKOp.S_CMPK_LT_U32: lambda a, b, scc: (0, int(a < (b & 0xffff))),
+  SOPK_BASE + SOPKOp.S_CMPK_LE_U32: lambda a, b, scc: (0, int(a <= (b & 0xffff))),
+}
+
+# VALU: op -> fn(s0, s1, s2) -> result
+VALU: dict[int, Callable] = {
+  # VOP2
+  VOP2_BASE + VOP2Op.V_ADD_F32: lambda a, b, c: i32(f32(a) + f32(b)),
+  VOP2_BASE + VOP2Op.V_SUB_F32: lambda a, b, c: i32(f32(a) - f32(b)),
+  VOP2_BASE + VOP2Op.V_SUBREV_F32: lambda a, b, c: i32(f32(b) - f32(a)),
+  VOP2_BASE + VOP2Op.V_MUL_F32: lambda a, b, c: i32(f32(a) * f32(b)),
+  VOP2_BASE + VOP2Op.V_MIN_F32: lambda a, b, c: i32(min(f32(a), f32(b))),
+  VOP2_BASE + VOP2Op.V_MAX_F32: lambda a, b, c: i32(max(f32(a), f32(b))),
+  VOP2_BASE + VOP2Op.V_ADD_NC_U32: lambda a, b, c: (a + b) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_SUB_NC_U32: lambda a, b, c: (a - b) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_SUBREV_NC_U32: lambda a, b, c: (b - a) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_AND_B32: lambda a, b, c: a & b,
+  VOP2_BASE + VOP2Op.V_OR_B32: lambda a, b, c: a | b,
+  VOP2_BASE + VOP2Op.V_XOR_B32: lambda a, b, c: a ^ b,
+  VOP2_BASE + VOP2Op.V_XNOR_B32: lambda a, b, c: (~(a ^ b)) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_LSHLREV_B32: lambda a, b, c: (b << (a & 0x1f)) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_LSHRREV_B32: lambda a, b, c: b >> (a & 0x1f),
+  VOP2_BASE + VOP2Op.V_ASHRREV_I32: lambda a, b, c: (sext(b, 32) >> (a & 0x1f)) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_MIN_I32: lambda a, b, c: a if sext(a, 32) < sext(b, 32) else b,
+  VOP2_BASE + VOP2Op.V_MAX_I32: lambda a, b, c: a if sext(a, 32) > sext(b, 32) else b,
+  VOP2_BASE + VOP2Op.V_MIN_U32: lambda a, b, c: min(a, b),
+  VOP2_BASE + VOP2Op.V_MAX_U32: lambda a, b, c: max(a, b),
+  VOP2_BASE + VOP2Op.V_MUL_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_MUL_HI_I32_I24: lambda a, b, c: ((sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24)) >> 32) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_MUL_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff)) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_MUL_HI_U32_U24: lambda a, b, c: (((a & 0xffffff) * (b & 0xffffff)) >> 32) & 0xffffffff,
+  VOP2_BASE + VOP2Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
+  VOP2_BASE + VOP2Op.V_LDEXP_F16: lambda a, b, c: i16(math.ldexp(f16(a), sext(b, 32))),
+  VOP2_BASE + VOP2Op.V_ADD_F16: lambda a, b, c: i16(f16(a) + f16(b)),
+  VOP2_BASE + VOP2Op.V_SUB_F16: lambda a, b, c: i16(f16(a) - f16(b)),
+  VOP2_BASE + VOP2Op.V_MUL_F16: lambda a, b, c: i16(f16(a) * f16(b)),
+  VOP2_BASE + VOP2Op.V_MIN_F16: lambda a, b, c: i16(min(f16(a), f16(b))),
+  VOP2_BASE + VOP2Op.V_MAX_F16: lambda a, b, c: i16(max(f16(a), f16(b))),
+  # VOP1
+  VOP1_BASE + VOP1Op.V_MOV_B32: lambda a, b, c: a,
+  VOP1_BASE + VOP1Op.V_NOT_B32: lambda a, b, c: (~a) & 0xffffffff,
+  VOP1_BASE + VOP1Op.V_BFREV_B32: lambda a, b, c: int(f'{a & 0xffffffff:032b}'[::-1], 2),
+  VOP1_BASE + VOP1Op.V_CLZ_I32_U32: lambda a, b, c: clz(a),
+  VOP1_BASE + VOP1Op.V_CLS_I32: lambda a, b, c: cls(a),
+  VOP1_BASE + VOP1Op.V_CVT_F32_I32: lambda a, b, c: i32(float(sext(a, 32))),
+  VOP1_BASE + VOP1Op.V_CVT_F32_U32: lambda a, b, c: i32(float(a)),
+  VOP1_BASE + VOP1Op.V_CVT_I32_F32: lambda a, b, c: _cvt_i32_f32(f32(a)),
+  VOP1_BASE + VOP1Op.V_CVT_U32_F32: lambda a, b, c: _cvt_u32_f32(f32(a)),
+  VOP1_BASE + VOP1Op.V_CVT_F16_F32: lambda a, b, c: i16(f32(a)),
+  VOP1_BASE + VOP1Op.V_CVT_F32_F16: lambda a, b, c: i32(f16(a)),
+  VOP1_BASE + VOP1Op.V_RCP_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
+  VOP1_BASE + VOP1Op.V_RCP_IFLAG_F32: lambda a, b, c: i32(1.0 / f32(a) if f32(a) != 0 else math.copysign(float('inf'), f32(a))),
+  VOP1_BASE + VOP1Op.V_RSQ_F32: lambda a, b, c: i32(1.0 / math.sqrt(f32(a)) if f32(a) > 0 else (float('nan') if f32(a) < 0 else float('inf'))),
+  VOP1_BASE + VOP1Op.V_SQRT_F32: lambda a, b, c: i32(math.sqrt(f32(a)) if f32(a) >= 0 else float('nan')),
+  VOP1_BASE + VOP1Op.V_LOG_F32: lambda a, b, c: i32(math.log2(f32(a)) if f32(a) > 0 else (float('-inf') if f32(a) == 0 else float('nan'))),
+  VOP1_BASE + VOP1Op.V_EXP_F32: lambda a, b, c: i32(float('inf') if f32(a) > 128 else (0.0 if f32(a) < -150 else math.pow(2.0, f32(a)))),
+  VOP1_BASE + VOP1Op.V_SIN_F32: lambda a, b, c: i32(math.sin(f32(a) * 2 * math.pi)),
+  VOP1_BASE + VOP1Op.V_COS_F32: lambda a, b, c: i32(math.cos(f32(a) * 2 * math.pi)),
+  VOP1_BASE + VOP1Op.V_FLOOR_F32: lambda a, b, c: i32(math.floor(f32(a))),
+  VOP1_BASE + VOP1Op.V_CEIL_F32: lambda a, b, c: i32(math.ceil(f32(a))),
+  VOP1_BASE + VOP1Op.V_TRUNC_F32: lambda a, b, c: i32(math.trunc(f32(a))),
+  VOP1_BASE + VOP1Op.V_RNDNE_F32: lambda a, b, c: i32(round(f32(a))),
+  VOP1_BASE + VOP1Op.V_FRACT_F32: lambda a, b, c: i32((v := f32(a)) - math.floor(v)),
+  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE0: lambda a, b, c: i32(float(a & 0xff)),
+  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE1: lambda a, b, c: i32(float((a >> 8) & 0xff)),
+  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE2: lambda a, b, c: i32(float((a >> 16) & 0xff)),
+  VOP1_BASE + VOP1Op.V_CVT_F32_UBYTE3: lambda a, b, c: i32(float((a >> 24) & 0xff)),
+  VOP1_BASE + VOP1Op.V_FREXP_MANT_F32: lambda a, b, c: i32(math.frexp(v)[0] if (v := f32(a)) != 0 else 0.0),
+  VOP1_BASE + VOP1Op.V_FREXP_EXP_I32_F32: lambda a, b, c: (math.frexp(v)[1] if (v := f32(a)) != 0 else 0) & 0xffffffff,
+  # VOP3
+  VOP3Op.V_FMA_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
+  VOP3Op.V_DIV_FMAS_F32: lambda a, b, c: i32(f32(a) * f32(b) + f32(c)),
+  VOP3Op.V_ADD3_U32: lambda a, b, c: (a + b + c) & 0xffffffff,
+  VOP3Op.V_LSHL_ADD_U32: lambda a, b, c: ((a << (b & 0x1f)) + c) & 0xffffffff,
+  VOP3Op.V_ADD_LSHL_U32: lambda a, b, c: ((a + b) << (c & 0x1f)) & 0xffffffff,
+  VOP3Op.V_XOR3_B32: lambda a, b, c: a ^ b ^ c,
+  VOP3Op.V_OR3_B32: lambda a, b, c: a | b | c,
+  VOP3Op.V_AND_OR_B32: lambda a, b, c: (a & b) | c,
+  VOP3Op.V_LSHL_OR_B32: lambda a, b, c: ((a << (b & 0x1f)) | c) & 0xffffffff,
+  VOP3Op.V_XAD_U32: lambda a, b, c: ((a ^ b) + c) & 0xffffffff,
+  VOP3Op.V_MAD_U32_U24: lambda a, b, c: ((a & 0xffffff) * (b & 0xffffff) + c) & 0xffffffff,
+  VOP3Op.V_MAD_I32_I24: lambda a, b, c: (sext(a & 0xffffff, 24) * sext(b & 0xffffff, 24) + sext(c, 32)) & 0xffffffff,
+  VOP3Op.V_BFE_U32: lambda a, b, c: (a >> (b & 0x1f)) & ((1 << (c & 0x1f)) - 1) if c & 0x1f else 0,
+  VOP3Op.V_BFE_I32: lambda a, b, c: sext((a >> (b & 0x1f)) & ((1 << w) - 1), w) & 0xffffffff if (w := c & 0x1f) else 0,
+  VOP3Op.V_ALIGNBIT_B32: lambda a, b, c: (((a << 32) | b) >> (c & 0x1f)) & 0xffffffff,
+  VOP3Op.V_MUL_LO_U32: lambda a, b, c: (a * b) & 0xffffffff,
+  VOP3Op.V_MUL_HI_U32: lambda a, b, c: ((a * b) >> 32) & 0xffffffff,
+  VOP3Op.V_MUL_HI_I32: lambda a, b, c: ((sext(a, 32) * sext(b, 32)) >> 32) & 0xffffffff,
+  VOP3Op.V_LDEXP_F32: lambda a, b, c: i32(math.ldexp(f32(a), sext(b, 32))),
+  VOP3Op.V_DIV_FIXUP_F32: lambda a, b, c: i32(math.copysign(float('inf'), f32(c)) if f32(b) == 0.0 else f32(c) / f32(b)),
+  VOP3Op.V_PACK_B32_F16: lambda a, b, c: (a & 0xffff) | ((b & 0xffff) << 16),
+  VOP3Op.V_CVT_PK_RTZ_F16_F32: lambda a, b, c: i16(f32(a)) | (i16(f32(b)) << 16),
+  VOP3Op.V_LSHLREV_B16: lambda a, b, c: ((b & 0xffff) << (a & 0xf)) & 0xffff,
+  VOP3Op.V_LSHRREV_B16: lambda a, b, c: (b & 0xffff) >> (a & 0xf),
+  VOP3Op.V_ASHRREV_I16: lambda a, b, c: (sext(b & 0xffff, 16) >> (a & 0xf)) & 0xffff,
+  VOP3Op.V_ADD_NC_U16: lambda a, b, c: ((a & 0xffff) + (b & 0xffff)) & 0xffff,
+  VOP3Op.V_SUB_NC_U16: lambda a, b, c: ((a & 0xffff) - (b & 0xffff)) & 0xffff,
+  VOP3Op.V_MUL_LO_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff)) & 0xffff,
+  VOP3Op.V_MIN_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff),
+  VOP3Op.V_MAX_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff),
+  VOP3Op.V_MIN_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) < sext(b & 0xffff, 16) else (b & 0xffff),
+  VOP3Op.V_MAX_I16: lambda a, b, c: (a & 0xffff) if sext(a & 0xffff, 16) > sext(b & 0xffff, 16) else (b & 0xffff),
+  VOP3Op.V_MAD_U16: lambda a, b, c: ((a & 0xffff) * (b & 0xffff) + (c & 0xffff)) & 0xffff,
+  VOP3Op.V_MAD_I16: lambda a, b, c: (sext(a & 0xffff, 16) * sext(b & 0xffff, 16) + sext(c & 0xffff, 16)) & 0xffff,
+  VOP3Op.V_FMA_F16: lambda a, b, c: i16(f16(a) * f16(b) + f16(c)),
+  VOP3Op.V_MIN3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[0] & 0xffffffff,
+  VOP3Op.V_MAX3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[2] & 0xffffffff,
+  VOP3Op.V_MED3_I32: lambda a, b, c: sorted([sext(a, 32), sext(b, 32), sext(c, 32)])[1] & 0xffffffff,
+  VOP3Op.V_MIN3_F16: lambda a, b, c: i16(min(f16(a), f16(b), f16(c))),
+  VOP3Op.V_MAX3_F16: lambda a, b, c: i16(max(f16(a), f16(b), f16(c))),
+  VOP3Op.V_MED3_F16: lambda a, b, c: i16(sorted([f16(a), f16(b), f16(c)])[1]),
+  VOP3Op.V_MIN3_U16: lambda a, b, c: min(a & 0xffff, b & 0xffff, c & 0xffff),
+  VOP3Op.V_MAX3_U16: lambda a, b, c: max(a & 0xffff, b & 0xffff, c & 0xffff),
+  VOP3Op.V_MED3_U16: lambda a, b, c: sorted([a & 0xffff, b & 0xffff, c & 0xffff])[1],
+  VOP3Op.V_MIN3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[0] & 0xffff,
+  VOP3Op.V_MAX3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[2] & 0xffff,
+  VOP3Op.V_MED3_I16: lambda a, b, c: sorted([sext(a & 0xffff, 16), sext(b & 0xffff, 16), sext(c & 0xffff, 16)])[1] & 0xffff,
+}
+
+def _cmp8(a, b): return [False, a < b, a == b, a <= b, a > b, a != b, a >= b, True]
+def _cmp6(a, b): return [a < b, a == b, a <= b, a > b, a != b, a >= b]
+
+def vopc(op: int, s0: int, s1: int, s0_hi: int = 0, s1_hi: int = 0) -> int:
+  base = op & 0x7f
+  if 16 <= base <= 31:  # F32
+    f0, f1, cmp, nan = f32(s0), f32(s1), base - 16, math.isnan(f32(s0)) or math.isnan(f32(s1))
+    return int([False, f0<f1, f0==f1, f0<=f1, f0>f1, f0!=f1, f0>=f1, not nan, nan, f0<f1 or nan, f0==f1 or nan, f0<=f1 or nan, f0>f1 or nan, f0!=f1 or nan, f0>=f1 or nan, True][cmp])
+  if 49 <= base <= 54: return int(_cmp6(sext(s0 & 0xffff, 16), sext(s1 & 0xffff, 16))[base - 49])  # I16
+  if 57 <= base <= 62: return int(_cmp6(s0 & 0xffff, s1 & 0xffff)[base - 57])  # U16
+  if 64 <= base <= 79:  # I32/U32
+    cmp = (base - 64) % 8
+    return int(_cmp8(sext(s0, 32), sext(s1, 32))[cmp] if base < 72 else _cmp8(s0, s1)[cmp])
+  if 80 <= base <= 95:  # I64/U64
+    s0_64, s1_64 = s0 | (s0_hi << 32), s1 | (s1_hi << 32)
+    return int(_cmp8(sext(s0_64, 64), sext(s1_64, 64))[(base - 80) % 8] if base < 88 else _cmp8(s0_64, s1_64)[(base - 80) % 8])
+  if base == 126:  # CLASS_F32
+    f, mask = f32(s0), s1
+    if math.isnan(f): return int(bool(mask & 0x3))
+    if math.isinf(f): return int(bool(mask & (0x4 if f < 0 else 0x200)))
+    if f == 0.0: return int(bool(mask & (0x20 if (s0 >> 31) & 1 else 0x40)))
+    exp, sign = (s0 >> 23) & 0xff, (s0 >> 31) & 1
+    return int(bool(mask & ((0x10 if sign else 0x80) if exp == 0 else (0x8 if sign else 0x100))))
+  raise NotImplementedError(f"VOPC op {op} (base {base})")

extra/assembly/rdna3/asm.py

@@ -0,0 +1,602 @@
+# RDNA3 assembler and disassembler
+from __future__ import annotations
+import re
+from extra.assembly.rdna3.lib import Inst, RawImm, Reg, SGPR, VGPR, TTMP, s, v, ttmp, _RegFactory, FLOAT_ENC, SRC_FIELDS, unwrap
+
+# Decoding helpers
+SPECIAL_GPRS = {106: "vcc_lo", 107: "vcc_hi", 124: "null", 125: "m0", 126: "exec_lo", 127: "exec_hi", 253: "scc"}
+SPECIAL_DEC = {**SPECIAL_GPRS, **{v: str(k) for k, v in FLOAT_ENC.items()}}
+SPECIAL_PAIRS = {106: "vcc", 126: "exec"}  # Special register pairs (for 64-bit ops)
+# GFX11 hwreg names (IDs 16-17 are TBA - not supported, IDs 18-19 are PERF_SNAPSHOT)
+HWREG_NAMES = {1: 'HW_REG_MODE', 2: 'HW_REG_STATUS', 3: 'HW_REG_TRAPSTS', 4: 'HW_REG_HW_ID', 5: 'HW_REG_GPR_ALLOC',
+               6: 'HW_REG_LDS_ALLOC', 7: 'HW_REG_IB_STS', 15: 'HW_REG_SH_MEM_BASES', 18: 'HW_REG_PERF_SNAPSHOT_PC_LO',
+               19: 'HW_REG_PERF_SNAPSHOT_PC_HI', 20: 'HW_REG_FLAT_SCR_LO', 21: 'HW_REG_FLAT_SCR_HI',
+               22: 'HW_REG_XNACK_MASK', 23: 'HW_REG_HW_ID1', 24: 'HW_REG_HW_ID2', 25: 'HW_REG_POPS_PACKER', 28: 'HW_REG_IB_STS2'}
+HWREG_IDS = {v.lower(): k for k, v in HWREG_NAMES.items()}  # Reverse map for assembler
+MSG_NAMES = {128: 'MSG_RTN_GET_DOORBELL', 129: 'MSG_RTN_GET_DDID', 130: 'MSG_RTN_GET_TMA',
+             131: 'MSG_RTN_GET_REALTIME', 132: 'MSG_RTN_SAVE_WAVE', 133: 'MSG_RTN_GET_TBA'}
+_16BIT_TYPES = ('f16', 'i16', 'u16', 'b16')
+def _is_16bit(s: str) -> bool: return any(s.endswith(x) for x in _16BIT_TYPES)
+
+def decode_src(val: int) -> str:
+  if val <= 105: return f"s{val}"
+  if val in SPECIAL_DEC: return SPECIAL_DEC[val]
+  if 108 <= val <= 123: return f"ttmp{val - 108}"
+  if 128 <= val <= 192: return str(val - 128)
+  if 193 <= val <= 208: return str(-(val - 192))
+  if 256 <= val <= 511: return f"v{val - 256}"
+  return "lit" if val == 255 else f"?{val}"
+
+def _reg(prefix: str, base: int, cnt: int = 1) -> str: return f"{prefix}{base}" if cnt == 1 else f"{prefix}[{base}:{base+cnt-1}]"
+def _sreg(base: int, cnt: int = 1) -> str: return _reg("s", base, cnt)
+def _vreg(base: int, cnt: int = 1) -> str: return _reg("v", base, cnt)
+
+def _fmt_sdst(v: int, cnt: int = 1) -> str:
+  """Format SGPR destination with special register names."""
+  if v == 124: return "null"
+  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
+  if cnt > 1 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
+  if cnt > 1: return _sreg(v, cnt)
+  return {126: "exec_lo", 127: "exec_hi", 106: "vcc_lo", 107: "vcc_hi", 125: "m0"}.get(v, f"s{v}")
+
+def _fmt_ssrc(v: int, cnt: int = 1) -> str:
+  """Format SGPR source with special register names and pairs."""
+  if cnt == 2:
+    if v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
+    if v <= 105: return _sreg(v, 2)
+    if 108 <= v <= 123: return _reg("ttmp", v - 108, 2)
+  return decode_src(v)
+
+def _fmt_src_n(v: int, cnt: int) -> str:
+  """Format source with given register count (1, 2, or 4)."""
+  if cnt == 1: return decode_src(v)
+  if v >= 256: return _vreg(v - 256, cnt)
+  if v <= 105: return _sreg(v, cnt)
+  if cnt == 2 and v in SPECIAL_PAIRS: return SPECIAL_PAIRS[v]
+  if 108 <= v <= 123: return _reg("ttmp", v - 108, cnt)
+  return decode_src(v)
+
+def _fmt_src64(v: int) -> str:
+  """Format 64-bit source (VGPR pair, SGPR pair, or special pair)."""
+  return _fmt_src_n(v, 2)
+
+def _parse_sop_sizes(op_name: str) -> tuple[int, ...]:
+  """Parse dst and src sizes from SOP instruction name. Returns (dst_cnt, src0_cnt) or (dst_cnt, src0_cnt, src1_cnt)."""
+  if op_name in ('s_bitset0_b64', 's_bitset1_b64'): return (2, 1)
+  if op_name in ('s_lshl_b64', 's_lshr_b64', 's_ashr_i64', 's_bfe_u64', 's_bfe_i64'): return (2, 2, 1)
+  if op_name in ('s_bfm_b64',): return (2, 1, 1)
+  # SOPC: s_bitcmp0_b64, s_bitcmp1_b64 - 64-bit src0, 32-bit src1 (bit index)
+  if op_name in ('s_bitcmp0_b64', 's_bitcmp1_b64'): return (1, 2, 1)
+  if m := re.search(r'_(b|i|u)(32|64)_(b|i|u)(32|64)$', op_name):
+    return (2 if m.group(2) == '64' else 1, 2 if m.group(4) == '64' else 1)
+  if m := re.search(r'_(b|i|u)(32|64)$', op_name):
+    sz = 2 if m.group(2) == '64' else 1
+    return (sz, sz)
+  return (1, 1)
+
+# Waitcnt helpers (RDNA3 format: bits 15:10=vmcnt, bits 9:4=lgkmcnt, bits 3:0=expcnt)
+def waitcnt(vmcnt: int = 0x3f, expcnt: int = 0x7, lgkmcnt: int = 0x3f) -> int:
+  return (expcnt & 0x7) | ((lgkmcnt & 0x3f) << 4) | ((vmcnt & 0x3f) << 10)
+def decode_waitcnt(val: int) -> tuple[int, int, int]:
+  return (val >> 10) & 0x3f, val & 0xf, (val >> 4) & 0x3f  # vmcnt, expcnt, lgkmcnt
+
+# VOP3SD opcodes (shared encoding with VOP3 but different field layout)
+# Note: opcodes 0-255 are VOPC promoted to VOP3 - never treat as VOP3SD
+VOP3SD_OPCODES = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
+
+# Disassembler
+def disasm(inst: Inst) -> str:
+  op_val = unwrap(inst._values.get('op', 0))
+  cls_name = inst.__class__.__name__
+  # VOP3 and VOP3SD share encoding - check opcode to determine which
+  is_vop3sd = cls_name == 'VOP3' and op_val in VOP3SD_OPCODES
+  try:
+    from extra.assembly.rdna3 import autogen
+    if is_vop3sd:
+      op_name = autogen.VOP3SDOp(op_val).name.lower()
+    else:
+      op_name = getattr(autogen, f"{cls_name}Op")(op_val).name.lower() if hasattr(autogen, f"{cls_name}Op") else f"op_{op_val}"
+  except (ValueError, KeyError): op_name = f"op_{op_val}"
+  def fmt_src(v): return f"0x{inst._literal:x}" if v == 255 and getattr(inst, '_literal', None) else decode_src(v)
+
+  # VOP1
+  if cls_name == 'VOP1':
+    vdst, src0 = unwrap(inst._values['vdst']), unwrap(inst._values['src0'])
+    if op_name == 'v_nop': return 'v_nop'
+    if op_name == 'v_pipeflush': return 'v_pipeflush'
+    parts = op_name.split('_')
+    is_16bit_dst = any(p in _16BIT_TYPES for p in parts[-2:-1]) or (len(parts) >= 2 and parts[-1] in _16BIT_TYPES and 'cvt' not in op_name)
+    is_16bit_src = parts[-1] in _16BIT_TYPES and 'sat_pk' not in op_name
+    _F64_OPS = ('v_ceil_f64', 'v_floor_f64', 'v_fract_f64', 'v_frexp_mant_f64', 'v_rcp_f64', 'v_rndne_f64', 'v_rsq_f64', 'v_sqrt_f64', 'v_trunc_f64')
+    is_f64_dst = op_name in _F64_OPS or op_name in ('v_cvt_f64_f32', 'v_cvt_f64_i32', 'v_cvt_f64_u32')
+    is_f64_src = op_name in _F64_OPS or op_name in ('v_cvt_f32_f64', 'v_cvt_i32_f64', 'v_cvt_u32_f64', 'v_frexp_exp_i32_f64')
+    if op_name == 'v_readfirstlane_b32':
+      return f"v_readfirstlane_b32 {decode_src(vdst)}, v{src0 - 256 if src0 >= 256 else src0}"
+    dst_str = _vreg(vdst, 2) if is_f64_dst else f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}" if is_16bit_dst else f"v{vdst}"
+    src_str = _fmt_src64(src0) if is_f64_src else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit_src and src0 >= 256 else fmt_src(src0)
+    return f"{op_name}_e32 {dst_str}, {src_str}"
+
+  # VOP2
+  if cls_name == 'VOP2':
+    vdst, src0_raw, vsrc1 = unwrap(inst._values['vdst']), unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
+    suffix = "" if op_name == "v_dot2acc_f32_f16" else "_e32"
+    is_16bit_op = ('_f16' in op_name or '_i16' in op_name or '_u16' in op_name) and '_f32' not in op_name and '_i32' not in op_name and 'pk_' not in op_name
+    if is_16bit_op:
+      dst_str = f"v{vdst & 0x7f}.{'h' if vdst >= 128 else 'l'}"
+      src0_str = f"v{(src0_raw - 256) & 0x7f}.{'h' if src0_raw >= 384 else 'l'}" if src0_raw >= 256 else fmt_src(src0_raw)
+      vsrc1_str = f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}"
+    else:
+      dst_str, src0_str, vsrc1_str = f"v{vdst}", fmt_src(src0_raw), f"v{vsrc1}"
+    return f"{op_name}{suffix} {dst_str}, {src0_str}, {vsrc1_str}" + (", vcc_lo" if op_name == "v_cndmask_b32" else "")
+
+  # VOPC
+  if cls_name == 'VOPC':
+    src0, vsrc1 = unwrap(inst._values['src0']), unwrap(inst._values['vsrc1'])
+    is_64bit = any(x in op_name for x in ('f64', 'i64', 'u64'))
+    is_64bit_vsrc1 = is_64bit and 'class' not in op_name
+    is_16bit = any(x in op_name for x in ('_f16', '_i16', '_u16')) and 'f32' not in op_name
+    is_cmpx = op_name.startswith('v_cmpx')  # VOPCX writes to exec, no vcc destination
+    src0_str = _fmt_src64(src0) if is_64bit else f"v{(src0 - 256) & 0x7f}.{'h' if src0 >= 384 else 'l'}" if is_16bit and src0 >= 256 else fmt_src(src0)
+    vsrc1_str = _vreg(vsrc1, 2) if is_64bit_vsrc1 else f"v{vsrc1 & 0x7f}.{'h' if vsrc1 >= 128 else 'l'}" if is_16bit else f"v{vsrc1}"
+    return f"{op_name}_e32 {src0_str}, {vsrc1_str}" if is_cmpx else f"{op_name}_e32 vcc_lo, {src0_str}, {vsrc1_str}"
+
+  # SOPP
+  if cls_name == 'SOPP':
+    simm16 = unwrap(inst._values.get('simm16', 0))
+    # No-operand instructions (simm16 is ignored)
+    no_imm_ops = ('s_endpgm', 's_barrier', 's_wakeup', 's_icache_inv', 's_ttracedata', 's_ttracedata_imm',
+                  's_wait_idle', 's_endpgm_saved', 's_code_end', 's_endpgm_ordered_ps_done')
+    if op_name in no_imm_ops: return op_name
+    if op_name == 's_waitcnt':
+      vmcnt, expcnt, lgkmcnt = decode_waitcnt(simm16)
+      parts = []
+      if vmcnt != 0x3f: parts.append(f"vmcnt({vmcnt})")
+      if expcnt != 0x7: parts.append(f"expcnt({expcnt})")
+      if lgkmcnt != 0x3f: parts.append(f"lgkmcnt({lgkmcnt})")
+      return f"s_waitcnt {' '.join(parts)}" if parts else "s_waitcnt 0"
+    if op_name == 's_delay_alu':
+      dep_names = ['VALU_DEP_1','VALU_DEP_2','VALU_DEP_3','VALU_DEP_4','TRANS32_DEP_1','TRANS32_DEP_2','TRANS32_DEP_3','FMA_ACCUM_CYCLE_1','SALU_CYCLE_1','SALU_CYCLE_2','SALU_CYCLE_3']
+      skip_names = ['SAME','NEXT','SKIP_1','SKIP_2','SKIP_3','SKIP_4']
+      id0, skip, id1 = simm16 & 0xf, (simm16 >> 4) & 0x7, (simm16 >> 7) & 0xf
+      def dep_name(v): return dep_names[v-1] if 0 < v <= len(dep_names) else str(v)
+      parts = [f"instid0({dep_name(id0)})"] if id0 else []
+      if skip: parts.append(f"instskip({skip_names[skip]})")
+      if id1: parts.append(f"instid1({dep_name(id1)})")
+      return f"s_delay_alu {' | '.join(p for p in parts if p)}" if parts else "s_delay_alu 0"
+    if op_name.startswith('s_cbranch') or op_name.startswith('s_branch'):
+      return f"{op_name} {simm16}"
+    # Most SOPP ops require immediate (s_nop, s_setkill, s_sethalt, s_sleep, s_setprio, s_sendmsg*, etc.)
+    return f"{op_name} 0x{simm16:x}"
+
+  # SMEM
+  if cls_name == 'SMEM':
+    if op_name in ('s_gl1_inv', 's_dcache_inv'): return op_name
+    sdata, sbase, soffset, offset = unwrap(inst._values['sdata']), unwrap(inst._values['sbase']), unwrap(inst._values['soffset']), unwrap(inst._values.get('offset', 0))
+    glc, dlc = unwrap(inst._values.get('glc', 0)), unwrap(inst._values.get('dlc', 0))
+    # Format offset: "soffset offset:X" if both, "0x{offset:x}" if only imm, or decode_src(soffset)
+    off_str = f"{decode_src(soffset)} offset:0x{offset:x}" if offset and soffset != 124 else f"0x{offset:x}" if offset else decode_src(soffset)
+    sbase_idx, sbase_cnt = sbase * 2, 4 if (8 <= op_val <= 12 or op_name == 's_atc_probe_buffer') else 2
+    sbase_str = _fmt_ssrc(sbase_idx, sbase_cnt) if sbase_cnt == 2 else _sreg(sbase_idx, sbase_cnt) if sbase_idx <= 105 else _reg("ttmp", sbase_idx - 108, sbase_cnt)
+    if op_name in ('s_atc_probe', 's_atc_probe_buffer'): return f"{op_name} {sdata}, {sbase_str}, {off_str}"
+    width = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val, 1)
+    mods = [m for m in ["glc" if glc else "", "dlc" if dlc else ""] if m]
+    return f"{op_name} {_fmt_sdst(sdata, width)}, {sbase_str}, {off_str}" + (" " + " ".join(mods) if mods else "")
+
+  # FLAT
+  if cls_name == 'FLAT':
+    vdst, addr, data, saddr, offset, seg = [unwrap(inst._values.get(f, 0)) for f in ['vdst', 'addr', 'data', 'saddr', 'offset', 'seg']]
+    instr = f"{['flat', 'scratch', 'global'][seg] if seg < 3 else 'flat'}_{op_name.split('_', 1)[1] if '_' in op_name else op_name}"
+    width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'u8':1, 'i8':1, 'u16':1, 'i16':1}.get(op_name.split('_')[-1], 1)
+    addr_str = _vreg(addr, 2) if saddr == 0x7F else _vreg(addr)
+    saddr_str = "" if saddr == 0x7F else f", {_sreg(saddr, 2)}" if saddr < 106 else ", off" if saddr == 124 else f", {decode_src(saddr)}"
+    off_str = f" offset:{offset}" if offset else ""
+    vdata_str = _vreg(data if 'store' in op_name else vdst, width)
+    return f"{instr} {addr_str}, {vdata_str}{saddr_str}{off_str}" if 'store' in op_name else f"{instr} {vdata_str}, {addr_str}{saddr_str}{off_str}"
+
+  # VOP3: vector ops with modifiers (can be 1, 2, or 3 sources depending on opcode range)
+  if cls_name == 'VOP3':
+    # Handle VOP3SD opcodes (same encoding, different field layout)
+    if is_vop3sd:
+      vdst = unwrap(inst._values.get('vdst', 0))
+      # VOP3SD: sdst is at bits [14:8], but VOP3 decodes opsel at [14:11], abs at [10:8], clmp at [15]
+      # We need to reconstruct sdst from these fields
+      opsel_raw = unwrap(inst._values.get('opsel', 0))
+      abs_raw = unwrap(inst._values.get('abs', 0))
+      clmp_raw = unwrap(inst._values.get('clmp', 0))
+      sdst = (clmp_raw << 7) | (opsel_raw << 3) | abs_raw
+      src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
+      neg = unwrap(inst._values.get('neg', 0))
+      omod = unwrap(inst._values.get('omod', 0))
+      omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
+      is_f64 = 'f64' in op_name
+      # v_mad_i64_i32/v_mad_u64_u32: 64-bit dst and src2, 32-bit src0/src1
+      is_mad64 = 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
+      def fmt_sd_src(v, neg_bit, is_64bit=False):
+        s = _fmt_src64(v) if (is_64bit or is_f64) else fmt_src(v)
+        return f"-{s}" if neg_bit else s
+      src0_str, src1_str = fmt_sd_src(src0, neg & 1), fmt_sd_src(src1, neg & 2)
+      src2_str = fmt_sd_src(src2, neg & 4, is_mad64)
+      dst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}"
+      sdst_str = _fmt_sdst(sdst, 1)
+      # v_add_co_u32, v_sub_co_u32, v_subrev_co_u32, v_add_co_ci_u32, etc. only use 2 sources
+      if op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32', 'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32'):
+        return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}"
+      # v_div_scale uses 3 sources
+      return f"{op_name} {dst_str}, {sdst_str}, {src0_str}, {src1_str}, {src2_str}" + omod_str
+
+    vdst = unwrap(inst._values.get('vdst', 0))
+    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
+    neg, abs_, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('abs', 0)), unwrap(inst._values.get('clmp', 0))
+    opsel = unwrap(inst._values.get('opsel', 0))
+    # Check if 64-bit op (needs register pairs)
+    is_f64 = 'f64' in op_name or 'i64' in op_name or 'u64' in op_name or 'b64' in op_name
+    # v_cmp_class_* has 64-bit src0 but 32-bit src1 (class mask)
+    is_class = 'class' in op_name
+    # Shift ops: v_*rev_*64 have 32-bit shift amount (src0), 64-bit value (src1)
+    is_shift64 = 'rev' in op_name and '64' in op_name and op_name.startswith('v_')
+    # v_ldexp_f64: 64-bit src0 (mantissa), 32-bit src1 (exponent)
+    is_ldexp64 = op_name == 'v_ldexp_f64'
+    # v_trig_preop_f64: 64-bit dst/src0, 32-bit src1 (exponent/scale)
+    is_trig_preop = op_name == 'v_trig_preop_f64'
+    # v_readlane_b32: destination is SGPR (despite vdst field)
+    is_readlane = op_name == 'v_readlane_b32'
+    # SAD/QSAD/MQSAD instructions have mixed sizes
+    # v_qsad_pk_u16_u8, v_mqsad_pk_u16_u8: 64-bit dst/src0/src2, 32-bit src1
+    # v_mqsad_u32_u8: 128-bit (4 reg) dst/src2, 64-bit src0, 32-bit src1
+    is_sad64 = any(x in op_name for x in ('qsad_pk', 'mqsad_pk'))
+    is_mqsad_u32 = 'mqsad_u32' in op_name
+    # Detect 16-bit and 64-bit operand sizes for various instruction patterns
+    if 'cvt_pk' in op_name:
+      is_f16_dst, is_f16_src, is_f16_src2 = False, op_name.endswith('16'), False
+    elif m := re.match(r'v_(?:cvt|frexp_exp)_([a-z0-9_]+)_([a-z0-9]+)', op_name):
+      dst_type, src_type = m.group(1), m.group(2)
+      is_f16_dst, is_f16_src, is_f16_src2 = _is_16bit(dst_type), _is_16bit(src_type), _is_16bit(src_type)
+      is_f64_dst, is_f64_src, is_f64 = '64' in dst_type, '64' in src_type, False
+    elif re.match(r'v_mad_[iu]32_[iu]16', op_name):
+      is_f16_dst, is_f16_src, is_f16_src2 = False, True, False  # 32-bit dst, 16-bit src0/src1, 32-bit src2
+    elif 'pack_b32' in op_name:
+      is_f16_dst, is_f16_src, is_f16_src2 = False, True, True  # 32-bit dst, 16-bit sources
+    else:
+      is_16bit_op = any(x in op_name for x in _16BIT_TYPES) and not any(x in op_name for x in ('dot2', 'pk_', 'sad', 'msad', 'qsad', 'mqsad'))
+      is_f16_dst = is_f16_src = is_f16_src2 = is_16bit_op
+    # Check if any opsel bit is set (any operand uses .h) - if so, we need explicit .l for low-half
+    any_hi = opsel != 0
+    def fmt_vop3_src(v, neg_bit, abs_bit, hi_bit=False, reg_cnt=1, is_16=False):
+      s = _fmt_src_n(v, reg_cnt) if reg_cnt > 1 else f"v{v - 256}.h" if is_16 and v >= 256 and hi_bit else f"v{v - 256}.l" if is_16 and v >= 256 and any_hi else fmt_src(v)
+      if abs_bit: s = f"|{s}|"
+      return f"-{s}" if neg_bit else s
+    # Determine register count for each source (check for cvt-specific 64-bit flags first)
+    is_src0_64 = locals().get('is_f64_src', is_f64 and not is_shift64) or is_sad64 or is_mqsad_u32
+    is_src1_64 = is_f64 and not is_class and not is_ldexp64 and not is_trig_preop
+    src0_cnt = 2 if is_src0_64 else 1
+    src1_cnt = 2 if is_src1_64 else 1
+    src2_cnt = 4 if is_mqsad_u32 else 2 if (is_f64 or is_sad64) else 1
+    src0_str = fmt_vop3_src(src0, neg & 1, abs_ & 1, opsel & 1, src0_cnt, is_f16_src)
+    src1_str = fmt_vop3_src(src1, neg & 2, abs_ & 2, opsel & 2, src1_cnt, is_f16_src)
+    src2_str = fmt_vop3_src(src2, neg & 4, abs_ & 4, opsel & 4, src2_cnt, is_f16_src2)
+    # Format destination - for 16-bit ops, use .h/.l suffix; readlane uses SGPR dest
+    is_dst_64 = locals().get('is_f64_dst', is_f64) or is_sad64
+    dst_cnt = 4 if is_mqsad_u32 else 2 if is_dst_64 else 1
+    if is_readlane:
+      dst_str = _fmt_sdst(vdst, 1)
+    elif dst_cnt > 1:
+      dst_str = _vreg(vdst, dst_cnt)
+    elif is_f16_dst:
+      dst_str = f"v{vdst}.h" if (opsel & 8) else f"v{vdst}.l" if any_hi else f"v{vdst}"
+    else:
+      dst_str = f"v{vdst}"
+    clamp_str = " clamp" if clmp else ""
+    omod = unwrap(inst._values.get('omod', 0))
+    omod_str = {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
+    # op_sel for non-VGPR sources (when opsel bits are set but source is not a VGPR)
+    # For 16-bit ops with VGPR sources, opsel is encoded in .h/.l suffix
+    # For non-VGPR sources or non-16-bit ops, we need explicit op_sel
+    has_nonvgpr_opsel = (src0 < 256 and (opsel & 1)) or (src1 < 256 and (opsel & 2)) or (src2 < 256 and (opsel & 4))
+    need_opsel = has_nonvgpr_opsel or (opsel and not is_f16_src)
+    # Helper to format opsel string based on source count
+    def fmt_opsel(num_src):
+      if not need_opsel: return ""
+      # When dst is .h (for 16-bit ops) and non-VGPR sources have opsel, use all 1s
+      if is_f16_dst and (opsel & 8):  # dst is .h
+        return f" op_sel:[1,1,1{',1' if num_src == 3 else ''}]"
+      # Otherwise output actual opsel values
+      if num_src == 3:
+        return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1},{(opsel >> 3) & 1}]"
+      return f" op_sel:[{opsel & 1},{(opsel >> 1) & 1},{(opsel >> 2) & 1}]"
+    # Determine number of sources based on opcode range:
+    # 0-255: VOPC promoted (comparison, 2 src, sdst)
+    # 256-383: VOP2 promoted (2 src)
+    # 384-511: VOP1 promoted (1 src)
+    # 512+: Native VOP3 (2 or 3 src depending on instruction)
+    if op_val < 256:  # VOPC promoted
+      # VOPCX (v_cmpx_*) writes to exec, no explicit destination
+      if op_name.startswith('v_cmpx'):
+        return f"{op_name}_e64 {src0_str}, {src1_str}"
+      return f"{op_name}_e64 {_fmt_sdst(vdst, 1)}, {src0_str}, {src1_str}"
+    elif op_val < 384:  # VOP2 promoted
+      # v_cndmask_b32 in VOP3 format has 3 sources (src2 is mask selector)
+      if 'cndmask' in op_name:
+        return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
+      return f"{op_name}_e64 {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
+    elif op_val < 512:  # VOP1 promoted
+      if op_name in ('v_nop', 'v_pipeflush'): return f"{op_name}_e64"
+      return f"{op_name}_e64 {dst_str}, {src0_str}" + fmt_opsel(1) + clamp_str + omod_str
+    else:  # Native VOP3 - determine 2 vs 3 sources based on instruction name
+      # 3-source ops: fma, mad, min3, max3, med3, div_fixup, div_fmas, sad, msad, qsad, mqsad, lerp, alignbit/byte, cubeid/sc/tc/ma, bfe, bfi, perm_b32, permlane, cndmask
+      # Note: v_writelane_b32 is 2-src (src0, src1 with vdst as 3rd operand - read-modify-write)
+      is_3src = any(x in op_name for x in ('fma', 'mad', 'min3', 'max3', 'med3', 'div_fix', 'div_fmas', 'sad', 'lerp', 'align', 'cube',
+                                            'bfe', 'bfi', 'perm_b32', 'permlane', 'cndmask', 'xor3', 'or3', 'add3', 'lshl_or', 'and_or', 'lshl_add',
+                                            'add_lshl', 'xad', 'maxmin', 'minmax', 'dot2', 'cvt_pk_u8', 'mullit'))
+      if is_3src:
+        return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}" + fmt_opsel(3) + clamp_str + omod_str
+      return f"{op_name} {dst_str}, {src0_str}, {src1_str}" + fmt_opsel(2) + clamp_str + omod_str
+
+  # VOP3SD: 3-source with scalar destination (v_div_scale_*, v_add_co_u32, v_mad_*64_*32, etc.)
+  if cls_name == 'VOP3SD':
+    vdst, sdst = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('sdst', 0))
+    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
+    neg, omod, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('omod', 0)), unwrap(inst._values.get('clmp', 0))
+    is_f64, is_mad64 = 'f64' in op_name, 'mad_i64_i32' in op_name or 'mad_u64_u32' in op_name
+    def fmt_neg(v, neg_bit, is_64=False): return f"-{_fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)}" if neg_bit else _fmt_src64(v) if (is_64 or is_f64) else fmt_src(v)
+    srcs = [fmt_neg(src0, neg & 1), fmt_neg(src1, neg & 2), fmt_neg(src2, neg & 4, is_mad64)]
+    dst_str, sdst_str = _vreg(vdst, 2) if (is_f64 or is_mad64) else f"v{vdst}", _fmt_sdst(sdst, 1)
+    clamp_str, omod_str = " clamp" if clmp else "", {1: " mul:2", 2: " mul:4", 3: " div:2"}.get(omod, "")
+    is_2src = op_name in ('v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32')
+    suffix = "_e64" if op_name.startswith('v_') and 'co_' in op_name else ""
+    return f"{op_name}{suffix} {dst_str}, {sdst_str}, {', '.join(srcs[:2] if is_2src else srcs)}" + clamp_str + omod_str
+
+  # VOPD: dual-issue instructions
+  if cls_name == 'VOPD':
+    from extra.assembly.rdna3 import autogen
+    opx, opy, vdstx, vdsty_enc = [unwrap(inst._values.get(f, 0)) for f in ('opx', 'opy', 'vdstx', 'vdsty')]
+    srcx0, vsrcx1, srcy0, vsrcy1 = [unwrap(inst._values.get(f, 0)) for f in ('srcx0', 'vsrcx1', 'srcy0', 'vsrcy1')]
+    vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1)  # Decode vdsty
+    def fmt_vopd(op, vdst, src0, vsrc1):
+      try: name = autogen.VOPDOp(op).name.lower()
+      except (ValueError, KeyError): name = f"op_{op}"
+      return f"{name} v{vdst}, {fmt_src(src0)}" if 'mov' in name else f"{name} v{vdst}, {fmt_src(src0)}, v{vsrc1}"
+    return f"{fmt_vopd(opx, vdstx, srcx0, vsrcx1)} :: {fmt_vopd(opy, vdsty, srcy0, vsrcy1)}"
+
+  # VOP3P: packed vector ops
+  if cls_name == 'VOP3P':
+    vdst, clmp = unwrap(inst._values.get('vdst', 0)), unwrap(inst._values.get('clmp', 0))
+    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
+    neg, neg_hi = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('neg_hi', 0))
+    opsel, opsel_hi, opsel_hi2 = unwrap(inst._values.get('opsel', 0)), unwrap(inst._values.get('opsel_hi', 0)), unwrap(inst._values.get('opsel_hi2', 0))
+    is_wmma, is_3src = 'wmma' in op_name, any(x in op_name for x in ('fma', 'mad', 'dot', 'wmma'))
+    def fmt_bits(name, val, n): return f"{name}:[{','.join(str((val >> i) & 1) for i in range(n))}]"
+    # WMMA: f16/bf16 use 8-reg sources, iu8 uses 4-reg, iu4 uses 2-reg; all have 8-reg dst
+    if is_wmma:
+      src_cnt = 2 if 'iu4' in op_name else 4 if 'iu8' in op_name else 8
+      src0_str, src1_str, src2_str = _fmt_src_n(src0, src_cnt), _fmt_src_n(src1, src_cnt), _fmt_src_n(src2, 8)
+      dst_str = _vreg(vdst, 8)
+    else:
+      src0_str, src1_str, src2_str = _fmt_src_n(src0, 1), _fmt_src_n(src1, 1), _fmt_src_n(src2, 1)
+      dst_str = f"v{vdst}"
+    n = 3 if is_3src else 2
+    full_opsel_hi = opsel_hi | (opsel_hi2 << 2)
+    mods = [fmt_bits("op_sel", opsel, n)] if opsel else []
+    if full_opsel_hi != (0b111 if is_3src else 0b11): mods.append(fmt_bits("op_sel_hi", full_opsel_hi, n))
+    if neg: mods.append(fmt_bits("neg_lo", neg, n))
+    if neg_hi: mods.append(fmt_bits("neg_hi", neg_hi, n))
+    if clmp: mods.append("clamp")
+    mod_str = " " + " ".join(mods) if mods else ""
+    return f"{op_name} {dst_str}, {src0_str}, {src1_str}, {src2_str}{mod_str}" if is_3src else f"{op_name} {dst_str}, {src0_str}, {src1_str}{mod_str}"
+
+  # VINTERP: interpolation instructions
+  if cls_name == 'VINTERP':
+    vdst = unwrap(inst._values.get('vdst', 0))
+    src0, src1, src2 = [unwrap(inst._values.get(f, 0)) for f in ('src0', 'src1', 'src2')]
+    neg, waitexp, clmp = unwrap(inst._values.get('neg', 0)), unwrap(inst._values.get('waitexp', 0)), unwrap(inst._values.get('clmp', 0))
+    def fmt_neg_vi(v, neg_bit): return f"-{v}" if neg_bit else v
+    srcs = [fmt_neg_vi(f"v{s - 256}" if s >= 256 else fmt_src(s), neg & (1 << i)) for i, s in enumerate([src0, src1, src2])]
+    mods = [m for m in [f"wait_exp:{waitexp}" if waitexp else "", "clamp" if clmp else ""] if m]
+    return f"{op_name} v{vdst}, {', '.join(srcs)}" + (" " + " ".join(mods) if mods else "")
+
+  # MUBUF/MTBUF helpers
+  def _buf_vaddr(vaddr, offen, idxen): return _vreg(vaddr, 2) if offen and idxen else f"v{vaddr}" if offen or idxen else "off"
+  def _buf_srsrc(srsrc): srsrc_base = srsrc * 4; return _reg("ttmp", srsrc_base - 108, 4) if 108 <= srsrc_base <= 123 else _sreg(srsrc_base, 4)
+
+  # MUBUF: buffer load/store
+  if cls_name == 'MUBUF':
+    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
+    offset, offen, idxen = unwrap(inst._values.get('offset', 0)), unwrap(inst._values.get('offen', 0)), unwrap(inst._values.get('idxen', 0))
+    glc, dlc, slc, tfe = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc', 'tfe')]
+    if op_name in ('buffer_gl0_inv', 'buffer_gl1_inv'): return op_name
+    # Determine data width from op name
+    if 'd16' in op_name: width = 2 if any(x in op_name for x in ('xyz', 'xyzw')) else 1
+    elif 'atomic' in op_name:
+      base_width = 2 if any(x in op_name for x in ('b64', 'u64', 'i64')) else 1
+      width = base_width * 2 if 'cmpswap' in op_name else base_width
+    else: width = {'b32':1, 'b64':2, 'b96':3, 'b128':4, 'b16':1, 'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
+    if tfe: width += 1
+    mods = [m for m in ["offen" if offen else "", "idxen" if idxen else "", f"offset:{offset}" if offset else "",
+                        "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else "", "tfe" if tfe else ""] if m]
+    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)}" + (" " + " ".join(mods) if mods else "")
+
+  # MTBUF: typed buffer load/store
+  if cls_name == 'MTBUF':
+    vdata, vaddr, srsrc, soffset = [unwrap(inst._values.get(f, 0)) for f in ('vdata', 'vaddr', 'srsrc', 'soffset')]
+    offset, tbuf_fmt, offen, idxen = [unwrap(inst._values.get(f, 0)) for f in ('offset', 'format', 'offen', 'idxen')]
+    glc, dlc, slc = [unwrap(inst._values.get(f, 0)) for f in ('glc', 'dlc', 'slc')]
+    mods = [f"format:{tbuf_fmt}"] + [m for m in ["idxen" if idxen else "", "offen" if offen else "", f"offset:{offset}" if offset else "",
+                                                  "glc" if glc else "", "dlc" if dlc else "", "slc" if slc else ""] if m]
+    width = 2 if 'd16' in op_name and any(x in op_name for x in ('xyz', 'xyzw')) else 1 if 'd16' in op_name else {'x':1, 'xy':2, 'xyz':3, 'xyzw':4}.get(op_name.split('_')[-1], 1)
+    return f"{op_name} {_vreg(vdata, width)}, {_buf_vaddr(vaddr, offen, idxen)}, {_buf_srsrc(srsrc)}, {decode_src(soffset)} {' '.join(mods)}"
+
+  # SOP1/SOP2/SOPC/SOPK
+  if cls_name in ('SOP1', 'SOP2', 'SOPC', 'SOPK'):
+    sizes = _parse_sop_sizes(op_name)
+    dst_cnt, src0_cnt = sizes[0], sizes[1]
+    src1_cnt = sizes[2] if len(sizes) > 2 else src0_cnt
+    if cls_name == 'SOP1':
+      sdst, ssrc0 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('ssrc0', 0))
+      if op_name == 's_getpc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}"
+      if op_name in ('s_setpc_b64', 's_rfe_b64'): return f"{op_name} {_fmt_ssrc(ssrc0, 2)}"
+      if op_name == 's_swappc_b64': return f"{op_name} {_fmt_sdst(sdst, 2)}, {_fmt_ssrc(ssrc0, 2)}"
+      if op_name in ('s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'):
+        return f"{op_name} {_fmt_sdst(sdst, 2 if 'b64' in op_name else 1)}, sendmsg({MSG_NAMES.get(ssrc0, str(ssrc0))})"
+      ssrc0_str = fmt_src(ssrc0) if src0_cnt == 1 else _fmt_ssrc(ssrc0, src0_cnt)
+      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {ssrc0_str}"
+    if cls_name == 'SOP2':
+      sdst, ssrc0, ssrc1 = [unwrap(inst._values.get(f, 0)) for f in ('sdst', 'ssrc0', 'ssrc1')]
+      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, {_fmt_ssrc(ssrc0, src0_cnt)}, {_fmt_ssrc(ssrc1, src1_cnt)}"
+    if cls_name == 'SOPC':
+      return f"{op_name} {_fmt_ssrc(unwrap(inst._values.get('ssrc0', 0)), src0_cnt)}, {_fmt_ssrc(unwrap(inst._values.get('ssrc1', 0)), src1_cnt)}"
+    if cls_name == 'SOPK':
+      sdst, simm16 = unwrap(inst._values.get('sdst', 0)), unwrap(inst._values.get('simm16', 0))
+      if op_name == 's_version': return f"{op_name} 0x{simm16:x}"
+      if op_name in ('s_setreg_b32', 's_getreg_b32'):
+        hwreg_id, hwreg_offset, hwreg_size = simm16 & 0x3f, (simm16 >> 6) & 0x1f, ((simm16 >> 11) & 0x1f) + 1
+        hwreg_str = f"0x{simm16:x}" if hwreg_id in (16, 17) else f"hwreg({HWREG_NAMES.get(hwreg_id, str(hwreg_id))}, {hwreg_offset}, {hwreg_size})"
+        return f"{op_name} {hwreg_str}, {_fmt_sdst(sdst, 1)}" if op_name == 's_setreg_b32' else f"{op_name} {_fmt_sdst(sdst, 1)}, {hwreg_str}"
+      return f"{op_name} {_fmt_sdst(sdst, dst_cnt)}, 0x{simm16:x}"
+
+  # Generic fallback
+  def fmt_field(n, v):
+    v = unwrap(v)
+    if n in SRC_FIELDS: return fmt_src(v) if v != 255 else "0xff"
+    if n in ('sdst', 'vdst'): return f"{'s' if n == 'sdst' else 'v'}{v}"
+    return f"v{v}" if n == 'vsrc1' else f"0x{v:x}" if n == 'simm16' else str(v)
+  ops = [fmt_field(n, inst._values.get(n, 0)) for n in inst._fields if n not in ('encoding', 'op')]
+  return f"{op_name} {', '.join(ops)}" if ops else op_name
+
+# Assembler
+SPECIAL_REGS = {'vcc_lo': RawImm(106), 'vcc_hi': RawImm(107), 'null': RawImm(124), 'off': RawImm(124), 'm0': RawImm(125), 'exec_lo': RawImm(126), 'exec_hi': RawImm(127), 'scc': RawImm(253)}
+FLOAT_CONSTS = {'0.5': 0.5, '-0.5': -0.5, '1.0': 1.0, '-1.0': -1.0, '2.0': 2.0, '-2.0': -2.0, '4.0': 4.0, '-4.0': -4.0}
+REG_MAP: dict[str, _RegFactory] = {'s': s, 'v': v, 't': ttmp, 'ttmp': ttmp}
+
+def parse_operand(op: str) -> tuple:
+  op = op.strip().lower()
+  neg = op.startswith('-') and not op[1:2].isdigit(); op = op[1:] if neg else op
+  abs_ = op.startswith('|') and op.endswith('|') or op.startswith('abs(') and op.endswith(')')
+  op = op[1:-1] if op.startswith('|') else op[4:-1] if op.startswith('abs(') else op
+  hi_half = op.endswith('.h')
+  op = re.sub(r'\.[lh]$', '', op)
+  if op in FLOAT_CONSTS: return (FLOAT_CONSTS[op], neg, abs_, hi_half)
+  if re.match(r'^-?\d+$', op): return (int(op), neg, abs_, hi_half)
+  if m := re.match(r'^-?0x([0-9a-f]+)$', op):
+    v = -int(m.group(1), 16) if op.startswith('-') else int(m.group(1), 16)
+    return (v, neg, abs_, hi_half)
+  if op in SPECIAL_REGS: return (SPECIAL_REGS[op], neg, abs_, hi_half)
+  if op == 'lit': return (RawImm(255), neg, abs_, hi_half)  # literal marker (actual value comes from literal word)
+  if m := re.match(r'^([svt](?:tmp)?)\[(\d+):(\d+)\]$', op): return (REG_MAP[m.group(1)][int(m.group(2)):int(m.group(3))], neg, abs_, hi_half)
+  if m := re.match(r'^([svt](?:tmp)?)(\d+)$', op):
+    reg = REG_MAP[m.group(1)][int(m.group(2))]
+    reg.hi = hi_half
+    return (reg, neg, abs_, hi_half)
+  # hwreg(name, offset, size) or hwreg(name) -> simm16 encoding
+  if m := re.match(r'^hwreg\((\w+)(?:,\s*(\d+),\s*(\d+))?\)$', op):
+    name_str = m.group(1).lower()
+    hwreg_id = HWREG_IDS.get(name_str, int(name_str) if name_str.isdigit() else None)
+    if hwreg_id is None: raise ValueError(f"unknown hwreg name: {name_str}")
+    offset, size = int(m.group(2)) if m.group(2) else 0, int(m.group(3)) if m.group(3) else 32
+    return (((size - 1) << 11) | (offset << 6) | hwreg_id, neg, abs_, hi_half)
+  raise ValueError(f"cannot parse operand: {op}")
+
+SMEM_OPS = {'s_load_b32', 's_load_b64', 's_load_b128', 's_load_b256', 's_load_b512',
+            's_buffer_load_b32', 's_buffer_load_b64', 's_buffer_load_b128', 's_buffer_load_b256', 's_buffer_load_b512'}
+SOP1_SRC_ONLY = {'s_setpc_b64', 's_rfe_b64'}
+SOP1_MSG_IMM = {'s_sendmsg_rtn_b32', 's_sendmsg_rtn_b64'}
+SOPK_IMM_ONLY = {'s_version'}
+SOPK_IMM_FIRST = {'s_setreg_b32'}
+SOPK_UNSUPPORTED = {'s_setreg_imm32_b32'}
+
+def asm(text: str) -> Inst:
+  from extra.assembly.rdna3 import autogen
+  text = text.strip()
+  clamp = 'clamp' in text.lower()
+  if clamp: text = re.sub(r'\s+clamp\s*$', '', text, flags=re.I)
+  modifiers = {}
+  if m := re.search(r'\s+wait_exp:(\d+)', text, re.I): modifiers['waitexp'] = int(m.group(1)); text = text[:m.start()] + text[m.end():]
+  parts = text.replace(',', ' ').split()
+  if not parts: raise ValueError("empty instruction")
+  mnemonic, op_str = parts[0].lower(), text[len(parts[0]):].strip()
+  # Handle s_waitcnt specially before operand parsing
+  if mnemonic == 's_waitcnt':
+    vmcnt, expcnt, lgkmcnt = 0x3f, 0x7, 0x3f
+    for part in op_str.replace(',', ' ').split():
+      if m := re.match(r'vmcnt\((\d+)\)', part): vmcnt = int(m.group(1))
+      elif m := re.match(r'expcnt\((\d+)\)', part): expcnt = int(m.group(1))
+      elif m := re.match(r'lgkmcnt\((\d+)\)', part): lgkmcnt = int(m.group(1))
+      elif re.match(r'^0x[0-9a-f]+$|^\d+$', part): return autogen.s_waitcnt(simm16=int(part, 0))
+    return autogen.s_waitcnt(simm16=waitcnt(vmcnt, expcnt, lgkmcnt))
+  # Handle VOPD dual-issue instructions: opx dst, src :: opy dst, src
+  if '::' in text:
+    x_part, y_part = text.split('::')
+    x_parts, y_parts = x_part.strip().replace(',', ' ').split(), y_part.strip().replace(',', ' ').split()
+    opx_name, opy_name = x_parts[0].upper(), y_parts[0].upper()
+    opx, opy = autogen.VOPDOp[opx_name], autogen.VOPDOp[opy_name]
+    x_ops, y_ops = [parse_operand(p)[0] for p in x_parts[1:]], [parse_operand(p)[0] for p in y_parts[1:]]
+    vdstx, srcx0 = x_ops[0], x_ops[1] if len(x_ops) > 1 else 0
+    vsrcx1 = x_ops[2] if len(x_ops) > 2 else VGPR(0)
+    vdsty, srcy0 = y_ops[0], y_ops[1] if len(y_ops) > 1 else 0
+    vsrcy1 = y_ops[2] if len(y_ops) > 2 else VGPR(0)
+    # Handle fmaak/fmamk literals (4th operand on x or y side)
+    lit = None
+    if 'fmaak' in opx_name.lower() and len(x_ops) > 3: lit = unwrap(x_ops[3])
+    elif 'fmamk' in opx_name.lower() and len(x_ops) > 3: lit, vsrcx1 = unwrap(x_ops[2]), x_ops[3]
+    elif 'fmaak' in opy_name.lower() and len(y_ops) > 3: lit = unwrap(y_ops[3])
+    elif 'fmamk' in opy_name.lower() and len(y_ops) > 3: lit, vsrcy1 = unwrap(y_ops[2]), y_ops[3]
+    return autogen.VOPD(opx, opy, vdstx=vdstx, vdsty=vdsty, srcx0=srcx0, vsrcx1=vsrcx1, srcy0=srcy0, vsrcy1=vsrcy1, literal=lit)
+  operands, current, depth, in_pipe = [], "", 0, False
+  for ch in op_str:
+    if ch in '[(': depth += 1
+    elif ch in '])': depth -= 1
+    elif ch == '|': in_pipe = not in_pipe
+    if ch == ',' and depth == 0 and not in_pipe: operands.append(current.strip()); current = ""
+    else: current += ch
+  if current.strip(): operands.append(current.strip())
+  parsed = [parse_operand(op) for op in operands]
+  values = [p[0] for p in parsed]
+  neg_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[1])
+  abs_bits = sum((1 << (i-1)) for i, p in enumerate(parsed) if i > 0 and p[2])
+  opsel_bits = (8 if len(parsed) > 0 and parsed[0][3] else 0) | sum((1 << i) for i, p in enumerate(parsed[1:4]) if p[3])
+  lit = None
+  if mnemonic in ('v_fmaak_f32', 'v_fmaak_f16') and len(values) == 4: lit, values = unwrap(values[3]), values[:3]
+  elif mnemonic in ('v_fmamk_f32', 'v_fmamk_f16') and len(values) == 4: lit, values = unwrap(values[2]), [values[0], values[1], values[3]]
+  vcc_ops = {'v_add_co_ci_u32', 'v_sub_co_ci_u32', 'v_subrev_co_ci_u32', 'v_add_co_u32', 'v_sub_co_u32', 'v_subrev_co_u32'}
+  if mnemonic.replace('_e32', '') in vcc_ops and len(values) >= 5: values = [values[0], values[2], values[3]]
+  if mnemonic.startswith('v_cmp') and len(values) >= 3 and operands[0].strip().lower() in ('vcc_lo', 'vcc_hi', 'vcc'):
+    values = values[1:]
+  # CMPX instructions with _e64 suffix: prepend implicit EXEC_LO destination (vdst=126)
+  if 'cmpx' in mnemonic and mnemonic.endswith('_e64') and len(values) == 2:
+    values = [VGPR(126, 1)] + values
+    # Recalculate modifiers: parsed[0]=src0, parsed[1]=src1 (no vdst in user input)
+    neg_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[1])
+    abs_bits = sum((1 << i) for i, p in enumerate(parsed[:3]) if p[2])
+    opsel_bits = sum((1 << i) for i, p in enumerate(parsed[:2]) if p[3])
+  vop3sd_ops = {'v_div_scale_f32', 'v_div_scale_f64'}
+  if mnemonic in vop3sd_ops and len(parsed) >= 5:
+    neg_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[1])
+    abs_bits = sum((1 << i) for i, p in enumerate(parsed[2:5]) if p[2])
+  if mnemonic in SOPK_UNSUPPORTED: raise ValueError(f"unsupported instruction: {mnemonic}")
+  elif mnemonic in SOP1_SRC_ONLY:
+    return getattr(autogen, mnemonic)(ssrc0=values[0])
+  elif mnemonic in SOP1_MSG_IMM:
+    return getattr(autogen, mnemonic)(sdst=values[0], ssrc0=RawImm(unwrap(values[1])))
+  elif mnemonic in SOPK_IMM_ONLY:
+    return getattr(autogen, mnemonic)(simm16=values[0])
+  elif mnemonic in SOPK_IMM_FIRST:
+    return getattr(autogen, mnemonic)(simm16=values[0], sdst=values[1])
+  elif mnemonic in SMEM_OPS and len(operands) >= 3 and re.match(r'^-?[0-9]|^-?0x', operands[2].strip().lower()):
+    return getattr(autogen, mnemonic)(sdata=values[0], sbase=values[1], offset=values[2], soffset=RawImm(124))
+  elif mnemonic.startswith('buffer_') and len(operands) >= 2 and operands[1].strip().lower() == 'off':
+    return getattr(autogen, mnemonic)(vdata=values[0], vaddr=0, srsrc=values[2], soffset=RawImm(unwrap(values[3])) if len(values) > 3 else RawImm(0))
+  elif (mnemonic.startswith('flat_load') or mnemonic.startswith('global_load') or mnemonic.startswith('scratch_load')) and len(values) >= 3:
+    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
+    return getattr(autogen, mnemonic)(vdst=values[0], addr=values[1], saddr=values[2], offset=offset)
+  elif (mnemonic.startswith('flat_store') or mnemonic.startswith('global_store') or mnemonic.startswith('scratch_store')) and len(values) >= 3:
+    offset = int(m.group(1)) if (m := re.search(r'offset:(-?\d+)', op_str)) else 0
+    return getattr(autogen, mnemonic)(addr=values[0], data=values[1], saddr=values[2], offset=offset)
+  for suffix in (['_e32', ''] if not (neg_bits or abs_bits or clamp) else ['', '_e32']):
+    if hasattr(autogen, name := mnemonic.replace('.', '_') + suffix):
+      use_opsel = 'opsel' in getattr(autogen, name).func._fields
+      vals = [type(v)(v.idx, v.count, False) if isinstance(v, Reg) and v.hi and use_opsel else v for v in values]
+      inst = getattr(autogen, name)(*vals, literal=lit, **modifiers)
+      if neg_bits and 'neg' in inst._fields: inst._values['neg'] = neg_bits
+      if opsel_bits and use_opsel: inst._values['opsel'] = opsel_bits
+      if abs_bits and 'abs' in inst._fields: inst._values['abs'] = abs_bits
+      if clamp and 'clmp' in inst._fields: inst._values['clmp'] = 1
+      return inst
+  raise ValueError(f"unknown instruction: {mnemonic}")

extra/assembly/rdna3/emu.py

@@ -0,0 +1,505 @@
+# RDNA3 emulator - pure Python implementation for testing
+from __future__ import annotations
+import ctypes, struct, math
+from typing import Callable
+from extra.assembly.rdna3.lib import Inst, Inst32, Inst64, RawImm
+
+Program = dict[int, Inst]  # pc (word offset) -> instruction
+from extra.assembly.rdna3.autogen import (
+  SOP1, SOP2, SOPC, SOPK, SOPP, SMEM, VOP1, VOP2, VOP3, VOP3SD, VOP3P, VOPC, DS, FLAT, VOPD, SrcEnum,
+  SOP1Op, SOP2Op, SOPCOp, SOPKOp, SOPPOp, SMEMOp, VOP1Op, VOP2Op, VOP3Op, VOP3SDOp, VOP3POp, VOPCOp, DSOp, FLATOp, GLOBALOp, VOPDOp
+)
+from extra.assembly.rdna3.alu import (
+  f32, i32, f16, i16, sext, vopc, SALU, VALU,
+  SOP1_BASE, SOP2_BASE, SOPC_BASE, SOPK_BASE, VOP1_BASE, VOP2_BASE
+)
+
+WAVE_SIZE, SGPR_COUNT, VGPR_COUNT = 32, 128, 256
+VCC_LO, VCC_HI, NULL, M0, EXEC_LO, EXEC_HI, SCC = SrcEnum.VCC_LO, SrcEnum.VCC_HI, SrcEnum.NULL, SrcEnum.M0, SrcEnum.EXEC_LO, SrcEnum.EXEC_HI, SrcEnum.SCC
+# Pre-computed inline constant table for src operands 128-254 (index = src - 128)
+_INLINE_CONSTS = [0] * 127
+for _i in range(65): _INLINE_CONSTS[_i] = _i  # 128-192 -> 0-64
+for _i in range(1, 17): _INLINE_CONSTS[64 + _i] = ((-_i) & 0xffffffff)  # 193-208 -> -1 to -16
+for _k, _v in {SrcEnum.POS_HALF: 0x3f000000, SrcEnum.NEG_HALF: 0xbf000000, SrcEnum.POS_ONE: 0x3f800000, SrcEnum.NEG_ONE: 0xbf800000,
+               SrcEnum.POS_TWO: 0x40000000, SrcEnum.NEG_TWO: 0xc0000000, SrcEnum.POS_FOUR: 0x40800000, SrcEnum.NEG_FOUR: 0xc0800000,
+               SrcEnum.INV_2PI: 0x3e22f983}.items(): _INLINE_CONSTS[_k - 128] = _v
+
+_valid_mem_ranges: list[tuple[int, int]] = []
+def set_valid_mem_ranges(ranges: set[tuple[int, int]]) -> None: global _valid_mem_ranges; _valid_mem_ranges = list(ranges)
+def _mem_valid(addr: int, size: int) -> bool:
+  for s, z in _valid_mem_ranges:
+    if s <= addr and addr + size <= s + z: return True
+  return not _valid_mem_ranges
+def _ctypes_at(addr: int, size: int): return (ctypes.c_uint8 if size == 1 else ctypes.c_uint16 if size == 2 else ctypes.c_uint32).from_address(addr)
+def mem_read(addr: int, size: int) -> int: return _ctypes_at(addr, size).value if _mem_valid(addr, size) else 0
+def mem_write(addr: int, size: int, val: int) -> None:
+  if _mem_valid(addr, size): _ctypes_at(addr, size).value = val
+
+# Memory op tables - (cnt, sz, sign) for loads, (cnt, sz) for stores
+def _mem_ops(ops, suffix_map):
+  return {getattr(e, f"{p}_{s}"): v for e in ops for s, v in suffix_map.items() for p in [e.__name__.replace("Op", "")]}
+_LOAD_MAP = {'LOAD_B32': (1,4,0), 'LOAD_B64': (2,4,0), 'LOAD_B96': (3,4,0), 'LOAD_B128': (4,4,0), 'LOAD_U8': (1,1,0), 'LOAD_I8': (1,1,1), 'LOAD_U16': (1,2,0), 'LOAD_I16': (1,2,1)}
+_STORE_MAP = {'STORE_B32': (1,4), 'STORE_B64': (2,4), 'STORE_B96': (3,4), 'STORE_B128': (4,4), 'STORE_B8': (1,1), 'STORE_B16': (1,2)}
+FLAT_LOAD = _mem_ops([GLOBALOp, FLATOp], _LOAD_MAP)
+FLAT_STORE = _mem_ops([GLOBALOp, FLATOp], _STORE_MAP)
+DS_LOAD: dict[int, tuple[int,int,int]] = {DSOp.DS_LOAD_B32: (1,4,0), DSOp.DS_LOAD_B64: (2,4,0), DSOp.DS_LOAD_B128: (4,4,0), DSOp.DS_LOAD_U8: (1,1,0), DSOp.DS_LOAD_I8: (1,1,1), DSOp.DS_LOAD_U16: (1,2,0), DSOp.DS_LOAD_I16: (1,2,1)}
+DS_STORE: dict[int, tuple[int,int]] = {DSOp.DS_STORE_B32: (1,4), DSOp.DS_STORE_B64: (2,4), DSOp.DS_STORE_B128: (4,4), DSOp.DS_STORE_B8: (1,1), DSOp.DS_STORE_B16: (1,2)}
+FLAT_D16_LO = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_U8', (1, 0)), ('LOAD_D16_I8', (1, 1)), ('LOAD_D16_B16', (2, 0))]}
+FLAT_D16_HI = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('LOAD_D16_HI_U8', (1, 0)), ('LOAD_D16_HI_I8', (1, 1)), ('LOAD_D16_HI_B16', (2, 0))]}
+FLAT_D16_STORE = {getattr(e, f"{e.__name__.replace('Op', '')}_{s}"): v for e in [FLATOp, GLOBALOp] for s, v in [('STORE_D16_HI_B8', 1), ('STORE_D16_HI_B16', 2)]}
+SMEM_LOAD: dict[int, int] = {SMEMOp.S_LOAD_B32: 1, SMEMOp.S_LOAD_B64: 2, SMEMOp.S_LOAD_B128: 4, SMEMOp.S_LOAD_B256: 8, SMEMOp.S_LOAD_B512: 16}
+SOPK_WAIT = {SOPKOp.S_WAITCNT_VSCNT, SOPKOp.S_WAITCNT_VMCNT, SOPKOp.S_WAITCNT_EXPCNT, SOPKOp.S_WAITCNT_LGKMCNT}
+
+class WaveState:
+  __slots__ = ('sgpr', 'vgpr', 'scc', 'pc', 'literal', '_pend_sgpr')
+  def __init__(self):
+    self.sgpr, self.vgpr = [0] * SGPR_COUNT, [[0] * VGPR_COUNT for _ in range(WAVE_SIZE)]
+    self.sgpr[EXEC_LO] = 0xffffffff  # wave32: all lanes active
+    self.scc = self.pc = self.literal = 0
+    self._pend_sgpr = {}
+
+  @property
+  def vcc(self) -> int: return self.sgpr[VCC_LO] | (self.sgpr[VCC_HI] << 32)
+  @vcc.setter
+  def vcc(self, v: int) -> None: self.sgpr[VCC_LO] = v & 0xffffffff; self.sgpr[VCC_HI] = (v >> 32) & 0xffffffff
+  @property
+  def exec_mask(self) -> int: return self.sgpr[EXEC_LO] | (self.sgpr[EXEC_HI] << 32)
+  @exec_mask.setter
+  def exec_mask(self, v: int) -> None: self.sgpr[EXEC_LO] = v & 0xffffffff; self.sgpr[EXEC_HI] = (v >> 32) & 0xffffffff
+
+  def rsgpr(self, i: int) -> int:
+    if i == NULL: return 0
+    if i == SCC: return self.scc
+    return self.sgpr[i] if i < SGPR_COUNT else 0
+  def wsgpr(self, i: int, v: int) -> None:
+    if i < SGPR_COUNT and i != NULL: self.sgpr[i] = v & 0xffffffff
+  def rsgpr64(self, i: int) -> int: return self.rsgpr(i) | (self.rsgpr(i+1) << 32)
+  def wsgpr64(self, i: int, v: int) -> None: self.wsgpr(i, v & 0xffffffff); self.wsgpr(i+1, (v >> 32) & 0xffffffff)
+
+  def rsrc(self, v: int, lane: int) -> int:
+    if v < SGPR_COUNT: return self.sgpr[v]
+    if v == SCC: return self.scc
+    if v < 255: return _INLINE_CONSTS[v - 128]
+    if v == 255: return self.literal
+    return self.vgpr[lane][v - 256] if v <= 511 else 0
+
+  def rsrc64(self, v: int, lane: int) -> int:
+    return self.rsrc(v, lane) | ((self.rsrc(v+1, lane) if v < VCC_LO or 256 <= v <= 511 else 0) << 32)
+
+  def pend_sgpr_lane(self, reg: int, lane: int, val: int) -> None:
+    if reg not in self._pend_sgpr: self._pend_sgpr[reg] = 0
+    if val: self._pend_sgpr[reg] |= (1 << lane)
+
+  def commit_pends(self) -> None:
+    for reg, val in self._pend_sgpr.items(): self.sgpr[reg] = val
+    self._pend_sgpr.clear()
+
+def decode_format(word: int) -> tuple[type[Inst] | None, bool]:
+  hi2 = (word >> 30) & 0x3
+  if hi2 == 0b11:
+    enc = (word >> 26) & 0xf
+    if enc == 0b1101: return SMEM, True
+    if enc == 0b0101:
+      op = (word >> 16) & 0x3ff
+      return (VOP3SD, True) if op in (288, 289, 290, 764, 765, 766, 767, 768, 769, 770) else (VOP3, True)
+    return {0b0011: (VOP3P, True), 0b0110: (DS, True), 0b0111: (FLAT, True), 0b0010: (VOPD, True)}.get(enc, (None, True))
+  if hi2 == 0b10:
+    enc = (word >> 23) & 0x7f
+    return {0b1111101: (SOP1, False), 0b1111110: (SOPC, False), 0b1111111: (SOPP, False)}.get(enc, (SOPK, False) if ((word >> 28) & 0xf) == 0b1011 else (SOP2, False))
+  enc = (word >> 25) & 0x7f
+  return (VOPC, False) if enc == 0b0111110 else (VOP1, False) if enc == 0b0111111 else (VOP2, False)
+
+def _unwrap(v) -> int: return v.val if isinstance(v, RawImm) else v.value if hasattr(v, 'value') else v
+
+def decode_program(data: bytes) -> Program:
+  result: Program = {}
+  i = 0
+  while i < len(data):
+    word = int.from_bytes(data[i:i+4], 'little')
+    inst_class, is_64 = decode_format(word)
+    if inst_class is None: i += 4; continue
+    base_size = 8 if is_64 else 4
+    inst = inst_class.from_bytes(data[i:i+base_size])
+    for name, val in inst._values.items(): setattr(inst, name, _unwrap(val))
+    has_literal = any(getattr(inst, fld, None) == 255 for fld in ('src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'srcx0', 'srcy0'))
+    if inst_class == VOP2 and inst.op in (44, 45, 55, 56): has_literal = True
+    if inst_class == VOPD and (inst.opx in (1, 2) or inst.opy in (1, 2)): has_literal = True
+    if inst_class == SOP2 and inst.op in (69, 70): has_literal = True
+    if has_literal: inst._literal = int.from_bytes(data[i+base_size:i+base_size+4], 'little')
+    inst._words = inst.size() // 4  # cache size for step_wave
+    result[i // 4] = inst
+    i += inst._words * 4
+  return result
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# SCALAR EXECUTION
+# ═══════════════════════════════════════════════════════════════════════════════
+def exec_sop1(st: WaveState, inst: SOP1) -> int:
+  s0, op = st.rsrc(inst.ssrc0, 0), inst.op
+  # 64-bit and special ops handled inline
+  if op == SOP1Op.S_MOV_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0)); return 0
+  if op == SOP1Op.S_NOT_B64: r = (~st.rsrc64(inst.ssrc0, 0)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
+  if op == SOP1Op.S_BITSET0_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) & ~(1 << (s0 & 0x1f))); return 0
+  if op == SOP1Op.S_BITSET1_B32: st.wsgpr(inst.sdst, st.rsgpr(inst.sdst) | (1 << (s0 & 0x1f))); return 0
+  if op == SOP1Op.S_AND_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
+  if op == SOP1Op.S_OR_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 | old; st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
+  if op == SOP1Op.S_AND_NOT1_SAVEEXEC_B32: old = st.exec_mask & 0xffffffff; st.exec_mask = s0 & (~old & 0xffffffff); st.scc = int(st.exec_mask != 0); st.wsgpr(inst.sdst, old); return 0
+  if op == SOP1Op.S_GETPC_B64: return -3
+  if op == SOP1Op.S_SETPC_B64: return -4
+  if op == SOP1Op.S_SWAPPC_B64: return -5
+  if (fn := SALU.get(SOP1_BASE + op)) is None: raise NotImplementedError(f"SOP1 op {op}")
+  r, scc = fn(s0, 0, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
+
+_SOP2_64: dict[int, Callable[[int, int], int]] = {SOP2Op.S_AND_B64: lambda a, b: a & b, SOP2Op.S_OR_B64: lambda a, b: a | b, SOP2Op.S_XOR_B64: lambda a, b: a ^ b}
+def exec_sop2(st: WaveState, inst: SOP2) -> int:
+  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
+  # 64-bit ops handled inline
+  if op == SOP2Op.S_LSHL_B64: r = (st.rsrc64(inst.ssrc0, 0) << (s1 & 0x3f)) & 0xffffffffffffffff; st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
+  if op == SOP2Op.S_LSHR_B64: r = st.rsrc64(inst.ssrc0, 0) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
+  if op == SOP2Op.S_ASHR_I64: r = sext(st.rsrc64(inst.ssrc0, 0), 64) >> (s1 & 0x3f); st.wsgpr64(inst.sdst, r & 0xffffffffffffffff); st.scc = int(r != 0); return 0
+  if (fn := _SOP2_64.get(op)): r = fn(st.rsrc64(inst.ssrc0, 0), st.rsrc64(inst.ssrc1, 0)); st.wsgpr64(inst.sdst, r); st.scc = int(r != 0); return 0
+  if op == SOP2Op.S_CSELECT_B64: st.wsgpr64(inst.sdst, st.rsrc64(inst.ssrc0, 0) if st.scc else st.rsrc64(inst.ssrc1, 0)); return 0
+  if op == SOP2Op.S_FMAC_F32: st.wsgpr(inst.sdst, i32(f32(st.rsgpr(inst.sdst)) + f32(s0) * f32(s1))); return 0
+  if op == SOP2Op.S_FMAAK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(s1) + f32(inst._literal or 0))); return 0
+  if op == SOP2Op.S_FMAMK_F32: st.wsgpr(inst.sdst, i32(f32(s0) * f32(inst._literal or 0) + f32(s1))); return 0
+  if (fn := SALU.get(SOP2_BASE + op)) is None: raise NotImplementedError(f"SOP2 op {op}")
+  r, scc = fn(s0, s1, st.scc); st.wsgpr(inst.sdst, r); st.scc = scc; return 0
+
+def exec_sopc(st: WaveState, inst: SOPC) -> int:
+  s0, s1, op = st.rsrc(inst.ssrc0, 0), st.rsrc(inst.ssrc1, 0), inst.op
+  if op == SOPCOp.S_CMP_EQ_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) == st.rsrc64(inst.ssrc1, 0)); return 0
+  if op == SOPCOp.S_CMP_LG_U64: st.scc = int(st.rsrc64(inst.ssrc0, 0) != st.rsrc64(inst.ssrc1, 0)); return 0
+  if (fn := SALU.get(SOPC_BASE + op)) is None: raise NotImplementedError(f"SOPC op {op}")
+  st.scc = fn(s0, s1, st.scc)[1]; return 0
+
+_SOPK_CMP = frozenset((SOPKOp.S_CMPK_EQ_I32, SOPKOp.S_CMPK_LG_I32, SOPKOp.S_CMPK_GT_I32, SOPKOp.S_CMPK_GE_I32,
+  SOPKOp.S_CMPK_LT_I32, SOPKOp.S_CMPK_LE_I32, SOPKOp.S_CMPK_EQ_U32, SOPKOp.S_CMPK_LG_U32,
+  SOPKOp.S_CMPK_GT_U32, SOPKOp.S_CMPK_GE_U32, SOPKOp.S_CMPK_LT_U32, SOPKOp.S_CMPK_LE_U32))
+def exec_sopk(st: WaveState, inst: SOPK) -> int:
+  simm, s0, op = inst.simm16, st.rsgpr(inst.sdst), inst.op
+  if op in SOPK_WAIT: return 0
+  if (fn := SALU.get(SOPK_BASE + op)) is None: raise NotImplementedError(f"SOPK op {op}")
+  r, scc = fn(s0, simm, st.scc)
+  if op not in _SOPK_CMP: st.wsgpr(inst.sdst, r)
+  st.scc = scc; return 0
+
+def exec_sopp(st: WaveState, inst: SOPP) -> int:
+  if inst.op == SOPPOp.S_ENDPGM: return -1
+  if inst.op == SOPPOp.S_BARRIER: return -2
+  if inst.op == SOPPOp.S_BRANCH: return sext(inst.simm16, 16)
+  if inst.op == SOPPOp.S_CBRANCH_SCC0: return sext(inst.simm16, 16) if st.scc == 0 else 0
+  if inst.op == SOPPOp.S_CBRANCH_SCC1: return sext(inst.simm16, 16) if st.scc == 1 else 0
+  # In wave32 mode, only VCC_LO is used for lane masks; VCC_HI is a free SGPR
+  if inst.op == SOPPOp.S_CBRANCH_VCCZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) == 0 else 0
+  if inst.op == SOPPOp.S_CBRANCH_VCCNZ: return sext(inst.simm16, 16) if (st.vcc & 0xffffffff) != 0 else 0
+  if inst.op == SOPPOp.S_CBRANCH_EXECZ: return sext(inst.simm16, 16) if st.exec_mask == 0 else 0
+  if inst.op == SOPPOp.S_CBRANCH_EXECNZ: return sext(inst.simm16, 16) if st.exec_mask != 0 else 0
+  # Scheduling hints and wait instructions are no-ops in emulation
+  if inst.op <= 31: return 0  # S_NOP, S_CLAUSE, S_DELAY_ALU, S_WAITCNT, etc.
+  # S_WAKEUP(52), S_SETPRIO(53), S_SENDMSG(54), S_SENDMSGHALT(55), perf counters, S_ICACHE_INV(60) are no-ops
+  if inst.op in (52, 53, 54, 55, 56, 57, 60): return 0
+  raise NotImplementedError(f"SOPP op {inst.op}")
+
+def exec_smem(st: WaveState, inst: SMEM) -> int:
+  addr = st.rsgpr64(inst.sbase * 2) + sext(inst.offset, 21)
+  if inst.soffset not in (NULL, 0x7f): addr += st.rsrc(inst.soffset, 0)
+  if (cnt := SMEM_LOAD.get(inst.op)) is None: raise NotImplementedError(f"SMEM op {inst.op}")
+  for i in range(cnt): st.wsgpr(inst.sdata + i, mem_read((addr + i * 4) & 0xffffffffffffffff, 4))
+  return 0
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# VECTOR EXECUTION
+# ═══════════════════════════════════════════════════════════════════════════════
+def f64(hi: int, lo: int) -> float: return struct.unpack('<d', struct.pack('<Q', (hi << 32) | lo))[0]
+def i64_parts(f: float) -> tuple[int, int]:
+  if math.isnan(f): val = 0x7ff8000000000000
+  elif math.isinf(f): val = 0x7ff0000000000000 if f > 0 else 0xfff0000000000000
+  else: val = struct.unpack('<Q', struct.pack('<d', f))[0]
+  return val & 0xffffffff, (val >> 32) & 0xffffffff
+
+def exec_vop1(st: WaveState, inst: VOP1, lane: int) -> None:
+  if inst.op == VOP1Op.V_NOP: return
+  V, s0 = st.vgpr[lane], st.rsrc(inst.src0, lane)
+  if inst.op == VOP1Op.V_READFIRSTLANE_B32:
+    first = (st.exec_mask & -st.exec_mask).bit_length() - 1 if st.exec_mask else 0
+    st.wsgpr(inst.vdst, st.rsrc(inst.src0, first) if inst.src0 >= 256 else s0); return
+  # F64 ops handled inline
+  if inst.op == VOP1Op.V_CVT_F64_F32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(f32(s0))); return
+  if inst.op == VOP1Op.V_CVT_F64_I32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(sext(s0, 32))); return
+  if inst.op == VOP1Op.V_CVT_F64_U32: V[inst.vdst], V[inst.vdst+1] = i64_parts(float(s0)); return
+  if inst.op in (VOP1Op.V_CVT_F32_F64, VOP1Op.V_CVT_I32_F64, VOP1Op.V_CVT_U32_F64):
+    src = inst.src0 - 256 if inst.src0 >= 256 else inst.src0
+    lo, hi = (V[src], V[src+1]) if inst.src0 >= 256 else (st.sgpr[src], st.sgpr[src+1])
+    v = f64(hi, lo)
+    if inst.op == VOP1Op.V_CVT_F32_F64: V[inst.vdst] = i32(v)
+    elif inst.op == VOP1Op.V_CVT_I32_F64: V[inst.vdst] = (max(-0x80000000, min(0x7fffffff, int(v))) & 0xffffffff) if math.isfinite(v) else 0
+    else: V[inst.vdst] = max(0, min(0xffffffff, int(v))) if math.isfinite(v) and v == v else 0
+    return
+  if (fn := VALU.get(VOP1_BASE + inst.op)): V[inst.vdst] = fn(s0, 0, 0); return
+  raise NotImplementedError(f"VOP1 op {inst.op}")
+
+def exec_vop2(st: WaveState, inst: VOP2, lane: int) -> None:
+  V, s0, s1, op = st.vgpr[lane], st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1], inst.op
+  if op == VOP2Op.V_CNDMASK_B32: V[inst.vdst] = s1 if (st.vcc >> lane) & 1 else s0; return
+  if op == VOP2Op.V_FMAC_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(V[inst.vdst])); return
+  if op == VOP2Op.V_FMAMK_F32: V[inst.vdst] = i32(f32(s0)*f32(st.literal)+f32(s1)); return
+  if op == VOP2Op.V_FMAAK_F32: V[inst.vdst] = i32(f32(s0)*f32(s1)+f32(st.literal)); return
+  if op == VOP2Op.V_FMAC_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(V[inst.vdst])); return
+  if op == VOP2Op.V_FMAMK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(st.literal)+f16(s1)); return
+  if op == VOP2Op.V_FMAAK_F16: V[inst.vdst] = (V[inst.vdst] & 0xffff0000) | i16(f16(s0)*f16(s1)+f16(st.literal)); return
+  if op == VOP2Op.V_PK_FMAC_F16:
+    lo = i16(f16(s0 & 0xffff) * f16(s1 & 0xffff) + f16(V[inst.vdst] & 0xffff))
+    hi = i16(f16((s0 >> 16) & 0xffff) * f16((s1 >> 16) & 0xffff) + f16((V[inst.vdst] >> 16) & 0xffff))
+    V[inst.vdst] = lo | (hi << 16); return
+  if op == VOP2Op.V_ADD_CO_CI_U32: r = s0+s1+((st.vcc>>lane)&1); st.pend_sgpr_lane(VCC_LO, lane, r >= 0x100000000); V[inst.vdst] = r & 0xffffffff; return
+  if op == VOP2Op.V_SUB_CO_CI_U32: b = (st.vcc>>lane)&1; st.pend_sgpr_lane(VCC_LO, lane, s1+b > s0); V[inst.vdst] = (s0-s1-b) & 0xffffffff; return
+  if (fn := VALU.get(VOP2_BASE + op)): V[inst.vdst] = fn(s0, s1, 0); return
+  raise NotImplementedError(f"VOP2 op {op}")
+
+def vop3_mod(val: int, neg: int, abs_: int, idx: int) -> int:
+  if (abs_ >> idx) & 1: val = i32(abs(f32(val)))
+  if (neg >> idx) & 1: val = i32(-f32(val))
+  return val
+
+def exec_vop3(st: WaveState, inst: VOP3, lane: int) -> None:
+  op, src0, src1, src2, vdst, neg, abs_ = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.neg, getattr(inst, 'abs', 0)
+  V = st.vgpr[lane]
+  # VOPC encoded in VOP3 (0-255)
+  if 0 <= op <= 255:
+    base = op & 0x7f
+    # For 64-bit comparisons (I64: 80-87, U64: 88-95), read raw 64-bit values (no float modifiers)
+    if 80 <= base <= 95:
+      s0_64, s1_64 = st.rsrc64(src0, lane), st.rsrc64(src1, lane)
+      result = vopc(op, s0_64 & 0xffffffff, s1_64 & 0xffffffff, (s0_64 >> 32) & 0xffffffff, (s1_64 >> 32) & 0xffffffff)
+    else:
+      s0, s1 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1)
+      result = vopc(op, s0, s1)
+    is_cmpx = op >= 128
+    st.pend_sgpr_lane(vdst, lane, result)
+    if is_cmpx: st.pend_sgpr_lane(EXEC_LO, lane, result)
+    return
+  s0, s1, s2 = vop3_mod(st.rsrc(src0, lane), neg, abs_, 0), vop3_mod(st.rsrc(src1, lane), neg, abs_, 1), vop3_mod(st.rsrc(src2, lane), neg, abs_, 2)
+  # Special ops
+  if op == VOP3Op.V_FMAC_F32: V[vdst] = i32(f32(s0)*f32(s1)+f32(V[vdst])); return
+  if op == VOP3Op.V_READLANE_B32: st.wsgpr(vdst, st.vgpr[s1 & 0x1f][src0 - 256] if src0 >= 256 else s0); return
+  if op == VOP3Op.V_WRITELANE_B32: st.vgpr[s1 & 0x1f][vdst] = s0; return
+  if op == VOP3Op.V_CNDMASK_B32:
+    mask = st.rsgpr(src2) if src2 < 256 else st.vcc
+    V[vdst] = s1 if (mask >> lane) & 1 else s0; return
+  if op in (VOP3Op.V_LSHLREV_B64, VOP3Op.V_LSHRREV_B64, VOP3Op.V_ASHRREV_I64):
+    v64 = st.rsrc64(src1, lane)
+    r = ((v64 << (s0 & 0x3f)) & 0xffffffffffffffff if op == VOP3Op.V_LSHLREV_B64 else
+         v64 >> (s0 & 0x3f) if op == VOP3Op.V_LSHRREV_B64 else sext(v64, 64) >> (s0 & 0x3f))
+    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff; return
+  if op in (VOP3Op.V_ADD_F64, VOP3Op.V_MUL_F64, VOP3Op.V_FMA_F64, VOP3Op.V_MAX_F64, VOP3Op.V_MIN_F64):
+    a, b = f64(st.rsrc(src0+1, lane), s0), f64(st.rsrc(src1+1, lane), s1)
+    c = f64(st.rsrc(src2+1, lane), s2) if op == VOP3Op.V_FMA_F64 else 0.0
+    rf = a + b if op == VOP3Op.V_ADD_F64 else a * b if op == VOP3Op.V_MUL_F64 else a * b + c if op == VOP3Op.V_FMA_F64 else max(a, b) if op == VOP3Op.V_MAX_F64 else min(a, b)
+    V[vdst], V[vdst+1] = i64_parts(rf); return
+  if (fn := VALU.get(op)): V[vdst] = fn(s0, s1, s2); return
+  raise NotImplementedError(f"VOP3 op {op}")
+
+def exec_vopc(st: WaveState, inst: VOPC, lane: int) -> None:
+  result, is_cmpx = vopc(inst.op, st.rsrc(inst.src0, lane), st.vgpr[lane][inst.vsrc1]), inst.op >= 128
+  st.pend_sgpr_lane(EXEC_LO if is_cmpx else VCC_LO, lane, result)
+
+def exec_vop3sd(st: WaveState, inst: VOP3SD, lane: int) -> None:
+  op, src0, src1, src2, vdst, sdst, neg = inst.op, inst.src0, inst.src1, inst.src2, inst.vdst, inst.sdst, inst.neg
+  s0, s1, s2 = st.rsrc(src0, lane), st.rsrc(src1, lane), st.rsrc(src2, lane)
+  if (neg >> 0) & 1: s0 = i32(-f32(s0))
+  if (neg >> 1) & 1: s1 = i32(-f32(s1))
+  if (neg >> 2) & 1: s2 = i32(-f32(s2))
+  V = st.vgpr[lane]
+  if op == VOP3SDOp.V_ADD_CO_U32: r = s0 + s1; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
+  elif op == VOP3SDOp.V_SUB_CO_U32: V[vdst] = (s0 - s1) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 > s0)
+  elif op == VOP3SDOp.V_SUBREV_CO_U32: V[vdst] = (s1 - s0) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s0 > s1)
+  elif op == VOP3SDOp.V_ADD_CO_CI_U32:
+    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
+    r = s0 + s1 + cin; V[vdst] = r & 0xffffffff; st.pend_sgpr_lane(sdst, lane, r >= 0x100000000)
+  elif op == VOP3SDOp.V_SUB_CO_CI_U32:
+    cin = (st.rsgpr(src2) >> lane) & 1 if src2 < 256 else (st.vcc >> lane) & 1
+    V[vdst] = (s0 - s1 - cin) & 0xffffffff; st.pend_sgpr_lane(sdst, lane, s1 + cin > s0)
+  elif op == VOP3SDOp.V_MAD_U64_U32:
+    s2_64 = s2 | (st.rsrc(src2+1, lane) << 32); r = s0 * s1 + s2_64
+    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
+  elif op == VOP3SDOp.V_MAD_I64_I32:
+    s2_64 = sext(s2 | (st.rsrc(src2+1, lane) << 32), 64)
+    r = (sext(s0, 32) * sext(s1, 32) + s2_64) & 0xffffffffffffffff
+    V[vdst], V[vdst+1] = r & 0xffffffff, (r >> 32) & 0xffffffff
+  elif op == VOP3SDOp.V_DIV_SCALE_F32: V[vdst] = 0; st.pend_sgpr_lane(sdst, lane, False)
+  elif op == VOP3SDOp.V_DIV_SCALE_F64: V[vdst], V[vdst+1] = s0, st.rsrc(src0+1, lane); st.pend_sgpr_lane(VCC_LO, lane, s0 == s2)
+  else: raise NotImplementedError(f"VOP3SD op {op}")
+
+def exec_flat(st: WaveState, inst: FLAT, lane: int) -> None:
+  op, addr_reg, data_reg, vdst, offset, saddr, V = inst.op, inst.addr, inst.data, inst.vdst, sext(inst.offset, 13), inst.saddr, st.vgpr[lane]
+  addr = V[addr_reg] | (V[addr_reg+1] << 32)
+  addr = (st.rsgpr64(saddr) + V[addr_reg] + offset) & 0xffffffffffffffff if saddr not in (NULL, 0x7f) else (addr + offset) & 0xffffffffffffffff
+  if op in FLAT_LOAD:
+    cnt, sz, sign = FLAT_LOAD[op]
+    for i in range(cnt): val = mem_read(addr + i * sz, sz); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
+  elif op in FLAT_STORE:
+    cnt, sz = FLAT_STORE[op]
+    for i in range(cnt): mem_write(addr + i * sz, sz, V[data_reg + i] & ((1 << (sz * 8)) - 1))
+  elif op in FLAT_D16_LO: sz, sign = FLAT_D16_LO[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0xffff0000) | ((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff))
+  elif op in FLAT_D16_HI: sz, sign = FLAT_D16_HI[op]; val = mem_read(addr, sz); V[vdst] = (V[vdst] & 0x0000ffff) | (((sext(val, sz * 8) & 0xffff) if sign else (val & 0xffff)) << 16)
+  elif op in FLAT_D16_STORE: mem_write(addr, FLAT_D16_STORE[op], (V[data_reg] >> 16) & ((1 << (FLAT_D16_STORE[op] * 8)) - 1))
+  else: raise NotImplementedError(f"FLAT op {op}")
+
+def exec_ds(st: WaveState, inst: DS, lane: int, lds: bytearray) -> None:
+  op, addr, vdst, V = inst.op, (st.vgpr[lane][inst.addr] + inst.offset0) & 0xffff, inst.vdst, st.vgpr[lane]
+  if op in DS_LOAD:
+    cnt, sz, sign = DS_LOAD[op]
+    for i in range(cnt): val = int.from_bytes(lds[addr+i*sz:addr+i*sz+sz], 'little'); V[vdst + i] = sext(val, sz * 8) & 0xffffffff if sign else val
+  elif op in DS_STORE:
+    cnt, sz = DS_STORE[op]
+    for i in range(cnt): lds[addr+i*sz:addr+i*sz+sz] = (V[inst.data0 + i] & ((1 << (sz * 8)) - 1)).to_bytes(sz, 'little')
+  else: raise NotImplementedError(f"DS op {op}")
+
+VOPD_OPS: dict[int, Callable[[int, int, int, int, int], int]] = {
+  VOPDOp.V_DUAL_MUL_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)), VOPDOp.V_DUAL_ADD_F32: lambda a, b, d, l, lit: i32(f32(a)+f32(b)),
+  VOPDOp.V_DUAL_SUB_F32: lambda a, b, d, l, lit: i32(f32(a)-f32(b)), VOPDOp.V_DUAL_SUBREV_F32: lambda a, b, d, l, lit: i32(f32(b)-f32(a)),
+  VOPDOp.V_DUAL_MAX_F32: lambda a, b, d, l, lit: i32(max(f32(a), f32(b))), VOPDOp.V_DUAL_MIN_F32: lambda a, b, d, l, lit: i32(min(f32(a), f32(b))),
+  VOPDOp.V_DUAL_MUL_DX9_ZERO_F32: lambda a, b, d, l, lit: i32(0.0 if f32(a) == 0.0 or f32(b) == 0.0 else f32(a)*f32(b)),
+  VOPDOp.V_DUAL_MOV_B32: lambda a, b, d, l, lit: a, VOPDOp.V_DUAL_ADD_NC_U32: lambda a, b, d, l, lit: (a + b) & 0xffffffff,
+  VOPDOp.V_DUAL_LSHLREV_B32: lambda a, b, d, l, lit: (b << (a & 0x1f)) & 0xffffffff, VOPDOp.V_DUAL_AND_B32: lambda a, b, d, l, lit: a & b,
+  VOPDOp.V_DUAL_FMAC_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(d)), VOPDOp.V_DUAL_FMAAK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(b)+f32(lit)),
+  VOPDOp.V_DUAL_FMAMK_F32: lambda a, b, d, l, lit: i32(f32(a)*f32(lit)+f32(b)), VOPDOp.V_DUAL_CNDMASK_B32: lambda a, b, d, l, lit: b if l else a,
+}
+def exec_vopd(st: WaveState, inst: VOPD, lane: int) -> None:
+  V, vdsty, vcc_lane = st.vgpr[lane], (inst.vdsty << 1) | ((inst.vdstx & 1) ^ 1), (st.vcc >> lane) & 1
+  sx0, sx1, sy0, sy1, dstx = st.rsrc(inst.srcx0, lane), V[inst.vsrcx1], st.rsrc(inst.srcy0, lane), V[inst.vsrcy1], inst.vdstx
+  if (fn := VOPD_OPS.get(inst.opx)): V[dstx] = fn(sx0, sx1, V[dstx], vcc_lane, st.literal)
+  else: raise NotImplementedError(f"VOPD opx {inst.opx}")
+  if (fn := VOPD_OPS.get(inst.opy)): V[vdsty] = fn(sy0, sy1, V[vdsty], vcc_lane, st.literal)
+  else: raise NotImplementedError(f"VOPD opy {inst.opy}")
+
+def exec_vop3p(st: WaveState, inst: VOP3P, lane: int) -> None:
+  op, vdst, V = inst.op, inst.vdst, st.vgpr[lane]
+  s0, s1, s2 = st.rsrc(inst.src0, lane), st.rsrc(inst.src1, lane), st.rsrc(inst.src2, lane)
+  opsel, opsel_hi = [(inst.opsel >> i) & 1 for i in range(3)], [(inst.opsel_hi >> i) & 1 for i in range(2)] + [inst.opsel_hi2]
+  neg, neg_hi = inst.neg, inst.neg_hi
+  def get_src(src: int, idx: int, for_mix: bool = False) -> float:
+    if for_mix:
+      if not opsel_hi[idx]: return abs(f32(src)) if (neg_hi >> idx) & 1 else f32(src)
+      return float(f16((src >> 16) & 0xffff) if opsel[idx] else f16(src & 0xffff))
+    use_hi = opsel[idx]
+    val = ((src >> 16) & 0xffff) if use_hi else (src & 0xffff)
+    f = f16(val)
+    if use_hi and (neg >> idx) & 1: f = -f
+    elif not use_hi and (neg_hi >> idx) & 1: f = -f
+    return f
+  if op == VOP3POp.V_FMA_MIX_F32: V[vdst] = i32(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
+  elif op == VOP3POp.V_FMA_MIXLO_F16: V[vdst] = (V[vdst] & 0xffff0000) | i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True))
+  elif op == VOP3POp.V_FMA_MIXHI_F16: V[vdst] = (V[vdst] & 0x0000ffff) | (i16(get_src(s0, 0, True) * get_src(s1, 1, True) + get_src(s2, 2, True)) << 16)
+  else: raise NotImplementedError(f"VOP3P op {op}")
+
+def exec_wmma_f32_16x16x16_f16(st: WaveState, inst: VOP3P, n_lanes: int) -> None:
+  src0_base, src1_base, src2_base = (inst.src0 - 256) if inst.src0 >= 256 else inst.src0, (inst.src1 - 256) if inst.src1 >= 256 else inst.src1, (inst.src2 - 256) if inst.src2 >= 256 else inst.src2
+  src0_is_vgpr, src1_is_vgpr, src2_is_vgpr, vdst = inst.src0 >= 256, inst.src1 >= 256, inst.src2 >= 256, inst.vdst
+  A, B, C = [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)], [[0.0] * 16 for _ in range(16)]
+  for lane in range(min(n_lanes, 16)):
+    V = st.vgpr[lane]
+    for reg in range(8):
+      val = V[src0_base + reg] if src0_is_vgpr else st.sgpr[src0_base + reg]
+      A[lane][reg * 2], A[lane][reg * 2 + 1] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
+      val = V[src1_base + reg] if src1_is_vgpr else st.sgpr[src1_base + reg]
+      B[reg * 2][lane], B[reg * 2 + 1][lane] = f16(val & 0xffff), f16((val >> 16) & 0xffff)
+  for row in range(16):
+    for col in range(16):
+      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
+      if lane_idx < n_lanes:
+        val = st.vgpr[lane_idx][src2_base + reg] if src2_is_vgpr else st.sgpr[src2_base + reg]
+        C[row][col] = f32(val)
+  for row in range(16):
+    for col in range(16):
+      for k in range(16): C[row][col] += A[row][k] * B[k][col]
+  for row in range(16):
+    for col in range(16):
+      idx, lane_idx, reg = row * 16 + col, (row * 16 + col) % 32, (row * 16 + col) // 32
+      if lane_idx < n_lanes and (st.exec_mask & (1 << lane_idx)): st.vgpr[lane_idx][vdst + reg] = i32(C[row][col])
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# MAIN EXECUTION LOOP
+# ═══════════════════════════════════════════════════════════════════════════════
+SCALAR: dict[type, Callable[..., int]] = {SOP1: exec_sop1, SOP2: exec_sop2, SOPC: exec_sopc, SOPK: exec_sopk, SOPP: exec_sopp, SMEM: exec_smem}
+VECTOR: dict[type, Callable[..., None]] = {VOP1: exec_vop1, VOP2: exec_vop2, VOP3: exec_vop3, VOP3SD: exec_vop3sd, VOPC: exec_vopc, FLAT: exec_flat, DS: exec_ds, VOPD: exec_vopd, VOP3P: exec_vop3p}
+
+_WMMA_OPS = frozenset((VOP3POp.V_WMMA_F32_16X16X16_F16, VOP3POp.V_WMMA_F32_16X16X16_BF16, VOP3POp.V_WMMA_F16_16X16X16_F16,
+                       VOP3POp.V_WMMA_BF16_16X16X16_BF16, VOP3POp.V_WMMA_I32_16X16X16_IU8, VOP3POp.V_WMMA_I32_16X16X16_IU4))
+
+def step_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int) -> int:
+  inst = program.get(st.pc)
+  if inst is None: return 1
+  inst_words, st.literal, inst_type = inst._words, inst._literal or 0, type(inst)
+  if (handler := SCALAR.get(inst_type)) is not None:
+    delta = handler(st, inst)
+    if delta == -1: return -1
+    if delta == -2: st.pc += inst_words; return -2
+    if delta == -3:  # S_GETPC_B64
+      sop1 = inst if isinstance(inst, SOP1) else None
+      assert sop1 is not None
+      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc += inst_words; return 0
+    if delta == -4:  # S_SETPC_B64
+      sop1 = inst if isinstance(inst, SOP1) else None
+      assert sop1 is not None
+      st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
+    if delta == -5:  # S_SWAPPC_B64
+      sop1 = inst if isinstance(inst, SOP1) else None
+      assert sop1 is not None
+      next_pc = (st.pc + inst_words) * 4; st.wsgpr(sop1.sdst, next_pc & 0xffffffff); st.wsgpr(sop1.sdst + 1, (next_pc >> 32) & 0xffffffff); st.pc = st.rsrc64(sop1.ssrc0, 0) // 4; return 0
+    st.pc += inst_words + delta
+  else:
+    vec_handler, exec_mask = VECTOR[inst_type], st.exec_mask
+    if inst_type is DS:
+      for lane in range(n_lanes):
+        if exec_mask & (1 << lane): vec_handler(st, inst, lane, lds)
+    elif inst_type is VOP3P:
+      vop3p = inst if isinstance(inst, VOP3P) else None
+      assert vop3p is not None
+      if vop3p.op in _WMMA_OPS:
+        exec_wmma_f32_16x16x16_f16(st, vop3p, n_lanes)
+      else:
+        for lane in range(n_lanes):
+          if exec_mask & (1 << lane): vec_handler(st, vop3p, lane)
+    else:
+      for lane in range(n_lanes):
+        if exec_mask & (1 << lane): vec_handler(st, inst, lane)
+    st.commit_pends(); st.pc += inst_words
+  return 0
+
+def exec_wave(program: Program, st: WaveState, lds: bytearray, n_lanes: int, wg_id: tuple[int,int,int]=(0,0,0), local_size: tuple[int,int,int]=(1,1,1), wave_start: int=0) -> int:
+  while st.pc in program:
+    result = step_wave(program, st, lds, n_lanes)
+    if result == -1: return 0
+    if result == -2: return -2
+  return 0
+
+def exec_workgroup(program: Program, workgroup_id: tuple[int, int, int], local_size: tuple[int, int, int], args_ptr: int, dispatch_dim: int) -> None:
+  lx, ly, lz = local_size
+  total_threads, lds = lx * ly * lz, bytearray(65536)
+  waves: list[tuple[WaveState, int, int]] = []
+  for wave_start in range(0, total_threads, WAVE_SIZE):
+    n_lanes, st = min(WAVE_SIZE, total_threads - wave_start), WaveState()
+    st.exec_mask = (1 << n_lanes) - 1
+    st.wsgpr64(0, args_ptr)
+    gx, gy, gz = workgroup_id
+    if dispatch_dim >= 3: st.sgpr[13], st.sgpr[14], st.sgpr[15] = gx, gy, gz
+    elif dispatch_dim == 2: st.sgpr[14], st.sgpr[15] = gx, gy
+    else: st.sgpr[15] = gx
+    for i in range(n_lanes):
+      tid = wave_start + i
+      st.vgpr[i][0] = tid if local_size == (lx, 1, 1) else ((tid // (lx * ly)) << 20) | (((tid // lx) % ly) << 10) | (tid % lx)
+    waves.append((st, n_lanes, wave_start))
+  has_barrier = any(isinstance(inst, SOPP) and inst.op == SOPPOp.S_BARRIER for inst in program.values())
+  for _ in range(2 if has_barrier else 1):
+    for st, n_lanes, wave_start in waves: exec_wave(program, st, lds, n_lanes, workgroup_id, local_size, wave_start)
+
+def run_asm(lib: int, lib_sz: int, gx: int, gy: int, gz: int, lx: int, ly: int, lz: int, args_ptr: int) -> int:
+  data = (ctypes.c_char * lib_sz).from_address(lib).raw
+  program = decode_program(data)
+  if not program: return -1
+  dispatch_dim = 3 if gz > 1 else (2 if gy > 1 else 1)
+  for gidz in range(gz):
+    for gidy in range(gy):
+      for gidx in range(gx): exec_workgroup(program, (gidx, gidy, gidz), (lx, ly, lz), args_ptr, dispatch_dim)
+  return 0

extra/assembly/rdna3/gen.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+# generates autogen/__init__.py by parsing the AMD RDNA3.5 ISA PDF
+import re, pdfplumber, pathlib
+from tinygrad.helpers import fetch
+
+PDF_URL = "https://docs.amd.com/api/khub/documents/UVVZM22UN7tMUeiW_4ShTQ/content"
+FIELD_TYPES = {'SSRC0': 'SSrc', 'SSRC1': 'SSrc', 'SOFFSET': 'SSrc', 'SADDR': 'SSrc', 'SRC0': 'Src', 'SRC1': 'Src', 'SRC2': 'Src',
+  'SDST': 'SGPRField', 'SBASE': 'SGPRField', 'SDATA': 'SGPRField', 'SRSRC': 'SGPRField', 'VDST': 'VGPRField', 'VSRC1': 'VGPRField', 'VDATA': 'VGPRField',
+  'VADDR': 'VGPRField', 'ADDR': 'VGPRField', 'DATA': 'VGPRField', 'DATA0': 'VGPRField', 'DATA1': 'VGPRField', 'SIMM16': 'SImm', 'OFFSET': 'Imm',
+  'OPX': 'VOPDOp', 'OPY': 'VOPDOp', 'SRCX0': 'Src', 'SRCY0': 'Src', 'VSRCX1': 'VGPRField', 'VSRCY1': 'VGPRField', 'VDSTX': 'VGPRField', 'VDSTY': 'VDSTYEnc'}
+FIELD_ORDER = {
+  'SOP2': ['op', 'sdst', 'ssrc0', 'ssrc1'], 'SOP1': ['op', 'sdst', 'ssrc0'], 'SOPC': ['op', 'ssrc0', 'ssrc1'],
+  'SOPK': ['op', 'sdst', 'simm16'], 'SOPP': ['op', 'simm16'], 'VOP1': ['op', 'vdst', 'src0'], 'VOPC': ['op', 'src0', 'vsrc1'],
+  'VOP2': ['op', 'vdst', 'src0', 'vsrc1'], 'VOP3SD': ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2', 'clmp'],
+  'SMEM': ['op', 'sdata', 'sbase', 'soffset', 'offset', 'glc', 'dlc'], 'DS': ['op', 'vdst', 'addr', 'data0', 'data1'],
+  'VOP3': ['op', 'vdst', 'src0', 'src1', 'src2', 'omod', 'neg', 'abs', 'clmp', 'opsel'],
+  'VOP3P': ['op', 'vdst', 'src0', 'src1', 'src2', 'neg', 'neg_hi', 'opsel', 'opsel_hi', 'clmp'],
+  'FLAT': ['op', 'vdst', 'addr', 'data', 'saddr', 'offset', 'seg', 'dlc', 'glc', 'slc'],
+  'MUBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
+  'MTBUF': ['op', 'vdata', 'vaddr', 'srsrc', 'soffset', 'offset', 'format', 'offen', 'idxen', 'glc', 'dlc', 'slc', 'tfe'],
+  'MIMG': ['op', 'vdata', 'vaddr', 'srsrc', 'ssamp', 'dmask', 'dim', 'unrm', 'dlc', 'glc', 'slc'],
+  'EXP': ['en', 'target', 'vsrc0', 'vsrc1', 'vsrc2', 'vsrc3', 'done', 'row'],
+  'VINTERP': ['op', 'vdst', 'src0', 'src1', 'src2', 'waitexp', 'clmp', 'opsel', 'neg'],
+  'VOPD': ['opx', 'opy', 'vdstx', 'vdsty', 'srcx0', 'vsrcx1', 'srcy0', 'vsrcy1'],
+  'LDSDIR': ['op', 'vdst', 'attr', 'attr_chan', 'wait_va']}
+SRC_EXTRAS = {233: 'DPP8', 234: 'DPP8FI', 250: 'DPP16', 251: 'VCCZ', 252: 'EXECZ', 254: 'LDS_DIRECT'}
+FLOAT_MAP = {'0.5': 'POS_HALF', '-0.5': 'NEG_HALF', '1.0': 'POS_ONE', '-1.0': 'NEG_ONE', '2.0': 'POS_TWO', '-2.0': 'NEG_TWO',
+  '4.0': 'POS_FOUR', '-4.0': 'NEG_FOUR', '1/(2*PI)': 'INV_2PI', '0': 'ZERO'}
+
+def parse_bits(s: str) -> tuple[int, int] | None:
+  return (int(m.group(1)), int(m.group(2) or m.group(1))) if (m := re.match(r'\[(\d+)(?::(\d+))?\]', s)) else None
+
+def parse_fields_table(table: list, fmt: str, enums: set[str]) -> list[tuple]:
+  fields = []
+  for row in table[1:]:
+    if not row or not row[0]: continue
+    name, bits_str = row[0].split('\n')[0].strip(), (row[1] or '').split('\n')[0].strip()
+    if not (bits := parse_bits(bits_str)): continue
+    enc_val, hi, lo = None, bits[0], bits[1]
+    if name == 'ENCODING' and row[2] and (m := re.search(r"'b([01_]+)", row[2])):
+      enc_bits = m.group(1).replace('_', '')
+      enc_val = int(enc_bits, 2)
+      declared_width, actual_width = hi - lo + 1, len(enc_bits)
+      if actual_width > declared_width: lo = hi - actual_width + 1
+    ftype = f"{fmt}Op" if name == 'OP' and f"{fmt}Op" in enums else FIELD_TYPES.get(name.upper())
+    fields.append((name, hi, lo, enc_val, ftype))
+  return fields
+
+def generate(output_path: pathlib.Path|str|None = None) -> dict:
+  """Generate RDNA3.5 instruction definitions from the AMD ISA PDF. Returns dict with formats for testing."""
+  pdf = pdfplumber.open(fetch(PDF_URL))
+  pages = pdf.pages[150:200]
+  page_texts = [p.extract_text() or '' for p in pages]
+  page_tables = [[t.extract() for t in p.find_tables()] for p in pages]
+  full_text = '\n'.join(page_texts)
+
+  # parse SSRC encoding from first page with VCC_LO
+  src_enum = dict(SRC_EXTRAS)
+  for text in page_texts[:10]:
+    if 'SSRC0' in text and 'VCC_LO' in text:
+      for m in re.finditer(r'^(\d+)\s+(\S+)', text, re.M):
+        val, name = int(m.group(1)), m.group(2).rstrip('.:')
+        if name in FLOAT_MAP: src_enum[val] = FLOAT_MAP[name]
+        elif re.match(r'^[A-Z][A-Z0-9_]*$', name): src_enum[val] = name
+      break
+
+  # parse opcode tables
+  enums: dict[str, dict[int, str]] = {}
+  for m in re.finditer(r'Table \d+\. (\w+) Opcodes(.*?)(?=Table \d+\.|\n\d+\.\d+\.\d+\.\s+\w+\s*\nDescription|$)', full_text, re.S):
+    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+([A-Z][A-Z0-9_]+)', m.group(2))}:
+      enums[m.group(1) + "Op"] = ops
+  if vopd_m := re.search(r'Table \d+\. VOPD Y-Opcodes\n(.*?)(?=Table \d+\.|15\.\d)', full_text, re.S):
+    if ops := {int(x.group(1)): x.group(2) for x in re.finditer(r'(\d+)\s+(V_DUAL_\w+)', vopd_m.group(1))}:
+      enums["VOPDOp"] = ops
+  enum_names = set(enums.keys())
+
+  def is_fields_table(t) -> bool: return t and len(t) > 1 and t[0] and 'Field' in str(t[0][0] or '')
+  def has_encoding(fields) -> bool: return any(f[0] == 'ENCODING' for f in fields)
+  def has_header_before_fields(text) -> bool:
+    return (pos := text.find('Field Name')) != -1 and bool(re.search(r'\d+\.\d+\.\d+\.\s+\w+\s*\n', text[:pos]))
+
+  # find format headers with their page indices
+  format_headers = []  # (fmt_name, page_idx)
+  for i, text in enumerate(page_texts):
+    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n?Description', text): format_headers.append((m.group(1), i, m.start()))
+    for m in re.finditer(r'\d+\.\d+\.\d+\.\s+(\w+)\s*\n', text):
+      if m.start() > len(text) - 200 and 'Description' not in text[m.end():] and i + 1 < len(page_texts):
+        next_text = page_texts[i + 1].lstrip()
+        if next_text.startswith('Description') or (next_text.startswith('"RDNA') and 'Description' in next_text[:200]):
+          format_headers.append((m.group(1), i, m.start()))
+
+  # parse instruction formats
+  formats: dict[str, list] = {}
+  for fmt_name, page_idx, header_pos in format_headers:
+    if fmt_name in formats: continue
+    text, tables = page_texts[page_idx], page_tables[page_idx]
+    field_pos = text.find('Field Name', header_pos)
+
+    # find fields table with ENCODING (same page or up to 2 pages ahead)
+    fields = None
+    for offset in range(3):
+      if page_idx + offset >= len(pages): break
+      if offset > 0 and has_header_before_fields(page_texts[page_idx + offset]): break
+      for t in page_tables[page_idx + offset] if offset > 0 or field_pos > header_pos else []:
+        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)) and has_encoding(f):
+          fields = f
+          break
+      if fields: break
+
+    # for modifier formats (no ENCODING), accept first fields table on same page
+    if not fields and field_pos > header_pos:
+      for t in tables:
+        if is_fields_table(t) and (f := parse_fields_table(t, fmt_name, enum_names)):
+          fields = f
+          break
+
+    if not fields: continue
+    field_names = {f[0] for f in fields}
+
+    # check next pages for continuation fields (tables without ENCODING)
+    for pg_offset in range(1, 3):
+      if page_idx + pg_offset >= len(pages) or has_header_before_fields(page_texts[page_idx + pg_offset]): break
+      for t in page_tables[page_idx + pg_offset]:
+        if is_fields_table(t) and (extra := parse_fields_table(t, fmt_name, enum_names)) and not has_encoding(extra):
+          for ef in extra:
+            if ef[0] not in field_names:
+              fields.append(ef)
+              field_names.add(ef[0])
+          break
+    formats[fmt_name] = fields
+
+  # fix known PDF errors (verified against LLVM test vectors)
+  # SMEM: PDF says DLC=bit14, GLC=bit16 but actual encoding is DLC=bit13, GLC=bit14
+  if 'SMEM' in formats:
+    formats['SMEM'] = [(n, 13 if n == 'DLC' else 14 if n == 'GLC' else h, 13 if n == 'DLC' else 14 if n == 'GLC' else l, e, t)
+                       for n, h, l, e, t in formats['SMEM']]
+
+  # generate output
+  def enum_lines(name, items):
+    return [f"class {name}(IntEnum):"] + [f"  {n} = {v}" for v, n in sorted(items.items())] + [""]
+  def field_key(f): return order.index(f[0].lower()) if f[0].lower() in order else 1000
+  lines = ["# autogenerated from AMD RDNA3.5 ISA PDF by gen.py - do not edit", "from enum import IntEnum",
+           "from typing import Annotated",
+           "from extra.assembly.rdna3.lib import bits, BitField, Inst32, Inst64, SGPR, VGPR, TTMP as TTMP, s as s, v as v, ttmp as ttmp, SSrc, Src, SImm, Imm, VDSTYEnc, SGPRField, VGPRField",
+           "import functools", ""]
+  lines += enum_lines("SrcEnum", src_enum) + sum([enum_lines(n, ops) for n, ops in sorted(enums.items())], [])
+  # Format-specific field defaults (verified against LLVM test vectors)
+  format_defaults = {'VOP3P': {'opsel_hi': 3, 'opsel_hi2': 1}}
+  lines.append("# instruction formats")
+  for fmt_name, fields in sorted(formats.items()):
+    base = "Inst64" if max(f[1] for f in fields) > 31 or fmt_name == 'VOP3SD' else "Inst32"
+    order = FIELD_ORDER.get(fmt_name, [])
+    lines.append(f"class {fmt_name}({base}):")
+    if enc := next((f for f in fields if f[0] == 'ENCODING'), None):
+      enc_str = f"bits[{enc[1]}:{enc[2]}] == 0b{enc[3]:b}" if enc[1] != enc[2] else f"bits[{enc[1]}] == {enc[3]}"
+      lines.append(f"  encoding = {enc_str}")
+    if defaults := format_defaults.get(fmt_name):
+      lines.append(f"  _defaults = {defaults}")
+    for name, hi, lo, _, ftype in sorted([f for f in fields if f[0] != 'ENCODING'], key=field_key):
+      # Wrap IntEnum types (ending in Op) with Annotated[BitField, ...] for correct typing
+      if ftype and ftype.endswith('Op'):
+        ann = f":Annotated[BitField, {ftype}]"
+      else:
+        ann = f":{ftype}" if ftype else ""
+      lines.append(f"  {name.lower()}{ann} = bits[{hi}]" if hi == lo else f"  {name.lower()}{ann} = bits[{hi}:{lo}]")
+    lines.append("")
+  lines.append("# instruction helpers")
+  for cls_name, ops in sorted(enums.items()):
+    fmt = cls_name[:-2]
+    for op_val, name in sorted(ops.items()):
+      seg = {"GLOBAL": ", seg=2", "SCRATCH": ", seg=2"}.get(fmt, "")
+      tgt = {"GLOBAL": "FLAT, GLOBALOp", "SCRATCH": "FLAT, SCRATCHOp"}.get(fmt, f"{fmt}, {cls_name}")
+      if fmt in formats or fmt in ("GLOBAL", "SCRATCH"):
+        # VOP1/VOP2/VOPC get _e32 suffix, VOP3 promoted ops (< 512) get _e64 suffix
+        if fmt in ("VOP1", "VOP2", "VOPC"):
+          suffix = "_e32"
+        elif fmt == "VOP3" and op_val < 512:
+          suffix = "_e64"
+        else:
+          suffix = ""
+        lines.append(f"{name.lower()}{suffix} = functools.partial({tgt}.{name}{seg})")
+  # export SrcEnum values, but skip DPP8/DPP16 which conflict with class names
+  skip_exports = {'DPP8', 'DPP16'}
+  lines += [""] + [f"{name} = SrcEnum.{name}" for _, name in sorted(src_enum.items()) if name not in skip_exports] + ["OFF = NULL\n"]
+
+  if output_path is not None: pathlib.Path(output_path).write_text('\n'.join(lines))
+  return {"formats": formats, "enums": enums, "src_enum": src_enum}
+
+if __name__ == "__main__":
+  result = generate("extra/assembly/rdna3/autogen/__init__.py")
+  print(f"generated SrcEnum ({len(result['src_enum'])}) + {len(result['enums'])} opcode enums + {len(result['formats'])} format classes")

extra/assembly/rdna3/generate.py

@@ -1,136 +0,0 @@
-import os, sys, struct
-sys.path.append(os.getcwd())
-# PROFILE=1 to use
-#os.environ["PROFILE"] = "1"
-os.environ["SQTT"] = "1"
-os.environ["SQTT_ITRACE_SE_MASK"] = "1"
-os.environ["SQTT_LIMIT_SE"] = "1"
-import xml.etree.ElementTree as ET
-
-from tinygrad import nn, Tensor, Device
-from tinygrad.helpers import get_single_element
-from tinygrad.engine.realize import lower_schedule
-from tinygrad.runtime.support.elf import elf_loader
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from extra.sqtt.attempt_sqtt_parse import parse_sqtt_print_packets
-
-def disassemble(text, root:ET.Element):
-  i = 0
-  while i < len(text):
-    ins = struct.unpack("I", text[i:i+4])[0]
-
-    # 1. Get the encoding
-    did_match = False
-    for enc_el in root.findall("./ISA/Encodings/Encoding"):
-      mask = enc_el.findtext("EncodingIdentifierMask")
-      assert len(mask)%32 == 0
-      bit_mask = int(mask, 2)
-      iden = [int(x.text, 2) for x in enc_el.find("EncodingIdentifiers").findall("EncodingIdentifier")]
-      for ide in iden:
-        if ins&bit_mask == ide:
-          did_match = True
-          break
-      if did_match: break
-    if not did_match: raise RuntimeError(f"unknown instruction {ins:08X}")
-    if len(mask) >= 64: ins = (struct.unpack("I", text[i+4:i+8])[0]<<32) | ins
-    if len(mask) >= 96: ins = (struct.unpack("I", text[i+8:i+12])[0]<<64) | ins
-    encoding_name = enc_el.findtext("EncodingName")
-
-    #print(ET.tostring(enc_el).decode())
-
-    # 2. Parse the Fields for this Encoding
-    field_data = {}
-    for field in enc_el.findall("MicrocodeFormat/BitMap/Field"):
-      # Fields can be split into multiple ranges (RangeCount > 1)
-      ranges = sorted(field.findall("BitLayout/Range"), key=lambda x: int(x.attrib.get('Order')))
-      val = 0
-      current_shift = 0
-      for rng in ranges:
-        width = int(rng.find("BitCount").text)
-        chunk = (ins >> int(rng.find("BitOffset").text)) & ((1 << width) - 1)
-        val |= (chunk << current_shift)
-        current_shift += width
-      field_data[field.find("FieldName").text] = val
-    # this is already used
-    del field_data["ENCODING"]
-
-    # 3. Extract the instruction
-    did_match = False
-    for ins_el in root.findall("./ISA/Instructions/Instruction"):
-      ins_name = ins_el.findtext("InstructionName")
-      for ins_enc in ins_el.findall("InstructionEncodings/InstructionEncoding"):
-        if ins_enc.findtext("EncodingName") == encoding_name:
-          opcode = int(ins_enc.findtext("Opcode"))
-          if "OP" in field_data and opcode == field_data["OP"]:
-            did_match = True
-            del field_data["OP"]
-            break
-        if did_match: break
-      if did_match: break
-
-    #print(ET.tostring(ins_enc).decode())
-    #print()
-    #print(field_data)
-    if not did_match:
-      print(f"{i:4X} : {ins:16x} -- {encoding_name}")
-    elif did_match:
-      params = []
-      #print(ET.tostring(ins_el).decode())
-
-      # 4. Extract the opcodes
-      for op_ins in ins_enc.findall("Operands/Operand"):
-        op_type = op_ins.findtext("OperandType")
-        op_size = op_ins.findtext("OperandSize")
-        op_fmt = op_ins.findtext("DataFormatName")
-        op_field_name = op_ins.findtext("FieldName")
-        if op_field_name is None: continue
-        assert op_field_name in field_data
-        # loop through operands for compare
-        for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-          test_op_type = op_el.findtext("OperandTypeName")
-          val_dict = {}
-          for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-            val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-          if op_type == test_op_type:
-            if field_data[op_field_name] in val_dict:
-              print(op_type, op_size, op_fmt)
-              params.append(val_dict[field_data[op_field_name]])
-            else:
-              params.append(f"{op_type}({field_data[op_field_name]})")
-            del field_data[op_field_name]
-            #print(op_type, op_size, op_fmt, op_el, op_field_name,
-            #      field_data[op_field_name],
-            #      val_dict.get(field_data[op_field_name], "<UNK>"))
-            #print(ET.tostring(op_el).decode())
-
-      print(f"{i:4X} : {ins:16x} -- {ins_name.lower()} {', '.join(params)}", field_data)
-
-    # advance
-    i += len(mask) // 8
-
-  #print(ET.tostring(root).decode())
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  with open("/tmp/rdna35.xml", "wb") as f: f.write(bytes(xml_str))
-  root = ET.fromstring(xml_str)
-
-  a = Tensor.empty(16)+1
-  for si, ei in lower_schedule(a.schedule()):
-    # get text
-    _, hdr, _ = elf_loader(ei.prg.lib)
-    text = get_single_element([x for x in hdr if x.name==".text"]).content
-
-    # llvm disassembler
-    Device["AMD"].compiler.disassemble(ei.prg.lib)
-
-    # run program
-    ei.run()
-
-  sqtt_events = [e for e in Device["AMD"].profile_events if isinstance(e, ProfileSQTTEvent)]
-  for e in sqtt_events[0:1]: # only the first SE
-    parse_sqtt_print_packets(e.blob)
-
-  disassemble(text[:0x40], root)

extra/assembly/rdna3/lib.py

@@ -0,0 +1,254 @@
+# library for RDNA3 assembly DSL
+from __future__ import annotations
+from enum import IntEnum
+from typing import overload, Annotated, TypeVar, Generic
+
+# Bit field DSL
+class BitField:
+  def __init__(self, hi: int, lo: int, name: str | None = None): self.hi, self.lo, self.name = hi, lo, name
+  def __set_name__(self, owner, name): self.name, self._owner = name, owner
+  def __eq__(self, val: int) -> tuple[BitField, int]: return (self, val)  # type: ignore
+  def mask(self) -> int: return (1 << (self.hi - self.lo + 1)) - 1
+  @property
+  def marker(self) -> type | None:
+    # Get marker from Annotated type hint if present
+    import typing
+    if hasattr(self, '_owner') and self.name:
+      hints = typing.get_type_hints(self._owner, include_extras=True)
+      if self.name in hints:
+        hint = hints[self.name]
+        if typing.get_origin(hint) is Annotated:
+          args = typing.get_args(hint)
+          return args[1] if len(args) > 1 else None
+    return None
+  @overload
+  def __get__(self, obj: None, objtype: type) -> BitField: ...
+  @overload
+  def __get__(self, obj: object, objtype: type | None = None) -> int: ...
+  def __get__(self, obj, objtype=None):
+    if obj is None: return self
+    val = unwrap(obj._values.get(self.name, 0))
+    # Convert to IntEnum if marker is an IntEnum subclass
+    if self.marker and isinstance(self.marker, type) and issubclass(self.marker, IntEnum):
+      try: return self.marker(val)
+      except ValueError: pass
+    return val
+
+class _Bits:
+  def __getitem__(self, key) -> BitField: return BitField(key.start, key.stop) if isinstance(key, slice) else BitField(key, key)
+bits = _Bits()
+
+# Register types
+class Reg:
+  def __init__(self, idx: int, count: int = 1, hi: bool = False, neg: bool = False): self.idx, self.count, self.hi, self.neg = idx, count, hi, neg
+  def __repr__(self): return f"{self.__class__.__name__.lower()[0]}[{self.idx}]" if self.count == 1 else f"{self.__class__.__name__.lower()[0]}[{self.idx}:{self.idx + self.count}]"
+  def __neg__(self): return self.__class__(self.idx, self.count, self.hi, neg=not self.neg)
+
+T = TypeVar('T', bound=Reg)
+class _RegFactory(Generic[T]):
+  def __init__(self, cls: type[T], name: str): self._cls, self._name = cls, name
+  @overload
+  def __getitem__(self, key: int) -> Reg: ...
+  @overload
+  def __getitem__(self, key: slice) -> Reg: ...
+  def __getitem__(self, key: int | slice) -> Reg:
+    return self._cls(key.start, key.stop - key.start + 1) if isinstance(key, slice) else self._cls(key)
+  def __repr__(self): return f"<{self._name} factory>"
+
+class SGPR(Reg): pass
+class VGPR(Reg): pass
+class TTMP(Reg): pass
+s: _RegFactory[SGPR] = _RegFactory(SGPR, "SGPR")
+v: _RegFactory[VGPR] = _RegFactory(VGPR, "VGPR")
+ttmp: _RegFactory[TTMP] = _RegFactory(TTMP, "TTMP")
+
+# Field type markers (runtime classes for validation)
+class _SSrc: pass
+class _Src: pass
+class _Imm: pass
+class _SImm: pass
+class _VDSTYEnc: pass  # VOPD vdsty: encoded = actual >> 1, actual = (encoded << 1) | ((vdstx & 1) ^ 1)
+class _SGPRField: pass
+class _VGPRField: pass
+
+# Type aliases for annotations - tells mypy it's a BitField while preserving marker info
+SSrc = Annotated[BitField, _SSrc]
+Src = Annotated[BitField, _Src]
+Imm = Annotated[BitField, _Imm]
+SImm = Annotated[BitField, _SImm]
+VDSTYEnc = Annotated[BitField, _VDSTYEnc]
+SGPRField = Annotated[BitField, _SGPRField]
+VGPRField = Annotated[BitField, _VGPRField]
+class RawImm:
+  def __init__(self, val: int): self.val = val
+  def __repr__(self): return f"RawImm({self.val})"
+  def __eq__(self, other): return isinstance(other, RawImm) and self.val == other.val
+
+def unwrap(val) -> int:
+  return val.val if isinstance(val, RawImm) else val.value if hasattr(val, 'value') else val.idx if hasattr(val, 'idx') else val
+
+# Encoding helpers
+FLOAT_ENC = {0.5: 240, -0.5: 241, 1.0: 242, -1.0: 243, 2.0: 244, -2.0: 245, 4.0: 246, -4.0: 247}
+SRC_FIELDS = {'src0', 'src1', 'src2', 'ssrc0', 'ssrc1', 'soffset', 'srcx0', 'srcy0'}
+RAW_FIELDS = {'vdata', 'vdst', 'vaddr', 'addr', 'data', 'data0', 'data1', 'sdst', 'sdata'}
+
+def _encode_reg(val) -> int:
+  if isinstance(val, TTMP): return 108 + val.idx
+  return val.idx | (0x80 if val.hi else 0)
+
+def encode_src(val) -> int:
+  if isinstance(val, VGPR): return 256 + _encode_reg(val)
+  if isinstance(val, Reg): return _encode_reg(val)
+  if hasattr(val, 'value'): return val.value
+  if isinstance(val, float): return 128 if val == 0.0 else FLOAT_ENC.get(val, 255)
+  return 128 + val if isinstance(val, int) and 0 <= val <= 64 else 192 + (-val) if isinstance(val, int) and -16 <= val <= -1 else 255
+
+# Instruction base class
+class Inst:
+  _fields: dict[str, BitField]
+  _encoding: tuple[BitField, int] | None = None
+  _defaults: dict[str, int] = {}
+  _values: dict[str, int | RawImm]
+  _words: int  # size in 32-bit words, set by decode_program
+  _literal: int | None
+
+  def __init_subclass__(cls, **kwargs):
+    super().__init_subclass__(**kwargs)
+    cls._fields = {n: v[0] if isinstance(v, tuple) else v for n, v in cls.__dict__.items() if isinstance(v, BitField) or (isinstance(v, tuple) and len(v) == 2 and isinstance(v[0], BitField))}
+    if 'encoding' in cls._fields and isinstance(cls.__dict__.get('encoding'), tuple): cls._encoding = cls.__dict__['encoding']
+
+  def __init__(self, *args, literal: int | None = None, **kwargs):
+    self._values, self._literal = dict(self._defaults), literal
+    # Map positional args to field names
+    field_names = [n for n in self._fields if n != 'encoding']
+    orig_args = dict(zip(field_names, args))
+    orig_args.update(kwargs)
+    self._values.update(orig_args)
+    # Validate register counts for SMEM instructions (before encoding)
+    if self.__class__.__name__ == 'SMEM':
+      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
+      if op_val is not None:
+        if hasattr(op_val, 'value'): op_val = op_val.value
+        expected_cnt = {0:1, 1:2, 2:4, 3:8, 4:16, 8:1, 9:2, 10:4, 11:8, 12:16}.get(op_val)
+        sdata_val = orig_args.get('sdata')
+        if expected_cnt is not None and isinstance(sdata_val, Reg) and sdata_val.count != expected_cnt:
+          raise ValueError(f"SMEM op {op_val} expects {expected_cnt} registers, got {sdata_val.count}")
+    # Validate register counts for SOP1 instructions (b32 = 1 reg, b64 = 2 regs)
+    if self.__class__.__name__ == 'SOP1':
+      op_val = orig_args.get(field_names[0]) if args else orig_args.get('op')
+      if op_val is not None and hasattr(op_val, 'name'):
+        expected = 2 if op_val.name.endswith('_B64') else 1
+        sdst_val, ssrc0_val = orig_args.get('sdst'), orig_args.get('ssrc0')
+        if isinstance(sdst_val, Reg) and sdst_val.count != expected:
+          raise ValueError(f"SOP1 {op_val.name} expects {expected} destination register(s), got {sdst_val.count}")
+        if isinstance(ssrc0_val, Reg) and ssrc0_val.count != expected:
+          raise ValueError(f"SOP1 {op_val.name} expects {expected} source register(s), got {ssrc0_val.count}")
+    # Type check and encode values
+    for name, val in list(self._values.items()):
+      if name == 'encoding': continue
+      # For RawImm, only process RAW_FIELDS to unwrap to int
+      if isinstance(val, RawImm):
+        if name in RAW_FIELDS: self._values[name] = val.val
+        continue
+      field = self._fields.get(name)
+      marker = field.marker if field else None
+      # Type validation
+      if marker is _SGPRField:
+        if isinstance(val, VGPR): raise TypeError(f"field '{name}' requires SGPR, got VGPR")
+        if not isinstance(val, (SGPR, TTMP, int, RawImm)): raise TypeError(f"field '{name}' requires SGPR, got {type(val).__name__}")
+      if marker is _VGPRField:
+        if not isinstance(val, VGPR): raise TypeError(f"field '{name}' requires VGPR, got {type(val).__name__}")
+      if marker is _SSrc and isinstance(val, VGPR): raise TypeError(f"field '{name}' requires scalar source, got VGPR")
+      # Encode source fields as RawImm for consistent disassembly
+      if name in SRC_FIELDS:
+        encoded = encode_src(val)
+        self._values[name] = RawImm(encoded)
+        # Handle negation modifier for VOP3 instructions
+        if isinstance(val, Reg) and val.neg and 'neg' in self._fields:
+          neg_bit = {'src0': 1, 'src1': 2, 'src2': 4}.get(name, 0)
+          cur_neg = self._values.get('neg', 0)
+          self._values['neg'] = (cur_neg.val if isinstance(cur_neg, RawImm) else cur_neg) | neg_bit
+        # Track literal value if needed (encoded as 255)
+        if encoded == 255 and self._literal is None and isinstance(val, int) and not isinstance(val, IntEnum):
+          self._literal = val
+        elif encoded == 255 and self._literal is None and isinstance(val, float):
+          import struct
+          self._literal = struct.unpack('<I', struct.pack('<f', val))[0]
+      # Encode raw register fields for consistent repr
+      elif name in RAW_FIELDS:
+        if isinstance(val, Reg): self._values[name] = _encode_reg(val)
+        elif hasattr(val, 'value'): self._values[name] = val.value  # IntEnum like SrcEnum.NULL
+      # Encode sbase (divided by 2) and srsrc/ssamp (divided by 4)
+      elif name == 'sbase' and isinstance(val, Reg):
+        self._values[name] = val.idx // 2
+      elif name in {'srsrc', 'ssamp'} and isinstance(val, Reg):
+        self._values[name] = val.idx // 4
+      # VOPD vdsty: encode as actual >> 1 (constraint: vdsty parity must be opposite of vdstx)
+      elif marker is _VDSTYEnc and isinstance(val, VGPR):
+        self._values[name] = val.idx >> 1
+
+  def _encode_field(self, name: str, val) -> int:
+    if isinstance(val, RawImm): return val.val
+    if name in {'srsrc', 'ssamp'}: return val.idx // 4 if isinstance(val, Reg) else val
+    if name == 'sbase': return val.idx // 2 if isinstance(val, Reg) else val
+    if name in RAW_FIELDS: return _encode_reg(val) if isinstance(val, Reg) else val
+    if isinstance(val, Reg) or name in SRC_FIELDS: return encode_src(val)
+    return val.value if hasattr(val, 'value') else val
+
+  def to_int(self) -> int:
+    word = (self._encoding[1] & self._encoding[0].mask()) << self._encoding[0].lo if self._encoding else 0
+    for n, bf in self._fields.items():
+      if n != 'encoding' and n in self._values: word |= (self._encode_field(n, self._values[n]) & bf.mask()) << bf.lo
+    return word
+
+  def _get_literal(self) -> int | None:
+    for n in SRC_FIELDS:
+      if n in self._values and not isinstance(v := self._values[n], RawImm) and isinstance(v, int) and not isinstance(v, IntEnum) and not (0 <= v <= 64 or -16 <= v <= -1): return v
+    return None
+
+  def to_bytes(self) -> bytes:
+    result = self.to_int().to_bytes(self._size(), 'little')
+    return result + (lit & 0xffffffff).to_bytes(4, 'little') if (lit := self._get_literal() or getattr(self, '_literal', None)) else result
+
+  @classmethod
+  def _size(cls) -> int: return 4 if issubclass(cls, Inst32) else 8
+  def size(self) -> int: return self._size() + (4 if self._literal is not None else 0)
+
+  @classmethod
+  def from_int(cls, word: int):
+    inst = object.__new__(cls)
+    inst._values = {n: RawImm(v) if n in SRC_FIELDS else v for n, bf in cls._fields.items() if n != 'encoding' for v in [(word >> bf.lo) & bf.mask()]}
+    inst._literal = None
+    return inst
+
+  @classmethod
+  def from_bytes(cls, data: bytes):
+    inst = cls.from_int(int.from_bytes(data[:cls._size()], 'little'))
+    op_val = inst._values.get('op', 0)
+    has_literal = cls.__name__ == 'VOP2' and op_val in (44, 45, 55, 56)
+    has_literal = has_literal or (cls.__name__ == 'SOP2' and op_val in (69, 70))
+    for n in SRC_FIELDS:
+      if n in inst._values and isinstance(inst._values[n], RawImm) and inst._values[n].val == 255: has_literal = True
+    if has_literal and len(data) >= cls._size() + 4: inst._literal = int.from_bytes(data[cls._size():cls._size()+4], 'little')
+    return inst
+
+  def __repr__(self):
+    # Use _fields order and exclude fields that are 0/default (for consistent repr after roundtrip)
+    def is_zero(v): return (isinstance(v, int) and v == 0) or (isinstance(v, VGPR) and v.idx == 0 and v.count == 1)
+    items = [(k, self._values[k]) for k in self._fields if k in self._values and k != 'encoding'
+             and not (is_zero(self._values[k]) and k not in {'op'})]
+    lit = f", literal={hex(self._literal)}" if self._literal is not None else ""
+    return f"{self.__class__.__name__}({', '.join(f'{k}={v}' for k, v in items)}{lit})"
+
+  def __eq__(self, other):
+    if not isinstance(other, Inst): return NotImplemented
+    return self.__class__ == other.__class__ and self._values == other._values and self._literal == other._literal
+
+  def __hash__(self): return hash((self.__class__.__name__, tuple(sorted((k, repr(v)) for k, v in self._values.items())), self._literal))
+
+  def disasm(self) -> str:
+    from extra.assembly.rdna3.asm import disasm
+    return disasm(self)
+
+class Inst32(Inst): pass
+class Inst64(Inst): pass

extra/assembly/rdna3/parse.py

@@ -1,15 +0,0 @@
-from tinygrad import Tensor, nn
-import xml.etree.ElementTree as ET
-
-if __name__ == "__main__":
-  # human readable manual at https://docs.amd.com/v/u/en-US/rdna35_instruction_set_architecture
-  fns = nn.state.zip_extract(Tensor.from_url("https://gpuopen.com/download/machine-readable-isa/latest/"))
-  xml_str = fns['amdgpu_isa_rdna3_5.xml'].to("CPU").data()
-  root = ET.fromstring(xml_str)
-
-  for op_el in root.findall("./ISA/OperandTypes/OperandType"):
-    op_name = op_el.findtext("OperandTypeName")
-    val_dict = {}
-    for op_val in op_el.findall("OperandPredefinedValues/PredefinedValue"):
-      val_dict[int(op_val.findtext("Value"))] = op_val.findtext("Name")
-    print(op_name, val_dict)

extra/assembly/rdna3/test/bench_emu.py

@@ -0,0 +1,294 @@
+#!/usr/bin/env python3
+"""Benchmark comparing Python vs Rust RDNA3 emulators on synthetic and real tinygrad kernels."""
+import ctypes, time, os, struct, cProfile, pstats, io
+from pathlib import Path
+from typing import Callable
+
+# Set AMD=1 before importing tinygrad
+os.environ["AMD"] = "1"
+
+from extra.assembly.rdna3.emu import run_asm as python_run_asm, set_valid_mem_ranges, decode_program, step_wave, WaveState, WAVE_SIZE
+
+REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
+if not REMU_PATH.exists():
+  REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.dylib"
+
+def get_rust_remu():
+  """Load the Rust libremu shared library."""
+  if not REMU_PATH.exists(): return None
+  remu = ctypes.CDLL(str(REMU_PATH))
+  remu.run_asm.restype = ctypes.c_int32
+  remu.run_asm.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32,
+                           ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_void_p]
+  return remu
+
+def count_instructions(kernel: bytes) -> int:
+  """Count instructions in a kernel."""
+  return len(decode_program(kernel))
+
+def setup_buffers(buf_sizes: list[int], init_data: dict[int, bytes] | None = None):
+  """Allocate buffers and return args pointer + valid ranges."""
+  if init_data is None: init_data = {}
+  buffers = []
+  for i, size in enumerate(buf_sizes):
+    padded = ((size + 15) // 16) * 16 + 16
+    data = init_data.get(i, b'\x00' * padded)
+    data_list = list(data) + [0] * (padded - len(data))
+    buf = (ctypes.c_uint8 * padded)(*data_list[:padded])
+    buffers.append(buf)
+  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
+  args_ptr = ctypes.addressof(args)
+  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
+  ranges.add((args_ptr, ctypes.sizeof(args)))
+  return buffers, args, args_ptr, ranges
+
+def benchmark_emulator(name: str, run_fn, kernel: bytes, global_size, local_size, args_ptr, iterations: int = 5):
+  """Benchmark an emulator and return average time."""
+  gx, gy, gz = global_size
+  lx, ly, lz = local_size
+  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+  lib_ptr = ctypes.addressof(kernel_buf)
+
+  # Warmup
+  run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
+
+  # Timed runs
+  times = []
+  for _ in range(iterations):
+    start = time.perf_counter()
+    result = run_fn(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
+    end = time.perf_counter()
+    if result != 0:
+      print(f"  {name} returned error: {result}")
+      return None
+    times.append(end - start)
+
+  return sum(times) / len(times)
+
+def create_synthetic_kernel(n_ops: int) -> bytes:
+  """Create a synthetic kernel with n_ops vector operations."""
+  instructions = []
+  # VOP2 instructions: v_add_f32, v_mul_f32, v_max_f32, v_min_f32
+  ops = [
+    (0b0000011 << 25) | (1 << 17) | (0 << 9) | 256,  # v_add_f32 v0, v0, v1
+    (0b0001000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_mul_f32 v0, v0, v1
+    (0b0010000 << 25) | (1 << 17) | (0 << 9) | 256,  # v_max_f32 v0, v0, v1
+    (0b0001111 << 25) | (1 << 17) | (0 << 9) | 256,  # v_min_f32 v0, v0, v1
+  ]
+  for i in range(n_ops):
+    instructions.append(ops[i % len(ops)])
+  # S_ENDPGM
+  instructions.append((0b101111111 << 23) | (48 << 16) | 0)
+  return b''.join(struct.pack('<I', inst) for inst in instructions)
+
+def get_tinygrad_kernel(op_name: str) -> tuple[bytes, tuple, tuple, list[int], dict[int, bytes]] | None:
+  """Get a real tinygrad kernel by operation name. Returns (code, global_size, local_size, buf_sizes, buf_data)."""
+  try:
+    from tinygrad import Tensor
+    from tinygrad.runtime.support.elf import elf_loader
+    import numpy as np
+    np.random.seed(42)
+
+    ops = {
+      "add": lambda: Tensor.empty(1024) + Tensor.empty(1024),
+      "mul": lambda: Tensor.empty(1024) * Tensor.empty(1024),
+      "matmul_small": lambda: Tensor.empty(16, 16) @ Tensor.empty(16, 16),
+      "matmul_medium": lambda: Tensor.empty(64, 64) @ Tensor.empty(64, 64),
+      "reduce_sum": lambda: Tensor.empty(4096).sum(),
+      "reduce_max": lambda: Tensor.empty(4096).max(),
+      "softmax": lambda: Tensor.empty(256).softmax(),
+      "layernorm": lambda: Tensor.empty(32, 64).layernorm(),
+      "conv2d": lambda: Tensor.empty(1, 4, 16, 16).conv2d(Tensor.empty(4, 4, 3, 3)),
+      "gelu": lambda: Tensor.empty(1024).gelu(),
+      "exp": lambda: Tensor.empty(1024).exp(),
+      "sin": lambda: Tensor.empty(1024).sin(),
+    }
+
+    if op_name not in ops: return None
+    out = ops[op_name]()
+    sched = out.schedule()
+
+    for ei in sched:
+      lowered = ei.lower()
+      if ei.ast.op.name == 'SINK' and lowered.prg and lowered.prg.p.lib:
+        lib = bytes(lowered.prg.p.lib)
+        _, sections, _ = elf_loader(lib)
+        for sec in sections:
+          if sec.name == '.text':
+            buf_sizes = [b.nbytes for b in lowered.bufs]
+            # Get initial data from numpy arrays if available
+            buf_data = {}
+            for i, buf in enumerate(lowered.bufs):
+              if hasattr(buf, 'base') and buf.base is not None and hasattr(buf.base, '_buf'):
+                try: buf_data[i] = bytes(buf.base._buf)
+                except: pass
+            return (bytes(sec.content), tuple(lowered.prg.p.global_size), tuple(lowered.prg.p.local_size), buf_sizes, buf_data)
+    return None
+  except Exception as e:
+    print(f"  Error getting kernel: {e}")
+    return None
+
+def profile_python_emu(kernel: bytes, global_size, local_size, args_ptr, n_runs: int = 1):
+  """Profile the Python emulator to find bottlenecks."""
+  gx, gy, gz = global_size
+  lx, ly, lz = local_size
+  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+  lib_ptr = ctypes.addressof(kernel_buf)
+
+  pr = cProfile.Profile()
+  pr.enable()
+  for _ in range(n_runs):
+    python_run_asm(lib_ptr, len(kernel), gx, gy, gz, lx, ly, lz, args_ptr)
+  pr.disable()
+
+  s = io.StringIO()
+  ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
+  ps.print_stats(20)
+  return s.getvalue()
+
+def measure_step_rate(kernel: bytes, n_steps: int = 10000) -> float:
+  """Measure raw step_wave() performance (steps per second)."""
+  program = decode_program(kernel)
+  if not program: return 0.0
+
+  st = WaveState()
+  st.exec_mask = 0xffffffff
+  lds = bytearray(65536)
+  n_lanes = 32
+
+  # Reset PC for each measurement
+  start = time.perf_counter()
+  for _ in range(n_steps):
+    st.pc = 0
+    while st.pc in program:
+      result = step_wave(program, st, lds, n_lanes)
+      if result == -1: break
+  elapsed = time.perf_counter() - start
+  return n_steps / elapsed if elapsed > 0 else 0
+
+# Test configurations
+SYNTHETIC_TESTS = [
+  ("synthetic_10ops", 10, (1, 1, 1), (32, 1, 1)),
+  ("synthetic_100ops", 100, (1, 1, 1), (32, 1, 1)),
+  ("synthetic_500ops", 500, (1, 1, 1), (32, 1, 1)),
+  ("synthetic_100ops_4wg", 100, (4, 1, 1), (32, 1, 1)),
+  ("synthetic_100ops_16wg", 100, (16, 1, 1), (32, 1, 1)),
+]
+
+TINYGRAD_TESTS = ["add", "mul", "reduce_sum", "softmax", "exp", "gelu", "matmul_small"]
+
+def main():
+  import argparse
+  parser = argparse.ArgumentParser(description="Benchmark RDNA3 emulators")
+  parser.add_argument("--profile", action="store_true", help="Profile Python emulator")
+  parser.add_argument("--synthetic-only", action="store_true", help="Only run synthetic tests")
+  parser.add_argument("--tinygrad-only", action="store_true", help="Only run tinygrad tests")
+  parser.add_argument("--iterations", type=int, default=3, help="Number of iterations per benchmark")
+  args = parser.parse_args()
+
+  rust_remu = get_rust_remu()
+  if rust_remu is None:
+    print("Rust libremu not found. Build with: cargo build --release --manifest-path extra/remu/Cargo.toml")
+    print("Running Python-only benchmarks...\n")
+
+  print("=" * 90)
+  print("RDNA3 Emulator Benchmark: Python vs Rust")
+  print("=" * 90)
+
+  results = []
+
+  # Synthetic workloads
+  if not args.tinygrad_only:
+    print("\n[SYNTHETIC WORKLOADS]")
+    print("-" * 90)
+
+    for name, n_ops, global_size, local_size in SYNTHETIC_TESTS:
+      kernel = create_synthetic_kernel(n_ops)
+      n_insts = count_instructions(kernel)
+      n_workgroups = global_size[0] * global_size[1] * global_size[2]
+      n_threads = local_size[0] * local_size[1] * local_size[2]
+      total_work = n_insts * n_workgroups * n_threads
+
+      print(f"\n{name}: {n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
+
+      buf_sizes = [4096]
+      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes)
+      set_valid_mem_ranges(ranges)
+
+      # Benchmark
+      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
+      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
+
+      if py_time:
+        py_rate = total_work / py_time / 1e6
+        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
+      if rust_time:
+        rust_rate = total_work / rust_time / 1e6
+        speedup = py_time / rust_time if py_time else 0
+        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
+
+      results.append(("synthetic", name, n_insts, n_workgroups, py_time, rust_time))
+
+  # Tinygrad kernels
+  if not args.synthetic_only:
+    print("\n[TINYGRAD KERNELS]")
+    print("-" * 90)
+
+    for op_name in TINYGRAD_TESTS:
+      print(f"\n{op_name}:", end=" ", flush=True)
+      kernel_info = get_tinygrad_kernel(op_name)
+      if kernel_info is None:
+        print("failed to compile")
+        continue
+
+      kernel, global_size, local_size, buf_sizes, buf_data = kernel_info
+      n_insts = count_instructions(kernel)
+      n_workgroups = global_size[0] * global_size[1] * global_size[2]
+      n_threads = local_size[0] * local_size[1] * local_size[2]
+      total_work = n_insts * n_workgroups * n_threads
+
+      print(f"{n_insts} insts × {n_workgroups} WGs × {n_threads} threads = {total_work:,} ops")
+
+      buffers, args_arr, args_ptr, ranges = setup_buffers(buf_sizes, buf_data)
+      set_valid_mem_ranges(ranges)
+
+      py_time = benchmark_emulator("Python", python_run_asm, kernel, global_size, local_size, args_ptr, args.iterations)
+      rust_time = benchmark_emulator("Rust", rust_remu.run_asm, kernel, global_size, local_size, args_ptr, args.iterations) if rust_remu else None
+
+      if py_time:
+        py_rate = total_work / py_time / 1e6
+        print(f"  Python: {py_time*1000:8.3f} ms  ({py_rate:7.2f} M ops/s)")
+      if rust_time:
+        rust_rate = total_work / rust_time / 1e6
+        speedup = py_time / rust_time if py_time else 0
+        print(f"  Rust:   {rust_time*1000:8.3f} ms  ({rust_rate:7.2f} M ops/s)  [{speedup:.1f}x faster]")
+
+      results.append(("tinygrad", op_name, n_insts, n_workgroups, py_time, rust_time))
+
+      # Optional profiling
+      if args.profile and py_time:
+        print("\n  [PROFILE - Top 10 functions]")
+        profile_output = profile_python_emu(kernel, global_size, local_size, args_ptr)
+        for line in profile_output.split('\n')[5:15]:
+          if line.strip(): print(f"    {line}")
+
+  # Summary table
+  print("\n" + "=" * 90)
+  print("SUMMARY")
+  print("=" * 90)
+  print(f"{'Type':<10} {'Name':<25} {'Insts':<8} {'WGs':<6} {'Python (ms)':<14} {'Rust (ms)':<14} {'Speedup':<10}")
+  print("-" * 90)
+
+  for test_type, name, n_insts, n_wgs, py_time, rust_time in results:
+    py_ms = f"{py_time*1000:.3f}" if py_time else "error"
+    if rust_time:
+      rust_ms = f"{rust_time*1000:.3f}"
+      speedup = f"{py_time/rust_time:.1f}x" if py_time else "N/A"
+    else:
+      rust_ms, speedup = "N/A", "N/A"
+    print(f"{test_type:<10} {name:<25} {n_insts:<8} {n_wgs:<6} {py_ms:<14} {rust_ms:<14} {speedup:<10}")
+
+
+
+if __name__ == "__main__":
+  main()

extra/assembly/rdna3/test/test_compare_emulators.py

@@ -0,0 +1,466 @@
+# Test to compare Python and Rust RDNA3 emulators by running real tinygrad kernels
+import unittest, ctypes, os
+from dataclasses import dataclass
+from pathlib import Path
+
+# Set environment before any tinygrad imports to use MOCKGPU
+# This allows generating AMD GPU kernels without requiring real hardware
+os.environ["AMD"] = "1"
+os.environ["MOCKGPU"] = "1"
+
+from extra.assembly.rdna3.emu import WaveState, decode_program, step_wave, WAVE_SIZE
+
+REMU_PATH = Path(__file__).parents[3] / "remu/target/release/libremu.so"
+
+@dataclass
+class KernelInfo:
+  code: bytes
+  global_size: tuple[int, int, int]
+  local_size: tuple[int, int, int]
+  buf_idxs: list[int]  # indices into shared buffer pool
+  buf_sizes: list[int]  # sizes for each buffer index
+
+@dataclass
+class StateSnapshot:
+  pc: int
+  scc: int
+  vcc: int
+  exec_mask: int
+  sgpr: list[int]
+  vgpr: list[list[int]]
+
+  def diff(self, other: 'StateSnapshot', n_lanes: int) -> list[str]:
+    """Return list of differences between two states."""
+    diffs = []
+    if self.pc != other.pc: diffs.append(f"pc: {self.pc} vs {other.pc}")
+    if self.scc != other.scc: diffs.append(f"scc: {self.scc} vs {other.scc}")
+    if self.vcc != other.vcc: diffs.append(f"vcc: 0x{self.vcc:08x} vs 0x{other.vcc:08x}")
+    if self.exec_mask != other.exec_mask: diffs.append(f"exec: 0x{self.exec_mask:08x} vs 0x{other.exec_mask:08x}")
+    for i, (a, b) in enumerate(zip(self.sgpr, other.sgpr)):
+      # Skip VCC_LO/HI (106/107) and EXEC_LO/HI (126/127) as they alias vcc/exec_mask which are compared separately
+      if i in (106, 107, 126, 127): continue
+      if a != b: diffs.append(f"sgpr[{i}]: 0x{a:08x} vs 0x{b:08x}")
+    for lane in range(n_lanes):
+      for i, (a, b) in enumerate(zip(self.vgpr[lane], other.vgpr[lane])):
+        if a != b: diffs.append(f"vgpr[{lane}][{i}]: 0x{a:08x} vs 0x{b:08x}")
+    return diffs
+
+class CStateSnapshot(ctypes.Structure):
+  _fields_ = [("pc", ctypes.c_uint32), ("scc", ctypes.c_uint32), ("vcc", ctypes.c_uint32), ("exec_mask", ctypes.c_uint32),
+              ("sgpr", ctypes.c_uint32 * 128), ("vgpr", (ctypes.c_uint32 * 256) * 32)]
+
+  def to_snapshot(self) -> StateSnapshot:
+    return StateSnapshot(pc=self.pc, scc=self.scc, vcc=self.vcc, exec_mask=self.exec_mask,
+                         sgpr=list(self.sgpr), vgpr=[list(self.vgpr[i]) for i in range(32)])
+
+class RustEmulator:
+  def __init__(self):
+    self.lib = ctypes.CDLL(str(REMU_PATH))
+    self.lib.wave_create.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
+    self.lib.wave_create.restype = ctypes.c_void_p
+    self.lib.wave_step.argtypes = [ctypes.c_void_p]
+    self.lib.wave_step.restype = ctypes.c_int32
+    self.lib.wave_get_snapshot.argtypes = [ctypes.c_void_p, ctypes.POINTER(CStateSnapshot)]
+    self.lib.wave_set_sgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32]
+    self.lib.wave_set_vgpr.argtypes = [ctypes.c_void_p, ctypes.c_uint32, ctypes.c_uint32, ctypes.c_uint32]
+    self.lib.wave_init_lds.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
+    self.lib.wave_free.argtypes = [ctypes.c_void_p]
+    self.ctx = None
+
+  def create(self, kernel: bytes, n_lanes: int):
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    self.ctx = self.lib.wave_create(ctypes.addressof(kernel_buf), len(kernel), n_lanes)
+    self._kernel_buf = kernel_buf
+
+  def step(self) -> int: return self.lib.wave_step(self.ctx)
+  def set_sgpr(self, idx: int, val: int): self.lib.wave_set_sgpr(self.ctx, idx, val)
+  def set_vgpr(self, lane: int, idx: int, val: int): self.lib.wave_set_vgpr(self.ctx, lane, idx, val)
+  def init_lds(self, size: int): self.lib.wave_init_lds(self.ctx, size)
+
+  def get_snapshot(self) -> StateSnapshot:
+    snap = CStateSnapshot()
+    self.lib.wave_get_snapshot(self.ctx, ctypes.byref(snap))
+    return snap.to_snapshot()
+
+  def free(self):
+    if self.ctx: self.lib.wave_free(self.ctx); self.ctx = None
+
+class PythonEmulator:
+  def __init__(self):
+    self.state: WaveState | None = None
+    self.program: dict | None = None
+    self.lds: bytearray | None = None
+    self.n_lanes = 0
+
+  def create(self, kernel: bytes, n_lanes: int):
+    self.program = decode_program(kernel)
+    self.state = WaveState()
+    self.state.exec_mask = (1 << n_lanes) - 1
+    self.lds = bytearray(65536)
+    self.n_lanes = n_lanes
+
+  def step(self) -> int:
+    assert self.program is not None and self.state is not None and self.lds is not None
+    return step_wave(self.program, self.state, self.lds, self.n_lanes)
+  def set_sgpr(self, idx: int, val: int):
+    assert self.state is not None
+    self.state.sgpr[idx] = val & 0xffffffff
+  def set_vgpr(self, lane: int, idx: int, val: int):
+    assert self.state is not None
+    self.state.vgpr[lane][idx] = val & 0xffffffff
+
+  def get_snapshot(self) -> StateSnapshot:
+    assert self.state is not None
+    return StateSnapshot(pc=self.state.pc, scc=self.state.scc, vcc=self.state.vcc & 0xffffffff,
+                         exec_mask=self.state.exec_mask & 0xffffffff, sgpr=list(self.state.sgpr),
+                         vgpr=[list(self.state.vgpr[i]) for i in range(WAVE_SIZE)])
+
+def run_single_kernel(kernel: bytes, n_lanes: int, args_ptr: int, global_size: tuple[int, int, int],
+                      program, max_steps: int, debug: bool, trace_len: int, kernel_idx: int = 0,
+                      max_workgroups: int = 64) -> tuple[bool, str, int]:
+  """Run a single kernel through both emulators. Returns (success, message, total_steps)."""
+  gx, gy, gz = global_size
+  total_steps = 0
+  wg_count = 0
+
+  for gidz in range(gz):
+    for gidy in range(gy):
+      for gidx in range(gx):
+        if wg_count >= max_workgroups: return True, f"Completed {wg_count} workgroups (limit reached)", total_steps
+        wg_count += 1
+        rust = RustEmulator()
+        python = PythonEmulator()
+        rust.create(kernel, n_lanes)
+        python.create(kernel, n_lanes)
+
+        # Initialize LDS (64KB, standard size for AMD GPUs)
+        rust.init_lds(65536)
+
+        for emu in (rust, python):
+          emu.set_sgpr(0, args_ptr & 0xffffffff)
+          emu.set_sgpr(1, (args_ptr >> 32) & 0xffffffff)
+          emu.set_sgpr(13, gidx)
+          emu.set_sgpr(14, gidy)
+          emu.set_sgpr(15, gidz)
+
+        step = 0
+        trace: list[tuple[int, int, str, StateSnapshot, StateSnapshot]] = []
+        try:
+          while step < max_steps:
+            rust_before = rust.get_snapshot()
+            python_before = python.get_snapshot()
+
+            inst = program.get(python_before.pc)
+            inst_str = inst.disasm() if inst else f"unknown at PC={python_before.pc}"
+            trace.append((step, python_before.pc, inst_str, rust_before, python_before))
+            if len(trace) > trace_len: trace.pop(0)
+
+            if debug: print(f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: PC={python_before.pc}, inst={inst_str}")
+
+            diffs = rust_before.diff(python_before, n_lanes)
+            if diffs:
+              trace_lines = []
+              for s, pc, d, rb, pb in trace[:-1]:
+                trace_lines.append(f"    step {s}: PC={pc:3d} {d}")
+                if trace.index((s, pc, d, rb, pb)) < len(trace) - 2:
+                  next_rb, next_pb = trace[trace.index((s, pc, d, rb, pb)) + 1][3:5]
+                  inst_diffs = rb.diff(next_rb, n_lanes)
+                  if inst_diffs: trace_lines.append(f"             rust changes: {', '.join(inst_diffs[:3])}")
+              trace_str = "\n".join(trace_lines)
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step} before inst '{inst_str}': states differ:\n  " + "\n  ".join(diffs[:10]) + f"\n  Recent instructions:\n{trace_str}", total_steps
+
+            rust_result = rust.step()
+            python_result = python.step()
+
+            if rust_result != python_result:
+              trace_str = "\n".join(f"    step {s}: PC={pc:3d} {d}" for s, pc, d, _, _ in trace)
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: different return codes: rust={rust_result}, python={python_result}, inst={inst_str}\n  Recent instructions:\n{trace_str}", total_steps
+
+            if rust_result == -1:
+              total_steps += step + 1
+              break
+            if rust_result == 1:
+              total_steps += step + 1
+              break
+            if rust_result < 0 and rust_result != -2:
+              return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Step {step}: error code {rust_result}", total_steps
+
+            step += 1
+          else:
+            return False, f"K{kernel_idx} WG({gidx},{gidy},{gidz}) Max steps ({max_steps}) reached", total_steps
+        finally:
+          rust.free()
+
+  return True, f"Completed {gx*gy*gz} workgroups", total_steps
+
+def compare_emulators_multi_kernel(kernels: list[KernelInfo], buf_pool: dict[int, int], max_steps: int = 1000,
+                                    debug: bool = False, trace_len: int = 10, buf_data: dict[int, bytes] | None = None) -> tuple[bool, str]:
+  """Run all kernels through both emulators with shared buffer pool."""
+  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
+  if buf_data is None: buf_data = {}
+
+  # Allocate shared buffer pool with padding for over-reads (GPU loads up to 16 bytes at once)
+  buf_id_to_ptr: dict[int, int] = {}
+  buffers = []
+  for buf_id, size in buf_pool.items():
+    padded_size = ((size + 15) // 16) * 16 + 16  # round up to 16 bytes + extra padding
+    # Initialize with data from COPY if available
+    init_data = buf_data.get(buf_id, b'\x00' * padded_size)
+    init_list = list(init_data) + [0] * (padded_size - len(init_data))
+    buf = (ctypes.c_uint8 * padded_size)(*init_list[:padded_size])
+    buffers.append((buf, padded_size))
+    buf_id_to_ptr[buf_id] = ctypes.addressof(buf)
+
+  # Set up valid memory ranges
+  ranges = {(ctypes.addressof(b), size) for b, size in buffers}
+
+  total_steps = 0
+  for ki, kernel in enumerate(kernels):
+    # Create args array for this kernel's buffers
+    args = (ctypes.c_uint64 * len(kernel.buf_idxs))(*[buf_id_to_ptr[bid] for bid in kernel.buf_idxs])
+    args_ptr = ctypes.addressof(args)
+
+    # Update valid ranges to include this args array
+    kernel_ranges = ranges | {(args_ptr, ctypes.sizeof(args))}
+    set_valid_mem_ranges(kernel_ranges)
+
+    program = decode_program(kernel.code)
+    n_lanes = kernel.local_size[0] * kernel.local_size[1] * kernel.local_size[2]
+
+    ok, msg, steps = run_single_kernel(
+      kernel.code, min(n_lanes, 32), args_ptr, kernel.global_size,
+      program, max_steps, debug, trace_len, ki
+    )
+    total_steps += steps
+    if not ok:
+      return False, msg
+
+  return True, f"Completed {len(kernels)} kernels, {total_steps} total steps"
+
+def compare_emulators_with_memory(kernel: bytes, n_lanes: int, buf_sizes: list, max_steps: int = 1000, debug: bool = False,
+                                   global_size: tuple[int, int, int] = (1, 1, 1), trace_len: int = 10) -> tuple[bool, str]:
+  """Run both emulators with memory set up for tinygrad kernels, executing all workgroups. Legacy wrapper."""
+  from extra.assembly.rdna3.emu import set_valid_mem_ranges, decode_program
+
+  # Allocate buffers
+  buffers = []
+  for size in buf_sizes:
+    buf = (ctypes.c_uint8 * size)(*[0] * size)
+    buffers.append(buf)
+
+  # Create args array with buffer pointers
+  args = (ctypes.c_uint64 * len(buffers))(*[ctypes.addressof(b) for b in buffers])
+  args_ptr = ctypes.addressof(args)
+
+  # Set up valid memory ranges for Python emulator
+  ranges = {(ctypes.addressof(b), len(b)) for b in buffers}
+  ranges.add((args_ptr, ctypes.sizeof(args)))
+  set_valid_mem_ranges(ranges)
+
+  program = decode_program(kernel)
+  ok, msg, _ = run_single_kernel(kernel, n_lanes, args_ptr, global_size, program, max_steps, debug, trace_len)
+  return ok, msg
+
+def get_kernels_from_tinygrad(op_fn) -> tuple[list[KernelInfo], dict[int, int], dict[int, bytes]]:
+  """Compile a tinygrad operation and extract all kernels with their buffer mappings."""
+  from tinygrad import Tensor
+  from tinygrad.runtime.support.elf import elf_loader
+
+  out = op_fn(Tensor)
+  sched = out.schedule()
+  kernels = []
+  buf_pool: dict[int, int] = {}  # buffer id -> size
+  buf_data: dict[int, bytes] = {}  # buffer id -> initial data from COPY
+
+  for ei in sched:
+    lowered = ei.lower()
+    if ei.ast.op.name == 'COPY':
+      # Handle COPY: extract source data to initialize destination buffer
+      if len(lowered.bufs) >= 2:
+        dst_buf, src_buf = lowered.bufs[0], lowered.bufs[1]
+        dst_id = id(dst_buf)
+        if dst_id not in buf_pool:
+          buf_pool[dst_id] = dst_buf.nbytes
+        # Get source data if it's from numpy/CPU
+        if hasattr(src_buf, 'base') and src_buf.base is not None and hasattr(src_buf.base, '_buf'):
+          src_data = bytes(src_buf.base._buf)
+          buf_data[dst_id] = src_data
+    elif ei.ast.op.name == 'SINK':
+      if lowered.prg and lowered.prg.p.lib:
+        lib = bytes(lowered.prg.p.lib)
+        _, sections, _ = elf_loader(lib)
+        for sec in sections:
+          if sec.name == '.text':
+            buf_idxs = []
+            buf_sizes = []
+            for b in lowered.bufs:
+              buf_id = id(b)
+              if buf_id not in buf_pool:
+                buf_pool[buf_id] = b.nbytes
+              buf_idxs.append(buf_id)
+              buf_sizes.append(b.nbytes)
+            kernels.append(KernelInfo(
+              code=bytes(sec.content),
+              global_size=tuple(lowered.prg.p.global_size),
+              local_size=tuple(lowered.prg.p.local_size),
+              buf_idxs=buf_idxs,
+              buf_sizes=buf_sizes
+            ))
+  if not kernels: raise RuntimeError("No kernel found")
+  return kernels, buf_pool, buf_data
+
+def get_kernel_from_tinygrad(op_fn) -> tuple[bytes, tuple[int, int, int], tuple[int, int, int], list]:
+  """Compile a tinygrad operation and extract the last (main) kernel binary. Legacy wrapper."""
+  kernels, _, _ = get_kernels_from_tinygrad(op_fn)
+  k = kernels[-1]
+  return k.code, k.global_size, k.local_size, k.buf_sizes
+
+@unittest.skipUnless(REMU_PATH.exists(), "libremu.so not found")
+class TestTinygradKernels(unittest.TestCase):
+  """Compare emulators on real tinygrad-compiled kernels."""
+
+  def _test_kernel(self, op_fn, max_steps=10000):
+    kernels, buf_pool, buf_data = get_kernels_from_tinygrad(op_fn)
+    ok, msg = compare_emulators_multi_kernel(kernels, buf_pool, max_steps=max_steps, buf_data=buf_data)
+    self.assertTrue(ok, msg)
+
+  # Basic unary ops
+  def test_neg(self): self._test_kernel(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
+  def test_relu(self): self._test_kernel(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
+  def test_exp(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).exp())
+  def test_log(self): self._test_kernel(lambda T: T([1.0, 2.0, 3.0]).log())
+  def test_sin(self): self._test_kernel(lambda T: T([0.0, 1.0, 2.0]).sin())
+  def test_sqrt(self): self._test_kernel(lambda T: T([1.0, 4.0, 9.0]).sqrt())
+  def test_recip(self): self._test_kernel(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
+
+  # Binary ops
+  def test_add(self): self._test_kernel(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
+  def test_sub(self): self._test_kernel(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
+  def test_mul(self): self._test_kernel(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
+  def test_div(self): self._test_kernel(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
+  def test_max_binary(self): self._test_kernel(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
+
+  # Reductions
+  def test_sum_reduce(self): self._test_kernel(lambda T: T.empty(64).sum())
+  def test_max_reduce(self): self._test_kernel(lambda T: T.empty(64).max())
+  def test_mean_reduce(self): self._test_kernel(lambda T: T.empty(32).mean())
+
+  # Matmul - various sizes
+  def test_gemm_4x4(self): self._test_kernel(lambda T: T.empty(4, 4) @ T.empty(4, 4), max_steps=100000)
+  def test_gemm_8x8(self): self._test_kernel(lambda T: T.empty(8, 8) @ T.empty(8, 8), max_steps=200000)
+  @unittest.skip("too slow")
+  def test_gemm_16x16(self): self._test_kernel(lambda T: T.empty(16, 16) @ T.empty(16, 16), max_steps=500000)
+  def test_gemv(self): self._test_kernel(lambda T: T.empty(1, 16) @ T.empty(16, 16), max_steps=100000)
+
+  # Complex ops
+  def test_softmax(self): self._test_kernel(lambda T: T.empty(16).softmax())
+  def test_layernorm(self): self._test_kernel(lambda T: T.empty(8, 8).layernorm())
+
+  # Memory patterns
+  def test_contiguous(self): self._test_kernel(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
+  def test_reshape(self): self._test_kernel(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
+  def test_expand(self): self._test_kernel(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
+
+  # Cast ops
+  def test_cast_int(self): self._test_kernel(lambda T: T.empty(16).int().float())
+  def test_cast_half(self): self._test_kernel(lambda T: T.empty(16).half().float())
+
+  # Min/max (uses comparison internally)
+  def test_min_binary(self): self._test_kernel(lambda T: T([1.0, 5.0, 3.0]).minimum(T([3.0, 2.0, 4.0])))
+
+  # Comparison ops (test VOPC instructions) - use 32+ elements to force vector instructions
+  def test_cmp_lt(self): self._test_kernel(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
+  def test_cmp_eq(self): self._test_kernel(lambda T: (T.empty(64) == T.empty(64)).where(T.empty(64), T.empty(64)))
+  def test_where(self): self._test_kernel(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
+
+  # Bitwise ops
+  def test_bitwise_and(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0xFF]).int() & T([0x0F, 0x0F, 0x00]).int())
+  def test_bitwise_or(self): self._test_kernel(lambda T: T([0xF0, 0x0F, 0x00]).int() | T([0x0F, 0x0F, 0xFF]).int())
+  def test_bitwise_xor(self): self._test_kernel(lambda T: T([0xFF, 0x0F, 0xF0]).int() ^ T([0x0F, 0xF0, 0xF0]).int())
+
+  # Integer ops - use 32+ elements to force vector instructions
+  def test_int_add(self): self._test_kernel(lambda T: (T.empty(64).int() + T.empty(64).int()).float())
+  def test_int_mul(self): self._test_kernel(lambda T: (T.empty(64).int() * T.empty(64).int()).float())
+  def test_int_mod(self): self._test_kernel(lambda T: (T.empty(64).int().abs() % (T.empty(64).int().abs() + 1)).float())
+
+  # More math ops - use 32+ elements to force vector instructions
+  def test_abs(self): self._test_kernel(lambda T: T.empty(64).abs())
+  def test_floor(self): self._test_kernel(lambda T: T.empty(64).floor())
+  def test_ceil(self): self._test_kernel(lambda T: T.empty(64).ceil())
+  def test_trunc(self): self._test_kernel(lambda T: T.empty(64).trunc())
+
+  # Fused ops
+  def test_fma(self): self._test_kernel(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
+
+  # Argmax/argmin (tests different reduction pattern) - use 32+ elements to force vector instructions
+  def test_argmax(self): self._test_kernel(lambda T: T.empty(64).argmax())
+  def test_argmin(self): self._test_kernel(lambda T: T.empty(64).argmin())
+
+  # Exact value tests - use 32+ elements to force vector instructions (small tensors use scalar ops which Rust emu doesn't fully support)
+  def test_abs_exact(self): self._test_kernel(lambda T: T([-1., 0., 1.]*11).abs())  # 33 elements
+  def test_neg_exact(self): self._test_kernel(lambda T: -T([-1., 0., 1.]*11))
+  def test_log_special(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).log())
+  def test_exp_exact(self): self._test_kernel(lambda T: T([0., 1., -1.]*11).exp())
+  def test_reciprocal_exact(self): self._test_kernel(lambda T: T([1., 2., 0.5]*11).reciprocal())
+
+  # Integer division and mod - use 32+ elements
+  def test_int_div(self): self._test_kernel(lambda T: (T([10, 20, 30]*11).int() // T([3, 4, 5]*11).int()).float())
+  def test_int_neg(self): self._test_kernel(lambda T: (-T([1, -2, 3]*11).int()).float())
+
+  # Mixed precision - use 32+ elements
+  def test_half_add(self): self._test_kernel(lambda T: (T([1., 2.]*16).half() + T([3., 4.]*16).half()).float())
+  def test_half_mul(self): self._test_kernel(lambda T: (T([2., 3.]*16).half() * T([4., 5.]*16).half()).float())
+
+  # Matrix ops - patterns from test_ops.py failures
+  def test_cat(self): self._test_kernel(lambda T: T.empty(32, 64).cat(T.empty(32, 64), dim=1))
+  def test_gather(self): self._test_kernel(lambda T: T.empty(64).gather(0, T.arange(32).int()))
+
+  # Tests from test_ops.py that are failing
+  def test_permute(self): self._test_kernel(lambda T: T.empty(3, 4, 5, 6).permute((3, 2, 1, 0)).contiguous())
+  def test_cat_large(self): self._test_kernel(lambda T: T.empty(45, 65, 9).cat(T.empty(45, 65, 9), T.empty(45, 65, 9), dim=1))
+  def test_gather_small(self): self._test_kernel(lambda T: T.empty(10).gather(0, T.arange(5).int()))
+  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
+  def test_cross_entropy(self): self._test_kernel(lambda T: T.randn(32, 10).softmax().log().sum())
+  def test_cross_entropy_class(self):
+    import numpy as np
+    np.random.seed(0)
+    classes = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
+    x_np = np.random.randn(32, 10).astype(np.float32)
+    self._test_kernel(lambda T: (T(x_np.tolist()).reshape(32,10) + 0).cross_entropy((T(classes).int().reshape(32) + 0)))
+
+  # Regression tests for BFE operations with width=0 (walrus operator bug)
+  def test_topk(self): self._test_kernel(lambda T: T.empty(64).topk(3)[0])
+  def test_interpolate_uint8(self): self._test_kernel(lambda T: T.empty(2,3,64,64).relu().cast('uint8').interpolate((10,10), mode="linear"))
+
+  # Regression test for 64-bit comparison (V_CMP_GT_I64, V_CMP_LT_U64, etc.) with rsrc64
+  def test_index_int64(self):
+    from tinygrad import dtypes
+    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), :])
+
+  @unittest.skip("only works with mock GPU")
+  def test_index_int64_2d(self):
+    from tinygrad import dtypes
+    # Tests 64-bit compare with inline constants (comparing against 0)
+    self._test_kernel(lambda T: T.empty(4, 4)[T.arange(4).cast(dtypes.int64), T.arange(4).cast(dtypes.int64)])
+
+  # Pooling operations - regression test for VCC wave32 mode (S_CBRANCH_VCCZ should only check VCC_LO)
+  def test_avg_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).avg_pool2d(kernel_size=(4,4), stride=2))
+  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
+  def test_avg_pool3d(self):
+    import numpy as np
+    np.random.seed(0)
+    self._test_kernel(lambda T: T(np.random.randn(1, 1, 16, 16, 16).astype(np.float32).tolist()).avg_pool2d(kernel_size=(8,8,8), stride=5, padding=1, count_include_pad=False))
+  def test_max_pool2d(self): self._test_kernel(lambda T: T.empty(1, 1, 8, 8).max_pool2d(kernel_size=(4,4), stride=2))
+
+  # Convolution operations - multi-kernel tests
+  def test_conv2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv2d(T.empty(4, 4, 3, 3)), max_steps=100000)
+  def test_conv_transpose2d(self): self._test_kernel(lambda T: T.empty(1, 4, 8, 8).conv_transpose2d(T.empty(4, 4, 3, 3)), max_steps=200000)
+  @unittest.skip("Rust emulator has S_ADD_I32 SCC bug - uses carry instead of signed overflow")
+  def test_conv_transpose3d(self):
+    import numpy as np
+    np.random.seed(0)
+    self._test_kernel(lambda T: T(np.random.randn(2, 4, 9, 9, 9).astype(np.float32).tolist()).conv_transpose2d(
+      T(np.random.randn(4, 4, 3, 3, 3).astype(np.float32).tolist())), max_steps=500000)
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_emu.py

@@ -0,0 +1,845 @@
+# Unit tests for RDNA3 Python emulator
+import unittest
+import ctypes
+import struct
+import math
+from extra.assembly.rdna3.emu import (
+  WaveState, decode_program, exec_wave, exec_workgroup, run_asm,
+  i32, f32, sext, WAVE_SIZE, set_valid_mem_ranges
+)
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.lib import RawImm
+
+def run_kernel(kernel: bytes, n_threads: int = 1, n_outputs: int = 1) -> list[int]:
+  """Helper to run a kernel and return output values."""
+  output = (ctypes.c_uint32 * (n_threads * n_outputs))(*[0xdead] * (n_threads * n_outputs))
+  output_ptr = ctypes.addressof(output)
+  args = (ctypes.c_uint64 * 1)(output_ptr)
+  args_ptr = ctypes.addressof(args)
+  kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+  kernel_ptr = ctypes.addressof(kernel_buf)
+  # Register valid memory ranges for bounds checking
+  set_valid_mem_ranges({
+    (output_ptr, ctypes.sizeof(output)),
+    (args_ptr, ctypes.sizeof(args)),
+    (kernel_ptr, len(kernel)),
+  })
+  result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
+  assert result == 0, f"run_asm failed with {result}"
+  return [output[i] for i in range(n_threads * n_outputs)]
+
+def make_store_kernel(setup_instrs: list, store_vreg: int = 1) -> bytes:
+  """Create a kernel that runs setup instructions then stores v[store_vreg] to output[tid]."""
+  kernel = b''
+  # Load output pointer
+  kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
+  kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+  # Run setup instructions
+  for instr in setup_instrs:
+    kernel += instr.to_bytes()
+  # Compute offset: v3 = tid * 4
+  kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
+  # Store result
+  kernel += global_store_b32(addr=v[3], data=v[store_vreg], saddr=s[2]).to_bytes()
+  kernel += s_endpgm().to_bytes()
+  return kernel
+
+class TestScalarOps(unittest.TestCase):
+  def test_s_mov_b32(self):
+    state = WaveState()
+    kernel = s_mov_b32(s[5], 42).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[5], 42)
+
+  def test_s_add_u32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 100, 50
+    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 150)
+    self.assertEqual(state.scc, 0)  # no carry
+
+  def test_s_add_u32_carry(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 0xffffffff, 1
+    kernel = s_add_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 0)
+    self.assertEqual(state.scc, 1)  # carry
+
+  def test_s_sub_u32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 100, 30
+    kernel = s_sub_u32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 70)
+    self.assertEqual(state.scc, 0)  # no borrow
+
+  def test_s_and_b32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 0xff00, 0x0ff0
+    kernel = s_and_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 0x0f00)
+
+  def test_s_or_b32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 0xff00, 0x00ff
+    kernel = s_or_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 0xffff)
+
+  def test_s_lshl_b32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 1, 4
+    kernel = s_lshl_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 16)
+
+  def test_s_lshr_b32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 256, 4
+    kernel = s_lshr_b32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 16)
+
+  def test_s_mul_i32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 7, 6
+    kernel = s_mul_i32(s[2], s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[2], 42)
+
+  def test_s_cmp_eq_u32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 42, 42
+    kernel = s_cmp_eq_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.scc, 1)
+
+  def test_s_cmp_lg_u32(self):
+    state = WaveState()
+    state.sgpr[0], state.sgpr[1] = 42, 43
+    kernel = s_cmp_lg_u32(s[0], s[1]).to_bytes() + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.scc, 1)
+
+class TestVectorOps(unittest.TestCase):
+  def test_v_mov_b32(self):
+    kernel = make_store_kernel([v_mov_b32_e32(v[1], 42)])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [42])
+
+  def test_v_add_nc_u32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 10),
+      v_mov_b32_e32(v[2], 32),
+      v_add_nc_u32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [42])
+
+  def test_v_sub_nc_u32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 50),
+      v_mov_b32_e32(v[2], 8),
+      v_sub_nc_u32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [42])
+
+  def test_v_mul_lo_u32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 6),
+      v_mov_b32_e32(v[2], 7),
+      v_mul_lo_u32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [42])
+
+  def test_v_and_b32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 0xff0f),
+      v_mov_b32_e32(v[2], 0x0fff),
+      v_and_b32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [0x0f0f])
+
+  def test_v_or_b32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 0xf000),
+      v_mov_b32_e32(v[2], 0x000f),
+      v_or_b32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [0xf00f])
+
+  def test_v_lshlrev_b32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 1),
+      v_lshlrev_b32_e32(v[1], 5, v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [32])
+
+  def test_v_lshrrev_b32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 128),
+      v_lshrrev_b32_e32(v[1], 3, v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out, [16])
+
+  def test_v_add_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(1.5)),
+      v_mov_b32_e32(v[2], i32(2.5)),
+      v_add_f32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 4.0)
+
+  def test_v_mul_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(3.0)),
+      v_mov_b32_e32(v[2], i32(4.0)),
+      v_mul_f32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 12.0)
+
+  def test_v_max_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(3.0)),
+      v_mov_b32_e32(v[2], i32(5.0)),
+      v_max_f32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 5.0)
+
+  def test_v_min_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(3.0)),
+      v_mov_b32_e32(v[2], i32(5.0)),
+      v_min_f32_e32(v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 3.0)
+
+class TestThreading(unittest.TestCase):
+  def test_thread_id(self):
+    """Each thread should get its own thread ID in v0."""
+    kernel = make_store_kernel([v_mov_b32_e32(v[1], v[0])], store_vreg=1)
+    out = run_kernel(kernel, n_threads=4)
+    self.assertEqual(out, [0, 1, 2, 3])
+
+  def test_thread_local_ops(self):
+    """Each thread computes tid * 10."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[2], 10),
+      v_mul_lo_u32(v[1], v[0], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=4)
+    self.assertEqual(out, [0, 10, 20, 30])
+
+  def test_exec_mask(self):
+    """Test that exec mask controls which lanes execute."""
+    kernel = b''
+    kernel += s_load_b64(s[2:3], s[0:1], 0, soffset=NULL).to_bytes()
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    kernel += v_mov_b32_e32(v[1], 100).to_bytes()  # default value
+    kernel += s_mov_b32(EXEC_LO, 0b0101).to_bytes()  # only lanes 0 and 2
+    kernel += v_mov_b32_e32(v[1], 42).to_bytes()   # only for active lanes
+    kernel += s_mov_b32(EXEC_LO, 0xf).to_bytes()   # restore all lanes
+    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
+    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[2]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    out = run_kernel(kernel, n_threads=4)
+    self.assertEqual(out, [42, 100, 42, 100])
+
+class TestBranching(unittest.TestCase):
+  def test_s_branch(self):
+    """Test unconditional branch."""
+    state = WaveState()
+    kernel = b''
+    kernel += s_mov_b32(s[0], 1).to_bytes()
+    kernel += s_branch(1).to_bytes()  # skip next instruction
+    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
+    kernel += s_mov_b32(s[1], 3).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[0], 1)  # not overwritten
+    self.assertEqual(state.sgpr[1], 3)
+
+  def test_s_cbranch_scc0(self):
+    """Test conditional branch on SCC=0."""
+    state = WaveState()
+    state.scc = 0
+    kernel = b''
+    kernel += s_mov_b32(s[0], 1).to_bytes()
+    kernel += s_cbranch_scc0(1).to_bytes()  # branch if scc=0
+    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[0], 1)
+
+  def test_s_cbranch_scc1(self):
+    """Test conditional branch on SCC=1."""
+    state = WaveState()
+    state.scc = 1
+    kernel = b''
+    kernel += s_mov_b32(s[0], 1).to_bytes()
+    kernel += s_cbranch_scc1(1).to_bytes()  # branch if scc=1
+    kernel += s_mov_b32(s[0], 2).to_bytes()  # should be skipped
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.sgpr[0], 1)
+
+  def test_unknown_sopp_opcode(self):
+    """Regression test: unknown SOPP opcodes should be ignored, not crash."""
+    state = WaveState()
+    # Create a raw SOPP instruction with opcode 8 (undefined in our enum)
+    # SOPP format: bits[31:23] = 0b101111111, bits[22:16] = op, bits[15:0] = simm16
+    unknown_sopp = (0b101111111 << 23) | (8 << 16) | 0  # op=8, simm16=0
+    kernel = unknown_sopp.to_bytes(4, 'little') + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    # Should not raise an exception
+    exec_wave(prog, state, bytearray(65536), 1)
+
+class TestMemory(unittest.TestCase):
+  def test_global_load_store(self):
+    """Test global load followed by store."""
+    # Create input buffer
+    input_buf = (ctypes.c_uint32 * 4)(10, 20, 30, 40)
+    input_ptr = ctypes.addressof(input_buf)
+    output_buf = (ctypes.c_uint32 * 4)(*[0]*4)
+    output_ptr = ctypes.addressof(output_buf)
+    args = (ctypes.c_uint64 * 2)(output_ptr, input_ptr)
+    args_ptr = ctypes.addressof(args)
+
+    # Kernel: load from input[tid], add 1, store to output[tid]
+    kernel = b''
+    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()  # output ptr
+    kernel += s_load_b64(s[4:5], s[0:1], soffset=NULL, offset=8).to_bytes()  # input ptr
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    kernel += v_lshlrev_b32_e32(v[2], 2, v[0]).to_bytes()  # offset = tid * 4
+    kernel += global_load_b32(vdst=v[1], addr=v[2], saddr=s[4]).to_bytes()
+    kernel += s_waitcnt(vmcnt=0).to_bytes()
+    kernel += v_add_nc_u32_e32(v[1], 1, v[1]).to_bytes()  # add 1
+    kernel += global_store_b32(addr=v[2], data=v[1], saddr=s[2]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    kernel_ptr = ctypes.addressof(kernel_buf)
+    set_valid_mem_ranges({
+      (input_ptr, ctypes.sizeof(input_buf)),
+      (output_ptr, ctypes.sizeof(output_buf)),
+      (args_ptr, ctypes.sizeof(args)),
+      (kernel_ptr, len(kernel)),
+    })
+    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, 4, 1, 1, args_ptr)
+    self.assertEqual(result, 0)
+    self.assertEqual([output_buf[i] for i in range(4)], [11, 21, 31, 41])
+
+class TestFloatOps(unittest.TestCase):
+  def test_v_rcp_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(4.0)),
+      v_rcp_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertAlmostEqual(f32(out[0]), 0.25, places=5)
+
+  def test_v_sqrt_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(16.0)),
+      v_sqrt_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertAlmostEqual(f32(out[0]), 4.0, places=5)
+
+  def test_v_floor_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(3.7)),
+      v_floor_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 3.0)
+
+  def test_v_ceil_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(3.2)),
+      v_ceil_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 4.0)
+
+  def test_v_cvt_f32_i32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 42),
+      v_cvt_f32_i32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 42.0)
+
+  def test_v_cvt_i32_f32(self):
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(42.9)),
+      v_cvt_i32_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0], 42)
+
+class TestVOP3(unittest.TestCase):
+  def test_v_fma_f32(self):
+    """Test fused multiply-add: a*b + c"""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(2.0)),
+      v_mov_b32_e32(v[2], i32(3.0)),
+      v_mov_b32_e32(v[4], i32(4.0)),
+      v_fma_f32(v[1], v[1], v[2], v[4]),  # 2*3+4 = 10
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 10.0)
+
+  def test_v_add3_u32(self):
+    """Test 3-operand add."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 10),
+      v_mov_b32_e32(v[2], 20),
+      v_mov_b32_e32(v[4], 12),
+      v_add3_u32(v[1], v[1], v[2], v[4]),  # 10+20+12 = 42
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0], 42)
+
+  def test_v_neg_modifier(self):
+    """Test VOP3 negation modifier."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(5.0)),
+      v_mov_b32_e32(v[2], i32(3.0)),
+      # v_add_f32 with neg on src1: 5 + (-3) = 2
+      v_add_f32_e64(v[1], v[1], v[2], neg=0b010),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 2.0)
+
+  def test_v_ldexp_f32(self):
+    """Regression test: V_LDEXP_F32 used by exp()."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(1.5)),
+      v_mov_b32_e32(v[2], 3),  # exponent
+      v_ldexp_f32(v[1], v[1], v[2]),  # 1.5 * 2^3 = 12.0
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(f32(out[0]), 12.0)
+
+  def test_v_xad_u32(self):
+    """Regression test: V_XAD_U32 (xor-add) used by random number generation."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 3),
+      v_mov_b32_e32(v[2], 4),
+      v_mov_b32_e32(v[4], 5),
+      v_xad_u32(v[1], v[1], v[2], v[4]),  # (3^4)+5 = 7+5 = 12
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0], 12)
+
+  def test_v_lshl_or_b32(self):
+    """Regression test: V_LSHL_OR_B32 operand order is (s0 << s1) | s2, not (s0 << s2) | s1."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 5),   # s0 = value to shift
+      v_mov_b32_e32(v[2], 2),   # s1 = shift amount
+      v_mov_b32_e32(v[4], 3),   # s2 = value to OR
+      v_lshl_or_b32(v[1], v[1], v[2], v[4]),  # (5 << 2) | 3 = 20 | 3 = 23
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0], 23)
+
+  def test_v_sqrt_f32_negative(self):
+    """Regression test: V_SQRT_F32 should return NaN for negative inputs, not 0."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(-1.0)),
+      v_sqrt_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertTrue(math.isnan(f32(out[0])))
+
+  def test_v_rsq_f32_negative(self):
+    """Regression test: V_RSQ_F32 should return NaN for negative inputs, not inf."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i32(-1.0)),
+      v_rsq_f32_e32(v[1], v[1]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertTrue(math.isnan(f32(out[0])))
+
+class TestVOPD(unittest.TestCase):
+  def test_vopd_add_nc_u32(self):
+    """Test VOPD V_DUAL_ADD_NC_U32."""
+    state = WaveState()
+    state.vgpr[0][1] = 100
+    state.vgpr[0][2] = 50
+    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
+    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
+                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.vgpr[0][3], 100)  # MOV result
+    self.assertEqual(state.vgpr[0][4], 150)  # 100 + 50
+
+  def test_vopd_lshlrev(self):
+    """Test VOPD V_DUAL_LSHLREV_B32."""
+    state = WaveState()
+    state.vgpr[0][1] = 0x10
+    state.vgpr[0][2] = 0
+    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
+    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
+                  opy=VOPDOp.V_DUAL_LSHLREV_B32, srcy0=4, vsrcy1=VGPR(1), vdsty=VGPR(4)).to_bytes()  # V4 = V1 << 4
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.vgpr[0][3], 0x10)  # MOV result
+    self.assertEqual(state.vgpr[0][4], 0x100)  # 0x10 << 4 = 0x100
+
+  def test_vopd_and(self):
+    """Test VOPD V_DUAL_AND_B32."""
+    state = WaveState()
+    state.vgpr[0][1] = 0xff
+    state.vgpr[0][2] = 0x0f
+    # vdsty = (vdsty_enc << 1) | ((vdstx & 1) ^ 1), so for vdstx=3 (odd), vdsty=4 requires VGPR(4)
+    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[1], vsrcx1=VGPR(0), vdstx=VGPR(3),
+                  opy=VOPDOp.V_DUAL_AND_B32, srcy0=v[1], vsrcy1=VGPR(2), vdsty=VGPR(4)).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.vgpr[0][3], 0xff)
+    self.assertEqual(state.vgpr[0][4], 0x0f)  # 0xff & 0x0f = 0x0f
+
+  def test_vopd_parallel_read(self):
+    """Regression: VOPD must read all inputs before writing - Y op reads register that X op writes."""
+    state = WaveState()
+    state.vgpr[0][4] = 0
+    state.vgpr[0][7] = 5  # Y op reads v7 as vsrcy1, X op writes to v7
+    # X: MOV v7, v0 (v0=0, so v7 becomes 0)
+    # Y: ADD v6, v4, v7 (should use original v7=5, not the overwritten 0)
+    # vdsty_enc=3 with vdstx=7 (odd) -> vdsty = (3 << 1) | (7&1)^1 = 6 | 0 = 6
+    kernel = VOPD(opx=VOPDOp.V_DUAL_MOV_B32, srcx0=v[0], vsrcx1=VGPR(0), vdstx=VGPR(7),
+                  opy=VOPDOp.V_DUAL_ADD_NC_U32, srcy0=v[4], vsrcy1=VGPR(7), vdsty=VGPR(6)).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.vgpr[0][7], 0)  # X op: v7 = v0 = 0
+    self.assertEqual(state.vgpr[0][6], 5)  # Y op: v6 = v4 + v7 = 0 + 5 (original v7)
+
+class TestDecoder(unittest.TestCase):
+  def test_vopd_literal_handling(self):
+    """Regression test: VOPD srcx0/srcy0 with literal (255) wasn't consuming the literal dword."""
+    state = WaveState()
+    # Create VOPD with srcx0=255 (literal), followed by literal value 0x12345678
+    vopd_bytes = VOPD(opx=8, srcx0=RawImm(255), vsrcx1=VGPR(0), vdstx=VGPR(1),  # MOV: V1 = literal
+                      opy=8, srcy0=RawImm(128), vsrcy1=VGPR(0), vdsty=VGPR(2)).to_bytes()  # MOV: V2 = 0
+    literal_bytes = (0x12345678).to_bytes(4, 'little')
+    kernel = vopd_bytes + literal_bytes + s_endpgm().to_bytes()
+    prog = decode_program(kernel)
+    # Should decode as 3 instructions: VOPD (with literal), then S_ENDPGM
+    # The literal should NOT be decoded as a separate instruction
+    self.assertEqual(len(prog), 2)  # VOPD + S_ENDPGM
+    exec_wave(prog, state, bytearray(65536), 1)
+    self.assertEqual(state.vgpr[0][1], 0x12345678)
+
+  def test_s_endpgm_stops_decode(self):
+    """Regression test: decoder should stop at S_ENDPGM, not read past into metadata."""
+    # Create a kernel followed by garbage that looks like an invalid instruction
+    kernel = s_mov_b32(s[0], 42).to_bytes() + s_endpgm().to_bytes()
+    garbage = bytes([0xff] * 16)  # garbage after kernel
+    prog = decode_program(kernel + garbage)
+    # Should only have 2 instructions (s_mov_b32 and s_endpgm)
+    self.assertEqual(len(prog), 2)
+
+class TestFloatConversion(unittest.TestCase):
+  """Unit tests for i32/i16/f32/f16 float conversion functions."""
+
+  def test_i32_preserves_nan_sign(self):
+    """NaN sign bit should be preserved when converting float to int bits."""
+    from extra.assembly.rdna3.emu import i32, f32
+    # 0 * -inf produces a negative NaN
+    neg_nan = 0.0 * float('-inf')
+    bits = i32(neg_nan)
+    # Should have sign bit set (0xffc00000), not canonical positive NaN (0x7fc00000)
+    self.assertEqual(bits & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{bits:08x}")
+    self.assertTrue(math.isnan(f32(bits)))
+
+  def test_i32_preserves_positive_nan(self):
+    """Positive NaN should remain positive."""
+    from extra.assembly.rdna3.emu import i32, f32
+    pos_nan = float('nan')
+    bits = i32(pos_nan)
+    # Standard Python NaN is positive (0x7fc00000)
+    self.assertEqual(bits & 0x80000000, 0, f"Expected positive NaN, got 0x{bits:08x}")
+    self.assertTrue(math.isnan(f32(bits)))
+
+  def test_i32_overflow_to_inf(self):
+    """Values too large for f32 should become inf."""
+    from extra.assembly.rdna3.emu import i32, f32
+    big = 2.0 ** 200
+    self.assertEqual(i32(big), 0x7f800000)  # +inf
+    self.assertEqual(i32(-big), 0xff800000)  # -inf
+
+  def test_i32_inf(self):
+    """Infinity should be preserved."""
+    from extra.assembly.rdna3.emu import i32
+    self.assertEqual(i32(float('inf')), 0x7f800000)
+    self.assertEqual(i32(float('-inf')), 0xff800000)
+
+  def test_i32_normal_values(self):
+    """Normal float values should round-trip correctly (within f32 precision)."""
+    from extra.assembly.rdna3.emu import i32, f32
+    # Use values exactly representable in float32
+    for val in [0.0, 1.0, -1.0, 0.5, -0.5, 100.0, -100.0, 1e10]:
+      bits = i32(val)
+      self.assertAlmostEqual(f32(bits), val, places=5)
+
+  def test_i16_overflow_to_inf(self):
+    """Values too large for f16 should become inf."""
+    from extra.assembly.rdna3.emu import i16
+    big = 100000.0  # way larger than f16 max (65504)
+    self.assertEqual(i16(big), 0x7c00)  # +inf
+    self.assertEqual(i16(-big), 0xfc00)  # -inf
+
+  def test_i16_inf(self):
+    """Infinity should be preserved."""
+    from extra.assembly.rdna3.emu import i16
+    self.assertEqual(i16(float('inf')), 0x7c00)
+    self.assertEqual(i16(float('-inf')), 0xfc00)
+
+  def test_fma_nan_sign_preserved(self):
+    """FMA producing NaN should preserve the correct sign bit."""
+    from extra.assembly.rdna3.emu import i32, f32
+    # 0 * (-inf) + 1.0 = NaN (from 0 * -inf)
+    a, b, c = 0.0, float('-inf'), 1.0
+    result = i32(a * b + c)
+    # The NaN should be negative since 0 * -inf produces negative NaN
+    self.assertEqual(result & 0x80000000, 0x80000000, f"Expected negative NaN, got 0x{result:08x}")
+
+class TestMultiWave(unittest.TestCase):
+  def test_all_waves_execute(self):
+    """Regression test: all waves in a workgroup must execute, not just the first."""
+    n_threads = 64  # 2 waves of 32 threads each
+    output = (ctypes.c_uint32 * n_threads)(*[0xdead] * n_threads)
+    output_ptr = ctypes.addressof(output)
+    args = (ctypes.c_uint64 * 1)(output_ptr)
+    args_ptr = ctypes.addressof(args)
+
+    # Simple kernel: store tid to output[tid]
+    kernel = b''
+    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    kernel += v_lshlrev_b32_e32(v[1], 2, v[0]).to_bytes()  # offset = tid * 4
+    kernel += global_store_b32(addr=v[1], data=v[0], saddr=s[2]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    kernel_ptr = ctypes.addressof(kernel_buf)
+    set_valid_mem_ranges({
+      (output_ptr, ctypes.sizeof(output)),
+      (args_ptr, ctypes.sizeof(args)),
+      (kernel_ptr, len(kernel)),
+    })
+    result = run_asm(kernel_ptr, len(kernel), 1, 1, 1, n_threads, 1, 1, args_ptr)
+    self.assertEqual(result, 0)
+    # All threads should have written their tid
+    for i in range(n_threads):
+      self.assertEqual(output[i], i, f"Thread {i} didn't execute")
+
+class TestRegressions(unittest.TestCase):
+  """Regression tests for bugs fixed in the emulator."""
+
+  def test_v_fmac_f16(self):
+    """V_FMAC_F16: fused multiply-add for FP16. Regression for VOP2 op 54."""
+    from extra.assembly.rdna3.emu import i16, f16
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], i16(2.0)),  # v1.lo = 2.0 (fp16)
+      v_mov_b32_e32(v[2], i16(3.0)),  # v2.lo = 3.0 (fp16)
+      # v1 = v1 * v2 + v1 = 2.0 * 3.0 + 2.0 = 8.0
+      VOP2(VOP2Op.V_FMAC_F16, v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertAlmostEqual(f16(out[0] & 0xffff), 8.0, places=2)
+
+  def test_v_cvt_f64_f32(self):
+    """V_CVT_F64_F32: convert float32 to float64. Regression for VOP1 op 16."""
+    kernel = b''
+    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    kernel += v_mov_b32_e32(v[1], i32(3.14159)).to_bytes()
+    kernel += VOP1(VOP1Op.V_CVT_F64_F32, v[4], v[1]).to_bytes()  # v4:v5 = f64(v1)
+    kernel += v_lshlrev_b32_e32(v[3], 3, v[0]).to_bytes()  # offset = tid * 8
+    kernel += global_store_b64(addr=v[3], data=v[4], saddr=s[2]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    output = (ctypes.c_double * 1)(0.0)
+    output_ptr = ctypes.addressof(output)
+    args = (ctypes.c_uint64 * 1)(output_ptr)
+    args_ptr = ctypes.addressof(args)
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    kernel_ptr = ctypes.addressof(kernel_buf)
+    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
+    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
+    self.assertAlmostEqual(output[0], 3.14159, places=4)
+
+  def test_v_add_f64(self):
+    """V_ADD_F64: add two float64 values. Regression for VOP3 op 807."""
+    from extra.assembly.rdna3.emu import i64_parts
+    kernel = b''
+    kernel += s_load_b64(s[2:3], s[0:1], soffset=NULL, offset=0).to_bytes()
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    # Load 1.5 into v1:v2
+    lo, hi = i64_parts(1.5)
+    kernel += v_mov_b32_e32(v[1], lo).to_bytes()
+    kernel += v_mov_b32_e32(v[2], hi).to_bytes()
+    # Load 2.5 into v3:v4
+    lo, hi = i64_parts(2.5)
+    kernel += v_mov_b32_e32(v[3], lo).to_bytes()
+    kernel += v_mov_b32_e32(v[4], hi).to_bytes()
+    # v5:v6 = v1:v2 + v3:v4 = 1.5 + 2.5 = 4.0
+    kernel += VOP3(VOP3Op.V_ADD_F64, v[5], v[1], v[3]).to_bytes()
+    kernel += v_lshlrev_b32_e32(v[7], 3, v[0]).to_bytes()
+    kernel += global_store_b64(addr=v[7], data=v[5], saddr=s[2]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+    output = (ctypes.c_double * 1)(0.0)
+    output_ptr = ctypes.addressof(output)
+    args = (ctypes.c_uint64 * 1)(output_ptr)
+    args_ptr = ctypes.addressof(args)
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    kernel_ptr = ctypes.addressof(kernel_buf)
+    set_valid_mem_ranges({(output_ptr, 8), (args_ptr, 8), (kernel_ptr, len(kernel))})
+    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
+    self.assertAlmostEqual(output[0], 4.0, places=10)
+
+  def test_flat_load_d16_hi_b16(self):
+    """FLAT_LOAD_D16_HI_B16: load 16-bit to high half. Regression for FLAT op 35."""
+    from extra.assembly.rdna3.emu import i16
+    # Create a buffer with test data
+    src_data = (ctypes.c_uint16 * 1)(0x1234)
+    src_ptr = ctypes.addressof(src_data)
+    output = (ctypes.c_uint32 * 1)(0xABCD0000)  # preset low bits
+    output_ptr = ctypes.addressof(output)
+    args = (ctypes.c_uint64 * 2)(output_ptr, src_ptr)
+    args_ptr = ctypes.addressof(args)
+
+    kernel = b''
+    kernel += s_load_b128(s[0:3], s[0:1], soffset=NULL, offset=0).to_bytes()
+    kernel += s_waitcnt(lgkmcnt=0).to_bytes()
+    kernel += v_mov_b32_e32(v[1], 0xDEAD).to_bytes()  # initial value with low bits set
+    kernel += v_mov_b32_e32(v[2], 0).to_bytes()  # offset = 0
+    kernel += FLAT(FLATOp.FLAT_LOAD_D16_HI_B16, v[1], v[2], saddr=s[2], offset=0).to_bytes()
+    kernel += s_waitcnt(vmcnt=0).to_bytes()
+    kernel += v_lshlrev_b32_e32(v[3], 2, v[0]).to_bytes()
+    kernel += global_store_b32(addr=v[3], data=v[1], saddr=s[0]).to_bytes()
+    kernel += s_endpgm().to_bytes()
+
+    kernel_buf = (ctypes.c_char * len(kernel)).from_buffer_copy(kernel)
+    kernel_ptr = ctypes.addressof(kernel_buf)
+    set_valid_mem_ranges({(output_ptr, 4), (src_ptr, 2), (args_ptr, 16), (kernel_ptr, len(kernel))})
+    run_asm(kernel_ptr, len(kernel), 1, 1, 1, 1, 1, 1, args_ptr)
+    # High 16 bits should be 0x1234, low 16 bits preserved as 0xDEAD
+    self.assertEqual(output[0], 0x1234DEAD)
+
+  def test_v_mad_u16(self):
+    """V_MAD_U16: multiply-add unsigned 16-bit. Regression for VOP3 op 577."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 10),  # a = 10
+      v_mov_b32_e32(v[2], 20),  # b = 20
+      v_mov_b32_e32(v[4], 5),   # c = 5
+      VOP3(VOP3Op.V_MAD_U16, v[1], v[1], v[2], v[4]),  # v1 = 10*20+5 = 205
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0] & 0xffff, 205)
+
+  def test_v_lshrrev_b16(self):
+    """V_LSHRREV_B16: logical shift right 16-bit. Regression for VOP3 op 825."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 0x8000),  # value to shift
+      v_mov_b32_e32(v[2], 4),       # shift amount
+      VOP3(VOP3Op.V_LSHRREV_B16, v[1], v[2], v[1]),  # v1 = 0x8000 >> 4 = 0x0800
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0] & 0xffff, 0x0800)
+
+  def test_v_min_u16(self):
+    """V_MIN_U16: minimum of two unsigned 16-bit values. Regression for VOP3 op 779."""
+    kernel = make_store_kernel([
+      v_mov_b32_e32(v[1], 100),
+      v_mov_b32_e32(v[2], 50),
+      VOP3(VOP3Op.V_MIN_U16, v[1], v[1], v[2]),
+    ])
+    out = run_kernel(kernel, n_threads=1)
+    self.assertEqual(out[0] & 0xffff, 50)
+
+class TestWMMA(unittest.TestCase):
+  """Tests for WMMA (Wave Matrix Multiply Accumulate) instructions."""
+
+  def test_wmma_f32_16x16x16_f16_identity(self):
+    """V_WMMA_F32_16X16X16_F16 with identity matrix. Regression for VOP3P op 64."""
+    from extra.assembly.rdna3.emu import i16, f16, exec_wmma_f32_16x16x16_f16, WaveState
+    # Test using direct emulator call rather than full kernel to simplify
+    st = WaveState()
+    st.exec_mask = 0xffffffff  # all 32 lanes active
+
+    # Set up A as identity matrix: A[i][i] = 1.0, rest = 0.0
+    # Lane i holds row i of A in 8 regs (2 fp16 per reg)
+    for lane in range(16):
+      for reg in range(8):
+        col0, col1 = reg * 2, reg * 2 + 1
+        val0 = i16(1.0) if col0 == lane else 0
+        val1 = i16(1.0) if col1 == lane else 0
+        st.vgpr[lane][0 + reg] = val0 | (val1 << 16)  # src0 = v0:v7
+
+    # Set up B as identity matrix: lane i holds column i of B
+    for lane in range(16):
+      for reg in range(8):
+        row0, row1 = reg * 2, reg * 2 + 1
+        val0 = i16(1.0) if row0 == lane else 0
+        val1 = i16(1.0) if row1 == lane else 0
+        st.vgpr[lane][8 + reg] = val0 | (val1 << 16)  # src1 = v8:v15
+
+    # Set up C as zeros
+    for lane in range(32):
+      for reg in range(8):
+        st.vgpr[lane][16 + reg] = 0  # src2 = v16:v23
+
+    # Create a fake VOP3P instruction
+    inst = VOP3P(VOP3POp.V_WMMA_F32_16X16X16_F16, v[24], src0=VGPR(0), src1=VGPR(8), src2=VGPR(16))
+
+    # Execute WMMA
+    exec_wmma_f32_16x16x16_f16(st, inst, 32)
+
+    # Check result: C should be identity (since A @ B where both are identity)
+    # Output i = row*16+col goes to lane (i%32), reg (i//32)
+    for row in range(16):
+      for col in range(16):
+        idx = row * 16 + col
+        lane, reg = idx % 32, idx // 32
+        result = st.vgpr[lane][24 + reg]
+        expected = 1.0 if row == col else 0.0
+        self.assertAlmostEqual(f32(result), expected, places=3,
+          msg=f"C[{row},{col}] = {f32(result)}, expected {expected}")
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_formats.py

@@ -0,0 +1,332 @@
+#!/usr/bin/env python3
+"""Test MUBUF, MTBUF, MIMG, EXP, DS formats against LLVM."""
+import unittest
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.lib import encode_src
+
+class TestMUBUF(unittest.TestCase):
+  """Test MUBUF (buffer) instructions."""
+
+  def test_buffer_load_b32_basic(self):
+    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_load_b32_idxen(self):
+    # buffer_load_b32 v5, v0, s[8:11], s3 idxen offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, idxen=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x82,0x03]))
+
+  def test_buffer_load_b32_offen(self):
+    # buffer_load_b32 v5, v0, s[8:11], s3 offen offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, offen=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x42,0x03]))
+
+  def test_buffer_load_b32_glc(self):
+    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc
+    # GFX11: encoding: [0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x4f,0x50,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_load_b32_slc(self):
+    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 slc
+    # GFX11: encoding: [0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, slc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x1f,0x50,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_load_b32_dlc(self):
+    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 dlc
+    # GFX11: encoding: [0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, dlc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x2f,0x50,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_load_b32_all_flags(self):
+    # buffer_load_b32 v5, off, s[8:11], s3 offset:4095 glc slc dlc
+    # GFX11: encoding: [0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1, slc=1, dlc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x7f,0x50,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_store_b32(self):
+    # buffer_store_b32 v1, off, s[12:15], s4 offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]
+    inst = buffer_store_b32(vdata=v[1], vaddr=v[0], srsrc=s[12:16], soffset=s[4], offset=4095)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x68,0xe0,0x00,0x01,0x03,0x04]))
+
+  def test_buffer_load_b64(self):
+    # buffer_load_b64 v[5:6], off, s[8:11], s3 offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]
+    inst = buffer_load_b64(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x54,0xe0,0x00,0x05,0x02,0x03]))
+
+  def test_buffer_load_soffset_m0(self):
+    # buffer_load_b32 v5, off, s[8:11], m0 offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=M0, offset=4095)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x7d]))
+
+  def test_buffer_load_soffset_inline_const(self):
+    # buffer_load_b32 v5, off, s[8:11], 0 offset:4095
+    # GFX11: encoding: [0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=0, offset=4095)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0x50,0xe0,0x00,0x05,0x02,0x80]))
+
+  def test_buffer_disasm_roundtrip(self):
+    inst = buffer_load_b32(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, glc=1)
+    decoded = MUBUF.from_bytes(inst.to_bytes())
+    self.assertEqual(decoded.to_bytes(), inst.to_bytes())
+
+
+class TestMTBUF(unittest.TestCase):
+  """Test MTBUF (typed buffer) instructions."""
+
+  def test_tbuffer_load_format_x(self):
+    # tbuffer_load_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
+    # BUF_FMT_32_FLOAT = 22
+    # GFX11: encoding: [0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]
+    inst = tbuffer_load_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb0,0xe8,0x00,0x05,0x02,0x03]))
+
+  def test_tbuffer_store_format_x(self):
+    # tbuffer_store_format_x v5, off, s[8:11], s3 format:[BUF_FMT_32_FLOAT] offset:4095
+    # BUF_FMT_32_FLOAT = 22
+    # GFX11: encoding: [0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]
+    inst = tbuffer_store_format_x(vdata=v[5], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=22)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x0f,0xb2,0xe8,0x00,0x05,0x02,0x03]))
+
+  def test_tbuffer_load_format_xy(self):
+    # tbuffer_load_format_xy v[5:6], off, s[8:11], s3 format:[BUF_FMT_32_32_FLOAT] offset:4095
+    # BUF_FMT_32_32_FLOAT = 50
+    # GFX11: encoding: [0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]
+    inst = tbuffer_load_format_xy(vdata=v[5:7], vaddr=v[0], srsrc=s[8:12], soffset=s[3], offset=4095, format=50)
+    self.assertEqual(inst.to_bytes(), bytes([0xff,0x8f,0x90,0xe9,0x00,0x05,0x02,0x03]))
+
+
+class TestMIMG(unittest.TestCase):
+  """Test MIMG (image) instructions."""
+
+  def test_image_load_2d(self):
+    # image_load v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
+    # GFX11: encoding: [0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
+    inst = image_load(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)  # dim=1 is SQ_RSRC_IMG_2D
+    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
+
+  def test_image_store_2d(self):
+    # image_store v[0:3], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D
+    # GFX11: encoding: [0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]
+    inst = image_store(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x18,0xf0,0x04,0x00,0x00,0x00]))
+
+  def test_image_load_1d(self):
+    # image_load v[0:3], v4, s[0:7] dmask:0xf dim:SQ_RSRC_IMG_1D
+    # GFX11: encoding: [0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]
+    inst = image_load(vdata=v[0:4], vaddr=v[4], srsrc=s[0:8], dmask=0xf, dim=0)  # dim=0 is SQ_RSRC_IMG_1D
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x0f,0x00,0xf0,0x04,0x00,0x00,0x00]))
+
+  def test_image_sample(self):
+    # image_sample v[0:3], v[4:5], s[0:7], s[8:11] dmask:0xf dim:SQ_RSRC_IMG_2D
+    # GFX11: encoding: [0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]
+    inst = image_sample(vdata=v[0:4], vaddr=v[4:6], srsrc=s[0:8], ssamp=s[8:12], dmask=0xf, dim=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x6c,0xf0,0x04,0x00,0x00,0x08]))
+
+  def test_image_load_d16(self):
+    # image_load v[0:1], v[4:5], s[0:7] dmask:0xf dim:SQ_RSRC_IMG_2D d16
+    # GFX11: encoding: [0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]
+    inst = image_load(vdata=v[0:2], vaddr=v[4:6], srsrc=s[0:8], dmask=0xf, dim=1, d16=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x04,0x0f,0x02,0xf0,0x04,0x00,0x00,0x00]))
+
+
+class TestEXP(unittest.TestCase):
+  """Test EXP (export) instructions."""
+
+  def test_exp_mrt0(self):
+    # exp mrt0 v0, v1, v2, v3
+    # GFX11: encoding: [0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]
+    inst = EXP(en=0xf, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[2], vsrc3=v[3])
+    self.assertEqual(inst.to_bytes(), bytes([0x0f,0x00,0x00,0xf8,0x00,0x01,0x02,0x03]))
+
+  def test_exp_mrtz(self):
+    # exp mrtz v4, v3, v2, v1
+    # GFX11: encoding: [0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]
+    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1])
+    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x00,0x00,0xf8,0x04,0x03,0x02,0x01]))
+
+  def test_exp_mrtz_done(self):
+    # exp mrtz v4, v3, v2, v1 done
+    # GFX11: encoding: [0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x01]
+    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[3], done=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x08,0x00,0xf8,0x04,0x03,0x02,0x03]))
+
+  def test_exp_partial_mask(self):
+    # exp mrt0 v0, v1, off, off (en=0x3, only first two components)
+    # GFX11: encoding: [0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]
+    inst = EXP(en=0x3, target=0, vsrc0=v[0], vsrc1=v[1], vsrc2=v[0], vsrc3=v[0])
+    self.assertEqual(inst.to_bytes(), bytes([0x03,0x00,0x00,0xf8,0x00,0x01,0x00,0x00]))
+
+  def test_exp_row_en(self):
+    # exp mrtz v4, v3, v2, v1 row_en
+    # GFX11: encoding: [0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]
+    inst = EXP(en=0xf, target=8, vsrc0=v[4], vsrc1=v[3], vsrc2=v[2], vsrc3=v[1], row=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x8f,0x20,0x00,0xf8,0x04,0x03,0x02,0x01]))
+
+
+class TestDS(unittest.TestCase):
+  """Test DS (data share / LDS) instructions."""
+
+  def test_ds_store_b32(self):
+    # ds_store_b32 v0, v1
+    # GFX11: encoding: [0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
+    inst = ds_store_b32(addr=v[0], data0=v[1])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
+
+  def test_ds_load_b32(self):
+    # ds_load_b32 v0, v1
+    # GFX11: encoding: [0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]
+    inst = ds_load_b32(vdst=v[0], addr=v[1])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd8,0x01,0x00,0x00,0x00]))
+
+  def test_ds_store_b32_offset(self):
+    # ds_store_b32 v0, v1 offset:64
+    # GFX11: encoding: [0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]
+    inst = ds_store_b32(addr=v[0], data0=v[1], offset0=64)
+    self.assertEqual(inst.to_bytes(), bytes([0x40,0x00,0x34,0xd8,0x00,0x01,0x00,0x00]))
+
+  def test_ds_load_b64(self):
+    # ds_load_b64 v[0:1], v2
+    # GFX11: encoding: [0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]
+    inst = ds_load_b64(vdst=v[0:2], addr=v[2])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0xd8,0xd9,0x02,0x00,0x00,0x00]))
+
+  def test_ds_add_u32(self):
+    # ds_add_u32 v0, v1
+    # GFX11: encoding: [0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]
+    inst = ds_add_u32(addr=v[0], data0=v[1])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x00,0xd8,0x00,0x01,0x00,0x00]))
+
+  def test_ds_store_b32_gds(self):
+    # ds_store_b32 v0, v1 gds
+    # GFX11: encoding: [0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]
+    inst = ds_store_b32(addr=v[0], data0=v[1], gds=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x36,0xd8,0x00,0x01,0x00,0x00]))
+
+
+class TestVOP3(unittest.TestCase):
+  """Test VOP3 (3-operand vector) instructions."""
+
+  def test_v_fma_f32(self):
+    # v_fma_f32 v0, v1, v2, v3
+    # GFX11: encoding: [0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]
+    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x13,0xd6,0x01,0x05,0x0e,0x04]))
+
+  def test_v_mad_f32(self):
+    # v_fmac_f32_e64 v0, v1, v2 (fmac is fma with implicit dst as src2)
+    # Use v_fma_f32 with vdst == src2
+    inst = v_fma_f32(vdst=v[0], src0=v[1], src1=v[2], src2=v[0])
+    self.assertEqual(inst.to_bytes()[:4], bytes([0x00,0x00,0x13,0xd6]))
+
+  def test_v_add3_u32(self):
+    # v_add3_u32 v0, v1, v2, v3
+    # GFX11: encoding: [0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]
+    inst = v_add3_u32(vdst=v[0], src0=v[1], src1=v[2], src2=v[3])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x55,0xd6,0x01,0x05,0x0e,0x04]))
+
+
+class TestFLAT(unittest.TestCase):
+  """Test FLAT/GLOBAL/SCRATCH memory instructions."""
+
+  def test_global_load_b32(self):
+    # global_load_b32 v0, v[1:2], off (seg=2 for global)
+    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]
+    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF)
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x7c,0x00]))
+
+  def test_global_store_b32(self):
+    # global_store_b32 v[0:1], v2, off (seg=2 for global)
+    # GFX11: encoding: [0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]
+    inst = global_store_b32(addr=v[0:2], data=v[2], saddr=OFF)
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x6a,0xdc,0x00,0x02,0x7c,0x00]))
+
+  def test_global_load_b32_saddr(self):
+    # global_load_b32 v0, v1, s[0:1] (seg=2 for global)
+    # GFX11: encoding: [0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]
+    inst = global_load_b32(vdst=v[0], addr=v[1], saddr=s[0:2])
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x52,0xdc,0x01,0x00,0x00,0x00]))
+
+  def test_global_load_b32_offset(self):
+    # global_load_b32 v0, v[1:2], off offset:256 (seg=2 for global)
+    # GFX11: encoding: [0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]
+    inst = global_load_b32(vdst=v[0], addr=v[1:3], saddr=OFF, offset=256)
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x01,0x52,0xdc,0x01,0x00,0x7c,0x00]))
+
+  def test_global_load_b64(self):
+    # global_load_b64 v[0:1], v[2:3], off (seg=2 for global)
+    # GFX11: encoding: [0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]
+    inst = global_load_b64(vdst=v[0:2], addr=v[2:4], saddr=OFF)
+    self.assertEqual(inst.to_bytes(), bytes([0x00,0x00,0x56,0xdc,0x02,0x00,0x7c,0x00]))
+
+
+class TestSMEM(unittest.TestCase):
+  """Test SMEM (scalar memory) instructions - regression tests for glc/dlc bit positions."""
+
+  def test_smem_dlc_bit_position(self):
+    # s_load_b32 s5, s[2:3], s0 dlc - tests that DLC is at bit 13 (not bit 14)
+    # GFX11: encoding: [0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]
+    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], dlc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00]))
+
+  def test_smem_glc_bit_position(self):
+    # s_load_b32 s5, s[2:3], s0 glc - tests that GLC is at bit 14 (not bit 16)
+    # GFX11: encoding: [0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]
+    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x41,0x41,0x00,0xf4,0x00,0x00,0x00,0x00]))
+
+  def test_smem_glc_dlc_combined(self):
+    # s_load_b32 s5, s[2:3], s0 glc dlc - tests both flags together
+    # GFX11: encoding: [0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]
+    inst = s_load_b32(sdata=s[5], sbase=s[2], soffset=s[0], glc=1, dlc=1)
+    self.assertEqual(inst.to_bytes(), bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00]))
+
+  def test_smem_disasm_roundtrip_dlc(self):
+    # Test that disassembly/reassembly preserves DLC bit correctly
+    data = bytes([0x41,0x21,0x00,0xf4,0x00,0x00,0x00,0x00])
+    decoded = SMEM.from_bytes(data)
+    self.assertEqual(decoded.to_bytes(), data)
+
+  def test_smem_disasm_roundtrip_glc_dlc(self):
+    # Test that disassembly/reassembly preserves GLC+DLC bits correctly
+    data = bytes([0x41,0x61,0x00,0xf4,0x00,0x00,0x00,0x00])
+    decoded = SMEM.from_bytes(data)
+    self.assertEqual(decoded.to_bytes(), data)
+
+
+class TestVOP3Literal(unittest.TestCase):
+  """Test VOP3 literal handling - regression tests for Inst64 literal encoding."""
+
+  def test_vop3_with_literal(self):
+    # v_add3_u32 v5, vcc_hi, 0xaf123456, v255
+    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf]
+    from extra.assembly.rdna3.lib import RawImm
+    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=RawImm(107), src1=0xaf123456, src2=v[255])
+    expected = bytes([0x05,0x00,0x55,0xd6,0x6b,0xfe,0xfd,0x07,0x56,0x34,0x12,0xaf])
+    self.assertEqual(inst.to_bytes(), expected)
+
+  def test_vop3_literal_null_operand(self):
+    # v_add3_u32 v5, null, exec_lo, 0xaf123456
+    # GFX11: encoding: [0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf]
+    from extra.assembly.rdna3.lib import RawImm
+    inst = VOP3(VOP3Op.V_ADD3_U32, vdst=v[5], src0=NULL, src1=RawImm(126), src2=0xaf123456)
+    expected = bytes([0x05,0x00,0x55,0xd6,0x7c,0xfc,0xfc,0x03,0x56,0x34,0x12,0xaf])
+    self.assertEqual(inst.to_bytes(), expected)
+
+  def test_vop3p_with_literal(self):
+    # Test VOP3P literal encoding (also uses Inst64)
+    from extra.assembly.rdna3.lib import RawImm
+    inst = VOP3P(VOP3POp.V_PK_ADD_F16, vdst=v[5], src0=RawImm(240), src1=0x12345678, src2=v[0])
+    self.assertEqual(len(inst.to_bytes()), 12)  # 8 bytes + 4 byte literal
+
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_handwritten.py

@@ -0,0 +1,178 @@
+# do not change these tests. we need to fix bugs to make them pass
+# the Inst constructor should be looking at the types of the fields to correctly set the value
+
+import unittest, struct
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.lib import Inst
+from extra.assembly.rdna3.asm import asm
+from extra.assembly.rdna3.test.test_roundtrip import compile_asm
+
+class TestIntegration(unittest.TestCase):
+  inst: Inst
+  def tearDown(self):
+    if not hasattr(self, 'inst'): return
+    b = self.inst.to_bytes()
+    st = self.inst.disasm()
+    reasm = asm(st)
+    desc = f"{st:25s} {self.inst} {b!r} {reasm}"
+    self.assertEqual(b, compile_asm(st), desc)
+    # TODO: this compare should work for valid things
+    #self.assertEqual(self.inst, reasm)
+    self.assertEqual(repr(self.inst), repr(reasm))
+    print(desc)
+
+  def test_load_b128(self):
+    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 0)
+
+  def test_load_b128_wrong_size(self):
+    # this should have to be 4 regs on the loaded to
+    with self.assertRaises(Exception):
+      self.inst = s_load_b128(s[4:6], s[0:1], NULL, 0)
+
+  def test_mov_b32(self):
+    self.inst = s_mov_b32(s[80], s[0])
+
+  def test_mov_b64(self):
+    self.inst = s_mov_b64(s[80:81], s[0:1])
+
+  def test_mov_b32_wrong(self):
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b32(s[80:81], s[0:1])
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b32(s[80:81], s[0])
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b32(s[80], s[0:1])
+
+  def test_mov_b64_wrong(self):
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b64(s[80], s[0])
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b64(s[80], s[0:1])
+    with self.assertRaises(Exception):
+      self.inst = s_mov_b64(s[80:81], s[0])
+
+  def test_load_b128_no_0(self):
+    self.inst = s_load_b128(s[4:7], s[0:1], NULL)
+
+  def test_load_b128_s(self):
+    self.inst = s_load_b128(s[4:7], s[0:1], s[8], 0)
+
+  def test_load_b128_v(self):
+    with self.assertRaises(TypeError):
+      self.inst = s_load_b128(s[4:7], s[0:1], v[8], 0)
+
+  def test_load_b128_off(self):
+    self.inst = s_load_b128(s[4:7], s[0:1], NULL, 3)
+
+  def test_simple_stos(self):
+    self.inst = s_mov_b32(s[0], s[1])
+
+  def test_simple_wrong(self):
+    with self.assertRaises(TypeError):
+      self.inst = s_mov_b32(v[0], s[1])
+
+  def test_simple_vtov(self):
+    self.inst = v_mov_b32_e32(v[0], v[1])
+
+  def test_simple_stov(self):
+    self.inst = v_mov_b32_e32(v[0], s[2])
+
+  def test_simple_float_to_v(self):
+    self.inst = v_mov_b32_e32(v[0], 1.0)
+
+  def test_simple_v_to_float(self):
+    with self.assertRaises(TypeError):
+      self.inst = v_mov_b32_e32(1, v[0])
+
+  def test_simple_int_to_v(self):
+    self.inst = v_mov_b32_e32(v[0], 1)
+
+  def test_three_add(self):
+    self.inst = v_add_co_ci_u32_e32(v[3], s[7], v[3])
+
+  def test_three_add_v(self):
+    self.inst = v_add_co_ci_u32_e32(v[3], v[7], v[3])
+
+  def test_three_add_const(self):
+    self.inst = v_add_co_ci_u32_e32(v[3], 2.0, v[3])
+
+  def test_swaitcnt_lgkm(self): self.inst = s_waitcnt(0xfc07)
+  def test_swaitcnt_vm(self): self.inst = s_waitcnt(0x03f7)
+
+  def test_vmad(self):
+    self.inst = v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2])
+
+  def test_large_imm(self):
+    self.inst = v_mov_b32_e32(v[0], 0x1234)
+
+  def test_dual_mov(self):
+    self.inst = VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_MOV_B32, vdstx=v[0], vdsty=v[1], srcx0=v[2], srcy0=v[4])
+
+  def test_dual_mul(self):
+    self.inst = v_dual_mul_f32(VOPDOp.V_DUAL_MUL_F32, vdstx=v[0], vdsty=v[1], srcx0=v[2], vsrcx1=v[3], srcy0=v[4], vsrcy1=v[5])
+
+  def test_simple_int_to_s(self):
+    self.inst = s_mov_b32(s[0], 3)
+
+  def test_complex_int_to_s(self):
+    self.inst = s_mov_b32(s[0], 0x235646)
+
+  def test_simple_float_to_s(self):
+    self.inst = s_mov_b32(s[0], 1.0)
+
+  def test_complex_float_to_s(self):
+    self.inst = s_mov_b32(s[0], 1337.0)
+    int_inst = s_mov_b32(s[0], struct.unpack("I", struct.pack("f", 1337.0))[0])
+    self.assertEqual(self.inst, int_inst)
+
+class TestRegisterSliceSyntax(unittest.TestCase):
+  """
+  Issue: Register slice syntax should use AMD assembly convention (inclusive end).
+
+  In AMD assembly, s[4:7] means registers s4, s5, s6, s7 (4 registers, inclusive).
+  The DSL should match this convention so that:
+  - s[4:7] gives 4 registers
+  - Disassembler output can be copied directly back into DSL code
+
+  Fix: Change _RegFactory.__getitem__ to use inclusive end:
+    key.stop - key.start + 1  (instead of key.stop - key.start)
+  """
+  def test_register_slice_count(self):
+    # s[4:7] should give 4 registers: s4, s5, s6, s7 (AMD convention, inclusive)
+    reg = s[4:7]
+    self.assertEqual(reg.count, 4, "s[4:7] should give 4 registers (s4, s5, s6, s7)")
+
+  def test_register_slice_roundtrip(self):
+    # Round-trip: DSL -> disasm -> DSL should preserve register count
+    reg = s[4:7]  # 4 registers in AMD convention
+    inst = s_load_b128(reg, s[0:1], NULL, 0)
+    disasm = inst.disasm()
+    # Disasm shows s[4:7] - user should be able to copy this back
+    self.assertIn("s[4:7]", disasm)
+    # And s[4:7] in DSL should give the same 4 registers
+    reg_from_disasm = s[4:7]
+    self.assertEqual(reg_from_disasm.count, 4, "s[4:7] from disasm should give 4 registers")
+
+class TestInstructionEquality(unittest.TestCase):
+  """
+  Issue: No __eq__ method - instruction comparison requires repr() workaround.
+
+  Two identical instructions should compare equal with ==, but currently:
+    inst1 == inst2 returns False
+
+  The test_handwritten.py works around this with:
+    self.assertEqual(repr(self.inst), repr(reasm))
+  """
+  def test_identical_instructions_equal(self):
+    inst1 = v_mov_b32_e32(v[0], v[1])
+    inst2 = v_mov_b32_e32(v[0], v[1])
+    self.assertEqual(inst1, inst2, "identical instructions should be equal")
+
+  def test_different_instructions_not_equal(self):
+    inst1 = v_mov_b32_e32(v[0], v[1])
+    inst2 = v_mov_b32_e32(v[0], v[2])
+    self.assertNotEqual(inst1, inst2, "different instructions should not be equal")
+
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_integration.py

@@ -0,0 +1,343 @@
+#!/usr/bin/env python3
+"""Integration test: round-trip RDNA3 assembly through AMD toolchain."""
+import unittest, re, io, sys
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.asm import waitcnt, asm
+
+def get_amd_toolchain():
+  """Check if AMD toolchain is available."""
+  try:
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    HIPCompiler("gfx1100").compile(".text\ns_endpgm")
+    return True
+  except Exception:
+    return False
+
+def disassemble(lib: bytes, arch: str = "gfx1100") -> str:
+  """Disassemble ELF binary using tinygrad's compiler, return raw output."""
+  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+  old_stdout = sys.stdout
+  sys.stdout = io.StringIO()
+  HIPCompiler(arch).disassemble(lib)
+  output = sys.stdout.getvalue()
+  sys.stdout = old_stdout
+  return output
+
+def parse_disassembly(raw: str) -> list[str]:
+  """Parse disassembly output to list of instruction mnemonics."""
+  lines = []
+  for line in raw.splitlines():
+    if line.startswith('\t'):
+      instr = line.split('//')[0].strip()
+      if instr: lines.append(instr)
+  return lines
+
+def assemble_and_disassemble(instructions: list, arch: str = "gfx1100") -> list[str]:
+  """Assemble instructions with our DSL, then disassemble with AMD toolchain."""
+  from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+  # Generate bytes from our DSL
+  code_bytes = b''.join(inst.to_bytes() for inst in instructions)
+
+  # Wrap in minimal ELF-compatible assembly with .byte directives
+  byte_str = ', '.join(f'0x{b:02x}' for b in code_bytes)
+  asm_src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n.byte {byte_str}\n"
+
+  # Assemble with AMD COMGR and disassemble
+  lib = HIPCompiler(arch).compile(asm_src)
+  return parse_disassembly(disassemble(lib, arch))
+
+@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
+class TestIntegration(unittest.TestCase):
+  """Test our assembler output matches LLVM disassembly."""
+
+  def test_simple_sop1(self):
+    """Test SOP1 instructions round-trip."""
+    instructions = [
+      s_mov_b32(s[0], s[1]),
+      s_mov_b32(s[2], 0),
+      s_not_b32(s[3], s[4]),
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    self.assertIn('s_mov_b32', disasm[0])
+    self.assertIn('s_mov_b32', disasm[1])
+    self.assertIn('s_not_b32', disasm[2])
+
+  def test_simple_sop2(self):
+    """Test SOP2 instructions round-trip."""
+    instructions = [
+      s_add_u32(s[0], s[1], s[2]),
+      s_sub_u32(s[3], s[4], 10),
+      s_and_b32(s[5], s[6], s[7]),
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    self.assertIn('s_add_u32', disasm[0])
+    self.assertIn('s_sub_u32', disasm[1])
+    self.assertIn('s_and_b32', disasm[2])
+
+  def test_simple_vop2(self):
+    """Test VOP2 instructions round-trip."""
+    instructions = [
+      v_add_f32_e32(v[0], v[1], v[2]),
+      v_mul_f32_e32(v[3], 1.0, v[4]),  # 1.0 is inline constant
+      v_and_b32_e32(v[5], 10, v[6]),  # small inline constant
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    self.assertIn('v_add_f32', disasm[0])
+    self.assertIn('v_mul_f32', disasm[1])
+
+  def test_control_flow(self):
+    """Test control flow instructions."""
+    instructions = [
+      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
+      s_endpgm(),
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    self.assertIn('s_waitcnt', disasm[0])
+    self.assertIn('s_endpgm', disasm[1])
+
+  def test_memory_ops(self):
+    """Test memory instructions."""
+    instructions = [
+      s_load_b32(s[0], s[0:2], NULL),
+      s_waitcnt(simm16=waitcnt(lgkmcnt=0)),
+      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
+      s_endpgm(),
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    self.assertIn('s_load_b32', disasm[0])
+    self.assertIn('s_waitcnt', disasm[1])
+    self.assertIn('global_store_b32', disasm[2])
+
+  def test_full_kernel(self):
+    """Test a complete kernel similar to tinygrad output."""
+    # Simple kernel: load value, add 1, store back
+    instructions = [
+      # Get thread ID
+      v_mov_b32_e32(v[0], s[0]),  # base addr low
+      v_mov_b32_e32(v[1], s[1]),  # base addr high
+      # Load value
+      global_load_b32(vdst=v[2], addr=v[0:2], saddr=OFF),
+      s_waitcnt(simm16=waitcnt(vmcnt=0)),
+      # Add 1.0
+      v_add_f32_e32(v[2], 1.0, v[2]),
+      # Store result
+      global_store_b32(addr=v[0:2], data=v[2], saddr=OFF),
+      s_endpgm(),
+    ]
+    disasm = assemble_and_disassemble(instructions)
+    # Verify key instructions are present
+    self.assertTrue(any('global_load' in d for d in disasm))
+    self.assertTrue(any('v_add_f32' in d for d in disasm))
+    self.assertTrue(any('global_store' in d for d in disasm))
+    self.assertTrue(any('s_endpgm' in d for d in disasm))
+
+  def test_bytes_roundtrip(self):
+    """Test that our bytes match what AMD assembler produces."""
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+    # Simple instruction
+    inst = s_mov_b32(s[0], s[1])
+    our_bytes = inst.to_bytes()
+
+    # Assemble same instruction with AMD toolchain
+    asm_src = ".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\ns_mov_b32 s0, s1\n"
+    compiler = HIPCompiler("gfx1100")
+    lib = compiler.compile(asm_src)
+    raw = disassemble(lib)
+
+    for line in raw.splitlines():
+      if 's_mov_b32' in line and '//' in line:
+        # Extract hex bytes from comment: "// 000000001300: BE800001"
+        comment = line.split('//')[1].strip()
+        hex_str = comment.split(':')[1].strip()
+        # Convert big-endian hex string to little-endian bytes
+        amd_bytes = bytes.fromhex(hex_str)[::-1]  # reverse for little-endian
+        self.assertEqual(our_bytes, amd_bytes, f"Bytes mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
+        return
+    self.fail("Could not find s_mov_b32 in disassembly")
+
+@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
+class TestAsm(unittest.TestCase):
+  """Test asm() string parsing."""
+
+  def test_asm_basic(self):
+    """Test basic instruction parsing."""
+    inst = asm('s_mov_b32 s0, s1')
+    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], s[1]).to_bytes())
+
+  def test_asm_with_immediates(self):
+    """Test parsing with immediate values."""
+    inst = asm('s_add_u32 s0, s1, 10')
+    self.assertEqual(inst.to_bytes(), s_add_u32(s[0], s[1], 10).to_bytes())
+
+  def test_asm_float_const(self):
+    """Test parsing float constants."""
+    inst = asm('v_mul_f32_e32 v0, 1.0, v1')
+    self.assertEqual(inst.to_bytes(), v_mul_f32_e32(v[0], 1.0, v[1]).to_bytes())
+
+  def test_asm_hex_immediate(self):
+    """Test parsing hex immediates."""
+    inst = asm('s_waitcnt 0xfc07')
+    self.assertEqual(inst.to_bytes(), s_waitcnt(simm16=0xfc07).to_bytes())
+
+  def test_asm_special_regs(self):
+    """Test parsing special registers."""
+    inst = asm('s_mov_b32 s0, vcc_lo')
+    self.assertEqual(inst.to_bytes(), s_mov_b32(s[0], VCC_LO).to_bytes())
+
+  def test_asm_register_range(self):
+    """Test parsing register ranges."""
+    inst = asm('s_load_b128 s[4:7], s[0:1], null')
+    self.assertEqual(inst.to_bytes(), s_load_b128(s[4:7], s[0:1], NULL).to_bytes())
+
+  def test_asm_matches_llvm(self):
+    """Test asm() output matches LLVM assembler."""
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    compiler = HIPCompiler('gfx1100')
+
+    def get_llvm_bytes(instr: str) -> bytes:
+      src = f'.text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n{instr}\n'
+      lib = compiler.compile(src)
+      raw = disassemble(lib)
+      for line in raw.splitlines():
+        if instr.split()[0] in line and '//' in line:
+          hex_str = line.split('//')[1].strip().split(':')[1].strip()
+          return bytes.fromhex(hex_str)[::-1]
+      return b''
+
+    tests = ['s_mov_b32 s0, s1', 's_endpgm', 'v_add_f32_e32 v0, v1, v2']
+    for t in tests:
+      self.assertEqual(asm(t).to_bytes(), get_llvm_bytes(t), f"mismatch for: {t}")
+
+  def test_asm_vop3_modifiers(self):
+    """Test asm() with VOP3 modifiers (neg, abs, clamp)."""
+    import subprocess, re
+
+    def get_llvm_encoding(instr: str) -> str:
+      result = subprocess.run(['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-show-encoding'],
+                              input=instr, capture_output=True, text=True)
+      if m := re.search(r'encoding:\s*\[(.*?)\]', result.stdout):
+        return m.group(1).replace('0x','').replace(',','').replace(' ','')
+      return ''
+
+    tests = [
+      'v_fma_f32 v0, -v1, v2, v3',       # neg on src0
+      'v_fma_f32 v0, v1, |v2|, v3',      # abs on src1
+      'v_fma_f32 v0, v1, v2, v3 clamp',  # clamp
+      'v_fma_f32 v0, -v1, |v2|, v3 clamp',  # all modifiers
+      'v_fma_f32 v0, -|v1|, v2, v3',     # neg+abs on same operand
+    ]
+    for t in tests:
+      our_hex = asm(t).to_bytes().hex()
+      llvm_hex = get_llvm_encoding(t)
+      self.assertEqual(our_hex, llvm_hex, f"mismatch for: {t}")
+
+@unittest.skipUnless(get_amd_toolchain(), "AMD toolchain not available")
+class TestTinygradIntegration(unittest.TestCase):
+  """Test that we can parse disassembled tinygrad kernels."""
+
+  def test_simple_add_kernel(self):
+    """Generate a simple add kernel from tinygrad and verify disassembly."""
+    from tinygrad import Tensor
+    from tinygrad.codegen import get_program
+    from tinygrad.renderer.cstyle import AMDHIPRenderer
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    from tinygrad.uop.ops import Ops
+
+    # Create a computation that generates a real kernel
+    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    b = Tensor([5.0, 6.0, 7.0, 8.0]).realize()
+    c = a + b
+
+    # Get schedule and find SINK
+    schedule = c.schedule()
+    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
+    self.assertTrue(len(sink_items) > 0, "No SINK in schedule")
+
+    # Generate program
+    renderer = AMDHIPRenderer('gfx1100')
+    prg = get_program(sink_items[0].ast, renderer)
+    self.assertIsNotNone(prg.src)
+
+    # Compile and disassemble
+    compiler = HIPCompiler('gfx1100')
+    lib = compiler.compile(prg.src)
+    raw_disasm = disassemble(lib)
+    instrs = parse_disassembly(raw_disasm)
+
+    # Verify we got some instructions
+    self.assertTrue(len(instrs) > 0, "No instructions in disassembly")
+    # Should have an endpgm
+    self.assertTrue(any('s_endpgm' in i for i in instrs), "Missing s_endpgm")
+
+  def test_matmul_kernel(self):
+    """Generate a matmul kernel and verify disassembly has expected patterns."""
+    from tinygrad import Tensor
+    from tinygrad.codegen import get_program
+    from tinygrad.renderer.cstyle import AMDHIPRenderer
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    from tinygrad.uop.ops import Ops
+
+    # Create a small matmul
+    a = Tensor.rand(4, 4).realize()
+    b = Tensor.rand(4, 4).realize()
+    c = a @ b
+
+    # Get schedule
+    schedule = c.schedule()
+    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
+    self.assertTrue(len(sink_items) > 0)
+
+    # Generate and compile
+    renderer = AMDHIPRenderer('gfx1100')
+    prg = get_program(sink_items[0].ast, renderer)
+    compiler = HIPCompiler('gfx1100')
+    lib = compiler.compile(prg.src)
+    raw_disasm = disassemble(lib)
+    instrs = parse_disassembly(raw_disasm)
+
+    # Matmul should have multiply and add instructions
+    has_mul = any('mul' in i.lower() for i in instrs)
+    has_add = any('add' in i.lower() for i in instrs)
+    self.assertTrue(has_mul or has_add, "Matmul should have mul/add ops")
+
+  def test_disasm_to_bytes_roundtrip(self):
+    """Parse disassembled instructions and verify we can re-encode some of them."""
+    from tinygrad import Tensor
+    from tinygrad.codegen import get_program
+    from tinygrad.renderer.cstyle import AMDHIPRenderer
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    from tinygrad.uop.ops import Ops
+
+    # Simple kernel
+    a = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
+    b = (a * 2.0)
+
+    schedule = b.schedule()
+    sink_items = [si for si in schedule if si.ast.op == Ops.SINK]
+    if not sink_items: return  # skip if no kernel
+
+    renderer = AMDHIPRenderer('gfx1100')
+    prg = get_program(sink_items[0].ast, renderer)
+    compiler = HIPCompiler('gfx1100')
+    lib = compiler.compile(prg.src)
+    raw_disasm = disassemble(lib)
+
+    # Find s_endpgm and verify we can encode it
+    for line in raw_disasm.splitlines():
+      if 's_endpgm' in line and '//' in line:
+        # Extract bytes from comment
+        comment = line.split('//')[1].strip()
+        hex_str = comment.split(':')[1].strip()
+        amd_bytes = bytes.fromhex(hex_str)[::-1]
+
+        # Our encoding
+        our_inst = s_endpgm()
+        our_bytes = our_inst.to_bytes()
+
+        self.assertEqual(our_bytes, amd_bytes, f"s_endpgm mismatch: ours={our_bytes.hex()} AMD={amd_bytes.hex()}")
+        return
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_llvm.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+"""Test RDNA3 assembler/disassembler against LLVM test vectors."""
+import unittest, re
+from tinygrad.helpers import fetch
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.asm import asm
+from extra.assembly.rdna3.test.test_roundtrip import compile_asm, disassemble_lib
+
+LLVM_BASE = "https://raw.githubusercontent.com/llvm/llvm-project/main/llvm/test/MC/AMDGPU"
+
+# Format info: (filename, format_class, op_enum)
+LLVM_TEST_FILES = {
+  # Scalar ALU
+  'sop1': ('gfx11_asm_sop1.s', SOP1, SOP1Op),
+  'sop2': ('gfx11_asm_sop2.s', SOP2, SOP2Op),
+  'sopp': ('gfx11_asm_sopp.s', SOPP, SOPPOp),
+  'sopk': ('gfx11_asm_sopk.s', SOPK, SOPKOp),
+  'sopc': ('gfx11_asm_sopc.s', SOPC, SOPCOp),
+  # Vector ALU
+  'vop1': ('gfx11_asm_vop1.s', VOP1, VOP1Op),
+  'vop2': ('gfx11_asm_vop2.s', VOP2, VOP2Op),
+  'vopc': ('gfx11_asm_vopc.s', VOPC, VOPCOp),
+  'vop3': ('gfx11_asm_vop3.s', VOP3, VOP3Op),
+  'vop3p': ('gfx11_asm_vop3p.s', VOP3P, VOP3POp),
+  'vop3sd': ('gfx11_asm_vop3.s', VOP3SD, VOP3SDOp),  # VOP3SD shares file with VOP3
+  'vinterp': ('gfx11_asm_vinterp.s', VINTERP, VINTERPOp),
+  'vopd': ('gfx11_asm_vopd.s', VOPD, VOPDOp),
+  'vopcx': ('gfx11_asm_vopcx.s', VOPC, VOPCOp),  # VOPCX uses VOPC format
+  # VOP3 promotions (VOP1/VOP2/VOPC promoted to VOP3 encoding)
+  'vop3_from_vop1': ('gfx11_asm_vop3_from_vop1.s', VOP3, VOP3Op),
+  'vop3_from_vop2': ('gfx11_asm_vop3_from_vop2.s', VOP3, VOP3Op),
+  'vop3_from_vopc': ('gfx11_asm_vop3_from_vopc.s', VOP3, VOP3Op),
+  'vop3_from_vopcx': ('gfx11_asm_vop3_from_vopcx.s', VOP3, VOP3Op),
+  # Memory
+  'ds': ('gfx11_asm_ds.s', DS, DSOp),
+  'smem': ('gfx11_asm_smem.s', SMEM, SMEMOp),
+  'flat': ('gfx11_asm_flat.s', FLAT, FLATOp),
+  'mubuf': ('gfx11_asm_mubuf.s', MUBUF, MUBUFOp),
+  'mtbuf': ('gfx11_asm_mtbuf.s', MTBUF, MTBUFOp),
+  'mimg': ('gfx11_asm_mimg.s', MIMG, MIMGOp),
+  # WMMA (matrix multiply)
+  'wmma': ('gfx11_asm_wmma.s', VOP3P, VOP3POp),
+  # Additional features
+  'vop3_features': ('gfx11_asm_vop3_features.s', VOP3, VOP3Op),
+  'vop3p_features': ('gfx11_asm_vop3p_features.s', VOP3P, VOP3POp),
+  'vopd_features': ('gfx11_asm_vopd_features.s', VOPD, VOPDOp),
+  # Alias files (alternative mnemonics)
+  'vop3_alias': ('gfx11_asm_vop3_alias.s', VOP3, VOP3Op),
+  'vop3p_alias': ('gfx11_asm_vop3p_alias.s', VOP3P, VOP3POp),
+  'vopc_alias': ('gfx11_asm_vopc_alias.s', VOPC, VOPCOp),
+  'vopcx_alias': ('gfx11_asm_vopcx_alias.s', VOPC, VOPCOp),
+  'vinterp_alias': ('gfx11_asm_vinterp_alias.s', VINTERP, VINTERPOp),
+  'smem_alias': ('gfx11_asm_smem_alias.s', SMEM, SMEMOp),
+  'mubuf_alias': ('gfx11_asm_mubuf_alias.s', MUBUF, MUBUFOp),
+  'mtbuf_alias': ('gfx11_asm_mtbuf_alias.s', MTBUF, MTBUFOp),
+}
+
+def parse_llvm_tests(text: str) -> list[tuple[str, bytes]]:
+  """Parse LLVM test format into (asm, expected_bytes) pairs."""
+  tests, lines = [], text.split('\n')
+  for i, line in enumerate(lines):
+    line = line.strip()
+    if not line or line.startswith(('//', '.', ';')): continue
+    asm_text = line.split('//')[0].strip()
+    if not asm_text: continue
+    for j in range(i, min(i + 3, len(lines))):
+      # Match GFX11, W32, or W64 encodings (all valid for gfx11)
+      if m := re.search(r'(?:GFX11|W32|W64)[^:]*:.*?encoding:\s*\[(.*?)\]', lines[j]):
+        hex_bytes = m.group(1).replace('0x', '').replace(',', '').replace(' ', '')
+        if hex_bytes:
+          try: tests.append((asm_text, bytes.fromhex(hex_bytes)))
+          except ValueError: pass
+        break
+  return tests
+
+def try_assemble(text: str):
+  """Try to assemble instruction text, return bytes or None on failure."""
+  try: return asm(text).to_bytes()
+  except: return None
+
+class TestLLVM(unittest.TestCase):
+  """Test assembler and disassembler against all LLVM test vectors."""
+  tests: dict[str, list[tuple[str, bytes]]] = {}
+
+  @classmethod
+  def setUpClass(cls):
+    for name, (filename, _, _) in LLVM_TEST_FILES.items():
+      try:
+        data = fetch(f"{LLVM_BASE}/{filename}").read_bytes()
+        cls.tests[name] = parse_llvm_tests(data.decode('utf-8', errors='ignore'))
+      except Exception as e:
+        print(f"Warning: couldn't fetch {filename}: {e}")
+        cls.tests[name] = []
+
+# Generate test methods dynamically for each format
+def _make_asm_test(name):
+  def test(self):
+    passed, failed, skipped = 0, 0, 0
+    for asm_text, expected in self.tests.get(name, []):
+      result = try_assemble(asm_text)
+      if result is None: skipped += 1
+      elif result == expected: passed += 1
+      else: failed += 1
+    print(f"{name.upper()} asm: {passed} passed, {failed} failed, {skipped} skipped")
+    self.assertEqual(failed, 0)
+  return test
+
+def _make_disasm_test(name):
+  def test(self):
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+    compiler = HIPCompiler('gfx1100')
+    _, fmt_cls, op_enum = LLVM_TEST_FILES[name]
+    passed, failed, skipped, failures = 0, 0, 0, []
+    # VOP3SD opcodes that share encoding with VOP3 (only for vop3sd test, not vopc promotions)
+    # Note: opcodes 0-255 are VOPC promoted to VOP3, never VOP3SD
+    vop3sd_opcodes = {288, 289, 290, 764, 765, 766, 767, 768, 769, 770}
+    # vop3_from_vopc/vopcx tests have VOPC opcodes 0-255, not VOP3SD - don't detect as VOP3SD
+    is_vopc_promotion = name in ('vop3_from_vopc', 'vop3_from_vopcx')
+    # Undocumented opcodes not in AMD ISA PDF - skip these
+    undocumented = {'smem': {34, 35}, 'sopk': {22, 23}, 'sopp': {8, 58, 59}}  # s_atc_probe*, s_subvector_loop*, s_waitcnt_depctr, unknown
+    for asm_text, data in self.tests.get(name, []):
+      if len(data) > fmt_cls._size(): continue  # skip literals (need different handling)
+      # Skip undocumented opcodes
+      temp_inst = fmt_cls.from_bytes(data)
+      temp_op = temp_inst._values.get('op', 0)
+      temp_op = temp_op.val if hasattr(temp_op, 'val') else temp_op
+      if temp_op in undocumented.get(name, set()): skipped += 1; continue
+      # Skip SOPP no-imm instructions with non-zero simm16 (can't roundtrip through LLVM)
+      if name == 'sopp':
+        simm16 = temp_inst._values.get('simm16', 0)
+        simm16 = simm16.val if hasattr(simm16, 'val') else simm16
+        sopp_no_imm = {48, 54, 53, 55, 60, 61, 62}  # s_endpgm, s_barrier, s_wakeup, s_icache_inv, s_wait_idle, s_endpgm_saved, s_code_end
+        if temp_op in sopp_no_imm and simm16 != 0: skipped += 1; continue
+      try:
+        # VOP3 and VOP3SD share encoding - peek at opcode to determine which class to use
+        if fmt_cls.__name__ in ('VOP3', 'VOP3SD'):
+          temp = VOP3.from_bytes(data)
+          op_val = temp._values.get('op', 0)
+          op_val = op_val.val if hasattr(op_val, 'val') else op_val
+          is_vop3sd = (op_val in vop3sd_opcodes) and not is_vopc_promotion
+          decoded = VOP3SD.from_bytes(data) if is_vop3sd else VOP3.from_bytes(data)
+          # Validate opcode with appropriate enum
+          if is_vop3sd:
+            VOP3SDOp(op_val)
+          else:
+            VOP3Op(op_val)
+        else:
+          decoded = fmt_cls.from_bytes(data)
+          op_val = decoded._values.get('op', 0)
+          op_val = op_val.val if hasattr(op_val, 'val') else op_val
+          op_enum(op_val)  # validate opcode
+        if decoded.to_bytes()[:len(data)] != data:
+          failed += 1; failures.append(f"decode roundtrip failed for {data.hex()}"); continue
+        disasm_str = decoded.disasm()
+        # Test: LLVM should assemble our disasm output to the same bytes
+        llvm_bytes = compile_asm(disasm_str, compiler)
+        if llvm_bytes is None:
+          failed += 1; failures.append(f"LLVM failed to assemble: '{disasm_str}' (from '{asm_text}')")
+        elif llvm_bytes == data: passed += 1
+        else: failed += 1; failures.append(f"'{disasm_str}': expected={data.hex()} got={llvm_bytes.hex()}")
+      except Exception as e:
+        failed += 1; failures.append(f"exception for {data.hex()}: {e}")
+    print(f"{name.upper()} disasm: {passed} passed, {failed} failed" + (f", {skipped} skipped" if skipped else ""))
+    if failures[:10]: print("  " + "\n  ".join(failures[:10]))
+    self.assertEqual(failed, 0)
+  return test
+
+for name in LLVM_TEST_FILES:
+  setattr(TestLLVM, f'test_{name}_asm', _make_asm_test(name))
+  setattr(TestLLVM, f'test_{name}_disasm', _make_disasm_test(name))
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_mockgpu_invalid.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+"""Test that invalid instructions raise exceptions through the mock GPU stack."""
+import unittest, subprocess, os, time
+
+class TestMockGPUInvalidInstruction(unittest.TestCase):
+  def test_unsupported_instruction_raises(self):
+    """Test that unsupported instructions raise immediately through the full MOCKGPU stack."""
+    test_code = '''
+import struct
+from tinygrad import Device, Tensor
+from tinygrad.engine.realize import get_runner
+from tinygrad.runtime.ops_amd import AMDProgram
+
+dev = Device["AMD"]
+a = Tensor([1.0]).realize()
+b = a + 1
+si = b.schedule()[-1]
+runner = get_runner(dev.device, si.ast)
+
+prg = runner._prg
+lib = bytearray(prg.lib)
+
+# Find s_endpgm (0xBFB00000) and replace with invalid SOPP op=127 (0xBFFF0000)
+found = False
+for i in range(0, len(lib) - 4, 4):
+  if struct.unpack("<I", lib[i:i+4])[0] == 0xBFB00000:
+    lib[i:i+4] = struct.pack("<I", 0xBFFF0000)
+    found = True
+    break
+assert found, "s_endpgm not found"
+
+patched_prg = AMDProgram(dev, "patched", bytes(lib))
+b.uop.buffer.allocate()
+patched_prg(b.uop.buffer._buf, a.uop.buffer._buf, global_size=(1,1,1), local_size=(1,1,1))
+dev.synchronize()
+'''
+
+    env = os.environ.copy()
+    env["AMD"] = "1"
+    env["MOCKGPU"] = "1"
+    env["PYTHON_REMU"] = "1"
+    env["HCQDEV_WAIT_TIMEOUT_MS"] = "10000"
+
+    st = time.perf_counter()
+    result = subprocess.run(["python", "-c", test_code], env=env, capture_output=True, text=True, timeout=60)
+    elapsed = time.perf_counter() - st
+
+    self.assertNotEqual(result.returncode, 0, "should have raised")
+    self.assertIn("NotImplementedError", result.stderr)
+    # Should exit immediately, not wait for the full timeout
+    self.assertLess(elapsed, 5.0, f"should exit immediately on emulator exception, took {elapsed:.1f}s")
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_pdf_parser.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+"""Test that PDF parser correctly extracts format fields."""
+import unittest
+from extra.assembly.rdna3.autogen import (
+  SOP1, SOP2, SOPK, SOPP, VOP1, VOP2, VOP3SD, VOPC, FLAT, VOPD,
+  SOP1Op, SOP2Op, VOP1Op, VOP3Op
+)
+
+# expected formats with key fields and whether they have ENCODING
+EXPECTED_FORMATS = {
+  'DPP16': (['SRC0', 'DPP_CTRL', 'BANK_MASK', 'ROW_MASK'], False),
+  'DPP8': (['SRC0', 'LANE_SEL0', 'LANE_SEL7'], False),
+  'DS': (['OP', 'ADDR', 'DATA0', 'DATA1', 'VDST'], True),
+  'EXP': (['EN', 'TARGET', 'VSRC0', 'VSRC1', 'VSRC2', 'VSRC3'], True),
+  'FLAT': (['OP', 'ADDR', 'DATA', 'SADDR', 'VDST', 'OFFSET'], True),
+  'LDSDIR': (['VDST', 'OP'], True),
+  'MIMG': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'DMASK'], True),
+  'MTBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'FORMAT', 'SOFFSET'], True),
+  'MUBUF': (['OP', 'VADDR', 'VDATA', 'SRSRC', 'SOFFSET'], True),
+  'SMEM': (['OP', 'SBASE', 'SDATA', 'OFFSET', 'SOFFSET'], True),
+  'SOP1': (['OP', 'SDST', 'SSRC0'], True),
+  'SOP2': (['OP', 'SDST', 'SSRC0', 'SSRC1'], True),
+  'SOPC': (['OP', 'SSRC0', 'SSRC1'], True),
+  'SOPK': (['OP', 'SDST', 'SIMM16'], True),
+  'SOPP': (['OP', 'SIMM16'], True),
+  'VINTERP': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
+  'VOP1': (['OP', 'VDST', 'SRC0'], True),
+  'VOP2': (['OP', 'VDST', 'SRC0', 'VSRC1'], True),
+  'VOP3': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
+  'VOP3P': (['OP', 'VDST', 'SRC0', 'SRC1', 'SRC2'], True),
+  'VOP3SD': (['OP', 'VDST', 'SDST', 'SRC0', 'SRC1', 'SRC2'], True),
+  'VOPC': (['OP', 'SRC0', 'VSRC1'], True),
+  'VOPD': (['OPX', 'OPY', 'SRCX0', 'SRCY0', 'VDSTX', 'VDSTY'], True),
+}
+
+class TestPDFParserGenerate(unittest.TestCase):
+  """Test the PDF parser by running generate() and checking results."""
+  result: dict
+
+  @classmethod
+  def setUpClass(cls):
+    from extra.assembly.rdna3.gen import generate
+    cls.result = generate()
+
+  def test_all_formats_present(self):
+    """All expected formats should be parsed."""
+    for fmt_name in EXPECTED_FORMATS:
+      self.assertIn(fmt_name, self.result["formats"], f"missing format {fmt_name}")
+
+  def test_format_count(self):
+    """Should have exactly 23 formats."""
+    self.assertEqual(len(self.result["formats"]), 23)
+
+  def test_no_duplicate_fields(self):
+    """No format should have duplicate field names."""
+    for fmt_name, fields in self.result["formats"].items():
+      field_names = [f[0] for f in fields]
+      self.assertEqual(len(field_names), len(set(field_names)), f"{fmt_name} has duplicate fields: {field_names}")
+
+  def test_expected_fields(self):
+    """Each format should have its expected key fields."""
+    for fmt_name, (expected_fields, has_encoding) in EXPECTED_FORMATS.items():
+      fields = {f[0] for f in self.result["formats"].get(fmt_name, [])}
+      for field in expected_fields:
+        self.assertIn(field, fields, f"{fmt_name} missing {field}")
+      if has_encoding:
+        self.assertIn("ENCODING", fields, f"{fmt_name} should have ENCODING")
+      else:
+        self.assertNotIn("ENCODING", fields, f"{fmt_name} should not have ENCODING")
+
+  def test_vopd_no_dpp16_fields(self):
+    """VOPD should not have DPP16-specific fields (parser boundary bug)."""
+    vopd_fields = {f[0] for f in self.result["formats"].get("VOPD", [])}
+    for field in ['DPP_CTRL', 'BANK_MASK', 'ROW_MASK']:
+      self.assertNotIn(field, vopd_fields, f"VOPD should not have {field}")
+
+  def test_dpp16_no_vinterp_fields(self):
+    """DPP16 should not have VINTERP-specific fields."""
+    dpp16_fields = {f[0] for f in self.result["formats"].get("DPP16", [])}
+    for field in ['VDST', 'WAITEXP']:
+      self.assertNotIn(field, dpp16_fields, f"DPP16 should not have {field}")
+
+  def test_sopp_no_smem_fields(self):
+    """SOPP should not have SMEM fields (page break bug)."""
+    sopp_fields = {f[0] for f in self.result["formats"].get("SOPP", [])}
+    for field in ['SBASE', 'SDATA']:
+      self.assertNotIn(field, sopp_fields, f"SOPP should not have {field}")
+
+class TestPDFParser(unittest.TestCase):
+  """Verify format classes have correct fields from PDF parsing."""
+
+  def test_sop2_fields(self):
+    """SOP2 should have op, sdst, ssrc0, ssrc1."""
+    for field in ['op', 'sdst', 'ssrc0', 'ssrc1']:
+      self.assertIn(field, SOP2._fields)
+    self.assertEqual(SOP2._fields['op'].hi, 29)
+    self.assertEqual(SOP2._fields['op'].lo, 23)
+
+  def test_sop1_fields(self):
+    """SOP1 should have op, sdst, ssrc0 with correct bit positions."""
+    for field in ['op', 'sdst', 'ssrc0']:
+      self.assertIn(field, SOP1._fields)
+    self.assertNotIn('simm16', SOP1._fields)
+    self.assertEqual(SOP1._fields['ssrc0'].hi, 7)
+    self.assertEqual(SOP1._fields['ssrc0'].lo, 0)
+    assert SOP1._encoding is not None
+    self.assertEqual(SOP1._encoding[0].hi, 31)
+    self.assertEqual(SOP1._encoding[1], 0b101111101)
+
+  def test_vop3sd_fields(self):
+    """VOP3SD should have all fields including src0/src1/src2 from page continuation."""
+    for field in ['op', 'vdst', 'sdst', 'src0', 'src1', 'src2']:
+      self.assertIn(field, VOP3SD._fields)
+    self.assertEqual(VOP3SD._fields['src0'].hi, 40)
+    self.assertEqual(VOP3SD._fields['src0'].lo, 32)
+    self.assertEqual(VOP3SD._size(), 8)
+
+  def test_flat_has_vdst(self):
+    """FLAT should have vdst field."""
+    self.assertIn('vdst', FLAT._fields)
+    self.assertEqual(FLAT._fields['vdst'].hi, 63)
+    self.assertEqual(FLAT._fields['vdst'].lo, 56)
+
+  def test_encoding_bits(self):
+    """Verify encoding bits are correct for major formats."""
+    tests = [
+      (SOP2, 31, 30, 0b10),
+      (SOPK, 31, 28, 0b1011),
+      (SOPP, 31, 23, 0b101111111),
+      (VOP1, 31, 25, 0b0111111),
+      (VOP2, 31, 31, 0b0),
+      (VOPC, 31, 25, 0b0111110),
+      (FLAT, 31, 26, 0b110111),
+    ]
+    for cls, hi, lo, val in tests:
+      assert cls._encoding is not None
+      self.assertEqual(cls._encoding[0].hi, hi, f"{cls.__name__} encoding hi")
+      self.assertEqual(cls._encoding[0].lo, lo, f"{cls.__name__} encoding lo")
+      self.assertEqual(cls._encoding[1], val, f"{cls.__name__} encoding val")
+
+  def test_opcode_enums_exist(self):
+    """Verify opcode enums are generated with expected counts."""
+    self.assertGreater(len(SOP1Op), 50)
+    self.assertGreater(len(SOP2Op), 50)
+    self.assertGreater(len(VOP1Op), 50)
+    self.assertGreater(len(VOP3Op), 200)
+
+  def test_vopd_no_duplicate_fields(self):
+    """VOPD should not have duplicate fields and should not include DPP16 fields."""
+    field_names = list(VOPD._fields.keys())
+    self.assertEqual(len(field_names), len(set(field_names)))
+    for field in ['srcx0', 'srcy0', 'opx', 'opy']:
+      self.assertIn(field, VOPD._fields)
+    for field in ['dpp_ctrl', 'bank_mask', 'row_mask']:
+      self.assertNotIn(field, VOPD._fields)
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_rdna3_asm.py

@@ -0,0 +1,94 @@
+#!/usr/bin/env python3
+import unittest, subprocess
+from extra.assembly.rdna3.autogen import *
+
+def llvm_assemble(asm: str) -> bytes:
+  """Assemble using llvm-mc and return bytes."""
+  result = subprocess.run(
+    ["llvm-mc", "-triple=amdgcn", "-mcpu=gfx1100", "-show-encoding"],
+    input=asm, capture_output=True, text=True
+  )
+  out = b''
+  for line in result.stdout.split('\n'):
+    if 'encoding:' in line:
+      enc = line.split('encoding:')[1].strip()
+      enc = enc.strip('[]').replace('0x', '').replace(',', '')
+      out += bytes.fromhex(enc)
+  if not out: raise ValueError(f"no encoding found: {result.stdout} {result.stderr}")
+  return out
+
+class TestRDNA3Asm(unittest.TestCase):
+  def test_full_program(self):
+    """Test the full program from rdna3fun.py matches llvm-mc output."""
+    program = [
+      v_bfe_u32(v[1], v[0], 10, 10),
+      s_load_b128(s[4:7], s[0:1], NULL),
+      v_and_b32_e32(v[0], 0x3FF, v[0]),
+      s_mulk_i32(s[3], 0x87),
+      v_mad_u64_u32(v[1:2], NULL, s[2], 3, v[1:2]),
+      v_mul_u32_u24_e32(v[0], 45, v[0]),
+      v_ashrrev_i32_e32(v[2], 31, v[1]),
+      v_add3_u32(v[0], v[0], s[3], v[1]),
+      v_lshlrev_b64(v[2:3], 2, v[1:2]),
+      v_ashrrev_i32_e32(v[1], 31, v[0]),
+      v_lshlrev_b64(v[0:1], 2, v[0:1]),
+      s_waitcnt(0xfc07),  # lgkmcnt(0)
+      v_add_co_u32(v[2], VCC_LO, s[6], v[2]),
+      v_add_co_ci_u32_e32(v[3], s[7], v[3]),
+      v_add_co_u32(v[0], VCC_LO, s[4], v[0]),
+      global_load_b32(vdst=v[2], addr=v[2], saddr=OFF),
+      v_add_co_ci_u32_e32(v[1], s[5], v[1]),
+      s_waitcnt(0x03f7),  # vmcnt(0)
+      global_store_b32(addr=v[0], data=v[2], saddr=OFF),
+      s_endpgm(),
+    ]
+
+    asm = """
+v_bfe_u32 v1, v0, 10, 10
+s_load_b128 s[4:7], s[0:1], null
+v_and_b32_e32 v0, 0x3FF, v0
+s_mulk_i32 s3, 0x87
+v_mad_u64_u32 v[1:2], null, s2, 3, v[1:2]
+v_mul_u32_u24_e32 v0, 45, v0
+v_ashrrev_i32_e32 v2, 31, v1
+v_add3_u32 v0, v0, s3, v1
+v_lshlrev_b64 v[2:3], 2, v[1:2]
+v_ashrrev_i32_e32 v1, 31, v0
+v_lshlrev_b64 v[0:1], 2, v[0:1]
+s_waitcnt lgkmcnt(0)
+v_add_co_u32 v2, vcc_lo, s6, v2
+v_add_co_ci_u32_e32 v3, vcc_lo, s7, v3, vcc_lo
+v_add_co_u32 v0, vcc_lo, s4, v0
+global_load_b32 v2, v[2:3], off
+v_add_co_ci_u32_e32 v1, vcc_lo, s5, v1, vcc_lo
+s_waitcnt vmcnt(0)
+global_store_b32 v[0:1], v2, off
+s_endpgm
+"""
+    expected = llvm_assemble(asm)
+    for inst,rt in zip(program, asm.strip().split("\n")): print(f"{inst.disasm():50s} {rt}")
+    actual = b''.join(inst.to_bytes() for inst in program)
+    self.assertEqual(actual, expected)
+
+  def test_sop2_s_add_u32(self):
+    inst = SOP2(SOP2Op.S_ADD_U32, s[3], s[0], s[1])
+    expected = llvm_assemble("s_add_u32 s3, s0, s1")
+    self.assertEqual(inst.to_bytes(), expected)
+
+  def test_vop2_v_and_b32_inline_const(self):
+    inst = v_and_b32_e32(v[0], 10, v[0])
+    expected = llvm_assemble("v_and_b32_e32 v0, 10, v0")
+    self.assertEqual(inst.to_bytes(), expected)
+
+  def test_sopp_s_endpgm(self):
+    inst = s_endpgm()
+    expected = llvm_assemble("s_endpgm")
+    self.assertEqual(inst.to_bytes(), expected)
+
+  def test_sop1_s_mov_b32(self):
+    inst = s_mov_b32(s[0], s[1])
+    expected = llvm_assemble("s_mov_b32 s0, s1")
+    self.assertEqual(inst.to_bytes(), expected)
+
+if __name__ == "__main__":
+  unittest.main()

extra/assembly/rdna3/test/test_roundtrip.py

@@ -0,0 +1,234 @@
+#!/usr/bin/env python3
+"""Roundtrip tests: generate tinygrad kernels, decode instructions, re-encode, verify match."""
+import unittest, io, sys, re
+from extra.assembly.rdna3.autogen import *
+from extra.assembly.rdna3.lib import Inst
+from extra.assembly.rdna3.asm import asm
+
+# Instruction format detection based on encoding bits
+def detect_format(data: bytes) -> type[Inst] | None:
+  """Detect instruction format from machine code bytes."""
+  if len(data) < 4: return None
+  word = int.from_bytes(data[:4], 'little')
+  enc_9bit = (word >> 23) & 0x1FF  # 9-bit encoding for SOP1/SOPC/SOPP
+  enc_8bit = (word >> 24) & 0xFF
+
+  # Check 9-bit encodings first (most specific)
+  if enc_9bit == 0x17D: return SOP1  # bits 31:23 = 101111101
+  if enc_9bit == 0x17E: return SOPC  # bits 31:23 = 101111110
+  if enc_9bit == 0x17F: return SOPP  # bits 31:23 = 101111111
+  # SOPK: bits 31:28 = 1011, bits 27:23 = opcode (check after SOP1/SOPC/SOPP)
+  if enc_8bit in range(0xB0, 0xC0): return SOPK
+  # SOP2: bits 31:23 in range 0x100-0x17C (0x80-0xBE in bits 31:24, but not SOPK)
+  if 0x80 <= enc_8bit <= 0x9F: return SOP2
+  # VOP1: bits 31:25 = 0111111 (0x3F)
+  if (word >> 25) == 0x3F: return VOP1
+  # VOPC: bits 31:25 = 0111110 (0x3E)
+  if (word >> 25) == 0x3E: return VOPC
+  # VOP2: bits 31:30 = 00
+  if (word >> 30) == 0: return VOP2
+
+  # Check 64-bit formats
+  if len(data) >= 8:
+    if enc_8bit in (0xD4, 0xD5, 0xD7): return VOP3
+    if enc_8bit == 0xD6: return VOP3SD
+    if enc_8bit == 0xCC: return VOP3P
+    if enc_8bit == 0xCD: return VINTERP
+    if enc_8bit in (0xC8, 0xC9): return VOPD
+    if enc_8bit == 0xF4: return SMEM
+    if enc_8bit == 0xD8: return DS
+    if enc_8bit in (0xDC, 0xDD, 0xDE, 0xDF): return FLAT
+    if enc_8bit in (0xE0, 0xE1, 0xE2, 0xE3): return MUBUF
+    if enc_8bit in (0xE8, 0xE9, 0xEA, 0xEB): return MTBUF
+
+  return None
+
+def disassemble_lib(lib: bytes, compiler) -> list[tuple[str, bytes]]:
+  """Disassemble ELF binary and return list of (instruction_text, machine_code_bytes)."""
+  old_stdout = sys.stdout
+  sys.stdout = io.StringIO()
+  compiler.disassemble(lib)
+  output = sys.stdout.getvalue()
+  sys.stdout = old_stdout
+
+  results = []
+  for line in output.splitlines():
+    if '//' not in line: continue
+    instr = line.split('//')[0].strip()
+    if not instr: continue
+    comment = line.split('//')[1].strip()
+    if ':' not in comment: continue
+    hex_str = comment.split(':')[1].strip().split()[0]
+    try:
+      machine_bytes = bytes.fromhex(hex_str)[::-1]  # big-endian to little-endian
+      results.append((instr, machine_bytes))
+    except ValueError:
+      continue
+  return results
+
+def compile_asm(instr: str, compiler=None) -> bytes | None:
+  """Compile a single instruction with llvm-mc and return the machine code bytes."""
+  import subprocess
+  try:
+    result = subprocess.run(
+      ['llvm-mc', '-triple=amdgcn', '-mcpu=gfx1100', '-mattr=+real-true16,+wavefrontsize32', '-show-encoding'],
+      input=f".text\n{instr}\n", capture_output=True, text=True)
+    if result.returncode != 0: return None
+    # Parse encoding: [0x01,0x39,0x0a,0x7e]
+    for line in result.stdout.split('\n'):
+      if 'encoding:' in line:
+        enc = line.split('encoding:')[1].strip()
+        if enc.startswith('[') and enc.endswith(']'):
+          hex_vals = enc[1:-1].replace('0x', '').replace(',', '').replace(' ', '')
+          return bytes.fromhex(hex_vals)
+  except Exception:
+    pass
+  return None
+
+class TestTinygradKernelRoundtrip(unittest.TestCase):
+  """Test roundtrip on real tinygrad-generated kernels using get_kernels_from_tinygrad pattern."""
+
+  def _test_kernel_roundtrip(self, op_fn):
+    """Generate kernel from op_fn, test:
+    1. decode -> reencode matches original bytes
+    2. asm(disasm()) matches LLVM output
+    3. our disasm() matches LLVM's disassembly string exactly
+    """
+    from extra.assembly.rdna3.test.test_compare_emulators import get_kernels_from_tinygrad
+    from tinygrad.runtime.support.compiler_amd import HIPCompiler
+
+    kernels, _, _ = get_kernels_from_tinygrad(op_fn)
+    compiler = HIPCompiler('gfx1100')
+
+    decode_passed, decode_failed, decode_skipped = 0, 0, 0
+    asm_passed, asm_failed, asm_skipped = 0, 0, 0
+    disasm_passed, disasm_failed, disasm_skipped = 0, 0, 0
+    decode_failures, asm_failures, disasm_failures = [], [], []
+
+    for ki, kernel in enumerate(kernels):
+      offset = 0
+      while offset < len(kernel.code):
+        remaining = kernel.code[offset:]
+        fmt = detect_format(remaining)
+        if fmt is None:
+          decode_skipped += 1
+          asm_skipped += 1
+          disasm_skipped += 1
+          offset += 4
+          continue
+
+        size = fmt._size()
+        if len(remaining) < size:
+          break
+
+        orig_bytes = remaining[:size]
+
+        # Test 1: decode -> reencode roundtrip
+        try:
+          decoded = fmt.from_bytes(orig_bytes)
+          reencoded = decoded.to_bytes()
+          if reencoded[:size] == orig_bytes:
+            decode_passed += 1
+          else:
+            decode_failed += 1
+            decode_failures.append(f"K{ki}@{offset}: {decoded.disasm()}: orig={orig_bytes.hex()} reenc={reencoded[:size].hex()}")
+
+          our_disasm = decoded.disasm()
+
+          # Test 2: asm(disasm()) matches LLVM output
+          try:
+            our_bytes = asm(our_disasm).to_bytes()
+            llvm_bytes = compile_asm(our_disasm, compiler)
+            if llvm_bytes is None:
+              asm_skipped += 1
+            elif our_bytes[:len(llvm_bytes)] == llvm_bytes:
+              asm_passed += 1
+            else:
+              asm_failed += 1
+              asm_failures.append(f"K{ki}@{offset}: '{our_disasm}': ours={our_bytes[:len(llvm_bytes)].hex()} llvm={llvm_bytes.hex()}")
+          except Exception:
+            asm_skipped += 1
+
+          # Test 3: our disasm() matches LLVM's disassembly string exactly
+          # Skip if instruction uses op_XX (unknown opcode) or looks malformed (many raw field values)
+          if our_disasm.startswith('op_') or re.search(r', \d+, \d+, \d+,', our_disasm):
+            disasm_skipped += 1
+          else:
+            try:
+              # Get LLVM's disassembly of our instruction
+              src = f".text\n.globl test\n.p2align 8\n.type test,@function\ntest:\n  {our_disasm}\n"
+              lib = compiler.compile(src)
+              llvm_instrs = disassemble_lib(lib, compiler)
+              if llvm_instrs:
+                llvm_disasm = llvm_instrs[0][0]
+                if our_disasm == llvm_disasm:
+                  disasm_passed += 1
+                else:
+                  disasm_failed += 1
+                  disasm_failures.append(f"K{ki}@{offset}: ours='{our_disasm}' llvm='{llvm_disasm}'")
+              else:
+                disasm_skipped += 1
+            except Exception:
+              disasm_skipped += 1
+
+        except Exception:
+          decode_skipped += 1
+          asm_skipped += 1
+          disasm_skipped += 1
+
+        offset += size
+
+    print(f"decode roundtrip: {decode_passed} passed, {decode_failed} failed, {decode_skipped} skipped")
+    print(f"asm vs llvm: {asm_passed} passed, {asm_failed} failed, {asm_skipped} skipped")
+    print(f"disasm vs llvm: {disasm_passed} passed, {disasm_failed} failed, {disasm_skipped} skipped")
+    self.assertEqual(decode_failed, 0, f"Decode failures:\n" + "\n".join(decode_failures[:20]))
+    self.assertEqual(asm_failed, 0, f"Asm failures:\n" + "\n".join(asm_failures[:20]))
+    self.assertEqual(disasm_failed, 0, f"Disasm failures:\n" + "\n".join(disasm_failures[:20]))
+
+  # Basic unary ops
+  def test_neg(self): self._test_kernel_roundtrip(lambda T: -T([1.0, -2.0, 3.0, -4.0]))
+  def test_relu(self): self._test_kernel_roundtrip(lambda T: T([-1.0, 0.0, 1.0, 2.0]).relu())
+  def test_exp(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).exp())
+  def test_log(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 3.0]).log())
+  def test_sin(self): self._test_kernel_roundtrip(lambda T: T([0.0, 1.0, 2.0]).sin())
+  def test_sqrt(self): self._test_kernel_roundtrip(lambda T: T([1.0, 4.0, 9.0]).sqrt())
+  def test_recip(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0, 4.0]).reciprocal())
+
+  # Binary ops
+  def test_add(self): self._test_kernel_roundtrip(lambda T: T([1.0, 2.0]) + T([3.0, 4.0]))
+  def test_sub(self): self._test_kernel_roundtrip(lambda T: T([5.0, 6.0]) - T([1.0, 2.0]))
+  def test_mul(self): self._test_kernel_roundtrip(lambda T: T([2.0, 3.0]) * T([4.0, 5.0]))
+  def test_div(self): self._test_kernel_roundtrip(lambda T: T([10.0, 20.0]) / T([2.0, 4.0]))
+  def test_max_binary(self): self._test_kernel_roundtrip(lambda T: T([1.0, 5.0]).maximum(T([3.0, 2.0])))
+
+  # Reductions
+  def test_sum_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).sum())
+  def test_max_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(64).max())
+  def test_mean_reduce(self): self._test_kernel_roundtrip(lambda T: T.empty(32).mean())
+
+  # Matmul
+  def test_gemm_4x4(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4) @ T.empty(4, 4))
+  def test_gemv(self): self._test_kernel_roundtrip(lambda T: T.empty(1, 16) @ T.empty(16, 16))
+
+  # Complex ops
+  def test_softmax(self): self._test_kernel_roundtrip(lambda T: T.empty(16).softmax())
+  def test_layernorm(self): self._test_kernel_roundtrip(lambda T: T.empty(8, 8).layernorm())
+
+  # Memory patterns
+  def test_contiguous(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 4).permute(1, 0).contiguous())
+  def test_reshape(self): self._test_kernel_roundtrip(lambda T: (T.empty(16) + 1).reshape(4, 4).contiguous())
+  def test_expand(self): self._test_kernel_roundtrip(lambda T: T.empty(4, 1).expand(4, 4).contiguous())
+
+  # Cast ops
+  def test_cast_int(self): self._test_kernel_roundtrip(lambda T: T.empty(16).int().float())
+  def test_cast_half(self): self._test_kernel_roundtrip(lambda T: T.empty(16).half().float())
+
+  # Comparison ops
+  def test_cmp_lt(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) < T.empty(64)).where(T.empty(64), T.empty(64)))
+  def test_where(self): self._test_kernel_roundtrip(lambda T: (T.empty(64) > 0).where(T.empty(64), T.empty(64)))
+
+  # Fused ops
+  def test_fma(self): self._test_kernel_roundtrip(lambda T: (T([1.0, 2.0]) * T([3.0, 4.0]) + T([5.0, 6.0])))
+
+if __name__ == "__main__":
+  unittest.main()

extra/augment.py

@@ -1,42 +0,0 @@
-import numpy as np
-from PIL import Image
-from pathlib import Path
-import sys
-cwd = Path.cwd()
-sys.path.append(cwd.as_posix())
-sys.path.append((cwd / 'test').as_posix())
-from extra.datasets import fetch_mnist
-from tqdm import trange
-
-def augment_img(X, rotate=10, px=3):
-  Xaug = np.zeros_like(X)
-  for i in trange(len(X)):
-    im = Image.fromarray(X[i])
-    im = im.rotate(np.random.randint(-rotate,rotate), resample=Image.BICUBIC)
-    w, h = X.shape[1:]
-    #upper left, lower left, lower right, upper right
-    quad = np.random.randint(-px,px,size=(8)) + np.array([0,0,0,h,w,h,w,0])
-    im = im.transform((w, h), Image.QUAD, quad, resample=Image.BICUBIC)
-    Xaug[i] = im
-  return Xaug
-
-if __name__ == "__main__":
-  import matplotlib.pyplot as plt
-  X_train, Y_train, X_test, Y_test = fetch_mnist()
-  X_train = X_train.reshape(-1, 28, 28).astype(np.uint8)
-  X_test = X_test.reshape(-1, 28, 28).astype(np.uint8)
-  X = np.vstack([X_train[:1]]*10+[X_train[1:2]]*10)
-  fig, a = plt.subplots(2,len(X))
-  Xaug = augment_img(X)
-  for i in range(len(X)):
-    a[0][i].imshow(X[i], cmap='gray')
-    a[1][i].imshow(Xaug[i],cmap='gray')
-    a[0][i].axis('off')
-    a[1][i].axis('off')
-  plt.show()
-
-  #create some nice gifs for doc?!
-  for i in range(10):
-    im = Image.fromarray(X_train[7353+i])
-    im_aug = [Image.fromarray(x) for x in augment_img(np.array([X_train[7353+i]]*100))]
-    im.save(f"aug{i}.gif", save_all=True, append_images=im_aug, duration=100, loop=0)

extra/backends/clang_graph.py

@@ -1,39 +0,0 @@
-from typing import List, Dict, cast
-import ctypes
-from tinygrad.helpers import dedup, cpu_time_execution, DEBUG
-from tinygrad.engine.jit import GraphRunner, GraphException
-from tinygrad.device import Buffer, Device
-from tinygrad.engine.realize import ExecItem, CompiledRunner
-from tinygrad.uop.ops import Variable
-from tinygrad.runtime.ops_cpu import ClangProgram
-from tinygrad.renderer.cstyle import ClangRenderer
-render_dtype = ClangRenderer().render_dtype
-
-class ClangGraph(GraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
-    super().__init__(jit_cache, input_rawbuffers, var_vals)
-    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException
-
-    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
-    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v}" for v in var_vals])
-    code = ["void batched("+','.join(args)+") {"]
-    for ji in jit_cache:
-      args = []
-      for buf in ji.bufs:
-        assert buf is not None
-        if buf in input_rawbuffers:
-          args.append(f"arg{input_rawbuffers.index(buf)}")
-        else:
-          args.append(f"({render_dtype(buf.dtype)}*)0x{ctypes.addressof(buf._buf):X}")
-      args += [x.expr for x in cast(CompiledRunner, ji.prg).p.vars]
-      code.append(f"  {cast(CompiledRunner, ji.prg).p.function_name}({','.join(args)});")
-    code.append("}")
-    if DEBUG >= 4: print("\n".join(code))
-    compiler = Device["CPU"].compiler
-    assert compiler is not None
-    self._prg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers
-
-  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
-    return cpu_time_execution(
-    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0])]), enable=wait)

extra/backends/graph_hip.py

@@ -1,27 +0,0 @@
-import ctypes
-from typing import Tuple
-import tinygrad.runtime.autogen.hip as hip
-from tinygrad.helpers import init_c_var, time_execution_cuda_style
-from tinygrad.runtime.ops_hip import check, hip_set_device
-from tinygrad.runtime.graph.cuda import CUDAGraph
-
-# TODO: this is only used in graph
-def hip_time_execution(cb, enable=False): return time_execution_cuda_style(cb, hip.hipEvent_t, hip.hipEventCreate, hip.hipEventRecord, hip.hipEventSynchronize, hip.hipEventDestroy, hip.hipEventElapsedTime, enable=enable)  # noqa: E501
-
-class HIPGraph(CUDAGraph):
-  def __del__(self):
-    if hasattr(self, 'graph'): check(hip.hipGraphDestroy(self.graph))
-    if hasattr(self, 'instance'): check(hip.hipGraphExecDestroy(self.instance))
-  def set_device(self): hip_set_device(self.dev)
-  def encode_args_info(self): return (hip.hipDeviceptr_t, (1,2,3))
-  def graph_create(self): return init_c_var(hip.hipGraph_t(), lambda x: check(hip.hipGraphCreate(ctypes.byref(x), 0)))
-  def graph_instantiate(self, graph):
-    return init_c_var(hip.hipGraphExec_t(), lambda x: check(hip.hipGraphInstantiate(ctypes.byref(x), graph, None, None, 0)))
-  def graph_add_kernel_node(self, graph, c_deps, c_params):
-    return init_c_var(hip.hipGraphNode_t(), lambda x: check(hip.hipGraphAddKernelNode(ctypes.byref(x), graph, c_deps, ctypes.sizeof(c_deps)//8 if c_deps else 0, ctypes.byref(c_params))))  # noqa: E501
-  def graph_launch(self, *args, wait=False): return hip_time_execution(lambda: check(hip.hipGraphLaunch(*args)), enable=wait)
-  def graph_exec_kernel_node_set_params(self, *args): return check(hip.hipGraphExecKernelNodeSetParams(*args))
-  def build_kernel_node_params(self, prg, global_size, local_size, c_config):
-    return hip.hipKernelNodeParams(hip.dim3(*local_size), c_config, ctypes.cast(prg.clprg.prg, ctypes.c_void_p), hip.dim3(*global_size), None, 0)
-  def set_kernel_node_launch_dims(self, node, global_size: Tuple[int, int, int], local_size: Tuple[int, int, int]):
-    node.blockDim.x, node.blockDim.y, node.blockDim.z, node.gridDim.x, node.gridDim.y, node.gridDim.z = *local_size, *global_size

extra/backends/hsa_driver.py

@@ -1,143 +0,0 @@
-import ctypes, collections
-import tinygrad.runtime.autogen.hsa as hsa
-from tinygrad.helpers import init_c_var
-
-def check(status):
-  if status != 0:
-    hsa.hsa_status_string(status, ctypes.byref(status_str := ctypes.POINTER(ctypes.c_char)()))
-    raise RuntimeError(f"HSA Error {status}: {ctypes.string_at(status_str).decode()}")
-
-# Precalulated AQL info
-AQL_PACKET_SIZE = ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t)
-EMPTY_SIGNAL = hsa.hsa_signal_t()
-
-DISPATCH_KERNEL_SETUP = 3 << hsa.HSA_KERNEL_DISPATCH_PACKET_SETUP_DIMENSIONS
-DISPATCH_KERNEL_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
-DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
-DISPATCH_KERNEL_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
-DISPATCH_KERNEL_HEADER |= hsa.HSA_PACKET_TYPE_KERNEL_DISPATCH << hsa.HSA_PACKET_HEADER_TYPE
-
-BARRIER_HEADER  = 1 << hsa.HSA_PACKET_HEADER_BARRIER
-BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCACQUIRE_FENCE_SCOPE
-BARRIER_HEADER |= hsa.HSA_FENCE_SCOPE_SYSTEM << hsa.HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE
-BARRIER_HEADER |= hsa.HSA_PACKET_TYPE_BARRIER_AND << hsa.HSA_PACKET_HEADER_TYPE
-
-class AQLQueue:
-  def __init__(self, device, sz=-1):
-    self.device = device
-
-    check(hsa.hsa_agent_get_info(self.device.agent, hsa.HSA_AGENT_INFO_QUEUE_MAX_SIZE, ctypes.byref(max_queue_size := ctypes.c_uint32())))
-    queue_size = min(max_queue_size.value, sz) if sz != -1 else max_queue_size.value
-
-    null_func = ctypes.CFUNCTYPE(None, hsa.hsa_status_t, ctypes.POINTER(hsa.struct_hsa_queue_s), ctypes.c_void_p)()
-    self.hw_queue = init_c_var(ctypes.POINTER(hsa.hsa_queue_t)(), lambda x: check(
-      hsa.hsa_queue_create(self.device.agent, queue_size, hsa.HSA_QUEUE_TYPE_SINGLE, null_func, None, (1<<32)-1, (1<<32)-1, ctypes.byref(x))))
-
-    self.next_doorbell_index = 0
-    self.queue_base = self.hw_queue.contents.base_address
-    self.queue_size = self.hw_queue.contents.size * AQL_PACKET_SIZE # in bytes
-    self.write_addr = self.queue_base
-    self.write_addr_end = self.queue_base + self.queue_size - 1 # precalc saves some time
-    self.available_packet_slots = self.hw_queue.contents.size
-
-    check(hsa.hsa_amd_queue_set_priority(self.hw_queue, hsa.HSA_AMD_QUEUE_PRIORITY_HIGH))
-    check(hsa.hsa_amd_profiling_set_profiler_enabled(self.hw_queue, 1))
-
-  def __del__(self):
-    if hasattr(self, 'hw_queue'): check(hsa.hsa_queue_destroy(self.hw_queue))
-
-  def submit_kernel(self, prg, global_size, local_size, kernargs, completion_signal=None):
-    if self.available_packet_slots == 0: self._wait_queue()
-
-    packet = hsa.hsa_kernel_dispatch_packet_t.from_address(self.write_addr)
-    packet.workgroup_size_x = local_size[0]
-    packet.workgroup_size_y = local_size[1]
-    packet.workgroup_size_z = local_size[2]
-    packet.reserved0 = 0
-    packet.grid_size_x = global_size[0] * local_size[0]
-    packet.grid_size_y = global_size[1] * local_size[1]
-    packet.grid_size_z = global_size[2] * local_size[2]
-    packet.private_segment_size = prg.private_segment_size
-    packet.group_segment_size = prg.group_segment_size
-    packet.kernel_object = prg.handle
-    packet.kernarg_address = kernargs
-    packet.reserved2 = 0
-    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
-    packet.setup = DISPATCH_KERNEL_SETUP
-    packet.header = DISPATCH_KERNEL_HEADER
-    self._submit_packet()
-
-  def submit_barrier(self, wait_signals=None, completion_signal=None):
-    assert wait_signals is None or len(wait_signals) <= 5
-    if self.available_packet_slots == 0: self._wait_queue()
-
-    packet = hsa.hsa_barrier_and_packet_t.from_address(self.write_addr)
-    packet.reserved0 = 0
-    packet.reserved1 = 0
-    for i in range(5):
-      packet.dep_signal[i] = wait_signals[i] if wait_signals and len(wait_signals) > i else EMPTY_SIGNAL
-    packet.reserved2 = 0
-    packet.completion_signal = completion_signal if completion_signal else EMPTY_SIGNAL
-    packet.header = BARRIER_HEADER
-    self._submit_packet()
-
-  def blit_packets(self, packet_addr, packet_cnt):
-    if self.available_packet_slots < packet_cnt: self._wait_queue(packet_cnt)
-
-    tail_blit_packets = min((self.queue_base + self.queue_size - self.write_addr) // AQL_PACKET_SIZE, packet_cnt)
-    rem_packet_cnt = packet_cnt - tail_blit_packets
-    ctypes.memmove(self.write_addr, packet_addr, AQL_PACKET_SIZE * tail_blit_packets)
-    if rem_packet_cnt > 0: ctypes.memmove(self.queue_base, packet_addr + AQL_PACKET_SIZE * tail_blit_packets, AQL_PACKET_SIZE * rem_packet_cnt)
-
-    self._submit_packet(packet_cnt)
-
-  def wait(self):
-    self.submit_barrier([], finish_signal := self.device.alloc_signal(reusable=True))
-    hsa.hsa_signal_wait_scacquire(finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-    self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE
-
-  def _wait_queue(self, need_packets=1):
-    while self.available_packet_slots < need_packets:
-      rindex = hsa.hsa_queue_load_read_index_relaxed(self.hw_queue)
-      self.available_packet_slots = self.queue_size // AQL_PACKET_SIZE - (self.next_doorbell_index - rindex)
-
-  def _submit_packet(self, cnt=1):
-    self.available_packet_slots -= cnt
-    self.next_doorbell_index += cnt
-    hsa.hsa_queue_store_write_index_relaxed(self.hw_queue, self.next_doorbell_index)
-    hsa.hsa_signal_store_screlease(self.hw_queue.contents.doorbell_signal, self.next_doorbell_index-1)
-
-    self.write_addr += AQL_PACKET_SIZE * cnt
-    if self.write_addr > self.write_addr_end:
-      self.write_addr = self.queue_base + (self.write_addr - self.queue_base) % self.queue_size
-
-def scan_agents():
-  agents = collections.defaultdict(list)
-
-  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_agent_t, ctypes.c_void_p)
-  def __scan_agents(agent, data):
-    status = hsa.hsa_agent_get_info(agent, hsa.HSA_AGENT_INFO_DEVICE, ctypes.byref(device_type := hsa.hsa_device_type_t()))
-    if status == 0: agents[device_type.value].append(agent)
-    return hsa.HSA_STATUS_SUCCESS
-
-  hsa.hsa_iterate_agents(__scan_agents, None)
-  return agents
-
-def find_memory_pool(agent, segtyp=-1, location=-1):
-  @ctypes.CFUNCTYPE(hsa.hsa_status_t, hsa.hsa_amd_memory_pool_t, ctypes.c_void_p)
-  def __filter_amd_memory_pools(mem_pool, data):
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SEGMENT, ctypes.byref(segment := hsa.hsa_amd_segment_t())))
-    if segtyp >= 0 and segment.value != segtyp: return hsa.HSA_STATUS_SUCCESS
-
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_LOCATION, ctypes.byref(loc:=hsa.hsa_amd_memory_pool_location_t())))
-    if location >= 0 and loc.value != location: return hsa.HSA_STATUS_SUCCESS
-
-    check(hsa.hsa_amd_memory_pool_get_info(mem_pool, hsa.HSA_AMD_MEMORY_POOL_INFO_SIZE, ctypes.byref(sz := ctypes.c_size_t())))
-    if sz.value == 0: return hsa.HSA_STATUS_SUCCESS
-
-    ret = ctypes.cast(data, ctypes.POINTER(hsa.hsa_amd_memory_pool_t))
-    ret[0] = mem_pool
-    return hsa.HSA_STATUS_INFO_BREAK
-
-  hsa.hsa_amd_agent_iterate_memory_pools(agent, __filter_amd_memory_pools, ctypes.byref(region := hsa.hsa_amd_memory_pool_t()))
-  return region

extra/backends/hsa_graph.py

@@ -1,171 +0,0 @@
-import ctypes, collections, time, itertools
-from typing import List, Any, Dict, cast, Optional, Tuple
-from tinygrad.helpers import init_c_var, round_up
-from tinygrad.device import Buffer, BufferSpec
-from tinygrad.device import Compiled, Device
-from tinygrad.uop.ops import Variable
-from tinygrad.runtime.ops_hsa import HSADevice, PROFILE, Profiler
-from tinygrad.engine.realize import ExecItem, BufferXfer, CompiledRunner
-from tinygrad.engine.jit import MultiGraphRunner, GraphException
-import tinygrad.runtime.autogen.hsa as hsa
-from tinygrad.runtime.support.hsa import check, AQLQueue, AQL_PACKET_SIZE, EMPTY_SIGNAL
-
-def dedup_signals(signals): return [hsa.hsa_signal_t(hndl) for hndl in set([x.handle for x in signals if isinstance(x, hsa.hsa_signal_t)])]
-
-class VirtAQLQueue(AQLQueue):
-  def __init__(self, device, sz):
-    self.device = device
-    self.virt_queue = (hsa.hsa_kernel_dispatch_packet_t * sz)()
-    self.queue_base = self.write_addr = ctypes.addressof(self.virt_queue)
-    self.packets_count = 0
-    self.available_packet_slots = sz
-  def _wait_queue(self, need_packets=1): assert False, f"VirtQueue is too small to handle {self.packets_count+need_packets} packets!"
-  def _submit_packet(self):
-    self.write_addr += AQL_PACKET_SIZE
-    self.packets_count += 1
-    self.available_packet_slots -= 1
-
-class HSAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
-    super().__init__(jit_cache, input_rawbuffers, var_vals)
-
-    # Check all jit items are compatible.
-    compiled_devices = set()
-    for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): compiled_devices.add(ji.prg.dev)
-      elif isinstance(ji.prg, BufferXfer):
-        for x in ji.bufs[0:2]: compiled_devices.add(Device[cast(Buffer, x).device])
-      else: raise GraphException
-    if any(not isinstance(d, HSADevice) for d in compiled_devices): raise GraphException
-
-    self.devices: List[HSADevice] = list(compiled_devices) #type:ignore
-
-    # Allocate kernel args.
-    kernargs_size: Dict[Compiled, int] = collections.defaultdict(int)
-    for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): kernargs_size[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
-    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferSpec()) for dev,sz in kernargs_size.items()}
-
-    # Fill initial arguments.
-    self.ji_kargs_structs: Dict[int, ctypes.Structure] = {}
-    for j,ji in enumerate(self.jit_cache):
-      if not isinstance(ji.prg, CompiledRunner): continue
-      self.ji_kargs_structs[j] = ji.prg._prg.args_struct_t.from_address(kernargs_ptrs[ji.prg.dev])
-      kernargs_ptrs[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
-      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
-      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])
-
-    # Build queues.
-    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
-    self.packets = {}
-    self.transfers = []
-    self.ji_to_transfer: Dict[int, int] = {} # faster to store transfers as list and update using this mapping table.
-    self.signals_to_reset: List[hsa.hsa_signal_t] = []
-    self.signals_to_devices: Dict[ctypes.c_uint64, List[HSADevice]] = {}
-    self.profile_info: Dict[Compiled, List[Tuple[Any, ...]]] = collections.defaultdict(list)
-
-    # Special packet to wait for the world.
-    self.kickoff_signals: Dict[HSADevice, hsa.hsa_signal_t] = {dev:self.alloc_signal(reset_on_start=True) for dev in self.devices}
-    for dev in self.devices: self.virt_aql_queues[dev].submit_barrier([], self.kickoff_signals[dev])
-
-    for j,ji in enumerate(self.jit_cache):
-      if isinstance(ji.prg, CompiledRunner):
-        wait_signals = self.access_resources(ji.bufs, ji.prg.p.outs, new_dependency=j, sync_with_aql_packets=False)
-        for i in range(0, len(wait_signals), 5):
-          self.virt_aql_queues[ji.prg.dev].submit_barrier(wait_signals[i:i+5])
-        self.packets[j] = hsa.hsa_kernel_dispatch_packet_t.from_address(self.virt_aql_queues[ji.prg.dev].write_addr)
-
-        sync_signal = self.alloc_signal(reset_on_start=True) if PROFILE else None
-        self.virt_aql_queues[ji.prg.dev].submit_kernel(ji.prg._prg, *ji.prg.p.launch_dims(var_vals), #type:ignore
-                                                          ctypes.addressof(self.ji_kargs_structs[j]), completion_signal=sync_signal)
-        if PROFILE: self.profile_info[ji.prg.dev].append((sync_signal, ji.prg._prg.name, False))
-      elif isinstance(ji.prg, BufferXfer):
-        dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
-        dest_dev, src_dev = cast(HSADevice, Device[dest.device]), cast(HSADevice, Device[src.device])
-        sync_signal = self.alloc_signal(reset_on_start=True, wait_on=[dest_dev, src_dev])
-
-        wait_signals = self.access_resources([dest, src], write=[0], new_dependency=sync_signal, sync_with_aql_packets=True)
-        self.transfers.append([dest._buf, dest_dev.agent, src._buf, src_dev.agent, dest.nbytes, len(wait_signals),
-                              (hsa.hsa_signal_t*len(wait_signals))(*wait_signals), sync_signal, hsa.HSA_AMD_SDMA_ENGINE_0, True])
-        self.ji_to_transfer[j] = len(self.transfers) - 1
-        if PROFILE: self.profile_info[src_dev].append((sync_signal, f"transfer: HSA:{src_dev.device_id} -> HSA:{dest_dev.device_id}", True))
-
-    # Wait for all active signals to finish the graph
-    wait_signals_to_finish: Dict[HSADevice, List[hsa.hsa_signal_t]] = collections.defaultdict(list)
-    for v in dedup_signals(list(self.w_dependency_map.values()) + list(itertools.chain.from_iterable(self.r_dependency_map.values()))):
-      for dev in self.signals_to_devices[v.handle]:
-        wait_signals_to_finish[dev].append(v)
-
-    self.finish_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
-    for dev in self.devices:
-      wait_signals = wait_signals_to_finish[dev]
-      for i in range(0, max(1, len(wait_signals)), 5):
-        self.virt_aql_queues[dev].submit_barrier(wait_signals[i:i+5], completion_signal=self.finish_signal if i+5>=len(wait_signals) else None)
-
-    # Zero signals to allow graph to start and execute.
-    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
-    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)
-
-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
-    # Wait and restore signals
-    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
-    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, len(self.devices))
-
-    # Update rawbuffers
-    for (j,i),input_idx in self.input_replace.items():
-      if j in self.ji_kargs_structs:
-        self.ji_kargs_structs[j].__setattr__(f'f{i}', input_rawbuffers[input_idx]._buf)
-      else:
-        if i == 0: self.transfers[self.ji_to_transfer[j]][0] = input_rawbuffers[input_idx]._buf # dest
-        elif i == 1: self.transfers[self.ji_to_transfer[j]][2] = input_rawbuffers[input_idx]._buf # src
-
-    # Update var_vals
-    for j in self.jc_idx_with_updatable_var_vals:
-      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
-        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])
-
-    # Update launch dims
-    for j in self.jc_idx_with_updatable_launch_dims:
-      gl, lc = cast(CompiledRunner, self.jit_cache[j].prg).p.launch_dims(var_vals)
-      self.packets[j].workgroup_size_x = lc[0]
-      self.packets[j].workgroup_size_y = lc[1]
-      self.packets[j].workgroup_size_z = lc[2]
-      self.packets[j].grid_size_x = gl[0] * lc[0]
-      self.packets[j].grid_size_y = gl[1] * lc[1]
-      self.packets[j].grid_size_z = gl[2] * lc[2]
-
-    for dev in self.devices:
-      dev.flush_hdp()
-      dev.hw_queue.blit_packets(self.virt_aql_queues[dev].queue_base, self.virt_aql_queues[dev].packets_count)
-
-    for transfer_data in self.transfers:
-      check(hsa.hsa_amd_memory_async_copy_on_engine(*transfer_data))
-
-    et = None
-    if wait:
-      st = time.perf_counter()
-      hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-      et = time.perf_counter() - st
-
-    for profdev,profdata in self.profile_info.items(): Profiler.tracked_signals[profdev] += profdata
-    return et
-
-  def alloc_signal(self, reset_on_start=False, wait_on=None):
-    sync_signal = init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(x))))
-    if reset_on_start: self.signals_to_reset.append(sync_signal)
-    if wait_on is not None: self.signals_to_devices[sync_signal.handle] = wait_on
-    return sync_signal
-
-  def dependency_as_signal(self, dep, sync_with_aql_packets) -> Optional[hsa.hsa_signal_t]:
-    if isinstance(dep, hsa.hsa_signal_t): return dep
-    elif sync_with_aql_packets and isinstance(packet := self.packets.get(dep), hsa.hsa_kernel_dispatch_packet_t):
-      if packet.completion_signal.handle == EMPTY_SIGNAL.handle: packet.completion_signal = self.alloc_signal(reset_on_start=True)
-      return packet.completion_signal
-    return None
-
-  def access_resources(self, rawbufs, write, new_dependency, sync_with_aql_packets=False):
-    rdeps = self._access_resources(rawbufs, write, new_dependency)
-    wait_signals = [self.dependency_as_signal(dep, sync_with_aql_packets=sync_with_aql_packets) for dep in rdeps]
-    if sync_with_aql_packets: wait_signals += [self.kickoff_signals[cast(HSADevice, Device[rawbuf.device])] for rawbuf in rawbufs]
-    return dedup_signals(wait_signals)

extra/backends/ops_hsa.py

@@ -1,275 +0,0 @@
-from __future__ import annotations
-import ctypes, functools, subprocess, io, atexit, collections, json
-from typing import Tuple, TypeVar, List, Dict, Any
-import tinygrad.runtime.autogen.hsa as hsa
-from tinygrad.helpers import DEBUG, init_c_var, from_mv, round_up, to_mv, init_c_struct_t, getenv, PROFILE
-from tinygrad.device import Compiled, Compiler, CompileError, BufferSpec, LRUAllocator
-from tinygrad.renderer.cstyle import HIPRenderer
-from tinygrad.runtime.support.hsa import check, scan_agents, find_memory_pool, AQLQueue
-from tinygrad.runtime.support.hip_comgr import compile_hip
-if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl  # noqa: F401
-
-class HSAProfiler:
-  def __init__(self):
-    self.tracked_signals = collections.defaultdict(list)
-    self.collected_events: List[Tuple[Any, ...]] = []
-    self.copy_timings = hsa.hsa_amd_profiling_async_copy_time_t()
-    self.disp_timings = hsa.hsa_amd_profiling_dispatch_time_t()
-
-  def track(self, signal, device, name, is_copy=False): self.tracked_signals[device].append((signal, name, is_copy))
-  def process(self, device):
-    # Process all tracked signals, should be called before any of tracked signals are reused.
-    for sig,name,is_copy in self.tracked_signals[device]:
-      if is_copy: check(hsa.hsa_amd_profiling_get_async_copy_time(sig, ctypes.byref(timings :=  self.copy_timings)))
-      else: check(hsa.hsa_amd_profiling_get_dispatch_time(device.agent, sig, ctypes.byref(timings := self.disp_timings))) #type:ignore
-      self.collected_events.append((device.device_id, 1 if is_copy else 0, name, timings.start, timings.end))
-    self.tracked_signals.pop(device)
-
-  def save(self, path):
-    mjson = []
-    for i in range(len(HSADevice.devices)):
-      mjson.append({"name": "process_name", "ph": "M", "pid": i, "args": {"name": "HSA"}})
-      mjson.append({"name": "thread_name", "ph": "M", "pid": i, "tid": 0, "args": {"name": "AQL"}})
-      mjson.append({"name": "thread_name", "ph": "M", "pid": i, "tid": 1, "args": {"name": "SDMA"}})
-
-    for dev_id,queue_id,name,st,et in self.collected_events:
-      mjson.append({"name": name, "ph": "B", "pid": dev_id, "tid": queue_id, "ts": st*1e-3})
-      mjson.append({"name": name, "ph": "E", "pid": dev_id, "tid": queue_id, "ts": et*1e-3})
-    with open(path, "w") as f: f.write(json.dumps({"traceEvents": mjson}))
-    print(f"Saved HSA profile to {path}")
-Profiler = HSAProfiler()
-
-class HSACompiler(Compiler):
-  def __init__(self, arch:str):
-    self.arch = arch
-    super().__init__(f"compile_hip_{self.arch}")
-  def compile(self, src:str) -> bytes:
-    try: return compile_hip(src, self.arch)
-    except RuntimeError as e: raise CompileError(e)
-
-class HSAProgram:
-  def __init__(self, device:HSADevice, name:str, lib:bytes):
-    self.device, self.name, self.lib = device, name, lib
-
-    if DEBUG >= 6:
-      asm = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib)
-      print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
-
-    self.exec = init_c_var(hsa.hsa_executable_t(), lambda x: check(hsa.hsa_executable_create_alt(hsa.HSA_PROFILE_FULL, hsa.HSA_DEFAULT_FLOAT_ROUNDING_MODE_DEFAULT, None, ctypes.byref(x)))) # noqa: E501
-    self.code_reader = init_c_var(hsa.hsa_code_object_reader_t(),
-                                  lambda x: check(hsa.hsa_code_object_reader_create_from_memory(lib, len(lib), ctypes.byref(x))))
-    check(hsa.hsa_executable_load_agent_code_object(self.exec, self.device.agent, self.code_reader, None, None))
-    check(hsa.hsa_executable_freeze(self.exec, None))
-
-    self.kernel = init_c_var(hsa.hsa_executable_symbol_t(), lambda x: check(hsa.hsa_executable_get_symbol_by_name(self.exec, (name+".kd").encode("utf-8"), ctypes.byref(self.device.agent), ctypes.byref(x)))) # noqa: E501
-    self.handle = init_c_var(ctypes.c_uint64(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_OBJECT, ctypes.byref(x)))) # noqa: E501
-    self.kernargs_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_KERNARG_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
-    self.group_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_GROUP_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
-    self.private_segment_size = init_c_var(ctypes.c_uint32(), lambda x: check(hsa.hsa_executable_symbol_get_info(self.kernel, hsa.HSA_EXECUTABLE_SYMBOL_INFO_KERNEL_PRIVATE_SEGMENT_SIZE, ctypes.byref(x)))).value # noqa: E501
-
-  def __del__(self):
-    self.device.synchronize()
-    if hasattr(self, 'code_reader'): check(hsa.hsa_code_object_reader_destroy(self.code_reader))
-    if hasattr(self, 'exec'): check(hsa.hsa_executable_destroy(self.exec))
-
-  def __call__(self, *args, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
-    if not hasattr(self, "args_struct_t"):
-      self.args_struct_t = init_c_struct_t(tuple([(f'f{i}', ctypes.c_void_p) for i in range(len(args))] +
-                                                 [(f'v{i}', ctypes.c_int) for i in range(len(vals))]))
-      if ctypes.sizeof(self.args_struct_t) != self.kernargs_segment_size:
-        raise RuntimeError(f"HSAProgram.__call__: incorrect args struct size {ctypes.sizeof(self.args_struct_t)} != {self.kernargs_segment_size}")
-
-    kernargs = None
-    if self.kernargs_segment_size > 0:
-      kernargs = self.device.alloc_kernargs(self.kernargs_segment_size)
-      args_st = self.args_struct_t.from_address(kernargs)
-      for i in range(len(args)): args_st.__setattr__(f'f{i}', args[i])
-      for i in range(len(vals)): args_st.__setattr__(f'v{i}', vals[i])
-      self.device.flush_hdp()
-
-    signal = self.device.alloc_signal(reusable=True) if wait or PROFILE else None
-    self.device.hw_queue.submit_kernel(self, global_size, local_size, kernargs, completion_signal=signal)
-    if PROFILE: Profiler.track(signal, self.device, self.name)
-    if wait:
-      hsa.hsa_signal_wait_scacquire(signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-      check(hsa.hsa_amd_profiling_get_dispatch_time(self.device.agent, signal, ctypes.byref(timings := hsa.hsa_amd_profiling_dispatch_time_t())))
-      return (timings.end - timings.start) * self.device.clocks_to_time
-
-T = TypeVar("T")
-CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
-class HSAAllocator(LRUAllocator):
-  def __init__(self, device:HSADevice):
-    self.device = device
-    super().__init__()
-
-  def _alloc(self, size:int, options:BufferSpec):
-    if options.host:
-      check(hsa.hsa_amd_memory_pool_allocate(HSADevice.cpu_mempool, size, 0, ctypes.byref(mem := ctypes.c_void_p())))
-      check(hsa.hsa_amd_agents_allow_access(2, (hsa.hsa_agent_t*2)(HSADevice.cpu_agent, self.device.agent), None, mem))
-      return mem.value
-    c_agents = (hsa.hsa_agent_t * len(HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU]))(*HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU])
-    check(hsa.hsa_amd_memory_pool_allocate(self.device.gpu_mempool, size, 0, ctypes.byref(buf := ctypes.c_void_p())))
-    check(hsa.hsa_amd_agents_allow_access(len(HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU]), c_agents, None, buf))
-    return buf.value
-
-  def _free(self, opaque:T, options:BufferSpec):
-    HSADevice.synchronize_system()
-    check(hsa.hsa_amd_memory_pool_free(opaque))
-
-  def _copyin(self, dest:T, src: memoryview):
-    # Async copyin sync model uses barriers on the main hw queue, since barriers are guaranteed to execute in order with all other packets.
-    self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
-    mem = self._alloc(src.nbytes, BufferSpec(host=True))
-    ctypes.memmove(mem, from_mv(src), src.nbytes)
-    check(hsa.hsa_amd_memory_async_copy_on_engine(dest, self.device.agent, mem, HSADevice.cpu_agent, src.nbytes, 1, ctypes.byref(sync_signal),
-                                                  copy_signal := self.device.alloc_signal(reusable=True), hsa.HSA_AMD_SDMA_ENGINE_0, True))
-    self.device.hw_queue.submit_barrier([copy_signal])
-    self.device.delayed_free.append(mem)
-    if PROFILE: Profiler.track(copy_signal, self.device, f"copyin: CPU -> HSA:{self.device.device_id}", is_copy=True)
-
-  def copy_from_fd(self, dest, fd, offset, size):
-    self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
-
-    if not hasattr(self, 'hb'):
-      self.hb = [self._alloc(CHUNK_SIZE, BufferSpec(host=True)) for _ in range(2)]
-      self.hb_signals = [self.device.alloc_signal(reusable=False) for _ in range(2)]
-      self.hb_polarity = 0
-      self.sdma = [hsa.HSA_AMD_SDMA_ENGINE_0, hsa.HSA_AMD_SDMA_ENGINE_1]
-      for sig in self.hb_signals: hsa.hsa_signal_store_relaxed(sig, 0)
-
-    fo = io.FileIO(fd, "a+b", closefd=False)
-    fo.seek(offset - (minor_offset:=offset % PAGE_SIZE))
-
-    copies_called = 0
-    copied_in = 0
-    for local_offset in range(0, size+minor_offset, CHUNK_SIZE):
-      local_size = min(round_up(size+minor_offset, PAGE_SIZE)-local_offset, CHUNK_SIZE)
-      copy_size = min(local_size-minor_offset, size-copied_in)
-      if copy_size == 0: break
-
-      hsa.hsa_signal_wait_scacquire(self.hb_signals[self.hb_polarity], hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-      self.device.reusable_signals.append(self.hb_signals[self.hb_polarity]) # it's free now and can be reused
-      self.hb_signals[self.hb_polarity] = self.device.alloc_signal(reusable=False)
-
-      fo.readinto(to_mv(self.hb[self.hb_polarity], local_size))
-      check(hsa.hsa_amd_memory_async_copy_on_engine(dest+copied_in, self.device.agent, self.hb[self.hb_polarity]+minor_offset, HSADevice.cpu_agent,
-                                                    copy_size, 1, ctypes.byref(sync_signal), self.hb_signals[self.hb_polarity],
-                                                    self.sdma[self.hb_polarity], True))
-      copied_in += copy_size
-      self.hb_polarity = (self.hb_polarity + 1) % len(self.hb)
-      minor_offset = 0 # only on the first
-      copies_called += 1
-
-    wait_signals = [self.hb_signals[self.hb_polarity - 1]]
-    if copies_called > 1: wait_signals.append(self.hb_signals[self.hb_polarity])
-    self.device.hw_queue.submit_barrier(wait_signals)
-
-  def _copyout(self, dest:memoryview, src:T):
-    HSADevice.synchronize_system()
-    copy_signal = self.device.alloc_signal(reusable=True)
-    c_agents = (hsa.hsa_agent_t*2)(self.device.agent, HSADevice.cpu_agent)
-    check(hsa.hsa_amd_memory_lock_to_pool(from_mv(dest), dest.nbytes, c_agents, 2, HSADevice.cpu_mempool, 0, ctypes.byref(addr:=ctypes.c_void_p())))
-    check(hsa.hsa_amd_memory_async_copy(addr, HSADevice.cpu_agent, src, self.device.agent, dest.nbytes, 0, None, copy_signal))
-    hsa.hsa_signal_wait_scacquire(copy_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
-    check(hsa.hsa_amd_memory_unlock(from_mv(dest)))
-    if PROFILE: Profiler.track(copy_signal, self.device, f"copyout: HSA:{self.device.device_id} -> CPU", is_copy=True)
-
-  def transfer(self, dest:T, src:T, sz:int, src_dev=None, dest_dev=None):
-    src_dev.hw_queue.submit_barrier([], sync_signal_1 := src_dev.alloc_signal(reusable=True))
-    dest_dev.hw_queue.submit_barrier([], sync_signal_2 := dest_dev.alloc_signal(reusable=True))
-    c_wait_signal = (hsa.hsa_signal_t*2)(sync_signal_1, sync_signal_2)
-    check(hsa.hsa_amd_memory_async_copy_on_engine(dest, dest_dev.agent, src, src_dev.agent, sz, 2, c_wait_signal,
-                                                  copy_signal := dest_dev.alloc_signal(reusable=False), hsa.HSA_AMD_SDMA_ENGINE_0, True))
-    src_dev.hw_queue.submit_barrier([copy_signal])
-    dest_dev.hw_queue.submit_barrier([copy_signal])
-    if PROFILE: Profiler.track(copy_signal, src_dev, f"transfer: HSA:{src_dev.device_id} -> HSA:{dest_dev.device_id}", is_copy=True)
-
-class HSADevice(Compiled):
-  devices: List[HSADevice] = []
-  agents: Dict[int, List[hsa.hsa_agent_t]] = {}
-  cpu_agent: hsa.hsa_agent_t
-  cpu_mempool: hsa.hsa_amd_memory_pool_t
-  def __init__(self, device:str=""):
-    if not HSADevice.agents:
-      check(hsa.hsa_init())
-      atexit.register(hsa_terminate)
-      HSADevice.agents = scan_agents()
-      HSADevice.cpu_agent = HSADevice.agents[hsa.HSA_DEVICE_TYPE_CPU][0]
-      HSADevice.cpu_mempool = find_memory_pool(HSADevice.cpu_agent, segtyp=hsa.HSA_AMD_SEGMENT_GLOBAL, location=hsa.HSA_AMD_MEMORY_POOL_LOCATION_CPU)
-      if PROFILE: check(hsa.hsa_amd_profiling_async_copy_enable(1))
-
-    self.device_id = int(device.split(":")[1]) if ":" in device else 0
-    self.agent = HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU][self.device_id]
-    self.gpu_mempool = find_memory_pool(self.agent, segtyp=hsa.HSA_AMD_SEGMENT_GLOBAL, location=hsa.HSA_AMD_MEMORY_POOL_LOCATION_GPU)
-    self.hw_queue = AQLQueue(self)
-    HSADevice.devices.append(self)
-
-    check(hsa.hsa_agent_get_info(self.agent, hsa.HSA_AGENT_INFO_NAME, ctypes.byref(agent_name_buf := ctypes.create_string_buffer(256))))
-    self.arch = ctypes.string_at(agent_name_buf).decode()
-
-    check(hsa.hsa_system_get_info(hsa.HSA_SYSTEM_INFO_TIMESTAMP_FREQUENCY, ctypes.byref(gpu_freq := ctypes.c_uint64())))
-    self.clocks_to_time: float = 1 / gpu_freq.value
-
-    check(hsa.hsa_agent_get_info(self.agent, hsa.HSA_AMD_AGENT_INFO_HDP_FLUSH, ctypes.byref(hdp_flush := hsa.hsa_amd_hdp_flush_t())))
-    self.hdp_flush = hdp_flush
-
-    self.delayed_free: List[int] = []
-    self.reusable_signals: List[hsa.hsa_signal_t] = []
-
-    from tinygrad.runtime.graph.hsa import HSAGraph
-    super().__init__(device, HSAAllocator(self), HIPRenderer(), HSACompiler(self.arch), functools.partial(HSAProgram, self), HSAGraph)
-
-    # Finish init: preallocate some signals + space for kernargs
-    self.signal_pool = [init_c_var(hsa.hsa_signal_t(), lambda x: check(hsa.hsa_signal_create(1, 0, None, ctypes.byref(x)))) for _ in range(4096)]
-    self._new_kernargs_region(16 << 20) # initial region size is 16mb
-
-  def synchronize(self):
-    self.hw_queue.wait()
-
-    for sig in self.reusable_signals: hsa.hsa_signal_silent_store_relaxed(sig, 1)
-    self.signal_pool.extend(self.reusable_signals)
-    self.reusable_signals.clear()
-
-    for opaque_to_free in self.delayed_free: check(hsa.hsa_amd_memory_pool_free(opaque_to_free))
-    self.delayed_free.clear()
-
-    self.kernarg_next_addr = self.kernarg_start_addr
-    Profiler.process(self)
-
-  @staticmethod
-  def synchronize_system():
-    for d in HSADevice.devices: d.synchronize()
-
-  def alloc_signal(self, reusable=False):
-    if len(self.signal_pool): signal = self.signal_pool.pop()
-    else: check(hsa.hsa_amd_signal_create(1, 0, None, 0, ctypes.byref(signal := hsa.hsa_signal_t())))
-
-    # reusable means a signal could be reused after synchronize for the device it's allocated from is called.
-    if reusable: self.reusable_signals.append(signal)
-    return signal
-
-  def alloc_kernargs(self, sz):
-    if self.kernarg_next_addr + sz >= self.kernarg_start_addr + self.kernarg_pool_sz: self._new_kernargs_region(int(self.kernarg_pool_sz * 2))
-    result = self.kernarg_next_addr
-    self.kernarg_next_addr = round_up(self.kernarg_next_addr + sz, 16)
-    return result
-
-  def _new_kernargs_region(self, sz:int):
-    if hasattr(self, 'kernarg_start_addr'): self.delayed_free.append(self.kernarg_start_addr)
-    self.kernarg_start_addr: int = self.allocator._alloc(sz, BufferSpec())
-    self.kernarg_next_addr = self.kernarg_start_addr
-    self.kernarg_pool_sz: int = sz
-
-  def flush_hdp(self): self.hdp_flush.HDP_MEM_FLUSH_CNTL[0] = 1
-
-def hsa_terminate():
-  # Need to stop/delete aql queue before hsa shut down, this leads to gpu hangs.
-  for dev in HSADevice.devices:
-    Profiler.process(dev)
-    del dev.hw_queue
-
-  # hsa_shut_down cleans up all hsa-related resources.
-  hsa.hsa_shut_down()
-  HSADevice.synchronize = lambda: None #type:ignore
-  HSAProgram.__del__ = lambda _: None #type:ignore
-  if Profiler.collected_events: Profiler.save("/tmp/profile.json")

extra/backends/rdna.py

@@ -1,127 +0,0 @@
-from typing import Dict, Set
-import yaml
-from tinygrad.codegen.uops import UOpGraph, UOps, UOp
-from tinygrad.uop.ops import BinaryOps
-from tinygrad.dtype import dtypes
-
-def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
-  replace: Dict[UOp, UOp] = {}
-  seen: Set[UOp] = set()
-  for u in uops:
-    if u in seen: continue
-    seen.add(u)
-    for o,n in replace.items():
-      if o in u.vin and u is not n:
-        u.vin = tuple(n if x == o else x for x in u.vin)
-    # pointer indexing
-    if u.uop in {UOps.LOAD, UOps.STORE} and u.vin[0].dtype.itemsize > 1:
-      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_at=uops.uops.index(u))
-      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_at=uops.uops.index(u))
-      u.vin = (u.vin[0], ptr) + u.vin[2:]
-  #uops.print()
-
-  args = []
-  ins = []
-
-  v_cnt = 3  # v[0:2] is local_xyz
-  s_cnt = 5  # s[0:1] is the address, s[2:4] is global_xyz
-
-  r: Dict[UOp, str] = {}
-  for u in uops:
-    if u.uop == UOps.SPECIAL:
-      if u.arg.startswith("lidx"):
-        r[u] = f'v{u.src[0].arg}'
-      elif u.arg.startswith("gidx"):
-        r[u] = f's{2+u.src[0].arg}'
-      else:
-        raise NotImplementedError
-    elif u.uop == UOps.CONST:
-      #r[u] = u.arg
-
-      # TODO: sometimes we can use s
-      #r[u] = f"s{s_cnt}"
-      #s_cnt += 1
-      #ins.append(f"s_mov_b32 {r[u]}, {u.arg}")
-
-      r[u] = f"v{v_cnt}"
-      v_cnt += 1
-      ins.append(f"v_mov_b32 {r[u]}, {u.arg}")
-    elif u.uop == UOps.ALU:
-      if u.arg == BinaryOps.ADD:
-        r[u] = f"v{v_cnt}"
-        v_cnt += 1
-        ins.append(f"v_add_f32_e32 {r[u]}, {r[u.vin[0]]}, {r[u.vin[1]]}")
-      elif u.arg == BinaryOps.MUL:
-        r[u] = f"v{v_cnt}"
-        v_cnt += 1
-        if dtypes.is_float(u.dtype):
-          ins.append(f"v_mul_f32_e32 {r[u]}, {r[u.vin[0]]}, {r[u.vin[1]]}")
-        else:
-          ins.append(f"v_mul_u32_u24 {r[u]}, {r[u.vin[0]]}, {r[u.vin[1]]}")
-      else:
-        raise NotImplementedError
-    elif u.uop == UOps.LOAD:
-      r[u] = f"v{v_cnt}"
-      v_cnt += 1
-      ins.append(f"global_load_b32 {r[u]}, {r[u.vin[1]]}, {r[u.vin[0]]}")
-      ins.append("s_waitcnt vmcnt(0)")
-    elif u.uop == UOps.STORE:
-      ins.append(f"global_store_b32 {r[u.vin[1]]}, {r[u.vin[2]]}, {r[u.vin[0]]}")
-    elif u.uop == UOps.DEFINE_GLOBAL:
-      i = u.arg[0]
-      args.append({'.address_space': 'global', '.name': f'buf_{i}', '.offset': i*8, '.size': 8,
-                   '.type_name': u.dtype.name+"*", '.value_kind': 'global_buffer'})
-      s_cnt += s_cnt%2  # skip
-      r[u] = f"s[{s_cnt}:{s_cnt+1}]"
-      s_cnt += 2
-      ins.append(f"s_load_b64 {r[u]}, s[0:1], {i*8}")
-      ins.append("s_waitcnt lgkmcnt(0)")
-    else:
-      raise NotImplementedError(f"can't render {u.uop}")
-
-  # *** boilerplate rendering ***
-
-  metadata = {
-    'amdhsa.kernels': [{'.args': args,
-      '.group_segment_fixed_size': 0, '.kernarg_segment_align': 8, '.kernarg_segment_size': args[-1][".offset"] + args[-1][".size"],
-      '.language': 'OpenCL C', '.language_version': [1, 2], '.max_flat_workgroup_size': 256,
-      '.name': function_name, '.private_segment_fixed_size': 0, '.sgpr_count': s_cnt, '.sgpr_spill_count': 0,
-      '.symbol': f'{function_name}.kd', '.uses_dynamic_stack': False, '.vgpr_count': v_cnt, '.vgpr_spill_count': 0,
-      '.wavefront_size': 32}],
-    'amdhsa.target': 'amdgcn-amd-amdhsa--gfx1100', 'amdhsa.version': [1, 2]}
-
-  boilerplate_start = f"""
-.rodata
-.global {function_name}.kd
-.type {function_name}.kd,STT_OBJECT
-.align 0x10
-.amdhsa_kernel {function_name}"""
-
-  kernel_desc = {
-    '.amdhsa_group_segment_fixed_size': 0, '.amdhsa_private_segment_fixed_size': 0, '.amdhsa_kernarg_size': 0,
-    '.amdhsa_next_free_vgpr': v_cnt,   # this matters!
-    '.amdhsa_reserve_vcc': 0, '.amdhsa_reserve_xnack_mask': 0,
-    '.amdhsa_next_free_sgpr': s_cnt,
-    '.amdhsa_float_round_mode_32': 0, '.amdhsa_float_round_mode_16_64': 0, '.amdhsa_float_denorm_mode_32': 3, '.amdhsa_float_denorm_mode_16_64': 3,
-    '.amdhsa_dx10_clamp': 1, '.amdhsa_ieee_mode': 1, '.amdhsa_fp16_overflow': 0,
-    '.amdhsa_workgroup_processor_mode': 1, '.amdhsa_memory_ordered': 1, '.amdhsa_forward_progress': 0, '.amdhsa_enable_private_segment': 0,
-    '.amdhsa_system_sgpr_workgroup_id_x': 1, '.amdhsa_system_sgpr_workgroup_id_y': 1, '.amdhsa_system_sgpr_workgroup_id_z': 1,
-    '.amdhsa_system_sgpr_workgroup_info': 0, '.amdhsa_system_vgpr_workitem_id': 2, # is amdhsa_system_vgpr_workitem_id real?
-    '.amdhsa_exception_fp_ieee_invalid_op': 0, '.amdhsa_exception_fp_denorm_src': 0,
-    '.amdhsa_exception_fp_ieee_div_zero': 0, '.amdhsa_exception_fp_ieee_overflow': 0, '.amdhsa_exception_fp_ieee_underflow': 0,
-    '.amdhsa_exception_fp_ieee_inexact': 0, '.amdhsa_exception_int_div_zero': 0,
-    '.amdhsa_user_sgpr_dispatch_ptr': 0, '.amdhsa_user_sgpr_queue_ptr': 0, '.amdhsa_user_sgpr_kernarg_segment_ptr': 1,
-    '.amdhsa_user_sgpr_dispatch_id': 0, '.amdhsa_user_sgpr_private_segment_size': 0, '.amdhsa_wavefront_size32': 1, '.amdhsa_uses_dynamic_stack': 0}
-
-  code_start = f""".end_amdhsa_kernel
-.text
-.global {function_name}
-.type {function_name},@function
-.p2align 8
-{function_name}:
-"""
-
-  ins += ['s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)', 's_endpgm', 's_code_end']
-  return ".amdgpu_metadata\n" + yaml.dump(metadata) + ".end_amdgpu_metadata" + \
-         boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) + "\n" + code_start + \
-         '\n'.join(ins) + f"\n.size {function_name}, .-{function_name}"

extra/backends/triton.py

@@ -1,131 +0,0 @@
-from typing import Dict, List, Final, Callable, DefaultDict
-from collections import defaultdict
-from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
-from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.codegen.opt.kernel import  UOp, Ops
-from triton.compiler import compile as triton_compile
-import linecache
-import math
-import re
-
-triton_dtypes = {dtypes.double: "tl.float64", dtypes.float32: "tl.float32", dtypes.float16: "tl.float16", dtypes.bool: "tl.int1", dtypes.int8: "tl.int8", dtypes.uint8: "tl.uint8", dtypes.int32: "tl.int32", dtypes.int64: "tl.int64", dtypes.uint32: "tl.uint32", dtypes.uint64: "tl.uint64", dtypes.int16: "tl.int16", dtypes.uint16: "tl.uint16"}
-signature_dtypes = {dtypes.double: "fp64",dtypes.float32: "fp32", dtypes.float16: "fp16", dtypes.bool: "i8", dtypes.int8: "i1", dtypes.uint8: "u8", dtypes.int32: "i32", dtypes.int64: "i64", dtypes.uint32: "u32", dtypes.uint64: "u64", dtypes.int16: "i16", dtypes.uint16: "u16"}
-
-def next_power_of_2(x):
-  return 1 << (x - 1).bit_length()
-
-def render_valid(valid):
-  return '(' * (len(valid) -1) + ') and '.join(valid) if len(valid) else 'True'
-
-#NOTE Triton requires matching dimensions for load/store, disable this and see TestOps::test_output_padded_conv_transpose2d fail to compile
-def fill_dims_for_idx(idx, dims):
-  return "(" + idx + "+ (" + (f"0*({'+'.join(d for d in dims)})))") if len(dims) else idx
-
-def get_max(var):
-  if isinstance(var, int): return var
-  return re.sub(r'\[(.*?)\]', '', str(var))[1:-1]
-
-#NOTE can be removed after https://github.com/gpuocelot/gpuocelot/issues/8 gets resolved
-def remove_single_scalar_curly_braces(ptx_code):
-  return '\n'.join([re.sub(r'\{\s*(%\w+)\s*\}', r'\1', line) for line in ptx_code.split('\n')])
-
-def render_const(args,dtype:DType):
-  return (('-' if args<0 else '') + 'tl.where(1,float("inf"),0)') if math.isinf(args) else ('tl.where(1,float("nan"),0)' if math.isnan(args) else f"{int(args)}" if dtypes.is_int(dtype) else str(args))
-
-def render_cast(x:str, dtype:DType, bitcast=False):
-  return f"{x}.to({triton_dtypes[dtype]}, bitcast={bitcast})"
-
-def define_scalar(local_size, dtype, args):
-  if len(local_size) > 0: return f"tl.full(({','.join([str(next_power_of_2(x)) for x in local_size])},),{render_const(args,dtype)}, dtype={triton_dtypes[dtype]})"
-  return render_const(args,dtype)
-
-def uops_to_triton(function_name:str, uops:List[UOp]):
-  local_size: List[int] = []
-  depth = 1
-  signatures, dims, bufs, kernel, valid = [], [], [], [], [] #type: ignore
-
-  c: DefaultDict[str, int] = defaultdict(int)
-  r: Dict[UOp, str] = {}
-  def ssa(u, prefix="t"):
-    nonlocal c, r
-    c[prefix] += 1
-    r[u]=f"{prefix}{c[prefix]-1}"
-    return r[u]
-
-  child_count: DefaultDict[UOp, int] = defaultdict(int)
-  for ru in uops:
-    for v in ru.vin:
-      child_count[v] += 1
-
-  def kk(s): kernel.append("  "*depth+s)
-  code_for_op: Final[Dict[Op, Callable]] = {
-    UnaryOps.EXP2: lambda x,dtype,: f"tl.math.exp2({x})",
-    UnaryOps.LOG2: lambda x,dtype,: f"tl.math.log2({x})",
-    UnaryOps.SIN: lambda x,dtype: f"tl.sin({x})",
-    UnaryOps.SQRT: lambda x,dtype: f"tl.sqrt({x})",
-    UnaryOps.NEG: lambda x,dtype: f"-{x}",
-    BinaryOps.ADD: lambda x,y,dtype: f"({x}+{y})", BinaryOps.SUB: lambda x,y,: f"({x}-{y})",
-    BinaryOps.MUL: lambda x,y,dtype: f"({x}*{y})", BinaryOps.DIV: lambda x,y,: f"({x}/{y})" if y != '0.0' else f"{x}*tl.where({x}==0.0, float('nan'), float('inf'))",
-    BinaryOps.MAX: lambda x,y,dtype: f"tl.maximum({x},{y})",
-    BinaryOps.CMPLT: lambda x,y,dtype: f"({x}<{y})",
-    BinaryOps.MOD: lambda x,y,dtype: f"tl.abs({x})%tl.abs({y})*tl.where({x}<0,-1,1)",
-    TernaryOps.MULACC: lambda x,y,z,dtype: f"(({x}*{y})+{z})",
-    TernaryOps.WHERE: lambda x,y,z,dtype: f"tl.where({x},{y},{z})",
-  }
-  def int_div(x,y): return f"({x}//{y})" if y != '0' else f"{x}*tl.where({x}==0, float('nan'), float('inf'))"
-  for u in uops:
-    uop,dtype,vin,args = u.uop,u.dtype,u.vin,u.arg
-    if uop == Ops.LOOP:
-      kk(f"for {ssa(u, 'ridx')} in range({vin[0].arg}, {r[vin[1]]}):")
-      depth += 1
-    elif uop == Ops.END: depth -= 1
-    elif uop == Ops.ALU:
-      assert dtype is not None
-      val = code_for_op[args](*[r[x] for x in vin])
-      if child_count[u] <=1 or dtypes.is_int(dtype): r[u] = int_div(*[r[x] for x in vin]) if args == BinaryOps.DIV and dtypes.is_int(dtype) else val
-      else: kk(f"{ssa(u, 'alu')} = ({val})")
-    elif uop == Ops.LOAD:
-      assert dtype is not None
-      if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}")
-      else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}")
-    elif uop == Ops.DEFINE_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
-    elif uop == Ops.CONST: r[u] = define_scalar([], dtype, args)
-    elif uop == Ops.ASSIGN:
-      kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}")
-      r[u] = r[vin[0]]
-    elif uop == Ops.STORE:
-      assert not isinstance(dtype, ImageDType), "unimplemented: image store"
-      kk(f"{'if '+r[vin[3]]+': ' if len(vin)>3 else ''}tl.store({r[vin[0]]} + {r[vin[1]]}, {r[vin[2]].replace('//', '/')}, mask = {render_valid(valid)}) ")
-    elif uop == Ops.DEFINE_GLOBAL:
-      bufs.append(args)
-      signatures.append("*" if isinstance(dtype, PtrDType) else "" +  signature_dtypes[dtype])
-      r[u] = args
-    elif uop == Ops.SPECIAL:
-      dims.append(args[1])
-      valid.append(f"{args[1]}<{get_max(args[2])}")
-      if args[1].startswith("g"): kk(f"{args[1]} = tl.program_id({args[0]}) # {args[2]}")
-      elif args[1].startswith("l"):
-        kk(f"{args[1]} = tl.arange({0}, {next_power_of_2(args[2])})")
-        local_size.append(args[2])
-      r[u] = args[1]
-    elif uop == Ops.CAST and dtype is not None: r[u] = render_cast(r[vin[0]], dtype, isinstance(args, tuple) and args[1])
-    else: raise NotImplementedError(f"unimplemented: {uop}")
-
-  prg = f"import triton\nimport triton.language as tl\ntl.core.TRITON_MAX_TENSOR_NUMEL = float('inf')\n@triton.jit\ndef {function_name}("+','.join(bufs)+"):\n"
-  for i, line in enumerate(list(filter(lambda line: "tl.arange" in line, kernel))): kernel[kernel.index(line)] +=  f"[{', '.join([':' if i == j else 'None' for j in range(len(local_size))])}]"
-  prg += "\n".join(kernel)
-
-  acc_local_size = 1
-  for x in local_size: acc_local_size *= next_power_of_2(x)
-  local_size = [acc_local_size] + [1] * (len(local_size) - 1)
-
-  if DEBUG >= 4: print(prg)
-  getlines = linecache.getlines
-  linecache.getlines = lambda filename, module_globals=None: prg.splitlines(keepends=True) if "<triton>" == filename else getlines(filename, module_globals)
-  exec(compile(prg, "<triton>", "exec"), globals()) # pylint: disable=W0122\
-  compiled = triton_compile(globals()[function_name], signature=",".join(signatures), device_type="cuda", debug=False, cc=(35 if getenv("CUDACPU", 0) else None))
-  prg = remove_single_scalar_curly_braces(compiled.asm["ptx"].split(".file")[0].split(".visible .func")[0])
-  max_local_size =  [int(x) for x in prg.split(".maxntid ")[1].split("\n")[0].split(", ")]
-  for i in range(len(local_size)): local_size[i] = min(local_size[i], max_local_size[i])
-
-  return prg, {"shared":compiled.metadata["shared"], "local_size":local_size + [1]*(3-len(local_size))}

extra/datasets/coco.py

@@ -1,199 +0,0 @@
-import json
-import pathlib
-import zipfile
-import numpy as np
-from tinygrad.helpers import fetch
-import pycocotools._mask as _mask
-from examples.mask_rcnn import Masker
-from pycocotools.coco import COCO
-from pycocotools.cocoeval import COCOeval
-
-iou         = _mask.iou
-merge       = _mask.merge
-frPyObjects = _mask.frPyObjects
-
-BASEDIR = pathlib.Path(__file__).parent / "COCO"
-BASEDIR.mkdir(exist_ok=True)
-
-def create_dict(key_row, val_row, rows): return {row[key_row]:row[val_row] for row in rows}
-
-
-if not pathlib.Path(BASEDIR/'val2017').is_dir():
-  fn = fetch('http://images.cocodataset.org/zips/val2017.zip')
-  with zipfile.ZipFile(fn, 'r') as zip_ref:
-    zip_ref.extractall(BASEDIR)
-  fn.unlink()
-
-
-if not pathlib.Path(BASEDIR/'annotations').is_dir():
-  fn = fetch('http://images.cocodataset.org/annotations/annotations_trainval2017.zip')
-  with zipfile.ZipFile(fn, 'r') as zip_ref:
-    zip_ref.extractall(BASEDIR)
-  fn.unlink()
-
-with open(BASEDIR/'annotations/instances_val2017.json', 'r') as f:
-  annotations_raw = json.loads(f.read())
-images = annotations_raw['images']
-categories = annotations_raw['categories']
-annotations = annotations_raw['annotations']
-file_name_to_id = create_dict('file_name', 'id', images)
-id_to_width = create_dict('id', 'width', images)
-id_to_height = create_dict('id', 'height', images)
-json_category_id_to_contiguous_id = {v['id']: i + 1 for i, v in enumerate(categories)}
-contiguous_category_id_to_json_id = {v:k for k,v in json_category_id_to_contiguous_id.items()}
-
-
-def encode(bimask):
-  if len(bimask.shape) == 3:
-    return _mask.encode(bimask)
-  elif len(bimask.shape) == 2:
-    h, w = bimask.shape
-    return _mask.encode(bimask.reshape((h, w, 1), order='F'))[0]
-
-def decode(rleObjs):
-  if type(rleObjs) == list:
-    return _mask.decode(rleObjs)
-  else:
-    return _mask.decode([rleObjs])[:,:,0]
-
-def area(rleObjs):
-  if type(rleObjs) == list:
-    return _mask.area(rleObjs)
-  else:
-    return _mask.area([rleObjs])[0]
-
-def toBbox(rleObjs):
-  if type(rleObjs) == list:
-    return _mask.toBbox(rleObjs)
-  else:
-    return _mask.toBbox([rleObjs])[0]
-
-
-def convert_prediction_to_coco_bbox(file_name, prediction):
-  coco_results = []
-  try:
-    original_id = file_name_to_id[file_name]
-    if len(prediction) == 0:
-      return coco_results
-
-    image_width = id_to_width[original_id]
-    image_height = id_to_height[original_id]
-    prediction = prediction.resize((image_width, image_height))
-    prediction = prediction.convert("xywh")
-
-    boxes = prediction.bbox.numpy().tolist()
-    scores = prediction.get_field("scores").numpy().tolist()
-    labels = prediction.get_field("labels").numpy().tolist()
-
-    mapped_labels = [contiguous_category_id_to_json_id[int(i)] for i in labels]
-
-    coco_results.extend(
-      [
-        {
-          "image_id": original_id,
-          "category_id": mapped_labels[k],
-          "bbox": box,
-          "score": scores[k],
-        }
-          for k, box in enumerate(boxes)
-      ]
-    )
-  except Exception as e:
-    print(file_name, e)
-  return coco_results
-
-masker = Masker(threshold=0.5, padding=1)
-
-def convert_prediction_to_coco_mask(file_name, prediction):
-  coco_results = []
-  try:
-    original_id = file_name_to_id[file_name]
-    if len(prediction) == 0:
-      return coco_results
-
-    image_width = id_to_width[original_id]
-    image_height = id_to_height[original_id]
-    prediction = prediction.resize((image_width, image_height))
-    masks = prediction.get_field("mask")
-
-    scores = prediction.get_field("scores").numpy().tolist()
-    labels = prediction.get_field("labels").numpy().tolist()
-
-    masks = masker([masks], [prediction])[0].numpy()
-
-    rles = [
-      encode(np.array(mask[0, :, :, np.newaxis], order="F"))[0]
-      for mask in masks
-    ]
-    for rle in rles:
-      rle["counts"] = rle["counts"].decode("utf-8")
-
-    mapped_labels = [contiguous_category_id_to_json_id[int(i)] for i in labels]
-
-    coco_results.extend(
-      [
-        {
-          "image_id": original_id,
-          "category_id": mapped_labels[k],
-          "segmentation": rle,
-          "score": scores[k],
-        }
-          for k, rle in enumerate(rles)
-      ]
-    )
-  except Exception as e:
-    print(file_name, e)
-  return coco_results
-
-
-
-def accumulate_predictions_for_coco(coco_results, json_result_file, rm=False):
-  path = pathlib.Path(json_result_file)
-  if rm and path.exists(): path.unlink()
-  with open(path, "a") as f:
-    for s in coco_results:
-      f.write(json.dumps(s))
-      f.write('\n')
-
-def remove_dup(l):
-  seen = set()
-  seen_add = seen.add
-  return [x for x in l if not (x in seen or seen_add(x))]
-
-class NpEncoder(json.JSONEncoder):
-  def default(self, obj):
-    if isinstance(obj, np.integer):
-      return int(obj)
-    if isinstance(obj, np.floating):
-      return float(obj)
-    if isinstance(obj, np.ndarray):
-      return obj.tolist()
-    return super(NpEncoder, self).default(obj)
-
-
-def evaluate_predictions_on_coco(json_result_file, iou_type="bbox"):
-  coco_results = []
-  with open(json_result_file, "r") as f:
-    for line in f:
-      coco_results.append(json.loads(line))
-
-  coco_gt = COCO(str(BASEDIR/'annotations/instances_val2017.json'))
-  set_of_json = remove_dup([json.dumps(d, cls=NpEncoder) for d in coco_results])
-  unique_list = [json.loads(s) for s in set_of_json]
-
-  with open(f'{json_result_file}.flattend', "w") as f:
-    json.dump(unique_list, f)
-
-  coco_dt = coco_gt.loadRes(str(f'{json_result_file}.flattend'))
-  coco_eval = COCOeval(coco_gt, coco_dt, iou_type)
-  coco_eval.evaluate()
-  coco_eval.accumulate()
-  coco_eval.summarize()
-  return coco_eval
-
-def iterate(files, bs=1):
-  batch = []
-  for file in files:
-    batch.append(file)
-    if len(batch) >= bs: yield batch; batch = []
-  if len(batch) > 0: yield batch; batch = []

extra/disassemblers/adreno/.gitignore

@@ -1 +0,0 @@
-disasm.so

extra/disassemblers/adreno/README

@@ -1,5 +0,0 @@
-From the Freedreno project
-
-https://gallium.readthedocs.io/en/latest/gallium/drivers/freedreno.html
-
-In Mesa3D, so licensed MIT.

extra/disassemblers/adreno/__init__.py

@@ -1,22 +0,0 @@
-import ctypes
-import os
-import pathlib
-import struct
-from hexdump import hexdump
-
-fxn = None
-def disasm_raw(buf):
-  global fxn
-  if fxn is None:
-    shared = pathlib.Path(__file__).parent / "disasm.so"
-    if not shared.is_file():
-      os.system(f'cd {pathlib.Path(__file__).parent} && gcc -shared disasm-a3xx.c -o disasm.so')
-    fxn = ctypes.CDLL(shared.as_posix())['disasm']
-  fxn(buf, len(buf))
-
-def disasm(buf):
-  def _read_lib(off): return struct.unpack("I", buf[off:off+4])[0]
-
-  image_offset = _read_lib(0xc0)
-  image_size = _read_lib(0x100)
-  disasm_raw(buf[image_offset:image_offset+image_size])

extra/disassemblers/adreno/shader_enums.h

@@ -1,906 +0,0 @@
-/*
- * Mesa 3-D graphics library
- *
- * Copyright (C) 1999-2008  Brian Paul   All Rights Reserved.
- * Copyright (C) 2009  VMware, Inc.  All Rights Reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included
- * in all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
- * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- */
-
-#ifndef SHADER_ENUMS_H
-#define SHADER_ENUMS_H
-
-#include <stdbool.h>
-
-/* Project-wide (GL and Vulkan) maximum. */
-#define MAX_DRAW_BUFFERS 8
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-/**
- * Shader stages.
- *
- * The order must match how shaders are ordered in the pipeline.
- * The GLSL linker assumes that if i<j, then the j-th shader is
- * executed later than the i-th shader.
- */
-typedef enum
-{
-   MESA_SHADER_NONE = -1,
-   MESA_SHADER_VERTEX = 0,
-   MESA_SHADER_TESS_CTRL = 1,
-   MESA_SHADER_TESS_EVAL = 2,
-   MESA_SHADER_GEOMETRY = 3,
-   MESA_SHADER_FRAGMENT = 4,
-   MESA_SHADER_COMPUTE = 5,
-   /* must be last so it doesn't affect the GL pipeline */
-   MESA_SHADER_KERNEL = 6,
-} gl_shader_stage;
-
-static inline bool
-gl_shader_stage_is_compute(gl_shader_stage stage)
-{
-   return stage == MESA_SHADER_COMPUTE || stage == MESA_SHADER_KERNEL;
-}
-
-/**
- * Number of STATE_* values we need to address any GL state.
- * Used to dimension arrays.
- */
-#define STATE_LENGTH 5
-
-typedef short gl_state_index16; /* see enum gl_state_index */
-
-const char *gl_shader_stage_name(gl_shader_stage stage);
-
-/**
- * Translate a gl_shader_stage to a short shader stage name for debug
- * printouts and error messages.
- */
-const char *_mesa_shader_stage_to_string(unsigned stage);
-
-/**
- * Translate a gl_shader_stage to a shader stage abbreviation (VS, GS, FS)
- * for debug printouts and error messages.
- */
-const char *_mesa_shader_stage_to_abbrev(unsigned stage);
-
-/**
- * GL related stages (not including CL)
- */
-#define MESA_SHADER_STAGES (MESA_SHADER_COMPUTE + 1)
-
-/**
- * All stages
- */
-#define MESA_ALL_SHADER_STAGES (MESA_SHADER_KERNEL + 1)
-
-
-/**
- * Indexes for vertex program attributes.
- * GL_NV_vertex_program aliases generic attributes over the conventional
- * attributes.  In GL_ARB_vertex_program shader the aliasing is optional.
- * In GL_ARB_vertex_shader / OpenGL 2.0 the aliasing is disallowed (the
- * generic attributes are distinct/separate).
- */
-typedef enum
-{
-   VERT_ATTRIB_POS,
-   VERT_ATTRIB_NORMAL,
-   VERT_ATTRIB_COLOR0,
-   VERT_ATTRIB_COLOR1,
-   VERT_ATTRIB_FOG,
-   VERT_ATTRIB_COLOR_INDEX,
-   VERT_ATTRIB_EDGEFLAG,
-   VERT_ATTRIB_TEX0,
-   VERT_ATTRIB_TEX1,
-   VERT_ATTRIB_TEX2,
-   VERT_ATTRIB_TEX3,
-   VERT_ATTRIB_TEX4,
-   VERT_ATTRIB_TEX5,
-   VERT_ATTRIB_TEX6,
-   VERT_ATTRIB_TEX7,
-   VERT_ATTRIB_POINT_SIZE,
-   VERT_ATTRIB_GENERIC0,
-   VERT_ATTRIB_GENERIC1,
-   VERT_ATTRIB_GENERIC2,
-   VERT_ATTRIB_GENERIC3,
-   VERT_ATTRIB_GENERIC4,
-   VERT_ATTRIB_GENERIC5,
-   VERT_ATTRIB_GENERIC6,
-   VERT_ATTRIB_GENERIC7,
-   VERT_ATTRIB_GENERIC8,
-   VERT_ATTRIB_GENERIC9,
-   VERT_ATTRIB_GENERIC10,
-   VERT_ATTRIB_GENERIC11,
-   VERT_ATTRIB_GENERIC12,
-   VERT_ATTRIB_GENERIC13,
-   VERT_ATTRIB_GENERIC14,
-   VERT_ATTRIB_GENERIC15,
-   VERT_ATTRIB_MAX
-} gl_vert_attrib;
-
-const char *gl_vert_attrib_name(gl_vert_attrib attrib);
-
-/**
- * Symbolic constats to help iterating over
- * specific blocks of vertex attributes.
- *
- * VERT_ATTRIB_FF
- *   includes all fixed function attributes as well as
- *   the aliased GL_NV_vertex_program shader attributes.
- * VERT_ATTRIB_TEX
- *   include the classic texture coordinate attributes.
- *   Is a subset of VERT_ATTRIB_FF.
- * VERT_ATTRIB_GENERIC
- *   include the OpenGL 2.0+ GLSL generic shader attributes.
- *   These alias the generic GL_ARB_vertex_shader attributes.
- * VERT_ATTRIB_MAT
- *   include the generic shader attributes used to alias
- *   varying material values for the TNL shader programs.
- *   They are located at the end of the generic attribute
- *   block not to overlap with the generic 0 attribute.
- */
-#define VERT_ATTRIB_FF(i)           (VERT_ATTRIB_POS + (i))
-#define VERT_ATTRIB_FF_MAX          VERT_ATTRIB_GENERIC0
-
-#define VERT_ATTRIB_TEX(i)          (VERT_ATTRIB_TEX0 + (i))
-#define VERT_ATTRIB_TEX_MAX         MAX_TEXTURE_COORD_UNITS
-
-#define VERT_ATTRIB_GENERIC(i)      (VERT_ATTRIB_GENERIC0 + (i))
-#define VERT_ATTRIB_GENERIC_MAX     MAX_VERTEX_GENERIC_ATTRIBS
-
-#define VERT_ATTRIB_MAT0            \
-   (VERT_ATTRIB_GENERIC_MAX - VERT_ATTRIB_MAT_MAX)
-#define VERT_ATTRIB_MAT(i)          \
-   VERT_ATTRIB_GENERIC((i) + VERT_ATTRIB_MAT0)
-#define VERT_ATTRIB_MAT_MAX         MAT_ATTRIB_MAX
-
-/**
- * Bitflags for vertex attributes.
- * These are used in bitfields in many places.
- */
-/*@{*/
-#define VERT_BIT_POS             BITFIELD_BIT(VERT_ATTRIB_POS)
-#define VERT_BIT_NORMAL          BITFIELD_BIT(VERT_ATTRIB_NORMAL)
-#define VERT_BIT_COLOR0          BITFIELD_BIT(VERT_ATTRIB_COLOR0)
-#define VERT_BIT_COLOR1          BITFIELD_BIT(VERT_ATTRIB_COLOR1)
-#define VERT_BIT_FOG             BITFIELD_BIT(VERT_ATTRIB_FOG)
-#define VERT_BIT_COLOR_INDEX     BITFIELD_BIT(VERT_ATTRIB_COLOR_INDEX)
-#define VERT_BIT_EDGEFLAG        BITFIELD_BIT(VERT_ATTRIB_EDGEFLAG)
-#define VERT_BIT_TEX0            BITFIELD_BIT(VERT_ATTRIB_TEX0)
-#define VERT_BIT_TEX1            BITFIELD_BIT(VERT_ATTRIB_TEX1)
-#define VERT_BIT_TEX2            BITFIELD_BIT(VERT_ATTRIB_TEX2)
-#define VERT_BIT_TEX3            BITFIELD_BIT(VERT_ATTRIB_TEX3)
-#define VERT_BIT_TEX4            BITFIELD_BIT(VERT_ATTRIB_TEX4)
-#define VERT_BIT_TEX5            BITFIELD_BIT(VERT_ATTRIB_TEX5)
-#define VERT_BIT_TEX6            BITFIELD_BIT(VERT_ATTRIB_TEX6)
-#define VERT_BIT_TEX7            BITFIELD_BIT(VERT_ATTRIB_TEX7)
-#define VERT_BIT_POINT_SIZE      BITFIELD_BIT(VERT_ATTRIB_POINT_SIZE)
-#define VERT_BIT_GENERIC0        BITFIELD_BIT(VERT_ATTRIB_GENERIC0)
-
-#define VERT_BIT(i)              BITFIELD_BIT(i)
-#define VERT_BIT_ALL             BITFIELD_RANGE(0, VERT_ATTRIB_MAX)
-
-#define VERT_BIT_FF(i)           VERT_BIT(i)
-#define VERT_BIT_FF_ALL          BITFIELD_RANGE(0, VERT_ATTRIB_FF_MAX)
-#define VERT_BIT_TEX(i)          VERT_BIT(VERT_ATTRIB_TEX(i))
-#define VERT_BIT_TEX_ALL         \
-   BITFIELD_RANGE(VERT_ATTRIB_TEX(0), VERT_ATTRIB_TEX_MAX)
-
-#define VERT_BIT_GENERIC(i)      VERT_BIT(VERT_ATTRIB_GENERIC(i))
-#define VERT_BIT_GENERIC_ALL     \
-   BITFIELD_RANGE(VERT_ATTRIB_GENERIC(0), VERT_ATTRIB_GENERIC_MAX)
-
-#define VERT_BIT_MAT(i)	         VERT_BIT(VERT_ATTRIB_MAT(i))
-#define VERT_BIT_MAT_ALL         \
-   BITFIELD_RANGE(VERT_ATTRIB_MAT(0), VERT_ATTRIB_MAT_MAX)
-/*@}*/
-
-#define MAX_VARYING 32 /**< number of float[4] vectors */
-
-/**
- * Indexes for vertex shader outputs, geometry shader inputs/outputs, and
- * fragment shader inputs.
- *
- * Note that some of these values are not available to all pipeline stages.
- *
- * When this enum is updated, the following code must be updated too:
- * - vertResults (in prog_print.c's arb_output_attrib_string())
- * - fragAttribs (in prog_print.c's arb_input_attrib_string())
- * - _mesa_varying_slot_in_fs()
- */
-typedef enum
-{
-   VARYING_SLOT_POS,
-   VARYING_SLOT_COL0, /* COL0 and COL1 must be contiguous */
-   VARYING_SLOT_COL1,
-   VARYING_SLOT_FOGC,
-   VARYING_SLOT_TEX0, /* TEX0-TEX7 must be contiguous */
-   VARYING_SLOT_TEX1,
-   VARYING_SLOT_TEX2,
-   VARYING_SLOT_TEX3,
-   VARYING_SLOT_TEX4,
-   VARYING_SLOT_TEX5,
-   VARYING_SLOT_TEX6,
-   VARYING_SLOT_TEX7,
-   VARYING_SLOT_PSIZ, /* Does not appear in FS */
-   VARYING_SLOT_BFC0, /* Does not appear in FS */
-   VARYING_SLOT_BFC1, /* Does not appear in FS */
-   VARYING_SLOT_EDGE, /* Does not appear in FS */
-   VARYING_SLOT_CLIP_VERTEX, /* Does not appear in FS */
-   VARYING_SLOT_CLIP_DIST0,
-   VARYING_SLOT_CLIP_DIST1,
-   VARYING_SLOT_CULL_DIST0,
-   VARYING_SLOT_CULL_DIST1,
-   VARYING_SLOT_PRIMITIVE_ID, /* Does not appear in VS */
-   VARYING_SLOT_LAYER, /* Appears as VS or GS output */
-   VARYING_SLOT_VIEWPORT, /* Appears as VS or GS output */
-   VARYING_SLOT_FACE, /* FS only */
-   VARYING_SLOT_PNTC, /* FS only */
-   VARYING_SLOT_TESS_LEVEL_OUTER, /* Only appears as TCS output. */
-   VARYING_SLOT_TESS_LEVEL_INNER, /* Only appears as TCS output. */
-   VARYING_SLOT_BOUNDING_BOX0, /* Only appears as TCS output. */
-   VARYING_SLOT_BOUNDING_BOX1, /* Only appears as TCS output. */
-   VARYING_SLOT_VIEW_INDEX,
-   VARYING_SLOT_VIEWPORT_MASK, /* Does not appear in FS */
-   VARYING_SLOT_VAR0, /* First generic varying slot */
-   /* the remaining are simply for the benefit of gl_varying_slot_name()
-    * and not to be construed as an upper bound:
-    */
-   VARYING_SLOT_VAR1,
-   VARYING_SLOT_VAR2,
-   VARYING_SLOT_VAR3,
-   VARYING_SLOT_VAR4,
-   VARYING_SLOT_VAR5,
-   VARYING_SLOT_VAR6,
-   VARYING_SLOT_VAR7,
-   VARYING_SLOT_VAR8,
-   VARYING_SLOT_VAR9,
-   VARYING_SLOT_VAR10,
-   VARYING_SLOT_VAR11,
-   VARYING_SLOT_VAR12,
-   VARYING_SLOT_VAR13,
-   VARYING_SLOT_VAR14,
-   VARYING_SLOT_VAR15,
-   VARYING_SLOT_VAR16,
-   VARYING_SLOT_VAR17,
-   VARYING_SLOT_VAR18,
-   VARYING_SLOT_VAR19,
-   VARYING_SLOT_VAR20,
-   VARYING_SLOT_VAR21,
-   VARYING_SLOT_VAR22,
-   VARYING_SLOT_VAR23,
-   VARYING_SLOT_VAR24,
-   VARYING_SLOT_VAR25,
-   VARYING_SLOT_VAR26,
-   VARYING_SLOT_VAR27,
-   VARYING_SLOT_VAR28,
-   VARYING_SLOT_VAR29,
-   VARYING_SLOT_VAR30,
-   VARYING_SLOT_VAR31,
-} gl_varying_slot;
-
-
-#define VARYING_SLOT_MAX	(VARYING_SLOT_VAR0 + MAX_VARYING)
-#define VARYING_SLOT_PATCH0	(VARYING_SLOT_MAX)
-#define VARYING_SLOT_TESS_MAX	(VARYING_SLOT_PATCH0 + MAX_VARYING)
-#define MAX_VARYINGS_INCL_PATCH (VARYING_SLOT_TESS_MAX - VARYING_SLOT_VAR0)
-
-const char *gl_varying_slot_name(gl_varying_slot slot);
-
-/**
- * Bitflags for varying slots.
- */
-/*@{*/
-#define VARYING_BIT_POS BITFIELD64_BIT(VARYING_SLOT_POS)
-#define VARYING_BIT_COL0 BITFIELD64_BIT(VARYING_SLOT_COL0)
-#define VARYING_BIT_COL1 BITFIELD64_BIT(VARYING_SLOT_COL1)
-#define VARYING_BIT_FOGC BITFIELD64_BIT(VARYING_SLOT_FOGC)
-#define VARYING_BIT_TEX0 BITFIELD64_BIT(VARYING_SLOT_TEX0)
-#define VARYING_BIT_TEX1 BITFIELD64_BIT(VARYING_SLOT_TEX1)
-#define VARYING_BIT_TEX2 BITFIELD64_BIT(VARYING_SLOT_TEX2)
-#define VARYING_BIT_TEX3 BITFIELD64_BIT(VARYING_SLOT_TEX3)
-#define VARYING_BIT_TEX4 BITFIELD64_BIT(VARYING_SLOT_TEX4)
-#define VARYING_BIT_TEX5 BITFIELD64_BIT(VARYING_SLOT_TEX5)
-#define VARYING_BIT_TEX6 BITFIELD64_BIT(VARYING_SLOT_TEX6)
-#define VARYING_BIT_TEX7 BITFIELD64_BIT(VARYING_SLOT_TEX7)
-#define VARYING_BIT_TEX(U) BITFIELD64_BIT(VARYING_SLOT_TEX0 + (U))
-#define VARYING_BITS_TEX_ANY BITFIELD64_RANGE(VARYING_SLOT_TEX0, \
-                                              MAX_TEXTURE_COORD_UNITS)
-#define VARYING_BIT_PSIZ BITFIELD64_BIT(VARYING_SLOT_PSIZ)
-#define VARYING_BIT_BFC0 BITFIELD64_BIT(VARYING_SLOT_BFC0)
-#define VARYING_BIT_BFC1 BITFIELD64_BIT(VARYING_SLOT_BFC1)
-#define VARYING_BITS_COLOR (VARYING_BIT_COL0 | \
-                            VARYING_BIT_COL1 |        \
-                            VARYING_BIT_BFC0 |        \
-                            VARYING_BIT_BFC1)
-#define VARYING_BIT_EDGE BITFIELD64_BIT(VARYING_SLOT_EDGE)
-#define VARYING_BIT_CLIP_VERTEX BITFIELD64_BIT(VARYING_SLOT_CLIP_VERTEX)
-#define VARYING_BIT_CLIP_DIST0 BITFIELD64_BIT(VARYING_SLOT_CLIP_DIST0)
-#define VARYING_BIT_CLIP_DIST1 BITFIELD64_BIT(VARYING_SLOT_CLIP_DIST1)
-#define VARYING_BIT_CULL_DIST0 BITFIELD64_BIT(VARYING_SLOT_CULL_DIST0)
-#define VARYING_BIT_CULL_DIST1 BITFIELD64_BIT(VARYING_SLOT_CULL_DIST1)
-#define VARYING_BIT_PRIMITIVE_ID BITFIELD64_BIT(VARYING_SLOT_PRIMITIVE_ID)
-#define VARYING_BIT_LAYER BITFIELD64_BIT(VARYING_SLOT_LAYER)
-#define VARYING_BIT_VIEWPORT BITFIELD64_BIT(VARYING_SLOT_VIEWPORT)
-#define VARYING_BIT_FACE BITFIELD64_BIT(VARYING_SLOT_FACE)
-#define VARYING_BIT_PNTC BITFIELD64_BIT(VARYING_SLOT_PNTC)
-#define VARYING_BIT_TESS_LEVEL_OUTER BITFIELD64_BIT(VARYING_SLOT_TESS_LEVEL_OUTER)
-#define VARYING_BIT_TESS_LEVEL_INNER BITFIELD64_BIT(VARYING_SLOT_TESS_LEVEL_INNER)
-#define VARYING_BIT_BOUNDING_BOX0 BITFIELD64_BIT(VARYING_SLOT_BOUNDING_BOX0)
-#define VARYING_BIT_BOUNDING_BOX1 BITFIELD64_BIT(VARYING_SLOT_BOUNDING_BOX1)
-#define VARYING_BIT_VIEWPORT_MASK BITFIELD64_BIT(VARYING_SLOT_VIEWPORT_MASK)
-#define VARYING_BIT_VAR(V) BITFIELD64_BIT(VARYING_SLOT_VAR0 + (V))
-/*@}*/
-
-/**
- * Bitflags for system values.
- */
-#define SYSTEM_BIT_SAMPLE_ID ((uint64_t)1 << SYSTEM_VALUE_SAMPLE_ID)
-#define SYSTEM_BIT_SAMPLE_POS ((uint64_t)1 << SYSTEM_VALUE_SAMPLE_POS)
-#define SYSTEM_BIT_SAMPLE_MASK_IN ((uint64_t)1 << SYSTEM_VALUE_SAMPLE_MASK_IN)
-#define SYSTEM_BIT_LOCAL_INVOCATION_ID ((uint64_t)1 << SYSTEM_VALUE_LOCAL_INVOCATION_ID)
-
-/**
- * If the gl_register_file is PROGRAM_SYSTEM_VALUE, the register index will be
- * one of these values.  If a NIR variable's mode is nir_var_system_value, it
- * will be one of these values.
- */
-typedef enum
-{
-   /**
-    * \name System values applicable to all shaders
-    */
-   /*@{*/
-
-   /**
-    * Builtin variables added by GL_ARB_shader_ballot.
-    */
-   /*@{*/
-
-   /**
-    * From the GL_ARB_shader-ballot spec:
-    *
-    *    "A sub-group is a collection of invocations which execute in lockstep.
-    *     The variable <gl_SubGroupSizeARB> is the maximum number of
-    *     invocations in a sub-group. The maximum <gl_SubGroupSizeARB>
-    *     supported in this extension is 64."
-    *
-    * The spec defines this as a uniform. However, it's highly unlikely that
-    * implementations actually treat it as a uniform (which is loaded from a
-    * constant buffer). Most likely, this is an implementation-wide constant,
-    * or perhaps something that depends on the shader stage.
-    */
-   SYSTEM_VALUE_SUBGROUP_SIZE,
-
-   /**
-    * From the GL_ARB_shader_ballot spec:
-    *
-    *    "The variable <gl_SubGroupInvocationARB> holds the index of the
-    *     invocation within sub-group. This variable is in the range 0 to
-    *     <gl_SubGroupSizeARB>-1, where <gl_SubGroupSizeARB> is the total
-    *     number of invocations in a sub-group."
-    */
-   SYSTEM_VALUE_SUBGROUP_INVOCATION,
-
-   /**
-    * From the GL_ARB_shader_ballot spec:
-    *
-    *    "The <gl_SubGroup??MaskARB> variables provide a bitmask for all
-    *     invocations, with one bit per invocation starting with the least
-    *     significant bit, according to the following table,
-    *
-    *       variable               equation for bit values
-    *       --------------------   ------------------------------------
-    *       gl_SubGroupEqMaskARB   bit index == gl_SubGroupInvocationARB
-    *       gl_SubGroupGeMaskARB   bit index >= gl_SubGroupInvocationARB
-    *       gl_SubGroupGtMaskARB   bit index >  gl_SubGroupInvocationARB
-    *       gl_SubGroupLeMaskARB   bit index <= gl_SubGroupInvocationARB
-    *       gl_SubGroupLtMaskARB   bit index <  gl_SubGroupInvocationARB
-    */
-   SYSTEM_VALUE_SUBGROUP_EQ_MASK,
-   SYSTEM_VALUE_SUBGROUP_GE_MASK,
-   SYSTEM_VALUE_SUBGROUP_GT_MASK,
-   SYSTEM_VALUE_SUBGROUP_LE_MASK,
-   SYSTEM_VALUE_SUBGROUP_LT_MASK,
-   /*@}*/
-
-   /**
-    * Builtin variables added by VK_KHR_subgroups
-    */
-   /*@{*/
-   SYSTEM_VALUE_NUM_SUBGROUPS,
-   SYSTEM_VALUE_SUBGROUP_ID,
-   /*@}*/
-
-   /*@}*/
-
-   /**
-    * \name Vertex shader system values
-    */
-   /*@{*/
-   /**
-    * OpenGL-style vertex ID.
-    *
-    * Section 2.11.7 (Shader Execution), subsection Shader Inputs, of the
-    * OpenGL 3.3 core profile spec says:
-    *
-    *     "gl_VertexID holds the integer index i implicitly passed by
-    *     DrawArrays or one of the other drawing commands defined in section
-    *     2.8.3."
-    *
-    * Section 2.8.3 (Drawing Commands) of the same spec says:
-    *
-    *     "The commands....are equivalent to the commands with the same base
-    *     name (without the BaseVertex suffix), except that the ith element
-    *     transferred by the corresponding draw call will be taken from
-    *     element indices[i] + basevertex of each enabled array."
-    *
-    * Additionally, the overview in the GL_ARB_shader_draw_parameters spec
-    * says:
-    *
-    *     "In unextended GL, vertex shaders have inputs named gl_VertexID and
-    *     gl_InstanceID, which contain, respectively the index of the vertex
-    *     and instance. The value of gl_VertexID is the implicitly passed
-    *     index of the vertex being processed, which includes the value of
-    *     baseVertex, for those commands that accept it."
-    *
-    * gl_VertexID gets basevertex added in.  This differs from DirectX where
-    * SV_VertexID does \b not get basevertex added in.
-    *
-    * \note
-    * If all system values are available, \c SYSTEM_VALUE_VERTEX_ID will be
-    * equal to \c SYSTEM_VALUE_VERTEX_ID_ZERO_BASE plus
-    * \c SYSTEM_VALUE_BASE_VERTEX.
-    *
-    * \sa SYSTEM_VALUE_VERTEX_ID_ZERO_BASE, SYSTEM_VALUE_BASE_VERTEX
-    */
-   SYSTEM_VALUE_VERTEX_ID,
-
-   /**
-    * Instanced ID as supplied to gl_InstanceID
-    *
-    * Values assigned to gl_InstanceID always begin with zero, regardless of
-    * the value of baseinstance.
-    *
-    * Section 11.1.3.9 (Shader Inputs) of the OpenGL 4.4 core profile spec
-    * says:
-    *
-    *     "gl_InstanceID holds the integer instance number of the current
-    *     primitive in an instanced draw call (see section 10.5)."
-    *
-    * Through a big chain of pseudocode, section 10.5 describes that
-    * baseinstance is not counted by gl_InstanceID.  In that section, notice
-    *
-    *     "If an enabled vertex attribute array is instanced (it has a
-    *     non-zero divisor as specified by VertexAttribDivisor), the element
-    *     index that is transferred to the GL, for all vertices, is given by
-    *
-    *         floor(instance/divisor) + baseinstance
-    *
-    *     If an array corresponding to an attribute required by a vertex
-    *     shader is not enabled, then the corresponding element is taken from
-    *     the current attribute state (see section 10.2)."
-    *
-    * Note that baseinstance is \b not included in the value of instance.
-    */
-   SYSTEM_VALUE_INSTANCE_ID,
-
-   /**
-    * Vulkan InstanceIndex.
-    *
-    * InstanceIndex = gl_InstanceID + gl_BaseInstance
-    */
-   SYSTEM_VALUE_INSTANCE_INDEX,
-
-   /**
-    * DirectX-style vertex ID.
-    *
-    * Unlike \c SYSTEM_VALUE_VERTEX_ID, this system value does \b not include
-    * the value of basevertex.
-    *
-    * \sa SYSTEM_VALUE_VERTEX_ID, SYSTEM_VALUE_BASE_VERTEX
-    */
-   SYSTEM_VALUE_VERTEX_ID_ZERO_BASE,
-
-   /**
-    * Value of \c basevertex passed to \c glDrawElementsBaseVertex and similar
-    * functions.
-    *
-    * \sa SYSTEM_VALUE_VERTEX_ID, SYSTEM_VALUE_VERTEX_ID_ZERO_BASE
-    */
-   SYSTEM_VALUE_BASE_VERTEX,
-
-   /**
-    * Depending on the type of the draw call (indexed or non-indexed),
-    * is the value of \c basevertex passed to \c glDrawElementsBaseVertex and
-    * similar, or is the value of \c first passed to \c glDrawArrays and
-    * similar.
-    *
-    * \note
-    * It can be used to calculate the \c SYSTEM_VALUE_VERTEX_ID as
-    * \c SYSTEM_VALUE_VERTEX_ID_ZERO_BASE plus \c SYSTEM_VALUE_FIRST_VERTEX.
-    *
-    * \sa SYSTEM_VALUE_VERTEX_ID_ZERO_BASE, SYSTEM_VALUE_VERTEX_ID
-    */
-   SYSTEM_VALUE_FIRST_VERTEX,
-
-   /**
-    * If the Draw command used to start the rendering was an indexed draw
-    * or not (~0/0). Useful to calculate \c SYSTEM_VALUE_BASE_VERTEX as
-    * \c SYSTEM_VALUE_IS_INDEXED_DRAW & \c SYSTEM_VALUE_FIRST_VERTEX.
-    */
-   SYSTEM_VALUE_IS_INDEXED_DRAW,
-
-   /**
-    * Value of \c baseinstance passed to instanced draw entry points
-    *
-    * \sa SYSTEM_VALUE_INSTANCE_ID
-    */
-   SYSTEM_VALUE_BASE_INSTANCE,
-
-   /**
-    * From _ARB_shader_draw_parameters:
-    *
-    *   "Additionally, this extension adds a further built-in variable,
-    *    gl_DrawID to the shading language. This variable contains the index
-    *    of the draw currently being processed by a Multi* variant of a
-    *    drawing command (such as MultiDrawElements or
-    *    MultiDrawArraysIndirect)."
-    *
-    * If GL_ARB_multi_draw_indirect is not supported, this is always 0.
-    */
-   SYSTEM_VALUE_DRAW_ID,
-   /*@}*/
-
-   /**
-    * \name Geometry shader system values
-    */
-   /*@{*/
-   SYSTEM_VALUE_INVOCATION_ID,  /**< (Also in Tessellation Control shader) */
-   /*@}*/
-
-   /**
-    * \name Fragment shader system values
-    */
-   /*@{*/
-   SYSTEM_VALUE_FRAG_COORD,
-   SYSTEM_VALUE_POINT_COORD,
-   SYSTEM_VALUE_FRONT_FACE,
-   SYSTEM_VALUE_SAMPLE_ID,
-   SYSTEM_VALUE_SAMPLE_POS,
-   SYSTEM_VALUE_SAMPLE_MASK_IN,
-   SYSTEM_VALUE_HELPER_INVOCATION,
-   SYSTEM_VALUE_COLOR0,
-   SYSTEM_VALUE_COLOR1,
-   /*@}*/
-
-   /**
-    * \name Tessellation Evaluation shader system values
-    */
-   /*@{*/
-   SYSTEM_VALUE_TESS_COORD,
-   SYSTEM_VALUE_VERTICES_IN,    /**< Tessellation vertices in input patch */
-   SYSTEM_VALUE_PRIMITIVE_ID,
-   SYSTEM_VALUE_TESS_LEVEL_OUTER, /**< TES input */
-   SYSTEM_VALUE_TESS_LEVEL_INNER, /**< TES input */
-   SYSTEM_VALUE_TESS_LEVEL_OUTER_DEFAULT, /**< TCS input for passthru TCS */
-   SYSTEM_VALUE_TESS_LEVEL_INNER_DEFAULT, /**< TCS input for passthru TCS */
-   /*@}*/
-
-   /**
-    * \name Compute shader system values
-    */
-   /*@{*/
-   SYSTEM_VALUE_LOCAL_INVOCATION_ID,
-   SYSTEM_VALUE_LOCAL_INVOCATION_INDEX,
-   SYSTEM_VALUE_GLOBAL_INVOCATION_ID,
-   SYSTEM_VALUE_GLOBAL_INVOCATION_INDEX,
-   SYSTEM_VALUE_WORK_GROUP_ID,
-   SYSTEM_VALUE_NUM_WORK_GROUPS,
-   SYSTEM_VALUE_LOCAL_GROUP_SIZE,
-   SYSTEM_VALUE_GLOBAL_GROUP_SIZE,
-   SYSTEM_VALUE_WORK_DIM,
-   SYSTEM_VALUE_USER_DATA_AMD,
-   /*@}*/
-
-   /** Required for VK_KHR_device_group */
-   SYSTEM_VALUE_DEVICE_INDEX,
-
-   /** Required for VK_KHX_multiview */
-   SYSTEM_VALUE_VIEW_INDEX,
-
-   /**
-    * Driver internal vertex-count, used (for example) for drivers to
-    * calculate stride for stream-out outputs.  Not externally visible.
-    */
-   SYSTEM_VALUE_VERTEX_CNT,
-
-   /**
-    * Required for AMD_shader_explicit_vertex_parameter and also used for
-    * varying-fetch instructions.
-    *
-    * The _SIZE value is "primitive size", used to scale i/j in primitive
-    * space to pixel space.
-    */
-   SYSTEM_VALUE_BARYCENTRIC_PERSP_PIXEL,
-   SYSTEM_VALUE_BARYCENTRIC_PERSP_SAMPLE,
-   SYSTEM_VALUE_BARYCENTRIC_PERSP_CENTROID,
-   SYSTEM_VALUE_BARYCENTRIC_PERSP_SIZE,
-   SYSTEM_VALUE_BARYCENTRIC_LINEAR_PIXEL,
-   SYSTEM_VALUE_BARYCENTRIC_LINEAR_CENTROID,
-   SYSTEM_VALUE_BARYCENTRIC_LINEAR_SAMPLE,
-   SYSTEM_VALUE_BARYCENTRIC_PULL_MODEL,
-
-   /**
-    * IR3 specific geometry shader and tesselation control shader system
-    * values that packs invocation id, thread id and vertex id.  Having this
-    * as a nir level system value lets us do the unpacking in nir.
-    */
-   SYSTEM_VALUE_GS_HEADER_IR3,
-   SYSTEM_VALUE_TCS_HEADER_IR3,
-
-   SYSTEM_VALUE_MAX             /**< Number of values */
-} gl_system_value;
-
-const char *gl_system_value_name(gl_system_value sysval);
-
-/**
- * The possible interpolation qualifiers that can be applied to a fragment
- * shader input in GLSL.
- *
- * Note: INTERP_MODE_NONE must be 0 so that memsetting the
- * ir_variable data structure to 0 causes the default behavior.
- */
-enum glsl_interp_mode
-{
-   INTERP_MODE_NONE = 0,
-   INTERP_MODE_SMOOTH,
-   INTERP_MODE_FLAT,
-   INTERP_MODE_NOPERSPECTIVE,
-   INTERP_MODE_EXPLICIT,
-   INTERP_MODE_COUNT /**< Number of interpolation qualifiers */
-};
-
-enum glsl_interface_packing {
-   GLSL_INTERFACE_PACKING_STD140,
-   GLSL_INTERFACE_PACKING_SHARED,
-   GLSL_INTERFACE_PACKING_PACKED,
-   GLSL_INTERFACE_PACKING_STD430
-};
-
-const char *glsl_interp_mode_name(enum glsl_interp_mode qual);
-
-/**
- * Fragment program results
- */
-typedef enum
-{
-   FRAG_RESULT_DEPTH = 0,
-   FRAG_RESULT_STENCIL = 1,
-   /* If a single color should be written to all render targets, this
-    * register is written.  No FRAG_RESULT_DATAn will be written.
-    */
-   FRAG_RESULT_COLOR = 2,
-   FRAG_RESULT_SAMPLE_MASK = 3,
-
-   /* FRAG_RESULT_DATAn are the per-render-target (GLSL gl_FragData[n]
-    * or ARB_fragment_program fragment.color[n]) color results.  If
-    * any are written, FRAG_RESULT_COLOR will not be written.
-    * FRAG_RESULT_DATA1 and up are simply for the benefit of
-    * gl_frag_result_name() and not to be construed as an upper bound
-    */
-   FRAG_RESULT_DATA0 = 4,
-   FRAG_RESULT_DATA1,
-   FRAG_RESULT_DATA2,
-   FRAG_RESULT_DATA3,
-   FRAG_RESULT_DATA4,
-   FRAG_RESULT_DATA5,
-   FRAG_RESULT_DATA6,
-   FRAG_RESULT_DATA7,
-} gl_frag_result;
-
-const char *gl_frag_result_name(gl_frag_result result);
-
-#define FRAG_RESULT_MAX		(FRAG_RESULT_DATA0 + MAX_DRAW_BUFFERS)
-
-/**
- * \brief Layout qualifiers for gl_FragDepth.
- *
- * Extension AMD_conservative_depth allows gl_FragDepth to be redeclared with
- * a layout qualifier.
- *
- * \see enum ir_depth_layout
- */
-enum gl_frag_depth_layout
-{
-   FRAG_DEPTH_LAYOUT_NONE, /**< No layout is specified. */
-   FRAG_DEPTH_LAYOUT_ANY,
-   FRAG_DEPTH_LAYOUT_GREATER,
-   FRAG_DEPTH_LAYOUT_LESS,
-   FRAG_DEPTH_LAYOUT_UNCHANGED
-};
-
-/**
- * \brief Buffer access qualifiers
- */
-enum gl_access_qualifier
-{
-   ACCESS_COHERENT      = (1 << 0),
-   ACCESS_RESTRICT      = (1 << 1),
-   ACCESS_VOLATILE      = (1 << 2),
-   ACCESS_NON_READABLE  = (1 << 3),
-   ACCESS_NON_WRITEABLE = (1 << 4),
-
-   /** The access may use a non-uniform buffer or image index */
-   ACCESS_NON_UNIFORM   = (1 << 5),
-
-   /* This has the same semantics as NIR_INTRINSIC_CAN_REORDER, only to be
-    * used with loads. In other words, it means that the load can be
-    * arbitrarily reordered, or combined with other loads to the same address.
-    * It is implied by ACCESS_NON_WRITEABLE together with ACCESS_RESTRICT, and
-    * a lack of ACCESS_COHERENT and ACCESS_VOLATILE.
-    */
-   ACCESS_CAN_REORDER = (1 << 6),
-
-   /** Use as little cache space as possible. */
-   ACCESS_STREAM_CACHE_POLICY = (1 << 7),
-};
-
-/**
- * \brief Blend support qualifiers
- */
-enum gl_advanced_blend_mode
-{
-   BLEND_NONE           = 0x0000,
-
-   BLEND_MULTIPLY       = 0x0001,
-   BLEND_SCREEN         = 0x0002,
-   BLEND_OVERLAY        = 0x0004,
-   BLEND_DARKEN         = 0x0008,
-   BLEND_LIGHTEN        = 0x0010,
-   BLEND_COLORDODGE     = 0x0020,
-   BLEND_COLORBURN      = 0x0040,
-   BLEND_HARDLIGHT      = 0x0080,
-   BLEND_SOFTLIGHT      = 0x0100,
-   BLEND_DIFFERENCE     = 0x0200,
-   BLEND_EXCLUSION      = 0x0400,
-   BLEND_HSL_HUE        = 0x0800,
-   BLEND_HSL_SATURATION = 0x1000,
-   BLEND_HSL_COLOR      = 0x2000,
-   BLEND_HSL_LUMINOSITY = 0x4000,
-
-   BLEND_ALL            = 0x7fff,
-};
-
-enum blend_func
-{
-   BLEND_FUNC_ADD,
-   BLEND_FUNC_SUBTRACT,
-   BLEND_FUNC_REVERSE_SUBTRACT,
-   BLEND_FUNC_MIN,
-   BLEND_FUNC_MAX,
-};
-
-enum blend_factor
-{
-   BLEND_FACTOR_ZERO,
-   BLEND_FACTOR_SRC_COLOR,
-   BLEND_FACTOR_DST_COLOR,
-   BLEND_FACTOR_SRC_ALPHA,
-   BLEND_FACTOR_DST_ALPHA,
-   BLEND_FACTOR_CONSTANT_COLOR,
-   BLEND_FACTOR_CONSTANT_ALPHA,
-   BLEND_FACTOR_SRC_ALPHA_SATURATE,
-};
-
-enum gl_tess_spacing
-{
-   TESS_SPACING_UNSPECIFIED,
-   TESS_SPACING_EQUAL,
-   TESS_SPACING_FRACTIONAL_ODD,
-   TESS_SPACING_FRACTIONAL_EVEN,
-};
-
-/**
- * A compare function enum for use in compiler lowering passes.  This is in
- * the same order as GL's compare functions (shifted down by GL_NEVER), and is
- * exactly the same as gallium's PIPE_FUNC_*.
- */
-enum compare_func
-{
-   COMPARE_FUNC_NEVER,
-   COMPARE_FUNC_LESS,
-   COMPARE_FUNC_EQUAL,
-   COMPARE_FUNC_LEQUAL,
-   COMPARE_FUNC_GREATER,
-   COMPARE_FUNC_NOTEQUAL,
-   COMPARE_FUNC_GEQUAL,
-   COMPARE_FUNC_ALWAYS,
-};
-
-/**
- * Arrangements for grouping invocations from NV_compute_shader_derivatives.
- *
- *   The extension provides new layout qualifiers that support two different
- *   arrangements of compute shader invocations for the purpose of derivative
- *   computation.  When specifying
- *
- *     layout(derivative_group_quadsNV) in;
- *
- *   compute shader invocations are grouped into 2x2x1 arrays whose four local
- *   invocation ID values follow the pattern:
- *
- *       +-----------------+------------------+
- *       | (2x+0, 2y+0, z) |  (2x+1, 2y+0, z) |
- *       +-----------------+------------------+
- *       | (2x+0, 2y+1, z) |  (2x+1, 2y+1, z) |
- *       +-----------------+------------------+
- *
- *   where Y increases from bottom to top.  When specifying
- *
- *     layout(derivative_group_linearNV) in;
- *
- *   compute shader invocations are grouped into 2x2x1 arrays whose four local
- *   invocation index values follow the pattern:
- *
- *       +------+------+
- *       | 4n+0 | 4n+1 |
- *       +------+------+
- *       | 4n+2 | 4n+3 |
- *       +------+------+
- *
- *   If neither layout qualifier is specified, derivatives in compute shaders
- *   return zero, which is consistent with the handling of built-in texture
- *   functions like texture() in GLSL 4.50 compute shaders.
- */
-enum gl_derivative_group {
-   DERIVATIVE_GROUP_NONE = 0,
-   DERIVATIVE_GROUP_QUADS,
-   DERIVATIVE_GROUP_LINEAR,
-};
-
-enum float_controls
-{
-   FLOAT_CONTROLS_DEFAULT_FLOAT_CONTROL_MODE        = 0x0000,
-   FLOAT_CONTROLS_DENORM_PRESERVE_FP16              = 0x0001,
-   FLOAT_CONTROLS_DENORM_PRESERVE_FP32              = 0x0002,
-   FLOAT_CONTROLS_DENORM_PRESERVE_FP64              = 0x0004,
-   FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP16         = 0x0008,
-   FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP32         = 0x0010,
-   FLOAT_CONTROLS_DENORM_FLUSH_TO_ZERO_FP64         = 0x0020,
-   FLOAT_CONTROLS_SIGNED_ZERO_INF_NAN_PRESERVE_FP16 = 0x0040,
-   FLOAT_CONTROLS_SIGNED_ZERO_INF_NAN_PRESERVE_FP32 = 0x0080,
-   FLOAT_CONTROLS_SIGNED_ZERO_INF_NAN_PRESERVE_FP64 = 0x0100,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTE_FP16            = 0x0200,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTE_FP32            = 0x0400,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTE_FP64            = 0x0800,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTZ_FP16            = 0x1000,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTZ_FP32            = 0x2000,
-   FLOAT_CONTROLS_ROUNDING_MODE_RTZ_FP64            = 0x4000,
-};
-
-#ifdef __cplusplus
-} /* extern "C" */
-#endif
-
-#endif /* SHADER_ENUMS_H */

extra/disassemblers/adreno/util/bitscan.h

@@ -1,326 +0,0 @@
-/**************************************************************************
- *
- * Copyright 2008 VMware, Inc.
- * All Rights Reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the
- * "Software"), to deal in the Software without restriction, including
- * without limitation the rights to use, copy, modify, merge, publish,
- * distribute, sub license, and/or sell copies of the Software, and to
- * permit persons to whom the Software is furnished to do so, subject to
- * the following conditions:
- *
- * The above copyright notice and this permission notice (including the
- * next paragraph) shall be included in all copies or substantial portions
- * of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
- * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
- * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT.
- * IN NO EVENT SHALL VMWARE AND/OR ITS SUPPLIERS BE LIABLE FOR
- * ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
- * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
- * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
- *
- **************************************************************************/
-
-
-#ifndef BITSCAN_H
-#define BITSCAN_H
-
-#include <assert.h>
-#include <stdint.h>
-#include <stdbool.h>
-#include <string.h>
-
-#if defined(_MSC_VER)
-#include <intrin.h>
-#endif
-
-#if defined(__POPCNT__)
-#include <popcntintrin.h>
-#endif
-
-//#include "c99_compat.h"
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-
-/**
- * Find first bit set in word.  Least significant bit is 1.
- * Return 0 if no bits set.
- */
-#ifdef HAVE___BUILTIN_FFS
-#define ffs __builtin_ffs
-#elif defined(_MSC_VER) && (_M_IX86 || _M_ARM || _M_AMD64 || _M_IA64)
-static inline
-int ffs(int i)
-{
-   unsigned long index;
-   if (_BitScanForward(&index, i))
-      return index + 1;
-   else
-      return 0;
-}
-#else
-extern
-int ffs(int i);
-#endif
-
-#ifdef HAVE___BUILTIN_FFSLL
-#define ffsll __builtin_ffsll
-#elif defined(_MSC_VER) && (_M_AMD64 || _M_ARM64 || _M_IA64)
-static inline int
-ffsll(long long int i)
-{
-   unsigned long index;
-   if (_BitScanForward64(&index, i))
-      return index + 1;
-   else
-      return 0;
-}
-#else
-extern int
-ffsll(long long int val);
-#endif
-
-
-/* Destructively loop over all of the bits in a mask as in:
- *
- * while (mymask) {
- *   int i = u_bit_scan(&mymask);
- *   ... process element i
- * }
- *
- */
-static inline int
-u_bit_scan(unsigned *mask)
-{
-   const int i = ffs(*mask) - 1;
-   *mask ^= (1u << i);
-   return i;
-}
-
-static inline int
-u_bit_scan64(uint64_t *mask)
-{
-   const int i = ffsll(*mask) - 1;
-   *mask ^= (((uint64_t)1) << i);
-   return i;
-}
-
-/* Determine if an unsigned value is a power of two.
- *
- * \note
- * Zero is treated as a power of two.
- */
-static inline bool
-util_is_power_of_two_or_zero(unsigned v)
-{
-   return (v & (v - 1)) == 0;
-}
-
-/* Determine if an uint64_t value is a power of two.
- *
- * \note
- * Zero is treated as a power of two.
- */
-static inline bool
-util_is_power_of_two_or_zero64(uint64_t v)
-{
-   return (v & (v - 1)) == 0;
-}
-
-/* Determine if an unsigned value is a power of two.
- *
- * \note
- * Zero is \b not treated as a power of two.
- */
-static inline bool
-util_is_power_of_two_nonzero(unsigned v)
-{
-   /* __POPCNT__ is different from HAVE___BUILTIN_POPCOUNT.  The latter
-    * indicates the existence of the __builtin_popcount function.  The former
-    * indicates that _mm_popcnt_u32 exists and is a native instruction.
-    *
-    * The other alternative is to use SSE 4.2 compile-time flags.  This has
-    * two drawbacks.  First, there is currently no build infrastructure for
-    * SSE 4.2 (only 4.1), so that would have to be added.  Second, some AMD
-    * CPUs support POPCNT but not SSE 4.2 (e.g., Barcelona).
-    */
-#ifdef __POPCNT__
-   return _mm_popcnt_u32(v) == 1;
-#else
-   return v != 0 && (v & (v - 1)) == 0;
-#endif
-}
-
-/* For looping over a bitmask when you want to loop over consecutive bits
- * manually, for example:
- *
- * while (mask) {
- *    int start, count, i;
- *
- *    u_bit_scan_consecutive_range(&mask, &start, &count);
- *
- *    for (i = 0; i < count; i++)
- *       ... process element (start+i)
- * }
- */
-static inline void
-u_bit_scan_consecutive_range(unsigned *mask, int *start, int *count)
-{
-   if (*mask == 0xffffffff) {
-      *start = 0;
-      *count = 32;
-      *mask = 0;
-      return;
-   }
-   *start = ffs(*mask) - 1;
-   *count = ffs(~(*mask >> *start)) - 1;
-   *mask &= ~(((1u << *count) - 1) << *start);
-}
-
-static inline void
-u_bit_scan_consecutive_range64(uint64_t *mask, int *start, int *count)
-{
-   if (*mask == ~0ull) {
-      *start = 0;
-      *count = 64;
-      *mask = 0;
-      return;
-   }
-   *start = ffsll(*mask) - 1;
-   *count = ffsll(~(*mask >> *start)) - 1;
-   *mask &= ~(((((uint64_t)1) << *count) - 1) << *start);
-}
-
-
-/**
- * Find last bit set in a word.  The least significant bit is 1.
- * Return 0 if no bits are set.
- * Essentially ffs() in the reverse direction.
- */
-static inline unsigned
-util_last_bit(unsigned u)
-{
-#if defined(HAVE___BUILTIN_CLZ)
-   return u == 0 ? 0 : 32 - __builtin_clz(u);
-#elif defined(_MSC_VER) && (_M_IX86 || _M_ARM || _M_AMD64 || _M_IA64)
-   unsigned long index;
-   if (_BitScanReverse(&index, u))
-      return index + 1;
-   else
-      return 0;
-#else
-   unsigned r = 0;
-   while (u) {
-      r++;
-      u >>= 1;
-   }
-   return r;
-#endif
-}
-
-/**
- * Find last bit set in a word.  The least significant bit is 1.
- * Return 0 if no bits are set.
- * Essentially ffsll() in the reverse direction.
- */
-static inline unsigned
-util_last_bit64(uint64_t u)
-{
-#if defined(HAVE___BUILTIN_CLZLL)
-   return u == 0 ? 0 : 64 - __builtin_clzll(u);
-#elif defined(_MSC_VER) && (_M_AMD64 || _M_ARM64 || _M_IA64)
-   unsigned long index;
-   if (_BitScanReverse64(&index, u))
-      return index + 1;
-   else
-      return 0;
-#else
-   unsigned r = 0;
-   while (u) {
-      r++;
-      u >>= 1;
-   }
-   return r;
-#endif
-}
-
-/**
- * Find last bit in a word that does not match the sign bit. The least
- * significant bit is 1.
- * Return 0 if no bits are set.
- */
-static inline unsigned
-util_last_bit_signed(int i)
-{
-   if (i >= 0)
-      return util_last_bit(i);
-   else
-      return util_last_bit(~(unsigned)i);
-}
-
-/* Returns a bitfield in which the first count bits starting at start are
- * set.
- */
-static inline unsigned
-u_bit_consecutive(unsigned start, unsigned count)
-{
-   assert(start + count <= 32);
-   if (count == 32)
-      return ~0;
-   return ((1u << count) - 1) << start;
-}
-
-static inline uint64_t
-u_bit_consecutive64(unsigned start, unsigned count)
-{
-   assert(start + count <= 64);
-   if (count == 64)
-      return ~(uint64_t)0;
-   return (((uint64_t)1 << count) - 1) << start;
-}
-
-/**
- * Return number of bits set in n.
- */
-static inline unsigned
-util_bitcount(unsigned n)
-{
-#if defined(HAVE___BUILTIN_POPCOUNT)
-   return __builtin_popcount(n);
-#else
-   /* K&R classic bitcount.
-    *
-    * For each iteration, clear the LSB from the bitfield.
-    * Requires only one iteration per set bit, instead of
-    * one iteration per bit less than highest set bit.
-    */
-   unsigned bits;
-   for (bits = 0; n; bits++) {
-      n &= n - 1;
-   }
-   return bits;
-#endif
-}
-
-static inline unsigned
-util_bitcount64(uint64_t n)
-{
-#ifdef HAVE___BUILTIN_POPCOUNTLL
-   return __builtin_popcountll(n);
-#else
-   return util_bitcount(n) + util_bitcount(n >> 32);
-#endif
-}
-
-#ifdef __cplusplus
-}
-#endif
-
-#endif /* BITSCAN_H */

extra/disassemblers/adreno/util/bitset.h

@@ -1,261 +0,0 @@
-/*
- * Mesa 3-D graphics library
- *
- * Copyright (C) 2006  Brian Paul   All Rights Reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included
- * in all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
- * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- */
-
-/**
- * \file bitset.h
- * \brief Bitset of arbitrary size definitions.
- * \author Michal Krol
- */
-
-#ifndef BITSET_H
-#define BITSET_H
-
-//#include "util/bitscan.h"
-//#include "util/macros.h"
-
-/****************************************************************************
- * generic bitset implementation
- */
-
-#define BITSET_WORD unsigned int
-#define BITSET_WORDBITS (sizeof (BITSET_WORD) * 8)
-
-/* bitset declarations
- */
-#define BITSET_WORDS(bits) (((bits) + BITSET_WORDBITS - 1) / BITSET_WORDBITS)
-#define BITSET_DECLARE(name, bits) BITSET_WORD name[BITSET_WORDS(bits)]
-
-/* bitset operations
- */
-#define BITSET_COPY(x, y) memcpy( (x), (y), sizeof (x) )
-#define BITSET_EQUAL(x, y) (memcmp( (x), (y), sizeof (x) ) == 0)
-#define BITSET_ZERO(x) memset( (x), 0, sizeof (x) )
-#define BITSET_ONES(x) memset( (x), 0xff, sizeof (x) )
-
-#define BITSET_BITWORD(b) ((b) / BITSET_WORDBITS)
-#define BITSET_BIT(b) (1u << ((b) % BITSET_WORDBITS))
-
-/* single bit operations
- */
-#define BITSET_TEST(x, b) (((x)[BITSET_BITWORD(b)] & BITSET_BIT(b)) != 0)
-#define BITSET_SET(x, b) ((x)[BITSET_BITWORD(b)] |= BITSET_BIT(b))
-#define BITSET_CLEAR(x, b) ((x)[BITSET_BITWORD(b)] &= ~BITSET_BIT(b))
-
-#define BITSET_MASK(b) (((b) % BITSET_WORDBITS == 0) ? ~0 : BITSET_BIT(b) - 1)
-#define BITSET_RANGE(b, e) ((BITSET_MASK((e) + 1)) & ~(BITSET_BIT(b) - 1))
-
-/* bit range operations
- */
-#define BITSET_TEST_RANGE(x, b, e) \
-   (BITSET_BITWORD(b) == BITSET_BITWORD(e) ? \
-   (((x)[BITSET_BITWORD(b)] & BITSET_RANGE(b, e)) != 0) : \
-   (assert (!"BITSET_TEST_RANGE: bit range crosses word boundary"), 0))
-#define BITSET_SET_RANGE(x, b, e) \
-   (BITSET_BITWORD(b) == BITSET_BITWORD(e) ? \
-   ((x)[BITSET_BITWORD(b)] |= BITSET_RANGE(b, e)) : \
-   (assert (!"BITSET_SET_RANGE: bit range crosses word boundary"), 0))
-#define BITSET_CLEAR_RANGE(x, b, e) \
-   (BITSET_BITWORD(b) == BITSET_BITWORD(e) ? \
-   ((x)[BITSET_BITWORD(b)] &= ~BITSET_RANGE(b, e)) : \
-   (assert (!"BITSET_CLEAR_RANGE: bit range crosses word boundary"), 0))
-
-/* Get first bit set in a bitset.
- */
-static inline int
-__bitset_ffs(const BITSET_WORD *x, int n)
-{
-   int i;
-
-   for (i = 0; i < n; i++) {
-      if (x[i])
-	 return ffs(x[i]) + BITSET_WORDBITS * i;
-   }
-
-   return 0;
-}
-
-#define BITSET_FFS(x) __bitset_ffs(x, ARRAY_SIZE(x))
-
-static inline unsigned
-__bitset_next_set(unsigned i, BITSET_WORD *tmp,
-                  const BITSET_WORD *set, unsigned size)
-{
-   unsigned bit, word;
-
-   /* NOTE: The initial conditions for this function are very specific.  At
-    * the start of the loop, the tmp variable must be set to *set and the
-    * initial i value set to 0.  This way, if there is a bit set in the first
-    * word, we ignore the i-value and just grab that bit (so 0 is ok, even
-    * though 0 may be returned).  If the first word is 0, then the value of
-    * `word` will be 0 and we will go on to look at the second word.
-    */
-   word = BITSET_BITWORD(i);
-   while (*tmp == 0) {
-      word++;
-
-      if (word >= BITSET_WORDS(size))
-         return size;
-
-      *tmp = set[word];
-   }
-
-   /* Find the next set bit in the non-zero word */
-   bit = ffs(*tmp) - 1;
-
-   /* Unset the bit */
-   *tmp &= ~(1ull << bit);
-
-   return word * BITSET_WORDBITS + bit;
-}
-
-/**
- * Iterates over each set bit in a set
- *
- * @param __i    iteration variable, bit number
- * @param __set  the bitset to iterate (will not be modified)
- * @param __size number of bits in the set to consider
- */
-#define BITSET_FOREACH_SET(__i, __set, __size) \
-   for (BITSET_WORD __tmp = *(__set), *__foo = &__tmp; __foo != NULL; __foo = NULL) \
-      for (__i = 0; \
-           (__i = __bitset_next_set(__i, &__tmp, __set, __size)) < __size;)
-
-#ifdef __cplusplus
-
-/**
- * Simple C++ wrapper of a bitset type of static size, with value semantics
- * and basic bitwise arithmetic operators.  The operators defined below are
- * expected to have the same semantics as the same operator applied to other
- * fundamental integer types.  T is the name of the struct to instantiate
- * it as, and N is the number of bits in the bitset.
- */
-#define DECLARE_BITSET_T(T, N) struct T {                       \
-      EXPLICIT_CONVERSION                                       \
-      operator bool() const                                     \
-      {                                                         \
-         for (unsigned i = 0; i < BITSET_WORDS(N); i++)         \
-            if (words[i])                                       \
-               return true;                                     \
-         return false;                                          \
-      }                                                         \
-                                                                \
-      T &                                                       \
-      operator=(int x)                                          \
-      {                                                         \
-         const T c = {{ (BITSET_WORD)x }};                      \
-         return *this = c;                                      \
-      }                                                         \
-                                                                \
-      friend bool                                               \
-      operator==(const T &b, const T &c)                        \
-      {                                                         \
-         return BITSET_EQUAL(b.words, c.words);                 \
-      }                                                         \
-                                                                \
-      friend bool                                               \
-      operator!=(const T &b, const T &c)                        \
-      {                                                         \
-         return !(b == c);                                      \
-      }                                                         \
-                                                                \
-      friend bool                                               \
-      operator==(const T &b, int x)                             \
-      {                                                         \
-         const T c = {{ (BITSET_WORD)x }};                      \
-         return b == c;                                         \
-      }                                                         \
-                                                                \
-      friend bool                                               \
-      operator!=(const T &b, int x)                             \
-      {                                                         \
-         return !(b == x);                                      \
-      }                                                         \
-                                                                \
-      friend T                                                  \
-      operator~(const T &b)                                     \
-      {                                                         \
-         T c;                                                   \
-         for (unsigned i = 0; i < BITSET_WORDS(N); i++)         \
-            c.words[i] = ~b.words[i];                           \
-         return c;                                              \
-      }                                                         \
-                                                                \
-      T &                                                       \
-      operator|=(const T &b)                                    \
-      {                                                         \
-         for (unsigned i = 0; i < BITSET_WORDS(N); i++)         \
-            words[i] |= b.words[i];                             \
-         return *this;                                          \
-      }                                                         \
-                                                                \
-      friend T                                                  \
-      operator|(const T &b, const T &c)                         \
-      {                                                         \
-         T d = b;                                               \
-         d |= c;                                                \
-         return d;                                              \
-      }                                                         \
-                                                                \
-      T &                                                       \
-      operator&=(const T &b)                                    \
-      {                                                         \
-         for (unsigned i = 0; i < BITSET_WORDS(N); i++)         \
-            words[i] &= b.words[i];                             \
-         return *this;                                          \
-      }                                                         \
-                                                                \
-      friend T                                                  \
-      operator&(const T &b, const T &c)                         \
-      {                                                         \
-         T d = b;                                               \
-         d &= c;                                                \
-         return d;                                              \
-      }                                                         \
-                                                                \
-      bool                                                      \
-      test(unsigned i) const                                    \
-      {                                                         \
-         return BITSET_TEST(words, i);                          \
-      }                                                         \
-                                                                \
-      T &                                                       \
-      set(unsigned i)                                           \
-      {                                                         \
-         BITSET_SET(words, i);                                  \
-         return *this;                                          \
-      }                                                         \
-                                                                \
-      T &                                                       \
-      clear(unsigned i)                                         \
-      {                                                         \
-         BITSET_CLEAR(words, i);                                \
-         return *this;                                          \
-      }                                                         \
-                                                                \
-      BITSET_WORD words[BITSET_WORDS(N)];                       \
-   }
-
-#endif
-
-#endif

extra/disassemblers/adreno/util/list.h

@@ -1,262 +0,0 @@
-/**************************************************************************
- *
- * Copyright 2006 VMware, Inc., Bismarck, ND. USA.
- * All Rights Reserved.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the
- * "Software"), to deal in the Software without restriction, including
- * without limitation the rights to use, copy, modify, merge, publish,
- * distribute, sub license, and/or sell copies of the Software, and to
- * permit persons to whom the Software is furnished to do so, subject to
- * the following conditions:
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL
- * THE COPYRIGHT HOLDERS, AUTHORS AND/OR ITS SUPPLIERS BE LIABLE FOR ANY CLAIM,
- * DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
- * OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
- * USE OR OTHER DEALINGS IN THE SOFTWARE.
- *
- * The above copyright notice and this permission notice (including the
- * next paragraph) shall be included in all copies or substantial portions
- * of the Software.
- *
- **************************************************************************/
-
-/**
- * \file
- * List macros heavily inspired by the Linux kernel
- * list handling. No list looping yet.
- *
- * Is not threadsafe, so common operations need to
- * be protected using an external mutex.
- */
-
-#ifndef _UTIL_LIST_H_
-#define _UTIL_LIST_H_
-
-
-#include <stdbool.h>
-#include <stddef.h>
-#include <assert.h>
-
-#ifdef DEBUG
-#  define list_assert(cond, msg)  assert(cond && msg)
-#else
-#  define list_assert(cond, msg)  (void)(0 && (cond))
-#endif
-
-struct list_head
-{
-    struct list_head *prev;
-    struct list_head *next;
-};
-
-static inline void list_inithead(struct list_head *item)
-{
-    item->prev = item;
-    item->next = item;
-}
-
-static inline void list_add(struct list_head *item, struct list_head *list)
-{
-    item->prev = list;
-    item->next = list->next;
-    list->next->prev = item;
-    list->next = item;
-}
-
-static inline void list_addtail(struct list_head *item, struct list_head *list)
-{
-    item->next = list;
-    item->prev = list->prev;
-    list->prev->next = item;
-    list->prev = item;
-}
-
-static inline bool list_is_empty(const struct list_head *list);
-
-static inline void list_replace(struct list_head *from, struct list_head *to)
-{
-    if (list_is_empty(from)) {
-        list_inithead(to);
-    } else {
-        to->prev = from->prev;
-        to->next = from->next;
-        from->next->prev = to;
-        from->prev->next = to;
-    }
-}
-
-static inline void list_del(struct list_head *item)
-{
-    item->prev->next = item->next;
-    item->next->prev = item->prev;
-    item->prev = item->next = NULL;
-}
-
-static inline void list_delinit(struct list_head *item)
-{
-    item->prev->next = item->next;
-    item->next->prev = item->prev;
-    item->next = item;
-    item->prev = item;
-}
-
-static inline bool list_is_empty(const struct list_head *list)
-{
-   return list->next == list;
-}
-
-/**
- * Returns whether the list has exactly one element.
- */
-static inline bool list_is_singular(const struct list_head *list)
-{
-   return list->next != NULL && list->next != list && list->next->next == list;
-}
-
-static inline unsigned list_length(const struct list_head *list)
-{
-   struct list_head *node;
-   unsigned length = 0;
-   for (node = list->next; node != list; node = node->next)
-      length++;
-   return length;
-}
-
-static inline void list_splice(struct list_head *src, struct list_head *dst)
-{
-   if (list_is_empty(src))
-      return;
-
-   src->next->prev = dst;
-   src->prev->next = dst->next;
-   dst->next->prev = src->prev;
-   dst->next = src->next;
-}
-
-static inline void list_splicetail(struct list_head *src, struct list_head *dst)
-{
-   if (list_is_empty(src))
-      return;
-
-   src->prev->next = dst;
-   src->next->prev = dst->prev;
-   dst->prev->next = src->next;
-   dst->prev = src->prev;
-}
-
-static inline void list_validate(const struct list_head *list)
-{
-   struct list_head *node;
-   assert(list->next->prev == list && list->prev->next == list);
-   for (node = list->next; node != list; node = node->next)
-      assert(node->next->prev == node && node->prev->next == node);
-}
-
-#define LIST_ENTRY(__type, __item, __field)   \
-    ((__type *)(((char *)(__item)) - offsetof(__type, __field)))
-
-/**
- * Cast from a pointer to a member of a struct back to the containing struct.
- *
- * 'sample' MUST be initialized, or else the result is undefined!
- */
-#ifndef container_of
-#define container_of(ptr, sample, member)				\
-    (void *)((char *)(ptr)						\
-	     - ((char *)&(sample)->member - (char *)(sample)))
-#endif
-
-#define list_first_entry(ptr, type, member) \
-        LIST_ENTRY(type, (ptr)->next, member)
-
-#define list_last_entry(ptr, type, member) \
-        LIST_ENTRY(type, (ptr)->prev, member)
-
-
-#define LIST_FOR_EACH_ENTRY(pos, head, member)				\
-   for (pos = NULL, pos = container_of((head)->next, pos, member);	\
-	&pos->member != (head);						\
-	pos = container_of(pos->member.next, pos, member))
-
-#define LIST_FOR_EACH_ENTRY_SAFE(pos, storage, head, member)	\
-   for (pos = NULL, pos = container_of((head)->next, pos, member),	\
-	storage = container_of(pos->member.next, pos, member);	\
-	&pos->member != (head);						\
-	pos = storage, storage = container_of(storage->member.next, storage, member))
-
-#define LIST_FOR_EACH_ENTRY_SAFE_REV(pos, storage, head, member)	\
-   for (pos = NULL, pos = container_of((head)->prev, pos, member),	\
-	storage = container_of(pos->member.prev, pos, member);		\
-	&pos->member != (head);						\
-	pos = storage, storage = container_of(storage->member.prev, storage, member))
-
-#define LIST_FOR_EACH_ENTRY_FROM(pos, start, head, member)		\
-   for (pos = NULL, pos = container_of((start), pos, member);		\
-	&pos->member != (head);						\
-	pos = container_of(pos->member.next, pos, member))
-
-#define LIST_FOR_EACH_ENTRY_FROM_REV(pos, start, head, member)		\
-   for (pos = NULL, pos = container_of((start), pos, member);		\
-	&pos->member != (head);						\
-	pos = container_of(pos->member.prev, pos, member))
-
-#define list_for_each_entry(type, pos, head, member)                    \
-   for (type *pos = LIST_ENTRY(type, (head)->next, member),             \
-	     *__next = LIST_ENTRY(type, pos->member.next, member);      \
-	&pos->member != (head);                                         \
-	pos = LIST_ENTRY(type, pos->member.next, member),               \
-	list_assert(pos == __next, "use _safe iterator"),               \
-	__next = LIST_ENTRY(type, __next->member.next, member))
-
-#define list_for_each_entry_safe(type, pos, head, member)               \
-   for (type *pos = LIST_ENTRY(type, (head)->next, member),             \
-	     *__next = LIST_ENTRY(type, pos->member.next, member);      \
-	&pos->member != (head);                                         \
-	pos = __next,                                                   \
-	__next = LIST_ENTRY(type, __next->member.next, member))
-
-#define list_for_each_entry_rev(type, pos, head, member)                \
-   for (type *pos = LIST_ENTRY(type, (head)->prev, member),             \
-	     *__prev = LIST_ENTRY(type, pos->member.prev, member);      \
-	&pos->member != (head);                                         \
-	pos = LIST_ENTRY(type, pos->member.prev, member),               \
-	list_assert(pos == __prev, "use _safe iterator"),               \
-	__prev = LIST_ENTRY(type, __prev->member.prev, member))
-
-#define list_for_each_entry_safe_rev(type, pos, head, member)           \
-   for (type *pos = LIST_ENTRY(type, (head)->prev, member),             \
-	     *__prev = LIST_ENTRY(type, pos->member.prev, member);      \
-	&pos->member != (head);                                         \
-	pos = __prev,                                                   \
-        __prev = LIST_ENTRY(type, __prev->member.prev, member))
-
-#define list_for_each_entry_from(type, pos, start, head, member)        \
-   for (type *pos = LIST_ENTRY(type, (start), member);                  \
-	&pos->member != (head);                                         \
-	pos = LIST_ENTRY(type, pos->member.next, member))
-
-#define list_for_each_entry_from_safe(type, pos, start, head, member)   \
-   for (type *pos = LIST_ENTRY(type, (start), member),                  \
-	     *__next = LIST_ENTRY(type, pos->member.next, member);      \
-	&pos->member != (head);                                         \
-	pos = __next,                                                   \
-	__next = LIST_ENTRY(type, __next->member.next, member))
-
-#define list_for_each_entry_from_rev(type, pos, start, head, member)    \
-   for (type *pos = LIST_ENTRY(type, (start), member);                  \
-	&pos->member != (head);                                         \
-	pos = LIST_ENTRY(type, pos->member.prev, member))
-
-#define list_pair_for_each_entry(type, pos1, pos2, head1, head2, member) \
-   for (type *pos1 = LIST_ENTRY(type, (head1)->next, member),           \
-             *pos2 = LIST_ENTRY(type, (head2)->next, member);           \
-        &pos1->member != (head1) && &pos2->member != (head2);           \
-	pos1 = LIST_ENTRY(type, pos1->member.next, member),               \
-	pos2 = LIST_ENTRY(type, pos2->member.next, member))
-
-#endif /*_UTIL_LIST_H_*/

extra/disassemblers/adreno/util/macros.h

@@ -1,346 +0,0 @@
-/*
- * Copyright © 2014 Intel Corporation
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice (including the next
- * paragraph) shall be included in all copies or substantial portions of the
- * Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
- * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
- * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
- * IN THE SOFTWARE.
- */
-
-#ifndef UTIL_MACROS_H
-#define UTIL_MACROS_H
-
-#include <assert.h>
-
-/* Compute the size of an array */
-#ifndef ARRAY_SIZE
-#  define ARRAY_SIZE(x) (sizeof(x) / sizeof((x)[0]))
-#endif
-
-/* For compatibility with Clang's __has_builtin() */
-#ifndef __has_builtin
-#  define __has_builtin(x) 0
-#endif
-
-/**
- * __builtin_expect macros
- */
-#if !defined(HAVE___BUILTIN_EXPECT)
-#  define __builtin_expect(x, y) (x)
-#endif
-
-#ifndef likely
-#  ifdef HAVE___BUILTIN_EXPECT
-#    define likely(x)   __builtin_expect(!!(x), 1)
-#    define unlikely(x) __builtin_expect(!!(x), 0)
-#  else
-#    define likely(x)   (x)
-#    define unlikely(x) (x)
-#  endif
-#endif
-
-
-/**
- * Static (compile-time) assertion.
- * Basically, use COND to dimension an array.  If COND is false/zero the
- * array size will be -1 and we'll get a compilation error.
- */
-#define STATIC_ASSERT(COND) \
-   do { \
-      (void) sizeof(char [1 - 2*!(COND)]); \
-   } while (0)
-
-
-/**
- * Unreachable macro. Useful for suppressing "control reaches end of non-void
- * function" warnings.
- */
-#if defined(HAVE___BUILTIN_UNREACHABLE) || __has_builtin(__builtin_unreachable)
-#define unreachable(str)    \
-do {                        \
-   assert(!str);            \
-   __builtin_unreachable(); \
-} while (0)
-#elif defined (_MSC_VER)
-#define unreachable(str)    \
-do {                        \
-   assert(!str);            \
-   __assume(0);             \
-} while (0)
-#else
-#define unreachable(str) assert(!str)
-#endif
-
-/**
- * Assume macro. Useful for expressing our assumptions to the compiler,
- * typically for purposes of silencing warnings.
- */
-#if __has_builtin(__builtin_assume)
-#define assume(expr)       \
-do {                       \
-   assert(expr);           \
-   __builtin_assume(expr); \
-} while (0)
-#elif defined HAVE___BUILTIN_UNREACHABLE
-#define assume(expr) ((expr) ? ((void) 0) \
-                             : (assert(!"assumption failed"), \
-                                __builtin_unreachable()))
-#elif defined (_MSC_VER)
-#define assume(expr) __assume(expr)
-#else
-#define assume(expr) assert(expr)
-#endif
-
-/* Attribute const is used for functions that have no effects other than their
- * return value, and only rely on the argument values to compute the return
- * value.  As a result, calls to it can be CSEed.  Note that using memory
- * pointed to by the arguments is not allowed for const functions.
- */
-#ifdef HAVE_FUNC_ATTRIBUTE_CONST
-#define ATTRIBUTE_CONST __attribute__((__const__))
-#else
-#define ATTRIBUTE_CONST
-#endif
-
-#ifdef HAVE_FUNC_ATTRIBUTE_FLATTEN
-#define FLATTEN __attribute__((__flatten__))
-#else
-#define FLATTEN
-#endif
-
-#ifdef HAVE_FUNC_ATTRIBUTE_FORMAT
-#define PRINTFLIKE(f, a) __attribute__ ((format(__printf__, f, a)))
-#else
-#define PRINTFLIKE(f, a)
-#endif
-
-#ifdef HAVE_FUNC_ATTRIBUTE_MALLOC
-#define MALLOCLIKE __attribute__((__malloc__))
-#else
-#define MALLOCLIKE
-#endif
-
-/* Forced function inlining */
-/* Note: Clang also sets __GNUC__ (see other cases below) */
-#ifndef ALWAYS_INLINE
-#  if defined(__GNUC__)
-#    define ALWAYS_INLINE inline __attribute__((always_inline))
-#  elif defined(_MSC_VER)
-#    define ALWAYS_INLINE __forceinline
-#  else
-#    define ALWAYS_INLINE inline
-#  endif
-#endif
-
-/* Used to optionally mark structures with misaligned elements or size as
- * packed, to trade off performance for space.
- */
-#ifdef HAVE_FUNC_ATTRIBUTE_PACKED
-#define PACKED __attribute__((__packed__))
-#else
-#define PACKED
-#endif
-
-/* Attribute pure is used for functions that have no effects other than their
- * return value.  As a result, calls to it can be dead code eliminated.
- */
-#ifdef HAVE_FUNC_ATTRIBUTE_PURE
-#define ATTRIBUTE_PURE __attribute__((__pure__))
-#else
-#define ATTRIBUTE_PURE
-#endif
-
-#ifdef HAVE_FUNC_ATTRIBUTE_RETURNS_NONNULL
-#define ATTRIBUTE_RETURNS_NONNULL __attribute__((__returns_nonnull__))
-#else
-#define ATTRIBUTE_RETURNS_NONNULL
-#endif
-
-#ifndef NORETURN
-#  ifdef _MSC_VER
-#    define NORETURN __declspec(noreturn)
-#  elif defined HAVE_FUNC_ATTRIBUTE_NORETURN
-#    define NORETURN __attribute__((__noreturn__))
-#  else
-#    define NORETURN
-#  endif
-#endif
-
-#ifdef __cplusplus
-/**
- * Macro function that evaluates to true if T is a trivially
- * destructible type -- that is, if its (non-virtual) destructor
- * performs no action and all member variables and base classes are
- * trivially destructible themselves.
- */
-#   if (defined(__clang__) && defined(__has_feature))
-#      if __has_feature(has_trivial_destructor)
-#         define HAS_TRIVIAL_DESTRUCTOR(T) __has_trivial_destructor(T)
-#      endif
-#   elif defined(__GNUC__)
-#      if ((__GNUC__ > 4) || ((__GNUC__ == 4) && (__GNUC_MINOR__ >= 3)))
-#         define HAS_TRIVIAL_DESTRUCTOR(T) __has_trivial_destructor(T)
-#      endif
-#   elif defined(_MSC_VER) && !defined(__INTEL_COMPILER)
-#      define HAS_TRIVIAL_DESTRUCTOR(T) __has_trivial_destructor(T)
-#   endif
-#   ifndef HAS_TRIVIAL_DESTRUCTOR
-       /* It's always safe (if inefficient) to assume that a
-        * destructor is non-trivial.
-        */
-#      define HAS_TRIVIAL_DESTRUCTOR(T) (false)
-#   endif
-#endif
-
-/**
- * PUBLIC/USED macros
- *
- * If we build the library with gcc's -fvisibility=hidden flag, we'll
- * use the PUBLIC macro to mark functions that are to be exported.
- *
- * We also need to define a USED attribute, so the optimizer doesn't
- * inline a static function that we later use in an alias. - ajax
- */
-#ifndef PUBLIC
-#  if defined(__GNUC__)
-#    define PUBLIC __attribute__((visibility("default")))
-#    define USED __attribute__((used))
-#  elif defined(_MSC_VER)
-#    define PUBLIC __declspec(dllexport)
-#    define USED
-#  else
-#    define PUBLIC
-#    define USED
-#  endif
-#endif
-
-/**
- * UNUSED marks variables (or sometimes functions) that have to be defined,
- * but are sometimes (or always) unused beyond that. A common case is for
- * a function parameter to be used in some build configurations but not others.
- * Another case is fallback vfuncs that don't do anything with their params.
- *
- * Note that this should not be used for identifiers used in `assert()`;
- * see ASSERTED below.
- */
-#ifdef HAVE_FUNC_ATTRIBUTE_UNUSED
-#define UNUSED __attribute__((unused))
-#else
-#define UNUSED
-#endif
-
-/**
- * Use ASSERTED to indicate that an identifier is unused outside of an `assert()`,
- * so that assert-free builds don't get "unused variable" warnings.
- */
-#ifdef NDEBUG
-#define ASSERTED UNUSED
-#else
-#define ASSERTED
-#endif
-
-#ifdef HAVE_FUNC_ATTRIBUTE_WARN_UNUSED_RESULT
-#define MUST_CHECK __attribute__((warn_unused_result))
-#else
-#define MUST_CHECK
-#endif
-
-#if defined(__GNUC__)
-#define ATTRIBUTE_NOINLINE __attribute__((noinline))
-#else
-#define ATTRIBUTE_NOINLINE
-#endif
-
-
-/**
- * Check that STRUCT::FIELD can hold MAXVAL.  We use a lot of bitfields
- * in Mesa/gallium.  We have to be sure they're of sufficient size to
- * hold the largest expected value.
- * Note that with MSVC, enums are signed and enum bitfields need one extra
- * high bit (always zero) to ensure the max value is handled correctly.
- * This macro will detect that with MSVC, but not GCC.
- */
-#define ASSERT_BITFIELD_SIZE(STRUCT, FIELD, MAXVAL) \
-   do { \
-      ASSERTED STRUCT s; \
-      s.FIELD = (MAXVAL); \
-      assert((int) s.FIELD == (MAXVAL) && "Insufficient bitfield size!"); \
-   } while (0)
-
-
-/** Compute ceiling of integer quotient of A divided by B. */
-#define DIV_ROUND_UP( A, B )  ( ((A) + (B) - 1) / (B) )
-
-/** Clamp X to [MIN,MAX].  Turn NaN into MIN, arbitrarily. */
-#define CLAMP( X, MIN, MAX )  ( (X)>(MIN) ? ((X)>(MAX) ? (MAX) : (X)) : (MIN) )
-
-/** Minimum of two values: */
-#define MIN2( A, B )   ( (A)<(B) ? (A) : (B) )
-
-/** Maximum of two values: */
-#define MAX2( A, B )   ( (A)>(B) ? (A) : (B) )
-
-/** Minimum and maximum of three values: */
-#define MIN3( A, B, C ) ((A) < (B) ? MIN2(A, C) : MIN2(B, C))
-#define MAX3( A, B, C ) ((A) > (B) ? MAX2(A, C) : MAX2(B, C))
-
-/** Align a value to a power of two */
-#define ALIGN_POT(x, pot_align) (((x) + (pot_align) - 1) & ~((pot_align) - 1))
-
-/**
- * Macro for declaring an explicit conversion operator.  Defaults to an
- * implicit conversion if C++11 is not supported.
- */
-#if __cplusplus >= 201103L
-#define EXPLICIT_CONVERSION explicit
-#elif defined(__cplusplus)
-#define EXPLICIT_CONVERSION
-#endif
-
-/** Set a single bit */
-#define BITFIELD_BIT(b)      (1u << (b))
-/** Set all bits up to excluding bit b */
-#define BITFIELD_MASK(b)      \
-   ((b) == 32 ? (~0u) : BITFIELD_BIT((b) % 32) - 1)
-/** Set count bits starting from bit b  */
-#define BITFIELD_RANGE(b, count) \
-   (BITFIELD_MASK((b) + (count)) & ~BITFIELD_MASK(b))
-
-/** Set a single bit */
-#define BITFIELD64_BIT(b)      (1ull << (b))
-/** Set all bits up to excluding bit b */
-#define BITFIELD64_MASK(b)      \
-   ((b) == 64 ? (~0ull) : BITFIELD64_BIT(b) - 1)
-/** Set count bits starting from bit b  */
-#define BITFIELD64_RANGE(b, count) \
-   (BITFIELD64_MASK((b) + (count)) & ~BITFIELD64_MASK(b))
-
-/* TODO: In future we should try to move this to u_debug.h once header
- * dependencies are reorganised to allow this.
- */
-enum pipe_debug_type
-{
-   PIPE_DEBUG_TYPE_OUT_OF_MEMORY = 1,
-   PIPE_DEBUG_TYPE_ERROR,
-   PIPE_DEBUG_TYPE_SHADER_INFO,
-   PIPE_DEBUG_TYPE_PERF_INFO,
-   PIPE_DEBUG_TYPE_INFO,
-   PIPE_DEBUG_TYPE_FALLBACK,
-   PIPE_DEBUG_TYPE_CONFORMANCE,
-};
-
-#endif /* UTIL_MACROS_H */

extra/disk_read_speed.py

@@ -1,120 +0,0 @@
-#!/usr/bin/env python3
-import os, ctypes, ctypes.util, io, mmap, pathlib
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import Timing, from_mv
-libc = ctypes.CDLL(ctypes.util.find_library("c"))
-
-#from extra.hip_gpu_driver import hip_ioctl
-
-# sudo su -c "echo 3 > /proc/sys/vm/drop_caches"
-
-# sudo su -c 'echo 8 > /proc/sys/kernel/printk'
-# sudo su -c "echo 'module amdgpu +p' > /sys/kernel/debug/dynamic_debug/control"
-
-libc.memcpy.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t]
-
-libc.read.argtypes = [ctypes.c_int, ctypes.c_void_p, ctypes.c_size_t]
-libc.read.restype = ctypes.c_size_t
-
-libc.malloc.argtypes = [ctypes.c_size_t]
-libc.malloc.restype = ctypes.c_void_p
-
-def read_direct(fd, sz):
-  with Timing("mmap: ", lambda x: f", {sz/x:.2f} GB/s"):
-    buf = mmap.mmap(-1, sz, flags=mmap.MAP_SHARED|mmap.MAP_POPULATE)
-  with Timing("read: ", lambda x: f", {sz/x:.2f} GB/s"):
-    ret = libc.read(fd, from_mv(buf), sz)
-  assert ret == sz
-
-def read_mmap(fd, sz):
-  with Timing("mmfd:       ", lambda x: f", {sz/x:.2f} GB/s"):
-    buf = mmap.mmap(fd, sz, flags=mmap.MAP_SHARED|mmap.MAP_POPULATE) #|MAP_LOCKED)
-    t = 0
-    for i in range(0, sz, 0x1000): t += buf[i]
-
-# def _copyin_async(self, dest:T, src:T, size:int): check(hip.hipMemcpyAsync(dest, src, size, hip.hipMemcpyHostToDevice, None))
-
-def read_to_gpu_mmap(fd, sz, gpubuf):
-  with Timing("gpu copyin: ", lambda x: f", {sz/x:.2f} GB/s"):
-    with Timing("mmfd:       ", lambda x: f", {sz/x:.2f} GB/s"):
-      buf = mmap.mmap(fd, sz, flags=mmap.MAP_SHARED|mmap.MAP_POPULATE) #|MAP_LOCKED)
-    dev.allocator._copyin_async(gpubuf, from_mv(buf), sz)
-    dev.synchronize()
-
-def read_to_gpu_single(fd, sz, gpubuf):
-  os.lseek(fd, 0, os.SEEK_SET)
-  with Timing("total: ", lambda x: f", {sz/x:.2f} GB/s"):
-    with Timing("gpu host alloc: ", lambda x: f", {sz/x:.2f} GB/s"):
-      hst = dev.allocator._hostalloc(sz)
-    with Timing("read to host:   ", lambda x: f", {sz/x:.2f} GB/s"):
-      ret = libc.read(fd, hst, sz)
-    with Timing("gpu host copy:  ", lambda x: f", {sz/x:.2f} GB/s"):
-      dev.allocator._copyin_async(gpubuf, hst, sz)
-      dev.synchronize()
-
-def read_to_gpu_pingpong(fd, sz, gpubuf):
-  psz = 256*1024*1024
-  print(f"piece size {psz/(1024*1024):.2f} MB")
-  with Timing("gpu host alloc: ", lambda x: f", {sz/x:.2f} GB/s"):
-    hst1 = dev.allocator._hostalloc(psz)
-    hst2 = dev.allocator._hostalloc(psz)
-
-  os.lseek(fd, 0, os.SEEK_SET)
-  with Timing("total: ", lambda x: f", {sz/x:.2f} GB/s"):
-    for i in range(sz//(psz*2)):
-      with Timing("tfer(0):           ", lambda x: f", {psz/x:.2f} GB/s"):
-        ret = libc.read(fd, hst1, psz)
-        dev.synchronize()
-        dev.allocator._copyin_async(gpubuf, hst1, psz)
-      with Timing("tfer(1):           ", lambda x: f", {psz/x:.2f} GB/s"):
-        ret = libc.read(fd, hst2, psz)
-        dev.synchronize()
-        dev.allocator._copyin_async(gpubuf, hst2, psz)
-    dev.synchronize()
-
-MAP_LOCKED = 0x2000
-MAP_HUGETLB = 0x40000
-
-if __name__ == "__main__":
-  dev = Device[Device.DEFAULT]
-
-  warm = (Tensor.ones(1024, device=Device.DEFAULT).contiguous() + Tensor.ones(1024, device=Device.DEFAULT).contiguous()).realize()
-  #fn = "/home/tiny/tinygrad/weights/rng"
-  fn = pathlib.Path(__file__).parents[1] / "weights/LLaMA-2/70B/consolidated.00.pth"
-  sz = os.stat(fn).st_size
-  t = Tensor.empty(sz, dtype=dtypes.uint8, device=f"disk:{fn}")
-  with Timing("copy:  ", lambda x: f", {sz/x:.2f} GB/s"):
-    on_dev = t.to(Device.DEFAULT).realize()
-
-  exit(0)
-
-  # 4GB of random numbers
-  #fd = os.open("/home/tiny/tinygrad/weights/rng", os.O_RDWR|os.O_DIRECT)
-  #sz = os.fstat(fd).st_size // 4
-  fd = os.open("/home/tiny/tinygrad/weights/LLaMA/7B/consolidated.00.pth", os.O_RDWR|os.O_DIRECT)
-  sz = os.fstat(fd).st_size
-  print(f"read {sz} from {fd}")
-
-  with Timing("gpu alloc:  ", lambda x: f", {sz/x:.2f} GB/s"):
-    gpubuf = dev.allocator._alloc(sz)
-  # warmup
-  dev.allocator._copyin_async(gpubuf, from_mv(bytearray(b"\x00\x00\x00\x00"*0x1000)), 0x4000)
-  print("copying, is warm")
-
-  print("****** read to gpu pingpong")
-  read_to_gpu_pingpong(fd, sz, gpubuf)
-  exit(0)
-
-  print("****** read direct")
-  read_direct(fd, sz)
-
-  print("****** read mmap")
-  read_mmap(fd, sz)
-
-  print("****** read to gpu single")
-  read_to_gpu_single(fd, sz, gpubuf)
-
-  print("****** read to gpu mmap")
-  read_to_gpu_mmap(fd, sz, gpubuf)
-
-  os._exit(0)

extra/dump_cache.py

@@ -1,21 +0,0 @@
-import sys, sqlite3, pickle
-from tinygrad.helpers import CACHEDB
-
-if __name__ == "__main__":
-  fn = sys.argv[1] if len(sys.argv) > 1 else CACHEDB
-  conn = sqlite3.connect(fn)
-  cur = conn.cursor()
-  cur.execute("SELECT name FROM sqlite_master WHERE type='table'")
-  for f in cur.fetchall():
-    table = f[0]
-    cur2 = conn.cursor()
-    cur2.execute(f"SELECT COUNT(*) FROM {table}")
-    cnt = cur2.fetchone()[0]
-    print(f"{table:20s} : {cnt}")
-
-    cur3 = conn.cursor()
-    cur3.execute(f"SELECT * FROM {table} LIMIT 10")
-    for f in cur3.fetchall():
-      v = pickle.loads(f[-1])
-      print("   ", len(f[0]) if isinstance(f[0], str) else f[0], f[1:-1], str(v)[0:50])
-      #print(f"{len(k):10d}, {sk} -> {v}")

extra/gemm/.gitignore

@@ -1,3 +1,2 @@
-*.s
 *.ll
 fp32_sgemm_amd

extra/gemm/amd_matmul.py

@@ -36,7 +36,7 @@ if __name__ == "__main__":
     for _ in range(run_count): tc = (a@b).realize()
 
   GlobalCounters.reset()
-  ei = ExecItem(runner, [a.uop.buffer, b.uop.buffer, c.uop.buffer])
+  ei = ExecItem(ast, [a.uop.buffer, b.uop.buffer, c.uop.buffer], prg=runner)
   with Context(DEBUG=2):
     for _ in range(run_count): ei.run(wait=True)
   print(f"custom  {(c-tc).square().mean().item()}")

extra/gemm/amd_uop_matmul.py

@@ -147,7 +147,7 @@ def test_matmul(sink:UOp, N=N):
   hc = Tensor.empty(N, N)
   Tensor.realize(a, b, hc)
 
-  ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
+  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))
 
   ets = []
   with Context(DEBUG=2):

extra/gemm/asm/test.py

@@ -0,0 +1,65 @@
+# Run assembly on the AMD runtime and check correctness
+# VIZ=2 to profile
+import pathlib
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.engine.realize import ExecItem, CompiledRunner
+from tinygrad.renderer import ProgramSpec
+from tinygrad.uop.ops import track_rewrites, UOp
+from tinygrad.helpers import TracingKey
+
+fp = pathlib.Path(__file__).parent/"gemm.s"
+
+# ** generate inputs on CPU
+
+N = 8192
+scale = 10.0
+
+import torch
+torch.manual_seed(0)
+A = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
+B = (torch.randn(N, N, dtype=torch.float32, device="cpu") / scale).to(torch.bfloat16).contiguous()
+Bt = B.t().contiguous() # transpose B for the baseline gemm
+C_torch = A@Bt
+
+# ** copy buffers to AMD
+
+# input creation and validation run on the copy engine for simpler tracing
+
+def from_torch(t:torch.Tensor) -> Tensor:
+  return Tensor.from_blob(t.data_ptr(), t.shape, dtype=dtypes.bfloat16, device="cpu").to(Device.DEFAULT).realize()
+
+C_tiny = Tensor.matmul(from_torch(A), from_torch(Bt), dtype=dtypes.float32).cast(dtypes.bfloat16)
+C_asm = Tensor.empty_like(C_tiny)
+C_asm.uop.buffer.allocate()
+
+# ** run gemms
+
+@track_rewrites(name=lambda *args,ret,**kwargs: TracingKey(ret.name, (ret.function_name,), ret=ret))
+def get_asm_gemm(ast:UOp, fp:pathlib.Path) -> ProgramSpec:
+  src = fp.read_text()
+  lib = Device[Device.DEFAULT].compiler.compile(src)
+  return ProgramSpec("gemm", src, Device.DEFAULT, ast, lib=lib, global_size=[1024, 1, 1], local_size=[256, 1, 1], globals=[0, 1, 2])
+
+sched = C_tiny.schedule()
+assert len(sched) == 1
+eis:list[ExecItem] = [sched[-1].lower()]
+ast = eis[0].ast
+prg = get_asm_gemm(ast, fp)
+eis.append(ExecItem(ast, [C_asm.uop.buffer, from_torch(B).uop.buffer, from_torch(A).uop.buffer], prg=CompiledRunner(prg)))
+
+
+for ei in eis:
+  et = ei.run(wait=True)
+  print(f"{(N*N*N*2 / et)*1e-12:.2f} REAL TFLOPS")
+
+# ** correctness
+
+import ctypes
+
+def torch_bf16(t:Tensor) -> torch.tensor:
+  asm_out = t.to("cpu").realize().uop.buffer._buf
+  buf = (ctypes.c_uint16*C_asm.uop.size).from_address(asm_out.va_addr)
+  return torch.frombuffer(buf, dtype=torch.bfloat16, count=C_asm.uop.size).reshape(C_asm.shape)
+
+assert torch.allclose(torch_bf16(C_asm), C_torch, rtol=1e-2, atol=1e-3)
+assert torch.allclose(torch_bf16(C_tiny), C_torch, rtol=1e-2, atol=1e-3)

extra/gemm/asm/unpack_kd.py

@@ -0,0 +1,179 @@
+# unpack the complete kernel descriptor of an amdgpu ELF of for gfx950
+# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
+import struct, pathlib
+from tinygrad.runtime.support.elf import elf_loader
+
+def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
+def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
+
+with open(fp:=pathlib.Path(__file__).parent/"lib", "rb") as f:
+  lib = f.read()
+
+image, sections, relocs = elf_loader(lib)
+rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
+
+# rodata is exactly 64 bytes
+kd = image[rodata_entry:rodata_entry+64]
+desc = int.from_bytes(kd, byteorder="little")
+
+group_segment_fixed_size = bits(desc, 0, 31)
+private_segment_fixed_size = bits(desc, 32, 63)
+kernarg_size = bits(desc, 64, 95)
+reserved_127_96 = bits(desc, 96, 127)
+assert reserved_127_96 == 0
+
+print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
+print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
+print("KERNARG_SIZE:", kernarg_size)
+print("RESERVED 127:96:", reserved_127_96)
+
+entry_off = bits(desc, 128, 191)
+
+# sign-extend manually if needed
+if entry_off & (1 << 63):
+  entry_off -= 1 << 64
+
+print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
+
+kd_addr = 0x1840
+entry_addr = kd_addr + entry_off
+
+print("Computed entry address: 0x%016x" % entry_addr)
+print("256B aligned:", entry_addr % 256 == 0)
+
+pgm_rsrc3 = bits(desc, 352, 383)
+pgm_rsrc1 = bits(desc, 384, 415)
+pgm_rsrc2 = bits(desc, 416, 447)
+
+print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
+print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
+print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
+
+# rsrc 3
+
+accum_offset_raw = bits(pgm_rsrc3, 0, 5)
+assert_zero(pgm_rsrc3, 6, 15)
+tg_split = bits(pgm_rsrc3, 16, 16)
+accum_offset_vgprs = (accum_offset_raw + 1) * 4
+print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
+print("RSRC3.TG_SPLIT:", tg_split)
+
+# rsrc 1
+
+vgpr_gran = bits(pgm_rsrc1, 0, 5)
+sgpr_gran = bits(pgm_rsrc1, 6, 9)
+assert_zero(pgm_rsrc1, 27, 28)
+
+# NOTE: this is vgprs + agprs
+vgprs_used = (vgpr_gran + 1) * 8
+assert 0 <= vgprs_used <= 512
+
+k = sgpr_gran // 2
+sgprs_used = (k + 1) * 16
+
+print("RSRC1.VGPRS:", vgprs_used)
+print("RSRC1.SGPRS:", sgprs_used)
+
+assert_zero(pgm_rsrc1, 10, 11)
+
+float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
+float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
+float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
+float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
+
+priv = bits(pgm_rsrc1, 20, 20)
+assert priv == 0
+enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
+debug_mode = bits(pgm_rsrc1, 22, 22)
+enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
+bulky = bits(pgm_rsrc1, 24, 24)
+assert bulky == 0
+cdbg_user = bits(pgm_rsrc1, 25, 25)
+assert cdbg_user == 0
+fp16_ovfl = bits(pgm_rsrc1, 26, 26)
+assert_zero(pgm_rsrc1, 27, 28)  # reserved
+assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
+assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
+assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
+
+# rsrc 2
+
+enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
+user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
+enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
+assert enable_trap_handler == 0
+
+enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
+enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
+enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
+enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
+
+enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
+
+enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
+assert enable_exception_address_watch == 0
+enable_exception_memory = bits(pgm_rsrc2, 14, 14)
+assert enable_exception_memory == 0
+
+granulated_lds_size = bits(pgm_rsrc2, 15, 23)
+assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
+
+enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
+enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
+enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
+enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
+enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
+enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
+enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
+
+assert_zero(pgm_rsrc2, 31, 31)
+
+print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
+print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
+print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
+
+print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
+print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
+print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
+print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
+print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
+print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
+print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
+
+# user sgprs
+
+enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
+enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
+enable_sgpr_queue_ptr = bits(desc, 450, 450)
+enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
+enable_sgpr_dispatch_id = bits(desc, 452, 452)
+enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
+enable_sgpr_private_segment_size = bits(desc, 454, 454)
+
+assert_zero(desc, 455, 457)
+
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
+print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
+print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
+print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
+print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
+print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
+
+assert_zero(desc, 458, 459)
+
+uses_dynamic_stack = bits(desc, 459, 460)
+print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
+
+assert_zero(desc, 460, 463)
+kernarg_preload_spec_length = bits(desc, 464, 470)
+print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
+
+kernarg_preload_spec_offset = bits(desc, 471, 479)
+print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
+
+assert_zero(desc, 480, 511)

extra/gemm/jax_pmatmul.py

@@ -1,27 +0,0 @@
-#!/usr/bin/env python3
-import time
-import jax
-import jax.numpy as jnp
-
-print(jax.devices())
-DEVICES = len(jax.devices())
-BS = 32
-N = 4096
-dtype = jnp.float16
-A = jnp.zeros((DEVICES, BS, N, N), dtype)
-B = jnp.zeros((1, 1, N, N), dtype)
-A = jax.device_put_sharded([A[i] for i in range(DEVICES)], jax.devices())
-B = jax.device_put_sharded([B for i in range(DEVICES)], jax.devices())
-
-OPS = DEVICES*BS*N*N*N*2
-def matmul(A,B): return jnp.matmul(A,B,preferred_element_type=jnp.float32)
-pmatmul = jax.pmap(matmul)
-
-MAX_TFLOPS = 123*DEVICES  # Peak FP16 Tensor TFLOPS with FP32 Acc (7900XTX)
-for i in range(10):
-  st = time.perf_counter()
-  C = pmatmul(A,B).block_until_ready()
-  et = time.perf_counter()-st
-  tflops = (OPS*1e-12)/et
-  print(f"time {et*1e3:.2f} ms, TFLOPS {tflops:6.2f}, MFU {(tflops/MAX_TFLOPS)*100:4.2f}% out shape {C.shape} dtype {C.dtype}")
-

extra/gemm/mlx_matmul.py

@@ -1,10 +0,0 @@
-import mlx.core as mx
-from tinygrad.helpers import Timing
-N = 4096
-x = mx.random.normal((N,N))
-w = mx.random.normal((N,N))
-
-FLOPS = N*N*N*2
-for i in range(10):
-  with Timing("", lambda x: f"  {FLOPS/x:.2f} GFLOPS"):
-    mx.eval(x@w)

extra/gemm/simple_matmul.py

@@ -3,7 +3,6 @@ from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
 from tinygrad.codegen.opt import OptOps
-from tinygrad.engine.realize import lower_schedule
 
 dtype_in = (dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else
             dtypes.fp8e4m3 if getenv("FP8E4M3") else dtypes.fp8e5m2 if getenv("FP8E5M2") else dtypes.float)
@@ -40,8 +39,8 @@ if __name__ == "__main__":
 
   if getenv("SHOULD_USE_TC"):
     sched = a.matmul(b, dtype=acc_dtype).schedule()
-    lowered = list(lower_schedule(sched))
-    ei = get_single_element(lowered)[1]
+    ei = get_single_element(sched)
+    ei.lower()
     assert any(opt.op is OptOps.TC for opt in ei.prg.p.applied_opts), f"TC not triggered, {ei.prg.p.applied_opts}"
 
   ref = a.numpy().astype(np.float32) @ b.numpy().astype(np.float32)

extra/gemm/tf_gemm.py

@@ -1,33 +0,0 @@
-import time
-import tensorflow as tf
-
-gpus = tf.config.list_physical_devices('GPU')
-if gpus:
-  try:
-    # Currently, memory growth needs to be the same across GPUs
-    for gpu in gpus:
-      tf.config.experimental.set_memory_growth(gpu, True)
-    logical_gpus = tf.config.list_logical_devices('GPU')
-    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
-  except RuntimeError as e:
-    # Memory growth must be set before GPUs have been initialized
-    print(e)
-
-for dtype in [tf.float16, tf.float32]:
-  for N in [256, 512, 1024, 2048, 4096, 8192]:
-    FLOPS = N*N*N*2
-
-    b = tf.random.uniform((N, N), dtype=dtype)
-    c = tf.random.uniform((N, N), dtype=dtype)
-
-    b = tf.Variable(b)
-    c = tf.Variable(c)
-
-    def tf_prog(b, c):
-      st = time.perf_counter()
-      a = tf.matmul(b, c)
-      tf.debugging.check_numerics(a, "Nan or Inf in result") # Ensures that the calculation is done.
-      return time.perf_counter() - st
-
-    tm = min([tf_prog(b, c) for _ in range(20)])
-    print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS {N:4d}x{N:4d}x{N:4d} matmul in {dtype}")

extra/gemm/tinygrad_nv_matmul.py

@@ -33,5 +33,5 @@ if __name__ == "__main__":
   new_src = prg.src
   # can mod source here
   prg = replace(prg, src=new_src)
-  ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)
+  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
   for i in range(5): ei.run(wait=True)

extra/gemm/triton_nv_matmul.py

@@ -88,7 +88,7 @@ if __name__ == "__main__":
   prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
                 global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
                 mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
-  ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)
+  ei = ExecItem(si.ast, [x.ensure_allocated() for x in si.bufs], si.metadata, prg=CompiledRunner(prg))
   tflops = []
   for i in range(5):
     tm = ei.run(wait=True)

extra/hip_events.py

@@ -1,12 +0,0 @@
-import ctypes
-import tinygrad.runtime.autogen.hip as hip
-from tinygrad.runtime.ops_hip import check
-from tinygrad.helpers import init_c_var
-
-if __name__ == "__main__":
-  check(hip.hipSetDevice(0))
-  evt = init_c_var(hip.hipEvent_t(), lambda x: check(hip.hipEventCreate(ctypes.byref(x))))
-  check(hip.hipSetDevice(1))
-  check(hip.hipStreamWaitEvent(None, evt, 0))
-  check(hip.hipSetDevice(0))
-  check(hip.hipEventRecord(evt, None))

extra/junk/sentencepiece_model_pb2.py

@@ -1,45 +0,0 @@
-# -*- coding: utf-8 -*-
-# Generated by the protocol buffer compiler.  DO NOT EDIT!
-# source: sentencepiece_model.proto
-# Protobuf Python Version: 4.25.1
-"""Generated protocol buffer code."""
-from google.protobuf import descriptor as _descriptor
-from google.protobuf import descriptor_pool as _descriptor_pool
-from google.protobuf import symbol_database as _symbol_database
-from google.protobuf.internal import builder as _builder
-# @@protoc_insertion_point(imports)
-
-_sym_db = _symbol_database.Default()
-
-
-
-
-DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19sentencepiece_model.proto\x12\rsentencepiece\"\x80\x0c\n\x0bTrainerSpec\x12\r\n\x05input\x18\x01 \x03(\t\x12\x14\n\x0cinput_format\x18\x07 \x01(\t\x12\x14\n\x0cmodel_prefix\x18\x02 \x01(\t\x12\x41\n\nmodel_type\x18\x03 \x01(\x0e\x32$.sentencepiece.TrainerSpec.ModelType:\x07UNIGRAM\x12\x18\n\nvocab_size\x18\x04 \x01(\x05:\x04\x38\x30\x30\x30\x12\x17\n\x0f\x61\x63\x63\x65pt_language\x18\x05 \x03(\t\x12 \n\x15self_test_sample_size\x18\x06 \x01(\x05:\x01\x30\x12*\n\x1b\x65nable_differential_privacy\x18\x32 \x01(\x08:\x05\x66\x61lse\x12+\n differential_privacy_noise_level\x18\x33 \x01(\x02:\x01\x30\x12\x32\n\'differential_privacy_clipping_threshold\x18\x34 \x01(\x04:\x01\x30\x12\"\n\x12\x63haracter_coverage\x18\n \x01(\x02:\x06\x30.9995\x12\x1e\n\x13input_sentence_size\x18\x0b \x01(\x04:\x01\x30\x12$\n\x16shuffle_input_sentence\x18\x13 \x01(\x08:\x04true\x12 \n\x14mining_sentence_size\x18\x0c \x01(\x05\x42\x02\x18\x01\x12\"\n\x16training_sentence_size\x18\r \x01(\x05\x42\x02\x18\x01\x12(\n\x17seed_sentencepiece_size\x18\x0e \x01(\x05:\x07\x31\x30\x30\x30\x30\x30\x30\x12\x1e\n\x10shrinking_factor\x18\x0f \x01(\x02:\x04\x30.75\x12!\n\x13max_sentence_length\x18\x12 \x01(\x05:\x04\x34\x31\x39\x32\x12\x17\n\x0bnum_threads\x18\x10 \x01(\x05:\x02\x31\x36\x12\x1d\n\x12num_sub_iterations\x18\x11 \x01(\x05:\x01\x32\x12$\n\x18max_sentencepiece_length\x18\x14 \x01(\x05:\x02\x31\x36\x12%\n\x17split_by_unicode_script\x18\x15 \x01(\x08:\x04true\x12\x1d\n\x0fsplit_by_number\x18\x17 \x01(\x08:\x04true\x12!\n\x13split_by_whitespace\x18\x16 \x01(\x08:\x04true\x12)\n\x1atreat_whitespace_as_suffix\x18\x18 \x01(\x08:\x05\x66\x61lse\x12+\n\x1c\x61llow_whitespace_only_pieces\x18\x1a \x01(\x08:\x05\x66\x61lse\x12\x1b\n\x0csplit_digits\x18\x19 \x01(\x08:\x05\x66\x61lse\x12#\n\x19pretokenization_delimiter\x18\x35 \x01(\t:\x00\x12\x17\n\x0f\x63ontrol_symbols\x18\x1e \x03(\t\x12\x1c\n\x14user_defined_symbols\x18\x1f \x03(\t\x12\x16\n\x0erequired_chars\x18$ \x01(\t\x12\x1c\n\rbyte_fallback\x18# \x01(\x08:\x05\x66\x61lse\x12+\n\x1dvocabulary_output_piece_score\x18  \x01(\x08:\x04true\x12\x1e\n\x10hard_vocab_limit\x18! \x01(\x08:\x04true\x12\x1c\n\ruse_all_vocab\x18\" \x01(\x08:\x05\x66\x61lse\x12\x11\n\x06unk_id\x18( \x01(\x05:\x01\x30\x12\x11\n\x06\x62os_id\x18) \x01(\x05:\x01\x31\x12\x11\n\x06\x65os_id\x18* \x01(\x05:\x01\x32\x12\x12\n\x06pad_id\x18+ \x01(\x05:\x02-1\x12\x18\n\tunk_piece\x18- \x01(\t:\x05<unk>\x12\x16\n\tbos_piece\x18. \x01(\t:\x03<s>\x12\x17\n\teos_piece\x18/ \x01(\t:\x04</s>\x12\x18\n\tpad_piece\x18\x30 \x01(\t:\x05<pad>\x12\x1a\n\x0bunk_surface\x18, \x01(\t:\x05 \xe2\x81\x87 \x12+\n\x1ctrain_extremely_large_corpus\x18\x31 \x01(\x08:\x05\x66\x61lse\"5\n\tModelType\x12\x0b\n\x07UNIGRAM\x10\x01\x12\x07\n\x03\x42PE\x10\x02\x12\x08\n\x04WORD\x10\x03\x12\x08\n\x04\x43HAR\x10\x04*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"\xd1\x01\n\x0eNormalizerSpec\x12\x0c\n\x04name\x18\x01 \x01(\t\x12\x1c\n\x14precompiled_charsmap\x18\x02 \x01(\x0c\x12\x1e\n\x10\x61\x64\x64_dummy_prefix\x18\x03 \x01(\x08:\x04true\x12&\n\x18remove_extra_whitespaces\x18\x04 \x01(\x08:\x04true\x12 \n\x12\x65scape_whitespaces\x18\x05 \x01(\x08:\x04true\x12\x1e\n\x16normalization_rule_tsv\x18\x06 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"y\n\x0cSelfTestData\x12\x33\n\x07samples\x18\x01 \x03(\x0b\x32\".sentencepiece.SelfTestData.Sample\x1a)\n\x06Sample\x12\r\n\x05input\x18\x01 \x01(\t\x12\x10\n\x08\x65xpected\x18\x02 \x01(\t*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\"\xfe\x03\n\nModelProto\x12\x37\n\x06pieces\x18\x01 \x03(\x0b\x32\'.sentencepiece.ModelProto.SentencePiece\x12\x30\n\x0ctrainer_spec\x18\x02 \x01(\x0b\x32\x1a.sentencepiece.TrainerSpec\x12\x36\n\x0fnormalizer_spec\x18\x03 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x12\x33\n\x0eself_test_data\x18\x04 \x01(\x0b\x32\x1b.sentencepiece.SelfTestData\x12\x38\n\x11\x64\x65normalizer_spec\x18\x05 \x01(\x0b\x32\x1d.sentencepiece.NormalizerSpec\x1a\xd2\x01\n\rSentencePiece\x12\r\n\x05piece\x18\x01 \x01(\t\x12\r\n\x05score\x18\x02 \x01(\x02\x12\x42\n\x04type\x18\x03 \x01(\x0e\x32,.sentencepiece.ModelProto.SentencePiece.Type:\x06NORMAL\"T\n\x04Type\x12\n\n\x06NORMAL\x10\x01\x12\x0b\n\x07UNKNOWN\x10\x02\x12\x0b\n\x07\x43ONTROL\x10\x03\x12\x10\n\x0cUSER_DEFINED\x10\x04\x12\x08\n\x04\x42YTE\x10\x06\x12\n\n\x06UNUSED\x10\x05*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02*\t\x08\xc8\x01\x10\x80\x80\x80\x80\x02\x42\x02H\x03')
-
-_globals = globals()
-_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
-_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'sentencepiece_model_pb2', _globals)
-if _descriptor._USE_C_DESCRIPTORS == False:
-  _globals['DESCRIPTOR']._options = None
-  _globals['DESCRIPTOR']._serialized_options = b'H\003'
-  _globals['_TRAINERSPEC'].fields_by_name['mining_sentence_size']._options = None
-  _globals['_TRAINERSPEC'].fields_by_name['mining_sentence_size']._serialized_options = b'\030\001'
-  _globals['_TRAINERSPEC'].fields_by_name['training_sentence_size']._options = None
-  _globals['_TRAINERSPEC'].fields_by_name['training_sentence_size']._serialized_options = b'\030\001'
-  _globals['_TRAINERSPEC']._serialized_start=45
-  _globals['_TRAINERSPEC']._serialized_end=1581
-  _globals['_TRAINERSPEC_MODELTYPE']._serialized_start=1517
-  _globals['_TRAINERSPEC_MODELTYPE']._serialized_end=1570
-  _globals['_NORMALIZERSPEC']._serialized_start=1584
-  _globals['_NORMALIZERSPEC']._serialized_end=1793
-  _globals['_SELFTESTDATA']._serialized_start=1795
-  _globals['_SELFTESTDATA']._serialized_end=1916
-  _globals['_SELFTESTDATA_SAMPLE']._serialized_start=1864
-  _globals['_SELFTESTDATA_SAMPLE']._serialized_end=1905
-  _globals['_MODELPROTO']._serialized_start=1919
-  _globals['_MODELPROTO']._serialized_end=2429
-  _globals['_MODELPROTO_SENTENCEPIECE']._serialized_start=2208
-  _globals['_MODELPROTO_SENTENCEPIECE']._serialized_end=2418
-  _globals['_MODELPROTO_SENTENCEPIECE_TYPE']._serialized_start=2323
-  _globals['_MODELPROTO_SENTENCEPIECE_TYPE']._serialized_end=2407
-# @@protoc_insertion_point(module_scope)

extra/mcts_search.py

@@ -1,176 +0,0 @@
-from __future__ import annotations
-from typing import List, Optional, Dict, cast
-import numpy as np
-np.set_printoptions(suppress=True)
-import math, functools, time, random, statistics
-from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.codegen.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
-from tinygrad.engine.realize import get_program
-
-class MCTSNode:
-  def __init__(self, kernel:Kernel, parent=None):
-    self.kernel:Kernel = kernel
-    self.t = math.inf
-    self.n = 0
-    self.tm = math.inf
-    self.i = -1
-    self.parents: List[MCTSNode] = [parent] if parent is not None else []
-    self.children: Optional[List[MCTSNode]] = None
-    self.removed_children: List[MCTSNode] = []
-
-def expand_node(node:MCTSNode):
-  assert node.children is None
-  node.children = [MCTSNode(x, node) for x in get_kernel_actions(node.kernel, include_0=False).values()]
-
-def remove_node(node:MCTSNode):
-  for parent in node.parents:
-    assert parent.children is not None
-    parent.children.remove(node)
-    parent.removed_children.append(node)
-
-C = math.sqrt(2)
-TEMP = 0.5
-def _sample_tree(node:MCTSNode, best_tm:float) -> MCTSNode:
-  if node.children is None or len(node.children) == 0: return node
-  unexplored_children = []
-  explored_children = []
-  ucb_explored_children: List[float] = []
-  for child in node.children:
-    if child.n == 0: unexplored_children.append(child)
-    else:
-      ucb = -child.t/best_tm + C*math.sqrt(math.log(node.n)/child.n)
-      if not math.isinf(ucb):
-        explored_children.append(child)
-        ucb_explored_children.append(ucb)
-  if len(unexplored_children): return random.choice(unexplored_children)
-  if not len(explored_children): return node
-  # safe softmax
-  ucb_exp = np.exp((np.array(ucb_explored_children)-max(ucb_explored_children))/TEMP)
-  return _sample_tree(explored_children[np.random.choice(len(ucb_exp), p=ucb_exp/np.sum(ucb_exp))], best_tm)
-
-# this will expand/remove sometimes
-def sample_tree(root:MCTSNode, best_tm:float) -> Optional[MCTSNode]:
-  if root.children is None: expand_node(root)
-  while root.children:
-    # tree traversal
-    node = _sample_tree(root, best_tm)
-
-    if node.children is not None and len(node.children) == 0:
-      remove_node(node)
-      continue
-
-    # node expansion
-    if node.n != 0:
-      if node.children is None: expand_node(node)
-      assert node.children is not None
-      if len(node.children) == 0:
-        remove_node(node)
-        continue
-      node = random.choice(node.children)
-    return node
-  return None
-
-def backprop(bnode:MCTSNode, tm, strength=1.0):
-  if bnode.t > tm: bnode.t = tm
-  bnode.n += strength
-  for parent in bnode.parents: backprop(parent, tm, strength/len(bnode.parents))
-
-graph_mcts_cnt = 0
-def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
-  global graph_mcts_cnt
-  # TODO: copied from BEAM
-  key = {"ast": lin.ast.key, "amt": amt, "device": lin.opts.device, "suffix": lin.opts.suffix}
-  if not getenv("IGNORE_MCTS_CACHE") and CACHELEVEL >= 1 and (val:=diskcache_get("mcts_search", key)) is not None:
-    ret = lin.copy()
-    for o in val[len(lin.applied_opts):]: ret.apply_opt(o)
-    return ret
-
-  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals = {k.expr:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
-  dev = Device[lin.opts.device]
-  root = MCTSNode(lin)
-
-  st = time.perf_counter()
-  best, best_idx, best_tm = lin, 0, math.inf
-  seen_libs: Dict[bytes, MCTSNode] = {}
-  seen_asts: Dict[bytes, MCTSNode] = {}
-  compile_time, runtime_time = 0.0, 0.0
-  for i in range(amt):
-    node = sample_tree(root, best_tm)  # sample and expand
-    if node is None: break  # finished the whole tree
-    node.i = i  # when was node explored
-
-    opt_ast = node.kernel.get_optimized_ast()
-    if (sibling_node:=seen_asts.get(opt_ast.key, None)) is not None:
-      # early check for same optimized AST hit
-      remove_node(node)
-      tm = sibling_node.t
-    else:
-      seen_asts[opt_ast.key] = node
-
-      # lowering (50% of the time)
-      p = get_program(node.kernel.get_optimized_ast(name_override="test"), node.kernel.opts)
-
-      # rollout
-      tm1 = time.perf_counter()
-      try:
-        lib = dev.compiler.compile(p.src)
-      except CompileError:
-        # NOTE: many of these "compiler errors" are caused by bad code output from the lowerer
-        lib = None
-      tm2 = time.perf_counter()
-      if lib is None:
-        tm = math.inf
-      else:
-        if (sibling_node:=seen_libs.get(lib, None)) is not None:
-          # NOTE: these should all be caught by the AST check, need to canonicalize
-          # remove this node, it's a duplicate
-          remove_node(node)
-          tm = sibling_node.t
-        else:
-          seen_libs[lib] = node
-          try: tm = statistics.median(_time_program(p, lib, var_vals, rawbufs, cnt=3, early_stop=best_tm*5/1e6))*1e6
-          except RuntimeError: tm = math.inf
-          node.tm = tm
-      tm3 = time.perf_counter()
-      compile_time += tm2-tm1
-      runtime_time += tm3-tm2
-
-      # mock rollout
-      #node.tm = tm = random.random() + 0.1
-
-    if tm < best_tm: best, best_idx, best_tm = node.kernel, i, tm
-    et = time.perf_counter() - st
-    if DEBUG>=2: print(f"\r{et:7.2f}s {colored(f'{compile_time*100/et:3.0f}%', 'cyan')} {colored(f'{runtime_time*100/et:3.0f}%', 'red')}: {tm:12.2f} us     best: {best_tm:12.2f} us @ {best_idx+1:4d}      {i+1:4d}/{amt:4d}  {int(round((i+1)/et)):4d}/s     {node.kernel.colored_shape()}\033[K", end="")  # noqa: E501
-
-    # backprop
-    backprop(node, tm)
-  if DEBUG>=2: print()
-
-  if getenv("MCTSGRAPH"):
-    import networkx as nx
-    import os
-    GRAPHPATH = "/tmp/net"
-    def save_graph(G, fn, opt=""):
-      print("saving", G, f"to {fn}.svg")
-      nx.drawing.nx_pydot.write_dot(G, f'{fn}.dot')
-      os.system(f'dot {opt} -Tsvg {fn}.dot -o {fn}.svg')
-
-    G = nx.DiGraph()
-    def add_node(node:MCTSNode):
-      if node.n == 0: return
-      for parent in node.parents: G.add_edge(parent, node)
-      gopts = node.kernel.applied_opts
-      edge_lbl = f"{str(gopts[-1].op)[7:]} {gopts[-1].axis} {gopts[-1].arg}" if len(gopts) else "ROOT"
-      G.add_node(node, label=f"{node.i+1}\n{node.tm:.2f} us\n{edge_lbl}\nt {node.t:.2f}\nn {node.n}",
-                 fillcolor="#80ff8080" if node.tm == best_tm else "#ffff8080", style='filled' if node.t == best_tm else '')
-      if node.children is not None:
-        for child in node.children+node.removed_children: add_node(child)
-    add_node(root)
-    save_graph(G, f"{GRAPHPATH}.{graph_mcts_cnt}.mcts", '-Grankdir=LR')
-    graph_mcts_cnt += 1
-
-  if CACHELEVEL >= 1: diskcache_put("mcts_search", key, best.applied_opts)
-  return best

extra/models/bert.py

@@ -242,7 +242,8 @@ class BertIntermediate:
   def __call__(self, hidden_states):
     x = self.dense(hidden_states)
     # tinygrad gelu is openai gelu but we need the original bert gelu
-    return gelu(x)
+    # NOTE: contiguous for speed
+    return gelu(x).contiguous()
 
 class BertAttention:
   def __init__(self, hidden_size, num_attention_heads, attention_probs_dropout_prob, hidden_dropout_prob):

extra/models/llama.py

@@ -245,6 +245,11 @@ def convert_from_huggingface(weights:dict[str, Tensor], n_layers: int, n_heads:
       continue
     sd[keymap[k]] = v
   for k,v in experts.items(): sd[k] = Tensor.stack(*[v[i] for i in range(len(v))])
+
+  # Handle tied embeddings (e.g., Llama 3.2 1B Instruct where lm_head shares weights with embed_tokens)
+  if "output.weight" not in sd and "tok_embeddings.weight" in sd:
+    sd["output.weight"] = sd["tok_embeddings.weight"]
+
   return sd
 
 def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):

extra/optimization/extract_policynet.py

@@ -1,114 +0,0 @@
-import os, sys, sqlite3, pickle, random
-from tqdm import tqdm, trange
-from copy import deepcopy
-from tinygrad.nn import Linear
-from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import Adam
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.codegen.opt.search import actions
-from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.helpers import getenv
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.uop.ops import Variable
-inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-
-INNER = 256
-class PolicyNet:
-  def __init__(self):
-    self.l1 = Linear(1021,INNER)
-    self.l2 = Linear(INNER,INNER)
-    self.l3 = Linear(INNER,1+len(actions))
-  def __call__(self, x):
-    x = self.l1(x).relu()
-    x = self.l2(x).relu().dropout(0.9)
-    return self.l3(x).log_softmax()
-
-def dataset_from_cache(fn):
-  conn = sqlite3.connect(fn)
-  cur = conn.cursor()
-  cur.execute("SELECT * FROM beam_search")
-  X,A = [], []
-  for f in tqdm(cur.fetchall()):
-    Xs,As = [], []
-    try:
-      lin = Kernel(eval(f[0]))
-      opts = pickle.loads(f[-1])
-      for o in opts:
-        Xs.append(lin_to_feats(lin, use_sts=True))
-        As.append(actions.index(o))
-        lin.apply_opt(o)
-      Xs.append(lin_to_feats(lin, use_sts=True))
-      As.append(0)
-    except Exception:
-      pass
-    X += Xs
-    A += As
-  return X,A
-
-if __name__ == "__main__":
-  if getenv("REGEN"):
-    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
-    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset_policy")
-  else:
-    ld = safe_load("/tmp/dataset_policy")
-    X,V = ld['X'].numpy(), ld['V'].numpy()
-
-  print(X.shape, V.shape)
-  order = list(range(X.shape[0]))
-  random.shuffle(order)
-  X, V = X[order], V[order]
-
-  ratio = -256
-  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
-  X,V = X[:ratio], V[:ratio]
-  print(X.shape, V.shape)
-
-  net = PolicyNet()
-  #if os.path.isfile("/tmp/policynet.safetensors"): load_state_dict(net, safe_load("/tmp/policynet.safetensors"))
-  optim = Adam(get_parameters(net))
-
-  def get_minibatch(X,Y,bs):
-    xs, ys = [], []
-    for _ in range(bs):
-      sel = random.randint(0, len(X)-1)
-      xs.append(X[sel])
-      ys.append(Y[sel])
-    return Tensor(xs), Tensor(ys)
-
-  Tensor.training = True
-  losses = []
-  test_losses = []
-  test_accuracy = 0
-  test_loss = float('inf')
-  for i in (t:=trange(500)):
-    x,y = get_minibatch(X,V,bs=256)
-    out = net(x)
-    loss = out.sparse_categorical_crossentropy(y)
-    optim.zero_grad()
-    loss.backward()
-    optim.step()
-    cat = out.argmax(axis=-1)
-    accuracy = (cat == y).mean()
-    t.set_description(f"loss {loss.numpy():7.2f} accuracy {accuracy.numpy()*100:7.2f}%, test loss {test_loss:7.2f} test accuracy {test_accuracy*100:7.2f}%")
-
-    losses.append(loss.numpy().item())
-    test_losses.append(test_loss)
-    if i % 10:
-      out = net(X_test)
-      test_loss = out.sparse_categorical_crossentropy(V_test).square().mean().numpy().item()
-      cat = out.argmax(axis=-1)
-      test_accuracy = (cat == y).mean().numpy()
-
-  safe_save(get_state_dict(net), "/tmp/policynet.safetensors")
-
-  import matplotlib.pyplot as plt
-  plt.plot(losses[10:])
-  plt.plot(test_losses[10:])
-  plt.show()

extra/optimization/extract_sa_pairs.py

@@ -1,129 +0,0 @@
-import sys, sqlite3, pickle, math
-from collections import defaultdict
-from tqdm import tqdm, trange
-import numpy as np
-
-# stuff needed to unpack a kernel
-from tinygrad.uop.ops import LazyOp, TernaryOps, BinaryOps, UnaryOps, ReduceOps, BufferOps, MemBuffer, ConstBuffer
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.uop.ops import Variable
-inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.opt.kernel import Opt, OptOps
-
-# more stuff
-from tinygrad.codegen.opt.kernel import Kernel
-from tinygrad.codegen.opt.search import actions
-from extra.optimization.helpers import lin_to_feats
-from extra.optimization.pretrain_valuenet import ValueNet
-from tinygrad.nn.optim import Adam
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-import random
-from tinygrad.tensor import Tensor
-from tinygrad.helpers import getenv
-
-def dataset_from_cache(fn):
-  conn = sqlite3.connect(fn)
-  cur = conn.cursor()
-  cur.execute("SELECT * FROM time_linearizer")
-  grouped = defaultdict(dict)
-  for f in tqdm(cur.fetchall()): grouped[f[0]][f[1:-1]] = pickle.loads(f[-1])
-
-  opts_to_outcome = {}
-
-  for ast,sk in grouped.items():
-    cnts = defaultdict(int)
-    for sks,tm in sk.items():
-      if sks[1] != 1: continue
-      opts = eval(sks[0])
-      cnts[(len(opts), sks[1])] += 1
-      opts_to_outcome[(ast, tuple(opts))] = tm
-    #print(cnts)
-
-  S,A,V = [], [], []
-  for ast,k in tqdm(opts_to_outcome):
-    if len(k) == 0: continue
-    old_tm = min(opts_to_outcome[(ast,k[:-1])])
-    new_tm = min(opts_to_outcome[(ast,k)])
-    if math.isinf(old_tm) or math.isinf(new_tm) or old_tm < 1e-9 or new_tm < 1e-9: continue
-    try:
-      lin = Kernel(eval(ast))
-    except Exception:
-      continue
-    lin.apply_opts(k[:-1])
-    act = k[-1]
-    log_ratio = math.log(old_tm/new_tm)
-    #print(f"ratio: {old_tm/new_tm:6.2f}x (log {log_ratio:5.2f}) from {str(act):50s} on {lin.colored_shape()}")
-    S.append(lin_to_feats(lin, use_sts=True))
-    A.append(actions.index(act))
-    V.append([log_ratio])  # NOTE: i have written the bug many times with this having the wrong dim
-
-  S, A, V = np.array(S), np.array(A), np.array(V, dtype=np.float32)
-  X = np.zeros((S.shape[0], S.shape[1]+len(actions)), dtype=np.float32)
-  X[:, :S.shape[1]] = S
-  X[range(S.shape[0]), S.shape[1]+A] = 1.0
-  return X, V
-
-def log_likelihood(x:Tensor, mu:Tensor, log_sigma:Tensor):
-  #print(x.shape, mu.shape, log_sigma.shape)
-  #return (x-mu).abs() * (-log_sigma).exp() + log_sigma
-  return (x-mu).square() * (-2*log_sigma).exp() / 2 + log_sigma
-
-if __name__ == "__main__":
-  if getenv("REGEN"):
-    X,V = dataset_from_cache(sys.argv[1] if len(sys.argv) > 1 else "/tmp/tinygrad_cache")
-    safe_save({"X": Tensor(X), "V": Tensor(V)}, "/tmp/dataset")
-  else:
-    ld = safe_load("/tmp/dataset")
-    X,V = ld['X'].numpy(), ld['V'].numpy()
-
-  print(X.shape, V.shape)
-  order = list(range(X.shape[0]))
-  random.shuffle(order)
-  X, V = X[order], V[order]
-
-  ratio = -512
-  X_test, V_test = Tensor(X[ratio:]), Tensor(V[ratio:])
-  X,V = X[:ratio], V[:ratio]
-  print(X.shape, V.shape)
-
-  #print(X[0], V[0])
-  #print(X[-1], V[-1])
-  print(X.shape)
-
-  net = ValueNet(X.shape[1], 2)
-  optim = Adam(get_parameters(net))
-
-  def get_minibatch(X,Y,bs):
-    xs, ys = [], []
-    #random.seed(1337)
-    for _ in range(bs):
-      sel = random.randint(0, len(X)-1)
-      xs.append(X[sel])
-      ys.append(Y[sel])
-    return Tensor(xs), Tensor(ys)
-
-  Tensor.training = True
-  losses = []
-  test_losses = []
-  test_loss = float('inf')
-  for i in (t:=trange(2000)):
-    x,y = get_minibatch(X,V,bs=256)
-    out = net(x)
-    #loss = (out-y).square().mean()
-    loss = log_likelihood(y, out[:, 0:1], out[:, 1:2]).mean()
-    optim.zero_grad()
-    loss.backward()
-    optim.step()
-    t.set_description(f"loss {loss.numpy():7.2f}, test loss {test_loss:7.2f}")
-    losses.append(loss.numpy().item())
-    test_losses.append(test_loss)
-    if i % 10: test_loss = (net(X_test)[:, 0:1]-V_test).square().mean().numpy().item()
-
-  safe_save(get_state_dict(net), "/tmp/qnet.safetensors")
-
-  import matplotlib.pyplot as plt
-  plt.plot(losses[20:])
-  plt.plot(test_losses[20:])
-  plt.show()

extra/optimization/helpers.py

@@ -1,124 +0,0 @@
-# stuff needed to unpack a kernel
-from tinygrad import Variable
-from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.dtype import dtypes, PtrDType
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.helpers import getenv
-from tinygrad.engine.realize import get_program
-inf, nan = float('inf'), float('nan')
-UOps = Ops
-
-# kernel unpacker
-from tinygrad.codegen.opt.kernel import Kernel
-def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
-def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
-def kern_str_to_lin(kern_str:str, opts=None):
-  (ast, applied_opts,) = eval(kern_str)
-  k = Kernel(ast, opts=opts)
-  k.apply_opts(applied_opts)
-  return k
-
-# load worlds, a dataset of about 12k kernels
-import gzip
-from pathlib import Path
-import random
-from tinygrad.helpers import dedup, DEBUG
-def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
-  fn = Path(__file__).parent.parent / "datasets/sops.gz"
-  ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
-  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
-  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
-  if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
-  if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
-  if filter_novariable: ast_strs = [x for x in ast_strs if "DEFINE_VAR" not in x]
-  if DEBUG >= 1: print(f"loaded {len(ast_strs)=} after filters {filter_reduce=}, {filter_noimage=}, {filter_novariable=}")
-  random.seed(1337)
-  random.shuffle(ast_strs)
-  return ast_strs
-
-def assert_same_lin(l1, l2):
-  assert l1.colored_shape() == l2.colored_shape()
-  assert all(x==y for x,y in zip(l1.sts, l2.sts))
-
-# get features
-import math
-
-MAX_DIMS = 16
-MAX_BUFS = 9
-def lin_to_feats(lin:Kernel, use_sts=True):
-  assert lin.shape_len < MAX_DIMS, "too many dims"
-
-  all_colors = ["blue", "cyan", "white", "green", "red", "magenta", "yellow"]
-  lc = [all_colors.index(x) for x in lin.colors()]
-
-  ret = []
-  # before, some generic linearizer stuff
-  ret.append(lin.upcasted)
-  ret.append(lin.local_dims)
-
-  # first, the full shape, including the colors
-  for s,os,c in zip(lin.full_shape,lin.output_shape,lc):
-    if isinstance(s, UOp):
-      ret.append(False)
-      ret += [0]*9
-    else:
-      ret.append(True)
-      ret.append(math.log2(s))
-      ret.append(min(33, s))
-      ret.append(math.log2(os))
-      ret.append(min(33, os))
-      ret.append(s%2 == 0)
-      ret.append(s%3 == 0)
-      ret.append(s%4 == 0)
-      ret.append(s%8 == 0)
-      ret.append(s%16 == 0)
-    cc = [0]*7
-    cc[c] = 1
-    ret += cc
-  ret += [0] * (17*(MAX_DIMS-len(lin.full_shape)))
-  ret = [float(x) for x in ret]
-
-  if use_sts:
-    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
-    assert len(my_sts) < MAX_BUFS
-    sts_len = 3 + 5*MAX_DIMS
-    for s in my_sts:
-      ret.append(s[0])  # reduce
-      ret.append(s[2])  # has mask
-      ret.append(s[3])  # len views
-      for d in s[1]:
-        ret.append(d is None)
-        ret.append(d == 0)
-        ret.append(d == 1)
-        ret.append(min(33, d) if d is not None else -1)
-        if d is not None and d >= 1: ret.append(math.log2(d))
-        else: ret.append(-1)
-      ret += [0] * (5*(MAX_DIMS - len(s[1])))
-    ret += [0] * (sts_len*(MAX_BUFS - len(my_sts)))
-    assert len(ret) == 1021, f"wrong len {len(ret)}"
-  else:
-    assert len(ret) == 274, f"wrong len {len(ret)}"
-  return ret
-
-from tinygrad.device import Device, Buffer
-from tinygrad.codegen.opt.search import _ensure_buffer_alloc, _time_program
-from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
-
-def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501
-  key = {"ast": lin.ast.key, "opts": str(lin.applied_opts), "allow_test_size": allow_test_size,
-         "max_global_size": max_global_size, "clear_l2": clear_l2, "device": lin.opts.device, "suffix": lin.opts.suffix}
-  if not disable_cache and CACHELEVEL >= 2 and (val:=diskcache_get("time_linearizer", key)) is not None: return min(val)
-
-  dev = Device[lin.opts.device]
-  assert dev.compiler is not None
-
-  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
-  p = get_program(lin.get_optimized_ast(), lin.opts)
-  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
-                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
-
-  if CACHELEVEL >= 2: diskcache_put("time_linearizer", key, tms)
-  return min(tms)