don't bufferize 0s

2026-08-14 17:58:28 +00:00 · 2025-08-20 21:02:09 -07:00
719 changed files with 209599 additions and 206699 deletions
@@ -1,3 +0,0 @@
-[run]
-source = tinygrad
-branch = True
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
-  mesa:
-    description: "Install mesa"
-    required: false
-    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -61,7 +57,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: ${{ github.workspace }}/.venv
-        key: venv-${{ runner.os }}-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 }}-${{ hashFiles('**/setup.py') }}-${{ env.PYTHON_CACHE_VERSION }}

    # **** Caching downloads ****

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

    # **** Python deps ****

@@ -187,7 +183,7 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_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')
@@ -221,7 +217,7 @@ runs:
        sudo mkdir -p /usr/local/lib
        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
-          sudo xargs curl -fL -o /usr/local/lib/libamd_comgr.dylib
+          sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
        cargo build --release --manifest-path ./extra/remu/Cargo.toml

    # **** gpuocelot ****
@@ -229,25 +225,16 @@ runs:
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
-      run: |
-        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
-        for f in "${pkgs[@]}"; do
-          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
-        done
-
-        # Fix boost 1.85 for gpuocelot
-        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
-        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
+      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true'
      id: cache-build
      uses: actions/cache@v4
      env:
-        cache-name: cache-gpuocelot-build-1
+        cache-name: cache-gpuocelot-build
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
    - name: Clone/compile gpuocelot
      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
      shell: bash
@@ -257,13 +244,7 @@ runs:
        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
        mkdir build
        cd build
-
-        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
-          CMAKE_ARGS="$CMAKE_ARGS -DBoost_INCLUDE_DIR=$(brew --prefix boost)/include -DBoost_LIBRARY_DIR=$(brew --prefix boost)/lib"
-        fi
-
-        cmake .. $CMAKE_ARGS
+        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
@@ -278,7 +259,7 @@ runs:
      if: inputs.webgpu == 'true' && runner.os == 'Linux'
      shell: bash
      run: |
-        sudo curl -fL https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
+        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
        sudo ldconfig
    - name: Install WebGPU dawn (macOS)
      if: inputs.webgpu == 'true' && runner.os == 'macOS'
@@ -293,13 +274,3 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
-
-    # **** mesa ****
-    - name: Install mesa (linux)
-      if: inputs.mesa == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: sudo curl -fL https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
-    - name: Install mesa (macOS)
-      if: inputs.mesa == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -1,178 +0,0 @@
-name: Autogen
-env:
-  # increment this when downloads substantially change to avoid the internet
-  CACHE_VERSION: '13'
-  CAPTURE_PROCESS_REPLAY: 1
-  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
-  PYTHONPATH: ${{ github.workspace }}
-
-on:
-  push:
-    branches:
-      - master
-  pull_request:
-    paths:
-    - 'tinygrad/runtime/autogen/**/*'
-    - 'tinygrad/runtime/support/autogen.py'
-  workflow_dispatch:
-    paths:
-    - 'tinygrad/runtime/autogen/**/*'
-    - 'tinygrad/runtime/support/autogen.py'
-
-jobs:
-  autogen:
-    name: In-tree Autogen
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        opencl: 'true'
-        amd: 'true'
-        cuda: 'true'
-        llvm: 'true'
-        webgpu: 'true'
-        mesa: 'true'
-        pydeps: 'pyyaml mako'
-    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev
-    - name: Verify OpenCL autogen
-      run: |
-        mv tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        python3 -c "from tinygrad.runtime.autogen import opencl"
-        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
-    - name: Verify CUDA autogen
-      run: |
-        mv tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
-        mv tinygrad/runtime/autogen/nvrtc.py /tmp/nvrtc.py.bak
-        mv tinygrad/runtime/autogen/nvjitlink.py /tmp/nvjitlink.py.bak
-        mv tinygrad/runtime/autogen/nv_570.py /tmp/nv_570.py.bak
-        mv tinygrad/runtime/autogen/nv.py /tmp/nv.py.bak
-        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv"
-        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
-        diff /tmp/nvrtc.py.bak tinygrad/runtime/autogen/nvrtc.py
-        diff /tmp/nvjitlink.py.bak tinygrad/runtime/autogen/nvjitlink.py
-        diff /tmp/nv_570.py.bak tinygrad/runtime/autogen/nv_570.py
-        diff /tmp/nv.py.bak tinygrad/runtime/autogen/nv.py
-    - name: Verify AMD autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
-        mv tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
-        mv tinygrad/runtime/autogen/hip.py /tmp/hip.py.bak
-        mv tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
-        mv tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
-        mv tinygrad/runtime/autogen/rocprof.py /tmp/rocprof.py.bak
-        mv tinygrad/runtime/autogen/am/am.py /tmp/am_am.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_soc15.py /tmp/am_pm4_soc15.py.bak
-        mv tinygrad/runtime/autogen/am/pm4_nv.py /tmp/am_pm4_nv.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_4_0_0.py /tmp/am_sdma_4_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_5_0_0.py /tmp/am_sdma_5_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/sdma_6_0_0.py /tmp/am_sdma_6_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v13_0_0.py /tmp/am_smu_v13_0_0.py.bak
-        mv tinygrad/runtime/autogen/am/smu_v14_0_2.py /tmp/am_smu_v14_0_2.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr, hsa, hip, amd_gpu, sqtt, rocprof; from tinygrad.runtime.autogen.am import am, pm4_soc15, pm4_nv, sdma_4_0_0, sdma_5_0_0, sdma_6_0_0, smu_v13_0_0, smu_v14_0_2"
-        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
-        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
-        diff /tmp/hip.py.bak tinygrad/runtime/autogen/hip.py
-        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
-        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
-        diff /tmp/rocprof.py.bak tinygrad/runtime/autogen/rocprof.py
-        diff /tmp/am_am.py.bak tinygrad/runtime/autogen/am/am.py
-        diff /tmp/am_pm4_soc15.py.bak tinygrad/runtime/autogen/am/pm4_soc15.py
-        diff /tmp/am_pm4_nv.py.bak tinygrad/runtime/autogen/am/pm4_nv.py
-        diff /tmp/am_sdma_4_0_0.py.bak tinygrad/runtime/autogen/am/sdma_4_0_0.py
-        diff /tmp/am_sdma_5_0_0.py.bak tinygrad/runtime/autogen/am/sdma_5_0_0.py
-        diff /tmp/am_sdma_6_0_0.py.bak tinygrad/runtime/autogen/am/sdma_6_0_0.py
-        diff /tmp/am_smu_v13_0_0.py.bak tinygrad/runtime/autogen/am/smu_v13_0_0.py
-        diff /tmp/am_smu_v14_0_2.py.bak tinygrad/runtime/autogen/am/smu_v14_0_2.py
-    - name: Verify Linux autogen
-      run: |
-        mv tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
-        mv tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
-        mv tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
-        mv tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
-        mv tinygrad/runtime/autogen/pci.py /tmp/pci.py.bak
-        mv tinygrad/runtime/autogen/vfio.py /tmp/vfio.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, ib, pci, vfio"
-        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
-        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
-        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
-        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
-        diff /tmp/pci.py.bak tinygrad/runtime/autogen/pci.py
-        diff /tmp/vfio.py.bak tinygrad/runtime/autogen/vfio.py
-    - name: Verify LLVM autogen
-      run: |
-        mv tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
-        python3 -c "from tinygrad.runtime.autogen import llvm"
-        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
-    - name: Verify WebGPU autogen
-      run: |
-        mv tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
-        python3 -c "from tinygrad.runtime.autogen import webgpu"
-        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
-    - name: Verify Qualcomm autogen
-      run: |
-        mv tinygrad/runtime/autogen/kgsl.py /tmp/kgsl.py.bak
-        mv tinygrad/runtime/autogen/adreno.py /tmp/adreno.py.bak
-        mv tinygrad/runtime/autogen/qcom_dsp.py /tmp/qcom_dsp.py.bak
-        python3 -c "from tinygrad.runtime.autogen import kgsl, adreno, qcom_dsp"
-        diff /tmp/kgsl.py.bak tinygrad/runtime/autogen/kgsl.py
-        diff /tmp/adreno.py.bak tinygrad/runtime/autogen/adreno.py
-        diff /tmp/qcom_dsp.py.bak tinygrad/runtime/autogen/qcom_dsp.py
-    - name: Verify libusb autogen
-      run: |
-        mv tinygrad/runtime/autogen/libusb.py /tmp/libusb.py.bak
-        python3 -c "from tinygrad.runtime.autogen import libusb"
-        diff /tmp/libusb.py.bak tinygrad/runtime/autogen/libusb.py
-    - name: Verify mesa autogen
-      run: |
-        mv tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
-        python3 -c "from tinygrad.runtime.autogen import mesa"
-        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
-    - name: Verify libclang autogen
-      run: |
-        cp tinygrad/runtime/autogen/libclang.py /tmp/libclang.py.bak
-        REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
-        diff /tmp/libclang.py.bak tinygrad/runtime/autogen/libclang.py
-  autogen-mac:
-    name: In-tree Autogen (macos)
-    runs-on: macos-14
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        llvm: 'true'
-    - name: Verify macos autogen
-      run: |
-        mv tinygrad/runtime/autogen/metal.py /tmp/metal.py.bak
-        LIBCLANG_PATH=/opt/homebrew/opt/llvm@20/lib/libclang.dylib python3 -c "from tinygrad.runtime.autogen import metal"
-        diff /tmp/metal.py.bak tinygrad/runtime/autogen/metal.py
-  autogen-comgr-3:
-    name: In-tree Autogen (comgr 3)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-    - name: Install autogen support packages
-      run: |
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.4 $(lsb_release -cs) main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt -qq update || true
-        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
-    - name: Verify comgr (3) autogen
-      run: |
-        mv tinygrad/runtime/autogen/comgr_3.py /tmp/comgr_3.py.bak
-        python3 -c "from tinygrad.runtime.autogen import comgr_3"
-        diff /tmp/comgr_3.py.bak tinygrad/runtime/autogen/comgr_3.py
@@ -28,7 +28,7 @@ jobs:
      # since sudo is required for usbgpu on macos, move the cache to a new location, as some of the files are owned by root
      PYTHONPYCACHEPREFIX: /tmp/tiny_python_pycache
    runs-on: [self-hosted, macOS]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -51,30 +51,25 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
-    - name: Print macOS version
-      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
-      run: METAL=1 NOCLANG=1 python3.11 test/external/external_model_benchmark.py
+      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
-      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
-        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - name: Run Tensor Core GEMM (float)
      run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
    - name: Run Tensor Core GEMM (half)
@@ -102,7 +97,7 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 JIT=1 ASSERT_MIN_STEP_TIME=13 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 JIT=1 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
      run: BENCHMARK_LOG=gpt2_half HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
@@ -111,27 +106,22 @@ jobs:
      run: BENCHMARK_LOG=olmoe python3.11 examples/olmoe.py
    - name: Train MNIST
      run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
-
-    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
-    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
-    #- name: Run 10 CIFAR training steps
-    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
-    #- name: Run 10 CIFAR training steps w HALF
-    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
-
+    - name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    - name: Run 10 CIFAR training steps w HALF
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #- name: Run 10 CIFAR training steps w BF16
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+    - name: UsbGPU openpilot test
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -168,7 +158,7 @@ jobs:
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -199,21 +189,20 @@ jobs:
    - name: Test speed vs torch
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test benchmark allreduce
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      run: |
-        NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
-        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Run Tensor Core GEMM (CUDA)
      run: |
        CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
        CUDA=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
        CUDA=1 SHOULD_USE_TC=1 ALLOW_TF32=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
-        CUDA=1 SHOULD_USE_TC=1 FP8E4M3=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_fp8.txt
    - name: Run Tensor Core GEMM (PTX)
-      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
      run: NV=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
    - name: Test NV=1
@@ -222,9 +211,8 @@ jobs:
      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
    - name: Run Stable Diffusion
      run: BENCHMARK_LOG=stable_diffusion NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA
      run: |
        BENCHMARK_LOG=llama_nojit NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -239,8 +227,6 @@ jobs:
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
      run: BENCHMARK_LOG=llama3_beam_4gpu NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
-    - name: Run quantized LLaMA3
-      run: BENCHMARK_LOG=llama3_fp8 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --temperature 0 --benchmark --quantize fp8 | tee llama3_fp8.txt
    # - name: Run LLaMA-3 8B on 6 GPUs
    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
    # - name: Run LLaMA-2 70B
@@ -250,9 +236,9 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 NV=1 JIT=1 ASSERT_MIN_STEP_TIME=4 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 NV=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 ASSERT_MIN_STEP_TIME=6 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
@@ -274,7 +260,6 @@ jobs:
          llama3_beam.txt
          llama3_four_gpu.txt
          llama3_six_gpu.txt
-          llama3_fp8.txt
          llama_2_70B.txt
          mixtral.txt
          gpt2_unjitted.txt
@@ -287,7 +272,7 @@ jobs:
  testmorenvidiabenchmark:
    name: tinybox green Training Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -312,28 +297,24 @@ 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
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (NV)
-    #   run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    # TODO: too slow
-    # - name: Fuzz Padded Tensor Core GEMM (PTX)
-    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    - name: HEVC Decode Benchmark
-      run: VALIDATE=1 MAX_FRAMES=100 NV=1 PYTHONPATH=. python3 extra/hevc/decode.py
+    - name: Fuzz Padded Tensor Core GEMM (NV)
+      run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    - name: Fuzz Padded Tensor Core GEMM (PTX)
+      run: NV=1 PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: Train MNIST
      run: time PYTHONPATH=. NV=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=110 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=120 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_half_wino NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run MLPerf resnet eval on training data
      run: time BENCHMARK_LOG=resnet_eval NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
@@ -363,7 +344,7 @@ jobs:
  testamdbenchmark:
    name: tinybox red Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -410,11 +391,11 @@ jobs:
    #    python3 -c "import torch; print(torch.__version__)"
    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 CCACHE=0 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
    - name: Test tensor cores
      run: |
-        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
-        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
@@ -432,9 +413,9 @@ jobs:
    - name: Test AM warm start time
      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=3200 CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run LLaMA 7B
      run: |
        BENCHMARK_LOG=llama_nojit AMD=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -460,9 +441,9 @@ jobs:
    - name: Run GPT2
      run: |
        BENCHMARK_LOG=gpt2_nojit AMD=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
-        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
+        BENCHMARK_LOG=gpt2 AMD=1 JIT=1 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.txt
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 ASSERT_MIN_STEP_TIME=5 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
    - name: Run GPT2 w HALF/BEAM
      run: BENCHMARK_LOG=gpt2_half_beam AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
    - uses: actions/upload-artifact@v4
@@ -493,7 +474,7 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -525,20 +506,19 @@ jobs:
    - name: Train MNIST
      run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
    - name: Run 10 CIFAR training steps
-      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=200 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    # - name: Run 10 CIFAR training steps w BF16
-    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
-    # - name: Run 10 CIFAR training steps w winograd
-    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_half AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    - name: Run 10 CIFAR training steps w BF16
+      run: BENCHMARK_LOG=cifar_10steps_bf16 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    - name: Run 10 CIFAR training steps w winograd
+      run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -557,7 +537,7 @@ jobs:
  testmlperfamdbenchmark:
    name: tinybox red MLPerf Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 60
+    timeout-minutes: 30
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -623,33 +603,47 @@ 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.10.0 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.10.0/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.0 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=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: 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=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_vision.onnx
-    - name: openpilot compile3 0.10.1 driving_policy
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/720392c9a5b986981fdbed1bb8c47a6c5573a50e/selfdrive/modeld/models/driving_policy.onnx
-    - name: openpilot compile3 0.10.1 dmonitoring
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=10 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 openpilot 0.9.9 driving_vision
+      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: benchmark openpilot 0.9.9 driving_policy
+      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: benchmark openpilot 0.9.9 dmonitoring
+      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 0.9.9 driving_vision
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+    - name: openpilot compile3 0.9.9 driving_policy
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot compile3 0.9.9 dmonitoring
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+    - name: openpilot compile3 Space Lab policy + vision
+      run: |
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
        ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
-        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 CPU=1 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
+    - uses: actions/upload-artifact@v4
+      with:
+        name: Speed (comma)
+        path: |
+          openpilot_compile_0_9_4.txt
+          openpilot_compile_0_9_7.txt
+          openpilot_0_9_4.txt
+          openpilot_0_9_7.txt
+          openpilot_image_0_9_4.txt
+          openpilot_image_0_9_7.txt

  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -685,8 +679,8 @@ jobs:
    # Fails on 9070
    # - name: Test tensor cores
    #   run: |
-    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
-    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
    #     AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: Run Tensor Core GEMM (AMD)
      run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee am_matmul_amd.txt
@@ -694,12 +688,9 @@ jobs:
      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
    - name: Test DISK copy time
      run: AMD=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        AMD=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+    # TODO: enable
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
@@ -719,7 +710,7 @@ jobs:
  testgreendriverbenchmark:
    name: NV Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -751,17 +742,13 @@ jobs:
    - name: Test driver start time
      run: time DEBUG=3 NV=1 python3 test/test_tiny.py TestTiny.test_plus
    - name: Test tensor cores
-      run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+      run: NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
    - name: Test DISK copy time
      run: NV=1 TESTFILE=/raid/downloads/llama3-8b-sfr/model-00001-of-00004.safetensors python3 test/external/external_benchmark_disk_raw.py
-    - name: Test CPU copy time
-      run: |
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
-        NV=1 GRAPH_ONE_KERNEL=1 PYTHONPATH=. NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
    - name: Test LLAMA-3
      run: BENCHMARK_LOG=llama3_beam NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --benchmark --temperature 0 | tee nv_llama3_beam.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
+      run: time BENCHMARK_LOG=cifar NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
@@ -22,13 +22,13 @@ jobs:
    - name: Run SDXL with new search
      # TODO: GCVM_L2_PROTECTION_FAULT_STATUS with llvm19
      run: |
-        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CCACHE=0 python examples/sdxl.py --noshow --timing --seed 0
+        BENCHMARK_LOG=search_sdxl PYTHONPATH=. AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run SDXL with cached search
      run: |
        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 CCACHE=0 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
    - name: Run winograd cifar with cached search
      run: |
        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 720
+    timeout-minutes: 360

    steps:
    - name: Checkout Code
@@ -27,4 +27,4 @@ jobs:
      run: |
        rm "~/.cache/tinygrad/cache_mlperf.db" || true
        BENCHMARK_LOG=mlpert_train_resnet LOGMLPERF=0 CACHEDB="~/.cache/tinygrad/cache_mlperf.db" examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-        rm "~/.cache/tinygrad/cache_mlperf.db"
+        rm "~/.cache/tinygrad/cache_mlperf.db"
@@ -20,11 +20,11 @@ jobs:
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
-        pip install setuptools wheel build twine
+        pip install setuptools wheel twine
    - name: Build and publish
      env:
        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
      run: |
-        python -m build
+        python setup.py sdist bdist_wheel
        twine upload dist/*
@@ -56,15 +56,15 @@ jobs:
        uses: actions/checkout@v4
        with:
          path: base
-      - name: Set up Python 3.12
+      - name: Set up Python 3.10
        uses: actions/setup-python@v5
        with:
-          python-version: '3.12'
+          python-version: '3.10'
      - name: Count Line Diff
        run: |
+          pip install tabulate
          BASE="$GITHUB_WORKSPACE/base"
          PR="$GITHUB_WORKSPACE/pr"
-          pip install tabulate $BASE
          cp "$BASE/sz.py" .
          echo "loc_content<<EOF" >> "$GITHUB_ENV"
          python sz.py "$BASE" "$PR" >> "$GITHUB_ENV"
@@ -38,7 +38,6 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
-venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -63,5 +62,3 @@ profile_stats
 *.log
 target
 .mypy_cache
-mutants
-.mutmut-cache
@@ -20,15 +20,27 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
-      - id: example
-        name: test all devices
-        entry: python3 test/external/external_test_example.py
+      - id: devicetests
+        name: select GPU tests
+        entry: env GPU=1 PYTHONPATH="." python3 -m pytest test/test_uops.py test/test_search.py
        language: system
        always_run: true
        pass_filenames: false
      - id: tests
        name: subset of tests
-        entry: env OMP_NUM_THREADS=1 PYTHONPATH="." python3 -m pytest -n=6 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
+        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
+      - id: example
+        name: multi device tests
+        entry: python3 test/external/external_test_example.py
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: pylint
+        name: pylint
+        entry: python3 -m pylint tinygrad/
+        language: system
+        always_run: true
+        pass_filenames: false
@@ -30,6 +30,10 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=

+# When enabled, pylint would attempt to guess common misconfiguration and emit
+# user-friendly hints instead of false-positive error messages
+suggestion-mode=yes
+
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
@@ -50,12 +54,11 @@ confidence=
 # --enable=similarities". If you want to run only the classes checker, but have
 # no Warning level messages displayed, use"--disable=all --enable=classes
 # --disable=W"
-disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method
 # E1101 for function binding
 # W0221 for Function class
 # W0105 for comment strings
 # E0401 for missing imports
-# W0707 for not reraising

 # Enable the message, report, category or checker with the given id(s). You can
 # either give multiple identifier separated by comma (,) or put this option
@@ -21,38 +21,17 @@ tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) an

 ---

-tinygrad is an end-to-end deep learning stack:
+Despite tinygrad's size, it is a fully featured deep learning framework.

- **Tensor library** with autograd
- **IR and compiler** that fuse and lower kernels
- **JIT + graph execution**
- **nn / optim / datasets** for real training
+Due to its extreme simplicity, it is the easiest framework to add new accelerators to, with support for both inference and training. If XLA is CISC, tinygrad is RISC.

-It’s inspired by PyTorch (ergonomics), JAX (functional transforms and IR-based AD), and TVM (scheduling and codegen), but stays intentionally tiny and hackable.
+tinygrad is now beta software, we [raised some money](https://geohot.github.io/blog/jekyll/update/2023/05/24/the-tiny-corp-raised-5M.html) to make it good. Someday, we will tape out chips.

---
+## Features

-## How tinygrad compares
+### LLaMA and Stable Diffusion

-**PyTorch**
-
- ✅ Similar: eager `Tensor` API, autograd, `optim`, basic datasets and layers.
- ✅ You can write familiar training loops.
- 🔁 Unlike PyTorch, the entire compiler and IR are visible and hackable.
-
-**JAX**
-
- ✅ IR-based autodiff over primitives (like JAXPR + XLA).
- ✅ Function-level JIT (`TinyJit`) that captures and replays kernels.
- 🔁 Fewer functional transforms (no full `vmap`/`pmap` yet), but far easier to read.
-
-**TVM**
-
- ✅ Multiple lowering passes, scheduling, and BEAM search over kernels.
- ✅ Device “graphs” for batched execution.
- 🔁 tinygrad also ships the **front-end framework** (tensors, nn, optim), not just the compiler.
-
---
+tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!

 ### Laziness

@@ -100,8 +79,9 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers

 tinygrad already supports numerous accelerators, including:

- [x] [OpenCL](tinygrad/runtime/ops_cl.py)
- [x] [CPU](tinygrad/runtime/ops_cpu.py)
+- [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
+- [x] [CPU (C Code)](tinygrad/runtime/ops_cpu.py)
+- [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
@@ -0,0 +1,470 @@
+#!/bin/bash -e
+
+# setup instructions for clang2py
+if [[ ! $(clang2py -V) ]]; then
+  pushd .
+  cd /tmp
+  sudo apt-get install -y --no-install-recommends clang
+  pip install --upgrade pip setuptools
+  pip install clang==14.0.6
+  git clone https://github.com/nimlgen/ctypeslib.git
+  cd ctypeslib
+  pip install .
+  clang2py -V
+  popd
+fi
+
+BASE=tinygrad/runtime/autogen/
+
+fixup() {
+  sed -i '1s/^/# mypy: ignore-errors\n/' $1
+  sed -i 's/ *$//' $1
+  grep FIXME_STUB $1 || true
+}
+
+patch_dlopen() {
+  path=$1; shift
+  name=$1; shift
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $path
+PATHS_TO_TRY = [
+$(for p in "$@"; do echo "  $p,"; done)
+]
+def _try_dlopen_$name():
+  library = ctypes.util.find_library("$name")
+  if library: return ctypes.CDLL(library)
+  for candidate in PATHS_TO_TRY:
+    try: return ctypes.CDLL(candidate)
+    except OSError: pass
+  return None
+EOF
+}
+
+generate_opencl() {
+  clang2py /usr/include/CL/cl.h -o $BASE/opencl.py -l /usr/lib/x86_64-linux-gnu/libOpenCL.so.1 -k cdefstum
+  fixup $BASE/opencl.py
+  # hot patches
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/opencl.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libOpenCL.so.1')\ctypes.CDLL(ctypes.util.find_library('OpenCL'))\g" $BASE/opencl.py
+  python3 -c "import tinygrad.runtime.autogen.opencl"
+}
+
+generate_hip() {
+  clang2py /opt/rocm/include/hip/hip_ext.h /opt/rocm/include/hip/hiprtc.h \
+  /opt/rocm/include/hip/hip_runtime_api.h /opt/rocm/include/hip/driver_types.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/hip.py -l /opt/rocm/lib/libamdhip64.so
+  echo "hipDeviceProp_t = hipDeviceProp_tR0600" >> $BASE/hip.py
+  echo "hipGetDeviceProperties = hipGetDevicePropertiesR0600" >> $BASE/hip.py
+  fixup $BASE/hip.py
+  # we can trust HIP is always at /opt/rocm/lib
+  #sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhiprtc.so')\ctypes.CDLL(ctypes.util.find_library('hiprtc'))\g" $BASE/hip.py
+  #sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamdhip64.so')\ctypes.CDLL(ctypes.util.find_library('amdhip64'))\g" $BASE/hip.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/hip.py
+  sed -i "s\'/opt/rocm/\os.getenv('ROCM_PATH', '/opt/rocm/')+'/\g" $BASE/hip.py
+  python3 -c "import tinygrad.runtime.autogen.hip"
+}
+
+generate_comgr() {
+  clang2py /opt/rocm/include/amd_comgr/amd_comgr.h \
+  --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
+  fixup $BASE/comgr.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
+  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
+  python3 -c "import tinygrad.runtime.autogen.comgr"
+}
+
+generate_kfd() {
+  clang2py /usr/include/linux/kfd_ioctl.h -o $BASE/kfd.py -k cdefstum
+
+  fixup $BASE/kfd.py
+  sed -i "s/import ctypes/import ctypes, os/g" $BASE/kfd.py
+  sed -i "s/import fcntl, functools/import functools/g" $BASE/kfd.py
+  sed -i "/import functools/a from tinygrad.runtime.support.hcq import FileIOInterface" $BASE/kfd.py
+  sed -i "s/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, \*\*kwargs):/def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:FileIOInterface, \*\*kwargs):/g" $BASE/kfd.py
+  sed -i "s/fcntl.ioctl(__fd, (__idir<<30)/__fd.ioctl((__idir<<30)/g" $BASE/kfd.py
+  sed -i "s/!!/not not /g" $BASE/kfd.py
+  python3 -c "import tinygrad.runtime.autogen.kfd"
+}
+
+generate_cuda() {
+  clang2py /usr/include/cuda.h --clang-args="-D__CUDA_API_VERSION_INTERNAL" -o $BASE/cuda.py -l /usr/lib/x86_64-linux-gnu/libcuda.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/cuda.py
+  sed -i "s\ctypes.CDLL('/usr/lib/x86_64-linux-gnu/libcuda.so')\ctypes.CDLL(ctypes.util.find_library('cuda'))\g" $BASE/cuda.py
+  fixup $BASE/cuda.py
+  python3 -c "import tinygrad.runtime.autogen.cuda"
+}
+
+generate_nvrtc() {
+  clang2py /usr/local/cuda/include/nvrtc.h /usr/local/cuda/include/nvJitLink.h -o $BASE/nvrtc.py -l /usr/local/cuda/lib64/libnvrtc.so -l /usr/local/cuda/lib64/libnvJitLink.so
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvrtc.so')\ctypes.CDLL(ctypes.util.find_library('nvrtc'))\g" $BASE/nvrtc.py
+  sed -i "s\ctypes.CDLL('/usr/local/cuda/lib64/libnvJitLink.so')\ctypes.CDLL(ctypes.util.find_library('nvJitLink'))\g" $BASE/nvrtc.py
+  fixup $BASE/nvrtc.py
+  python3 -c "import tinygrad.runtime.autogen.nvrtc"
+}
+
+generate_nv() {
+  NVKERN_COMMIT_HASH=81fe4fb417c8ac3b9bdcc1d56827d116743892a5
+  NVKERN_SRC=/tmp/open-gpu-kernel-modules-$NVKERN_COMMIT_HASH
+  if [ ! -d "$NVKERN_SRC" ]; then
+    git clone https://github.com/NVIDIA/open-gpu-kernel-modules $NVKERN_SRC
+    pushd .
+    cd $NVKERN_SRC
+    git reset --hard $NVKERN_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    extra/nv_gpu_driver/clc6c0qmd.h \
+    extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc96f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc761.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl83de.h \
+    $NVKERN_SRC/src/nvidia/generated/g_allclasses.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc6c0.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clcdc0.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc6b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/clc9b5.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/uvm_linux_ioctl.h \
+    $NVKERN_SRC/kernel-open/nvidia-uvm/hwref/ampere/ga100/dev_fault.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv_escape.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numbers.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-ioctl-numa.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include/nv-unix-nvos-params-wrappers.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/alloc/alloc_channel.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/nvos.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0000/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl0080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl2080/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl83de/*.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
+    --clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv_gpu.py
+  fixup $BASE/nv_gpu.py
+  sed -i "s\(0000000001)\1\g" $BASE/nv_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/nv_gpu.py
+  sed -i 's/#\?\s\([A-Za-z0-9_]\+\) = MW ( \([0-9]\+\) : \([0-9]\+\) )/\1 = (\2 , \3)/' $BASE/nv_gpu.py # NVC6C0_QMDV03_00 processing
+  sed -i 's/#\sdef NVC6C0_QMD\([A-Za-z0-9_()]\+\):/def NVC6C0_QMD\1:/' $BASE/nv_gpu.py
+  sed -i 's/#\sdef NVCEC0_QMD\([A-Za-z0-9_()]\+\):/def NVCEC0_QMD\1:/' $BASE/nv_gpu.py
+  sed -E -i -n '/^def (NVCEC0_QMDV05_00_RELEASE)(_ENABLE)\(i\):/{p;s//\1'"0"'\2=\1\2(0)\n\1'"1"'\2=\1\2(1)/;H;b};p;${x;s/^\n//;p}' "$BASE/nv_gpu.py"
+  sed -i 's/#\s*return MW(\([0-9i()*+]\+\):\([0-9i()*+]\+\))/    return (\1 , \2)/' $BASE/nv_gpu.py
+  sed -i 's/#\?\s*\(.*\)\s*=\s*\(NV\)\?BIT\(32\)\?\s*(\s*\([0-9]\+\)\s*)/\1 = (1 << \4)/' $BASE/nv_gpu.py # name = BIT(x) -> name = (1 << x)
+  sed -i "s/UVM_\([A-Za-z0-9_]\+\) = \['i', '(', '\([0-9]\+\)', ')'\]/UVM_\1 = \2/" $BASE/nv_gpu.py # UVM_name = ['i', '(', '<num>', ')'] -> UVM_name = <num>
+
+  # Parse status codes
+  sed -n '1i\
+nv_status_codes = {}
+/^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.nv_gpu"
+
+  clang2py -k cdefstum \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
+    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.h \
+    extra/nv_gpu_driver/g_rpc-message-header.h \
+    extra/nv_gpu_driver/gsp_static_config.h \
+    extra/nv_gpu_driver/vbios.h \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv/nv.py
+
+  fixup $BASE/nv/nv.py
+  python3 -c "import tinygrad.runtime.autogen.nv.nv"
+}
+
+generate_amd() {
+  # clang2py broken when pass -x c++ to prev headers
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    extra/hip_gpu_driver/nvd.h \
+    extra/hip_gpu_driver/gc_11_0_0_offset.h \
+    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/amd_gpu.py
+
+  fixup $BASE/amd_gpu.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/amd_gpu.py
+  python3 -c "import tinygrad.runtime.autogen.amd_gpu"
+}
+
+generate_hsa() {
+  clang2py \
+    /opt/rocm/include/hsa/hsa.h \
+    /opt/rocm/include/hsa/hsa_ext_amd.h \
+    /opt/rocm/include/hsa/amd_hsa_signal.h \
+    /opt/rocm/include/hsa/amd_hsa_queue.h \
+    /opt/rocm/include/hsa/amd_hsa_kernel_code.h \
+    /opt/rocm/include/hsa/hsa_ext_finalize.h /opt/rocm/include/hsa/hsa_ext_image.h \
+    /opt/rocm/include/hsa/hsa_ven_amd_aqlprofile.h \
+    --clang-args="-I/opt/rocm/include" \
+    -o $BASE/hsa.py -l /opt/rocm/lib/libhsa-runtime64.so
+
+  fixup $BASE/hsa.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/hsa.py
+  sed -i "s\ctypes.CDLL('/opt/rocm/lib/libhsa-runtime64.so')\ctypes.CDLL(os.getenv('ROCM_PATH')+'/lib/libhsa-runtime64.so' if os.getenv('ROCM_PATH') else ctypes.util.find_library('hsa-runtime64'))\g" $BASE/hsa.py
+  python3 -c "import tinygrad.runtime.autogen.hsa"
+}
+
+generate_io_uring() {
+  clang2py -k cdefstum \
+    /usr/include/liburing.h \
+    /usr/include/linux/io_uring.h \
+    -o $BASE/io_uring.py
+
+  sed -r '/^#define __NR_io_uring/ s/^#define __(NR_io_uring[^ ]+) (.*)$/\1 = \2/; t; d' /usr/include/asm-generic/unistd.h >> $BASE/io_uring.py # io_uring syscalls numbers
+  fixup $BASE/io_uring.py
+}
+
+generate_ib() {
+  clang2py -k cdefstum \
+    /usr/include/infiniband/verbs.h \
+    /usr/include/infiniband/verbs_api.h \
+    /usr/include/infiniband/ib_user_ioctl_verbs.h \
+    /usr/include/rdma/ib_user_verbs.h \
+    -o $BASE/ib.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util\g" "$BASE/ib.py"
+  sed -i "s\FIXME_STUB\libibverbs\g" "$BASE/ib.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(ctypes.util.find_library('ibverbs'), use_errno=True)\g" "$BASE/ib.py"
+
+  fixup $BASE/ib.py
+}
+
+generate_libc() {
+  clang2py -k cdefstum \
+    $(dpkg -L libc6-dev | grep sys/mman.h) \
+    $(dpkg -L libc6-dev | grep sys/syscall.h) \
+    /usr/include/string.h \
+    /usr/include/elf.h \
+    /usr/include/unistd.h \
+    /usr/include/asm-generic/mman-common.h \
+    -o $BASE/libc.py
+
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
+  sed -i "s\FIXME_STUB\libc\g" $BASE/libc.py
+  sed -i "s\FunctionFactoryStub()\None if (libc_path := ctypes.util.find_library('c')) is None else ctypes.CDLL(libc_path, use_errno=True)\g" $BASE/libc.py
+
+  fixup $BASE/libc.py
+}
+
+generate_llvm() {
+  INC="$(llvm-config-14 --includedir)"
+  clang2py -k cdefstum \
+    $(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
+    "$INC/llvm/Config/Targets.def" \
+    "$INC/llvm/Config/AsmPrinters.def" \
+    "$INC/llvm/Config/AsmParsers.def" \
+    "$INC/llvm/Config/Disassemblers.def" \
+    --clang-args="$(llvm-config-14 --cflags)" \
+    -o "$BASE/llvm.py"
+
+  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
+  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
+  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
+
+  fixup "$BASE/llvm.py"
+}
+
+generate_kgsl() {
+  clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
+  fixup $BASE/kgsl.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
+  sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
+  python3 -c "import tinygrad.runtime.autogen.kgsl"
+}
+
+generate_adreno() {
+  clang2py extra/qcom_gpu_driver/a6xx.xml.h -o $BASE/adreno.py -k cestum
+  sed -nE 's/#define ([A-Za-z0-9_]+)__SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1__SHIFT) \& \1__MASK/p' extra/qcom_gpu_driver/a6xx.xml.h >> $BASE/adreno.py
+  fixup $BASE/adreno.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/adreno.py
+  python3 -c "import tinygrad.runtime.autogen.adreno"
+}
+
+generate_qcom() {
+  clang2py -k cdefstum \
+    extra/dsp/include/ion.h \
+    extra/dsp/include/msm_ion.h \
+    extra/dsp/include/adsprpc_shared.h \
+    extra/dsp/include/remote_default.h \
+    extra/dsp/include/apps_std.h \
+    -o $BASE/qcom_dsp.py
+
+  fixup $BASE/qcom_dsp.py
+  python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
+}
+
+generate_pci() {
+  clang2py -k cdefstum \
+    /usr/include/linux/pci_regs.h \
+    -o $BASE/pci.py
+  fixup $BASE/pci.py
+}
+
+generate_vfio() {
+  clang2py -k cdefstum \
+    /usr/include/linux/vfio.h \
+    -o $BASE/vfio.py
+  fixup $BASE/vfio.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
+  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
+  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import FileIOInterface\g" $BASE/vfio.py
+  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
+}
+
+generate_am() {
+  AMKERN_COMMIT_HASH=ceb12c04e2b5b53ec0779362831f5ee40c4921e4
+  AMKERN_SRC=/tmp/ROCK-Kernel-Driver-$AMKERN_COMMIT_HASH
+  if [ ! -d "$AMKERN_SRC" ]; then
+    git clone https://github.com/ROCm/ROCK-Kernel-Driver $AMKERN_SRC --depth 1
+  fi
+  AMKERN_AMD=$AMKERN_SRC/drivers/gpu/drm/amd/
+  AMKERN_INC=$AMKERN_AMD/include/
+
+  clang2py -k cdefstum \
+    extra/amdpci/headers/v11_structs.h \
+    extra/amdpci/headers/v12_structs.h \
+    extra/amdpci/headers/amdgpu_vm.h \
+    extra/amdpci/headers/discovery.h \
+    extra/amdpci/headers/amdgpu_ucode.h \
+    extra/amdpci/headers/psp_gfx_if.h \
+    extra/amdpci/headers/amdgpu_psp.h \
+    extra/amdpci/headers/amdgpu_irq.h \
+    extra/amdpci/headers/amdgpu_doorbell.h \
+    $AMKERN_INC/soc15_ih_clientid.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/am.py
+  fixup $BASE/am/am.py
+  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
+  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/soc15d.h \
+    -o $BASE/am/pm4_soc15.py
+  fixup $BASE/am/pm4_soc15.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/amdkfd/kfd_pm4_headers_ai.h \
+    $AMKERN_AMD/amdgpu/nvd.h \
+    -o $BASE/am/pm4_nv.py
+  fixup $BASE/am/pm4_nv.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_4_0_0.py
+  fixup $BASE/am/sdma_4_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/navi10_sdma_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_5_0_0.py
+  fixup $BASE/am/sdma_5_0_0.py
+
+  clang2py -k cdefstum \
+    extra/hip_gpu_driver/sdma_registers.h \
+    $AMKERN_AMD/amdgpu/sdma_v6_0_0_pkt_open.h \
+    --clang-args="-I/opt/rocm/include -x c++" \
+    -o $BASE/am/sdma_6_0_0.py
+  fixup $BASE/am/sdma_6_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v13_0_0_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu13_driver_if_v13_0_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    -o $BASE/am/smu_v13_0_0.py
+  fixup $BASE/am/smu_v13_0_0.py
+
+  clang2py -k cdefstum \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_0_pmfw.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu_v14_0_2_ppsmc.h \
+    $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
+    extra/amdpci/headers/amdgpu_smu.h \
+    --clang-args="-include stdint.h" \
+    -o $BASE/am/smu_v14_0_2.py
+  fixup $BASE/am/smu_v14_0_2.py
+}
+
+generate_sqtt() {
+  clang2py -k cdefstum \
+    extra/sqtt/sqtt.h \
+    -o $BASE/sqtt.py
+
+  fixup $BASE/sqtt.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
+  python3 -c "import tinygrad.runtime.autogen.sqtt"
+}
+
+generate_webgpu() {
+  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
+  fixup $BASE/webgpu.py
+  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
+  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
+  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
+  python3 -c "import tinygrad.runtime.autogen.webgpu"
+}
+
+generate_libusb() {
+  clang2py -k cdefstum \
+    /usr/include/libusb-1.0/libusb.h \
+    -o $BASE/libusb.py
+
+  fixup $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
+  sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
+  sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
+  sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
+  python3 -c "import tinygrad.runtime.autogen.libusb"
+}
+
+if [ "$1" == "opencl" ]; then generate_opencl
+elif [ "$1" == "hip" ]; then generate_hip
+elif [ "$1" == "comgr" ]; then generate_comgr
+elif [ "$1" == "cuda" ]; then generate_cuda
+elif [ "$1" == "nvrtc" ]; then generate_nvrtc
+elif [ "$1" == "hsa" ]; then generate_hsa
+elif [ "$1" == "kfd" ]; then generate_kfd
+elif [ "$1" == "nv" ]; then generate_nv
+elif [ "$1" == "amd" ]; then generate_amd
+elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
+elif [ "$1" == "sqtt" ]; then generate_sqtt
+elif [ "$1" == "qcom" ]; then generate_qcom
+elif [ "$1" == "io_uring" ]; then generate_io_uring
+elif [ "$1" == "ib" ]; then generate_ib
+elif [ "$1" == "libc" ]; then generate_libc
+elif [ "$1" == "llvm" ]; then generate_llvm
+elif [ "$1" == "kgsl" ]; then generate_kgsl
+elif [ "$1" == "adreno" ]; then generate_adreno
+elif [ "$1" == "pci" ]; then generate_pci
+elif [ "$1" == "vfio" ]; then generate_vfio
+elif [ "$1" == "webgpu" ]; then generate_webgpu
+elif [ "$1" == "libusb" ]; then generate_libusb
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+else echo "usage: $0 <type>"
+fi
@@ -0,0 +1,137 @@
+# tinygrad is a tensor library, and as a tensor library it has multiple parts
+# 1. a "runtime". this allows buffer management, compilation, and running programs
+# 2. a "Device" that uses the runtime but specifies compute in an abstract way for all
+# 3. a "UOp" that fuses the compute into kernels, using memory only when needed
+# 4. a "Tensor" that provides an easy to use frontend with autograd ".backward()"
+
+
+print("******** first, the runtime ***********")
+
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
+
+cpu = CPUDevice()
+
+# allocate some buffers
+out = cpu.allocator.alloc(4)
+a = cpu.allocator.alloc(4)
+b = cpu.allocator.alloc(4)
+
+# load in some values (little endian)
+cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+
+# compile a program to a binary
+lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+
+# create a runtime for the program
+fxn = cpu.runtime("add", lib)
+
+# run the program
+fxn(out, a, b)
+
+# check the data out
+print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
+assert val == 5
+
+
+print("******** second, the Device ***********")
+
+DEVICE = "CPU"   # NOTE: you can change this!
+
+import struct
+from tinygrad.dtype import dtypes
+from tinygrad.device import Buffer, Device
+from tinygrad.uop.ops import UOp, Ops
+from tinygrad.shape.shapetracker import ShapeTracker
+
+# allocate some buffers + load in values
+out = Buffer(DEVICE, 1, dtypes.int32).allocate()
+a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+# NOTE: a._buf is the same as the return from cpu.allocator.alloc
+
+# describe the computation
+buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
+buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
+alu = ld_1 + ld_2
+output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
+s = UOp(Ops.SINK, dtypes.void, (st_0,))
+
+# convert the computation to a "linearized" format (print the format)
+from tinygrad.engine.realize import get_program, CompiledRunner
+program = get_program(s, Device[DEVICE].renderer)
+
+# compile a program (and print the source)
+fxn = CompiledRunner(program)
+print(fxn.p.src)
+# NOTE: fxn.clprg is the CPUProgram
+
+# run the program
+fxn.exec([out, a, b])
+
+# check the data out
+assert out.as_buffer().cast('I')[0] == 5
+
+
+print("******** third, the UOp ***********")
+
+from tinygrad.engine.realize import run_schedule
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.schedule.kernelize import get_kernelize_map
+
+# allocate some values + load in values
+a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+b = UOp.new_buffer(DEVICE, 1, dtypes.int32)
+a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
+b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+
+# describe the computation
+out = a + b
+s = UOp(Ops.SINK, dtypes.void, (out,))
+
+# group the computation into kernels
+becomes_map = get_kernelize_map(s)
+
+# the compute maps to an assign
+assign = becomes_map[a+b]
+
+# the first source is the output buffer (data)
+assert assign.src[0].op is Ops.BUFFER
+# the second source is the kernel (compute)
+assert assign.src[1].op is Ops.KERNEL
+
+# schedule the kernel graph in a linear list
+s = UOp(Ops.SINK, dtypes.void, (assign,))
+sched, _ = create_schedule_with_vars(s)
+assert len(sched) == 1
+
+# DEBUGGING: print the compute ast
+print(sched[-1].ast)
+# NOTE: sched[-1].ast is the same as st_0 above
+
+# the output will be stored in a new buffer
+out = assign.buf_uop
+assert out.op is Ops.BUFFER and not out.buffer.is_allocated()
+print(out)
+
+# run that schedule
+run_schedule(sched)
+
+# check the data out
+assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
+
+
+print("******** fourth, the Tensor ***********")
+
+from tinygrad import Tensor
+
+a = Tensor([2], dtype=dtypes.int32, device=DEVICE)
+b = Tensor([3], dtype=dtypes.int32, device=DEVICE)
+out = a + b
+
+# check the data out
+print(val:=out.item())
+assert val == 5
@@ -0,0 +1,109 @@
+# Kernel Creation
+
+Tinygrad lazily builds up a graph of Tensor operations. The Tensor graph includes a mix of:
+
+- Buffer and Assignment Ops: `BUFFER`, `BUFFER_VIEW`, `COPY`, `ASSIGN`
+- Movement Ops: `RESHAPE`, `EXPAND`, `PERMUTE`, `PAD`, `SHRINK`, `FLIP`
+- Compute Ops: `ADD`, `MUL`, `REDUCE_AXIS`, ...
+
+`Tensor.kernelize` creates the kernels and buffers needed to realize the output Tensor(s).
+
+## Kernelize flow
+
+Let's see how a multiply add Tensor graph becomes a fused elementwise kernel.
+
+```py
+# initialize 3 input buffers on the device
+a = Tensor([1]).realize()
+b = Tensor([2]).realize()
+c = Tensor([3]).realize()
+
+# create the Tensor graph
+mul = a*b
+out = mul+c
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ADD: 52>, None)> on METAL with grad None>
+
+out.kernelize()
+
+print(mul) # <Tensor <UOp METAL (1,) int (<Ops.MUL: 48>, None)> on METAL with grad None>
+print(out) # <Tensor <UOp METAL (1,) int (<Ops.ASSIGN: 66>, None)> on METAL with grad None>
+```
+
+The multiply Tensor stays the same because it is fused. The output Tensor's UOp becomes a new ASSIGN UOp:
+
+```py
+print(out.uop)
+```
+
+The first source is the output BUFFER:
+
+```
+UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+  UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),))
+```
+
+And the second source is the KERNEL and its 4 buffer edges (output_buffer, a, b, c):
+
+```
+UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 45>,) (__add__, __mul__)>, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+    x1:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),
+    UOp(Ops.UNIQUE, dtypes.void, arg=6, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=3, src=()),)),
+  UOp(Ops.BUFFER, dtypes.int, arg=1, src=(
+     x1,
+    UOp(Ops.UNIQUE, dtypes.void, arg=5, src=()),)),))
+```
+
+KERNEL describes the compute AST, metadata and memory dependencies.
+
+BUFFER holds a reference to the device memory where the output will be stored.
+
+Once a Tensor is kernelized, all children will LOAD its BUFFER, instead of fusing it:
+
+```py
+child = out+2
+child.kernelize()
+print(child.uop.src[1].arg.ast)
+```
+
+```
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=0, src=()),
+    x2:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=()),
+         x2,)),
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+         x2,)),)),)),))
+```
+
+`Tensor.realize` will execute the kernels and write outputs to memory:
+
+```py
+Tensor.realize(out)
+print(out)        # <Tensor <UOp METAL (1,) int (<Ops.BUFFER: 23>, <buf real:True device:METAL size:1 dtype:dtypes.int offset:0>)> on METAL with grad None>
+print(out.item()) # 5
+```
+
+<hr />
+
+**Summary**
+
+- The large Tensor graph is built from a mix of data, compute and movement Ops.
+
+- `Tensor.kernelize` splits the Tensor graph into data (BUFFER), compute (KERNEL) and links dependencies with ASSIGN.
+
+- `Tensor.realize` executes KERNELs on device and replaces the Tensor graph with just a BUFFER.
+
+- Kernelize can be called multiple times on a Tensor. This allows for incrementally building the kernel fusion layout of a large Tensor graph, without having to call `realize` or `schedule`.
@@ -10,7 +10,7 @@ Directories are listed in order of how they are processed.

 Group UOps into kernels.

-::: tinygrad.schedule.rangeify.get_rangeify_map
+::: tinygrad.schedule.kernelize.get_kernelize_map
    options:
        members: false
        show_labels: false
@@ -22,6 +22,12 @@ Group UOps into kernels.

 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.

+::: tinygrad.codegen.opt.get_optimized_ast
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
 ---

 ## tinygrad/codegen
@@ -3,7 +3,7 @@
 This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
 Most of these are self-explanatory, and are usually used to set an option at runtime.

-Example: `CL=1 DEBUG=4 python3 -m pytest`
+Example: `GPU=1 DEBUG=4 python3 -m pytest`

 However you can also decorate a function to set a value only inside that function.

@@ -31,17 +31,20 @@ These control the behavior of core tinygrad even when used as a library.
 Variable | Possible Value(s) | Description
 ---|---|---
 DEBUG               | [1-7]      | enable debugging output (operations, timings, speed, generated code and more)
-CL                  | [1]        | enable OpenCL backend
+GPU                 | [1]        | enable the GPU (OpenCL) backend
 CUDA                | [1]        | enable CUDA backend
 AMD                 | [1]        | enable AMD backend
 NV                  | [1]        | enable NV backend
 METAL               | [1]        | enable Metal backend (for Mac M1 and after)
-CPU                 | [1]        | enable CPU backend
+CPU                 | [1]        | enable CPU (Clang) backend
+LLVM                | [1]        | enable LLVM backend
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
-HCQ_VISIBLE_DEVICES | [list[int]]| restricts the HCQ devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
+PTX                 | [1]        | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
+PROFILE             | [1]        | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
+VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
 VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
 ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
@@ -131,7 +131,7 @@ timeit.repeat(jit_step, repeat=5, number=1)

 1.0 ms is 75x faster! Note that we aren't syncing the GPU, so GPU time may be slower.

-The first two runs of the function execute normally, with the JIT capturing the kernels. Starting from the third run, only the tinygrad operations are replayed, removing the overhead by skipping Python code execution. So be aware that any non-tinygrad Python values affecting the kernels will be "frozen" from the second run. Note that `Tensor` randomness functions work as expected.
+The slowness the first two times is the JIT capturing the kernels. And this JIT will not run any Python in the function, it will just replay the tinygrad kernels that were run, so be aware that non tinygrad Python operations won't work. Randomness functions work as expected.

 Unlike other JITs, we JIT everything, including the optimizer. Think of it as a dumb replay on different data.

@@ -0,0 +1,293 @@
+#!/usr/bin/env python3
+
+# this file is a "ramp" for people new to tinygrad to think about how to approach it
+# it is runnable and editable.
+# whenever you see stuff like DEBUG=2 or CPU=1 discussed, these are environment variables
+# in a unix shell like bash `DEBUG=2 CPU=1 python docs/ramp.py`
+
+# this pip installs tinygrad master for the system
+# the -e allows you to edit the tinygrad folder and update system tinygrad
+# tinygrad is pure Python, so you are encouraged to do this
+# git pull in the tinygrad directory will also get you the latest
+"""
+git clone https://github.com/tinygrad/tinygrad.git
+cd tinygrad
+python3 -m pip install -e .
+"""
+
+# %% ********
+print("******* PART 1 *******")
+
+# we start with a Device.
+# a Device is where Tensors are stored and compute is run
+# tinygrad autodetects the best device on your system and makes it the DEFAULT
+from tinygrad import Device
+print(Device.DEFAULT)  # on Mac, you can see this prints METAL
+
+# now, lets create a Tensor
+from tinygrad import Tensor, dtypes
+t = Tensor([1,2,3,4])
+
+# you can see this Tensor is on the DEFAULT device with int dtype and shape (4,)
+assert t.device == Device.DEFAULT
+assert t.dtype == dtypes.int
+assert t.shape == (4,)
+
+# unlike in torch, if we print it, it doesn't print the contents
+# this is because tinygrad is lazy
+# this Tensor has not been computed yet
+print(t)
+# <Tensor <UOp METAL (4,) int (<Ops.COPY: 7>, None)> on METAL with grad None>
+
+# the ".uop" property on Tensor contains the specification of how to compute it
+print(t.uop)
+"""
+UOp(Ops.COPY, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=0, src=()),
+    UOp(Ops.DEVICE, dtypes.void, arg='PYTHON', src=()),)),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# as you can see, it's specifying a copy from PYTHON device
+# which is where the [1,2,3,4] array lives
+
+# UOps are the specification language in tinygrad
+# they are immutable and form a DAG
+# they have a "Ops", a "dtype", a tuple of srcs (parents), and an arg
+
+t.realize()
+# if we want to "realize" a tensor, we can with the "realize" method
+# now when we look at the uop, it's changed
+print(t.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# the copy was actually run, and now the "uop" of the Tensor is just a BUFFER
+# if you run this script with DEBUG=2 in the environment, you can see the copy happen
+# *** METAL      1 copy       16,   METAL <- PYTHON ...
+
+# now let's do some compute
+# we look at the uop to see the specification of the compute
+t_times_2 = t * 2
+print(t_times_2.uop)
+"""
+UOp(Ops.MUL, dtypes.int, arg=None, src=(
+  UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=2, src=(
+        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+           x2,)),)),)),)),))
+"""
+# the BUFFER from above is being multiplied by a CONST 2
+# it's RESHAPEd and EXPANDed to broadcast the CONST to the BUFFER
+
+# we can check the result with
+assert t_times_2.tolist() == [2, 4, 6, 8]
+
+# UOps are both immutable and globally unique
+# if i multiply the Tensor by 4 twice, these result Tensors will have the same uop specification
+t_times_4_try_1 = t * 4
+t_times_4_try_2 = t * 4
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# the specification isn't just the same, it's the exact same Python object
+assert t_times_4_try_1 is not t_times_4_try_2
+# the Tensor is a different Python object
+
+# if we realize `t_times_4_try_1` ...
+t_times_4_try_1.realize()
+print(t_times_4_try_2.uop)
+"""
+UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+  UOp(Ops.UNIQUE, dtypes.void, arg=4, src=()),
+  UOp(Ops.DEVICE, dtypes.void, arg='METAL', src=()),))
+"""
+# ... `t_times_4_try_2` also becomes the same BUFFER
+assert t_times_4_try_1.uop is t_times_4_try_2.uop
+# so this print doesn't require any computation, just a copy back to the CPU so we can print it
+print("** only the copy start")
+print(t_times_4_try_2.tolist())  # [4, 8, 12, 16]
+print("** only the copy end")
+# you can confirm this with DEBUG=2, seeing what's printed in between the "**" prints
+
+# tinygrad has an auto differentiation engine that operates according to these same principles
+# the derivative of "log(x)" is "1/x", and you can see this on line 20 of gradient.py
+t_float = Tensor([3.0])
+t_log = t_float.log()
+t_log_grad, = t_log.sum().gradient(t_float)
+# due to how log is implemented, this gradient contains a lot of UOps
+print(t_log_grad.uop)
+# ...not shown here...
+# but if you run with DEBUG=4 (CPU=1 used here for simpler code), you can see the generated code
+"""
+void E_(float* restrict data0, float* restrict data1) {
+  float val0 = *(data1+0);
+  *(data0+0) = (1/val0);
+}
+"""
+# the derivative is close to 1/3
+assert (t_log_grad.item() - 1/3) < 1e-6
+
+# %% ********
+print("******* PART 2 *******")
+
+# we redefine the same t here so this cell can run on it's own
+from tinygrad import Tensor
+t = Tensor([1,2,3,4])
+
+# what's above gives you enough of an understanding to go use tinygrad as a library
+# however, a lot of the beauty of tinygrad is in how easy it is to interact with the internals
+# NOTE: the APIs here are subject to change
+
+t_plus_3_plus_4 = t + 3 + 4
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  UOp(Ops.ADD, dtypes.int, arg=None, src=(
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+      x3:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+    UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+      UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+        UOp(Ops.CONST, dtypes.int, arg=3, src=(
+          x7:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(), strides=(), offset=0, mask=None, contiguous=True),)), src=(
+             x3,)),)),)),)),)),
+  UOp(Ops.EXPAND, dtypes.int, arg=(4,), src=(
+    UOp(Ops.RESHAPE, dtypes.int, arg=(1,), src=(
+      UOp(Ops.CONST, dtypes.int, arg=4, src=(
+         x7,)),)),)),))
+"""
+# you can see it's adding both 3 and 4
+
+# but by the time we are actually running the code, it's adding 7
+# `kernelize` will simplify and group the operations in the graph into kernels
+t_plus_3_plus_4.kernelize()
+print(t_plus_3_plus_4.uop)
+"""
+UOp(Ops.ASSIGN, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+    UOp(Ops.UNIQUE, dtypes.void, arg=7, src=()),
+    x2:=UOp(Ops.DEVICE, dtypes.void, arg='CPU', src=()),)),
+  UOp(Ops.KERNEL, dtypes.void, arg=<Kernel 12 SINK(<Ops.STORE: 48>,) (__add__,)>, src=(
+     x0,
+    UOp(Ops.BUFFER, dtypes.int, arg=4, src=(
+      UOp(Ops.UNIQUE, dtypes.void, arg=1, src=()),
+       x2,)),)),))
+"""
+# ASSIGN has two srcs, src[0] is the BUFFER that's assigned to, and src[1] is the thing to assign
+# src[1] is the GPU Kernel that's going to be run
+# we can get the ast of the Kernel as follows
+kernel_ast = t_plus_3_plus_4.uop.src[1].arg.ast
+
+# almost everything in tinygrad functions as a rewrite of the UOps
+# the codegen rewrites the ast to a simplified form ready for "rendering"
+from tinygrad.codegen import full_rewrite_to_sink
+rewritten_ast = full_rewrite_to_sink(kernel_ast)
+print(rewritten_ast)
+"""
+UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+      UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=0, src=()),
+      x3:=UOp(Ops.SPECIAL, dtypes.int, arg=('gidx0', 4), src=()),)),
+    UOp(Ops.ADD, dtypes.int, arg=None, src=(
+      UOp(Ops.LOAD, dtypes.int, arg=None, src=(
+        UOp(Ops.INDEX, dtypes.int.ptr(4), arg=None, src=(
+          UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(4), arg=1, src=()),
+           x3,)),)),
+      UOp(Ops.CONST, dtypes.int, arg=7, src=()),)),)),))
+"""
+# you can see at this point we are adding 7, not 3 and 4
+
+# with DEBUG=4, we can see the code.
+# since optimizations are on, it UPCASTed the operation, explicitly writing out all 4 +7s
+t_plus_3_plus_4.realize()
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  int val0 = *(data1+0);
+  int val1 = *(data1+1);
+  int val2 = *(data1+2);
+  int val3 = *(data1+3);
+  *(data0+0) = (val0+7);
+  *(data0+1) = (val1+7);
+  *(data0+2) = (val2+7);
+  *(data0+3) = (val3+7);
+}
+"""
+# the function name E_4n2 is "E" for elementwise op (as opposed to "r" for reduce op)
+# "4" for the size, and "n2" for name deduping (it's the 3rd function with the same E and 4 in this session)
+# when you print the name with DEBUG=2, you'll see the 4 is yellow, meaning that it's upcasted
+# if you run with NOOPT=1 ...
+"""
+void E_4n2(int* restrict data0, int* restrict data1) {
+  for (int ridx0 = 0; ridx0 < 4; ridx0++) {
+    int val0 = *(data1+ridx0);
+    *(data0+ridx0) = (val0+7);
+  }
+}
+"""
+# ... you get this unoptimized code with a loop and the 4 is blue (for global). the color code is in kernel.py
+
+# %% ********
+print("******* PART 3 *******")
+
+# now, we go even lower and understand UOps better and how the graph rewrite engine works.
+# it's much simpler than what's in LLVM or MLIR
+
+from tinygrad import dtypes
+from tinygrad.uop.ops import UOp, Ops
+
+# first, we'll construct some const UOps
+a = UOp(Ops.CONST, dtypes.int, arg=2)
+b = UOp(Ops.CONST, dtypes.int, arg=2)
+
+# if you have been paying attention, you should know these are the same Python object
+assert a is b
+
+# UOps support normal Python math operations, so a_plus_b expresses the spec for 2 + 2
+a_plus_b = a + b
+print(a_plus_b)
+"""
+UOp(Ops.ADD, dtypes.int, arg=None, src=(
+  x0:=UOp(Ops.CONST, dtypes.int, arg=2, src=()),
+   x0,))
+"""
+
+# we could actually render this 2+2 into a language like c and run it
+# or, we can use tinygrad's graph rewrite engine to "constant fold"
+
+from tinygrad.uop.ops import graph_rewrite, UPat, PatternMatcher
+
+# a `PatternMatcher` is a list of tuples. for each element in the list:
+# [0] is the pattern to match, and [1] is the function to run.
+# this function can return either a UOp to replace the pattern with, or None to not replace
+simple_pm = PatternMatcher([
+  (UPat(Ops.ADD, src=(UPat(Ops.CONST, name="c1"), UPat(Ops.CONST, name="c2"))),
+   lambda c1,c2: UOp(Ops.CONST, dtype=c1.dtype, arg=c1.arg+c2.arg)),
+])
+# this pattern matches the addition of two CONST and rewrites it into a single CONST UOp
+
+# to actually apply the pattern to a_plus_b, we use graph_rewrite
+a_plus_b_simplified = graph_rewrite(a_plus_b, simple_pm)
+print(a_plus_b_simplified)
+"""
+UOp(Ops.CONST, dtypes.int, arg=4, src=())
+"""
+# 2+2 is in fact, 4
+
+# we can also use syntactic sugar to write the pattern nicer
+simpler_pm = PatternMatcher([
+  (UPat.cvar("c1")+UPat.cvar("c2"), lambda c1,c2: c1.const_like(c1.arg+c2.arg))
+])
+assert graph_rewrite(a_plus_b, simple_pm) is graph_rewrite(a_plus_b, simpler_pm)
+# note again the use of is, UOps are immutable and globally unique
+
+# %% ********
+
+# that brings you to an understanding of the most core concepts in tinygrad
+# you can run this with VIZ=1 to use the web based graph rewrite explorer
+# hopefully now you understand it. the nodes in the graph are just UOps
@@ -2,17 +2,17 @@

 tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).

-| Runtime | Description | Compiler Options | Requirements |
-|---------|-------------|------------------|--------------|
-| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | nvrtc (default)<br>PTX (`NV_PTX=1`) | Ampere/Ada/Blackwell series GPUs.<br>You can select an interface via `NV_IFACE=(NVK\|PCI)`. See [NV interfaces](#nv-interfaces) for details. |
-| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | LLVM (`AMD_LLVM=1`)<br>HIP/COMGR (`AMD_HIP=1`) | RDNA2 or newer GPUs.<br>You can select an interface via `AMD_IFACE=(KFD\|PCI\|USB)`. See [AMD interfaces](#amd-interfaces) for details. |
-| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | - | 6xx series GPUs |
-| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | - | M1+ Macs; Metal 3.0+ for `bfloat` support |
-| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | nvrtc (default)<br> PTX (`CUDA_PTX=1`) | NVIDIA GPU with CUDA support |
-| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
-| [CPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang or llvm compiler | Clang JIT (default)<br>LLVM IR (`CPU_LLVM=1`) | `clang` compiler in system `PATH` |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | - | Dawn library installed and discoverable. Binaries: [pydawn v0.3.0](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0) |
-
+| Runtime | Description | Requirements |
+|---------|-------------|--------------|
+| [NV](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_nv.py) | Provides acceleration for NVIDIA GPUs | Ampere/Ada series GPUs |
+| [AMD](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_amd.py) | Provides acceleration for AMD GPUs | RDNA2/RDNA3/RDNA4 series GPUs. You can select one of the interfaces for communication by setting `AMD_IFACE=(KFD|PCI)`. See [AMD interfaces](#amd-interfaces) for more details. |
+| [QCOM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_qcom.py) | Provides acceleration for QCOM GPUs | 6xx series GPUs |
+| [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
+| [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
+| [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
+| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
+| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
+| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |

 ## Interoperability

@@ -70,12 +70,5 @@ AMD backend supports several interfaces for communicating with devices:

 * `KFD`: uses the amdgpu driver
 * `PCI`: uses the [AM driver](developer/am.md)
-* `USB`: USB3 interafce for asm24xx chips.

 You can force an interface by setting `AMD_IFACE` to one of these values. In the case of `AMD_IFACE=PCI`, this may unbind your GPU from the amdgpu driver.
-
-## NV Interfaces
-NV backend supports several interfaces for communicating with devices:
-
-* `NVK`: uses the nvidia driver
-* `PCI`: uses the [NV driver](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/support/nv/nvdev.py)
@@ -78,7 +78,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
 ::: tinygrad.Tensor.copysign
-::: tinygrad.Tensor.logaddexp

 ## Casting Ops

@@ -6,7 +6,7 @@ If you don't have a tinybox and you want one, see [tinygrad.org](https://tinygra

 ## Welcome

-Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, the green box includes six 4090 GPUs, and the green v2 box includes four 5090 GPUs. Whether you bought a red one or a green one, we want you to love it.
+Welcome to your tinybox! The tinybox is the universal system purpose-built for all AI infrastructure and workloads, from training to inference. The red box includes six 7900XTX GPUs, and the green box includes six 4090 GPUs. Whether you bought a red one or a green one, we want you to love it.

 We don't have a stupid cloud service, you don't have to create a tiny account to set it up, and we aren't tracking how you use the box. We're just happy you bought one. This petaflop is your petaflop.

@@ -41,7 +41,7 @@ The BMC also has a web interface you can use if you find that easier.
 It is recommended that you change the BMC password after setting up the box, as the password on the screen is only the initial password.

 If you do decide to change the BMC password and no longer want the initial password to be displayed, remove the `/root/.bmc_password` file.
-Reboot after making these changes or restart the `tinybox-display.service` service.
+Reboot after making these changes or restart the `displayservice.service` service.

 ## What do I use it for?

@@ -2,6 +2,7 @@ import time
 start_tm = time.perf_counter()
 import math
 from typing import Tuple, cast
+import numpy as np
 from tinygrad import Tensor, nn, GlobalCounters, TinyJit, dtypes, Device
 from tinygrad.helpers import partition, trange, getenv, Context
 from extra.lr_scheduler import OneCycleLR
@@ -10,7 +11,7 @@ GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]

 # override tinygrad defaults
 dtypes.default_float = dtypes.half
-Context(FUSE_OPTIM=1).__enter__()
+Context(FUSE_ARANGE=1, FUSE_OPTIM=1).__enter__()

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -149,12 +150,13 @@ if __name__ == "__main__":
      acc.append((out.argmax(-1) == Y).sum() / eval_batchsize)
    return Tensor.stack(*loss).mean() / (batchsize*loss_batchsize_scaler), Tensor.stack(*acc).mean()

-  Tensor.manual_seed(1337)
-  num_train_samples = X_train.shape[0]
-
+  np.random.seed(1337)
  for epoch in range(math.ceil(hyp['misc']['train_epochs'])):
+    # TODO: move to tinygrad
    gst = time.perf_counter()
-    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
+    idxs = np.arange(X_train.shape[0])
+    np.random.shuffle(idxs)
+    tidxs = Tensor(idxs, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)  # NOTE: long doesn't fold
    train_loss:float = 0
    for epoch_step in (t:=trange(num_steps_per_epoch)):
      st = time.perf_counter()
@@ -21,7 +21,7 @@ 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))
+  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))

  @TinyJit
  @Tensor.train()
@@ -1,6 +1,6 @@
 import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs, validate

 def load_onnx_model(onnx_file):
@@ -8,7 +8,7 @@ import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
 from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model
@@ -26,8 +26,8 @@ class Attention:
      start_pos = start_pos.val

    if HALF: x = x.half()
-    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
-    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
+    xqkv = self.c_attn(x)
+    xq, xk, xv = [xqkv.shrink((None, None, (i*self.dim, (i+1)*self.dim))).reshape(None, None, self.n_heads, self.head_dim) for i in range(3)]
    bsz, seqlen, _, _ = xq.shape

    # create kv cache
@@ -35,11 +35,11 @@ class Attention:
      self.cache_kv = Tensor.zeros(2, bsz, MAX_CONTEXT, self.n_heads, self.head_dim, dtype=x.dtype).contiguous().realize()

    # update the cache
-    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()
+    self.cache_kv.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()

    if start_pos > 0:
-      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
+      values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
    else:
      keys = xk
      values = xv
@@ -64,7 +64,7 @@ class TransformerBlock:

  def __call__(self, x:Tensor, start_pos:Variable, mask:Optional[Tensor]):
    h = x + self.attn(self.ln_1(x), start_pos, mask).float()
-    return (h + self.mlp(self.ln_2(h))).contiguous()
+    return (h + self.mlp(self.ln_2(h)))

 class Transformer:
  def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
@@ -181,7 +181,6 @@ class GPT2:
    self.tokenizer = tokenizer

  def generate(self, prompt:str, max_length:int, temperature:float, timing:bool=False, batch_size:int=1):
-    step_times = []
    prompt_tokens = self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})
    toks = [prompt_tokens[:] for _ in range(batch_size)]
    start_pos = 0
@@ -189,7 +188,7 @@ class GPT2:
      GlobalCounters.reset()
      if timing: print("")
      st = GlobalCounters.time_sum_s
-      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+      with Timing("ran model in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                  f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                  (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=timing):
        with WallTimeEvent(BenchEvent.STEP):
@@ -198,13 +197,8 @@ class GPT2:
          else:
            tokens = Tensor([x[start_pos:] for x in toks])
          tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT-1).bind(start_pos), temperature).tolist()
-      step_times.append((GlobalCounters.time_sum_s-st)*1e3)
      start_pos = len(toks[0])
      for i,t in enumerate(tok): toks[i].append(t)
-
-    if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-      min_time = min(step_times)
-      assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
    return [self.tokenizer.decode(x) for x in toks]

 # **** main code ****
@@ -232,7 +226,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -1,108 +0,0 @@
-import itertools
-from typing import Callable
-from tinygrad import nn, Tensor, dtypes, Device, TinyJit
-from tinygrad.helpers import getenv, trange, partition
-
-class Model:
-  def __init__(self):
-    self.layers: list[Callable[[Tensor], Tensor]] = [
-      nn.Conv2d(1, 32, 5), Tensor.relu,
-      nn.Conv2d(32, 32, 5), Tensor.relu,
-      nn.BatchNorm(32), Tensor.max_pool2d,
-      nn.Conv2d(32, 64, 3), Tensor.relu,
-      nn.Conv2d(64, 64, 3), Tensor.relu,
-      nn.BatchNorm(64), Tensor.max_pool2d,
-      lambda x: x.flatten(1), nn.Linear(576, 10)]
-
-  def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers)
-
-# TODO: refactor this into optim/onnx
-def functional_adam(g:Tensor, m:Tensor, v:Tensor, b1_t:Tensor, b2_t:Tensor, lr=0.001, b1=0.9, b2=0.999, eps=1e-6) -> Tensor:
-  b1_t *= b1
-  b2_t *= b2
-  m.assign(b1 * m + (1.0 - b1) * g)
-  v.assign(b2 * v + (1.0 - b2) * (g * g))
-  m_hat = m / (1.0 - b1_t)
-  v_hat = v / (1.0 - b2_t)
-  return lr * (m_hat / (v_hat.sqrt() + eps))
-
-if __name__ == "__main__":
-  BS = getenv("BS", 512)
-  ACC_STEPS = getenv("ACC_STEPS", 8)
-
-  X_train, Y_train, X_test, Y_test = nn.datasets.mnist()
-  model = Model()
-
-  params = nn.state.get_parameters(model)
-
-  # init params, set requires grad on the ones we need gradients of
-  for x in params:
-    if x.requires_grad is None: x.requires_grad_()
-    x.replace(x.contiguous())
-  Tensor.realize(*params)
-
-  # split params (with grads) and buffers (without)
-  params, buffers = partition(params, lambda x: x.requires_grad)
-  print(f"params: {len(params)} buffers: {len(buffers)}")
-
-  # optim params
-  pos_params = list(itertools.accumulate(params, lambda x,y: x+y.numel(), initial=0))
-  adam_m = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_v = Tensor.zeros(pos_params[-1], device="CPU").contiguous()
-  adam_b1_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_b2_t = Tensor.ones((1,), dtype=dtypes.float32, device="CPU", requires_grad=False).contiguous()
-  adam_params = [adam_m, adam_v, adam_b1_t, adam_b2_t]
-
-  # create loss and grads. init all state so the JIT works on microbatch
-  for x in params: x.assign(x.detach())
-  loss = Tensor.zeros(tuple()).contiguous()
-  grads = Tensor.zeros(pos_params[-1]).contiguous()
-  Tensor.realize(*params, *buffers, *adam_params, loss, grads)
-
-  @TinyJit
-  @Tensor.train()
-  def microbatch():
-    samples = Tensor.randint(BS // ACC_STEPS, high=X_train.shape[0])
-    for t in params: t.grad = None
-    # divide by ACC_STEPS at the loss
-    uloss = (model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]) / ACC_STEPS).backward()
-    ugrads = Tensor.cat(*[t.grad.contiguous().flatten() for t in params], dim=0)
-    for t in params: t.grad = None
-    # concat the grads and assign them
-    loss.assign(loss + uloss)
-    grads.assign(grads + ugrads)
-    Tensor.realize(*params, *buffers, loss, grads)
-
-  @TinyJit
-  def optimizer():
-    # run optimizer (on CPU, where adam params live)
-    delta = functional_adam(grads.to("CPU"), adam_m, adam_v, adam_b1_t, adam_b2_t)
-
-    # update the params, copying back the delta one at a time to avoid OOM
-    # NOTE: the scheduler is ordering things poorly, all the copies are happening before the adds
-    for j,tt in enumerate(params):
-      tt.assign(tt.detach() - delta[pos_params[j]:pos_params[j+1]].reshape(tt.shape).to(Device.DEFAULT))
-
-    # realize everything, zero out loss and grads
-    loss.assign(Tensor.zeros_like(loss))
-    grads.assign(Tensor.zeros_like(grads))
-    Tensor.realize(*params, *adam_params, loss, grads)
-
-  @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))):
-    # microbatch sets the gradients
-    for _ in range(ACC_STEPS): microbatch()
-
-    # get the loss before the optimizer clears it
-    # this is already realized so this isn't a schedule
-    loss_item = loss.item()
-
-    # run the optimizer
-    optimizer()
-
-    # eval
-    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}%")
@@ -118,7 +118,7 @@ class SpeedyResNet:
 # hyper-parameters were exactly the same as the original repo
 bias_scaler = 58
 hyp = {
-  'seed' : 201,
+  'seed' : 200,
  'opt': {
    'bias_lr':            1.76 * bias_scaler/512,
    'non_bias_lr':        1.76 / 512,
@@ -145,6 +145,7 @@ hyp = {
  },
 }

+@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():

  def set_seed(seed):
@@ -228,8 +229,7 @@ def train_cifar():
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
-      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
-      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
    X, Y = X[perms], Y[perms]
    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])

@@ -355,7 +355,7 @@ def train_cifar():

  # https://www.anandtech.com/show/16727/nvidia-announces-geforce-rtx-3080-ti-3070-ti-upgraded-cards-coming-in-june
  # 136 TFLOPS is the theoretical max w float16 on 3080 Ti
-  step_times = []
+
  model_ema: Optional[modelEMA] = None
  projected_ema_decay_val = hyp['ema']['decay_base'] ** hyp['ema']['every_n_steps']
  i = 0
@@ -413,17 +413,12 @@ def train_cifar():
          model_ema.update(model, Tensor([projected_ema_decay_val*(i/STEPS)**hyp['ema']['decay_pow']]))

      cl = time.monotonic()
-      step_times.append((cl-st)*1000.0)
      device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
      #  53  221.74 ms run,    2.22 ms python,  219.52 ms CL,  803.39 loss, 0.000807 LR, 4.66 GB used,   3042.49 GFLOPS,    674.65 GOPS
      print(f"{i:3d} {(cl-st)*1000.0:7.2f} ms run, {(et-st)*1000.0:7.2f} ms python, {(cl-et)*1000.0:7.2f} ms {device_str}, {loss_cpu:7.2f} loss, {opt_non_bias.lr.numpy()[0]:.6f} LR, {GlobalCounters.mem_used/1e9:.2f} GB used, {GlobalCounters.global_ops*1e-9/(cl-st):9.2f} GFLOPS, {GlobalCounters.global_ops*1e-9:9.2f} GOPS")
      st = cl
      i += 1

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
  # verify eval acc
  if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
    if eval_acc_pct >= target:
@@ -478,7 +478,7 @@ After you are done speaking, output [EOS]. You are not Chad.
      with Profiling(enabled=args.profile):
        with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
          with WallTimeEvent(BenchEvent.STEP):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
              tok_tensor = llama.model(next_tok, start_pos, args.temperature)
@@ -1,6 +1,8 @@
 from pathlib import Path
 from typing import List
 import json, argparse, random, time, os
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
 from extra.models.llama import Transformer, convert_from_huggingface, convert_from_gguf, fix_bf16
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters, gguf_load
 from tinygrad import Tensor, dtypes, nn, Context, Device, GlobalCounters
@@ -10,8 +12,6 @@ from extra.bench_log import BenchEvent, WallTimeEvent
 class Tokenizer:
  pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
  def __init__(self, model_path: str):
-    import tiktoken
-    from tiktoken.load import load_tiktoken_bpe
    mergeable_ranks = load_tiktoken_bpe(model_path)
    self.num_base_tokens = len(mergeable_ranks)
    special_tokens = [
@@ -145,41 +145,6 @@ def NF4Linear(block_size):
      return new_state_dict
  return _NF4Linear

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

@@ -278,7 +242,7 @@ if __name__ == "__main__":
  parser.add_argument("--model", type=Path, help="Model path")
  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
  parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4", "float16", "fp8"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
  parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
  parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
  parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -477,7 +441,7 @@ if __name__ == "__main__":
      with Profiling(enabled=args.profile):
        with Timing("total ", on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
          with WallTimeEvent(BenchEvent.STEP):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None):
              tok = model(Tensor([[last_tok]], device=device), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P)
@@ -515,7 +479,7 @@ if __name__ == "__main__":
        st = GlobalCounters.time_sum_s
        with Profiling(enabled=args.profile):
          with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "")+
                        f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB"+
                        (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):

@@ -279,15 +279,9 @@ def generate(model, tokenizer, prompt: str, n_tokens_to_gen: int = 10, temp: boo
  # Loading in the prompt tokens
  logits = model.forward(Tensor([tks]))[:, -1, :]
  for _ in tqdm(range(n_tokens_to_gen), desc="Speed Gen"):
+    # TODO: topk
    if sample:
-      scaled_logits = logits / temp
-      if top_k is not None:
-        topk_values, topk_indices = scaled_logits.topk(top_k)
-        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
-        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
-        tok_Tens = filtered_logits.softmax().multinomial()
-      else:
-        tok_Tens = scaled_logits.softmax().multinomial()
+      tok_Tens = (logits/temp).softmax().multinomial()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
@@ -304,7 +298,6 @@ if __name__ == "__main__":
  parser.add_argument("--size", type=str, default="370m",
                      help=f"Size of model to use [{', '.join([k for k in MODELS.keys()])}]")
  parser.add_argument("--n_tokens", type=int, default=10, help="Number of tokens to generate")
-  parser.add_argument("--top_k", type=int, help="Limit sampling to the top k most likely tokens")
  parser.add_argument("--sample", dest="sample", action="store_true", help="Sample flag")
  parser.add_argument("--temp", type=float, default=1.0, help="Sampling temp has to be <=1.0")
  args = parser.parse_args()
@@ -315,9 +308,8 @@ if __name__ == "__main__":
  num_toks = args.n_tokens
  sample = args.sample
  temp = args.temp
-  top_k = args.top_k
  s = time.time()
-  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp)
  print(tinyoutput)
  print('TIME: ', time.time() - s)
  TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
@@ -511,33 +511,6 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
      # happens with BENCHMARK set
      pass

-# stable diffusion callbacks to match mlperf ref; declared here because they're pickled
-def filter_dataset(sample:dict): return {k:v for k,v in sample.items() if k in {'npy', 'txt'}}
-def collate(batch:list[dict]):
-  ret = {"npy": [], "txt": [], "__key__": []}
-  for sample in batch:
-    for k,v in sample.items():
-      ret[k].append(v)
-  return ret
-def collate_fn(batch): return batch
-
-# Reference (code): https://github.com/mlcommons/training/blob/2f4a93fb4888180755a8ef55f4b977ef8f60a89e/stable_diffusion/ldm/data/webdatasets.py, Line 55
-# Reference (params): https://github.com/mlcommons/training/blob/ab4ae1ca718d7fe62c369710a316dff18768d04b/stable_diffusion/configs/train_01x08x08.yaml, Line 107
-def batch_load_train_stable_diffusion(urls:str, BS:int):
-  import webdataset
-  dataset = webdataset.WebDataset(urls=urls, resampled=True, cache_size=-1, cache_dir=None)
-  dataset = dataset.shuffle(size=1000)
-  dataset = dataset.decode()
-  dataset = dataset.map(filter_dataset)
-  dataset = dataset.batched(BS, partial=False, collation_fn=collate)
-  dataset = webdataset.WebLoader(dataset, batch_size=None, shuffle=False, num_workers=1, persistent_workers=True, collate_fn=collate_fn)
-
-  for x in dataset:
-    assert isinstance(x, dict) and all(isinstance(k, str) for k in x.keys()) and all(isinstance(v, list) for v in x.values())
-    assert all(isinstance(moment_mean_logvar, np.ndarray) and moment_mean_logvar.shape==(1,8,64,64) for moment_mean_logvar in x["npy"])
-    assert all(isinstance(caption, str) for caption in x["txt"])
-    yield x
-
 # llama3

 class BinIdxDataset:
@@ -785,27 +758,6 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
      batch.append(tokens)
    yield Tensor.stack(batch, dim=0)

-def batch_load_llama3_small(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
-  if val:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-validation-91205-samples.en_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, False)
-  else:
-    dataset = BlendedGPTDataset([
-      base_dir / "c4-train.en_6_text_document",
-    ], [
-      1.0
-    ], samples, seqlen, seed, True)
-
-  for b in range(math.ceil(samples / bs)):
-    batch = []
-    for i in range(bs):
-      tokens = dataset.get(b * bs + i)
-      batch.append(tokens)
-    yield Tensor.stack(batch, dim=0)
-
 if __name__ == "__main__":
  def load_unet3d(val):
    assert not val, "validation set is not supported due to different sizes on inputs"
@@ -2,9 +2,7 @@ import math
 from typing import Union

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix, Context
-from tinygrad.nn.state import get_parameters
-from extra.models.unet import UNetModel
+from tinygrad.helpers import prod, argfix

 # rejection sampling truncated randn
 def rand_truncn(*shape, dtype=None, truncstds=2, **kwargs) -> Tensor:
@@ -19,10 +17,6 @@ def he_normal(*shape, a: float = 0.00, **kwargs) -> Tensor:
  std = math.sqrt(2.0 / (1 + a ** 2)) / math.sqrt(prod(argfix(*shape)[1:])) / 0.87962566103423978
  return std * rand_truncn(*shape, **kwargs)

-# Stable Diffusion v2 training uses default torch gelu, which doesn't use tanh approximation
-def gelu_erf(x:Tensor) -> Tensor:
-  return 0.5 * x * (1.0 + (x / 1.4142135623730951).erf())
-
 class Conv2dHeNormal(nn.Conv2d):
  def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
    super().__init__(in_channels, out_channels, kernel_size, stride=stride, padding=padding, dilation=dilation, groups=groups, bias=bias)
@@ -133,59 +127,3 @@ class Conv2dRetinaNet(nn.Conv2d):
  def __call__(self, x:Tensor) -> Tensor:
    return x.conv2d(self.weight.cast(dtypes.default_float), self.bias.cast(dtypes.default_float) if self.bias is not None else None,
                    groups=self.groups, stride=self.stride, dilation=self.dilation, padding=self.padding)
-
-# copy torch AMP: isolate mixed precision to just the below autocast ops, instead of using dtypes.default_float which affects all new Tensors
-class AutocastLinear(nn.Linear):
-  cast_dtype=dtypes.bfloat16 # enable monkeypatching of the mixed precision dtype
-  def __call__(self, x:Tensor) -> Tensor:
-    dtype = type(self).cast_dtype
-    return x.cast(dtype).linear(self.weight.cast(dtype).transpose(), self.bias.cast(dtype) if self.bias is not None else None)
-
-class AutocastConv2d(nn.Conv2d):
-  cast_dtype=dtypes.bfloat16
-  def __call__(self, x:Tensor) -> Tensor:
-    dtype = type(self).cast_dtype
-    return x.cast(dtype).conv2d(self.weight.cast(dtype), self.bias.cast(dtype), self.groups, self.stride, self.dilation, self.padding)
-
-# copy torch AMP: upcast to float32 before GroupNorm and LayerNorm
-class AutocastGroupNorm(nn.GroupNorm):
-  def __call__(self, x:Tensor) -> Tensor:
-    return super().__call__(x.cast(dtypes.float32))
-
-class AutocastLayerNorm(nn.LayerNorm):
-  def __call__(self, x:Tensor) -> Tensor:
-    return super().__call__(x.cast(dtypes.float32))
-
-def zero_module(module):
-  for p in get_parameters(module): p.assign(Tensor.zeros_like(p).contiguous())
-
-# Stable Diffusion mlperf reference doesn't call scaled_dot_product_attention
-# copy torch AMP: upcast to float32 before softmax on CUDA
-def attn_f32_softmax(q:Tensor, k:Tensor, v:Tensor) -> Tensor:
-  return (q.matmul(k.transpose(-2,-1), dtype=dtypes.float32) / math.sqrt(q.shape[-1])).softmax(-1).cast(q.dtype) @ v
-
-def init_stable_diffusion(version:str, pretrained:str, devices:list[str]):
-  from examples.stable_diffusion import StableDiffusion
-  from tinygrad.nn.state import safe_load, safe_save, load_state_dict, get_state_dict
-  from tempfile import TemporaryDirectory
-  model = StableDiffusion(version=version, pretrained=pretrained)
-  unet:UNetModel = model.model.diffusion_model
-
-  # this prevents extra consumption of memory, enabling much larger BS
-  Tensor.realize(*get_parameters(unet))
-  with TemporaryDirectory(prefix="unet_init") as tmp:
-    safe_save(get_state_dict(unet), init_fn:=f"{tmp}/init_model.safetensors")
-    load_state_dict(unet, safe_load(init_fn))
-
-  sqrt_alphas_cumprod = model.alphas_cumprod.sqrt().realize()
-  sqrt_one_minus_alphas_cumprod = (1 - model.alphas_cumprod).sqrt().realize()
-
-  if len(devices) > 1:
-    to_move = [sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod]
-    if version == "v2-mlperf-train": to_move += get_parameters(unet) + get_parameters(model.cond_stage_model)
-    for p in to_move:
-      p.to_(devices)
-    with Context(BEAM=0):
-      Tensor.realize(*to_move)
-
-  return model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod
@@ -1,9 +1,8 @@
 import math
-from tinygrad import dtypes, Tensor
+from tinygrad import dtypes
 from tinygrad.nn.optim import Optimizer

 from extra.lr_scheduler import LR_Scheduler
-from typing import Callable

 # https://github.com/mlcommons/training/blob/e237206991d10449d9675d95606459a3cb6c21ad/image_classification/tensorflow2/lars_util.py
 class PolynomialDecayWithWarmup(LR_Scheduler):
@@ -37,24 +36,4 @@ class CosineAnnealingLRWithWarmup(LR_Scheduler):
  def get_lr(self):
    warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
    decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
-    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
-
-# Reference: https://github.com/mlcommons/training/blob/64b14a9abc74e08779a175abca7d291f8c957632/stable_diffusion/ldm/lr_scheduler.py, Lines 36-97
-class LambdaLinearScheduler:
-  def __init__(self, warm_up_steps:int, f_min:float, f_max:float, f_start:float, cycle_lengths:int):
-    self.lr_warm_up_steps, self.f_min, self.f_max, self.f_start, self.cycle_lengths = warm_up_steps, f_min, f_max, f_start, cycle_lengths
-
-  def schedule(self, n:Tensor) -> Tensor:
-    warm_up = (n < self.lr_warm_up_steps)
-    f_warm_up = (self.f_max - self.f_start) / self.lr_warm_up_steps * n + self.f_start
-    return warm_up.where(f_warm_up, self.f_min + (self.f_max - self.f_min) * (self.cycle_lengths - n) / (self.cycle_lengths))
-
-# based on torch.optim.lr_scheduler.LambdaLR
-class LambdaLR(LR_Scheduler):
-  def __init__(self, optimizer:Optimizer, base_lr:Tensor, lr_lambda:Callable):
-    super().__init__(optimizer)
-    self.base_lr, self.lr_lambda = base_lr, lr_lambda
-    self.step()
-
-  def get_lr(self):
-    return self.base_lr * self.lr_lambda(self.epoch_counter - 1)
+    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)
@@ -1,10 +1,10 @@
-import time, math, os
+import time, math
 start = time.perf_counter()
 from pathlib import Path
 import numpy as np
 from tinygrad import Tensor, Device, dtypes, GlobalCounters, TinyJit
 from tinygrad.nn.state import get_parameters, load_state_dict, safe_load
-from tinygrad.helpers import getenv, Context, prod
+from tinygrad.helpers import getenv
 from extra.bench_log import BenchEvent, WallTimeEvent
 def tlog(x): print(f"{x:25s}  @ {time.perf_counter()-start:5.2f}s")

@@ -243,299 +243,31 @@ def eval_mrcnn():

 def eval_llama3():
  from extra.models.llama import Transformer
-  from examples.llama3 import MODEL_PARAMS, load, convert_from_huggingface
+  from examples.llama3 import MODEL_PARAMS
  from tinygrad.helpers import tqdm

-  BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
-  BS = getenv("BS", 4)
-  SMALL = getenv("SMALL", 0)
-  SEQLEN = getenv("SEQLEN", 8192)
-  MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
+  bs = 4
+  sequence_length = 512

-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
-  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
-  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
-
-  # load weights
-  weights = load(str(MODEL_PATH / "model.safetensors.index.json"))
-  if "model.embed_tokens.weight" in weights:
-    print("converting from huggingface format")
-    weights = convert_from_huggingface(weights, params["n_layers"], params["n_heads"], params["n_kv_heads"])
-
-  load_state_dict(model, weights, strict=False, consume=True)
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)

  @TinyJit
  def eval_step(model, tokens):
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
-    return loss.flatten().float()
+    return loss.flatten()

-  if SMALL:
-    from examples.mlperf.dataloader import batch_load_llama3_small
-    iter = batch_load_llama3_small(BS, 5760, SEQLEN, BASEDIR, val=True)
-  else:
-    from examples.mlperf.dataloader import batch_load_llama3
-    iter = batch_load_llama3(BS, 5760, SEQLEN, BASEDIR, val=True)
+  from examples.mlperf.dataloader import batch_load_llama3
+  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)

  losses = []
-  for tokens in tqdm(iter, total=5760//BS):
+  for tokens in tqdm(iter, total=5760//bs):
    GlobalCounters.reset()
    losses += eval_step(model, tokens).tolist()
    tqdm.write(f"loss: {np.mean(losses)}")

-  log_perplexity = np.mean(losses)
-  print(f"Log Perplexity: {log_perplexity}")
-
-# NOTE: BEAM hangs on 8xmi300x with DECODE_BS=384 in final realize below; function is declared here for external testing
-@TinyJit
-def vae_decode(x:Tensor, vae, disable_beam=False) -> Tensor:
-  from examples.stable_diffusion import AutoencoderKL
-  assert isinstance(vae, AutoencoderKL)
-  x = vae.post_quant_conv(1./0.18215 * x)
-
-  x = vae.decoder.conv_in(x)
-  x = vae.decoder.mid(x)
-  for i, l in enumerate(vae.decoder.up[::-1]):
-    print("decode", x.shape)
-    for b in l['block']: x = b(x)
-    if 'upsample' in l:
-      bs,c,py,px = x.shape
-      x = x.reshape(bs, c, py, 1, px, 1).expand(bs, c, py, 2, px, 2).reshape(bs, c, py*2, px*2)
-      x = l['upsample']['conv'](x)
-    if i == len(vae.decoder.up) - 1 and disable_beam:
-      with Context(BEAM=0): x.realize()
-    else: x.realize()
-  x = vae.decoder.conv_out(vae.decoder.norm_out(x).swish())
-
-  x = ((x + 1.0) / 2.0).clip(0.0, 1.0)
-  return x
-
-def eval_stable_diffusion():
-  import csv, PIL, sys
-  from tqdm import tqdm
-  from examples.mlperf.initializers import init_stable_diffusion, gelu_erf
-  from examples.stable_diffusion import AutoencoderKL
-  from extra.models.unet import UNetModel
-  from tinygrad.nn.state import load_state_dict, torch_load
-  from tinygrad.helpers import BEAM
-  from extra.models import clip
-  from extra.models.clip import FrozenOpenClipEmbedder
-  from extra.models.clip import OpenClipEncoder
-  from extra.models.inception import FidInceptionV3
-
-  config = {}
-  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
-  for x in GPUS: Device[x]
-  print(f"running eval on {GPUS}")
-  seed               = config["seed"]                   = getenv("SEED", 12345)
-  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
-  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
-  CONTEXT_BS         = config["CONTEXT_BS"]             = getenv("CONTEXT_BS", 1 * len(GPUS))
-  DENOISE_BS         = config["DENOISE_BS"]             = getenv("DENOISE_BS", 1 * len(GPUS))
-  DECODE_BS          = config["DECODE_BS"]              = getenv("DECODE_BS", 1 * len(GPUS))
-  INCEPTION_BS       = config["INCEPTION_BS"]           = getenv("INCEPTION_BS", 1 * len(GPUS))
-  CLIP_BS            = config["CLIP_BS"]                = getenv("CLIP_BS", 1 * len(GPUS))
-  EVAL_CKPT_DIR      = config["EVAL_CKPT_DIR"]          = getenv("EVAL_CKPT_DIR", "")
-  STOP_IF_CONVERGED  = config["STOP_IF_CONVERGED"]      = getenv("STOP_IF_CONVERGED", 0)
-
-  if (WANDB := getenv("WANDB", "")):
-    import wandb
-    wandb.init(config=config, project="MLPerf-Stable-Diffusion")
-
-  assert EVAL_CKPT_DIR != "", "provide a directory with checkpoints to be evaluated"
-  print(f"running eval on checkpoints in {EVAL_CKPT_DIR}\nSEED={seed}")
-  eval_queue:list[tuple[int, Path]] = []
-  for p in Path(EVAL_CKPT_DIR).iterdir():
-    if p.name.endswith(".safetensors"):
-      ckpt_iteration = p.name.split(".safetensors")[0]
-      assert ckpt_iteration.isdigit(), f"invalid checkpoint name: {p.name}, expected <digits>.safetensors"
-      eval_queue.append((int(ckpt_iteration), p))
-  assert len(eval_queue), f'no files ending with ".safetensors" were found in {EVAL_CKPT_DIR}'
-  print(sorted(eval_queue, reverse=True))
-
-  Tensor.manual_seed(seed)  # seed for weight initialization
-  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-eval", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
-
-  # load prompts for generating images for validation; 2 MB of data total
-  with open(DATADIR / "coco2014" / "val2014_30k.tsv") as f:
-    reader = csv.DictReader(f, delimiter="\t")
-    eval_inputs:list[dict] = [{"image_id": int(row["image_id"]), "id": int(row["id"]), "caption": row["caption"]} for row in reader]
-  assert len(eval_inputs) == 30_000
-  # NOTE: the clip weights are the same between model.cond_stage_model and clip_encoder
-  eval_timesteps = list(reversed(range(1, 1000, 20)))
-
-  original_device, Device.DEFAULT = Device.DEFAULT, "CPU"
-  # The choice of alphas_prev[0] = alphas_cumprod[0] seems arbitrary, but it's how the mlperf ref does it:
-  #   alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
-  eval_alphas_prev = model.alphas_cumprod[0:1].cat(model.alphas_cumprod[list(range(1, 1000, 20))[:-1]]).to(GPUS).realize()
-  inception = FidInceptionV3().load_from_pretrained(CKPTDIR / "inception" / "pt_inception-2015-12-05-6726825d.pth")
-  vision_cfg = {'width': 1280, 'layers': 32, 'd_head': 80, 'image_size': 224, 'patch_size': 14}
-  text_cfg = {'width': 1024, 'n_heads': 16, 'layers': 24, 'vocab_size': 49408, 'ctx_length': 77}
-  clip.gelu = gelu_erf
-  clip_encoder = OpenClipEncoder(1024, text_cfg, vision_cfg)
-  loaded = torch_load(CKPTDIR / "clip" / "open_clip_pytorch_model.bin")
-  loaded.update({"attn_mask": clip_encoder.attn_mask, "mean": clip_encoder.mean, "std": clip_encoder.std})
-  load_state_dict(clip_encoder, loaded)
-  Device.DEFAULT=original_device
-
-  @TinyJit
-  def denoise_step(x:Tensor, x_x:Tensor, t_t:Tensor, uc_c:Tensor, sqrt_alphas_cumprod_t:Tensor, sqrt_one_minus_alphas_cumprod_t:Tensor,
-                    alpha_prev:Tensor, unet:UNetModel, GPUS) -> Tensor:
-    out_uncond, out = unet(x_x, t_t, uc_c).to("CPU").reshape(-1, 2, 4, 64, 64).chunk(2, dim=1)
-    out_uncond = out_uncond.squeeze(1).shard(GPUS,axis=0)
-    out = out.squeeze(1).shard(GPUS,axis=0)
-    v_t = out_uncond + 8.0 * (out - out_uncond)
-    e_t = sqrt_alphas_cumprod_t * v_t + sqrt_one_minus_alphas_cumprod_t * x
-    pred_x0 = sqrt_alphas_cumprod_t * x - sqrt_one_minus_alphas_cumprod_t * v_t
-    dir_xt = (1. - alpha_prev).sqrt() * e_t
-    x_prev = alpha_prev.sqrt() * pred_x0 + dir_xt
-    return x_prev.realize()
-
-  def shard_tensor(t:Tensor) -> Tensor: return t.shard(GPUS, axis=0) if len(GPUS) > 1 else t.to(GPUS[0])
-  def get_batch(whole:Tensor, i:int, bs:int) -> tuple[Tensor, int]:
-    batch = whole[i: i + bs].to("CPU")
-    if (unpadded_bs:=batch.shape[0]) < bs:
-      batch = batch.cat(batch[-1:].expand(bs - unpadded_bs, *batch[-1].shape))
-    return batch, unpadded_bs 
-
-  @Tensor.train(mode=False)
-  def eval_unet(eval_inputs:list[dict], unet:UNetModel, cond_stage:FrozenOpenClipEmbedder, first_stage:AutoencoderKL,
-                inception:FidInceptionV3, clip:OpenClipEncoder) -> tuple[float, float]:
-    # Eval is divided into 5 jits, one per model
-    # It doesn't make sense to merge these jits, e.g. unet repeats 50 times in isolation; images fork to separate inception/clip
-    # We're generating and scoring 30,000 images per eval, and all the data can flow through one jit at a time
-    # To maximize throughput for each jit, we have only one model/jit on the GPU at a time, and pool outputs from each jit off-GPU
-    for model in (unet, first_stage, inception, clip):
-      Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
-
-    uc_written = False
-    models = (cond_stage, unet, first_stage, inception, clip)
-    jits = (jit_context:=TinyJit(cond_stage.embed_tokens), denoise_step, vae_decode, jit_inception:=TinyJit(inception),
-            jit_clip:=TinyJit(clip.get_clip_score))
-    all_bs = (CONTEXT_BS, DENOISE_BS, DECODE_BS, INCEPTION_BS, CLIP_BS)
-    if (EVAL_SAMPLES:=getenv("EVAL_SAMPLES", 0)) and EVAL_SAMPLES > 0:
-      eval_inputs = eval_inputs[0:EVAL_SAMPLES]
-    output_shapes = [(ns:=len(eval_inputs),77), (ns,77,1024), (ns,4,64,64), (ns,3,512,512), (ns,2048), (ns,)]
-    # Writing progress to disk lets us resume eval if we crash
-    stages = ["tokens", "embeds", "latents", "imgs", "inception", "clip"]
-    disk_tensor_names, disk_tensor_shapes = stages + ["end", "uc"], output_shapes + [(6,), (1,77,1024)]
-    if not all(os.path.exists(f"{EVAL_CKPT_DIR}/{name}.bytes") for name in disk_tensor_names):
-      for name, shape in zip(disk_tensor_names, disk_tensor_shapes):
-        file = Path(f"{EVAL_CKPT_DIR}/{name}.bytes")
-        file.unlink(missing_ok=True)
-        with file.open("wb") as f: f.truncate(prod(shape) * 4)
-    progress = {name: Tensor.empty(*shape, device=f"disk:{EVAL_CKPT_DIR}/{name}.bytes", dtype=dtypes.int if name in {"tokens", "end"} else dtypes.float)
-                for name, shape in zip(disk_tensor_names, disk_tensor_shapes)}
-
-    def embed_tokens(tokens:Tensor) -> Tensor:
-      nonlocal uc_written
-      if not uc_written:
-        with Context(BEAM=0): progress["uc"].assign(cond_stage.embed_tokens(cond_stage.tokenize("").to(GPUS)).to("CPU").realize()).realize()
-        uc_written = True
-      return jit_context(shard_tensor(tokens))
-
-    def generate_latents(embeds:Tensor) -> Tensor:
-      uc_c = Tensor.stack(progress["uc"].to("CPU").expand(bs, 77, 1024), embeds, dim=1).reshape(-1, 77, 1024)
-      uc_c = shard_tensor(uc_c)
-      x = shard_tensor(Tensor.randn(bs,4,64,64))
-      for step_idx, timestep in enumerate(tqdm(eval_timesteps)):
-        reversed_idx = Tensor([50 - step_idx - 1], device=GPUS)
-        alpha_prev = eval_alphas_prev[reversed_idx]
-        ts = Tensor.full(bs, fill_value=timestep, dtype=dtypes.int, device="CPU")
-        ts_ts = shard_tensor(ts.cat(ts))
-        ts = shard_tensor(ts)
-        sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
-        sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[ts].reshape(bs, 1, 1, 1)
-        x_x = shard_tensor(Tensor.stack(x.to("CPU"), x.to("CPU"), dim=1).reshape(-1, 4, 64, 64))
-        x.assign(denoise_step(x, x_x, ts_ts, uc_c, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t, alpha_prev, unet, GPUS)).realize()
-      return x
-
-    def decode_latents(latents:Tensor) -> Tensor: return vae_decode(shard_tensor(latents), first_stage, disable_beam=True)
-    def generate_inception(imgs:Tensor) -> Tensor: return jit_inception(shard_tensor(imgs))[:,:,0,0]
-
-    def calc_clip_scores(batch:Tensor, batch_tokens:Tensor) -> Tensor:
-      # Tensor.interpolate does not yet support bicubic, so we use PIL
-      batch = (batch.to(GPUS[0]).permute(0,2,3,1) * 255).clip(0, 255).cast(dtypes.uint8).numpy()
-      batch = [np.array(PIL.Image.fromarray(batch[i]).resize((224,224), PIL.Image.BICUBIC)) for i in range(bs)]
-      batch = shard_tensor(Tensor(np.stack(batch, axis=0).transpose(0,3,1,2), device="CPU").realize())
-      batch = batch.cast(dtypes.float) / 255
-      batch = (batch - model.mean) / model.std
-      batch = jit_clip(shard_tensor(batch_tokens), batch)
-      return batch
-
-    callbacks = (embed_tokens, generate_latents, decode_latents, generate_inception, calc_clip_scores)
-
-    # save every forward pass output to disk; NOTE: this needs ~100 GB disk space because 30k images are large
-    def stage_progress(stage_idx:int) -> int: return progress["end"].to("CPU")[stage_idx].item()
-    if stage_progress(0) < len(eval_inputs):
-      tokens = []
-      for i in tqdm(range(0, len(eval_inputs), CONTEXT_BS)):
-        subset = [cond_stage.tokenize(row["caption"], device="CPU") for row in eval_inputs[i: i+CONTEXT_BS]]
-        tokens.append(Tensor.cat(*subset, dim=0).realize())
-      progress["tokens"].assign(Tensor.cat(*tokens, dim=0).realize()).realize()
-      progress["end"][0:1].assign(Tensor([len(eval_inputs)], dtype=dtypes.int)).realize()
-    prev_stage = "tokens"
-    tokens = progress["tokens"]
-
-    # wrapper code for every model
-    for stage_idx, model, jit, bs, callback in zip(range(1,6), models, jits, all_bs, callbacks):
-      stage = stages[stage_idx]
-      if stage_progress(stage_idx) >= len(eval_inputs):
-        prev_stage = stage
-        continue # use cache
-      t0 = time.perf_counter()
-      print(f"starting eval with model: {model}")
-      if stage_idx == 1: inputs = tokens
-      elif stage_idx == 5: inputs = progress["imgs"]
-      else: inputs = progress[prev_stage]
-
-      Tensor.realize(*[p.to_(GPUS) for p in get_parameters(model)])
-      for batch_idx in tqdm(range(stage_progress(stage_idx), inputs.shape[0], bs)):
-        t1 = time.perf_counter()
-        batch, unpadded_bs = get_batch(inputs, batch_idx, bs)
-        if isinstance(model, OpenClipEncoder): batch = callback(batch, get_batch(tokens, batch_idx, bs)[0].realize())
-        else: batch = callback(batch)
-        # to(GPUS[0]) is necessary for this to work, without that the result is still on GPUS, probably due to a bug
-        batch = batch.to(GPUS[0]).to("CPU")[0:unpadded_bs].realize()
-        progress[stage][batch_idx: batch_idx + bs].assign(batch).realize()
-        # keep track of what our last output was, so we can resume from there if we crash in this loop
-        progress["end"][stage_idx: stage_idx + 1].assign(Tensor([batch_idx + bs], dtype=dtypes.int)).realize()
-        print(f"model: {model}, batch_idx: {batch_idx}, elapsed: {(time.perf_counter() - t1):.2f}")
-      del batch
-        
-      jit.reset()
-      Tensor.realize(*[p.to_("CPU") for p in get_parameters(model)])
-      print(f"done with model: {model}, elapsed: {(time.perf_counter() - t0):.2f}")
-      prev_stage = stage
-
-    inception_stats_fn = str(DATADIR / "coco2014" / "val2014_30k_stats.npz")
-    fid_score = inception.compute_score(progress["inception"].to("CPU"), inception_stats_fn)
-    clip_score = progress["clip"].to(GPUS[0]).mean().item()
-    for name in disk_tensor_names:
-      Path(f"{EVAL_CKPT_DIR}/{name}.bytes").unlink(missing_ok=True)
-    
-    if EVAL_SAMPLES and BEAM:
-      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      sys.exit() # Don't eval additional models; we don't care about clip/fid scores when running BEAM on eval sample subset
-
-    return clip_score, fid_score
-
-  # evaluate checkpoints in reverse chronological order
-  for ckpt_iteration, p in sorted(eval_queue, reverse=True):
-    unet_ckpt = safe_load(p)
-    load_state_dict(unet, unet_ckpt)
-    clip_score, fid_score = eval_unet(eval_inputs, unet, model.cond_stage_model, model.first_stage_model, inception, clip_encoder)
-    converged = True if clip_score >= 0.15 and fid_score <= 90 else False
-    print(f"eval results for {EVAL_CKPT_DIR}/{p.name}: clip={clip_score}, fid={fid_score}, converged={converged}")
-    if WANDB:
-      wandb.log({"eval/ckpt_iteration": ckpt_iteration, "eval/clip_score": clip_score, "eval/fid_score": fid_score})
-    if converged and STOP_IF_CONVERGED:
-      print(f"Convergence detected, exiting early before evaluating other checkpoints due to STOP_IF_CONVERGED={STOP_IF_CONVERGED}")
-      sys.exit()
-
-  # for testing
-  return clip_score, fid_score, ckpt_iteration
+  log_perplexity = Tensor(losses).mean()
+  print(f"Log Perplexity: {log_perplexity.item()}")

 if __name__ == "__main__":
  # inference only
@@ -3,8 +3,8 @@ from pathlib import Path
 import multiprocessing

 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
-from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
-from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
+from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

 from extra.lr_scheduler import LRSchedulerGroup
@@ -252,10 +252,6 @@ def train_resnet():
        print(f"epoch global_ops: {steps_in_train_epoch * GlobalCounters.global_ops:_}, "
              f"epoch global_mem: {steps_in_train_epoch * GlobalCounters.global_mem:_}")
        # if we are doing beam search, run the first eval too
-        if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-          min_time = min(step_times)
-          assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
        if (TRAIN_BEAM or EVAL_BEAM) and e == start_epoch: break
        return
    if MLLOGGER and RUNMLPERF:
@@ -348,8 +344,6 @@ def train_resnet():
        print(f"saving ckpt to {fn}")
        safe_save(get_training_state(model, optimizer_group, scheduler_group), fn)

-
-
 def train_retinanet():
  from contextlib import redirect_stdout
  from examples.mlperf.dataloader import batch_load_retinanet
@@ -707,7 +701,7 @@ def train_unet3d():
  ```BASEDIR=<folder_path> ./examples/mlperf/scripts/setup_kits19_dataset.sh```

  2) To start training the model, run the following:
-  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
+  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
  """
  from examples.mlperf.losses import dice_ce_loss
  from examples.mlperf.metrics import dice_score
@@ -749,6 +743,7 @@ def train_unet3d():
    "train_beam": TRAIN_BEAM,
    "eval_beam": EVAL_BEAM,
    "wino": WINO.value,
+    "fuse_conv_bw": FUSE_CONV_BW.value,
    "gpus": GPUS,
    "default_float": dtypes.default_float.name
  }
@@ -1188,9 +1183,7 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
-        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
-        train_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1224,7 +1217,7 @@ def train_bert():
          return

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

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -1297,20 +1290,18 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

  config = {}
-  BASEDIR            = config["BASEDIR"]                = Path(getenv("BASEDIR", "/raid/datasets/c4/"))
  BS                 = config["BS"]                     = getenv("BS", 16)
  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
  SEED               = config["SEED"]                   = getenv("SEED", 5760)
  SEQLEN             = config["SEQLEN"]                 = getenv("SEQLEN", 8192)
  TRAIN_ON_VAL       = config["TRAIN_ON_VAL"]           = getenv("TRAIN_ON_VAL", 0)
-  SMALL              = config["SMALL"]                  = getenv("SMALL", 0)
  SAMPLES            = config["SAMPLES"]                = getenv("SAMPLES", 5_760 if TRAIN_ON_VAL else 1_200_000 * 1152)
  EVAL_FREQ          = config["EVAL_FREQ"]              = getenv("EVAL_FREQ", 46080)
  EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 16)
  EVAL_TARGET        = config["EVAL_TARGET"]            = getenv("EVAL_TARGET", 5.6)

-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
+  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 FUSE_ARANGE=1 JITBEAM=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B WARMUP_STEPS=36 DECAY_STEPS=360 SEQLEN=512 PYTHONPATH=. AMD=1 AMD_LLVM=0 MODEL=llama3 python3 examples/mlperf/model_train.py
  # trains to 7

  opt_adamw_beta_1 = 0.9
@@ -1320,14 +1311,13 @@ def train_llama3():

  opt_gradient_clip_norm = 1.0
  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
-  opt_learning_rate_decay_steps = getenv("MAX_STEPS", math.ceil(1_200_000 * 1152 / GBS)) - opt_learning_rate_warmup_steps
+  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
-  opt_end_learning_rate = getenv("END_LR", 8e-7)
+  opt_end_learning_rate = 8e-7

  # TODO: confirm weights are in bf16
  # vocab_size from the mixtral tokenizer
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
-  params = params | {"vocab_size": 32000} if not SMALL else params
+  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)

@@ -1363,15 +1353,6 @@ def train_llama3():
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)

-  if resume_ckpt := getenv("RESUME_CKPT"):
-    fn = f"./ckpts/llama3_{resume_ckpt}.safe"
-    print(f"loading initial checkpoint from {fn}")
-    load_state_dict(model, safe_load(fn), realize=False)
-
-    fn = f"./ckpts/llama3_{resume_ckpt}_optim.safe"
-    print(f"loading optim checkpoint from {fn}")
-    load_state_dict(scheduler, safe_load(fn), realize=False)
-
  @TinyJit
  @Tensor.train()
  def train_step(model, tokens:Tensor, grad_acc:int):
@@ -1422,55 +1403,43 @@ def train_llama3():
  # ** data iters **
  def fake_data(bs, samples):
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=params["vocab_size"], dtype=dtypes.int32, device=Device.DEFAULT)
+      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)

  def get_train_iter():
    if getenv("FAKEDATA", 0):
      return fake_data(GBS, SAMPLES)
    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(GBS, SAMPLES, SEQLEN, BASEDIR, seed=SEED, val=bool(TRAIN_ON_VAL))
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))

  def get_eval_iter():
    if getenv("FAKEDATA", 0):
      return fake_data(EVAL_BS, 5760)
    else:
-      if SMALL:
-        from examples.mlperf.dataloader import batch_load_llama3_small
-        return batch_load_llama3_small(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
-      else:
-        from examples.mlperf.dataloader import batch_load_llama3
-        return batch_load_llama3(EVAL_BS, 5760, SEQLEN, BASEDIR, val=True)
+      from examples.mlperf.dataloader import batch_load_llama3
+      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)

  iter = get_train_iter()
-  i, sequences_seen = resume_ckpt, 0
+  i, sequences_seen = 0, 0
  for tokens in tqdm(iter, total=SAMPLES//GBS):
    t = time.perf_counter()
    GlobalCounters.reset()
    loss, lr = train_step(model, tokens, grad_acc)
    loss = loss.float().item()
-
-    i += 1
-    sequences_seen += tokens.shape[0]
-
+    # above as tqdm.write f-string
    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
    if (fname:=getenv("LOSS_FILE", "")):
      with open(fname, "a") as f:
        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")

-    if (ckpt_freq := getenv("CKPT")) and (i % ckpt_freq == 0 and (i != 1 or ckpt_freq == 1)):
+    if getenv("CKPT") and (i % 200 == 0 or i == 10):
      tqdm.write("saving checkpoint")
      if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
      fn = f"{ckpt_dir}/llama3_{i}.safe"
      safe_save(get_state_dict(model), fn)

-      tqdm.write("saving optim checkpoint")
-      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
-      safe_save(get_state_dict(scheduler), fn)
+    i += 1
+    sequences_seen += tokens.shape[0]

    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
      tqdm.write(f"evaluating after {sequences_seen} sequences")
@@ -1494,144 +1463,6 @@ def train_llama3():
          safe_save(get_state_dict(model), fn)
        break

-def train_stable_diffusion():
-  from extra.models.unet import UNetModel
-  from examples.mlperf.dataloader import batch_load_train_stable_diffusion
-  from examples.mlperf.lr_schedulers import LambdaLR, LambdaLinearScheduler
-  from examples.mlperf.initializers import init_stable_diffusion
-  from examples.mlperf.helpers import get_training_state
-  import numpy as np
-
-  config = {}
-  GPUS               = config["GPUS"]                   = [f"{Device.DEFAULT}:{i}" for i in range(getenv("GPUS", 1))]
-  seed               = config["seed"]                   = getenv("SEED", 12345)
-  # ** hyperparameters **
-  BS                 = config["BS"]                     = getenv("BS", 1 * len(GPUS))
-  BASE_LR            = config["LEARNING_RATE"]          = getenv("LEARNING_RATE", 2.5e-7)
-  # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
-  # "Checkpoint must be collected every 512,000 images. CEIL(512000 / global_batch_size) if 512000 is not divisible by GBS."
-  # NOTE: It's inferred that "steps" is the unit for the output of the CEIL formula, based on all other cases of CEIL in the rules
-  CKPT_STEP_INTERVAL = config["CKPT_STEP_INTERVAL"]     = getenv("CKPT_STEP_INTERVAL", math.ceil(512_000 / BS))
-  CKPTDIR            = config["CKPTDIR"]                = Path(getenv("CKPTDIR", "./checkpoints"))
-  DATADIR            = config["DATADIR"]                = Path(getenv("DATADIR", "./datasets"))
-  UNET_CKPTDIR       = config["UNET_CKPTDIR"]           = Path(getenv("UNET_CKPTDIR", "./checkpoints"))
-  TOTAL_CKPTS        = config["TOTAL_CKPTS"]            = getenv("TOTAL_CKPTS", 0)
-
-  print(f"training on {GPUS}")
-  lr = BS * BASE_LR
-  print(f"BS={BS}, BASE_LR={BASE_LR}, lr={lr}")
-  print(f"CKPT_STEP_INTERVAL = {CKPT_STEP_INTERVAL}")
-  for x in GPUS: Device[x]
-  if (WANDB := getenv("WANDB", "")):
-    import wandb
-    wandb.init(config=config, project="MLPerf-Stable-Diffusion")
-
-  Tensor.manual_seed(seed)  # seed for weight initialization
-  model, unet, sqrt_alphas_cumprod, sqrt_one_minus_alphas_cumprod = init_stable_diffusion("v2-mlperf-train", CKPTDIR / "sd" / "512-base-ema.ckpt", GPUS)
-
-  optimizer = AdamW(get_parameters(unet))
-  lambda_lr_callback = LambdaLinearScheduler(1000, 1.0, 1.0, 1e-06, 10000000000000).schedule
-  lr_scheduler = LambdaLR(optimizer, Tensor(lr, dtype=dtypes.float, device=optimizer.device), lambda_lr_callback)
-
-  @TinyJit
-  def train_step(mean:Tensor, logvar:Tensor, tokens:Tensor, unet:UNetModel, optimizer:LAMB, lr_scheduler:LambdaLR) -> Tensor:
-    optimizer.zero_grad()
-
-    timestep = Tensor.randint(BS, low=0, high=model.alphas_cumprod.shape[0], dtype=dtypes.int, device=GPUS[0])
-    latent_randn = Tensor.randn(*mean.shape, device=GPUS[0])
-    noise = Tensor.randn(*mean.shape, device=GPUS[0])
-    for t in (mean, logvar, tokens, timestep, latent_randn, noise):
-      t.shard_(GPUS, axis=0)
-
-    std = Tensor.exp(0.5 * logvar.clamp(-30.0, 20.0))
-    latent = (mean + std * latent_randn) * 0.18215
-
-    sqrt_alphas_cumprod_t = sqrt_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
-    sqrt_one_minus_alphas_cumprod_t = sqrt_one_minus_alphas_cumprod[timestep].reshape(timestep.shape[0], 1, 1, 1)
-    latent_with_noise = sqrt_alphas_cumprod_t * latent + sqrt_one_minus_alphas_cumprod_t * noise
-    v_true = sqrt_alphas_cumprod_t * noise - sqrt_one_minus_alphas_cumprod_t * latent
-
-    context = model.cond_stage_model.embed_tokens(tokens)
-
-    out = unet(latent_with_noise, timestep, context)
-    loss = ((out - v_true) ** 2).mean()
-    del mean, logvar, std, latent, noise, sqrt_alphas_cumprod_t, sqrt_one_minus_alphas_cumprod_t
-    del out, v_true, context, latent_randn, tokens, timestep
-    loss.backward()
-
-    optimizer.step()
-    lr_scheduler.step()
-    loss, out_lr = loss.detach().to("CPU"), optimizer.lr.to("CPU")
-    Tensor.realize(loss, out_lr)
-    return loss, out_lr
-    
-  # checkpointing takes ~9 minutes without this, and ~1 minute with this
-  @TinyJit
-  def ckpt_to_cpu():
-    ckpt = get_training_state(unet, optimizer, lr_scheduler)
-    # move to CPU first so more GPU bufs aren't created (can trigger OOM)
-    for k,v in ckpt.items(): ckpt[k] = v.detach().to("CPU")
-    Tensor.realize(*[v for v in ckpt.values()])
-    for k,v in ckpt.items(): ckpt[k] = v.cast(v.dtype.base).contiguous()
-    Tensor.realize(*[v for v in ckpt.values()])
-    return ckpt
-
-  # training loop
-  dl = batch_load_train_stable_diffusion(f'{DATADIR}/laion-400m/webdataset-moments-filtered/{{00000..00831}}.tar', BS)
-  # for tests
-  saved_checkpoints = []
-
-  train_start_time = time.perf_counter()
-  t0 = t6 = time.perf_counter()
-  for i, batch in enumerate(dl, start=1):
-    loop_time = time.perf_counter() - t0
-    t0 = time.perf_counter()
-    dl_time = t0 - t6
-    GlobalCounters.reset()
-
-    mean, logvar = np.split(np.concatenate(batch["npy"], axis=0), 2, axis=1)
-    mean, logvar = Tensor(mean, dtype=dtypes.float32, device="CPU"), Tensor(logvar, dtype=dtypes.float32, device="CPU")
-    tokens = []
-    for text in batch['txt']: tokens += model.cond_stage_model.tokenizer.encode(text, pad_with_zeros=True)
-    tokens = Tensor(tokens, dtype=dtypes.int32, device="CPU").reshape(-1, 77)
-
-    t1 = time.perf_counter()
-    loss, lr = train_step(mean, logvar, tokens, unet, optimizer, lr_scheduler)
-    loss_item, lr_item = loss.item(), lr.item()
-    t2 = time.perf_counter()
-
-    if i == 3:
-      for _ in range(3): ckpt_to_cpu() # do this at the beginning of run to prevent OOM surprises when checkpointing
-      print("BEAM COMPLETE", flush=True) # allows wrapper script to detect BEAM search completion and retry if it failed
-      
-    total_train_time = time.perf_counter() - train_start_time
-    if WANDB:
-      wandb.log({"train/loss": loss_item, "train/lr": lr_item, "train/loop_time_prev": loop_time, "train/dl_time": dl_time, "train/step": i,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (t2-t1), "train/input_prep_time": t1-t0,
-                    "train/train_step_time": t2-t1, "train/total_time": total_train_time})
-
-      if i == 1 and wandb.run is not None:
-        with open(f"{UNET_CKPTDIR}/wandb_run_id_{wandb.run.id}", "w") as f:
-          f.write(f"wandb.run.id = {wandb.run.id}")
-
-    if i % CKPT_STEP_INTERVAL == 0:
-      # https://github.com/mlcommons/training_policies/blob/cfa99da479b8d5931f7a3c67612d021dfb47510a/training_rules.adoc#benchmark_specific_rules
-      # "evaluation is done offline, the time is not counted towards the submission time."
-      fn = f"{UNET_CKPTDIR}/{i}.safetensors"
-      print(f"saving unet checkpoint at {fn}")
-      saved_checkpoints.append(fn)
-      safe_save({k.replace("model.", ""):v for k,v in ckpt_to_cpu().items() if k.startswith("model.")}, fn)
-      if TOTAL_CKPTS and i == TOTAL_CKPTS * CKPT_STEP_INTERVAL:
-        print(f"ending run after {i} steps ({TOTAL_CKPTS} checkpoints collected)")
-        return saved_checkpoints
-
-    t3 = time.perf_counter()
-    print(f"""step {i}: {GlobalCounters.global_ops * 1e-9 / (t2-t1):9.2f} GFLOPS, mem_used: {GlobalCounters.mem_used / 1e9:.2f} GB,
-  loop_time_prev: {loop_time:.2f}, dl_time: {dl_time:.2f}, input_prep_time: {t1-t0:.2f}, train_step_time: {t2-t1:.2f},
-  t3-t2: {t3-t2:.4f}, loss:{loss_item:.5f}, lr:{lr_item:.3e}, total_train_time:{total_train_time:.2f}
-  """)
-    t6 = time.perf_counter()
-
 if __name__ == "__main__":
  multiprocessing.set_start_method('spawn')

@@ -1640,7 +1471,7 @@ if __name__ == "__main__":
  else: bench_log_manager = contextlib.nullcontext()

  with Tensor.train():
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
@@ -1,57 +0,0 @@
-#!/usr/bin/env bash
-# adapted from https://github.com/mlcommons/training/blob/4bdf5c8ed218ad76565a2ba1ac27c919ccc6d689/stable_diffusion/README.md
-
-# setup dirs
-
-DATA=/raid/datasets/stable_diffusion
-
-LAION=$DATA/laion-400m/webdataset-moments-filtered 
-COCO=$DATA/coco2014
-mkdir -p $LAION $COCO
-
-CKPT=/raid/weights/stable_diffusion
-mkdir -p $CKPT/clip $CKPT/sd $CKPT/inception
-
-# download data
-
-# if rclone isn't installed system-wide / in your PATH, put the executable path in quotes below
-#RCLONE=""
-RCLONE="rclone"
-
-## VAE-encoded image latents, from 6.1M image subset of laion-400m
-## about 1 TB for whole download
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/ ${LAION} --include="*.tar" -P
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/laion-400m/moments-webdataset-filtered/sha512sums.txt ${LAION} -P
-cd $LAION && grep -E '\.tar$' sha512sums.txt | sha512sum -c --quiet - && \
-  echo "All .tar files verified" || { echo "Checksum failure when validating downloaded Laion moments"; exit 1; }
-
-## prompts and FID statistics from 30k image subset of coco2014
-## 33 MB
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k.tsv ${COCO} -P
-
-$RCLONE config create mlc-training s3 provider=Cloudflare access_key_id=76ea42eadb867e854061a1806220ee1e secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
-$RCLONE copy mlc-training:mlcommons-training-wg-public/stable_diffusion/datasets/coco2014/val2014_30k_stats.npz ${COCO} -P
-
-# download checkpoints
-
-## clip (needed for text and vision encoders for validation)
-CLIP_WEIGHTS_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/resolve/main/open_clip_pytorch_model.bin"
-CLIP_WEIGHTS_SHA256="9a78ef8e8c73fd0df621682e7a8e8eb36c6916cb3c16b291a082ecd52ab79cc4"
-CLIP_CONFIG_URL="https://huggingface.co/laion/CLIP-ViT-H-14-laion2B-s32B-b79K/raw/main/open_clip_config.json"
-wget -N -P ${CKPT}/clip ${CLIP_WEIGHTS_URL}
-wget -N -P ${CKPT}/clip ${CLIP_CONFIG_URL}
-echo "${CLIP_WEIGHTS_SHA256}  ${CKPT}/clip/open_clip_pytorch_model.bin"                    | sha256sum -c
-
-## sd (needed for latent->image decoder for validation, also has clip text encoder for training)
-SD_WEIGHTS_URL='https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt'
-SD_WEIGHTS_SHA256="d635794c1fedfdfa261e065370bea59c651fc9bfa65dc6d67ad29e11869a1824"
-wget -N -P ${CKPT}/sd ${SD_WEIGHTS_URL}
-echo "${SD_WEIGHTS_SHA256}  ${CKPT}/sd/512-base-ema.ckpt"                    | sha256sum -c
-
-## inception (needed for validation)
-FID_WEIGHTS_URL='https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'
-FID_WEIGHTS_SHA1="bd836944fd6db519dfd8d924aa457f5b3c8357ff"
-wget -N -P ${CKPT}/inception ${FID_WEIGHTS_URL}
-echo "${FID_WEIGHTS_SHA1}  ${CKPT}/inception/pt_inception-2015-12-05-6726825d.pth"                    | sha1sum -c
@@ -1,72 +0,0 @@
-#!/usr/bin/env bash
-
-DATETIME=${2:-$(date "+%m%d%H%M")}
-LOGFILE="${HOME}/logs/sd_mi300x_${DATETIME}.log"
-# UNET_CKPTDIR must be set: training saves checkpoints to this path, then a separate eval process scans this path to know which checkpoints to eval
-export UNET_CKPTDIR="${HOME}/stable_diffusion/training_checkpoints/${DATETIME}"
-mkdir -p "${HOME}/logs" "$UNET_CKPTDIR"
-
-# run this script in isolation when using the --bg flag
-if [[ "${1:-}" == "--bg" ]]; then
-  echo "logging output to $LOGFILE"
-  echo "saving UNet checkpoints to $UNET_CKPTDIR"
-  script_path="$(readlink -f "${BASH_SOURCE[0]}")"
-  nohup bash "$script_path" run "$DATETIME" >"$LOGFILE" 2>&1 & disown $!
-  exit 0
-fi
-
-# venv management
-if [[ -d .venv-sd-mlperf ]]; then
-  . .venv-sd-mlperf/bin/activate
-else
-  python3 -m venv .venv-sd-mlperf && . .venv-sd-mlperf/bin/activate
-  pip install --index-url https://download.pytorch.org/whl/cpu torch && pip install tqdm numpy ftfy regex pillow scipy wandb webdataset
-fi
-pip list
-apt list --installed | grep amdgpu
-rocm-smi --version
-modinfo amdgpu | grep version
-
-export BEAM=2 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 IGNORE_JIT_FIRST_BEAM=1 HCQDEV_WAIT_TIMEOUT_MS=300000
-export AMD_LLVM=0 # bf16 seems to require this
-export DATADIR="/raid/datasets/stable_diffusion"
-export CKPTDIR="/raid/weights/stable_diffusion"
-export EVAL_CKPT_DIR=$UNET_CKPTDIR
-export MODEL="stable_diffusion" PYTHONPATH="."
-export GPUS=8 BS=304
-export CONTEXT_BS=816 DENOISE_BS=600 DECODE_BS=384 INCEPTION_BS=560 CLIP_BS=240
-export WANDB=1
-export PARALLEL=4
-export PYTHONUNBUFFERED=1
-sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setperfdeterminism 1500 || exit 1
-
-# Retry BEAM search if script fails before BEAM COMPLETE is printed, but don't retry after that
-run_retry(){ local try=0 max=5 code tmp py pgid kids
-  while :; do
-    tmp=$(mktemp)
-    setsid bash -c 'exec env "$@"' _ "$@" > >(tee -a "$LOGFILE" | tee "$tmp") 2>&1 &
-    py=$!; pgid=$(ps -o pgid= -p "$py" | tr -d ' ')
-    wait "$py"; code=$?
-    [[ -n "$pgid" ]] && { kill -TERM -"$pgid" 2>/dev/null; sleep 1; kill -KILL -"$pgid" 2>/dev/null; }
-    kids=$(pgrep -P "$py" || true)
-    while [[ -n "$kids" ]]; do
-      kill -TERM $kids 2>/dev/null; sleep 0.5
-      kids=$(for k in $kids; do pgrep -P "$k" || true; done)
-    done
-    grep -q 'BEAM COMPLETE' "$tmp" && { rm -f "$tmp"; return 1; }
-    rm -f "$tmp"
-    ((code==0)) && return 0
-    ((try>=max)) && return 2
-    ((try++)); sleep 90; echo "try = ${try}"
-  done
-}
-
-# Power limiting to 400W is only needed if GPUs fall out of sync (causing 2.2x increased train time) at higher power, which has been observed at 450W
-sudo rocm-smi -d 0 1 2 3 4 5 6 7 --setpoweroverdrive 750 && \
-run_retry TOTAL_CKPTS=7 python3 examples/mlperf/model_train.py; (( $? == 2 )) && { echo "training failed before BEAM completion"; exit 2; }
-sleep 90
-
-run_retry EVAL_SAMPLES=600 python3 examples/mlperf/model_eval.py; (( $? == 2 )) && { echo "eval failed before BEAM completion"; exit 2; }
-# Checkpoints will be evaluated in reverse chronological order, even if above training crashed early
-# STOP_IF_CONVERGED=1: Stop the eval after the first time convergence is detected; no more checkpoints will be evaluated after that.
-STOP_IF_CONVERGED=1 python3 examples/mlperf/model_eval.py
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
-
-export IGNORE_OOB=1
-
-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-# export BEAM_LOG_SURPASS_MAX=1
-# export BASEDIR="/raid/datasets/wiki"
-
-export RESET_STEP=1
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
-export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
-export BASEDIR="/raid/datasets/wiki"
-
-export BENCHMARK=10 BERT_LAYERS=2
-
-python3 examples/mlperf/model_train.py
@@ -1,20 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
-
-# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
-export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
-export TRAIN_STEPS=3900
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
-export BASEDIR="/raid/datasets/wiki"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export SUBMISSION_PLATFORM="tinybox_8xMI300X"
-export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
-
-# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
-export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
-export TRAIN_STEPS=3900
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
-
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,17 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
-export BASEDIR="/raid/datasets/wiki"
-
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,28 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="bert"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_green_${DATETIME}_${SEED}.log"
-
-# init
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,69 +0,0 @@
-# 1. Problem
-
-This problem uses BERT for NLP.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-Also install gdown (for dataset), numpy, tqdm and tensorflow.
-```
-pip install gdown numpy tqdm tensorflow
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download and verify data
-
-### 1. Download raw data
-
-```
-BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
-```
-
-### 2. Preprocess train and validation data
-
-Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
-
-#### Training:
-```
-BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
-```
-
-Generating a specific topic (Between 0 and 499)
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
-```
-
-#### Validation:
-```
-BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
-```
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
-```
-
-### tinybox_red
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
-```
-### tinybox_8xMI300X
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
-```
@@ -1,18 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
-export BASEDIR="/raid/datasets/wiki"
-
-export RESET_STEP=1
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,16 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-export WANDB=1 PARALLEL=0
-
-RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -1,31 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." AMD=1
-export MODEL="bert"
-export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export IGNORE_OOB=1
-export REWRITE_STACK_LIMIT=500000
-
-export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/wiki"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="bert_red_${DATETIME}_${SEED}.log"
-
-export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
-
-# init
-sleep 5 && sudo rmmod amdgpu || true
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
-# 1. Problem
-
-This problem uses the ResNet-50 CNN to do image classification.
-
-## Requirements
-
-Install tinygrad and mlperf-logging from master.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-### tinybox_red
-Disable cwsr
-This is the default on production tinybox red.
-```
-sudo vi /etc/modprobe.d/amdgpu.conf
-cat <<EOF > /etc/modprobe.d/amdgpu.conf
-options amdgpu cwsr_enable=0
-EOF
-sudo update-initramfs -u
-sudo reboot
-
-# validate
-sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
-```
-
-# 2. Directions
-
-## Steps to download and verify data
-
-```
-IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
-```
-
-## Steps for one time setup
-
-### tinybox_red
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
-```
-
-## Steps to run benchmark
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-```
@@ -1,13 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-export BENCHMARK=10 DEBUG=2
-
-python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-export EVAL_START_EPOCH=3 EVAL_FREQ=4
-
-export WANDB=1 PARALLEL=0
-
-python3 examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="resnet"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
-
-# init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,50 +0,0 @@
-# 1. Problem
-
-This problem uses the ResNet-50 CNN to do image classification.
-
-## Requirements
-
-Install tinygrad and mlperf-logging from master.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-### tinybox_red
-Disable cwsr
-This is the default on production tinybox red.
-```
-sudo vi /etc/modprobe.d/amdgpu.conf
-cat <<EOF > /etc/modprobe.d/amdgpu.conf
-options amdgpu cwsr_enable=0
-EOF
-sudo update-initramfs -u
-sudo reboot
-
-# validate
-sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
-```
-
-# 2. Directions
-
-## Steps to download and verify data
-
-```
-IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
-```
-
-## Steps for one time setup
-
-### tinybox_red
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
-```
-
-## Steps to run benchmark
-```
-examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
-```
@@ -1,13 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=10 DEBUG=${DEBUG:-2}
-
-python3 examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export EVAL_START_EPOCH=3 EVAL_FREQ=4
-
-export WANDB=1 PARALLEL=0
-
-python3 examples/mlperf/model_train.py
@@ -1,26 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." AMD=1
-export MODEL="resnet"
-export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
-
-export RESET_STEP=0
-
-export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=${LOGMLPERF:-1}
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
-
-# init
-sleep 5 && sudo rmmod amdgpu || true
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,8 +0,0 @@
-#!/bin/bash
-
-rocm-smi --setprofile compute
-rocm-smi --setmclk 3
-rocm-smi --setperflevel high
-
-# power cap to 350W
-echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -1,38 +0,0 @@
-# 1. Problem
-
-This problem uses RetinaNet for SSD.
-
-## Requirements
-
-Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
-```
-git clone https://github.com/tinygrad/tinygrad.git
-python3 -m pip install -e ".[mlperf]"
-```
-
-Also install the following dependencies:
-```
-pip install tqdm numpy pycocotools boto3 pandas torch torchvision
-```
-
-### tinybox_green
-Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
-This is the default on production tinybox green.
-
-# 2. Directions
-
-## Steps to download data
-
-Run the following:
-```
-BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
-```
-
-## Running
-
-### tinybox_green
-
-#### Steps to run benchmark
-```
-examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
-```
@@ -1,14 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=5 DEBUG=2
-
-python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export WANDB=1 PARALLEL=0
-export RUNMLPERF=1
-
-python examples/mlperf/model_train.py
@@ -1,25 +0,0 @@
-#!/bin/bash
-set -e  # Exit on any error
-set -o pipefail  # Make pipeline fail if any command fails
-
-export PYTHONPATH="." NV=1
-export MODEL="retinanet"
-export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-
-export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-export IGNORE_JIT_FIRST_BEAM=1
-export BASEDIR="/raid/datasets/openimages"
-
-# pip install -e ".[mlperf]"
-export LOGMLPERF=1
-
-export SEED=$RANDOM
-DATETIME=$(date "+%m%d%H%M")
-LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
-
-# init
-BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
-
-# run
-PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -1,14 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export BENCHMARK=5 DEBUG=2
-
-python examples/mlperf/model_train.py
@@ -1,15 +0,0 @@
-#!/bin/bash
-
-export PYTHONPATH="." AMD=1
-export MODEL="retinanet"
-export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
-export BASEDIR="/raid/datasets/openimages"
-
-# export RESET_STEP=0
-
-export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
-
-export WANDB=1 PARALLEL=0
-export RUNMLPERF=1
-
-python examples/mlperf/model_train.py
@@ -1,38 +0,0 @@
-{
-    "submitter": "tinycorp",
-    "division": "closed",
-    "status": "Available on-premise",
-    "system_name": "tinybox 8xMI300X",
-    "number_of_nodes": "1",
-    "host_processors_per_node": "2",
-    "host_processor_model_name": "AMD EPYC 9354",
-    "host_processor_core_count": "32",
-    "host_processor_vcpu_count": "64",
-    "host_processor_frequency": "",
-    "host_processor_caches": "",
-    "host_processor_interconnect": "",
-    "host_memory_capacity": "2304GB",
-    "host_storage_type": "NVMe SSD",
-    "host_storage_capacity": "3x 4TB raid array",
-    "host_networking": "",
-    "host_networking_topology": "",
-    "host_memory_configuration": "24x 96GB DDR5",
-    "accelerators_per_node": "8",
-    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
-    "accelerator_host_interconnect": "PCIe 5.0 x16",
-    "accelerator_frequency": "",
-    "accelerator_on-chip_memories": "",
-    "accelerator_memory_configuration": "HBM3",
-    "accelerator_memory_capacity": "192GB",
-    "accelerator_interconnect": "",
-    "accelerator_interconnect_topology": "",
-    "cooling": "air",
-    "hw_notes": "",
-    "framework": "tinygrad, branch mlperf_training_v5.0",
-    "other_software_stack": {
-        "python": "3.10.16",
-        "ROCm": "3.0.0+94441cb"
-    },
-    "operating_system": "Ubuntu 24.04.1 LTS",
-    "sw_notes": ""
-  }
@@ -1,38 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox green",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6X",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12",
-    "CUDA": "12.4"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}
@@ -1,37 +0,0 @@
-{
-  "submitter": "tinycorp",
-  "division": "closed",
-  "status": "Available on-premise",
-  "system_name": "tinybox red",
-  "number_of_nodes": "1",
-  "host_processors_per_node": "1",
-  "host_processor_model_name": "AMD EPYC 7532",
-  "host_processor_core_count": "32",
-  "host_processor_vcpu_count": "64",
-  "host_processor_frequency": "",
-  "host_processor_caches": "",
-  "host_processor_interconnect": "",
-  "host_memory_capacity": "128GB",
-  "host_storage_type": "NVMe SSD",
-  "host_storage_capacity": "4 TB raid array + 1 TB boot",
-  "host_networking": "",
-  "host_networking_topology": "",
-  "host_memory_configuration": "8x 16GB DDR4",
-  "accelerators_per_node": "6",
-  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
-  "accelerator_host_interconnect": "PCIe 4.0 x16",
-  "accelerator_frequency": "",
-  "accelerator_on-chip_memories": "",
-  "accelerator_memory_configuration": "GDDR6",
-  "accelerator_memory_capacity": "24GB",
-  "accelerator_interconnect": "",
-  "accelerator_interconnect_topology": "",
-  "cooling": "air",
-  "hw_notes": "",
-  "framework": "tinygrad, branch mlperf_training_v5.0",
-  "other_software_stack": {
-    "python": "3.10.12"
-  },
-  "operating_system": "Ubuntu 22.04.4",
-  "sw_notes": ""
-}
@@ -1,11 +1,16 @@
-import os, sys, pickle, time, re
+import os, sys, pickle, time
 import numpy as np
+if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
+if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
+if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"

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

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

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

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

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

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

  # run 20 times
-  step_times = []
  for _ in range(20):
    st = time.perf_counter()
    out = run(**inputs)
    mt = time.perf_counter()
    val = out.numpy()
    et = time.perf_counter()
-    step_times.append((et-st)*1e3)
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
-
-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
+  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")

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

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

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

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

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

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

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

-  if getenv("BENCHMARK_LOG", ""):
-    bench(pickle_loaded, inputs)
@@ -0,0 +1,47 @@
+import sys
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp
+from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.schedule.kernelize import get_kernelize_map
+from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.realize import run_schedule
+
+# NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+
+OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
+OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
+
+if __name__ == "__main__":
+  onnx_file = fetch(OPENPILOT_MODEL)
+  run_onnx = OnnxRunner(onnx_file)
+
+  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
+  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
+  root = out.uop
+  targets = [x.uop for x in inputs.values()]
+  print(targets)
+
+  # TODO: abstract this from gradient?
+
+  # compute the target path (top down)
+  in_target_path: dict[UOp, bool] = {}
+  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
+  independent_set = {}
+  for u in root.toposort():
+    if in_target_path[u]:
+      for s in u.src:
+        if not in_target_path[s]:
+          independent_set[s] = None
+  independent = UOp.sink(*independent_set.keys())
+  kernelized = get_kernelize_map(independent)
+  independent = independent.substitute(kernelized)
+  schedule, var_vals = create_schedule_with_vars(independent)
+  run_schedule(schedule)
+
+  print("**** real ****")
+  GlobalCounters.reset()
+  out.uop = root.substitute(kernelized)
+  out.kernelize()
+
+  # realize
+  out.realize()
@@ -27,7 +27,7 @@ class Model(nn.Module):

 if __name__ == "__main__":
  if getenv("TINY_BACKEND"):
-    import tinygrad.nn.torch  # noqa: F401
+    import tinygrad.frontend.torch  # noqa: F401
    device = torch.device("tiny")
  else:
    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
@@ -8,7 +8,7 @@ from typing import Dict, Union

 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from examples.llama3 import load
-from tinygrad import nn, Tensor, Device
+from tinygrad import nn, Tensor
 from tinygrad.helpers import fetch, colored, GlobalCounters, Timing, DEBUG
 from tinygrad.nn.state import load_state_dict, get_parameters

@@ -80,7 +80,7 @@ if __name__ == "__main__":
    st = GlobalCounters.time_sum_s
    next_tok = Tensor([toks[start_pos:]]) if tok_tensor is None or (len(toks)-start_pos) > 1 else tok_tensor.reshape(1, 1)
    with Timing("total ", enabled=args.timing, on_exit=lambda x: f", {1e9/x:.2f} tok/s, {GlobalCounters.global_mem/x:.2f} GB/s, param {param_bytes/x:.2f} GB/s"):
-      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" if DEBUG>=2 else "") +
+      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on GPU" if DEBUG>=2 else "") +
                  f", {GlobalCounters.global_ops*1e-9:.2f} GOPS, {GlobalCounters.global_mem*1e-9:.2f} GB" +
                  (f", {GlobalCounters.global_mem*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s, param {param_bytes*1e-9/(GlobalCounters.time_sum_s-st):.2f} GB/s" if DEBUG>=2 else "")) if DEBUG else None, enabled=args.timing):
        tok_tensor = transformer(next_tok, start_pos, args.temperature)
@@ -99,7 +99,6 @@ if __name__ == "__main__":
  parser.add_argument('--timing',      action='store_true', help="Print timing per step")
  parser.add_argument('--noshow',      action='store_true', help="Don't show the image")
  parser.add_argument('--fp16',        action='store_true', help="Cast the weights to float16")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

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

  model = StableDiffusionV2(**params)

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

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

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

  c  = { "crossattn": model.cond_stage_model(args.prompt) }
  uc = { "crossattn": model.cond_stage_model("") }
@@ -6,7 +6,7 @@
 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import safe_load, load_state_dict
-from tinygrad.helpers import fetch, trange, colored, Timing, getenv
+from tinygrad.helpers import fetch, trange, colored, Timing
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -14,7 +14,7 @@ from examples.stable_diffusion import ResnetBlock, Mid
 import numpy as np

 from typing import Dict, List, Callable, Optional, Any, Set, Tuple, Union, Type
-import argparse, tempfile, time
+import argparse, tempfile
 from abc import ABC, abstractmethod
 from pathlib import Path
 from PIL import Image
@@ -342,13 +342,11 @@ class DPMPP2MSampler:
    sigmas = self.discretization(num_steps).to(x.device)
    x *= Tensor.sqrt(1.0 + sigmas[0] ** 2.0)
    num_sigmas = len(sigmas)
-    step_times = []

    old_denoised = None
    for i in trange(num_sigmas - 1):
      with Timing("step in ", enabled=timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
        GlobalCounters.reset()
-        st = time.perf_counter_ns()
        with WallTimeEvent(BenchEvent.STEP):
          x, old_denoised = self.sampler_step(
            old_denoised=old_denoised,
@@ -360,13 +358,8 @@ class DPMPP2MSampler:
            c=c,
            uc=uc,
          )
-          step_times.append(t:=(time.perf_counter_ns() - st)*1e-6)
          x.realize(old_denoised)

-    if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-      min_time = min(step_times)
-      assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-
    return x


@@ -437,8 +430,8 @@ if __name__ == "__main__":
    im.show()

  # validation!
-  is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024
-  if is_default and not args.weights and not args.fakeweights:
+  if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \
+    and not args.weights:
    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "sdxl_seed0.png")))
    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
    assert distance < 4e-3, colored(f"validation failed with {distance=}", "red")
@@ -2,39 +2,40 @@ import os, pathlib, argparse
 from examples.llama3 import Tokenizer
 from tabulate import tabulate
 from tinygrad import fetch
-from tinygrad.helpers import flatten, getenv
-from sz import NONCORE_DIRS
+from tinygrad.helpers import flatten

 # llama 3 tokenizer
 tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())

-def read_code(base_path, full=False):
+def read_code(base_path):
  ret = []
  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
-    if not full and any(path.split("./")[1].startswith(x) for x in NONCORE_DIRS): continue
    for name in files:
      if not name.endswith(".py"): continue
      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
      fullpath = os.path.join(path, name)
      code = pathlib.Path(fullpath).read_text()
-      ret.append((fullpath.split("tinygrad/", 1)[1], code))
+      ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
  return ret

+def write_code_to_file(filename, code_list):
+  """Writes the combined code to a specified file."""
+  with open(filename, 'w') as f:
+    f.write('\n'.join(flatten(code_list)))
+
 if __name__ == "__main__":
  parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
  parser.add_argument("--output", help="Output file to write the combined code to.")
-  parser.add_argument("--full", action="store_true", help="All directories")
  args = parser.parse_args()

-  ret = read_code(".", args.full)
+  ret = read_code(".")

  table = []
  for name,code in ret:
-    table.append([name, len(tokenizer.encode(code))])
+    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))

-  banner = "#"*40
-  code_str = ''.join([f"{banner}\n# {name}\n{banner}\n\n{code}\n" for name,code in ret])
+  code_str = '\x00'.join(flatten(ret))
  print(f"code has {len(code_str)} chars")
  newline_count = code_str.count('\n')
  print(f"code has {newline_count} newlines")
@@ -43,5 +44,5 @@ if __name__ == "__main__":
  print(f"code has {len(encoded)} tokens")

  if args.output:
-    with open(args.output, 'w') as f: f.write(code_str)
-    print(f"Combined code written to {args.output}")
+    write_code_to_file(args.output, ret)
+    print(f"Combined code written to {args.output}")
@@ -2,20 +2,18 @@
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
 import tempfile
 from pathlib import Path
-import argparse, time
+import argparse
 from collections import namedtuple
 from typing import Dict, Any

 from PIL import Image
 import numpy as np
 from tinygrad import Device, GlobalCounters, dtypes, Tensor, TinyJit
-from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten, profile_marker
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import torch_load, load_state_dict, get_state_dict
-from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
-from extra.models import unet, clip
+from extra.models.clip import Closed, Tokenizer
 from extra.models.unet import UNetModel
-from examples.mlperf.initializers import AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm, zero_module, attn_f32_softmax, gelu_erf
 from extra.bench_log import BenchEvent, WallTimeEvent

 class AttnBlock:
@@ -156,46 +154,12 @@ unet_params: Dict[str,Any] = {
  "use_linear": False,
 }

-mlperf_params: Dict[str,Any] = {"adm_in_ch": None, "in_ch": 4, "out_ch": 4, "model_ch": 320, "attention_resolutions": [4, 2, 1], "num_res_blocks": 2,
-                                "channel_mult": [1, 2, 4, 4], "d_head": 64, "transformer_depth": [1, 1, 1, 1], "ctx_dim": 1024, "use_linear": True,
-                                "num_groups":16, "st_norm_eps":1e-6}
-
 class StableDiffusion:
-  def __init__(self, version:str|None=None, pretrained:str|None=None):
+  def __init__(self):
    self.alphas_cumprod = get_alphas_cumprod()
-    if version != "v2-mlperf-train":
-      self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments
-
-    if not version:
-      self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
-      unet_init_params = unet_params
-    elif version in {"v2-mlperf-train", "v2-mlperf-eval"}:
-      unet_init_params = mlperf_params
-      clip.gelu = gelu_erf
-      self.cond_stage_model = FrozenOpenClipEmbedder(**{"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True,
-                                                        "clip_tokenizer_version": "sd_mlperf_v5_0"})
-      unet.Linear, unet.Conv2d, unet.GroupNorm, unet.LayerNorm = AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm
-      unet.attention, unet.gelu, unet.mixed_precision_dtype = attn_f32_softmax, gelu_erf, dtypes.bfloat16
-      if pretrained:
-        print("loading text encoder")
-        weights: dict[str,Tensor] = {k.replace("cond_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("cond_stage_model.")}
-        weights["model.attn_mask"] = Tensor.full((77, 77), fill_value=float("-inf")).triu(1)
-        load_state_dict(self.cond_stage_model, weights)
-        # only the eval model needs the decoder
-        if version == "v2-mlperf-eval":
-          print("loading image latent encoder")
-          weights = {k.replace("first_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("first_stage_model.")}
-          load_state_dict(self.first_stage_model, weights)
-
-    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_init_params))
-    if version == "v2-mlperf-train":
-      # the mlperf reference inits certain weights as zeroes
-      for bb in flatten(self.model.diffusion_model.input_blocks) + self.model.diffusion_model.middle_block + flatten(self.model.diffusion_model.output_blocks):
-        if isinstance(bb, unet.ResBlock):
-          zero_module(bb.out_layers[3])
-        elif isinstance(bb, unet.SpatialTransformer):
-          zero_module(bb.proj_out)
-      zero_module(self.model.diffusion_model.out[2])
+    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_params))
+    self.first_stage_model = AutoencoderKL()
+    self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))

  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
    temperature = 1
@@ -263,34 +227,25 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
-  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()

-  profile_marker("create model")
  model = StableDiffusion()

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

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

-    Tensor.realize(*get_state_dict(model).values())
-
-  profile_marker("run clip (conditional)")
+  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
  prompt = Tensor([tokenizer.encode(args.prompt)])
  context = model.cond_stage_model.transformer.text_model(prompt).realize()
  print("got CLIP context", context.shape)

-  profile_marker("run clip (unconditional)")
  prompt = Tensor([tokenizer.encode("")])
  unconditional_context = model.cond_stage_model.transformer.text_model(prompt).realize()
  print("got unconditional CLIP context", unconditional_context.shape)
@@ -311,39 +266,30 @@ if __name__ == "__main__":
  def run(model, *x): return model(*x).realize()

  # this is diffusion
-  step_times = []
  with Context(BEAM=getenv("LATEBEAM")):
    for index, timestep in (t:=tqdm(list(enumerate(timesteps))[::-1])):
-      profile_marker(f"step {len(timesteps)-index-1}")
      GlobalCounters.reset()
-      st = time.perf_counter_ns()
      t.set_description("%3d %3d" % (index, timestep))
      with Timing("step in ", enabled=args.timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
        with WallTimeEvent(BenchEvent.STEP):
          tid = Tensor([index])
          latent = run(model, unconditional_context, context, latent, Tensor([timestep]), alphas[tid], alphas_prev[tid], Tensor([args.guidance]))
          if args.timing: Device[Device.DEFAULT].synchronize()
-      step_times.append((time.perf_counter_ns() - st)*1e-6)
-    # done with diffusion model
    del run
-    del model.model

-  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
-    min_time = min(step_times)
-    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
-  profile_marker("run decoder") # upsample latent space to image with autoencoder
-  x = model.decode(latent).realize()
+  # upsample latent space to image with autoencoder
+  x = model.decode(latent)
  print(x.shape)

-  profile_marker("save image")
+  # save image
  im = Image.fromarray(x.numpy())
  print(f"saving {args.out}")
  im.save(args.out)
  # Open image.
  if not args.noshow: im.show()

+  # validation!
  if args.prompt == default_prompt and args.steps == 6 and args.seed == 0 and args.guidance == 7.5:
-    profile_marker("validate")
    ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "stable_diffusion_seed0.png")))
    distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
    assert distance < 3e-3, colored(f"validation failed with {distance=}", "red")  # higher distance with WINO
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

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

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -32,7 +32,7 @@ if __name__ == "__main__":

  lr = 5e-3
  transform = ComposeTransforms([
-    lambda x: [Image.fromarray(xx).resize((64, 64)) for xx in x],
+    lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
    lambda x: np.stack([np.asarray(xx) for xx in x], 0),
    lambda x: x / 255.0,
    lambda x: np.tile(np.expand_dims(x, 1), (1, 3, 1, 1)).astype(np.float32),
@@ -3,7 +3,7 @@
 import sys, base64, multiprocessing, itertools, collections
 from typing import Optional, Union, Literal, List

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

@@ -109,7 +109,7 @@ class TextDecoder:

  def forward(self, x:Tensor, pos:Union[Variable, Literal[0]], encoded_audio:Tensor):
    seqlen = x.shape[-1]
-    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
+    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None, None))
    for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
    return self.output_tok(x)

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

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

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

  def gettexttoks(line): return [tok for tok in line if tok < eot or tok > enc._special_tokens["<|notimestamps|>"]][-nsample+len(start_tokens):]
@@ -2,7 +2,7 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs

 os.chdir("/tmp")
@@ -37,7 +37,7 @@ def main():
  dev = PCIIface(None, 0)
  for x, y in dev.dev_impl.__dict__.items():
    if isinstance(y, AMRegister):
-      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+      for inst, addr in y.addr.keys(): reg_names[addr] = f"{x}, xcc={inst}"

  with open(sys.argv[1], 'r') as f:
    log_content = log_content_them = f.read()
@@ -1,7 +1,7 @@
 # copying the kernels from https://github.com/microsoft/ArchProbe into Python
 import numpy as np
 import pickle
-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
 from tinygrad import dtypes
 from tqdm import trange, tqdm
 from matplotlib import pyplot as plt
@@ -4,7 +4,7 @@ from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
 from tinygrad.codegen.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_cl import ROCM_LLVM_PATH
+from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH

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

-from tinygrad.runtime.ops_cl import CLProgram, CLBuffer, ROCM_LLVM_PATH
+from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH

 ENABLE_NON_ASM = False

@@ -10,13 +10,13 @@ from tinygrad.renderer.cstyle import ClangRenderer
 render_dtype = ClangRenderer().render_dtype

 class ClangGraph(GraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v}" for v in var_vals])
+    args += sorted([f"int {v.expr}" for v in var_vals])
    code = ["void batched("+','.join(args)+") {"]
    for ji in jit_cache:
      args = []
@@ -34,6 +34,6 @@ class ClangGraph(GraphRunner):
    assert compiler is not None
    self._prg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers

-  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[Variable, int], wait=False):
    return cpu_time_execution(
-    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0])]), enable=wait)
+    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), enable=wait)
@@ -26,7 +26,7 @@ class VirtAQLQueue(AQLQueue):
    self.available_packet_slots -= 1

 class HSAGraph(MultiGraphRunner):
-  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)

    # Check all jit items are compatible.
@@ -53,7 +53,7 @@ class HSAGraph(MultiGraphRunner):
      self.ji_kargs_structs[j] = ji.prg._prg.args_struct_t.from_address(kernargs_ptrs[ji.prg.dev])
      kernargs_ptrs[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
      for i in range(len(ji.bufs)): self.ji_kargs_structs[j].__setattr__(f'f{i}', cast(Buffer, ji.bufs[i])._buf)
-      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i]])

    # Build queues.
    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
@@ -106,7 +106,7 @@ class HSAGraph(MultiGraphRunner):
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 0)
    hsa.hsa_signal_silent_store_relaxed(self.finish_signal, 0)

-  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[Variable, int], wait=False) -> Optional[float]:
    # Wait and restore signals
    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
@@ -123,7 +123,7 @@ class HSAGraph(MultiGraphRunner):
    # Update var_vals
    for j in self.jc_idx_with_updatable_var_vals:
      for i,v in enumerate(cast(CompiledRunner, self.jit_cache[j].prg).p.vars):
-        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v])

    # Update launch dims
    for j in self.jc_idx_with_updatable_launch_dims:
@@ -29,10 +29,10 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
  r: Dict[UOp, str] = {}
  for u in uops:
    if u.uop == UOps.SPECIAL:
-      if u.arg.startswith("lidx"):
-        r[u] = f'v{u.src[0].arg}'
-      elif u.arg.startswith("gidx"):
-        r[u] = f's{2+u.src[0].arg}'
+      if u.arg[1].startswith("lidx"):
+        r[u] = f'v{u.arg[0]}'
+      elif u.arg[1].startswith("gidx"):
+        r[u] = f's{2+u.arg[0]}'
      else:
        raise NotImplementedError
    elif u.uop == UOps.CONST:
@@ -1,34 +0,0 @@
-#!/usr/bin/env python3
-from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
-from tinygrad.helpers import getenv, colored
-
-SZ = 8_000_000_000
-GPUS = getenv("GPUS", 4) # TODO: expose a way in tinygrad to access this
-
-if __name__ == "__main__":
-  # create tensors
-  tens = [Tensor.ones(SZ, dtype=dtypes.uint8, device=f"{Device.DEFAULT}:{i}").contiguous() for i in range(GPUS)]
-  Tensor.realize(*tens)
-
-  bw = [[0.0]*GPUS for _ in range(GPUS)]
-  for i in range(GPUS):
-    for j in range(GPUS):
-      GlobalCounters.reset()
-      with Context(DEBUG=2):
-        if i == j:
-          # this copy would be optimized out, just add 1
-          (tens[i]+1).realize()
-        else:
-          tens[i].to(f"{Device.DEFAULT}:{j}").realize()
-      t = max(GlobalCounters.time_sum_s, 1e-9)
-      bw[i][j] = SZ / t / 1e9  # GB/s
-
-  def fmt(x):
-    c = "green" if x > 50 else "yellow" if x > 20 else "red"
-    return colored(f"{x:6.1f}", c)
-
-  # header
-  print(" " * 8 + " ".join(f"{'d'+str(j):>6}" for j in range(GPUS)))
-  # rows
-  for i in range(GPUS):
-    print(f"{'s'+str(i):>6} -> " + " ".join(fmt(x) for x in bw[i]))
@@ -1,4 +0,0 @@
-# source extra/cl_android.sh
-export LD_LIBRARY_PATH=/data/data/com.termux/files/usr/lib:/system/vendor/lib64
-export LD_PRELOAD=/system/vendor/lib64/libOpenCL.so
-
@@ -10,7 +10,7 @@ from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict

-EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "CL"]
+EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CPU", "CUDA", "GPU"]

 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]]:
  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
@@ -67,12 +67,11 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
  forward_args = ",".join(f"{dtype}{'*' if name not in symbolic_vars.values() else ''} {name}" for name,dtype,_ in (outputs+inputs if wasm else inputs+outputs))

  if not wasm:
-    thread_id = 0 # NOTE: export does not support threading, thread_id is always 0
    for name,cl in bufs_to_save.items():
      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
      cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
    cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
-    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -240,9 +239,7 @@ export default {model_name};

 def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
-
-  # NOTE: CPU_COUNT=1, since export does not support threading
-  with Context(JIT=2, CPU_COUNT=1): run,special_names = jit_model(model, *inputs)
+  with Context(JIT=2): run,special_names = jit_model(model, *inputs)
  functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
  state = get_state_dict(model)
  weight_names = {id(x.uop.base.realized): name for name, x in state.items()}
@@ -1,169 +1,335 @@
-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.uop.ops import UOp, Ops, KernelInfo, graph_rewrite, AxisType, PatternMatcher, UPat
+from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
 from tinygrad.dtype import AddrSpace
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, colored, prod, unwrap
+from tinygrad.shape.shapetracker import ShapeTracker, View
+from tinygrad.shape.view import strides_for_shape
+from tinygrad.codegen.opt.kernel import axis_colors
+from tinygrad.codegen.opt.swizzler import merge_views, view_left

-N = getenv("N", 4096)
-M = K = N
-run_count = getenv("CNT", 5)
+def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])

-# ---------------------------
-# launch/config constants
-# ---------------------------
+N = 4096
+run_count = 5

-WARP_SIZE = 32
+BN = 128
+BM = 128
+BK = 8

-# 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
+TN = 4
+TM = 4

-# Register tile sizes (per-thread accumulator tile of C)
-TN = 4     # columns per thread
-TM = 4     # rows per thread
+# NOTE: this is from testgrad
+# change reduceop axes and input ShapeTrackers, view gets replaced with a reshape.
+# src->r->view  -->   src->view->r
+def swizzle_reduceop(src:UOp, r:UOp, view:UOp):
+  if r.tag is not None: return None
+  # confirm the input is in order
+  # TODO: replace this with a UOp that allows for nothing else then remove this
+  permute = tuple(i for i in range(len(src.shape)) if i not in r.axis_arg)+r.axis_arg
+  assert permute == tuple(range(len(permute))), f"reduce axis must already be in order, {permute} isn't"

-is_kernel5 = getenv("K5", 0)
-THREADS_PER_BLOCK = 128 if is_kernel5 else 256
-assert THREADS_PER_BLOCK % BLOCK_N == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_N"
-assert THREADS_PER_BLOCK % BLOCK_K == 0, "THREADS_PER_BLOCK must be divisible by BLOCK_K"
-assert (BLOCK_N * BLOCK_K) % THREADS_PER_BLOCK == 0
-assert (BLOCK_M * BLOCK_K) % THREADS_PER_BLOCK == 0
+  # append the reduce shape to each of the views
+  prshape = prod(rshape:=src.shape[-len(r.axis_arg):])
+  rstrides = strides_for_shape(rshape)
+  nv = [View.create(v.shape+rshape, tuple(x*prshape for x in v.strides)+rstrides, v.offset*prshape,
+                    v.mask+tuple((0,s) for s in rshape) if v.mask is not None else None) for v in unwrap(view.st).views]

-WARPS_PER_BLOCK = THREADS_PER_BLOCK // WARP_SIZE
-WAVE_TILE_N = 128 if is_kernel5 else 64
-WAVE_TILE_M = BLOCK_N * BLOCK_M // WARPS_PER_BLOCK // WAVE_TILE_N
-assert BLOCK_N % WAVE_TILE_N == 0, "BN must be a multiple of WN"
-assert BLOCK_M % WAVE_TILE_M == 0, "BM must be a multiple of WM"
-WAVES_IN_BLOCK_X = BLOCK_N // WAVE_TILE_N
-WAVES_IN_BLOCK_Y = BLOCK_M // WAVE_TILE_M
-assert WAVES_IN_BLOCK_X * WAVES_IN_BLOCK_Y == WARPS_PER_BLOCK, "wave grid must match warps/block"
+  # no reshape required with shrinking REDUCE_AXIS
+  return UOp(Ops.REDUCE_AXIS, r.dtype, (src.view(ShapeTracker(tuple(nv))),),
+             (r.arg[0], tuple(range(len(view.shape), len(view.shape) + len(r.axis_arg)))))

-LANES_PER_WAVE_X = 8
-LANES_PER_WAVE_Y = 4
-ITERS_PER_WAVE_N = WAVE_TILE_N // (LANES_PER_WAVE_X * TN)
-ITERS_PER_WAVE_M = WAVE_TILE_M // (LANES_PER_WAVE_Y * TM)
-assert WAVE_TILE_N % (LANES_PER_WAVE_X * TN) == 0, "WAVE_TILE_N must be divisible by LANES_PER_WAVE_X*TN"
-assert WAVE_TILE_M % (LANES_PER_WAVE_Y * TM) == 0, "WAVE_TILE_M must be divisible by LANES_PER_WAVE_Y*TM"
+pm = PatternMatcher([
+  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
+])

-def rngs_for_shape(shape:tuple[sint, ...], rng:int, axis_type=AxisType.LOOP): return [UOp.range(s, rng+i, axis_type) for i,s in enumerate(shape)]
-def copy(dest:UOp, src:UOp, rng:int, set=False, upcast=False):
-  assert dest.shape == src.shape
-  rngs = rngs_for_shape(src.shape, rng, AxisType.UPCAST if upcast else AxisType.LOOP)
-  copy = dest[*rngs].store(src[*rngs]).end(*rngs)
-  return dest.after(copy) if set else copy
+def top_spec_kernel3():
+  a = Tensor.empty(N,N)
+  b = Tensor.empty(N,N)
+  c = a@b
+  sink = c.schedule()[-1].ast
+  L = 16
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
+  sink = graph_rewrite(sink, view_left+pm)
+  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
+  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))

-def hand_spec_kernel3():
-  # ---------------------------
-  # block indices & placeholders
-  # ---------------------------
-  blockIdx_x = UOp.special(N // BLOCK_N, "gidx0")
-  blockIdx_y = UOp.special(N // BLOCK_M, "gidx1")
+def hl_spec_kernel3():
+  nbIterWaveM = 2
+  nbIterWaveN = 2

-  a = UOp.placeholder((N, N), dtypes.float, slot=1)
-  b = UOp.placeholder((N, N), dtypes.float, slot=2)
-  c = UOp.placeholder((N, N), dtypes.float, slot=0)
+  # define buffers
+  # TODO: remove these views once the defines have a shape
+  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1).view(ShapeTracker.from_shape((N,N)))
+  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2).view(ShapeTracker.from_shape((N,N))).permute((1,0))
+  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0).view(ShapeTracker.from_shape((N,N)))
+  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BM, AddrSpace.LOCAL), arg=0).view(ShapeTracker.from_shape((BK, BM))).permute((1,0))
+  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1).view(ShapeTracker.from_shape((BK, BN))).permute((1,0))
+  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), arg=0).view(ShapeTracker.from_shape((nbIterWaveM * TM,)))
+  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), arg=1).view(ShapeTracker.from_shape((nbIterWaveN * TN,)))

-  # index the output with the globals
-  c = c.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_N, BLOCK_N)[blockIdx_y, :, blockIdx_x, :]
+  # shape buffers. TODO: permutes
+  full_shape = (N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)
+  a = a.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, N//BK, BK)).expand(full_shape)
+  b = b.reshape((1, 1, 1, 1, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, N//BK, BK)).expand(full_shape)
+  c = c.reshape((N//BM, nbIterWaveM, BM//(nbIterWaveM * TM), TM, N//BN, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, 1))
+  As = As.reshape((1, nbIterWaveM, BM//(nbIterWaveM * TM), TM, 1, 1, 1, 1, 1, BK)).expand(full_shape)
+  Bs = Bs.reshape((1, 1, 1, 1, 1, nbIterWaveN, BN//(nbIterWaveN * TN), TN, 1, BK)).expand(full_shape)
+  A_col = A_col.reshape((1, nbIterWaveM, 1, TM, 1, 1, 1, 1, 1, 1)).expand(full_shape)
+  B_row = B_row.reshape((1, 1, 1, 1, 1, nbIterWaveN, 1, TN, 1, 1)).expand(full_shape)

-  # open the main reduction range
-  k_tile_range = UOp.range(N // BLOCK_K, 0, AxisType.REDUCE)
-  a = a.reshape(M // BLOCK_M, BLOCK_M, N // BLOCK_K, BLOCK_K)[blockIdx_y, :, k_tile_range, :]
-  b = b.reshape(N // BLOCK_K, BLOCK_K, N // BLOCK_N, BLOCK_N)[k_tile_range, :, blockIdx_x, :]
+  #                     U1   L2 L3 L4 L5   U6 U7      U9   L10 L11 L12 L13   U14 U15      U17  U18  U19
+  expanded_shape = (32, 2,   2, 2, 2, 2,   2, 2,  32, 2,   2,  2,  2,  2,    2,  2,  512, 2,   2,   2)
+  assert len(expanded_shape) == 20
+  permute_a = list(range(len(expanded_shape)))
+  permute_b = permute_a[:]

-  # globals are no longer used, they are already in the indexes
-  del blockIdx_y, blockIdx_x
+  # this makes all the global loads match
+  # this can also be more simply done by rebinding the RANGEs
+  # but sadly, rebinding the RANGEs doesn't work to change the order of the local axes
+  permute_a[17:20] = [11,12,13]
+  permute_a[11:14] = [17,18,19]
+  permute_a[7], permute_a[10] = permute_a[10], permute_a[7]
+  permute_a[2:7] = [3,4,5,6,2]

-  # ---------------------------
-  # GLOBAL -> LOCAL (As, Bs)
-  # ---------------------------
-  tid = UOp.special(THREADS_PER_BLOCK, "lidx0")
+  permute_b[2:16] = [19,9,10,11,17,18,8,2,12,13,14,15,3,4]
+  permute_b[17:20] = [5,6,7]

-  # A: read BM x BK tiles (permute on store into locals)
-  BM_As_stride = (BLOCK_M + 4) if is_kernel5 else BLOCK_M
-  As = UOp.placeholder((BLOCK_K, BM_As_stride), dtypes.float, slot=0, addrspace=AddrSpace.LOCAL).shrink_to((BLOCK_K, BLOCK_M))
-  As_store = copy(As.permute((1,0)).reshape(-1, THREADS_PER_BLOCK)[:, tid], a.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=100)
+  a_permute   = a.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)
+  As_permute = As.reshape(expanded_shape).permute(tuple(permute_a)).reshape(full_shape)

-  # B: read BK x BN tiles
-  Bs = UOp.placeholder((BLOCK_K, BLOCK_N), dtypes.float, slot=1, addrspace=AddrSpace.LOCAL)
-  Bs_store = copy(Bs.reshape(-1, THREADS_PER_BLOCK)[:, tid], b.reshape(-1, THREADS_PER_BLOCK)[:, tid], rng=200)
+  b_permute   = b.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)
+  Bs_permute = Bs.reshape(expanded_shape).permute(tuple(permute_b)).reshape(full_shape)

-  # TODO: can we automate barrier?
-  barrier = UOp.barrier(As_store, Bs_store)
-  As, Bs = As.after(barrier), Bs.after(barrier)
+  #out = (a.load() * b.load()).r(Ops.ADD, (8, 9))
+  out = (As.load(As_permute.store(a_permute.load())) * Bs.load(Bs_permute.store(b_permute.load()))).r(Ops.ADD, (8, 9))
+  #out = (A_col.load(A_col.store(As.load(As.store(a.load())))) * B_row.load(B_row.store(Bs.load(Bs.store(b.load()))))).r(Ops.ADD, (8, 9))

-  # open inner k range
-  k = UOp.range(BLOCK_K, 3, AxisType.REDUCE)
+  axis_types = (
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.GLOBAL, AxisType.UPCAST, AxisType.LOCAL, AxisType.UPCAST,
+    AxisType.REDUCE, AxisType.REDUCE)

-  # ---------------------------
-  # LOCAL -> REG (per-wave tiles)
-  # ---------------------------
-  waveIdx = (tid // WARP_SIZE) % WAVES_IN_BLOCK_X
-  waveIdy = (tid // WARP_SIZE) // WAVES_IN_BLOCK_X
-  assert waveIdy.vmax+1 == WAVES_IN_BLOCK_Y
+  sink = c.store(out).sink(arg=KernelInfo(name="tg_"+to_colored(full_shape, axis_types), axis_types=axis_types))
+  sink = graph_rewrite(sink, merge_views)
+  return sink

-  laneIdx = (tid % WARP_SIZE) % LANES_PER_WAVE_X
-  laneIdy = (tid % WARP_SIZE) // LANES_PER_WAVE_X
-  assert laneIdy.vmax+1 == LANES_PER_WAVE_Y
+def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
+  BLOCK_SIZE = 128 if kernel5 else 256

-  A_col = UOp.placeholder((ITERS_PER_WAVE_M, TM), dtypes.float, slot=0, addrspace=AddrSpace.REG)
-  A_col = copy(A_col, As[k, :].reshape(WAVES_IN_BLOCK_Y, ITERS_PER_WAVE_M, LANES_PER_WAVE_Y, TM)[waveIdy, :, laneIdy, :], 300, set=True, upcast=True)
+  nbWaves = BLOCK_SIZE // 32
+  WN = 128 if kernel5 else 64
+  WM = BN * BM // nbWaves // WN

-  B_row = UOp.placeholder((ITERS_PER_WAVE_N, TN), dtypes.float, slot=1, addrspace=AddrSpace.REG)
-  B_row = copy(B_row, Bs[k, :].reshape(WAVES_IN_BLOCK_X, ITERS_PER_WAVE_N, LANES_PER_WAVE_X, TN)[waveIdx, :, laneIdx, :], 400, set=True, upcast=True)
+  nbWaveX = BN // WN
+  nbWaveY = BM // WM

-  # ---------------------------
-  # FMA: c_regs += A_col * B_row
-  # ---------------------------
-  c_regs = UOp.placeholder((ITERS_PER_WAVE_M, TM, ITERS_PER_WAVE_N, TN), dtypes.float, slot=2, addrspace=AddrSpace.REG)
-  i = UOp.range(c_regs.size, 16)
-  c_regs = c_regs.after(c_regs.flatten()[i].store(0.0).end(i))
+  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
+  waveIndex = threadIdx_x // 32
+  waveIdx = waveIndex % nbWaveX
+  waveIdy = waveIndex // nbWaveX
+  indexInWave = threadIdx_x % 32

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

-  # Close k, sync, and close K tiles
-  sink = sink.end(k).barrier().end(k_tile_range)
+  idxInWave = indexInWave % nbThreadXPerWave
+  idyInWave = indexInWave // nbThreadXPerWave

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

-  return sink.sink(arg=KernelInfo(opts_to_apply=())).simplify()
+  SUBWN = WN // nbIterWaveN
+  SUBWM = WM // nbIterWaveM

-def test_matmul(sink:UOp, N=N):
-  rng = np.random.default_rng()
-  a = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  b = Tensor(rng.random((N, N), dtype=np.float32)-0.5)
-  hc = Tensor.empty(N, N)
-  Tensor.realize(a, b, hc)
+  # Thread mapping to read BKxBN block from A
+  rAIdx = threadIdx_x % BK
+  rAIdy = threadIdx_x // BK
+  # Thread mapping to read BNxBK block from B
+  rBIdx = threadIdx_x % BN
+  rBIdy = threadIdx_x // BN

-  ei = ExecItem(get_runner(Device.DEFAULT, sink), [t.uop.buffer for t in [hc, a, b]])
+  strideReadB = BLOCK_SIZE // BN
+  strideReadA = BLOCK_SIZE // BK
+  nbReadsB = BN * BK // BLOCK_SIZE
+  nbReadsA = BM * BK // BLOCK_SIZE

-  ets = []
-  with Context(DEBUG=2):
-    for _ in range(run_count):
-      ets.append(ei.run(wait=True))
-  print(f"REAL TFLOPS {N * N * N * 2 / min(ets) * 1e-12:.2f}")
+  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
+  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))

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

 if __name__ == "__main__":
-  test_matmul(hand_spec_kernel3(), N=N)
+  HL = getenv("HL")
+  if HL == 2: hprg = top_spec_kernel3()
+  elif HL == 1: hprg = hl_spec_kernel3()
+  else: hprg = hand_spec_kernel3()
+  prg = get_program(hprg, Device.default.renderer)
+  print(prg.src)
+  if getenv("SRC"): exit(0)
+  hrunner = CompiledRunner(prg)
+
+  a = Tensor.randn(N, N).realize()
+  b = Tensor.randn(N, N).realize()
+  hc = Tensor.zeros(N, N).contiguous().realize()
+
+  GlobalCounters.reset()
+  with Context(DEBUG=2):
+    for _ in range(run_count): tc = (a@b).realize()
+
+  GlobalCounters.reset()
+  buffers = [hc.uop.buffer, a.uop.buffer, b.uop.buffer]
+  ei = ExecItem(hrunner, buffers)
+  with Context(DEBUG=2):
+    for _ in range(run_count): ei.run(wait=True)
+  err = (hc-tc).square().mean().item()
+  print(f"hrunner {err}")
+  if err > 1e-06: raise RuntimeError("matmul is wrong!")
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import numpy as np
-from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
+from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler
 from tinygrad import Device, dtypes
 from tinygrad.device import Buffer
 from hexdump import hexdump
@@ -11,7 +11,7 @@ from hexdump import hexdump
 # https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
 # https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf

-device = Device["CL"]
+device = Device["GPU"]

 # NOTE: only the subgroup type 8 ones work
 prog = CLProgram(device, "test", CLCompiler(device, "test").compile(f"""
@@ -26,9 +26,9 @@ __kernel void test(__global float* data0, const __global int* data1, const __glo
 """))
 #with open("/tmp/test.elf", "wb") as f: f.write(prog.lib)

-a = Buffer("CL", 8, dtypes.float32).allocate()
-b = Buffer("CL", 0x10, dtypes.float16).allocate()
-c = Buffer("CL", 8*0x10, dtypes.float16).allocate()
+a = Buffer("GPU", 8, dtypes.float32).allocate()
+b = Buffer("GPU", 0x10, dtypes.float16).allocate()
+c = Buffer("GPU", 8*0x10, dtypes.float16).allocate()

 row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
 mat = np.random.random((8, 0x10)).astype(np.float16)
@@ -1,11 +1,17 @@
 import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
+from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 from tinygrad.engine.realize import get_program

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

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

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

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

    if GEMM_VARIATION == "max" and (M%64)==0 and (N%128)==0 and (K%64)==0 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.float and DTYPE_ACC == dtypes.float:
      print("Using CUDA and triton-generated kernel")
-      # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py`
+      # See nv_triton_gemm.annotated.ptx for PTX code which was generated from `PYTHONPATH=. DEBUG=6 CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py`
      # this kernel with M=N=K=4096 does 162TFLOPS, vs torch at 144TFLOPS and BEAM=8 tinygrad at 138TFLOPS.  theo max is 165TFLOPS.

      # WMMA element size is (M, N, K) = (16, 8, 16)
@@ -177,11 +189,11 @@ if __name__ == "__main__":

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

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