fine for ptx

that hack broke things
Revert "fix tests"
2025-08-05 18:26:06 -07:00 · 2025-08-05 18:24:58 -07:00 · 2025-08-05 18:06:27 -07:00 · 2025-08-05 17:50:14 -07:00 · 2025-08-05 17:45:35 -07:00 · 2025-08-05 17:07:12 -07:00
643 changed files with 234496 additions and 76612 deletions
@@ -41,10 +41,6 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
  mesa:
    description: "Install mesa"
    required: false
    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -125,7 +121,7 @@ runs:
        echo 'Acquire::GzipIndexes "true";' | sudo tee /etc/apt/apt.conf.d/gzip
        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
        echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' | sudo tee -a /etc/apt/apt.conf.d/99keep-debs
-
+    
    - name: Add OpenCL Repo
      if: inputs.opencl == 'true' && runner.os == 'Linux'
      shell: bash
@@ -178,7 +174,7 @@ runs:
        if [[ "${{ inputs.llvm }}" == "true" ]]; then
          pkgs+=" libllvm20 clang-20 lld-20"
        fi
-
+        
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
@@ -187,21 +183,21 @@ runs:
      uses: actions/cache@v4
      with:
        path: /var/cache/apt/archives/
-        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.APT_CACHE_VERSION }}
+        key: ${{ runner.os }}-apt-${{ steps.apt-pkgs.outputs.hash }}
    - name: Run apt Update + Install
      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.cuda == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
-
+        
        # ******** do install ********
        if [[ -n "${{ steps.apt-pkgs.outputs.pkgs }}" ]]; then
          sudo apt-get -y --allow-unauthenticated --no-install-recommends install ${{ steps.apt-pkgs.outputs.pkgs }}
        fi
-
+        
        sudo chown -R $USER:$USER /var/cache/apt/archives/
-
+    
    # **** AMD ****
    - name: Setup AMD (Linux)
      if: inputs.amd == 'true' && runner.os == 'Linux'
@@ -229,25 +225,16 @@ runs:
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
-      run: |
+      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
        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
    - 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.BUILD_CACHE_VERSION }}
+        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
    - 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 .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
        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
        ninja
    - name: Install gpuocelot
      if: inputs.ocelot == 'true'
@@ -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 -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
    - name: Install mesa (macOS)
      if: inputs.mesa == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install sirhcm/tinymesa/tinymesa_cpu
@@ -1,97 +0,0 @@
 name: Autogen
 env:
  # increment this when downloads substantially change to avoid the internet
  DOWNLOAD_CACHE_VERSION: '12'
  PYTHON_CACHE_VERSION: '3'
  APT_CACHE_VERSION: '1'
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
  PYTHONPATH: ${{ github.workspace }}
 on:
  push:
    branches:
      - master
  pull_request:
    paths: 
    - 'tinygrad/runtime/autogen/**/*'
  workflow_dispatch:
    paths: 
    - 'tinygrad/runtime/autogen/**/*'
 jobs:
  autogen:
    name: 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'
        webgpu: 'true'
        llvm: 'true'
        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
    - name: Verify OpenCL autogen
      run: |
        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
        ./autogen_stubs.sh opencl
        diff /tmp/opencl.py.bak tinygrad/runtime/autogen/opencl.py
    - name: Verify CUDA autogen
      run: |
        cp tinygrad/runtime/autogen/cuda.py /tmp/cuda.py.bak
        cp tinygrad/runtime/autogen/nv_gpu.py /tmp/nv_gpu.py.bak
        ./autogen_stubs.sh cuda
        ./autogen_stubs.sh nv
        diff /tmp/cuda.py.bak tinygrad/runtime/autogen/cuda.py
        diff /tmp/nv_gpu.py.bak tinygrad/runtime/autogen/nv_gpu.py
    - name: Verify AMD autogen
      run: |
        cp tinygrad/runtime/autogen/hsa.py /tmp/hsa.py.bak
        cp tinygrad/runtime/autogen/kfd.py /tmp/kfd.py.bak
        cp tinygrad/runtime/autogen/comgr.py /tmp/comgr.py.bak
        cp tinygrad/runtime/autogen/amd_gpu.py /tmp/amd_gpu.py.bak
        cp tinygrad/runtime/autogen/sqtt.py /tmp/sqtt.py.bak
        ./autogen_stubs.sh hsa
        ./autogen_stubs.sh kfd
        ./autogen_stubs.sh comgr
        ./autogen_stubs.sh amd
        ./autogen_stubs.sh sqtt
        diff /tmp/hsa.py.bak tinygrad/runtime/autogen/hsa.py
        diff /tmp/kfd.py.bak tinygrad/runtime/autogen/kfd.py
        diff /tmp/comgr.py.bak tinygrad/runtime/autogen/comgr.py
        diff /tmp/amd_gpu.py.bak tinygrad/runtime/autogen/amd_gpu.py
        diff /tmp/sqtt.py.bak tinygrad/runtime/autogen/sqtt.py
    - name: Verify Linux autogen
      run: |
        cp tinygrad/runtime/autogen/libc.py /tmp/libc.py.bak
        cp tinygrad/runtime/autogen/io_uring.py /tmp/io_uring.py.bak
        cp tinygrad/runtime/autogen/ib.py /tmp/ib.py.bak
        ./autogen_stubs.sh libc
        ./autogen_stubs.sh io_uring
        ./autogen_stubs.sh ib
        diff /tmp/libc.py.bak tinygrad/runtime/autogen/libc.py
        diff /tmp/io_uring.py.bak tinygrad/runtime/autogen/io_uring.py
        diff /tmp/ib.py.bak tinygrad/runtime/autogen/ib.py
    - name: Verify WebGPU autogen
      run: |
        cp tinygrad/runtime/autogen/webgpu.py /tmp/webgpu.py.bak
        ./autogen_stubs.sh webgpu
        diff /tmp/webgpu.py.bak tinygrad/runtime/autogen/webgpu.py
    - name: Verify LLVM autogen
      run: |
        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
        ./autogen_stubs.sh llvm
        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
    - name: Verify mesa autogen
      run: |
        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
        ./autogen_stubs.sh mesa
        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
@@ -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,27 @@ 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=800 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 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 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 python3.11 test/external/external_model_benchmark.py
    - name: Run huggingface_onnx test
      run: METAL=1 python3.11 extra/huggingface_onnx/run_models.py test --debug FacebookAI/xlm-roberta-large
    - name: Test speed vs torch
-      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
+      run: BIG=2 MPS=1 python3.11 test/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 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 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        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
        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
    - 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 +99,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 +108,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
-
+    - name: Run 10 CIFAR training steps
-    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
+      run: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
-    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
+    - name: Run 10 CIFAR training steps w HALF
-    #- name: Run 10 CIFAR training steps
+      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
    #  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 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
-    # - 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
    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    #- name: UsbGPU openpilot test
+    - name: UsbGPU openpilot test
-    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -168,7 +160,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}
@@ -197,22 +189,22 @@ jobs:
    - name: Run model inference benchmark
      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
    - 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
+      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      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 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        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
    - 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
@@ -221,9 +213,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
-    # - 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
    #   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 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
@@ -247,9 +238,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
@@ -283,7 +274,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}
@@ -308,36 +299,33 @@ 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)
-    # - 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
-    #   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)
-    # TODO: too slow
+      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: 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: 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=270 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=240 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=270 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    # TODO: too slow
+    - name: Run 10 CIFAR training steps w winograd
-    # - 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
    #   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
    - 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)
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+      run: BENCHMARK_LOG=resnet_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NVIDIA Training)
@@ -358,7 +346,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}
@@ -403,13 +391,13 @@ jobs:
    #- name: Test speed vs torch
    #  run: |
    #    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
+    #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
    - name: Test speed vs theoretical
      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=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
+        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
@@ -427,10 +415,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
-    # TODO: too slow
+    - name: Run SDXL
-    # - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    #   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
    - 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
@@ -456,9 +443,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
@@ -489,7 +476,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}
@@ -521,20 +508,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=330 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=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.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
+    - 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
+      run: BENCHMARK_LOG=cifar_10steps_bf16 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
-    # - 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
    #   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
    - 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
+    - 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)
+    - 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)
@@ -553,7 +539,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}
@@ -584,13 +570,13 @@ jobs:
      run: test/external/process_replay/reset.py
    - name: Run MLPerf resnet eval
      run: time BENCHMARK_LOG=resnet_eval AMD=1 MODEL=resnet python3 examples/mlperf/model_eval.py
-    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
+      run: BENCHMARK_LOG=resnet_10steps AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
-    #- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
+    - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-    #  run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
+      run: BENCHMARK_LOG=resnet_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
    - name: Run 10 MLPerf Bert training steps (6 gpu)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: BENCHMARK_LOG=bert_10steps_6gpu AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD MLPerf)
@@ -619,28 +605,30 @@ 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.9.9 driving_vision
+    - name: validate openpilot 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_9_vision PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: openpilot compile3 0.9.9 driving_policy
+    - name: benchmark openpilot 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_9_policy PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+      run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: openpilot compile3 0.9.9 dmonitoring
+    - name: benchmark openpilot w IMAGE=2 0.9.7
-      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: BENCHMARK_LOG=openpilot_0_9_7_image 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.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: openpilot compile3 0.10.1 driving_vision
+    - name: openpilot compile3 0.9.7
-      # TODO: ASSERT_MIN_STEP_TIME=17
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx
-      run: BENCHMARK_LOG=openpilot_0_10_1_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=25 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/cf6376aa9a090f0da26c280ef69eabf9bbdd51d1faac9ed392919c3db69be916
+    - name: openpilot compile3 0.9.7+ tomb raider
-    - name: openpilot compile3 0.10.1 driving_policy
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/e8bea2c78ffa92685ece511e9b554122aaf1a79d/selfdrive/modeld/models/supercombo.onnx
-      run: BENCHMARK_LOG=openpilot_0_10_1_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/driving_policy.onnx
+    - name: openpilot dmonitoring compile3 0.9.7
-    - name: openpilot compile3 0.10.1 dmonitoring
+      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/dmonitoring_model.onnx
-      # TODO: ASSERT_MIN_STEP_TIME=10
+    - name: openpilot compile3 Space Lab policy + vision
-      run: BENCHMARK_LOG=openpilot_0_10_1_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=13 DEV=QCOM FLOAT16=1 IMAGE=2 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/refs/heads/master/selfdrive/modeld/models/dmonitoring_model.onnx
+      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
@@ -657,7 +645,7 @@ jobs:
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 20
+    timeout-minutes: 15
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -693,8 +681,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=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     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
@@ -702,18 +690,13 @@ 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)
-    # - 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
    #   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)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee am_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AM Driver)
@@ -728,7 +711,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}
@@ -760,22 +743,18 @@ 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 AMD=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)
+    - 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
+      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)
      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
+      run: BENCHMARK_LOG=bert_10steps NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 FUSE_ARANGE=1 FUSE_ARANGE_UINT=0 MODEL=bert python3 examples/mlperf/model_train.py | tee nv_train_bert_one_gpu.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (NV Driver)
@@ -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
@@ -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
@@ -20,15 +20,27 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
-      - id: example
+      - id: devicetests
-        name: test all devices
+        name: select GPU tests
-        entry: python3 test/external/external_test_example.py
+        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 PYTHONPATH="." python3 -m pytest -n=8 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
@@ -79,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] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
- [x] [CPU](tinygrad/runtime/ops_cpu.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)
@@ -198,7 +198,11 @@ generate_amd() {
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    extra/hip_gpu_driver/nvd.h \
    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
    extra/hip_gpu_driver/soc21_enum.h \
    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
    extra/hip_gpu_driver/gc_11_0_0_offset.h \
    extra/hip_gpu_driver/gc_10_3_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
@@ -372,6 +376,26 @@ generate_am() {
    -o $BASE/am/pm4_nv.py
  fixup $BASE/am/pm4_nv.py
  clang2py -k cdefstum \
    $AMKERN_INC/vega10_enum.h \
    -o $BASE/am/vega10.py
  fixup $BASE/am/vega10.py
  clang2py -k cdefstum \
    $AMKERN_INC/navi10_enum.h \
    -o $BASE/am/navi10.py
  fixup $BASE/am/navi10.py
  clang2py -k cdefstum \
    $AMKERN_INC/soc21_enum.h \
    -o $BASE/am/soc21.py
  fixup $BASE/am/soc21.py
  clang2py -k cdefstum \
    $AMKERN_INC/soc24_enum.h \
    -o $BASE/am/soc24.py
  fixup $BASE/am/soc24.py
  clang2py -k cdefstum \
    extra/hip_gpu_driver/sdma_registers.h \
    $AMKERN_AMD/amdgpu/vega10_sdma_pkt_open.h \
@@ -414,29 +438,10 @@ generate_sqtt() {
  clang2py -k cdefstum \
    extra/sqtt/sqtt.h \
    -o $BASE/sqtt.py
  fixup $BASE/sqtt.py
  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
  python3 -c "import tinygrad.runtime.autogen.sqtt"
  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
  if [ ! -d "$ROCPROF_SRC" ]; then
    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
    pushd .
    cd $ROCPROF_SRC
    git reset --hard $ROCPROF_COMMIT_HASH
    popd
  fi
  clang2py -k cdefstum \
    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
    $ROCPROF_SRC/include/trace_decoder_instrument.h \
    $ROCPROF_SRC/include/trace_decoder_types.h \
    -o extra/sqtt/rocprof/rocprof.py
  fixup extra/sqtt/rocprof/rocprof.py
  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
 }
 generate_webgpu() {
@@ -461,80 +466,6 @@ generate_libusb() {
  python3 -c "import tinygrad.runtime.autogen.libusb"
 }
 generate_mesa() {
  MESA_TAG="mesa-25.2.4"
  MESA_SRC=/tmp/mesa-$MESA_TAG
  TINYMESA_TAG=tinymesa-32dc66c
  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
  if [ ! -d "$MESA_SRC" ]; then
    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
    pushd .
    cd $MESA_SRC
    git reset --hard $MESA_COMMIT_HASH
    # clang 14 doesn't support packed enums
    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
    mkdir -p gen/util/format
    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
    popd
  fi
  if [ ! -d "$TINYMESA_DIR" ]; then
    mkdir $TINYMESA_DIR
    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
  fi
  clang2py -k cdefstu \
    $MESA_SRC/src/compiler/nir/nir.h \
    $MESA_SRC/src/compiler/nir/nir_builder.h \
    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
    $MESA_SRC/src/compiler/nir/nir_serialize.h \
    $MESA_SRC/gen/nir_intrinsics.h \
    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
    $MESA_SRC/src/nouveau/compiler/nak.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
    $MESA_SRC/src/compiler/glsl_types.h \
    $MESA_SRC/src/util/blob.h \
    $MESA_SRC/src/util/ralloc.h \
    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
    -l $TINYMESA_SO \
    -o $BASE/mesa.py
  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
  fixup $BASE/mesa.py
  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "'/opt/homebrew/lib/libtinymesa_cpu.dylib'" "'/opt/homebrew/lib/libtinymesa.dylib'"
  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
  sed -i "/in_dll/s/.*/try: &\nexcept (AttributeError, ValueError): pass/" $BASE/mesa.py
  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
  echo "def __getattr__(nm): raise AttributeError('LLVMpipe requires tinymesa_cpu' if 'tinymesa_cpu' not in dll._name else f'attribute {nm} not found') if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
  # bitfield bug in clang2py
  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
  python3 -c "import tinygrad.runtime.autogen.mesa"
 }
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -558,7 +489,6 @@ elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
-elif [ "$1" == "mesa" ]; then generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
 elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
 else echo "usage: $0 <type>"
 fi
@@ -42,6 +42,7 @@ 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()
@@ -50,14 +51,13 @@ b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struc
 # NOTE: a._buf is the same as the return from cpu.allocator.alloc
 # describe the computation
 idx = UOp.const(dtypes.index, 0)
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
 buf_2 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 2)
-ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.view(ShapeTracker.from_shape((1,))),))
-ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
+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.index(idx), alu))
+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)
@@ -80,7 +80,7 @@ print("******** third, the UOp ***********")
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.schedule.rangeify import get_rangeify_map
+from tinygrad.schedule.kernelize import get_kernelize_map
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +93,10 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))
 # group the computation into kernels
-becomes_map = get_rangeify_map(s)
+becomes_map = get_kernelize_map(s)
 # the compute maps to an assign
-assign = becomes_map[a+b].base
+assign = becomes_map[a+b]
 # the first source is the output buffer (data)
 assert assign.src[0].op is Ops.BUFFER
@@ -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
@@ -18,10 +18,16 @@ Group UOps into kernels.
 ---
-## tinygrad/codegen/opt
+## tinygrad/opt
 Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
 ::: tinygrad.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,16 +31,19 @@ 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
 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)
@@ -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 |
+| Runtime | Description | 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. |
+| [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 | 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. |
+| [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 |
+| [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 |
+| [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 |
+| [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 |
-| [CL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cl.py) | Accelerates computations using OpenCL on GPUs | - | OpenCL 2.0 compatible device |
+| [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](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` |
+| [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` |
-| [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) |
+| [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.
@@ -47,8 +47,8 @@ Reboot after making these changes or restart the `displayservice.service` servic
 The [default tinybox image](https://github.com/tinygrad/tinyos) ships with tinygrad and PyTorch. While we develop tinygrad, the box is universal hardware. Use whatever framework you desire, run notebooks, download demos, install more things, train, inference, live, laugh, love, you aren't paying per hour for this box so the only limit is your imagination.
-## Building the OS image
+## tinychat
-The OS image is built using `ubuntu-image` from <https://github.com/tinygrad/tinyos>.
+Since LLMs are so popular, we ship with a built in tinygrad based chatbot using a LLaMA-3 finetune. Visit the IP (not the BMC IP) of your tinybox in a web browser on your computer or phone, and you'll find a friendly looking chat interface. This chatbot also provides an OpenAI compatible LLM API on that port, so you can script it.
-After cloning, run `make green` or `make red` to build a tinybox green or tinybox red image respectively.
+The conversations you have with this chatbot are between you and your tinybox. Also, the history in the web app is saved on the client, not the tinybox.
@@ -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)
+  np.random.seed(1337)
  num_train_samples = X_train.shape[0]
  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()
@@ -1,12 +1,12 @@
 # model based off https://medium.com/data-science/going-beyond-99-mnist-handwritten-digits-recognition-cfff96337392
-from typing import Callable
+from typing import List, Callable
 from tinygrad import Tensor, TinyJit, nn, GlobalCounters
 from tinygrad.helpers import getenv, colored, trange
 from tinygrad.nn.datasets import mnist
 class Model:
  def __init__(self):
-    self.layers: list[Callable[[Tensor], Tensor]] = [
+    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,
@@ -21,15 +21,17 @@ if __name__ == "__main__":
  X_train, Y_train, X_test, Y_test = mnist(fashion=getenv("FASHION"))
  model = Model()
-  opt = (nn.optim.Adam if not getenv("MUON") else nn.optim.Muon)(nn.state.get_parameters(model))
+  opt = nn.optim.Adam(nn.state.get_parameters(model))
  @TinyJit
  @Tensor.train()
  def train_step() -> Tensor:
    opt.zero_grad()
    samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0])
    # TODO: this "gather" of samples is very slow. will be under 5s when this is fixed
    loss = model(X_train[samples]).sparse_categorical_crossentropy(Y_train[samples]).backward()
-    return loss.realize(*opt.schedule_step())
+    opt.step()
    return loss
  @TinyJit
  def get_test_acc() -> Tensor: return (model(X_test).argmax(axis=1) == Y_test).mean()*100
@@ -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)
+    xqkv = self.c_attn(x)
-    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
+    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, :, :]
+      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
-      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
+      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:
@@ -0,0 +1,134 @@
 from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
 from tinygrad.opt.kernel import Kernel
 from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
 from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
 from tinygrad.engine.realize import get_program
 def get_sched_resnet():
  mdl = ResNet50()
  optim = (nn.optim.LARS if getenv("LARS") else nn.optim.SGD)(nn.state.get_parameters(mdl))
  BS = getenv("BS", 64)
  # run model twice to get only what changes, these are the kernels of the model
  for _ in range(2):
    out = mdl(Tensor.empty(BS, 3, 224, 224))
    targets = [out]
    if getenv("BACKWARD"):
      optim.zero_grad()
      out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
      targets += [x for x in optim.schedule_step()]
    sched = Tensor.schedule(*targets)
    print(f"schedule length {len(sched)}")
  return sched
 def get_sched_bert():
  mdl = get_mlperf_bert_model()
  optim = nn.optim.LAMB(nn.state.get_parameters(mdl))
  # fake data
  BS = getenv("BS", 9)
  input_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
  segment_ids = Tensor.empty((BS, 512), dtype=dtypes.float32)
  attention_mask = Tensor.empty((BS, 512), dtype=dtypes.default_float)
  masked_positions = Tensor.empty((BS, 76), dtype=dtypes.float32)
  masked_lm_ids = Tensor.empty((BS, 76), dtype=dtypes.float32)
  masked_lm_weights = Tensor.empty((BS, 76), dtype=dtypes.float32)
  next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
  # run model twice to get only what changes, these are the kernels of the model
  for _ in range(2):
    lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
    targets = [lm_logits, seq_relationship_logits]
    if getenv("BACKWARD"):
      optim.zero_grad()
      loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
      # ignore grad norm and loss scaler for now
      loss.backward()
      targets += [x for x in optim.schedule_step()]
    sched = Tensor.schedule(*targets)
    print(f"schedule length {len(sched)}")
  return sched
 if __name__ == "__main__":
  if getenv("HALF", 1):
    dtypes.default_float = dtypes.half
  # the device we are optimizing for
  device: Compiled = Device[Device.DEFAULT]
  if getenv("BACKWARD"): Tensor.training = True
  print(f"optimizing for {Device.DEFAULT}")
  sched = globals()[f"get_sched_{getenv('MODEL', 'resnet')}"]()
  sched = [x for x in sched if x.ast.op is Ops.SINK]
  # focus on one kernel
  if getenv("KERNEL", -1) >= 0: sched = sched[getenv("KERNEL", -1):getenv("KERNEL", -1)+1]
  # work with the schedule
  total_tm = 0
  running_gflops = 0
  usage = {}
  for i,si in enumerate(sched):
    if DEBUG >= 3: print(si.ast)
    rawbufs = bufs_from_lin(Kernel(si.ast))
    # "linearize" the op into uops in different ways
    lins: list[tuple[Kernel, str]] = []
    # always try hand coded opt
    lin = Kernel(si.ast, opts=device.renderer)
    lin.apply_opts(hand_coded_optimizations(lin))
    lins.append((lin, "HC"))
    # maybe try tensor cores
    lin = Kernel(si.ast, opts=device.renderer)
    if lin.apply_tensor_cores():
      lins.append((lin, "TC"))
    # try a beam search
    if beam:=getenv("BEAM"):
      lin = Kernel(si.ast, opts=device.renderer)
      lin = beam_search(lin, rawbufs, beam, bool(getenv("BEAM_ESTIMATE", 1)))
      lins.append((lin, "BEAM"))
    # try MCTS
    if mcts:=getenv("MCTS"):
      lin = Kernel(si.ast, opts=device.renderer)
      lin = mcts_search(lin, rawbufs, mcts)
      lins.append((lin, "MCTS"))
    # benchmark the programs
    choices = []
    for lin, nm in lins:
      tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
      ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
      gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
      choices.append((tm, gflops, lin, prg, nm))
    sorted_choices = sorted(choices, key=lambda x: x[0])
    if DEBUG >= 1: # print all kernels
      for tm, gflops, lin, prg, nm in choices:
        print(f"                 kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS -- {colored(nm, 'green') if lin is sorted_choices[0][2] else nm}")
    tm, gflops, lin, prg, nm = sorted_choices[0]
    if getenv("SRC"):
      print(si.ast)
      print(lin.applied_opts)
      print(get_program(lin.get_optimized_ast(), lin.opts).src)
    total_tm += tm
    running_gflops += gflops * tm
    if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)
    usage[key] = (usage[key][0] + tm, usage[key][1] + 1)
    print(f"*** {total_tm*1000:7.2f} ms : kernel {i:2d} {lin.name+' '*(37-ansilen(lin.name))} {str(prg.global_size):18s} {str(prg.local_size):12s} takes {tm*1000:7.2f} ms, {gflops:6.0f} GFLOPS {[repr(m) if TRACEMETA >= 2 else str(m) for m in si.metadata]}")
  print(f"******* total {total_tm*1000:.2f} ms, {running_gflops/total_tm:6.0f} GFLOPS")
  print("usage:")
  for k in sorted(usage, key=lambda x: -usage[x][0])[:10]:
    print(f"{usage[k][0]*1000:.2f} ms: {k} ({usage[k][1]} times)")
@@ -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) # flip LR
      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # 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)
@@ -441,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)
@@ -479,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
+      tok_Tens = (logits/temp).softmax().multinomial()
      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()
    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"
@@ -1,4 +1,6 @@
-import os, random, pickle, queue, struct, math, functools, hashlib, time
+import functools
 import hashlib
 import os, random, pickle, queue, struct, math
 from typing import List
 from pathlib import Path
 from multiprocessing import Queue, Process, shared_memory, connection, Lock, cpu_count
@@ -511,33 +513,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:
@@ -557,21 +532,21 @@ class BinIdxDataset:
    start = self.idx.tell()
    end = start + self.count * dtypes.int32.itemsize
-    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32).numpy()
+    self.sizes = self.idx_t[start:end].bitcast(dtypes.int32)
    start = end
    end = start + self.count * dtypes.int64.itemsize
-    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+    self.pointers = self.idx_t[start:end].bitcast(dtypes.int64)
    start = end
    end = start + doc_count * dtypes.int64.itemsize
-    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64).numpy()
+    self.doc_idx = self.idx_t[start:end].bitcast(dtypes.int64)
    # bin file
    self.bin_t = Tensor(base_path.with_name(f"{base_path.name}.bin"))
  def _index(self, idx) -> tuple[int, int]:
-    return int(self.pointers[idx]), int(self.sizes[idx])
+    return self.pointers[idx].item(), self.sizes[idx].item()
  def get(self, idx, offset:int=0, length:int|None=None):
    ptr, size = self._index(idx)
@@ -591,13 +566,10 @@ class GPTDataset:
    # check for cache
    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
    cache_path = base_path.with_name(f"{base_path.name}.{cache_hash}.index_cache")
    print(f"try loading GPTDataset from {cache_path}...")
    if cache_path.exists():
      print("cache found, loading...")
      with open(cache_path, "rb") as f:
        self.doc_idx, self.sample_idx, self.shuffle_idx = pickle.load(f)
    else:
      print("cache not found, building index...")
      self.doc_idx = self._build_doc_idx()
      self.sample_idx = self._build_sample_idx()
      self.shuffle_idx = self._build_shuffle_idx()
@@ -656,20 +628,14 @@ class GPTDataset:
  # https://github.com/NVIDIA/Megatron-LM/blob/94bd476bd840c2fd4c3ebfc7448c2af220f4832b/megatron/core/datasets/gpt_dataset.py#L558
  def _build_doc_idx(self):
-    print(f"building doc_idx for {self.num_epochs=}, {self.indexed_dataset.count=}")
+    doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
-    st = time.perf_counter()
+    doc_idx = doc_idx.reshape(-1)
    # doc_idx = np.mgrid[:self.num_epochs, :self.indexed_dataset.count][1]
    doc_idx = np.arange(self.indexed_dataset.count).reshape(1, -1).repeat(self.num_epochs, axis=0).flatten()
    doc_idx = doc_idx.astype(np.int32)
    at = time.perf_counter()
    if self.shuffle: self.rng.shuffle(doc_idx)
    print(f"doc_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
    return doc_idx
  def _build_sample_idx(self):
-    print(f"building sample_idx for {self.samples=}, {self.seqlen=}, {self.doc_idx.shape[0]=}")
+    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int32)
    sample_idx_max = max(self.doc_idx.shape[0], self.indexed_dataset.sizes.max())
    sample_idx = np.empty((self.samples + 1, 2), dtype=np.int64 if sample_idx_max > dtypes.int32.max else np.int32)
    sample_idx_idx, doc_idx_idx, doc_offset = 0, 0, 0
    sample_idx[sample_idx_idx, 0], sample_idx[sample_idx_idx, 1] = doc_idx_idx, doc_offset
@@ -679,7 +645,7 @@ class GPTDataset:
      remaining_seqlen = self.seqlen + 1
      while remaining_seqlen > 0:
        doc_idx = int(self.doc_idx[doc_idx_idx])
-        doc_len = int(self.indexed_dataset.sizes[doc_idx]) - doc_offset
+        doc_len = self.indexed_dataset.sizes[doc_idx].item() - doc_offset
        remaining_seqlen -= doc_len
        if remaining_seqlen <= 0:
          doc_offset += remaining_seqlen + doc_len - 1
@@ -688,7 +654,7 @@ class GPTDataset:
          if doc_idx_idx == len(self.doc_idx) - 1:
            assert sample_idx_idx == self.samples
            doc_idx = int(self.doc_idx[doc_idx_idx])
-            doc_offset = int(self.indexed_dataset.sizes[doc_idx]) - 1
+            doc_offset = self.indexed_dataset.sizes[doc_idx].item() - 1
            break
          doc_idx_idx += 1
          doc_offset = 0
@@ -699,18 +665,13 @@ class GPTDataset:
    return sample_idx
  def _build_shuffle_idx(self):
    print(f"building shuffle_idx for {self.samples=}")
    st = time.perf_counter()
    shuffle_idx = np.arange(self.samples, dtype=np.int32)
    at = time.perf_counter()
    if self.shuffle: self.rng.shuffle(shuffle_idx)
    print(f"shuffle_idx built in {at - st:.3f}s, shuffled in {time.perf_counter() - at:.3f}s")
    return shuffle_idx
 class BlendedGPTDataset:
  def __init__(self, paths:list[Path], weights:list[float], samples:int, seqlen:int, seed:int, shuffle:bool):
-    self.shuffle = shuffle
+    self.seed = seed
    self.rng = np.random.RandomState(seed)
    # normalize weights
    total_weight = sum(weights)
@@ -722,47 +683,10 @@ class BlendedGPTDataset:
    self.datasets = [GPTDataset(path, samples_per_blend[i], seqlen, seed + i, shuffle) for i,path in enumerate(paths)]
    # check for cache
    cache_hash = hashlib.sha256(f"{samples}:{seqlen}:{seed}:{shuffle}".encode()).hexdigest()
    cache_path = paths[0].with_name(f"{paths[0].name}.{cache_hash}.blend_cache")
    print(f"try loading BlendedGPTDataset from {cache_path}...")
    if cache_path.exists():
      print("cache found, loading...")
      with open(cache_path, "rb") as f:
        self.dataset_idx, self.dataset_sample_idx = pickle.load(f)
    else:
      print("cache not found, building index...")
      self.dataset_idx, self.dataset_sample_idx = self._build_blend_idx()
      # save cache
      with open(cache_path, "wb") as f:
        pickle.dump((self.dataset_idx, self.dataset_sample_idx), f)
  def get(self, idx:int):
-    tokens = self.datasets[self.dataset_idx[idx]][self.dataset_sample_idx[idx]]
+    tokens = self.datasets[0][idx]
    return tokens
  def _build_blend_idx(self):
    dataset_idx = np.zeros(self.samples, dtype=np.int16)
    dataset_sample_idx = np.zeros(self.samples, dtype=np.int64)
    unspent_datasets = set(range(len(self.datasets)))
    dataset_sample_counts = [0] * len(self.datasets)
    for i in tqdm(range(self.samples)):
      error_argmax, error_max = 0, 0.0
      for di in unspent_datasets:
        error = self.weights[di] * max(i, 1) - dataset_sample_counts[di]
        if error > error_max:
          error_max = error
          error_argmax = di
      dataset_idx[i] = error_argmax
      dataset_sample_idx[i] = dataset_sample_counts[error_argmax]
      dataset_sample_counts[error_argmax] += 1
    return dataset_idx, dataset_sample_idx
 def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0, val:bool=True):
  if val:
    dataset = BlendedGPTDataset([
@@ -785,27 +709,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"
@@ -836,8 +739,8 @@ if __name__ == "__main__":
  def load_llama3(val):
    bs = 24
-    samples = 5760 if val else 1_200_000 * 1152
+    samples = 5760 if val else 1_200_000
-    seqlen = 8192
+    seqlen = 512
    max_, min_ = 0, math.inf
    for tokens in tqdm(batch_load_llama3(bs, samples, seqlen, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=5760, val=bool(val)), total=samples//bs):
@@ -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.helpers import prod, argfix
 from tinygrad.nn.state import get_parameters
 from extra.models.unet import UNetModel
 # 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)
+    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)
@@ -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 = 4
-  BS = getenv("BS", 4)
+  sequence_length = 512
  SMALL = getenv("SMALL", 0)
  SEQLEN = getenv("SEQLEN", 8192)
  MODEL_PATH = Path(getenv("MODEL_PATH", "/raid/weights/llama31_8b/"))
-  params = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
  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)
  @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
-    from examples.mlperf.dataloader import batch_load_llama3_small
+  iter = batch_load_llama3(bs, 5760, sequence_length, Path(getenv("BASEDIR", "/raid/datasets/c4/")), True)
    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)
  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)
+  log_perplexity = Tensor(losses).mean()
-  print(f"Log Perplexity: {log_perplexity}")
+  print(f"Log Perplexity: {log_perplexity.item()}")
 # 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
 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.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
-from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
+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:
+      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
        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,15 @@ 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)
  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,58 +1308,44 @@ 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 = MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}
  params = params | {"vocab_size": 32000} if not SMALL else params
  if (llama_layers:=getenv("LLAMA_LAYERS")) != 0: params['n_layers'] = llama_layers
  model = Transformer(**params, max_context=SEQLEN, jit=False, disable_kv_cache=True)
  if getenv("FAKEDATA"):
    for v in get_parameters(model):
      v = v.assign(Tensor.empty(v.shape))
  if (DP := getenv("DP", 1)) > 1:
    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
    for v in get_parameters(model):
      v.shard_(device, axis=None)
-  if (MP := getenv("MP", 1)) > 1:
+  # TODO: MP
-    device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+  # if (GPUS := getenv("GPUS", 1)) > 1:
-    for k,v in get_state_dict(model).items():
+  #   device = tuple(f"{Device.DEFAULT}:{i}" for i in range(GPUS))
-      if 'scale' in k: v.shard_(device, axis=None)  # from quantized
+  #   for k,v in get_state_dict(model).items():
-      elif '.attention.wq' in k: v.shard_(device, axis=0)
+  #     if 'scale' in k: v.shard_(device, axis=None)  # from quantized
-      elif '.attention.wk' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wq' in k: v.shard_(device, axis=0)
-      elif '.attention.wv' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wk' in k: v.shard_(device, axis=0)
-      elif '.attention.wo' in k: v.shard_(device, axis=1)
+  #     # elif '.attention.wv' in k: v.shard_(device, axis=0)
-      elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
+  #     # elif '.attention.wo' in k: v.shard_(device, axis=1)
-      elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
+  #     # elif '.feed_forward.w1.' in k: v.shard_(device, axis=0)
-      elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
+  #     # elif '.feed_forward.w2.' in k: v.shard_(device, axis=1)
-      elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
+  #     # elif '.feed_forward.w3.' in k: v.shard_(device, axis=0)
-      elif 'output.weight' in k: v.shard_(device, axis=0)
+  #     # elif 'tok_embeddings.weight' in k: v.shard_(device, axis=0)
-      else:
+  #     elif 'output.weight' in k: v.shard_(device, axis=0)  # 243.32
-        # attention_norm, ffn_norm, norm
+  #     else:
-        v.shard_(device, axis=None)
+  #       # print(k)
-      # prevents memory spike on device 0
+  #       # attention_norm, ffn_norm, norm
-      v.realize()
+  #       v.shard_(device, axis=None)
  optim = AdamW(get_parameters(model), lr=0.0,
                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
  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):
@@ -1381,9 +1355,6 @@ def train_llama3():
      if (DP := getenv("DP", 1)) > 1:
        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
        batch = batch.shard(device, 0)
      if (MP := getenv("MP", 1)) > 1:
        device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
        batch = batch.shard(device)
      logits:Tensor = model(batch[:, :-1], start_pos=0, temperature=math.nan)
      loss = logits.sparse_categorical_crossentropy(batch[:, 1:])
      loss.backward()
@@ -1397,7 +1368,7 @@ def train_llama3():
        total_norm += p.grad.float().square().sum()
      total_norm = total_norm.sqrt().contiguous()
      for p in optim.params:
-        p.grad = p.grad * (opt_gradient_clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
    optim.step()
    scheduler.step()
@@ -1406,231 +1377,32 @@ def train_llama3():
    loss.realize(lr)
    return loss, lr
-  @TinyJit
+  if getenv("FAKEDATA", 0):
-  @Tensor.train(False)
+    def fake_data():
-  def eval_step(model, tokens:Tensor):
+      for _ in range(SAMPLES // GBS):
-    if (DP := getenv("DP", 1)) > 1:
+        yield Tensor.randint(GBS, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(DP))
+    iter = fake_data()
-      tokens = tokens.shard(device, 0)
+  else:
-    if (MP := getenv("MP", 1)) > 1:
+    from examples.mlperf.dataloader import batch_load_llama3
-      device = tuple(f"{Device.DEFAULT}:{i}" for i in range(MP))
+    iter = batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
      tokens = tokens.shard(device)
    logits:Tensor = model(tokens[:, :-1], start_pos=0, temperature=math.nan)
    loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
    return loss.flatten().float()
-  # ** data iters **
+  i = 0
  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)
  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))
  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)
  iter = get_train_iter()
  i, sequences_seen = resume_ckpt, 0
  for tokens in tqdm(iter, total=SAMPLES//GBS):
    t = time.perf_counter()
    GlobalCounters.reset()
    loss, lr = train_step(model, tokens, grad_acc)
-    loss = loss.float().item()
+    # above as tqdm.write f-string
-
+    tqdm.write(f"{loss.item():.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
    i += 1
    sequences_seen += tokens.shape[0]
    tqdm.write(f"{loss:.4f} loss, {lr.item():.12f} LR, {GlobalCounters.mem_used / 1e9:.2f} GB used, {time.perf_counter()-t:.2f} s")
    if (fname:=getenv("LOSS_FILE", "")):
      with open(fname, "a") as f:
-        f.write(f"{i} {loss:.4f} {lr.item():.12f} {GlobalCounters.mem_used / 1e9:.2f}\n")
+        f.write(f"{i} {loss.item():.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"
+      fn = f"{ckpt_dir}/{i}.safe"
      safe_save(get_state_dict(model), fn)
-
+    i += 1
      tqdm.write("saving optim checkpoint")
      fn = f"{ckpt_dir}/llama3_{i}_optim.safe"
      safe_save(get_state_dict(scheduler), fn)
    if sequences_seen % EVAL_FREQ == 0 and (i != 1 or EVAL_FREQ == 1):
      tqdm.write(f"evaluating after {sequences_seen} sequences")
      # run eval
      eval_losses = []
      eval_iter = get_eval_iter()
      tqdm.write(f"evaluating {5760//EVAL_BS} batches of {EVAL_BS} sequences")
      for tokens in tqdm(eval_iter, total=5760//EVAL_BS):
        eval_losses += eval_step(model, tokens).tolist()
      log_perplexity = Tensor(eval_losses).mean().float().item()
      tqdm.write(f"eval log perplexity: {log_perplexity:.4f}")
      if log_perplexity < EVAL_TARGET:
        tqdm.write(f"target achieved after {sequences_seen} sequences")
        if getenv("CKPT"):
          if not os.path.exists(ckpt_dir := "./ckpts"): os.mkdir(ckpt_dir)
          fn = f"{ckpt_dir}/llama3.safe"
          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 +1412,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
@@ -4,8 +4,6 @@ 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
@@ -5,8 +5,6 @@ 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 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"
@@ -8,8 +8,6 @@ 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 TRAIN_STEPS=3900
 export IGNORE_OOB=1
 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"
@@ -11,8 +11,6 @@ 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 TRAIN_STEPS=3900
 export IGNORE_OOB=1
 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"
@@ -2,9 +2,9 @@
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -2,9 +2,9 @@
 export PYTHONPATH="." NV=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -5,9 +5,9 @@ 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=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -2,9 +2,9 @@
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -2,9 +2,9 @@
 export PYTHONPATH="." AMD=1
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -5,9 +5,9 @@ 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=90 EVAL_BS=90
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
-export IGNORE_OOB=1
+export FUSE_ARANGE=1 FUSE_ARANGE_UINT=0
 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
@@ -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=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,7 +1,8 @@
 # https://arxiv.org/pdf/2409.02060
-import time, functools
+import time
 import numpy as np
 np.set_printoptions(suppress=True, linewidth=1000)
 import functools
 from tinygrad import Tensor, nn, Device, GlobalCounters
 from tinygrad.helpers import Timing, getenv
 from extra.models.llama import Transformer, convert_from_huggingface
@@ -16,7 +17,7 @@ class MixtureFeedForward:
  def __call__(self, x:Tensor) -> Tensor:
    assert x.shape[0] == 1, "only BS=1"
    assert x.shape[1] == 1, "only length=1"
-    g = self.gate(x).softmax(-1)
+    g = self.gate(x).float().softmax(-1)
    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
    probs, sel = g.topk(self.activated_experts)
@@ -24,7 +25,7 @@ class MixtureFeedForward:
    # run MoE
    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
-    return (x_down * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
+    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
 # model is bf16, 1.3B active, 6.9B total
 # M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
@@ -1,12 +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
 import onnx
-from tinygrad.nn.onnx import OnnxRunner
+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"
@@ -17,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())}
+  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
-  if not getenv("NPY_IMG"):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
    inputs = {k:Tensor(v.numpy(), device=Device.DEFAULT).realize() if 'img' in k else v for k,v in inputs.items()}
  print("created tensors")
  run_onnx_jit = TinyJit(lambda **kwargs:
@@ -32,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
@@ -49,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=}"
@@ -66,76 +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 {(et-st)*1e3:6.2f} ms")
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+  print(out, val.shape, val.dtype)
  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 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()}
+  if any([x.device == 'NPY' for x in inputs.values()]):
-  out = run(**inputs_2x)
+    for v in new_inputs_numpy.values(): v *= 2
-  changed_val = out.numpy()
+    out = run(**inputs)
-  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+    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):
+def test_vs_onnx(new_inputs, test_val, onnx_file, ort=False):
-  import onnxruntime as ort
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
  onnx_inputs = {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)
+  if ort:
-  onnx_types = {x.name: ORT_TO_NP_DTYPES[x.type] for x in onnx_session.get_inputs()}
+    # test with onnxruntime
-  onnx_inputs = {k:onnx_inputs[k].astype(onnx_types[k]) for k in onnx_inputs}
+    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):
+  if test_val is not None:
-    st = time.perf_counter()
+    np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
-    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], onnx_inputs)
+    print("test vs onnx passed")
    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")
  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)
+  # same randomness as compile
-  if getenv("SELFTEST"):
+  Tensor.manual_seed(100)
-    test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
+  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):
+  default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
-    if not args.fakeweights:
+  weights_fn = args.weights_fn
-      default_weights_url = 'https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors'
+  if not weights_fn:
-      weights_fn = args.weights_fn
+    weights_url = args.weights_url if args.weights_url else default_weights_url
-      if not weights_fn:
+    weights_fn  = fetch(weights_url, os.path.basename(str(weights_url)))
        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))
+          v.replace(v.cast(dtypes.float16).realize())
    Tensor.realize(*get_state_dict(model).values())
  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 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:
+    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,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
+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.clip import Closed, Tokenizer
 from extra.models import unet, clip
 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.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_params))
-      self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments
+    self.first_stage_model = AutoencoderKL()
-
+    self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
    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])
  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
    temperature = 1
@@ -263,23 +227,18 @@ if __name__ == "__main__":
  parser.add_argument('--timing', action='store_true', help="Print timing per step")
  parser.add_argument('--seed', type=int, help="Set the random latent seed")
  parser.add_argument('--guidance', type=float, default=7.5, help="Prompt strength")
  parser.add_argument('--fakeweights', action='store_true', help="Skip loading checkpoints and use fake weights")
  args = parser.parse_args()
  model = StableDiffusion()
  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    if not args.fakeweights:
+    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)
      model_bin = fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt')
      load_state_dict(model, torch_load(model_bin)['state_dict'], verbose=False, strict=False, realize=False)
    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())
  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
@@ -307,23 +266,17 @@ 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])):
      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)
    del run
  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"
  # upsample latent space to image with autoencoder
  x = model.decode(latent)
  print(x.shape)
@@ -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
+# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 python3 examples/benchmark_onnx.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),
@@ -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)
@@ -71,8 +71,8 @@ def bbox_iou(box1, box2):
  # get the coordinates of the intersection rectangle
  inter_rect_x1 = np.maximum(b1_x1, b2_x1)
  inter_rect_y1 = np.maximum(b1_y1, b2_y1)
-  inter_rect_x2 = np.minimum(b1_x2, b2_x2)
+  inter_rect_x2 = np.maximum(b1_x2, b2_x2)
-  inter_rect_y2 = np.minimum(b1_y2, b2_y2)
+  inter_rect_y2 = np.maximum(b1_y2, b2_y2)
  #Intersection area
  inter_area = np.clip(inter_rect_x2 - inter_rect_x1 + 1, 0, 99999) * np.clip(inter_rect_y2 - inter_rect_y1 + 1, 0, 99999)
  #Union Area
@@ -297,13 +297,13 @@ class Darknet:
          # Get the number of weights of batchnorm
          num_bn_biases = math.prod(bn.bias.shape)
          # Load weights
-          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases].astype(np.float32))
+          bn_biases = Tensor(weights[ptr:ptr + num_bn_biases])
          ptr += num_bn_biases
-          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
+          bn_weights = Tensor(weights[ptr:ptr+num_bn_biases])
          ptr += num_bn_biases
-          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
+          bn_running_mean = Tensor(weights[ptr:ptr+num_bn_biases])
          ptr += num_bn_biases
-          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases].astype(np.float32))
+          bn_running_var = Tensor(weights[ptr:ptr+num_bn_biases])
          ptr += num_bn_biases
          # Cast the loaded weights into dims of model weights
          bn_biases = bn_biases.reshape(shape=tuple(bn.bias.shape))
@@ -319,7 +319,7 @@ class Darknet:
          # load biases of the conv layer
          num_biases = math.prod(conv.bias.shape)
          # Load weights
-          conv_biases = Tensor(weights[ptr: ptr+num_biases].astype(np.float32))
+          conv_biases = Tensor(weights[ptr: ptr+num_biases])
          ptr += num_biases
          # Reshape
          conv_biases = conv_biases.reshape(shape=tuple(conv.bias.shape))
@@ -327,7 +327,7 @@ class Darknet:
          conv.bias = conv_biases
        # Load weighys for conv layers
        num_weights = math.prod(conv.weight.shape)
-        conv_weights = Tensor(weights[ptr:ptr+num_weights].astype(np.float32))
+        conv_weights = Tensor(weights[ptr:ptr+num_weights])
        ptr += num_weights
        conv_weights = conv_weights.reshape(shape=tuple(conv.weight.shape))
        conv.weight = conv_weights
@@ -371,7 +371,7 @@ class Darknet:
 if __name__ == "__main__":
  model = Darknet(fetch('https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg').read_bytes())
  print("Loading weights file (237MB). This might take a while…")
-  model.load_weights('https://github.com/shadiakiki1986/yolov3.weights/releases/download/3.0.1/yolov3.weights')
+  model.load_weights('https://pjreddie.com/media/files/yolov3.weights')
  if len(sys.argv) > 1:
    url = sys.argv[1]
  else:
@@ -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")
@@ -1,16 +1,6 @@
-import re, ctypes, sys, importlib
+import re, ctypes, sys
-from tinygrad.runtime.support.am.amdev import AMDev, AMRegister
+from tinygrad.runtime.autogen.am import am, mp_11_0, mp_13_0_0, nbio_4_3_0, mmhub_3_0_0, gc_11_0_0, osssys_6_0_0
 class AMDFake(AMDev):
  def __init__(self, devfmt, vram, doorbell, mmio, dma_regions=None):
    self.devfmt, self.vram, self.doorbell64, self.mmio, self.dma_regions = devfmt, vram, doorbell, mmio, dma_regions
    self._run_discovery()
    self._build_regs()
 amdev = importlib.import_module("tinygrad.runtime.support.am.amdev")
 amdev.AMDev = AMDFake
 from tinygrad.runtime.ops_amd import PCIIface
 def parse_amdgpu_logs(log_content, register_names=None):
  register_map = register_names
@@ -33,11 +23,16 @@ def parse_amdgpu_logs(log_content, register_names=None):
  return processed_log
 def main():
  regs_offset = {13: {0: [3072, 37784576]}, 28: {0: [93184, 37754880], 1: [201327616, 201461760], 2: [209716224, 209850368], 3: [218104832, 218238976], 4: [226493440, 226627584], 5: [234882048, 235016192], 6: [243270656, 243404800]}, 21: {0: [28672, 12582912, 37795840, 130023424, 306184192], 1: [201326592, 201463808, 201465856, 204210176, 204472320], 2: [209715200, 209852416, 209854464, 212598784, 212860928], 3: [218103808, 218241024, 218243072, 220987392, 221249536], 4: [226492416, 226629632, 226631680, 229376000, 229638144], 5: [234881024, 235018240, 235020288, 237764608, 238026752], 6: [243269632, 243406848, 243408896, 246153216, 246415360]}, 22: {0: [18, 192, 13504, 36864, 37764096]}, 1: {0: [4704, 40960, 114688, 37760000]}, 2: {0: [3872, 37790720]}, 11: {0: [70656, 38103040]}, 12: {0: [106496, 37783552]}, 15: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 16: {0: [90112, 14417920, 14680064, 14942208, 38009856]}, 14: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 26: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 23: {0: [4256, 37789696]}, 33: {0: [0, 20, 3360, 66560, 37859328, 67371008]}, 25: {0: []}, 3: {0: [4704, 40960, 114688, 37760000]}, 4: {0: [4704, 40960, 114688, 37760000]}, 24: {0: [92160, 92672, 37752832, 54788096]}, 27: {0: [91648, 37751808], 1: [201339904, 201458176], 2: [209728512, 209846784], 3: [218117120, 218235392], 4: [226505728, 226624000], 5: [234894336, 235012608], 6: [243282944, 243401216]}, 29: {0: [201342976, 201344000, 205520896, 205537280], 1: [209731584, 209732608, 213909504, 213925888], 2: [218120192, 218121216, 222298112, 222314496], 3: [226508800, 226509824, 230686720, 230703104], 4: [234897408, 234898432, 239075328, 239091712], 5: [243286016, 243287040, 247463936, 247480320]}, 17: {0: [30720, 32256], 1: [31488, 73728]}}
  reg_names = {}
-  dev = PCIIface(None, 0)
+  def _prepare_registers(modules):
-  for x, y in dev.dev_impl.__dict__.items():
+    for base, m in modules:
-    if isinstance(y, AMRegister):
+      for k, regval in m.__dict__.items():
-      for inst, addr in y.addr.items(): reg_names[addr] = f"{x}, xcc={inst}"
+        if k.startswith("reg") and not k.endswith("_BASE_IDX") and (base_idx:=getattr(m, f"{k}_BASE_IDX", None)) is not None:
          reg_names[regs_offset[am.__dict__.get(f"{base}_HWIP")][0][base_idx] + regval] = k
  _prepare_registers([("MP0", mp_13_0_0), ("NBIO", nbio_4_3_0), ("MMHUB", mmhub_3_0_0), ("GC", gc_11_0_0), ("OSSSYS", osssys_6_0_0)])
  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
@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.opt.kernel import Ops, MemOp, UOp
+from tinygrad.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG
@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.opt.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.opt.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch
@@ -2,9 +2,9 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.codegen.opt.kernel import Ops
+from tinygrad.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
@@ -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"):
+      if u.arg[1].startswith("lidx"):
-        r[u] = f'v{u.src[0].arg}'
+        r[u] = f'v{u.arg[0]}'
-      elif u.arg.startswith("gidx"):
+      elif u.arg[1].startswith("gidx"):
-        r[u] = f's{2+u.src[0].arg}'
+        r[u] = f's{2+u.arg[0]}'
      else:
        raise NotImplementedError
    elif u.uop == UOps.CONST:
@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.codegen.opt.kernel import  UOp, Ops
+from tinygrad.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math
@@ -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"
-
+  with Context(JIT=2): run,special_names = jit_model(model, *inputs)
  # NOTE: CPU_COUNT=1, since export does not support threading
  with Context(JIT=2, CPU_COUNT=1): 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()}
@@ -2,11 +2,11 @@ from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
 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.schedule.kernelize import merge_views, view_left
 from tinygrad.helpers import getenv, colored, prod, unwrap
 from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.shape.view import strides_for_shape
-from tinygrad.codegen.opt.kernel import axis_colors, Opt, OptOps
+from tinygrad.opt.kernel import axis_colors
 from tinygrad.codegen.opt.swizzler import merge_views, view_left
 def to_colored(full_shape, axis_types): return '_'.join([colored(str(s), axis_colors[at]) for s,at in zip(full_shape, axis_types)])
@@ -44,27 +44,13 @@ pm = PatternMatcher([
  (UPat(Ops.VIEW, src=(UPat(Ops.REDUCE_AXIS, src=(UPat.var("src"),), name="r"),), name="view"), swizzle_reduceop),
 ])
 def rangeify_kernel3():
  a = Tensor.empty(N,N)
  b = Tensor.empty(N,N)
  c = a@b
  #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
  sink = c.schedule()[-1].ast
  #print(sink)
  opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
  opts += [Opt(OptOps.UPCAST, 1, 4), Opt(OptOps.LOCAL, 1, 16), Opt(OptOps.UPCAST, 1, 2)]
  opts += [Opt(OptOps.UNROLL, 0, 8)]
  return sink.replace(arg=KernelInfo(opts_to_apply=tuple(opts)))
 def top_spec_kernel3():
  a = Tensor.empty(N,N)
  b = Tensor.empty(N,N)
  c = a@b
  sink = c.schedule()[-1].ast
  L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
+  sink = 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))
@@ -185,7 +171,7 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), arg=2)
-  i = UOp.range(c_regs.dtype.size, 16)
+  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:
@@ -196,53 +182,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    kId = 0
    # load from globals into locals
-    i = UOp.range(nbReadsB, 0)
+    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(nbReadsA, 1)
+    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(N//BK-1, 2)
+    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(nbReadsB, 3)
+    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(nbReadsA, 4)
+    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(BK, first_range+0)
+      k = UOp.range(dtypes.int, BK, first_range+0)
      # load from locals into registers
-      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
-      i = UOp.range(TN, first_range+2)
+      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(nbIterWaveM, first_range+3)
+      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
-      i = UOp.range(TM, first_range+4)
+      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(nbIterWaveM, first_range+5)
+      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
-      yt = UOp.range(TM, first_range+6)
+      yt = UOp.range(dtypes.int, TM, first_range+6)
-      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
+      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
-      xt = UOp.range(TN, first_range+8)
+      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]
@@ -255,12 +241,12 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    sink = inner_loop(5, (barrier, regB_store, regA_store)).barrier()
    # load from registers into locals
-    i = UOp.range(nbReadsB, 14)
+    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(nbReadsA, 15)
+    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)
@@ -268,40 +254,40 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    # final iteration without the copy
    sink = inner_loop(16, (UOp.barrier(Bs_store, As_store),))
  else:
-    kId_range = UOp.range(N//BK, 0)
+    kId_range = UOp.range(dtypes.int, N//BK, 0)
    kId = kId_range*BK
    # load from globals into locals
-    i = UOp.range(nbReadsB, 1)
+    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(nbReadsA, 2)
+    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(BK, 3)
+    k = UOp.range(dtypes.int, BK, 3)
    # load from locals into registers
-    iterWave = UOp.range(nbIterWaveN, 4)
+    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-    i = UOp.range(TN, 5)
+    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(nbIterWaveM, 6)
+    iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
-    i = UOp.range(TM, 7)
+    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(nbIterWaveM, 8)
+    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-    yt = UOp.range(TM, 9)
+    yt = UOp.range(dtypes.int, TM, 9)
-    iterWaveN = UOp.range(nbIterWaveN, 10)
+    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
-    xt = UOp.range(TN, 12)
+    xt = UOp.range(dtypes.int, TN, 12)
    x = iterWaveN * TN + xt
    y = iterWaveM * TM + yt
    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -309,10 +295,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
                            iterWaveM, iterWaveN, yt, xt, k, kId_range)
  # store c_regs into c
-  iterWaveM = UOp.range(nbIterWaveM, 1000)
+  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 1000)
-  yt = UOp.range(TM, 1001)
+  yt = UOp.range(dtypes.int, TM, 1001)
-  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
-  xt = UOp.range(TN, 1003)
+  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
@@ -323,14 +309,10 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
 if __name__ == "__main__":
  HL = getenv("HL")
-  if HL == 3: hprg = rangeify_kernel3()
+  if HL == 2: hprg = top_spec_kernel3()
  elif HL == 2: hprg = top_spec_kernel3()
  elif HL == 1: hprg = hl_spec_kernel3()
  else: hprg = hand_spec_kernel3()
-  if HL == 3:
+  prg = get_program(hprg, Device.default.renderer)
    prg = get_program(hprg, Device.default.renderer)
  else:
    prg = get_program(hprg, Device.default.renderer)
  print(prg.src)
  if getenv("SRC"): exit(0)
  hrunner = CompiledRunner(prg)
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	29262a7543	fine for ptx	2025-08-05 18:26:06 -07:00
geohot	dcc6ddf0eb	that hack broke things	2025-08-05 18:24:58 -07:00
geohot	d1d935242b	Revert "fix tests" This reverts commit `a27019383d`.	2025-08-05 18:06:27 -07:00
geohot	6b330f302d	remote metal was flaky	2025-08-05 17:50:14 -07:00
geohot	a27019383d	fix tests	2025-08-05 17:45:35 -07:00
geohot	8b285e193a	move those to fix_kernel_ops	2025-08-05 17:07:12 -07:00
geohot	f0c9b11b9e	early meta ops	2025-08-05 17:00:49 -07:00
geohot	1902a85ac1	early load buffer	2025-08-05 16:56:10 -07:00
geohot	b2fc111e3f	cleanup fix_kernel	2025-08-05 16:46:13 -07:00