ln

cleanups
better uop matmul
2026-03-12 08:23:06 +08:00 · 2026-03-12 08:18:31 +08:00 · 2026-03-12 08:13:09 +08:00 · 2026-03-11 21:50:31 +08:00 · 2026-03-11 15:22:46 +08:00
201 changed files with 6496 additions and 6216 deletions
@@ -66,14 +66,14 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
    - name: Cache Python packages
      if: github.event_name != 'pull_request'
      id: restore-venv
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-${{ runner.arch }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}
+        key: venv-${{ runner.os }}-python-${{ steps.setup-python.outputs.python-version }}-${{ inputs.deps }}-${{ inputs.pydeps }}-${{ env.CACHE_VERSION }}

    # **** Caching downloads ****

@@ -199,13 +199,13 @@ runs:
      uses: actions/cache/restore@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
@@ -287,7 +287,6 @@ runs:

        CMAKE_ARGS="-Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5"
        if [[ "${{ runner.os }}" == "macOS" ]]; then
-          sudo xcode-select -s /Applications/Xcode_16.2.app/Contents/Developer
          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
        fi

@@ -56,7 +56,6 @@ jobs:
        python3 -c "from tinygrad.runtime.autogen import libusb"
        python3 -c "from tinygrad.runtime.autogen import mesa"
        python3 -c "from tinygrad.runtime.autogen import avcodec"
-        python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
    - name: Check for differences
      run: |
@@ -49,6 +49,8 @@ jobs:
        source /tmp/tinygrad_pytest_ci/bin/activate
        pytest -nauto --durations=20
    - name: openpilot compile3 0.10.1 driving_vision
+      run: FLOAT16=1 CL=1 IMAGE=2 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: FLOAT16=1 CL=1 IMAGE=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx

  testmacbenchmark:
@@ -185,13 +187,13 @@ jobs:
        PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
        PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
    - name: UsbGPU boot time
-      run: sudo -E PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
+      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
-      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
-      run: sudo -E PYTHONPATH=. GMMU=0 AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
+      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=. GMMU=0 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 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - name: UsbGPU (USB4/TB) boot time
      run: PYTHONPATH=. DEBUG=3 NV=1 NV_IFACE=PCI NV_NAK=1 time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU (USB4/TB) tiny tests
@@ -588,22 +590,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: openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: IR3 openpilot compile3 0.11.0 driving_vision
-      run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM QCOM_IR3=1 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.11.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.11.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.10.0 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 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=11 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: DEBUG=2 openpilot compile3 0.10.1 driving_vision
+      run: PYTHONPATH="." DEBUG=2 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: DEBUG=2 IMAGE=1 openpilot compile3 0.10.1 driving_vision
      run: PYTHONPATH="." DEBUG=2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+    - name: IMAGE=1 openpilot compile3 0.10.1 driving_vision
+      run: BENCHMARK_LOG=image_1_openpilot_0_10_1_vision PYTHONPATH="." DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 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=3 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3 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
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 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
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
@@ -632,9 +634,9 @@ jobs:
        echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: openpilot compile3 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." GMMU=0 DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision PYTHONPATH="." DEV=AMD AMD_LLVM=1 AMD_IFACE=USB ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot load_pickle 0.10.1 driving_vision
-      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." GMMU=0 DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py
+      run: BENCHMARK_LOG=usbgpu_openpilot_0_10_1_vision_load_pickle PYTHONPATH="." DEV=AMD AMD_IFACE=USB ASSERT_MIN_LOAD_TIME=15 python3 examples/openpilot/load_pickle.py

  testreddriverbenchmark:
    name: AM Benchmark
@@ -697,14 +699,6 @@ jobs:
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps AMD=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
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6482 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_IFACE=PCI python3 test/test_tiny.py
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6482 AM_RESET=1 AMD=1 AMD_AQL=1 AMD_IFACE=PCI python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
    - 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

@@ -761,12 +755,5 @@ jobs:
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
      run: BENCHMARK_LOG=bert_10steps NV=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
-    - name: Remote
-      run: |
-        pkill -f 'extra/remote/serve.py' || true
-        PYTHONPATH=. python3 extra/remote/serve.py 6483 &
-        sleep 1
-        DEBUG=2 PYTHONPATH=. REMOTE=127.0.0.1:6483 NV=1 python3 test/test_tiny.py
-        pkill -f 'extra/remote/serve.py' || true
    - 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
@@ -70,7 +70,7 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        BS=2 STEPS=10 MAX_BUFFER_SIZE=0 python beautiful_mnist.py
+        BS=2 STEPS=10 python beautiful_mnist.py
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
@@ -141,7 +141,7 @@ jobs:
        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 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      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

@@ -167,10 +167,10 @@ jobs:
      run: PYTHON=1 python3 test/backend/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
      run: PYTHON=1 python3 -m pytest -rA test/backend/test_renderer_failures.py::TestRendererFailures
-    - name: Test IMAGE support
+    - name: Test IMAGE=2 support
      run: |
-        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
-        IMAGE=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_simple_conv2d
    - name: Test emulated METAL tensor cores
      run: |
        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/backend/test_ops.py TestOps.test_big_gemm
@@ -369,11 +369,11 @@ jobs:
          key: gpu-image
          deps: testing_unit
          opencl: 'true'
-      - name: Test CL IMAGE=1 ops
+      - name: Test CL IMAGE=2 ops
        run: |
-          CL=1 IMAGE=1 python -m pytest -n=auto test/backend/test_ops.py --durations=20
+          CL=1 IMAGE=2 python -m pytest -n=auto test/backend/test_ops.py --durations=20
        # TODO: training is broken
-        # CL=1 IMAGE=1 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
+        # CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -418,13 +418,13 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1486 ALLOWED_GATED_READ_IMAGE=17 FLOAT16=1 CL=1 IMAGE=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+          ALLOWED_KERNEL_COUNT=123 ALLOWED_READ_IMAGE=1397 ALLOWED_GATED_READ_IMAGE=94 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=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        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: CL=1 IMAGE=1 SELFTEST=1 python examples/openpilot/compile3.py https://github.com/haraschax/filedump/raw/refs/heads/master/driving_vision_fp32.onnx
+        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: IMAGE=1 FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+        run: FLOAT16=1 CPU=1 CPU_LLVM=1 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -505,13 +505,8 @@ jobs:
        uses: ./.github/actions/setup-tinygrad
        with:
          key: apps_llm
-      - name: Test 1B LLM (llama)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (llama q4)
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster
-      - name: Test 1B LLM (qwen)
-        # NOTE: qwen is dumb and only knows about female chickens
-        run: echo "What's a female chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen
+      - name: Test 1B LLM
+        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster

 # ****** Models Tests ******

@@ -704,7 +699,7 @@ jobs:
      - name: Run test_tiny on MOCKAM
        run: python test/test_tiny.py
      - name: Run test_tiny on MOCKAM USB
-        run: GMMU=0 AMD_IFACE=USB python test/test_tiny.py
+        run: AMD_IFACE=USB python test/test_tiny.py
      - name: Run test_hcq on MOCKAM
        run: python -m pytest test/device/test_hcq.py

@@ -713,7 +708,8 @@ jobs:
      fail-fast: false
      matrix:
        backend: [amd, amdllvm]
-        arch: [rdna3, rdna4, cdna4]
+        arch: [rdna3, rdna4]
+        #arch: [rdna3, rdna4, cdna4]

    name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
    runs-on: ubuntu-22.04
@@ -739,7 +735,7 @@ jobs:
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['AMD'], Device.DEFAULT"
          DEBUG=5 FORWARD_ONLY=1 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
+        run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/testextra/test_cfg_viz.py test/external/external_test_am.py --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/backend/test_ops.py::TestOps::test_sin test/backend/test_ops.py::TestOps::test_cos test/backend/test_ops.py::TestOps::test_tan test/backend/test_ops.py::TestOps::test_exp test/backend/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
@@ -1,6 +1,6 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
 from typing import Callable
-from tinygrad import Tensor, TinyJit, nn, GlobalCounters, function
+from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist

@@ -15,31 +15,30 @@ class Model:
      nn.BatchNorm(64), Tensor.max_pool2d,
      lambda x: x.flatten(1), nn.Linear(576, 10)]

-  @function
  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)

-  @TinyJit
-  @Tensor.train()
-  def train_step(self, X_train:Tensor, Y_train:Tensor) -> Tensor:
-    opt.zero_grad()
-    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
-    loss = self(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
-
-  @TinyJit
-  def get_test_acc(self, X_test:Tensor, Y_test:Tensor) -> Tensor: return (self(X_test).argmax(axis=1) == Y_test).mean()*100
-
 if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))

  model = Model()
  opt = (nn.optim.Muon if getenv("MUON") else nn.optim.SGD if getenv("SGD") else nn.optim.Adam)(nn.state.get_parameters(model))

+  @TinyJit
+  @Tensor.train()
+  def train_step() -> Tensor:
+    opt.zero_grad()
+    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
+    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
+    return loss.realize(*opt.schedule_step())
+
+  @TinyJit
+  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
+
  test_acc = float('nan')
  for i in (t:=trange(getenv("STEPS", 70))):
    GlobalCounters.reset()   # NOTE: this makes it nice for DEBUG=2 timing
-    loss = model.train_step(X_train, Y_train)
-    if i%10 == 9: test_acc = model.get_test_acc(X_test, Y_test).item()
+    loss = train_step()
+    if i%10 == 9: test_acc = get_test_acc().item()
    t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")

  # verify eval acc
@@ -5,7 +5,7 @@ from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
  run_onnx = OnnxRunner(onnx_file)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
  return run_onnx_jit, run_onnx.graph_inputs

 if __name__ == "__main__":
@@ -13,8 +13,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 # TODO: fix benchmark logging and use tinygrad tqdm
 from tqdm import tqdm

-from tinygrad.uop.ops import UOp
-
 def train_resnet():
  from extra.models import resnet
  from examples.mlperf.dataloader import batch_load_resnet
@@ -1284,7 +1282,7 @@ def train_bert():
        previous_step = i

 def train_llama3():
-  from examples.mlperf.models.flat_llama import FlatTransformer, apply_grad
+  from extra.models.llama import Transformer
  from examples.llama3 import MODEL_PARAMS
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
  from examples.mlperf.optim import GradAccClipAdamW
@@ -1345,8 +1343,7 @@ def train_llama3():
  if (MP := getenv("MP", 1)) > 1: model_params['vocab_size'] = round_up(model_params['vocab_size'], 256 * MP)
  vocab_mask:Tensor = Tensor.arange(model_params['vocab_size']).reshape(1, 1, -1) >= real_vocab_size

-  model = FlatTransformer(**model_params, max_context=SEQLEN)
-
+  model = Transformer(**model_params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  params = get_parameters(model)
  # weights are all bfloat16 for now
  assert params and all(p.dtype == dtypes.bfloat16 for p in params)
@@ -1355,25 +1352,45 @@ def train_llama3():
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape))

-  is_dp = (DP := getenv("DP", 1)) > 1
-  is_mp = (MP := getenv("MP", 1)) > 1
-  is_sharding = is_dp or is_mp
-  device_count = max(DP, MP)
-  device = tuple(f"{Device.DEFAULT}:{i}" for i in range(device_count))
+  if (DP := getenv("DP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    for v in get_parameters(model):
+      v.shard_(device, axis=None)

-  model.shard(device, is_mp)
+    vocab_mask.shard_(device, axis=None)

-  if is_dp: vocab_mask.shard_(device, axis=None).realize()
-  if is_mp: vocab_mask.shard_(device, axis=2).realize()
+  if (MP := getenv("MP", 1)) > 1:
+    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    for k,v in get_state_dict(model).items():
+      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+      elif '.attention.wq' in k: v.shard_(device, axis=0)
+      elif '.attention.wk' in k: v.shard_(device, axis=0)
+      elif '.attention.wv' in k: v.shard_(device, axis=0)
+      elif '.attention.wqkv' in k: v.shard_(device, axis=0)
+      elif '.attention.wo' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+      elif 'output.weight' in k: v.shard_(device, axis=0)
+      else:
+        # attention_norm, ffn_norm, norm
+        v.shard_(device, axis=None)
+      # prevents memory spike on device 0
+      v.realize()

-  is_offload_optim = bool(getenv("OFFLOAD_OPTIM"))
-  is_fake_offload = Device.DEFAULT == "NULL"
-  optim_device = ("CPU" if not is_fake_offload else "NULL:99") if is_offload_optim else None
-  optim = GradAccClipAdamW(params, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
+    vocab_mask.shard_(device, axis=2).realize()
+
+  optim_device = "CPU" if getenv("OFFLOAD_OPTIM") else None
+  optim = GradAccClipAdamW(get_parameters(model), lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2,
                           eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device)

  # init grads
-  grads = [Tensor.zeros_like(p).contiguous() for p in optim.params]
+  for p in optim.params:
+    p.grad = p.empty_like().realize()
+  grads: list[Tensor] = [p.grad for p in optim.params]
+  for p in optim.params:
+    p.grad.assign(p.grad.zeros_like()).realize()

  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

@@ -1388,24 +1405,28 @@ def train_llama3():

  @TinyJit
  def minibatch(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if is_mp: tokens = tokens.shard(device)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.to(None).shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
-
-    for i,(t,g) in enumerate(zip(optim.params, loss.gradient(*optim.params))):
-      grads[i].replace(Tensor(grads[i].uop.after(UOp.group(*apply_grad(grads[i].uop, g.uop))), device=t.device))
-
+    loss.backward()
+    assert all(p.grad is g for p,g in zip(optim.params, grads))
    loss_cpu = loss.flatten().float().to("CPU")
-    return loss_cpu.realize(*grads)
+    Tensor.realize(loss_cpu, *grads)
+    return loss_cpu

  @TinyJit
  def optim_step():
    grad_norm = optim.fstep(grads)
    scheduler.step()

-    for g in grads: g.assign(g.zeros_like())
+    for g in grads:
+      g.assign(g.zeros_like()).realize()

    lr_cpu = optim.lr.float().to("CPU")
    grad_norm_cpu = grad_norm.float().to("CPU")
@@ -1416,10 +1437,14 @@ def train_llama3():
  @TinyJit
  @Tensor.train(False)
  def eval_step(tokens:Tensor):
-    if is_dp: tokens = tokens.to(None).shard(device, 0)
-    if is_mp: tokens = tokens.shard(device)
-    if not is_sharding: tokens = tokens.to(None)
-    logits:Tensor = model(tokens[:, :-1])
+    if (DP := getenv("DP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+      tokens = tokens.to(None).shard(device, 0)
+    if (MP := getenv("MP", 1)) > 1:
+      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+      tokens = tokens.shard(device)
+    if DP == 1 and MP == 1: tokens = tokens.to(None)
+    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = vocab_mask.where(-1e9, logits).sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float().to("CPU")

@@ -1,193 +0,0 @@
-import math, os
-if __name__ == "__main__":
-  os.environ["DEFAULT_FLOAT"] = "bfloat16"
-  os.environ["OPTIM_DTYPE"] = "bfloat16"
-  if "DEV" not in os.environ: os.environ["DEV"] = "NULL"
-  # CDNA
-  os.environ["EMULATE"] = "AMD_CDNA4"
-  os.environ["DEVICE_IN_FUNCTION_BUG"] = "1"
-  os.environ["ALL2ALL"] = "1"
-  os.environ["USE_ATOMICS"] = "1"
-  if "HK_FLASH_ATTENTION" not in os.environ:
-    os.environ["HK_FLASH_ATTENTION"] = "1"
-    if "ASM_GEMM" not in os.environ:
-      os.environ["ASM_GEMM"] = "1"
-from tinygrad import Tensor, nn, function, getenv, dtypes, TinyJit
-from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
-from tinygrad.uop.ops import Ops, UOp
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-FP8 = getenv("FP8", 0)
-
-FP8_DTYPE = dtypes.fp8e4m3
-FP8_MAX = 448.0
-
-def quantize_fp8(x:Tensor):
-  scale = FP8_MAX / (x.abs().max().detach() + 1e-8)
-  x_scaled = x * scale
-  x_clamped = x_scaled + (x_scaled.detach().clamp(-FP8_MAX, FP8_MAX) - x_scaled.detach())  # STE
-  return x_clamped.cast(FP8_DTYPE), scale.float().reciprocal()
-
-def matmul(x:Tensor, w:Tensor) -> Tensor:
-  if not FP8: return x @ w.T
-  # weights are already FP8, just quantize activations
-  x_fp8, x_scale = quantize_fp8(x)
-  return x_fp8.dot(w.T, dtype=dtypes.float) * x_scale
-
-def rmsnorm(x_in:Tensor, eps:float):
-  x = x_in.float()
-  x = x * (x.square().mean(-1, keepdim=True) + eps).rsqrt()
-  return x.cast(x_in.dtype)
-
-class FlatTransformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024):
-    self.vocab_size = vocab_size
-    self.n_layers = n_layers
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    # Attention
-    self.wqkv = self.lin_per_layer(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2)
-    self.wo = self.lin_per_layer(self.n_heads * self.head_dim, dim)
-
-    # FeedForward
-    self.w1 = self.lin_per_layer(dim, hidden_dim)
-    self.w2 = self.lin_per_layer(hidden_dim, dim)
-    self.w3 = self.lin_per_layer(dim, hidden_dim)
-
-    self.norm_eps = norm_eps
-    self.attention_norm = Tensor.ones(n_layers, dim).contiguous()
-    self.ffn_norm = Tensor.ones(n_layers, dim).contiguous()
-
-    # output
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = nn.Embedding(vocab_size, dim)
-    self.output = nn.Linear(dim, vocab_size, bias=False)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def lin_per_layer(self, in_features:int, out_features:int):
-    bound = 1 / math.sqrt(in_features)
-    dt = FP8_DTYPE if FP8 else None
-    if getenv("ZEROS"): return Tensor.zeros(self.n_layers, out_features, in_features, dtype=dt)
-    return Tensor.uniform(self.n_layers, out_features, in_features, low=-bound, high=bound, dtype=dt)
-
-  def attention(self, x:Tensor, freqs_cis:Tensor, attention_norm:Tensor, wqkv:Tensor, wo:Tensor):
-    x = rmsnorm(x, self.norm_eps) * attention_norm
-    xqkv = matmul(x, wqkv)
-
-    bsz, seqlen, _ = xqkv.shape
-    # interleaved layout: each kv group has [n_rep q heads, 1 k head, 1 v head] for clean MP sharding
-    xqkv = xqkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
-    xq = xqkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
-    xk = xqkv[:, :, :, self.n_rep].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-    xv = xqkv[:, :, :, self.n_rep+1].reshape(bsz, seqlen, self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-    attn = attn.reshape(bsz, seqlen, -1)
-    return matmul(attn, wo)
-
-  def feed_forward(self, x:Tensor, ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
-    x = rmsnorm(x, self.norm_eps) * ffn_norm
-    x_w1 = matmul(x, w1).silu()
-    x_w3 = matmul(x.contiguous_backward(), w3)
-    return matmul(x_w1 * x_w3, w2)
-
-  @function(precompile=True, precompile_backward=True)
-  def run_layer(self, x:Tensor, freqs_cis:Tensor,
-                attention_norm:Tensor, wqkv:Tensor, wo:Tensor,
-                ffn_norm:Tensor, w1:Tensor, w2:Tensor, w3:Tensor):
-    h = x + self.attention(x, freqs_cis, attention_norm, wqkv, wo)
-    return h + self.feed_forward(h, ffn_norm, w1, w2, w3)
-
-  def shard(self, device:tuple[str, ...], mp:bool=False):
-    from tinygrad.nn.state import get_parameters
-    if not mp:
-      for v in get_parameters(self): v.shard_(device, axis=None)
-    else:
-      # flat per-layer weights: axis 0 is n_layers, so shard axes are +1 vs per-layer Transformer
-      self.wqkv.shard_(device, axis=1).realize()          # (n_layers, out, dim) shard out
-      self.wo.shard_(device, axis=2).realize()             # (n_layers, dim, in) shard in
-      self.w1.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      self.w2.shard_(device, axis=2).realize()             # (n_layers, dim, hidden) shard in
-      self.w3.shard_(device, axis=1).realize()             # (n_layers, hidden, dim) shard out
-      self.attention_norm.shard_(device, axis=None).realize()
-      self.ffn_norm.shard_(device, axis=None).realize()
-      self.norm.weight.shard_(device, axis=None).realize()
-      self.tok_embeddings.weight.shard_(device, axis=0).realize()
-      self.output.weight.shard_(device, axis=0).realize()
-      self.freqs_cis.shard_(device, axis=None).realize()
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for i in range(self.n_layers):
-      h = self.run_layer(h, freqs_cis,
-                         self.attention_norm[i], self.wqkv[i], self.wo[i],
-                         self.ffn_norm[i], self.w1[i], self.w2[i], self.w3[i])
-    logits = self.output(self.norm(h))
-    return logits
-
-# TODO: this shouldn't be needed, but it prevents a copy of the grads. CAT can help
-def apply_grad(old_grad:UOp, new_grad:UOp) -> list[UOp]:
-  if new_grad.op == Ops.ADD:
-    return apply_grad(old_grad, new_grad.src[0])+apply_grad(old_grad, new_grad.src[1])
-  elif new_grad.op == Ops.PAD:
-    grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(new_grad.src[0].shape, new_grad.marg)])
-    return apply_grad(old_grad.shrink(grad_shrink), new_grad.src[0])
-  else:
-    return [old_grad.store(old_grad + new_grad)]
-
-if __name__ == "__main__":
-  config = {}
-  BS                 = config["BS"]                     = getenv("BS", 16)
-  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
-
-  from examples.llama3 import MODEL_PARAMS
-  model_params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: model_params['n_layers'] = llama_layers
-  model = FlatTransformer(**model_params, max_context=SEQLEN)
-  state = nn.state.get_state_dict(model)
-  print("tensor count:", len(state))
-
-  # shard the model
-  from tinygrad import Device
-  if (DP := getenv("DP", 1)) > 1:
-    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)))
-  if (MP := getenv("MP", 1)) > 1:
-    model.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)), mp=True)
-
-  # preallocate all the grad buffers and zero them out
-  grads = {x:Tensor.zeros_like(x).contiguous() for x in state.values() if x.requires_grad is None}
-
-  # print model size
-  sz = 0
-  for k,v in state.items():
-    print(f"{colored(k, 'green' if v in grads else 'white'):30s} {str(v.shape):30s} {str(v.dtype):20s} {v.device}  {v.nbytes()/1e9:.2f} GB")
-    sz += v.nbytes()
-  print(f"total sz: {sz/1e9:.2f} GB")
-
-  with Timing("fake data: "): tokens = Tensor.randint(BS, SEQLEN+1, low=0, high=model.vocab_size, dtype=dtypes.int)
-  with Timing("realize weights/grads/data: "): Tensor.realize(*state.values(), *grads.values(), tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
-  if DP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(DP)), axis=0)
-  if MP > 1: tokens = tokens.shard(tuple(f"{Device.DEFAULT}:{i}" for i in range(MP)))
-
-  @TinyJit
-  def jit_step(tokens:Tensor):
-    with Timing("python forward: "): loss = model(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    with Timing("python backward: "):
-      for t,g in zip(grads, loss.gradient(*grads)):
-        grads[t] = Tensor(grads[t].uop.after(UOp.group(*apply_grad(grads[t].uop, g.uop))), device=t.device)
-    with Timing("run step: "): loss.realize(*grads.values())
-
-  for i in range(6):
-    GlobalCounters.reset()
-    profile_marker(f"step {i}")
-    with Timing(colored(f"*** step {i}: ", "red")):
-      jit_step(tokens)
-  print("mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -1,80 +0,0 @@
-from tinygrad import Tensor, nn
-from tinygrad.helpers import getenv
-from extra.models.llama import apply_rotary_emb, precompute_freqs_cis
-
-class Attention:
-  def __init__(self, dim:int, n_heads:int, n_kv_heads:int|None=None, linear=nn.Linear):
-    self.n_heads = n_heads
-    self.n_kv_heads = n_kv_heads if n_kv_heads is not None else n_heads # n_kv_heads != n_heads implies MQA [arxiv/2307.09288, A.2.1]
-    self.head_dim = dim // n_heads
-    self.n_rep = self.n_heads // self.n_kv_heads
-
-    if getenv("WQKV"):
-      self.wqkv = linear(dim, self.n_heads * self.head_dim + self.n_kv_heads * self.head_dim * 2, bias=False)
-    else:
-      self.wq = linear(dim, self.n_heads * self.head_dim, bias=False)
-      self.wk = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-      self.wv = linear(dim, self.n_kv_heads * self.head_dim, bias=False)
-
-    self.wo = linear(self.n_heads * self.head_dim, dim, bias=False)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor) -> Tensor:
-    if getenv("WQKV"):
-      xqkv = self.wqkv(x)
-      xqkv = xqkv.reshape(xqkv.shape[0], xqkv.shape[1], self.n_kv_heads, self.n_rep + 2, self.head_dim)
-      xq = xqkv[:, :, :, :self.n_rep].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xk = xqkv[:, :, :, self.n_rep:self.n_rep+1].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-      xv = xqkv[:, :, :, self.n_rep+1:self.n_rep+2].reshape(xqkv.shape[0], xqkv.shape[1], -1)
-    else:
-      xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
-
-    xq = xq.reshape(xq.shape[0], xq.shape[1], self.n_heads, self.head_dim)
-    xk = xk.reshape(xk.shape[0], xk.shape[1], self.n_kv_heads, self.head_dim)
-    xv = xv.reshape(xv.shape[0], xv.shape[1], self.n_kv_heads, self.head_dim)
-
-    xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
-    bsz, seqlen, _, _ = xq.shape
-
-    xq, xk, xv = xq.transpose(1, 2), xk.transpose(1, 2), xv.transpose(1, 2)
-    attn = xq.scaled_dot_product_attention(xk, xv, is_causal=True, enable_gqa=True).transpose(1, 2)
-
-    attn = attn.reshape(bsz, seqlen, -1)
-    return self.wo(attn)
-
-class FeedForward:
-  def __init__(self, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.w1 = linear(dim, hidden_dim, bias=False)
-    self.w2 = linear(hidden_dim, dim, bias=False)
-    self.w3 = linear(dim, hidden_dim, bias=False) # the gate in Gated Linear Unit
-
-  def __call__(self, x:Tensor) -> Tensor:
-    w1 = self.w1(x).silu()
-    w3 = self.w3(x)
-    return self.w2(w1 * w3)
-
-class TransformerBlock:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_kv_heads:int|None, norm_eps:float, linear=nn.Linear):
-    self.attention = Attention(dim, n_heads, n_kv_heads, linear)
-    self.feed_forward = FeedForward(dim, hidden_dim, linear)
-    self.attention_norm = nn.RMSNorm(dim, norm_eps)
-    self.ffn_norm = nn.RMSNorm(dim, norm_eps)
-
-  def __call__(self, x:Tensor, freqs_cis:Tensor):
-    h = x + self.attention(self.attention_norm(x), freqs_cis)
-    return h + self.feed_forward(self.ffn_norm(h))
-
-class Transformer:
-  def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size:int, n_kv_heads:int|None=None,
-               rope_theta:int=10000, max_context:int=1024, linear=nn.Linear, embedding=nn.Embedding):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, linear) for _ in range(n_layers)]
-    self.norm = nn.RMSNorm(dim, norm_eps)
-    self.tok_embeddings = embedding(vocab_size, dim)
-    self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, max_context * 2, rope_theta).contiguous().requires_grad_(False)
-
-  def __call__(self, tokens:Tensor):
-    h = self.tok_embeddings(tokens)
-    freqs_cis = self.freqs_cis.cast(h.dtype)[:, :tokens.shape[1], :, :, :]
-    for layer in self.layers: h = layer(h, freqs_cis)
-    logits = self.output(self.norm(h))
-    return logits
@@ -1,140 +0,0 @@
-import os
-os.environ["WQKV"] = "1"
-import unittest
-import numpy as np
-from tinygrad import Tensor, nn, dtypes
-from tinygrad.nn.state import get_parameters
-from tinygrad.device import is_dtype_supported
-from examples.mlperf.models.llama import Transformer
-from examples.mlperf.models.flat_llama import FlatTransformer
-
-def copy_weights(flat:FlatTransformer, ref:Transformer):
-  n_layers = flat.n_layers
-  Tensor.realize(*nn.state.get_state_dict(ref).values())
-  flat.wqkv.assign(Tensor(np.stack([ref.layers[i].attention.wqkv.weight.numpy() for i in range(n_layers)])))
-  flat.wo.assign(Tensor(np.stack([ref.layers[i].attention.wo.weight.numpy() for i in range(n_layers)])))
-  flat.w1.assign(Tensor(np.stack([ref.layers[i].feed_forward.w1.weight.numpy() for i in range(n_layers)])))
-  flat.w2.assign(Tensor(np.stack([ref.layers[i].feed_forward.w2.weight.numpy() for i in range(n_layers)])))
-  flat.w3.assign(Tensor(np.stack([ref.layers[i].feed_forward.w3.weight.numpy() for i in range(n_layers)])))
-  flat.attention_norm.assign(Tensor(np.stack([ref.layers[i].attention_norm.weight.numpy() for i in range(n_layers)])))
-  flat.ffn_norm.assign(Tensor(np.stack([ref.layers[i].ffn_norm.weight.numpy() for i in range(n_layers)])))
-  flat.norm.weight.assign(Tensor(ref.norm.weight.numpy()))
-  flat.tok_embeddings.weight.assign(Tensor(ref.tok_embeddings.weight.numpy()))
-  flat.output.weight.assign(Tensor(ref.output.weight.numpy()))
-
-class TestFlatLlama(unittest.TestCase):
-  def test_forward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-    tokens = Tensor([[1, 50, 100, 999, 2]])
-    ref_logits = ref(tokens).realize()
-    flat_logits = flat(tokens).realize()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    diff = (ref_logits - flat_logits).abs().max().item()
-    self.assertLess(diff, 1e-5, f"forward mismatch: max abs diff {diff}")
-
-  def test_backward_match(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-
-    for p in get_parameters(ref): p.requires_grad_(True)
-    for p in get_parameters(flat): p.requires_grad_(True)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-    tokens = Tensor([[1, 50, 100, 999, 2, 10]])
-
-    ref_loss = ref(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    ref_loss.backward()
-    ref_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(ref).items() if v.grad is not None}
-
-    flat_loss = flat(tokens[:, :-1]).sparse_categorical_crossentropy(tokens[:, 1:])
-    flat_loss.backward()
-    flat_grads = {k: v.grad.numpy() for k, v in nn.state.get_state_dict(flat).items() if v.grad is not None}
-
-    # check loss matches
-    self.assertAlmostEqual(ref_loss.item(), flat_loss.item(), places=4)
-
-    # check output weight grad matches
-    diff = abs(ref_grads["output.weight"] - flat_grads["output.weight"]).max()
-    self.assertLess(diff, 1e-4, f"output.weight grad mismatch: max abs diff {diff}")
-
-    # check per-layer weight grads match
-    for i in range(params["n_layers"]):
-      for flat_key, ref_key in [
-        ("wqkv", f"layers.{i}.attention.wqkv.weight"),
-        ("wo", f"layers.{i}.attention.wo.weight"),
-        ("w1", f"layers.{i}.feed_forward.w1.weight"),
-        ("w2", f"layers.{i}.feed_forward.w2.weight"),
-        ("w3", f"layers.{i}.feed_forward.w3.weight"),
-      ]:
-        diff = abs(ref_grads[ref_key] - flat_grads[flat_key][i]).max()
-        self.assertLess(diff, 1e-4, f"layer {i} {flat_key} grad mismatch: max abs diff {diff}")
-
-  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
-  def test_forward_match_mp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices, mp=True)
-
-    tokens = Tensor([[1, 50, 100, 999, 2]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @unittest.skipUnless(os.getenv("CPU", "") == "1", "multi-device CPU test")
-  def test_forward_match_dp(self):
-    Tensor.manual_seed(42)
-    params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-    from tinygrad import Device
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    ref = Transformer(**params)
-    flat = FlatTransformer(**params)
-    copy_weights(flat, ref)
-    Tensor.realize(*nn.state.get_state_dict(flat).values())
-    flat.shard(devices)
-
-    tokens = Tensor([[1, 50, 100, 999, 2], [2, 100, 50, 1, 999]], device=devices[0])
-    ref_logits = ref(tokens.to(devices[0])).numpy()
-    flat_logits = flat(tokens.shard(devices, axis=0)).numpy()
-    self.assertEqual(ref_logits.shape, flat_logits.shape)
-    np.testing.assert_allclose(flat_logits, ref_logits, atol=1e-4, rtol=1e-4)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3), "fp8 not supported on this device")
-  def test_forward_fp8(self):
-    import examples.mlperf.models.flat_llama as flat_llama_mod
-    old_fp8 = flat_llama_mod.FP8
-    try:
-      flat_llama_mod.FP8 = 1
-      Tensor.manual_seed(42)
-      params = dict(dim=128, hidden_dim=256, n_heads=4, n_kv_heads=2, n_layers=2, norm_eps=1e-5, vocab_size=1024, rope_theta=10000, max_context=64)
-      ref = Transformer(**params)
-      flat = FlatTransformer(**params)
-      copy_weights(flat, ref)
-      Tensor.realize(*nn.state.get_state_dict(flat).values())
-
-      tokens = Tensor([[1, 50, 100, 999, 2]])
-      ref_logits = ref(tokens).numpy()
-      flat_logits = flat(tokens).numpy()
-      self.assertEqual(ref_logits.shape, flat_logits.shape)
-      # FP8 has lower precision, allow larger tolerance
-      np.testing.assert_allclose(flat_logits, ref_logits, atol=1.0, rtol=0.1)
-    finally:
-      flat_llama_mod.FP8 = old_fp8
-
-if __name__ == "__main__":
-  unittest.main()
@@ -25,10 +25,8 @@ class GradAccClipAdamW(Optimizer):
    return extra[-1]

  def _step(self, params:list[Tensor], grads:list[Tensor]) -> tuple[list[Tensor], list[Tensor]]:
-    grads = list(grads)
-
    for i in range(len(grads)):
-      if grads[i].device != self.m[i].device: grads[i] = grads[i].to(self.m[i].device)
+      if grads[i].device != self.m[i].device: grads[i].assign(grads[i].to(self.m[i].device))

    if self.fused:
      grads[0].assign(grads[0] / self.grad_acc)
@@ -36,10 +34,10 @@ class GradAccClipAdamW(Optimizer):
      grads[0].assign((grads[0] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[0].dtype))
    else:
      for i in range(len(grads)):
-        grads[i].assign(grads[i] / self.grad_acc)
-      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
+        grads[i].assign(grads[i] / self.grad_acc).realize()
+      total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous().realize()
      for i in range(len(grads)):
-        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype))
+        grads[i].assign((grads[i] * (self.clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(grads[i].dtype)).realize()

    ret = []
    self.b1_t *= self.b1
@@ -47,13 +45,13 @@ class GradAccClipAdamW(Optimizer):
    for i, g in enumerate(grads):
      self.m[i].assign((self.b1 * self.m[i] + (1.0 - self.b1) * g).cast(self.m[i].dtype))
      self.v[i].assign((self.b2 * self.v[i] + (1.0 - self.b2) * (g * g)).cast(self.v[i].dtype))
-      m_hat = (self.m[i] / (1.0 - self.b1_t)).cast(self.m[i].dtype)
-      v_hat = (self.v[i] / (1.0 - self.b2_t)).cast(self.v[i].dtype)
+      m_hat = self.m[i] / (1.0 - self.b1_t)
+      v_hat = self.v[i] / (1.0 - self.b2_t)
      up = m_hat / (v_hat.sqrt() + self.eps)
      ret.append((self.lr * up).cast(g.dtype))
    return ret, [self.b1_t, self.b2_t] + self.m + self.v + [total_norm]

  def _apply_update(self, t:Tensor, up:Tensor) -> Tensor:
-    wd = self.wd if t.ndim >= 3 else 0.0
+    wd = self.wd if t.ndim >= 2 else 0.0
    up = up.shard_like(t) + self.lr.to(t.device) * wd * t.detach()
    return t.detach() - up.cast(t.dtype)
@@ -5,7 +5,6 @@ export DEV=${DEV:-AMD}
 export EMULATE="AMD_CDNA4"
 export CHECK_OOB=0
 export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
-export DEVICE_IN_FUNCTION_BUG=1

 export DEBUG=${DEBUG:-2}
 export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
@@ -13,7 +12,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}
@@ -12,7 +12,6 @@ export ALL2ALL=${ALL2ALL:-1}
 export USE_ATOMICS=${USE_ATOMICS:-0}
 export ASM_GEMM=${ASM_GEMM:-1}
 export WQKV=${WQKV:-1}
-export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-1}

 export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
 export DP=${DP:-1} MP=${MP:-8}
@@ -1,5 +1,6 @@
 #!/bin/bash
 export BENCHMARK=5
 export EVAL_BS=0
-VIZ=${VIZ:--1} examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
-extra/viz/cli.py --profile --device "AMD" --limit 20
+export VIZ=${VIZ:--1}
+examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama8b/implementations/tinybox_8xMI350X/dev_run.sh
+extra/viz/cli.py --profile --device "AMD" --top 20
@@ -1,11 +1,12 @@
 from pathlib import Path
+import numpy as np
 import torch
 from torchvision.utils import make_grid, save_image
 from tinygrad.nn.state import get_parameters
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import trange
 from tinygrad.nn import optim
-from tinygrad.nn.datasets import mnist
+from extra.datasets import fetch_mnist

 class LinearGen:
  def __init__(self):
@@ -37,14 +38,14 @@ class LinearDisc:
    return x

 def make_batch(images):
-  sample = Tensor.randint(batch_size, low=0, high=images.shape[0])
-  return images[sample].reshape(batch_size, 28*28).cast('float').div(127.5).sub(1.0)
+  sample = np.random.randint(0, len(images), size=(batch_size))
+  image_b = images[sample].reshape(-1, 28*28).astype(np.float32) / 127.5 - 1.0
+  return Tensor(image_b)

 def make_labels(bs, col, val=-2.0):
-  y = Tensor.zeros(bs, 2)
-  if col == 0: y = y + Tensor([val, 0.0])
-  else: y = y + Tensor([0.0, val])
-  return y
+  y = np.zeros((bs, 2), np.float32)
+  y[range(bs), [col] * bs] = val  # Can we do label smoothing? i.e -2.0 changed to -1.98789.
+  return Tensor(y)

 def train_discriminator(optimizer, data_real, data_fake):
  real_labels = make_labels(batch_size, 1)
@@ -70,12 +71,12 @@ def train_generator(optimizer, data_fake):

 if __name__ == "__main__":
  # data for training and validation
-  X_train, _, _, _ = mnist()
+  images_real = np.vstack(fetch_mnist()[::2])
  ds_noise = Tensor.randn(64, 128, requires_grad=False)
  # parameters
  epochs, batch_size, k = 300, 512, 1
  sample_interval = epochs // 10
-  n_steps = X_train.shape[0] // batch_size
+  n_steps = len(images_real) // batch_size
  # models and optimizer
  generator = LinearGen()
  discriminator = LinearDisc()
@@ -83,24 +84,24 @@ if __name__ == "__main__":
  output_dir = Path(".").resolve() / "outputs"
  output_dir.mkdir(exist_ok=True)
  # optimizers
-  optim_g = optim.Adam(get_parameters(generator), lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
-  optim_d = optim.Adam(get_parameters(discriminator), lr=0.0002, b1=0.5)
+  optim_g = optim.Adam(get_parameters(generator),lr=0.0002, b1=0.5)  # 0.0002 for equilibrium!
+  optim_d = optim.Adam(get_parameters(discriminator),lr=0.0002, b1=0.5)
  # training loop
-  with Tensor.train():
-    for epoch in (t := trange(epochs)):
-      loss_g, loss_d = 0.0, 0.0
-      for _ in range(n_steps):
-        data_real = make_batch(X_train)
-        for step in range(k):  # Try with k = 5 or 7.
-          noise = Tensor.randn(batch_size, 128)
-          data_fake = generator.forward(noise).detach()
-          loss_d += train_discriminator(optim_d, data_real, data_fake)
+  Tensor.training = True
+  for epoch in (t := trange(epochs)):
+    loss_g, loss_d = 0.0, 0.0
+    for _ in range(n_steps):
+      data_real = make_batch(images_real)
+      for step in range(k):  # Try with k = 5 or 7.
        noise = Tensor.randn(batch_size, 128)
-        data_fake = generator.forward(noise)
-        loss_g += train_generator(optim_g, data_fake)
-      if (epoch + 1) % sample_interval == 0:
-        fake_images = generator.forward(ds_noise).detach().numpy()
-        fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
-        save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
-      t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
+        data_fake = generator.forward(noise).detach()
+        loss_d += train_discriminator(optim_d, data_real, data_fake)
+      noise = Tensor.randn(batch_size, 128)
+      data_fake = generator.forward(noise)
+      loss_g += train_generator(optim_g, data_fake)
+    if (epoch + 1) % sample_interval == 0:
+      fake_images = generator.forward(ds_noise).detach().numpy()
+      fake_images = (fake_images.reshape(-1, 1, 28, 28) + 1) / 2  # 0 - 1 range.
+      save_image(make_grid(torch.tensor(fake_images)), output_dir / f"image_{epoch+1}.jpg")
+    t.set_description(f"Generator loss: {loss_g/n_steps}, Discriminator loss: {loss_d/n_steps}")
  print("Training Completed!")
@@ -65,7 +65,7 @@ def get_bar0_size(pcibus):
 class AMSMI(AMDev):
  def __init__(self, pcibus, vram_bar:MMIOInterface, doorbell_bar:MMIOInterface, mmio_bar:MMIOInterface):
    self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
+    self.vram, self.doorbell64, self.mmio, self.dma_regions = vram_bar, doorbell_bar, mmio_bar, None
    self.pci_state = self.read_pci_state()
    if self.pci_state == "D0": self._init_from_d0()

@@ -2,7 +2,7 @@

 import os
 from tinygrad.helpers import Context
-from tinygrad.runtime.support.system import System, PCIDevice
+from tinygrad.runtime.support.system import System, PCIDevice, PCIDevImplBase
 from tinygrad.runtime.support.hcq import FileIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev

@@ -12,7 +12,7 @@ if __name__ == "__main__":
    drv_path = f"/sys/bus/pci/devices/{gpu}/driver"
    if FileIOInterface.exists(drv_path) and os.path.basename(os.readlink(drv_path)) == "amdgpu":
      raise RuntimeError(f"amdgpu is bound to {gpu}. Stopping...")
-  pcidevs = [PCIDevice("AM", gpu) for gpu in gpus]
+  pcidevs = [PCIDevice("AM", gpu, bars=[0, 2, 5]) for gpu in gpus]
  amdevs = []
  with Context(DEBUG=2):
    for pcidev in pcidevs:
@@ -7,8 +7,8 @@ class GFXFake:
  def __init__(self): self.xccs = 8

 class AMDFake(AMDev):
-  def __init__(self, pci_dev):
-    self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
+  def __init__(self, pci_dev, dma_regions=None):
+    self.pci_dev, self.devfmt, self.dma_regions = pci_dev, pci_dev.pcibus, dma_regions
    self.vram, self.doorbell64, self.mmio = self.pci_dev.map_bar(0), self.pci_dev.map_bar(2, fmt='Q'), self.pci_dev.map_bar(5, fmt='I')
    self._run_discovery()
    self._build_regs()
@@ -0,0 +1,139 @@
+from typing import Callable
+from tinygrad import UOp, dtypes, Device, Tensor, getenv, function
+from tinygrad.uop.ops import AxisType, AddrSpace
+
+def simple_function(fxn:Callable[..., UOp]) -> Callable[..., UOp]:
+  def wrapper(*args:UOp) -> UOp:
+    params:list[UOp] = [x.param_like(i) for i,x in enumerate(args)]
+    return fxn(*params).call(*args)
+  return wrapper
+
+THREADS_PER_BLOCK = 128
+WARP_SIZE = 32
+
+# Register tile sizes (per-thread accumulator tile of C)
+TN = 4     # columns per thread
+TM = 4     # rows per thread
+
+WAVE_TILE_N = 128
+WAVE_TILE_M = 32
+
+LANES_PER_WAVE_X = 8
+LANES_PER_WAVE_Y = 4
+ITERS_PER_WAVE_N = 4 #WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
+ITERS_PER_WAVE_M = 2 #WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
+
+WAVES_IN_BLOCK_Y = 4
+WAVES_IN_BLOCK_X = 1
+
+
+N = getenv("N", 4096)
+M = K = N
+
+# 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
+
+@simple_function
+def slice_matmul(c_regs, a_local, b_local):
+  # 2x
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+
+  pass
+
+@simple_function
+def compute_local(c:UOp, a_local:UOp, b_local:UOp) -> UOp:
+  # this is the LID level on the GPU, here we can define regs
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M*TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N*TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+  # do the math
+  A_col = A_col.assign(a_local[k_tile].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :].flatten())
+  B_row = B_row.assign(b_local[k_tile].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :].flatten())
+  c_regs += A_col.reshape(-1, 1) * B_row.reshape(1, -1)  #
+  c_regs
+
+
+
+@simple_function
+def load_local(a_local, b_local, a_global, b_global):
+  # NOTE: it ends this range, so there's a BARRIER
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  return UOp.group(
+    a_local[:, tid].store(a_global[tid, :]),
+    b_local[:, tid].store(b_global[:, tid]))
+
+@simple_function
+def reg_matmul(c_regs, a_local, b_local):
+  A_col = UOp.placeholder((ITERS_PER_WAVE_M*TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
+  B_row = UOp.placeholder((ITERS_PER_WAVE_N*TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
+
+
+
+@simple_function
+def local_matmul(c:UOp, a:UOp, b:UOp, a_local:UOp, b_local:UOp):
+  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
+  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
+
+  # this is the LID level on the GPU, this (and below) is where we define REGs
+  c_regs = UOp.placeholder((ITERS_PER_WAVE_M*TM, ITERS_PER_WAVE_N*TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
+
+  # 128x128, Kx128, Kx128
+  k_tile = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)*BLOCK_K
+  fxn = reg_matmul(c_regs.assign(0),
+                   a_local[:, tid].assign(a[k_tile:k_tile+BLOCK_K, tid]),
+                   b_local[:, tid].assign(b[k_tile:k_tile+BLOCK_K, tid]))
+
+  # do math
+  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)
+  return c[waveIdy, :, laneIdy, :, waveIdx, :, laneIdx, :].store(c_regs.after(fxn))
+
+@simple_function
+def global_matmul(c:UOp, a:UOp, b:UOp):
+  # this is the GID level on the GPU, this is where we define LOCAL buffers shared across lids
+  gx = UOp.range(N//BLOCK_N, 0, AxisType.GLOBAL) * BLOCK_N
+  gy = UOp.range(M//BLOCK_M, 1, AxisType.GLOBAL) * BLOCK_M
+  a_local = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL)
+  b_local = UOp.placeholder((BLOCK_K, BLOCK_M), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
+  return local_matmul(c[gx:gx+BLOCK_N, gy:gy+BLOCK_M], a.permute(1,0)[:, gx:gx+BLOCK_N], b[:, gy:gy+BLOCK_M], a_local, b_local)
+
+  #ll = load_local(a_local, b_local, a.permute(1,0)[:, gx:gx+BLOCK_N], b[:, gy:gy+BLOCK_M])
+  #return compute_local(c[gx:gx+BLOCK_N, gy:gy+BLOCK_M], a_local.after(ll), b_local.after(ll))
+
+if __name__ == "__main__":
+  # this is the outer lvel on the GPU, this is where we define GLOBAL buffers
+  C = Tensor.empty(N, M)
+  A = Tensor.randn(N, K)
+  B = Tensor.randn(K, M)
+  c_out = C.call(A, B, fxn=global_matmul).numpy()
+
+  #C = UOp.new_buffer(Device.DEFAULT, N*M, dtypes.float).reshape(N,M)
+  #A = UOp.new_buffer(Device.DEFAULT, N*K, dtypes.float).reshape(N,K)
+  #B = UOp.new_buffer(Device.DEFAULT, K*M, dtypes.float).reshape(K,M)
+  #global_matmul(C, A, B).realize()
+
+  # input matmuls
+  #c = UOp.param(0, dtypes.float, (N, M))
+  #a = UOp.param(1, dtypes.float, (N, K))
+  #b = UOp.param(2, dtypes.float, (K, M))
+
+
+  #ba = a.rearrange("(n bn) (k bk) -> n k bn bk", bn=BLOCK_N, bk=BLOCK_K)[gx, k_tile_range]
+  #bb = b.rearrange("(k bk) (m bm) -> k m bk bm", bk=BLOCK_K, bm=BLOCK_M)[k_tile_range, gy]
+  #bc = c.rearrange("(n bn) (m bm) -> n m bn bm", bn=BLOCK_N, bm=BLOCK_M)[gx, gy]
@@ -0,0 +1,85 @@
+from tinygrad import UOp, dtypes, Device, Tensor
+
+if __name__ == "__main__":
+  B0 = UOp.new_buffer(Device.DEFAULT, 100, dtypes.float).reshape(10,10)
+  B1 = UOp.new_buffer(Device.DEFAULT, 100, dtypes.float).reshape(10,10)
+
+
+  b0 = UOp.param(0, dtypes.float, (10,10))
+  b1 = UOp.param(1, dtypes.float, (10,10))
+  r0 = UOp.range(10, axis_id=0)
+  r1 = UOp.range(10, axis_id=1)
+
+  fxn = (b0[r0, r1] + b1[r0, r1]).call(B0, B1)
+  t = Tensor(fxn)
+  t.realize()
+
+# gemm (N,N)
+
+# (N//k, k, N//k, k)
+
+
+# what if call just implicitly ends all ranges and you don't need to connect them?
+# you do have to connect them, and it does end the ranges
+
+# if assign (store+after) is on call, we move the store into the call (indexed with the ranges) and replace the assign with an after
+
+
+def gemm(A, B):
+  N = 4096
+  k = 128
+
+  ia = UOp.param(0, dtypes.float, (k, k)).reshape(k, 1, k)
+  ib = UOp.param(1, dtypes.float, (k, k)).reshape(1, k, k)
+  gemm_fxn = (ia * ib).sum(2) # <-- rangeify this
+
+  a = UOp.param(0, dtypes.float, (N, N))
+  b = UOp.param(1, dtypes.float, (N, N))
+  r0 = UOp.range(N//k, 0)
+  r1 = UOp.range(N//k, 1)
+  local_fxn = gemm_fxn.call(a.reshape(N//k, k, N//k, k)[r0, :, r1, :], b.reshape(N//k, k, N//k, k)[r0, :, r1, :], r0, r1).permute(0,2,1,3).reshape(N,N)
+
+  fxn = local_fxn.call(A,B)
+
+
+
+  return
+
+
+
+
+
+  a = UOp.param(0, dtypes.float, (N//k, k, N//k, k))
+  b = UOp.param(1, dtypes.float, (N//k, k, N//k, k))
+
+
+
+  # inner kxk GEMM (are WMMAs calls?)
+  ia = UOp.param(0, dtypes.float, (k,k)).reshape(k, 1, k)
+  ib = UOp.param(1, dtypes.float, (k,k)).reshape(1, k, k)
+  r0 = UOp.range(N//k, 0)
+  r1 = UOp.range(N//k, 1)
+  fxn = (ia * ib).sum(2).call(a[:, r0, :, r1], b[:, r0, :, r1])   # this call ends these ranges implicitly
+  assert fxn.shape == (N//k, N//k, k, k)
+
+
+  #.call(A, B, UOp.range(N//k), UOp.range(N//k))
+
+  #r0 = UOp.param(2, dtypes.index, (), vmin_vmax=(0, N//k-1))
+  #r1 = UOp.param(3, dtypes.index, (), vmin_vmax=(0, N//k-1))
+
+
+
+
+
+
+
+
+
+
+# Q = [batch, seq_len, heads, dim]
+# K = [batch, seq_len, head_kv, dim]
+# V = [batch, seq_len, head_kv, dim]
+
+
+
@@ -1,11 +1,11 @@
 from typing import Tuple, Dict, List, Optional
-from tinygrad.dtype import DType, dtypes
+from tinygrad.dtype import DType
 from tinygrad.renderer import ProgramSpec
-from tinygrad.tensor import Tensor
-from tinygrad.device import Device
+from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
 from tinygrad.helpers import Context, to_mv
+from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict
@@ -13,20 +13,12 @@ from collections import OrderedDict
 EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]

 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
-  # memory-planned subbuffers can have multiple Buffer objects for the same memory region
-  canon, _seen = {}, {}
-  for ji in run.jit_cache:
-    for b in ji.bufs:
-      if b is not None: canon[id(b)] = _seen.setdefault((id(b.base._buf), b.offset, b.size, b.dtype), b)
-  special_names = {id(canon[k]): v for k, v in special_names.items() if k in canon}
-
  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
  for ji in run.jit_cache:
    fxn: ProgramSpec = ji.prg.p
    functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
    cargs = []
    for i,arg in enumerate(ji.bufs):
-      arg = canon[id(arg)]
      key = id(arg)
      if key not in bufs:
        if key in special_names:
@@ -291,7 +291,7 @@ def build_kernel(N, arch='gfx1100'):
  # MAIN GEMM LOOP
  # ===========================================================================

-  NO_ALU, NO_DS, NO_GLOBAL = getenv("NO_ALU", 0), getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)
+  NO_DS, NO_GLOBAL = getenv("NO_DS", 0), getenv("NO_GLOBAL", 0)

  k.label('LOOP_INC')
  k.emit(s_add_i32(s[S_LOOP_CTR], s[S_LOOP_CTR], 8))
@@ -350,11 +350,10 @@ def build_kernel(N, arch='gfx1100'):

    # 64 dual FMACs
    k.waitcnt(lgkm=0)
-    if not NO_ALU:
-      k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
-      for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
-        k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
-                    vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))
+    k.emit(s_clause(simm16=len(FMAC_PATTERN)-1))
+    for i, (vdst_x, vdst_y, ax, bx, ay, by) in enumerate(FMAC_PATTERN):
+      k.emit(VOPD(VOPDOp.V_DUAL_FMAC_F32, VOPDOp.V_DUAL_FMAC_F32,
+                  vdstx=v[vdst_x], vdsty=v[vdst_y], srcx0=v[ax], vsrcx1=v[bx], srcy0=v[ay], vsrcy1=v[by]))

  # wait for all global loads to finish
  # then sync the warp so it's safe to store local
@@ -1,110 +0,0 @@
-from tinygrad import UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-from tinygrad.dtype import AddrSpace, dtypes
-
-N = getenv("N", 4096)
-M = getenv("M", N)
-K = getenv("K", N)
-
-WARP_SIZE = 32
-BLOCK_M, BLOCK_N = 128, 128
-BLOCK_K = getenv("BK", 16)
-assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
-
-use_wmma = getenv("WMMA")
-if use_wmma:
-  WAVES_M, WAVES_N = 2, 2
-  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-  UNROLL_M, UNROLL_N = 1, 1
-
-  # wmma params
-  WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-  WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
-else:
-  WAVES_M, WAVES_N = 4, 1
-  LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
-  UNROLL_M, UNROLL_N = 4, 4
-
-# WARP_SIZE * total waves
-THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
-
-# accumulator size
-TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
-TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
-
-def block_128x128_gemm(c:UOp, a:UOp, b:UOp) -> UOp:
-  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
-  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
-  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-
-  # -- GLOBAL -> LOCAL --
-  # wmma: spatial outer, k inner (k contiguous for vectorized WMMA tile loads)
-  # gemm: k outer, spatial inner
-  A_local = UOp.placeholder((BLOCK_M, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_M), a.dtype.base, slot=0, addrspace=AddrSpace.LOCAL)
-  B_local = UOp.placeholder((BLOCK_N, BLOCK_K) if use_wmma else (BLOCK_K, BLOCK_N), b.dtype.base, slot=1, addrspace=AddrSpace.LOCAL)
-
-  a = a.reshape(K // BLOCK_K, BLOCK_K, BLOCK_M)
-  b = b.reshape(K // BLOCK_K, BLOCK_K, BLOCK_N)
-  k_tile = UOp.range(K // BLOCK_K, 100, AxisType.REDUCE)
-
-  # copy with transpose for wmma (input is k×spatial, LDS is spatial×k)
-  A_copy = A_local.permute((1,0)) if use_wmma else A_local
-  B_copy = B_local.permute((1,0)) if use_wmma else B_local
-  A_store = A_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(a[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  B_store = B_copy.reshape(-1, THREADS_PER_BLOCK)[:, tid].store(b[k_tile].reshape(-1, THREADS_PER_BLOCK)[:, tid])
-  barrier = UOp.barrier(A_store, B_store)
-  A_local, B_local = A_local.after(barrier), B_local.after(barrier)
-
-  # -- COMPUTE --
-  lane_m, lane_n = lane // LANES_PER_WAVE_N, lane % LANES_PER_WAVE_N
-
-  # accumulator (unified: both paths use (TM, TN) with scalar dtypes.float)
-  acc = UOp.placeholder((TM, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(UOp.const(dtypes.float, 0).reshape((1,)*len(acc.shape)).expand(acc.shape)))
-
-  if use_wmma:
-    k = UOp.range(BLOCK_K // WMMA_K, 101, AxisType.REDUCE)
-    tile_m = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-    tile_n = UOp.range(TN, 201, AxisType.LOOP)
-
-    acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0,2,1)[tile_m, tile_n]
-    a_frag = A_local.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_K // WMMA_K, WMMA_K)[wave_m, tile_m, lane_n, k]
-    b_frag = B_local.reshape(WAVES_N, TN, WMMA_N, BLOCK_K // WMMA_K, WMMA_K)[wave_n, tile_n, lane_n, k]
-
-    wmma = UOp(Ops.SHAPED_WMMA, dtypes.float, (a_frag, b_frag, acc_frag.after(k)), arg=((16, 16, 16), 'AMD', 32))
-    acc_store = acc_frag.store(wmma).end(tile_m, tile_n)
-  else:
-    # registers for LOCAL -> REG
-    a_frag = UOp.placeholder((TM//UNROLL_M, UNROLL_M), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-    b_frag = UOp.placeholder((TN//UNROLL_N, UNROLL_N), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-
-    k = UOp.range(BLOCK_K, 101, AxisType.REDUCE)
-    a_frag = a_frag.after(a_frag.store(A_local[k].reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M)[wave_m, :, lane_m, :]))
-    b_frag = b_frag.after(b_frag.store(B_local[k].reshape(WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)[wave_n, :, lane_n, :]))
-
-    # FMA
-    a_frag = a_frag.reshape(TM, 1).expand(TM, TN)
-    b_frag = b_frag.reshape(1, TN).expand(TM, TN)
-    acc_store = acc.store(acc.after(k) + (a_frag * b_frag))
-
-  # store accumulator and loop
-  acc = acc.after(acc_store.end(k).barrier().end(k_tile))
-
-  # store accumulator to output (unified)
-  c = c.reshape(WAVES_M, TM//UNROLL_M, LANES_PER_WAVE_M, UNROLL_M,
-                WAVES_N, TN//UNROLL_N, LANES_PER_WAVE_N, UNROLL_N)
-  c = c.permute((0,4,2,6, 1,3,5,7)).reshape(THREADS_PER_BLOCK, TM, TN)
-  return c[tid].store(acc).end(wave_m, wave_n, lane)
-
-def amd_copy_matmul(c:UOp, a:UOp, b:UOp) -> UOp:
-  block_id_m = UOp.range(M // BLOCK_M, 0, AxisType.GLOBAL)
-  block_id_n = UOp.range(N // BLOCK_N, 1, AxisType.GLOBAL)
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]
-  a = a.T.reshape(K, M // BLOCK_M, BLOCK_M)[:, block_id_m, :]
-  b = b.reshape(K, N // BLOCK_N, BLOCK_N)[:, block_id_n, :]
-  return block_128x128_gemm(c, a, b).end(block_id_n, block_id_m).sink(arg=KernelInfo(opts_to_apply=()))
-
-if __name__ == "__main__":
-  from amd_uop_matmul import eval_custom_matmul
-  eval_custom_matmul(amd_copy_matmul, dtypes.half if use_wmma else dtypes.float)
@@ -1,205 +0,0 @@
-from tinygrad import Tensor, UOp, getenv
-from tinygrad.uop.ops import AxisType, KernelInfo, Ops
-from tinygrad.dtype import AddrSpace, dtypes
-from tinygrad.helpers import DEBUG, GlobalCounters, Context
-import math
-
-BLOCK_M, BLOCK_N = 64, 64
-WARP_SIZE = 32
-WMMA_M, WMMA_N, WMMA_K = 16, 16, 16
-WAVES_M, WAVES_N = 4, 1
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 2, 16
-WMMA_ACC = WMMA_M // LANES_PER_WAVE_M
-THREADS_PER_BLOCK = WARP_SIZE * WAVES_M * WAVES_N
-LDS_PAD = 4  # pad LDS rows to reduce bank conflicts
-
-WMMA_ARG = ((WMMA_M, WMMA_N, WMMA_K), 'AMD', 32)
-LOG2E = math.log2(math.e)
-
-def warp_shfl_xor(val, offset, lane):
-  """Read val from lane ^ offset using ds_bpermute."""
-  idx = ((lane ^ offset) * 4).cast(dtypes.int)
-  return UOp(Ops.CUSTOM, dtypes.float, (idx, val),
-             arg="__builtin_bit_cast(float, __builtin_amdgcn_ds_bpermute({0}, __builtin_bit_cast(int, {1})))")
-
-def warp_reduce_max(val, lane):
-  """Tree reduce MAX across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = UOp(Ops.MAX, dtypes.float, (val, warp_shfl_xor(val, offset, lane)))
-  return val
-
-def warp_reduce_sum(val, lane):
-  """Tree reduce SUM across LANES_PER_WAVE_N=16 lanes."""
-  for offset in [8, 4, 2, 1]:
-    val = val + warp_shfl_xor(val, offset, lane)
-  return val
-
-def amd_flash_attention(o:UOp, q:UOp, k:UOp, v:UOp) -> UOp:
-  # inputs are (B*H, N, D)
-  BH, N, D = q.shape
-  assert N % BLOCK_M == 0 and N % BLOCK_N == 0, f"N={N} must be divisible by BLOCK_M={BLOCK_M} and BLOCK_N={BLOCK_N}"
-  assert D % WMMA_K == 0 and D % LANES_PER_WAVE_N == 0, f"D={D} must be divisible by WMMA_K={WMMA_K} and LANES_PER_WAVE_N={LANES_PER_WAVE_N}"
-  assert BLOCK_M % (WAVES_M * WMMA_M) == 0 and BLOCK_N % LANES_PER_WAVE_N == 0
-  TM = BLOCK_M // (WAVES_M * LANES_PER_WAVE_M)
-  TN = BLOCK_N // (WAVES_N * LANES_PER_WAVE_N)
-  TD = D // (WAVES_N * LANES_PER_WAVE_N)
-  SCALE = 1.0 / math.sqrt(D)
-
-  block_bh = UOp.range(BH, 0, AxisType.GLOBAL)
-  block_m = UOp.range(N // BLOCK_M, 1, AxisType.GLOBAL)
-
-  q = q.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-  k = k.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  v = v.reshape(BH, N//BLOCK_N, BLOCK_N, D)[block_bh]
-  o = o.reshape(BH, N//BLOCK_M, BLOCK_M, D)[block_bh, block_m]
-
-  wave_m = UOp.range(WAVES_M, 2, AxisType.LOCAL)
-  wave_n = UOp.range(WAVES_N, 3, AxisType.LOCAL)
-  lane = UOp.range(WARP_SIZE, -1, AxisType.WARP)
-  tid = (wave_m * WAVES_N + wave_n) * WARP_SIZE + lane
-  lane_m = lane // LANES_PER_WAVE_N
-  lane_n = lane % LANES_PER_WAVE_N
-
-  # LDS allocation: slot 0 = Q then P (shared), slot 1 = K then V
-  # TODO: the memory planner should be able to find this reuse
-  ELEMS_PER_THREAD = BLOCK_M * D // THREADS_PER_BLOCK
-  QP_lds = UOp.placeholder((BLOCK_M, D + LDS_PAD), dtypes.half, slot=0, addrspace=AddrSpace.LOCAL)
-  KV_lds = UOp.placeholder((BLOCK_N, D + LDS_PAD), dtypes.half, slot=1, addrspace=AddrSpace.LOCAL)[:, :D]
-
-  # register state
-  acc = UOp.placeholder((TM, TD), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  m_i = UOp.placeholder((TM,), dtypes.float, slot=3, addrspace=AddrSpace.REG)
-  l_i = UOp.placeholder((TM,), dtypes.float, slot=4, addrspace=AddrSpace.REG)
-  acc = acc.after(acc.store(acc.const_like(0)))
-  m_i = m_i.after(m_i.store(m_i.const_like(-math.inf)))
-  l_i = l_i.after(l_i.store(l_i.const_like(0)))
-
-  # ====== KV tile loop ======
-  n_tile = UOp.range(N // BLOCK_N, 100, AxisType.REDUCE)
-
-  # load Q + K into LDS (Q reloaded each iteration since P overwrites slot 0)
-  Q_lds = QP_lds[:, :D]
-  Q_store = Q_lds.after(n_tile).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    q.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  K_store = KV_lds.reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    k[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  qk_load_barrier = UOp.barrier(UOp.group(Q_store, K_store))
-  Q_lds = Q_lds.after(qk_load_barrier)
-  KV_lds_k = KV_lds.after(qk_load_barrier)
-
-  # -- S = Q @ K^T via WMMA (re-init each n_tile) --
-  S_reg = UOp.placeholder((TM, TN), dtypes.float, slot=6, addrspace=AddrSpace.REG)
-  S_reg = S_reg.after(S_reg.after(n_tile).store(S_reg.const_like(0)))
-  k_qk = UOp.range(D // WMMA_K, 101, AxisType.REDUCE)
-  tm1 = UOp.range(TM // WMMA_ACC, 200, AxisType.LOOP)
-  tn1 = UOp.range(TN, 201, AxisType.LOOP)
-  S_frag = S_reg.reshape(TM // WMMA_ACC, WMMA_ACC, TN).permute(0, 2, 1)[tm1, tn1]
-  q_frag = Q_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, D // WMMA_K, WMMA_K)[wave_m, tm1, lane_n, k_qk]
-  k_frag = KV_lds_k.reshape(WAVES_N, TN, WMMA_N, D // WMMA_K, WMMA_K)[wave_n, tn1, lane_n, k_qk]
-  qk = UOp(Ops.SHAPED_WMMA, dtypes.float, (q_frag, k_frag, S_frag.after(k_qk)), arg=WMMA_ARG)
-  qk_done = S_frag.store(qk).end(tm1, tn1).end(k_qk)
-  S_reg = S_reg.after(qk_done)
-
-  # -- softmax in registers with warp shuffles --
-  S_reg = S_reg.after(S_reg.store(S_reg * SCALE))
-
-  # per-thread local row max over TN=4 elements, then warp reduce across 16 lanes
-  m_ij = UOp.placeholder((TM,), dtypes.float, slot=7, addrspace=AddrSpace.REG)
-  m_ij = m_ij.after(m_ij.after(n_tile).store(m_ij.const_like(-math.inf)))
-  rm2 = UOp.range(TN, 261, AxisType.REDUCE)
-  m_ij = m_ij.after(m_ij.store(m_ij.after(rm2).maximum(S_reg[:, rm2])).end(rm2))
-  # warp reduce max (in-place)
-  ri_w = UOp.range(TM, 270, AxisType.LOOP)
-  m_ij = m_ij.after(m_ij[ri_w].store(warp_reduce_max(m_ij[ri_w], lane)).end(ri_w))
-
-  # compute P = exp(S - m_ij) in S_reg
-  S_reg = S_reg.after(S_reg.store(((S_reg - m_ij.reshape(TM, 1).expand(TM, TN)) * LOG2E).exp2()))
-
-  p_local = UOp.placeholder((TM,), dtypes.float, slot=8, addrspace=AddrSpace.REG)
-  p_local = p_local.after(p_local.after(n_tile).store(p_local.const_like(0)))
-  rp2 = UOp.range(TN, 291, AxisType.REDUCE)
-  p_local = p_local.after(p_local.store(p_local.after(rp2) + S_reg[:, rp2]).end(rp2))
-  ri_ws = UOp.range(TM, 295, AxisType.LOOP)
-  p_sum = p_local.after(p_local[ri_ws].store(warp_reduce_sum(p_local[ri_ws], lane)).end(ri_ws))
-
-  # write P = exp(S - m_ij) to P_lds (reuses slot 0, Q no longer needed)
-  P_lds = QP_lds[:, :BLOCK_N]
-  P_write = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TN, LANES_PER_WAVE_N)
-  P_write = P_write.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TN)
-  # TODO: P_write[tid].store(S_reg.cast(dtypes.half)) — shaped store fails due to RESHAPE(DEFINE_LOCAL) surviving linearization
-  rw1 = UOp.range(TM, 296, AxisType.LOOP)
-  rw2 = UOp.range(TN, 297, AxisType.LOOP)
-  P_store = P_write[tid, rw1, rw2].store(S_reg[rw1, rw2].cast(dtypes.half)).end(rw1, rw2)
-
-  # -- online softmax correction --
-  ri4 = UOp.range(TM, 330, AxisType.LOOP)
-  m_new_val = m_i[ri4].maximum(m_ij[ri4])
-  alpha_val = ((m_i[ri4] - m_new_val) * LOG2E).exp2()
-  beta_val = ((m_ij[ri4] - m_new_val) * LOG2E).exp2()
-  rj4 = UOp.range(TD, 331, AxisType.LOOP)
-  correction = UOp.group(
-    acc[ri4, rj4].store(alpha_val * acc[ri4, rj4]).end(rj4),
-    l_i[ri4].store(alpha_val * l_i[ri4] + beta_val * p_sum[ri4]),
-    m_i[ri4].store(m_new_val),
-  ).end(ri4)
-  acc = acc.after(correction)
-  l_i = l_i.after(correction)
-  m_i = m_i.after(correction)
-
-  # load V into KV_lds (must wait for QK WMMA to finish reading K from KV_lds)
-  V_store = KV_lds.after(qk_done).reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid].store(
-    v[n_tile].reshape(THREADS_PER_BLOCK, ELEMS_PER_THREAD)[tid])
-  pv_barrier = UOp.barrier(UOp.group(P_store, V_store))
-  P_lds = P_lds.after(pv_barrier)
-  KV_lds_v = KV_lds.after(pv_barrier)
-
-  # -- acc += P @ V via WMMA --
-  k_pv = UOp.range(BLOCK_N // WMMA_K, 400, AxisType.REDUCE)
-  tm2 = UOp.range(TM // WMMA_ACC, 401, AxisType.LOOP)
-  tn2 = UOp.range(TD, 402, AxisType.LOOP)
-  acc_frag = acc.reshape(TM // WMMA_ACC, WMMA_ACC, TD).permute(0, 2, 1)[tm2, tn2]
-  p_frag = P_lds.reshape(WAVES_M, TM // WMMA_ACC, WMMA_M, BLOCK_N // WMMA_K, WMMA_K)[wave_m, tm2, lane_n, k_pv]
-  v_frag = KV_lds_v.reshape(WAVES_N, TD, WMMA_N, BLOCK_N // WMMA_K, WMMA_K)[wave_n, tn2, lane_n, k_pv]
-  pv = UOp(Ops.SHAPED_WMMA, dtypes.float, (p_frag, v_frag, acc_frag.after(k_pv)), arg=WMMA_ARG)
-
-  # end KV tile loop
-  n_tile_end = acc_frag.store(pv).end(tm2, tn2).end(k_pv).barrier().end(n_tile)
-  acc = acc.after(n_tile_end)
-  l_i = l_i.after(n_tile_end)
-  m_i = m_i.after(n_tile_end)
-
-  # normalize: acc /= l_i
-  acc = acc.after(acc.store(acc * (1 / l_i).reshape(TM, 1).expand(TM, TD)))
-
-  # store output
-  o = o.reshape(WAVES_M, TM // WMMA_ACC, WMMA_ACC, LANES_PER_WAVE_M, WAVES_N, TD, LANES_PER_WAVE_N)
-  o = o.permute((0, 4, 3, 6, 1, 2, 5)).reshape(THREADS_PER_BLOCK, TM, TD)
-  return o[tid].store(acc).end(wave_m, wave_n, lane).end(block_m, block_bh).sink(arg=KernelInfo(opts_to_apply=()))
-
-if __name__ == "__main__":
-  B, H, N, D = getenv("B", 1), getenv("H", 32), getenv("N", 1024), getenv("D", 64)
-  q = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  k = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  v = Tensor.rand(B, H, N, D).cast(dtypes.half)
-  o = Tensor.empty(B, H, N, D, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(q, k, v)
-
-  q_flat, k_flat, v_flat, o_flat = q.reshape(B*H, N, D), k.reshape(B*H, N, D), v.reshape(B*H, N, D), o.reshape(B*H, N, D)
-  NUM_RUNS = getenv("CNT", 5)
-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(NUM_RUNS):
-      GlobalCounters.reset()
-      tst = Tensor.custom_kernel(o_flat, q_flat, k_flat, v_flat, fxn=amd_flash_attention)[0].realize()
-      ets.append(GlobalCounters.time_sum_s)
-  print(f"best time: {min(ets)*1e3:.2f}ms")
-
-  if getenv("VERIFY", 1):
-    with Context(DEBUG=0):
-      ref = q.float().scaled_dot_product_attention(k.float(), v.float()).reshape(B*H, N, D).realize()
-      err = (ref - tst).square().mean().item()
-    print(f"mean squared error {err}")
-    if err > 1e-2:
-      raise RuntimeError("flash attention is wrong!")
-    else:
-      print("flash attention is correct!")
@@ -1,7 +1,9 @@
-from tinygrad import Tensor, Context, GlobalCounters, dtypes
+import numpy as np
+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.dtype import AddrSpace
-from tinygrad.helpers import DEBUG, getenv
+from tinygrad.helpers import getenv

 N = getenv("N", 4096)
 M = getenv("M", N)
@@ -13,34 +15,60 @@ NUM_RUNS = getenv("CNT", 5)
 # ---------------------------

 WARP_SIZE = 32
-BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 8
-TM, TN = 4, 4
-LANES_PER_WAVE_M, LANES_PER_WAVE_N = 4, 8
-assert N % BLOCK_N == 0 and M % BLOCK_M == 0 and K % BLOCK_K == 0
+
+# 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
+assert N % BLOCK_N == 0, f"N ({N}) must be a multiple of BLOCK_N ({BLOCK_N})"
+assert M % BLOCK_M == 0, f"M ({M}) must be a multiple of BLOCK_M ({BLOCK_M})"
+assert K % BLOCK_K == 0, f"K ({K}) must be a multiple of BLOCK_K ({BLOCK_K})"
+
+# Register tile sizes (per-thread accumulator tile of C)
+TN = 4     # columns per thread
+TM = 4     # rows per thread

 is_kernel5 = getenv("K5", 0)
 THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-WAVES_PER_BLOCK_N = 1 if is_kernel5 else 2
-WAVES_PER_BLOCK_M = THREADS_PER_BLOCK // WARP_SIZE // WAVES_PER_BLOCK_N
-REG_TILES_PER_WAVE_N = BLOCK_N // (WAVES_PER_BLOCK_N * LANES_PER_WAVE_N * TN)
-REG_TILES_PER_WAVE_M = BLOCK_M // (WAVES_PER_BLOCK_M * LANES_PER_WAVE_M * TM)
+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

-assert WAVES_PER_BLOCK_M*REG_TILES_PER_WAVE_M*LANES_PER_WAVE_M*TM == BLOCK_M, "M reshape is wrong"
-assert WAVES_PER_BLOCK_N*REG_TILES_PER_WAVE_N*LANES_PER_WAVE_N*TN == BLOCK_N, "N reshape is wrong"
+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_PER_BLOCK_N = BLOCK_N // WAVE_TILE_N
+WAVES_PER_BLOCK_M = BLOCK_M // WAVE_TILE_M
+assert WAVES_PER_BLOCK_N * WAVES_PER_BLOCK_M == WARPS_PER_BLOCK, "wave grid must match warps/block"
+
+LANES_PER_WAVE_N = 8
+LANES_PER_WAVE_M = 4
+REG_TILES_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_N * TN)
+REG_TILES_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_M * TM)
+assert WAVE_TILE_N % (LANES_PER_WAVE_N * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_N*TN"
+assert WAVE_TILE_M % (LANES_PER_WAVE_M * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_M*TM"

 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, upcast=False):
+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)
-  return dest[*rngs].store(src[*rngs]).end(*rngs)
+  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
+  return dest.after(copy) if set else copy

-def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
+def hand_spec_kernel3():
  # ---------------------------
-  # block indices
+  # block indices & placeholders
  # ---------------------------
  block_id_n = UOp.special(N // BLOCK_N, "gidx0")
  block_id_m = UOp.special(M // BLOCK_M, "gidx1")

+  a = UOp.placeholder((M, K), dtypes.float, slot=1)
+  b = UOp.placeholder((K, N), dtypes.float, slot=2)
+  c = UOp.placeholder((M, N), dtypes.float, slot=0)
+
  # index the output with the globals
  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[block_id_m, :, block_id_n, :]

@@ -76,18 +104,21 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:
  # ---------------------------
  # LOCAL -> REG (per-wave tiles)
  # ---------------------------
-  warp, lane = tid // WARP_SIZE, tid % WARP_SIZE
-  waveIdx, waveIdy = warp % WAVES_PER_BLOCK_N, warp // WAVES_PER_BLOCK_N
-  laneIdx, laneIdy = lane % LANES_PER_WAVE_N, lane // LANES_PER_WAVE_N
-  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M and laneIdy.vmax+1 == LANES_PER_WAVE_M
+  waveIdx = (tid // WARP_SIZE) % WAVES_PER_BLOCK_N
+  waveIdy = (tid // WARP_SIZE) // WAVES_PER_BLOCK_N
+  assert waveIdy.vmax+1 == WAVES_PER_BLOCK_M
+
+  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_N
+  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_N
+  assert laneIdy.vmax+1 == LANES_PER_WAVE_M

  A_col = UOp.placeholder((REG_TILES_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
  A_local_slice = A_local[k, :].reshape(WAVES_PER_BLOCK_M, REG_TILES_PER_WAVE_M, LANES_PER_WAVE_M, TM)[waveIdy, :, laneIdy, :]
-  A_col = A_col.after(copy(A_col, A_local_slice, 300, upcast=True))
+  A_col = copy(A_col, A_local_slice , 300, set=True, upcast=True)

  B_row = UOp.placeholder((REG_TILES_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
  B_local_slice = B_local[k, :].reshape(WAVES_PER_BLOCK_N, REG_TILES_PER_WAVE_N, LANES_PER_WAVE_N, TN)[waveIdx, :, laneIdx, :]
-  B_row = B_row.after(copy(B_row, B_local_slice, 400, upcast=True))
+  B_row = copy(B_row, B_local_slice, 400, set=True, upcast=True)

  # ---------------------------
  # FMA: c_regs += A_col * B_row
@@ -115,29 +146,30 @@ def hand_spec_kernel3(c:UOp, a:UOp, b:UOp) -> UOp:

  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()

-def eval_custom_matmul(fxn, dt=dtypes.float):
-  a = Tensor.randn(M, K, dtype=dt)
-  b = Tensor.randn(K, N, dtype=dt)
-  c = Tensor.empty(M, N, dtype=dtypes.float)
-  with Context(DEBUG=0): Tensor.realize(a, b)
+def test_matmul(sink:UOp, dtype=dtypes.float32, M=M, N=N, K=K):
+  rng = np.random.default_rng()
+  a = Tensor(rng.random((M, K), dtype=np.float32)-0.5, dtype=dtype)
+  b = Tensor(rng.random((K, N), dtype=np.float32)-0.5, dtype=dtype)
+  hc = Tensor.empty(M, N, dtype=dtype)
+  Tensor.realize(a, b, hc)
+
+  ei = ExecItem(sink, [t.uop.buffer for t in [hc, a, b]], prg=get_runner(Device.DEFAULT, sink))

  ets = []
-  with Context(DEBUG=max(2, DEBUG.value), DEVECTORIZE=2 if dt == dtypes.half else 0):
+  with Context(DEBUG=2):
    for _ in range(NUM_RUNS):
-      GlobalCounters.reset()
-      tst = Tensor.custom_kernel(c, a, b, fxn=fxn)[0].realize()
-      ets.append(GlobalCounters.time_sum_s)
+      ets.append(ei.run(wait=True))
  print(f"REAL TFLOPS {M * N * K * 2 / min(ets) * 1e-12:.2f}")

  if getenv("VERIFY", 1):
    GlobalCounters.reset()
    with Context(DEBUG=2):
-      tc = (a.float() @ b.float()).realize()
+      tc = (a @ b).realize()
    with Context(DEBUG=0):
-      err = (tc - tst).square().mean().item()
+      err = (hc - tc).square().mean().item()
    print(f"mean squared error {err}")
-    if err > (1e-2 if dt == dtypes.half else 1e-6):
+    if err > 1e-06:
      raise RuntimeError("matmul is wrong!")

 if __name__ == "__main__":
-  eval_custom_matmul(hand_spec_kernel3)
+  test_matmul(hand_spec_kernel3(), N=N)
@@ -1,14 +1,5 @@
-import atexit, functools
-from tinygrad import Tensor, Device, dtypes
-from tinygrad.dtype import AddrSpace
-from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
-from tinygrad.renderer import Estimates
-from tinygrad.helpers import getenv, all_same, DEBUG
 from tinygrad.runtime.autogen.amd.cdna.ins import *
-
-# ** CDNA4 assembly gemm
-
-WORKGROUP_SIZE = 256
+from tinygrad.dtype import dtypes

 # M0 is encoded with 124 (NULL in RDNA) in CDNA
 M0 = NULL
@@ -2606,121 +2597,3 @@ def build_kernel(batch, M, N, K, dtype):
  k.label('KernelEnd')
  k.emit(s_endpgm())
  return k.finalize()
-
-# ** ASM_GEMM custom kernel
-
-@functools.cache
-def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
-  batch, M, K = A.shape
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
-  gidx = UOp.special(NUM_WG, "gidx0")
-  insts = build_kernel(batch, M, N, K, A.dtype.base)
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
-  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
-                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
-                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-counters = {"used":0, "todos":[]}
-def todo(msg:str) -> bool: counters["todos"].append(msg); return False
-def _asm_gemm_report():
-  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
-  if DEBUG >= 2 and counters["todos"]:
-    from collections import Counter
-    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
-atexit.register(_asm_gemm_report)
-
-def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
-  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
-  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
-  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
-  N = b.shape[1]
-  if isinstance(a.device, tuple):
-    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
-    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
-    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
-    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
-    dname = a.device[0]
-  else: dname = a.device
-  arch = getattr(Device[dname].renderer, "arch", "")
-  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
-  # blacklist slow matmul
-  # TODO: why is this slow?
-  if (M,N,K) == (8192, 2304, 16384): return todo("blacklisted slow matmul")
-  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
-    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
-  return True
-
-# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
-# note: this can be removed after we have GEMM on mixins
-
-def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
-  M, K = A.shape[0]*A.shape[1], A.shape[2]
-  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
-  assert K == K2
-  m = UOp.range(M, 1, AxisType.LOOP)
-  n = UOp.range(N, 2, AxisType.LOOP)
-  k = UOp.range(K, 0, AxisType.REDUCE)
-  mul = (A.flatten().index((m*UOp.const(dtypes.weakint, K)+k))*
-         B.flatten().index((k*UOp.const(dtypes.weakint, N)+n))).cast(dtypes.float32)
-  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
-  store = C.flatten().index((m*UOp.const(dtypes.weakint, N)+n), ptr=True).store(red).end(m, n)
-  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
-
-# ** backward gemm, might use the asm gemm
-
-def custom_gemm_bw(gradient:UOp, kernel:UOp):
-  out, a, b = kernel.src[1:]
-  assert all_same([gradient.device, a.device, b.device, out.device])
-  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
-  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
-  g_t = g_t[:a.shape[0]]
-  grad_a = (g_t @ b_t.T).uop
-  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
-  return (None, grad_a, grad_b)
-
-# ** main gemm function
-
-def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
-  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
-  counters["used"] += 1
-  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
-  if unfold_batch:
-    orig_batch = a.shape[0]
-    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
-  squeeze = a.ndim == 2
-  if squeeze: a = a.unsqueeze(0)
-
-  batch, M, K = a.shape
-  N = b.shape[1]
-  is_multi = isinstance(a.device, tuple)
-  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
-  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
-  n_sharded = is_multi and b.uop.axis == 1
-
-  if is_multi:
-    if n_sharded:
-      out = Tensor(Tensor.invalid(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
-    elif m_sharded:
-      out = Tensor(Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
-    else:
-      out = Tensor(Tensor.invalid(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0),
-                   device=a.device)
-  else:
-    out = Tensor.invalid(batch, M, N, dtype=a.dtype, device=a.device)
-
-  renderer = Device[a.device[0] if is_multi else a.device].renderer
-  dname, arch = renderer.device, getattr(renderer, "arch", "")
-  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
-    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
-  else:
-    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
-  if k_sharded: out = out.sum(0)
-  out = out.squeeze(0) if squeeze else out
-  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
-  return out
@@ -0,0 +1,122 @@
+import atexit, functools
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.dtype import AddrSpace
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
+from tinygrad.renderer import Estimates
+from tinygrad.helpers import getenv, all_same, DEBUG
+from extra.gemm.asm.cdna.asm import build_kernel, TILE_M, TILE_N, TILE_K, NUM_WG
+
+# ** CDNA4 assembly gemm
+
+WORKGROUP_SIZE = 256
+
+@functools.cache
+def custom_asm_gemm(C:UOp, A:UOp, B:UOp, dname:str) -> UOp:
+  batch, M, K = A.shape
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  lidx = UOp.special(WORKGROUP_SIZE, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+  insts = build_kernel(batch, M, N, K, A.dtype.base)
+  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=133_120, addrspace=AddrSpace.LOCAL), (), 'lds')
+  sink = UOp.sink(C.base, A.base, B.base, lds, lidx, gidx,
+                  arg=KernelInfo(name=f"gemm_{batch}_{M}_{N}_{K}", estimates=Estimates(ops=2*batch*M*N*K, mem=(batch*M*K + K*N + batch*M*N)*2)))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname),
+                                UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
+
+counters = {"used":0, "todos":[]}
+def todo(msg:str) -> bool: counters["todos"].append(msg); return False
+def _asm_gemm_report():
+  print(f'asm_gemm: {counters["used"]} used, {len(counters["todos"])} not used')
+  if DEBUG >= 2 and counters["todos"]:
+    from collections import Counter
+    for msg, cnt in Counter(counters["todos"]).most_common(): print(f'  {cnt:3d}x {msg}')
+atexit.register(_asm_gemm_report)
+
+def can_use_asm_gemm(a:Tensor, b:Tensor) -> bool:
+  if a.dtype != b.dtype: return todo(f"dtypes must match {a.dtype} != {b.dtype}")
+  if a.dtype not in {dtypes.bfloat16, dtypes.float16}: return todo(f"only bfloat16/float16, got {a.dtype}")
+  batch, M, K = (1, *a.shape) if a.ndim == 2 else a.shape
+  N = b.shape[1]
+  if isinstance(a.device, tuple):
+    if a.ndim == 2 and a.uop.axis == 0 and b.uop.axis is None: M //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis == 1 and b.uop.axis == 0: K //= len(a.device)
+    elif a.ndim == 2 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 0 and b.uop.axis is None: batch //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis is None and b.uop.axis == 1: N //= len(a.device)
+    elif a.ndim == 3 and a.uop.axis == 2 and b.uop.axis == 0: K //= len(a.device)
+    else: return todo(f"sharding mismatch a.ndim={a.ndim} a.uop.axis={a.uop.axis} b.uop.axis={b.uop.axis}")
+    dname = a.device[0]
+  else: dname = a.device
+  arch = getattr(Device[dname].renderer, "arch", "")
+  if batch not in {1, 2}: return todo(f"GEMM batch size {batch}")
+  if (M % TILE_M != 0 or N % TILE_N != 0 or K % TILE_K != 0) and arch == "gfx950":
+    return todo(f"GEMM shape ({M},{N},{K}) not a multiple of ({TILE_M},{TILE_N},{TILE_K})")
+  return True
+
+# ** UOp gemm to test Tensor.custom_kernel multi and backward correctness on non cdna4
+# note: this can be removed after we have GEMM on mixins
+
+def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
+  M, K = A.shape[0]*A.shape[1], A.shape[2]
+  K2, N = B.shape[(1 if B.ndim == 3 else 0):]
+  assert K == K2
+  m = UOp.range(M, 1, AxisType.LOOP)
+  n = UOp.range(N, 2, AxisType.LOOP)
+  k = UOp.range(K, 0, AxisType.REDUCE)
+  mul = (A.index((m*UOp.const(dtypes.index, K)+k))*B.index((k*UOp.const(dtypes.index, N)+n))).cast(dtypes.float32)
+  red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype.base)
+  store = C.index((m*UOp.const(dtypes.index, N)+n), ptr=True).store(red).end(m, n)
+  return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
+
+# ** backward gemm, might use the asm gemm
+
+def custom_gemm_bw(gradient:UOp, kernel:UOp):
+  out, a, b = kernel.src[1:]
+  assert all_same([gradient.device, a.device, b.device, out.device])
+  a_t, b_t, g_t = Tensor(a, device=a.device), Tensor(b, device=a.device), Tensor(gradient, device=a.device)
+  # TODO: this needs to be cleaned up and done properly, the batch dim of grad and a multi need to align
+  g_t = g_t[:a.shape[0]]
+  grad_a = (g_t @ b_t.T).uop
+  grad_b = (a_t.permute(2, 0, 1).reshape(a_t.shape[2], -1) @ g_t.reshape(-1, g_t.shape[-1])).uop
+  return (None, grad_a, grad_b)
+
+# ** main gemm function
+
+def asm_gemm(a:Tensor, b:Tensor) -> Tensor:
+  assert can_use_asm_gemm(a, b), f"{counters['todos'][-1]}"
+  counters["used"] += 1
+  unfold_batch = a.ndim == 3 and isinstance(a.device, tuple) and a.uop.axis == 2 and b.uop.axis == 0
+  if unfold_batch:
+    orig_batch = a.shape[0]
+    a = a.reshape(a.shape[0]*a.shape[1], a.shape[2])
+  squeeze = a.ndim == 2
+  if squeeze: a = a.unsqueeze(0)
+
+  batch, M, K = a.shape
+  N = b.shape[1]
+  is_multi = isinstance(a.device, tuple)
+  if (k_sharded:=is_multi and a.uop.axis == 2): K //= len(a.device)
+  if (m_sharded:=is_multi and a.uop.axis == 1): M //= len(a.device)
+  n_sharded = is_multi and b.uop.axis == 1
+
+  if is_multi:
+    if n_sharded:
+      out = Tensor(Tensor.empty(batch, M, N//len(a.device), dtype=a.dtype, device=a.device).uop.multi(2), device=a.device)
+    elif m_sharded:
+      out = Tensor(Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device).uop.multi(1), device=a.device)
+    else:
+      out = Tensor(Tensor.empty(batch//len(a.device) if a.uop.axis==0 else batch, M, N, dtype=a.dtype, device=a.device).uop.multi(0), device=a.device)
+  else:
+    out = Tensor.empty(batch, M, N, dtype=a.dtype, device=a.device)
+
+  renderer = Device[a.device[0] if is_multi else a.device].renderer
+  dname, arch = renderer.device, getattr(renderer, "arch", "")
+  if arch.startswith("gfx950") and getenv("USE_ASM", 1):
+    out = Tensor.custom_kernel(out, a, b, fxn=functools.partial(custom_asm_gemm, dname=dname), grad_fxn=custom_gemm_bw)[0]
+  else:
+    out = Tensor.custom_kernel(out, a, b, fxn=custom_uop_gemm, grad_fxn=custom_gemm_bw)[0]
+  if k_sharded: out = out.sum(0)
+  out = out.squeeze(0) if squeeze else out
+  if unfold_batch: out = out.reshape(orig_batch, -1, out.shape[-1])
+  return out
@@ -0,0 +1,76 @@
+.text
+.section	.text.
+.global	gemm
+.p2align	8
+.type	gemm,@function
+
+gemm:
+INSTRUCTIONS
+
+.section .rodata,"a",@progbits
+.p2align 6, 0x0
+.amdhsa_kernel gemm
+  # basic memory requirements
+  .amdhsa_group_segment_fixed_size 30336
+  .amdhsa_private_segment_fixed_size 0
+  .amdhsa_kernarg_size 32
+  # register usage (RSRC1)
+  .amdhsa_next_free_vgpr 256
+  .amdhsa_next_free_sgpr 100
+  # workgroup / workitem IDs (RSRC2)
+  .amdhsa_system_sgpr_workgroup_id_x 1
+  .amdhsa_system_sgpr_workgroup_id_y 1
+  .amdhsa_system_sgpr_workgroup_id_z 1
+  # user SGPRs: kernarg ptr in s[0:1]
+  .amdhsa_user_sgpr_kernarg_segment_ptr 1
+  .amdhsa_user_sgpr_count 2
+  # gfx10+ / gfx11 specifics (RSRC1[29..31])
+  .amdhsa_wavefront_size32 1
+  .amdhsa_workgroup_processor_mode 1
+  .amdhsa_memory_ordered 1
+  .amdhsa_forward_progress 1
+  # misc for gfx11
+  .amdhsa_dx10_clamp 1
+  .amdhsa_ieee_mode 1
+  .amdhsa_uses_dynamic_stack 0
+.end_amdhsa_kernel
+
+.amdgpu_metadata
+---
+amdhsa.kernels:
+  - .args:
+      - .address_space:  generic
+        .name:           C
+        .offset:         0
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           A
+        .offset:         8
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+      - .address_space:  generic
+        .name:           B
+        .offset:         16
+        .size:           8
+        .value_kind:     global_buffer
+        .value_type:     f16
+    .group_segment_fixed_size: 30336
+    .kernarg_segment_align: 8
+    .kernarg_segment_size: 32
+    .max_flat_workgroup_size: 128
+    .name:           gemm
+    .private_segment_fixed_size: 0
+    .sgpr_count:     70
+    .sgpr_spill_count: 0
+    .symbol:         gemm.kd
+    .vgpr_count:     256
+    .vgpr_spill_count: 0
+    .wavefront_size: 32
+amdhsa.version:
+  - 1
+  - 1
+...
+.end_amdgpu_metadata
@@ -0,0 +1,30 @@
+import math, pathlib
+
+from tinygrad import Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo
+
+from extra.gemm.amd_uop_matmul import test_matmul
+
+N = 4096
+TN = 96
+THREADS_PER_WG = 128
+NUM_WG = math.ceil(N / TN) * math.ceil(N / TN)
+
+dname:str = Device.DEFAULT
+template:str = (pathlib.Path(__file__).parent/"template.s").read_text()
+
+def asm_kernel() -> UOp:
+  lidx = UOp.special(THREADS_PER_WG, "lidx0")
+  gidx = UOp.special(NUM_WG, "gidx0")
+
+  a = UOp.placeholder((N*N,), dtypes.half, slot=1)
+  b = UOp.placeholder((N*N,), dtypes.half, slot=2)
+  c = UOp.placeholder((N*N,), dtypes.half, slot=0)
+
+  src = template.replace("INSTRUCTIONS", (pathlib.Path(__file__).parent/"gemm.s").read_text())
+
+  sink = UOp.sink(a, b, c, lidx, gidx, arg=KernelInfo(name="gemm"))
+  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg=dname), UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=src)))
+
+if __name__ == "__main__":
+  test_matmul(asm_kernel(), dtype=dtypes.half, N=N)
@@ -0,0 +1,179 @@
+# unpack the complete kernel descriptor of an amdgpu ELF
+# https://rocm.docs.amd.com/projects/llvm-project/en/latest/LLVM/llvm/html/AMDGPUUsage.html#code-object-v3-kernel-descriptor
+import struct, pathlib, sys
+from tinygrad.runtime.support.elf import elf_loader
+
+def bits(x, lo, hi): return (x >> lo) & ((1 << (hi - lo + 1)) - 1)
+def assert_zero(x, lo, hi): assert bits(x, lo, hi) == 0
+
+with open(sys.argv[1], "rb") as f:
+  lib = f.read()
+
+image, sections, relocs = elf_loader(lib)
+rodata_entry = next((sh.header.sh_addr for sh in sections if sh.name == ".rodata"))
+
+# rodata is exactly 64 bytes
+kd = image[rodata_entry:rodata_entry+64]
+desc = int.from_bytes(kd, byteorder="little")
+
+group_segment_fixed_size = bits(desc, 0, 31)
+private_segment_fixed_size = bits(desc, 32, 63)
+kernarg_size = bits(desc, 64, 95)
+reserved_127_96 = bits(desc, 96, 127)
+assert reserved_127_96 == 0
+
+print("GROUP_SEGMENT_FIXED_SIZE:", group_segment_fixed_size)
+print("PRIVATE_SEGMENT_FIXED_SIZE:", private_segment_fixed_size)
+print("KERNARG_SIZE:", kernarg_size)
+print("RESERVED 127:96:", reserved_127_96)
+
+entry_off = bits(desc, 128, 191)
+
+# sign-extend manually if needed
+if entry_off & (1 << 63):
+  entry_off -= 1 << 64
+
+print("KERNEL_CODE_ENTRY_BYTE_OFFSET:", entry_off)
+
+kd_addr = 0x1840
+entry_addr = kd_addr + entry_off
+
+print("Computed entry address: 0x%016x" % entry_addr)
+print("256B aligned:", entry_addr % 256 == 0)
+
+pgm_rsrc3 = bits(desc, 352, 383)
+pgm_rsrc1 = bits(desc, 384, 415)
+pgm_rsrc2 = bits(desc, 416, 447)
+
+print("COMPUTE_PGM_RSRC3: 0x%08x" % pgm_rsrc3)
+print("COMPUTE_PGM_RSRC1: 0x%08x" % pgm_rsrc1)
+print("COMPUTE_PGM_RSRC2: 0x%08x" % pgm_rsrc2)
+
+# rsrc 3 (gfx950)
+
+accum_offset_raw = bits(pgm_rsrc3, 0, 5)
+assert_zero(pgm_rsrc3, 6, 15)
+tg_split = bits(pgm_rsrc3, 16, 16)
+accum_offset_vgprs = (accum_offset_raw + 1) * 4
+print("RSRC3.ACCUM_OFFSET (AccVGPR index):", accum_offset_vgprs)
+print("RSRC3.TG_SPLIT:", tg_split)
+
+# rsrc 1
+
+vgpr_gran = bits(pgm_rsrc1, 0, 5)
+sgpr_gran = bits(pgm_rsrc1, 6, 9)
+assert_zero(pgm_rsrc1, 27, 28)
+
+# NOTE: this is vgprs + agprs
+vgprs_used = (vgpr_gran + 1) * 8
+assert 0 <= vgprs_used <= 512
+
+k = sgpr_gran // 2
+sgprs_used = (k + 1) * 16
+
+print("RSRC1.VGPRS:", vgprs_used)
+print("RSRC1.SGPRS:", sgprs_used)
+
+assert_zero(pgm_rsrc1, 10, 11)
+
+float_round_mode_32 = bits(pgm_rsrc1, 12, 13)
+float_round_mode_16_64 = bits(pgm_rsrc1, 15, 14)
+float_denorm_mode_32 = bits(pgm_rsrc1, 16, 17)
+float_denorm_mode_16_64 = bits(pgm_rsrc1, 18, 19)
+
+priv = bits(pgm_rsrc1, 20, 20)
+assert priv == 0
+enable_dx10_clamp_wg_rr_en = bits(pgm_rsrc1, 21, 21)
+debug_mode = bits(pgm_rsrc1, 22, 22)
+enable_ieee_mode = bits(pgm_rsrc1, 23, 23)
+bulky = bits(pgm_rsrc1, 24, 24)
+assert bulky == 0
+cdbg_user = bits(pgm_rsrc1, 25, 25)
+assert cdbg_user == 0
+fp16_ovfl = bits(pgm_rsrc1, 26, 26)
+assert_zero(pgm_rsrc1, 27, 28)  # reserved
+assert_zero(pgm_rsrc1, 29, 29)  # WGP_MODE (reserved on gfx9)
+assert_zero(pgm_rsrc1, 30, 30)  # MEM_ORDERED (reserved on gfx9)
+assert_zero(pgm_rsrc1, 31, 31)  # FWD_PROGRESS (reserved on gfx9)
+
+# rsrc 2
+
+enable_private_segment = bits(pgm_rsrc2, 0, 0)  # SCRATCH_EN
+user_sgpr_count = bits(pgm_rsrc2, 1, 5)         # USER_SGPR
+enable_trap_handler = bits(pgm_rsrc2, 6, 6)     # TRAP_PRESENT (must be 0 here)
+assert enable_trap_handler == 0
+
+enable_sgpr_workgroup_id_x = bits(pgm_rsrc2, 7, 7)
+enable_sgpr_workgroup_id_y = bits(pgm_rsrc2, 8, 8)
+enable_sgpr_workgroup_id_z = bits(pgm_rsrc2, 9, 9)
+enable_sgpr_workgroup_info = bits(pgm_rsrc2, 10, 10)
+
+enable_vgpr_workitem_id = bits(pgm_rsrc2, 11, 12)  # TIDIG_CMP_CNT enum (0..3)
+
+enable_exception_address_watch = bits(pgm_rsrc2, 13, 13)
+assert enable_exception_address_watch == 0
+enable_exception_memory = bits(pgm_rsrc2, 14, 14)
+assert enable_exception_memory == 0
+
+granulated_lds_size = bits(pgm_rsrc2, 15, 23)
+assert granulated_lds_size == 0  # spec: must be 0; CP uses dispatch packet rounding
+
+enable_exception_fp_invalid = bits(pgm_rsrc2, 24, 24)
+enable_exception_fp_denorm_src = bits(pgm_rsrc2, 25, 25)
+enable_exception_fp_div0 = bits(pgm_rsrc2, 26, 26)
+enable_exception_fp_overflow = bits(pgm_rsrc2, 27, 27)
+enable_exception_fp_underflow = bits(pgm_rsrc2, 28, 28)
+enable_exception_fp_inexact = bits(pgm_rsrc2, 29, 29)
+enable_exception_int_div0 = bits(pgm_rsrc2, 30, 30)
+
+assert_zero(pgm_rsrc2, 31, 31)
+
+print("RSRC2.ENABLE_PRIVATE_SEGMENT:", enable_private_segment)
+print("RSRC2.USER_SGPR_COUNT:", user_sgpr_count)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_X:", enable_sgpr_workgroup_id_x)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Y:", enable_sgpr_workgroup_id_y)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_ID_Z:", enable_sgpr_workgroup_id_z)
+print("RSRC2.ENABLE_SGPR_WORKGROUP_INFO:", enable_sgpr_workgroup_info)
+print("RSRC2.ENABLE_VGPR_WORKITEM_ID (enum):", enable_vgpr_workitem_id)
+
+print("RSRC2.EXC_FP_INVALID:", enable_exception_fp_invalid)
+print("RSRC2.EXC_FP_DENORM_SRC:", enable_exception_fp_denorm_src)
+print("RSRC2.EXC_FP_DIV0:", enable_exception_fp_div0)
+print("RSRC2.EXC_FP_OVERFLOW:", enable_exception_fp_overflow)
+print("RSRC2.EXC_FP_UNDERFLOW:", enable_exception_fp_underflow)
+print("RSRC2.EXC_FP_INEXACT:", enable_exception_fp_inexact)
+print("RSRC2.EXC_INT_DIV0:", enable_exception_int_div0)
+
+# user sgprs
+
+enable_sgpr_private_segment_buffer = bits(desc, 448, 448)
+enable_sgpr_dispatch_ptr = bits(desc, 449, 449)
+enable_sgpr_queue_ptr = bits(desc, 450, 450)
+enable_sgpr_kernarg_segment_ptr = bits(desc, 451, 451)
+enable_sgpr_dispatch_id = bits(desc, 452, 452)
+enable_sgpr_flat_scratch_init = bits(desc, 453, 453)
+enable_sgpr_private_segment_size = bits(desc, 454, 454)
+
+assert_zero(desc, 455, 457)
+
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER:", enable_sgpr_private_segment_buffer)
+print("DESC.ENABLE_SGPR_DISPATCH_PTR:", enable_sgpr_dispatch_ptr)
+print("DESC.ENABLE_SGPR_QUEUE_PTR:", enable_sgpr_queue_ptr)
+print("DESC.ENABLE_SGPR_KERNARG_SEGMENT_PTR:", enable_sgpr_kernarg_segment_ptr)
+print("DESC.ENABLE_SGPR_DISPATCH_ID:", enable_sgpr_dispatch_id)
+print("DESC.ENABLE_SGPR_FLAT_SCRATCH_INIT:", enable_sgpr_flat_scratch_init)
+print("DESC.ENABLE_SGPR_PRIVATE_SEGMENT_SIZE:", enable_sgpr_private_segment_size)
+
+assert_zero(desc, 458, 459)
+
+uses_dynamic_stack = bits(desc, 459, 460)
+print("DESC.USES_DYNAMIC_STACK:", uses_dynamic_stack)
+
+# gfx950 only
+assert_zero(desc, 460, 463)
+kernarg_preload_spec_length = bits(desc, 464, 470)
+print("DESC.KERNARG_PRELOAD_SPEC_LENGTH:", kernarg_preload_spec_length)
+kernarg_preload_spec_offset = bits(desc, 471, 479)
+print("DESC.KERNARG_PRELOAD_SPEC_OFFSET:", kernarg_preload_spec_offset)
+
+assert_zero(desc, 480, 511)
@@ -1,46 +0,0 @@
-#!/usr/bin/env python3
-import os, sys, time
-from tinygrad.runtime.support.system import RemotePCIDevice
-
-LAT_N_RUNS = 500
-THROUGHPUT_N_RUNS = 8
-SIZES = [4, 1 << 10, 8 << 20]
-
-if __name__ == "__main__":
-  os.environ["REMOTE"] = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("REMOTE", "127.0.0.1:6667")
-
-  # choose any amd/nv gpu.
-  devs = RemotePCIDevice.remote_list(0x1002, ((0, (0,)),), 0) or RemotePCIDevice.remote_list(0x10de, ((0, (0,)),), 0x03)
-  if not devs: raise RuntimeError("no GPU found on remote")
-
-  pci = RemotePCIDevice("BN", devs[0])
-  print(f"connected to {os.environ['REMOTE']}, device: {devs[0]}\n")
-
-  # ping (minimal server round-trip, no device I/O)
-  from tinygrad.runtime.support.system import RemoteCmd
-  sock = pci.sock
-  for _ in range(10): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
-  st = time.perf_counter()
-  for _ in range(LAT_N_RUNS): RemotePCIDevice._rpc(sock, 0, RemoteCmd.PING)
-  ping_lat = (time.perf_counter() - st) / LAT_N_RUNS
-  print(f"PING latency: {ping_lat*1e6:.1f} us ({1/ping_lat:,.0f} ops/sec)\n")
-
-  # throughput
-  sysmem, _ = pci.alloc_sysmem(max(SIZES))
-  print(f"{'size':>10s}  {'write MB/s':>10s}  {'read MB/s':>10s}")
-  for sz in SIZES:
-    data = b'\x01' * sz
-
-    for _ in range(5): sysmem[0:sz] = data
-    st = time.perf_counter()
-    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz] = data
-    pci.read_config(0, 4) # flush, since writes are posted
-    w = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
-
-    for _ in range(5): sysmem[0:sz]
-    st = time.perf_counter()
-    for _ in range(THROUGHPUT_N_RUNS): sysmem[0:sz]
-    r = (time.perf_counter() - st) / THROUGHPUT_N_RUNS
-
-    sfx, div = [('B',1),('K',1<<10),('M',1<<20)][[sz>=1<<10,sz>=1<<20,sz>=1<<30].count(True)]
-    print(f"{sz/div:>9.4g}{sfx}  {sz/w/1e6:>10.1f}  {sz/r/1e6:>10.1f}")
@@ -1,103 +0,0 @@
-#!/usr/bin/env python3
-import socket, struct, sys
-from tinygrad.runtime.support.system import PCIDevice, RemoteCmd, System
-from tinygrad.helpers import DEBUG, OSX
-
-def resp(resp0=0, resp1=0, status=0): return struct.pack('<BQQ', status, resp0, resp1)
-def resp_err(msg): return struct.pack('<BQQ', 1, len(err:=msg.encode()), 0) + err
-
-discovered_devices: list[str] = []
-opened_devices: dict[int, PCIDevice] = {}
-mapped_bars: dict[tuple[int, int], object] = {}
-sysmem_allocs: list[tuple] = []
-
-def handle(conn, cmd, dev_id, bar, arg0, arg1, arg2):
-  if cmd == RemoteCmd.PING:
-    return conn.sendall(resp())
-
-  if cmd == RemoteCmd.PROBE:
-    payload = conn.recv(arg1, socket.MSG_WAITALL) if arg1 > 0 else b""
-    filter_devices: dict[int, list[int]] = {}
-    for i in range(0, len(payload), 8):
-      mask, dev = struct.unpack('<II', payload[i:i+8])
-      filter_devices.setdefault(mask, []).append(dev)
-    base_class = None if arg0 == 0 else int(arg0)
-    devs = System.list_devices(arg2, tuple([(x, tuple(y)) for x,y in filter_devices.items()]), base_class)
-    for p in devs:
-      if p not in discovered_devices: discovered_devices.append(p)
-    data = "\n".join(f"{p[1]}:{discovered_devices.index(p)}" for p in devs).encode()
-    return conn.sendall(resp(len(data), len(devs)) + data)
-
-  # lazy device open
-  if dev_id not in opened_devices:
-    if dev_id >= len(discovered_devices): raise RuntimeError(f"device {dev_id} not probed")
-    cl, pcibus = discovered_devices[dev_id]
-    opened_devices[dev_id] = cl("SV", pcibus)
-  pci_dev = opened_devices[dev_id]
-
-  if cmd == RemoteCmd.MAP_BAR:
-    if (dev_id, bar) not in mapped_bars: mapped_bars[(dev_id, bar)] = pci_dev.map_bar(bar)
-    conn.sendall(resp(*pci_dev.bar_info(bar)))
-  elif cmd == RemoteCmd.CFG_READ:
-    conn.sendall(resp(pci_dev.read_config(arg0, arg1)))
-  elif cmd == RemoteCmd.CFG_WRITE:
-    pci_dev.write_config(arg0, arg2, arg1)
-    conn.sendall(resp())
-  elif cmd == RemoteCmd.RESIZE_BAR:
-    pci_dev.resize_bar(bar)
-    conn.sendall(resp())
-  elif cmd == RemoteCmd.RESET:
-    pci_dev.reset()
-    conn.sendall(resp())
-  elif cmd == RemoteCmd.MMIO_READ:
-    bar_view = mapped_bars[(dev_id, bar)]
-    if arg0 % 4 == 0 and arg1 == 4: conn.sendmsg([resp(arg1), struct.pack(f'<{arg1 // 4}I', bar_view.view(arg0, arg1, fmt='I')[0])])
-    else: conn.sendmsg([resp(arg1), bar_view[arg0:arg0+arg1]])
-  elif cmd == RemoteCmd.MMIO_WRITE:
-    data = conn.recv(arg1, socket.MSG_WAITALL)
-    bar_view = mapped_bars[(dev_id, bar)]
-    if arg0 % 4 == 0 and arg1 == 4: bar_view.view(arg0, arg1, fmt='I')[0] = struct.unpack(f'<{arg1 // 4}I', data)[0]
-    else: bar_view[arg0:arg0+arg1] = data
-  elif cmd == RemoteCmd.MAP_SYSMEM:
-    memview, paddrs = pci_dev.alloc_sysmem(arg0, contiguous=bool(arg1))
-    sysmem_allocs.append((memview, paddrs))
-    paddrs_bytes = struct.pack(f'<{len(paddrs)}Q', *paddrs)
-    conn.sendall(resp(len(paddrs_bytes), len(sysmem_allocs) - 1) + paddrs_bytes)
-  elif cmd == RemoteCmd.SYSMEM_READ:
-    conn.sendmsg([resp(arg1), sysmem_allocs[bar][0][arg0:arg0+arg1]])
-  elif cmd == RemoteCmd.SYSMEM_WRITE:
-    sysmem_allocs[bar][0][arg0:arg0+arg1] = conn.recv(arg1, socket.MSG_WAITALL)
-  else: raise RuntimeError(f"unknown command {cmd}")
-
-def serve(conn:socket.socket):
-  REQ = '<BIIQQQ'
-  while True:
-    hdr = conn.recv(struct.calcsize(REQ), socket.MSG_WAITALL)
-    if len(hdr) < struct.calcsize(REQ): raise ConnectionError("client disconnected")
-    cmd, dev_id, bar, arg0, arg1, arg2 = struct.unpack(REQ, hdr)
-    if DEBUG >= 4: print(f"cmd={RemoteCmd(cmd).name} dev={dev_id} bar={bar} arg0={arg0:#x} arg1={arg1:#x} arg2={arg2:#x}")
-    try: handle(conn, cmd, dev_id, bar, arg0, arg1, arg2)
-    except ConnectionError: raise
-    except Exception as e:
-      if cmd in {RemoteCmd.MMIO_WRITE, RemoteCmd.SYSMEM_WRITE}: raise ConnectionError(f"write failed: {e}")
-      print(f"ERROR: {e}")
-      conn.sendall(resp_err(str(e)))
-
-if __name__ == "__main__":
-  if not OSX: System.reserve_hugepages(128) # for sysmem allocations
-
-  port = int(sys.argv[1]) if len(sys.argv) > 1 else 6667
-  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-  server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-  server.bind(("0.0.0.0", port))
-  server.listen(1)
-  s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-  try: s.connect(("8.8.8.8", 80)); ip = s.getsockname()[0]
-  finally: s.close()
-  print(f"listening on {ip}:{port}")
-  while True:
-    conn, addr = server.accept()
-    conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
-    for bt in [socket.SO_SNDBUF, socket.SO_RCVBUF]: conn.setsockopt(socket.SOL_SOCKET, bt, 64 << 20)
-    try: serve(conn)
-    except ConnectionError: print("disconnected")
@@ -1,31 +0,0 @@
-#!/bin/sh
-install_loc="$HOME/.local/bin"
-docker build --platform=linux/amd64 -t rocm-hipcc:7.2 - <<'EOF'
-FROM ubuntu:22.04
-ENV DEBIAN_FRONTEND=noninteractive
-ENV TZ=Etc/UTC
-RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates gnupg tzdata && \
-  wget https://repo.radeon.com/amdgpu-install/7.2/ubuntu/jammy/amdgpu-install_7.2.70200-1_all.deb && \
-  apt-get install -y ./amdgpu-install_7.2.70200-1_all.deb && \
-  amdgpu-install -y --usecase=rocm --no-dkms --no-32 && \
-  rm -rf /var/lib/apt/lists/*
-ENV PATH=/opt/rocm/bin:$PATH
-EOF
-
-mkdir -p "$install_loc"
-tee "$install_loc/hipccshim" >/dev/null <<'EOF'
-#!/bin/sh
-set -eu
-cname="rocm-hipcc-persistent"
-if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
-  docker rm -f "$cname" 2>/dev/null || true
-  docker run -d --platform=linux/amd64 --name "$cname" \
-    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
-    rocm-hipcc:7.2 sleep 300 >/dev/null
-fi
-exec docker exec "$cname" "$(basename "$0")" "$@"
-EOF
-chmod +x "$install_loc/hipccshim"
-for t in hipcc hipconfig; do
-  ln -sf "$install_loc/hipccshim" "$install_loc/$t"
-done
@@ -1,9 +1,9 @@
 #!/bin/sh
 install_loc="$HOME/.local/bin"
-docker build --platform=linux/arm64 -t cuda-nvcc:12.8 - <<'EOF'
+docker build --platform=linux/amd64 -t cuda-nvcc:12.8 - <<'EOF'
 FROM ubuntu:22.04
 RUN apt-get update && apt-get install -y --no-install-recommends wget ca-certificates && \
-  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/sbsa/cuda-keyring_1.1-1_all.deb && \
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
  dpkg -i cuda-keyring_1.1-1_all.deb && \
  apt-get update && apt-get install -y --no-install-recommends cuda-nvcc-12-8 cuda-nvdisasm-12-8 cuda-cuobjdump-12-8 && rm -rf /var/lib/apt/lists/*
 ENV PATH=/usr/local/cuda/bin:$PATH
@@ -13,14 +13,10 @@ mkdir -p "$install_loc"
 tee "$install_loc/nvccshim" >/dev/null <<'EOF'
 #!/bin/sh
 set -eu
-cname="cuda-nvcc-persistent"
-if ! docker inspect --format='{{.State.Running}}' "$cname" 2>/dev/null | grep -q true; then
-  docker rm -f "$cname" 2>/dev/null || true
-  docker run -d --platform=linux/arm64 --name "$cname" \
-    -v /var/folders:/var/folders -v "$HOME":"$HOME" \
-    cuda-nvcc:12.8 sleep 300 >/dev/null
-fi
-exec docker exec "$cname" "$(basename "$0")" "$@"
+# assume the final arg is the input path
+# mount it so that container can read it
+dir=$(dirname "${@: -1}")
+exec docker run --rm --platform=linux/amd64 -v "$dir":"$dir" cuda-nvcc:12.8 "$(basename "$0")" "$@"
 EOF
 chmod +x "$install_loc/nvccshim"
 for t in nvcc nvdisasm; do
@@ -0,0 +1,148 @@
+#!/usr/bin/env python3
+# Run all ALU and memory instructions in the ISA
+import functools, inspect
+from enum import Enum
+from tinygrad import Tensor, Device, dtypes
+from tinygrad.uop.ops import UOp, Ops, KernelInfo, AddrSpace
+from tinygrad.renderer.amd.dsl import Inst, Reg, OPERANDS, SrcField, VGPRField, SGPRField, SSrcField, SBaseField, AlignedSGPRField, BitField
+from tinygrad.renderer.amd.dsl import FixedBitField, EnumBitField, s, v, NULL, VCC_LO
+from extra.gemm.amd_asm_matmul import Kernel
+
+# skip instructions that mutate wave state (PC, EXEC, allocations, signals)
+SKIP = {"S_SETPC_B64", "S_SWAPPC_B64", "S_RFE_B64", "S_BARRIER_SIGNAL_ISFIRST", "S_GET_BARRIER_STATE", "S_ALLOC_VGPR", "S_SLEEP_VAR", "S_GETPC_B64",
+        "S_SENDMSG_RTN_B32", "S_SENDMSG_RTN_B64"}
+# skip barriers, s_waits, wrap level atomics, and ray tracing (bvh)
+SKIP_SUBSTR = ["SAVEEXEC", "CMPX", "WREXEC", "MOVREL", "ATOMIC", "S_BUFFER_", "S_ATC_PROBE", "BARRIER", "S_WAITCNT", "BVH",
+               "DS_CMPSTORE_RTN", "DS_WRAP_RTN_B32", "DS_ORDERED_COUNT", "DS_GWS", "GS_REG", "GLOBAL_LOAD_LDS", "GLOBAL_STORE_BLOCK"]
+
+ALU_FORMATS = {"VOP1", "VOP1_LIT", "VOP1_SDST", "VOP2", "VOP2_LIT", "VOP3", "VOP3_SDST", "VOP3SD", "VOP3P", "VOP3P_MFMA", "VOP3PX2",
+               "VOPC", "SOP1", "SOP1_LIT", "SOP2", "SOP2_LIT", "SOPC", "SOPC_LIT", "SOPK", "SOPK_LIT", "VINTERP"}
+# intentionally not testing scratch memory ops
+MEM_FORMATS = {"VGLOBAL", "GLOBAL", "SMEM", "DS"}
+
+def should_skip(op:Enum) -> bool: return (name:=op.name) in SKIP or any(sub in name for sub in SKIP_SUBSTR)
+
+# ** named register assignments
+
+# ALU operands
+ALU_VGPR_STRIDE = 16          # v[0], v[16], v[32], ... per ALU operand slot
+ALU_SGPR_STRIDE = 4           # s[0], s[4], s[8], ... per ALU operand slot
+
+# memory address registers
+S_KERNARG_PTR = (0, 1)
+S_BUF_PTR = (2, 3)
+V_VADDR = (0, 1)
+V_DS_ADDR = 0
+
+# memory data registers
+MEM_VGPR_BASE = 32            # v[32], v[48], ... for vdst/vdata/vsrc
+MEM_VGPR_STRIDE = 16          # spacing between memory data vgpr slots
+MEM_SGPR_BASE = 8             # s[8], s[10], ... for SMEM sdata
+MEM_SGPR_STRIDE = 2           # spacing between memory data sgpr slots
+
+# ** create an ALU instruction based on the operands
+
+def create_alu_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
+  inst_cls, operands, slot = builder.func, OPERANDS[op], 0
+  kwargs:dict[str, Reg|int] = {}
+  for name, field in inst_cls._fields:
+    if isinstance(field, (FixedBitField, EnumBitField)): continue
+    nregs = max(1, operands[name][1] // 32) if name in operands else 1
+    is_sreg = name in operands and "SREG" in str(operands[name][2])
+    base_v, base_s = slot * ALU_VGPR_STRIDE, slot * ALU_SGPR_STRIDE
+    if name == "sdst" and isinstance(field, SGPRField): reg = VCC_LO
+    elif is_sreg and not isinstance(field, VGPRField): reg = VCC_LO
+    elif isinstance(field, VGPRField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
+    elif isinstance(field, SSrcField): reg = VCC_LO if nregs <= 2 else s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
+    elif isinstance(field, SGPRField): reg = s[base_s:base_s+nregs-1] if nregs > 1 else s[base_s]
+    elif isinstance(field, SrcField): reg = v[base_v:base_v+nregs-1] if nregs > 1 else v[base_v]
+    else: reg = None
+    if reg is not None: kwargs[name] = reg; slot += 1
+    elif isinstance(field, BitField): kwargs[name] = field.default
+  return builder(**kwargs)
+
+# ** create a memory instruction with pre set address registers
+
+MEM_PRESET_REGS:dict[str, dict[str, Reg]] = {
+  "VGLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "vaddr":v[V_VADDR[0]:V_VADDR[1]]},
+  "GLOBAL":{"saddr":s[S_BUF_PTR[0]:S_BUF_PTR[1]], "addr":v[V_DS_ADDR]},  # addr is 32-bit offset when saddr is valid SGPR
+  "DS":{"addr":v[V_DS_ADDR]},
+  "SMEM":{"sbase":s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], "soffset":NULL},
+}
+
+def create_mem_inst(op:Enum, builder:functools.partial[Inst]) -> Inst:
+  inst_cls, operands, field_map = builder.func, OPERANDS.get(op, {}), MEM_PRESET_REGS.get(builder.func.__name__, {})
+  kwargs:dict[str, Reg|int] = {}
+  vslot, sslot = 0, 0
+  for name, field in inst_cls._fields:
+    if isinstance(field, (FixedBitField, EnumBitField)): continue
+    if name in field_map:
+      kwargs[name] = field_map[name]
+      continue
+    nregs = max(1, operands[name][1] // 32) if name in operands else 1
+    if isinstance(field, VGPRField):
+      vi = MEM_VGPR_BASE + vslot * MEM_VGPR_STRIDE
+      kwargs[name] = v[vi:vi+nregs-1] if nregs > 1 else v[vi]
+      vslot += 1
+    elif isinstance(field, (SGPRField, AlignedSGPRField, SBaseField)):
+      si = MEM_SGPR_BASE + sslot * MEM_SGPR_STRIDE
+      kwargs[name] = s[si:si+nregs-1] if nregs > 1 else s[si]
+      sslot += 1
+    elif isinstance(field, BitField): kwargs[name] = field.default
+  return builder(**kwargs)
+
+# ** collect all memory and ALU instructions from the ISA autogen
+
+def collect_instructions() -> tuple[list[Inst], list[Inst], list[str]]:
+  op_map:dict[Enum, functools.partial[Inst]] = {}
+  for name, obj in inspect.getmembers(all_insts):
+    if isinstance(obj, functools.partial) and len(obj.args) == 1: op_map[obj.args[0]] = obj
+  alu_insts:list[Inst] = []
+  mem_insts:list[Inst] = []
+  skipped:list[str] = []
+  for op_enum, builder in op_map.items():
+    if should_skip(op_enum) or op_enum not in OPERANDS: skipped.append(op_enum.name); continue
+    fmt = builder.func.__name__
+    if fmt in ALU_FORMATS: alu_insts.append(create_alu_inst(op_enum, builder))
+    elif fmt in MEM_FORMATS: mem_insts.append(create_mem_inst(op_enum, builder))
+  return alu_insts, mem_insts, skipped
+
+def exec_insts(insts:list):
+  k = Kernel(arch)
+  # ** prologue for global memory
+  k.emit(s_load_b64(sdata=s[S_BUF_PTR[0]:S_BUF_PTR[1]], sbase=s[S_KERNARG_PTR[0]:S_KERNARG_PTR[1]], soffset=NULL))
+  k.waitcnt(lgkm=0)
+  k.emit(v_mov_b32_e32(v[V_VADDR[0]], 0))
+  k.emit(v_mov_b32_e32(v[V_VADDR[1]], 0))
+  # ** emit
+  for inst in insts: k.emit(inst)
+  k.emit(s_endpgm())
+  # ** run
+  NUM_THREADS, NUM_GRIDS, BUF_SIZE = 32, 1, 1024*1024
+  def fxn(A:UOp, B:UOp, C:UOp) -> UOp:
+    lidx, gidx = UOp.special(NUM_THREADS, "lidx0"), UOp.special(NUM_GRIDS, "gidx0")
+    lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=BUF_SIZE, addrspace=AddrSpace.LOCAL), (), "lds")
+    sink = UOp.sink(A.base, B.base, C.base, lds, lidx, gidx, arg=KernelInfo(name="discover_ops"))
+    return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in k.finalize()))))
+  A = Tensor.empty(BUF_SIZE, dtype=dtypes.uint8)
+  B = Tensor.empty(1, dtype=dtypes.uint8)
+  C = Tensor.empty(1, dtype=dtypes.uint8)
+  Tensor.custom_kernel(A, B, C, fxn=fxn)[0].realize()
+
+if __name__ == "__main__":
+  import sys
+  arch = Device[Device.DEFAULT].renderer.arch
+  if arch.startswith("gfx12"):
+    from tinygrad.runtime.autogen.amd.rdna4.ins import *
+    import tinygrad.runtime.autogen.amd.rdna4.ins as all_insts
+  elif arch.startswith("gfx11"):
+    from tinygrad.runtime.autogen.amd.rdna3.ins import *
+    import tinygrad.runtime.autogen.amd.rdna3.ins as all_insts
+    # these don"t exist in RDNA3, only RDNA3.5 and above
+    SKIP.update(["S_FMAAK_F32", "S_FMAMK_F32"])
+  else:
+    print(f"{arch} not supported yet")
+    sys.exit(0)
+  alu_insts, mem_insts, skipped = collect_instructions()
+  print(f"collected {len(alu_insts)} ALU + {len(mem_insts)} memory instructions ({len(skipped)} skipped)")
+  exec_insts(mem_insts+alu_insts)
@@ -8,8 +8,8 @@ PROFILE_PATH = Path(temp("profile.pkl", append_user=True))
 EXAMPLES = {
  "empty":"test/backend/test_custom_kernel.py TestCustomKernel.test_empty",
  "plus":"test/test_tiny.py TestTiny.test_plus",
-  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=32, N)@Tensor.empty(N, N)).realize()\"",
-  "sync":"test/amd/test_custom_kernel.py TestCustomKernel.test_lds_sync",
+  "gemm":"-c \"from tinygrad import Tensor; (Tensor.empty(N:=64, N)@Tensor.empty(N, N)).realize()\"",
+  "ops":"extra/sqtt/examples/discover_ops.py"
 }

 if __name__ == "__main__":
@@ -10,11 +10,11 @@ from tinygrad.uop.ops import UOp, Ops, KernelInfo

 def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None=None) -> Tensor:
  dtype = dtype or ref.dtype
-  if not isinstance(ref.device, tuple): return Tensor.invalid(*shape, dtype=dtype, device=ref.device)
+  if not isinstance(ref.device, tuple): return Tensor.empty(*shape, dtype=dtype, device=ref.device)
  shard_axis = ref.uop.axis if axis is None else axis
  shape = tuple(s // len(ref.device) if i == shard_axis else s for i, s in enumerate(shape))
  axis = ref.uop.axis if axis is None else axis
-  return Tensor(Tensor.invalid(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)
+  return Tensor(Tensor.empty(*shape, dtype=dtype, device=ref.device).uop.multi(axis), dtype=dtype, device=ref.device)

 def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
  return _sharded_empty(ref.shape, ref, axis)
@@ -1,113 +0,0 @@
-#ifndef _TINYDRENO_H
-#define _TINYDRENO_H
-
-#include <stdint.h>
-#include <stddef.h>
-
-typedef void * cl_llvm_instance;
-
-cl_llvm_instance cl_compiler_create_llvm_instance(void);
-void cl_compiler_destroy_llvm_instance(cl_llvm_instance inst);
-
-enum cl_handle_type {
-  CL_HANDLE_COMPILED = 1,
-  CL_HANDLE_LIBRARY,
-  CL_HANDLE_LINKED
-};
-
-// handle->data for CL_HANDLE_COMPILED and CL_HANDLE_LIBRARY
-struct cl_compiled_data {
-  uint64_t chip_id;
-  uint32_t mode;
-  void    *llvm_bitcode;
-  uint64_t llvm_bitcode_size;
-  char    *build_log;
-  uint32_t build_log_len;
-  uint32_t error_code;
-};
-
-// handle->data for CL_HANDLE_LINKED
-struct cl_executable_data {
-  int32_t  num_kernels;
-  void    *kernel_props;
-  uint32_t error_code;
-  char    *build_log;
-  char     _unk0[0x20];
-  uint64_t chip_id;
-  uint32_t mode;
-};
-
-typedef struct {
-  enum cl_handle_type type;
-  union {
-    struct cl_compiled_data *compiled;
-    struct cl_executable_data *executable;
-  };
-} cl_handle;
-
-
-#define CL_MODE_32BIT 0
-#define CL_MODE_64BIT 1
-
-#define CL_SRC_STR  0
-#define CL_SRC_BLOB 1
-
-cl_handle *cl_compiler_compile_source(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int p5, uint64_t p6, uint64_t p7,
-                                      const char *source, uint64_t source_len, uint64_t source_type, void *p11);
-
-cl_handle *cl_compiler_link_program(cl_llvm_instance inst, uint64_t chip_id, int mode, const char *options, int num_handles,
-                                    cl_handle **input_handles);
-
-
-void cl_compiler_handle_create_binary(cl_handle *handle, void **out_ptr, size_t *out_size);
-
-// lib binary format (output of handle_create_binary for type 3)
-// layout: cl_lib_header, then cl_lib_section[num_sections], then data
-
-#define CL_LIB_PROGRAM     0
-#define CL_LIB_CONSTS      6
-#define CL_LIB_IMAGE       7
-#define CL_LIB_CODE       10
-#define CL_LIB_IMAGE_DESC 11
-
-typedef struct {
-  uint32_t id;
-  uint32_t offset;
-  uint32_t size;
-  uint32_t count;
-  uint32_t entry_size;
-} cl_lib_section;
-
-typedef struct {
-  uint32_t _unk0[6];
-  uint32_t num_sections;
-  uint32_t _unk1[5];
-  cl_lib_section sections[];
-} cl_lib_header;
-
-// at sections[CL_LIB_PROGRAM].offset
-typedef struct {
-  char     name[8];
-  uint32_t _unk0[3];
-  uint32_t fregs;
-  uint32_t hregs;
-} cl_lib_prog;
-
-// at sections[CL_LIB_IMAGE_DESC].offset
-typedef struct {
-  char     _unk0[0xc4];
-  uint32_t prg_offset;
-  uint32_t pvtmem;
-  char     _unk1[0x0c];
-  uint32_t shmem;
-  uint32_t samp_cnt;
-  char     _unk2[0x28];
-  uint32_t brnchstck;
-  char     _unk4[0x4c];
-  char     kernel_name[];
-} cl_lib_img_desc;
-
-void cl_compiler_free_handle(cl_handle *handle);
-void cl_compiler_free_assembly(void *ptr);
-
-#endif
@@ -590,7 +590,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
  "aten.repeat": lambda x,*repeats: Tensor.repeat(x,*repeats).contiguous(), # not a view
  "aten._softmax": lambda self,dim,half_to_float: self.softmax(dim),
  "aten._log_softmax": lambda self,dim,half_to_float: self.log_softmax(dim),
-  "aten.random_": lambda self: Tensor.randint(*self.shape, low=self.dtype.min, high=self.dtype.max, device=self.device, dtype=self.dtype),
+  "aten.random_": lambda self: Tensor.randint(*self.shape, low=dtypes.min(self.dtype), high=dtypes.max(self.dtype), device=self.device, dtype=self.dtype),
  "aten.random_.from": lambda self, from_, to: Tensor.randint(*self.shape, low=from_, high=to, device=self.device, dtype=self.dtype),
  "aten.uniform_": lambda self, low=0, high=1: Tensor.uniform(*self.shape, low=low, high=high, dtype=self.dtype),
  "aten.normal_": lambda self, mean=0, std=1: Tensor.normal(*self.shape, mean=mean, std=std, dtype=self.dtype),
@@ -24,7 +24,7 @@ if __name__ == "__main__":
  kernel_count = GlobalCounters.kernel_count
  assert kernel_count > 0, "No kernels, test failed"
  # NOTE: this is 124 on torch 2.10.0
-  expected_kernels = 334
+  expected_kernels = 332
  expectation = f"ResNet18 kernels are {kernel_count} vs {expected_kernels} expected."
  if kernel_count < expected_kernels: warnings.warn(f"{expectation} Expectation can be lowered.", UserWarning)
-  assert kernel_count <= expected_kernels, f"{expectation}"
+  assert kernel_count <= expected_kernels, f"{expectation}"
@@ -23,7 +23,7 @@ class TestKernelFusionRegression(unittest.TestCase):
    def fn():
      x = torch.randn(128, 128, device=device)
      return (x + 1.0) * 2.0 - 0.5
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 5)

  def test_relu_fusion(self):
    def fn():
@@ -31,7 +31,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      conv = torch.nn.Conv2d(3, 16, 3, padding=1).to(device)
      with torch.no_grad():
        return torch.nn.functional.relu(conv(x))
-    self._check_kernel_count(fn, 8)
+    self._check_kernel_count(fn, 6)

  def test_batchnorm_fusion(self):
    def fn():
@@ -41,26 +41,26 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.eval()
      with torch.no_grad():
        return torch.nn.functional.relu(bn(conv(x)))
-    self._check_kernel_count(fn, 12)
+    self._check_kernel_count(fn, 10)

  def test_reduce_fusion(self):
    def fn():
      x = torch.randn(64, 64, device=device)
      return (x * 2.0).sum()
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 5)

  def test_matmul_elementwise_fusion(self):
    def fn():
      x = torch.randn(32, 32, device=device)
      w = torch.randn(32, 32, device=device)
      return torch.nn.functional.relu(x @ w + 1.0)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 7)

  def test_pooling_fusion(self):
    def fn():
      x = torch.randn(1, 8, 16, 16, device=device)
      return torch.nn.functional.max_pool2d(x * 2.0, 2)
-    self._check_kernel_count(fn, 7)
+    self._check_kernel_count(fn, 5)

  def test_residual_add_relu_fusion(self):
    def fn():
@@ -68,7 +68,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      identity = torch.randn(1, 8, 16, 16, device=device)
      out = x + identity
      return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 7)

  def test_inplace_add_relu_fusion(self):
    def fn():
@@ -76,7 +76,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      y = torch.randn(1, 16, 32, 32, device=device)
      x += y
      return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 9)
+    self._check_kernel_count(fn, 7)

  def test_conv_bn_add_relu_fusion(self):
    def fn():
@@ -89,7 +89,7 @@ class TestKernelFusionRegression(unittest.TestCase):
        out = bn(conv(x))
        out += identity
        return torch.nn.functional.relu(out)
-    self._check_kernel_count(fn, 14)
+    self._check_kernel_count(fn, 12)

  def test_multiple_inplace_ops_fusion(self):
    def fn():
@@ -97,7 +97,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      x += 1.0
      x *= 2.0
      return torch.nn.functional.relu(x)
-    self._check_kernel_count(fn, 6)
+    self._check_kernel_count(fn, 4)

  def test_view_inplace_no_fusion_break(self):
    def fn():
@@ -105,7 +105,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      view = x[1:3]
      view += 1.0
      return x.sum()
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)

  def test_batchnorm_running_stats_update(self):
    def fn():
@@ -114,7 +114,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      bn.train()
      with torch.no_grad():
        return bn(x)
-    self._check_kernel_count(fn, 10)
+    self._check_kernel_count(fn, 8)

  # this is a minimal extra/other_mnist/beautiful_mnist_torch.py to cover fusion for training with optimizer
  def test_mnist_training_fusion(self):
@@ -135,7 +135,7 @@ class TestKernelFusionRegression(unittest.TestCase):
      loss.backward()
      optimizer.step()
      return loss
-    self._check_kernel_count(fn, 26)
+    self._check_kernel_count(fn, 24)

 if __name__ == "__main__":
  unittest.main()
@@ -17,7 +17,6 @@
 // Protocol

 enum {
-  CMD_PROBE = 0,          // probe devices, returns count
  CMD_MAP_BAR = 1,        // map PCI BAR, returns size
  CMD_MAP_SYSMEM_FD = 2,  // alloc DMA memory, returns fd via SCM_RIGHTS
  CMD_CFG_READ = 3,       // read PCI config space
@@ -25,15 +24,11 @@ enum {
  CMD_RESET = 5,          // reset device
  CMD_MMIO_READ = 6,      // bulk read from BAR
  CMD_MMIO_WRITE = 7,     // bulk write to BAR
-  CMD_MAP_SYSMEM = 8,     // map system memory
-  CMD_SYSMEM_READ = 9,    // bulk read from system memory
-  CMD_SYSMEM_WRITE = 10,  // bulk write to system memory
-  CMD_RESIZE_BAR = 11,    // resize bar (noop)
  RESP_OK = 0, RESP_ERR = 1,
 };

-typedef struct { uint8_t cmd; uint32_t dev_id, bar; uint64_t arg0, arg1, arg2; } __attribute__((packed)) request_t;
-typedef struct { uint8_t status; uint64_t resp0, resp1; } __attribute__((packed)) response_t;
+typedef struct { uint8_t cmd, bar; uint64_t offset, size, value; } __attribute__((packed)) request_t;
+typedef struct { uint8_t status; uint64_t value, addr; } __attribute__((packed)) response_t;

 // Constants and state

@@ -82,7 +77,7 @@ static int send_response(int fd, response_t *resp, int send_fd) {
 }

 static void send_error(int fd, const char *msg) {
-  response_t resp = {.status = RESP_ERR, .resp0 = strlen(msg)};
+  response_t resp = {.status = RESP_ERR, .value = strlen(msg)};
  send_response(fd, &resp, -1);
  send(fd, msg, strlen(msg), 0);
 }
@@ -121,12 +116,12 @@ static int dext_rpc(uint32_t sel, uint64_t *in, uint32_t in_cnt, uint64_t *out_v
 static int map_bar(uint32_t bar, response_t *resp) {
  if (bar >= MAX_BARS) return -1;
  if (!g_bars[bar].addr && IOConnectMapMemory64(g_conn, bar, mach_task_self(), &g_bars[bar].addr, &g_bars[bar].size, kIOMapAnywhere)) return -1;
-  resp->resp0 = g_bars[bar].addr;
-  resp->resp1 = g_bars[bar].size;
+  resp->addr = g_bars[bar].addr;
+  resp->value = g_bars[bar].size;
  return 0;
 }

-static int map_sysmem_fd(uint64_t size, int contiguous, response_t *resp, int *out_fd) {
+static int map_sysmem_fd(uint64_t size, response_t *resp, int *out_fd) {
  if (g_sysmem_count >= MAX_SYSMEM) return -1;
  int idx = g_sysmem_count;
  int fd = -1;
@@ -153,7 +148,7 @@ static int map_sysmem_fd(uint64_t size, int contiguous, response_t *resp, int *o
  strncpy(g_sysmem[idx].shm_name, shm_name, sizeof(g_sysmem[idx].shm_name));
  g_sysmem_count++;

-  *resp = (response_t){.resp0 = alloc_sz, .resp1 = idx};
+  *resp = (response_t){.addr = idx, .value = alloc_sz};
  *out_fd = fd;
  return 0;

@@ -208,42 +203,39 @@ static void handle_client(int fd) {

    case CMD_MAP_SYSMEM_FD: {
      int shm_fd = -1;
-      resp.status = map_sysmem_fd(req.arg0, (int)req.arg1, &resp, &shm_fd) ? 1 : 0;
+      resp.status = map_sysmem_fd(req.size, &resp, &shm_fd) ? 1 : 0;
      send_response(fd, &resp, shm_fd);
      continue;
    }

    case CMD_CFG_READ: {
-      uint64_t in[2] = {req.arg0, req.arg1};
-      resp.status = dext_rpc(0, in, 2, &resp.resp0) ? 1 : 0;
+      uint64_t in[2] = {req.offset, req.size};
+      resp.status = dext_rpc(0, in, 2, &resp.value) ? 1 : 0;
      break;
    }

    case CMD_CFG_WRITE: {
-      uint64_t in[3] = {req.arg0, req.arg1, req.arg2};
+      uint64_t in[3] = {req.offset, req.size, req.value};
      resp.status = dext_rpc(1, in, 3, NULL) ? 1 : 0;
      break;
    }

-    case CMD_RESIZE_BAR:
-      break;
-
    case CMD_RESET:
      resp.status = dext_rpc(2, NULL, 0, NULL) ? 1 : 0;
      break;

    case CMD_MMIO_READ:
-      if (validate_bar(req.bar, req.arg0, req.arg1)) { resp.status = 1; break; }
-      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.arg0), req.arg1);
-      resp.resp0 = req.arg1;
+      if (validate_bar(req.bar, req.offset, req.size)) { resp.status = 1; break; }
+      mmio_copy(g_bulk_buf, (void*)(g_bars[req.bar].addr + req.offset), req.size);
+      resp.value = req.size;
      send_response(fd, &resp, -1);
-      send(fd, g_bulk_buf, req.arg1, 0);
+      send(fd, g_bulk_buf, req.size, 0);
      continue;

    case CMD_MMIO_WRITE:
-      recvall(fd, g_bulk_buf, req.arg1);
-      if (!validate_bar(req.bar, req.arg0, req.arg1))
-        mmio_copy((void*)(g_bars[req.bar].addr + req.arg0), g_bulk_buf, req.arg1);
+      recvall(fd, g_bulk_buf, req.size);
+      if (!validate_bar(req.bar, req.offset, req.size))
+        mmio_copy((void*)(g_bars[req.bar].addr + req.offset), g_bulk_buf, req.size);
      continue;

    default:
@@ -1,7 +1,6 @@
 A command line tool for exploring the VIZ trace.

-1. Set VIZ to -1 to save the trace.
-2. Use `extra/viz/cli.py` to inspect the trace files.
+After running with VIZ=-1, use `extra/viz/cli.py` to explore the saved trace files.

 ## Inspect runtime profiling

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
-import argparse, pathlib, sys, struct, json, itertools
+import os
+os.environ["VIZ"] = "0"
+import argparse, pathlib, sys, struct, json
 from typing import Iterator
 from tinygrad.viz import serve as viz
 from tinygrad.uop.ops import RewriteTrace
@@ -46,16 +48,11 @@ def decode_profile(data:bytes) -> dict:
        name, ref, key, st, dur, fmt = u("<IIIIfI")
        v["events"].append({"name":strings[name], "ref":option(ref), "key":option(key), "st":st, "dur":dur, "fmt":strings[fmt]})
    else:
-      v["linear"] = u("<B")[0]
      v["peak"] = u("<Q")[0]
      for _ in range(event_count):
-        if v["linear"]:
-          ts, value = u("<IQ")
-          v["events"].append({"event":"freq", "ts":ts, "value":value})
-        else:
-          alloc, ts, key = u("<BII")
-          if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
-          else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
+        alloc, ts, key = u("<BII")
+        if alloc: v["events"].append({"event":"alloc", "ts":ts, "key":key, "arg": {"dtype":strings[u("<I")[0]], "sz":u("<Q")[0]}})
+        else: v["events"].append({"event":"free", "ts":ts, "key":key, "arg": {"users":[u("<IIIB") for _ in range(u("<I")[0])]}})
  return {"dur":total_dur, "peak":global_peak, "layout":layout, "markers":markers}

 if __name__ == "__main__":
@@ -65,12 +62,9 @@ if __name__ == "__main__":
  g_mode.add_argument("--rewrites", action="store_true", help="View rewrites trace")
  g_common = parser.add_argument_group("common options")
  g_common.add_argument("--kernel", type=str, default=None, metavar="NAME", help="Select a kernel by name (optional name, default: only list names)")
-  g_common.add_argument("--no-color", action="store_true", default=not (sys.stdin.isatty() and sys.stdout.isatty()),
-                        help="Disable colored output (default: true in non-interactive mode)")
  g_profile = parser.add_argument_group("profile options")
  g_profile.add_argument("--device", type=str, default=None, metavar="NAME", help="Select a device (optional name, default: only list names)")
-  g_profile.add_argument("--offset", type=int, default=0, metavar="N", help="event offset (default: 0)")
-  g_profile.add_argument("--limit", type=int, default=10, metavar="N", help="events to display (-1 for all, default: 10)")
+  g_profile.add_argument("--top", type=int, default=10, metavar="N", help="Number of top kernels to show (-1 for all, default: 10)")
  g_rewrites = parser.add_argument_group("rewrites options")
  g_rewrites.add_argument("--select", type=str, default=None, metavar="NAME",
                          help="Select an item within the chosen kernel (optional name, default: only list names)")
@@ -86,64 +80,14 @@ if __name__ == "__main__":
  viz.trace = viz.load_pickle(args.rewrites_path, default=RewriteTrace([], [], {}))
  viz.ctxs = viz.get_rewrites(viz.trace)

-  def format_colored(s:str) -> str: return ansistrip(s) if args.no_color else s
-
  if args.profile:
    from tabulate import tabulate
-    profile = decode_profile(viz.get_profile(profile_data:=viz.load_pickle(args.profile_path, default=[])))
-    viz.load_amd_counters(viz.ctxs, profile_data)
-    counters = {f'{c["name"]} SQTT {s["name"]}': s["data"] for c in viz.ctxs if c["name"].startswith("Exec") for s in c["steps"]
-                if s["name"].startswith("PKTS")}
-    if args.device is None:
-      print("Select a device:")
-      for k in (*profile["layout"], *counters):
-        print(f"  {format_colored(k)}")
-      sys.exit(0)
-
-    # ** SQTT printer
-    if args.device is not None and (sqtt_data:=next((v for k,v in counters.items() if ansistrip(k) == args.device), None)) is not None:
-      assert args.limit > 1, f"SQTT limit must be greater than 1, got {args.limit}"
-      sqtt_events, has_more = viz.sqtt_timeline(*sqtt_data, max_pkts=args.offset+args.limit)
-      sqtt_pkts = [e for e in sqtt_events if type(e).__name__ == "ProfileRangeEvent"]
-      pc_map = next((e.arg for e in sqtt_events if type(e).__name__ == "ProfilePointEvent" and e.key == 'pcMap'), None)
-      if pc_map is None:
-        print(f"No SQTT instruction trace data for {args.device}")
-        sys.exit(0)
-      # modern terminals support 24-bit color
-      def hex_colored(st:str, color:str) -> str: return f"\x1b[38;2;{int(color[1:3],16)};{int(color[3:5],16)};{int(color[5:7],16)}m{st}\x1b[0m"
-      WAVE_COLORS = ((('VALU', 'VINTERP'), '#ffffc0'), (('SALU',), '#cef263'), (('VMEM',), '#b2b7c9'), (('LOAD', 'SMEM'), '#ffc0c0'),
-                     (('STORE',), '#4fa3cc'), (('IMMEDIATE',), '#f3b44a'), (('BARRIER',), '#d00000'), (('LDS',), '#9fb4a6'), (('JUMP',), '#ffb703'),
-                     (('JUMP_NO',), '#fb8500'), (('MESSAGE',), '#90dbf4'), (('WAVERDY',), '#1a2a2a'))
-      print(f"{'Clk':<12} {'Unit':<20} {'Op':<22} {'Dur':<4} {'Info'}")
-      print("-" * 90)
-      # start from the first packet in trace, prepare packet indexes and map dispatches
-      pkt_idxs:dict[str, itertools.count] = {}
-      dispatch_to_pc:dict[str, int] = {}
-      for e in sqtt_pkts[:-args.limit]:
-        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
-        if e.name.ret is not None and e.name.ret.startswith("PC:"): dispatch_to_pc[f"{e.device}-{idx}"] = int(e.name.ret.replace("PC:", ""))
-      # start printing from the offset point
-      for e in sqtt_pkts[-args.limit:]:
-        op_name, info = e.name.display_name, e.name.ret or ""
-        color = next((c for p, c in WAVE_COLORS if any(x in op_name for x in p)), None)
-        op_str = hex_colored(op_name, color) if color and not args.no_color else op_name
-        phase, pc = None, None
-        idx = next(pkt_idxs.setdefault(e.device, itertools.count()))
-        if info.startswith("PC:"):
-          dispatch_to_pc[f"{e.device}-{idx}"] = pc = int(info.replace("PC:", ""))
-          phase = "DISPATCH"
-        if info.startswith("LINK:"): phase, pc = "EXEC", dispatch_to_pc[info.replace("LINK:", "")]
-        if pc and phase: info = f"{phase:<8} 0x{pc:05x} {pc_map[pc]}"
-        print(f"{int(e.st):<12} {e.device:<20} {op_str}{' '*(22-ansilen(op_str))} {int(e.en-e.st):<4} {info}")
-      # note: we only print the important packets and skip the rest
-      if has_more: print(f"Selected packets {args.offset:,}-{args.offset + args.limit:,}. Use --offset and --limit to see others")
-      sys.exit(0)
-
-    # ** Profiler printer
+    profile = decode_profile(viz.get_profile(viz.load_pickle(args.profile_path, default=[])))
    agg, total, n = {}, 0, 0
+    if args.device is None: print("Select a device:")
    for k,v in profile["layout"].items():
      if not optional_eq({"name":k}, args.device): continue
-      print(f"  {format_colored(k)}")
+      print(f"  {k}")
      if args.device is None: continue
      for e in v.get("events", []):
        et = e["dur"]*1e-6
@@ -151,7 +95,7 @@ if __name__ == "__main__":
          if optional_eq(e, args.kernel) and n < 10:
            ptm = colored(time_to_str(et, w=9), "yellow" if et > 0.01 else None) if et is not None else ""
            name = e["name"]+(" " * (46 - ansilen(e["name"])))
-            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e.get('fmt', '').replace('\n', ' | ')+"  ")
+            print(f"{name} {ptm}/{(et or 0)*1e3:9.2f}ms  "+e['fmt'].replace('\n', ' | ')+"  ")
            n += 1
        else:
          a = agg.setdefault(e["name"], [0.0, 0])
@@ -160,15 +104,14 @@ if __name__ == "__main__":
          total += et
    if agg and total > 0:
      items = sorted(agg.items(), key=lambda kv:kv[1][0], reverse=True)
-      sel = items if args.limit == -1 else items[args.offset:args.offset+args.limit]
+      sel = items if args.top == -1 else items[:args.top]
      table = [[name, time_to_str(t, w=9), c, f"{(t/total*100.0):.2f}%"] for name,(t,c) in sel]
-      if args.limit != -1 and (other:=items[len(sel):]):
+      if args.top != -1 and (other:=items[len(sel):]):
        other_t = total-sum(t for _, (t, _) in sel)
        table.append([f"Other ({len(other)} unique)", time_to_str(other_t, w=9), sum(c for _,(_,c) in other), f"{other_t/total*100.0:.2f}%"])
      print(tabulate(table, headers=["name", "total", "count", "pct"], tablefmt="github"))
    sys.exit(0)

-  # ** Graph rewrites printer
  for k in viz.ctxs:
    if not optional_eq(k, args.kernel): continue
    print(k["name"])
@@ -5,8 +5,8 @@ from tinygrad.runtime.autogen import llvm
 from tinygrad.runtime.support.elf import elf_loader

 ARCH_TO_TARGET:dict[str, list[str]] = {
-  "rdna3":["gfx1100", "gfx1151"],
-  "rdna4":["gfx1200", "gfx1201"],
+  "rdna3":["gfx1100"],
+  "rdna4":["gfx1200"],
  "cdna":["gfx950", "gfx942"],
 }

@@ -117,58 +117,6 @@ class TestDS2AddrMore(unittest.TestCase):
    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB)
    self.assertEqual(st.vgpr[0][4], 0xDEADBEEF, "v4 should be untouched")

-  def test_ds_load_2addr_b64_addr_overlaps_vdst(self):
-    """DS_LOAD_2ADDR_B64 where addr register overlaps vdst range.
-
-    Hardware reads the address before writing any results, so addr=v[4]
-    with vdst=v[4:7] must load all 4 dwords using the original v[4] value.
-    """
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[2], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
-      s_mov_b32(s[2], 0xBBBBBBBB),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
-      s_mov_b32(s[2], 0xCCCCCCCC),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=8),
-      s_mov_b32(s[2], 0xDDDDDDDD),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=12),
-      s_waitcnt(lgkmcnt=0),
-      # addr=v[4] overlaps vdst=v[4:7]
-      v_mov_b32_e32(v[4], 0),
-      DS(DSOp.DS_LOAD_2ADDR_B64, addr=v[4], vdst=v[4:7], offset0=0, offset1=1),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][4], 0xAAAAAAAA, "v4 = LDS[0:4]")
-    self.assertEqual(st.vgpr[0][5], 0xBBBBBBBB, "v5 = LDS[4:8]")
-    self.assertEqual(st.vgpr[0][6], 0xCCCCCCCC, "v6 = LDS[8:12]")
-    self.assertEqual(st.vgpr[0][7], 0xDDDDDDDD, "v7 = LDS[12:16]")
-
-  def test_ds_load_2addr_b32_addr_overlaps_vdst(self):
-    """DS_LOAD_2ADDR_B32 where addr register overlaps vdst range."""
-    instructions = [
-      v_mov_b32_e32(v[10], 0),
-      s_mov_b32(s[2], 0xAAAAAAAA),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=0),
-      s_mov_b32(s[2], 0xBBBBBBBB),
-      v_mov_b32_e32(v[0], s[2]),
-      ds_store_b32(addr=v[10], data0=v[0], offset0=4),
-      s_waitcnt(lgkmcnt=0),
-      # addr=v[2] overlaps vdst=v[2:3]
-      v_mov_b32_e32(v[2], 0),
-      DS(DSOp.DS_LOAD_2ADDR_B32, addr=v[2], vdst=v[2:3], offset0=0, offset1=1),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xAAAAAAAA, "v2 = LDS[0:4]")
-    self.assertEqual(st.vgpr[0][3], 0xBBBBBBBB, "v3 = LDS[4:8]")
-
  def test_ds_load_b64_no_overwrite(self):
    """DS_LOAD_B64 should only write 2 VGPRs."""
    instructions = [
@@ -811,47 +759,6 @@ class TestDsPermute(unittest.TestCase):
    # Lane 0 receives data (highest numbered active lane wins)
    self.assertEqual(st.vgpr[0][2], 0x11111111)

-  def test_ds_bpermute_b32_xor_swap(self):
-    """DS_BPERMUTE_B32 with XOR-1 pattern — each lane reads from lane^1.
-
-    This is the pattern used by warp_shfl_xor in flash attention for reduce_max/reduce_sum.
-    Each lane has a unique value (lane_id + 100), and reads from the adjacent lane.
-    """
-    instructions = [
-      # v[0] = (lane_id ^ 1) * 4  (byte offset for bpermute)
-      v_xor_b32_e32(v[0], 1, v[255]),
-      v_lshlrev_b32_e32(v[0], 2, v[0]),
-      # v[1] = lane_id + 100 (unique per-lane value)
-      s_mov_b32(s[0], 100),
-      v_add_nc_u32_e32(v[1], s[0], v[255]),
-      # ds_bpermute: v[2] = v[1] from lane (lane_id ^ 1)
-      ds_bpermute_b32(vdst=v[2], addr=v[0], data0=v[1]),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=32)
-    for lane in range(32):
-      src_lane = lane ^ 1
-      expected = src_lane + 100
-      self.assertEqual(st.vgpr[lane][2], expected, f"lane {lane}: expected v[1] from lane {src_lane} = {expected}, got {st.vgpr[lane][2]}")
-
-
-class TestDSSubDword(unittest.TestCase):
-  """Tests for sub-dword DS operations (ds_store_b16, ds_store_b16_d16_hi)."""
-
-  def test_ds_store_b16_and_d16_hi(self):
-    """DS_STORE_B16 stores low 16 bits, DS_STORE_B16_D16_HI stores high 16 bits to adjacent LDS half-words."""
-    instructions = [
-      v_mov_b32_e32(v[0], 0),
-      v_mov_b32_e32(v[1], 0xBEEF1234),
-      DS(DSOp.DS_STORE_B16, addr=v[0], data0=v[1], offset0=0),
-      DS(DSOp.DS_STORE_B16_D16_HI, addr=v[0], data0=v[1], offset0=2),
-      s_waitcnt(lgkmcnt=0),
-      ds_load_b32(vdst=v[2], addr=v[0], offset0=0),
-      s_waitcnt(lgkmcnt=0),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertEqual(st.vgpr[0][2], 0xBEEF1234, "lo=0x1234 at byte 0, hi=0xBEEF at byte 2")
-

 class TestDSLargeOffset(unittest.TestCase):
  """Tests for DS instructions with offsets > 255 (offset1 > 0).
@@ -408,23 +408,6 @@ class TestVOP3P(unittest.TestCase):
    self.assertEqual(lo, 0x0005, f"lo: expected 0x0005, got 0x{lo:04x}")
    self.assertEqual(hi, 0x4003, f"hi: expected 0x4003, got 0x{hi:04x}")

-  def test_v_pk_add_u16_literal_constant(self):
-    """V_PK_ADD_U16 with a literal constant (value > 64, requires VOP3P_LIT encoding).
-    Regression test: VOP3P literal constants were not passed to rsrc_dyn, so literal src read as 0.
-    """
-    instructions = [
-      s_mov_b32(s[0], 0x1C001C00),  # packed u16: hi=0x1C00, lo=0x1C00 (f16 for 2^-8)
-      v_mov_b32_e32(v[0], s[0]),
-      v_pk_add_u16(v[1], 0x2000, v[0], opsel_hi=2, opsel_hi2=1),  # add 0x2000 bias to both halves
-    ]
-    st = run_program(instructions, n_lanes=1)
-    result = st.vgpr[0][1]
-    lo = result & 0xffff
-    hi = (result >> 16) & 0xffff
-    # lo = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0), hi = 0x1C00 + 0x2000 = 0x3C00 (f16 1.0)
-    self.assertEqual(lo, 0x3C00, f"lo: expected 0x3C00, got 0x{lo:04x}")
-    self.assertEqual(hi, 0x3C00, f"hi: expected 0x3C00, got 0x{hi:04x}")
-

 class TestWMMAF16(unittest.TestCase):
  """Tests for WMMA F16 output variant (V_WMMA_F16_16X16X16_F16).
@@ -138,47 +138,6 @@ class TestVOPDLiterals(unittest.TestCase):
    self.assertAlmostEqual(i2f(st.vgpr[0][2]), 13.0, places=5, msg="fma(2.0, 5.0, 3.0) should be 13.0")


-class TestVOPDDot2Acc(unittest.TestCase):
-  """Tests for V_DUAL_DOT2ACC_F32_F16 - packed f16 dot product accumulate."""
-
-  def test_vopd_dot2acc_f32_f16_basic(self):
-    """V_DUAL_DOT2ACC_F32_F16: D += lo(S0)*lo(S1) + hi(S0)*hi(S1).
-
-    S0 = pack(1.0h, 2.0h), S1 = pack(3.0h, 4.0h), D = 10.0f
-    result = 10.0 + 1.0*3.0 + 2.0*4.0 = 10.0 + 3.0 + 8.0 = 21.0
-    """
-    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
-    pk_s0 = f32_to_f16(1.0) | (f32_to_f16(2.0) << 16)  # lo=1.0h, hi=2.0h
-    pk_s1 = f32_to_f16(3.0) | (f32_to_f16(4.0) << 16)  # lo=3.0h, hi=4.0h
-    instructions = [
-      v_mov_b32_e32(v[0], pk_s0),
-      v_mov_b32_e32(v[1], pk_s1),
-      v_mov_b32_e32(v[3], f2i(10.0)),  # accumulator in v[3] (vdsty with vdstx=v[4])
-      # X: v[4] = MOV v[0] (don't care), Y: v[3] += dot2(v[0], v[1])
-      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][3]), 21.0, places=2, msg="10.0 + 1.0*3.0 + 2.0*4.0 = 21.0")
-
-  def test_vopd_dot2acc_f32_f16_zero_accum(self):
-    """V_DUAL_DOT2ACC_F32_F16 with zero accumulator — pure dot product.
-
-    S0 = pack(0.5h, -1.0h), S1 = pack(2.0h, 3.0h), D = 0.0f
-    result = 0.0 + 0.5*2.0 + (-1.0)*3.0 = 1.0 - 3.0 = -2.0
-    """
-    from test.amd.hw.helpers import f2i, i2f, f32_to_f16
-    pk_s0 = f32_to_f16(0.5) | (f32_to_f16(-1.0) << 16)
-    pk_s1 = f32_to_f16(2.0) | (f32_to_f16(3.0) << 16)
-    instructions = [
-      v_mov_b32_e32(v[0], pk_s0),
-      v_mov_b32_e32(v[1], pk_s1),
-      v_mov_b32_e32(v[3], f2i(0.0)),  # zero accumulator in v[3]
-      VOPD(VOPDOp.V_DUAL_MOV_B32, VOPDOp.V_DUAL_DOT2ACC_F32_F16, v[4], v[3], v[0], v[0], v[0], v[1]),
-    ]
-    st = run_program(instructions, n_lanes=1)
-    self.assertAlmostEqual(i2f(st.vgpr[0][3]), -2.0, places=2, msg="0.5*2.0 + (-1.0)*3.0 = -2.0")
-
-
 class TestVOPDMultilane(unittest.TestCase):
  """Tests for VOPD with multiple lanes."""

@@ -1,14 +1,10 @@
 import unittest
-import functools
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.renderer import Estimates
-from tinygrad.dtype import AddrSpace
+
 from tinygrad.runtime.autogen.amd.rdna3.ins import *
-import tinygrad.runtime.autogen.amd.rdna3.ins as r3
-import tinygrad.runtime.autogen.amd.rdna4.ins as r4
 from tinygrad.renderer.amd.dsl import s, v
-from test.amd.helpers import TARGET_TO_ARCH

 def custom_add_one(A:UOp) -> UOp:
  A = A.flatten()
@@ -47,60 +43,9 @@ def custom_add_var(A:UOp, B:UOp) -> UOp:
  sink = UOp.sink(A.base, B.base, var, threads, arg=KernelInfo(f"custom_add_var_{A.size}"))
  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))

-def custom_wave_sync(A:UOp, arch:str) -> UOp:
-  # 4 waves across 1024 WG — enough to saturate a SIMD with many concurrent WGs
-  # s_sleep yields the SIMD so waves from different WGs interleave, causing barrier packet reordering
-  threads = UOp.special(128, "lidx0")
-  wg = UOp.special(1024, "gidx0")
-  insts = []
-  for _ in range(4):
-    insts.append(s_sleep(4))
-    insts += [s_barrier()] if arch == "rdna3" else [r4.s_barrier_signal(), r4.s_barrier_wait()]
-    insts += [s_nop(0)]*4
-  insts.append(s_endpgm())
-  sink = UOp.sink(A.base, threads, wg, arg=KernelInfo("custom_wave_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
-def custom_lds_sync(A:UOp, arch:str) -> UOp:
-  A = A.flatten()
-  num_threads = A.shape[0]
-  threads = UOp.special(num_threads, "lidx0")
-  wg = UOp.special(1, "gidx0")
-  lds = UOp(Ops.DEFINE_LOCAL, dtypes.uint8.ptr(size=512, addrspace=AddrSpace.LOCAL), (), 'lds')  # 128 * 4 bytes
-  isa = r4 if arch == "rdna4" else r3
-  wait_kmcnt = [isa.s_wait_kmcnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
-  wait_dscnt = [isa.s_wait_dscnt(simm16=0)] if arch == "rdna4" else [isa.s_waitcnt(lgkmcnt=0)]
-  barrier = [isa.s_barrier_signal(ssrc0=-1), isa.s_barrier_wait(simm16=-1)] if arch == "rdna4" else [isa.s_barrier()]
-  global_store = [isa.global_store_b32(vaddr=v[6:7], saddr=s[0:1], vsrc=v[5])] if arch == "rdna4" \
-      else [isa.global_store_b32(addr=v[6], data=v[5], saddr=s[0:1])]
-  insts = [
-    isa.s_load_b64(s[0:1], s[0:1], soffset=NULL),
-    *wait_kmcnt,
-    isa.v_lshlrev_b32_e32(v[1], 2, v[0]),
-    # lds[thread_idx] = thread_idx
-    isa.ds_store_b32(addr=v[1], data0=v[0]),
-    *wait_dscnt,
-    *barrier,
-    # out[threaed_idx] = thread_idx == num_threads ? -1 : lds[thread_idx + 1]
-    isa.v_add_nc_u32_e32(v[2], 4, v[1]),
-    isa.v_cmp_gt_u32_e32(num_threads-1, v[0]),
-    isa.ds_load_b32(vdst=v[3], addr=v[2]),
-    *wait_dscnt,
-    isa.v_mov_b32_e32(v[4], -1),
-    isa.v_cndmask_b32_e32(v[5], v[4], v[3]),
-    isa.v_lshlrev_b32_e32(v[6], 2, v[0]),
-    *global_store,
-    isa.s_endpgm(),
-  ]
-  sink = UOp.sink(A.base, lds, threads, wg, arg=KernelInfo("custom_lds_sync"))
-  return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.DEVICE, arg="AMD"), UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
-
@unittest.skipUnless(Device.DEFAULT == "AMD", "requires AMD device")
 class TestCustomKernel(unittest.TestCase):
-  def setUp(self): self.arch = TARGET_TO_ARCH[Device["AMD"].arch]
-
  def test_simple(self):
-    if self.arch != "rdna3": self.skipTest("only rdna3")
    a = Tensor.full((16, 16), 1.).contiguous().realize()
    a = Tensor.custom_kernel(a, fxn=custom_add_one)[0]
    ei = a.schedule()[-1].lower()
@@ -110,7 +55,6 @@ class TestCustomKernel(unittest.TestCase):
    self.assertTrue((a.numpy() == 2.).all())

  def test_variable(self):
-    if self.arch != "rdna3": self.skipTest("only rdna3")
    b = Tensor.full((16, 16), 1, dtype=dtypes.uint32).contiguous().realize()
    a = Tensor.zeros_like(b).contiguous().realize()
    a = Tensor.custom_kernel(a, b, fxn=custom_add_var)[0]
@@ -119,14 +63,5 @@ class TestCustomKernel(unittest.TestCase):
      ei.run({"var":i})
      self.assertTrue((a.numpy() == 1+i).all())

-  def test_lds_sync(self):
-    if self.arch not in ("rdna3", "rdna4"): self.skipTest("only rdna3/rdna4")
-    a = Tensor.empty(128, dtype=dtypes.int32).contiguous().realize()
-    a = Tensor.custom_kernel(a, fxn=functools.partial(custom_lds_sync, arch=self.arch))[0]
-    a.realize()
-    ref = Tensor.arange(1, 129, dtype=dtypes.int32)
-    ref[127] = -1
-    self.assertListEqual(a.tolist(), ref.tolist())
-
 if __name__ == "__main__":
  unittest.main()
@@ -10,11 +10,18 @@ from tinygrad.runtime.autogen.amd.rdna3.ins import SOPP
 from tinygrad.runtime.autogen.amd.rdna3.enum import SOPPOp
 from tinygrad.renderer.amd.sqtt import (decode, LAYOUT_HEADER, WAVESTART, WAVESTART_RDNA4, WAVEEND, INST, INST_RDNA4, VALUINST,
                                     IMMEDIATE, IMMEDIATE_MASK, PACKET_TYPES_RDNA3, PACKET_TYPES_RDNA4, PACKET_TYPES_CDNA, CDNA_WAVESTART,
-                                     print_packets, CDNA_WAVEEND, CDNA_INST)
+                                     InstOp, InstOpRDNA4, print_packets, CDNA_WAVEEND)
 from test.amd.helpers import TARGET_TO_ARCH

 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"
+# INST ops for non-traced SIMDs (excluded from instruction count)
+OTHER_SIMD_OPS = {InstOp.OTHER_LDS_LOAD, InstOp.OTHER_LDS_STORE, InstOp.OTHER_LDS_STORE_64, InstOp.OTHER_LDS_STORE_128,
+                  InstOp.OTHER_FLAT_LOAD, InstOp.OTHER_FLAT_STORE, InstOp.OTHER_FLAT_STORE_64, InstOp.OTHER_FLAT_STORE_96,
+                  InstOp.OTHER_FLAT_STORE_128, InstOp.OTHER_GLOBAL_LOAD, InstOp.OTHER_GLOBAL_LOAD_VADDR,
+                  InstOp.OTHER_GLOBAL_STORE_64, InstOp.OTHER_GLOBAL_STORE_96, InstOp.OTHER_GLOBAL_STORE_128,
+                  InstOp.OTHER_GLOBAL_STORE_VADDR_128}
+OTHER_SIMD_OPS_RDNA4 = {InstOpRDNA4.OTHER_VMEM, InstOpRDNA4.OTHER_VMEM_5}

 # ═══════════════════════════════════════════════════════════════════════════════
 # ROCPROF DECODER
@@ -146,10 +153,9 @@ class SQTTExamplesTestBase(unittest.TestCase):
      if "gemm" not in name: continue
      with self.subTest(example=name):
        all_packets = [p for e in events for p in decode(e.blob)]
-        inst_packets = [p for p in all_packets if isinstance(p, (INST, INST_RDNA4, CDNA_INST))]
-        self.assertGreater(len(inst_packets), 0, f"no INST packets in {name}")
-        if isinstance(inst_packets[0], (INST, INST_RDNA4)):
-          self.assertGreater(len([p for p in inst_packets if p.op.name.startswith("JUMP")]), 0, f"no JUMP packets in {name}")
+        inst_names = [p.op.name for p in all_packets if isinstance(p, (INST, INST_RDNA4))]
+        self.assertGreater(len(inst_names), 0, f"no INST packets in {name}")
+        self.assertGreater(len([n for n in inst_names if n.startswith("JUMP")]), 0, f"no JUMP packets in {name}")

  expected: dict[str, list[int]] = {}  # override in subclasses
  def test_packet_counts(self):
@@ -176,18 +182,11 @@ class SQTTExamplesTestBase(unittest.TestCase):
        our_waves: list[tuple[int, int]] = []
        for event in events:
          wave_starts: dict[tuple[int, int, int], int] = {}
-          first_timestamp:int|None = None
          for p in decode(event.blob):
-            if first_timestamp is None: first_timestamp = p._time
            if isinstance(p, (WAVESTART, CDNA_WAVESTART, WAVESTART_RDNA4)): wave_starts[(p.wave, p.simd, p.cu)] = p._time
            elif isinstance(p, (WAVEEND, CDNA_WAVEEND)) and (key := (p.wave, p.simd, p.cu)) in wave_starts:
              our_waves.append((wave_starts[key], p._time))
-          for st in wave_starts.values():
-            self.assertGreater(st, first_timestamp, "wave start must be after the first packet")
-        # rocprof fails non deterministically and gives inaccurate timestamps.
-        #self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")
-        for st, et in our_waves:
-          self.assertGreater(et, st, "wave end must be after start")
+        self.assertEqual(sorted(our_waves), sorted(roc_waves), f"wave times mismatch in {name}")

  def test_rocprof_inst_times_match(self):
    """Instruction times must match rocprof exactly (excluding s_endpgm)."""
@@ -200,8 +199,8 @@ class SQTTExamplesTestBase(unittest.TestCase):
        our_insts: list[int] = []
        for event in events:
          for p in decode(event.blob):
-            # INST ops for non-traced SIMDs (excluded from instruction count)
-            if isinstance(p, (INST, INST_RDNA4)) and not p.op.name.startswith("OTHER_"): our_insts.append(p._time)
+            if isinstance(p, INST) and p.op not in OTHER_SIMD_OPS: our_insts.append(p._time)
+            elif isinstance(p, INST_RDNA4) and p.op not in OTHER_SIMD_OPS_RDNA4: our_insts.append(p._time)
            elif isinstance(p, VALUINST): our_insts.append(p._time)
            elif isinstance(p, IMMEDIATE): our_insts.append(p._time)
            elif isinstance(p, IMMEDIATE_MASK):
@@ -211,22 +210,23 @@ class SQTTExamplesTestBase(unittest.TestCase):
 class TestSQTTExamplesRDNA3(SQTTExamplesTestBase):
  target = "gfx1100"
  expected = {
-    "profile_empty_run_0": [1880, 1867, 1920, 1971, 1998, 1904],
-    "profile_empty_run_1": [1880, 1867, 1920, 1971, 1998, 1904],
-    "profile_gemm_run_0": [3275, 3278, 2426, 2475, 2511, 2431],
-    "profile_gemm_run_1": [3264, 3268, 2420, 2469, 2504, 2401],
-    "profile_ops_run_0": [1944, 4903, 1984, 2035, 2062, 1968],
-    "profile_ops_run_1": [1944, 4918, 1984, 2035, 2062, 1968],
-    "profile_plus_run_0": [1938, 1932, 1978, 2029, 2056, 1962],
-    "profile_plus_run_1": [1891, 1874, 1931, 1982, 2009, 1915],
+    "profile_empty_run_0": [1974, 1961, 2014, 2065, 2092, 1998],
+    "profile_empty_run_1": [1979, 1972, 2019, 2070, 2097, 2003],
+    "profile_gemm_run_0": [2038, 11076, 2324, 2129, 2156, 2062],
+    "profile_gemm_run_1": [2038, 11037, 2318, 2129, 2156, 2062],
+    "profile_ops_run_0": [2038, 5070, 2078, 2129, 2156, 2062],
+    "profile_ops_run_1": [2038, 5007, 2078, 2129, 2156, 2062],
+    "profile_plus_run_0": [1979, 1979, 2030, 2070, 2097, 2003],
+    "profile_plus_run_1": [1979, 2043, 2030, 2070, 2097, 2003],
  }

 class TestSQTTExamplesRDNA4(SQTTExamplesTestBase): target = "gfx1200"

 class TestSQTTExamplesCDNA(SQTTExamplesTestBase):
  target = "gfx950"
+  def test_decode_all_examples(self): self.skipTest("TODO: correct deltas in the timestamp packet types, first packet is REGCS_CDNA")
+  def test_gemm_has_instructions(self): self.skipTest("TODO: decode CDNA inst packets")
  def test_rocprof_wave_times_match(self): self.skipTest("TODO: requires timestamp patching")
-  def test_rocprof_inst_times_match(self): self.skipTest("TODO: requires timestamp patching")

 if __name__ == "__main__":
  unittest.main()
@@ -2,7 +2,7 @@
 import unittest, pickle
 from typing import Iterator
 from pathlib import Path
-from tinygrad.helpers import DEBUG, getenv, temp
+from tinygrad.helpers import DEBUG, OSX, getenv, temp
 from tinygrad.renderer.amd.sqtt import print_packets, map_insts
 from tinygrad.runtime.autogen.amd.rdna3.ins import s_endpgm
 from tinygrad.viz.serve import sqtt_timeline
@@ -11,7 +11,7 @@ from test.amd.disasm import disasm
 import tinygrad
 EXAMPLES_DIR = Path(tinygrad.__file__).parent.parent / "extra/sqtt/examples"

-def rocprof_inst_traces_match(sqtt, prg, target):
+def rocprof_inst_traces_match(sqtt, prg, target, pass_rocprof_err=False):
  from tinygrad.viz.serve import amd_decode
  from extra.sqtt.roc import decode as roc_decode, InstExec
  addr_table = amd_decode(prg.lib, target)
@@ -31,7 +31,7 @@ def rocprof_inst_traces_match(sqtt, prg, target):
    rocprof_inst = next(rwaves_iter[info.wave][0])
    ref_pc = rocprof_inst.pc-prg.base
    # always check pc matches
-    assert ref_pc == info.pc, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
+    assert ref_pc == info.pc or pass_rocprof_err, f"pc mismatch {ref_pc}:{disasm_map[rocprof_inst.pc]} != {info.pc}:{disasm(info.inst)}"
    # special handling for s_endpgm, it marks the wave completion.
    if info.inst == s_endpgm():
      completed_wave = list(rwaves_iter[info.wave].pop(0))
@@ -64,13 +64,13 @@ class TestSQTTMapBase(unittest.TestCase):

  def test_rocprof_inst_traces_match(self):
    for name, (events, kern_events, target) in self.examples.items():
-      if "sync" in name and self.target.startswith("gfx12"):
-        self.skipTest("our timestamps are off by a few cycles because rocprof patches timestamps for rdna4 barriers")
      for event in events:
        if not event.itrace: continue
        if event.kern not in kern_events: continue
        with self.subTest(example=name, kern=event.kern):
-          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target)
+          # rocprof OSX has a bug for sopk decoding, linux rocprof works
+          pass_rocprof_err = OSX and target == "gfx1200" and name.startswith("profile_ops")
+          passed_insts, n_waves, n_units = rocprof_inst_traces_match(event, kern_events[event.kern], target, pass_rocprof_err)
          if n_waves: print(f"{name}: passed for {passed_insts} instructions across {n_waves} waves scheduled on {n_units} wave units")

  def test_sqtt_timeline(self):
@@ -78,43 +78,20 @@ class TestSQTTMapBase(unittest.TestCase):
      for event in events:
        if (p:=kern_events.get(event.kern)) is None: continue
        with self.subTest(example=name, kern=event.kern):
-          # skip if there's no SQTT frequency data
-          if not (timeline:=sqtt_timeline(event.blob, p.lib, target)[0]): continue
-          if not (frequency:=[e.key for e in timeline if type(e).__name__ == "ProfilePointEvent" and e.name == "freq_hz"]): continue
-          mean = sum(frequency) / len(frequency)
-          variance = sum((v - mean) ** 2 for v in frequency) / len(frequency)
-          self.assertGreater(mean, 0)
-          if DEBUG >= 2: print(f"{name:20s} SE:{event.se} {mean/1e9:.2f} GHz mean, {variance/1e18:.2f} GHz^2 variance")
-          events = [e for e in timeline if type(e).__name__ == "ProfileRangeEvent"]
+          events = [e for e in sqtt_timeline(event.blob, p.lib, target) if type(e).__name__ == "ProfileRangeEvent"]
          insts, execs = 0, 0
          for e in events:
            if "EXEC" in e.device:
              if "ALT" not in e.name.display_name: execs += 1
            elif "WAVE" in e.device:
              # sopk/immediates don't get ALU/MEM EXEC
-              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE", "BARRIER", "BARRIER_SIGNAL",
-                                             "WAVEEND", "WAVERDY"}: insts += 1
+              if e.name.display_name not in {"IMMEDIATE", "IMMEDIATE_MASK", "JUMP", "JUMP_NO", "MESSAGE"}: insts += 1
            else: raise Exception(f"timeline row must be INST or EXEC, got {e.device}")
          self.assertEqual(execs, insts)

-  def test_wave_sync(self):
-    for name, (events, kern_events, target) in self.examples.items():
-      for event in events:
-        wave_barriers = {}
-        for e in sqtt_timeline(event.blob, kern_events[event.kern].lib, target)[0]:
-          if type(e).__name__ == "ProfileRangeEvent" and e.name.display_name == "BARRIER": wave_barriers.setdefault(e.device, []).append(e)
-        if not wave_barriers: continue
-        for row, events in wave_barriers.items():
-          for e in events:
-            assert e.en-e.st > 1, f"all barriers must have a duration greater than 1, got {e}"
-
 class TestSQTTMapRDNA3(TestSQTTMapBase): target = "gfx1100"

 class TestSQTTMapRDNA4(TestSQTTMapBase): target = "gfx1200"

-class TestSQTTMapCDNA(TestSQTTMapBase):
-  target = "gfx950"
-  def test_rocprof_inst_traces_match(self): self.skipTest("requires timestamp patching to match rocprof, currently it's off by a few cycles")
-
 if __name__ == "__main__":
  unittest.main()
@@ -8,7 +8,6 @@ from tinygrad.engine.schedule import ExecItem
 from tinygrad.uop.ops import Ops
 from tinygrad.renderer import Estimates
 from tinygrad.renderer.ptx import PTXRenderer
-from test.helpers import needs_second_gpu

 class TestArange(unittest.TestCase):
  def _get_flops(self, tensor, desired):
@@ -189,33 +188,6 @@ class TestIndexing(unittest.TestCase):
    for i in idx.flatten().numpy(): expected_grad[i] += 2
    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)

-  @needs_second_gpu
-  @unittest.skipIf(Device.DEFAULT not in ("CPU", "AMD"), "atomics only on AMD/CPU")
-  @Context(USE_ATOMICS=1, SPEC=1)
-  def test_embedding_backward_vocab_sharded(self):
-    from tinygrad.renderer.cstyle import CStyleLanguage
-    if Device.DEFAULT == "CPU" and not isinstance(Device["CPU"].renderer, CStyleLanguage): self.skipTest("CPU needs Clang renderer")
-    devices = (f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1")
-    vocab_size, embed_size = 1000, 128
-    bs, seqlen = 4, 256
-    idx = Tensor.randint(bs, seqlen, high=vocab_size)
-    emb = nn.Embedding(vocab_size, embed_size)
-    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
-    gt = Tensor.zeros(bs, seqlen, embed_size)
-    Tensor.realize(idx, emb.weight, gt)
-    # compute expected grad on single device
-    expected_grad = np.zeros((vocab_size, embed_size), dtype=np.float32)
-    for i in idx.flatten().numpy(): expected_grad[i] += 2
-    # now shard the embedding weight on vocab axis and recompute
-    emb.weight = Tensor.ones(vocab_size, embed_size, requires_grad=True)
-    emb.weight.shard_(devices, axis=0)
-    idx = idx.shard(devices, axis=None)
-    gt = gt.shard(devices, axis=None)
-    Tensor.realize(idx, emb.weight, gt)
-    loss = (emb(idx)-gt).square().sum()
-    loss.backward()
-    np.testing.assert_allclose(emb.weight.grad.numpy(), expected_grad, rtol=1e-5, atol=1e-5)
-
  @unittest.skipUnless(Device.DEFAULT == "AMD" or (Device.DEFAULT == "NULL" and EMULATE.value.startswith("AMD")), "tests AMD bf16 cast overhead")
  def base_test_llama_8b_rope_backward(self, dtype):
    from extra.models.llama import precompute_freqs_cis, apply_rotary_emb
@@ -2,7 +2,7 @@ import unittest
 from tinygrad import Tensor, Device, dtypes, Context
 from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv
-from extra.gemm.cdna_asm_gemm import asm_gemm
+from extra.gemm.asm.cdna.gemm import asm_gemm
 from test.helpers import needs_second_gpu

 # On non CDNA4 it will only validate the Tensor.custom_kernel integration
@@ -81,37 +81,13 @@ class TestGemm(unittest.TestCase):
  @needs_second_gpu
  def test_gemm_k_sharded_3d(self): verify_asm_gemm_k_sharded_3d(1, 64, 32, 2*64, gpus=2)

-# uses the smallest size for the cdna assembly gemm
-class TestAsmGEMM(unittest.TestCase):
+# uses the Asm GEMM on CDNA4 only for speed reasons
+class TestGemmLarge(unittest.TestCase):
  def setUp(self):
    if not is_cdna4():
-      self.skipTest("assembly gemm is only for cdna4")
-
-  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
-
-  def test_verify_with_numpy(self):
-    import numpy as np
-    M, N, K = 256, 256, 64
-    rng = np.random.default_rng(0)
-    a_np = (rng.random((M, K), dtype=np.float32) - 0.5).astype(np.half)
-    b_np = (rng.random((K, N), dtype=np.float32) - 0.5).astype(np.half)
-    c_np = a_np @ b_np
-    a, b = Tensor(a_np), Tensor(b_np)
-    c = asm_gemm(a, b)
-    c.realize()
-    # no validation on the NULL device
-    if a.device.startswith("NULL"): return None
-    np.testing.assert_allclose(c.numpy(), c_np, atol=2e-3, rtol=5e-2)
-
-# test the Asm GEMM with Llama shapes, only run on the real machine for speed
-class TestGemmLlama(unittest.TestCase):
-  def setUp(self):
-    if not is_cdna4() or getenv("MOCKGPU"):
      self.skipTest("very slow on non mi350x")

-  @Context(ASM_GEMM=1)
-  def test_empty(self): (Tensor.empty(N:=getenv("N", 4096), N, dtype=dtypes.half)@Tensor.empty(N, N, dtype=dtypes.half)).realize()
-
+  def test_tiny(self): verify_asm_gemm(1, 256, 256, 64)
  def test_simple(self): verify_asm_gemm(1, N:=getenv("N", 4096), N, N, dtype=dtypes.half)
  def test_gemm(self): verify_asm_gemm(1, 8192, 4096, 14336)
  def test_gemm_batched(self): verify_asm_gemm(2, 8192, 4096, 4096)
@@ -181,7 +157,7 @@ class TestGemmLlama(unittest.TestCase):

 class TestMagicGu(unittest.TestCase):
  def test_magicgu_matches_old(self):
-    from extra.gemm.cdna_asm_gemm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
+    from extra.gemm.asm.cdna.asm import _magicgu_mulhi, TILE_M, TILE_N, TILE_K
    old_iters_args = {64: (67108864, 0), 128: (33554432, 0), 224: (613566757, 2147483656)}
    old_gemm_shapes = [
      (8192, 4096, 4096), (8192, 14336, 4096), (8192, 4096, 14336),
@@ -283,30 +283,5 @@ class TestCustomKernel(unittest.TestCase):
    self.assertIsNotNone(custom_idx, "custom_addmul kernel not found in schedule")
    self.assertEqual(custom_idx, 3, f"custom_addmul should be at index 3, got {custom_idx}")

-  def test_anonymous_buffers_in_function(self):
-    """Test that custom kernels with anonymous output buffers work inside @function."""
-    a = Tensor.full((4, 4), 3.).contiguous()
-    b = Tensor.full((4, 4), 2.).contiguous()
-    Tensor.realize(a, b)
-
-    def custom_add_with_tmp(o1:UOp, o2:UOp, A:UOp, B:UOp) -> UOp:
-      o1,o2,A,B = o1.flatten(), o2.flatten(), A.flatten(), B.flatten()
-      i = UOp.range(o1.size, 0)
-      store_o1 = o1[i].store(A[i]+B[i])
-      store_o2 = o2[i].store(A[i]+B[i]+2)
-      return UOp.group(store_o1, store_o2).end(i).sink(arg=KernelInfo(name=f"add_with_tmp_{o1.size}")).simplify()
-
-    from tinygrad import function
-    @function(precompile=True)
-    def run(x:Tensor, w:Tensor) -> Tensor:
-      out = Tensor.invalid(*x.shape, dtype=x.dtype)
-      tmp = Tensor.invalid(*x.shape, dtype=x.dtype)
-      out, tmp = Tensor.custom_kernel(out, tmp, x, w, fxn=custom_add_with_tmp)[:2]
-      return out+tmp
-
-    result = run(a, b).flatten().tolist()
-    expected = (3+2)*2+2
-    assert all(x == expected for x in result), f"expected all {expected}, got {result}"
-
 if __name__ == '__main__':
  unittest.main()
@@ -10,7 +10,7 @@ from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Context, Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
-from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX, FP8E4M3FNUZ_MAX, FP8E5M2FNUZ_MAX
+from test.unit.test_dtype_spec import _assert_eq, core_dtypes, dtype_ints, dtype_floats, FP8E4M3_MAX, FP8E5M2_MAX
 import pytest
 pytestmark = pytest.mark.filterwarnings("ignore")

@@ -18,10 +18,11 @@ settings.register_profile("my_profile", max_examples=200, deadline=None, derando
 settings.load_profile("my_profile")

 def get_available_cast_dtypes(dtype: DType) -> List[DType]:
-  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) or v in dtypes.fp8s+(dtypes.half,dtypes.bfloat16,dtypes.long)]
-  if dtype in (dtypes.long, dtypes.ulong) and (not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes)):
-    return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
-  if not is_dtype_supported(dtype) and dtype not in dtypes.fp8s+(dtypes.half,dtypes.bfloat16): return []
+  # dont cast internal dtypes
+  dts = [v for k, v in DTYPES_DICT.items() if v != dtype and is_dtype_supported(v) and not k.startswith("_")]
+  if not is_dtype_supported(dtype) or dtypes.long in EMULATED_DTYPES.tolist(dtypes):
+    if dtype in (dtypes.long, dtypes.ulong): return [dt for dt in dts if dt != dtypes.double] # can't bitcast with no 64-bit support
+    else: return []
  return dts

 def _to_torch_storage_type(dtype:DType):
@@ -59,8 +60,10 @@ class TestDType(unittest.TestCase):
  DATA: Any = None
  @classmethod
  def setUpClass(cls):
-    if cls.DTYPE is None: raise unittest.SkipTest("base class")
-    cls.DATA = rand_for_dtype(cls.DTYPE, 0x10, allow_subnormal=is_dtype_supported(cls.DTYPE))
+    if not cls.DTYPE or not is_dtype_supported(cls.DTYPE): raise unittest.SkipTest("dtype not supported")
+    cls.DATA = rand_for_dtype(cls.DTYPE, 10)
+  def setUp(self):
+    if self.DTYPE is None: raise unittest.SkipTest("base class")

  def test_to_np(self):
    _test_to_np(Tensor(self.DATA, dtype=self.DTYPE), _to_np_dtype(self.DTYPE), np.array(self.DATA, dtype=_to_np_dtype(self.DTYPE)))
@@ -98,14 +101,14 @@ class TestDType(unittest.TestCase):
  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
  def test_uint_overflow(self):
    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
-    v = self.DTYPE.max
+    v = dtypes.max(self.DTYPE)
    _test_to_np(Tensor(v, dtype=self.DTYPE)+2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))+2)
    _test_to_np(Tensor(v, dtype=self.DTYPE)*2, _to_np_dtype(self.DTYPE), np.array(v, dtype=_to_np_dtype(self.DTYPE))*2)

  def test_dtypes_DTYPES_DICT(self):
    self.assertIn("float", DTYPES_DICT)
    self.assertIn("float32", DTYPES_DICT)
-    self.assertEqual(len(DTYPES_DICT), 28)
+    self.assertEqual(len(DTYPES_DICT), 26)
    self.assertTrue(all(isinstance(value, DType) for value in DTYPES_DICT.values()))
    self.assertTrue(all(issubclass(_to_np_dtype(value), np.generic) for value in DTYPES_DICT.values() if _to_np_dtype(value) is not None))

@@ -129,6 +132,7 @@ class TestDType(unittest.TestCase):

 def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
  target_dtype = target_dtype or least_upper_dtype(a_dtype, b_dtype)
+  if not is_dtype_supported(a_dtype) or not is_dtype_supported(b_dtype) or not is_dtype_supported(target_dtype): return
  if a_dtype == dtypes.bool or b_dtype == dtypes.bool: return
  _assert_eq(Tensor([1,2,3,4], dtype=a_dtype)+Tensor([1,2,3,4], dtype=b_dtype), target_dtype, [2,4,6,8])
  _assert_eq((Tensor([1], dtype=a_dtype).cast(b_dtype)+Tensor([1], dtype=a_dtype).cast(b_dtype)).cast(a_dtype), a_dtype, [2])
@@ -139,8 +143,6 @@ def _test_ops(a_dtype:DType, b_dtype:DType, target_dtype=None):
 class TestFp8s(unittest.TestCase):
  def test_fp8e4m3_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3).dtype == dtypes.fp8e4m3
  def test_fp8e5m2_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2).dtype == dtypes.fp8e5m2
-  def test_fp8e4m3fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e4m3fnuz).dtype == dtypes.fp8e4m3fnuz
-  def test_fp8e5m2fnuz_creation(self): assert Tensor([-1, 1, 2], dtype=dtypes.fp8e5m2fnuz).dtype == dtypes.fp8e5m2fnuz

 class TestFp8sConversions(unittest.TestCase):
  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3_MAX, max_value=FP8E4M3_MAX))
@@ -167,30 +169,7 @@ class TestFp8sConversions(unittest.TestCase):
  def test_fp8e5m2_to_float(self, x):
    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2).float().item())

-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E4M3FNUZ_MAX, max_value=FP8E4M3FNUZ_MAX))
-  def test_float_to_fp8e4m3fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e4m3fnuz_extreme_values(self):
-    for x in [FP8E4M3FNUZ_MAX, FP8E4M3FNUZ_MAX*1.01, -FP8E4M3FNUZ_MAX, -FP8E4M3FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.float8_e4m3fnuz).view(torch.uint8).item())
-
-  @given(strat.floats(width=32, allow_subnormal=True, allow_nan=False, allow_infinity=False, min_value=-FP8E5M2FNUZ_MAX, max_value=FP8E5M2FNUZ_MAX))
-  def test_float_to_fp8e5m2fnuz(self, x):
-    np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  def test_float_to_fp8e5m2fnuz_extreme_values(self):
-    for x in [FP8E5M2FNUZ_MAX, FP8E5M2FNUZ_MAX*1.01, -FP8E5M2FNUZ_MAX, -FP8E5M2FNUZ_MAX*1.01, math.inf, -math.inf, math.nan, 0.0, -0.0]:
-      np.testing.assert_equal(float_to_fp8(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.float8_e5m2fnuz).view(torch.uint8).item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e4m3fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e4m3fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e4m3fnuz).float().item())
-
-  @given(strat.integers(min_value=0, max_value=255))
-  def test_fp8e5m2fnuz_to_float(self, x):
-    np.testing.assert_equal(fp8_to_float(x, dtypes.fp8e5m2fnuz), torch.tensor(x, dtype=torch.uint8).view(torch.float8_e5m2fnuz).float().item())
-
+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16(unittest.TestCase):
  def test_bf16_creation_numpy(self):
    data = [-1, 1, 2]
@@ -210,6 +189,7 @@ class TestBFloat16(unittest.TestCase):
    assert t.dtype == dtypes.bfloat16
    np.testing.assert_allclose(t.numpy(), np.eye(3))

+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16), "bfloat16 not supported")
 class TestBFloat16DType(unittest.TestCase):
  def test_bf16_to_float(self):
    _test_cast(Tensor([100000], dtype=dtypes.bfloat16), dtypes.float32)
@@ -223,6 +203,7 @@ class TestBFloat16DType(unittest.TestCase):
    back = t.cast(dtypes.float32)
    assert tuple(back.numpy().tolist()) == (9984., -1, -1000, -9984, 20)

+@unittest.skipUnless(is_dtype_supported(dtypes.bfloat16) and is_dtype_supported(dtypes.float16), "bfloat16 or float16 not supported")
 class TestBFloat16DTypeCast(unittest.TestCase):
  def test_f16_to_bf16_conversion(self):
    original_tensor = Tensor([1.0, 2.0, 3.0], dtype=dtypes.float16)
@@ -271,7 +252,6 @@ class TestFloatDType(TestDType):
    _test_op(lambda: Tensor([-0.9, -0.3, 1.2], dtype=dtypes.float32).cast(dtypes.uint32), dtypes.uint32,
             [0, 0, 1])

-@unittest.skipUnless(is_dtype_supported(dtypes.double), f"no double on {Device.DEFAULT}")
 class TestDoubleDType(TestDType):
  DTYPE = dtypes.double
  @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
@@ -510,3 +490,4 @@ class TestOpsBFloat16(unittest.TestCase):

 if __name__ == '__main__':
  unittest.main()
+
@@ -52,8 +52,6 @@ class ht:
 ht.bfloat16 = ht.uint16.filter(lambda x: ((x >> 7) & 0xFF) != 0)  # filter subnormal bfloat16
 ht.fp8e4m3 = ht.uint8
 ht.fp8e5m2 = ht.uint8
-ht.fp8e4m3fnuz = ht.uint8
-ht.fp8e5m2fnuz = ht.uint8

 def universal_test(a, b, dtype, op):
  if not isinstance(op, tuple): op = (op, op)
@@ -69,8 +67,7 @@ def universal_test(a, b, dtype, op):
    if not is_dtype_supported(dtype) or dtype in EMULATED_DTYPES.tolist(dtypes): # denormals are zero
      fe, fm = dtypes.finfo(dtype)
      atol, rtol = 2 ** (2 - (1 << (fe - 1))), 2 ** (-fm)
-    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1),
-                        dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz:(5e-1, 5e-1)}.get(dtype, (1e-10, 1e-7))
+    else: atol, rtol = {dtypes.bfloat16:(1e-3, 1e-2), dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2:(1.0, 5e-1)}.get(dtype, (1e-10, 1e-7))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -81,8 +78,8 @@ def universal_test_unary(a, dtype, op):
  if op[0] == Tensor.cos and abs(a) > 30: return
  if op[0] == Tensor.log and a <= 0: return
  if dtype in dtypes.fp8s:
-    # denormals are zero
-    if dtype in EMULATED_DTYPES.tolist(dtypes) or not is_dtype_supported(dtype) and abs(ta.numpy().item()) < 0.015625: return
+    # normals are zero
+    if dtype in EMULATED_DTYPES.tolist(dtypes) and abs(ta.numpy().item()) < 0.015625: return
    tensor_value = fp8_to_float(op[0](ta.realize()).bitcast(dtypes.uint8).item(), dtype)
    numpy_value = truncate[dtype](v:=op[1](ta.numpy()).item())
    # cuda cast f32 inf to f8 MAX, amd cast it to nan(E4M3)/inf(E5M2)
@@ -90,8 +87,7 @@ def universal_test_unary(a, dtype, op):
  else: tensor_value, numpy_value = op[0](ta).numpy(), op[1](ta.numpy())
  if dtype in dtypes.floats:
    atol, rtol = { dtypes.float16:(1e-3, 1e-2), dtypes.bfloat16:(1e-3, 2e-2),
-      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1),
-      dtypes.fp8e4m3fnuz:(1e-1, 1e-1), dtypes.fp8e5m2fnuz: (5e-1, 5e-1)}.get(dtype, (1e-6, 1e-5))
+      dtypes.fp8e4m3:(1e-1, 1e-1), dtypes.fp8e5m2: (1.0, 5e-1)}.get(dtype, (1e-6, 1e-5))
    np.testing.assert_allclose(tensor_value, numpy_value, atol=atol, rtol=rtol)
  else: np.testing.assert_equal(tensor_value, numpy_value)

@@ -159,26 +155,6 @@ class TestDTypeALU(unittest.TestCase):
  def test_emulated_fp8e5m2(self, a, b, op):
    universal_test(from_storage_scalar(a, dtypes.fp8e5m2), from_storage_scalar(b, dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  def test_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, ht.fp8e4m3fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e4m3fnuz), from_storage_scalar(b, dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, ht.fp8e5m2fnuz, strat.sampled_from(binary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz(self, a, b, op):
-    universal_test(from_storage_scalar(a, dtypes.fp8e5m2fnuz), from_storage_scalar(b, dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
  @given(ht.float32, strat.sampled_from(unary_operations))
  def test_float32_unary(self, a, op): universal_test_unary(a, dtypes.float32, op)

@@ -222,30 +198,6 @@ class TestDTypeALU(unittest.TestCase):
    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2) != 0.0)
    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2), dtypes.fp8e5m2, op)

-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e4m3fnuz), f"no fp8e4m3fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.fp8e5m2fnuz), f"no fp8e5m2fnuz on {Device.DEFAULT}")
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  def test_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
-  @given(ht.fp8e4m3fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e4m3fnuz")
-  def test_emulated_fp8e4m3fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e4m3fnuz), dtypes.fp8e4m3fnuz, op)
-
-  @given(ht.fp8e5m2fnuz, strat.sampled_from(unary_operations))
-  @Context(EMULATED_DTYPES="fp8e5m2fnuz")
-  def test_emulated_fp8e5m2fnuz_unary(self, a, op):
-    if op[1] == np.reciprocal: assume(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz) != 0.0)
-    universal_test_unary(from_storage_scalar(a, dtype=dtypes.fp8e5m2fnuz), dtypes.fp8e5m2fnuz, op)
-
  @given(ht.uint8, ht.uint8, strat.sampled_from(integer_binary_operations))
  def test_uint8(self, a, b, op): universal_test(a, b, dtypes.uint8, op)

@@ -366,7 +318,7 @@ class TestDTypeALU(unittest.TestCase):

  @unittest.expectedFailure
  def test_unsafe_cast_float_to_int_failure(self):
-    val = float(dtypes.int32.max - 1)
+    val = float(dtypes.max(dtypes.int32) - 1)
    t1 = Tensor([val], dtype=dtypes.float32).cast(dtypes.int32)
    t2 = Tensor(val, dtype=dtypes.float32).cast(dtypes.int32)
    np.testing.assert_equal(t1.item(), t2.item())
@@ -1,17 +1,186 @@
 import unittest
 import numpy as np
 from tinygrad import Device, dtypes, Tensor, Context
-from tinygrad.helpers import unwrap
+from tinygrad.device import LRUAllocator, is_dtype_supported
+from tinygrad.dtype import ImageDType
+from tinygrad.helpers import prod, unwrap

 IMAGE_SUPPORTED_DEVICES = ("QCOM", "CL")

+@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
+class TestImageCopy(unittest.TestCase):
+  def test_image_copyout_1x8(self, img_type=dtypes.imagef):
+    it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
+    buf = it.uop.buffer
+    out = buf.as_memoryview()
+    np.testing.assert_equal(out.cast(it.dtype.fmt).tolist(), np.arange(32))
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.half, device="PYTHON"), "need half")
+  def test_imageh_copyout_1x8(self): self.test_image_copyout_1x8(img_type=dtypes.imageh)
+
+  def test_image_numpy_1x8(self, img_type=dtypes.imagef):
+    it = Tensor.arange(32).cast(img_type((1,8,4))).realize()
+    np.testing.assert_equal(it.numpy(), np.arange(32))
+  def test_imageh_numpy_1x8(self): self.test_image_numpy_1x8(img_type=dtypes.imageh)
+
+  def test_image_copyout_2x4(self):
+    it = Tensor.arange(2*4*4).cast(dtypes.imagef((2,4,4))).realize()
+    buf = it.uop.buffer
+    out = buf.as_memoryview()
+    np.testing.assert_equal(out.cast('f').tolist(), np.arange(2*4*4))
+
+  def test_image_roundtrip(self):
+    sz = (4,2,4)
+    it = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
+    buf = it.uop.buffer
+    out = buf.as_memoryview()
+
+    it2 = Tensor.rand(prod(sz)).cast(dtypes.imagef(sz)).realize()
+    buf2 = it2.uop.buffer
+    buf2.copyin(out)
+
+    assert (it == it2).sum().item() == prod(sz)
+
@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
 class TestImageDType(unittest.TestCase):
+  def test_image_pitch(self):
+    def __validate(imgdt, expected_pitch):
+      assert imgdt.pitch == expected_pitch, f"Failed pitch for image: {imgdt}. Got 0x{imgdt.pitch:X}, expected 0x{expected_pitch:X}"
+
+    # Match opencl pitches for perf
+    __validate(dtypes.imageh((1, 201)), 0x680)
+    __validate(dtypes.imageh((16, 216)), 0x700)
+    __validate(dtypes.imageh((16, 9)), 0x80)
+    __validate(dtypes.imageh((48, 64)), 0x200)
+    __validate(dtypes.imageh((32, 128)), 0x400)
+    __validate(dtypes.imageh((96, 128)), 0x400)
+    __validate(dtypes.imageh((64, 256)), 0x840)
+    __validate(dtypes.imageh((64, 9)), 0x80)
+    __validate(dtypes.imageh((192, 256)), 0x840)
+    __validate(dtypes.imageh((64, 768)), 0x1840)
+    __validate(dtypes.imageh((256, 49)), 0x1C0)
+    __validate(dtypes.imageh((128, 9)), 0x80)
+    __validate(dtypes.imageh((16, 1024)), 0x2080)
+    __validate(dtypes.imageh((64, 512)), 0x1040)
+    __validate(dtypes.imageh((16, 512)), 0x1080)
+    __validate(dtypes.imageh((132, 64)), 0x200)
+    __validate(dtypes.imageh((4, 512)), 0x1200)
+    __validate(dtypes.imageh((8, 512)), 0x1100)
+    __validate(dtypes.imageh((128, 128)), 0x400)
+    __validate(dtypes.imageh((32, 512)), 0x1040)
+    __validate(dtypes.imageh((26, 64)), 0x200)
+    __validate(dtypes.imageh((32, 516)), 0x1040)
+    __validate(dtypes.imageh((32, 1024)), 0x2040)
+    __validate(dtypes.imageh((16, 2048)), 0x4080)
+    __validate(dtypes.imageh((8, 2048)), 0x4100)
+    __validate(dtypes.imageh((4, 4096)), 0x8200)
+
+    __validate(dtypes.imagef((16, 49)), 0x380)
+    __validate(dtypes.imagef((16, 1024)), 0x4080)
+    __validate(dtypes.imagef((256, 64)), 0x400)
+    __validate(dtypes.imagef((64, 512)), 0x2040)
+    __validate(dtypes.imagef((16, 512)), 0x2080)
+    __validate(dtypes.imagef((132, 64)), 0x400)
+    __validate(dtypes.imagef((4, 512)), 0x2200)
+    __validate(dtypes.imagef((4, 16)), 0x200)
+    __validate(dtypes.imagef((2, 16)), 0x400)
+    __validate(dtypes.imagef((8, 512)), 0x2100)
+    __validate(dtypes.imagef((12, 64)), 0x400)
+    __validate(dtypes.imagef((3, 32)), 0x400)
+    __validate(dtypes.imagef((128, 128)), 0x840)
+    __validate(dtypes.imagef((32, 512)), 0x2040)
+    __validate(dtypes.imagef((8, 3072)), 0xC100)
+    __validate(dtypes.imagef((4, 2048)), 0x8200)
+    __validate(dtypes.imagef((4, 1024)), 0x4200)
+    __validate(dtypes.imagef((4, 4096)), 0x10200)
+    __validate(dtypes.imagef((10, 384)), 0x1900)
+    __validate(dtypes.imagef((24, 64)), 0x400)
+    __validate(dtypes.imagef((128, 12)), 0xC0)
+    __validate(dtypes.imagef((10, 24)), 0x200)
+    __validate(dtypes.imagef((1, 129)), 0x840)
+    __validate(dtypes.imagef((1, 32)), 0x200)
+    __validate(dtypes.imagef((1, 64)), 0x400)
+    __validate(dtypes.imagef((1, 1239)), 0x4D80)
+    __validate(dtypes.imagef((1, 1)), 0x40)
+
+  def test_image_and_back(self):
+    data = Tensor.randn(9*32*4).realize()
+    tst = data.numpy()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    assert isinstance(it.uop.base.realized.dtype, ImageDType)
+    np.testing.assert_equal(tst, it.numpy())
+
+  def test_image_cast_and_back_collapses(self):
+    data = Tensor.randn(9*27*4).realize()
+    tst = data.numpy()
+    it = data.cast(dtypes.imagef((9,27,4))).realize()
+    # the underlying UOp is identical
+    #self.assertIs(it.uop.base.realized, data.uop.base.realized)
+    np.testing.assert_equal(tst, it.numpy())
+
+  def test_image_and_back_wrong_shape(self):
+    data = Tensor.randn(9*27*4).realize()
+    tst = data.numpy()
+    it = data.cast(dtypes.imagef((9,12,4))).realize()
+    assert not isinstance(it.uop.base.realized.dtype, ImageDType)
+    np.testing.assert_equal(tst, it.numpy())
+
+  def test_shrink_load_float(self):
+    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
+    imgv = it.numpy()
+    np.testing.assert_equal(imgv[0:2], it[0:2].numpy())
+
+  def test_mul_stays_image(self):
+    # NOTE: contiguous is needed otherwise this folds
+    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).contiguous().realize()
+    out = (it*2).realize()
+    assert isinstance(out.uop.base.realized.dtype, ImageDType)
+
+  def test_sum(self):
+    it = Tensor.rand(8).cast(dtypes.imagef((1,2,4))).realize()
+    itn = it.numpy()
+    np.testing.assert_allclose(np.sum(itn), it.sum().numpy(), rtol=1e-6)
+
+  def test_shrink_max(self):
+    it = Tensor.randn(16).cast(dtypes.imagef((1,4,4))).realize()
+    imgv = it.numpy()
+    np.testing.assert_equal(np.maximum(imgv[0:3], 0), it[0:3].relu().numpy())
+
+  def test_shrink_to_float(self):
+    it = Tensor.randn(4, 4).cast(dtypes.imagef((1,4,4))).realize()
+    imgv = it.numpy()
+    np.testing.assert_equal(np.maximum(imgv[:, 0], 0), it[:, 0].relu().numpy())
+
+  @unittest.skipUnless(isinstance(Device.default.allocator, LRUAllocator), "Requires LRU")
+  def test_lru_alloc(self):
+    data = Tensor.randn(9*27*4).realize()
+    it = data.cast(dtypes.imagef((9,27,4))).realize()
+    b1 = it.uop.base.realized._buf
+    del it
+    it = data.cast(dtypes.imagef((9,27,4))).realize()
+    assert it.uop.base.realized._buf == b1
+
+  def test_no_lru_alloc(self):
+    data = Tensor.randn(9*32*4).realize()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    b1 = it.uop.base.realized._buf
+    del it
+    it = data.reshape(9,32,4).pad_to(10, None, None).cast(dtypes.imagef((10,32,4))).contiguous().realize()
+    assert it.uop.base.realized._buf != b1
+
+  def test_no_lru_alloc_dtype(self):
+    data = Tensor.randn(9*32*4).realize()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    b1 = it.uop.base.realized._buf
+    del it
+    it = data.cast(dtypes.imageh((9,32,4))).realize()
+    assert it.uop.base.realized._buf != b1
+
  # issue caused by: don't realize image to image casts. this is part of a larger problem
  #@unittest.expectedFailure
  # update: passing after tensor_map
  def test_lil_model(self):
-    with Context(IMAGE=1):
+    with Context(IMAGE=2):
      x = Tensor.zeros(1, 1)
      w1 = Tensor.zeros(1, 8, requires_grad=True)
      w2 = Tensor.zeros(8, 2)
@@ -29,5 +198,39 @@ class TestImageDType(unittest.TestCase):
      self.assertEqual(w1.grad.uop.base.buffer.dtype, dtypes.float32)
      self.assertEqual(len(sched), 9)

+@unittest.skipUnless(Device.DEFAULT in IMAGE_SUPPORTED_DEVICES, "Images not supported")
+class TestImageRealization(unittest.TestCase):
+  def test_image_dtype_expand(self):
+    data = Tensor.randn(9*32*4).realize()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
+    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous().realize()
+    self.assertEqual(it_expanded.dtype, dtypes.float32)
+
+  def test_image_dtype_expand_and_back(self):
+    data = Tensor.randn(9*32*4).realize()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
+    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4))
+    it2 = it_expanded.sum(3).realize()
+    self.assertEqual(it2.dtype, dtypes.imagef((9,32,4)))
+
+  def test_image_alu_children(self):
+    data = Tensor.randn(9*32*4).realize()
+    it = data.cast(dtypes.imagef((9,32,4))).contiguous().realize()
+    self.assertEqual(it.dtype, dtypes.imagef((9,32,4)))
+    it_expanded = it.reshape((9,32,4,1)).expand((9,32,4,4)).contiguous()
+    alu1 = it_expanded+1
+    alu2 = it_expanded.sum(3)
+    it_expanded.realize()
+    # NOTE: the parent becomes float, but the alu child will stay image until its output cannot fit the image
+    self.assertEqual(alu1.dtype, dtypes.imagef((9,32,4)))
+    alu1.realize()
+    self.assertEqual(alu1.dtype, dtypes.float32)
+    # alu2 is back in image because it fits the dtype again
+    self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
+    alu2.realize()
+    self.assertEqual(alu2.dtype, dtypes.imagef((9,32,4)))
+
 if __name__ == '__main__':
  unittest.main()
@@ -39,18 +39,6 @@ class TestJit(unittest.TestCase):
    def add(a, b): return (a+b).realize()
    _simple_test(add)

-  def test_jitbeam_triggers_beam(self):
-    from unittest.mock import patch
-    from tinygrad.helpers import getenv as _getenv
-    @TinyJit
-    def add(a, b): return (a+b).realize()
-    a, b = Tensor.ones(10, 10).contiguous().realize(), Tensor.ones(10, 10).contiguous().realize()
-    with patch("tinygrad.codegen.opt.search.beam_search", wraps=lambda k,*a,**kw: k) as mock_beam:
-      add(a, b)
-      assert mock_beam.call_count == 0
-      with patch("tinygrad.engine.jit.getenv", side_effect=lambda k, d=0: 1 if k == "JITBEAM" else _getenv(k, d)): add(a, b)
-      assert mock_beam.call_count == 1
-
  def test_simple_jit_reset(self):
    @TinyJit
    def add(a, b): return (a+b).realize()
@@ -630,7 +618,7 @@ class TestJitFree(unittest.TestCase):

    expected_savings = (len(inp) * inp.dtype.itemsize * 2) + dtypes.float32.itemsize # (t1 and t2) + out

-    self.assertGreaterEqual(savings_after_free, expected_savings)
+    self.assertEqual(savings_after_free, expected_savings)
    out = fxn(Tensor([11,1,2,3,4]))
    self.assertEqual(out.item(), 136)

@@ -640,7 +628,7 @@ class TestJitFree(unittest.TestCase):
    fxn.captured.free_intermediates() # 2nd time to validate
    savings_after_free = pre_free - GlobalCounters.mem_used

-    self.assertGreaterEqual(savings_after_free, expected_savings)
+    self.assertEqual(savings_after_free, expected_savings)
    out = fxn(Tensor([11,1,2,3,4]))
    self.assertEqual(out.item(), 136)

@@ -660,6 +648,25 @@ class TestJitFree(unittest.TestCase):
    fxn(Tensor([2]))
    self.assertEqual(x.item(), 8)

+  def test_replan_buffers_memory_layout(self):
+    if not hasattr(Device[Device.DEFAULT].allocator, '_offset'): raise unittest.SkipTest("replan_buffers_memory_layout useless")
+
+    ext_tensor = Tensor([1,24,23,45,1]).contiguous()
+    ext_tensor_2 = Tensor([2,2,2,2,2]).contiguous()
+    @TinyJit
+    def fxn(x:Tensor):
+      out = (x*ext_tensor_2+ext_tensor).reshape(5,1).expand(5, 100).contiguous()
+      return out.sum()
+    for i in range(5):
+      out = fxn(Tensor([i,1,2,3,4]))
+      self.assertEqual(out.item(), 11400+200*i)
+    self.assertEqual(len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])), 4)
+    fxn.captured.replan_buffers_memory_layout()
+    self.assertEqual(len(set([b.base for item in fxn.captured.jit_cache for b in item.bufs if b is not None])), 2)
+
+    out = fxn(Tensor([11,1,2,3,4]))
+    self.assertEqual(out.item(), 13600)
+
 class TestJitGraphSplit(unittest.TestCase):
  def compute(self, device, inp):
    assert inp.device == device, f"Input device {inp.device} does not match expected {device}"
@@ -12,18 +12,18 @@ class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
    c0 = UOp(Ops.PARAM, dtypes.uchar.ptr(4014080), arg=0, src=())
-    c1 = UOp.range(UOp.const(dtypes.weakint, 512), 0, AxisType.GLOBAL)
-    c2 = UOp.range(UOp.const(dtypes.weakint, 784), 1, AxisType.GLOBAL)
-    c3 = UOp.range(UOp.const(dtypes.weakint, 10), 3, AxisType.GLOBAL)
+    c1 = UOp.range(UOp.const(dtypes.index, 512), 0, AxisType.GLOBAL)
+    c2 = UOp.range(UOp.const(dtypes.index, 784), 1, AxisType.GLOBAL)
+    c3 = UOp.range(UOp.const(dtypes.index, 10), 3, AxisType.GLOBAL)
    c4 = UOp(Ops.PARAM, dtypes.int.ptr(512), arg=1, src=())
    c5 = c4.index(c1.valid(UOp.const(dtypes.bool, True)))
-    c6 = UOp.range(UOp.const(dtypes.weakint, 6000), 1004, AxisType.REDUCE)
-    c7 = UOp.range(UOp.const(dtypes.weakint, 3750), 2006, AxisType.REDUCE)
-    c8 = UOp.range(UOp.const(dtypes.weakint, 16), 2007, AxisType.GROUP_REDUCE)
+    c6 = UOp.range(UOp.const(dtypes.index, 6000), 1004, AxisType.REDUCE)
+    c7 = UOp.range(UOp.const(dtypes.index, 3750), 2006, AxisType.REDUCE)
+    c8 = UOp.range(UOp.const(dtypes.index, 16), 2007, AxisType.GROUP_REDUCE)
    c9 = UOp(Ops.PARAM, dtypes.uchar.ptr(47040000), arg=2, src=())
-    c10 = c9.index((((c3*UOp.const(dtypes.weakint, 4704000))+c2)+(c6*UOp.const(dtypes.weakint, 784))).valid(UOp.const(dtypes.bool, True)))
-    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.weakint, 6000))+c6)+((c7*UOp.const(dtypes.weakint, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.weakint, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
-    c12 = c0.index((((c1*UOp.const(dtypes.weakint, 7840))+(c2*UOp.const(dtypes.weakint, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
+    c10 = c9.index((((c3*UOp.const(dtypes.index, 4704000))+c2)+(c6*UOp.const(dtypes.index, 784))).valid(UOp.const(dtypes.bool, True)))
+    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(dtypes.index, 6000))+c6)+((c7*UOp.const(dtypes.index, 16))+c8)).alu(Ops.CMPLT, UOp.const(dtypes.index, 59999)).where(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(dtypes.int, -1))).where(UOp.const(dtypes.uchar, 0), c10).reduce(c6, arg=Ops.ADD)
+    c12 = c0.index((((c1*UOp.const(dtypes.index, 7840))+(c2*UOp.const(dtypes.index, 10)))+c3).valid(UOp.const(dtypes.bool, True))).store(c11).end(c1, c2, c3)
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = get_program(ast, Device["METAL"].renderer)

@@ -110,7 +110,7 @@ class TestMultiTensor(unittest.TestCase):
  def test_tensor_from_multi(self):
    X = Tensor([1, 2], dtype=dtypes.int).shard_(devices_2, 0)
    Y = Tensor(X.uop)
-    self.assertEqual(Y.device, devices_2)
+    self.assertEqual(Y.device, Device.DEFAULT)
    np.testing.assert_equal(X.numpy(), Y.numpy())

    with self.assertRaises(AssertionError):
@@ -645,7 +645,6 @@ class TestMultiTensor(unittest.TestCase):
    out = t0.flip(0) + 1
    self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)))

-  @unittest.skip("flaky")
  def test_reshape_on_axis(self):
    t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)

@@ -1130,51 +1129,6 @@ class TestTensorOps(unittest.TestCase):
  def test_bitcast(self):
    helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))

-@unittest.skipIf(not_support_multi_device(), "need multi")
-class TestMultiBufferView(unittest.TestCase):
-  @needs_second_gpu
-  def setUp(self): pass
-
-  def _check(self, a_ref:Tensor, a_multi:Tensor, view_fn):
-    """Apply view_fn to both, verify zero compiled kernels and matching values."""
-    b_ref = view_fn(a_ref)
-    b_multi = view_fn(a_multi).contiguous()
-    sched = b_multi.schedule()
-    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0, f"expected zero compiled kernels, got {len(compiled)}")
-    run_schedule(sched)
-    np.testing.assert_equal(b_multi.numpy(), b_ref.numpy())
-
-  @unittest.skip("flaky on LLVM")
-  def test_shrink_non_shard_axis(self):
-    ref = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().realize()
-    a = Tensor.arange(8*4*10).reshape(8, 4, 10).contiguous().shard(devices_2, axis=1).realize()
-    self._check(ref, a, lambda t: t[3])
-
-  def test_shrink_2d(self):
-    ref = Tensor.arange(6*4).reshape(6, 4).contiguous().realize()
-    a = Tensor.arange(6*4).reshape(6, 4).contiguous().shard(devices_2, axis=1).realize()
-    self._check(ref, a, lambda t: t.shrink(((1, 4), None)))
-
-  def test_reshape_then_shrink(self):
-    ref = Tensor.arange(8*6).reshape(8, 6).contiguous().realize()
-    a = Tensor.arange(8*6).reshape(8, 6).contiguous().shard(devices_2, axis=1).realize()
-    self._check(ref, a, lambda t: t.reshape(4, 2, 6)[1])
-
-  def test_chained_shrink(self):
-    ref = Tensor.arange(10*8).reshape(10, 8).contiguous().realize()
-    a = Tensor.arange(10*8).reshape(10, 8).contiguous().shard(devices_2, axis=1).realize()
-    self._check(ref, a, lambda t: t.shrink(((2, 8), None)).shrink(((1, 4), None)))
-
-  def test_4_devices(self):
-    ref = Tensor.arange(8*12).reshape(8, 12).contiguous().realize()
-    a = Tensor.arange(8*12).reshape(8, 12).contiguous().shard(devices_4, axis=1).realize()
-    sched = a[5].contiguous().schedule()
-    compiled = [si for si in sched if isinstance(si.prg, CompiledRunner)]
-    self.assertEqual(len(compiled), 0)
-    run_schedule(sched)
-    np.testing.assert_equal(a[5].contiguous().numpy(), ref[5].numpy())
-
@unittest.skipIf(not_support_multi_device(), "need multi")
 class TestMultiFromUnrenderable(unittest.TestCase):
  @needs_second_gpu
@@ -479,9 +479,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -496,9 +496,9 @@ class TestOps(unittest.TestCase):
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], 3.])
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[1., 0., 3., -4.], [-1., -2., 3., 0.]])
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.max], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.max(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum,
-                   vals=[[-1234, 0, 1234, dtypes.int.max, dtypes.int.min], dtypes.int.min], forward_only=True)
+                   vals=[[-1234, 0, 1234, dtypes.max(dtypes.int), dtypes.min(dtypes.int)], dtypes.min(dtypes.int)], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], True], forward_only=True)
    helper_test_op(None, torch.minimum, Tensor.minimum, vals=[[True, False, False], [True, True, False]], forward_only=True)

@@ -204,13 +204,13 @@ class TestQuantizeOnnx(unittest.TestCase):
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
    out = (X.int().matmul(W.int())//1000)
-    if clip: out = out.clip(tg_dtype.min, tg_dtype.max)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
    out = out.cast(tg_dtype)
    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
-    if clip: mout = mout.clip(tg_dtype.min, tg_dtype.max)
+    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
    mout = mout.astype(xi)
    print(tout)
    print(mout)
@@ -62,7 +62,7 @@ def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 4
  return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \
    (torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha))

-def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
+def normal_test(func, shape=(20, 23), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)

 class TestRandomness(unittest.TestCase):
  def test_rand(self):
@@ -131,32 +131,29 @@ class TestRandomness(unittest.TestCase):
    """
    key0 = 1337
    key1 = int.from_bytes(hashlib.sha256(int(0).to_bytes(4)).digest(), "big") & 0xffffffff
-    # derive new key for the counter offset (c_low=0, c_high=0 for first call)
-    new_key_values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.array([0, 0], dtype=np.uint32))
-    new_key = (np.uint32(new_key_values[0]), np.uint32(new_key_values[1]))
-    values = jax.extend.random.threefry_2x32(new_key, np.arange(20, dtype=np.uint32))
+    values = jax.extend.random.threefry_2x32((np.uint32(key1), np.uint32(key0)), np.arange(20, dtype=np.uint32))
    values = (values >> (32 - 23)) | np.array(1, dtype=np.float32).view(np.uint32)
-    values = values.view(np.float32) - 1
+    values =  values.view(np.float32) - 1
    print(f"[{', '.join(f'{v}' for v in values)}]")
    """
-    jr = np.array([0.45735931396484375, 0.6311527490615845, 0.15571284294128418, 0.8149417638778687, 0.7862188816070557,
-                   0.8008807897567749, 0.568588376045227, 0.9852620363235474, 0.42314577102661133, 0.9811755418777466,
-                   0.38059568405151367, 0.09186363220214844, 0.9497315883636475, 0.5826880931854248, 0.3796330690383911,
-                   0.5610522031784058, 0.16122901439666748, 0.3732343912124634, 0.9795231819152832, 0.3280656337738037], dtype=np.float32)
+    jr = np.array([0.9073467254638672, 0.8235964775085449, 0.6872662305831909, 0.9920015335083008, 0.4941047430038452,
+                   0.3108327388763428, 0.09639489650726318, 0.004686474800109863, 0.8435229063034058, 0.824237585067749,
+                   0.5873836278915405, 0.4232727289199829, 0.2530076503753662, 0.40300023555755615, 0.03966474533081055,
+                   0.27904558181762695, 0.9150195121765137, 0.48057758808135986, 0.23821306228637695, 0.7676635980606079], dtype=np.float32)
    r = Tensor.rand(20).numpy()
    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)

-    # next 20 (c_low=20, c_high=0)
-    jr = np.array([0.09199333190917969, 0.9130761623382568, 0.7048608064651489, 0.22254979610443115, 0.0014830827713012695,
-                   0.37023448944091797, 0.7790107727050781, 0.7484984397888184, 0.7524604797363281, 0.19875383377075195,
-                   0.48537540435791016, 0.10002851486206055, 0.5369305610656738, 0.3294715881347656, 0.5246957540512085,
-                   0.7659651041030884, 0.7949080467224121, 0.34988296031951904, 0.9798505306243896, 0.2599533796310425], dtype=np.float32)
+    # next 20, np.arange(20, 40, dtype=np.uint32)
+    jr = np.array([0.7444133758544922, 0.7713677883148193, 0.8233780860900879, 0.43871235847473145, 0.517757773399353,
+                   0.6437174081802368, 0.967403769493103, 0.26167726516723633, 0.6825339794158936, 0.14966607093811035,
+                   0.28920769691467285, 0.017063498497009277, 0.2627382278442383, 0.9525482654571533, 0.9351049661636353,
+                   0.43904995918273926, 0.043945908546447754, 0.6616791486740112, 0.6667773723602295, 0.5228077173233032], dtype=np.float32)
    r = Tensor.rand(20).numpy()
    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)

-    # next 10 (c_low=40, c_high=0)
-    jr = np.array([0.3198714256286621, 0.7984923124313354, 0.320881724357605, 0.4716068506240845, 0.7323365211486816,
-                   0.9663800001144409, 0.13873648643493652, 0.16062307357788086, 0.49300849437713623, 0.10077548027038574], dtype=np.float32)
+    # next 10, np.arange(40, 50, dtype=np.uint32)
+    jr = np.array([0.9614430665969849, 0.059279561042785645, 0.01909029483795166, 0.47882091999053955, 0.9677121639251709,
+                   0.36863112449645996, 0.3102607727050781, 0.06608951091766357, 0.35329878330230713, 0.26518797874450684], dtype=np.float32)
    r = Tensor.rand(10).numpy()
    np.testing.assert_allclose(r, jr, atol=1e-5, rtol=1e-5)

@@ -327,7 +324,7 @@ class TestRandomness(unittest.TestCase):
                                                              lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))

  def test_kaiming_uniform(self):
-    for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]:
+    for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]:
      self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape))

  def test_kaiming_normal(self):
@@ -387,21 +384,6 @@ class TestRandomness(unittest.TestCase):
    for _ in range(833): Tensor.rand(1)
    Tensor.rand(1).realize()

-  def test_random_counter_overflow(self):
-    device = Device.DEFAULT
-    Tensor.manual_seed(1337)
-    Tensor.rand(1).realize()
-
-    Tensor._device_rng_counters[device].assign(Tensor([dtypes.uint32.max - 5, 0], device=device, dtype=dtypes.uint32)).realize()
-
-    Tensor.rand(10).realize()
-    c = Tensor._device_rng_counters[device].numpy()
-    np.testing.assert_allclose(c, [4, 1])
-
-    Tensor.rand(10).realize()
-    c = Tensor._device_rng_counters[device].numpy()
-    np.testing.assert_allclose(c, [14, 1])
-
 # TODO: still fails with MAX_KERNEL_BUFFERS
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
 class TestSample(unittest.TestCase):
@@ -62,7 +62,7 @@ class TestRendererFailures(unittest.TestCase):
 class TestCStyleFailures(unittest.TestCase):
  def test_inline_const_alu(self):
    # CPU doesn't use the max function
-    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.int.min+1))
+    ret = _setup_and_test_alu(Ops.MAX, 1, UOp.const(dtypes.int, dtypes.min(dtypes.int)+1))
    self.assertEqual(ret[0], 1)

  def _test_src_strip_paren(self, op: Ops, should_strip_paren:bool=True):
@@ -9,7 +9,7 @@ from hypothesis import assume, given, settings, strategies as strat

 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.device import is_dtype_supported
-from tinygrad.dtype import DType
+from tinygrad.dtype import DType, ImageDType
 from tinygrad.uop.ops import UOp, Ops, UPat
 from tinygrad.helpers import CI, DEBUG, OSX, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.engine.realize import CompiledRunner, run_schedule
@@ -783,6 +783,18 @@ class TestSchedule(unittest.TestCase):
    gc.collect()
    self.assertEqual(GlobalCounters.mem_used-base, 1024)

+  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  def test_image_matmul(self):
+    with Context(IMAGE=2):
+      x = Tensor.randn((9, 9)).realize()
+      y = Tensor.randn((9, 9)).realize()
+      out = x@y
+      run_schedule(check_schedule(out, 3))
+      np.testing.assert_allclose(out.numpy(), x.numpy()@y.numpy(), atol=1e-4, rtol=1e-4)
+      self.assertIsInstance(out.dtype, ImageDType)
+      self.assertIsNotNone(out.uop.base.realized)
+      self.assertIsInstance(out.uop.base.realized.dtype, ImageDType)
+
  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_dot_f16_fusion(self):
    with Context(FLOAT16=1, OPENPILOT_HACKS=1):
@@ -793,8 +805,11 @@ class TestSchedule(unittest.TestCase):
        for si in sched: si.lower()
        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])

-      with Context(IMAGE=1):
-        self.assertEqual(cnt(), 5)
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, 5)
+      self.assertEqual(cnt2, 5)

  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
  def test_image_f16_residual_fusion(self):
@@ -810,10 +825,14 @@ class TestSchedule(unittest.TestCase):
        for si in sched: si.lower()
        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])

-      with Context(IMAGE=1):
-        self.assertEqual(cnt(), 9)
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, 9)
+      self.assertEqual(cnt2, 9)

  @unittest.skipIf(Device.DEFAULT != "CL", "image only supported on CL")
+  @unittest.expectedFailure
  def test_image_conv_fusion(self):
    with Context(OPENPILOT_HACKS=1):
      def cnt():
@@ -824,8 +843,10 @@ class TestSchedule(unittest.TestCase):
        for si in sched: si.lower()
        return len([si for si in sched if isinstance(si.prg, CompiledRunner)])

-      with Context(IMAGE=1):
-        self.assertEqual(cnt(), 5)
+      with Context(IMAGE=1): cnt1 = cnt()
+      with Context(IMAGE=2): cnt2 = cnt()
+
+      self.assertEqual(cnt1, cnt2)

  def _test_fusion(self, shapes, f, cnt):
    with Context(DEBUG=0, TRACK_MATCH_STATS=0): args = [Tensor.randn(s).realize() for s in shapes]
@@ -292,13 +292,6 @@ class TestSetitem(unittest.TestCase):
    np.testing.assert_allclose(a.numpy(), [4, 6, 8, 10])
    np.testing.assert_allclose(b.numpy(), [0, 2, 4, 6])

-  def test_setitem_multiple_disjoint_on_invalid(self):
-    z = Tensor.invalid(10, dtype="int").realize()
-    z[2:5] = 2
-    z[6:7] = 3
-    z.realize()
-    self.assertListEqual(z[2:5].tolist(), [2, 2, 2])
-    self.assertListEqual(z[6:7].tolist(), [3])

 class TestWithGrad(unittest.TestCase):
  def test_no_requires_grad_works(self):
@@ -256,7 +256,7 @@ class TestTinygrad(unittest.TestCase):
  def test_randperm(self):
    Tensor.manual_seed(0)
    a = Tensor.randperm(10).realize()
-    np.testing.assert_equal(a.numpy(), [8, 9, 4, 3, 6, 1, 7, 5, 2, 0])
+    np.testing.assert_equal(a.numpy(), [5, 2, 8, 1, 3, 7, 9, 6, 0, 4])
    b = Tensor.randperm(1000).realize()
    np.testing.assert_equal(set(b.numpy()), set(range(1000)))

@@ -493,17 +493,6 @@ class TestTinygrad(unittest.TestCase):
    dev = a.to(Device.DEFAULT)
    np.testing.assert_allclose(a.numpy(), dev.numpy())

-  def test_copy_from_numpy_dtype(self):
-    data = np.array([1.0, 2, 3], dtype=np.float32)
-    t = Tensor(data, dtype=dtypes.bfloat16)
-    try:
-      # TODO: fix dtype in tinygrad space
-      assert t.dtype == dtypes.bfloat16
-    except AssertionError:
-      assert t.dtype == dtypes.float32
-    np.testing.assert_equal(t.tolist(), data)
-    np.testing.assert_equal((t+1).tolist(), data+1)
-
  # Regression test for https://github.com/tinygrad/tinygrad/issues/1751
  def test_copy_from_numpy_unaligned(self):
    # 2**15 is the minimum for repro
@@ -77,15 +77,15 @@ class TestUOps(unittest.TestCase):
  def _test_uop_fxn(self, op, fxn, dts=(dtypes.float32, )):
    for f in [_test_single_value, _test_single_value_const]:
      for a in [-2.0, 0.0, 1.0]:
-        a = dts[0].const(a)
+        a = dtypes.as_const(a, dts[0])
        self._equal(f([a], op, dts), fxn(a))

  def _test_bop_fxn(self, op, fxn, dts=(dtypes.float32, )*2, no_b_zero=False, no_b_neg=False):
    for f in [_test_single_value, _test_single_value_const]:
      for a in [-2.0, 0.0, 1.0]:
        for b in [-3.0, 1.0] + ([] if no_b_zero else [0.0]):
-          a = dts[0].const(a)
-          b = dts[1].const(abs(b) if no_b_neg else b)
+          a = dtypes.as_const(a, dts[0])
+          b = dtypes.as_const(abs(b) if no_b_neg else b, dts[1])
          self._equal(f([a,b], op, dts), fxn(a,b))

  def _test_top_fxn(self, op, fxn, dts=(dtypes.float32, )*3):
@@ -93,9 +93,9 @@ class TestUOps(unittest.TestCase):
      for a in [-2.0, 0, 1]:
        for b in [-3.0, 3.0]:
          for c in [-4.0, 4.0]:
-            a = dts[0].const(a)
-            b = dts[1].const(b)
-            c = dts[2].const(c)
+            a = dtypes.as_const(a, dts[0])
+            b = dtypes.as_const(b, dts[1])
+            c = dtypes.as_const(c, dts[2])
            self._equal(f([a,b,c], op, dts), fxn(a,b,c))

 class TestFloatUOps(TestUOps):
@@ -117,7 +117,7 @@ class TestFloatUOps(TestUOps):
  def test_cmpne_nan(self):  # NaN != x for any x (IEEE 754)
    for a, b in [(math.nan, 1.0), (1.0, math.nan), (math.nan, math.nan)]:
      self.assertTrue(_test_single_value(
-        [dtypes.float32.const(a), dtypes.float32.const(b)],
+        [dtypes.as_const(a, dtypes.float32), dtypes.as_const(b, dtypes.float32)],
        Ops.CMPNE, (dtypes.float32, dtypes.float32)))
  # MOD isn't tested on floats

@@ -23,7 +23,7 @@ def vision_conv_143():
  c32 = ((c27<3)!=True)&(c27<67)
  c34 = UOp(Ops.PARAM, dtypes.imageh((32, 1024, 4)), (), 1)
  c38 = c5//2
-  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.weakint, Invalid))
+  c45 = (c32&c24).where((c27*64+c38+c17*4096+-12480), UOp.const(dtypes.index, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
  c49 = UOp(Ops.PARAM, dtypes.imageh((64, 49, 4)), (), 2)
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
@@ -49,7 +49,7 @@ def vision_conv_153():
  c32 = ((c27<3)!=True)&(c27<35)
  c34 = UOp(Ops.PARAM, dtypes.imageh((16, 1024, 4)), (), 1)
  c38 = c5//2
-  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.weakint, Invalid))
+  c45 = (c32&c24).where((c27*128+c38+c17*4096+-12672), UOp.const(dtypes.index, Invalid))
  c48 = (c24&c32).where(c34.index(c45), UOp.const(dtypes.float, 0.0))
  c49 = UOp(Ops.PARAM, dtypes.imageh((128, 49, 4)), (), 2)
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	b384c27314	ln	2026-03-12 08:23:06 +08:00
geohot	5ed68aeed5	cleanups	2026-03-12 08:18:31 +08:00
geohot	f682af2a31	better uop matmul	2026-03-12 08:13:09 +08:00
geohot	62dbf12655	call matmul	2026-03-11 21:50:31 +08:00
geohot	b17e15d1aa	support ranges on call	2026-03-11 15:22:46 +08:00