add more uop gc test

reduce_unparented only checks ranges (#12548 )
faster backward_slice (#12515 )
2025-10-09 14:27:22 +08:00 · 2025-10-09 14:14:03 +08:00 · 2025-10-09 14:12:20 +08:00 · 2025-10-09 13:53:08 +08:00 · 2025-10-09 01:52:30 -04:00 · 2025-10-09 07:14:27 +03:00
366 changed files with 19415 additions and 7900 deletions
@@ -225,13 +225,22 @@ runs:
    - name: Install gpuocelot dependencies (MacOS)
      if: inputs.ocelot == 'true' && runner.os == 'macOS'
      shell: bash
-      run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
+      run: |
+        pkgs=(cmake ninja llvm@15 zlib glew flex bison [email protected] zstd ncurses)
+        for f in "${pkgs[@]}"; do
+          brew ls --versions "$f" >/dev/null 2>&1 || brew install --quiet "$f"
+        done
+
+        # Fix boost 1.85 for gpuocelot
+        ln -s /opt/homebrew/opt/[email protected] /opt/homebrew/opt/boost || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_atomic-mt.dylib /opt/homebrew/opt/boost/lib/libboost_atomic.dylib || true
+        ln -s /opt/homebrew/opt/boost/lib/libboost_thread-mt.dylib /opt/homebrew/opt/boost/lib/libboost_thread.dylib || true
    - name: Cache gpuocelot
      if: inputs.ocelot == 'true'
      id: cache-build
      uses: actions/cache@v4
      env:
-        cache-name: cache-gpuocelot-build
+        cache-name: cache-gpuocelot-build-1
      with:
        path: ${{ github.workspace }}/gpuocelot/ocelot
        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-${{ env.BUILD_CACHE_VERSION }}
@@ -244,7 +253,13 @@ 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'
@@ -0,0 +1,91 @@
+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'
+    - name: Install autogen support packages
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+    - name: Verify OpenCL autogen
+      run: |
+        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
+        ./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
@@ -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: 20
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -52,24 +52,28 @@ jobs:
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      # TODO: very slow step time
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
-    - name: Run SDXL
-      run: BENCHMARK_LOG=stable_diffusion_xl CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
      run: BIG=2 MPS=1 python3.11 test/speed/external_test_speed_v_torch.py | tee torch_speed.txt
    - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: METAL=1 python3.11 test/opt/test_tensor_cores.py
    - name: Test AMX tensor cores
      run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_tensor_cores.py
+        DEBUG=2 CPU=1 CPU_LLVM=0 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 CPU_LLVM=1 AMX=1 python3.11 test/opt/test_gen_float4.py TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
    - 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)
@@ -97,7 +101,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 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.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
    - 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
@@ -106,22 +110,27 @@ 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
-      run: BENCHMARK_LOG=cifar_10steps JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
-    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+
+    # NOTE: this is failing in CI. it is not failing on my machine and I don't really have a way to debug it
+    # the error is "RuntimeError: Internal Error (0000000e:Internal Error)"
+    #- name: Run 10 CIFAR training steps
+    #  run: BENCHMARK_LOG=cifar_10steps JIT=1 ASSERT_MIN_STEP_TIME=3000 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
+    #- name: Run 10 CIFAR training steps w HALF
+    #  run: BENCHMARK_LOG=cifar_10steps_half JIT=2 ASSERT_MIN_STEP_TIME=3000 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+
    #- name: Run 10 CIFAR training steps w BF16
    #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    - name: Run 10 CIFAR training steps w winograd
-      run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+    # TODO: too slow
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_wino JIT=1 ASSERT_MIN_STEP_TIME=150 WINO=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: UsbGPU boot time
      run: sudo -E PYTHONPATH=. DEBUG=2 AM_RESET=1 AMD=1 AMD_IFACE=USB time python3.11 test/test_tiny.py TestTiny.test_plus
    - name: UsbGPU tiny tests
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/test_tiny.py
    - name: UsbGPU copy speeds
      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
-    - name: UsbGPU openpilot test
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB 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
+    #- name: UsbGPU openpilot test
+    #  run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (Mac)
@@ -158,7 +167,7 @@ jobs:
  testnvidiabenchmark:
    name: tinybox green Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 30
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -194,15 +203,15 @@ jobs:
      run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
    - name: Test tensor cores
      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
-        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
+        NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
+        NV=1 NV_PTX=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
    - 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 PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
+      run: NV=1 NV_PTX=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
    - name: Run Tensor Core GEMM (NV)
      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
@@ -211,8 +220,9 @@ 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
-    - 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
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=2000 CAPTURE_PROCESS_REPLAY=0 NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run 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
@@ -236,9 +246,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 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      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
    - 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
@@ -272,7 +282,7 @@ jobs:
  testmorenvidiabenchmark:
    name: tinybox green Training Benchmark
    runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 20
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -297,30 +307,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
-    - name: Fuzz Padded Tensor Core GEMM (NV)
-      run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
-    - name: Fuzz Padded Tensor Core GEMM (PTX)
-      run: NV=1 PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    # TODO: too slow
+    # - name: Fuzz Padded Tensor Core GEMM (NV)
+    #   run: NV=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
+    # TODO: too slow
+    # - name: Fuzz Padded Tensor Core GEMM (PTX)
+    #   run: NV=1 NV_PTX=1 M_START=12 M_STOP=20 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 python3 ./extra/gemm/fuzz_matmul.py
    - name: 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 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=270 NV=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    - name: Run 10 CIFAR training steps w winograd
-      run: BENCHMARK_LOG=cifar_10steps_half_wino NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=310 NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    # TODO: too slow
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=350 NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      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
+      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
    - 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.2 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.0 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)
-      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)
-      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 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
+    #- 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
    - 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
@@ -344,7 +357,7 @@ jobs:
  testamdbenchmark:
    name: tinybox red Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 20
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -394,8 +407,8 @@ jobs:
      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/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 AMD_LLVM=0 python3 test/opt/test_tensor_cores.py
+        AMD=1 AMD_LLVM=1 python3 test/opt/test_tensor_cores.py
        AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
    - name: 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
@@ -413,9 +426,10 @@ 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 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
-    - 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 ASSERT_MIN_STEP_TIME=550 AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+    # TODO: too slow
+    # - name: Run SDXL
+    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=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
@@ -441,9 +455,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 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_jitted.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
    - name: Run GPT2 w HALF
-      run: BENCHMARK_LOG=gpt2_half AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
+      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
    - 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
@@ -474,7 +488,7 @@ jobs:
  testmoreamdbenchmark:
    name: tinybox red Training Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 30
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -506,19 +520,20 @@ 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 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
+      run: BENCHMARK_LOG=cifar_10steps ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar.txt
    - name: Run 10 CIFAR training steps w HALF
-      run: BENCHMARK_LOG=cifar_10steps_half AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
-    - name: Run 10 CIFAR training steps w BF16
-      run: BENCHMARK_LOG=cifar_10steps_bf16 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
-    - name: Run 10 CIFAR training steps w winograd
-      run: BENCHMARK_LOG=cifar_10steps_half_wino AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: BENCHMARK_LOG=cifar_10steps_half ASSERT_MIN_STEP_TIME=330 AMD=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_half.txt
+    # - name: Run 10 CIFAR training steps w BF16
+    #   run: BENCHMARK_LOG=cifar_10steps_bf16 ASSERT_MIN_STEP_TIME=288 AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
+    # TODO: too slow
+    # - name: Run 10 CIFAR training steps w winograd
+    #   run: BENCHMARK_LOG=cifar_10steps_half_wino ASSERT_MIN_STEP_TIME=66 AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
    - name: Run full CIFAR training w 1 GPU
-      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS
-      run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
-    - name: Run full CIFAR training steps w 6 GPUS (REMOTE)
-      run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
+      run: time BENCHMARK_LOG=cifar AMD=1 DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS
+    #  run: time BENCHMARK_LOG=cifar_6gpu AMD=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+    #- name: Run full CIFAR training steps w 6 GPUS (REMOTE)
+    #  run: time BENCHMARK_LOG=cifar_6gpu_remote REMOTE=1 REMOTEDEV=AMD DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu_remote.txt
    - uses: actions/upload-artifact@v4
      with:
        name: Speed (AMD Training)
@@ -537,7 +552,7 @@ jobs:
  testmlperfamdbenchmark:
    name: tinybox red MLPerf Benchmark
    runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 30
+    timeout-minutes: 60
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -568,10 +583,10 @@ 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)
-      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)
-      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 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
+    #- 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
    - 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
@@ -610,21 +625,21 @@ jobs:
    - name: benchmark openpilot 0.9.9 dmonitoring
      run: BENCHMARK_LOG=openpilot_0_9_9_dmonitoring PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 0.9.9 driving_vision
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=22 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_vision.onnx
    - name: openpilot compile3 0.9.9 driving_policy
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=7 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/driving_policy.onnx
    - name: openpilot compile3 0.9.9 dmonitoring
-      run: PYTHONPATH="." QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
+      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=15 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
-        PYTHONPATH="." 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
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
        # generate quantized weights
        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 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 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
        PYTHONPATH=. CC=clang-19 DSP=1 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
@@ -643,7 +658,7 @@ jobs:
  testreddriverbenchmark:
    name: AM Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 15
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -679,8 +694,8 @@ jobs:
    # Fails on 9070
    # - name: Test tensor cores
    #   run: |
-    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    #     AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+    #     AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
+    #     AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py test/opt/test_tensor_cores.py
    #     AMD=1 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
@@ -688,8 +703,12 @@ 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 LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee am_train_cifar_one_gpu.txt
+      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
    # TODO: enable
    # - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
    #   run: BENCHMARK_LOG=resnet_10steps AMD=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee am_train_resnet_one_gpu.txt
@@ -710,7 +729,7 @@ jobs:
  testgreendriverbenchmark:
    name: NV Benchmark
    runs-on: [self-hosted, Linux, tinyboxrandom]
-    timeout-minutes: 15
+    timeout-minutes: 20
    defaults:
      run:
        shell: bash -e -o pipefail {0}
@@ -742,15 +761,19 @@ 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/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: NV=1 ALLOW_TF32=1 python3 test/opt/test_tensor_cores.py
    - 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 LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee nv_train_cifar_one_gpu.txt
-    - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
-      run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
+      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
+    #- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
+    #  run: BENCHMARK_LOG=resnet_10steps NV=1 MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee nv_train_resnet_one_gpu.txt
    - name: Run 10 MLPerf Bert training steps (1 gpu)
      # 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
@@ -7,6 +7,7 @@ env:
  BUILD_CACHE_VERSION: '1'
  CAPTURE_PROCESS_REPLAY: 1
  GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+  PYTHONPATH: ${{ github.workspace }}

 on:
  push:
@@ -29,12 +30,10 @@ jobs:
        key: llvm-speed
        deps: testing_minimal
        llvm: 'true'
-    - name: External Benchmark Schedule
-      run: PYTHONPATH="." python3 test/external/external_benchmark_schedule.py
    - name: Speed Test
-      run: LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py
    - name: Speed Test (BEAM=2)
-      run: BEAM=2 LLVM=1 python3 test/speed/external_test_speed_v_torch.py
+      run: BEAM=2 CPU=1 CPU_LLVM=1 python3 test/speed/external_test_speed_v_torch.py

  docs:
    name: Docs
@@ -47,7 +46,7 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        deps: docs
-        pydeps: "capstone"
+        pydeps: "capstone torch"
    - name: Build wheel and show size
      run: |
        pip install build
@@ -71,98 +70,29 @@ jobs:
        source venv/bin/activate
        pip install $GITHUB_WORKSPACE
        cp $GITHUB_WORKSPACE/examples/beautiful_mnist.py .
-        PYTHONPATH=$GITHUB_WORKSPACE BS=2 STEPS=10 python beautiful_mnist.py
+        BS=2 STEPS=10 python beautiful_mnist.py
    - name: Test Docs Build
      run: python -m mkdocs build --strict
    - name: Test Docs
      run: |
        python docs/abstractions2.py
        python docs/abstractions3.py
+    - name: Test README
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  python README.py
    - name: Test Quickstart
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  PYTHONPATH=. python quickstart.py
+      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py &&  python quickstart.py
    - name: Test DEBUG
      run: DEBUG=100 python3 -c "from tinygrad import Tensor; N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N); c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2); print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
    - name: Compile EfficientNet to C and test it
      run: |
-        CPU=1 PYTHONPATH="." python examples/compile_efficientnet.py > recognize.c
+        CPU=1 CPU_LLVM=0 python examples/compile_efficientnet.py > recognize.c
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  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'
-    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
-    - name: Verify OpenCL autogen
-      run: |
-        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
-        ./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
-
  torchbackend:
    name: Torch Backend Tests
    runs-on: ubuntu-latest
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -182,26 +112,24 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check extra/torch_backend/backend.py
    - name: Test one op
-      run: PYTHONPATH=. FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
    - name: Test ResNet-18
-      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/example.py
+      run: DEBUG=2 python3 extra/torch_backend/example.py
    - name: My (custom) tests
-      run: PYTHONPATH=. python3 extra/torch_backend/test.py
+      run: python3 extra/torch_backend/test.py
    - name: Test one op in torch tests
-      run: PYTHONPATH=. DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
    - name: Test Ops with TINY_BACKEND
-      run: PYTHONPATH=. LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
    - name: Test in-place operations on views
-      run: PYTHONPATH=. TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
    - name: Test multi-gpu
-      run: PYTHONPATH=. LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py

  torchbackendmore:
    name: Torch Backend Tests More
    runs-on: ubuntu-latest
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -216,86 +144,14 @@ jobs:
        sudo apt update || true
        sudo apt install -y --no-install-recommends ninja-build
    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: SPLIT_REDUCEOP=0 FUSE_ARANGE=1 PYTHONPATH=. LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
    - name: Test some torch tests (expect failure)
-      run: PYTHONPATH=. python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
-
-  tc:
-    name: Tensor Core tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 0
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: uops-minimal
-        deps: testing_minimal
-    - name: Test IMAGE=2 support
-      run: |
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
-    - name: Test emulated METAL tensor cores
-      run: |
-        DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test emulated AMX tensor cores
-      run: PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-    - name: Test emulated AMD tensor cores
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test emulated AMD MFMA tensor cores
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_MFMA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test emulated AMD RDNA4 tensor cores
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test emulated CUDA tensor cores
-      run: |
-        DEBUG=2 EMULATE_CUDA=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
-        DEBUG=2 EMULATE_CUDA_SM75=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH="." DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-    - name: Test emulated INTEL OpenCL tensor cores
-      run: DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
-    - name: Full test tensor cores
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-    - name: Test device flop counts
-      run: |
-        PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
-        PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul
+      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

  bepython:
    name: Python Backend
    runs-on: ubuntu-latest
-    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 0
+    timeout-minutes: 15
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -305,15 +161,60 @@ jobs:
        key: be-minimal
        deps: testing_minimal
    - name: Test dtype with Python emulator
-      run: DEBUG=1 PYTHONPATH=. PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
+      run: DEBUG=1 PYTHON=1 python3 -m pytest -n=auto test/test_dtype.py test/test_dtype_alu.py
    - name: Test ops with Python emulator
-      run: DEBUG=2 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py -k "not (test_split or test_simple_cumsum or test_cumsum or test_einsum or test_dot or test_dot_1d or test_big_gemm or test_broadcastdot or test_multidot or test_var_axis or test_std_axis or test_broadcast_full or test_broadcast_partial or test_simple_conv3d or test_dilated_conv_transpose2d or test_simple_conv_transpose3d or test_large_input_conv2d or test_max_pool2d or test_max_pool2d_simple or test_max_pool2d_bigger_stride or test_avg_pool2d or test_cat or test_scaled_product_attention or test_scaled_product_attention_causal or test_slice_fancy_indexing_dim_inject_none or test_slice_fancy_indexing_list_indices or test_slice_fancy_indexing_no_dim_collapse or test_slice_fancy_indexing_tuple_indices or test_slice_fancy_indexing_list_with_tensors or test_slice_fancy_indexing_dim_collapse_int or test_interpolate_bilinear or test_interpolate_bilinear_corners_aligned or test_scaled_dot_product_attention or test_cummax or test_simple_cummax or test_logcumsumexp or test_sort or test_cumprod)" --durations=20
+      run: DEBUG=2 SKIP_SLOW_TEST=1 PYTHON=1 python3 -m pytest -n=auto test/test_ops.py --durations=20
    - name: Test uops with Python emulator
      run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
    - name: Test symbolic with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
+      run: PYTHON=1 python3 test/test_symbolic_ops.py
    - name: test_renderer_failures with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+      run: PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
+    - name: Test IMAGE=2 support
+      run: |
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        IMAGE=2 PYTHON=1 python3 test/test_ops.py TestOps.test_simple_conv2d
+    - name: Test emulated METAL tensor cores
+      run: |
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_big_gemm
+        DEBUG=2 EMULATE=METAL FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test emulated AMX tensor cores
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+    - name: Test emulated AMD tensor cores
+      run: |
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test emulated AMD MFMA tensor cores
+      run: |
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_MFMA FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test emulated AMD RDNA4 tensor cores
+      run: |
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        DEBUG=2 EMULATE=AMD_RDNA4 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test emulated CUDA tensor cores
+      run: |
+        DEBUG=2 EMULATE=CUDA FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm
+        DEBUG=2 EMULATE=CUDA_SM75 FORWARD_ONLY=1 PYTHON=1 python3 test/test_ops.py TestOps.test_gemm_fp16
+        DEBUG=2 EMULATE=CUDA ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test emulated INTEL OpenCL tensor cores
+      run: DEBUG=2 EMULATE=INTEL FORWARD_ONLY=1 PYTHON=1 HALF=1 N=64 python3 ./extra/gemm/simple_matmul.py
+    - name: Test emulated AMX tensor cores
+      run: DEBUG=2 AMX=1 EMULATE=AMX FORWARD_ONLY=1 PYTHON=1 python3 test/opt/test_tensor_cores.py
+    - name: Test device flop counts
+      run: |
+        DEBUG=2 EMULATE=METAL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=AMD PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=CUDA PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 EMULATE=INTEL PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
+        DEBUG=2 AMX=1 EMULATE=AMX PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStats.test_simple_matmul

  linter:
    name: Linters
@@ -358,32 +259,39 @@ jobs:
      uses: ./.github/actions/setup-tinygrad
      with:
        key: unittest-12
-        pydeps: "pillow"
+        pydeps: "pillow numpy ftfy regex"
        deps: testing_unit
-    - name: Test README
-      run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
+    - name: Check Device.DEFAULT
+      run: python -c "from tinygrad import Device; assert Device.DEFAULT == 'CPU', Device.DEFAULT"
    - name: Run unit tests
-      run: PYTHONPATH="." python -m pytest -n=auto test/unit/ --durations=20
+      run: CPU=1 python -m pytest -n=auto test/unit/ --durations=20
+    - name: Check SPEC=1
+      run: SPEC=1 python3 test/test_tiny.py
    - name: Run targetted tests on NULL backend
-      run: PYTHONPATH="." NULL=1 python3 test/test_multitensor.py TestMultiTensor.test_data_parallel_resnet_train_step
-    - name: Run SDXL on NULL backend
-      run: MAX_BUFFER_SIZE=0 PYTHONPATH="." NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+      run: NULL=1 python3 -m unittest test.test_multitensor.TestMultiTensor.test_data_parallel_resnet_train_step test/device/test_null.py
+    # TODO: too slow
+    # - name: Run SDXL on NULL backend
+    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
+    - name: Run Clip tests for SD MLPerf on NULL backend
+      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
    - name: Run GC tests
-      run: PYTHONPATH="." python test/external/external_uop_gc.py
+      run: python test/external/external_uop_gc.py
+    - name: External Benchmark Schedule
+      run: python3 test/external/external_benchmark_schedule.py
    - name: Run process replay tests
      uses: ./.github/actions/process-replay
    - name: Regen dataset on test_tiny
      run: |
        test/external/process_replay/reset.py
        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
-        PYTHONPATH=. python extra/optimization/extract_dataset.py
+        python extra/optimization/extract_dataset.py
        gzip -c /tmp/sops > extra/datasets/sops.gz
-        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
-    - name: Repo line count < 17500 lines
-      run: MAX_LINE_COUNT=17500 python sz.py
+        #DEBUG=1 MIN_ASTS=1 python extra/optimization/get_action_space.py
+    - name: Repo line count < 18000 lines
+      run: MAX_LINE_COUNT=18000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -403,17 +311,15 @@ jobs:
      run: python test/external/fuzz_fast_idiv.py
    - name: Fuzz Test shapetracker
      run: |
-        PYTHONPATH="." python test/external/fuzz_shapetracker.py
-        PYTHONPATH="." python test/external/fuzz_shapetracker_math.py
+        python test/external/fuzz_shapetracker.py
+        python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

-  testgpuimage:
-    name: 'GPU IMAGE Tests'
+  testopenclimage:
+    name: CL IMAGE Tests
    runs-on: ubuntu-22.04
-    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 0
+    timeout-minutes: 15
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -423,25 +329,17 @@ jobs:
          key: gpu-image
          deps: testing_minimal
          opencl: 'true'
-      - name: Run Kernel Count Test
-        run: PYTHONPATH="." GPU=1 python -m pytest -n=auto test/external/external_test_opt.py
-      - name: Test WINO=1
-        run: GPU=1 DEBUG=2 WINO=1 python3 test/test_ops.py TestOps.test_simple_conv2d
-      - name: Test GPU IMAGE=2 ops + training
+      - name: Test CL IMAGE=2 ops + training
        run: |
-          PYTHONPATH="." GPU=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
-          PYTHONPATH="." GPU=1 IMAGE=2 python3 test/models/test_end2end.py TestEnd2End.test_linear_mnist
-      - name: Run fused optimizer tests
-        run: PYTHONPATH="." GPU=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
+          CL=1 IMAGE=2 python -m pytest -n=auto test/test_ops.py --durations=20
+          CL=1 IMAGE=2 python test/models/test_end2end.py TestEnd2End.test_linear_mnist
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testgendataset:
-    name: 'GPU Generate Kernel Dataset'
+  testgpumisc:
+    name: CL Misc tests
    runs-on: ubuntu-22.04
    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -452,7 +350,11 @@ jobs:
          deps: testing_minimal
          opencl: 'true'
      - name: Generate Dataset
-        run: PYTHONPATH="." extra/optimization/generate_dataset.sh
+        run: CL=1 extra/optimization/generate_dataset.sh
+      - name: Run Kernel Count Test
+        run: CL=1 python -m pytest -n=auto test/external/external_test_opt.py
+      - name: Run fused optimizer tests
+        run: CL=1 FUSE_OPTIM=1 python -m pytest -n=auto test/models/test_mnist.py
      - name: Upload artifact
        uses: actions/upload-artifact@v4
        with:
@@ -460,11 +362,9 @@ jobs:
          path: /tmp/sops.gz

  testopenpilot:
-    name: 'openpilot Compile Tests'
+    name: openpilot Compile Tests
    runs-on: ubuntu-22.04
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -477,26 +377,26 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=33 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
-        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
-        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      # - name: Test openpilot simple_plan vision model correctness (float32)
-      #   run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
+      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
-        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot compile4
-        run: PYTHONPATH="." NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
+        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+# ****** ONNX Tests ******
+
  testonnxcpu:
-    name: 'ONNX (CPU) Tests'
+    name: ONNX (CPU) Tests
    runs-on: ubuntu-22.04
    timeout-minutes: 20
-    env:
-      IGNORE_OOB: 0

    steps:
      - name: Checkout Code
@@ -509,25 +409,22 @@ jobs:
          python-version: '3.11'
          llvm: 'true'
      - name: Test ONNX (CPU)
-        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
-        run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX Runner (CPU)
-        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_runner.py
      - name: Test Additional ONNX Ops (CPU)
-        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
+        run: CPU=1 CPU_LLVM=0 python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
-        run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+        run: CPU=1 CPU_LLVM=0 python3 test/test_quantize_onnx.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

  testopencl:
-    name: 'ONNX (GPU)+Optimization Tests'
+    name: ONNX (CL)+Optimization Tests
    runs-on: ubuntu-22.04
    timeout-minutes: 20
-    env:
-      IGNORE_OOB: 0
-
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -539,18 +436,23 @@ jobs:
          pydeps: "tensorflow==2.15.1 tensorflow_addons"
          python-version: '3.11'
          opencl: 'true'
-      - name: Test ONNX (GPU)
-        run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test Optimization Helpers
-        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
+      - name: Test ONNX (CL)
+        run: CL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      #- name: Test Optimization Helpers
+      #  run: DEBUG=1 python3 extra/optimization/test_helpers.py
      #- name: Test Action Space
-      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
+      #  run: DEBUG=1 CL=1 python3 extra/optimization/get_action_space.py
      - name: Test Beam Search
-        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+        run: CL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
-        run: GPU=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
+        run: CL=1 python -m pytest -n=auto test/external/external_test_optim.py test/external/external_test_losses.py test/external/external_test_metrics.py test/external/external_test_datasets.py --durations=20
+      - name: NULL=1 beautiful_mnist_multigpu
+        run: NULL=1 python examples/beautiful_mnist_multigpu.py
+      - name: Test Bert training
+        run: NULL=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=24 GPUS=4 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
      - name: Test llama 3 training
-        run: MAX_BUFFER_SIZE=0 PYTHONPATH="." DEV=NULL SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
+        # TODO: remove LLAMA_LAYERS once it's fast
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 LLAMA_LAYERS=4 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -568,12 +470,12 @@ jobs:
      - name: Test 1B LLM
        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster

+# ****** Models Tests ******
+
  testmodels:
    name: Models (llvm+cpu+gpu)
    runs-on: ubuntu-22.04
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -585,47 +487,38 @@ jobs:
          opencl: 'true'
          llvm: 'true'
      - name: Test models (llvm)
-        run: LLVM=1 python -m pytest -n=auto test/models --durations=20
-      - name: Test models (gpu)
-        run: GPU=1 python -m pytest -n=auto test/models --durations=20
+        run: CPU=1 CPU_LLVM=1 python -m pytest -n=auto test/models --durations=20
+      - name: Test models (opencl)
+        run: CL=1 python -m pytest -n=auto test/models --durations=20
      - name: Test models (cpu)
-        run: CPU=1 python -m pytest -n=auto test/models --durations=20
+        run: CPU=1 CPU_LLVM=0 python -m pytest -n=auto test/models --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  testrangeify:
-    name: Linux (rangeify)
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
+  testmetalmodels:
+    name: Models (metal)
+    runs-on: macos-14
+    timeout-minutes: 20
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: rangeify-minimal-llvm
-        deps: testing_minimal
-        llvm: "true"
-    - name: Test CPU=1 RANGEIFY=1
-      # TODO: add more passing tests here
-      # test_symbolic_arange_sym_step is passing now
-      # test_threefry_doesnt_use_long is because there's a contig after the long now
-      run: |
-        CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-          -k "not test_symbolic_arange_sym_step and not test_threefry_doesnt_use_long" \
-          test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
-          test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_tensor_data.py
-    - name: Test CPU=1 RANGEIFY=2
-      run: CPU=1 RANGEIFY=2 python3 -m pytest -n auto test/test_tiny.py test/test_rangeify.py test/test_ops.py --durations 20
-    - name: Test LLVM=1 RANGEIFY=1 (slow tests)
-      run: LLVM=1 RANGEIFY=1 python3 -m pytest -n auto test/models/test_mnist.py --durations 20
+        key: metal
+        deps: testing
+        python-version: '3.11'
+    - name: Test models (Metal)
+      run: METAL=1 python -m pytest -n=auto test/models --durations=20
+    - name: Test LLaMA compile speed
+      run: METAL=1 python test/external/external_test_speed_llama.py
+
+# ****** Feature Tests ******

  testdevectorize:
    name: Linux (devectorize)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -637,18 +530,16 @@ jobs:
        pydeps: "pillow"
        llvm: "true"
    - name: Test LLVM=1 DEVECTORIZE=0
-      run: LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
    - name: Test LLVM=1 DEVECTORIZE=0 for model
-      run: PYTHONPATH="." LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
+      run: CPU=1 CPU_LLVM=1 DEVECTORIZE=0 python3 test/models/test_efficientnet.py
    - name: Test CPU=1 DEVECTORIZE=0
-      run: CPU=1 DEVECTORIZE=0 FUSE_ARANGE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"
+      run: CPU=1 CPU_LLVM=0 DEVECTORIZE=0 python3 -m pytest -n auto test/test_tiny.py test/test_ops.py -k "not test_avg_pool3d_failure"

  testdsp:
    name: Linux (DSP)
    runs-on: ubuntu-24.04
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
@@ -675,9 +566,9 @@ jobs:
    - name: Run test_tiny on DSP
      run: DEBUG=2 DSP=1 python test/test_tiny.py
    - name: Test transcendentals
-      run: CC=clang-20 PYTHONPATH="." DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
+      run: CC=clang-20 DEBUG=2 DSP=1 python test/test_transcendental.py TestTranscendentalVectorized
    - name: Test quantize onnx
-      run: PYTHONPATH="." DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py
+      run: DEBUG=2 DSP=1 python3 test/test_quantize_onnx.py

  testwebgpu:
    name: Linux (WebGPU)
@@ -715,12 +606,10 @@ jobs:
    runs-on: ubuntu-22.04
    timeout-minutes: 20
    env:
-      IGNORE_OOB: 0
      AMD: 1
      MOCKGPU: 1
      FORWARD_ONLY: 1
      AMD_LLVM: ${{ matrix.backend == 'amdllvm' && '1' || matrix.backend != 'amdllvm' && '0' }}
-      PYTHONPATH: ${{ github.workspace }}
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -746,7 +635,7 @@ jobs:
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run TestOps.test_add with SQTT
        run: |
-          PROFILE=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
+          VIZ=1 SQTT=1 DEBUG=5 python3 test/test_ops.py TestOps.test_add
          extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
      - name: Run process replay tests
        uses: ./.github/actions/process-replay
@@ -760,7 +649,9 @@ jobs:
    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
    timeout-minutes: 20
-
+    env:
+      MOCKGPU: 1
+      FORWARD_ONLY: 1
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -772,29 +663,26 @@ jobs:
          cuda: 'true'
          ocelot: 'true'
      - name: Set env
-        run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'PTX' && 'CUDA=1\nCUDA_PTX=1' || matrix.backend == 'nv' && 'NV=1\nSKIP_SLOW_TEST=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
-          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CUDA','NV'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
      - name: Run pytest (cuda)
        # skip multitensor because it's slow
        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --ignore test/test_gc.py --ignore test/test_multitensor.py --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

-  tests:
+  testcpuopencl:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, gpu]
+        backend: [llvm, cpu, opencl]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
    timeout-minutes: 20
-    env:
-      IGNORE_OOB: 0
-
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -803,65 +691,124 @@ jobs:
        with:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
-          opencl: ${{ matrix.backend == 'gpu' && 'true' }}
+          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'gpu' && 'GPU=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          PYTHONPATH=${{ github.workspace }} python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['LLVM','CPU','GPU'], Device.DEFAULT"
-          DEBUG=5 PYTHONPATH=${{ github.workspace }} FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-      - name: Run pytest (not cuda)
+          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
+          DEBUG=5 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+      - name: Run pytest (${{ matrix.backend }})
        run: python -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
      - name: Run TRANSCENDENTAL math
        run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

+  amdremote:
+    name: Linux (remote)
+    runs-on: ubuntu-22.04
+    timeout-minutes: 20
+    env:
+      REMOTE: 1
+    steps:
+      - name: Checkout Code
+        uses: actions/checkout@v4
+      - name: Setup Environment
+        uses: ./.github/actions/setup-tinygrad
+        with:
+          key: linux-remote
+          deps: testing_minimal
+          amd: 'true'
+          llvm: 'true'
+          opencl: 'true'
+      - name: Start remote server
+        run: |
+          start_server() {
+            systemd-run --user \
+              --unit="$1" \
+              --setenv=REMOTEDEV="$2" \
+              --setenv=MOCKGPU=1 \
+              --setenv=PYTHONPATH=. \
+              --setenv=PORT="$3" \
+              --working-directory="$(pwd)" \
+              python tinygrad/runtime/ops_remote.py
+          }
+
+          start_server "remote-server-amd-1" "AMD" 6667
+          start_server "remote-server-amd-2" "AMD" 6668
+          start_server "remote-server-gpu" "CL" 7667
+          start_server "remote-server-cpu" "CPU" 8667
+      - name: Check Device.DEFAULT and print some source
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
+          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
+      - name: Run REMOTE=1 Test (AMD)
+        env:
+          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
+      - name: Run REMOTE=1 Test (CL)
+        env:
+          HOST: 127.0.0.1:7667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
+          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
+      - name: Run REMOTE=1 Test (CPU)
+        env:
+          HOST: 127.0.0.1:8667*6
+        run: |
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
+      - name: Show remote server logs
+        if: always()
+        run: |
+          journalctl --user -u remote-server-amd-1 --no-pager
+          journalctl --user -u remote-server-amd-2 --no-pager
+          journalctl --user -u remote-server-gpu --no-pager
+          journalctl --user -u remote-server-cpu --no-pager
+
 # ****** OSX Tests ******

-  testmetal2:
+  testmetal:
    name: MacOS (unit)
    runs-on: macos-14
    timeout-minutes: 20
-    env:
-      IGNORE_OOB: 0
-
    steps:
    - name: Checkout Code
      uses: actions/checkout@v4
    - name: Setup Environment
      uses: ./.github/actions/setup-tinygrad
      with:
-        key: metal2
+        key: metal
        deps: testing
        python-version: '3.11'
        amd: 'true'
        cuda: 'true'
        ocelot: 'true'
        llvm: 'true'
-    - name: Run real world test
-      run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
-    - name: Test models (Metal)
-      run: METAL=1 python -m pytest -n=auto test/models -v --durations=20
+    - name: Run unit tests
+      run: METAL=1 python -m pytest -n=auto test/unit/ --durations=20
    - name: Run ONNX
      run: METAL=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
    - name: Test tensor core ops (fake)
-      run: TC=2 METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_gemm
+      run: METAL=1 DEBUG=3 TC=2 python test/test_ops.py TestOps.test_gemm
    - name: Test tensor core ops (real)
      run: METAL=1 DEBUG=3 python test/test_ops.py TestOps.test_big_gemm
-    - name: Test LLaMA compile speed
-      run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
    - name: Test Beam Search
-      run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
+      run: METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
    #- name: Fuzz Test linearizer
-    #  run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #  run: METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
    - name: Run TRANSCENDENTAL math
-      run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
+      run: METAL=1 TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
    - name: Run pytest (amd)
      env:
        MOCKGPU: 1
        AMD: 1
+        AMD_LLVM: 0
        FORWARD_ONLY: 1
      run: |
        python3 -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py --durations=20
@@ -869,13 +816,14 @@ jobs:
      env:
        MOCKGPU: 1
        AMD: 1
+        AMD_LLVM: 1
        FORWARD_ONLY: 1
      run: |
        python -m pytest -n=auto test/device/test_hcq.py test/test_tiny.py test/device/test_amd_llvm.py --durations=20
    - name: Run pytest (ptx)
      env:
        MOCKGPU: 1
-        PTX: 1
+        NV_PTX: 1
        NV: 1
        FORWARD_ONLY: 1
      run: |
@@ -918,7 +866,7 @@ jobs:
    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
    #    node test_viz.js
    - name: Test ONNX Runner (WEBGPU)
-      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+      run: WEBGPU=1 python3 test/external/external_test_onnx_runner.py

  osxremote:
   name: MacOS (remote metal)
@@ -944,72 +892,6 @@ jobs:
       run: |
         python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py

-  amdremote:
-    name: Linux (remote)
-    runs-on: ubuntu-22.04
-    timeout-minutes: 20
-    env:
-      REMOTE: 1
-      PYTHONPATH: ${{ github.workspace }}
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: linux-remote
-          deps: testing_minimal
-          amd: 'true'
-          llvm: 'true'
-          opencl: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "GPU" 7667
-          start_server "remote-server-cpu" "CPU" 8667
-      - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
-          python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
-          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py --durations 20
-      - name: Run REMOTE=1 Test (GPU)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py --durations 20
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py --durations 20
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
-
  osxtests:
    strategy:
      fail-fast: false
@@ -1018,8 +900,6 @@ jobs:
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
-    env:
-      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -1031,10 +911,10 @@ jobs:
          pydeps: "capstone"
          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -1055,8 +935,6 @@ jobs:
    name: Windows (${{ matrix.backend }})
    runs-on: windows-latest
    timeout-minutes: 15
-    env:
-      IGNORE_OOB: 0
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -1068,12 +946,13 @@ jobs:
          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
      - name: Set env
        shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
      - name: Run unit tests
        if: matrix.backend=='llvm'
-        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py
+        # test_newton_schulz hits RecursionError
+        run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py --ignore=test/unit/test_linalg.py --durations=20
      - name: Run pytest (${{ matrix.backend }})
        shell: bash
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == '${{ matrix.backend }}'.upper(), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          python -m pytest -n=auto test/test_tiny.py test/test_ops.py --durations=20
@@ -20,12 +20,6 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
-      - id: devicetests
-        name: select GPU tests
-        entry: env GPU=1 PYTHONPATH="." python3 -m pytest test/test_uops.py test/test_search.py
-        language: system
-        always_run: true
-        pass_filenames: false
      - id: tests
        name: subset of tests
        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
@@ -54,11 +54,12 @@ 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
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method,W0707
 # 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,9 +79,8 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers

 tinygrad already supports numerous accelerators, including:

- [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
- [x] [CPU (C Code)](tinygrad/runtime/ops_cpu.py)
- [x] [LLVM](tinygrad/runtime/ops_llvm.py)
+- [x] [OpenCL](tinygrad/runtime/ops_cl.py)
+- [x] [CPU](tinygrad/runtime/ops_cpu.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
@@ -42,7 +42,6 @@ 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()
@@ -51,13 +50,14 @@ 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.view(ShapeTracker.from_shape((1,))),))
-ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.view(ShapeTracker.from_shape((1,))),))
+ld_1 = UOp(Ops.LOAD, dtypes.int32, (buf_1.index(idx),))
+ld_2 = UOp(Ops.LOAD, dtypes.int32, (buf_2.index(idx),))
 alu = ld_1 + ld_2
 output_buf = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 0)
-st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,))), alu))
+st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.index(idx), alu))
 s = UOp(Ops.SINK, dtypes.void, (st_0,))

 # convert the computation to a "linearized" format (print the format)
@@ -80,7 +80,8 @@ print("******** third, the UOp ***********")

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

 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +94,10 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))

 # group the computation into kernels
-becomes_map = get_kernelize_map(s)
+becomes_map = get_rangeify_map(s)

 # the compute maps to an assign
-assign = becomes_map[a+b]
+assign = becomes_map[a+b].base

 # 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.kernelize.get_kernelize_map
+::: tinygrad.schedule.rangeify.get_rangeify_map
    options:
        members: false
        show_labels: false
@@ -22,12 +22,6 @@ Group UOps into kernels.

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

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

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

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

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

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

 ## Interoperability

@@ -70,5 +70,12 @@ 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)
@@ -2,7 +2,6 @@ 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
@@ -11,7 +10,7 @@ GPUS = [f'{Device.DEFAULT}:{i}' for i in range(getenv("GPUS", 1))]

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

 # from https://github.com/tysam-code/hlb-CIFAR10/blob/main/main.py
 batchsize = getenv("BS", 1024)
@@ -150,13 +149,12 @@ 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()

-  np.random.seed(1337)
+  Tensor.manual_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()
-    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
+    tidxs = Tensor.randperm(num_train_samples, dtype='int')[:num_steps_per_epoch*batchsize].reshape(num_steps_per_epoch, batchsize)
    train_loss:float = 0
    for epoch_step in (t:=trange(num_steps_per_epoch)):
      st = time.perf_counter()
@@ -1,6 +1,6 @@
 import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.nn.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.frontend.onnx import OnnxRunner
+from tinygrad.nn.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)
-    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)]
+    xqkv = self.c_attn(x).reshape(None, None, 3, self.n_heads, self.head_dim)
+    xq, xk, xv = [xqkv[:, :, i, :, :] for i in range(3)]
    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.shrink((None, None,(start_pos,start_pos+seqlen),None,None)).assign(Tensor.stack(xk, xv)).realize()
+    self.cache_kv[:, :, start_pos:start_pos+seqlen, :, :].assign(Tensor.stack(xk, xv)).realize()

    if start_pos > 0:
-      keys = self.cache_kv[0].shrink((None, (0, start_pos+seqlen), None, None))
-      values = self.cache_kv[1].shrink((None, (0, start_pos+seqlen), None, None))
+      keys = self.cache_kv[0][:, :start_pos+seqlen, :, :]
+      values = self.cache_kv[1][:, :start_pos+seqlen, :, :]
    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)))
+    return (h + self.mlp(self.ln_2(h))).contiguous()

 class Transformer:
  def __init__(self, dim, n_heads, n_layers, norm_eps, vocab_size, max_seq_len=1024):
@@ -181,6 +181,7 @@ 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
@@ -188,7 +189,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 GPU" if DEBUG>=2 else "")+
+      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 "")+
                  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):
@@ -197,8 +198,13 @@ 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 ****
@@ -145,7 +145,6 @@ hyp = {
  },
 }

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

  def set_seed(seed):
@@ -229,7 +228,8 @@ def train_cifar():
    if getenv("RANDOM_CROP", 1):
      X = random_crop(X, crop_size=32)
    if getenv("RANDOM_FLIP", 1):
-      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
+      # NOTE: RANGEIFY=1 needs this contiguous or the X[perms] is very slow
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X).contiguous() # flip LR
    X, 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,12 +413,17 @@ 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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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 GPU" if DEBUG>=2 else "")+
+            with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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,9 +279,15 @@ 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:
-      tok_Tens = (logits/temp).softmax().multinomial()
+      scaled_logits = logits / temp
+      if top_k is not None:
+        topk_values, topk_indices = scaled_logits.topk(top_k)
+        filtered_logits = Tensor.full_like(scaled_logits, -float("inf"))
+        filtered_logits = filtered_logits.scatter(dim=-1, index=topk_indices, src=topk_values)
+        tok_Tens = filtered_logits.softmax().multinomial()
+      else:
+        tok_Tens = scaled_logits.softmax().multinomial()
    else:
      tok_Tens = logits.argmax(axis=-1).unsqueeze(0)
    tok = tok_Tens.item()
@@ -298,6 +304,7 @@ 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()
@@ -308,8 +315,9 @@ 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)
+  tinyoutput = generate(model, tokenizer, prompt, n_tokens_to_gen=num_toks, sample=sample, temp=temp, top_k=top_k)
  print(tinyoutput)
  print('TIME: ', time.time() - s)
  TORCHOUTPUT = "Why is gravity \nso important?\nBecause it's the only"
@@ -511,6 +511,33 @@ 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:
@@ -758,6 +785,27 @@ def batch_load_llama3(bs:int, samples:int, seqlen:int, base_dir:Path, seed:int=0
      batch.append(tokens)
    yield Tensor.stack(batch, dim=0)

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

 from tinygrad import Tensor, nn, dtypes
-from tinygrad.helpers import prod, argfix
+from tinygrad.helpers import prod, argfix, Context
+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:
@@ -17,6 +19,10 @@ 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)
@@ -127,3 +133,59 @@ 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,8 +1,9 @@
 import math
-from tinygrad import dtypes
+from tinygrad import dtypes, Tensor
 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):
@@ -36,4 +37,24 @@ 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
+import time, math, os
 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
+from tinygrad.helpers import getenv, Context, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
 def tlog(x): print(f"{x:25s}  @ {time.perf_counter()-start:5.2f}s")

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

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

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

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

  @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()
+    return loss.flatten().float()

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

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

 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, FUSE_CONV_BW, Profiling
-from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, Profiling
+from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW

 from extra.lr_scheduler import LRSchedulerGroup
@@ -252,6 +252,10 @@ 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:
@@ -344,6 +348,8 @@ 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
@@ -701,7 +707,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 FUSE_CONV_BW=1 GPUS=6 BS=6 MODEL=unet3d python3 examples/mlperf/model_train.py```
+  ```time PYTHONPATH=. WANDB=1 TRAIN_BEAM=3 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
@@ -743,7 +749,6 @@ 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
  }
@@ -1290,18 +1295,20 @@ def train_llama3():
  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup

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

-  # LR=1e-4 TRAIN_ON_VAL=1 DEFAULT_FLOAT=bfloat16 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
+  # 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
  # trains to 7

  opt_adamw_beta_1 = 0.9
@@ -1311,13 +1318,14 @@ 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("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
+  opt_learning_rate_decay_steps = getenv("MAX_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 = 8e-7
+  opt_end_learning_rate = getenv("END_LR", 8e-7)

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

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

+  if resume_ckpt := getenv("RESUME_CKPT"):
+    fn = f"./ckpts/llama3_{resume_ckpt}.safe"
+    print(f"loading initial checkpoint from {fn}")
+    load_state_dict(model, safe_load(fn), realize=False)
+
+    fn = f"./ckpts/llama3_{resume_ckpt}_optim.safe"
+    print(f"loading optim checkpoint from {fn}")
+    load_state_dict(scheduler, safe_load(fn), realize=False)
+
  @TinyJit
  @Tensor.train()
  def train_step(model, tokens:Tensor, grad_acc:int):
@@ -1403,43 +1420,55 @@ def train_llama3():
  # ** data iters **
  def fake_data(bs, samples):
    for _ in range(samples // bs):
-      yield Tensor.randint(bs, SEQLEN + 1, low=0, high=32000, dtype=dtypes.int32, device=Device.DEFAULT)
+      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:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(GBS, SAMPLES, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=bool(TRAIN_ON_VAL))
+      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:
-      from examples.mlperf.dataloader import batch_load_llama3
-      return batch_load_llama3(EVAL_BS, 5760, SEQLEN, Path(getenv("BASEDIR", "/raid/datasets/c4/")), seed=SEED, val=True)
+      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 = 0, 0
+  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
+
+    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")

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

-    i += 1
-    sequences_seen += tokens.shape[0]
+      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")
@@ -1463,6 +1492,144 @@ def train_llama3():
          safe_save(get_state_dict(model), fn)
        break

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

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

  with Tensor.train():
-    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn").split(","):
+    for m in getenv("MODEL", "resnet,retinanet,unet3d,rnnt,bert,maskrcnn,stable_diffusion").split(","):
      nm = f"train_{m}"
      if nm in globals():
        print(f"training {m}")
@@ -0,0 +1,57 @@
+#!/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
@@ -0,0 +1,72 @@
+#!/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
@@ -10,7 +10,7 @@ from tinygrad.helpers import DEBUG, getenv
 from tinygrad.engine.realize import CompiledRunner

 import onnx
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.nn.onnx import OnnxRunner

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -77,13 +77,20 @@ def test_vs_compile(run, new_inputs, test_val=None):
            **{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()
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
+    step_times.append((et-st)*1e3)
+    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
+
+  if (assert_time:=getenv("ASSERT_MIN_STEP_TIME")):
+    min_time = min(step_times)
+    assert min_time < assert_time, f"Speed regression, expected min step time of < {assert_time} ms but took: {min_time} ms"
+
  print(out, val.shape, val.dtype)
  if test_val is not None: np.testing.assert_equal(test_val, val)
  print("**** test done ****")
@@ -1,12 +1,12 @@
 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.nn.onnx import OnnxRunner
+from tinygrad.schedule.rangeify import get_rangeify_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule

-# NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
+# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py

 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
@@ -33,7 +33,7 @@ if __name__ == "__main__":
        if not in_target_path[s]:
          independent_set[s] = None
  independent = UOp.sink(*independent_set.keys())
-  kernelized = get_kernelize_map(independent)
+  kernelized = get_rangeify_map(independent)
  independent = independent.substitute(kernelized)
  schedule, var_vals = create_schedule_with_vars(independent)
  run_schedule(schedule)
@@ -27,7 +27,7 @@ class Model(nn.Module):

 if __name__ == "__main__":
  if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch  # noqa: F401
+    import tinygrad.nn.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
+from tinygrad import nn, Tensor, Device
 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 GPU" if DEBUG>=2 else "") +
+      with Timing("enqueue in ", on_exit=(lambda et: (f", {(GlobalCounters.time_sum_s-st)*1e3:.2f} ms on {Device.DEFAULT}" 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)
@@ -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
+from tinygrad.helpers import fetch, trange, colored, Timing, getenv
 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
+import argparse, tempfile, time
 from abc import ABC, abstractmethod
 from pathlib import Path
 from PIL import Image
@@ -342,11 +342,13 @@ 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,
@@ -358,8 +360,13 @@ 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


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

  # validation!
-  if args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024 \
-    and not args.weights:
+  is_default = args.prompt == default_prompt and args.steps == 10 and args.seed == 0 and args.guidance == 6.0 and args.width == args.height == 1024
+  if is_default and not args.weights and not args.fakeweights:
    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,18 +2,20 @@
 # https://github.com/ekagra-ranjan/huggingface-blog/blob/main/stable_diffusion.md
 import tempfile
 from pathlib import Path
-import argparse
+import argparse, time
 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
+from tinygrad.helpers import Timing, Context, getenv, fetch, colored, tqdm, flatten
 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
+from extra.models.clip import Closed, Tokenizer, FrozenOpenClipEmbedder
+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:
@@ -154,12 +156,46 @@ 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):
+  def __init__(self, version:str|None=None, pretrained:str|None=None):
    self.alphas_cumprod = get_alphas_cumprod()
-    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_params))
-    self.first_stage_model = AutoencoderKL()
-    self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
+    if version != "v2-mlperf-train":
+      self.first_stage_model = AutoencoderKL() # only needed for decoding generated latents to images; not needed in mlperf training from preprocessed moments
+
+    if not version:
+      self.cond_stage_model = namedtuple("CondStageModel", ["transformer"])(transformer = namedtuple("Transformer", ["text_model"])(text_model = Closed.ClipTextTransformer()))
+      unet_init_params = unet_params
+    elif version in {"v2-mlperf-train", "v2-mlperf-eval"}:
+      unet_init_params = mlperf_params
+      clip.gelu = gelu_erf
+      self.cond_stage_model = FrozenOpenClipEmbedder(**{"dims": 1024, "n_heads": 16, "layers": 24, "return_pooled": False, "ln_penultimate": True,
+                                                        "clip_tokenizer_version": "sd_mlperf_v5_0"})
+      unet.Linear, unet.Conv2d, unet.GroupNorm, unet.LayerNorm = AutocastLinear, AutocastConv2d, AutocastGroupNorm, AutocastLayerNorm
+      unet.attention, unet.gelu, unet.mixed_precision_dtype = attn_f32_softmax, gelu_erf, dtypes.bfloat16
+      if pretrained:
+        print("loading text encoder")
+        weights: dict[str,Tensor] = {k.replace("cond_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("cond_stage_model.")}
+        weights["model.attn_mask"] = Tensor.full((77, 77), fill_value=float("-inf")).triu(1)
+        load_state_dict(self.cond_stage_model, weights)
+        # only the eval model needs the decoder
+        if version == "v2-mlperf-eval":
+          print("loading image latent encoder")
+          weights = {k.replace("first_stage_model.", "", 1):v for k,v in torch_load(pretrained)["state_dict"].items() if k.startswith("first_stage_model.")}
+          load_state_dict(self.first_stage_model, weights)
+
+    self.model = namedtuple("DiffusionModel", ["diffusion_model"])(diffusion_model = UNetModel(**unet_init_params))
+    if version == "v2-mlperf-train":
+      # the mlperf reference inits certain weights as zeroes
+      for bb in flatten(self.model.diffusion_model.input_blocks) + self.model.diffusion_model.middle_block + flatten(self.model.diffusion_model.output_blocks):
+        if isinstance(bb, unet.ResBlock):
+          zero_module(bb.out_layers[3])
+        elif isinstance(bb, unet.SpatialTransformer):
+          zero_module(bb.proj_out)
+      zero_module(self.model.diffusion_model.out[2])

  def get_x_prev_and_pred_x0(self, x, e_t, a_t, a_prev):
    temperature = 1
@@ -233,12 +269,14 @@ if __name__ == "__main__":

  # load in weights
  with WallTimeEvent(BenchEvent.LOAD_WEIGHTS):
-    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)
+    load_state_dict(model, torch_load(fetch('https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt', 'sd-v1-4.ckpt'))['state_dict'], verbose=False, strict=False, realize=False)

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

  # run through CLIP to get context
  tokenizer = Tokenizer.ClipTokenizer()
@@ -266,17 +304,23 @@ 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)
@@ -32,7 +32,7 @@ if __name__ == "__main__":

  lr = 5e-3
  transform = ComposeTransforms([
-    lambda x: [Image.fromarray(xx, mode='L').resize((64, 64)) for xx in x],
+    lambda x: [Image.fromarray(xx).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, None))
+    x = self.token_embedding(x) + self.positional_embedding.shrink(((pos, pos+seqlen), None))
    for block in self.blocks: x = block(x, xa=encoded_audio, mask=self.mask, len=pos)
    return self.output_tok(x)

@@ -2,7 +2,7 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.nn.onnx import OnnxRunner
 from extra.onnx_helpers import get_example_inputs

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

  with open(sys.argv[1], 'r') as f:
    log_content = log_content_them = f.read()
@@ -1,7 +1,7 @@
 # copying the kernels from https://github.com/microsoft/ArchProbe into Python
 import numpy as np
 import pickle
-from tinygrad.runtime.ops_gpu import CLProgram, CLBuffer
+from tinygrad.runtime.ops_cl import CLProgram, CLBuffer
 from tinygrad import dtypes
 from tqdm import trange, tqdm
 from matplotlib import pyplot as plt
@@ -4,7 +4,7 @@ from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
 from tinygrad.codegen.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
+from tinygrad.runtime.ops_cl 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_gpu import CLProgram, CLBuffer, ROCM_LLVM_PATH
+from tinygrad.runtime.ops_cl 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[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)
    if not all(isinstance(ji.prg, CompiledRunner) for ji in jit_cache): raise GraphException

    prgs = '\n'.join(dedup([cast(CompiledRunner, ji.prg).p.src for ji in jit_cache]))
    args = [f"{render_dtype(x.dtype)}* arg{i}" for i,x in enumerate(input_rawbuffers)]
-    args += sorted([f"int {v.expr}" for v in var_vals])
+    args += sorted([f"int {v}" 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[Variable, int], wait=False):
+  def __call__(self, rawbufs: List[Buffer], var_vals: Dict[str, int], wait=False):
    return cpu_time_execution(
-    lambda: self._prg(*[x._buf for x in rawbufs], *[x[1] for x in sorted(var_vals.items(), key=lambda x: x[0].expr)]), 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])]), 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[Variable, int]):
+  def __init__(self, jit_cache: List[ExecItem], input_rawbuffers: List[Buffer], var_vals: Dict[str, int]):
    super().__init__(jit_cache, input_rawbuffers, var_vals)

    # Check all jit items are compatible.
@@ -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]])
+      for i in range(len(ji.prg.p.vars)): self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[ji.prg.p.vars[i].expr])

    # Build queues.
    self.virt_aql_queues: Dict[Compiled, VirtAQLQueue] = {dev:VirtAQLQueue(dev, 2*len(self.jit_cache)+16) for dev in self.devices}
@@ -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[Variable, int], wait=False) -> Optional[float]:
+  def __call__(self, input_rawbuffers: List[Buffer], var_vals: Dict[str, int], wait=False) -> Optional[float]:
    # Wait and restore signals
    hsa.hsa_signal_wait_scacquire(self.finish_signal, hsa.HSA_SIGNAL_CONDITION_LT, 1, (1 << 64) - 1, hsa.HSA_WAIT_STATE_ACTIVE)
    for sig in self.signals_to_reset: hsa.hsa_signal_silent_store_relaxed(sig, 1)
@@ -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])
+        self.ji_kargs_structs[j].__setattr__(f'v{i}', var_vals[v.expr])

    # Update launch dims
    for j in self.jc_idx_with_updatable_launch_dims:
@@ -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[1].startswith("lidx"):
-        r[u] = f'v{u.arg[0]}'
-      elif u.arg[1].startswith("gidx"):
-        r[u] = f's{2+u.arg[0]}'
+      if u.arg.startswith("lidx"):
+        r[u] = f'v{u.src[0].arg}'
+      elif u.arg.startswith("gidx"):
+        r[u] = f's{2+u.src[0].arg}'
      else:
        raise NotImplementedError
    elif u.uop == UOps.CONST:
@@ -10,7 +10,7 @@ from tinygrad.uop.ops import Ops
 import json
 from collections import OrderedDict

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

 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
  functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
@@ -67,11 +67,12 @@ 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)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
+    cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)}, {thread_id});" for (name, args, _global_size, _local_size) in statements] + ["}"]
    return '\n'.join(headers + cprog)
  else:
    if bufs_to_save:
@@ -239,7 +240,9 @@ 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()}
@@ -65,7 +65,7 @@ def top_spec_kernel3():
  c = a@b
  sink = c.schedule()[-1].ast
  L = 16
-  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(dtypes.int, N//BM, 0), 2:UOp.range(dtypes.int, N//BN, 1)})
+  sink = sink.reshape((N//L, L, N//L, L)) #.lift({0:UOp.range(N//BM, 0), 2:UOp.range(N//BN, 1)})
  sink = graph_rewrite(sink, view_left+pm)
  axis_types = (AxisType.GLOBAL, AxisType.LOCAL, AxisType.GLOBAL, AxisType.LOCAL, AxisType.REDUCE)
  return sink.replace(arg=KernelInfo(name="top_"+to_colored(sink.full_shape, axis_types), axis_types=axis_types))
@@ -186,7 +186,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(dtypes.int, c_regs.dtype.size, 16)
+  i = UOp.range(c_regs.dtype.size, 16)
  init_store = c_regs[i].store(UOp.const(dtypes.float, 0.0), i)

  if kernel4:
@@ -197,53 +197,53 @@ def hand_spec_kernel3(kernel4=getenv("K4", 0), kernel5=getenv("K5", 0)):
    kId = 0

    # load from globals into locals
-    i = UOp.range(dtypes.int, nbReadsB, 0)
+    i = UOp.range(nbReadsB, 0)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(dtypes.int, nbReadsA, 1)
+    i = UOp.range(nbReadsA, 1)
    index_x = rAIdx + kId
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)

    # iterate over the middle chunk
-    kId_range = UOp.range(dtypes.int, N//BK-1, 2)
+    kId_range = UOp.range(N//BK-1, 2)
    kId = kId_range*BK

    barrier = UOp.barrier(As_store, Bs_store)

    # load from globals into registers (next round)
-    i = UOp.range(dtypes.int, nbReadsB, 3)
+    i = UOp.range(nbReadsB, 3)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId + BK
    regB_store = regB[i].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(dtypes.int, nbReadsA, 4)
+    i = UOp.range(nbReadsA, 4)
    index_x = rAIdx + kId + BK
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    regA_store = regA[i].store(a[N * index_y + index_x].load(), i)

    def inner_loop(first_range, inp_dep=()):
      # inner unroll
-      k = UOp.range(dtypes.int, BK, first_range+0)
+      k = UOp.range(BK, first_range+0)

      # load from locals into registers
-      iterWave = UOp.range(dtypes.int, nbIterWaveN, first_range+1)
-      i = UOp.range(dtypes.int, TN, first_range+2)
+      iterWave = UOp.range(nbIterWaveN, first_range+1)
+      i = UOp.range(TN, first_range+2)
      index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
      B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(*inp_dep), iterWave, i)

-      iterWave = UOp.range(dtypes.int, nbIterWaveM, first_range+3)
-      i = UOp.range(dtypes.int, TM, first_range+4)
+      iterWave = UOp.range(nbIterWaveM, first_range+3)
+      i = UOp.range(TM, first_range+4)
      index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
      A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(*inp_dep), iterWave, i)

      # do the GEMM math
-      iterWaveM = UOp.range(dtypes.int, nbIterWaveM, first_range+5)
-      yt = UOp.range(dtypes.int, TM, first_range+6)
-      iterWaveN = UOp.range(dtypes.int, nbIterWaveN, first_range+7)
-      xt = UOp.range(dtypes.int, TN, first_range+8)
+      iterWaveM = UOp.range(nbIterWaveM, first_range+5)
+      yt = UOp.range(TM, first_range+6)
+      iterWaveN = UOp.range(nbIterWaveN, first_range+7)
+      xt = UOp.range(TN, first_range+8)
      x = iterWaveN * TN + xt
      y = iterWaveM * TM + yt
      c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -256,12 +256,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(dtypes.int, nbReadsB, 14)
+    i = UOp.range(nbReadsB, 14)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId + BK
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(regB[i].load(sink), i, kId_range)

-    i = UOp.range(dtypes.int, nbReadsA, 15)
+    i = UOp.range(nbReadsA, 15)
    index_x = rAIdx + kId + BK
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(regA[i].load(sink), i, kId_range)
@@ -269,40 +269,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(dtypes.int, N//BK, 0)
+    kId_range = UOp.range(N//BK, 0)
    kId = kId_range*BK

    # load from globals into locals
-    i = UOp.range(dtypes.int, nbReadsB, 1)
+    i = UOp.range(nbReadsB, 1)
    index_x = BN * blockIdx_x + rBIdx
    index_y = rBIdy + i * strideReadB + kId
    Bs_store = Bs[(index_y % BK) * BN + index_x % BN].store(b[N * index_y + index_x].load(), i)

-    i = UOp.range(dtypes.int, nbReadsA, 2)
+    i = UOp.range(nbReadsA, 2)
    index_x = rAIdx + kId
    index_y = BM * blockIdx_y + rAIdy + i * strideReadA
    As_store = As[(index_x % BK) * BM_As_stride + index_y % BM].store(a[N * index_y + index_x].load(), i)

    barrier = UOp.barrier(As_store, Bs_store)

-    k = UOp.range(dtypes.int, BK, 3)
+    k = UOp.range(BK, 3)

    # load from locals into registers
-    iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-    i = UOp.range(dtypes.int, TN, 5)
+    iterWave = UOp.range(nbIterWaveN, 4)
+    i = UOp.range(TN, 5)
    index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
    B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)

-    iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
-    i = UOp.range(dtypes.int, TM, 7)
+    iterWave = UOp.range(nbIterWaveM, 6)
+    i = UOp.range(TM, 7)
    index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
    A_col_store = A_col[iterWave*TM + i].store(As[k*BM_As_stride + index].load(barrier), iterWave, i)

    # do the GEMM math
-    iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-    yt = UOp.range(dtypes.int, TM, 9)
-    iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 10)
-    xt = UOp.range(dtypes.int, TN, 12)
+    iterWaveM = UOp.range(nbIterWaveM, 8)
+    yt = UOp.range(TM, 9)
+    iterWaveN = UOp.range(nbIterWaveN, 10)
+    xt = UOp.range(TN, 12)
    x = iterWaveN * TN + xt
    y = iterWaveM * TM + yt
    c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
@@ -310,10 +310,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(dtypes.int, nbIterWaveM, 1000)
-  yt = UOp.range(dtypes.int, TM, 1001)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 1002)
-  xt = UOp.range(dtypes.int, TN, 1003)
+  iterWaveM = UOp.range(nbIterWaveM, 1000)
+  yt = UOp.range(TM, 1001)
+  iterWaveN = UOp.range(nbIterWaveN, 1002)
+  xt = UOp.range(TN, 1003)
  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
  indexC = N * (yOut + yt) + xOut + xt
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 import numpy as np
-from tinygrad.runtime.ops_gpu import CLProgram, CLCompiler
+from tinygrad.runtime.ops_cl import CLProgram, CLCompiler
 from tinygrad import Device, dtypes
 from tinygrad.device import Buffer
 from hexdump import hexdump
@@ -11,7 +11,7 @@ from hexdump import hexdump
 # https://registry.khronos.org/OpenCL/extensions/intel/cl_intel_subgroup_split_matrix_multiply_accumulate.html
 # https://hc34.hotchips.org/assets/program/conference/day1/GPU%20HPC/Intel_s%20Ponte%20Vecchio%20GPU%20-%20Architecture%20Systems%20and%20Software%20FINAL.pdf

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

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

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

 row = np.array([1,2,3,4,5,6,7,8,1,2,3,4,5,6,7,8], np.float16)
 mat = np.random.random((8, 0x10)).astype(np.float16)
@@ -56,7 +56,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
  k = Kernel(ast)
  k.apply_opts(opts)
-  p = get_program(k.get_optimized_ast(), k.opts)
+  p = get_program(k.ast, k.opts, k.applied_opts)
  return CUDAProgram(device, p.function_name, compiler.compile(p.src))

 if __name__ == "__main__":
@@ -75,7 +75,7 @@ if __name__ == "__main__":

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

      # WMMA element size is (M, N, K) = (16, 8, 16)
@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.opt.kernel import OptOps
+from tinygrad.codegen.opt import OptOps
 from tinygrad.engine.realize import lower_schedule

 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
@@ -29,7 +29,7 @@ if __name__ == "__main__":
      Opt(op=OptOps.LOCAL, axis=0, amt=2),
    ]
  k.apply_opts(opts)
-  prg = get_program(k.get_optimized_ast(), k.opts)
+  prg = get_program(k.ast, k.opts, k.applied_opts)
  new_src = prg.src
  # can mod source here
  prg = replace(prg, src=new_src)
@@ -43,7 +43,7 @@ def matmul_kernel(c_ptr, a_ptr, b_ptr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N:
  c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
  tl.store(c_ptrs, c)

-# CUDA=1 PTX=1 python3 extra/gemm/triton_nv_matmul.py
+# CUDA=1 CUDA_PTX=1 python3 extra/gemm/triton_nv_matmul.py
 if __name__ == "__main__":
  BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
  M, N, K = 4096, 4096, 4096
@@ -7,7 +7,6 @@ bert_train_params = {
  "GPUS": 6,
  "BS": 96,
  "EVAL_BS": 96,
-  "FUSE_ARANGE": 1,
  "BASEDIR": "/raid/datasets/wiki",
 }

@@ -50,7 +50,7 @@ def ioctls_from_header():
  hdr = (pathlib.Path(__file__).parent / "kfd_ioctl.h").read_text().replace("\\\n", "")
  pattern = r'#define\s+(AMDKFD_IOC_[A-Z0-9_]+)\s+AMDKFD_IOW?R?\((0x[0-9a-fA-F]+),\s+struct\s([A-Za-z0-9_]+)\)'
  matches = re.findall(pattern, hdr, re.MULTILINE)
-  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname)) for name, nr, sname in matches}
+  return {int(nr, 0x10):(name, getattr(kfd_ioctl, "struct_"+sname, None)) for name, nr, sname in matches}
 nrs = ioctls_from_header()

@ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_int, ctypes.c_ulong, ctypes.c_void_p)
@@ -1,7 +1,7 @@
 import onnx, yaml, tempfile, time, argparse, json
 from pathlib import Path
 from typing import Any
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.nn.onnx import OnnxRunner
 from extra.onnx_helpers import validate, get_example_inputs
 from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry

@@ -88,7 +88,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
    return ret

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals = {k:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals = {k.expr:(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  dev = Device[lin.opts.device]
  root = MCTSNode(lin)

@@ -9,6 +9,9 @@ from PIL import Image
 import numpy as np
 import re, gzip

+# Allow for monkeypatching for mlperf.
+gelu = Tensor.gelu
+
@lru_cache()
 def default_bpe():
  # Clip tokenizer, taken from https://github.com/openai/CLIP/blob/main/clip/simple_tokenizer.py (MIT license)
@@ -53,8 +56,8 @@ class Tokenizer:
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))
  class ClipTokenizer:
-    def __init__(self):
-      self.byte_encoder = Tokenizer.bytes_to_unicode()
+    def __init__(self, version=None):
+      self.byte_encoder, self.version = Tokenizer.bytes_to_unicode(), version
      merges = gzip.open(default_bpe()).read().decode("utf-8").split('\n')
      merges = merges[1:49152-256-2+1]
      merges = [tuple(merge.split()) for merge in merges]
@@ -62,11 +65,17 @@ class Tokenizer:
      vocab = vocab + [v+'</w>' for v in vocab]
      for merge in merges:
        vocab.append(''.join(merge))
-      vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+      if self.version == "sd_mlperf_v5_0":
+        import regex
+        vocab.extend(['<start_of_text>', '<end_of_text>'])
+        self.cache = {'<start_of_text>': '<start_of_text>', '<end_of_text>': '<end_of_text>'}
+        self.pat = regex.compile(r"""<start_of_text>|<end_of_text>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", regex.IGNORECASE)
+      else:
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)
      self.encoder = dict(zip(vocab, range(len(vocab))))
      self.bpe_ranks = dict(zip(merges, range(len(merges))))
-      self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
-      self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[^\s]+""", re.IGNORECASE)

    def bpe(self, token):
      if token in self.cache:
@@ -110,8 +119,17 @@ class Tokenizer:

    def encode(self, text:str, pad_with_zeros:bool=False) -> List[int]:
      bpe_tokens: List[int] = []
-      text = Tokenizer.whitespace_clean(text.strip()).lower()
-      for token in re.findall(self.pat, text):
+      if self.version == "sd_mlperf_v5_0":
+        import regex, ftfy, html
+        text = ftfy.fix_text(text)
+        text = html.unescape(html.unescape(text)).strip()
+        text = Tokenizer.whitespace_clean(text).lower()
+        re_module = regex
+      else:
+        text = Tokenizer.whitespace_clean(text.strip()).lower()
+        re_module = re
+
+      for token in re_module.findall(self.pat, text):
        token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
        bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
      # Truncation, keeping two slots for start and end tokens.
@@ -252,10 +270,8 @@ class Open:
      q,k,v = [y.reshape(T, B*self.n_heads, self.d_head).transpose(0, 1).reshape(B, self.n_heads, T, self.d_head) for y in proj.chunk(3)]

      attn_output = Tensor.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
-      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T*B, C)
-
+      attn_output = attn_output.permute(2, 0, 1, 3).reshape(T, B, C)
      attn_output = self.out_proj(attn_output)
-      attn_output = attn_output.reshape(T, B, C)

      return attn_output

@@ -263,9 +279,10 @@ class Open:
    def __init__(self, dims, hidden_dims):
      self.c_fc   = Linear(dims, hidden_dims)
      self.c_proj = Linear(hidden_dims, dims)
+      self.gelu = gelu

    def __call__(self, x:Tensor) -> Tensor:
-      return x.sequential([self.c_fc, Tensor.gelu, self.c_proj])
+      return x.sequential([self.c_fc, self.gelu, self.c_proj])

  # https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/src/open_clip/transformer.py#L210
  class ResidualAttentionBlock:
@@ -350,15 +367,15 @@ class Open:
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L396
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/encoders/modules.py#L498
 class FrozenOpenClipEmbedder(Embedder):
-  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False):
-    self.tokenizer = Tokenizer.ClipTokenizer()
+  def __init__(self, dims:int, n_heads:int, layers:int, return_pooled:bool, ln_penultimate:bool=False, clip_tokenizer_version=None):
+    self.tokenizer = Tokenizer.ClipTokenizer(version=clip_tokenizer_version)
    self.model = Open.ClipTextTransformer(dims, n_heads, layers)
    self.return_pooled = return_pooled
    self.input_key = "txt"
    self.ln_penultimate = ln_penultimate

  def tokenize(self, text:str, device:Optional[str]=None) -> Tensor:
-    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int64, device=device).reshape(1,-1)
+    return Tensor(self.tokenizer.encode(text, pad_with_zeros=True), dtype=dtypes.int32, device=device).reshape(1,-1)

  def text_transformer_forward(self, x:Tensor, attn_mask:Optional[Tensor]=None):
    for r in self.model.transformer.resblocks:
@@ -449,7 +466,7 @@ class OpenClipEncoder:
    x = x + self.positional_embedding
    x = self.transformer(x, attn_mask=self.attn_mask)
    x = self.ln_final(x)
-    x = x[:, tokens.argmax(axis=-1)]
+    x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
    x = x @ self.text_projection
    return x

@@ -270,8 +270,10 @@ class FidInceptionV3:
    self.Mixed_7b = inception.Mixed_7b
    self.Mixed_7c = inception.Mixed_7c

-  def load_from_pretrained(self):
-    state_dict = torch_load(str(fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")))
+  def load_from_pretrained(self, path=None):
+    if path is None:
+      path = fetch("https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth", "pt_inception-2015-12-05-6726825d.pth")
+    state_dict = torch_load(str(path))
    for k,v in state_dict.items():
      if k.endswith(".num_batches_tracked"):
        state_dict[k] = v.reshape(1)
@@ -249,8 +249,5 @@ def convert_from_gguf(weights:dict[str, Tensor], n_layers:int):
  return sd

 def fix_bf16(weights:dict[Any, Tensor]):
-  if getenv("SUPPORT_BF16", 1):
-    # TODO: without casting to float16, 70B llama OOM on tinybox.
-    return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
-  # TODO: check if device supports bf16
-  return {k:v.llvm_bf16_cast(dtypes.half).to(v.device) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
+  # TODO: without casting to float16, 70B llama OOM on tinybox.
+  return {k:v.cast(dtypes.float32).cast(dtypes.float16) if v.dtype == dtypes.bfloat16 else v for k,v in weights.items()}
@@ -1,21 +1,24 @@
-from tinygrad import Tensor, dtypes
-from tinygrad.nn import Linear, Conv2d, GroupNorm, LayerNorm
+from tinygrad import Tensor, dtypes, nn
 from tinygrad.device import is_dtype_supported
-from typing import Optional, Union, List, Any, Tuple
+from typing import Optional, Union, List, Any, Tuple, Callable
 import math

+# allow for monkeypatching
+Linear, Conv2d, GroupNorm, LayerNorm = nn.Linear, nn.Conv2d, nn.GroupNorm, nn.LayerNorm
+attention, gelu, mixed_precision_dtype = Tensor.scaled_dot_product_attention, Tensor.gelu, dtypes.float16
+
 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/util.py#L207
 def timestep_embedding(timesteps:Tensor, dim:int, max_period=10000):
  half = dim // 2
  freqs = (-math.log(max_period) * Tensor.arange(half, device=timesteps.device) / half).exp()
  args = timesteps.unsqueeze(1) * freqs.unsqueeze(0)
  out = Tensor.cat(args.cos(), args.sin(), dim=-1)
-  return out.cast(dtypes.float16) if is_dtype_supported(dtypes.float16) else out
+  return out.cast(mixed_precision_dtype) if is_dtype_supported(mixed_precision_dtype) else out

 class ResBlock:
-  def __init__(self, channels:int, emb_channels:int, out_channels:int):
+  def __init__(self, channels:int, emb_channels:int, out_channels:int, num_groups:int=32):
    self.in_layers = [
-      GroupNorm(32, channels),
+      GroupNorm(num_groups, channels),
      Tensor.silu,
      Conv2d(channels, out_channels, 3, padding=1),
    ]
@@ -24,7 +27,7 @@ class ResBlock:
      Linear(emb_channels, out_channels),
    ]
    self.out_layers = [
-      GroupNorm(32, out_channels),
+      GroupNorm(num_groups, out_channels),
      Tensor.silu,
      lambda x: x,  # needed for weights loading code to work
      Conv2d(out_channels, out_channels, 3, padding=1),
@@ -45,35 +48,37 @@ class CrossAttention:
    self.to_v = Linear(ctx_dim,   n_heads*d_head, bias=False)
    self.num_heads = n_heads
    self.head_size = d_head
+    self.attn = attention
    self.to_out = [Linear(n_heads*d_head, query_dim)]

  def __call__(self, x:Tensor, ctx:Optional[Tensor]=None) -> Tensor:
    ctx = x if ctx is None else ctx
    q,k,v = self.to_q(x), self.to_k(ctx), self.to_v(ctx)
    q,k,v = [y.reshape(x.shape[0], -1, self.num_heads, self.head_size).transpose(1,2) for y in (q,k,v)]
-    attention = Tensor.scaled_dot_product_attention(q, k, v).transpose(1,2)
+    attention = self.attn(q, k, v).transpose(1,2)
    h_ = attention.reshape(x.shape[0], -1, self.num_heads * self.head_size)
    return h_.sequential(self.to_out)

 class GEGLU:
  def __init__(self, dim_in:int, dim_out:int):
    self.proj = Linear(dim_in, dim_out * 2)
+    self.gelu = gelu
    self.dim_out = dim_out

  def __call__(self, x:Tensor) -> Tensor:
    x, gate = self.proj(x).chunk(2, dim=-1)
-    return x * gate.gelu()
+    return x * self.gelu(gate)

 class FeedForward:
  def __init__(self, dim:int, mult:int=4):
-    self.net = [
+    self.net: tuple[GEGLU, Callable, nn.Linear] = (
      GEGLU(dim, dim*mult),
      lambda x: x,  # needed for weights loading code to work
      Linear(dim*mult, dim)
-    ]
+    )

  def __call__(self, x:Tensor) -> Tensor:
-    return x.sequential(self.net)
+    return x.sequential(list(self.net))

 class BasicTransformerBlock:
  def __init__(self, dim:int, ctx_dim:int, n_heads:int, d_head:int):
@@ -92,12 +97,13 @@ class BasicTransformerBlock:

 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/attention.py#L619
 class SpatialTransformer:
-  def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1):
+  def __init__(self, channels:int, n_heads:int, d_head:int, ctx_dim:Union[int,List[int]], use_linear:bool, depth:int=1,
+               norm_eps:float=1e-5):
    if isinstance(ctx_dim, int):
      ctx_dim = [ctx_dim]*depth
    else:
      assert isinstance(ctx_dim, list) and depth == len(ctx_dim)
-    self.norm = GroupNorm(32, channels)
+    self.norm = GroupNorm(32, channels, eps=norm_eps)
    assert channels == n_heads * d_head
    self.proj_in  = Linear(channels, channels) if use_linear else Conv2d(channels, channels, 1)
    self.transformer_blocks = [BasicTransformerBlock(channels, ctx_dim[d], n_heads, d_head) for d in range(depth)]
@@ -134,7 +140,9 @@ class Upsample:

 # https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/sgm/modules/diffusionmodules/openaimodel.py#L472
 class UNetModel:
-  def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int, channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None, n_heads:Optional[int]=None):
+  def __init__(self, adm_in_ch:Optional[int], in_ch:int, out_ch:int, model_ch:int, attention_resolutions:List[int], num_res_blocks:int,
+               channel_mult:List[int], transformer_depth:List[int], ctx_dim:Union[int,List[int]], use_linear:bool=False, d_head:Optional[int]=None,
+               n_heads:Optional[int]=None, num_groups:int=32, st_norm_eps:float=1e-5):
    self.model_ch = model_ch
    self.num_res_blocks = [num_res_blocks] * len(channel_mult)

@@ -174,12 +182,12 @@ class UNetModel:
    for idx, mult in enumerate(channel_mult):
      for _ in range(self.num_res_blocks[idx]):
        layers: List[Any] = [
-          ResBlock(ch, time_embed_dim, model_ch*mult),
+          ResBlock(ch, time_embed_dim, model_ch*mult, num_groups),
        ]
        ch = mult * model_ch
        if ds in attention_resolutions:
          d_head, n_heads = get_d_and_n_heads(ch)
-          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))
+          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))

        self.input_blocks.append(layers)
        input_block_channels.append(ch)
@@ -193,9 +201,9 @@ class UNetModel:

    d_head, n_heads = get_d_and_n_heads(ch)
    self.middle_block: List = [
-      ResBlock(ch, time_embed_dim, ch),
-      SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1]),
-      ResBlock(ch, time_embed_dim, ch),
+      ResBlock(ch, time_embed_dim, ch, num_groups),
+      SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[-1], norm_eps=st_norm_eps),
+      ResBlock(ch, time_embed_dim, ch, num_groups),
    ]

    self.output_blocks = []
@@ -203,13 +211,13 @@ class UNetModel:
      for i in range(self.num_res_blocks[idx] + 1):
        ich = input_block_channels.pop()
        layers = [
-          ResBlock(ch + ich, time_embed_dim, model_ch*mult),
+          ResBlock(ch + ich, time_embed_dim, model_ch*mult, num_groups),
        ]
        ch = model_ch * mult

        if ds in attention_resolutions:
          d_head, n_heads = get_d_and_n_heads(ch)
-          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx]))
+          layers.append(SpatialTransformer(ch, n_heads, d_head, ctx_dim, use_linear, depth=transformer_depth[idx], norm_eps=st_norm_eps))

        if idx > 0 and i == self.num_res_blocks[idx]:
          layers.append(Upsample(ch))
@@ -217,7 +225,7 @@ class UNetModel:
        self.output_blocks.append(layers)

    self.out = [
-      GroupNorm(32, ch),
+      GroupNorm(num_groups, ch),
      Tensor.silu,
      Conv2d(model_ch, out_ch, 3, padding=1),
    ]
@@ -230,10 +238,10 @@ class UNetModel:
      assert y.shape[0] == x.shape[0]
      emb = emb + y.sequential(self.label_emb[0])

-    if is_dtype_supported(dtypes.float16):
-      emb = emb.cast(dtypes.float16)
-      ctx = ctx.cast(dtypes.float16)
-      x   = x  .cast(dtypes.float16)
+    if is_dtype_supported(mixed_precision_dtype):
+      emb = emb.cast(mixed_precision_dtype)
+      ctx = ctx.cast(mixed_precision_dtype)
+      x   = x  .cast(mixed_precision_dtype)

    def run(x:Tensor, bb) -> Tensor:
      if isinstance(bb, ResBlock): x = bb(x, emb)
@@ -272,4 +272,4 @@ def compare_launch_state(states, good_states):

  return True, "PASS"

-# IOCTL=1 PTX=1 CUDA=1 python3 test/test_ops.py TestOps.test_tiny_add
+# IOCTL=1 CUDA=1 CUDA_PTX=1 python3 test/test_ops.py TestOps.test_tiny_add
@@ -1,6 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner, OnnxValue
+from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort

@@ -7,7 +7,7 @@ rm $LOGOPS
 test/external/process_replay/reset.py

 CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
-GPU=1 python3 -m pytest test/test_tiny.py
+CL=1 python3 -m pytest test/test_tiny.py

 # extract, sort and uniq
 extra/optimization/extract_dataset.py
@@ -1,6 +1,6 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.opt.kernel import Opt, OptOps
+from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -115,7 +115,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_
  assert dev.compiler is not None

  rawbufs = _ensure_buffer_alloc(rawbufs)
-  var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
+  var_vals: dict[str, int] = {k.expr:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
  p = get_program(lin.get_optimized_ast(), lin.opts)
  tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
                      max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
@@ -16,9 +16,9 @@ class TestBeamSearch(unittest.TestCase):
    BEAM.value = self.old_beam

  def test_variable_ast_beam(self):
-    with Context(IGNORE_OOB=1):
-      a = rand(3, 3).reshape((Variable("a", 1, 10).bind(3), 3))
-      a = (a+1).realize()
+    vi = Variable("a", 1, 10).bind(3)
+    a = rand(10, 3)[:vi]
+    a = (a+1).realize()

  def test_big_prime_number(self):
    a = rand(367, 367)
@@ -42,18 +42,16 @@ class TestBeamSearch(unittest.TestCase):

  def test_variable_big_prime_number(self):
    v = Variable("v", 1, 400).bind(367)
-    a = rand(367, 367)
-    b = rand(367, 367)
-    with Context(IGNORE_OOB=1):
-      c = (a.reshape(367, v) @ b.reshape(v, 367)).realize()
-      np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), atol=1e-4, rtol=1e-4)
+    a = rand(367, 400)
+    b = rand(400, 367)
+    c = (a[:, :v] @ b[:v, :]).realize()
+    np.testing.assert_allclose(c.numpy(), a[:, :367].numpy() @ b[:367, :].numpy(), atol=1e-4, rtol=1e-4)

  def test_variable_shrink_prime_number(self):
    v = Variable("v", 1, 400).bind(367)
    a = rand(400, 367)
-    with Context(IGNORE_OOB=1):
-      b = (a.shrink(((0,v), None))+1).reshape(367,367).realize()
-      np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)
+    b = (a.shrink(((0,v), None))+1)[:367,:367].realize()
+    np.testing.assert_allclose(b.numpy(), a.numpy()[:367]+1, atol=1e-4, rtol=1e-4)

  def test_no_mutate_rawbuffers(self):
    a = rand(3, 3).realize()
@@ -1,6 +1,6 @@
 import ctypes, array
 from hexdump import hexdump
-from tinygrad.runtime.ops_gpu import GPUDevice
+from tinygrad.runtime.ops_cl import CLDevice
 from tinygrad.helpers import getenv, to_mv, mv_address
 from tinygrad.dtype import dtypes
 from tinygrad import Tensor, TinyJit
@@ -8,7 +8,7 @@ from tinygrad.runtime.autogen import opencl as cl
 if getenv("IOCTL"): import extra.qcom_gpu_driver.opencl_ioctl  # noqa: F401  # pylint: disable=unused-import

 # create raw opencl buffer.
-gdev = GPUDevice()
+gdev = CLDevice()
 cl_buf = cl.clCreateBuffer(gdev.context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
 assert status.value == 0

@@ -673,6 +673,7 @@ impl<'a> Thread<'a> {
                                        39 => f32::log2(s0),
                                        42 => 1.0 / s0,
                                        43 => 1.0 / s0,
+                                        46 => 1.0 / f32::sqrt(s0),
                                        51 => f32::sqrt(s0),
                                        _ => todo_instr!(instruction)?,
                                    }
@@ -929,7 +930,7 @@ impl<'a> Thread<'a> {

            let op = ((instr >> 16) & 0x3ff) as u32;
            match op {
-                764 | 765 | 288 | 289 | 290 | 766 | 768 | 769 => {
+                764 | 765 | 288 | 289 | 290 | 766 | 767 | 768 | 769 => {
                    let vdst = (instr & 0xff) as usize;
                    let sdst = ((instr >> 8) & 0x7f) as usize;
                    let f = |i: u32| -> usize { ((instr >> i) & 0x1ff) as usize };
@@ -943,6 +944,16 @@ impl<'a> Thread<'a> {
                    assert_eq!(clmp, 0);

                    let vcc = match op {
+                        767 => {
+                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
+                            let (mul_result, overflow_mul) = (s0 as i64).overflowing_mul(s1 as i64);
+                            let (ret, overflow_add) = mul_result.overflowing_add(s2 as i64);
+                            let overflowed = overflow_mul || overflow_add;
+                            if self.exec.read() {
+                                self.vec_reg.write64(vdst, ret as u64);
+                            }
+                            overflowed
+                        },
                        766 => {
                            let (s0, s1, s2): (u32, u32, u64) = (self.val(s0), self.val(s1), self.val(s2));
                            let (mul_result, overflow_mul) = (s0 as u64).overflowing_mul(s1 as u64);
@@ -1246,7 +1257,7 @@ impl<'a> Thread<'a> {
                            }

                            let ret = match op {
-                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 531 | 537 | 540 | 551 | 567 | 796 => {
+                                257 | 259 | 299 | 260 | 261 | 264 | 272 | 392 | 426 | 430 | 531 | 537 | 540 | 551 | 567 | 796 => {
                                    let s0 = f32::from_bits(s0).negate(0, neg).absolute(0, abs);
                                    let s1 = f32::from_bits(s1).negate(1, neg).absolute(1, abs);
                                    let s2 = f32::from_bits(s2).negate(2, neg).absolute(2, abs);
@@ -1258,6 +1269,7 @@ impl<'a> Thread<'a> {
                                        272 => f32::max(s0, s1),
                                        299 => f32::mul_add(s0, s1, f32::from_bits(self.vec_reg[vdst])),
                                        426 => s0.recip(),
+                                        430 => 1.0 / f32::sqrt(s0),
                                        531 => f32::mul_add(s0, s1, s2),
                                        537 => f32::min(f32::min(s0, s1), s2),
                                        540 => f32::max(f32::max(s0, s1), s2),
@@ -2625,6 +2637,14 @@ mod test_vop1 {
        assert_eq!(thread.vec_reg[3], 1071644672);
    }

+    #[test]
+    fn test_v_rsq_f32() {
+        let mut thread = _helper_test_thread();
+        thread.vec_reg[0] = f32::to_bits(4.0);
+        r(&vec![0x7E005D00, END_PRG], &mut thread);
+        assert_eq!(f32::from_bits(thread.vec_reg[0]), 0.5);
+    }
+
    #[test]
    fn test_v_frexp_exp_i32_f64() {
        [(3573412790272.0, 42), (69.0, 7), (2.0, 2), (f64::NEG_INFINITY, 0)]
@@ -58,7 +58,7 @@ if __name__ == "__main__":
            GlobalCounters.kernel_count -= 1

        if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = get_program(k.get_optimized_ast(), k.opts)
+        p2 = get_program(k.ast, k.opts, k.applied_opts)
        new_ei = replace(ei, prg=CompiledRunner(p2))
        new_ei.run()
        new_jit.append(new_ei)
@@ -4,7 +4,7 @@

 Only supported on 7900XTX, requires either AM (`rmmod amdgpu`) or disabling power gating on AMD (`ppfeaturemask=0xffff3fff`, don't forget to rebuild initramfs)

-SQTT is implemented on top of normal tinygrad PROFILE=1, `PROFILE=1 SQTT=1` to get profile pickle with sqtt data embedded in it.
+SQTT is implemented on top of normal tinygrad profiling, `VIZ=1 SQTT=1` to get profile pickle with sqtt data embedded in it.

 `SQTT_BUFFER_SIZE=X` to change size of SQTT buffer (per shader engine, 6 SEs on 7900xtx) in megabytes, default 256.

@@ -0,0 +1,20 @@
+from extra.optimization.helpers import load_worlds, ast_str_to_ast
+from tinygrad.helpers import tqdm
+from tinygrad.uop.ops import pyrender, UOp, Ops
+from tinygrad import dtypes
+from tinygrad.shape.shapetracker import ShapeTracker, View
+inf, nan = float('inf'), float('nan')
+
+if __name__ == "__main__":
+  ast_strs = load_worlds()
+  for i, ast_str in enumerate(tqdm(ast_strs)):
+    good_ast = ast_str_to_ast(ast_str)
+    code = '\n'.join(pyrender(good_ast))
+    print("\n***************\n\n"+code)
+    exec(code)
+    if str(good_ast) != str(ast):
+      print(code)
+      print("MISMATCH")
+      print(good_ast)
+      print(ast)
+      break
@@ -4,13 +4,13 @@ import struct
 import json
 import traceback
 import numpy as np
-from tinygrad.runtime.ops_gpu import CLProgram, compile_gpu
+from tinygrad.runtime.ops_cl import CLProgram, compile_gpu
 from tinygrad.device import Device
 from tinygrad.helpers import DEBUG, getenv
 from collections import defaultdict
 import pyopencl as cl
-from tinygrad.runtime.ops_gpu import OSX_TIMING_RATIO
-CL = Device["GPU"]
+from tinygrad.runtime.ops_cl import OSX_TIMING_RATIO
+CL = Device["CL"]

 DEBUGCL = getenv("DEBUGCL", 0)
 FLOAT16 = getenv("FLOAT16", 0)
@@ -110,7 +110,7 @@ class Thneed:
    prgs = {}
    for o in jdat['binaries']:
      nptr = ptr + o['length']
-      prgs[o['name']] = CLProgram(Device["GPU"], o['name'], weights[ptr:nptr])
+      prgs[o['name']] = CLProgram(Device["CL"], o['name'], weights[ptr:nptr])
      ptr = nptr

    # populate the cl_cache
@@ -267,7 +267,7 @@ class Thneed:
    for prg, args in self.cl_cache:
      events.append(prg.clprg(CL.queue, *args))
    mt = time.monotonic()
-    Device["GPU"].synchronize()
+    Device["CL"].synchronize()
    et = time.monotonic() - st
    print(f"submit in {(mt-st)*1000.0:.2f} ms, total runtime is {et*1000.0:.2f} ms")

@@ -0,0 +1,74 @@
+# include directory copied from https://github.com/HazyResearch/ThunderMittens
+# https://hazyresearch.stanford.edu/blog/2024-11-28-tk-mlx
+
+gemm = """
+#include <metal_stdlib>
+#include "include/tk.metal"
+using namespace mittens;
+
+#define GEMM_PARAMS_DEF(T) \
+    device T* D [[buffer(0)]], \
+    device T* A [[buffer(1)]], \
+    device T* B [[buffer(2)]], \
+    const constant int &N [[buffer(3)]], \
+    const constant int &K [[buffer(4)]], \
+    const constant int &M [[buffer(5)]], \
+    uint3 tg_id [[threadgroup_position_in_grid]], \
+    uint simd_lane_id [[thread_index_in_simdgroup]]
+
+template<typename T, unsigned N_BLOCK, unsigned K_BLOCK, unsigned M_BLOCK>
+kernel void matmul_naive(GEMM_PARAMS_DEF(T)) {
+  using global_layout = gl<T, 1, 1, -1, -1>;
+  global_layout gl_a(A, nullptr, nullptr, N, K);
+  global_layout gl_b(B, nullptr, nullptr, K, M);
+  global_layout gl_d(D, nullptr, nullptr, N, M);
+  rt<T,     N_BLOCK * TILE_DIM, K_BLOCK * TILE_DIM> a_reg;
+  rt<T,     K_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> b_reg;
+  rt<float, N_BLOCK * TILE_DIM, M_BLOCK * TILE_DIM> d_reg;
+  zero(d_reg);
+  #pragma clang loop unroll(full)
+  for (int k = 0; k < K / (K_BLOCK * TILE_DIM); k++) {
+    load(a_reg, gl_a, {0, 0, (int)tg_id.y, k}, simd_lane_id);
+    load(b_reg, gl_b, {0, 0, k, (int)tg_id.x}, simd_lane_id);
+    mma_AB(d_reg, a_reg, b_reg, d_reg);
+  }
+  store(gl_d, d_reg, {0, 0, (int)tg_id.y, (int)tg_id.x}, simd_lane_id);
+}
+
+#define instantiate_matmul_custom(type_name, T) \
+   template [[host_name("matmul_custom_" #type_name)]] [[kernel]] \
+   void matmul_naive<T, 4, 2, 4>(GEMM_PARAMS_DEF(T)); \
+
+instantiate_matmul_custom(float32, float);
+"""
+
+from tinygrad import Device, Tensor, Context
+
+if __name__ == "__main__":
+  device = Device["METAL"]
+  lib = device.compiler.compile(gemm)
+  prg = device.runtime("matmul_custom_float32", lib)
+
+  N = 4096
+  a = Tensor.randn(N, N)
+  b = Tensor.randn(N, N)
+  c = Tensor.empty(N, N)
+  Tensor.realize(a, b, c)
+
+  TILE_DIM = 8
+  N_BLOCK = 4
+  M_BLOCK = 4
+
+  gsz = (N // (M_BLOCK * TILE_DIM), N // (N_BLOCK * TILE_DIM), 1)
+  for _ in range(5):
+    et = prg(c.uop.buffer.ensure_allocated()._buf, a.uop.buffer._buf, b.uop.buffer._buf,
+            global_size=gsz, local_size=(32,1,1), vals=(N, N, N), wait=True)
+    print(f"{N*N*N*2/(et*1e9):2f} GFLOPS")
+
+  for _ in range(5):
+    with Context(DEBUG=2):
+      ref = (a@b).realize()
+
+  print((ref-c).mean().item())
+
+
@@ -0,0 +1,392 @@
+/**
+ * @file
+ * @brief Basic operations on generic types.
+ */
+#pragma once
+#include "base_types.metal"
+#include <metal_math>
+
+namespace mittens {
+/**
+ * @namespace base_ops
+ *
+ * @brief A namespace for operations on basic data types.
+ */
+namespace base_ops {
+#define TEMPLATE_OPS_SINGLE(func_contents) \
+    template<typename T> static METAL_FUNC T op(device const T &x)      { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &x) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread const T &x)      { func_contents }
+
+#define TEMPLATE_OPS_OVERRIDE_SINGLE(T, op_name, func_contents) \
+    template<> METAL_FUNC T op_name::op<T>(device const T &x)      { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &x) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread const T &x)      { func_contents }
+
+#define TEMPLATE_OPS_DOUBLE(func_contents) \
+    template<typename T> static METAL_FUNC T op(device      const T &a, device      const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, threadgroup const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, thread      const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, device      const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread      const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, device      const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, threadgroup const T &b) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, thread      const T &b) { func_contents }
+
+#define TEMPLATE_OPS_OVERRIDE_DOUBLE(T, op_name, func_contents) \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, device      const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, threadgroup const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, thread      const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device      const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread      const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, device      const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, threadgroup const T &b) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, thread      const T &b) { func_contents }
+    
+#define TEMPLATE_OPS_TRIPLE(func_contents) \
+    template<typename T> static METAL_FUNC T op(device      const T &a, device      const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device 	    const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device 	    const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(device      const T &a, thread      const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, device      const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &a, thread      const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, device      const T &b, device 	    const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread 	    const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<typename T> static METAL_FUNC T op(thread      const T &a, thread      const T &b, thread      const T &c) { func_contents }
+
+#define TEMPLATE_OPS_OVERRIDE_TRIPLE(T, op_name, func_contents) \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, device      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device 	   const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(device      const T &a, thread      const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(threadgroup const T &a, thread      const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, device      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, device      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, device      const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, threadgroup const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, threadgroup const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, threadgroup const T &b, thread      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, thread      const T &b, device      const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, thread      const T &b, threadgroup const T &c) { func_contents } \
+    template<> METAL_FUNC T op_name::op<T>(thread      const T &a, thread      const T &b, thread      const T &c) { func_contents }
+
+
+
+/* ----------  CONST OPS  ---------- */
+
+/**
+ * @brief Represents the zero constant operation.
+ *
+ * This operation returns the zero value of the specified type.
+ *
+ * @tparam T The data type for which to return the zero value.
+ * @return The zero value of type T.
+ */
+struct zero {
+    template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::zero();      }
+};
+/**
+ * @brief Represents the one constant operation.
+ *
+ * This operation returns the one value of the specified type.
+ *
+ * @tparam T The data type for which to return the one value.
+ * @return The one value of type T.
+ */
+struct one {
+    template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::one();       }
+};
+    
+/**
+ * @brief Represents the positive infinity constant operation.
+ *
+ * This operation returns the positive infinity value of the specified type.
+ *
+ * @tparam T The data type for which to return the positive infinity value.
+ * @return The positive infinity value of type T.
+ */
+struct pos_infty {
+    template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
+};
+/**
+ * @brief Represents the negative infinity constant operation.
+ *
+ * This operation returns the negative infinity value of the specified type.
+ *
+ * @tparam T The data type for which to return the negative infinity value.
+ * @return The negative infinity value of type T.
+ */
+struct neg_infty {
+    template<typename T, typename... args> static METAL_FUNC constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
+};
+    
+
+/* ----------  UNARY OPS  ---------- */
+/**
+ * @brief Exponential function operation.
+ *
+ * This operation calculates the exponential of the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The exponential of the input value.
+ */
+struct exp {
+    TEMPLATE_OPS_SINGLE(return metal::exp(x);)
+};
+
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp, return bf16(metal::exp((float)x));)
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp, return bf16_2(metal::exp(float2(x)));)
+    
+    /**
+ * @brief Exponential function operation, in base 2
+ *
+ * This operation calculates the exponential of the input value, in base 2.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The exponential of the input value.
+ */
+struct exp2 {
+    template<typename T> static METAL_FUNC T op(device const T &x)      { return metal::exp2(x); } \
+    template<typename T> static METAL_FUNC T op(threadgroup const T &x) { return metal::exp2(x); } \
+    template<typename T> static METAL_FUNC T op(thread const T &x)      { return metal::exp2(x); }
+};
+    
+//template<> METAL_FUNC bf16 exp2::op<bf16>(device const bf16 &x)      { return bf16(metal::exp2(x)); } \
+//template<> METAL_FUNC bf16 exp2::op<bf16>(threadgroup const bf16 &x) { return bf16(metal::exp2(x)); } \
+//template<> METAL_FUNC bf16 exp2::op<bf16>(thread const bf16 &x)      { return bf16(metal::exp2(x)); }
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, exp2, return bf16(metal::exp2(x));)
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, exp2, return bf16_2(metal::exp2((float2)x));)
+
+/**
+ * @brief Natural log function operation.
+ *
+ * This operation calculates the natural logarithm of the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The natural logarithm of the input value.
+ */
+struct log {
+    TEMPLATE_OPS_SINGLE(return metal::log(x);)
+};
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16, log, return bf16(metal::log(x));)
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, log, return bf16_2(metal::log((float2)x));)
+
+/**
+ * @brief Absolute value operation.
+ *
+ * This operation calculates the absolute value of the input.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The absolute value of the input.
+ */
+struct abs {
+    TEMPLATE_OPS_SINGLE(return metal::abs(x);)
+};
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16  , abs, return bf16(metal::abs((float)x));)
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, abs, return bf16_2(metal::abs((float2)x));)
+/**
+ * @brief Rectified Linear Unit (ReLU) operation.
+ *
+ * This operation applies the ReLU function to the input, which is the
+ * maximum of zero and the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The result of ReLU function applied to the input.
+ */
+struct relu {
+    TEMPLATE_OPS_SINGLE(return max(x, base_types::constants<T>::zero());)
+};
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16  , relu, return bf16(metal::max((float)x, base_types::constants<float>::zero()));)
+TEMPLATE_OPS_OVERRIDE_SINGLE(bf16_2, relu, return bf16_2(metal::max((float2)x, base_types::constants<float2>::zero()));)
+/**
+ * @brief Copy operation.
+ *
+ * This operation returns the input value unchanged.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The input value.
+ * @return The same value as the input.
+ */
+struct copy { // for non-compile-time setters.
+    TEMPLATE_OPS_SINGLE(return x;)
+};
+
+/* ----------  BINARY OPS  ---------- */
+  
+    
+/**
+ * @brief Copy2 operation.
+ *
+ * This operation returns the second input value unchanged.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value (ignored).
+ * @param b[in] The second input value.
+ * @return The same value as the second input.
+ */
+struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
+    TEMPLATE_OPS_DOUBLE(return b;)
+};
+/**
+ * @brief Sum operation.
+ *
+ * This operation calculates the sum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The sum of the input values.
+ */
+struct sum {
+    TEMPLATE_OPS_DOUBLE(return a+b;)
+};
+    
+/**
+ * @brief Subtraction operation.
+ *
+ * This operation calculates the difference between two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The difference between the input values.
+ */
+struct sub {
+    TEMPLATE_OPS_DOUBLE(return a-b;)
+};
+/**
+ * @brief Multiplication operation.
+ *
+ * This operation calculates the product of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The product of the input values.
+ */
+struct mul {
+    TEMPLATE_OPS_DOUBLE(return a*b;)
+};
+/**
+ * @brief Division operation.
+ *
+ * This operation calculates the quotient of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The quotient of the input values.
+ */
+struct div {
+    TEMPLATE_OPS_DOUBLE(return a/b;)
+};
+/**
+ * @brief Maximum operation.
+ *
+ * This operation calculates the maximum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The maximum of the input values.
+ */
+struct max {
+    TEMPLATE_OPS_DOUBLE(return metal::max(a,b);)
+};
+TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16  , max, return (bf16)metal::max((float)a, (float)b);)
+TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, max, return (bf16_2)metal::max((float2)a, (float2)b);)
+/**
+ * @brief Minimum operation.
+ *
+ * This operation calculates the minimum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The minimum of the input values.
+ */
+struct min {
+    TEMPLATE_OPS_DOUBLE(return metal::min(a,b);)
+};
+TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16  , min, return (bf16)metal::min((float)a, (float)b);)
+TEMPLATE_OPS_OVERRIDE_DOUBLE(bf16_2, min, return (bf16_2)metal::min((float2)a, (float2)b);)
+
+
+/* ----------  TERNARY OPS  ---------- */
+/**
+ * @brief Fused multiply-add operation A * B + C.
+ *
+ * This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @param c[in] The third input value to be added.
+ * @return The result of the fused multiply-add operation.
+ */
+struct fma_AxBtC {
+    TEMPLATE_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, b), c);)
+};
+    
+/**
+ * @brief Fused multiply-add operation A * C + B.
+ *
+ * This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
+ * This is particularly useful for attention mechanisms in neural networks.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value. 
+ * @param b[in] The third input value to be added.
+ * @param c[in] The second input value.
+ * @return The result of the fused multiply-add operation.
+ */
+struct fma_AxCtB { // this is the one needed for attention
+    TEMPLATE_OPS_TRIPLE(return sum::op<T>(mul::op<T>(a, c), b);)
+};
+
+#undef TEMPLATE_OPS_SINGLE
+#undef TEMPLATE_OPS_OVERRIDE_SINGLE
+#undef TEMPLATE_OPS_DOUBLE
+#undef TEMPLATE_OPS_OVERRIDE_DOUBLE
+#undef TEMPLATE_OPS_TRIPLE
+#undef TEMPLATE_OPS_OVERRIDE_TRIPLE
+} // base_ops
+} // mittens
@@ -0,0 +1,321 @@
+
+#pragma once
+
+namespace mittens {
+    
+using bf16 = bfloat;
+using bf16_2 = bfloat2;
+using bf16_4 = bfloat4;
+//using half_2 = half2;
+     
+namespace ducks {
+namespace base_types {
+template <typename T>
+static METAL_FUNC constexpr const bool isT1() {
+    return metal::is_same<typename T::dtype, float>::value ||
+           metal::is_same<typename T::dtype, bf16 >::value ||
+           metal::is_same<typename T::dtype, half>::value;
+}
+template <typename T>
+static METAL_FUNC constexpr const bool isT2() {
+    return metal::is_same<typename T::dtype, float2>::value ||
+           metal::is_same<typename T::dtype, bf16_2>::value ||
+           metal::is_same<typename T::dtype, half2>::value;
+}
+    
+template <typename T>
+static METAL_FUNC constexpr const bool isT1Type() {
+    return metal::is_same<T, float>::value ||
+           metal::is_same<T, bf16 >::value ||
+           metal::is_same<T, half>::value;
+}
+template <typename T>
+static METAL_FUNC constexpr const bool isT2Type() {
+    return metal::is_same<T, float2>::value ||
+           metal::is_same<T, bf16_2>::value ||
+           metal::is_same<T, half2>::value;
+}
+    
+template <typename T>
+static METAL_FUNC constexpr const bool isT1Ptr() {
+    return metal::is_same<T, device      float*>::value ||
+           metal::is_same<T, threadgroup float*>::value ||
+           metal::is_same<T, thread      float*>::value ||
+           metal::is_same<T, device      bf16*>::value ||
+           metal::is_same<T, threadgroup bf16*>::value ||
+           metal::is_same<T, thread      bf16*>::value ||
+           metal::is_same<T, device      half*>::value ||
+           metal::is_same<T, threadgroup half*>::value ||
+           metal::is_same<T, thread      half*>::value;
+}
+template <typename T>
+static METAL_FUNC constexpr const bool isT2Ptr() {
+    return metal::is_same<T, device      float2*>::value ||
+           metal::is_same<T, threadgroup float2*>::value ||
+           metal::is_same<T, thread      float2*>::value ||
+           metal::is_same<T, device      bf16_2*>::value ||
+           metal::is_same<T, threadgroup bf16_2*>::value ||
+           metal::is_same<T, thread      bf16_2*>::value ||
+           metal::is_same<T, device      half2*>::value ||
+           metal::is_same<T, threadgroup half2*>::value ||
+           metal::is_same<T, thread      half2*>::value;
+}
+
+template <typename T>
+static METAL_FUNC constexpr const bool isTKType() { // good enough
+    return !isT1Type<T>() && !isT2Type<T>() && !isT1Ptr<T>() && !isT2Ptr<T>();
+}
+    
+} // namespace base_types
+} // namespace ducks
+    
+/**
+ * @namespace base_types
+ *
+ * @brief A namespace for Thundermittens basic data types.
+ */
+namespace base_types {
+/**
+ * @brief Provides compile-time constants for different types.
+ *
+ * @tparam T The type for which to provide constants.
+ */
+template<typename T> struct constants {
+    /**
+     * @brief Zero
+     * @return Constexpr zero with type T
+     */
+    static METAL_FUNC constexpr T zero()      { return T{0}; }
+    /**
+     * @brief One
+     * @return Constexpr one with type T
+     */
+    static METAL_FUNC constexpr T one()       { return T{1}; }
+    /**
+     * @brief Positive infinity. Particularly useful for initializing before a min op.
+     * @return Constexpr positive infinity with type T
+     */
+    static METAL_FUNC constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
+    /**
+     * @brief Negative infinity. Particularly useful for initializing before a max op.
+     * @return Constexpr negative infinity with type T
+     */
+    static METAL_FUNC constexpr T neg_infty() { return T{-INFINITY}; }
+};
+template<> struct constants<float> {
+    static METAL_FUNC constexpr float zero()      { return 0.f; }
+    static METAL_FUNC constexpr float one()       { return 1.f; }
+    static METAL_FUNC constexpr float pos_infty() { return INFINITY; }
+    static METAL_FUNC constexpr float neg_infty() { return -INFINITY; }
+};
+template<> struct constants<float2> {
+    static METAL_FUNC constexpr float2 zero()      { return float2(0.f, 0.f); }
+    static METAL_FUNC constexpr float2 one()       { return float2(1.f, 1.f); }
+    static METAL_FUNC constexpr float2 pos_infty() { return float2(constants<float>::pos_infty(), constants<float>::pos_infty()); }
+    static METAL_FUNC constexpr float2 neg_infty() { return float2(constants<float>::neg_infty(), constants<float>::neg_infty()); }
+};
+template<> struct constants<bf16> {
+    static METAL_FUNC constexpr bf16 zero()      { return 0.bf; }
+    static METAL_FUNC constexpr bf16 one()       { return 1.bf; }
+    static METAL_FUNC constexpr bf16 pos_infty() { return HUGE_VALBF; }
+    static METAL_FUNC constexpr bf16 neg_infty() { return -HUGE_VALBF; }
+};
+template<> struct constants<bf16_2> {
+    static METAL_FUNC constexpr bf16_2 zero()      { return bf16_2(constants<bf16>::zero(),      constants<bf16>::zero());      }
+    static METAL_FUNC constexpr bf16_2 one()       { return bf16_2(constants<bf16>::one(),       constants<bf16>::one());       }
+    static METAL_FUNC constexpr bf16_2 pos_infty() { return bf16_2(constants<bf16>::pos_infty(), constants<bf16>::pos_infty()); }
+    static METAL_FUNC constexpr bf16_2 neg_infty() { return bf16_2(constants<bf16>::neg_infty(), constants<bf16>::neg_infty()); }
+};
+template<> struct constants<half> {
+    static METAL_FUNC constexpr half zero()      { return half(0.h); }
+    static METAL_FUNC constexpr half one()       { return half(1.h); }
+    static METAL_FUNC constexpr half pos_infty() { return HUGE_VALH; }
+    static METAL_FUNC constexpr half neg_infty() { return -HUGE_VALH; }
+};
+    
+template<> struct constants<half2> {
+    static METAL_FUNC constexpr half2 zero()      { return half2(constants<half>::zero(),      constants<half>::zero());      }
+    static METAL_FUNC constexpr half2 one()       { return half2(constants<half>::one(),       constants<half>::one());       }
+    static METAL_FUNC constexpr half2 pos_infty() { return half2(constants<half>::pos_infty(), constants<half>::pos_infty()); }
+    static METAL_FUNC constexpr half2 neg_infty() { return half2(constants<half>::neg_infty(), constants<half>::neg_infty()); }
+};
+
+    
+
+/**
+ * @brief Provides information about packing of elements for a given type.
+ *
+ * @tparam T The type for which to provide packing information.
+ */
+template<typename T> struct packing {
+//    /**
+//     * @brief The number of elements packed together.
+//     *
+//     * @return constexpr int representing number of elements within the type.
+//     */
+//    static METAL_FUNC constexpr int num() { return 1; }
+//    /**
+//     * @brief Packs a single T element twice (replicated) into its packed type.
+//     *
+//     * @param i[in] The element to pack.
+//     * @return The packed type.
+//     */
+//    static METAL_FUNC constexpr T pack(device const bf16 &i);
+//    static METAL_FUNC constexpr T pack(threadgroup const bf16 &i);
+//    static METAL_FUNC constexpr T pack(thread const bf16 &i);
+};
+
+#define PACK_FUNCTIONS(T1, T2) \
+    static METAL_FUNC constexpr T2 pack(device const T1 &i) { return T2{i, i};      } \
+    static METAL_FUNC constexpr T2 pack(threadgroup const T1 &i) { return T2{i, i}; } \
+    static METAL_FUNC constexpr T2 pack(thread const T1 &i) { return T2{i, i};      }
+
+template<> struct packing<bf16> {
+    static METAL_FUNC constexpr int num() { return 1; }
+    using unpacked_type = bf16;
+    using packed_type = bf16_2;
+    using packed_four = bf16_4;
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<half> {
+    static METAL_FUNC constexpr int num() { return 1; }
+    using unpacked_type = half;
+    using packed_type = half2;
+    using packed_four = half4;
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<float> {
+    static METAL_FUNC constexpr int num() { return 1; }
+    using unpacked_type = float;
+    using packed_type = float2;
+    using packed_four = float4;
+    
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<bf16_2> {
+    static METAL_FUNC constexpr int num() { return 2; }
+    using unpacked_type = bf16;
+    using packed_type = bf16_2;
+    using packed_four = bf16_4;
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<half2> {
+    static METAL_FUNC constexpr int num() { return 2; }
+    using unpacked_type = half;
+    using packed_type = half2;
+    using packed_four = half4;
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<float2> {
+    static METAL_FUNC constexpr int num() { return 2; }
+    using unpacked_type = float;
+    using packed_type = float2;
+    using packed_four = float4;
+    PACK_FUNCTIONS(unpacked_type, packed_type)
+};
+template<> struct packing<int2> {
+    static METAL_FUNC constexpr int num() { return 2; }
+};
+template<> struct packing<float4> {
+    static METAL_FUNC constexpr int num() { return 4; }
+};
+template<> struct packing<int4> {
+    static METAL_FUNC constexpr int num() { return 4; }
+};
+
+    
+/**
+ * @brief Provides templated functionality to convert between different types.
+ *
+ * @tparam T The target type for conversion.
+ * @tparam U The source type for conversion.
+ */
+template<typename T, typename U> struct convertor {
+    /**
+     * @brief Converts a value of type U to type T.
+     *
+     * @param u[in] The value of type U to convert.
+     * @return T The converted value of type T.
+     */
+    static METAL_FUNC T convert(device const U & u)      { return (T)u; }
+    static METAL_FUNC T convert(threadgroup const U & u) { return (T)u; }
+    static METAL_FUNC T convert(thread const U & u)      { return (T)u; }
+};
+    
+template<> struct convertor<float, bf16> {
+    // fptrunc float %_ to bfloat
+    static METAL_FUNC float convert(device const bf16 & u)      { return float(u);}
+    static METAL_FUNC float convert(threadgroup const bf16 & u) { return float(u);}
+    static METAL_FUNC float convert(thread const bf16 & u)      { return float(u);}
+};
+template<> struct convertor<bf16, float> {
+    // fpext bfloat %_ to float
+    static METAL_FUNC bf16 convert(device const float & u)      { return bf16(u); }
+    static METAL_FUNC bf16 convert(threadgroup const float & u) { return bf16(u); }
+    static METAL_FUNC bf16 convert(thread const float & u)      { return bf16(u); }
+};
+template<> struct convertor<float2, bf16_2> {
+    // tail call fast <2 x float> @air.convert.f.v2f32.f.v2bf16(<2 x bfloat> %_)
+    static METAL_FUNC float2 convert(device const bf16_2 & u)      { return float2(u); }
+    static METAL_FUNC float2 convert(threadgroup const bf16_2 & u) { return float2(u); }
+    static METAL_FUNC float2 convert(thread const bf16_2 & u)      { return float2(u); }
+};
+template<> struct convertor<bf16_2, float2> {
+    // tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f32(<2 x float> %_)
+    static METAL_FUNC bf16_2 convert(device const float2 & u)      { return bf16_2(u); }
+    static METAL_FUNC bf16_2 convert(threadgroup const float2 & u) { return bf16_2(u); }
+    static METAL_FUNC bf16_2 convert(thread const float2 & u)      { return bf16_2(u); }
+};
+    
+template<> struct convertor<float, half> {
+    // fptrunc float %_ to half
+    static METAL_FUNC float convert(device const half & u)      { return float(u); }
+    static METAL_FUNC float convert(threadgroup const half & u) { return float(u); }
+    static METAL_FUNC float convert(thread const half & u)      { return float(u); }
+};
+template<> struct convertor<half, float> {
+    //fpext half %_ to float
+    static METAL_FUNC half convert(device const float & u)      { return half(u); }
+    static METAL_FUNC half convert(threadgroup const float & u) { return half(u); }
+    static METAL_FUNC half convert(thread const float & u)      { return half(u); }
+};
+template<> struct convertor<float2, half2> {
+    // tail call fast <2 x float> @air.convert.f.v2f32.f.v2f16(<2 x half> %_)
+    static METAL_FUNC float2 convert(device const half2 & u)      { return float2(u); }
+    static METAL_FUNC float2 convert(threadgroup const half2 & u) { return float2(u); }
+    static METAL_FUNC float2 convert(thread const half2 & u)      { return float2(u); }
+};
+template<> struct convertor<half2, float2> {
+    // tail call fast <2 x half> @air.convert.f.v2f16.f.v2f32(<2 x float> %_)
+    static METAL_FUNC half2 convert(device const float2 & u)      { return half2(u); }
+    static METAL_FUNC half2 convert(threadgroup const float2 & u) { return half2(u); }
+    static METAL_FUNC half2 convert(thread const float2 & u)      { return half2(u); }
+};
+template<> struct convertor<bf16, half> {
+    static METAL_FUNC bf16 convert(device const half & u)      { return bf16(u); }
+    static METAL_FUNC bf16 convert(threadgroup const half & u) { return bf16(u); }
+    static METAL_FUNC bf16 convert(thread const half & u)      { return bf16(u); }
+};
+template<> struct convertor<half, bf16> {
+    static METAL_FUNC half convert(device const bf16 & u)      { return half(u); }
+    static METAL_FUNC half convert(threadgroup const bf16 & u) { return half(u); }
+    static METAL_FUNC half convert(thread const bf16 & u)      { return half(u); }
+};
+template<> struct convertor<bf16_2, half2> {
+    // tail call fast <2 x bfloat> @air.convert.f.v2bf16.f.v2f16(<2 x half> %_)
+    static METAL_FUNC bf16_2 convert(device const half2 & u)      { return bf16_2(u); }
+    static METAL_FUNC bf16_2 convert(threadgroup const half2 & u) { return bf16_2(u); }
+    static METAL_FUNC bf16_2 convert(thread const half2 & u)      { return bf16_2(u); }
+};
+template<> struct convertor<half2, bf16_2> {
+    // tail call fast <2 x half> @air.convert.f.v2f16.f.v2bf16(<2 x bfloat> %_)
+    static METAL_FUNC half2 convert(device const bf16_2 & u)      { return half2(u); }
+    static METAL_FUNC half2 convert(threadgroup const bf16_2 & u) { return half2(u); }
+    static METAL_FUNC half2 convert(thread const bf16_2 & u)      { return half2(u); }
+};
+    
+    
+    
+} // base_types
+
+} // mittens
@@ -0,0 +1,10 @@
+/**
+ * @file
+ * @brief A collection of common resources on which Thundermittens depends.
+ */
+
+
+#pragma once
+#include "base_types.metal"
+#include "base_ops.metal"
+#include "utils.metal"
@@ -0,0 +1,225 @@
+/**
+ * @file
+ * @brief General utilities for Thundermittens.
+ */
+#pragma once // not done
+/*
+ TODO:
+   shared allocator
+   max shared mem for other hardware
+ */
+
+#include <metal_stdlib>
+#include "base_types.metal"
+/**
+ * @namespace mittens
+ *
+ * @brief The main namespace of Thundermittens.
+ */
+namespace mittens {
+/**
+ * @namespace ore
+ *
+ * @brief The main namespace of Thundermittens Metal.
+ */
+
+/* ----------  GENERAL CONSTANTS FOR mittens  ---------- */
+
+/**
+ * @brief Tile dimension constant.
+ */
+constant constexpr const int TILE_DIM{8};
+constant constexpr const int TILE_ELEMENTS{TILE_DIM*TILE_DIM};
+constant constexpr const int SIMD_THREADS{32};
+    
+    
+#ifdef M2_PRO
+constant constexpr int MAX_SHARED_MEMORY = 32768;
+#else
+constant constexpr int MAX_SHARED_MEMORY = 32768;
+#endif
+/* ----------  TYPE HELPERS  ---------- */
+/**
+ * @namespace ducks
+ *
+ * @brief Thundermittens' namespace for template metaprogramming..
+ *
+ * This includes primarily dummy types and concept wrappers, along
+ * with a few additional utilities.
+ */
+namespace ducks {
+
+/**
+ * @brief A type representing an empty default for a template.
+ */
+struct default_type {};
+
+// This macro can't be done as a template, so it doesn't really have a location in mittens.
+#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
+
+    
+}
+    
+/* ----------  SHUFFLE UTILS  ---------- */
+/**
+ * @brief Mask constant for all active threads in a warp.
+ */
+constant static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
+    
+template<typename T>
+static METAL_FUNC T shfl_sync(thread const T &f, const ushort laneid) {
+    return metal::simd_shuffle(f, laneid);
+}
+    
+template<>
+METAL_FUNC bfloat shfl_sync<bfloat>(thread const bf16 &f, const ushort laneid) {
+//    return as_type<bf16>(metal::simd_shuffle(*(thread half*)(&f), laneid));
+    float f_val = (float)f;
+    float shfl_val = metal::simd_shuffle(f_val, laneid);
+    return (bf16)shfl_val;
+}
+    
+template<>
+METAL_FUNC bfloat2 shfl_sync<bfloat2>(thread const bf16_2 &f, const ushort laneid) {
+//    return as_type<bf16_2>(metal::simd_shuffle(*(thread half2*)(&f), laneid));
+    float2 f_val = (float2)f;
+    float2 shfl_val = metal::simd_shuffle(f_val, laneid);
+    return (bf16_2)shfl_val;
+}
+    
+template<typename T>
+static METAL_FUNC T shfl_down_fill_sync(thread const T &f, thread const T& fill_data, const ushort laneid) {
+    return metal::simd_shuffle_and_fill_down(f, laneid, fill_data);
+}
+    
+template<>
+METAL_FUNC bfloat shfl_down_fill_sync<bfloat>(thread const bfloat &f, thread const bfloat &fill_data, const ushort laneid) {
+//    return as_type<bf16>(metal::simd_shuffle_and_fill_down(*(thread half*)(&f), *(thread half*)(&fill_data), laneid));
+    float f_val = (float)f;
+    float fill_data_f = (float)fill_data;
+    float shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
+    return (bf16)shfl_val;
+}
+template<>
+METAL_FUNC bfloat2 shfl_down_fill_sync<bfloat2>(thread const bfloat2 &f, thread const bfloat2 &fill_data, const ushort laneid) {
+//    return as_type<bf16_2>(metal::simd_shuffle_and_fill_down(*(thread half2*)(&f), *(thread half2*)(&fill_data), laneid));
+    float2 f_val = (float2)f;
+    float2 fill_data_f = (float2)fill_data;
+    float2 shfl_val = metal::simd_shuffle_and_fill_down(f_val, fill_data_f, laneid);
+    return (bf16_2)shfl_val;
+}
+/**
+ * @brief Perform a shuffle down operation on a packed type synchronously across a warp.
+ * @tparam T The type of the value to be shuffled.
+ * @param mask[in] The mask of active threads.
+ * @param f[in] The value to be shuffled.
+ * @param delta[in] The number of positions to shuffle down.
+ * @return The result of the shuffle operation.
+ */
+template<typename T>
+static METAL_FUNC T shfl_down_sync(thread const T &f, int delta) {
+    return metal::simd_shuffle_rotate_down(f, delta);
+}
+
+template<>
+METAL_FUNC bfloat shfl_down_sync<bfloat>(thread const bf16 &f, int delta) {
+//    return base_types::convertor<bf16, float>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float, bf16>::convert(f), delta));
+//    return as_type<bf16>(metal::simd_shuffle_rotate_down(*(thread half*)(&f), delta));
+    float f_val = (float)f;
+    float shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
+    return (bf16)shfl_val;
+}
+
+template<>
+METAL_FUNC bfloat2 shfl_down_sync<bfloat2>(thread const bf16_2 &f, int delta) {
+//    return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread const half2*)(&f), delta));
+//    return base_types::convertor<bf16_2, float2>::convert(metal::simd_shuffle_rotate_down(base_types::convertor<float2, bf16_2>::convert(f), delta));
+    
+    float2 f_val = (float2)f;
+    float2 shfl_val = metal::simd_shuffle_rotate_down(f_val, delta);
+    return (bf16_2)shfl_val;
+//    return as_type<bf16_2>(metal::simd_shuffle_rotate_down(*(thread half2*)(&f), delta));
+}
+    
+    
+/* ----------  LOOP UNROLLING UTILS  ---------- */
+    
+namespace meta {
+template <int Start, int End, int Stride, bool = (Start < End)>
+struct unroll_i_in_range {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F f, Args... args) {
+        f(Start, args...);
+        unroll_i_in_range<Start + Stride, End, Stride>::run(f, args...);
+    }
+};
+
+template <int Start, int End, int Stride>
+struct unroll_i_in_range<Start, End, Stride, false> {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F, Args...) {
+    }
+};
+
+
+template <int Start, int End, int Stride, bool = (Start < End)>
+struct unroll_i_j_in_range_inner {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F f, int outerIndex, Args... args) {
+        f(outerIndex, Start, args...);
+        unroll_i_j_in_range_inner<Start + Stride, End, Stride>::run(f, outerIndex, args...);
+    }
+};
+
+template <int Start, int End, int Stride>
+struct unroll_i_j_in_range_inner<Start, End, Stride, false> {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F, int, Args...) {
+    }
+};
+
+template <int StartOuter, int EndOuter, int StrideOuter,
+          int StartInner, int EndInner, int StrideInner,
+          bool = (StartOuter < EndOuter)>
+struct unroll_i_j_in_range {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F f, Args... args) {
+        unroll_i_j_in_range_inner<StartInner, EndInner, StrideInner>::run(
+            f, StartOuter, args...
+        );
+        unroll_i_j_in_range<
+            StartOuter + StrideOuter, EndOuter, StrideOuter,
+            StartInner, EndInner, StrideInner
+        >::run(f, args...);
+    }
+};
+
+template <int StartOuter, int EndOuter, int StrideOuter,
+          int StartInner, int EndInner, int StrideInner>
+struct unroll_i_j_in_range<StartOuter, EndOuter, StrideOuter,
+                          StartInner, EndInner, StrideInner, false> {
+    template<class F, typename... Args>
+    static METAL_FUNC void run(F, Args...) {
+    }
+};
+    
+}
+
+
+template <int N>
+struct ReadVector {
+    float _[N];
+};
+
+/* ----------  SHARED MEMORY UTILS  ---------- */
+
+#define mittens_ALIGN_AS(n) alignas(n)
+#define mittens_DEFAULT_ALIGN mittens_ALIGN_AS(16)
+    
+/**
+ * @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
+ */
+struct mittens_DEFAULT_ALIGN alignment_dummy { int dummy; };
+}
+
+
@@ -0,0 +1,24 @@
+/**
+ * @file
+ * @brief An aggregate header of all group (multi-warp) operations defined by Thundermittens
+ */
+
+#pragma once
+#include "../../common/common.metal"
+#include "../../types/types.metal"
+#include "../warp/warp.metal" // several group memory ops rely on underlying warp-scope ops
+namespace mittens {
+template<int N_WARPS>
+struct group {
+    constant static constexpr int GROUP_WARPS = N_WARPS; // This alias produces nice parallelism.
+    constant static constexpr int GROUP_THREADS = N_WARPS * mittens::SIMD_THREADS; // This alias produces nice parallelism.
+    static METAL_FUNC int simd_laneid(const unsigned threadIdx) { return threadIdx % mittens::SIMD_THREADS;         }
+    static METAL_FUNC int laneid     (const unsigned threadIdx) { return threadIdx % GROUP_THREADS;                      }
+    static METAL_FUNC int warpid     (const unsigned threadIdx) { return laneid(threadIdx) / mittens::SIMD_THREADS; }
+    static METAL_FUNC int groupid    (const unsigned threadIdx) { return threadIdx / GROUP_THREADS;                      }
+    #include "memory/memory.metal"
+    #include "shared/shared.metal"
+};
+    
+
+}
@@ -0,0 +1,2 @@
+#include "tile/tile.metal"
+#include "vec/vec.metal"
@@ -0,0 +1,132 @@
+
+/**
+ * @file
+ * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively loads data from a source array into row-major layout tiles.
+ *
+ * @tparam RT The row-major layout tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from. 
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<typename RT, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const device U *src = (device U*)&_src.template get<RT>(idx);
+    const int row_stride = _src.row_stride();
+    
+    int warp_laneid = threadIdx % 32;
+    const int row_offset = dst.rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < dst.height; i++) {
+        int row = simd_y + i * RT::tile_size;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < dst.width; j++) {
+            int col = simd_x + j * RT::tile_size;
+            T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
+            dst.tiles[i][j].data.thread_elements()[0] = src2[0];
+            dst.tiles[i][j].data.thread_elements()[1] = src2[1];
+        }
+    }
+}
+
+template<typename RT, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+load(thread RT &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const device U *src = (device U*)&_src.template get<RT>(idx);
+    const int row_stride = _src.row_stride();
+    
+    int warp_laneid = threadIdx % 32;
+    const int row_offset = dst.rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;;
+    const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < dst.height; i++) {
+        int row = simd_y + i * RT::tile_size;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < dst.width; j++) {
+            int col = simd_x + j * RT::tile_size;
+            T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src[row * row_stride + col])));
+            dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[row * row_stride + col]);
+            dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[(row + 1) * row_stride + col]);
+        }
+    }
+}
+/**
+ * @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
+ *
+ * @tparam RT The register tile type with a row-major layout.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<typename RT, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
+store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    device U *dst = (device U*)&(_dst.template get<RT>(idx));
+    const int row_stride = _dst.row_stride();
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = src.rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < src.height; i++) {
+        int row = simd_y + i * RT::tile_size;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < src.width; j++) {
+            int col = simd_x + j * RT::tile_size;
+            U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]));
+            *(device U2*)(&dst[row*row_stride + col]) = src2;
+        }
+    }
+}
+
+template<typename RT, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
+store(thread GL &_dst, thread const RT &src, thread const coord &idx, const int threadIdx) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    device U *dst = (device U*)&(_dst.template get<RT>(idx));
+    const int row_stride = _dst.row_stride();
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = src.rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+//    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+//    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < src.height; i++) {
+        int row = simd_y + i * RT::tile_size;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < src.width; j++) {
+            int col = simd_x + j * RT::tile_size;
+            dst[row*row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
+            dst[(row + 1) * row_stride + col] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
+        }
+    }
+}
@@ -0,0 +1,144 @@
+/**
+ * @file
+ * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory.
+ */
+
+
+//template<typename ST, typename U>
+//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+//load(int i,
+//     threadgroup ST *dst, device U* src,
+//     thread const int& group_laneid,
+//     thread const int& memcpy_per_row,
+//     thread const int& elem_per_memcpy,
+//     thread const int& row_stride)
+//{
+//    int idx = i * GROUP_THREADS + group_laneid;
+//    int row = idx / memcpy_per_row;
+//    int col = (idx*elem_per_memcpy) % ST::cols;
+//    if (row < ST::rows) {
+//        *(threadgroup float4*)(&(*dst)[{row, col}]) = *(device float4*)(&src[row*row_stride + col]);
+//    }
+//}
+
+
+template<typename ST, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+    int group_laneid = threadIdx % GROUP_THREADS;
+    using T = typename ST::T;
+    using U = typename GL::dtype;
+    device U *src = (device U*)&_src.template get<ST>(idx);
+    const int row_stride = _src.row_stride();
+    using read_vector = ReadVector<1>;
+    // we can handle this many rows each time we run a memcpy_async
+    constexpr const int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype);
+    constexpr const int memcpy_per_row = ST::cols / elem_per_memcpy;
+    int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < total_calls; i++) {
+
+        int idx = i * GROUP_THREADS + group_laneid;
+        int row = idx / memcpy_per_row;
+        int col = (idx*elem_per_memcpy) % dst.cols;
+        if (row<dst.rows && col < dst.cols) {
+            *(threadgroup read_vector*)(&dst[{row, col}]) = *(device read_vector*)(&src[row*row_stride + col]);
+//            *(threadgroup float*)(&dst[{row, col}]) = 1.0f;
+        }
+    }
+//    dst[{0, 0}] = base_types::convertor<T, float>::convert(1.f);
+//    dst[{0, 0}] = total_calls;
+//    meta::unroll_i_in_range<0, total_calls, 1>::run(load<ST, typename GL::dtype>, &dst, src, group_laneid, memcpy_per_row, elem_per_memcpy, row_stride);
+}
+
+
+//template<typename ST, typename GL>
+//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+//    int group_laneid = threadIdx % GROUP_THREADS;
+//    int groupid = threadIdx / GROUP_THREADS;
+//    int laneid = threadIdx % SIMD_THREADS;
+//    
+//    using U = typename GL::dtype;
+//    device U *src = (device U*)&_src.template get<ST>(idx);
+//    const int row_stride = _src.row_stride();
+//    
+//    int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
+//    int memcpy_per_row = ST::cols / elem_per_memcpy;
+//    int total_calls = ((ST::height * ST::width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
+//    /*
+//     1x16 or 8 x 128
+//     */
+//    int offset = ST::num_elements / (GROUP_WARPS);
+////    int offset = group_laneid
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < total_calls; i++) {
+//        int idx = i * SIMD_THREADS + laneid;
+////        int idx = i * () + group_laneid;
+//        int row = idx / memcpy_per_row;
+//        int col = (idx*elem_per_memcpy) % dst.cols;
+//        if (row<dst.rows) {
+//            *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row*row_stride + col]);
+//        }
+//    }
+//}
+//
+//template<typename ST, typename GL>
+//static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+//load(threadgroup ST &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+//    int warp_id = threadIdx / SIMD_THREADS;
+//    int lane_id = threadIdx % SIMD_THREADS;
+////    int N_WARPS = /* number of warps in your group */;
+//    
+//    using U = typename GL::dtype;
+//    device U *src = (device U*)&_src.template get<ST>(idx);
+//    const int row_stride = _src.row_stride();
+//    
+//    int elem_per_memcpy = sizeof(float)/sizeof(typename ST::dtype);
+//    int memcpy_per_row = ST::cols / elem_per_memcpy;
+//    int total_memcpy_elems = (ST::height * ST::cols) / elem_per_memcpy;
+//    int elems_per_warp = (total_memcpy_elems + N_WARPS - 1) / N_WARPS;  // Ceiling division
+//    
+//    int start_idx = warp_id * elems_per_warp;
+//    int end_idx = metal::min(start_idx + elems_per_warp, total_memcpy_elems);
+//    
+//    #pragma clang loop unroll(full)
+//    for (int idx = start_idx + lane_id; idx < end_idx; idx += SIMD_THREADS) {
+//        int row = idx / memcpy_per_row;
+//        int col = (idx % memcpy_per_row) * elem_per_memcpy;
+//        if (row < ST::height) {
+//            *(threadgroup float*)(&dst[{row, col}]) = *(device float*)(&src[row * row_stride + col]);
+//        }
+//    }
+//}
+
+template<typename ST, typename GL>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+store(thread const GL &_dst, threadgroup const ST &src, thread const coord &idx, const int threadIdx) {
+    int group_laneid = threadIdx % GROUP_THREADS;
+    using U = typename GL::dtype;
+    device U *dst = (device U*)&_dst.template get<ST>(idx);
+    const int row_stride = _dst.row_stride();
+    using read_vector = ReadVector<1>;
+    // we can handle this many rows each time we run a memcpy_async
+    int elem_per_memcpy = sizeof(read_vector)/sizeof(typename ST::dtype); // float/float -> 1
+    int memcpy_per_row = ST::cols / elem_per_memcpy; // 240 memcpy per row
+    int total_calls = ((src.height * src.width + (N_WARPS-1))) * TILE_DIM*TILE_DIM / (N_WARPS*SIMD_THREADS*elem_per_memcpy); // round up
+    
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < total_calls; i++) {
+
+        int idx = i * GROUP_THREADS + group_laneid;
+        
+        int row = idx / memcpy_per_row;
+        int col = (idx*elem_per_memcpy) % src.cols;
+        if (row<src.rows && col < src.cols) {
+            *(device read_vector*)(&dst[row*row_stride + col]) = *(threadgroup read_vector*)(&src[{row, col}]);
+//            *(device float*)(&dst[row*row_stride + col]) = 1.f;
+        }
+    }
+//    dst[0] = src[{0,0}];
+//    dst[0] = total_calls;
+//    dst[0] = base_types::convertor<U, float>::convert(1);
+}
+
@@ -0,0 +1,152 @@
+/**
+ * @file
+ * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = RT::rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < dst.height; i++) {
+        int row = simd_y + i * mittens::TILE_DIM;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < dst.width; j++) {
+            int col = simd_x + j * mittens::TILE_DIM;
+            T2 src2 = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[{row, col}])));
+            dst.tiles[i][j].data.thread_elements()[0] = src2[0];
+            dst.tiles[i][j].data.thread_elements()[1] = src2[1];
+        }
+    } 
+}
+
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, threadgroup const ST &src, const int threadIdx) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = RT::rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < dst.height; i++) {
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < dst.width; j++) {
+            int row = simd_y + i * mittens::TILE_DIM;
+            int col = simd_x + j * mittens::TILE_DIM;
+            dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[{row + 0, col}]);
+            dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[{row + 1, col}]);
+        }
+    }
+}
+
+/**
+ * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ */
+template<typename ST, typename RT>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = RT::rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < RT::height; i++) {
+        int row = simd_y + i * mittens::TILE_DIM;
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < RT::width; j++) {
+            int col = simd_x + j * mittens::TILE_DIM;
+            U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
+                                                                src.tiles[i][j].data.thread_elements()[1]));
+            *(threadgroup U2*)(&dst[{row, col}]) = src2;
+        }
+    }
+}
+
+
+template<typename ST, typename RT>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(threadgroup ST &dst, thread const RT &src, const int threadIdx) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%N_WARPS == 0, "Group load / store requires tile height to be a multiple of N_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(warp_height * N_WARPS == height, "RT height * N_WARPS must = ST height");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int warp_laneid = simd_laneid(threadIdx);
+    const int row_offset = RT::rows * warpid(threadIdx);
+    const short qid = warp_laneid / 4;
+//    const short simd_y = row_offset + (qid & 4) + (warp_laneid / 2) % 4;
+//    const short simd_x = (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    const short simd_y = row_offset + (qid & 2) * 2 + (warp_laneid % 2) * 2;
+    const short simd_x = (qid & 4) + (warp_laneid / 2) % 4;
+    #pragma clang loop unroll(full)
+    for(int i = 0; i < RT::height; i++) {
+        
+        #pragma clang loop unroll(full)
+        for(int j = 0; j < RT::width; j++) {
+            int row = simd_y + i * mittens::TILE_DIM;
+            int col = simd_x + j * mittens::TILE_DIM;
+//            U2 src2 = base_types::convertor<U2, T2>::convert(T2(src.tiles[i][j].data.thread_elements()[0],
+//                                                                src.tiles[i][j].data.thread_elements()[1]));
+//            *(threadgroup U2*)(&dst[{row, col}]) = src2;
+            
+            dst[{row + 0, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
+            dst[{row + 1, col}] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
+        }
+    }
+}
@@ -0,0 +1,8 @@
+/**
+ * @file
+ * @brief An aggregate header of group memory operations on tiles.
+ */
+
+#include "shared_to_register.metal"
+#include "global_to_register.metal"
+#include "global_to_shared.metal"
@@ -0,0 +1,47 @@
+
+/**
+ * @file
+ * @brief Functions for a warpgroup to collaboratively transfer  data directly between global memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively loads data into register vectors from a source array in global memory.
+ *
+ * @tparam RV The register vector type.
+ * @tparam U The data type of the source array.
+ * @param[out] dst The destination register vector to load data into.
+ * @param[in] src The source array in global memory to load data from.
+ */
+template<typename RV, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
+load(thread RV &dst, thread const GL &_src, thread coord idx, const int threadIdx) {
+    using T  = typename RV::dtype;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using T2 = typename base_types::packing<T>::packed_type;
+    
+    idx.c += warpid(threadIdx);
+    // Call warp level store
+    ::mittens::load(dst, _src, idx, simd_laneid(threadIdx));
+}
+
+/**
+ * @brief Collaboratively stores data from register vectors to a destination array in global memory.
+ *
+ * @tparam RV The register vector type.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register vector to store data from.
+ */
+template<typename RV, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>(), void>::type
+store(thread GL &_dst, thread const RV &src, thread coord idx, const int threadIdx) {
+    using T  = typename RV::dtype;
+//    using U2 = typename base_types::packing<U>::packed_type;
+    using T2 = typename base_types::packing<T>::packed_type;
+    
+    idx.c += warpid(threadIdx);
+
+    // Call warp level store
+    ::mittens::store(_dst, src, idx, simd_laneid(threadIdx));
+}
@@ -0,0 +1,59 @@
+/**
+ * @file
+ * @brief Group (collaborative warp) ops for loading shared vectors from and storing to global memory.
+ */
+
+/**
+ * @brief Loads data from global memory into shared memory vector.
+ *
+ * This function loads data from a global memory location pointed to by `src` into a shared memory vector `dst`.
+ * It calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
+ * The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
+ *
+ * @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
+ * @param dst Reference to the shared vector where the data will be loaded.
+ * @param src Pointer to the global memory location from where the data will be loaded.
+ */
+template<typename SV, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+load(threadgroup SV &dst, thread const GL &_src, thread const coord &idx, const int threadIdx) {
+    using U = typename GL::dtype;
+    using read_vector = ReadVector<1>;
+    constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
+    constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
+    device U *src = (device U*)&_src.template get<SV>(idx);
+    
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < total_calls; i+=GROUP_THREADS) {
+        if(i * elem_per_transfer < dst.length)
+            *(threadgroup read_vector*)&dst[i*elem_per_transfer] = *(device read_vector*)&src[i*elem_per_transfer];
+    }
+}
+ 
+/**
+ * @brief Stores data from a shared memory vector to global memory.
+ *
+ * This function stores data from a shared memory vector `src` to a global memory location pointed to by `dst`.
+ * Similar to the load function, it calculates the number of elements that can be transferred in one operation based on the size ratio of `float4` to the data type of `SV`.
+ * The function ensures coalesced memory access and efficient use of bandwidth by dividing the work among threads in a warp.
+ *
+ * @tparam SV Shared vector type, must satisfy ducks::sv::all concept.
+ * @param dst Pointer to the global memory location where the data will be stored.
+ * @param src Reference to the shared vector from where the data will be stored.
+ */
+template<typename SV, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+store(thread const GL &_dst, threadgroup const SV &src, thread const coord &idx, const int threadIdx) {
+    using read_vector = ReadVector<1>;
+    using U = typename GL::dtype;
+    constexpr int elem_per_transfer = sizeof(read_vector) / sizeof(typename SV::dtype);
+    constexpr int total_calls = SV::length / elem_per_transfer; // guaranteed to divide
+    device U *dst = (device U*)&_dst.template get<SV>(idx);
+    
+    metal::simdgroup_barrier(metal::mem_flags::mem_none);
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < total_calls; i+= GROUP_THREADS) {
+        if(i * elem_per_transfer < src.length)
+            *(device read_vector*)&dst[i*elem_per_transfer] = *(threadgroup read_vector*)&src[i*elem_per_transfer]; // lmao it's identical
+    }
+}
@@ -0,0 +1,60 @@
+/**
+ * @file
+ * @brief Functions for a group to collaboratively transfer data directly between shared memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively load data from a shared vector into register vectors split across a warpgroup.
+ *
+ * @tparam RV The register vector type
+ * @tparam SV The shared vector type
+ * @param dst[out] The destination register vector.
+ * @param src[in]  The source shared vector.
+ */
+template<typename RV, typename SV>
+METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
+load(thread RV &dst, threadgroup const SV &_src, const int threadIdx) {
+    using T  = typename RV::dtype;
+    using U  = typename SV::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using T2 = typename base_types::packing<T>::packed_type;
+
+    static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
+//    threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
+    //    threadgroup subvec &src = subvec_inplace<RV::outer_dim, SV>(_src, warpid(threadIdx));
+    unsigned warpId = warpid(threadIdx);
+    using subvec = typename SV::template subvec<RV::length>;
+
+    threadgroup subvec& src = *(threadgroup subvec*)(&_src[warpId *RV::length]);
+    
+    ::mittens::load<RV, subvec>(dst, src, simd_laneid(threadIdx)); // warp-level
+}
+
+/**
+ * @brief Collaboratively store data into a shared vector from register vectors split across a warpgroup.
+ *
+ * @tparam RV The register vector type
+ * @tparam SV The shared vector type
+ * @param dst[out] The destination shared vector.
+ * @param src[in]  The source register vector.
+ */
+template<typename SV, typename RV>
+METAL_FUNC static typename metal::enable_if<ducks::is_register_vector<RV>() && ducks::is_shared_vector<SV>(), void>::type
+store(threadgroup SV &_dst, thread const RV &src, const int threadIdx) {
+    using T  = typename RV::dtype;
+    using U  = typename SV::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U2 = typename base_types::packing<U>::packed_type;
+    
+
+    static_assert(SV::length == RV::length*N_WARPS, "rv and sv dimensions do not match");// confirm size correct
+    
+//    threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim> &dst = subvec_inplace<RV::outer_dim, SV>(_dst, warpid(threadIdx));
+//    ::mittens::store<threadgroup typename SV::template subvec<typename SV::dtype, RV::outer_dim>, RV>(dst, src, simd_laneid(threadIdx)); // warp-level
+    
+    unsigned warpId = warpid(threadIdx);
+    using subvec = typename SV::template subvec<RV::length>;
+    threadgroup subvec& dst = *(threadgroup subvec*)(&_dst[warpId * RV::length]);
+    
+    ::mittens::store(dst, src, simd_laneid(threadIdx)); // warp-level
+}
@@ -0,0 +1,8 @@
+/**
+ * @file
+ * @brief An aggregate header of group memory operations on vectors.
+ */
+
+#include "shared_to_register.metal"
+#include "global_to_register.metal"
+#include "global_to_shared.metal"
@@ -0,0 +1,3 @@
+
+#include "tile/tile.metal"
+#include "vec/vec.metal"
@@ -0,0 +1,27 @@
+/**
+ * @file
+ * @brief Group conversions between different shared memory tile types.
+ */
+
+/* ----------  COPIES  ---------- */
+
+/**
+ * @brief Copies data from one shared memory tile to another, potentially with different data types and layouts.
+ *
+ * @tparam T The data type of the destination tile.
+ * @tparam U The data type of the source tile.
+ * @tparam _height The height of the tile.
+ * @tparam _width The width of the tile.
+ * @tparam L1 The layout of the destination tile.
+ * @tparam L2 The layout of the source tile.
+ * @param[out] dst The destination tile.
+ * @param[in] src The source tile.
+ */
+template<typename T, typename U, int _height, int _width>
+static METAL_FUNC void copy(threadgroup st<T, _height, _width> &dst, threadgroup const st<U, _height, _width> &src, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i+=GROUP_THREADS) {
+        int row = i/dst.cols, col = i%dst.cols;
+        dst[{row, col}] = base_types::convertor<T, U>::convert(src[{row, col}]);
+    }
+}
@@ -0,0 +1,475 @@
+/**  
+ * @file
+ * @brief Group maps on shared tiles.
+ */
+
+/**
+ * @brief Performs a uniform unary operation on a tile.
+ *
+ * This function applies a given unary operation to each element of the source tile and stores the result in the destination tile.
+ * The operation is applied independently to each element, without considering its position or the values of neighboring elements.
+ *
+ * @tparam op The unary operation to be applied. Must be specialized to support operation on the data type of T.
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the unary operation is applied.
+ */
+template<typename op, typename ST> // T2, w, h can be inferred from dst as long as op is specialized
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ unary_map(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
+        dst.data[i] = op::template op<typename ST::dtype>(src.data[i]);
+    }
+}
+
+/**
+ * @brief Performs a uniform binary operation on a tile with a scalar parameter.
+ *
+ * This function applies a given binary operation to each element of the source tile and a scalar parameter, then stores the result in the destination tile.
+ * The operation is applied independently to each element, treating the scalar parameter as the second operand for each operation.
+ *
+ * @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the scalar parameter.
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the binary operation is applied.
+ * @param[in] param The scalar parameter to be used as the second operand in the binary operation.
+ */
+template<typename op, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ bin_map(threadgroup ST &dst, threadgroup const ST &src, thread const typename ST::dtype &param, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
+        dst.data[i] = op::template op<typename ST::dtype>(src.data[i], param);
+    }
+}
+
+/**
+ * @brief Performs a uniform binary operation on two tiles.
+ *
+ * This function applies a given binary operation to corresponding elements of two source tiles and stores the result in the destination tile.
+ * The operation is applied independently to each pair of elements, without considering their positions or the values of neighboring elements.
+ *
+ * @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T.
+ * @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile to which the binary operation is applied.
+ * @param[in] rhs The second source tile to which the binary operation is applied.
+ */
+template<typename op, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ bin_map(threadgroup ST &dst, threadgroup const ST &lhs, threadgroup const ST &rhs, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
+        dst.data[i] = op::template op<typename ST::dtype>(lhs.data[i], rhs.data[i]);
+    }
+}
+
+/**
+ * @brief Performs a row-wise binary operation on a tile with a vector.
+ *
+ * This function applies a given binary operation to each row of the source tile and the corresponding element of the source vector,
+ * then stores the result in the destination tile. The operation is applied independently to each row, using the vector element as
+ * the second operand for each element in the row.
+ *
+ * @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
+ * @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
+ * @tparam V The type of the vector. Must have the same data type as T.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the binary operation is applied.
+ * @param[in] vec The source vector containing the second operand for each row operation.
+ */
+template<typename op, typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+row_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const int threadIdx) {
+    static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
+    static_assert(SV::length == ST::rows, "Vector length must match the number of rows in the tile");
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
+        int row = i/dst.cols, col = i%dst.cols;
+        dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[row]);
+    }
+}
+
+/**
+ * @brief Performs a column-wise binary operation on a tile with a vector.
+ *
+ * This function applies a given binary operation to each column of the source tile and the corresponding element of the source vector,
+ * then stores the result in the destination tile. The operation is applied independently to each column, using the vector element as
+ * the second operand for each element in the column.
+ *
+ * @tparam op The binary operation to be applied. Must be specialized to support operation on the data type of T and the vector elements.
+ * @tparam T The type of the tiles. Must satisfy the `ducks::st::all` concept.
+ * @tparam V The type of the vector. Must have the same data type as T.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the binary operation is applied.
+ * @param[in] vec The source vector containing the second operand for each column operation.
+ */
+template<typename op, typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_map(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &vec, const int threadIdx) {
+    static_assert(metal::is_same<typename ST::dtype, typename SV::dtype>::value, "Tile and vector must have the same data type");
+    static_assert(SV::length == ST::cols, "Vector length must match the number of columns in the tile");
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.num_elements; i += GROUP_THREADS) {
+        int row = i/dst.cols, col = i%dst.cols;
+        dst[{row, col}] = op::template op<typename ST::dtype>(src[{row, col}], vec[col]);
+    }
+}
+
+
+/* ----------  WRAPPERS FOR PRETTINESS  ---------- */
+
+// All of the annoying qualifiers *should* be automatically inferred during compile-time.
+// So, syntax should just be mittens::add_row(tile, colvec);
+
+// const maps
+/**
+ * @brief Sets all elements of the destination tile to zero.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ zero(threadgroup ST &dst, const int threadIdx) {
+    unary_map<base_ops::zero, ST>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of the destination tile to one.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ one(threadgroup ST &dst, const int threadIdx) {
+    unary_map<base_ops::one, ST>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of the destination tile to positive infinity.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ pos_infty(threadgroup ST &dst, const int threadIdx) {
+    unary_map<base_ops::pos_infty, ST>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of the destination tile to negative infinity.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ neg_infty(threadgroup ST &dst, const int threadIdx) {
+    unary_map<base_ops::neg_infty, ST>(dst, dst, threadIdx);
+}
+
+// unary maps
+/**
+ * @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the exponential function is applied.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ exp(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    unary_map<base_ops::exp, ST>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the exponential function to each element of the source tile and stores the result in the destination tile, in base 2.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the exponential function is applied.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ exp2(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    unary_map<base_ops::exp2, ST>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the natural logarithm function to each element of the source tile and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the natural logarithm function is applied.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+log(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    unary_map<base_ops::log, ST>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the absolute function to each element of the source tile and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the absolute function is applied.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+abs(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    unary_map<base_ops::abs, ST>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the rectified linear unit function to each element of the source tile and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source tile to which the rectified linear unit function is applied.
+ */
+template<typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+relu(threadgroup ST &dst, threadgroup const ST &src, const int threadIdx) {
+    unary_map<base_ops::relu, ST>(dst, src, threadIdx);
+}
+/**
+ * @brief Copies the elements of the source tile to the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] src The source data to be copied.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ copy(threadgroup ST &dst, thread const U &src, const int threadIdx) {
+    bin_map<base_ops::copy, ST>(dst, src, threadIdx);
+}
+
+// uniform binary maps
+/**
+ * @brief Finds the maximum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ max(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::max, ST>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Finds the minimum of each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ min(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::min, ST>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Adds each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ add(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::sum, ST>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Subtracts each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ sub(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::sub, ST>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Multiplies each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+ mul(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::mul, ST>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Divides each pair of corresponding elements in the two source tiles and stores the result in the destination tile.
+ *
+ * @tparam T The type of the tile. Must satisfy the `ducks::st::all` concept.
+ * @tparam U The type of the second source data. Must be convertible to the data type of the destination tile.
+ * @param[out] dst The destination tile where the results are stored.
+ * @param[in] lhs The first source tile.
+ * @param[in] rhs The second source data.
+ */
+template<typename ST, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+div(threadgroup ST &dst, threadgroup const ST &lhs, thread const U &rhs, const int threadIdx) {
+    bin_map<base_ops::div, ST>(dst, lhs, rhs, threadIdx);
+}
+
+// Row and col maps
+
+/**
+ * @brief Adds row values to each row of a tile.
+ *
+ * @tparam T Tile type.
+ * @tparam V Column vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the addition on.
+ * @param row_values[in] Column vector containing values to add to each row.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ add_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
+    row_map<base_ops::sum, ST, SV>(dst, src, row_values, threadIdx);
+}
+/**
+ * @brief Subtracts row values from each row of a tile.
+ *
+ * @tparam T Tile type.
+ * @tparam V Column vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the subtraction on.
+ * @param row_values[in] Column vector containing values to subtract from each row.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ sub_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
+    row_map<base_ops::sub, ST, SV>(dst, src, row_values, threadIdx);
+}
+/**
+ * @brief Multiplies each row of a tile by row values.
+ *
+ * @tparam T Tile type.
+ * @tparam V Column vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the multiplication on.
+ * @param row_values[in] Column vector containing values to multiply each row by.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ mul_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
+    row_map<base_ops::mul, ST, SV>(dst, src, row_values, threadIdx);
+}
+/**
+ * @brief Divides each row of a tile by row values.
+ *
+ * @tparam T Tile type.
+ * @tparam V Column vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the division on.
+ * @param row_values[in] Column vector containing values to divide each row by.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ div_row(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &row_values, const int threadIdx) {
+    row_map<base_ops::div, ST, SV>(dst, src, row_values, threadIdx);
+}
+/**
+ * @brief Broadcast a vector into into a tile's rows.
+ *
+ * @tparam T Tile type.
+ * @tparam V Column vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param row_values[in] Column vector containing values to broadcast into rows.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ broadcast_row(threadgroup ST &dst, threadgroup const SV &row_values, const int threadIdx) {
+    row_map<base_ops::copy2, ST, SV>(dst, dst, row_values, threadIdx);
+}
+
+
+// col maps
+/**
+ * @brief Adds column values to each column of a tile.
+ *
+ * @tparam T Tile type.
+ * @tparam V Row vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the addition on.
+ * @param col_values[in] Row vector containing values to add to each column.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ add_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
+    col_map<base_ops::sum, ST, SV>(dst, src, col_values, threadIdx);
+}
+/**
+ * @brief Subtracts column values from each column of a tile.
+ *
+ * @tparam T Tile type.
+ * @tparam V Row vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the subtraction on.
+ * @param col_values[in] Row vector containing values to subtract from each column.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ sub_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
+    col_map<base_ops::sub, ST, SV>(dst, src, col_values, threadIdx);
+}
+/**
+ * @brief Multiplies each column of a tile by column values.
+ *
+ * @tparam T Tile type.
+ * @tparam V Row vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the multiplication on.
+ * @param col_values[in] Row vector containing values to multiply each column by.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ mul_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
+     col_map<base_ops::mul, ST, SV>(dst, src, col_values, threadIdx);
+}
+/**
+ * @brief Divides each column of a tile by column values.
+ *
+ * @tparam T Tile type.
+ * @tparam V Row vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param src[in] Source tile to apply the division on.
+ * @param col_values[in] Row vector containing values to divide each column by.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ div_col(threadgroup ST &dst, threadgroup const ST &src, threadgroup const SV &col_values, const int threadIdx) {
+    col_map<base_ops::div, ST, SV>(dst, src, col_values, threadIdx);
+}
+/**
+ * @brief Broadcast a vector into into a tile's columns.
+ *
+ * @tparam T Tile type.
+ * @tparam V Row vector type.
+ * @param dst[out] Destination tile where the result is stored.
+ * @param row_values[in] Row vector containing values to broadcast into cols.
+ */
+template<typename ST, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ broadcast_col(threadgroup ST &dst, threadgroup const SV &col_values, const int threadIdx) {
+    col_map<base_ops::copy2, ST, SV>(dst, dst, col_values, threadIdx);
+}
@@ -0,0 +1,284 @@
+/**
+ * @file
+ * @brief Group reductions on shared tiles.
+ */
+
+/**
+ * Performs row-wise reduction on a matrix using a specified operation.
+ *
+ * @tparam op The operation to be applied for reduction.
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type with row layout.
+ * @param row_accum The accumulator where the result of the reduction is stored.
+ * @param src The source matrix on which to perform the reduction.
+ * @param src_accum The initial value of the accumulator, used when reset is false.
+ * @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
+ */
+template<typename op, typename SV, typename ST, bool reset>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_reduce(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
+    using dtype = typename SV::dtype;
+    for (int row = laneid(threadIdx); row < src.rows; row += GROUP_THREADS) {
+        dtype accum = src[{row, 0}];
+        #pragma clang loop unroll(full)
+        for (int col = 1; col < src.cols; col++) {
+            accum = op::template op<dtype>(accum, src[{row, col}]);
+        }
+        if (reset) {
+            row_accum[row] = accum;
+        } else {
+            row_accum[row] = op::template op<dtype>(src_accum[row], accum);
+        }
+    }
+}
+
+/**
+ * Performs column-wise reduction on a matrix using a specified operation.
+ *
+ * @tparam op The operation to be applied for reduction.
+ * @tparam V The shared vector type for the column accumulator.
+ * @tparam T The shared matrix type with column layout.
+ * @param col_accum The accumulator where the result of the reduction is stored.
+ * @param src The source matrix on which to perform the reduction.
+ * @param src_accum The initial value of the accumulator, used when reset is false.
+ * @param reset A boolean flag indicating whether to reset the accumulator (ignore src_accum) or not.
+ */
+template<typename op, typename SV, typename ST, bool reset>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_reduce(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
+    using dtype = typename SV::dtype;
+    for (int col = laneid(threadIdx); col < src.cols; col += GROUP_THREADS) {
+        dtype accum = src[{0, col}];
+        #pragma clang loop unroll(full)
+        for (int row = 1; row < src.rows; row++) {
+            accum = op::template op<dtype>(accum, src[{row, col}]);
+        }
+        if (reset) {
+            col_accum[col] = accum;
+        } else {
+            col_accum[col] = op::template op<dtype>(src_accum[col], accum);
+        }
+    }
+}
+
+/* ----------  WRAPPERS FOR PRETTINESS  ---------- */
+
+/**
+ * @brief Store the maximum of each row of the src shared matrix in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_max(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx)  {
+    row_reduce<base_ops::max, SV, ST, true>(row_accum, src, row_accum, threadIdx);
+}
+/**
+ * @brief Store the minimum of each row of the src shared matrix in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_min(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx)  {
+    row_reduce<base_ops::min, SV, ST, true>(row_accum, src, row_accum, threadIdx);
+}
+/**
+ * @brief Store the sum of each row of the src shared matrix in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_sum(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx)  {
+    row_reduce<base_ops::sum, SV, ST, true>(row_accum, src, row_accum, threadIdx);
+}
+/**
+ * @brief Store the product of each row of the src shared matrix in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_prod(threadgroup SV &row_accum, threadgroup const ST &src, const int threadIdx) {
+    row_reduce<base_ops::mul, SV, ST, true>(row_accum, src, row_accum, threadIdx);
+}
+
+/**
+ * @brief Store the maximum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_max(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    row_reduce<base_ops::max, SV, ST, false>(row_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the minimum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_min(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    row_reduce<base_ops::min, SV, ST, false>(row_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the sum of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_sum(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    row_reduce<base_ops::sum, SV, ST, false>(row_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the product of each row of the src shared matrix, as well as the src_accum shared vector, in the row_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] row_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ row_prod(threadgroup SV &row_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
+    row_reduce<base_ops::mul, SV, ST, false>(row_accum, src, src_accum, threadIdx);
+}
+
+/**
+ * @brief Store the maximum of each column of the src shared matrix in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_max(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx)  {
+    col_reduce<base_ops::max, SV, ST, true>(col_accum, src, col_accum, threadIdx);
+}
+/**
+ * @brief Store the minimum of each column of the src shared matrix in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_min(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx)  {
+    col_reduce<base_ops::min, threadgroup SV, threadgroup ST, true>(col_accum, src, col_accum, threadIdx);
+}
+/**
+ * @brief Store the sum of each column of the src shared matrix in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_sum(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx)  {
+    col_reduce<base_ops::sum, SV, ST, true>(col_accum, src, col_accum, threadIdx);
+}
+/**
+ * @brief Store the product of each column of the src shared matrix in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_prod(threadgroup SV &col_accum, threadgroup const ST &src, const int threadIdx) {
+    col_reduce<base_ops::mul, SV, ST, true>(col_accum, src, col_accum, threadIdx);
+}
+
+/**
+ * @brief Store the maximum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_max(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    col_reduce<base_ops::max, SV, ST, false>(col_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the minimum of each column of the src shared matrix, as well as the src_accum shared vector, in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_min(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    col_reduce<base_ops::min, SV, ST, false>(col_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the sum of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_sum(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx)  {
+    col_reduce<base_ops::sum, SV, ST, false>(col_accum, src, src_accum, threadIdx);
+}
+/**
+ * @brief Store the product of each column of the src shared tile, as well as the src_accum row vector, in the col_accum shared vector.
+ *
+ * @tparam V The shared vector type for the row accumulator.
+ * @tparam T The shared matrix type.
+ * @param[out] col_accum The accumulator where the result of the reduction is stored.
+ * @param[in] src The source matrix on which to perform the reduction.
+ * @param[in] src_accum The initial value of the accumulator, used when accumulating onto an existing value.
+ */
+template<typename SV, typename ST>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_shared_vector<SV>(), void>::type
+ col_prod(threadgroup SV &col_accum, threadgroup const ST &src, threadgroup const SV &src_accum, const int threadIdx) {
+    col_reduce<base_ops::mul, SV, ST, false>(col_accum, src, src_accum, threadIdx);
+}
@@ -0,0 +1,3 @@
+#include "conversions.metal"
+#include "maps.metal"
+#include "reductions.metal"
@@ -0,0 +1,29 @@
+/**
+ * @file
+ * @brief Group conversions on shared vectors.
+ */
+
+/**
+ * @brief Copies data from one shared vector to another, converting data types if necessary.
+ *
+ * This function copies data from the source shared vector `src` to the destination shared vector `dst`.
+ * If the data types of `src` and `dst` are the same, it performs a direct memory copy. Otherwise, it
+ * converts each element from the source data type to the destination data type using the appropriate
+ * converter before copying.
+ *
+ * @tparam SV1 The type of the destination shared vector, must satisfy the ducks::sv::all concept.
+ * @tparam SV2 The type of the source shared vector, must satisfy the ducks::sv::all concept.
+ * @param[out] dst The destination shared vector.
+ * @param[in] src The source shared vector.
+ * @note The lengths of `src` and `dst` must be equal. This is enforced at compile time.
+ */
+template<typename SV1, typename SV2>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV1>() && ducks::is_shared_vector<SV2>(), void>::type
+copy(threadgroup SV1 &dst, threadgroup const SV2 &src, const int threadIdx) {
+    static_assert(SV1::length == SV2::length, "Source and destination vectors must have the same length.");
+    #pragma clang loop unroll(full)
+    for(int i = laneid(threadIdx); i < dst.length; i+=GROUP_THREADS) {
+        dst[i] = base_types::convertor<typename SV1::dtype, typename SV2::dtype>::convert(src[i]);
+    }
+}
+
@@ -0,0 +1,267 @@
+/**
+ * @file
+ * @brief Group maps on shared vectors.
+ */
+
+/**
+ * @brief Applies a unary operation to each element of a shared memory vector.
+ *
+ * @tparam op Unary operation type.
+ * @tparam T Shared memory vector type.
+ * @param dst[out] Destination vector in which to store the result.
+ * @param src[in] Source vector to apply the unary operation.
+ */
+template<typename op, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+unary_op(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
+        dst[cur] = op::template op<typename SV::dtype>(src[cur]);
+    }
+}
+/**
+ * @brief Perform a binary operation on two shared vectors.
+ *
+ * @tparam op The binary operation to perform.
+ * @tparam T The type of the vectors.
+ * @param dst[out] The destination vector where the result is stored.
+ * @param lhs[in] The left-hand side vector for the operation.
+ * @param rhs[in] The right-hand side vector for the operation.
+ */
+template<typename op, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+bin_op(threadgroup SV &dst, threadgroup const SV &lhs, threadgroup const SV &rhs, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
+        dst[cur] = op::template op<typename SV::dtype>(lhs[cur], rhs[cur]);
+    }
+}
+/**
+ * @brief Perform a binary operation on a shared vector and a scalar.
+ *
+ * @tparam op The binary operation to perform.
+ * @tparam T The type of the vector.
+ * @param dst[out] The destination vector where the result is stored.
+ * @param src[in] The source vector for the operation.
+ * @param param[in] The scalar parameter for the operation.
+ */
+template<typename op, typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+bin_op(threadgroup SV &dst, threadgroup const SV &src, thread const typename SV::dtype &param, const int threadIdx) {
+    #pragma clang loop unroll(full)
+    for(auto cur = laneid(threadIdx); cur < SV::length; cur+=GROUP_THREADS) {
+        dst[cur] = op::template op<typename SV::dtype>(src[cur], param);
+    }
+}
+
+/* ----------  WRAPPERS FOR PRETTINESS  ---------- */
+
+// ---- const ops ----
+
+/**
+ * @brief Sets all elements of a shared memory vector to zero.
+ *
+ * @tparam T Shared memory vector type.
+ * @param dst[out] Destination vector to be set to zero.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+zero(threadgroup SV &dst, const int threadIdx) {
+    unary_op<base_ops::zero, SV>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of a shared memory vector to one.
+ *
+ * @tparam T Shared memory vector type.
+ * @param dst[out] Destination vector to be set to one.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+one(threadgroup SV &dst, const int threadIdx) {
+    unary_op<base_ops::one, SV>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of a shared memory vector to positive infinity.
+ *
+ * @tparam T Shared memory vector type.
+ * @param dst[out] Destination vector to be set to positive infinity.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+pos_infty(threadgroup SV &dst, const int threadIdx) {
+    unary_op<base_ops::pos_infty, SV>(dst, dst, threadIdx);
+}
+/**
+ * @brief Sets all elements of a shared memory vector to negative infinity.
+ *
+ * @tparam T Shared memory vector type.
+ * @param dst[out] Destination vector to be set to negative infinity.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+neg_infty(threadgroup SV &dst, const int threadIdx) {
+    unary_op<base_ops::neg_infty, SV>(dst, dst, threadIdx);
+}
+
+// ---- unary ops ----
+
+/**
+ * @brief Copies the elements from one shared vector to another.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the source vector.
+ * @param dst[out] Destination vector where the elements will be copied to.
+ * @param src[in] Source vector to copy the elements from.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+copy(threadgroup SV &dst, thread const U &src, const int threadIdx) {
+    bin_op<base_ops::copy2, SV>(dst, dst, src, threadIdx); // the second arg is ignored here.
+}
+/**
+ * @brief Applies the exponential function element-wise to a shared vector.
+ *
+ * @tparam T Shared vector type.
+ * @param dst[out] Destination vector where the exponential values will be stored.
+ * @param src[in] Source vector to apply the exponential function to.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+exp(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    unary_op<base_ops::exp, SV>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the exponential function element-wise to a shared vector, in base 2.
+ *
+ * @tparam T Shared vector type.
+ * @param dst[out] Destination vector where the exponential values will be stored.
+ * @param src[in] Source vector to apply the exponential function to.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+exp2(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    unary_op<base_ops::exp2, SV>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the natural logarithm function element-wise to a shared vector.
+ *
+ * @tparam T Shared vector type.
+ * @param dst[out] Destination vector where the exponential values will be stored.
+ * @param src[in] Source vector to apply the logarithm function to.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+log(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    unary_op<base_ops::log, SV>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the absolute value function element-wise to a shared vector.
+ *
+ * @tparam T Shared vector type.
+ * @param dst[out] Destination vector where the absolute values will be stored.
+ * @param src[in] Source vector to apply the absolute value function to.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+abs(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    unary_op<base_ops::abs, SV>(dst, src, threadIdx);
+}
+/**
+ * @brief Applies the rectified linear unit (ReLU) function element-wise to a shared vector.
+ *
+ * @tparam T Shared vector type.
+ * @param dst[out] Destination vector where the ReLU values will be stored.
+ * @param src[in] Source vector to apply the ReLU function to.
+ */
+template<typename SV>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+relu(threadgroup SV &dst, threadgroup const SV &src, const int threadIdx) {
+    unary_op<base_ops::relu, SV>(dst, src, threadIdx);
+}
+
+// ---- binary ops ----
+
+/**
+ * @brief Computes the element-wise maximum of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the maximum values will be stored.
+ * @param lhs[in] First vector for the maximum operation.
+ * @param rhs[in] Second vector for the maximum operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+max(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::max, SV>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Computes the element-wise minimum of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the minimum values will be stored.
+ * @param lhs[in] First vector for the minimum operation.
+ * @param rhs[in] Second vector for the minimum operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+min(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::min, SV>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Computes the element-wise sum of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the sum values will be stored.
+ * @param lhs[in] First vector for the sum operation.
+ * @param rhs[in] Second vector for the sum operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+add(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::sum, SV>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Computes the element-wise difference of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the difference values will be stored.
+ * @param lhs[in] First vector for the difference operation.
+ * @param rhs[in] Second vector for the difference operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+sub(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::sub, SV>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Computes the element-wise product of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the product values will be stored.
+ * @param lhs[in] First vector for the product operation.
+ * @param rhs[in] Second vector for the product operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+mul(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::mul, SV>(dst, lhs, rhs, threadIdx);
+}
+/**
+ * @brief Computes the element-wise division of two shared vectors.
+ *
+ * @tparam T Shared vector type.
+ * @tparam U Type of the second vector.
+ * @param dst[out] Destination vector where the division values will be stored.
+ * @param lhs[in] First vector for the division operation.
+ * @param rhs[in] Second vector for the division operation.
+ */
+template<typename SV, typename U>
+static METAL_FUNC typename metal::enable_if<ducks::is_shared_vector<SV>(), void>::type
+div(threadgroup SV &dst, threadgroup const SV &lhs, thread const U &rhs, const int threadIdx) {
+    bin_op<base_ops::div, SV>(dst, lhs, rhs, threadIdx);
+}
@@ -0,0 +1,3 @@
+#include "conversions.metal"
+#include "maps.metal"
+
@@ -0,0 +1,3 @@
+#pragma once
+#include "group/group.metal"
+#include "warp/warp.metal"
@@ -0,0 +1,4 @@
+#pragma once
+#include "tile/tile.metal"
+#include "util/util.metal"
+#include "vec/vec.metal"
@@ -0,0 +1,51 @@
+/**
+* @file
+* @brief Functions for transferring data directly between global memory and registers and back.
+*/
+
+#pragma once
+
+#include "../../../../../common/common.metal"
+#include "../../../../../types/types.metal"
+
+#include "../global_to_register.metal"
+
+namespace mittens {
+/**
+ * @brief Load data from source arrays into a complex-type tile.
+ *
+ * @tparam CRT The complex tile type.
+ * @tparam U The data type of the source arrays.
+ * @param dst[out] The destination tile to load data into.
+ * @param resrc[in] The source array to load the real component data from.
+ * @param imsrc[in] The source array to load the imaginary component data from.
+ * @param re_row_stride[in] The stride in elements between rows in the real component source array.
+ * @param im_row_stride[in] The stride in elements between rows in the imaginary component source array.
+ */
+template<typename CRT, typename CGL>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_register_tile<CRT>() && ducks::is_complex_global_layout<CGL>(), void>::type
+load(thread CRT &dst, thread const CGL &src, thread const coord &idx, const short laneid) {
+    // Internally will use the correct load() method for row and column types
+    load(dst.real, src.real, idx);
+    load(dst.imag, src.imag, idx);
+}
+
+/**
+ * @brief Store data from a complex register tile to destination arrays in global memory.
+ *
+ * @tparam CRT The complex tile type.
+ * @tparam U The data type of the destination arrays.
+ * @param redst[out] The destination array in global memory to store the real component data into.
+ * @param imdst[out] The destination array in global memory to store the imaginary component data into.
+ * @param src[in] The source register tile to store data from.
+ * @param re_row_stride[in] The stride in elements between rows in the real component destination array.
+ * @param im_row_stride[in] The stride in elements between rows in the imaginary component destination array.
+ */
+template<typename CRT, typename CGL>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_register_tile<CRT>() && ducks::is_complex_global_layout<CGL>(), void>::type
+store(thread CGL &dst, thread const CRT &src, thread const coord &idx) {
+    // Internally will use the correct load() method for row and column types
+    store(dst.real, src.real, idx);
+    store(dst.imag, src.imag, idx);
+}
+}
@@ -0,0 +1,48 @@
+/**
+* @file
+* @brief Functions for transferring data directly between global and shared memory and back.
+*/
+
+#pragma once
+
+#include "../../../../../common/common.metal"
+#include "../../../../../types/types.metal"
+
+#include "../global_to_shared.metal"
+
+namespace mittens {
+/**
+ * @brief Loads data from global memory into a complex shared memory tile with a row layout.
+ *
+ * @tparam CST The type of the complex shared tile.
+ * @param[out] dst The destination complex shared memory tile.
+ * @param[in] resrc The source global memory array for the real component.
+ * @param[in] imsrc The source global memory array for the imaginary component.
+ * @param re_row_stride[in] The stride between rows in the source real component array.
+ * @param im_row_stride[in] The stride between rows in the source imaginary component array.
+ */
+template<typename CST, typename CGL>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_global_layout<CGL>(), void>::type
+load(threadgroup CST &dst, thread const CGL &src, thread const coord &idx) {
+    load(dst.real, src.real, idx);
+    load(dst.imag, src.imag, idx);
+}
+
+/**
+ * @brief Stores bf16 data from a complex shared memory tile with a row layout into global memory.
+ *
+ * @tparam CST The type of the complex shared tile.
+ * @param[out] redst The destination global memory array for the real component.
+ * @param[out] imdst The destination global memory array for the imaginary component.
+ * @param[in] src The source complex shared memory tile.
+ * @param re_row_stride[in] The stride between rows in the destination real component array.
+ * @param im_row_stride[in] The stride between rows in the destination imaginary component array.
+ */
+template<typename CST, typename CGL>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_global_layout<CGL>(), void>::type
+store(thread const CGL &dst, threadgroup CST &src, thread const coord &idx) {
+    store(dst.real, src.real, idx);
+    store(dst.imag, src.imag, idx);
+}
+
+}
@@ -0,0 +1,47 @@
+/**
+* @file
+* @brief Functions for transferring data directly between shared memory and registers and back.
+*/
+
+#pragma once
+
+
+#include "../../../../../common/common.metal"
+#include "../../../../../types/types.metal"
+
+#include "../shared_to_register.metal"
+
+namespace mittens {
+/**
+ * @brief Load data from a complex shared tile into a complex register tile.
+ *
+ * @tparam CRT The complex register tile type
+ * @tparam CST The complex shared tile type
+ * @param dst[out] The destination complex register tile.
+ * @param src[in]  The source complex shared tile.
+ */
+template<typename CRT, typename CST>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_register_tile<CRT>(), void>::type
+load(thread CRT &dst, threadgroup const CST &src) {
+    load(dst.real, src.real);
+    load(dst.imag, src.imag);
+}
+
+/**
+ * @brief Store data into a complex shared tile from a complex register tile.
+ *
+ * @tparam RT The complex register tile type
+ * @tparam ST The complex shared tile type
+ * @param dst[out] The destination complex shared tile.
+ * @param src[in]  The source complex register tile.
+ */
+template<typename CRT, typename CST>
+METAL_FUNC static typename metal::enable_if<ducks::is_complex_shared_tile<CST>() && ducks::is_complex_register_tile<CRT>(), void>::type
+store(threadgroup CST &dst, thread const CRT &src) {
+    store(dst.real, src.real);
+    store(dst.imag, src.imag);
+}
+
+
+}
+
@@ -0,0 +1,217 @@
+/**
+ * @file
+ * @brief Functions for transferring data directly between global memory and registers and back.
+ */
+
+#pragma once // done!
+#include "../../../../types/types.metal"
+#include "../../../../common/common.metal"
+#include <metal_stdlib>
+namespace mittens{
+
+namespace meta {
+template<typename RT, typename U>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
+load(int i, int j, thread RT *dst, const device U *src_ptr, const short simd_y, const short simd_x, const int row_stride) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+    T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src_ptr[offset])));
+    dst->tiles[i][j].data.thread_elements()[0] = src2[0];
+    dst->tiles[i][j].data.thread_elements()[1] = src2[1];
+}
+    
+template<typename RT, typename U>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
+load(int i, int j, thread RT *dst, const device U *src_ptr, const short simd_y, const short simd_x, const int row_stride) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+    dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src_ptr[offset]);
+    offset += row_stride;
+    dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src_ptr[offset]);
+}
+    
+template<typename RT, typename U>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>(), void>::type
+store(int i, int j, device U *dst_ptr, const thread RT *src, const short simd_y, const short simd_x, const int row_stride) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    unsigned offset = (simd_y + i * TILE_DIM) * row_stride + (simd_x + j * TILE_DIM);
+    U2 src2 = base_types::convertor<U2, T2>::convert(
+                                                     T2(src->tiles[i][j].data.thread_elements()[0],
+                                                        src->tiles[i][j].data.thread_elements()[1])
+                                                     );
+    *((device U2*)&dst_ptr[offset]) = src2;
+}
+
+template<typename RT, typename U>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>(), void>::type
+store(int i, int j, device U *dst_ptr, const thread RT *src, const short simd_y, const short simd_x, const int row_stride) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+    dst_ptr[offset] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
+    offset += row_stride;
+    dst_ptr[offset] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
+}
+
+}
+
+/**
+ * @brief Load data from a source array into a row-major layout tile.
+ *
+ * @tparam RT The row-major layout tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from.
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<typename RT, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+load(thread RT &dst, thread const GL &src, thread const coord &idx, const short laneid) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    const device U *src_ptr = (device U*)&src.template get<RT>(idx);
+    const int row_stride = src.row_stride();
+    
+    const short qid = laneid / 4;
+    const short simd_y = (qid & 4) + (laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for (int i = 0; i < RT::height; i++) {
+//        #pragma clang loop unroll(full)
+//        for (int j = 0; j < RT::width; j++) {
+//            unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+//            T2 src2 = base_types::convertor<T2, U2>::convert(*((device U2*)(&src_ptr[offset])));
+//            dst.tiles[i][j].data.thread_elements()[0] = src2[0];
+//            dst.tiles[i][j].data.thread_elements()[1] = src2[1];
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::load<RT, U>, &dst, src_ptr, simd_y, simd_x, row_stride);
+}
+/**
+ * @brief Load data from a source array into a col-major layout tile.
+ *
+ * @tparam RT The row-major layout tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from.
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<typename RT, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+load(thread RT &dst, thread const GL &src, thread const coord &idx, const short laneid) {
+    using T = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U = typename GL::dtype;
+    using layout = typename RT::layout;
+    const device U *src_ptr = (device U*)&(src.template get<RT>(idx));
+    const int row_stride = src.row_stride();
+    
+    const short qid = laneid / 4;
+    const short simd_x = (qid & 4) + (laneid / 2) % 4;
+    const short simd_y = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for (int i = 0; i < RT::height; i++) {
+//        #pragma clang loop unroll(full)
+//        for (int j = 0; j < RT::width; j++) {
+//            unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+//            dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src_ptr[offset]);
+//            offset += row_stride;
+//            dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src_ptr[offset]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::load<RT, U>, &dst, src_ptr, simd_y, simd_x, row_stride);
+}
+    
+/**
+ * @brief Store data from a register tile to a destination array in global memory with a row-major layout.
+ *
+ * @tparam RT The register tile type with a row-major layout.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<typename RT, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+store(thread GL &dst, thread const RT &src, thread const coord &idx, const short laneid) {
+    using T  = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    device U *dst_ptr = (device U*)&(dst.template get<RT>(idx));
+//    device U* dst_ptr = dst.raw_ptr;
+    const int row_stride = dst.row_stride();
+    const short qid = laneid / 4;
+    const short simd_y = (qid & 4) + (laneid / 2) % 4;
+    const short simd_x = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for (int i = 0; i < RT::height; i++) {
+//        #pragma clang loop unroll(full)
+//        for (int j = 0; j < RT::width; j++) {
+//            unsigned offset = (simd_y + i * TILE_DIM) * row_stride + (simd_x + j * TILE_DIM);
+//            U2 src2 = base_types::convertor<U2, T2>::convert(
+//                                                             T2(src.tiles[i][j].data.thread_elements()[0],
+//                                                                src.tiles[i][j].data.thread_elements()[1])
+//                                                             );
+//            *((device U2*)&dst_ptr[offset]) = src2;
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::store<RT, U>, dst_ptr, &src, simd_y, simd_x, row_stride);
+}
+
+/**
+ * @brief Store data from a register tile to a destination array in global memory with a col-major layout.
+ *
+ * @tparam RT The register tile type with a row-major layout.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<typename RT, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_global_layout<GL>(), void>::type
+store(thread GL &dst, thread const RT &src, thread const coord &idx, const short laneid) {
+    using T  = typename RT::dtype;
+    using T2 = typename RT::T2;
+    using U  = typename GL::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    using layout = typename RT::layout;
+    device U *dst_ptr = (device U*)&(dst.template get<RT>(idx));
+    const int row_stride = dst.row_stride();
+    const short qid = laneid / 4;
+    const short simd_x = (qid & 4) + (laneid / 2) % 4;
+    const short simd_y = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for (int i = 0; i < RT::height; i++) {
+//        #pragma clang loop unroll(full)
+//        for (int j = 0; j < RT::width; j++) {
+//            unsigned offset = (simd_y + i * rt_base<T, layout>::tile_size) * row_stride + (simd_x + j * rt_base<T, layout>::tile_size);
+//            dst_ptr[offset] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
+//            offset += row_stride;
+//            dst_ptr[offset] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::store<RT, U>, dst_ptr, &src, simd_y, simd_x, row_stride);
+}
+
+
+}
@@ -0,0 +1,192 @@
+/**
+ * @file
+ * @brief Functions for transferring data directly between global and shared memory and back.
+ */
+
+#pragma once // not done!
+#include "../../../../types/types.metal"
+#include "../../../../common/common.metal"
+#include <metal_stdlib>
+namespace mittens {
+
+//    
+namespace meta {
+template<typename ST, int memcpy_per_row, int elem_per_memcpy, int READ_FLOATS>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+load(int i, threadgroup ST *dst, device const typename ST::dtype *src, thread const int& row_stride, thread const short& laneid) {
+    {
+        unsigned idx = i + laneid;
+        unsigned row = idx / memcpy_per_row;
+        unsigned col = (idx*elem_per_memcpy) % ST::cols;
+        *(threadgroup ReadVector<READ_FLOATS>*)(&(*dst)[int2(row, col)]) = *(device ReadVector<READ_FLOATS>*)(&src[row*row_stride + col]);
+    }
+}
+
+template<typename ST, int memcpy_per_row, int elem_per_memcpy, int READ_FLOATS>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>(), void>::type
+store(int i, device typename ST::dtype *dst, threadgroup const ST *src, thread const int& row_stride, thread const short& laneid) {
+    {
+        unsigned idx = i + laneid;
+        unsigned row = idx / memcpy_per_row;
+        unsigned col = (idx*elem_per_memcpy) % ST::cols;
+        *(device ReadVector<READ_FLOATS>*)(&dst[row*row_stride + col]) = *(threadgroup ReadVector<READ_FLOATS>*)(&(*src)[int2(row, col)]);
+    }
+}
+    
+} // namespace meta
+
+//
+///**
+// * @brief Loads data from global memory into a shared memory tile with a row layout.
+// *
+// * @tparam ST The type of the shared tile.
+// * @param[out] dst The destination shared memory tile.
+// * @param[in] src The source global memory array.
+// * @param row_stride[in] The stride between rows in the source array.
+// * @param laneid[in] Thread's index in SIMD group
+// */
+//template<typename ST>
+//static METAL_FUNC void load(threadgroup ST &dst, device const typename ST::dtype *src, const int row_stride, short laneid) {
+//    using read_type = float;
+//    ducks::assert_shared_tile<ST>();
+//    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
+//    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;                    // 32/2=16 not power of 2
+//    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
+////    #pragma clang loop unroll_count(1)
+////    #pragma clang loop unroll(disable)
+//    #pragma clang loop unroll(full)
+//    for(unsigned i = 0; i < total_calls; i++) {
+//        unsigned idx = i * 32 + laneid;
+//        unsigned row = idx / memcpy_per_row;
+//        unsigned col = (idx*elem_per_memcpy) % ST::cols;
+//        *(threadgroup read_type*)(&dst[int2(row, col)]) = *(device read_type*)(&src[row*row_stride + col]);
+//    }
+//    
+////    ducks::assert_shared_tile<ST>();
+////    const constexpr int read_size = 1;
+////    using read_type = ReadVector<read_size>;
+////    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
+////    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;                    // 32/2=16 not power of 2
+////    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
+////    
+////    
+////    meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::load<ST, memcpy_per_row, elem_per_memcpy, read_size>, &dst, src, row_stride, laneid);
+//}
+//    
+//    
+///**
+// * @brief Stores data from a shared memory tile with a row layout into global memory.
+// *
+// * @tparam ST The type of the shared tile.
+// * @param[out] dst The destination global memory array.
+// * @param[in] src The source shared memory tile.
+// * @param row_stride[in] The stride between rows in the destination array.
+// * @param laneid[in] Thread's index in SIMD group
+// */
+//template<typename ST>
+//static METAL_FUNC void store(device typename ST::dtype *dst, threadgroup const ST &src, const int row_stride, short laneid) {
+//    using read_type = float4;
+//    ducks::assert_shared_tile<ST>();
+//    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
+//    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
+//    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
+////    #pragma clang loop unroll_count(READ_SIZE)
+////#pragma clang loop unroll(disable)
+//    #pragma clang loop unroll(full)
+//    for(unsigned i = 0; i < total_calls; i++) {
+//        unsigned idx = i * 32 + laneid;
+//        unsigned row = idx / memcpy_per_row;
+//        unsigned col = (idx*elem_per_memcpy) % src.cols;
+//        *(device read_type*)(&dst[row*row_stride + col]) = *(threadgroup read_type*)(&src[int2(row, col)]);
+//    }
+//    
+////    
+////    ducks::assert_shared_tile<ST>();
+////    const constexpr int read_size = 1;
+////    using read_type = ReadVector<read_size>;
+////
+////    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
+////    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
+////    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
+////
+////
+////    meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::store<ST, memcpy_per_row, elem_per_memcpy, read_size>, dst, &src, row_stride, laneid);
+//}
+    
+
+
+/**
+ * @brief Loads data from global memory into a shared memory tile with a row layout.
+ *
+ * @tparam ST The type of the shared tile.
+ * @param[out] dst The destination shared memory tile.
+ * @param[in] src The source global memory array.
+ * @param row_stride[in] The stride between rows in the source array.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename ST, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+load(threadgroup ST &dst, thread const GL &src, thread const coord &idx, short laneid) {
+    using U = typename GL::dtype;
+    constexpr const int read_size = 1;
+    using read_type = ReadVector<read_size>;
+    device U *src_ptr = (device U*)&src.template get<ST>(idx);
+    const int row_stride = src.row_stride();
+    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype); // 2
+    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;                    // 32/2=16 not power of 2
+    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy); // 1024/(32*2)=16
+//    #pragma clang loop unroll_count(1)
+//    #pragma clang loop unroll(disable)
+//    #pragma clang loop unroll(full)
+//    for(unsigned i = 0; i < total_calls; i++) {
+//        unsigned idx = i * 32 + laneid;
+//        unsigned row = idx / memcpy_per_row;
+//        unsigned col = (idx*elem_per_memcpy) % ST::cols;
+//        *(threadgroup read_type*)(&dst[int2(row, col)]) = *(device read_type*)(&src_ptr[row*row_stride + col]);
+//    }
+    meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::load<ST, memcpy_per_row, elem_per_memcpy, read_size>, &dst, src_ptr, row_stride, laneid);
+}
+    /*
+     
+     */
+    
+    
+/**
+ * @brief Stores data from a shared memory tile with a row layout into global memory.
+ *
+ * @tparam ST The type of the shared tile.
+ * @param[out] dst The destination global memory array.
+ * @param[in] src The source shared memory tile.
+ * @param row_stride[in] The stride between rows in the destination array.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename ST, typename GL>
+METAL_FUNC static typename metal::enable_if<ducks::is_shared_tile<ST>() && ducks::is_global_layout<GL>(), void>::type
+store(thread GL &dst, threadgroup const ST &src, thread const coord &idx, short laneid) {
+    using U = typename GL::dtype;
+    constexpr const int read_size = 1;
+    using read_type = ReadVector<read_size>;
+    device U *dst_ptr = (device U*)&dst.template get<ST>(idx);
+    const int row_stride = dst.row_stride();
+    
+    constexpr const unsigned elem_per_memcpy = sizeof(read_type)/sizeof(typename ST::dtype);
+    constexpr const unsigned memcpy_per_row = ST::cols / elem_per_memcpy;
+    constexpr const unsigned total_calls = ST::num_elements / (SIMD_THREADS*elem_per_memcpy);
+//    #pragma clang loop unroll_count(READ_SIZE)
+//#pragma clang loop unroll(disable)
+//    #pragma clang loop unroll(full)
+//    for(unsigned i = 0; i < total_calls; i++) {
+//        unsigned idx = i * 32 + laneid;
+//        unsigned row = idx / memcpy_per_row;
+//        unsigned col = (idx*elem_per_memcpy) % src.cols;
+//        *(device read_type*)(&dst_ptr[row*row_stride + col]) = *(threadgroup read_type*)(&src[int2(row, col)]);
+//    }
+    
+    meta::unroll_i_in_range<0, total_calls * SIMD_THREADS, SIMD_THREADS>::run(meta::store<ST, memcpy_per_row, elem_per_memcpy, read_size>, dst_ptr, &src, row_stride, laneid);
+}
+    
+
+    
+}
+
+
@@ -0,0 +1,461 @@
+/**
+ * @file
+ * @brief Functions for transferring data directly between shared memory and registers and back.
+ */
+#pragma once // done!
+
+#include "../../../../types/types.metal"
+#include "../../../../common/common.metal"
+#include <metal_stdlib>
+namespace mittens {
+
+// These probably need to be redone to reduce bank conflicts.
+// They currently work fine with xor layout but it should be
+// possible to reduce their bank conflicts with other layouts too.
+//
+namespace meta {
+
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+loadStR(int i, int j, thread RT *dst, threadgroup const ST *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+    T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&(*src)[int2(y, x)])));
+    dst->tiles[i][j].data.thread_elements()[0] = values[0];
+    dst->tiles[i][j].data.thread_elements()[1] = values[1];
+//    
+//    simdgroup_load(dst->tiles[i][j].data,
+//                   (threadgroup T*)(src->data),
+//                   src->cols,
+//                   {i * mittens::TILE_DIM, j * mittens::TILE_DIM},
+//                   
+}
+    
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+storeStR(int i, int j, threadgroup ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+    U2 values = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
+    *((threadgroup U2*)(&(*dst)[int2(y, x)])) = values;
+}
+
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+loadStR(int i, int j, thread RT *dst, threadgroup const ST *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+//    dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert((*src)[int2(y  , x)]);
+//    dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert((*src)[int2(y+1, x)]);
+    T2 vals = base_types::convertor<T2, U2>::convert({(*src)[int2(y  , x)], (*src)[int2(y+1, x)]});
+    dst->tiles[i][j].data.thread_elements()[0] = vals[0];
+    dst->tiles[i][j].data.thread_elements()[1] = vals[1];
+}
+    
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+storeStR(int i, int j, threadgroup ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+//    (*dst)[int2(y  , x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
+//    (*dst)[int2(y+1, x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
+    
+    U2 vals = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
+    (*dst)[int2(y  , x)] = vals[0];
+    (*dst)[int2(y+1, x)] = vals[1];
+}
+
+}
+    
+/**
+ * @brief Load data from a shared tile into a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, threadgroup const ST &src, short laneid) {
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const short qid = laneid / 4;
+    int offsetY = (qid & 4) + (laneid / 2) % 4;
+    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < dst.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < dst.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[int2(y, x)])));
+//            dst.tiles[i][j].data.thread_elements()[0] = values[0];
+//            dst.tiles[i][j].data.thread_elements()[1] = values[1];
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+/**
+ * @brief Load data from a shared tile into a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, threadgroup const ST &src, short laneid) {
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const short qid = laneid / 4;
+//    int offsetY = (qid & 4) + (laneid / 2) % 4;
+//    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    int offsetX = (qid & 4) + (laneid / 2) % 4;
+    int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < dst.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < dst.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[int2(y  , x)]);
+//            dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[int2(y+1, x)]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+
+/**
+ * @brief Store data into a shared tile from a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(threadgroup ST &dst, thread const RT &src, short laneid) {
+    ducks::assert_register_tile<RT>();
+    ducks::assert_shared_tile<ST>();
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+
+    const short qid = laneid / 4;
+    int offsetY = (qid & 4) + (laneid / 2) % 4;
+    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < src.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < src.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            U2 values = base_types::convertor<U2, T2>::convert({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
+//            *((threadgroup U2*)(&dst[int2(y, x)])) = values;
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+/**
+ * @brief Store data into a shared tile from a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(threadgroup ST &dst, thread const RT &src, short laneid) {
+    ducks::assert_register_tile<RT>();
+    ducks::assert_shared_tile<ST>();
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+
+    const short qid = laneid / 4;
+//    int offsetY = (qid & 4) + (laneid / 2) % 4;
+//    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    int offsetX = (qid & 4) + (laneid / 2) % 4;
+    int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < src.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < src.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            dst[int2(y  , x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
+//            dst[int2(y+1, x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+/*---------------------------------------------------------------------------------*/
+// These probably need to be redone to reduce bank conflicts.
+// They currently work fine with xor layout but it should be
+// possible to reduce their bank conflicts with other layouts too.
+//
+namespace meta {
+
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+loadStR_r(int i, int j, thread RT *dst, thread const ST *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+    T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&(*src)[int2(y, x)])));
+    dst->tiles[i][j].data.thread_elements()[0] = values[0];
+    dst->tiles[i][j].data.thread_elements()[1] = values[1];
+//
+//    simdgroup_load(dst->tiles[i][j].data,
+//                   (threadgroup T*)(src->data),
+//                   src->cols,
+//                   {i * mittens::TILE_DIM, j * mittens::TILE_DIM},
+//
+}
+    
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+storeStR_r(int i, int j, thread ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+    U2 values = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
+    *((threadgroup U2*)(&(*dst)[int2(y, x)])) = values;
+}
+
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+loadStR_c(int i, int j, thread RT *dst, thread const ST *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+//    dst->tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert((*src)[int2(y  , x)]);
+//    dst->tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert((*src)[int2(y+1, x)]);
+    T2 vals = base_types::convertor<T2, U2>::convert({(*src)[int2(y  , x)], (*src)[int2(y+1, x)]});
+    dst->tiles[i][j].data.thread_elements()[0] = vals[0];
+    dst->tiles[i][j].data.thread_elements()[1] = vals[1];
+}
+    
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+storeStR_c(int i, int j, thread ST *dst, thread const RT *src, short laneid, int offsetY, int offsetX) {
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    int y = offsetY + i * mittens::TILE_DIM;
+    int x = offsetX + j * mittens::TILE_DIM;
+//    (*dst)[int2(y  , x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[0]);
+//    (*dst)[int2(y+1, x)] = base_types::convertor<U, T>::convert(src->tiles[i][j].data.thread_elements()[1]);
+    
+    U2 vals = base_types::convertor<U2, T2>::convert({src->tiles[i][j].data.thread_elements()[0], src->tiles[i][j].data.thread_elements()[1]});
+    (*dst)[int2(y  , x)] = vals[0];
+    (*dst)[int2(y+1, x)] = vals[1];
+}
+
+}
+    
+/**
+ * @brief Load data from a shared tile into a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, thread const ST &src, short laneid) {
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const short qid = laneid / 4;
+    int offsetY = (qid & 4) + (laneid / 2) % 4;
+    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < dst.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < dst.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            T2 values = base_types::convertor<T2, U2>::convert(*((threadgroup U2*)(&src[int2(y, x)])));
+//            dst.tiles[i][j].data.thread_elements()[0] = values[0];
+//            dst.tiles[i][j].data.thread_elements()[1] = values[1];
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR_r<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+/**
+ * @brief Load data from a shared tile into a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+load(thread RT &dst, thread const ST &src, short laneid) {
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+    const short qid = laneid / 4;
+//    int offsetY = (qid & 4) + (laneid / 2) % 4;
+//    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    int offsetX = (qid & 4) + (laneid / 2) % 4;
+    int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < dst.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < dst.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            dst.tiles[i][j].data.thread_elements()[0] = base_types::convertor<T, U>::convert(src[int2(y  , x)]);
+//            dst.tiles[i][j].data.thread_elements()[1] = base_types::convertor<T, U>::convert(src[int2(y+1, x)]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::loadStR_c<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+
+/**
+ * @brief Store data into a shared tile from a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_row_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(thread ST &dst, thread const RT &src, short laneid) {
+    ducks::assert_register_tile<RT>();
+    ducks::assert_shared_tile<ST>();
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+
+    const short qid = laneid / 4;
+    int offsetY = (qid & 4) + (laneid / 2) % 4;
+    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < src.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < src.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            U2 values = base_types::convertor<U2, T2>::convert({src.tiles[i][j].data.thread_elements()[0], src.tiles[i][j].data.thread_elements()[1]});
+//            *((threadgroup U2*)(&dst[int2(y, x)])) = values;
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR_r<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+/**
+ * @brief Store data into a shared tile from a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ * @param laneid[in] Thread's index in SIMD group
+ */
+template<typename RT, typename ST>
+METAL_FUNC static typename metal::enable_if<ducks::is_col_register_tile<RT>() && ducks::is_shared_tile<ST>(), void>::type
+store(thread ST &dst, thread const RT &src, short laneid) {
+    ducks::assert_register_tile<RT>();
+    ducks::assert_shared_tile<ST>();
+    static_assert(RT::height == ST::height, "register tile and shared tile must match height");
+    static_assert(RT::width  == ST::width,  "register tile and shared tile must match width");
+    using T  = typename RT::dtype;
+    using T2 = typename base_types::packing<T>::packed_type;
+    using U  = typename ST::dtype;
+    using U2 = typename base_types::packing<U>::packed_type;
+
+    const short qid = laneid / 4;
+//    int offsetY = (qid & 4) + (laneid / 2) % 4;
+//    int offsetX = (qid & 2) * 2 + (laneid % 2) * 2;
+    int offsetX = (qid & 4) + (laneid / 2) % 4;
+    int offsetY = (qid & 2) * 2 + (laneid % 2) * 2;
+    
+//    #pragma clang loop unroll(full)
+//    for(int i = 0; i < src.height; i++) {
+//        #pragma clang loop unroll(full)
+//        for(int j = 0; j < src.width; j++) {
+//            int y = offsetY + i * mittens::TILE_DIM;
+//            int x = offsetX + j * mittens::TILE_DIM;
+//            dst[int2(y  , x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[0]);
+//            dst[int2(y+1, x)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data.thread_elements()[1]);
+//        }
+//    }
+    meta::unroll_i_j_in_range<0, RT::height, 1, 0, RT::width, 1>::run(meta::storeStR_c<RT, ST>, &dst, &src, laneid, offsetY, offsetX);
+}
+    
+}
+
+
--- a/Show More
+++ b/Show More