remove on_stack from graph_rewrite

2026-08-14 08:18:26 +00:00 · 2025-10-09 14:14:25 +08:00
511 changed files with 7885 additions and 60025 deletions
@@ -1,3 +0,0 @@
-[run]
-source = tinygrad
-branch = True
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  mesa:
-    description: "Install mesa"
-    required: false
-    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -293,13 +289,3 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
-
-    # **** mesa ****
-    - name: Install mesa (linux)
-      if: inputs.mesa == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
-    - name: Install mesa (macOS)
-      if: inputs.mesa == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -2,7 +2,7 @@ name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
+  PYTHON_CACHE_VERSION: '3'
  APT_CACHE_VERSION: '1'
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
@@ -36,9 +36,8 @@ jobs:
        cuda: 'true'
        webgpu: 'true'
        llvm: 'true'
-        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
    - name: Verify OpenCL autogen
      run: |
        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -90,8 +89,3 @@ jobs:
        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
        ./autogen_stubs.sh llvm
        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-    - name: Verify mesa autogen
-      run: |
-        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
-        ./autogen_stubs.sh mesa
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
@@ -51,18 +51,17 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
-    - name: Print macOS version
-      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
@@ -131,7 +130,7 @@ jobs:
    - 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
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -199,7 +198,7 @@ jobs:
    - name: Test speed vs torch
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
@@ -211,7 +210,6 @@ jobs:
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
-        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
    - name: Run Tensor Core GEMM (PTX)
      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
@@ -239,8 +237,6 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
-    - name: Run quantized LLaMA3
-      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -274,7 +270,6 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
-          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -323,9 +318,9 @@ jobs:
    - name: Run 10 CIFAR training steps
      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=240 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
    # - name: Run 10 CIFAR training steps w winograd
    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
@@ -409,7 +404,7 @@ jobs:
    #    python3 -c "import torch; print(torch.__version__)"
    #    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 DISABLE_COMPILER_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 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
@@ -527,7 +522,7 @@ jobs:
    - name: Run 10 CIFAR training steps
      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=390 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    # - name: Run 10 CIFAR training steps w BF16
    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
    # TODO: too slow
@@ -623,24 +618,22 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: test/external/process_replay/reset.py
+    - name: benchmark openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.9.9 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.9.9 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 0.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/dmonitoring_model.onnx
-    - name: openpilot compile3 0.10.1 driving_vision
-      # TODO: ASSERT_MIN_STEP_TIME=17
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=21 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.1 dmonitoring
-      # TODO: ASSERT_MIN_STEP_TIME=10
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 Space Lab policy + vision
+      run: |
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -648,9 +641,19 @@ jobs:
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - 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
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (comma)
+        path: |
+          openpilot_compile_0_9_4.txt
+          openpilot_compile_0_9_7.txt
+          openpilot_0_9_4.txt
+          openpilot_0_9_7.txt
+          openpilot_image_0_9_4.txt
+          openpilot_image_0_9_7.txt

  testreddriverbenchmark:
    name: AM Benchmark
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 720
+    timeout-minutes: 360

    steps:
    - name: Checkout Code
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -2,7 +2,7 @@ name: Unit Tests
 env:
  # increment this when downloads substantially change to avoid the internet
  DOWNLOAD_CACHE_VERSION: '12'
-  PYTHON_CACHE_VERSION: '4'
+  PYTHON_CACHE_VERSION: '3'
  APT_CACHE_VERSION: '1'
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
@@ -89,65 +89,64 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  # TODO: fix the torch backend and reenable
-  # torchbackend:
-  #   name: Torch Backend Tests
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       pydeps: "pillow torchvision expecttest"
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Lint with ruff
-  #     run: |
-  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
-  #       python3 -m ruff check extra/torch_backend/backend.py
-  #   - name: Test one op
-  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-  #   - name: Test ResNet-18
-  #     run: DEBUG=2 python3 extra/torch_backend/example.py
-  #   - name: My (custom) tests
-  #     run: python3 extra/torch_backend/test.py
-  #   - name: Test one op in torch tests
-  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-  #   - name: Test Ops with TINY_BACKEND
-  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-  #   - name: Test in-place operations on views
-  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-  #   - name: Test multi-gpu
-  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  torchbackend:
+    name: Torch Backend Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        pydeps: "pillow torchvision expecttest"
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Lint with ruff
+      run: |
+        pip3 install --upgrade --force-reinstall ruff==0.11.0
+        python3 -m ruff check extra/torch_backend/backend.py
+    - name: Test one op
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+    - name: Test ResNet-18
+      run: DEBUG=2 python3 extra/torch_backend/example.py
+    - name: My (custom) tests
+      run: python3 extra/torch_backend/test.py
+    - name: Test one op in torch tests
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+    - name: Test Ops with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+    - name: Test in-place operations on views
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+    - name: Test multi-gpu
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py

-  # torchbackendmore:
-  #   name: Torch Backend Tests More
-  #   runs-on: ubuntu-latest
-  #   timeout-minutes: 15
-  #   steps:
-  #   - name: Checkout Code
-  #     uses: actions/checkout@v4
-  #   - name: Setup Environment
-  #     uses: ./.github/actions/setup-tinygrad
-  #     with:
-  #       key: torch-backend-pillow-torchvision-et-pt
-  #       deps: testing_minimal
-  #       llvm: 'true'
-  #   - name: Install ninja
-  #     run: |
-  #       sudo apt update || true
-  #       sudo apt install -y --no-install-recommends ninja-build
-  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
-  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-  #   - name: Test some torch tests (expect failure)
-  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+  torchbackendmore:
+    name: Torch Backend Tests More
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: torch-backend-pillow-torchvision-et-pt
+        deps: testing_minimal
+        llvm: 'true'
+    - name: Install ninja
+      run: |
+        sudo apt update || true
+        sudo apt install -y --no-install-recommends ninja-build
+    - name: Test beautiful_mnist in torch with TINY_BACKEND
+      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+    - name: Test some torch tests (expect failure)
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

  bepython:
    name: Python Backend
@@ -204,7 +203,7 @@ jobs:
        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE=CUDA_SM89 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
    - name: Test emulated INTEL OpenCL tensor cores
      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
    - name: Test emulated AMX tensor cores
@@ -230,7 +229,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: linting-only
-        python-version: '3.11'
+        python-version: '3.10'
        deps: linting
    - name: Lint bad-indentation and trailing-whitespace with pylint
      run: python -m pylint --disable=all -e W0311 -e C0303 --jobs=0 --indent-string='  ' --recursive=y .
@@ -239,13 +238,14 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
+    - name: Lint tinygrad with pylint
+      run: python -m pylint tinygrad/
    - 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
@@ -265,6 +265,8 @@ jobs:
      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Check SPEC=1
+      run: SPEC=1 python3 test/test_tiny.py
    - name: Run targetted tests on NULL backend
      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
    # TODO: too slow
@@ -272,8 +274,6 @@ jobs:
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
-    - name: Run AMD emulated BERT training on NULL backend
-      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -290,27 +290,8 @@ jobs:
        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
-    - name: Repo line count < 18500 lines
-      run: MAX_LINE_COUNT=18500 python sz.py
-
-  spec:
-    strategy:
-      fail-fast: false
-      matrix:
-        group: [1, 2]
-    name: SPEC=2 (${{ matrix.group }})
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: spec-unit
-        deps: testing_unit
-    - 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 }}
+    - name: Repo line count < 18000 lines
+      run: MAX_LINE_COUNT=18000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -328,6 +309,10 @@ jobs:
      run: python test/external/fuzz_symbolic.py
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
+    - name: Fuzz Test shapetracker
+      run: |
+        python test/external/fuzz_shapetracker.py
+        python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

@@ -369,7 +354,7 @@ jobs:
      - name: Run Kernel Count Test
        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
      - name: Run fused optimizer tests
-        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py test/test_optim.py -k "not muon"
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -392,13 +377,17 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1452 ALLOWED_GATED_READ_IMAGE=122 FLOAT16=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
-      - name: Test openpilot CL compile fp16
-        run: FLOAT16=1 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
-      - name: Test openpilot CL compile fp32 (test correctness)
-        run: DEBUGCL=1 CL=1 IMAGE=2 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
-      - name: Test openpilot LLVM compile fp16
-        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot alt model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot fastvits model correctness (float32)
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      # - name: Test openpilot simple_plan vision model correctness (float32)
+      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
+      - name: Test openpilot LLVM compile
+        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot compile4
+        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -462,7 +451,8 @@ jobs:
      - name: Test Bert training
        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        # TODO: remove LLAMA_LAYERS once it's fast
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 LLAMA_LAYERS=4 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -540,11 +530,11 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test LLVM=1 DEVECTORIZE=0 for model
      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testdsp:
    name: Linux (DSP)
@@ -645,7 +635,6 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          VIZ=1 PMC=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
@@ -689,7 +678,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl, lvp]
+        backend: [llvm, cpu, opencl]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -703,10 +692,9 @@ jobs:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -908,7 +896,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [metal, llvm, cpu, lvp]
+        backend: [metal, llvm, cpu]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
@@ -921,13 +909,12 @@ jobs:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
          pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
-          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -38,7 +38,6 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
-venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -63,5 +62,3 @@ profile_stats
 *.log
 target
 .mypy_cache
-mutants
-.mutmut-cache
@@ -20,15 +20,21 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
+      - id: tests
+        name: subset of tests
+        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        language: system
+        always_run: true
+        pass_filenames: false
      - id: example
-        name: test all devices
+        name: multi device tests
        entry: python3 test/external/external_test_example.py
        language: system
        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
+      - id: pylint
+        name: pylint
+        entry: python3 -m pylint tinygrad/
        language: system
        always_run: true
-        pass_filenames: false
+        pass_filenames: false
@@ -30,6 +30,10 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=

+# When enabled, pylint would attempt to guess common misconfiguration and emit
+# user-friendly hints instead of false-positive error messages
+suggestion-mode=yes
+
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
@@ -31,9 +31,7 @@ $(for p in "$@"; do echo "  $p,"; done)
 ]
 def _try_dlopen_$name():
  library = ctypes.util.find_library("$name")
-  if library:
-    try: return ctypes.CDLL(library)
-    except OSError: pass
+  if library: return ctypes.CDLL(library)
  for candidate in PATHS_TO_TRY:
    try: return ctypes.CDLL(candidate)
    except OSError: pass
@@ -188,7 +186,6 @@ nv_status_codes = {}
    extra/nv_gpu_driver/g_rpc-message-header.h \
    extra/nv_gpu_driver/gsp_static_config.h \
    extra/nv_gpu_driver/vbios.h \
-    extra/nv_gpu_driver/pci_exp_table.h \
    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
    -o $BASE/nv/nv.py

@@ -417,31 +414,10 @@ generate_sqtt() {
  clang2py -k cdefstum \
    extra/sqtt/sqtt.h \
    -o $BASE/sqtt.py
+
  fixup $BASE/sqtt.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
  python3 -c "import tinygrad.runtime.autogen.sqtt"
-
-  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
-  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
-  if [ ! -d "$ROCPROF_SRC" ]; then
-    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
-    pushd .
-    cd $ROCPROF_SRC
-    git reset --hard $ROCPROF_COMMIT_HASH
-    popd
-  fi
-
-  clang2py -k cdefstum \
-    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
-    $ROCPROF_SRC/include/trace_decoder_instrument.h \
-    $ROCPROF_SRC/include/trace_decoder_types.h \
-    -o $BASE/rocprof.py
-  fixup $BASE/rocprof.py
-  sed -i '1s/^/# pylint: skip-file\n/' $BASE/rocprof.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util/g" $BASE/rocprof.py
-  patch_dlopen $BASE/rocprof.py rocprof-trace-decoder "'/usr/local/lib/librocprof-trace-decoder.so'" "'/usr/local/lib/librocprof-trace-decoder.dylib'"
-  sed -i "s/def _try_dlopen_rocprof-trace-decoder():/def _try_dlopen_rocprof_trace_decoder():/g" $BASE/rocprof.py
-  sed -i "s|FunctionFactoryStub()|_try_dlopen_rocprof_trace_decoder()|g" $BASE/rocprof.py
 }

 generate_webgpu() {
@@ -466,80 +442,6 @@ generate_libusb() {
  python3 -c "import tinygrad.runtime.autogen.libusb"
 }

-generate_mesa() {
-  MESA_TAG="mesa-25.2.4"
-  MESA_SRC=/tmp/mesa-$MESA_TAG
-  TINYMESA_TAG=tinymesa-32dc66c
-  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
-  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
-  if [ ! -d "$MESA_SRC" ]; then
-    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
-    pushd .
-    cd $MESA_SRC
-    git reset --hard $MESA_COMMIT_HASH
-    # clang 14 doesn't support packed enums
-    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
-    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
-    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
-    mkdir -p gen/util/format
-    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
-    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
-    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
-    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
-    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
-    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
-    popd
-  fi
-
-  if [ ! -d "$TINYMESA_DIR" ]; then
-    mkdir $TINYMESA_DIR
-    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
-  fi
-
-  clang2py -k cdefstu \
-    $MESA_SRC/src/compiler/nir/nir.h \
-    $MESA_SRC/src/compiler/nir/nir_builder.h \
-    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
-    $MESA_SRC/src/compiler/nir/nir_serialize.h \
-    $MESA_SRC/gen/nir_intrinsics.h \
-    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
-    $MESA_SRC/src/nouveau/compiler/nak.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
-    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
-    $MESA_SRC/src/compiler/glsl_types.h \
-    $MESA_SRC/src/util/blob.h \
-    $MESA_SRC/src/util/ralloc.h \
-    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
-    -l $TINYMESA_SO \
-    -o $BASE/mesa.py
-
-  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
-
-  fixup $BASE/mesa.py
-  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
-  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
-  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
-  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
-  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
-  echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
-  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
-  # bitfield bug in clang2py
-  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
-  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
-  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
-  sed -i '/struct_nir_builder._pack_ = 1 # source:False/d' "$BASE/mesa.py"
-  python3 -c "import tinygrad.runtime.autogen.mesa"
-}
-
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -550,6 +452,7 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
@@ -562,7 +465,6 @@ elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
-elif [ "$1" == "mesa" ]; then generate_mesa
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
 else echo "usage: $0 <type>"
 fi
@@ -53,7 +53,9 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 idx = UOp.const(dtypes.index, 0)
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-alu = buf_1.index(idx) + buf_2.index(idx)
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
+alu = ld_1 + ld_2
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
 st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))
@@ -78,6 +80,7 @@ print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.helpers import RANGEIFY
 from tinygrad.schedule.rangeify import get_rangeify_map

 # allocate some values + load in values
@@ -0,0 +1,109 @@
+# Kernel Creation
+
+Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
+
+- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
+- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
+- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
+
+`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
+
+## Kernelize flow
+
+Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
+
+```py
+# initialize 3 input buffers on the device
+a = Tensor([1]).realize()
+b = Tensor([2]).realize()
+c = Tensor([3]).realize()
+
+# create the Tensor graph
+mul = a*b
+out = mul+c
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
+
+out.kernelize()
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
+```
+
+The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
+
+```py
+print(out.uop)
+```
+
+The first source is the output BUFFER:
+
+```
+UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
+```
+
+And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
+
+```
+UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
+```
+
+KERNEL describes the compute AST, metadata and memory dependencies.
+
+BUFFER holds a reference to the device memory where the output will be stored.
+
+Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
+
+```py
+child = out+2
+child.kernelize()
+print(child.uop.src[1].arg.ast)
+```
+
+```
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
+    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
+         x2,)),
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+         x2,)),)),)),))
+```
+
+`Tensor.realize` will execute the kernels and write outputs to memory:
+
+```py
+Tensor.realize(out)
+print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
+print(out.item()) # 5
+```
+
+<hr />
+
+**Summary**
+
+- The large Tensor graph is built from a mix of data, compute and movement Ops.
+
+- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
+
+- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
+
+- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -41,7 +41,7 @@ BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -145,41 +145,6 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear

-def quantize_to_fp8(x: Tensor, dtype=dtypes.fp8e4m3):
-  fp8_min = -448.0 if dtype == dtypes.fp8e4m3 else -57344.0
-  fp8_max = 448.0 if dtype == dtypes.fp8e4m3 else 57344.0
-  scale = fp8_max / x.abs().max()
-  x_scl_sat = (x * scale).clamp(fp8_min, fp8_max)
-  return x_scl_sat.cast(dtype), scale.float().reciprocal()
-
-class FP8Linear:
-  def __init__(self, in_features, out_features, bias=True):
-    self.weight = Tensor.empty(out_features, in_features, dtype=dtypes.fp8e4m3)
-    self.bias = Tensor.empty(out_features, dtype=dtypes.float16) if bias else None
-    self.weight_scale = Tensor.empty((), dtype=dtypes.float16)
-
-  def __call__(self, x:Tensor):
-    y = x.dot(self.weight.T.cast(dtypes.float32)) * self.weight_scale
-    if self.bias is not None: y = y + self.bias.cast(y.dtype)
-    return y.cast(x.dtype)
-
-  @staticmethod
-  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
-    assert not quantize_embeds
-    new_tensors = {}
-    for name,v in tensors.items():
-      if "feed_forward" in name or "attention.w" in name:
-        assert "weight" in name, name
-        fp8_weight, scale = quantize_to_fp8(v)
-        new_tensors[name] = fp8_weight
-        new_tensors[name.replace('weight', 'weight_scale')] = scale.cast(scale_dtype)
-        if isinstance(device, tuple):
-          new_tensors[name].shard_(device, axis=-1)
-          new_tensors[name.replace('weight', 'weight_scale')].shard_(device, axis=None)
-      else:
-        new_tensors[name] = v
-    return new_tensors
-
 MODEL_PARAMS = {
  "1B": {
    "args": {"dim": 2048, "n_heads": 32, "n_kv_heads": 8, "n_layers": 16, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256, "hidden_dim": 8192},
@@ -202,7 +167,6 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
  # build model
  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
-  elif quantize == "fp8": linear, embedding, quantize_embeds = FP8Linear, nn.Embedding, False
  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)

@@ -278,7 +242,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -1188,9 +1188,7 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
-        train_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1224,7 +1222,7 @@ def train_bert():
          return

      if getenv("RESET_STEP"): eval_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
-
-export IGNORE_OOB=1
-
-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-# export BEAM_LOG_SURPASS_MAX=1
-# export BASEDIR="/raid/datasets/wiki"
-
-export RESET_STEP=1
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
-export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-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"
-
-export BENCHMARK=10 BERT_LAYERS=2
-
-python3 examples/mlperf/model_train.py
@@ -1,20 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-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=500000
-
-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"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
-#!/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_8xMI300X"
-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=500000
-
-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_8xMI300x_${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
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
-export BASEDIR="/raid/datasets/wiki"
-
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,28 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_green_${DATETIME}_${SEED}.log"
-
-# init
-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
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
-export BASEDIR="/raid/datasets/wiki"
-
-export RESET_STEP=1
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
-#!/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_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_red_${DATETIME}_${SEED}.log"
-
-export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
-
-# init
-sleep 5 && sudo rmmod amdgpu || true
-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
@@ -1,50 +0,0 @@
-# 1. Problem
-
-This problem uses the ResNet-50 CNN to do image classification.
-
-## Requirements
-
-Install tinygrad and mlperf-logging from master.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-### tinybox_red
-Disable cwsr
-This is the default on production tinybox red.
-```
-sudo vi /etc/modprobe.d/amdgpu.conf
-cat <<EOF > /etc/modprobe.d/amdgpu.conf
-options amdgpu cwsr_enable=0
-EOF
-sudo update-initramfs -u
-sudo reboot
-
-# validate
-sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
-```
-
-# 2. Directions
-
-## Steps to download and verify data
-
-```
-IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
-```
-
-## Steps for one time setup
-
-### tinybox_red
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
-```
-
-## Steps to run benchmark
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-```
@@ -1,13 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-export BENCHMARK=10 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-export EVAL_START_EPOCH=3 EVAL_FREQ=4
-
-export WANDB=1 PARALLEL=0
-
-python3 examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
-
-# init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
-# 1. Problem
-
-This problem uses the ResNet-50 CNN to do image classification.
-
-## Requirements
-
-Install tinygrad and mlperf-logging from master.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-### tinybox_red
-Disable cwsr
-This is the default on production tinybox red.
-```
-sudo vi /etc/modprobe.d/amdgpu.conf
-cat <<EOF > /etc/modprobe.d/amdgpu.conf
-options amdgpu cwsr_enable=0
-EOF
-sudo update-initramfs -u
-sudo reboot
-
-# validate
-sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
-```
-
-# 2. Directions
-
-## Steps to download and verify data
-
-```
-IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
-```
-
-## Steps for one time setup
-
-### tinybox_red
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
-```
-
-## Steps to run benchmark
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-```
@@ -1,13 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=10 DEBUG=${DEBUG:-2}
-
-python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export EVAL_START_EPOCH=3 EVAL_FREQ=4
-
-export WANDB=1 PARALLEL=0
-
-python3 examples/mlperf/model_train.py
@@ -1,26 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
-
-# init
-sleep 5 && sudo rmmod amdgpu || true
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,8 +0,0 @@
-#!/bin/bash
-
-rocm-smi --setprofile compute
-rocm-smi --setmclk 3
-rocm-smi --setperflevel high
-
-# power cap to 350W
-echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -1,38 +0,0 @@
-# 1. Problem
-
-This problem uses RetinaNet for SSD.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-Also install the following dependencies:
-```
-pip install tqdm numpy pycocotools boto3 pandas torch torchvision
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download data
-
-Run the following:
-```
-BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
-```
-
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
-```
@@ -1,14 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=5 DEBUG=2
-
-python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export WANDB=1 PARALLEL=0
-export RUNMLPERF=1
-
-python examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/openimages"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
-
-# init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,14 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=5 DEBUG=2
-
-python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export WANDB=1 PARALLEL=0
-export RUNMLPERF=1
-
-python examples/mlperf/model_train.py
@@ -1,38 +0,0 @@
-{
-    "submitter": "tinycorp",
-    "division": "closed",
-    "status": "Available on-premise",
-    "system_name": "tinybox 8xMI300X",
-    "number_of_nodes": "1",
-    "host_processors_per_node": "2",
-    "host_processor_model_name": "AMD EPYC 9354",
-    "host_processor_core_count": "32",
-    "host_processor_vcpu_count": "64",
-    "host_processor_frequency": "",
-    "host_processor_caches": "",
-    "host_processor_interconnect": "",
-    "host_memory_capacity": "2304GB",
-    "host_storage_type": "NVMe SSD",
-    "host_storage_capacity": "3x 4TB raid array",
-    "host_networking": "",
-    "host_networking_topology": "",
-    "host_memory_configuration": "24x 96GB DDR5",
-    "accelerators_per_node": "8",
-    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
-    "accelerator_host_interconnect": "PCIe 5.0 x16",
-    "accelerator_frequency": "",
-    "accelerator_on-chip_memories": "",
-    "accelerator_memory_configuration": "HBM3",
-    "accelerator_memory_capacity": "192GB",
-    "accelerator_interconnect": "",
-    "accelerator_interconnect_topology": "",
-    "cooling": "air",
-    "hw_notes": "",
-    "framework": "tinygrad, branch mlperf_training_v5.0",
-    "other_software_stack": {
-        "python": "3.10.16",
-        "ROCm": "3.0.0+94441cb"
-    },
-    "operating_system": "Ubuntu 24.04.1 LTS",
-    "sw_notes": ""
-  }
@@ -1,38 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox green",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6X",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12",
-    "CUDA": "12.4"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}
@@ -1,37 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox red",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}
@@ -1,9 +1,15 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time
 import numpy as np
+if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
+if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
+if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
+if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

 from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner
+
+import onnx
 from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
@@ -15,14 +21,11 @@ def compile(onnx_file):

  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
-
  # Float inputs and outputs to tinyjits for openpilot are always float32
-  # TODO this seems dumb
  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
  Tensor.manual_seed(100)
-  inputs = {k:Tensor(Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize().numpy(), device='NPY') for k,shp in sorted(input_shapes.items())}
-  if not getenv("NPY_IMG"):
-    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  print("created tensors")

  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -30,6 +33,8 @@ def compile(onnx_file):
  for i in range(3):
    GlobalCounters.reset()
    print(f"run {i}")
+    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
+              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
    with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
      ret = run_onnx_jit(**inputs).numpy()
    # copy i == 1 so use of JITBEAM is okay
@@ -38,7 +43,7 @@ def compile(onnx_file):
  np.testing.assert_equal(test_val, ret, "JIT run failed")
  print("jit run validated")

-  # check gated read_image usage
+  # checks from compile2
  kernel_count = 0
  read_image_count = 0
  gated_read_image_count = 0
@@ -47,8 +52,6 @@ def compile(onnx_file):
      kernel_count += 1
      read_image_count += ei.prg.p.src.count("read_image")
      gated_read_image_count += ei.prg.p.src.count("?read_image")
-      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
-        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
  print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
  if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -64,9 +67,14 @@ def compile(onnx_file):
  print(f"mdl size is {mdl_sz/1e6:.2f}M")
  print(f"pkl size is {pkl_sz/1e6:.2f}M")
  print("**** compile done ****")
-  return inputs, test_val
+  return test_val

-def test_vs_compile(run, inputs, test_val=None):
+def test_vs_compile(run, new_inputs, test_val=None):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+
+  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
+  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
+            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}

  # run 20 times
  step_times = []
@@ -83,58 +91,68 @@ def test_vs_compile(run, inputs, test_val=None):
    min_time = min(step_times)
    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"

+  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

  # test that changing the numpy changes the model outputs
-  inputs_2x = {k: Tensor(v.numpy()*2, device=v.device) for k,v in inputs.items()}
-  out = run(**inputs_2x)
-  changed_val = out.numpy()
-  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  if any([x.device == 'NPY' for x in inputs.values()]):
+    for v in new_inputs_numpy.values(): v *= 2
+    out = run(**inputs)
+    changed_val = out.numpy()
+    np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
  return val

-def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
-  import onnx
-  import onnxruntime as ort
-  
-  onnx_inputs = {k:v.numpy() for k,v in new_inputs.items()}
+def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_model = onnx.load(onnx_file)

-  ORT_TO_NP_DTYPES: dict[str, np.dtype] = {
-    'tensor(float)': np.dtype('float32'),
-    'tensor(float16)': np.dtype('float16'),
-    'tensor(uint8)': np.dtype('uint8'),
-  }
-
  timings = []
-  onnx_session = ort.InferenceSession(onnx_file)
-  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
-  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+  if ort:
+    # test with onnxruntime
+    import onnxruntime as ort
+    onnx_session = ort.InferenceSession(onnx_file)
+    for _ in range(1 if test_val is not None else 5):
+      st = time.perf_counter()
+      onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
+      timings.append(time.perf_counter() - st)
+    new_torch_out = onnx_output[0]
+  else:
+    # test with torch
+    import torch
+    from onnx2torch import convert
+    inputs = {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}
+    torch_model = convert(onnx_model).float()
+    with torch.no_grad():
+      for _ in range(1 if test_val is not None else 5):
+        st = time.perf_counter()
+        torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
+        timings.append(time.perf_counter() - st)
+      new_torch_out = torch_out.numpy()

-  for _ in range(1 if test_val is not None else 5):
-    st = time.perf_counter()
-    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
-    timings.append(time.perf_counter() - st)
-
-  np.testing.assert_allclose(onnx_output[0].reshape(test_val.shape), test_val, atol=tol, rtol=tol)
-  print("test vs onnx passed")
+  if test_val is not None:
+    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
+    print("test vs onnx passed")
  return timings

-def bench(run, inputs):
-  from extra.bench_log import WallTimeEvent, BenchEvent
-  for _ in range(10):
-    with WallTimeEvent(BenchEvent.STEP):
-      run(**inputs).numpy()
-
 if __name__ == "__main__":
  onnx_file = fetch(OPENPILOT_MODEL)
-  inputs, outputs = compile(onnx_file)
+  test_val = compile(onnx_file) if not getenv("RUN") else None

  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)

-  test_vs_compile(pickle_loaded, inputs, outputs)
-  if getenv("SELFTEST"):
-    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+  # same randomness as compile
+  Tensor.manual_seed(100)
+  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
+                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
+
+  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
+  if getenv("BENCHMARK"):
+    for be in ["torch", "ort"]:
+      try:
+        timings = test_vs_onnx(new_inputs, None, onnx_file, be=="ort")
+        print(f"timing {be}: {min(timings)*1000:.2f} ms")
+      except Exception as e:
+        print(f"{be} fail with {e}")
+  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file, getenv("ORT"))

-  if getenv("BENCHMARK_LOG", ""):
-    bench(pickle_loaded, inputs)
@@ -0,0 +1,47 @@
+import sys
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.realize import run_schedule
+
+# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+
+OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
+OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+
+if __name__ == "__main__":
+  onnx_file = fetch(OPENPILOT_MODEL)
+  run_onnx = OnnxRunner(onnx_file)
+
+  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
+  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
+  root = out.uop
+  targets = [x.uop for x in inputs.values()]
+  print(targets)
+
+  # TODO: abstract this from gradient?
+
+  # compute the target path (top down)
+  in_target_path: dict[UOp, bool] = {}
+  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
+  independent_set = {}
+  for u in root.toposort():
+    if in_target_path[u]:
+      for s in u.src:
+        if not in_target_path[s]:
+          independent_set[s] = None
+  independent = UOp.sink(*independent_set.keys())
+  kernelized = get_rangeify_map(independent)
+  independent = independent.substitute(kernelized)
+  schedule, var_vals = create_schedule_with_vars(independent)
+  run_schedule(schedule)
+
+  print("**** real ****")
+  GlobalCounters.reset()
+  out.uop = root.substitute(kernelized)
+  out.kernelize()
+
+  # realize
+  out.realize()
@@ -99,7 +99,6 @@ if __name__ == "__main__":
  parser.add_argument('--timing',      action='store_true', help="Print timing per step")
  parser.add_argument('--noshow',      action='store_true', help="Don't show the image")
  parser.add_argument('--fp16',        action='store_true', help="Cast the weights to float16")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

  N = 1
@@ -113,22 +112,19 @@ if __name__ == "__main__":

  model = StableDiffusionV2(**params)

-  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
-      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
-      weights_fn = args.weights_fn
-      if not weights_fn:
-        weights_url = args.weights_url if args.weights_url else default_weights_url
-        weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
+  default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+  weights_fn = args.weights_fn
+  if not weights_fn:
+    weights_url = args.weights_url if args.weights_url else default_weights_url
+    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))

-        load_state_dict(model, safe_load(weights_fn), strict=False)
+  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
+    load_state_dict(model, safe_load(weights_fn), strict=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
        if k.startswith("model"):
-          v.replace(v.cast(dtypes.float16))
-
-    Tensor.realize(*get_state_dict(model).values())
+          v.replace(v.cast(dtypes.float16).realize())

  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -263,16 +263,13 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

  model = StableDiffusion()

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
-      model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
-      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
+    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

 # QUANT=1 python3 examples/test_onnx_imagenet.py
 #   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
+# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

-from tinygrad import Tensor, TinyJit, Variable, nn, dtypes
+from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
 from tinygrad.helpers import getenv, fetch

@@ -244,16 +244,15 @@ def transcribe_waveform(model: Whisper, enc, waveforms, truncate=False):

  log_spec = prep_audio(waveforms, model.batch_size, truncate)
  nsample = model.decoder.max_tokens_to_sample
-  nctx = model.decoder.max_self_attn_cache_len

  def inferloop(ctx: Union[np.ndarray, List[np.ndarray]], encoded_audio):
    pos, next_tokens = 0, ctx
-    for i in range(nsample):
-      next_tokens = model.decoder(Tensor(next_tokens, dtype=dtypes.int32), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
+    for i in range((nsample-len(start_tokens))*2):
+      next_tokens = model.decoder(Tensor(next_tokens), pos, encoded_audio)[:, -1].argmax(axis=-1).numpy().astype(np.int32).reshape(-1, 1)
      next_tokens[ctx[:, -1] == eot] = eot
      ctx = np.concatenate((ctx, next_tokens), axis=1)
      pos = ctx.shape[-1] - 1
-      if (next_tokens == eot).all() or pos == nctx: break
+      if (next_tokens == eot).all(): break
    return ctx

  def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]
@@ -1,4 +0,0 @@
-# source extra/cl_android.sh
-export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
-export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
-
@@ -1,168 +1,355 @@
 from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, KernelInfo, sint, AxisType
-from tinygrad.engine.realize import ExecItem, get_runner
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, colored, prod, unwrap
+from tinygrad.shape.shapetracker import ShapeTracker, View
+from tinygrad.shape.view import strides_for_shape
+from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.codegen.opt.swizzler import merge_views, view_left
+
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])

 N = 4096
-M = K = N
 run_count = 5

-# ---------------------------
-# launch/config constants
-# ---------------------------
+BN = 128
+BM = 128
+BK = 8

-WARP_SIZE = 32
+TN = 4
+TM = 4

-# Threadblock tile sizes (block-level tile of C that a block computes)
-BLOCK_N = 128   # columns of C (N-dim) per block
-BLOCK_M = 128   # rows of C (M-dim) per block
-BLOCK_K = 8     # K-slice per block iteration
+# NOTE: this is from testgrad
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
+# src->r->view  -->   src->view->r
+def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
+  if r.tag is not None: return None
+  # confirm the input is in order
+  # TODO: replace this with a UOp that allows for nothing else then remove this
+  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
+  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"

-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+  # append the reduce shape to each of the views
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]

-is_kernel5 = getenv("K5", 0)
-THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+  # no reshape required with shrinking REDUCE_AXIS
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
+             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))

-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
+pm = PatternMatcher([
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])

-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+def rangeify_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
+  with Context(RANGEIFY=1):
+    sink = c.schedule()[-1].ast
+  #print(sink)

-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
-  assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
+  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
+  opts += [Opt(OptOps.UNROLL, 0, 8)]

-def hand_spec_kernel3():
-  # ---------------------------
-  # block indices & placeholders
-  # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
+  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))

-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+def top_spec_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  sink = c.schedule()[-1].ast
+  L = 16
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
+  sink = graph_rewrite(sink, view_left+pm)
+  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
+  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))

-  # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+def hl_spec_kernel3():
+  nbIterWaveM = 2
+  nbIterWaveN = 2

-  # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  # define buffers
+  # TODO: remove these views once the defines have a shape
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))

-  # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  # shape buffers. TODO: permutes
+  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
+  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
+  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
+  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
+  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
+  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
+  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
+  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)

-  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
-  # ---------------------------
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
+  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
+  assert len(expanded_shape) == 20
+  permute_a = list(range(len(expanded_shape)))
+  permute_b = permute_a[:]

-  # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  # this makes all the global loads match
+  # this can also be more simply done by rebinding the RANGEs
+  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
+  permute_a[17:20] = [11,12,13]
+  permute_a[11:14] = [17,18,19]
+  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
+  permute_a[2:7] = [3,4,5,6,2]

-  # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
+  permute_b[17:20] = [5,6,7]

-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)

-  # open inner k range
-  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
+  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)

-  # ---------------------------
-  # LOCAL -> REG (per-wave tiles)
-  # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
+  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
+  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))

-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+  axis_types = (
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.REDUCE, AxisType.REDUCE)

-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
+  sink = graph_rewrite(sink, merge_views)
+  return sink

-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256

-  # ---------------------------
-  # FMA: c_regs += A_col * B_row
-  # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  i = UOp.range(c_regs.size, 16)
-  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
+  nbWaves = BLOCK_SIZE // 32
+  WN = 128 if kernel5 else 64
+  WM = BN * BM // nbWaves // WN

-  # TODO: why don't these work as upcast?
-  # why if the ranges merge is it slow?!? (if you change the order on end, they will merge. big slowdown on METAL)
-  iterWaveM, yt, iterWaveN, xt = rngs = rngs_for_shape(c_regs.shape, 500)
-  sink = c_regs[*rngs].store(c_regs.after(k)[*rngs] + A_col[iterWaveM, yt] * B_row[iterWaveN, xt]).end(iterWaveM, iterWaveN, yt, xt)
+  nbWaveX = BN // WN
+  nbWaveY = BM // WM

-  # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
+  waveIndex = threadIdx_x // 32
+  waveIdx = waveIndex % nbWaveX
+  waveIdy = waveIndex // nbWaveX
+  indexInWave = threadIdx_x % 32

-  # ---------------------------
-  # REG -> GLOBAL (epilogue)
-  # ---------------------------
-  c = c.reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM,
-                WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)
-  c = c[waveIdy, :, laneIdy, :,
-        waveIdx, :, laneIdx, :]
-  sink = copy(c, c_regs.after(sink), rng=600)
+  nbThreadXPerWave = 8
+  nbThreadYPerWave = 4

-  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+  idxInWave = indexInWave % nbThreadXPerWave
+  idyInWave = indexInWave // nbThreadXPerWave

-def test_matmul(sink:UOp, N=N):
-  with Context(DEBUG=0):
-    a = Tensor.randn(N, N)
-    b = Tensor.randn(N, N)
-    hc = Tensor.empty(N, N)
-    Tensor.realize(a, b, hc)
+  nbIterWaveN = WN // (nbThreadXPerWave * TN)
+  nbIterWaveM = WM // (nbThreadYPerWave * TM)

-  ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
+  SUBWN = WN // nbIterWaveN
+  SUBWM = WM // nbIterWaveM

-  GlobalCounters.reset()
-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  # Thread mapping to read BKxBN block from A
+  rAIdx = threadIdx_x % BK
+  rAIdy = threadIdx_x // BK
+  # Thread mapping to read BNxBK block from B
+  rBIdx = threadIdx_x % BN
+  rBIdy = threadIdx_x // BN

-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    tc = (a @ b).realize()
-  with Context(DEBUG=0):
-    err = (hc - tc).square().mean().item()
-  print(f"mean squared error {err}")
-  if err > 1e-06:
-    raise RuntimeError("matmul is wrong!")
+  strideReadB = BLOCK_SIZE // BN
+  strideReadA = BLOCK_SIZE // BK
+  nbReadsB = BN * BK // BLOCK_SIZE
+  nbReadsA = BM * BK // BLOCK_SIZE
+
+  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
+  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
+
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
+
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0)
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1)
+
+  BM_As_stride = (BM+4) if kernel5 else BM
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM_As_stride, AddrSpace.LOCAL), arg=0)
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
+
+  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
+
+  i = UOp.range(c_regs.dtype.size, 16)
+  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)
+
+  if kernel4:
+    regA = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsA, AddrSpace.REG), arg=3)
+    regB = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbReadsB, AddrSpace.REG), arg=4)
+
+    # initial load from globals into locals (0)
+    kId = 0
+
+    # load from globals into locals
+    i = UOp.range(nbReadsB, 0)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId
+    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
+
+    i = UOp.range(nbReadsA, 1)
+    index_x = rAIdx + kId
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
+
+    # iterate over the middle chunk
+    kId_range = UOp.range(N//BK-1, 2)
+    kId = kId_range*BK
+
+    barrier = UOp.barrier(As_store, Bs_store)
+
+    # load from globals into registers (next round)
+    i = UOp.range(nbReadsB, 3)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)
+
+    i = UOp.range(nbReadsA, 4)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)
+
+    def inner_loop(first_range, inp_dep=()):
+      # inner unroll
+      k = UOp.range(BK, first_range+0)
+
+      # load from locals into registers
+      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      i = UOp.range(TN, first_range+2)
+      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
+      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)
+
+      iterWave = UOp.range(nbIterWaveM, first_range+3)
+      i = UOp.range(TM, first_range+4)
+      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
+      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)
+
+      # do the GEMM math
+      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
+      yt = UOp.range(TM, first_range+6)
+      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
+      xt = UOp.range(TN, first_range+8)
+      x = iterWaveN * TN + xt
+      y = iterWaveM * TM + yt
+      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+      # sketchy, this should end the kId_range but it doesn't
+      sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                              iterWaveM, iterWaveN, yt, xt, k)
+      return sink
+
+    # TODO: kId_range should endrange after a barrier
+    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
+
+    # load from registers into locals
+    i = UOp.range(nbReadsB, 14)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId + BK
+    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)
+
+    i = UOp.range(nbReadsA, 15)
+    index_x = rAIdx + kId + BK
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
+
+    # final iteration without the copy
+    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
+  else:
+    kId_range = UOp.range(N//BK, 0)
+    kId = kId_range*BK
+
+    # load from globals into locals
+    i = UOp.range(nbReadsB, 1)
+    index_x = BN * blockIdx_x + rBIdx
+    index_y = rBIdy + i * strideReadB + kId
+    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)
+
+    i = UOp.range(nbReadsA, 2)
+    index_x = rAIdx + kId
+    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
+    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)
+
+    barrier = UOp.barrier(As_store, Bs_store)
+
+    k = UOp.range(BK, 3)
+
+    # load from locals into registers
+    iterWave = UOp.range(nbIterWaveN, 4)
+    i = UOp.range(TN, 5)
+    index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
+    B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
+
+    iterWave = UOp.range(nbIterWaveM, 6)
+    i = UOp.range(TM, 7)
+    index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
+    A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)
+
+    # do the GEMM math
+    iterWaveM = UOp.range(nbIterWaveM, 8)
+    yt = UOp.range(TM, 9)
+    iterWaveN = UOp.range(nbIterWaveN, 10)
+    xt = UOp.range(TN, 12)
+    x = iterWaveN * TN + xt
+    y = iterWaveM * TM + yt
+    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
+    sink = c_regs_idx.store(c_regs_idx.load(init_store) + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
+                            iterWaveM, iterWaveN, yt, xt, k, kId_range)
+
+  # store c_regs into c
+  iterWaveM = UOp.range(nbIterWaveM, 1000)
+  yt = UOp.range(TM, 1001)
+  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  xt = UOp.range(TN, 1003)
+  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
+  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
+  indexC = N * (yOut + yt) + xOut + xt
+  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink),
+                         iterWaveM, iterWaveN, yt, xt)
+
+  return sink.sink(arg=KernelInfo(name="tinygemm"))

 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  HL = getenv("HL")
+  if HL == 3: hprg = rangeify_kernel3()
+  elif HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
+  if HL == 3:
+    with Context(RANGEIFY=1, BLOCK_REORDER=0):
+      prg = get_program(hprg, Device.default.renderer)
+  else:
+    prg = get_program(hprg, Device.default.renderer)
+  print(prg.src)
+  if getenv("SRC"): exit(0)
+  hrunner = CompiledRunner(prg)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  hc = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
+  ei = ExecItem(hrunner, buffers)
+  with Context(DEBUG=2):
+    for _ in range(run_count): ei.run(wait=True)
+  err = (hc-tc).square().mean().item()
+  print(f"hrunner {err}")
+  if err > 1e-06: raise RuntimeError("matmul is wrong!")
@@ -1,11 +1,17 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
+from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program

 # for copied uops
-from tinygrad import dtypes
-from tinygrad.dtype import DTYPES_DICT
+from tinygrad.codegen.opt.kernel import Kernel, KernelOptError
+from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
+from tinygrad.codegen.opt.search import Opt, OptOps
+from tinygrad import Device, dtypes, Tensor
+from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View

 script_dir = os.path.dirname(os.path.abspath(__file__))

@@ -47,6 +53,12 @@ def randoms():
    nc = nc.astype(np.bfloat16 if DTYPE_IN == dtypes.bfloat16 else np.float16)
  return na, nb, nc

+def ast_to_cuda_prog(compiler, ast, opts):
+  k = Kernel(ast)
+  k.apply_opts(opts)
+  p = get_program(k.ast, k.opts, k.applied_opts)
+  return CUDAProgram(device, p.function_name, compiler.compile(p.src))
+
 if __name__ == "__main__":
  print(f"gemm variation: {GEMM_VARIATION=} {M=} {N=} {K=} {DTYPE_IN=} {DTYPE_OUT=} {DTYPE_ACC=}")
  prog, global_size, local_size = None, None, None
@@ -177,11 +189,11 @@ if __name__ == "__main__":

    tms = []
    na, nb, nc = randoms()
-    cudaalloc._copyin(a, memoryview(bytearray(na)))
-    cudaalloc._copyin(b, memoryview(bytearray(nb)))
+    cudaalloc.copyin(a, bytearray(na))
+    cudaalloc.copyin(b, bytearray(nb))
    for i in range(CNT):
      tms.append(prog(*args, **kwargs))
-    cudaalloc._copyout(flat_mv(nc.data), c)
+    cudaalloc.copyout(flat_mv(nc.data), c)
    comp = na.astype(np.float32) @ nb.astype(np.float32)
    result = nc.reshape(M, N).astype(np.float32)

@@ -1,42 +0,0 @@
-from tinygrad import UOp, dtypes
-from tinygrad.uop.ops import AxisType, Ops, KernelInfo, AddrSpace
-from extra.gemm.amd_uop_matmul import test_matmul
-
-N = 2048
-
-# metal has an 8x8 tensor core. this is the indexing
-def mat_idx(buf, g0, g1, warp, u):
-  l = [(warp//2**i)%2 for i in range(5)]
-  return buf[g0, l[4]*4 + l[2]*2 + l[1], g1, l[3]*4 + l[0]*2 + u]
-
-def hand_spec_tc_cores():
-  gx = UOp.special(N // 8, "gidx0")
-  gy = UOp.special(N // 8, "gidx1")
-  warp = UOp.special(32, "lidx0")
-
-  c = UOp.placeholder((N, N), dtypes.float, slot=0).reshape((N//8, 8, N//8, 8))
-  a = UOp.placeholder((N, N), dtypes.float, slot=1).reshape((N//8, 8, N//8, 8))
-  b = UOp.placeholder((N, N), dtypes.float, slot=2).reshape((N//8, 8, N//8, 8))
-
-  gk = UOp.range(N // 8, 0, AxisType.REDUCE)
-
-  a_tc = UOp.vectorize(*[mat_idx(a, gx, gk, warp, i) for i in range(2)])
-  b_tc = UOp.vectorize(*[mat_idx(b, gk, gy, warp, i) for i in range(2)])
-
-  acc = UOp.placeholder((2,), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  acc = acc[0].set(0.0)
-  acc = acc[1].set(0.0)
-
-  # TODO: make this simple
-  wmma_arg = ('WMMA_8_8_8_float_float', (8, 8, 8), dtypes.float, dtypes.float, 'METAL', 32, (((3, 2),), ((3, 2),), ((3, 2),)), ())
-
-  acc_load = UOp.vectorize(acc.after(gk)[0], acc.after(gk)[1])
-  out = UOp(Ops.WMMA, dtypes.float.vec(2), (a_tc, b_tc, acc_load), arg=wmma_arg)
-
-  end_loop = UOp.group(*[acc[i].store(out.gep(i)) for i in range(2)]).end(gk)
-
-  sink = UOp.group(*[mat_idx(c.after(end_loop), gx, gy, warp, i).store(acc[i]) for i in range(2)])
-  return sink.sink(arg=KernelInfo(name="custom_metal_matmul", opts_to_apply=())).simplify()
-
-if __name__ == "__main__":
-  test_matmul(hand_spec_tc_cores(), N=N)
@@ -1,229 +0,0 @@
-import os
-import numpy as np
-np.set_printoptions(linewidth=1000000)
-os.environ["AMD_LLVM"] = "0"
-
-from tinygrad import Tensor, Context, dtypes, UOp, GlobalCounters
-from tinygrad.helpers import DEBUG, getenv
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-
-WARP_SIZE = 64
-
-# Reg tile sizes (tensor cores)
-TC_M = 16
-TC_N = 16
-TC_K = 32
-
-# 1024 matrix cores
-# 16 cycle mfma
-# 2.2 GHz
-# 16x16x32x2 FLOPS/mma = 16384
-# 2.2*1e9*16384*1024/16*1e-12 TFLOPS = 2306 TFLOPS
-
-#N,M,K = 256,256,64
-N,M,K = 4096,4096,4096
-
-# Threadblock tile sizes (block-level tile of C that a block computes)
-#BLOCK_M = 128   # rows of C (M-dim) per block
-#BLOCK_N = 128   # columns of C (N-dim) per block
-#BLOCK_K = 128   # K-slice per block iteration
-
-BLOCK_M = 64
-BLOCK_N = 64
-BLOCK_K = 128
-
-WARPGROUP_SIZE = 1
-BLOCK_M = BLOCK_M * WARPGROUP_SIZE
-
-# TODO: improve the syntax of this. better syntax, faster iteration
-#  -- DONE: add working slice a[gx, :, i] -> shape of the : (aka (16,16,32) becomes (16,))
-#  -- DONE(ish): add argfix to movement (traits shared with Tensor)
-#  -- fix WMMA to not require all the junk
-#  -- improve syntax for vectorized loads/stores (both with DEVECTORIZE and without)
-#  -- DONE: be able to use CONTRACT on a range
-#  -- fix upcasted RANGE on an already vectorized buffer
-#  -- improve "all ranges not ended error" / fix the bug with after on ended ranges (if you are after end of range, range is closed)
-
-CUS_PER_GPU = 256
-assert ((M//BLOCK_M) * (N//BLOCK_N)) >= CUS_PER_GPU, "not enough globals"
-
-def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  # A = (M x K)
-  # B = (K x N)
-  # C = (M x N)
-
-  # check it's proper matmul
-  assert C.shape[0] == A.shape[0]
-  assert C.shape[1] == B.shape[1]
-  assert A.shape[1] == B.shape[0]
-
-  gx, gy = UOp.special(M//BLOCK_M, "gidx0"), UOp.special(N//BLOCK_N, "gidx1")
-  warp = UOp.special(WARP_SIZE, "lidx0")
-  warpgroup = UOp.special(WARPGROUP_SIZE, "lidx1")
-
-  # generic copy logic (not good)
-  def generic_copy(glbl, gargs, lcl, rng):
-    # Fully coalesced 128-bit loads/stores.
-    INNER_SIZE = 8
-    cp_i = UOp.range(lcl.size//(WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE), rng)
-    cp_inner = UOp.range(INNER_SIZE, rng+1, AxisType.UPCAST)
-    idx_i = cp_i*WARPGROUP_SIZE*WARP_SIZE*INNER_SIZE + warpgroup*WARP_SIZE*INNER_SIZE + warp*INNER_SIZE + cp_inner
-    return lcl[idx_i].store(glbl[*gargs, idx_i]).end(cp_i, cp_inner)
-
-  # split out the globals into blocks
-  C = C.reshape((M//BLOCK_M, BLOCK_M, N//BLOCK_N, BLOCK_N))
-  A = A.reshape((M//BLOCK_M, BLOCK_M, K//BLOCK_K, BLOCK_K))
-  B = B.reshape((K//BLOCK_K, BLOCK_K, N//BLOCK_N, BLOCK_N))
-
-  # this is the big accumulator
-  acc = UOp.placeholder((BLOCK_N//TC_N, BLOCK_M//TC_M//WARPGROUP_SIZE), dtypes.float.vec(4), 0, AddrSpace.REG)
-  assert acc.size*WARP_SIZE*WARPGROUP_SIZE*4 == BLOCK_M*BLOCK_N
-  acc = acc[init_l:=UOp.range(acc.size, 500)].set(UOp.const(dtypes.float.vec(4), 0.0), end=init_l)
-
-  # create locals (note A is permuted, and the stride is changed to avoid bank conflicts)
-  def make_locals(slot) -> tuple[UOp, UOp]:
-    BM_As_stride = (BLOCK_M + 1)
-    BN_Bs_stride = (BLOCK_N + 0)
-    INNER_SLICE = 8
-    As = UOp.placeholder((BLOCK_K//INNER_SLICE, BM_As_stride, INNER_SLICE), dtypes.half, slot=slot, addrspace=AddrSpace.LOCAL)
-    INNER_SLICE = 1
-    Bs = UOp.placeholder((BLOCK_K//INNER_SLICE, BN_Bs_stride, INNER_SLICE), dtypes.half, slot=slot+1, addrspace=AddrSpace.LOCAL)
-    As = As.permute((0,2,1)).reshape((BLOCK_K, BM_As_stride)).shrink_to((BLOCK_K, BLOCK_M))
-    Bs = Bs.permute((0,2,1)).reshape((BLOCK_K, BN_Bs_stride)).shrink_to((BLOCK_K, BLOCK_N))
-    return As, Bs
-
-  # load from globals into locals (TODO: use the warpgroup)
-
-  def load_to_locals(l_K_outer_loop:UOp, Asl:UOp, Bsl:UOp, rng:int, barrier=True) -> tuple[UOp, UOp]:
-    if getenv("FAKE"):
-      return Asl[0].set(0), Bsl[0].set(0)
-    else:
-      pA = A.permute((0,2,1,3)).reshape((M//BLOCK_M, K//BLOCK_K, BLOCK_M*BLOCK_K))
-      pas = Asl.permute((1,0)).reshape((BLOCK_M*BLOCK_K,))
-      As_store = generic_copy(pA, (gx, l_K_outer_loop), pas, rng)
-
-      pB = B.permute((0,2,1,3)).reshape((K//BLOCK_K, N//BLOCK_N, BLOCK_K*BLOCK_N))
-      pbs = Bsl.reshape((BLOCK_K*BLOCK_N,))
-      Bs_store = generic_copy(pB, (l_K_outer_loop, gy), pbs, rng+2)
-
-      barrier = UOp.barrier(As_store, Bs_store) if barrier else UOp.group(As_store, Bs_store)
-      return Asl.after(barrier), Bsl.after(barrier)
-
-  def compute_on_locals(acc:UOp, Asl:UOp, Bsl:UOp, rng:int, afters:tuple[UOp, ...]=()) -> UOp:
-    K_inner_loop = UOp.range(BLOCK_K//TC_K, rng, AxisType.REDUCE)
-
-    # load from locals into registers
-    Ar = UOp.placeholder((BLOCK_M//TC_M//WARPGROUP_SIZE,), dtypes.half.vec(8), slot=1, addrspace=AddrSpace.REG)
-    Br = UOp.placeholder((BLOCK_N//TC_N,), dtypes.half.vec(8), slot=2, addrspace=AddrSpace.REG)
-
-    M_load_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+10)
-    Asl = Asl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_M//TC_M//WARPGROUP_SIZE, WARPGROUP_SIZE, TC_M))
-    load_rng = UOp.range(8, rng+11, axis_type=AxisType.UPCAST)
-    A_in = Asl[K_inner_loop, (warp//16)*8+load_rng, M_load_loop, warpgroup, warp%16].contract(load_rng)
-    Ar = Ar[M_load_loop].set(A_in, end=M_load_loop)
-
-    N_load_loop = UOp.range(BLOCK_N//TC_N, rng+20)
-    Bsl = Bsl.reshape((BLOCK_K//TC_K, TC_K, BLOCK_N//TC_N, TC_N))
-    load_rng = UOp.range(8, rng+21, axis_type=AxisType.UPCAST)
-    B_in = Bsl[K_inner_loop, (warp//16)*8+load_rng, N_load_loop, warp%16].contract(load_rng)
-    Br = Br[N_load_loop].set(B_in, end=N_load_loop)
-
-    M_inner_loop = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, rng+30)
-    N_inner_loop = UOp.range(BLOCK_N//TC_N, rng+31)
-
-    # load values
-    acc_after = acc.after(*afters, M_inner_loop, N_inner_loop, K_inner_loop)
-    acc_load = acc_after[N_inner_loop, M_inner_loop]
-
-    # do WMMA
-    wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-    out = UOp(Ops.WMMA, dtypes.float.vec(4), (Ar[M_inner_loop], Br[N_inner_loop], acc_load), arg=wmma_arg)
-
-    # store back the acc
-    acc_store = acc[N_inner_loop, M_inner_loop].store(out)
-    return acc_store.end(M_inner_loop, N_inner_loop, K_inner_loop)
-
-  # **** START INNER LOOP *****
-  # inner loop -- locals -> regs
-
-  # no pipeline
-  if not getenv("PIPELINE"):
-    As, Bs = make_locals(slot=0)
-
-    K_outer_loop = UOp.range(K//BLOCK_K, 0, AxisType.REDUCE)
-    As, Bs = load_to_locals(K_outer_loop, As, Bs, 1000, barrier=True)
-    acc_store = compute_on_locals(acc, As, Bs, 1500, afters=(K_outer_loop,))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-  else:
-    # this doesn't work
-    As0, Bs0 = make_locals(slot=0)
-    As1, Bs1 = make_locals(slot=2)
-    As0, Bs0 = load_to_locals(0, As0, Bs0, 1000)
-
-    K_outer_loop = UOp.range((K//BLOCK_K-2)//2, 0, AxisType.REDUCE)
-    As1, Bs1 = load_to_locals(K_outer_loop+1, As1, Bs1, 2000, barrier=False)
-    acc_store = compute_on_locals(acc, As0, Bs0, 1500, afters=(K_outer_loop,))
-    As0, Bs0 = load_to_locals(K_outer_loop+2, As0, Bs0, 3000, barrier=False)
-    acc_store = compute_on_locals(acc, As1, Bs1, 2500, afters=(acc_store, As0, Bs0))
-    acc = acc.after(acc_store.barrier().end(K_outer_loop))
-
-    #acc_store = compute_on_locals(acc, As0, Bs0, 3500, afters=(acc_store.barrier().end(K_outer_loop)))
-    """
-    As1, Bs1 = load_to_locals(K//BLOCK_K-1, As1, Bs1, 4000)
-    acc_store = compute_on_locals(acc, As1, Bs1, 4500, afters=(acc_store))
-    """
-    #acc = acc.after(acc_store)
-
-  # **** END LOOPS *****
-
-  # store the acc into gmem
-  cp_i, cp_j = UOp.range(BLOCK_M//TC_M//WARPGROUP_SIZE, 10004), UOp.range(BLOCK_N//TC_N, 10005)
-  c_load = lambda i: C[gx, cp_i*TC_M*WARPGROUP_SIZE + warpgroup*TC_M + (warp//16)*4+i, gy, cp_j*TC_N + warp%16]
-  store = UOp.group(*[c_load(i).store(acc[cp_j, cp_i].gep(i)) for i in range(4)])
-  store = store.end(cp_i, cp_j)
-
-  return store.sink(arg=KernelInfo(name="custom_gemm", opts_to_apply=())).simplify()
-
-# simplest WMMA
-"""
-# init the acc
-acc = UOp.placeholder((4,), dtypes.float, 0, AddrSpace.REG)
-acc = acc[init_l:=UOp.range(4, 1)].set(0.0, end=init_l)
-
-# do the wmma
-acc_load = UOp.vectorize(*[acc.after(K_loop)[i] for i in range(4)])
-wmma_arg = ('WMMA_16_16_32_half_float', (16, 16, 32), dtypes.half, dtypes.float, 'AMD', 64, ((), (), ((3, 2), (2, 2))), ())
-out = UOp(Ops.WMMA, dtypes.float.vec(4), (A_in, B_in, acc_load), arg=wmma_arg)
-
-# store back the acc
-acc = acc.after(UOp.group(*[acc[i].store(out.gep(i)) for i in range(4)]).end(K_loop))
-
-# store the acc into gmem
-store = UOp.group(*[C[gx, (warp//16)*4+i, gy, warp%16].store(acc[i]) for i in range(4)])
-"""
-
-if __name__ == "__main__":
-  a = Tensor.randn(M, K, dtype=dtypes.half)
-  b = Tensor.randn(K, N, dtype=dtypes.half)
-
-  #a = Tensor.zeros(M, K, dtype=dtypes.half).contiguous()
-  #a[0,16] = 1
-  #b = Tensor.ones(K, N, dtype=dtypes.half).contiguous()
-
-  c = Tensor.empty(M, N, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(a,b)
-
-  ref = a.dot(b, dtype=dtypes.float)
-  ref.realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2):
-    tst = Tensor.custom_kernel(c, a, b, fxn=custom_gemm)[0]
-    tst.realize()
-  print(f"{(N*M*K*2 / GlobalCounters.time_sum_s)*1e-12:.2f} REAL TFLOPS")
-
-  with Context(DEBUG=0):
-    #print(ref.numpy())
-    #print(tst.numpy())
-    assert Tensor.isclose(ref, tst, atol=1e-2).all().item(), "matrix not close"
@@ -5,10 +5,8 @@ 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)
-acc_dtype = (dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else
-            dtypes.fp8e4m3 if getenv("ACC_FP8E4M3") else dtypes.fp8e5m2 if getenv("ACC_FP8E5M2") else None)
+dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
+acc_dtype = dtypes.half if getenv("ACC_HALF") else dtypes.bfloat16 if getenv("ACC_BFLOAT16") else None
 if getenv("INT"):  dtype_in, acc_dtype = dtypes.int8, dtypes.int32
 if getenv("UINT"): dtype_in, acc_dtype = dtypes.uint8, dtypes.int32

@@ -16,10 +14,8 @@ N = getenv("N", 4096)
 M = getenv("M", N)
 K = getenv("K", N)
 CNT = getenv("CNT", 10)
-
-atol, rtol = {dtypes.half:{1e-3, 1e-2}, dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype_in, (1e-4, 3e-2))
-ATOL, RTOL = getenv("ATOL", atol), getenv("RTOL", rtol)
-
+ATOL = getenv("ATOL", 1e-4)
+RTOL = getenv("RTOL", 3e-2)
 INT_LOW = getenv("INT_LOW", 0)
 INT_HIGH = getenv("INT_HIGH", 10)

@@ -8,23 +8,19 @@ import torch
 torch.set_num_threads(1)
 from tinygrad.helpers import getenv
 CUDA = getenv("CUDA", 1)
-MPS = getenv("MPS", 0)
-if getenv("FP16_ACC"): torch.backends.cuda.matmul.allow_fp16_accumulation = True

-for dtype in [torch.float32, torch.float16, torch.bfloat16]:
+for dtype in [torch.float32, torch.float16]:
  for N in [256, 512, 1024, 2048, 4096]:
    FLOPS = N*N*N*2

    b = torch.rand((N,N), dtype=dtype)
    c = torch.rand((N,N), dtype=dtype)
    if CUDA: b,c = b.cuda(),c.cuda()
-    if MPS: b,c = b.to('mps'),c.to('mps')

    def torch_prog(b, c):
      st = time.perf_counter()
      a = b@c
      if CUDA: torch.cuda.synchronize()
-      if MPS: torch.mps.synchronize()
      return time.perf_counter() - st
    tm = min([torch_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}")
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3

-import argparse, glob, os, time, subprocess, sys
+from tinygrad.runtime.support.system import System
+import argparse, glob, os, re, time, subprocess, sys

 def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
  target_dev = args.pci_bus if 'pci_bus' in args.__dir__() else ""
@@ -11,7 +12,7 @@ def scan_devs_based_on_lock(prefix:str, args) -> list[str]:
    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}") and dev_id.startswith(target_dev): devs.append(dev_id)
  return devs

-def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/devices/{pci_bus}/reset'")
+def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
 def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")

 def cmd_remove_module(args):
@@ -1,23 +0,0 @@
-#!/bin/sh
-
-if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
-  echo "usage: $0 MESA_PREFIX"
-  exit 1
-fi
-
-TMP=$(mktemp)
-trap 'rm -f "$TMP"' EXIT
-
-(
-  cat <<EOF
-#define HAVE_ENDIAN_H
-#define HAVE_STRUCT_TIMESPEC
-#define HAVE_PTHREAD
-#include <unistd.h>
-#include "nir_shader_compiler_options.h"
-#include "compiler/shader_enums.h"
-EOF
-  sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
-  echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
-) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
-
@@ -7,34 +7,31 @@ import os
 NUM_WORKGROUPS = 96
 WAVE_SIZE = 32
 NUM_WAVES = 2
-FLOPS_PER_MATMUL = 16*16*16*2
-INTERNAL_LOOP = 1_000_00
-INSTRUCTIONS_PER_LOOP = 200
-DIRECTIVE = ".amdhsa_wavefront_size32 1"
+FLOPS_PER_MATMUL =  16*16*16*2
+INTERNAL_LOOP  = 1_000_000
+INSTRUCTIONS_PER_LOOP = 1_000

 assemblyTemplate = (pathlib.Path(__file__).parent / "template.s").read_text()

-def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, extra=""):
-  if accum:
-    instructions = "{} a[0:{}], v[{}:{}], v[{}:{}], 1{}\n".format(instruction, vgprIndices[0],
-                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2], extra)
-  elif dense:
+def launchBenchmark(instruction, vgprIndices, dense = True):
+  if dense:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], 1\n".format(instruction, vgprIndices[0],
                                                                vgprIndices[1], vgprIndices[2],
-                                                                vgprIndices[1], vgprIndices[2])
+                                                                vgprIndices[1], vgprIndices[2]) * INSTRUCTIONS_PER_LOOP
  else:
    instructions = "{} v[0:{}], v[{}:{}], v[{}:{}], v{}\n".format(instruction, vgprIndices[0],
-                                                                  vgprIndices[1], vgprIndices[2],
-                                                                  vgprIndices[3], vgprIndices[4],
-                                                                  vgprIndices[5])
-  src = assemblyTemplate.replace("INTERNAL_LOOP", str(INTERNAL_LOOP)).replace("INSTRUCTION", instructions*INSTRUCTIONS_PER_LOOP)
-  src = src.replace("DIRECTIVE", DIRECTIVE)
+                                                                vgprIndices[1], vgprIndices[2],
+                                                                vgprIndices[3], vgprIndices[4],
+                                                                vgprIndices[5]) * INSTRUCTIONS_PER_LOOP
+  src = assemblyTemplate.replace("INSTRUCTION", instructions)
  lib = COMPILER.compile(src)
  fxn = AMDProgram(DEV, "matmul", lib)
-  elapsed = fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True)
+  start = time.perf_counter()
+  fxn(global_size=(NUM_WORKGROUPS,1,1), local_size=(WAVE_SIZE*NUM_WAVES,1,1), wait=True) #For some reason the returned time is very small after the first kernel execution
+  end = time.perf_counter()
+  elapsed = end-start
  FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
-  print(f"{instruction:<29} : {FLOPs/elapsed/10**12:.2f} T(FL)OPS")
+  print("{:<29} : {} T(FL)OPS".format(instruction, round(FLOPs/elapsed/10**12, 2)))

 if __name__=="__main__":
  DEVICENUM = os.getenv("DEVICENUM", "0")
@@ -43,17 +40,18 @@ if __name__=="__main__":
  except:
    raise RuntimeError("Error while initiating AMD device")

-  COMPILER = HIPCompiler(DEV.arch)
-  if DEV.arch in {'gfx1100', 'gfx1103'}:
-    if DEV.arch == 'gfx1103':
-      NUM_WORKGROUPS = 8
+  if (ARCH := DEV.arch) not in ['gfx1100', 'gfx1201']:
+    raise RuntimeError("only gfx1100 and gfx1201 supported")
+  COMPILER = HIPCompiler(ARCH)
+
+  if ARCH == 'gfx1100':
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_bf16", (7,8,15))
    launchBenchmark("v_wmma_f32_16x16x16_f16", (7,8,15))
    launchBenchmark("v_wmma_i32_16x16x16_iu4", (7,8,9))
    launchBenchmark("v_wmma_i32_16x16x16_iu8", (7,8,11))
-  elif DEV.arch == 'gfx1201':
+  if ARCH == 'gfx1201':
    NUM_WORKGROUPS = 64
    launchBenchmark("v_wmma_bf16_16x16x16_bf16", (3,4,7))
    launchBenchmark("v_wmma_f16_16x16x16_f16", (3,4,7))
@@ -78,20 +76,4 @@ if __name__=="__main__":
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_fp8", (7,8,9,10,13,14), False)
    launchBenchmark("v_swmmac_f32_16x16x32_bf8_bf8", (7,8,9,10,13,14), False)
    FLOPS_PER_MATMUL = 16*16*64*2
-    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
-  elif DEV.arch == 'gfx950':
-    DIRECTIVE = ".amdhsa_accum_offset 4"
-    NUM_WORKGROUPS = 256
-    WAVE_SIZE = 64
-    NUM_WAVES = 4
-    launchBenchmark("v_mfma_f32_16x16x16_f16", (3,0,1), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x16_bf16", (3,0,1), accum=True)
-    FLOPS_PER_MATMUL = 16*16*32*2
-    launchBenchmark("v_mfma_f32_16x16x32_f16", (3,0,3), accum=True)
-    launchBenchmark("v_mfma_f32_16x16x32_bf16", (3,0,3), accum=True)
-    FLOPS_PER_MATMUL = 16*16*128*2
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,7), accum=True) # fp8
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,5), accum=True, extra=", cbsz:2 blgp:2") # fp6
-    launchBenchmark("v_mfma_f32_16x16x128_f8f6f4", (3,0,3), accum=True, extra=", cbsz:4 blgp:4") # fp4
-  else:
-    raise RuntimeError(f"arch {DEV.arch} not supported.")
+    launchBenchmark("v_swmmac_i32_16x16x64_iu4", (7,8,9,10,13,14), False)
@@ -1,9 +1,9 @@
    .text
    .globl matmul
-    .p2align 8
+    .p2align 8 
    .type matmul,@function
 matmul:
-    s_mov_b32 s1, INTERNAL_LOOP
+    s_mov_b32 s1, 1000000
    s_mov_b32 s2, 0
    inner_loop:
        INSTRUCTION
@@ -17,7 +17,7 @@ matmul:
 .amdhsa_kernel matmul
  .amdhsa_next_free_vgpr .amdgcn.next_free_vgpr
  .amdhsa_next_free_sgpr .amdgcn.next_free_sgpr
-  DIRECTIVE
+  .amdhsa_wavefront_size32 1
 .end_amdhsa_kernel

 .amdgpu_metadata
@@ -28,7 +28,7 @@ amdhsa.version:
 amdhsa.kernels:
  - .name: matmul
    .symbol: matmul.kd
-    .kernarg_segment_size: 0
+    .kernarg_segment_size:  0
    .group_segment_fixed_size: 0
    .private_segment_fixed_size: 0
    .kernarg_segment_align: 4
@@ -36,5 +36,6 @@ amdhsa.kernels:
    .sgpr_count: 8
    .vgpr_count: 32
    .max_flat_workgroup_size: 1024
+    .args:
 ...
 .end_amdgpu_metadata
@@ -89,20 +89,6 @@ class Attention:
    keys, values = repeat_kv(keys, self.n_rep), repeat_kv(values, self.n_rep)
    xq, keys, values = xq.transpose(1, 2), keys.transpose(1, 2), values.transpose(1, 2)
    attn = xq.scaled_dot_product_attention(keys, values, mask).transpose(1, 2)
-    if getenv("STUB_ATTENTION"):
-      # TODO: do we need mask?
-      from tinygrad.uop.ops import UOp, KernelInfo
-      def fa_custom_forward(attn:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-        return UOp.sink(arg=KernelInfo(name="fa_custom_forward"))
-      def fa_custom_backward(out_q:UOp, out_k:UOp, out_v:UOp, grad:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-        return UOp.sink(arg=KernelInfo(name="fa_custom_backward"))
-      def fa_backward(grad:UOp, kernel:UOp) -> tuple[None, UOp, UOp, UOp]:
-        grad_q = Tensor.empty_like(q:=Tensor(kernel.src[1]))
-        grad_k = Tensor.empty_like(k:=Tensor(kernel.src[2]))
-        grad_v = Tensor.empty_like(v:=Tensor(kernel.src[3]))
-        ck = Tensor.custom_kernel(grad_q, grad_k, grad_v, Tensor(grad), q, k, v, fxn=fa_custom_backward)[:3]
-        return (None, ck[0].uop, ck[1].uop, ck[2].uop)
-      attn = Tensor.empty_like(attn).custom_kernel(xq, keys, values, fxn=fa_custom_forward, grad_fxn=fa_backward)[0]
    attn = attn.reshape(bsz, seqlen, -1)
    return self.wo(attn)

@@ -1,134 +0,0 @@
-/*
- * SPDX-FileCopyrightText: Copyright (c) 1993-2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
- * SPDX-License-Identifier: MIT
- *
- * 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 PCIEXPTBL_H
-#define PCIEXPTBL_H
-
-#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_BASE   0x00
-#define NV_BCRT_HASH_INFO_BASE_CODE_TYPE_VBIOS_EXT    0xE0
-
-//
-// The VBIOS object comes from walking the PCI expansion code block
-// The following structure holds the expansion code format.
-//
-#define PCI_EXP_ROM_SIGNATURE     0xaa55
-#define PCI_EXP_ROM_SIGNATURE_NV  0x4e56 // "VN" in word format
-#define PCI_EXP_ROM_SIGNATURE_NV2 0xbb77
-#define IS_VALID_PCI_ROM_SIG(sig) ((sig == PCI_EXP_ROM_SIGNATURE) ||    \
-                                   (sig == PCI_EXP_ROM_SIGNATURE_NV) || \
-                                   (sig == PCI_EXP_ROM_SIGNATURE_NV2))
-
-#define OFFSETOF_PCI_EXP_ROM_SIG                    0x0
-#define OFFSETOF_PCI_EXP_ROM_NBSI_DATA_OFFSET       0x16
-#define OFFSETOF_PCI_EXP_ROM_PCI_DATA_STRUCT_PTR    0x18
-
-#pragma pack(1)
-typedef struct _PCI_EXP_ROM_STANDARD
-{
-    NvU16       sig;                //  00h: ROM Signature 0xaa55
-    NvU8        reserved [0x16];    //  02h: Reserved (processor architecture unique data)
-    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
-    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
-} PCI_EXP_ROM_STANDARD, *PPCI_EXP_ROM_STANDARD;
-#pragma pack()
-
-#pragma pack(1)
-typedef struct _PCI_EXP_ROM_NBSI
-{
-    NvU16       sig;                //  00h: ROM Signature 0xaa55
-    NvU8        reserved [0x14];    //  02h: Reserved (processor architecture unique data)
-    NvU16       nbsiDataOffset;     //  16h: Offset from header to NBSI image
-    NvU16       pciDataStrucPtr;    //  18h: Pointer to PCI Data Structure
-    NvU32       sizeOfBlock;        //  1Ah: <NBSI-specific appendage>
-} PCI_EXP_ROM_NBSI, *PPCI_EXP_ROM_NBSI;
-#pragma pack()
-
-typedef union _PCI_EXP_ROM {
-    PCI_EXP_ROM_STANDARD standard;
-    PCI_EXP_ROM_NBSI nbsi;
-} PCI_EXP_ROM, *PPCI_EXP_ROM;
-
-#define PCI_DATA_STRUCT_SIGNATURE     0x52494350 // "PCIR" in dword format
-#define PCI_DATA_STRUCT_SIGNATURE_NV  0x5344504E // "NPDS" in dword format
-#define PCI_DATA_STRUCT_SIGNATURE_NV2 0x53494752 // "RGIS" in dword format
-#define IS_VALID_PCI_DATA_SIG(sig) ((sig == PCI_DATA_STRUCT_SIGNATURE) ||    \
-                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV) || \
-                                    (sig == PCI_DATA_STRUCT_SIGNATURE_NV2))
-
-#define PCI_LAST_IMAGE NVBIT(7)
-#define PCI_ROM_IMAGE_BLOCK_SIZE 512U
-
-#define OFFSETOF_PCI_DATA_STRUCT_SIG        0x0
-#define OFFSETOF_PCI_DATA_STRUCT_VENDOR_ID  0x4
-#define OFFSETOF_PCI_DATA_STRUCT_LEN        0xa
-#define OFFSETOF_PCI_DATA_STRUCT_CLASS_CODE 0xd
-#define OFFSETOF_PCI_DATA_STRUCT_CODE_TYPE  0x14
-#define OFFSETOF_PCI_DATA_STRUCT_IMAGE_LEN  0x10
-#define OFFSETOF_PCI_DATA_STRUCT_LAST_IMAGE 0x15
-
-#pragma pack(1)
-typedef struct _PCI_DATA_STRUCT
-{
-    NvU32       sig;                //  00h: Signature, the string "PCIR" or NVIDIA's alternate "NPDS"
-    NvU16       vendorID;           //  04h: Vendor Identification
-    NvU16       deviceID;           //  06h: Device Identification
-    NvU16       deviceListPtr;      //  08h: Device List Pointer
-    NvU16       pciDataStructLen;   //  0Ah: PCI Data Structure Length
-    NvU8        pciDataStructRev;   //  0Ch: PCI Data Structure Revision
-    NvU8        classCode[3];       //  0Dh: Class Code
-    NvU16       imageLen;           //  10h: Image Length (units of 512 bytes)
-    NvU16       vendorRomRev;       //  12h: Revision Level of the Vendor's ROM
-    NvU8        codeType;           //  14h: holds NBSI_OBJ_CODE_TYPE (0x70) and others
-    NvU8        lastImage;          //  15h: Last Image Indicator: bit7=1 is lastImage
-    NvU16       maxRunTimeImageLen; //  16h: Maximum Run-time Image Length (units of 512 bytes)
-} PCI_DATA_STRUCT, *PPCI_DATA_STRUCT;
-#pragma pack()
-
-#define NV_PCI_DATA_EXT_SIG 0x4544504E // "NPDE" in dword format
-#define NV_PCI_DATA_EXT_REV_10 0x100      // 1.0
-#define NV_PCI_DATA_EXT_REV_11 0x101      // 1.1
-
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_SIG            0x0
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_LEN            0x6
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_REV            0x4
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_SUBIMAGE_LEN   0x8
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_LAST_IMAGE     0xa
-#define OFFSETOF_PCI_DATA_EXT_STRUCT_FLAGS          0xb
-
-#define PCI_DATA_EXT_STRUCT_FLAGS_CHECKSUM_DISABLED 0x04
-
-#pragma pack(1)
-typedef struct _NV_PCI_DATA_EXT_STRUCT
-{
-    NvU32   signature;          //  00h: Signature, the string "NPDE"
-    NvU16   nvPciDataExtRev;    //  04h: NVIDIA PCI Data Extension Revision
-    NvU16   nvPciDataExtLen;    //  06h: NVIDIA PCI Data Extension Length
-    NvU16   subimageLen;        //  08h: Sub-image Length
-    NvU8    privLastImage;      //  0Ah: Private Last Image Indicator
-    NvU8    flags;              //  0Bh: Private images enabled if bit0=1
-} NV_PCI_DATA_EXT_STRUCT, *PNV_PCI_DATA_EXT_STRUCT;
-#pragma pack()
-
-#endif // PCIEXPTBL_H
-
-
@@ -3,21 +3,8 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
-ort_options = ort.SessionOptions()
-ort_options.log_severity_level = 3

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
-  """
-  Generate example input tensors based on the provided ONNX graph input specifications.
-
-  NOTE: This is not guaranteed to be reliable. It's a best-effort helper
-  that uses heuristics to guess input shapes and values.
-
-  Example:
-    from tinygrad.nn.onnx import OnnxRunner
-    from extra.onnx_helpers import get_example_inputs
-    inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
-  """
  def _get_shape(onnx_shape: tuple[str|int]):
    shape = []
    for onnx_dim in onnx_shape:
@@ -57,9 +44,11 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
    ret.update({name:value})
  return ret

-def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
+def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  run_onnx = OnnxRunner(onnx_file)

+  ort_options = ort.SessionOptions()
+  ort_options.log_severity_level = 3
  ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
  np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
  out_names = list(run_onnx.graph_outputs)
@@ -67,101 +56,9 @@ def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
  ort_out = dict(zip(out_names, out_values))

  tinygrad_out = run_onnx(inputs)
-  Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
-  tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
-  return tinygrad_out, ort_out
-
-def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
-  """
-  Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
-  """
-  tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)

  assert tinygrad_out.keys() == ort_out.keys()
  for k in tinygrad_out.keys():
    tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
    if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
-    else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
-
-def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
-  """
-  Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
-  """
-  report = generate_node_output_report(onnx_file, inputs)
-  for i, node in enumerate(report):
-    node_name = node["node"]
-    op = node["op"]
-    outputs = node["outputs"]
-    for output in outputs:
-      output_name = output["name"]
-      tinygrad_out = output["tinygrad"]
-      ort_out = output["onnxruntime"]
-      try:
-        if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
-        else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
-        print(f"Validated {i}: {op=} {node_name=} {output_name=}")
-      except AssertionError as e:
-        print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
-        print(str(e).strip() + "\n")
-
-def generate_node_output_report(onnx_file, inputs):
-  """
-  Build a report of all ONNX node outputs from tinygrad and onnxruntime
-
-  Returns:
-    A list of dictionaries, where each entry corresponds to one
-    node in the ONNX graph. The structure is as follows:
-    [
-      {
-        "node": str,  # The name of the ONNX node.
-        "op": str,    # The operation type of the ONNX node.
-        "outputs": [
-          {
-            "name": str,                       # The name of the output tensor.
-            "tinygrad": np.ndarray | None,     # The output value from tinygrad.
-            "onnxruntime": np.ndarray | None,  # The output value from onnxruntime.
-          },
-          ...
-        ]
-      },
-      ...
-    ]
-  """
-  import onnx_graphsurgeon as gs
-  import onnx
-  import tempfile
-
-  # rewrite the model to output all the node outputs
-  # `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
-  inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
-  model = gs.import_onnx(inferred_model)
-  model_nodes = model.nodes
-  node_outputs = [n.outputs for n in model.nodes]
-  model.outputs = [
-      each_output for outputs in node_outputs for each_output in outputs
-      if not (each_output.dtype is None and each_output.shape is None)  # output with None dtype and None shape is likely a `None` value
-  ]
-  rewritten_model = gs.export_onnx(model)
-
-  # TODO: remove this once ORT supports 1.18.0
-  if getattr(rewritten_model, "ir_version", 0) > 10:
-    rewritten_model.ir_version = 10
-
-  with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
-    onnx.save(rewritten_model, f.name)
-    rewritten_model_path = f.name
-    tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
-
-  report = []
-  for node in model_nodes:
-    outputs = []
-    for each_output in node.outputs:
-      if each_output.dtype is None and each_output.shape is None:
-        continue
-      name = each_output.name
-      tinygrad_output = tinygrad_out[name]
-      ort_output = ort_out[name]
-      outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
-    report.append({"node": node.name, "op": node.op, "outputs": outputs})
-
-  return report
+    else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
  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])
+    my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), 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:
@@ -1,8 +1,7 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <database xmlns="http://nouveau.freedesktop.org/"
 xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
-<import file="freedreno_copyright.xml"/>
+xsi:schemaLocation="http://nouveau.freedesktop.org/ rules-ng.xsd">
 <import file="adreno/adreno_common.xml"/>

 <enum name="vgt_event_type" varset="chip">
@@ -21,9 +20,9 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="HLSQ_FLUSH" value="7" variants="A3XX-A4XX"/>
 	<value name="VIZQUERY_END" value="8" variants="A2XX"/>
 	<value name="SC_WAIT_WC" value="9" variants="A2XX"/>
-	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX-"/>
-	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX-"/>
-	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX-"/>
+	<value name="WRITE_PRIMITIVE_COUNTS" value="9" variants="A6XX"/>
+	<value name="START_PRIMITIVE_CTRS" value="11" variants="A6XX"/>
+	<value name="STOP_PRIMITIVE_CTRS" value="12" variants="A6XX"/>
 	<!-- Not sure that these 4 events don't have the same meaning as on A5XX+ -->
 	<value name="RST_PIX_CNT" value="13" variants="A2XX-A4XX"/>
 	<value name="RST_VTX_CNT" value="14" variants="A2XX-A4XX"/>
@@ -31,8 +30,8 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="STAT_EVENT" value="16" variants="A2XX-A4XX"/>
 	<value name="CACHE_FLUSH_AND_INV_TS_EVENT" value="20" variants="A2XX-A4XX"/>
 	<doc>
-		If A6XX_RB_SAMPLE_COUNTER_CNTL.copy is true, writes OQ Z passed
-		sample counts to RB_SAMPLE_COUNTER_BASE.  This writes to main
+		If A6XX_RB_SAMPLE_COUNT_CONTROL.copy is true, writes OQ Z passed
+		sample counts to RB_SAMPLE_COUNT_ADDR.  This writes to main
 		memory, skipping UCHE.
 	</doc>
 	<value name="ZPASS_DONE" value="21"/>
@@ -97,13 +96,6 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	</doc>
 	<value name="BLIT" value="30" variants="A5XX-"/>

-	<doc>
-	Flip between the primary and secondary LRZ buffers. This is used
-	for concurrent binning, so that BV can write to one buffer while
-	BR reads from the other.
-	</doc>
-	<value name="LRZ_FLIP_BUFFER" value="36" variants="A7XX-"/>
-
 	<doc>
 		Clears based on GRAS_LRZ_CNTL configuration, could clear
 		fast-clear buffer or LRZ direction.
@@ -120,12 +112,11 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="LRZ_FLUSH" value="38" variants="A5XX-"/>
 	<value name="BLIT_OP_FILL_2D" value="39" variants="A5XX-"/>
 	<value name="BLIT_OP_COPY_2D" value="40" variants="A5XX-A6XX"/>
-	<value name="LRZ_CACHE_INVALIDATE" value="40" variants="A7XX-"/>
-	<value name="LRZ_Q_CACHE_INVALIDATE" value="41" variants="A7XX-"/>
+	<value name="UNK_40" value="40" variants="A7XX"/>
 	<value name="BLIT_OP_SCALE_2D" value="42" variants="A5XX-"/>
 	<value name="CONTEXT_DONE_2D" value="43" variants="A5XX-"/>
-	<value name="VSC_BINNING_START" value="44" variants="A5XX-"/>
-	<value name="VSC_BINNING_END" value="45" variants="A5XX-"/>
+	<value name="UNK_2C" value="44" variants="A5XX-"/>
+	<value name="UNK_2D" value="45" variants="A5XX-"/>

 	<!-- a6xx events -->
 	<doc>
@@ -138,22 +129,21 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<!-- note, some of these are the same as a6xx, just named differently -->

 	<doc> Doesn't seem to do anything </doc>
-	<value name="DUMMY_EVENT" value="1" variants="A7XX-"/>
-	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX-"/>
-	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX-"/>
-	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX-"/>
-	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX-"/>
-	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX-"/>
-	<value name="CCU_RESOLVE" value="30" variants="A7XX-"/>
-	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX-"/>
-	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX-"/>
-	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX-"/>
-	<value name="CACHE_RESET" value="48" variants="A7XX-"/>
-	<value name="CACHE_CLEAN" value="49" variants="A7XX-"/>
+	<value name="DUMMY_EVENT" value="1" variants="A7XX"/>
+	<value name="CCU_INVALIDATE_DEPTH" value="24" variants="A7XX"/>
+	<value name="CCU_INVALIDATE_COLOR" value="25" variants="A7XX"/>
+	<value name="CCU_RESOLVE_CLEAN" value="26" variants="A7XX"/>
+	<value name="CCU_FLUSH_DEPTH" value="28" variants="A7XX"/>
+	<value name="CCU_FLUSH_COLOR" value="29" variants="A7XX"/>
+	<value name="CCU_RESOLVE" value="30" variants="A7XX"/>
+	<value name="CCU_END_RESOLVE_GROUP" value="31" variants="A7XX"/>
+	<value name="CCU_CLEAN_DEPTH" value="32" variants="A7XX"/>
+	<value name="CCU_CLEAN_COLOR" value="33" variants="A7XX"/>
+	<value name="CACHE_RESET" value="48" variants="A7XX"/>
+	<value name="CACHE_CLEAN" value="49" variants="A7XX"/>
 	<!-- TODO: deal with name conflicts with other gens -->
-	<value name="CACHE_FLUSH7" value="50" variants="A7XX-"/>
-	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX-"/>
-	<value name="DEPTH_BUFFER_FLIP" value="0x3d" variants="A8XX-"/>
+	<value name="CACHE_FLUSH7" value="50" variants="A7XX"/>
+	<value name="CACHE_INVALIDATE7" value="51" variants="A7XX"/>
 </enum>

 <enum name="pc_di_primtype">
@@ -334,7 +324,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<doc>fetch state sub-blocks and initiate shader code DMAs</doc>
 	<value name="CP_SET_STATE" value="0x25"/>
 	<doc>load constant into chip and to memory</doc>
-	<value name="CP_SET_CONSTANT" value="0x2d" variants="A2XX"/>
+	<value name="CP_SET_CONSTANT" value="0x2d"/>
 	<doc>load sequencer instruction memory (pointer-based)</doc>
 	<value name="CP_IM_LOAD" value="0x27"/>
 	<doc>load sequencer instruction memory (code embedded in packet)</doc>
@@ -381,7 +371,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_LOAD_STATE" value="0x30" variants="A3XX"/>
 	<value name="CP_LOAD_STATE4" value="0x30" variants="A4XX-A5XX"/>
 	<doc>Conditionally load a IB based on a flag, prefetch enabled</doc>
-	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a" variants="A3XX-A5XX"/>
+	<value name="CP_COND_INDIRECT_BUFFER_PFE" value="0x3a"/>
 	<doc>Conditionally load a IB based on a flag, prefetch disabled</doc>
 	<value name="CP_COND_INDIRECT_BUFFER_PFD" value="0x32" variants="A3XX"/>
 	<doc>Load a buffer with pre-fetch enabled</doc>
@@ -524,7 +514,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<!--
 	Seems to set the mode flags which control which CP_SET_DRAW_STATE
 	packets are executed, based on their ENABLE_MASK values
-
+	
 	CP_SET_MODE w/ payload of 0x1 seems to cause CP_SET_DRAW_STATE
 	packets w/ ENABLE_MASK & 0x6 to execute immediately
 	 -->
@@ -547,7 +537,7 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_LOAD_STATE6_GEOM" value="0x32" variants="A6XX-"/>
 	<value name="CP_LOAD_STATE6_FRAG" value="0x34" variants="A6XX-"/>
 	<!--
-	Note: For UAV state (Image/SSBOs) which have shared state across
+	Note: For IBO state (Image/SSBOs) which have shared state across
 	shader stages, for 3d pipeline CP_LOAD_STATE6 is used.  But for
 	compute shaders, CP_LOAD_STATE6_FRAG is used.  Possibly they are
 	interchangable.
@@ -576,21 +566,20 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="IN_PREEMPT" value="0x0f" variants="A6XX-"/>

 	<!-- TODO do these exist on A5xx? -->
-	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX-"/>
+	<value name="CP_SCRATCH_WRITE" value="0x4c" variants="A6XX"/>
 	<value name="CP_REG_TO_MEM_OFFSET_MEM" value="0x74" variants="A6XX-"/>
 	<value name="CP_REG_TO_MEM_OFFSET_REG" value="0x72" variants="A6XX-"/>
 	<value name="CP_WAIT_MEM_GTE" value="0x14" variants="A6XX"/>
 	<value name="CP_WAIT_TWO_REGS" value="0x70" variants="A6XX"/>
 	<value name="CP_MEMCPY" value="0x75" variants="A6XX-"/>
 	<value name="CP_SET_BIN_DATA5_OFFSET" value="0x2e" variants="A6XX-"/>
-	<!-- A750+, set in place of CP_SET_BIN_DATA5_OFFSET but has different values -->
-	<value name="CP_SET_UNK_BIN_DATA" value="0x2d" variants="A7XX-"/>
 	<doc>
                Write CP_CONTEXT_SWITCH_*_INFO from CP to the following dwords,
                and forcibly switch to the indicated context.
 	</doc>
 	<value name="CP_CONTEXT_SWITCH" value="0x54" variants="A6XX"/>
-	<value name="CP_SET_AMBLE" value="0x55" variants="A6XX-"/>
+	<!-- Note, kgsl calls this CP_SET_AMBLE: -->
+	<value name="CP_SET_CTXSWITCH_IB" value="0x55" variants="A6XX-"/>

 	<!--
 	Seems to always have the payload:
@@ -641,7 +630,8 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">
 	<value name="CP_BV_BR_COUNT_OPS" value="0x1b" variants="A7XX-"/>
 	<doc> Clears, adds to local, or adds to global timestamp </doc>
 	<value name="CP_MODIFY_TIMESTAMP" value="0x1c" variants="A7XX-"/>
-	<value name="CP_NON_CONTEXT_REG_BUNCH" value="0x5d" variants="A7XX-"/>
+	<!-- similar to CP_CONTEXT_REG_BUNCH, but discards first two dwords?? -->
+	<value name="CP_CONTEXT_REG_BUNCH2" value="0x5d" variants="A7XX-"/>
 	<doc>
 		Write to a scratch memory that is read by CP_REG_TEST with
 		SOURCE_SCRATCH_MEM set. It's not the same scratch as scratch registers.
@@ -658,11 +648,6 @@ xsi:schemaLocation="https://gitlab.freedesktop.org/freedreno/ rules-fd.xsd">

 	<doc>Reset various on-chip state used for synchronization</doc>
 	<value name="CP_RESET_CONTEXT_STATE" value="0x1f" variants="A7XX-"/>
-
-	<doc>Invalidates the "CCHE" introduced on a740</doc>
-	<value name="CP_CCHE_INVALIDATE" value="0x3a" variants="A7XX-"/>
-
-	<value name="CP_SCOPE_CNTL" value="0x6c" variants="A7XX-"/>
 </enum>


@@ -805,14 +790,14 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="SB6_GS_SHADER" value="0xb"/>
 		<value name="SB6_FS_SHADER" value="0xc"/>
 		<value name="SB6_CS_SHADER" value="0xd"/>
-		<value name="SB6_UAV"       value="0xe"/>
-		<value name="SB6_CS_UAV"    value="0xf"/>
+		<value name="SB6_IBO"       value="0xe"/>
+		<value name="SB6_CS_IBO"    value="0xf"/>
 	</enum>
 	<enum name="a6xx_state_type">
 		<value name="ST6_SHADER" value="0"/>
 		<value name="ST6_CONSTANTS" value="1"/>
 		<value name="ST6_UBO" value="2"/>
-		<value name="ST6_UAV" value="3"/>
+		<value name="ST6_IBO" value="3"/>
 	</enum>
 	<enum name="a6xx_state_src">
 		<value name="SS6_DIRECT" value="0"/>
@@ -918,6 +903,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>

 	<stripe varset="chip" variants="A5XX-">
+		<reg32 offset="4" name="4">
+			<bitfield name="INDX_BASE_LO" low="0" high="31"/>
+		</reg32>
+		<reg32 offset="5" name="5">
+			<bitfield name="INDX_BASE_HI" low="0" high="31"/>
+		</reg32>
 		<reg64 offset="4" name="INDX_BASE" type="address"/>
 		<reg32 offset="6" name="6">
 			<!-- max # of elements in index buffer -->
@@ -1093,10 +1084,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			<bitfield name="BINNING" pos="20" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="GMEM" pos="21" varset="chip" variants="A6XX-" type="boolean"/>
 			<bitfield name="SYSMEM" pos="22" varset="chip" variants="A6XX-" type="boolean"/>
-			<!-- high bit is 28 until a750: -->
-			<bitfield name="GROUP_ID" low="24" high="29" type="uint"/>
+			<bitfield name="GROUP_ID" low="24" high="28" type="uint"/>
 		</reg32>
-		<reg64 offset="1" name="ADDR" type="address"/>
 		<reg32 offset="1" name="1">
 			<bitfield name="ADDR_LO" low="0" high="31" type="hex"/>
 		</reg32>
@@ -1130,63 +1119,39 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<enum name="a7xx_abs_mask_mode">
-	<value name="ABS_MASK" value="0x1"/>
-	<value name="NO_ABS_MASK" value="0x0"/>
-</enum>
-
 <domain name="CP_SET_BIN_DATA5" width="32">
 	<reg32 offset="0" name="0">
-		<bitfield name="VSC_MASK" low="0" high="15" type="hex">
-			<doc>
-				A mask of bins, starting at VSC_N, whose
-				visibility is OR'd together. A value of 0 is
-				interpreted as 1 (i.e. just use VSC_N for
-				visbility) for backwards compatibility. Only
-				exists on a7xx.
-			</doc>
-		</bitfield>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
-		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes">
-			<doc>
-				If this field is 1, VSC_MASK and VSC_N are
-				ignored and instead a new ordinal immediately
-				after specifies the full 32-bit mask of bins
-				to use. The mask is "absolute" instead of
-				relative to VSC_N.
-			</doc>
-		</bitfield>
 	</reg32>
-	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<reg64 offset="1" name="BIN_DATA_ADDR" type="address"/>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-		<reg64 offset="3" name="BIN_SIZE_ADDR" type="address"/>
-		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
-		<reg64 offset="5" name="BIN_PRIM_STRM" type="address"/>
-		<!--
-			a7xx adds a few more addresses to the end of the pkt
-		 -->
-		<reg64 offset="7" name="7"/>
-		<reg64 offset="9" name="9"/>
-	</stripe>
-	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
-		<reg32 offset="1" name="ABS_MASK"/>
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<reg64 offset="2" name="BIN_DATA_ADDR" type="address"/>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-		<reg64 offset="4" name="BIN_SIZE_ADDR" type="address"/>
-		<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
-		<reg64 offset="6" name="BIN_PRIM_STRM" type="address"/>
-		<!--
-			a7xx adds a few more addresses to the end of the pkt
-		 -->
-		<reg64 offset="8" name="8"/>
-		<reg64 offset="10" name="10"/>
-	</stripe>
+	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+	<reg32 offset="1" name="1">
+		<bitfield name="BIN_DATA_ADDR_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="2" name="2">
+		<bitfield name="BIN_DATA_ADDR_HI" low="0" high="31" type="hex"/>
+	</reg32>
+	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+	<reg32 offset="3" name="3">
+		<bitfield name="BIN_SIZE_ADDRESS_LO" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="4" name="4">
+		<bitfield name="BIN_SIZE_ADDRESS_HI" low="0" high="31"/>
+	</reg32>
+	<!-- new on a6xx, where BIN_DATA_ADDR is the DRAW_STRM: -->
+	<reg32 offset="5" name="5">
+		<bitfield name="BIN_PRIM_STRM_LO" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="6" name="6">
+		<bitfield name="BIN_PRIM_STRM_HI" low="0" high="31"/>
+	</reg32>
+	<!--
+		a7xx adds a few more addresses to the end of the pkt
+	 -->
+	<reg64 offset="7" name="7"/>
+	<reg64 offset="9" name="9"/>
 </domain>

 <domain name="CP_SET_BIN_DATA5_OFFSET" width="32">
@@ -1197,42 +1162,23 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
                stream is recorded.
 	</doc>
 	<reg32 offset="0" name="0">
-		<bitfield name="VSC_MASK" low="0" high="15" type="hex"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.SIZE on a3xx/a4xx: -->
 		<bitfield name="VSC_SIZE" low="16" high="21" type="uint"/>
 		<!-- equiv to PC_VSTREAM_CONTROL.N on a3xx/a4xx: -->
 		<bitfield name="VSC_N" low="22" high="26" type="uint"/>
-		<bitfield name="ABS_MASK" pos="28" type="a7xx_abs_mask_mode" addvariant="yes"/>
 	</reg32>
-	<stripe varset="a7xx_abs_mask_mode" variants="NO_ABS_MASK">
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<reg32 offset="1" name="1">
-			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-		<reg32 offset="2" name="2">
-			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
-		<reg32 offset="3" name="3">
-			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-	</stripe>
-	<stripe varset="a7xx_abs_mask_mode" variants="ABS_MASK">
-		<reg32 offset="1" name="ABS_MASK"/>
-		<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
-		<reg32 offset="2" name="2">
-			<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-		<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
-		<reg32 offset="3" name="3">
-			<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-		<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
-		<reg32 offset="4" name="4">
-			<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
-		</reg32>
-	</stripe>
+	<!-- BIN_DATA_ADDR -> VSC_PIPE[p].DATA_ADDRESS -->
+	<reg32 offset="1" name="1">
+		<bitfield name="BIN_DATA_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
+	<!-- BIN_SIZE_ADDRESS -> VSC_SIZE_ADDRESS + (p * 4)-->
+	<reg32 offset="2" name="2">
+		<bitfield name="BIN_SIZE_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
+	<!-- BIN_DATA2_ADDR -> VSC_PIPE[p].DATA2_ADDRESS -->
+	<reg32 offset="3" name="3">
+		<bitfield name="BIN_DATA2_OFFSET" low="0" high="31" type="uint"/>
+	</reg32>
 </domain>

 <domain name="CP_REG_RMW" width="32">
@@ -1250,9 +1196,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</doc>
 	<reg32 offset="0" name="0">
 		<bitfield name="DST_REG" low="0" high="17" type="hex"/>
-		<bitfield name="DST_SCRATCH" pos="19" type="boolean" varset="chip" variants="A7XX-"/>
-		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
-		<bitfield name="SKIP_WAIT_FOR_ME" pos="23" type="boolean" varset="chip" variants="A7XX-"/>
 		<bitfield name="ROTATE" low="24" high="28" type="uint"/>
 		<bitfield name="SRC1_ADD" pos="29" type="boolean"/>
 		<bitfield name="SRC1_IS_REG" pos="30" type="boolean"/>
@@ -1266,7 +1209,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_REG_TO_MEM" width="32" prefix="chip">
+<domain name="CP_REG_TO_MEM" width="32">
 	<reg32 offset="0" name="0">
 		<bitfield name="REG" low="0" high="17" type="hex"/>
 		<!-- number of registers/dwords copied is max(CNT, 1). -->
@@ -1274,12 +1217,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<stripe varset="chip" variants="A2XX-A4XX">
-		<reg32 offset="1" name="DEST" type="address"/>
-	</stripe>
-	<stripe varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="DEST" type="address"/>
-	</stripe>
+	<reg32 offset="1" name="1">
+		<bitfield name="DEST" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
+		<bitfield name="DEST_HI" low="0" high="31"/>
+	</reg32>
 </domain>

 <domain name="CP_REG_TO_MEM_OFFSET_REG" width="32">
@@ -1295,7 +1238,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="DEST" type="waddress"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="DEST" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
+		<bitfield name="DEST_HI" low="0" high="31"/>
+	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="OFFSET0" low="0" high="17" type="hex"/>
 		<bitfield name="OFFSET0_SCRATCH" pos="19" type="boolean"/>
@@ -1315,8 +1263,18 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="64B" pos="30" type="boolean"/>
 		<bitfield name="ACCUMULATE" pos="31" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="DEST" type="waddress"/>
-	<reg64 offset="3" name="OFFSET" type="waddress"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="DEST" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
+		<bitfield name="DEST_HI" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="3" name="3">
+		<bitfield name="OFFSET_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="4" name="4">
+		<bitfield name="OFFSET_HI" low="0" high="31" type="hex"/>
+	</reg32>
 </domain>

 <domain name="CP_MEM_TO_REG" width="32">
@@ -1329,12 +1287,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- does the same thing as CP_MEM_TO_MEM::UNK31 -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
-	<stripe varset="chip" variants="A2XX-A4XX">
-		<reg32 offset="1" name="SRC" type="address"/>
-	</stripe>
-	<stripe varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="SRC" type="address"/>
-	</stripe>
+	<reg32 offset="1" name="1">
+		<bitfield name="SRC" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="2" name="2" varset="chip" variants="A5XX-">
+		<bitfield name="SRC_HI" low="0" high="31"/>
+	</reg32>
 </domain>

 <domain name="CP_MEM_TO_MEM" width="32">
@@ -1354,10 +1312,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- some other kind of wait -->
 		<bitfield name="UNK31" pos="31" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="DST" type="waddress"/>
-	<reg64 offset="3" name="SRC_A" type="address"/>
-	<reg64 offset="5" name="SRC_B" type="address"/>
-	<reg64 offset="7" name="SRC_C" type="address"/>
 	<!--
 	followed by sequence of addresses.. the first is the
 	destination and the rest are N src addresses which are
@@ -1392,8 +1346,6 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="SCRATCH" low="20" high="22" type="uint"/>
 		<!-- number of registers/dwords copied is CNT + 1. -->
 		<bitfield name="CNT" low="24" high="26" type="uint"/>
-		<!-- skip implied CP_WAIT_FOR_IDLE + CP_WAIT_FOR_ME -->
-		<bitfield name="SKIP_WAIT_FOR_ME" pos="27" type="boolean" varset="chip" variants="A7XX-"/>
 	</reg32>
 </domain>

@@ -1416,12 +1368,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 </domain>

 <domain name="CP_MEM_WRITE" width="32">
-	<stripe varset="chip" variants="A2XX-A4XX">
-		<reg32 offset="0" name="ADDR" type="address"/>
-	</stripe>
-	<stripe varset="chip" variants="A5XX-">
-		<reg64 offset="0" name="ADDR" type="address"/>
-	</stripe>
+	<reg32 offset="0" name="0">
+		<bitfield name="ADDR_LO" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="1" name="1">
+		<bitfield name="ADDR_HI" low="0" high="31"/>
+	</reg32>
 	<!-- followed by the DWORDs to write -->
 </domain>

@@ -1473,14 +1425,24 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="2" name="2">
+		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
+	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
 	<reg32 offset="4" name="4">
 		<bitfield name="MASK" low="0" high="31"/>
 	</reg32>
-	<reg64 offset="5" name="WRITE_ADDR" type="waddress"/>
+	<reg32 offset="5" name="5">
+		<bitfield name="WRITE_ADDR_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="6" name="6">
+		<bitfield name="WRITE_ADDR_HI" low="0" high="31" type="hex"/>
+	</reg32>
 	<reg32 offset="7" name="7">
 		<bitfield name="WRITE_DATA" low="0" high="31"/>
 	</reg32>
@@ -1495,7 +1457,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<!-- Reserved for flags, presumably? Unused in FW -->
 		<bitfield name="RESERVED" low="0" high="31" type="hex"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="2" name="2">
+		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
+	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1513,7 +1480,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="POLL" low="4" high="5" type="poll_memory_type"/>
 		<bitfield name="WRITE_MEMORY" pos="8" type="boolean"/>
 	</reg32>
-	<reg64 offset="1" name="POLL_ADDR" type="address"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="POLL_ADDR_LO" low="0" high="31" type="hex"/>
+	</reg32>
+	<reg32 offset="2" name="2">
+		<bitfield name="POLL_ADDR_HI" low="0" high="31" type="hex"/>
+	</reg32>
 	<reg32 offset="3" name="3">
 		<bitfield name="REF" low="0" high="31"/>
 	</reg32>
@@ -1647,7 +1619,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	TODO what is gpuaddr for, seems to be all 0's.. maybe needed for
 	context switch?
 	 -->
-	<reg64 offset="1" name="ADDR" type="waddress"/>
+	<reg32 offset="1" name="1">
+		<bitfield name="ADDR_0_LO" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="2" name="2">
+		<bitfield name="ADDR_0_HI" low="0" high="31"/>
+	</reg32>
 	<reg32 offset="3" name="3">
 		<!-- ??? -->
 	</reg32>
@@ -1676,8 +1653,8 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_SAMPLE_COUNT" pos="12" type="boolean"/>
 		<!-- Write sample count at (iova + 16) -->
 		<bitfield name="SAMPLE_COUNT_END_OFFSET" pos="13" type="boolean"/>
-		<!-- *(iova + 8) += *(iova + 16) - *iova -->
-		<bitfield name="WRITE_ACCUM_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>
+		<!-- *(iova + 8) = *(iova + 16) - *iova -->
+		<bitfield name="WRITE_SAMPLE_COUNT_DIFF" pos="14" type="boolean"/>

 		<!-- Next 4 flags are valid to set only when concurrent binning is enabled -->
 		<!-- Increment 16b BV counter. Valid only in BV pipe -->
@@ -1691,11 +1668,15 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<bitfield name="WRITE_DST" pos="24" type="event_write_dst" addvariant="yes"/>
 		<!-- Writes into WRITE_DST from WRITE_SRC. RB_DONE_TS requires WRITE_ENABLED. -->
 		<bitfield name="WRITE_ENABLED" pos="27" type="boolean"/>
-		<bitfield name="IRQ" pos="31" type="boolean"/>
 	</reg32>

 	<stripe varset="event_write_dst" variants="EV_DST_RAM">
-		<reg64 offset="1" name="1" type="waddress"/>
+		<reg32 offset="1" name="1">
+			<bitfield name="ADDR_0_LO" low="0" high="31"/>
+		</reg32>
+		<reg32 offset="2" name="2">
+			<bitfield name="ADDR_0_HI" low="0" high="31"/>
+		</reg32>
 		<reg32 offset="3" name="3">
 			<bitfield name="PAYLOAD_0" low="0" high="31"/>
 		</reg32>
@@ -1762,7 +1743,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 	</reg32>
 	<stripe varset="chip" variants="A4XX">
-		<reg32 offset="1" name="ADDR" type="address"/>
+		<reg32 offset="1" name="1">
+			<bitfield name="ADDR" low="0" high="31"/>
+		</reg32>
 		<reg32 offset="2" name="2">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1771,7 +1754,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		</reg32>
 	</stripe>
 	<stripe varset="chip" variants="A5XX-">
-		<reg64 offset="1" name="ADDR" type="address"/>
+		<reg32 offset="1" name="1">
+			<bitfield name="ADDR_LO" low="0" high="31"/>
+		</reg32>
+		<reg32 offset="2" name="2">
+			<bitfield name="ADDR_HI" low="0" high="31"/>
+		</reg32>
 		<reg32 offset="3" name="3">
 			<!-- localsize is value minus one: -->
 			<bitfield name="LOCALSIZEX" low="2" high="11" type="uint"/>
@@ -1783,88 +1771,40 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_SET_MARKER" width="32" varset="chip" prefix="chip" variants="A6XX-">
 	<doc>Tell CP the current operation mode, indicates save and restore procedure</doc>
-	<enum name="set_marker_mode">
-		<value value="0" name="SET_RENDER_MODE"/>
-		<!-- IFPC - inter-frame power collapse -->
-		<value value="1" name="SET_IFPC_MODE"/>
-	</enum>
-	<enum name="a6xx_ifpc_mode">
-		<value value="0" name="IFPC_ENABLE"/>
-		<value value="1" name="IFPC_DISABLE"/>
-	</enum>
 	<enum name="a6xx_marker">
-		<value value="1" name="RM6_DIRECT_RENDER"/>
-		<value value="2" name="RM6_BIN_VISIBILITY"/>
-		<value value="3" name="RM6_BIN_DIRECT"/>
-		<value value="4" name="RM6_BIN_RENDER_START"/>
-		<value value="5" name="RM6_BIN_END_OF_DRAWS"/>
-		<value value="6" name="RM6_BIN_RESOLVE"/>
-		<value value="7" name="RM6_BIN_RENDER_END"/>
+		<value value="1" name="RM6_BYPASS"/>
+		<value value="2" name="RM6_BINNING"/>
+		<value value="4" name="RM6_GMEM"/>
+		<value value="5" name="RM6_ENDVIS"/>
+		<value value="6" name="RM6_RESOLVE"/>
+		<value value="7" name="RM6_YIELD"/>
 		<value value="8" name="RM6_COMPUTE"/>
-		<value value="12" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->
+		<value value="0xc" name="RM6_BLIT2DSCALE"/>  <!-- no-op (at least on current sqe fw) -->

 		<!--
 			These values come from a6xx_set_marker() in the
 			downstream kernel, and they can only be set by the kernel
 		-->
-		<value value="13" name="RM6_IB1LIST_START"/>
-		<value value="14" name="RM6_IB1LIST_END"/>
-		<value value="15" name="RM7_BIN_VISIBILITY_END"/>
-
-		<!-- new in a8xx: -->
-		<value value="32" name="RM8_DEPTH_PASS_START"/>
-		<value value="33" name="RM8_DEPTH_PASS_END"/>
+		<value value="0xd" name="RM6_IB1LIST_START"/>
+		<value value="0xe" name="RM6_IB1LIST_END"/>
+		<!-- IFPC - inter-frame power collapse -->
+		<value value="0x100" name="RM6_IFPC_ENABLE"/>
+		<value value="0x101" name="RM6_IFPC_DISABLE"/>
 	</enum>
-	<stripe varset="chip" variants="A6XX-A7XX">
-		<reg32 offset="0" name="0">
-			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
-			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
-
-
-			<bitfield name="MODE" low="0" high="3" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
-			<!-- used by preemption to determine if GMEM needs to be saved or not -->
-			<bitfield name="USES_GMEM" pos="4" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
-
-
-			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
-
-
-			<!--
-				CP_SET_MARKER is used with these bits to create a
-				critical section around a workaround for ray tracing.
-				The workaround happens after BVH building, and appears
-				to invalidate the RTU's BVH node cache. It makes sure
-				that only one of BR/BV/LPAC is executing the
-				workaround at a time, and no draws using RT on BV/LPAC
-				are executing while the workaround is executed on BR (or
-				vice versa, that no draws on BV/BR using RT are executed
-				while the workaround executes on LPAC), by
-				hooking subsequent CP_EVENT_WRITE/CP_DRAW_*/CP_EXEC_CS.
-				The blob usage is:
-
-
-				CP_SET_MARKER(RT_WA_START)
-				... workaround here ...
-				CP_SET_MARKER(RT_WA_END)
-				...
-				CP_SET_MARKER(SHADER_USES_RT)
-				CP_DRAW_INDX(...) or CP_EXEC_CS(...)
-			-->
-			<bitfield name="SHADER_USES_RT" pos="9" type="boolean" variants="A7XX-"/>
-			<bitfield name="RT_WA_START" pos="10" type="boolean" variants="A7XX-"/>
-			<bitfield name="RT_WA_END" pos="11" type="boolean" variants="A7XX-"/>
-		</reg32>
-	</stripe>
-	<stripe varset="chip" variants="A8XX-">
-		<reg32 offset="0" name="0">
-			<!-- if b8 is set, the low bits are interpreted differently (and b4 ignored) -->
-			<bitfield name="MARKER_MODE" pos="8" type="set_marker_mode" addvariant="yes"/>
-			<bitfield name="USES_GMEM" pos="7" type="boolean" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
-			<bitfield name="MODE" low="0" high="6" type="a6xx_marker" varset="set_marker_mode" variants="SET_RENDER_MODE"/>
-			<bitfield name="IFPC_MODE" pos="0" type="a6xx_ifpc_mode" varset="set_marker_mode" variants="SET_IFPC_MODE"/>
-			<!-- idk if the RT w/a fields apply to a8xx as well -->
-		</reg32>
-	</stripe>
+	<reg32 offset="0" name="0">
+		<!--
+			NOTE: blob driver and some versions of freedreno/turnip set
+			b4, which is unused (at least by current sqe fw), but interferes
+			with parsing if we extend the size of the bitfield to include
+			b8 (only sent by kernel mode driver).  Really, the way the
+			parsing works in the firmware, only b0-b3 are considered, but
+			if b8 is set, the low bits are interpreted differently.  To
+			model this, without getting confused by spurious b4, this is
+			described as two overlapping bitfields:
+		 -->
+		<bitfield name="MODE" low="0" high="8" type="a6xx_marker"/>
+		<bitfield name="MARKER" low="0" high="3" type="a6xx_marker"/>
+	</reg32>
 </domain>

 <domain name="CP_SET_PSEUDO_REG" width="32" varset="chip" prefix="chip" variants="A6XX-">
@@ -1890,9 +1830,9 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			If concurrent binning is disabled then BR also does binning so it will also
 			write the "real" registers in BR.
 		-->
-		<value value="8" name="VSC_PIPE_DATA_DRAW_BASE"/>
-		<value value="9" name="VSC_SIZE_BASE"/>
-		<value value="10" name="VSC_PIPE_DATA_PRIM_BASE"/>
+		<value value="8" name="DRAW_STRM_ADDRESS"/>
+		<value value="9" name="DRAW_STRM_SIZE_ADDRESS"/>
+		<value value="10" name="PRIM_STRM_ADDRESS"/>
 		<value value="11" name="UNK_STRM_ADDRESS"/>
 		<value value="12" name="UNK_STRM_SIZE_ADDRESS"/>

@@ -1993,11 +1933,11 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 			a bitmask of which modes pass the test.
 		-->

-		<!-- RM6_BIN_VISIBILITY -->
+		<!-- RM6_BINNING -->
 		<bitfield name="BINNING" pos="25" variants="RENDER_MODE" type="boolean"/>
 		<!-- all others -->
 		<bitfield name="GMEM" pos="26" variants="RENDER_MODE" type="boolean"/>
-		<!-- RM6_DIRECT_RENDER -->
+		<!-- RM6_BYPASS -->
 		<bitfield name="SYSMEM" pos="27" variants="RENDER_MODE" type="boolean"/>

 		<bitfield name="BV" pos="25" variants="THREAD_MODE" type="boolean"/>
@@ -2070,45 +2010,54 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	</reg32>
 </domain>

-<domain name="CP_SET_AMBLE" width="32">
+<domain name="CP_SET_CTXSWITCH_IB" width="32">
 	<doc>
-		Used by the userspace and kernel drivers to set various IB's
-		which are executed during context save/restore for handling
-		state that isn't restored by the context switch routine itself.
+                Used by the userspace driver to set various IB's which are
+                executed during context save/restore for handling
+                state that isn't restored by the
+                context switch routine itself.
 	</doc>
-	<enum name="amble_type">
-		<value name="PREAMBLE_AMBLE_TYPE" value="0">
+	<enum name="ctxswitch_ib">
+		<value name="RESTORE_IB" value="0">
 			<doc>Executed unconditionally when switching back to the context.</doc>
 		</value>
-		<value name="BIN_PREAMBLE_AMBLE_TYPE" value="1">
+		<value name="YIELD_RESTORE_IB" value="1">
                        <doc>
 				Executed when switching back after switching
 				away during execution of
-				a CP_SET_MARKER packet with RM6_BIN_RENDER_END as the
-				payload *and* skipsaverestore is set. This is
-				expected to restore static register values not
-				saved when skipsaverestore is set.
+				a CP_SET_MARKER packet with RM6_YIELD as the
+				payload *and* the normal save routine was
+				bypassed for a shorter one. I think this is
+				connected to the "skipsaverestore" bit set by
+				the kernel when preempting.
 			</doc>
 		</value>
-		<value name="POSTAMBLE_AMBLE_TYPE" value="2">
+		<value name="SAVE_IB" value="2">
                        <doc>
 				Executed when switching away from the context,
 				except for context switches initiated via
 				CP_YIELD.
                        </doc>
 		</value>
-		<value name="KMD_AMBLE_TYPE" value="3">
+		<value name="RB_SAVE_IB" value="3">
 			<doc>
 				This can only be set by the RB (i.e. the kernel)
 				and executes with protected mode off, but
-				is otherwise similar to POSTAMBLE_AMBLE_TYPE.
+				is otherwise similar to SAVE_IB.
+
+				Note, kgsl calls this CP_KMD_AMBLE_TYPE
 			</doc>
 		</value>
 	</enum>
-	<reg64 offset="0" name="ADDR" type="address"/>
+	<reg32 offset="0" name="0">
+		<bitfield name="ADDR_LO" low="0" high="31"/>
+	</reg32>
+	<reg32 offset="1" name="1">
+		<bitfield name="ADDR_HI" low="0" high="31"/>
+	</reg32>
 	<reg32 offset="2" name="2">
 		<bitfield name="DWORDS" low="0" high="19" type="uint"/>
-		<bitfield name="TYPE" low="20" high="21" type="amble_type"/>
+		<bitfield name="TYPE" low="20" high="21" type="ctxswitch_ib"/>
 	</reg32>
 </domain>

@@ -2140,12 +2089,12 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 		<value name="UNK_EVENT_WRITE" value="0x4"/>
 		<doc>
 			Tracks GRAS_LRZ_CNTL::GREATER, GRAS_LRZ_CNTL::DIR, and
-			GRAS_LRZ_VIEW_INFO with previous values, and if one of
+			GRAS_LRZ_DEPTH_VIEW with previous values, and if one of
 			the following is true:
 			- GRAS_LRZ_CNTL::GREATER has changed
 			- GRAS_LRZ_CNTL::DIR has changed, the old value is not
 			  CUR_DIR_GE, and the new value is not CUR_DIR_DISABLED
-			- GRAS_LRZ_VIEW_INFO has changed
+			- GRAS_LRZ_DEPTH_VIEW has changed
 			then it does a LRZ_FLUSH with GRAS_LRZ_CNTL::ENABLE
 			forced to 1.
 			Only exists in a650_sqe.fw.
@@ -2260,7 +2209,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)

 <domain name="CP_MEM_TO_SCRATCH_MEM" width="32">
 	<doc>
-		Best guess is that it is a faster way to fetch all the VSC_CHANNEL_VISIBILITY registers
+		Best guess is that it is a faster way to fetch all the VSC_STATE registers
 		and keep them in a local scratch memory instead of fetching every time
 		when skipping IBs.
 	</doc>
@@ -2308,25 +2257,7 @@ opcode: CP_LOAD_STATE4 (30) (4 dwords)
 	<reg32 offset="0" name="0">
 		<bitfield name="CLEAR_ON_CHIP_TS" pos="0" type="boolean"/>
 		<bitfield name="CLEAR_RESOURCE_TABLE" pos="1" type="boolean"/>
-		<bitfield name="CLEAR_BV_BR_COUNTER" pos="2" type="boolean"/>
-		<bitfield name="RESET_GLOBAL_LOCAL_TS" pos="3" type="boolean"/>
-	</reg32>
-</domain>
-
-<domain name="CP_SCOPE_CNTL" width="32">
-	<enum name="cp_scope">
-		<value value="0" name="INTERRUPTS"/>
-	</enum>
-	<reg32 offset="0" name="0">
-		<bitfield name="DISABLE_PREEMPTION" pos="0" type="boolean"/>
-		<bitfield low="28" high="31" name="SCOPE" type="cp_scope"/>
-	</reg32>
-</domain>
-
-<domain name="CP_INDIRECT_BUFFER" width="32" varset="chip" prefix="chip" variants="A5XX-">
-	<reg64 offset="0" name="IB_BASE" type="address"/>
-	<reg32 offset="2" name="2">
-		<bitfield name="IB_SIZE" low="0" high="19"/>
+		<bitfield name="CLEAR_GLOBAL_LOCAL_TS" pos="2" type="boolean"/>
 	</reg32>
 </domain>

@@ -10,8 +10,6 @@ import ctypes


 class AsDictMixin:
-    import sys
-    if sys.version_info >= (3, 14): _layout_ = 'ms'
    @classmethod
    def as_dict(cls, self):
        result = {}
@@ -97,7 +97,7 @@ def parse_cmd_buf(dat):

        if state_block == SB6_CS_SHADER:
          from extra.disassemblers.adreno import disasm_raw
-          if state_type == ST6_SHADER and IOCTL > 3:
+          if state_type == ST6_SHADER and IOCTL > 2:
            disasm_raw(get_mem(((vals[2] << 32) | vals[1]), num_unit * 128))
          if state_type == ST6_CONSTANTS:
            x = get_mem(((vals[2] << 32) | vals[1]), num_unit*4)
@@ -106,30 +106,25 @@ def parse_cmd_buf(dat):
              print('constants')
              hexdump(x)
          if state_type == ST6_IBO:
-            if state_src == 0x1:
-              ibos_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
-            else: ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
+            ibos_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 16 * 4)
            CAPTURED_STATE['ibos'] = ibos_bytes[:]
            if IOCTL > 1:
              print('texture ibos')
              hexdump(ibos_bytes)
        elif state_block == SB6_CS_TEX:
          if state_type == ST6_SHADER:
-            if state_src == 0x1:
-              samplers_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
-            else: samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
+            samplers_bytes = get_mem((vals[2] << 32) | vals[1], num_unit * 4 * 4)
            CAPTURED_STATE['samplers'] = samplers_bytes[:]
            if IOCTL > 1:
              print('texture samplers')
              hexdump(samplers_bytes)
          if state_type == ST6_CONSTANTS:
-            if state_src == 0x1:
-              descriptors_bytes = get_mem(CAPTURED_STATE['bindless_base'] + ((vals[2] << 32) | vals[1]) * 4, num_unit * 64)
-            else: descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
+            descriptors_bytes = get_mem((vals[2] << 32) | vals[1], 1600)
            CAPTURED_STATE['descriptors'] = descriptors_bytes[:]
            if IOCTL > 1:
              print('texture descriptors')
              hexdump(descriptors_bytes)
+
      elif ops[opcode] == "CP_REG_TO_MEM":
        reg, cnt, b64, accum = vals[0] & 0x3FFFF, (vals[0] >> 18) & 0xFFF, (vals[0] >> 30) & 0x1, (vals[0] >> 31) & 0x1
        dest = vals[1] | (vals[2] << 32)
@@ -157,10 +152,6 @@ def parse_cmd_buf(dat):
        if IOCTL > 0:
          print(f'THREADSIZE-{(vals[0] >> 20)&0x1}\nEARLYPREAMBLE-{(vals[0] >> 23) & 0x1}\nMERGEDREGS-{(vals[0] >> 3) & 0x1}\nTHREADMODE-{vals[0] & 0x1}\nHALFREGFOOTPRINT-{(vals[0] >> 1) & 0x3f}\nFULLREGFOOTPRINT-{(vals[0] >> 7) & 0x3f}\nBRANCHSTACK-{(vals[0] >> 14) & 0x3f}\n')
          print(f'SP_CS_UNKNOWN_A9B1-{vals[1]}\nSP_CS_BRANCH_COND-{vals[2]}\nSP_CS_OBJ_FIRST_EXEC_OFFSET-{vals[3]}\nSP_CS_OBJ_START-{vals[4] | (vals[5] << 32)}\nSP_CS_PVT_MEM_PARAM-{vals[6]}\nSP_CS_PVT_MEM_ADDR-{vals[7] | (vals[8] << 32)}\nSP_CS_PVT_MEM_SIZE-{vals[9]}')
-      if offset == 0xa9e8:
-        CAPTURED_STATE['bindless_base'] = (vals[0] | (vals[1] << 32)) & ~0b11
-        # print(hex(CAPTURED_STATE['bindless_base']))
-        # hexdump(get_mem(CAPTURED_STATE['bindless_base'], 0x200))
      if offset == 0xb180:
        if IOCTL > 0:
          print('border color offset', hex(vals[1] << 32 | vals[0]))
@@ -180,8 +171,8 @@ def ioctl(fd, request, argp):
    name, stype = nrs[nr]
    s = get_struct(argp, stype)
    if IOCTL > 0: print(f"{ret:2d} = {name:40s}", ' '.join(format_struct(s)))
-    if name == "IOCTL_KGSL_GPUOBJ_INFO":
-      mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
+    if name == "IOCTL_KGSL_GPUOBJ_INFO": pass
+      # mmaped[s.gpuaddr] = mmap.mmap(fd, s.size, offset=s.id*0x1000)
    if name == "IOCTL_KGSL_GPU_COMMAND":
      for i in range(s.numcmds):
        cmd = get_struct(s.cmdlist+ctypes.sizeof(msm_kgsl.struct_kgsl_command_object)*i, msm_kgsl.struct_kgsl_command_object)
@@ -882,11 +882,6 @@ impl<'a> Thread<'a> {
                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i32;
                            (s0 * s1) as u32
                        }
-                        10 => {
-                            let s0 = sign_ext((s0 & 0xffffff) as u64, 24) as i64;
-                            let s1 = sign_ext((s1 & 0xffffff) as u64, 24) as i64;
-                            ((s0 * s1) >> 32) as u32
-                        }
                        17 | 18 | 26 => {
                            let (s0, s1) = (s0 as i32, s1 as i32);
                            (match op {
@@ -935,7 +930,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 | 770 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -1001,10 +996,6 @@ impl<'a> Thread<'a> {
                                    let ret = s0.wrapping_sub(s1);
                                    (ret as u32, s1 > s0)
                                }
-                                770 => {
-                                    let ret = s1.wrapping_sub(s0);
-                                    (ret as u32, s0 > s1)
-                                }
                                _ => todo_instr!(instruction)?,
                            };
                            if self.exec.read() {
@@ -1,32 +1,98 @@
 import numpy as np
 import unittest
 import subprocess, struct, math
-from tinygrad import Tensor, dtypes, Device, UOp
-from tinygrad.helpers import getenv
-from tinygrad.runtime.support.compiler_amd import amdgpu_disassemble
-from tinygrad.renderer import ProgramSpec
-from tinygrad.engine.realize import CompiledRunner
+from typing import cast
+from tinygrad.runtime.ops_amd import AMDProgram, AMDDevice
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.helpers import diskcache, OSX, getenv

-def get_output(asm:str, n_threads:int=1):
-  input_asm = "\n".join([ln if ln.strip().startswith('asm volatile') else f'asm volatile("{ln.strip().lstrip()}" : "+v"(a), "+v"(b));'
-                         for ln in asm.strip().splitlines() if ln.strip()])
-  src = f"""
-  typedef long unsigned int size_t;
-  extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
-  extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, {n_threads}))) test(unsigned int* data0_1) {{
-    int l = __ockl_get_local_id(0);
-    unsigned a = 0, b = 0, c = 0;
-    {input_asm}
-    unsigned res;
-    asm volatile("v_mov_b32 %0, %1" : "=v"(res) : "v"(a));
-    *(data0_1+l) = res;
-  }}"""
-  t = Tensor.zeros(n_threads, dtype=dtypes.uint32).contiguous().realize()
-  prg = ProgramSpec("test", src, Device.DEFAULT, UOp.sink(t), global_size=[1, 1, 1], local_size=[n_threads, 1, 1])
-  car = CompiledRunner(prg)
-  if getenv("PRINT_ASM"): amdgpu_disassemble(car.lib)
-  car([t.uop.buffer], {}, wait=True)
-  return t.numpy()
+@diskcache
+def assemble(code:str) -> bytes:
+  try:
+    LLVM_MC = "llvm-mc" if OSX else "/opt/rocm/llvm/bin/llvm-mc"
+    return subprocess.run([LLVM_MC, "--arch=amdgcn", "--mcpu=gfx1100", "--triple=amdgcn-amd-amdhsa", "-filetype=obj", "-o", "-"],
+                           input=code.encode("utf-8"), stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True).stdout
+  except subprocess.CalledProcessError as e:
+    print("stderr:")
+    print(e.stderr.decode())
+    raise
+
+# copied from extra/rdna
+def get_prg(code:str, v_cnt:int, s_cnt:int):
+  function_name = "test"
+  metadata = f"""
+amdhsa.kernels:
+- .args:
+  - .address_space: global
+    .name: buf_0
+    .offset: 0
+    .size: 8
+    .type_name: unsigned int*
+    .value_kind: global_buffer
+  .group_segment_fixed_size: 0
+  .kernarg_segment_align: 8
+  .kernarg_segment_size: 8
+  .language: OpenCL C
+  .language_version:
+  - 1
+  - 2
+  .max_flat_workgroup_size: 256
+  .name: test
+  .private_segment_fixed_size: 0
+  .sgpr_count: {s_cnt}
+  .sgpr_spill_count: 0
+  .symbol: test.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}:
+  """
+  ret = ".amdgpu_metadata\n" + metadata + ".end_amdgpu_metadata" + boilerplate_start + "\n" + '\n'.join("%s %d" % x for x in kernel_desc.items()) \
+      + "\n" + code_start + code + f"\n.size {function_name}, .-{function_name}"
+  return AMDProgram(cast(AMDDevice, Device["AMD"]), function_name, assemble(ret))
+
+def get_output(s:str, n_threads:int=1):
+  assert n_threads <= 32
+  code = "\n".join(["s_load_b64 s[0:1], s[0:1], null", "v_lshlrev_b32_e32 v0, 2, v0", s,
+    "s_waitcnt 0",
+    "global_store_b32 v0, v1, s[0:1]",
+    "s_nop 0", "s_sendmsg sendmsg(MSG_DEALLOC_VGPRS)", "s_endpgm"])
+  test = Tensor.zeros((n_threads,), dtype=dtypes.uint32).contiguous().realize().uop.buffer
+  prg = get_prg(code, 32, 32)
+  prg(test._buf, global_size=(1, 1, 1), local_size=(n_threads, 1, 1), wait=True)
+  return test.numpy()

 def f16_to_bits(x:float) -> int: return struct.unpack('<H', struct.pack('<e', x))[0]
 def f32_from_bits(x:int) -> float: return struct.unpack('<f', struct.pack('<I', x))[0]
@@ -39,57 +105,54 @@ class TestHW(unittest.TestCase):

  def test_simple(self):
    out = get_output("""
-    v_mov_b32_e32 %1 42
-    v_mov_b32_e32 %2 %1
-    """)[0]
+    v_mov_b32_e32 v10 42
+    v_mov_b32_e32 v1 v10
+    """, n_threads=2)
    np.testing.assert_equal(out, 42)

  def test_exec_mov(self):
    out = get_output("""
-    v_mov_b32_e32 %1 42
+    v_mov_b32_e32 v10 42
    s_mov_b32_e32 exec_lo 0b10
-    v_mov_b32_e32 %1 10
+    v_mov_b32_e32 v10 10
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 %1
+    v_mov_b32_e32 v1 v10
    """, n_threads=2)
    np.testing.assert_equal(out, [42, 10])

  def test_exec_cmp_vopc(self):
    out = get_output("""
    s_mov_b32 vcc_lo 0 // reset vcc
-    v_mov_b32_e32 %1 42
-    v_mov_b32_e32 %2 10
+    v_mov_b32_e32 v10 42
+    v_mov_b32_e32 v11 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmp_ne_u32 %1 %2
+    v_cmp_ne_u32 v10 v11
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 vcc_lo
+    v_mov_b32_e32 v1 vcc_lo
    """, n_threads=2)
    np.testing.assert_equal(out, 0b01)

  def test_exec_cmpx_vop3(self):
    out = get_output("""
-    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %1 42
-    v_mov_b32_e32 %2 10
+    v_mov_b32_e32 v10 42
+    v_mov_b32_e32 v11 10
    s_mov_b32_e32 exec_lo 0b01
-    v_cmpx_ne_u32 %1 %2
+    v_cmpx_ne_u32 v10 v11
    s_mov_b32_e32 s10 exec_lo
    s_mov_b32_e32 exec_lo 0b11
-    v_mov_b32_e32 %2 s10
-    """, n_threads=2)[0]
-    np.testing.assert_equal(out & 0b11, 0b01)
+    v_mov_b32_e32 v1 s10
+    """, n_threads=2)
+    np.testing.assert_equal(out, 0b01)

  def test_fmac_vop3_modifier(self):
    init_state = f"""
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(4.0)}" : "+v"(a));
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(3.0)}" : "+v"(b));
-    asm volatile("v_mov_b32_e32 %1, {f16_to_bits(2.0)}" : "+v"(c));
+    v_mov_b32_e32 v10 {f16_to_bits(4.0)}
+    v_mov_b32_e32 v11 {f16_to_bits(3.0)}
+    v_mov_b32_e32 v1 {f16_to_bits(2.0)}
    """
-    mov = """asm volatile("v_mov_b32_e32 %1, %2" : "+v"(c), "+v"(a));"""
-    def fmac(a, b, c): return f"""asm volatile("v_fmac_f16_e64 {c}, {a}, {b}" : "+v"(c) : "v"(a), "v"(b));"""+"\n"+mov
-    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "%2", "%3")), f16_to_bits(14.))
-    self.assertEqual(get_output(init_state+"\n"+fmac("%1", "-%2", "%3")), f16_to_bits(-10.))
-    self.assertEqual(get_output(init_state+"\n"+fmac("-%1", "-%2", "%3")), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 v11 v10"), f16_to_bits(14.))
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 v10"), f16_to_bits(-10.))
+    self.assertEqual(get_output(init_state+"\n"+"v_fmac_f16_e64 v1 -v11 -v10"), f16_to_bits(14.))

  def test_s_abs_i32(self):
    def s_abs_i32(x, y, dst="s10", scc=0):
@@ -97,7 +160,7 @@ class TestHW(unittest.TestCase):
        self.assertEqual(get_output(f"""
        s_mov_b32_e32 {dst} {x}
        s_abs_i32 {dst} {dst}
-        v_mov_b32_e32 %2 {reg}
+        v_mov_b32_e32 v1 {reg}
        """)[0], val)
    s_abs_i32(0x00000001, 0x00000001, scc=1)
    s_abs_i32(0x7fffffff, 0x7fffffff, scc=1)
@@ -110,8 +173,8 @@ class TestHW(unittest.TestCase):
  def test_v_rcp_f32_neg_vop3(self):
    def v_neg_rcp_f32(x:float, y:float):
      out = get_output(f"""
-      v_mov_b32_e32 %2 {f32_to_bits(x)}
-      v_rcp_f32_e64 %2, -%2
+      v_mov_b32_e32 v1 {f32_to_bits(x)}
+      v_rcp_f32_e64 v1, -v1
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg_rcp_f32(math.inf, -0.0)
@@ -123,11 +186,10 @@ class TestHW(unittest.TestCase):

  def test_v_cndmask_b32_neg(self):
    def v_neg(x:int|float, y:float):
-      # always pick -v1
      out = get_output(f"""
-      v_mov_b32_e32 %2 {f32_to_bits(x)}
-      s_mov_b32_e32 s10 1
-      v_cndmask_b32 %2, %2, -%2 s10
+      v_mov_b32_e32 v1 {f32_to_bits(x)}
+      s_mov_b32_e32 s10 1 // always pick -v1
+      v_cndmask_b32 v1, v1, -v1 s10
      """)[0]
      assert out == f32_to_bits(y), f"{f32_from_bits(out)} != {y} / {out} != {f32_to_bits(y)}"
    v_neg(-0.0, 0.0)
@@ -136,12 +198,5 @@ class TestHW(unittest.TestCase):
    v_neg(math.inf, -math.inf)
    v_neg(-math.inf, math.inf)

-  def test_v_subrev_wrap(self):
-    out = get_output("""
-    v_dual_mov_b32 %1, 0xffffffff :: v_dual_mov_b32 %2, 0x0
-    v_subrev_co_u32 %2, vcc_lo, %2, %1
-    """)[0]
-    self.assertEqual(out, 0xffff_ffff)
-
 if __name__ == "__main__":
  unittest.main()
@@ -2,6 +2,8 @@

 ## Getting SQ Thread Trace

+Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)
+
 SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.
@@ -1,18 +0,0 @@
-#!/usr/bin/env python3
-import os, shutil
-from pathlib import Path
-from tinygrad.helpers import fetch, OSX
-
-DEST = Path("/usr/local/lib")
-DEST.mkdir(exist_ok=True)
-
-if __name__ == "__main__":
-  if OSX:
-    fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
-    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
-    os.chmod(fp, 0o755)
-    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
-  else:
-    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/43bf0fef74a83c3c25badfc5a09c0bd39ed8c6f9/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
-  shutil.copy2(lib, DEST)
-  print(f"Installed {lib.name} to", DEST)
@@ -4,7 +4,7 @@ import argparse, ctypes, struct, hashlib, pickle, code, typing, functools
 import tinygrad.runtime.autogen.sqtt as sqtt
 from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
 from tinygrad.runtime.ops_amd import ProfileSQTTEvent
-from tinygrad.helpers import round_up, flatten, all_same, temp
+from tinygrad.helpers import round_up, flatten, all_same
 from dataclasses import dataclass

 CHUNK_CLASSES = {
@@ -154,25 +154,7 @@ class RGP:
      if device not in device_events: raise RuntimeError(f"Device {device} not found in profile, devices in profile: {', '.join(device_events.keys())} ")
      device_event = device_events[device]
    sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
-    device_props = device_event.props
-    # merge events per SE
-    merged_sqtt_events:dict[int, ProfileSQTTEvent] = {}
-    for ev in sqtt_events:
-      if ev.se not in merged_sqtt_events: merged_sqtt_events[ev.se] = ev
-      else:
-        merged_sqtt_events[ev.se] = ProfileSQTTEvent(
-          device=ev.device,
-          kern=ev.kern,
-          se=ev.se,
-          itrace=merged_sqtt_events[ev.se].itrace or ev.itrace,
-          blob=merged_sqtt_events[ev.se].blob + ev.blob,
-        )
-    sqtt_events = list(merged_sqtt_events.values())
-
    if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
-    gfx_ver = device_props['gfx_target_version'] // 10000
-    gfx_iplvl = getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}_{(device_props['gfx_target_version']//100)%100}",
-                        getattr(sqtt, f"SQTT_GFXIP_LEVEL_GFXIP_{device_props['gfx_target_version']//10000}", None))
    sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
    sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
    sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -210,21 +192,21 @@ class RGP:
        flags=0,
        trace_shader_core_clock=0x93f05080,
        trace_memory_clock=0x4a723a40,
-        device_id={110000: 0x744c, 110003: 0x7480, 120001: 0x7550, 120000: 0x7550}[device_props['gfx_target_version']],
+        device_id=0x744c,
        device_revision_id=0xc8,
        vgprs_per_simd=1536,
        sgprs_per_simd=128*16,
-        shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
-        compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
-        simd_per_compute_unit=device_props['simd_per_cu'],
-        wavefronts_per_simd=device_props['max_waves_per_simd'],
+        shader_engines=6,
+        compute_unit_per_shader_engine=16,
+        simd_per_compute_unit=2,
+        wavefronts_per_simd=16,
        minimum_vgpr_alloc=4,
        vgpr_alloc_granularity=8,
        minimum_sgpr_alloc=128,
        sgpr_alloc_granularity=128,
        hardware_contexts=8,
        gpu_type=sqtt.SQTT_GPU_TYPE_DISCRETE,
-        gfxip_level=gfx_iplvl,
+        gfxip_level=sqtt.SQTT_GFXIP_LEVEL_GFXIP_11_0,
        gpu_index=0,
        gds_size=0,
        gds_per_shader_engine=0,
@@ -236,7 +218,7 @@ class RGP:
        vram_bus_width=384, # 384-bit
        l2_cache_size=6 * 1024 * 1024, # 6 MB
        l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
-        lds_size=device_props['lds_size_in_kb'] * 1024,
+        lds_size=65536, # 64 KB per CU
        gpu_name=b'NAVI31',
        alu_per_clock=0,
        texture_per_clock=0,
@@ -275,7 +257,7 @@ class RGP:
            major_version=0, minor_version=2,
          ),
          shader_engine_index=sqtt_event.se,
-          sqtt_version={11: sqtt.SQTT_VERSION_3_2, 12: sqtt.SQTT_VERSION_3_3}.get(gfx_ver),
+          sqtt_version=sqtt.SQTT_VERSION_3_2,
          _0=sqtt.union_sqtt_file_chunk_sqtt_desc_0(
            v1=sqtt.struct_sqtt_file_chunk_sqtt_desc_0_v1(
              instrumentation_spec_version=1,
@@ -324,7 +306,7 @@ class RGP:
 if __name__ == '__main__':
  parser = argparse.ArgumentParser(prog='rgptool', description='A tool to create (from pickled tinygrad profile), inspect and modify Radeon GPU Profiler files')
  parser.add_argument('command')
-  parser.add_argument('input', nargs='?', default=temp("profile.pkl", append_user=True))
+  parser.add_argument('input')
  parser.add_argument('-d', '--device')
  parser.add_argument('-o', '--output')
  args = parser.parse_args()
@@ -346,4 +328,3 @@ if __name__ == '__main__':

  if args.output is not None:
    with open(args.output, 'wb+') as fd: fd.write(rgp.to_bytes())
-    print(f"Saved to {args.output}")
@@ -1,171 +0,0 @@
-import ctypes, pathlib, argparse, pickle, re, functools, dataclasses, itertools
-from tinygrad.helpers import temp, unwrap, DEBUG
-from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileProgramEvent
-from tinygrad.runtime.ops_amd import ProfileSQTTEvent, ProfilePMCEvent
-from tinygrad.runtime.autogen import llvm, rocprof
-from tinygrad.runtime.support.elf import elf_loader
-
-# to pass NULL to callbacks
-llvm.LLVMCreateDisasmCPUFeatures.argtypes = tuple(llvm.LLVMCreateDisasmCPUFeatures.argtypes[:5]) + (ctypes.c_void_p, ctypes.c_void_p)
-def llvm_disasm(arch:str, lib:bytes) -> dict[int, tuple[str, int]]:
-  llvm.LLVMInitializeAMDGPUTargetInfo()
-  llvm.LLVMInitializeAMDGPUTargetMC()
-  llvm.LLVMInitializeAMDGPUAsmParser()
-  llvm.LLVMInitializeAMDGPUDisassembler()
-  ctx = llvm.LLVMCreateDisasmCPUFeatures("amdgcn-amd-amdhsa".encode(), arch.encode(), "".encode(), None, 0, None, None)
-
-  image, sections, relocs = elf_loader(lib)
-  text = next((sh.header for sh in sections if sh.name == ".text"), None)
-  off, sz = unwrap(text).sh_addr, unwrap(text).sh_size
-
-  addr_table:dict[int, tuple[str, int]] = {}
-  out = ctypes.create_string_buffer(128)
-  cur_off = off
-  while cur_off < sz + off:
-    view = (ctypes.c_ubyte * ((sz + off) - cur_off)).from_buffer_copy(memoryview(image)[cur_off:])
-    instr_sz = llvm.LLVMDisasmInstruction(ctx, view, ctypes.c_uint64(len(view)), ctypes.c_uint64(0), out, ctypes.c_size_t(128))
-    addr_table[cur_off] = (out.value.decode("utf-8", "replace").strip(), instr_sz)
-    cur_off += instr_sz
-  return addr_table
-
-@dataclasses.dataclass
-class InstInfo:
-  typ:str=""
-  inst:str=""
-  hit:int=0
-  lat:int=0
-  stall:int=0
-  def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
-
-  def on_ev(self, ev):
-    self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
-
-@dataclasses.dataclass(frozen=True)
-class InstExec:
-  typ:str
-  inst:str
-  stall:int
-  dur:int
-  time:int
-
-@dataclasses.dataclass(frozen=True)
-class PrgExec:
-  name:str
-  wave:int
-  cu:int
-  simd:int
-  def __str__(self): return f"{self.name},{self.wave},{self.cu},{self.simd}"
-
-@dataclasses.dataclass(frozen=True)
-class WaveExec:
-  wave_id:int
-  cu:int
-  simd:int
-  insts:list[InstExec]
-
-class _ROCParseCtx:
-  def __init__(self, dev_evs:dict[str, ProfileDeviceEvent], sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
-    self.dev_evs, self.sqtt_evs, self.prog_evs = dev_evs, iter(sqtt_evs), prog_evs
-    self.wave_events:dict[PrgExec, dict[int, InstInfo]] = {}
-    self.disasms:dict[tuple[str, int], tuple[str, int]] = {}
-    self.inst_execs:dict[str, list[WaveExec]] = {}
-
-    for prog in prog_evs:
-      arch = "gfx%d%x%x" % ((trgt:=unwrap(dev_evs[prog.device].props)['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
-      for addr, info in llvm_disasm(arch, unwrap(prog.lib)).items():
-        self.disasms[(prog.name, unwrap(prog.base) + addr)] = info
-
-  def next_sqtt(self):
-    x = next(self.sqtt_evs, None)
-    self.active_kern = x.kern if x is not None else None
-    self.active_se = x.se if x is not None else None
-    return x
-
-  def on_occupancy_ev(self, ev):
-    if DEBUG >= 5: print("OCC", ev.time, self.active_se, ev.cu, ev.simd, ev.wave_id, ev.start)
-
-  def on_wave_ev(self, ev):
-    if DEBUG >= 5: print("WAVE", ev.wave_id, self.active_se, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
-
-    asm:dict[int, InstInfo] = {}
-    inst_execs:list[InstExec] = []
-    for j in range(ev.instructions_size):
-      inst_ev = ev.instructions_array[j]
-      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
-      inst_disasm = self.disasms[(unwrap(self.active_kern), unwrap(inst_ev.pc.address))][0]
-      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=inst_disasm))
-      asm[inst_ev.pc.address].on_ev(inst_ev)
-      inst_execs.append(InstExec(inst_typ, inst_disasm, inst_ev.stall, inst_ev.duration, inst_ev.time))
-
-    if ev.instructions_size > 0:
-      self.wave_events[key:=PrgExec(unwrap(self.active_kern), ev.wave_id, ev.cu, ev.simd)] = asm
-      self.inst_execs.setdefault(key.name, []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, inst_execs))
-
-def decode(profile:list[ProfileEvent]) -> _ROCParseCtx:
-  dev_events:dict[str, ProfileDeviceEvent] = {}
-  sqtt_events:list[ProfileSQTTEvent] = []
-  prog_events:list[ProfileProgramEvent] = []
-  for e in profile:
-    if isinstance(e, ProfileDeviceEvent): dev_events[e.device] = e
-    if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
-    if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
-
-  ROCParseCtx = _ROCParseCtx(dev_events, sqtt_events, prog_events)
-
-  @rocprof.rocprof_trace_decoder_se_data_callback_t
-  def copy_cb(buf, buf_size, data_ptr):
-    if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
-    buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
-    buf_size[0] = len(prof.blob)
-    return len(prof.blob)
-
-  @rocprof.rocprof_trace_decoder_trace_callback_t
-  def trace_cb(record_type, events_ptr, n, data_ptr):
-    match record_type:
-      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
-        for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
-      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
-        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
-      case _:
-        if DEBUG >= 5: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
-    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
-
-  @rocprof.rocprof_trace_decoder_isa_callback_t
-  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
-    instr, mem_size_ptr[0] = ROCParseCtx.disasms[(unwrap(ROCParseCtx.active_kern), pc.address)]
-
-    # this is the number of bytes to next instruction, set to 0 for end_pgm
-    if instr == "s_endpgm": mem_size_ptr[0] = 0
-    if (max_sz:=size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
-
-    # truncate the instr if it doesn't fit
-    if (str_sz:=len(instr_bytes:=instr.encode()))+1 > max_sz: str_sz = max_sz
-    ctypes.memmove(instr_ptr, instr_bytes, str_sz)
-    size_ptr[0] = str_sz
-
-    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
-
-  try:
-    rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
-  except AttributeError as e: raise RuntimeError("Failed to find rocprof-trace-decoder. Run ./extra/sqtt/install_sqtt_decoder.py to install") from e
-  return ROCParseCtx
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
-  args = parser.parse_args()
-
-  with args.profile.open("rb") as f: profile = pickle.load(f)
-  rctx = decode(profile)
-  print('SQTT:', rctx.wave_events.keys())
-
-  for ev in profile:
-    if not isinstance(ev, ProfilePMCEvent): continue
-    print(f"PMC Event: dev={ev.device} kern={ev.kern}")
-    ptr = 0
-    for s in ev.sched:
-      view = memoryview(ev.blob).cast('Q')
-      print(f"\t{s.name}")
-      for xcc, inst, se_idx, sa_idx, wgp_idx in itertools.product(range(s.xcc), range(s.inst), range(s.se), range(s.sa), range(s.wgp)):
-        print(f"\t\tXCC {xcc} Inst {inst} SE {se_idx} SA {sa_idx} WGP {wgp_idx}: {view[ptr]:#x}")
-        ptr += 1
@@ -43,7 +43,6 @@ enum sqtt_version
   SQTT_VERSION_2_3 = 0x6, /* GFX9 */
   SQTT_VERSION_2_4 = 0x7, /* GFX10+ */
   SQTT_VERSION_3_2 = 0xb, /* GFX11+ */
-   SQTT_VERSION_3_3 = 0xc, /* GFX12+ */
 };

 enum sqtt_file_chunk_type
@@ -145,8 +144,6 @@ enum sqtt_gfxip_level
   SQTT_GFXIP_LEVEL_GFXIP_10_1 = 0x7,
   SQTT_GFXIP_LEVEL_GFXIP_10_3 = 0x9,
   SQTT_GFXIP_LEVEL_GFXIP_11_0 = 0xc,
-   SQTT_GFXIP_LEVEL_GFXIP_11_5 = 0xd,
-   SQTT_GFXIP_LEVEL_GFXIP_12 = 0x10,
 };

 enum sqtt_memory_type
@@ -430,8 +427,6 @@ enum elf_gfxip_level
   EF_AMDGPU_MACH_AMDGCN_GFX1010 = 0x033,
   EF_AMDGPU_MACH_AMDGCN_GFX1030 = 0x036,
   EF_AMDGPU_MACH_AMDGCN_GFX1100 = 0x041,
-   EF_AMDGPU_MACH_AMDGCN_GFX1150 = 0x043,
-   EF_AMDGPU_MACH_AMDGCN_GFX1200 = 0x04e,
 };

 struct sqtt_file_chunk_spm_db {
@@ -1,105 +0,0 @@
-import os
-os.environ["PYTHONPATH"] = "."
-os.environ["SQTT"] = "1"
-os.environ["AMD"] = "1"
-os.environ["VIZ"] = "1"
-os.environ["AMD_LLVM"] = "0"
-
-import unittest
-import sys, contextlib
-from tinygrad import Tensor
-from tinygrad.dtype import dtypes
-from tinygrad.renderer import ProgramSpec
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.device import Device, ProfileDeviceEvent
-
-from extra.sqtt.roc import decode, InstExec, PrgExec
-
-dev = Device["AMD"]
-
-def custom(arg:str, s:UOp|None=None) -> UOp: return UOp(Ops.CUSTOM, src=(s,) if s is not None else (), arg=arg)
-
-def asm_kernel(instrs:list[str], l:int=1, g:int=1) -> Tensor:
-  name = sys._getframe(1).f_code.co_name
-  def fxn(_):
-    L = UOp.special(l, "lidx0")
-    G = UOp.special(g, "gidx0")
-    op = custom("asm volatile (")
-    for inst in instrs: op = custom(f'  "{inst}\\n\\t"', op)
-    op = custom(");", op)
-    return UOp.sink(op, L, G, arg=KernelInfo(name=name))
-  k = Tensor.custom_kernel(Tensor.empty(1), fxn=fxn)[0]
-  return k
-
-@contextlib.contextmanager
-def save_sqtt():
-  # clear the old traces
-  dev.profile_events.clear()
-  sqtt:dict[PrgExec, list[InstExec]] = {}
-  yield sqtt
-  # decode sqtt
-  rctx = decode(dev.profile_events+[ProfileDeviceEvent("AMD", props=dev.device_props())])
-  assert len(rctx.inst_execs) > 0, "empty sqtt output"
-  sqtt.update(rctx.inst_execs)
-
-class TestTiming(unittest.TestCase):
-  def test_v_add(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([f"v_add_f32 v{10+i} v{10+i+1} {10+i}" for i in range(3)]).realize()
-    wave = list(sqtt.values())[0][:-1]
-    assert all(s.dur == 1 for s in wave)
-    assert all(s.stall == 0 for s in wave)
-
-  def test_chain_v_add_1l(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_add_f32_e32 v1 v0 v0",
-        "v_add_f32_e32 v2 v1 v1",
-      ]).realize()
-    wave = list(sqtt.values())[0][:-1]
-    assert all(s.dur == 1 for s in wave)
-    assert all(s.stall == 0 for s in wave)
-
-  def test_multi_cycle_inst(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_mov_b32_e32 v4 0x3f800000",
-        "v_rcp_f32_e32 v5 v4",
-        "v_mul_f32_e32 v6 v5 v4",
-      ]).realize()
-    w = list(sqtt.values())[0]
-    rcp, mul = w[1], w[2]
-    self.assertGreater(rcp.dur, 1) # 4 cycles on gfx11
-    self.assertEqual(mul.dur, 1)
-    # mul depends on v5, how can it run before rcp is done?
-    self.assertGreaterEqual(mul.time, rcp.time+rcp.dur)
-
-  def test_wmma(self):
-    with save_sqtt() as sqtt:
-      asm_kernel([
-        "v_wmma_f32_16x16x16_f16 v[16:23], v[0:7], v[8:15], v[16:23]",
-        "v_add_f32_e32 v0 v16 v0",
-      ], l=32*4).realize()
-    assert len(sqtt) == 2, f"expected two waves, got {len(sqtt)} {list(sqtt.keys())}"
-    wmma = list(sqtt.values())[0][0]
-    self.assertGreater(wmma.dur, 1) # rgp says 32 clocks
-
-  def test_sleep(self):
-    n = 1
-    def sleep_kernel(data0):
-      assert data0.dtype.base == dtypes.ulong
-      op = custom("unsigned long long t0 = __builtin_readcyclecounter();")
-      op = custom(f"__builtin_amdgcn_s_sleep({n});", op)
-      op = custom(f"unsigned long long t1 = __builtin_readcyclecounter();", op)
-      op = custom(f"data0_{data0.size}[0] = t1 - t0;", op)
-      return UOp.sink(data0, op, arg=KernelInfo(name=f"sleep_{n}"))
-    diff_hw_reg = Tensor.empty(1, dtype=dtypes.ulong)
-    diff_hw_reg = Tensor.custom_kernel(diff_hw_reg, fxn=sleep_kernel)[0]
-    with save_sqtt() as sqtt:
-      diff_hw_reg.realize()
-    diff_sqtt = list(sqtt.values())[0][2]
-    self.assertEqual(diff_sqtt.dur, diff_hw_reg.item()-1) # 1 cycle for reading the counter register
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,106 +0,0 @@
-#include "kittens.cuh"
-
-using namespace kittens;
-
-constexpr int NUM_WORKERS = 4;
-constexpr int PIPE_STAGES = 3;
-
-constexpr int ATTN_B = 16;
-constexpr int ATTN_N = 1024;
-constexpr int ATTN_H = 16;
-constexpr int ATTN_D = 64;
-
-template<int D> constexpr size_t ROWS = 16*(64/D); // height of each worker tile (rows)
-template<int D, typename T=bf16, typename L=row_l> using qkvo_tile = rt<T, ROWS<D>, D, L>;
-template<int D, typename T=float> using attn_tile = rt<T, ROWS<D>, ROWS<D>>;
-template<int D> using shared_tile = st_bf<ROWS<D>, D>;
-template<int D> using global_layout = gl<bf16, -1, -1, -1, D>; // B, N, H, specified at runtime, D known at compile time for this kernel
-template<int D> struct globals { global_layout<D> Qg, Kg, Vg, Og; };
-
-__launch_bounds__(NUM_WORKERS*WARP_THREADS, 1)
-__global__ void attend_ker(bf16 *O_ptr, bf16 *Q_ptr, bf16 *K_ptr, bf16 *V_ptr) {
-    constexpr int D = ATTN_D;
-    global_layout<D> Qg{Q_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
-    global_layout<D> Kg{K_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
-    global_layout<D> Vg{V_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
-    global_layout<D> Og{O_ptr, ATTN_B, ATTN_N, ATTN_H, nullptr};
-    globals<D> g(Qg, Kg, Vg, Og);
-
-    using load_group = kittens::group<2>; // pairs of workers collaboratively load k, v tiles
-    int loadid = load_group::groupid(), workerid = kittens::warpid(); // which worker am I?
-    constexpr int LOAD_BLOCKS = NUM_WORKERS / load_group::GROUP_WARPS;
-    const int batch = blockIdx.z, head = blockIdx.y, q_seq = blockIdx.x * NUM_WORKERS + workerid;
-
-    extern __shared__ alignment_dummy __shm[];
-    shared_allocator al((int*)&__shm[0]);
-
-    shared_tile<D> (&k_smem)[LOAD_BLOCKS][PIPE_STAGES] = al.allocate<shared_tile<D>, LOAD_BLOCKS, PIPE_STAGES>();
-    shared_tile<D> (&v_smem)[LOAD_BLOCKS][PIPE_STAGES] = al.allocate<shared_tile<D>, LOAD_BLOCKS, PIPE_STAGES>();
-
-    shared_tile<D> (&qo_smem)[NUM_WORKERS] = reinterpret_cast<shared_tile<D>(&)[NUM_WORKERS]>(k_smem);
-    // Initialize all of the register tiles.
-    qkvo_tile<D, bf16> q_reg, k_reg; // Q and K are both row layout, as we use mma_ABt.
-    qkvo_tile<D, bf16, col_l> v_reg; // V is column layout, as we use mma_AB.
-    qkvo_tile<D, float> o_reg; // Output tile.
-    attn_tile<D, float> att_block; // attention tile, in float. (We want to use float wherever possible.)
-    attn_tile<D, bf16> att_block_mma; // bf16 attention tile for the second mma_AB. We cast right before that op.
-    typename attn_tile<D, float>::col_vec max_vec_last, max_vec, norm_vec; // these are column vectors for the in-place softmax.
-    // each warp loads its own Q tile of 16x64
-    if (q_seq*ROWS<D> < g.Qg.depth()) {
-      warp::load<1, false>(qo_smem[workerid], g.Qg, {batch, q_seq, head, 0});  // going through shared memory improves coalescing of dram reads.
-      __syncwarp();
-      warp::load(q_reg, qo_smem[workerid]);
-    }
-    __syncthreads();
-
-    if constexpr(D == 64) q_reg *= __float2bfloat16(0.125f * 1.44269504089f);
-    else if constexpr(D == 128) q_reg *= __float2bfloat16(0.08838834764f * 1.44269504089f);
-
-    max_vec = base_types::constants<float>::neg_infty();
-    norm_vec = 0.f;
-    o_reg = 0.f;
-    // launch the load of the first k, v tiles
-    int kv_blocks = (g.Kg.depth() + LOAD_BLOCKS*ROWS<D>-1) / (LOAD_BLOCKS*ROWS<D>), tic = 0;
-    load_group::load_async<1, false>(k_smem[loadid][0], g.Kg, {batch, loadid, head, 0});
-    load_group::load_async<1, false>(v_smem[loadid][0], g.Vg, {batch, loadid, head, 0});
-    // iterate over k, v for these q's that have been loaded
-    for(auto kv_idx = 0; kv_idx < kv_blocks; kv_idx++, tic=(tic+1)%3) {
-        int next_load_idx = (kv_idx+1)*LOAD_BLOCKS + loadid;
-        if(next_load_idx*ROWS<D> < g.Kg.depth()) {
-            int next_tic = (tic+1)%3;
-            load_group::load_async<1, false>(k_smem[loadid][next_tic], g.Kg, {batch, next_load_idx, head, 0});
-            load_group::load_async<1, false>(v_smem[loadid][next_tic], g.Vg, {batch, next_load_idx, head, 0});
-            load_async_wait<1>(); // next k, v can stay in flight.
-        }
-        else load_async_wait();
-        __syncthreads();
-
-        #pragma unroll LOAD_BLOCKS
-        for(int subtile = 0; subtile < LOAD_BLOCKS && (kv_idx*LOAD_BLOCKS + subtile)*ROWS<D> < g.Kg.depth(); subtile++) {
-          warp::load(k_reg, k_smem[subtile][tic]); // load k from shared into registers
-          att_block = 0.f; // zero 16x16 attention tile
-          warp::mma<transpose::N, transpose::T>(att_block, q_reg, k_reg, att_block); // [email protected]
-          // int first_index = (kv_idx*LOAD_BLOCKS + subtile)*ROWS<D>; // one past the last KV index of this tile
-          // int start_fill = g.Kg.depth()-first_index < ROWS<D> ? g.Kg.depth()-first_index : ROWS<D>;
-          // right_fill(att_block, att_block, start_fill, base_types::constants<float>::neg_infty());
-          max_vec_last = max_vec;
-          max_vec = warp::max<axis::COL>(att_block, max_vec); 
-          att_block = warp::exp2(att_block - max_vec); 
-          max_vec_last = warp::exp2(max_vec_last - max_vec); 
-          norm_vec *= max_vec_last; 
-          norm_vec = warp::sum<axis::COL>(att_block, norm_vec); 
-          att_block_mma = att_block; // copy to bf16 tile
-          warp::load(v_reg, v_smem[subtile][tic]);
-          o_reg *= max_vec_last;
-          warp::mma<transpose::N, transpose::N>(o_reg, att_block_mma, v_reg, o_reg);
-        }
-    }
-
-    o_reg /= norm_vec;
-    __syncthreads();
-    if (q_seq*ROWS<D> < g.Og.depth()) { // write out o.
-      warp::store(qo_smem[workerid], o_reg); // going through shared memory improves coalescing of dram writes.
-      __syncwarp();
-      warp::store<1, false>(g.Og, qo_smem[workerid], {batch, q_seq, head, 0});
-    }
-}
@@ -1,43 +0,0 @@
-import pathlib
-from tinygrad import Device, Tensor
-from tinygrad.helpers import Context
-from tinygrad.runtime.support.compiler_cuda import pretty_ptx, NVCCCompiler
-
-if __name__ == "__main__":
-  code = (pathlib.Path(__file__).parent / "fa.cu").read_text()
-  device = Device["CUDA"]
-  kitten_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "--expt-relaxed-constexpr", "-DKITTENS_4090"]
-  lib = NVCCCompiler(device.compiler.arch, kitten_args).compile(code)
-  kernel_name = lib.decode().split(".globl\t")[1].split("\n")[0]
-  print("kernel name", kernel_name)
-  print(pretty_ptx(lib.decode()))
-
-  prg = device.runtime(kernel_name, lib)
-  prg.smem = 16384 * 3
-
-  B, N, H, D = 16, 1024, 16, 64
-  q = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
-  k = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
-  v = Tensor.randn(B, N, H, D, device='CUDA', dtype="bfloat16")
-  out = Tensor.empty(B, N, H, D, device='CUDA', dtype="bfloat16")
-  Tensor.realize(q, k, v, out)
-
-  NUM_WORKERS = 4
-  ROWS = 16 * (64 // D)
-
-  gsz = (N // (ROWS*NUM_WORKERS), H, B)
-  for _ in range(5):
-    et = prg(out.uop.buffer.ensure_allocated()._buf, q.uop.buffer._buf, k.uop.buffer._buf, v.uop.buffer._buf,
-             global_size=gsz, local_size=(ROWS*NUM_WORKERS,1,1), wait=True)
-
-    attn_flops = 2 * B * H * N * N * D + \
-                 4 * B * H * N * N + \
-                 2 * B * H * N * N * D
-    print(f"{attn_flops/(et*1e9):2f} GFLOPS")
-
-  for _ in range(5):
-    with Context(DEBUG=2):
-      ref = q.scaled_dot_product_attention(k, v)
-
-  ref, out = ref.float(), out.float()
-  print((ref-out).mean().item(), (ref-out).max().item())
@@ -1,400 +0,0 @@
-/**
- * @file
- * @brief Basic operations on generic types.
- */
-
-#pragma once
-
-#include <cuda_bf16.h>
-#include <limits>
-#include "base_types.cuh"
-
-namespace kittens {
-
-/**
- * @namespace base_ops
- *
- * @brief A namespace for operations on basic data types.
- */
-namespace base_ops {
-
-/* ----------  CONST OPS  ---------- */
-
-/**
- * @brief Represents the zero constant operation.
- *
- * This operation returns the zero value of the specified type.
- *
- * @tparam T The data type for which to return the zero value.
- * @return The zero value of type T.
- */
-struct zero {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::zero();      }
-};
-/**
- * @brief Represents the one constant operation.
- *
- * This operation returns the one value of the specified type.
- *
- * @tparam T The data type for which to return the one value.
- * @return The one value of type T.
- */
-struct one {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::one();       }
-};
-/**
- * @brief Represents the positive infinity constant operation.
- *
- * This operation returns the positive infinity value of the specified type.
- *
- * @tparam T The data type for which to return the positive infinity value.
- * @return The positive infinity value of type T.
- */
-struct pos_infty {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
-};
-/**
- * @brief Represents the negative infinity constant operation.
- *
- * This operation returns the negative infinity value of the specified type.
- *
- * @tparam T The data type for which to return the negative infinity value.
- * @return The negative infinity value of type T.
- */
-struct neg_infty {
-    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
-};
-
-
-/* ----------  UNARY OPS  ---------- */
-
-/**
- * @brief Exponential function operation.
- *
- * This operation calculates the exponential of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The exponential of the input value.
- */
-struct exp {
-    template<typename T> static __device__ inline T op(const T &x) { return exp(x); }
-};
-template<> __device__ inline float  exp::op<float> (const float &x ) { return __expf(x);                        }
-template<> __device__ inline float2 exp::op<float2>(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; }
-template<> __device__ inline bf16   exp::op<bf16>  (const bf16 &x  ) { return hexp(x);                          }
-template<> __device__ inline bf16_2 exp::op<bf16_2>(const bf16_2 &x) { return h2exp(x);                         }
-template<> __device__ inline half   exp::op<half>  (const half &x  ) { return hexp(x);                          }
-template<> __device__ inline half_2 exp::op<half_2>(const half_2 &x) { return h2exp(x);                         }
-
-/**
- * @brief Exponential function operation, in base 2
- *
- * This operation calculates the exponential of the input value, in base 2.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The exponential of the input value.
- */
-struct exp2 {
-    template<typename T> static __device__ inline T op(const T &x) { return exp2f(x); }
-};
-template<> __device__ inline float  exp2::op<float> (const float &x ) { return exp2f(x);                        }
-template<> __device__ inline float2 exp2::op<float2>(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; }
-template<> __device__ inline bf16   exp2::op<bf16>  (const bf16 &x  ) { return hexp2(x);                          }
-template<> __device__ inline bf16_2 exp2::op<bf16_2>(const bf16_2 &x) { return h2exp2(x);                         }
-template<> __device__ inline half   exp2::op<half>  (const half &x  ) { return hexp2(x);                          }
-template<> __device__ inline half_2 exp2::op<half_2>(const half_2 &x) { return h2exp2(x);                         }
-/**
- * @brief Natural log function operation.
- *
- * This operation calculates the natural logarithm of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The natural logarithm of the input value.
- */
-struct log {
-    template<typename T> static __device__ inline T op(const T &x) { return log(x); }
-};
-template<> __device__ inline float  log::op<float> (const float &x ) { return __logf(x);                        }
-template<> __device__ inline float2 log::op<float2>(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; }
-template<> __device__ inline bf16   log::op<bf16>  (const bf16 &x  ) { return hlog(x);                          }
-template<> __device__ inline bf16_2 log::op<bf16_2>(const bf16_2 &x) { return h2log(x);                         }
-template<> __device__ inline half   log::op<half>  (const half &x  ) { return hlog(x);                          }
-template<> __device__ inline half_2 log::op<half_2>(const half_2 &x) { return h2log(x);                         }
-/**
- * @brief Logarithm base 2 operation.
- *
- * This operation calculates the logarithm base 2 of the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The logarithm base 2 of the input value.
- */
-struct log2 {
-    template<typename T> static __device__ inline T op(const T &x) { return log2(x); }
-};
-template<> __device__ inline float  log2::op<float> (const float &x ) { return __log2f(x);                        }
-template<> __device__ inline float2 log2::op<float2>(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; }
-template<> __device__ inline bf16   log2::op<bf16>  (const bf16 &x  ) { return hlog2(x);                          }
-template<> __device__ inline bf16_2 log2::op<bf16_2>(const bf16_2 &x) { return h2log2(x);                         }
-template<> __device__ inline half   log2::op<half>  (const half &x  ) { return hlog2(x);                          }
-template<> __device__ inline half_2 log2::op<half_2>(const half_2 &x) { return h2log2(x);                         }
-/**
- * @brief Absolute value operation.
- *
- * This operation calculates the absolute value of the input.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The absolute value of the input.
- */
-struct abs {
-    template<typename T> static __device__ inline T op(const T &x) { return abs(x); }
-};
-template<> __device__ inline float  abs::op<float> (const float &x ) { return fabsf(x);                       }
-template<> __device__ inline float2 abs::op<float2>(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; }
-template<> __device__ inline bf16   abs::op<bf16>  (const bf16 &x  ) { return __habs(x);                      }
-template<> __device__ inline bf16_2 abs::op<bf16_2>(const bf16_2 &x) { return __habs2(x);                     }
-template<> __device__ inline half   abs::op<half>  (const half &x  ) { return __habs(x);                      }
-template<> __device__ inline half_2 abs::op<half_2>(const half_2 &x) { return __habs2(x);                     }
-/**
- * @brief Rectified Linear Unit (ReLU) operation.
- *
- * This operation applies the ReLU function to the input, which is the
- * maximum of zero and the input value.
- *
- * @tparam T The data type of the input and output values.
- * @param x[in] The input value.
- * @return The result of ReLU function applied to the input.
- */
-struct relu {
-    template<typename T> static __device__ inline T op(const T &x) { return max(x, base_types::constants<T>::zero()); }
-};
-template<> __device__ inline float  relu::op<float> (const float &x ) { return max(x, 0.f);                                  }
-template<> __device__ inline float2 relu::op<float2>(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)};         }
-template<> __device__ inline bf16   relu::op<bf16>  (const bf16 &x  ) { return __hmax(x, base_types::constants<bf16>::zero());    }
-template<> __device__ inline bf16_2 relu::op<bf16_2>(const bf16_2 &x) { return __hmax2(x, base_types::constants<bf16_2>::zero()); }
-template<> __device__ inline half   relu::op<half>  (const half &x  ) { return __hmax(x, base_types::constants<half>::zero());    }
-template<> __device__ inline half_2 relu::op<half_2>(const half_2 &x) { return __hmax2(x, base_types::constants<half_2>::zero()); }
-/**
- * @brief Copy operation.
- *
- * This operation returns the input value unchanged.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The input value.
- * @return The same value as the input.
- */
-struct copy { // for non-compile-time setters.
-    template<typename T> static __device__ inline T op(const T &a) { return a; }
-};
-
-
-/* ----------  BINARY OPS  ---------- */
-
-/**
- * @brief Copy2 operation.
- *
- * This operation returns the second input value unchanged.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value (ignored).
- * @param b[in] The second input value.
- * @return The same value as the second input.
- */
-struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return b; }
-};
-/**
- * @brief Sum operation.
- *
- * This operation calculates the sum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The sum of the input values.
- */
-struct sum {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a+b; }
-};
-template<> __device__ inline float2 sum::op<float2>(const float2 &a, const float2 &b) {
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x+b.x, a.y+b.y};
-#endif
-}
-template<> __device__ inline bf16   sum::op<bf16>  (const bf16   &a, const bf16   &b) { return __hadd(a, b);             }
-template<> __device__ inline bf16_2 sum::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b);            }
-template<> __device__ inline half   sum::op<half>  (const half   &a, const half   &b) { return __hadd(a, b);             }
-template<> __device__ inline half_2 sum::op<half_2>(const half_2 &a, const half_2 &b) { return __hadd2(a, b);            }
-/**
- * @brief Subtraction operation.
- *
- * This operation calculates the difference between two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The difference between the input values.
- */
-struct sub {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a-b; }
-};
-template<> __device__ inline float2 sub::op<float2>(const float2 &a, const float2 &b) { 
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x-b.x, a.y-b.y}; 
-#endif
-}
-template<> __device__ inline bf16   sub::op<bf16>  (const bf16   &a, const bf16   &b) { return __hsub(a, b);             }
-template<> __device__ inline bf16_2 sub::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b);            }
-template<> __device__ inline half   sub::op<half>  (const half   &a, const half   &b) { return __hsub(a, b);             }
-template<> __device__ inline half_2 sub::op<half_2>(const half_2 &a, const half_2 &b) { return __hsub2(a, b);            }
-/**
- * @brief Multiplication operation.
- *
- * This operation calculates the product of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The product of the input values.
- */
-struct mul {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a*b; }
-};
-template<> __device__ inline float2 mul::op<float2>(const float2 &a, const float2 &b) { 
-#ifdef KITTENS_BLACKWELL
-    float2 c;
-    asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
-    return c;
-#else
-    return float2{a.x*b.x, a.y*b.y}; 
-#endif
-}
-template<> __device__ inline bf16   mul::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmul(a, b);             }
-template<> __device__ inline bf16_2 mul::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b);            }
-template<> __device__ inline half   mul::op<half>  (const half   &a, const half   &b) { return __hmul(a, b);             }
-template<> __device__ inline half_2 mul::op<half_2>(const half_2 &a, const half_2 &b) { return __hmul2(a, b);            }
-/**
- * @brief Division operation.
- *
- * This operation calculates the quotient of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The quotient of the input values.
- */
-struct div {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a/b; }
-};
-template<> __device__ inline float2 div::op<float2>(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; }
-template<> __device__ inline bf16   div::op<bf16>  (const bf16   &a, const bf16   &b) { return __hdiv(a, b);             }
-template<> __device__ inline bf16_2 div::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b);            } // this op is a special snowflake
-template<> __device__ inline half   div::op<half>  (const half   &a, const half   &b) { return __hdiv(a, b);             }
-template<> __device__ inline half_2 div::op<half_2>(const half_2 &a, const half_2 &b) { return __h2div(a, b);            }
-/**
- * @brief Maximum operation.
- *
- * This operation calculates the maximum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The maximum of the input values.
- */
- struct max {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); }
-};
-template<>  __device__ inline float2 max::op<float2>(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; }
-template<>  __device__ inline bf16   max::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmax(a, b);                             }
-template<>  __device__ inline bf16_2 max::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b);                            }
-template<>  __device__ inline half   max::op<half>  (const half   &a, const half   &b) { return __hmax(a, b);                             }
-template<>  __device__ inline half_2 max::op<half_2>(const half_2 &a, const half_2 &b) { return __hmax2(a, b);                            }
-/**
- * @brief Minimum operation.
- *
- * This operation calculates the minimum of two input values.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @return The minimum of the input values.
- */
-struct min {
-    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); }
-};
-template<>  __device__ inline float2 min::op<float2>(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; }
-template<>  __device__ inline bf16   min::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmin(a, b);                         }
-template<>  __device__ inline bf16_2 min::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b);                        }
-template<>  __device__ inline half   min::op<half>  (const half   &a, const half   &b) { return __hmin(a, b);                         }
-template<>  __device__ inline half_2 min::op<half_2>(const half_2 &a, const half_2 &b) { return __hmin2(a, b);                        }
-
-
-/* ----------  TERNARY OPS  ---------- */
-
-/**
- * @brief Fused multiply-add operation A * B + C.
- *
- * This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The second input value.
- * @param c[in] The third input value to be added.
- * @return The result of the fused multiply-add operation.
- */
-struct fma_AxBtC {
-    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
-        return sum::op<T>(mul::op<T>(a, b), c);
-    }
-};
-template<> __device__ inline float2 fma_AxBtC::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
-#ifdef KITTENS_BLACKWELL
-    float2 d;
-    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c));
-    return d;
-#else
-    return float2{a.x*b.x+c.x, a.y*b.y+c.y};
-#endif
-}
-/**
- * @brief Fused multiply-add operation A * C + B.
- *
- * This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
- * This is particularly useful for attention mechanisms in neural networks.
- *
- * @tparam T The data type of the input and output values.
- * @param a[in] The first input value.
- * @param b[in] The third input value to be added.
- * @param c[in] The second input value.
- * @return The result of the fused multiply-add operation.
- */
-struct fma_AxCtB { // this is the one needed for attention
-    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
-        return sum::op<T>(mul::op<T>(a, c), b);
-    }
-};
-template<> __device__ inline float2 fma_AxCtB::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
-#ifdef KITTENS_BLACKWELL
-    float2 d;
-    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b));
-    return d;
-#else
-    return float2{a.x*c.x+b.x, a.y*c.y+b.y};
-#endif
-}
-
-} // namespace base_ops
-
-} // namespace kittens
@@ -1,519 +0,0 @@
-/**
- * @file
- * @brief Declarations, manipulations, and wrappers for basic types.
- * 
- * This file is a bunch of utilities for going back and forth between different types.
- * 
- * Many of them are for the compiler, so as to clean up the code. It unfortunately
- * seems necessary when we have types we really care about that are less than word width.
- */
-
-#pragma once
-
-#ifdef KITTENS_HOPPER
-#include <cuda_fp8.h>
-#endif
-
-#include <cuda_bf16.h>
-#include <cuda_fp16.h>
-#include <string>
-#include <bit>
-
-
-namespace kittens {
-
-/**
- * @brief Bfloat16 floating-point type.
- */
-using bf16 = __nv_bfloat16;
-/**
- * @brief Half-precision floating-point type.
- */
-using half = __half;
-/**
- * @brief Packed word of two bfloat16 floating-point values.
- */
-using bf16_2 = __nv_bfloat162;
-/**
- * @brief Packed word of two half-precision floating-point values.
- */
-using half_2 = __half2;
-#ifdef KITTENS_HOPPER
-/**
- * @brief float8 floating-point type.
- */
-using fp8e4m3 = __nv_fp8_e4m3;
-using fp8e5m2 = __nv_fp8_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0 = __nv_fp8_e8m0;
-#endif
-/**
- * @brief 2-packed float8 floating-point type.
- */
-using fp8e4m3_2 = __nv_fp8x2_e4m3;
-using fp8e5m2_2 = __nv_fp8x2_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0_2 = __nv_fp8x2_e8m0;
-#endif
-/**
- * @brief 4-packed float8 floating-point type.
- */
-using fp8e4m3_4 = __nv_fp8x4_e4m3;
-using fp8e5m2_4 = __nv_fp8x4_e5m2;
-#ifdef KITTENS_BLACKWELL
-using fp8e8m0_4 = __nv_fp8x4_e8m0;
-#endif
-#endif
-
-namespace ducks {
-/**
- * @namespace base_types
- *
- * @brief A namespace for concepts for basic data types.
- */
-namespace base_types {
-
-#ifdef KITTENS_HOPPER
-#ifdef KITTENS_BLACKWELL
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4> || std::is_same_v<T, fp8e8m0_4>; // could add half_2 later if implemented.
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2> || std::is_same_v<T, fp8e8m0>; // could add half_2 later if implemented.
-#else
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4>;
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2>;
-#endif
-#else
-template<typename T>
-concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2>;
-template<typename T>
-concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half>;
-#endif
-
-} // namespace base_types
-} // namespace ducks
-
-/**
- * @namespace base_types
- *
- * @brief A namespace for ThunderKittens basic data types.
- */
-namespace base_types {
-
-/**
- * @brief Provides compile-time constants for different types.
- *
- * @tparam T The type for which to provide constants.
- */
-template<typename T> struct constants {
-    /**
-     * @brief Zero
-     * @return Constexpr zero with type T
-     */
-    static __device__ inline constexpr T zero()      { return T{0}; }
-    /**
-     * @brief One
-     * @return Constexpr one with type T
-     */
-    static __device__ inline constexpr T one()       { return T{1}; }
-    /**
-     * @brief Positive infinity. Particularly useful for initializing before a min op.
-     * @return Constexpr positive infinity with type T
-     */
-    static __device__ inline constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
-    /**
-     * @brief Negative infinity. Particularly useful for initializing before a max op.
-     * @return Constexpr negative infinity with type T
-     */
-    static __device__ inline constexpr T neg_infty() { return T{-INFINITY}; }
-};
-template<> struct constants<float2> {
-    static __device__ inline constexpr float2 zero()      { return float2{0.f, 0.f}; }
-    static __device__ inline constexpr float2 one()       { return float2{1.f, 1.f}; }
-    static __device__ inline constexpr float2 pos_infty() { return float2{constants<float>::pos_infty(), constants<float>::pos_infty()}; }
-    static __device__ inline constexpr float2 neg_infty() { return float2{constants<float>::neg_infty(), constants<float>::neg_infty()}; }
-};
-template<> struct constants<bf16> {
-    static __device__ inline constexpr bf16 zero()      { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } // unfortunately __float2bf16_rn is not constexpr
-    static __device__ inline constexpr bf16 one()       { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); }
-    static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); }
-    static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); }
-};
-template<> struct constants<bf16_2> {
-    static __device__ inline constexpr bf16_2 zero()      { return bf16_2{constants<bf16>::zero(),      constants<bf16>::zero()};      }
-    static __device__ inline constexpr bf16_2 one()       { return bf16_2{constants<bf16>::one(),       constants<bf16>::one()};       }
-    static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants<bf16>::pos_infty(), constants<bf16>::pos_infty()}; }
-    static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants<bf16>::neg_infty(), constants<bf16>::neg_infty()}; }
-};
-template<> struct constants<half> {
-    static __device__ inline constexpr half zero()      { return std::bit_cast<__half>(uint16_t(0x0000)); }
-    static __device__ inline constexpr half one()       { return std::bit_cast<__half>(uint16_t(0x3C00)); }
-    static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); }
-    static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); }
-};
-template<> struct constants<half_2> {
-    static __device__ inline constexpr half_2 zero()      { return half_2{constants<half>::zero(),      constants<half>::zero()};      }
-    static __device__ inline constexpr half_2 one()       { return half_2{constants<half>::one(),       constants<half>::one()};       }
-    static __device__ inline constexpr half_2 pos_infty() { return half_2{constants<half>::pos_infty(), constants<half>::pos_infty()}; }
-    static __device__ inline constexpr half_2 neg_infty() { return half_2{constants<half>::neg_infty(), constants<half>::neg_infty()}; }
-};
-#ifdef KITTENS_HOPPER
-template<> struct constants<fp8e4m3> {
-    static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); }
-    static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); }
-};
-template<> struct constants<fp8e4m3_2> {
-    static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x0000)); }
-    static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x3838)); }
-};
-template<> struct constants<fp8e4m3_4> {
-    static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x00000000)); }
-    static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x38383838)); }
-};
-template<> struct constants<fp8e5m2> {
-    static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); }
-    static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); }
-};
-template<> struct constants<fp8e5m2_2> {
-    static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x0000)); }
-    static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x3C3C)); }
-};
-template<> struct constants<fp8e5m2_4> {
-    static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x00000000)); }
-    static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x3C3C3C3C)); }
-};
-#endif
-
-template<> struct constants<int> {
-    static __device__ inline constexpr int zero()      { return 0; }
-    static __device__ inline constexpr int one()       { return 1; }
-};
-template<> struct constants<int2> {
-    static __device__ inline constexpr int2 zero()      { return int2{0, 0}; }
-    static __device__ inline constexpr int2 one()       { return int2{1, 1}; }
-};
-
-/**
- * @brief Provides information about packing of elements for a given type.
- *
- * @tparam T The type for which to provide packing information.
- */
-template<typename T> struct packing {
-    /**
-     * @brief The number of elements packed together.
-     *
-     * @return constexpr int representing number of elements within the type.
-     */
-    static __device__ inline constexpr int num() { return 1; }
-    /**
-     * @brief Packs a single T element twice (replicated) into its packed type.
-     *
-     * @param i[in] The element to pack.
-     * @return The packed type.
-     */
-    static __device__ inline constexpr T pack(const bf16 &i);
-};
-template<> struct packing<bf16> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = bf16;
-    using packed_type = bf16_2;
-    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; }
-};
-template<> struct packing<bf16_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = bf16;
-    using packed_type = bf16_2;
-    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<half> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = half;
-    using packed_type = half_2;
-    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; }
-};
-template<> struct packing<half_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = half;
-    using packed_type = half_2;
-    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<float> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = float;
-    using packed_type = float2;
-    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; }
-};
-template<> struct packing<float2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = float;
-    using packed_type = float2;
-    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<char> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = char;
-    using packed_type = char2;
-    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<char2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = char;
-    using packed_type = char2;
-    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<int> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = int;
-    using packed_type = int2;
-    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<int2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = int;
-    using packed_type = int2;
-    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = uint;
-    using packed_type = uint2;
-    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = uint;
-    using packed_type = uint2;
-    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
-};
-struct uint64_2 { uint64_t x, y; };
-template<> struct packing<uint64_t> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = uint64_t;
-    using packed_type = uint64_2;
-    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<uint64_2> {
-    static __device__ inline constexpr int num() { return 2; }
-    using unpacked_type = uint64_t;
-    using packed_type = uint64_2;
-    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
-};
-template<> struct packing<float4> {
-    static __device__ inline constexpr int num() { return 4; }
-};
-template<> struct packing<int4> {
-    static __device__ inline constexpr int num() { return 4; }
-};
-#ifdef KITTENS_HOPPER
-template<> struct packing<fp8e4m3> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e4m3;
-    using packed_type = fp8e4m3_4;
-};
-template<> struct packing<fp8e4m3_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e4m3;
-    using packed_type = fp8e4m3_4;
-};
-template<> struct packing<fp8e5m2> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e5m2;
-    using packed_type = fp8e5m2_4;
-};
-template<> struct packing<fp8e5m2_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e5m2;
-    using packed_type = fp8e5m2_4;
-};
-#ifdef KITTENS_BLACKWELL
-template<> struct packing<fp8e8m0> {
-    static __device__ inline constexpr int num() { return 1; }
-    using unpacked_type = fp8e8m0;
-    using packed_type = fp8e8m0_4;
-};
-template<> struct packing<fp8e8m0_4> {
-    static __device__ inline constexpr int num() { return 4; }
-    using unpacked_type = fp8e8m0;
-    using packed_type = fp8e8m0_4;
-};
-#endif
-#endif
-
-
-/**
- * @brief Provides templated functionality to convert between different types.
- *
- * @tparam T The target type for conversion.
- * @tparam U The source type for conversion.
- */
-template<typename T, typename U> struct convertor {
-    /**
-     * @brief Converts a value of type U to type T.
-     *
-     * @param u[in] The value of type U to convert.
-     * @return T The converted value of type T.
-     */
-    static __host__ __device__ inline T convert(const U & u) {
-        return (T)u;
-    }
-};
-template<> struct convertor<float, bf16> {
-    static __host__ __device__ inline float convert(const bf16 & u) {
-        return 	__bfloat162float(u);
-    }
-};
-template<> struct convertor<bf16, float> {
-    static __host__ __device__ inline bf16 convert(const float & u) {
-        return 	__float2bfloat16_rn(u);
-    }
-};
-template<> struct convertor<float2, bf16_2> {
-    static __host__ __device__ inline float2 convert(const bf16_2 & u) {
-        return 	__bfloat1622float2(u);
-    }
-};
-template<> struct convertor<bf16_2, float2> {
-    static __host__ __device__ inline bf16_2 convert(const float2 & u) {
-        return 	__float22bfloat162_rn(u);
-    }
-};
-template<> struct convertor<float, half> {
-    static __host__ __device__ inline float convert(const half & u) {
-        return __half2float(u);
-    }
-};
-template<> struct convertor<half, float> {
-    static __host__ __device__ inline half convert(const float & u) {
-        return __float2half(u);
-    }
-};
-template<> struct convertor<float2, half_2> {
-    static __host__ __device__ inline float2 convert(const half_2 & u) {
-        return __half22float2(u);
-    }
-};
-template<> struct convertor<half_2, float2> {
-    static __host__ __device__ inline half_2 convert(const float2 & u) {
-        return __float22half2_rn(u);
-    }
-};
-template<> struct convertor<bf16, half> {
-    static __host__ __device__ inline bf16 convert(const half & u) {
-        return __float2bfloat16_rn(__half2float(u));
-    }
-};
-template<> struct convertor<half, bf16> {
-    static __host__ __device__ inline half convert(const bf16 & u) {
-        return __float2half(__bfloat162float(u));
-    }
-};
-template<> struct convertor<bf16_2, half_2> {
-    static __host__ __device__ inline bf16_2 convert(const half_2 & u) {
-        return __float22bfloat162_rn(__half22float2(u));
-    }
-};
-template<> struct convertor<half_2, bf16_2> {
-    static __host__ __device__ inline half_2 convert(const bf16_2 & u) {
-        return __float22half2_rn(__bfloat1622float2(u));
-    }
-};
-#ifdef KITTENS_HOPPER
-// fp8e4m3
-template<> struct convertor<fp8e4m3_4, float4> {
-    static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) {
-        return __nv_fp8x4_e4m3(u); 
-    }
-};
-template<> struct convertor<float4, fp8e4m3_4> {
-    static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) {
-        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u));
-        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
-    }
-};
-template<> struct convertor<fp8e4m3_2, float2> {
-    static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) {
-        return __nv_fp8x2_e4m3(u); 
-    }
-};
-template<> struct convertor<float2, fp8e4m3_2> {
-    static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) {
-        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u));
-        return make_float2(float(vals[0]), float(vals[1]));
-    }
-};
-template<> struct convertor<fp8e4m3, float> {
-    static __host__ __device__ inline fp8e4m3 convert(const float & u) {
-        return __nv_fp8_e4m3(u);
-    }
-};
-template<> struct convertor<float, fp8e4m3> {
-    static __host__ __device__ inline float convert(const fp8e4m3 & u) {
-        return float(u);
-    }
-};
-template<> struct convertor<bf16_2, fp8e4m3_4> {
-    static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) {
-        float4 f4 = convertor<float4, fp8e4m3_4>::convert(u);
-        float2 f2 = make_float2(f4.x, f4.y);
-        return __float22bfloat162_rn(f2);
-    }
-};
-template<> struct convertor<fp8e4m3_4, bf16_2> {
-    static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) {
-        float2 f2 = __bfloat1622float2(u);
-        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
-        return __nv_fp8x4_e4m3(f4);
-    }
-};
-// fp8e5m2
-template<> struct convertor<fp8e5m2_4, float4> {
-    static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) {
-        return __nv_fp8x4_e5m2(u); 
-    }
-};
-template<> struct convertor<float4, fp8e5m2_4> {
-    static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) {
-        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u));
-        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
-    }
-};
-template<> struct convertor<fp8e5m2_2, float2> {
-    static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) {
-        return __nv_fp8x2_e5m2(u); 
-    }
-};
-template<> struct convertor<float2, fp8e5m2_2> {
-    static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) {
-        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u));
-        return make_float2(float(vals[0]), float(vals[1]));
-    }
-};
-template<> struct convertor<fp8e5m2, float> {
-    static __host__ __device__ inline fp8e5m2 convert(const float & u) {
-        return __nv_fp8_e5m2(u);
-    }
-};
-template<> struct convertor<float, fp8e5m2> {
-    static __host__ __device__ inline float convert(const fp8e5m2 & u) {
-        return float(u);
-    }
-};
-template<> struct convertor<bf16_2, fp8e5m2_4> {
-    static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) {
-        float4 f4 = convertor<float4, fp8e5m2_4>::convert(u);
-        float2 f2 = make_float2(f4.x, f4.y);
-        return __float22bfloat162_rn(f2);
-    }
-};
-template<> struct convertor<fp8e5m2_4, bf16_2> {
-    static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) {
-        float2 f2 = __bfloat1622float2(u);
-        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
-        return __nv_fp8x4_e5m2(f4);
-    }
-};
-#endif
-}
-}
@@ -1,11 +0,0 @@
-/**
- * @file
- * @brief A collection of common resources on which ThunderKittens depends.
- */
- 
-
-#pragma once
-
-#include "util.cuh"
-#include "base_types.cuh"
-#include "base_ops.cuh"
@@ -1,56 +0,0 @@
-#pragma once
-
-// Reset
-#define TK_RESET "\033[0m"
-
-// Foreground colors
-#define TK_FG_BLACK "\033[30m"
-#define TK_FG_RED "\033[31m"
-#define TK_FG_GREEN "\033[32m"
-#define TK_FG_YELLOW "\033[33m"
-#define TK_FG_BLUE "\033[34m"
-#define TK_FG_MAGENTA "\033[35m"
-#define TK_FG_CYAN "\033[36m"
-#define TK_FG_WHITE "\033[37m"
-
-// Background colors
-#define TK_BG_BLACK "\033[40m"
-#define TK_BG_RED "\033[41m"
-#define TK_BG_GREEN "\033[42m"
-#define TK_BG_YELLOW "\033[43m"
-#define TK_BG_BLUE "\033[44m"
-#define TK_BG_MAGENTA "\033[45m"
-#define TK_BG_CYAN "\033[46m"
-#define TK_BG_WHITE "\033[47m"
-
-// Bright foreground colors
-#define TK_FG_BRIGHT_BLACK "\033[90m"
-#define TK_FG_BRIGHT_RED "\033[91m"
-#define TK_FG_BRIGHT_GREEN "\033[92m"
-#define TK_FG_BRIGHT_YELLOW "\033[93m"
-#define TK_FG_BRIGHT_BLUE "\033[94m"
-#define TK_FG_BRIGHT_MAGENTA "\033[95m"
-#define TK_FG_BRIGHT_CYAN "\033[96m"
-#define TK_FG_BRIGHT_WHITE "\033[97m"
-
-// Bright background colors
-#define TK_BG_BRIGHT_BLACK "\033[100m"
-#define TK_BG_BRIGHT_RED "\033[101m"
-#define TK_BG_BRIGHT_GREEN "\033[102m"
-#define TK_BG_BRIGHT_YELLOW "\033[103m"
-#define TK_BG_BRIGHT_BLUE "\033[104m"
-#define TK_BG_BRIGHT_MAGENTA "\033[105m"
-#define TK_BG_BRIGHT_CYAN "\033[106m"
-#define TK_BG_BRIGHT_WHITE "\033[107m"
-
-// Text styles
-#define TK_BOLD "\033[1m"
-#define TK_DIM "\033[2m"
-#define TK_ITALIC "\033[3m"
-#define TK_UNDERLINE "\033[4m"
-#define TK_BLINK "\033[5m"
-#define TK_REVERSE "\033[7m"
-#define TK_HIDDEN "\033[8m"
-
-// Macro to combine styles
-#define TK_STYLE(...)  "\033[" #__VA_ARGS__ "m"
@@ -1,314 +0,0 @@
-/**
- * @file
- * @brief General utilities for ThunderKittens.
- */
-
-#pragma once
-
-#include <stdint.h>
-#include <type_traits>
-#include <concepts>
-#include <memory>
-
-// CUDA driver API
-#define CUCHECK(cmd) do {                                     \
-    CUresult err = cmd;                                       \
-    if (err != CUDA_SUCCESS) {                                \
-        const char *errStr;                                   \
-        cuGetErrorString(err, &errStr);                       \
-        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
-            __FILE__, __LINE__, errStr);                      \
-        exit(EXIT_FAILURE);                                   \
-    }                                                         \
-} while(0)
-
-// CUDA runtime API
-#define CUDACHECK(cmd) do {                                   \
-    cudaError_t err = cmd;                                    \
-    if (err != cudaSuccess) {                                 \
-        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
-            __FILE__, __LINE__, cudaGetErrorString(err));     \
-        exit(EXIT_FAILURE);                                   \
-    }                                                         \
-} while(0)
-
-/**
- * @namespace kittens
- *
- * @brief The main namespace of ThunderKittens.
- */
-namespace kittens {
-
-/* ----------  GENERAL CONSTANTS FOR KITTENS  ---------- */
-
-/**
- * @brief Tile dimension constant.
- */
-template<typename T> constexpr int TILE_COL_DIM = sizeof(T) == 1 ? 32 : 16;
-template<typename T> constexpr int TILE_ROW_DIM = 16;
-/**
- * @brief Tile num elements constant calculated as TILE_DIM squared.
- */
-template<typename T> constexpr int TILE_ELEMENTS{TILE_COL_DIM<T>*TILE_ROW_DIM<T>};
-/**
- * @brief Constant representing number of threads in a warp.
- */
-constexpr int WARP_THREADS{32};
-/**
- * @brief Constant representing number of threads in a warpgroup of four warps.
- */
-constexpr int WARPGROUP_THREADS{128};
-/**
-
- * @brief Constant representing number of warps in a warpgroup of four warps.
- */
-constexpr int WARPGROUP_WARPS{4};
-/**
-
- * @brief Get the warp ID of the current thread.
- * @return The warp ID.
- */
-__device__ static __forceinline__ int warpid() {
-    // uint32_t wid;
-    // asm volatile("mov.u32 %0, %warpid;" : "=r"(wid));
-    // return wid;
-    return threadIdx.x >> 5;
-}
-/**
- * @brief Get the warpgroup ID of the current thread.
- * @return The warpgroup ID.
- */
-__device__ static __forceinline__ int warpgroupid() { return warpid() >> 2; }
-/**
- * @brief Get the lane ID of the current thread within its warp.
- * @return The lane ID.
- */
-__device__ static __forceinline__ int laneid() {
-    // uint32_t lid;
-    // asm volatile("mov.u32 %0, %laneid;" : "=r"(lid));
-    // return lid;
-    return threadIdx.x & 31;
-}
-
-#if defined(KITTENS_HOPPER)
-constexpr int MAX_SHARED_MEMORY = 227000;
-#elif defined(KITTENS_A100)
-constexpr int MAX_SHARED_MEMORY = 164000;
-#elif defined(KITTENS_4090)
-constexpr int MAX_SHARED_MEMORY = 100000;
-#endif
-
-struct transpose {
-    static constexpr int N = 0; // not transposed
-    static constexpr int T = 1; // transposed
-};
-struct axis {
-    static constexpr int ROW = 0; // row axis of a tile
-    static constexpr int COL = 1; // column axis of a tile
-};
-
-/* ----------  TYPE HELPERS  ---------- */
-
-/**
- * @namespace ducks
- *
- * @brief ThunderKittens' namespace for template metaprogramming..
- * 
- * This includes primarily dummy types and concept wrappers, along
- * with a few additional utilities.
- */
-namespace ducks {
-
-/**
- * @brief A type representing an empty default for a template.
- */
-struct default_type {};
-
-// This macro can't be done as a template, so it doesn't really have a location in kittens.
-#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
-
-}
-
-/* ----------  SHUFFLE UTILS  ---------- */
-
-/**
- * @brief Mask constant for all active threads in a warp.
- */
-static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
-
-/**
- * @brief Perform a shuffle down operation on a packed type synchronously across a warp.
- * @tparam T The type of the value to be shuffled.
- * @param mask[in] The mask of active threads.
- * @param f[in] The value to be shuffled.
- * @param delta[in] The number of positions to shuffle down.
- * @return The result of the shuffle operation.
- */
-template<typename T>
-__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) {
-    return __shfl_down_sync(mask, f, delta);
-}
-template<>
-__device__ inline float2 packed_shfl_down_sync<float2>(uint32_t mask, const float2 &f, int delta) {
-    float2 r;
-    r.x = __shfl_down_sync(mask, f.x, delta);
-    r.y = __shfl_down_sync(mask, f.y, delta);
-    return r;
-}
-/**
- * @brief Perform a packed shuffle operation synchronously across a warp.
- * @tparam T The type of the value to be shuffled.
- * @param mask[in] The mask of active threads.
- * @param f[in] The value to be shuffled.
- * @param src[in] The source lane from which to shuffle.
- * @return The result of the shuffle operation.
- */
-template<typename T>
-__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) {
-    return __shfl_sync(mask, f, src);
-}
-template<>
-__device__ inline float2 packed_shfl_sync<float2>(uint32_t mask, const float2 &f, int src) {
-    float2 r;
-    r.x = __shfl_sync(mask, f.x, src);
-    r.y = __shfl_sync(mask, f.y, src);
-    return r;
-}
-
-/* ----------  SHARED MEMORY UTILS  ---------- */
-
-// namespace ducks {
-// namespace sb {
-// struct identifier {};
-// }
-// }
-
-// template<typename Args...>
-// struct sb {
-//     using identifier = ducks::sb::identifier;
-//     Args... args;
-// };
-
-// namespace ducks {
-// namespace sb {
-// template<typename T> concept all = requires {
-//     typename T::identifier;
-// } && std::is_same_v<T::identifier, identifier>;
-// }
-// }
-
-// Joyously stolen from https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cute/container/alignment.hpp#L51
-#if defined(__CUDACC__)
-#define KITTENS_ALIGN_AS(n) __align__(n)
-#else
-#define KITTENS_ALIGN_AS(n) alignas(n)
-#endif
-
-#ifdef KITTENS_HOPPER
-#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128)
-#else
-#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16)
-#endif
-
-/**
- * @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
- */
-struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; };
-/**
- * @brief Very simple allocator for dynamic shared memory. Advances pointer and tracks alignments.
- * @tparam default_alignment The default alignment this allocator will enforce. If <=0 (default -1) it will not align.
- */
-#ifdef KITTENS_HOPPER
-template<int default_alignment=1024> 
-#else
-template<int default_alignment=16> 
-#endif
-struct shared_allocator {
-    int *ptr;
-
-    private:
-        // Recursive template to generate N-dimensional array type
-        template<typename A, size_t... dims>
-        struct variadic_array;
-        template<typename A, size_t first_dim, size_t... rest_dims>
-        struct variadic_array<A, first_dim, rest_dims...> {
-            using type = typename variadic_array<A, rest_dims...>::type[first_dim];
-        };
-        template<typename A>
-        struct variadic_array<A> {
-            using type = A;
-        };
-        template<typename A, size_t... dims> 
-        using variadic_array_t = typename variadic_array<A, dims...>::type;
-
-        template<int alignment>
-        __device__ inline void align_ptr() {
-            if constexpr (alignment > 0) {
-                uint64_t p = reinterpret_cast<uint64_t>(ptr);
-                if(p % alignment != 0) {
-                    ptr = (int*)(p + (alignment-(p%alignment)));
-                }
-            }
-        }
-
-    public:
-        /**
-        * @brief Construct a new shared allocator using a pointer to extern shared memory.
-        * @param[in] _ptr Pointer to the start of the extern shared memory.
-        */
-        __device__ shared_allocator(int *_ptr): ptr(_ptr) {}
-        /**
-        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
-        * @tparam A The type of the object to allocate.
-        * @tparam dims... A list of dimensions for the N-dimensional array.
-        * @return Reference to the allocated object.
-        */
-        template<typename A, size_t... dims> 
-        __device__ inline variadic_array_t<A, dims...>& allocate() {
-            // static_assert(sizeof(A) % default_alignment == 0, "Type is not aligned properly for array allocation");
-            align_ptr<default_alignment>();
-            using at = variadic_array_t<A, dims...>;
-            at*p = reinterpret_cast<at*>(ptr);
-            ptr += sizeof(at)/sizeof(int);
-            return *p;
-        }
-        /**
-        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
-        * @tparam alignment An alignment to enforce for this particular object.
-        * @tparam A The type of the object to allocate.
-        * @tparam dims... A list of dimensions for the N-dimensional array.
-        * @return Reference to the allocated object.
-        */
-        template<int alignment, typename A, size_t... dims> 
-        __device__ inline variadic_array_t<A, dims...>& allocate() {
-            // static_assert(sizeof(A) % alignment == 0, "Type is not aligned properly for array allocation");
-            align_ptr<alignment>();
-            using at = variadic_array_t<A, dims...>;
-            at*p = reinterpret_cast<at*>(ptr);
-            ptr += sizeof(at)/sizeof(int);
-            return *p;
-        }
-};
-#if (defined(KITTENS_HOPPER) || defined(KITTENS_BLACKWELL))
-/**
- * @brief A wrapper for an allocator that enforces sufficient alignment to be used for TMA loads and stores.
- */
-using tma_allocator = shared_allocator<1024>;
-using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
-
-/* Get CTA ID within a cluster */
-__device__ static inline int3 clusterIdx() {
-    int3 cluster_idx;
-    asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
-    asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
-    asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
-    return cluster_idx;
-}
-__device__ static inline int cluster_ctarank() {
-    uint32_t ctarank;
-    asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
-    return ctarank;
-}
-#endif
-
-} // namespace kittens
@@ -1,12 +0,0 @@
-/**
- * @file
- * @brief The master header file of ThunderKittens. This file includes everything you need!
- */
-
-#pragma once
-
-#include "common/common.cuh"
-#include "types/types.cuh"
-#include "ops/ops.cuh"
-#include "pyutils/util.cuh"
-// #include "pyutils/pyutils.cuh" // for simple binding without including torch
@@ -1,51 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of all device (multi-GPU) operations defined by ThunderKittens
- */
-
-#pragma once
-
-#include "../../types/types.cuh"
-
-namespace kittens {
-
-template<int _NUM_DEVICES>
-struct device {
-
-static_assert(_NUM_DEVICES >= 0 && _NUM_DEVICES <= 72, "Invalid number of devices");
-static constexpr int NUM_DEVICES = _NUM_DEVICES;
-
-#ifdef KITTENS_HOPPER
-
-using barrier_t = pgl<gl<int, 1, 1, 1, -1>, NUM_DEVICES, true>;
-
-/**
- * @brief Multi-GPU synchronization barrier for coordinated kernel exit
- * 
- * Performs a synchronization across all devices to ensure all GPUs complete 
- * their work before any kernel exits. Does not synchronize intra-node threads
- * or threadblocks.
- * 
- * @param barrier Pre-allocated barrier structure, must be initialized to 0
- * @param dev_idx Current device index (0 to NUM_DEVICES - 1)
- * @param id Synchronization point identifier (default: 0). 0 is fine for most cases
- * 
- */
-__device__ static inline void sync_on_exit(const barrier_t &barrier, const int dev_idx, const int id = 0) {
-    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
-        threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
-        cuda::atomic_ref<int, cuda::thread_scope_system> barrier_uc(barrier[dev_idx][{id}]);
-
-        // Inter-note check-in
-        multimem<int>::red<reduce_op::ADD>(barrier.mc_ptr_at({id}), 1);
-        asm volatile ("{fence.proxy.alias;}" ::: "memory");
-        while (barrier_uc.load(cuda::memory_order_acquire) < NUM_DEVICES);
-        barrier_uc.fetch_sub(NUM_DEVICES, cuda::memory_order_release);
-    }
-}
-
-#endif
-
-};
-
-} // namespace kittens
@@ -1,96 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of all group (multi-warp) operations defined by ThunderKittens
- */
-
-#pragma once
-
-#include <cuda/pipeline>
-
-#include "../../common/common.cuh"
-#include "../../types/types.cuh"
-#include "../thread/thread.cuh" // several group memory ops rely on underlying warp-scope ops
-
-#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group.");
-// A "warpgroup" is a special group of 4 consecutive warps defined by NVIDIA for certain SM_90+ operations.
-#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group.");
-
-// WGMMA relies on some template structures that cannot be specialized within the group struct, so we declare them in advance.
-#ifdef KITTENS_HOPPER
-#include "mma/warpgroup/base/base.cuh"
-#endif
-
-namespace kittens {
-/*
-This is meant to be used with a `using group_N = kittens::group<NUM_WORKERS>;` at the start of every kernel.
-*/
-template<int _GROUP_WARPS>
-struct group {
-static constexpr int GROUP_WARPS = _GROUP_WARPS; // This alias produces nice parallelism.
-static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; // This alias produces nice parallelism.
-__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; }
-__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; }
-__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; }
-
-__device__ static inline void sync(int id) {
-    asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
-}
-template<uint32_t MASK=0xFFFFFFFF> __device__ static inline void sync() {
-    static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!");
-    asm volatile("bar.warp.sync %0;\n" :: "n"(MASK));
-}
-__device__ static inline void arrive(int id) {
-    asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
-}
-
-#include "memory/memory.cuh"
-#include "shared/shared.cuh"
-#include "register/register.cuh"
-#include "mma/mma.cuh"
-
-#ifdef KITTENS_HOPPER
-
-template<int n_reg> __device__ static inline void increase_registers() {
-    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
-    asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg));
-}
-template<int n_reg> __device__ static inline void decrease_registers() {
-    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
-    asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg));
-}
-__device__ static inline void producer_registers() { decrease_registers<24>(); }
-template<int NCWG> __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); }
-
-#endif
-
-};
-
-namespace everyone {
-
-// Block-level synchronization
-__device__ static inline void sync(int id) {
-    asm volatile("bar.sync %0;\n" :: "r"(id));
-}
-
-// Cluster-level synchronization functions
-namespace tma {
-namespace cluster {
-__device__ static inline void arrive_aligned() { // All threads in the cluster must call this
-    asm volatile ("barrier.cluster.arrive.release.aligned;\n");
-}
-__device__ static inline void wait_aligned() {
-    asm volatile ("barrier.cluster.wait.acquire.aligned;\n");
-}
-__device__ static inline void sync() {
-    arrive_aligned();
-    wait_aligned();
-}
-}
-}
-
-};
-
-using warp = group<1>;      // scope used by most pre-Hopper GPUs, and also for most register operations.
-using warpgroup = group<4>; // special scope commonly used by Hopper and later.
-
-}
@@ -1,21 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of colaborative group memory movement operations
- */
-
-#include "util/util.cuh"
-#include "tile/tile.cuh"
-#include "vec/vec.cuh"
-
-#ifdef KITTENS_HOPPER
-struct tma {
-#include "util/tma.cuh"
-#include "tile/tma.cuh"
-#include "vec/tma.cuh"
-struct cluster {
-#include "util/tma_cluster.cuh"
-#include "tile/tma_cluster.cuh"
-#include "vec/tma_cluster.cuh"
-};
-};
-#endif
@@ -1,42 +0,0 @@
-/**
- * @file
- * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
- */
-
-/**
- * @brief Collaboratively loads data from a source array into register tiles.
- *
- * @tparam RT The register tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
-    load<axis, CRT::component, CGL::component, COORD>(dst.real, src.real, idx);
-    load<axis, CRT::component, CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
-    load<2, CRT, CGL>(dst, src, idx);
-}
-
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory.
- *
- * @tparam RT The register tile type.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
-    store<axis, typename CRT::component, typename CGL::component>(dst.real, src.real, idx);
-    store<axis, typename CRT::component, typename CGL::component>(dst.imag, src.imag, idx);
-}
-template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
-__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
-    store<2, CRT, CGL>(dst, src, idx);
-}
@@ -1,37 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
- */
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
-    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
-    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
-    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
-    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-
-template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
-    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
-template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
-__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
-    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
-    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
-}
@@ -1,34 +0,0 @@
-/**
- * @file
- * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
- */
-
-/**
- * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source shared tile.
- */
-template<ducks::crt::all RT, ducks::cst::all ST>
-__device__ inline static void load(RT &dst, const ST &src) {
-    load(dst.real, src.real);
-    load(dst.imag, src.imag);
-}
-
-
-/**
- * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination shared tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::cst::all ST, ducks::crt::all RT>
-__device__ inline static void store(ST &dst, const RT &src) {
-    store(dst.real, src.real);
-    store(dst.imag, src.imag);
-}
-
@@ -1,207 +0,0 @@
-/**
- * @file
- * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
- */
-
-/**
- * @brief Collaboratively loads data from a source array into row-major layout tiles.
- *
- * @tparam RT The row-major layout tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    using T2 = RT::dtype;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-
-    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = src.template stride<axis>();
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = dst.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        int row = row_offset + i*dst.tile_size_row + (warp_laneid / 4);
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-            dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)]));
-            dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+8)]));
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-            dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)]));
-            dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+8)]));
-        }
-    }
-}
-/**
- * @brief Collaboratively loads data from a source array into column-major layout tiles.
- *
- * @tparam RT The column-major layout tile type.
- * @tparam U The data type of the source array.
- * @param dst[out] The destination tile to load data into.
- * @param src[in] The source array to load data from.
- * @param row_stride[in] The stride in elements between rows in the source array.
- */
-template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    using T = typename RT::T;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T, fp8e4m3> && !std::is_same_v<T, fp8e5m2>, "Unsupported type for load/store");
-    #endif
-
-    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = src.template stride<axis>();
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = dst.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        int row = row_offset + i*dst.tile_size_row + 2*(warp_laneid % 4);
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[0].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[1].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[0].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[1].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[2].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[3].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+8)]);
-        }
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            int col = j*dst.tile_size_col + (warp_laneid / 4);
-            dst.tiles[i][j].data[2].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+0)]);
-            dst.tiles[i][j].data[3].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+8)]);
-        }
-    }
-}
-template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
-    load<2>(dst, src, idx);
-}
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
- *
- * @tparam RT The register tile type with a row-major layout.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    using T2 = RT::dtype;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-
-    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = dst.template stride<axis>();
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = src.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < src.height; i++) {
-        int row = row_offset + i*src.tile_size_row + (warp_laneid / 4);
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-        }
-    }
-}
-/**
- * @brief Collaboratively stores data from register tiles to a destination array in global memory with a column-major layout.
- *
- * @tparam RT The register tile type with a column-major layout.
- * @tparam U The data type of the destination array.
- * @param[out] dst The destination array in global memory to store data into.
- * @param[in] src The source register tile to store data from.
- * @param row_stride[in] The stride in elements between rows in the destination array.
- */
-template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    using T = base_types::packing<typename RT::dtype>::unpacked_type;
-    using U = typename GL::dtype;
-
-    #ifdef KITTENS_HOPPER
-    static_assert(!std::is_same_v<T, fp8e4m3_4> && !std::is_same_v<T, fp8e5m2_4>, "Unsupported type for load/store");
-    #endif
-    
-    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
-    const int row_stride = dst.template stride<axis>();
-    int warp_laneid = threadIdx.x % WARP_THREADS;
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    const int row_offset = src.rows*local_warpid;
-    #pragma unroll
-    for(int i = 0; i < src.height; i++) {
-        int row = row_offset + i*src.tile_size_row + 2*(warp_laneid % 4);
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].x);
-            dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].x);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].y);
-            dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].y);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].x);
-            dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].x);
-        }
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            int col = j*src.tile_size_col + (warp_laneid / 4);
-            dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].y);
-            dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].y);
-        }
-    }
-}
-template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
-__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
-    store<2>(dst, src, idx);
-}
@@ -1,168 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
- */
- 
-
-/**
- * @brief Loads data from global memory into a shared memory tile.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination shared memory tile.
- * @param[in] src The source global memory array.
- * @param[in] idx The coordinate of the tile in the global memory array.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = src.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
-    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-    constexpr int total_rows = dst.height*dst.width;
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
-    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % dst.cols;
-
-        if constexpr (assume_aligned) {
-            float4 tmp;
-            move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
-            move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
-                float4 tmp;
-                move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
-            }
-            else {
-                float4 zeros = {0.f,0.f,0.f,0.f};
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
-            }
-        }
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
-    load<2, false, ST, GL, COORD>(dst, src, idx);
-}
-
-/**
- * @brief Stores data from a shared memory tile into global memory.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination global memory array.
- * @param[in] src The source shared memory tile.
- * @param row_stride[in] The stride between rows in the destination array.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = dst.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = src.cols / elem_per_memcpy;
-    constexpr int total_calls = (src.height*src.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord];
-    uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % src.cols;
-
-        if constexpr (assume_aligned) {
-            float4 tmp;
-            move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
-            move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < dst.template shape<axis>()) {
-                float4 tmp;
-                move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
-                move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
-            }
-        }
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
-    store<2, false, ST, GL, COORD>(dst, src, idx);
-}
-
-/**
- * @brief Asynchronously loads data from global memory into a shared memory tile.
- *
- * @tparam ST The type of the shared tile.
- * @param[out] dst The destination shared memory tile.
- * @param[in] src The source global memory array.
- *
- * @note This function expects 16-byte alignments. Otherwise, behavior is undefined.
- */
-template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
-    using T = typename ST::dtype;
-    const int row_stride = src.template stride<axis>();
-    // we can handle this many rows each time we run a memcpy_async
-    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
-    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
-    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
-
-    coord<> unit_coord = idx.template unit_coord<axis, 3>();
-    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
-    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-    int laneid = threadIdx.x % GROUP_THREADS;
-
-    #pragma unroll
-    for(int i = 0; i < total_calls; i++) {
-
-        int load_idx = i * GROUP_THREADS + laneid;
-        
-        int row = load_idx / memcpy_per_row;
-        int col = (load_idx*elem_per_memcpy) % dst.cols;
-
-        if constexpr (assume_aligned) {
-            asm volatile(
-                "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
-                :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
-                : "memory"
-            );
-        }
-        else {
-            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
-                asm volatile(
-                    "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
-                    :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
-                    : "memory"
-                );
-            }
-            else {
-                // printf("thread %d skipping async load on row %d, col %d\n", threadIdx.x, row + unit_coord.template dim<axis>(), col);
-                float4 zeros = {0.f,0.f,0.f,0.f};
-                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
-            }
-        }
-    }
-    asm volatile("cp.async.commit_group;\n" ::: "memory");
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
-    load_async<2, false, ST, GL, COORD>(dst, src, idx);
-}
@@ -1,323 +0,0 @@
-/**
- * @file
- * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
- */
-
-/**
- * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source shared tile.
- */
-template<ducks::rt::all RT, ducks::st::all ST>
-__device__ inline static void load(RT &dst, const ST &src) {
-    constexpr int height = ST::height;
-    constexpr int warp_height = RT::height;
-    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
-    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
-    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    using T2 = RT::dtype;
-    using U  = ST::dtype;
-    using T  = base_types::packing<T2>::unpacked_type;
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-
-    #pragma unroll
-    for(int i = 0; i < dst.height; i++) {
-        #pragma unroll
-        for(int j = 0; j < dst.width; j++) {
-            if constexpr (sizeof(typename ST::dtype) == 2) {
-                // handle the row-major layout for 16-bit types
-                U2 tmp[4];
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid % 16);
-                int col = j*dst.tile_size_col + (warp_laneid / 16) * 8;
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                else {
-                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) {
-                // handle the row-major layout for 8-bit types
-                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
-                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
-                int col = j*dst.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
-
-                U2 tmp[4];
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                else {
-                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
-                }
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
-                // handle the row-major layout for 32-bit types
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid / 4);
-                int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
-                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
-                    row += src.row_offset;
-                    col += src.col_offset;
-                }
-                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
-                U2 tmp[4];
-                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
-                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
-                const int outer_idx = col/subtile_cols;
-                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
-                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
-                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
-                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
-                move<U>::lds(tmp[0].x, (addr_1+ 0)^swizzle_1);
-                move<U>::lds(tmp[0].y, (addr_1+ 4)^swizzle_1);
-                move<U>::lds(tmp[2].x, (addr_1+32)^swizzle_1);
-                move<U>::lds(tmp[2].y, (addr_1+36)^swizzle_1);
-                move<U>::lds(tmp[1].x, (addr_2+ 0)^swizzle_2);
-                move<U>::lds(tmp[1].y, (addr_2+ 4)^swizzle_2);
-                move<U>::lds(tmp[3].x, (addr_2+32)^swizzle_2);
-                move<U>::lds(tmp[3].y, (addr_2+36)^swizzle_2);
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-                if(blit) {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x};
-                    }
-                }
-            }
-            else {
-                // handle the column-major layout
-                int row = (local_warpid*warp_height + i)*dst.tile_size_row + 2*(warp_laneid % 4);
-                int col = j*dst.tile_size_col + (warp_laneid / 4);
-                U2 tmp[4];
-                move<U>::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0}));
-                move<U>::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0}));
-                move<U>::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8}));
-                move<U>::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8}));
-                move<U>::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0}));
-                move<U>::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0}));
-                move<U>::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8}));
-                move<U>::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8}));
-                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
-                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
-                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
-                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
-            }
-        }
-    }
-}
-
-
-/**
- * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
- *
- * @tparam RT The register tile type
- * @tparam ST The shared tile type
- * @param dst[out] The destination shared tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::st::all ST, ducks::rt::all RT>
-__device__ inline static void store(ST &dst, const RT &src) {
-    constexpr int height = ST::height;
-    constexpr int warp_height = RT::height;
-    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
-    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
-    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
-    int local_warpid;
-    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
-    else local_warpid = warpid();
-    using T2 = RT::dtype;
-    using U  = ST::dtype;
-    using T  = base_types::packing<T2>::unpacked_type;
-    using U2 = base_types::packing<U>::packed_type;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-
-    #pragma unroll
-    for(int i = 0; i < warp_height; i++) {
-        #pragma unroll
-        for(int j = 0; j < src.width; j++) {
-            if constexpr (sizeof(typename ST::dtype) == 2) {
-                // handle the row-major layout
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-#ifdef KITTENS_HOPPER
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid % 16);
-                int col = j*src.tile_size_col + (warp_laneid / 16) * 8;
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
-                }
-                else {
-                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
-                }
-#else
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
-                    int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]);
-                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]);
-                }
-                else {
-                    int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
-                    int col = j*src.tile_size_col + (warp_laneid / 4);
-                    move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
-                    move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
-                    move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
-                }
-#endif
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) { 
-                // handle the row-major layout for 8-bit types
-
-                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
-                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
-                int col = j*src.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
-
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
-                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
-                }
-                else {
-                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
-                }
-            }
-            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
-                // handle the row-major layout for 32-bit types
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
-                int col = j*src.tile_size_col + 2*(warp_laneid % 4);
-                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
-                    row += dst.row_offset;
-                    col += dst.col_offset;
-                }
-                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
-                T2 reg_tmp[4];
-                if(blit) {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x};
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int k = 0; k < 4; k++) {
-                        reg_tmp[k] = src.tiles[i][j].data[k];
-                    }
-                }
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(reg_tmp[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(reg_tmp[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(reg_tmp[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(reg_tmp[3]);
-                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
-                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
-                const int outer_idx = col/subtile_cols;
-                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
-                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
-                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
-                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
-                move<U>::sts((addr_1+ 0)^swizzle_1, tmp[0].x);
-                move<U>::sts((addr_1+ 4)^swizzle_1, tmp[0].y);
-                move<U>::sts((addr_1+32)^swizzle_1, tmp[2].x);
-                move<U>::sts((addr_1+36)^swizzle_1, tmp[2].y);
-                move<U>::sts((addr_2+ 0)^swizzle_2, tmp[1].x);
-                move<U>::sts((addr_2+ 4)^swizzle_2, tmp[1].y);
-                move<U>::sts((addr_2+32)^swizzle_2, tmp[3].x);
-                move<U>::sts((addr_2+36)^swizzle_2, tmp[3].y);
-            }
-            else {
-                // handle the column-major layout
-                int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
-                int col = j*src.tile_size_col + (warp_laneid / 4);
-                U2 tmp[4];
-                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
-                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
-                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
-                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
-                move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
-                move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
-                move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
-                move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
-                move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
-                move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
-                move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
-                move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
-            }
-        }
-    }
-}
-
-// Load and store of vectors from/to shared tiles.
-
-template<ducks::rv::naive_layout RV, ducks::st::all ST>
-__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) {
-    KITTENS_CHECK_WARP;
-    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
-    using T = RV::T;
-    using U = ST::T;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
-
-    #pragma unroll
-    for(int col = warp_laneid; col < dst.length; col+=WARP_THREADS) {
-        U tmp;
-        move<U>::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col}));
-        dst.data[col/WARP_THREADS][0] = base_types::convertor<T, U>::convert(tmp);
-    }
-}
-
-template<ducks::rv::naive_layout RV, ducks::st::all ST>
-__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) {
-    KITTENS_CHECK_WARP;
-    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
-    using T = RV::T;
-    using U = ST::T;
-    int warp_laneid = ::kittens::laneid();
-
-    // convert to shared state space
-    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
-
-    #pragma unroll
-    for(int col = warp_laneid; col < src.length; col+=WARP_THREADS) {
-        U tmp = base_types::convertor<U, T>::convert(src.data[col/WARP_THREADS][0]);
-        move<U>::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp);
-    }
-}
@@ -1,325 +0,0 @@
-/**
- * @file
- * @brief Group (collaborative warp) ops for loading tensor tiles into register tiles.
- */
-
-/**
- * @brief Load data from a tensor tile into a register tile.
- *
- * @tparam RT The register tile type
- * @tparam TM The tensor memory tile type
- * @param dst[out] The destination register tile.
- * @param src[in]  The source tensor tile.
- */
-template<ducks::rt::row_layout RT, ducks::tt::all TM>
-__device__ inline static void load_async(RT &dst, const TM &src) {
-    if constexpr (GROUP_WARPS == 1) {
-        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
-        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
-
-        using T2 = RT::dtype;
-        using U  = typename TM::dtype;
-        using U2 = base_types::packing<typename TM::dtype>::packed_type;
-
-        if constexpr (sizeof(typename TM::dtype) == 1) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                #pragma unroll
-                for(int j = 0; j < dst.width; j++) {
-                    asm volatile(
-                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
-                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
-                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                    );
-                }
-            }
-        } else if constexpr (sizeof(typename TM::dtype) == 2) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                #pragma unroll
-                for(int j = 0; j < dst.width; j++) {
-                    asm volatile(
-                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
-                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
-                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
-                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col))
-                    );
-                }
-            }
-        }
-        else if constexpr (sizeof(typename TM::dtype) == 4) {
-            #pragma unroll
-            for(int i = 0; i < dst.height; i++) {
-                if constexpr (dst.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j+=4) {
-                        U2 data[16];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y),
-                            "=f"(data[4].x), "=f"(data[4].y),
-                            "=f"(data[5].x), "=f"(data[5].y),
-                            "=f"(data[6].x), "=f"(data[6].y),
-                            "=f"(data[7].x), "=f"(data[7].y),
-                            "=f"(data[8].x), "=f"(data[8].y),
-                            "=f"(data[9].x), "=f"(data[9].y),
-                            "=f"(data[10].x), "=f"(data[10].y),
-                            "=f"(data[11].x), "=f"(data[11].y),
-                            "=f"(data[12].x), "=f"(data[12].y),
-                            "=f"(data[13].x), "=f"(data[13].y),
-                            "=f"(data[14].x), "=f"(data[14].y),
-                            "=f"(data[15].x), "=f"(data[15].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
-                            dst.tiles[i][j+2].data[k] = base_types::convertor<T2, U2>::convert(data[k+8]);
-                            dst.tiles[i][j+3].data[k] = base_types::convertor<T2, U2>::convert(data[k+12]);
-                        }
-                    }
-                }
-                else if constexpr (dst.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j+=2) {
-                        U2 data[8];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x4.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y),
-                            "=f"(data[4].x), "=f"(data[4].y),
-                            "=f"(data[5].x), "=f"(data[5].y),
-                            "=f"(data[6].x), "=f"(data[6].y),
-                            "=f"(data[7].x), "=f"(data[7].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
-                        }
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < dst.width; j++) {
-                        U2 data[4];
-                        asm volatile(
-                            "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
-                            : "=f"(data[0].x), "=f"(data[0].y),
-                            "=f"(data[1].x), "=f"(data[1].y),
-                            "=f"(data[2].x), "=f"(data[2].y),
-                            "=f"(data[3].x), "=f"(data[3].y)
-                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
-                        );
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            dst.tiles[i][j].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
-                        }
-                    }
-                }
-            }
-        }
-    }
-    else {
-        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
-        constexpr int warp_rows = TM::rows/GROUP_WARPS;
-        static_assert(TM::cols==RT::cols);
-        static_assert(warp_rows==RT::rows);
-        if constexpr (GROUP_WARPS == 4) {
-            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
-            ::kittens::group<1>::load_async(dst, src_subtile);
-        }
-        else {
-            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
-            ::kittens::group<1>::load_async(dst, src_subtile);
-        }
-    }
-}
-
-
-/**
- * @brief Store data into a tensor tile from a register tile.
- *
- * @tparam RT The register tile type
- * @tparam TM The tensor memory tile type
- * @param dst[out] The destination tensor tile.
- * @param src[in]  The source register tile.
- */
-template<ducks::rt::all RT, ducks::tt::all TM>
-__device__ inline static void store_async(TM &dst, const RT &src) {
-    if constexpr (GROUP_WARPS == 1) {
-        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
-        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
-
-        using T2 = RT::dtype;
-        using T = base_types::packing<T2>::unpacked_type;
-        using U = TM::dtype;
-        using U2 = base_types::packing<U>::packed_type;
-
-        if constexpr (sizeof(typename TM::dtype) == 2) {
-            #pragma unroll
-            for(int i = 0; i < src.height; i++) {
-                if constexpr (src.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=4) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[3])
-                        );
-                    }
-                }
-                else if constexpr (src.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=2) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3])
-                        );
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j++) {
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[0]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[1]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[2]),
-                            "r"(*(uint32_t*)&src.tiles[i][j].data[3])
-                        );
-                    }
-                }
-            }
-        }
-        else if constexpr (sizeof(typename TM::dtype) == 4) {
-            #pragma unroll
-            for(int i = 0; i < src.height; i++) {
-                if constexpr(src.width%4 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=4) {
-                        U2 data[16];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
-                            data[k+8] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+2].data[k]);
-                            data[k+12] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+3].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y),
-                            "f"(data[4].x), "f"(data[4].y),
-                            "f"(data[5].x), "f"(data[5].y),
-                            "f"(data[6].x), "f"(data[6].y),
-                            "f"(data[7].x), "f"(data[7].y),
-                            "f"(data[8].x), "f"(data[8].y),
-                            "f"(data[9].x), "f"(data[9].y),
-                            "f"(data[10].x), "f"(data[10].y),
-                            "f"(data[11].x), "f"(data[11].y),
-                            "f"(data[12].x), "f"(data[12].y),
-                            "f"(data[13].x), "f"(data[13].y),
-                            "f"(data[14].x), "f"(data[14].y),
-                            "f"(data[15].x), "f"(data[15].y)
-                        );
-                    }
-                }
-                else if constexpr(src.width%2 == 0) {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j+=2) {
-                        U2 data[8];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y),
-                            "f"(data[4].x), "f"(data[4].y),
-                            "f"(data[5].x), "f"(data[5].y),
-                            "f"(data[6].x), "f"(data[6].y),
-                            "f"(data[7].x), "f"(data[7].y)
-                        );
-                    }
-                }
-                else {
-                    #pragma unroll
-                    for(int j = 0; j < src.width; j++) {
-                        U2 data[4];
-                        #pragma unroll
-                        for(int k = 0; k < 4; k++) {
-                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
-                        }
-                        asm volatile(
-                            "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
-                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
-                            "f"(data[0].x), "f"(data[0].y),
-                            "f"(data[1].x), "f"(data[1].y),
-                            "f"(data[2].x), "f"(data[2].y),
-                            "f"(data[3].x), "f"(data[3].y)
-                        );
-                    }
-                }
-            }
-        }
-    }
-    else {
-        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
-        constexpr int warp_rows = TM::rows/GROUP_WARPS;
-        static_assert(TM::cols==RT::cols);
-        static_assert(warp_rows==RT::rows);
-        if constexpr (GROUP_WARPS == 4) {
-            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
-            ::kittens::group<1>::store_async(dst_subtile, src);
-        }
-        else {
-            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
-            ::kittens::group<1>::store_async(dst_subtile, src);
-        }
-    }
-}
@@ -1,16 +0,0 @@
-/**
- * @file
- * @brief An aggregate header of group memory operations on tiles.
- */
-
-#include "shared_to_register.cuh"
-#include "global_to_register.cuh"
-#include "global_to_shared.cuh"
-#ifdef KITTENS_BLACKWELL
-#include "tensor_to_register.cuh"
-#endif
-
-#include "complex/complex_shared_to_register.cuh"
-#include "complex/complex_global_to_register.cuh"
-#include "complex/complex_global_to_shared.cuh"
-
@@ -1,134 +0,0 @@
-/**
- * @file
- * @brief Functions for a group scope to call tile TMA functions.
- */
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::prefetch<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::prefetch<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
-    if(laneid() == 0) {
-        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
-    }
-}
-
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
-    if(laneid() == 0) {
-        ::kittens::tma::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar); // Don't do the mask
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
-    if(laneid() == 0) {
-        ::kittens::tma::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar);
-    }
-}
@@ -1,33 +0,0 @@
-/**
- * @file
- * @brief Functions for a group scope to call tile TMA cluster functions.
- */
-
-
-#ifdef KITTENS_BLACKWELL
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
-    }
-}
-#else
-template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
-    }
-}
-template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
-__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
-    if(laneid() == 0) {
-        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
-    }
-}
-#endif
--- a/Show More
+++ b/Show More