might end if

* rocprof parser

* viewer

* vw

* skip

.github/actions/setup-tinygrad/action.yml

@@ -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'

.github/workflows/autogen.yml

@@ -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

.github/workflows/benchmark.yml

@@ -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

.github/workflows/test.yml

@@ -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=41 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,22 @@ 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
+        run: NULL=1 SAMPLES=300 BS=8 SEQLEN=512 GRADIENT_ACC_STEPS=8 FAKEDATA=1 DEFAULT_FLOAT=bfloat16 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=1B MODEL=llama3 python3 examples/mlperf/model_train.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -568,12 +469,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 +486,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 +529,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 +565,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 +605,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 +634,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 +648,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 +662,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 +690,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 +815,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 +865,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 +891,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 +899,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 +910,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 +934,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 +945,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

.pre-commit-config.yaml

@@ -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

.pylintrc

@@ -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

README.md

@@ -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)

autogen_stubs.sh

@@ -414,10 +414,29 @@ generate_sqtt() {
   clang2py -k cdefstum \
     extra/sqtt/sqtt.h \
     -o $BASE/sqtt.py
-
   fixup $BASE/sqtt.py
   sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
   python3 -c "import tinygrad.runtime.autogen.sqtt"
+
+  ROCPROF_COMMIT_HASH=dd0485100971522cc4cd8ae136bdda431061a04d
+  ROCPROF_SRC=/tmp/rocprof-trace-decoder-$ROCPROF_COMMIT_HASH
+  if [ ! -d "$ROCPROF_SRC" ]; then
+    git clone https://github.com/ROCm/rocprof-trace-decoder $ROCPROF_SRC
+    pushd .
+    cd $ROCPROF_SRC
+    git reset --hard $ROCPROF_COMMIT_HASH
+    popd
+  fi
+
+  clang2py -k cdefstum \
+    $ROCPROF_SRC/include/rocprof_trace_decoder.h \
+    $ROCPROF_SRC/include/trace_decoder_instrument.h \
+    $ROCPROF_SRC/include/trace_decoder_types.h \
+    -o extra/sqtt/rocprof/rocprof.py
+  fixup extra/sqtt/rocprof/rocprof.py
+  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes, tinygrad.helpers.fetch as tgfetch/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(tgfetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
 }
 
 generate_webgpu() {

docs/abstractions2.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,7 @@ 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.schedule.rangeify import get_rangeify_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +93,10 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))
 
 # group the computation into kernels
-becomes_map = get_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

docs/developer/layout.md

@@ -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

docs/env_vars.md

@@ -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)

docs/runtime.md

@@ -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)

examples/beautiful_cifar.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()

examples/benchmark_onnx.py

@@ -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):

examples/compile_tensorflow.py

@@ -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

examples/gpt2.py

@@ -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 ****

examples/hlb_cifar10.py

@@ -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:

examples/llama.py

@@ -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)

examples/llama3.py

@@ -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):
 

examples/mamba.py

@@ -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"

examples/mlperf/dataloader.py

@@ -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"

examples/mlperf/initializers.py

@@ -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

examples/mlperf/lr_schedulers.py

@@ -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)

examples/mlperf/model_eval.py

@@ -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

examples/mlperf/model_train.py

@@ -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}")

examples/mlperf/scripts/stable_diffusion_downloads.sh

@@ -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

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/stable_diffusion/implementations/tinybox_8xMI300X/dev_run.sh

@@ -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

examples/openpilot/compile3.py

@@ -1,4 +1,4 @@
-import os, sys, pickle, time
+import os, sys, pickle, time, re
 import numpy as np
 if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
 if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
@@ -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"
@@ -52,6 +52,8 @@ def compile(onnx_file):
       kernel_count += 1
       read_image_count += ei.prg.p.src.count("read_image")
       gated_read_image_count += ei.prg.p.src.count("?read_image")
+      for v in [m.group(1) for m in re.finditer(r'(val\d+)\s*=\s*read_imagef\(', ei.prg.p.src)]:
+        if len(re.findall(fr'[\?\:]{v}\.[xyzw]', ei.prg.p.src)) > 0: gated_read_image_count += 1
   print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
   if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
     assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
@@ -77,13 +79,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 ****")

examples/openpilot/compile4.py

@@ -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)

examples/other_mnist/beautiful_mnist_torch.py

@@ -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"))

examples/qwq.py

@@ -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)

examples/sdxl.py

@@ -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")

examples/stable_diffusion.py

@@ -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)

examples/train_resnet.py

@@ -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),

examples/whisper.py

@@ -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)
 

examples/yolov8-onnx.py

@@ -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")

extra/amdpci/proclogs.py

@@ -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()

extra/archprobe.py

@@ -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

extra/assembly/assembly_rdna.py

@@ -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

extra/assembly/rocm/rdna3/asm.py

@@ -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
 

extra/backends/clang_graph.py

@@ -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)

extra/backends/hsa_graph.py

@@ -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:

extra/backends/rdna.py

@@ -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:

extra/export_model.py

@@ -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()}

extra/gemm/amd_uop_matmul.py

@@ -49,8 +49,7 @@ def rangeify_kernel3():
   b = Tensor.empty(N,N)
   c = a@b
   #c = c.reshape((32,2,16,4,32,2,16,4)).contiguous()
-  with Context(RANGEIFY=1):
-    sink = c.schedule()[-1].ast
+  sink = c.schedule()[-1].ast
   #print(sink)
 
   opts  = [Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.LOCAL, 0, 16), Opt(OptOps.UPCAST, 0, 2)]
@@ -65,7 +64,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 +185,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 +196,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 +255,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 +268,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 +309,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
@@ -329,7 +328,7 @@ if __name__ == "__main__":
   elif HL == 1: hprg = hl_spec_kernel3()
   else: hprg = hand_spec_kernel3()
   if HL == 3:
-    with Context(RANGEIFY=1, BLOCK_REORDER=0):
+    with Context(BLOCK_REORDER=0):
       prg = get_program(hprg, Device.default.renderer)
   else:
     prg = get_program(hprg, Device.default.renderer)

extra/gemm/intel_xmx.py

@@ -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)

extra/gemm/max_matmul.py

@@ -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)

extra/gemm/simple_matmul.py

@@ -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

extra/gemm/tinygrad_nv_matmul.py

@@ -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)

extra/gemm/triton_nv_matmul.py

@@ -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

extra/hcqfuzz/tests/bert.py

@@ -7,7 +7,6 @@ bert_train_params = {
   "GPUS": 6,
   "BS": 96,
   "EVAL_BS": 96,
-  "FUSE_ARANGE": 1,
   "BASEDIR": "/raid/datasets/wiki",
 }
 

extra/hip_gpu_driver/hip_ioctl.py

@@ -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)

extra/huggingface_onnx/run_models.py

@@ -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
 

extra/mcts_search.py

@@ -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)
 

extra/models/clip.py

@@ -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
 

extra/models/inception.py

@@ -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)

extra/models/llama.py

@@ -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()}

extra/models/unet.py

@@ -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)

extra/nv_gpu_driver/nv_ioctl.py

@@ -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

extra/onnx_helpers.py

@@ -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
 

extra/optimization/generate_dataset.sh

@@ -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

extra/optimization/helpers.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
@@ -81,7 +81,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   ret = [float(x) for x in ret]
 
   if use_sts:
-    my_sts = dedup([(x.shape == lin.full_shape, x.real_strides(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
+    my_sts = dedup([(x.shape == lin.full_shape, x.is_expanded(), any(v.mask is not None for v in x.views), len(x.views)) for x in lin.sts])
     assert len(my_sts) < MAX_BUFS
     sts_len = 3 + 5*MAX_DIMS
     for s in my_sts:
@@ -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))

extra/optimization/test_beam_search.py

@@ -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()

extra/qcom_gpu_driver/qcom_opencl_interop.py

@@ -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
 

extra/remu/src/thread.rs

@@ -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)]

extra/replay_pkl.py

@@ -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)

extra/sqtt/README.md

@@ -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.
 

extra/sqtt/disasm.py

@@ -0,0 +1,68 @@
+import ctypes
+from dataclasses import dataclass
+import tinygrad.runtime.autogen.comgr as comgr
+from tinygrad.runtime.support.compiler_amd import check
+
+@dataclass
+class InstrCtx:
+  pc:int=0
+  inst:str=""
+
+@comgr.amd_comgr_create_disassembly_info.argtypes[2]
+def instr_cb(text, user_data):
+  c = ctypes.cast(user_data, ctypes.POINTER(ctypes.py_object)).contents.value
+  c.inst = ctypes.string_at(text).decode("utf-8","replace").strip()
+  return comgr.AMD_COMGR_STATUS_SUCCESS
+
+# nop callback
+@comgr.amd_comgr_create_disassembly_info.argtypes[3]
+def addr_cb(*args): return comgr.AMD_COMGR_STATUS_SUCCESS
+
+def comgr_get_address_table(lib:bytes) -> dict[int, tuple[str, int]]:
+  check(comgr.amd_comgr_create_data(comgr.AMD_COMGR_DATA_KIND_EXECUTABLE, ctypes.byref(data_src:=comgr.amd_comgr_data_t())))
+  lib_buf = ctypes.create_string_buffer(lib, len(lib))
+  check(comgr.amd_comgr_set_data(data_src, len(lib), lib_buf))
+  check(comgr.amd_comgr_get_data_isa_name(data_src, isa_sz:=ctypes.c_size_t(128), isa:=(ctypes.c_char*isa_sz.value)()))
+
+  @comgr.amd_comgr_create_disassembly_info.argtypes[1]
+  def memory_cb(from_addr, to, size, _):
+    base, buf_len = ctypes.addressof(lib_buf), len(lib_buf)
+    start = int(from_addr) - base
+    if start < 0 or start >= buf_len: return 0
+    ctypes.memmove(to, base + start, n:=min(int(size), buf_len - start))
+    return n
+
+  info_src = comgr.amd_comgr_disassembly_info_t()
+  check(comgr.amd_comgr_create_disassembly_info(ctypes.cast(isa, ctypes.POINTER(ctypes.c_char)), memory_cb, instr_cb, addr_cb, info_src))
+
+  @comgr.amd_comgr_iterate_symbols.argtypes[1]
+  def sym_callback(sym, udata):
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_TYPE, ctypes.byref(sym_type:=ctypes.c_int())))
+    if sym_type.value != comgr.AMD_COMGR_SYMBOL_TYPE_FUNC: return comgr.AMD_COMGR_STATUS_SUCCESS
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_VALUE, ctypes.byref(vaddr:=ctypes.c_uint64())))
+    check(comgr.amd_comgr_symbol_get_info(sym, comgr.AMD_COMGR_SYMBOL_INFO_SIZE, ctypes.byref(size:=ctypes.c_uint64())))
+    check(comgr.amd_comgr_map_elf_virtual_address_to_code_object_offset(data_src, vaddr.value, ctypes.byref(offset:=ctypes.c_uint64()),
+                                                                        ctypes.byref(ctypes.c_uint64()), ctypes.byref(nobits:=ctypes.c_bool())))
+    check(nobits.value)
+    base = ctypes.addressof(lib_buf)
+    pc = base + offset.value
+    end = pc + size.value
+    addr_table = ctypes.cast(udata, ctypes.POINTER(ctypes.py_object)).contents.value
+    instr_ref = ctypes.py_object(ctx:=InstrCtx())
+    instr_ptr = ctypes.cast(ctypes.pointer(instr_ref), ctypes.c_void_p)
+    while pc < end:
+      size_read = ctypes.c_uint64(0)
+      ctx.pc = pc
+      st = comgr.amd_comgr_disassemble_instruction(info_src, ctypes.c_uint64(pc), instr_ptr, ctypes.byref(size_read))
+      if st == comgr.AMD_COMGR_STATUS_SUCCESS and size_read.value:
+        rel = (pc - base) - offset.value
+        addr_table[vaddr.value + rel] = (ctx.inst, int(size_read.value))
+        pc += size_read.value
+      else: # don't inf loop if comgr fails
+        b = ctypes.c_ubyte.from_buffer(lib_buf, pc - base).value
+        addr_table[vaddr.value + (pc - base - offset.value)] = (f"DISASSEMBLER ISSUE 0x{b:02x}", 1)
+        pc += 1
+    return comgr.AMD_COMGR_STATUS_SUCCESS
+  addr_table:dict[int, tuple[str, int]] = {}
+  check(comgr.amd_comgr_iterate_symbols(data_src, sym_callback, ctypes.cast(ctypes.pointer(ctypes.py_object(addr_table)), ctypes.c_void_p)))
+  return addr_table

extra/sqtt/rgptool.py

@@ -155,6 +155,7 @@ class RGP:
       device_event = device_events[device]
     sqtt_events = [x for x in profile if isinstance(x, ProfileSQTTEvent) and x.device == device_event.device]
     if len(sqtt_events) == 0: raise RuntimeError(f"Device {device_event.device} doesn't contain SQTT data")
+    device_props = sqtt_events[0].props
     sqtt_itrace_enabled = any([event.itrace for event in sqtt_events])
     sqtt_itrace_masked = not all_same([event.itrace for event in sqtt_events])
     sqtt_itrace_se_mask = functools.reduce(lambda a,b: a|b, [int(event.itrace) << event.se for event in sqtt_events], 0) if sqtt_itrace_masked else 0
@@ -192,14 +193,14 @@ class RGP:
         flags=0,
         trace_shader_core_clock=0x93f05080,
         trace_memory_clock=0x4a723a40,
-        device_id=0x744c,
+        device_id={110000: 0x744c, 110003: 0x7480}[device_props['gfx_target_version']],
         device_revision_id=0xc8,
         vgprs_per_simd=1536,
         sgprs_per_simd=128*16,
-        shader_engines=6,
-        compute_unit_per_shader_engine=16,
-        simd_per_compute_unit=2,
-        wavefronts_per_simd=16,
+        shader_engines=device_props['array_count'] // device_props['simd_arrays_per_engine'],
+        compute_unit_per_shader_engine=device_props['simd_count'] // device_props['simd_per_cu'] // (device_props['array_count'] // device_props['simd_arrays_per_engine']),
+        simd_per_compute_unit=device_props['simd_per_cu'],
+        wavefronts_per_simd=device_props['max_waves_per_simd'],
         minimum_vgpr_alloc=4,
         vgpr_alloc_granularity=8,
         minimum_sgpr_alloc=128,
@@ -218,7 +219,7 @@ class RGP:
         vram_bus_width=384, # 384-bit
         l2_cache_size=6 * 1024 * 1024, # 6 MB
         l1_cache_size=32 * 1024, # 32 KB per SIMD (?)
-        lds_size=65536, # 64 KB per CU
+        lds_size=device_props['lds_size_in_kb'] * 1024,
         gpu_name=b'NAVI31',
         alu_per_clock=0,
         texture_per_clock=0,

extra/sqtt/roc.py

@@ -0,0 +1,95 @@
+import ctypes, pathlib, argparse, pickle, re, functools, dataclasses
+from extra.sqtt.rocprof import rocprof
+from extra.sqtt.disasm import comgr_get_address_table
+from tinygrad.helpers import temp, DEBUG
+from tinygrad.device import ProfileEvent, ProfileProgramEvent
+from tinygrad.runtime.ops_amd import ProfileSQTTEvent
+
+@dataclasses.dataclass
+class InstInfo:
+  typ:str=""
+  inst:str=""
+  hit:int=0
+  lat:int=0
+  stall:int=0
+  def __str__(self): return f"{self.inst:>20} hits:{self.typ:>6} hits:{self.hit:>6} latency:{self.lat:>6} stall:{self.stall:>6}"
+
+  def on_ev(self, ev):
+    self.hit, self.lat, self.stall = self.hit + 1, self.lat + ev.duration, self.stall + ev.stall
+
+class _ROCParseCtx:
+  def __init__(self, sqtt_evs:list[ProfileSQTTEvent], prog_evs:list[ProfileProgramEvent]):
+    self.sqtt_evs, self.prog_evs = iter(sqtt_evs), prog_evs
+    self.wave_events = {}
+
+  def next_sqtt(self): return next(self.sqtt_evs, None)
+  def find_program(self, idx): return self.prog_evs[idx]
+  def get_instr_info(self, idx, exec_addr): return self.disasm_program(idx)[exec_addr - self.find_program(idx).base]
+
+  @functools.lru_cache(None)
+  def disasm_program(self, idx): return comgr_get_address_table(self.find_program(idx).lib)
+
+  def on_occupancy_ev(self, ev):
+    if DEBUG >= 4: print("OCC", ev.time, ev.cu, ev.simd, ev.wave_id, ev.start)
+
+  def on_wave_ev(self, ev):
+    if DEBUG >= 4: print("WAVE", ev.wave_id, ev.cu, ev.simd, ev.contexts, ev.begin_time, ev.end_time)
+
+    asm = {}
+    for j in range(ev.instructions_size):
+      inst_ev = ev.instructions_array[j]
+      inst_typ = rocprof.rocprofiler_thread_trace_decoder_inst_category_t__enumvalues[inst_ev.category]
+      asm.setdefault(inst_ev.pc.address, InstInfo(typ=inst_typ, inst=self.get_instr_info(inst_ev.pc.code_object_id, inst_ev.pc.address)[0]))
+      asm[inst_ev.pc.address].on_ev(inst_ev)
+
+    self.wave_events[(self.find_program(ev.instructions_array[0].pc.code_object_id).name, ev.wave_id, ev.cu, ev.simd)] = asm
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument('--profile', type=pathlib.Path, help='Path to profile', default=pathlib.Path(temp("profile.pkl", append_user=True)))
+  args = parser.parse_args()
+
+  with args.profile.open("rb") as f: profile = pickle.load(f)
+  sqtt_events:list[ProfileSQTTEvent] = []
+  prog_events:list[ProfileProgramEvent] = []
+  for e in profile:
+    if isinstance(e, ProfileSQTTEvent): sqtt_events.append(e)
+    if isinstance(e, ProfileProgramEvent) and e.device.startswith("AMD"): prog_events.append(e)
+
+  ROCParseCtx = _ROCParseCtx(sqtt_events, prog_events)
+
+  @rocprof.rocprof_trace_decoder_se_data_callback_t
+  def copy_cb(buf, buf_size, data_ptr):
+    if (prof:=ROCParseCtx.next_sqtt()) is None: return 0
+    buf[0] = ctypes.cast((ctypes.c_ubyte * len(prof.blob)).from_buffer_copy(prof.blob), ctypes.POINTER(ctypes.c_ubyte))
+    buf_size[0] = len(prof.blob)
+    return len(prof.blob)
+
+  @rocprof.rocprof_trace_decoder_trace_callback_t
+  def trace_cb(record_type, events_ptr, n, data_ptr):
+    match record_type:
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY:
+        for ev in (rocprof.rocprofiler_thread_trace_decoder_occupancy_t * n).from_address(events_ptr): ROCParseCtx.on_occupancy_ev(ev)
+      case rocprof.ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE:
+        for ev in (rocprof.rocprofiler_thread_trace_decoder_wave_t * n).from_address(events_ptr): ROCParseCtx.on_wave_ev(ev)
+      case _:
+        if DEBUG >= 2: print(rocprof.rocprofiler_thread_trace_decoder_record_type_t__enumvalues[record_type], events_ptr, n)
+    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
+
+  @rocprof.rocprof_trace_decoder_isa_callback_t
+  def isa_cb(instr_ptr, mem_size_ptr, size_ptr, pc, data_ptr):
+    instr, mem_size_ptr[0] = ROCParseCtx.get_instr_info(pc.code_object_id, pc.address)
+
+    # this is the number of bytes to next instruction, set to 0 for end_pgm
+    if instr == "s_endpgm": mem_size_ptr[0] = 0
+    if (max_sz:=size_ptr[0]) == 0: return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES
+
+    # truncate the instr if it doesn't fit
+    if (str_sz:=len(instr_bytes:=instr.encode()))+1 > max_sz: str_sz = max_sz
+    ctypes.memmove(instr_ptr, instr_bytes, str_sz)
+    size_ptr[0] = str_sz
+
+    return rocprof.ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS
+
+  rocprof.rocprof_trace_decoder_parse_data(copy_cb, trace_cb, isa_cb, None)
+  print(ROCParseCtx.wave_events.keys())

extra/sqtt/rocprof/rocprof.py

@@ -0,0 +1,656 @@
+# pylint: skip-file
+# mypy: ignore-errors
+# -*- coding: utf-8 -*-
+#
+# TARGET arch is: []
+# WORD_SIZE is: 8
+# POINTER_SIZE is: 8
+# LONGDOUBLE_SIZE is: 16
+#
+import ctypes, tinygrad.helpers.fetch as tgfetch
+
+
+class AsDictMixin:
+    @classmethod
+    def as_dict(cls, self):
+        result = {}
+        if not isinstance(self, AsDictMixin):
+            # not a structure, assume it's already a python object
+            return self
+        if not hasattr(cls, "_fields_"):
+            return result
+        # sys.version_info >= (3, 5)
+        # for (field, *_) in cls._fields_:  # noqa
+        for field_tuple in cls._fields_:  # noqa
+            field = field_tuple[0]
+            if field.startswith('PADDING_'):
+                continue
+            value = getattr(self, field)
+            type_ = type(value)
+            if hasattr(value, "_length_") and hasattr(value, "_type_"):
+                # array
+                if not hasattr(type_, "as_dict"):
+                    value = [v for v in value]
+                else:
+                    type_ = type_._type_
+                    value = [type_.as_dict(v) for v in value]
+            elif hasattr(value, "contents") and hasattr(value, "_type_"):
+                # pointer
+                try:
+                    if not hasattr(type_, "as_dict"):
+                        value = value.contents
+                    else:
+                        type_ = type_._type_
+                        value = type_.as_dict(value.contents)
+                except ValueError:
+                    # nullptr
+                    value = None
+            elif isinstance(value, AsDictMixin):
+                # other structure
+                value = type_.as_dict(value)
+            result[field] = value
+        return result
+
+
+class Structure(ctypes.Structure, AsDictMixin):
+
+    def __init__(self, *args, **kwds):
+        # We don't want to use positional arguments fill PADDING_* fields
+
+        args = dict(zip(self.__class__._field_names_(), args))
+        args.update(kwds)
+        super(Structure, self).__init__(**args)
+
+    @classmethod
+    def _field_names_(cls):
+        if hasattr(cls, '_fields_'):
+            return (f[0] for f in cls._fields_ if not f[0].startswith('PADDING'))
+        else:
+            return ()
+
+    @classmethod
+    def get_type(cls, field):
+        for f in cls._fields_:
+            if f[0] == field:
+                return f[1]
+        return None
+
+    @classmethod
+    def bind(cls, bound_fields):
+        fields = {}
+        for name, type_ in cls._fields_:
+            if hasattr(type_, "restype"):
+                if name in bound_fields:
+                    if bound_fields[name] is None:
+                        fields[name] = type_()
+                    else:
+                        # use a closure to capture the callback from the loop scope
+                        fields[name] = (
+                            type_((lambda callback: lambda *args: callback(*args))(
+                                bound_fields[name]))
+                        )
+                    del bound_fields[name]
+                else:
+                    # default callback implementation (does nothing)
+                    try:
+                        default_ = type_(0).restype().value
+                    except TypeError:
+                        default_ = None
+                    fields[name] = type_((
+                        lambda default_: lambda *args: default_)(default_))
+            else:
+                # not a callback function, use default initialization
+                if name in bound_fields:
+                    fields[name] = bound_fields[name]
+                    del bound_fields[name]
+                else:
+                    fields[name] = type_()
+        if len(bound_fields) != 0:
+            raise ValueError(
+                "Cannot bind the following unknown callback(s) {}.{}".format(
+                    cls.__name__, bound_fields.keys()
+            ))
+        return cls(**fields)
+
+
+class Union(ctypes.Union, AsDictMixin):
+    pass
+
+
+
+c_int128 = ctypes.c_ubyte*16
+c_uint128 = c_int128
+void = None
+if ctypes.sizeof(ctypes.c_longdouble) == 16:
+    c_long_double_t = ctypes.c_longdouble
+else:
+    c_long_double_t = ctypes.c_ubyte*16
+
+def string_cast(char_pointer, encoding='utf-8', errors='strict'):
+    value = ctypes.cast(char_pointer, ctypes.c_char_p).value
+    if value is not None and encoding is not None:
+        value = value.decode(encoding, errors=errors)
+    return value
+
+
+def char_pointer_cast(string, encoding='utf-8'):
+    if encoding is not None:
+        try:
+            string = string.encode(encoding)
+        except AttributeError:
+            # In Python3, bytes has no encode attribute
+            pass
+    string = ctypes.c_char_p(string)
+    return ctypes.cast(string, ctypes.POINTER(ctypes.c_char))
+
+
+
+class FunctionFactoryStub:
+    def __getattr__(self, _):
+      return ctypes.CFUNCTYPE(lambda y:y)
+
+# libraries['FIXME_STUB'] explanation
+# As you did not list (-l libraryname.so) a library that exports this function
+# This is a non-working stub instead.
+# You can either re-run clan2py with -l /path/to/library.so
+# Or manually fix this by comment the ctypes.CDLL loading
+_libraries = {}
+_libraries['FIXME_STUB'] = ctypes.CDLL(str(tgfetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
+
+
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_info_t'
+rocprofiler_thread_trace_decoder_info_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE = 0
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST = 1
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE = 2
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE = 3
+ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST = 4
+rocprofiler_thread_trace_decoder_info_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_pc_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_pc_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_pc_t._fields_ = [
+    ('address', ctypes.c_uint64),
+    ('code_object_id', ctypes.c_uint64),
+]
+
+rocprofiler_thread_trace_decoder_pc_t = struct_rocprofiler_thread_trace_decoder_pc_t
+class struct_rocprofiler_thread_trace_decoder_perfevent_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_perfevent_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_perfevent_t._fields_ = [
+    ('time', ctypes.c_int64),
+    ('events0', ctypes.c_uint16),
+    ('events1', ctypes.c_uint16),
+    ('events2', ctypes.c_uint16),
+    ('events3', ctypes.c_uint16),
+    ('CU', ctypes.c_ubyte),
+    ('bank', ctypes.c_ubyte),
+    ('PADDING_0', ctypes.c_ubyte * 6),
+]
+
+rocprofiler_thread_trace_decoder_perfevent_t = struct_rocprofiler_thread_trace_decoder_perfevent_t
+class struct_rocprofiler_thread_trace_decoder_occupancy_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_occupancy_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_occupancy_t._fields_ = [
+    ('pc', rocprofiler_thread_trace_decoder_pc_t),
+    ('time', ctypes.c_uint64),
+    ('reserved', ctypes.c_ubyte),
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('start', ctypes.c_uint32, 1),
+    ('_rsvd', ctypes.c_uint32, 31),
+]
+
+rocprofiler_thread_trace_decoder_occupancy_t = struct_rocprofiler_thread_trace_decoder_occupancy_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_wstate_type_t'
+rocprofiler_thread_trace_decoder_wstate_type_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY = 0
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE = 1
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC = 2
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT = 3
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL = 4
+ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST = 5
+rocprofiler_thread_trace_decoder_wstate_type_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_wave_state_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_wave_state_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_wave_state_t._fields_ = [
+    ('type', ctypes.c_int32),
+    ('duration', ctypes.c_int32),
+]
+
+rocprofiler_thread_trace_decoder_wave_state_t = struct_rocprofiler_thread_trace_decoder_wave_state_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_inst_category_t'
+rocprofiler_thread_trace_decoder_inst_category_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
+    6: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
+    7: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
+    8: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
+    9: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
+    10: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
+    11: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
+    12: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
+    13: 'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE = 0
+ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM = 1
+ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU = 2
+ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM = 3
+ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT = 4
+ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS = 5
+ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU = 6
+ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP = 7
+ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT = 8
+ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED = 9
+ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT = 10
+ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE = 11
+ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH = 12
+ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST = 13
+rocprofiler_thread_trace_decoder_inst_category_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_inst_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_inst_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_inst_t._fields_ = [
+    ('category', ctypes.c_uint32, 8),
+    ('stall', ctypes.c_uint32, 24),
+    ('duration', ctypes.c_int32),
+    ('time', ctypes.c_int64),
+    ('pc', rocprofiler_thread_trace_decoder_pc_t),
+]
+
+rocprofiler_thread_trace_decoder_inst_t = struct_rocprofiler_thread_trace_decoder_inst_t
+class struct_rocprofiler_thread_trace_decoder_wave_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_wave_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_wave_t._fields_ = [
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('contexts', ctypes.c_ubyte),
+    ('_rsvd1', ctypes.c_uint32),
+    ('_rsvd2', ctypes.c_uint32),
+    ('_rsvd3', ctypes.c_uint32),
+    ('begin_time', ctypes.c_int64),
+    ('end_time', ctypes.c_int64),
+    ('timeline_size', ctypes.c_uint64),
+    ('instructions_size', ctypes.c_uint64),
+    ('timeline_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_wave_state_t)),
+    ('instructions_array', ctypes.POINTER(struct_rocprofiler_thread_trace_decoder_inst_t)),
+]
+
+rocprofiler_thread_trace_decoder_wave_t = struct_rocprofiler_thread_trace_decoder_wave_t
+class struct_rocprofiler_thread_trace_decoder_realtime_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_realtime_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_realtime_t._fields_ = [
+    ('shader_clock', ctypes.c_int64),
+    ('realtime_clock', ctypes.c_uint64),
+    ('reserved', ctypes.c_uint64),
+]
+
+rocprofiler_thread_trace_decoder_realtime_t = struct_rocprofiler_thread_trace_decoder_realtime_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_shaderdata_flags_t'
+rocprofiler_thread_trace_decoder_shaderdata_flags_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM = 0
+ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV = 1
+rocprofiler_thread_trace_decoder_shaderdata_flags_t = ctypes.c_uint32 # enum
+class struct_rocprofiler_thread_trace_decoder_shaderdata_t(Structure):
+    pass
+
+struct_rocprofiler_thread_trace_decoder_shaderdata_t._pack_ = 1 # source:False
+struct_rocprofiler_thread_trace_decoder_shaderdata_t._fields_ = [
+    ('time', ctypes.c_int64),
+    ('value', ctypes.c_uint64),
+    ('cu', ctypes.c_ubyte),
+    ('simd', ctypes.c_ubyte),
+    ('wave_id', ctypes.c_ubyte),
+    ('flags', ctypes.c_ubyte),
+    ('reserved', ctypes.c_uint32),
+]
+
+rocprofiler_thread_trace_decoder_shaderdata_t = struct_rocprofiler_thread_trace_decoder_shaderdata_t
+
+# values for enumeration 'rocprofiler_thread_trace_decoder_record_type_t'
+rocprofiler_thread_trace_decoder_record_type_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
+    6: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
+    7: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
+    8: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
+    9: 'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP = 0
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY = 1
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT = 2
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE = 3
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO = 4
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG = 5
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA = 6
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME = 7
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY = 8
+ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST = 9
+rocprofiler_thread_trace_decoder_record_type_t = ctypes.c_uint32 # enum
+
+# values for enumeration 'c__EA_rocprofiler_thread_trace_decoder_status_t'
+c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues = {
+    0: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
+    1: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
+    2: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
+    3: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
+    4: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
+    5: 'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
+}
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS = 0
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR = 1
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES = 2
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT = 3
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA = 4
+ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST = 5
+c__EA_rocprofiler_thread_trace_decoder_status_t = ctypes.c_uint32 # enum
+rocprofiler_thread_trace_decoder_status_t = c__EA_rocprofiler_thread_trace_decoder_status_t
+rocprofiler_thread_trace_decoder_status_t__enumvalues = c__EA_rocprofiler_thread_trace_decoder_status_t__enumvalues
+rocprof_trace_decoder_trace_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, rocprofiler_thread_trace_decoder_record_type_t, ctypes.POINTER(None), ctypes.c_uint64, ctypes.POINTER(None))
+rocprof_trace_decoder_isa_callback_t = ctypes.CFUNCTYPE(c__EA_rocprofiler_thread_trace_decoder_status_t, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint64), struct_rocprofiler_thread_trace_decoder_pc_t, ctypes.POINTER(None))
+rocprof_trace_decoder_se_data_callback_t = ctypes.CFUNCTYPE(ctypes.c_uint64, ctypes.POINTER(ctypes.POINTER(ctypes.c_ubyte)), ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(None))
+try:
+    rocprof_trace_decoder_parse_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_parse_data
+    rocprof_trace_decoder_parse_data.restype = rocprofiler_thread_trace_decoder_status_t
+    rocprof_trace_decoder_parse_data.argtypes = [rocprof_trace_decoder_se_data_callback_t, rocprof_trace_decoder_trace_callback_t, rocprof_trace_decoder_isa_callback_t, ctypes.POINTER(None)]
+except AttributeError:
+    pass
+try:
+    rocprof_trace_decoder_get_info_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_info_string
+    rocprof_trace_decoder_get_info_string.restype = ctypes.POINTER(ctypes.c_char)
+    rocprof_trace_decoder_get_info_string.argtypes = [rocprofiler_thread_trace_decoder_info_t]
+except AttributeError:
+    pass
+try:
+    rocprof_trace_decoder_get_status_string = _libraries['FIXME_STUB'].rocprof_trace_decoder_get_status_string
+    rocprof_trace_decoder_get_status_string.restype = ctypes.POINTER(ctypes.c_char)
+    rocprof_trace_decoder_get_status_string.argtypes = [rocprofiler_thread_trace_decoder_status_t]
+except AttributeError:
+    pass
+rocprofiler_thread_trace_decoder_debug_callback_t = ctypes.CFUNCTYPE(None, ctypes.c_int64, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char), ctypes.POINTER(None))
+uint64_t = ctypes.c_uint64
+try:
+    rocprof_trace_decoder_dump_data = _libraries['FIXME_STUB'].rocprof_trace_decoder_dump_data
+    rocprof_trace_decoder_dump_data.restype = rocprofiler_thread_trace_decoder_status_t
+    rocprof_trace_decoder_dump_data.argtypes = [ctypes.POINTER(ctypes.c_char), uint64_t, rocprofiler_thread_trace_decoder_debug_callback_t, ctypes.POINTER(None)]
+except AttributeError:
+    pass
+class union_rocprof_trace_decoder_gfx9_header_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_gfx9_header_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_gfx9_header_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_gfx9_header_t_0._fields_ = [
+    ('legacy_version', ctypes.c_uint64, 13),
+    ('gfx9_version2', ctypes.c_uint64, 3),
+    ('DSIMDM', ctypes.c_uint64, 4),
+    ('DCU', ctypes.c_uint64, 5),
+    ('reserved1', ctypes.c_uint64, 1),
+    ('SEID', ctypes.c_uint64, 6),
+    ('reserved2', ctypes.c_uint64, 32),
+]
+
+union_rocprof_trace_decoder_gfx9_header_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_gfx9_header_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_gfx9_header_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_gfx9_header_t_0),
+    ('raw', ctypes.c_uint64),
+]
+
+rocprof_trace_decoder_gfx9_header_t = union_rocprof_trace_decoder_gfx9_header_t
+class union_rocprof_trace_decoder_instrument_enable_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_instrument_enable_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_instrument_enable_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_instrument_enable_t_0._fields_ = [
+    ('char1', ctypes.c_uint32, 8),
+    ('char2', ctypes.c_uint32, 8),
+    ('char3', ctypes.c_uint32, 8),
+    ('char4', ctypes.c_uint32, 8),
+]
+
+union_rocprof_trace_decoder_instrument_enable_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_instrument_enable_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_instrument_enable_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_instrument_enable_t_0),
+    ('u32All', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_instrument_enable_t = union_rocprof_trace_decoder_instrument_enable_t
+class union_rocprof_trace_decoder_packet_header_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_packet_header_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_packet_header_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_packet_header_t_0._fields_ = [
+    ('opcode', ctypes.c_uint32, 8),
+    ('type', ctypes.c_uint32, 4),
+    ('data20', ctypes.c_uint32, 20),
+]
+
+union_rocprof_trace_decoder_packet_header_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_packet_header_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_packet_header_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_packet_header_t_0),
+    ('u32All', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_packet_header_t = union_rocprof_trace_decoder_packet_header_t
+
+# values for enumeration 'rocprof_trace_decoder_packet_opcode_t'
+rocprof_trace_decoder_packet_opcode_t__enumvalues = {
+    4: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
+    5: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
+    6: 'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
+}
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ = 4
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP = 5
+ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO = 6
+rocprof_trace_decoder_packet_opcode_t = ctypes.c_uint32 # enum
+
+# values for enumeration 'rocprof_trace_decoder_agent_info_type_t'
+rocprof_trace_decoder_agent_info_type_t__enumvalues = {
+    0: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
+    1: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
+    2: 'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
+}
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ = 0
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL = 1
+ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST = 2
+rocprof_trace_decoder_agent_info_type_t = ctypes.c_uint32 # enum
+class union_rocprof_trace_decoder_codeobj_marker_tail_t(Union):
+    pass
+
+class struct_rocprof_trace_decoder_codeobj_marker_tail_t_0(Structure):
+    pass
+
+struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._pack_ = 1 # source:False
+struct_rocprof_trace_decoder_codeobj_marker_tail_t_0._fields_ = [
+    ('isUnload', ctypes.c_uint32, 1),
+    ('bFromStart', ctypes.c_uint32, 1),
+    ('legacy_id', ctypes.c_uint32, 30),
+]
+
+union_rocprof_trace_decoder_codeobj_marker_tail_t._pack_ = 1 # source:False
+union_rocprof_trace_decoder_codeobj_marker_tail_t._anonymous_ = ('_0',)
+union_rocprof_trace_decoder_codeobj_marker_tail_t._fields_ = [
+    ('_0', struct_rocprof_trace_decoder_codeobj_marker_tail_t_0),
+    ('raw', ctypes.c_uint32),
+]
+
+rocprof_trace_decoder_codeobj_marker_tail_t = union_rocprof_trace_decoder_codeobj_marker_tail_t
+
+# values for enumeration 'rocprof_trace_decoder_codeobj_marker_type_t'
+rocprof_trace_decoder_codeobj_marker_type_t__enumvalues = {
+    0: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
+    1: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
+    2: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
+    3: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
+    4: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
+    5: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
+    6: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
+    7: 'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
+}
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL = 0
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO = 1
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO = 2
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI = 3
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI = 4
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO = 5
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI = 6
+ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST = 7
+rocprof_trace_decoder_codeobj_marker_type_t = ctypes.c_uint32 # enum
+__all__ = \
+    ['ROCPROFILER_THREAD_TRACE_DECODER_INFO_DATA_LOST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_NONE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_STITCH_INCOMPLETE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INFO_WAVE_INCOMPLETE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_BVH',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_CONTEXT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_FLAT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_IMMED',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_JUMP',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_LDS',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_MESSAGE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NEXT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_NONE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SALU',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_SMEM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VALU',
+    'ROCPROFILER_THREAD_TRACE_DECODER_INST_VMEM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_DEBUG',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_GFXIP',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_INFO',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_OCCUPANCY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_PERFEVENT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_REALTIME',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_RT_FREQUENCY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_SHADERDATA',
+    'ROCPROFILER_THREAD_TRACE_DECODER_RECORD_WAVE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_IMM',
+    'ROCPROFILER_THREAD_TRACE_DECODER_SHADERDATA_FLAGS_PRIV',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_ARGUMENT',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_INVALID_SHADER_DATA',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_ERROR_OUT_OF_RESOURCES',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_STATUS_SUCCESS',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EMPTY',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_EXEC',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_IDLE',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_LAST',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_STALL',
+    'ROCPROFILER_THREAD_TRACE_DECODER_WSTATE_WAIT',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_COUNTER_INTERVAL',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_LAST',
+    'ROCPROF_TRACE_DECODER_AGENT_INFO_TYPE_RT_FREQUENCY_KHZ',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ADDR_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_ID_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_LAST',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_HI',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_SIZE_LO',
+    'ROCPROF_TRACE_DECODER_CODEOBJ_MARKER_TYPE_TAIL',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_AGENT_INFO',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_CODEOBJ',
+    'ROCPROF_TRACE_DECODER_PACKET_OPCODE_RT_TIMESTAMP',
+    'c__EA_rocprofiler_thread_trace_decoder_status_t',
+    'rocprof_trace_decoder_agent_info_type_t',
+    'rocprof_trace_decoder_codeobj_marker_tail_t',
+    'rocprof_trace_decoder_codeobj_marker_type_t',
+    'rocprof_trace_decoder_dump_data',
+    'rocprof_trace_decoder_get_info_string',
+    'rocprof_trace_decoder_get_status_string',
+    'rocprof_trace_decoder_gfx9_header_t',
+    'rocprof_trace_decoder_instrument_enable_t',
+    'rocprof_trace_decoder_isa_callback_t',
+    'rocprof_trace_decoder_packet_header_t',
+    'rocprof_trace_decoder_packet_opcode_t',
+    'rocprof_trace_decoder_parse_data',
+    'rocprof_trace_decoder_se_data_callback_t',
+    'rocprof_trace_decoder_trace_callback_t',
+    'rocprofiler_thread_trace_decoder_debug_callback_t',
+    'rocprofiler_thread_trace_decoder_info_t',
+    'rocprofiler_thread_trace_decoder_inst_category_t',
+    'rocprofiler_thread_trace_decoder_inst_t',
+    'rocprofiler_thread_trace_decoder_occupancy_t',
+    'rocprofiler_thread_trace_decoder_pc_t',
+    'rocprofiler_thread_trace_decoder_perfevent_t',
+    'rocprofiler_thread_trace_decoder_realtime_t',
+    'rocprofiler_thread_trace_decoder_record_type_t',
+    'rocprofiler_thread_trace_decoder_shaderdata_flags_t',
+    'rocprofiler_thread_trace_decoder_shaderdata_t',
+    'rocprofiler_thread_trace_decoder_status_t',
+    'rocprofiler_thread_trace_decoder_status_t__enumvalues',
+    'rocprofiler_thread_trace_decoder_wave_state_t',
+    'rocprofiler_thread_trace_decoder_wave_t',
+    'rocprofiler_thread_trace_decoder_wstate_type_t',
+    'struct_rocprof_trace_decoder_codeobj_marker_tail_t_0',
+    'struct_rocprof_trace_decoder_gfx9_header_t_0',
+    'struct_rocprof_trace_decoder_instrument_enable_t_0',
+    'struct_rocprof_trace_decoder_packet_header_t_0',
+    'struct_rocprofiler_thread_trace_decoder_inst_t',
+    'struct_rocprofiler_thread_trace_decoder_occupancy_t',
+    'struct_rocprofiler_thread_trace_decoder_pc_t',
+    'struct_rocprofiler_thread_trace_decoder_perfevent_t',
+    'struct_rocprofiler_thread_trace_decoder_realtime_t',
+    'struct_rocprofiler_thread_trace_decoder_shaderdata_t',
+    'struct_rocprofiler_thread_trace_decoder_wave_state_t',
+    'struct_rocprofiler_thread_trace_decoder_wave_t', 'uint64_t',
+    'union_rocprof_trace_decoder_codeobj_marker_tail_t',
+    'union_rocprof_trace_decoder_gfx9_header_t',
+    'union_rocprof_trace_decoder_instrument_enable_t',
+    'union_rocprof_trace_decoder_packet_header_t']

extra/test_pyrender.py

@@ -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

extra/thneed.py

@@ -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")
 

extra/thunder/cuda/include/common/base_ops.cuh

@@ -0,0 +1,400 @@
+/**
+ * @file
+ * @brief Basic operations on generic types.
+ */
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <limits>
+#include "base_types.cuh"
+
+namespace kittens {
+
+/**
+ * @namespace base_ops
+ *
+ * @brief A namespace for operations on basic data types.
+ */
+namespace base_ops {
+
+/* ----------  CONST OPS  ---------- */
+
+/**
+ * @brief Represents the zero constant operation.
+ *
+ * This operation returns the zero value of the specified type.
+ *
+ * @tparam T The data type for which to return the zero value.
+ * @return The zero value of type T.
+ */
+struct zero {
+    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::zero();      }
+};
+/**
+ * @brief Represents the one constant operation.
+ *
+ * This operation returns the one value of the specified type.
+ *
+ * @tparam T The data type for which to return the one value.
+ * @return The one value of type T.
+ */
+struct one {
+    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::one();       }
+};
+/**
+ * @brief Represents the positive infinity constant operation.
+ *
+ * This operation returns the positive infinity value of the specified type.
+ *
+ * @tparam T The data type for which to return the positive infinity value.
+ * @return The positive infinity value of type T.
+ */
+struct pos_infty {
+    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::pos_infty(); }
+};
+/**
+ * @brief Represents the negative infinity constant operation.
+ *
+ * This operation returns the negative infinity value of the specified type.
+ *
+ * @tparam T The data type for which to return the negative infinity value.
+ * @return The negative infinity value of type T.
+ */
+struct neg_infty {
+    template<typename T, typename... args> __device__ static inline constexpr T op(args... _) { return base_types::constants<T>::neg_infty(); }
+};
+
+
+/* ----------  UNARY OPS  ---------- */
+
+/**
+ * @brief Exponential function operation.
+ *
+ * This operation calculates the exponential of the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The exponential of the input value.
+ */
+struct exp {
+    template<typename T> static __device__ inline T op(const T &x) { return exp(x); }
+};
+template<> __device__ inline float  exp::op<float> (const float &x ) { return __expf(x);                        }
+template<> __device__ inline float2 exp::op<float2>(const float2 &x) { return float2{__expf(x.x), __expf(x.y)}; }
+template<> __device__ inline bf16   exp::op<bf16>  (const bf16 &x  ) { return hexp(x);                          }
+template<> __device__ inline bf16_2 exp::op<bf16_2>(const bf16_2 &x) { return h2exp(x);                         }
+template<> __device__ inline half   exp::op<half>  (const half &x  ) { return hexp(x);                          }
+template<> __device__ inline half_2 exp::op<half_2>(const half_2 &x) { return h2exp(x);                         }
+
+/**
+ * @brief Exponential function operation, in base 2
+ *
+ * This operation calculates the exponential of the input value, in base 2.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The exponential of the input value.
+ */
+struct exp2 {
+    template<typename T> static __device__ inline T op(const T &x) { return exp2f(x); }
+};
+template<> __device__ inline float  exp2::op<float> (const float &x ) { return exp2f(x);                        }
+template<> __device__ inline float2 exp2::op<float2>(const float2 &x) { return float2{exp2f(x.x), exp2f(x.y)}; }
+template<> __device__ inline bf16   exp2::op<bf16>  (const bf16 &x  ) { return hexp2(x);                          }
+template<> __device__ inline bf16_2 exp2::op<bf16_2>(const bf16_2 &x) { return h2exp2(x);                         }
+template<> __device__ inline half   exp2::op<half>  (const half &x  ) { return hexp2(x);                          }
+template<> __device__ inline half_2 exp2::op<half_2>(const half_2 &x) { return h2exp2(x);                         }
+/**
+ * @brief Natural log function operation.
+ *
+ * This operation calculates the natural logarithm of the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The natural logarithm of the input value.
+ */
+struct log {
+    template<typename T> static __device__ inline T op(const T &x) { return log(x); }
+};
+template<> __device__ inline float  log::op<float> (const float &x ) { return __logf(x);                        }
+template<> __device__ inline float2 log::op<float2>(const float2 &x) { return float2{__logf(x.x), __logf(x.y)}; }
+template<> __device__ inline bf16   log::op<bf16>  (const bf16 &x  ) { return hlog(x);                          }
+template<> __device__ inline bf16_2 log::op<bf16_2>(const bf16_2 &x) { return h2log(x);                         }
+template<> __device__ inline half   log::op<half>  (const half &x  ) { return hlog(x);                          }
+template<> __device__ inline half_2 log::op<half_2>(const half_2 &x) { return h2log(x);                         }
+/**
+ * @brief Logarithm base 2 operation.
+ *
+ * This operation calculates the logarithm base 2 of the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The logarithm base 2 of the input value.
+ */
+struct log2 {
+    template<typename T> static __device__ inline T op(const T &x) { return log2(x); }
+};
+template<> __device__ inline float  log2::op<float> (const float &x ) { return __log2f(x);                        }
+template<> __device__ inline float2 log2::op<float2>(const float2 &x) { return float2{__log2f(x.x), __log2f(x.y)}; }
+template<> __device__ inline bf16   log2::op<bf16>  (const bf16 &x  ) { return hlog2(x);                          }
+template<> __device__ inline bf16_2 log2::op<bf16_2>(const bf16_2 &x) { return h2log2(x);                         }
+template<> __device__ inline half   log2::op<half>  (const half &x  ) { return hlog2(x);                          }
+template<> __device__ inline half_2 log2::op<half_2>(const half_2 &x) { return h2log2(x);                         }
+/**
+ * @brief Absolute value operation.
+ *
+ * This operation calculates the absolute value of the input.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The absolute value of the input.
+ */
+struct abs {
+    template<typename T> static __device__ inline T op(const T &x) { return abs(x); }
+};
+template<> __device__ inline float  abs::op<float> (const float &x ) { return fabsf(x);                       }
+template<> __device__ inline float2 abs::op<float2>(const float2 &x) { return float2{fabsf(x.x), fabsf(x.y)}; }
+template<> __device__ inline bf16   abs::op<bf16>  (const bf16 &x  ) { return __habs(x);                      }
+template<> __device__ inline bf16_2 abs::op<bf16_2>(const bf16_2 &x) { return __habs2(x);                     }
+template<> __device__ inline half   abs::op<half>  (const half &x  ) { return __habs(x);                      }
+template<> __device__ inline half_2 abs::op<half_2>(const half_2 &x) { return __habs2(x);                     }
+/**
+ * @brief Rectified Linear Unit (ReLU) operation.
+ *
+ * This operation applies the ReLU function to the input, which is the
+ * maximum of zero and the input value.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param x[in] The input value.
+ * @return The result of ReLU function applied to the input.
+ */
+struct relu {
+    template<typename T> static __device__ inline T op(const T &x) { return max(x, base_types::constants<T>::zero()); }
+};
+template<> __device__ inline float  relu::op<float> (const float &x ) { return max(x, 0.f);                                  }
+template<> __device__ inline float2 relu::op<float2>(const float2 &x) { return float2{max(x.x, 0.f), max(x.y, 0.f)};         }
+template<> __device__ inline bf16   relu::op<bf16>  (const bf16 &x  ) { return __hmax(x, base_types::constants<bf16>::zero());    }
+template<> __device__ inline bf16_2 relu::op<bf16_2>(const bf16_2 &x) { return __hmax2(x, base_types::constants<bf16_2>::zero()); }
+template<> __device__ inline half   relu::op<half>  (const half &x  ) { return __hmax(x, base_types::constants<half>::zero());    }
+template<> __device__ inline half_2 relu::op<half_2>(const half_2 &x) { return __hmax2(x, base_types::constants<half_2>::zero()); }
+/**
+ * @brief Copy operation.
+ *
+ * This operation returns the input value unchanged.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The input value.
+ * @return The same value as the input.
+ */
+struct copy { // for non-compile-time setters.
+    template<typename T> static __device__ inline T op(const T &a) { return a; }
+};
+
+
+/* ----------  BINARY OPS  ---------- */
+
+/**
+ * @brief Copy2 operation.
+ *
+ * This operation returns the second input value unchanged.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value (ignored).
+ * @param b[in] The second input value.
+ * @return The same value as the second input.
+ */
+struct copy2 { // this turns out to be a slightly hacky op that makes some code cleaner :/
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return b; }
+};
+/**
+ * @brief Sum operation.
+ *
+ * This operation calculates the sum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The sum of the input values.
+ */
+struct sum {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a+b; }
+};
+template<> __device__ inline float2 sum::op<float2>(const float2 &a, const float2 &b) {
+#ifdef KITTENS_BLACKWELL
+    float2 c;
+    asm volatile("add.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
+    return c;
+#else
+    return float2{a.x+b.x, a.y+b.y};
+#endif
+}
+template<> __device__ inline bf16   sum::op<bf16>  (const bf16   &a, const bf16   &b) { return __hadd(a, b);             }
+template<> __device__ inline bf16_2 sum::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hadd2(a, b);            }
+template<> __device__ inline half   sum::op<half>  (const half   &a, const half   &b) { return __hadd(a, b);             }
+template<> __device__ inline half_2 sum::op<half_2>(const half_2 &a, const half_2 &b) { return __hadd2(a, b);            }
+/**
+ * @brief Subtraction operation.
+ *
+ * This operation calculates the difference between two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The difference between the input values.
+ */
+struct sub {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a-b; }
+};
+template<> __device__ inline float2 sub::op<float2>(const float2 &a, const float2 &b) { 
+#ifdef KITTENS_BLACKWELL
+    float2 c;
+    asm volatile("sub.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
+    return c;
+#else
+    return float2{a.x-b.x, a.y-b.y}; 
+#endif
+}
+template<> __device__ inline bf16   sub::op<bf16>  (const bf16   &a, const bf16   &b) { return __hsub(a, b);             }
+template<> __device__ inline bf16_2 sub::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hsub2(a, b);            }
+template<> __device__ inline half   sub::op<half>  (const half   &a, const half   &b) { return __hsub(a, b);             }
+template<> __device__ inline half_2 sub::op<half_2>(const half_2 &a, const half_2 &b) { return __hsub2(a, b);            }
+/**
+ * @brief Multiplication operation.
+ *
+ * This operation calculates the product of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The product of the input values.
+ */
+struct mul {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a*b; }
+};
+template<> __device__ inline float2 mul::op<float2>(const float2 &a, const float2 &b) { 
+#ifdef KITTENS_BLACKWELL
+    float2 c;
+    asm volatile("mul.f32x2 %0, %1, %2;" : "=l"(*(uint64_t*)&c) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b));
+    return c;
+#else
+    return float2{a.x*b.x, a.y*b.y}; 
+#endif
+}
+template<> __device__ inline bf16   mul::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmul(a, b);             }
+template<> __device__ inline bf16_2 mul::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmul2(a, b);            }
+template<> __device__ inline half   mul::op<half>  (const half   &a, const half   &b) { return __hmul(a, b);             }
+template<> __device__ inline half_2 mul::op<half_2>(const half_2 &a, const half_2 &b) { return __hmul2(a, b);            }
+/**
+ * @brief Division operation.
+ *
+ * This operation calculates the quotient of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The quotient of the input values.
+ */
+struct div {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return a/b; }
+};
+template<> __device__ inline float2 div::op<float2>(const float2 &a, const float2 &b) { return float2{a.x/b.x, a.y/b.y}; }
+template<> __device__ inline bf16   div::op<bf16>  (const bf16   &a, const bf16   &b) { return __hdiv(a, b);             }
+template<> __device__ inline bf16_2 div::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __h2div(a, b);            } // this op is a special snowflake
+template<> __device__ inline half   div::op<half>  (const half   &a, const half   &b) { return __hdiv(a, b);             }
+template<> __device__ inline half_2 div::op<half_2>(const half_2 &a, const half_2 &b) { return __h2div(a, b);            }
+/**
+ * @brief Maximum operation.
+ *
+ * This operation calculates the maximum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The maximum of the input values.
+ */
+ struct max {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::max(a, b); }
+};
+template<>  __device__ inline float2 max::op<float2>(const float2 &a, const float2 &b) { return float2{::max(a.x, b.x), ::max(a.y, b.y)}; }
+template<>  __device__ inline bf16   max::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmax(a, b);                             }
+template<>  __device__ inline bf16_2 max::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmax2(a, b);                            }
+template<>  __device__ inline half   max::op<half>  (const half   &a, const half   &b) { return __hmax(a, b);                             }
+template<>  __device__ inline half_2 max::op<half_2>(const half_2 &a, const half_2 &b) { return __hmax2(a, b);                            }
+/**
+ * @brief Minimum operation.
+ *
+ * This operation calculates the minimum of two input values.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @return The minimum of the input values.
+ */
+struct min {
+    template<typename T> static __device__ inline T op(const T &a, const T &b) { return ::min(a, b); }
+};
+template<>  __device__ inline float2 min::op<float2>(const float2 &a, const float2 &b) { return float2{::min(a.x, b.x), ::min(a.y, b.y)}; }
+template<>  __device__ inline bf16   min::op<bf16>  (const bf16   &a, const bf16   &b) { return __hmin(a, b);                         }
+template<>  __device__ inline bf16_2 min::op<bf16_2>(const bf16_2 &a, const bf16_2 &b) { return __hmin2(a, b);                        }
+template<>  __device__ inline half   min::op<half>  (const half   &a, const half   &b) { return __hmin(a, b);                         }
+template<>  __device__ inline half_2 min::op<half_2>(const half_2 &a, const half_2 &b) { return __hmin2(a, b);                        }
+
+
+/* ----------  TERNARY OPS  ---------- */
+
+/**
+ * @brief Fused multiply-add operation A * B + C.
+ *
+ * This operation performs a fused multiply-add, computing (A * B) + C with only one rounding.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The second input value.
+ * @param c[in] The third input value to be added.
+ * @return The result of the fused multiply-add operation.
+ */
+struct fma_AxBtC {
+    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
+        return sum::op<T>(mul::op<T>(a, b), c);
+    }
+};
+template<> __device__ inline float2 fma_AxBtC::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
+#ifdef KITTENS_BLACKWELL
+    float2 d;
+    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&b), "l"(*(uint64_t*)&c));
+    return d;
+#else
+    return float2{a.x*b.x+c.x, a.y*b.y+c.y};
+#endif
+}
+/**
+ * @brief Fused multiply-add operation A * C + B.
+ *
+ * This operation performs a fused multiply-add, computing (A * C) + B with only one rounding.
+ * This is particularly useful for attention mechanisms in neural networks.
+ *
+ * @tparam T The data type of the input and output values.
+ * @param a[in] The first input value.
+ * @param b[in] The third input value to be added.
+ * @param c[in] The second input value.
+ * @return The result of the fused multiply-add operation.
+ */
+struct fma_AxCtB { // this is the one needed for attention
+    template<typename T> static __device__ inline T op(const T &a, const T &b, const T &c) {
+        return sum::op<T>(mul::op<T>(a, c), b);
+    }
+};
+template<> __device__ inline float2 fma_AxCtB::op<float2>(const float2 &a, const float2 &b, const float2 &c) {
+#ifdef KITTENS_BLACKWELL
+    float2 d;
+    asm volatile("fma.rn.f32x2 %0, %1, %2, %3;" : "=l"(*(uint64_t*)&d) : "l"(*(uint64_t*)&a), "l"(*(uint64_t*)&c), "l"(*(uint64_t*)&b));
+    return d;
+#else
+    return float2{a.x*c.x+b.x, a.y*c.y+b.y};
+#endif
+}
+
+} // namespace base_ops
+
+} // namespace kittens

extra/thunder/cuda/include/common/base_types.cuh

@@ -0,0 +1,519 @@
+/**
+ * @file
+ * @brief Declarations, manipulations, and wrappers for basic types.
+ * 
+ * This file is a bunch of utilities for going back and forth between different types.
+ * 
+ * Many of them are for the compiler, so as to clean up the code. It unfortunately
+ * seems necessary when we have types we really care about that are less than word width.
+ */
+
+#pragma once
+
+#ifdef KITTENS_HOPPER
+#include <cuda_fp8.h>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <string>
+#include <bit>
+
+
+namespace kittens {
+
+/**
+ * @brief Bfloat16 floating-point type.
+ */
+using bf16 = __nv_bfloat16;
+/**
+ * @brief Half-precision floating-point type.
+ */
+using half = __half;
+/**
+ * @brief Packed word of two bfloat16 floating-point values.
+ */
+using bf16_2 = __nv_bfloat162;
+/**
+ * @brief Packed word of two half-precision floating-point values.
+ */
+using half_2 = __half2;
+#ifdef KITTENS_HOPPER
+/**
+ * @brief float8 floating-point type.
+ */
+using fp8e4m3 = __nv_fp8_e4m3;
+using fp8e5m2 = __nv_fp8_e5m2;
+#ifdef KITTENS_BLACKWELL
+using fp8e8m0 = __nv_fp8_e8m0;
+#endif
+/**
+ * @brief 2-packed float8 floating-point type.
+ */
+using fp8e4m3_2 = __nv_fp8x2_e4m3;
+using fp8e5m2_2 = __nv_fp8x2_e5m2;
+#ifdef KITTENS_BLACKWELL
+using fp8e8m0_2 = __nv_fp8x2_e8m0;
+#endif
+/**
+ * @brief 4-packed float8 floating-point type.
+ */
+using fp8e4m3_4 = __nv_fp8x4_e4m3;
+using fp8e5m2_4 = __nv_fp8x4_e5m2;
+#ifdef KITTENS_BLACKWELL
+using fp8e8m0_4 = __nv_fp8x4_e8m0;
+#endif
+#endif
+
+namespace ducks {
+/**
+ * @namespace base_types
+ *
+ * @brief A namespace for concepts for basic data types.
+ */
+namespace base_types {
+
+#ifdef KITTENS_HOPPER
+#ifdef KITTENS_BLACKWELL
+template<typename T>
+concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4> || std::is_same_v<T, fp8e8m0_4>; // could add half_2 later if implemented.
+template<typename T>
+concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2> || std::is_same_v<T, fp8e8m0>; // could add half_2 later if implemented.
+#else
+template<typename T>
+concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2> || std::is_same_v<T, fp8e4m3_4> || std::is_same_v<T, fp8e5m2_4>;
+template<typename T>
+concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half> || std::is_same_v<T, fp8e4m3> || std::is_same_v<T, fp8e5m2>;
+#endif
+#else
+template<typename T>
+concept T2 = std::is_same_v<T, float2> || std::is_same_v<T, bf16_2> || std::is_same_v<T, half_2>;
+template<typename T>
+concept T1 = std::is_same_v<T, float>  || std::is_same_v<T, bf16  > || std::is_same_v<T, half>;
+#endif
+
+} // namespace base_types
+} // namespace ducks
+
+/**
+ * @namespace base_types
+ *
+ * @brief A namespace for ThunderKittens basic data types.
+ */
+namespace base_types {
+
+/**
+ * @brief Provides compile-time constants for different types.
+ *
+ * @tparam T The type for which to provide constants.
+ */
+template<typename T> struct constants {
+    /**
+     * @brief Zero
+     * @return Constexpr zero with type T
+     */
+    static __device__ inline constexpr T zero()      { return T{0}; }
+    /**
+     * @brief One
+     * @return Constexpr one with type T
+     */
+    static __device__ inline constexpr T one()       { return T{1}; }
+    /**
+     * @brief Positive infinity. Particularly useful for initializing before a min op.
+     * @return Constexpr positive infinity with type T
+     */
+    static __device__ inline constexpr T pos_infty() { return T{INFINITY}; } // I'll find a better way at some point but this appears to work.
+    /**
+     * @brief Negative infinity. Particularly useful for initializing before a max op.
+     * @return Constexpr negative infinity with type T
+     */
+    static __device__ inline constexpr T neg_infty() { return T{-INFINITY}; }
+};
+template<> struct constants<float2> {
+    static __device__ inline constexpr float2 zero()      { return float2{0.f, 0.f}; }
+    static __device__ inline constexpr float2 one()       { return float2{1.f, 1.f}; }
+    static __device__ inline constexpr float2 pos_infty() { return float2{constants<float>::pos_infty(), constants<float>::pos_infty()}; }
+    static __device__ inline constexpr float2 neg_infty() { return float2{constants<float>::neg_infty(), constants<float>::neg_infty()}; }
+};
+template<> struct constants<bf16> {
+    static __device__ inline constexpr bf16 zero()      { return std::bit_cast<__nv_bfloat16>(uint16_t(0x0000)); } // unfortunately __float2bf16_rn is not constexpr
+    static __device__ inline constexpr bf16 one()       { return std::bit_cast<__nv_bfloat16>(uint16_t(0x3F80)); }
+    static __device__ inline constexpr bf16 pos_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0x7F80)); }
+    static __device__ inline constexpr bf16 neg_infty() { return std::bit_cast<__nv_bfloat16>(uint16_t(0xFF80)); }
+};
+template<> struct constants<bf16_2> {
+    static __device__ inline constexpr bf16_2 zero()      { return bf16_2{constants<bf16>::zero(),      constants<bf16>::zero()};      }
+    static __device__ inline constexpr bf16_2 one()       { return bf16_2{constants<bf16>::one(),       constants<bf16>::one()};       }
+    static __device__ inline constexpr bf16_2 pos_infty() { return bf16_2{constants<bf16>::pos_infty(), constants<bf16>::pos_infty()}; }
+    static __device__ inline constexpr bf16_2 neg_infty() { return bf16_2{constants<bf16>::neg_infty(), constants<bf16>::neg_infty()}; }
+};
+template<> struct constants<half> {
+    static __device__ inline constexpr half zero()      { return std::bit_cast<__half>(uint16_t(0x0000)); }
+    static __device__ inline constexpr half one()       { return std::bit_cast<__half>(uint16_t(0x3C00)); }
+    static __device__ inline constexpr half pos_infty() { return std::bit_cast<__half>(uint16_t(0x7C00)); }
+    static __device__ inline constexpr half neg_infty() { return std::bit_cast<__half>(uint16_t(0xFC00)); }
+};
+template<> struct constants<half_2> {
+    static __device__ inline constexpr half_2 zero()      { return half_2{constants<half>::zero(),      constants<half>::zero()};      }
+    static __device__ inline constexpr half_2 one()       { return half_2{constants<half>::one(),       constants<half>::one()};       }
+    static __device__ inline constexpr half_2 pos_infty() { return half_2{constants<half>::pos_infty(), constants<half>::pos_infty()}; }
+    static __device__ inline constexpr half_2 neg_infty() { return half_2{constants<half>::neg_infty(), constants<half>::neg_infty()}; }
+};
+#ifdef KITTENS_HOPPER
+template<> struct constants<fp8e4m3> {
+    static __device__ inline constexpr fp8e4m3 zero() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x00)); }
+    static __device__ inline constexpr fp8e4m3 one() { return std::bit_cast<__nv_fp8_e4m3>(uint8_t(0x38)); }
+};
+template<> struct constants<fp8e4m3_2> {
+    static __device__ inline constexpr fp8e4m3_2 zero() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x0000)); }
+    static __device__ inline constexpr fp8e4m3_2 one() { return std::bit_cast<fp8e4m3_2>(uint16_t(0x3838)); }
+};
+template<> struct constants<fp8e4m3_4> {
+    static __device__ inline constexpr fp8e4m3_4 zero() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x00000000)); }
+    static __device__ inline constexpr fp8e4m3_4 one() { return std::bit_cast<fp8e4m3_4>(uint32_t(0x38383838)); }
+};
+template<> struct constants<fp8e5m2> {
+    static __device__ inline constexpr fp8e5m2 zero() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x00)); }
+    static __device__ inline constexpr fp8e5m2 one() { return std::bit_cast<__nv_fp8_e5m2>(uint8_t(0x3C)); }
+};
+template<> struct constants<fp8e5m2_2> {
+    static __device__ inline constexpr fp8e5m2_2 zero() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x0000)); }
+    static __device__ inline constexpr fp8e5m2_2 one() { return std::bit_cast<fp8e5m2_2>(uint16_t(0x3C3C)); }
+};
+template<> struct constants<fp8e5m2_4> {
+    static __device__ inline constexpr fp8e5m2_4 zero() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x00000000)); }
+    static __device__ inline constexpr fp8e5m2_4 one() { return std::bit_cast<fp8e5m2_4>(uint32_t(0x3C3C3C3C)); }
+};
+#endif
+
+template<> struct constants<int> {
+    static __device__ inline constexpr int zero()      { return 0; }
+    static __device__ inline constexpr int one()       { return 1; }
+};
+template<> struct constants<int2> {
+    static __device__ inline constexpr int2 zero()      { return int2{0, 0}; }
+    static __device__ inline constexpr int2 one()       { return int2{1, 1}; }
+};
+
+/**
+ * @brief Provides information about packing of elements for a given type.
+ *
+ * @tparam T The type for which to provide packing information.
+ */
+template<typename T> struct packing {
+    /**
+     * @brief The number of elements packed together.
+     *
+     * @return constexpr int representing number of elements within the type.
+     */
+    static __device__ inline constexpr int num() { return 1; }
+    /**
+     * @brief Packs a single T element twice (replicated) into its packed type.
+     *
+     * @param i[in] The element to pack.
+     * @return The packed type.
+     */
+    static __device__ inline constexpr T pack(const bf16 &i);
+};
+template<> struct packing<bf16> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = bf16;
+    using packed_type = bf16_2;
+    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; }
+};
+template<> struct packing<bf16_2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = bf16;
+    using packed_type = bf16_2;
+    static __device__ inline constexpr bf16_2 pack(const bf16 &i) { return bf16_2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<half> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = half;
+    using packed_type = half_2;
+    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; }
+};
+template<> struct packing<half_2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = half;
+    using packed_type = half_2;
+    static __device__ inline constexpr half_2 pack(const half &i) { return half_2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<float> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = float;
+    using packed_type = float2;
+    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; }
+};
+template<> struct packing<float2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = float;
+    using packed_type = float2;
+    static __device__ inline constexpr float2 pack(const float &i) { return float2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<char> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = char;
+    using packed_type = char2;
+    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<char2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = char;
+    using packed_type = char2;
+    static __device__ inline constexpr char2 pack(const char &i) { return char2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<int> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = int;
+    using packed_type = int2;
+    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<int2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = int;
+    using packed_type = int2;
+    static __device__ inline constexpr int2 pack(const int &i) { return int2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<uint> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = uint;
+    using packed_type = uint2;
+    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<uint2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = uint;
+    using packed_type = uint2;
+    static __device__ inline constexpr uint2 pack(const uint &i) { return uint2{i, i}; } // this replication makes code cleaner later.
+};
+struct uint64_2 { uint64_t x, y; };
+template<> struct packing<uint64_t> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = uint64_t;
+    using packed_type = uint64_2;
+    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<uint64_2> {
+    static __device__ inline constexpr int num() { return 2; }
+    using unpacked_type = uint64_t;
+    using packed_type = uint64_2;
+    static __device__ inline constexpr uint64_2 pack(const uint64_t &i) { return uint64_2{i, i}; } // this replication makes code cleaner later.
+};
+template<> struct packing<float4> {
+    static __device__ inline constexpr int num() { return 4; }
+};
+template<> struct packing<int4> {
+    static __device__ inline constexpr int num() { return 4; }
+};
+#ifdef KITTENS_HOPPER
+template<> struct packing<fp8e4m3> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = fp8e4m3;
+    using packed_type = fp8e4m3_4;
+};
+template<> struct packing<fp8e4m3_4> {
+    static __device__ inline constexpr int num() { return 4; }
+    using unpacked_type = fp8e4m3;
+    using packed_type = fp8e4m3_4;
+};
+template<> struct packing<fp8e5m2> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = fp8e5m2;
+    using packed_type = fp8e5m2_4;
+};
+template<> struct packing<fp8e5m2_4> {
+    static __device__ inline constexpr int num() { return 4; }
+    using unpacked_type = fp8e5m2;
+    using packed_type = fp8e5m2_4;
+};
+#ifdef KITTENS_BLACKWELL
+template<> struct packing<fp8e8m0> {
+    static __device__ inline constexpr int num() { return 1; }
+    using unpacked_type = fp8e8m0;
+    using packed_type = fp8e8m0_4;
+};
+template<> struct packing<fp8e8m0_4> {
+    static __device__ inline constexpr int num() { return 4; }
+    using unpacked_type = fp8e8m0;
+    using packed_type = fp8e8m0_4;
+};
+#endif
+#endif
+
+
+/**
+ * @brief Provides templated functionality to convert between different types.
+ *
+ * @tparam T The target type for conversion.
+ * @tparam U The source type for conversion.
+ */
+template<typename T, typename U> struct convertor {
+    /**
+     * @brief Converts a value of type U to type T.
+     *
+     * @param u[in] The value of type U to convert.
+     * @return T The converted value of type T.
+     */
+    static __host__ __device__ inline T convert(const U & u) {
+        return (T)u;
+    }
+};
+template<> struct convertor<float, bf16> {
+    static __host__ __device__ inline float convert(const bf16 & u) {
+        return 	__bfloat162float(u);
+    }
+};
+template<> struct convertor<bf16, float> {
+    static __host__ __device__ inline bf16 convert(const float & u) {
+        return 	__float2bfloat16_rn(u);
+    }
+};
+template<> struct convertor<float2, bf16_2> {
+    static __host__ __device__ inline float2 convert(const bf16_2 & u) {
+        return 	__bfloat1622float2(u);
+    }
+};
+template<> struct convertor<bf16_2, float2> {
+    static __host__ __device__ inline bf16_2 convert(const float2 & u) {
+        return 	__float22bfloat162_rn(u);
+    }
+};
+template<> struct convertor<float, half> {
+    static __host__ __device__ inline float convert(const half & u) {
+        return __half2float(u);
+    }
+};
+template<> struct convertor<half, float> {
+    static __host__ __device__ inline half convert(const float & u) {
+        return __float2half(u);
+    }
+};
+template<> struct convertor<float2, half_2> {
+    static __host__ __device__ inline float2 convert(const half_2 & u) {
+        return __half22float2(u);
+    }
+};
+template<> struct convertor<half_2, float2> {
+    static __host__ __device__ inline half_2 convert(const float2 & u) {
+        return __float22half2_rn(u);
+    }
+};
+template<> struct convertor<bf16, half> {
+    static __host__ __device__ inline bf16 convert(const half & u) {
+        return __float2bfloat16_rn(__half2float(u));
+    }
+};
+template<> struct convertor<half, bf16> {
+    static __host__ __device__ inline half convert(const bf16 & u) {
+        return __float2half(__bfloat162float(u));
+    }
+};
+template<> struct convertor<bf16_2, half_2> {
+    static __host__ __device__ inline bf16_2 convert(const half_2 & u) {
+        return __float22bfloat162_rn(__half22float2(u));
+    }
+};
+template<> struct convertor<half_2, bf16_2> {
+    static __host__ __device__ inline half_2 convert(const bf16_2 & u) {
+        return __float22half2_rn(__bfloat1622float2(u));
+    }
+};
+#ifdef KITTENS_HOPPER
+// fp8e4m3
+template<> struct convertor<fp8e4m3_4, float4> {
+    static __host__ __device__ inline fp8e4m3_4 convert(const float4& u) {
+        return __nv_fp8x4_e4m3(u); 
+    }
+};
+template<> struct convertor<float4, fp8e4m3_4> {
+    static __host__ __device__ inline float4 convert(const fp8e4m3_4& u) {
+        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x4_e4m3*>(&u));
+        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
+    }
+};
+template<> struct convertor<fp8e4m3_2, float2> {
+    static __host__ __device__ inline fp8e4m3_2 convert(const float2& u) {
+        return __nv_fp8x2_e4m3(u); 
+    }
+};
+template<> struct convertor<float2, fp8e4m3_2> {
+    static __host__ __device__ inline float2 convert(const fp8e4m3_2& u) {
+        __nv_fp8_e4m3 *vals = reinterpret_cast<__nv_fp8_e4m3*>(const_cast<__nv_fp8x2_e4m3*>(&u));
+        return make_float2(float(vals[0]), float(vals[1]));
+    }
+};
+template<> struct convertor<fp8e4m3, float> {
+    static __host__ __device__ inline fp8e4m3 convert(const float & u) {
+        return __nv_fp8_e4m3(u);
+    }
+};
+template<> struct convertor<float, fp8e4m3> {
+    static __host__ __device__ inline float convert(const fp8e4m3 & u) {
+        return float(u);
+    }
+};
+template<> struct convertor<bf16_2, fp8e4m3_4> {
+    static __host__ __device__ inline bf16_2 convert(const fp8e4m3_4 & u) {
+        float4 f4 = convertor<float4, fp8e4m3_4>::convert(u);
+        float2 f2 = make_float2(f4.x, f4.y);
+        return __float22bfloat162_rn(f2);
+    }
+};
+template<> struct convertor<fp8e4m3_4, bf16_2> {
+    static __host__ __device__ inline fp8e4m3_4 convert(const bf16_2 & u) {
+        float2 f2 = __bfloat1622float2(u);
+        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
+        return __nv_fp8x4_e4m3(f4);
+    }
+};
+// fp8e5m2
+template<> struct convertor<fp8e5m2_4, float4> {
+    static __host__ __device__ inline fp8e5m2_4 convert(const float4& u) {
+        return __nv_fp8x4_e5m2(u); 
+    }
+};
+template<> struct convertor<float4, fp8e5m2_4> {
+    static __host__ __device__ inline float4 convert(const fp8e5m2_4& u) {
+        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x4_e5m2*>(&u));
+        return make_float4(float(vals[0]), float(vals[1]), float(vals[2]), float(vals[3]));
+    }
+};
+template<> struct convertor<fp8e5m2_2, float2> {
+    static __host__ __device__ inline fp8e5m2_2 convert(const float2& u) {
+        return __nv_fp8x2_e5m2(u); 
+    }
+};
+template<> struct convertor<float2, fp8e5m2_2> {
+    static __host__ __device__ inline float2 convert(const fp8e5m2_2& u) {
+        __nv_fp8_e5m2 *vals = reinterpret_cast<__nv_fp8_e5m2*>(const_cast<__nv_fp8x2_e5m2*>(&u));
+        return make_float2(float(vals[0]), float(vals[1]));
+    }
+};
+template<> struct convertor<fp8e5m2, float> {
+    static __host__ __device__ inline fp8e5m2 convert(const float & u) {
+        return __nv_fp8_e5m2(u);
+    }
+};
+template<> struct convertor<float, fp8e5m2> {
+    static __host__ __device__ inline float convert(const fp8e5m2 & u) {
+        return float(u);
+    }
+};
+template<> struct convertor<bf16_2, fp8e5m2_4> {
+    static __host__ __device__ inline bf16_2 convert(const fp8e5m2_4 & u) {
+        float4 f4 = convertor<float4, fp8e5m2_4>::convert(u);
+        float2 f2 = make_float2(f4.x, f4.y);
+        return __float22bfloat162_rn(f2);
+    }
+};
+template<> struct convertor<fp8e5m2_4, bf16_2> {
+    static __host__ __device__ inline fp8e5m2_4 convert(const bf16_2 & u) {
+        float2 f2 = __bfloat1622float2(u);
+        float4 f4 = make_float4(f2.x, f2.y, 0.0f, 0.0f);
+        return __nv_fp8x4_e5m2(f4);
+    }
+};
+#endif
+}
+}

extra/thunder/cuda/include/common/common.cuh

@@ -0,0 +1,11 @@
+/**
+ * @file
+ * @brief A collection of common resources on which ThunderKittens depends.
+ */
+ 
+
+#pragma once
+
+#include "util.cuh"
+#include "base_types.cuh"
+#include "base_ops.cuh"

extra/thunder/cuda/include/common/debug.cuh

@@ -0,0 +1,56 @@
+#pragma once
+
+// Reset
+#define TK_RESET "\033[0m"
+
+// Foreground colors
+#define TK_FG_BLACK "\033[30m"
+#define TK_FG_RED "\033[31m"
+#define TK_FG_GREEN "\033[32m"
+#define TK_FG_YELLOW "\033[33m"
+#define TK_FG_BLUE "\033[34m"
+#define TK_FG_MAGENTA "\033[35m"
+#define TK_FG_CYAN "\033[36m"
+#define TK_FG_WHITE "\033[37m"
+
+// Background colors
+#define TK_BG_BLACK "\033[40m"
+#define TK_BG_RED "\033[41m"
+#define TK_BG_GREEN "\033[42m"
+#define TK_BG_YELLOW "\033[43m"
+#define TK_BG_BLUE "\033[44m"
+#define TK_BG_MAGENTA "\033[45m"
+#define TK_BG_CYAN "\033[46m"
+#define TK_BG_WHITE "\033[47m"
+
+// Bright foreground colors
+#define TK_FG_BRIGHT_BLACK "\033[90m"
+#define TK_FG_BRIGHT_RED "\033[91m"
+#define TK_FG_BRIGHT_GREEN "\033[92m"
+#define TK_FG_BRIGHT_YELLOW "\033[93m"
+#define TK_FG_BRIGHT_BLUE "\033[94m"
+#define TK_FG_BRIGHT_MAGENTA "\033[95m"
+#define TK_FG_BRIGHT_CYAN "\033[96m"
+#define TK_FG_BRIGHT_WHITE "\033[97m"
+
+// Bright background colors
+#define TK_BG_BRIGHT_BLACK "\033[100m"
+#define TK_BG_BRIGHT_RED "\033[101m"
+#define TK_BG_BRIGHT_GREEN "\033[102m"
+#define TK_BG_BRIGHT_YELLOW "\033[103m"
+#define TK_BG_BRIGHT_BLUE "\033[104m"
+#define TK_BG_BRIGHT_MAGENTA "\033[105m"
+#define TK_BG_BRIGHT_CYAN "\033[106m"
+#define TK_BG_BRIGHT_WHITE "\033[107m"
+
+// Text styles
+#define TK_BOLD "\033[1m"
+#define TK_DIM "\033[2m"
+#define TK_ITALIC "\033[3m"
+#define TK_UNDERLINE "\033[4m"
+#define TK_BLINK "\033[5m"
+#define TK_REVERSE "\033[7m"
+#define TK_HIDDEN "\033[8m"
+
+// Macro to combine styles
+#define TK_STYLE(...)  "\033[" #__VA_ARGS__ "m"

extra/thunder/cuda/include/common/util.cuh

@@ -0,0 +1,314 @@
+/**
+ * @file
+ * @brief General utilities for ThunderKittens.
+ */
+
+#pragma once
+
+#include <stdint.h>
+#include <type_traits>
+#include <concepts>
+#include <memory>
+
+// CUDA driver API
+#define CUCHECK(cmd) do {                                     \
+    CUresult err = cmd;                                       \
+    if (err != CUDA_SUCCESS) {                                \
+        const char *errStr;                                   \
+        cuGetErrorString(err, &errStr);                       \
+        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
+            __FILE__, __LINE__, errStr);                      \
+        exit(EXIT_FAILURE);                                   \
+    }                                                         \
+} while(0)
+
+// CUDA runtime API
+#define CUDACHECK(cmd) do {                                   \
+    cudaError_t err = cmd;                                    \
+    if (err != cudaSuccess) {                                 \
+        fprintf(stderr, "Failed: CUDA error %s:%d '%s'\n",    \
+            __FILE__, __LINE__, cudaGetErrorString(err));     \
+        exit(EXIT_FAILURE);                                   \
+    }                                                         \
+} while(0)
+
+/**
+ * @namespace kittens
+ *
+ * @brief The main namespace of ThunderKittens.
+ */
+namespace kittens {
+
+/* ----------  GENERAL CONSTANTS FOR KITTENS  ---------- */
+
+/**
+ * @brief Tile dimension constant.
+ */
+template<typename T> constexpr int TILE_COL_DIM = sizeof(T) == 1 ? 32 : 16;
+template<typename T> constexpr int TILE_ROW_DIM = 16;
+/**
+ * @brief Tile num elements constant calculated as TILE_DIM squared.
+ */
+template<typename T> constexpr int TILE_ELEMENTS{TILE_COL_DIM<T>*TILE_ROW_DIM<T>};
+/**
+ * @brief Constant representing number of threads in a warp.
+ */
+constexpr int WARP_THREADS{32};
+/**
+ * @brief Constant representing number of threads in a warpgroup of four warps.
+ */
+constexpr int WARPGROUP_THREADS{128};
+/**
+
+ * @brief Constant representing number of warps in a warpgroup of four warps.
+ */
+constexpr int WARPGROUP_WARPS{4};
+/**
+
+ * @brief Get the warp ID of the current thread.
+ * @return The warp ID.
+ */
+__device__ static __forceinline__ int warpid() {
+    // uint32_t wid;
+    // asm volatile("mov.u32 %0, %warpid;" : "=r"(wid));
+    // return wid;
+    return threadIdx.x >> 5;
+}
+/**
+ * @brief Get the warpgroup ID of the current thread.
+ * @return The warpgroup ID.
+ */
+__device__ static __forceinline__ int warpgroupid() { return warpid() >> 2; }
+/**
+ * @brief Get the lane ID of the current thread within its warp.
+ * @return The lane ID.
+ */
+__device__ static __forceinline__ int laneid() {
+    // uint32_t lid;
+    // asm volatile("mov.u32 %0, %laneid;" : "=r"(lid));
+    // return lid;
+    return threadIdx.x & 31;
+}
+
+#if defined(KITTENS_HOPPER)
+constexpr int MAX_SHARED_MEMORY = 227000;
+#elif defined(KITTENS_A100)
+constexpr int MAX_SHARED_MEMORY = 164000;
+#elif defined(KITTENS_4090)
+constexpr int MAX_SHARED_MEMORY = 100000;
+#endif
+
+struct transpose {
+    static constexpr int N = 0; // not transposed
+    static constexpr int T = 1; // transposed
+};
+struct axis {
+    static constexpr int ROW = 0; // row axis of a tile
+    static constexpr int COL = 1; // column axis of a tile
+};
+
+/* ----------  TYPE HELPERS  ---------- */
+
+/**
+ * @namespace ducks
+ *
+ * @brief ThunderKittens' namespace for template metaprogramming..
+ * 
+ * This includes primarily dummy types and concept wrappers, along
+ * with a few additional utilities.
+ */
+namespace ducks {
+
+/**
+ * @brief A type representing an empty default for a template.
+ */
+struct default_type {};
+
+// This macro can't be done as a template, so it doesn't really have a location in kittens.
+#define typeof(A) typename std::remove_const<typename std::remove_reference<decltype(A)>::type>::type
+
+}
+
+/* ----------  SHUFFLE UTILS  ---------- */
+
+/**
+ * @brief Mask constant for all active threads in a warp.
+ */
+static constexpr uint32_t MASK_ALL = 0xFFFFFFFF;
+
+/**
+ * @brief Perform a shuffle down operation on a packed type synchronously across a warp.
+ * @tparam T The type of the value to be shuffled.
+ * @param mask[in] The mask of active threads.
+ * @param f[in] The value to be shuffled.
+ * @param delta[in] The number of positions to shuffle down.
+ * @return The result of the shuffle operation.
+ */
+template<typename T>
+__device__ static inline T packed_shfl_down_sync(uint32_t mask, const T &f, int delta) {
+    return __shfl_down_sync(mask, f, delta);
+}
+template<>
+__device__ inline float2 packed_shfl_down_sync<float2>(uint32_t mask, const float2 &f, int delta) {
+    float2 r;
+    r.x = __shfl_down_sync(mask, f.x, delta);
+    r.y = __shfl_down_sync(mask, f.y, delta);
+    return r;
+}
+/**
+ * @brief Perform a packed shuffle operation synchronously across a warp.
+ * @tparam T The type of the value to be shuffled.
+ * @param mask[in] The mask of active threads.
+ * @param f[in] The value to be shuffled.
+ * @param src[in] The source lane from which to shuffle.
+ * @return The result of the shuffle operation.
+ */
+template<typename T>
+__device__ static inline T packed_shfl_sync(uint32_t mask, const T &f, int src) {
+    return __shfl_sync(mask, f, src);
+}
+template<>
+__device__ inline float2 packed_shfl_sync<float2>(uint32_t mask, const float2 &f, int src) {
+    float2 r;
+    r.x = __shfl_sync(mask, f.x, src);
+    r.y = __shfl_sync(mask, f.y, src);
+    return r;
+}
+
+/* ----------  SHARED MEMORY UTILS  ---------- */
+
+// namespace ducks {
+// namespace sb {
+// struct identifier {};
+// }
+// }
+
+// template<typename Args...>
+// struct sb {
+//     using identifier = ducks::sb::identifier;
+//     Args... args;
+// };
+
+// namespace ducks {
+// namespace sb {
+// template<typename T> concept all = requires {
+//     typename T::identifier;
+// } && std::is_same_v<T::identifier, identifier>;
+// }
+// }
+
+// Joyously stolen from https://github.com/NVIDIA/cutlass/blob/5c447dd84f8ae0e1d48ff9a2eae26ce8c4958101/include/cute/container/alignment.hpp#L51
+#if defined(__CUDACC__)
+#define KITTENS_ALIGN_AS(n) __align__(n)
+#else
+#define KITTENS_ALIGN_AS(n) alignas(n)
+#endif
+
+#ifdef KITTENS_HOPPER
+#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(128)
+#else
+#define KITTENS_DEFAULT_ALIGN KITTENS_ALIGN_AS(16)
+#endif
+
+/**
+ * @brief Dummy structure for alignment purposes. Needed for WGMMA and TMA calls.
+ */
+struct KITTENS_DEFAULT_ALIGN alignment_dummy { int dummy; };
+/**
+ * @brief Very simple allocator for dynamic shared memory. Advances pointer and tracks alignments.
+ * @tparam default_alignment The default alignment this allocator will enforce. If <=0 (default -1) it will not align.
+ */
+#ifdef KITTENS_HOPPER
+template<int default_alignment=1024> 
+#else
+template<int default_alignment=16> 
+#endif
+struct shared_allocator {
+    int *ptr;
+
+    private:
+        // Recursive template to generate N-dimensional array type
+        template<typename A, size_t... dims>
+        struct variadic_array;
+        template<typename A, size_t first_dim, size_t... rest_dims>
+        struct variadic_array<A, first_dim, rest_dims...> {
+            using type = typename variadic_array<A, rest_dims...>::type[first_dim];
+        };
+        template<typename A>
+        struct variadic_array<A> {
+            using type = A;
+        };
+        template<typename A, size_t... dims> 
+        using variadic_array_t = typename variadic_array<A, dims...>::type;
+
+        template<int alignment>
+        __device__ inline void align_ptr() {
+            if constexpr (alignment > 0) {
+                uint64_t p = reinterpret_cast<uint64_t>(ptr);
+                if(p % alignment != 0) {
+                    ptr = (int*)(p + (alignment-(p%alignment)));
+                }
+            }
+        }
+
+    public:
+        /**
+        * @brief Construct a new shared allocator using a pointer to extern shared memory.
+        * @param[in] _ptr Pointer to the start of the extern shared memory.
+        */
+        __device__ shared_allocator(int *_ptr): ptr(_ptr) {}
+        /**
+        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
+        * @tparam A The type of the object to allocate.
+        * @tparam dims... A list of dimensions for the N-dimensional array.
+        * @return Reference to the allocated object.
+        */
+        template<typename A, size_t... dims> 
+        __device__ inline variadic_array_t<A, dims...>& allocate() {
+            // static_assert(sizeof(A) % default_alignment == 0, "Type is not aligned properly for array allocation");
+            align_ptr<default_alignment>();
+            using at = variadic_array_t<A, dims...>;
+            at*p = reinterpret_cast<at*>(ptr);
+            ptr += sizeof(at)/sizeof(int);
+            return *p;
+        }
+        /**
+        * @brief Allocate shared memory for a single instance or N-dimensional array of type A.
+        * @tparam alignment An alignment to enforce for this particular object.
+        * @tparam A The type of the object to allocate.
+        * @tparam dims... A list of dimensions for the N-dimensional array.
+        * @return Reference to the allocated object.
+        */
+        template<int alignment, typename A, size_t... dims> 
+        __device__ inline variadic_array_t<A, dims...>& allocate() {
+            // static_assert(sizeof(A) % alignment == 0, "Type is not aligned properly for array allocation");
+            align_ptr<alignment>();
+            using at = variadic_array_t<A, dims...>;
+            at*p = reinterpret_cast<at*>(ptr);
+            ptr += sizeof(at)/sizeof(int);
+            return *p;
+        }
+};
+#if (defined(KITTENS_HOPPER) || defined(KITTENS_BLACKWELL))
+/**
+ * @brief A wrapper for an allocator that enforces sufficient alignment to be used for TMA loads and stores.
+ */
+using tma_allocator = shared_allocator<1024>;
+using tma_swizzle_allocator = tma_allocator; // swizzled TMA modes require up to 1024 byte alignments :/
+
+/* Get CTA ID within a cluster */
+__device__ static inline int3 clusterIdx() {
+    int3 cluster_idx;
+    asm volatile("mov.u32 %0, %clusterid.x;\n" : "=r"(cluster_idx.x));
+    asm volatile("mov.u32 %0, %clusterid.y;\n" : "=r"(cluster_idx.y));
+    asm volatile("mov.u32 %0, %clusterid.z;\n" : "=r"(cluster_idx.z));
+    return cluster_idx;
+}
+__device__ static inline int cluster_ctarank() {
+    uint32_t ctarank;
+    asm volatile("mov.u32 %0, %cluster_ctarank;\n" : "=r"(ctarank));
+    return ctarank;
+}
+#endif
+
+} // namespace kittens

extra/thunder/cuda/include/kittens.cuh

@@ -0,0 +1,12 @@
+/**
+ * @file
+ * @brief The master header file of ThunderKittens. This file includes everything you need!
+ */
+
+#pragma once
+
+#include "common/common.cuh"
+#include "types/types.cuh"
+#include "ops/ops.cuh"
+#include "pyutils/util.cuh"
+// #include "pyutils/pyutils.cuh" // for simple binding without including torch

extra/thunder/cuda/include/ops/device/device.cuh

@@ -0,0 +1,51 @@
+/**
+ * @file
+ * @brief An aggregate header of all device (multi-GPU) operations defined by ThunderKittens
+ */
+
+#pragma once
+
+#include "../../types/types.cuh"
+
+namespace kittens {
+
+template<int _NUM_DEVICES>
+struct device {
+
+static_assert(_NUM_DEVICES >= 0 && _NUM_DEVICES <= 72, "Invalid number of devices");
+static constexpr int NUM_DEVICES = _NUM_DEVICES;
+
+#ifdef KITTENS_HOPPER
+
+using barrier_t = pgl<gl<int, 1, 1, 1, -1>, NUM_DEVICES, true>;
+
+/**
+ * @brief Multi-GPU synchronization barrier for coordinated kernel exit
+ * 
+ * Performs a synchronization across all devices to ensure all GPUs complete 
+ * their work before any kernel exits. Does not synchronize intra-node threads
+ * or threadblocks.
+ * 
+ * @param barrier Pre-allocated barrier structure, must be initialized to 0
+ * @param dev_idx Current device index (0 to NUM_DEVICES - 1)
+ * @param id Synchronization point identifier (default: 0). 0 is fine for most cases
+ * 
+ */
+__device__ static inline void sync_on_exit(const barrier_t &barrier, const int dev_idx, const int id = 0) {
+    if (blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0 &&
+        threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
+        cuda::atomic_ref<int, cuda::thread_scope_system> barrier_uc(barrier[dev_idx][{id}]);
+
+        // Inter-note check-in
+        multimem<int>::red<reduce_op::ADD>(barrier.mc_ptr_at({id}), 1);
+        asm volatile ("{fence.proxy.alias;}" ::: "memory");
+        while (barrier_uc.load(cuda::memory_order_acquire) < NUM_DEVICES);
+        barrier_uc.fetch_sub(NUM_DEVICES, cuda::memory_order_release);
+    }
+}
+
+#endif
+
+};
+
+} // namespace kittens

extra/thunder/cuda/include/ops/group/group.cuh

@@ -0,0 +1,96 @@
+/**
+ * @file
+ * @brief An aggregate header of all group (multi-warp) operations defined by ThunderKittens
+ */
+
+#pragma once
+
+#include <cuda/pipeline>
+
+#include "../../common/common.cuh"
+#include "../../types/types.cuh"
+#include "../thread/thread.cuh" // several group memory ops rely on underlying warp-scope ops
+
+#define KITTENS_CHECK_WARP static_assert(GROUP_WARPS==1, "Warp (GROUP_WARPS=1) function called from a non-warp group.");
+// A "warpgroup" is a special group of 4 consecutive warps defined by NVIDIA for certain SM_90+ operations.
+#define KITTENS_CHECK_WARPGROUP static_assert(GROUP_WARPS==4, "Warpgroup (GROUP_WARPS=4) function called from a non-warpgroup group.");
+
+// WGMMA relies on some template structures that cannot be specialized within the group struct, so we declare them in advance.
+#ifdef KITTENS_HOPPER
+#include "mma/warpgroup/base/base.cuh"
+#endif
+
+namespace kittens {
+/*
+This is meant to be used with a `using group_N = kittens::group<NUM_WORKERS>;` at the start of every kernel.
+*/
+template<int _GROUP_WARPS>
+struct group {
+static constexpr int GROUP_WARPS = _GROUP_WARPS; // This alias produces nice parallelism.
+static constexpr int GROUP_THREADS = GROUP_WARPS * kittens::WARP_THREADS; // This alias produces nice parallelism.
+__device__ static inline int laneid() { return threadIdx.x % GROUP_THREADS; }
+__device__ static inline int warpid() { return laneid() / kittens::WARP_THREADS; }
+__device__ static inline int groupid() { return threadIdx.x / GROUP_THREADS; }
+
+__device__ static inline void sync(int id) {
+    asm volatile("bar.sync %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
+}
+template<uint32_t MASK=0xFFFFFFFF> __device__ static inline void sync() {
+    static_assert(GROUP_WARPS==1, "barrier-less sync() can only be called by a single warp!");
+    asm volatile("bar.warp.sync %0;\n" :: "n"(MASK));
+}
+__device__ static inline void arrive(int id) {
+    asm volatile("bar.arrive %0, %1;\n" :: "r"(id), "n"(GROUP_THREADS));
+}
+
+#include "memory/memory.cuh"
+#include "shared/shared.cuh"
+#include "register/register.cuh"
+
+#ifdef KITTENS_HOPPER
+#include "mma/mma.cuh"
+
+template<int n_reg> __device__ static inline void increase_registers() {
+    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
+    asm volatile("setmaxnreg.inc.sync.aligned.u32 %0;\n" :: "n"(n_reg));
+}
+template<int n_reg> __device__ static inline void decrease_registers() {
+    static_assert(n_reg % 8 == 0, "n_reg must be a multiple of 8");
+    asm volatile("setmaxnreg.dec.sync.aligned.u32 %0;\n" :: "n"(n_reg));
+}
+__device__ static inline void producer_registers() { decrease_registers<24>(); }
+template<int NCWG> __device__ static inline void consumer_registers() { increase_registers<480/NCWG - 8*(NCWG>3) - 224*(NCWG==1)>(); }
+
+#endif
+
+};
+
+namespace everyone {
+
+// Block-level synchronization
+__device__ static inline void sync(int id) {
+    asm volatile("bar.sync %0;\n" :: "r"(id));
+}
+
+// Cluster-level synchronization functions
+namespace tma {
+namespace cluster {
+__device__ static inline void arrive_aligned() { // All threads in the cluster must call this
+    asm volatile ("barrier.cluster.arrive.release.aligned;\n");
+}
+__device__ static inline void wait_aligned() {
+    asm volatile ("barrier.cluster.wait.acquire.aligned;\n");
+}
+__device__ static inline void sync() {
+    arrive_aligned();
+    wait_aligned();
+}
+}
+}
+
+};
+
+using warp = group<1>;      // scope used by most pre-Hopper GPUs, and also for most register operations.
+using warpgroup = group<4>; // special scope commonly used by Hopper and later.
+
+}

extra/thunder/cuda/include/ops/group/memory/memory.cuh

@@ -0,0 +1,21 @@
+/**
+ * @file
+ * @brief An aggregate header of colaborative group memory movement operations
+ */
+
+#include "util/util.cuh"
+#include "tile/tile.cuh"
+#include "vec/vec.cuh"
+
+#ifdef KITTENS_HOPPER
+struct tma {
+#include "util/tma.cuh"
+#include "tile/tma.cuh"
+#include "vec/tma.cuh"
+struct cluster {
+#include "util/tma_cluster.cuh"
+#include "tile/tma_cluster.cuh"
+#include "vec/tma_cluster.cuh"
+};
+};
+#endif

extra/thunder/cuda/include/ops/group/memory/tile/complex/complex_global_to_register.cuh

@@ -0,0 +1,42 @@
+/**
+ * @file
+ * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively loads data from a source array into register tiles.
+ *
+ * @tparam RT The register tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from.
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
+__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
+    load<axis, CRT::component, CGL::component, COORD>(dst.real, src.real, idx);
+    load<axis, CRT::component, CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
+__device__ inline static void load(CRT &dst, const CGL &src, const COORD &idx) {
+    load<2, CRT, CGL>(dst, src, idx);
+}
+
+/**
+ * @brief Collaboratively stores data from register tiles to a destination array in global memory.
+ *
+ * @tparam RT The register tile type.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<int axis, ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
+__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
+    store<axis, typename CRT::component, typename CGL::component>(dst.real, src.real, idx);
+    store<axis, typename CRT::component, typename CGL::component>(dst.imag, src.imag, idx);
+}
+template<ducks::crt::all CRT, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<crt<typename CRT::T, GROUP_WARPS*CRT::rows, CRT::cols, typename CRT::layout>>>
+__device__ inline static void store(CGL &dst, const CRT &src, const COORD &idx) {
+    store<2, CRT, CGL>(dst, src, idx);
+}

extra/thunder/cuda/include/ops/group/memory/tile/complex/complex_global_to_shared.cuh

@@ -0,0 +1,37 @@
+/**
+ * @file
+ * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
+ */
+
+template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
+    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    load<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void load(CST &dst, const CGL &src, const COORD &idx) {
+    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    load<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+
+template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
+    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    store<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void store(CGL &dst, const CST &src, const COORD &idx) {
+    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    store<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+
+template<int axis, bool assume_aligned, ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
+    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    load_async<axis, assume_aligned, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}
+template<ducks::cst::all CST, ducks::cgl::all CGL, ducks::coord::tile COORD=coord<CST>>
+__device__ static inline void load_async(CST &dst, const CGL &src, const COORD &idx) {
+    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.real, src.real, idx);
+    load_async<2, false, typename CST::component, typename CGL::component, COORD>(dst.imag, src.imag, idx);
+}

extra/thunder/cuda/include/ops/group/memory/tile/complex/complex_shared_to_register.cuh

@@ -0,0 +1,34 @@
+/**
+ * @file
+ * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ */
+template<ducks::crt::all RT, ducks::cst::all ST>
+__device__ inline static void load(RT &dst, const ST &src) {
+    load(dst.real, src.real);
+    load(dst.imag, src.imag);
+}
+
+
+/**
+ * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ */
+template<ducks::cst::all ST, ducks::crt::all RT>
+__device__ inline static void store(ST &dst, const RT &src) {
+    store(dst.real, src.real);
+    store(dst.imag, src.imag);
+}
+

extra/thunder/cuda/include/ops/group/memory/tile/global_to_register.cuh

@@ -0,0 +1,207 @@
+/**
+ * @file
+ * @brief Functions for a group to collaboratively transfer data directly between global memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively loads data from a source array into row-major layout tiles.
+ *
+ * @tparam RT The row-major layout tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from.
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
+    using T2 = RT::dtype;
+    using U = typename GL::dtype;
+
+    #ifdef KITTENS_HOPPER
+    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
+    #endif
+
+    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
+    const int row_stride = src.template stride<axis>();
+    using U2 = base_types::packing<U>::packed_type;
+    int warp_laneid = threadIdx.x % WARP_THREADS;
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    const int row_offset = dst.rows*local_warpid;
+    #pragma unroll
+    for(int i = 0; i < dst.height; i++) {
+        int row = row_offset + i*dst.tile_size_row + (warp_laneid / 4);
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
+            dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+0)]));
+            dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+0)*row_stride + (col+8)]));
+        }
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
+            dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+0)]));
+            dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(*(U2*)(&src_ptr[(row+8)*row_stride + (col+8)]));
+        }
+    }
+}
+/**
+ * @brief Collaboratively loads data from a source array into column-major layout tiles.
+ *
+ * @tparam RT The column-major layout tile type.
+ * @tparam U The data type of the source array.
+ * @param dst[out] The destination tile to load data into.
+ * @param src[in] The source array to load data from.
+ * @param row_stride[in] The stride in elements between rows in the source array.
+ */
+template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
+    using T = typename RT::T;
+    using U = typename GL::dtype;
+
+    #ifdef KITTENS_HOPPER
+    static_assert(!std::is_same_v<T, fp8e4m3> && !std::is_same_v<T, fp8e5m2>, "Unsupported type for load/store");
+    #endif
+
+    U *src_ptr = (U*)&src[(idx.template unit_coord<axis, 3>())];
+    const int row_stride = src.template stride<axis>();
+    int warp_laneid = threadIdx.x % WARP_THREADS;
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    const int row_offset = dst.rows*local_warpid;
+    #pragma unroll
+    for(int i = 0; i < dst.height; i++) {
+        int row = row_offset + i*dst.tile_size_row + 2*(warp_laneid % 4);
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + (warp_laneid / 4);
+            dst.tiles[i][j].data[0].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+0)]);
+            dst.tiles[i][j].data[1].x = base_types::convertor<T, U>::convert(src_ptr[(row+0)*row_stride + (col+8)]);
+        }
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + (warp_laneid / 4);
+            dst.tiles[i][j].data[0].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+0)]);
+            dst.tiles[i][j].data[1].y = base_types::convertor<T, U>::convert(src_ptr[(row+1)*row_stride + (col+8)]);
+        }
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + (warp_laneid / 4);
+            dst.tiles[i][j].data[2].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+0)]);
+            dst.tiles[i][j].data[3].x = base_types::convertor<T, U>::convert(src_ptr[(row+8)*row_stride + (col+8)]);
+        }
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            int col = j*dst.tile_size_col + (warp_laneid / 4);
+            dst.tiles[i][j].data[2].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+0)]);
+            dst.tiles[i][j].data[3].y = base_types::convertor<T, U>::convert(src_ptr[(row+9)*row_stride + (col+8)]);
+        }
+    }
+}
+template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void load(RT &dst, const GL &src, const COORD &idx) {
+    load<2>(dst, src, idx);
+}
+/**
+ * @brief Collaboratively stores data from register tiles to a destination array in global memory with a row-major layout.
+ *
+ * @tparam RT The register tile type with a row-major layout.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<int axis, ducks::rt::row_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
+    using T2 = RT::dtype;
+    using U = typename GL::dtype;
+
+    #ifdef KITTENS_HOPPER
+    static_assert(!std::is_same_v<T2, fp8e4m3_4> && !std::is_same_v<T2, fp8e5m2_4>, "Unsupported type for load/store");
+    #endif
+
+    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
+    const int row_stride = dst.template stride<axis>();
+    using U2 = base_types::packing<U>::packed_type;
+    int warp_laneid = threadIdx.x % WARP_THREADS;
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    const int row_offset = src.rows*local_warpid;
+    #pragma unroll
+    for(int i = 0; i < src.height; i++) {
+        int row = row_offset + i*src.tile_size_row + (warp_laneid / 4);
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
+            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
+            *(U2*)(&dst_ptr[(row+0)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
+        }
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + 2*(warp_laneid % 4);
+            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+0)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
+            *(U2*)(&dst_ptr[(row+8)*row_stride + (col+8)]) = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
+        }
+    }
+}
+/**
+ * @brief Collaboratively stores data from register tiles to a destination array in global memory with a column-major layout.
+ *
+ * @tparam RT The register tile type with a column-major layout.
+ * @tparam U The data type of the destination array.
+ * @param[out] dst The destination array in global memory to store data into.
+ * @param[in] src The source register tile to store data from.
+ * @param row_stride[in] The stride in elements between rows in the destination array.
+ */
+template<int axis, ducks::rt::col_layout RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
+    using T = base_types::packing<typename RT::dtype>::unpacked_type;
+    using U = typename GL::dtype;
+
+    #ifdef KITTENS_HOPPER
+    static_assert(!std::is_same_v<T, fp8e4m3_4> && !std::is_same_v<T, fp8e5m2_4>, "Unsupported type for load/store");
+    #endif
+    
+    U *dst_ptr = (U*)&dst[(idx.template unit_coord<axis, 3>())];
+    const int row_stride = dst.template stride<axis>();
+    int warp_laneid = threadIdx.x % WARP_THREADS;
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    const int row_offset = src.rows*local_warpid;
+    #pragma unroll
+    for(int i = 0; i < src.height; i++) {
+        int row = row_offset + i*src.tile_size_row + 2*(warp_laneid % 4);
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + (warp_laneid / 4);
+            dst_ptr[(row+0)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].x);
+            dst_ptr[(row+0)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].x);
+        }
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + (warp_laneid / 4);
+            dst_ptr[(row+1)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[0].y);
+            dst_ptr[(row+1)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[1].y);
+        }
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + (warp_laneid / 4);
+            dst_ptr[(row+8)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].x);
+            dst_ptr[(row+8)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].x);
+        }
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            int col = j*src.tile_size_col + (warp_laneid / 4);
+            dst_ptr[(row+9)*row_stride + (col+0)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[2].y);
+            dst_ptr[(row+9)*row_stride + (col+8)] = base_types::convertor<U, T>::convert(src.tiles[i][j].data[3].y);
+        }
+    }
+}
+template<ducks::rt::all RT, ducks::gl::all GL, ducks::coord::tile COORD=coord<rt<typename RT::T, GROUP_WARPS*RT::rows, RT::cols, typename RT::layout>>>
+__device__ inline static void store(const GL &dst, const RT &src, const COORD &idx) {
+    store<2>(dst, src, idx);
+}

extra/thunder/cuda/include/ops/group/memory/tile/global_to_shared.cuh

@@ -0,0 +1,168 @@
+/**
+ * @file
+ * @brief Group (collaborative warp) ops for loading shared tiles from and storing to global memory. 
+ */
+ 
+
+/**
+ * @brief Loads data from global memory into a shared memory tile.
+ *
+ * @tparam ST The type of the shared tile.
+ * @param[out] dst The destination shared memory tile.
+ * @param[in] src The source global memory array.
+ * @param[in] idx The coordinate of the tile in the global memory array.
+ */
+template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
+    using T = typename ST::dtype;
+    const int row_stride = src.template stride<axis>();
+    // we can handle this many rows each time we run a memcpy_async
+    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
+    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
+    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
+    constexpr int total_rows = dst.height*dst.width;
+
+    coord<> unit_coord = idx.template unit_coord<axis, 3>();
+    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
+    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+    int laneid = threadIdx.x % GROUP_THREADS;
+
+    #pragma unroll
+    for(int i = 0; i < total_calls; i++) {
+
+        int load_idx = i * GROUP_THREADS + laneid;
+        
+        int row = load_idx / memcpy_per_row;
+        int col = (load_idx*elem_per_memcpy) % dst.cols;
+
+        if constexpr (assume_aligned) {
+            float4 tmp;
+            move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
+            move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
+        }
+        else {
+            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
+                float4 tmp;
+                move<float4>::ldg(tmp, (float4*)&src_ptr[row*row_stride + col]);
+                move<float4>::sts(dst.idx(dst_ptr, {row, col}), tmp);
+            }
+            else {
+                float4 zeros = {0.f,0.f,0.f,0.f};
+                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
+            }
+        }
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load(ST &dst, const GL &src, const COORD &idx) {
+    load<2, false, ST, GL, COORD>(dst, src, idx);
+}
+
+/**
+ * @brief Stores data from a shared memory tile into global memory.
+ *
+ * @tparam ST The type of the shared tile.
+ * @param[out] dst The destination global memory array.
+ * @param[in] src The source shared memory tile.
+ * @param row_stride[in] The stride between rows in the destination array.
+ */
+template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
+    using T = typename ST::dtype;
+    const int row_stride = dst.template stride<axis>();
+    // we can handle this many rows each time we run a memcpy_async
+    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
+    constexpr int memcpy_per_row = src.cols / elem_per_memcpy;
+    constexpr int total_calls = (src.height*src.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
+
+    coord<> unit_coord = idx.template unit_coord<axis, 3>();
+    typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[unit_coord];
+    uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
+    int laneid = threadIdx.x % GROUP_THREADS;
+
+    #pragma unroll
+    for(int i = 0; i < total_calls; i++) {
+
+        int load_idx = i * GROUP_THREADS + laneid;
+        
+        int row = load_idx / memcpy_per_row;
+        int col = (load_idx*elem_per_memcpy) % src.cols;
+
+        if constexpr (assume_aligned) {
+            float4 tmp;
+            move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
+            move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
+        }
+        else {
+            if (row + unit_coord.template dim<axis>() < dst.template shape<axis>()) {
+                float4 tmp;
+                move<float4>::lds(tmp, src.idx(src_ptr, {row, col}));
+                move<float4>::stg((float4*)&dst_ptr[row*row_stride + col], tmp);
+            }
+        }
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store(const GL &dst, const ST &src, const COORD &idx) {
+    store<2, false, ST, GL, COORD>(dst, src, idx);
+}
+
+/**
+ * @brief Asynchronously loads data from global memory into a shared memory tile.
+ *
+ * @tparam ST The type of the shared tile.
+ * @param[out] dst The destination shared memory tile.
+ * @param[in] src The source global memory array.
+ *
+ * @note This function expects 16-byte alignments. Otherwise, behavior is undefined.
+ */
+template<int axis, bool assume_aligned, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
+    using T = typename ST::dtype;
+    const int row_stride = src.template stride<axis>();
+    // we can handle this many rows each time we run a memcpy_async
+    constexpr int elem_per_memcpy = sizeof(float4)/sizeof(typename ST::dtype);
+    constexpr int memcpy_per_row = dst.cols / elem_per_memcpy;
+    constexpr int total_calls = (dst.height*dst.width * kittens::TILE_ROW_DIM<T>*kittens::TILE_COL_DIM<T> + GROUP_THREADS*elem_per_memcpy-1) / (GROUP_THREADS*elem_per_memcpy); // round up
+
+    coord<> unit_coord = idx.template unit_coord<axis, 3>();
+    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[unit_coord];
+    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+    int laneid = threadIdx.x % GROUP_THREADS;
+
+    #pragma unroll
+    for(int i = 0; i < total_calls; i++) {
+
+        int load_idx = i * GROUP_THREADS + laneid;
+        
+        int row = load_idx / memcpy_per_row;
+        int col = (load_idx*elem_per_memcpy) % dst.cols;
+
+        if constexpr (assume_aligned) {
+            asm volatile(
+                "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
+                :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
+                : "memory"
+            );
+        }
+        else {
+            if (row + unit_coord.template dim<axis>() < src.template shape<axis>()) {
+                asm volatile(
+                    "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
+                    :: "r"(dst.idx(dst_ptr, {row, col})), "l"(&src_ptr[row*row_stride + col])
+                    : "memory"
+                );
+            }
+            else {
+                // printf("thread %d skipping async load on row %d, col %d\n", threadIdx.x, row + unit_coord.template dim<axis>(), col);
+                float4 zeros = {0.f,0.f,0.f,0.f};
+                move<float4>::sts(dst.idx(dst_ptr, {row, col}), zeros); // use the default value
+            }
+        }
+    }
+    asm volatile("cp.async.commit_group;\n" ::: "memory");
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx) {
+    load_async<2, false, ST, GL, COORD>(dst, src, idx);
+}

extra/thunder/cuda/include/ops/group/memory/tile/shared_to_register.cuh

@@ -0,0 +1,323 @@
+/**
+ * @file
+ * @brief Functions for a warpgroup to collaboratively transfer data directly between shared memory and registers and back.
+ */
+
+/**
+ * @brief Collaboratively load data from a shared tile into register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source shared tile.
+ */
+template<ducks::rt::all RT, ducks::st::all ST>
+__device__ inline static void load(RT &dst, const ST &src) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    using T2 = RT::dtype;
+    using U  = ST::dtype;
+    using T  = base_types::packing<T2>::unpacked_type;
+    using U2 = base_types::packing<U>::packed_type;
+    int warp_laneid = ::kittens::laneid();
+
+    // convert to shared state space
+    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
+
+    #pragma unroll
+    for(int i = 0; i < dst.height; i++) {
+        #pragma unroll
+        for(int j = 0; j < dst.width; j++) {
+            if constexpr (sizeof(typename ST::dtype) == 2) {
+                // handle the row-major layout for 16-bit types
+                U2 tmp[4];
+                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid % 16);
+                int col = j*dst.tile_size_col + (warp_laneid / 16) * 8;
+                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
+                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
+                }
+                else {
+                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
+                }
+                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
+                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
+                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
+                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
+            }
+            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) {
+                // handle the row-major layout for 8-bit types
+                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
+                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
+                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
+                int col = j*dst.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
+
+                U2 tmp[4];
+                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
+                    move<U2>::ldsm4(tmp[0], tmp[1], tmp[2], tmp[3], src.idx(shared_addr, {row, col}));
+                }
+                else {
+                    move<U2>::ldsm4t(tmp[0], tmp[2], tmp[1], tmp[3], src.idx(shared_addr, {row, col}));
+                }
+                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
+                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
+                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
+                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
+            }
+            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
+                // handle the row-major layout for 32-bit types
+                int row = (local_warpid*warp_height + i)*dst.tile_size_row + (warp_laneid / 4);
+                int col = j*dst.tile_size_col + 2*(warp_laneid % 4);
+                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
+                    row += src.row_offset;
+                    col += src.col_offset;
+                }
+                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
+                U2 tmp[4];
+                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
+                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
+                const int outer_idx = col/subtile_cols;
+                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
+                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
+                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
+                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
+                move<U>::lds(tmp[0].x, (addr_1+ 0)^swizzle_1);
+                move<U>::lds(tmp[0].y, (addr_1+ 4)^swizzle_1);
+                move<U>::lds(tmp[2].x, (addr_1+32)^swizzle_1);
+                move<U>::lds(tmp[2].y, (addr_1+36)^swizzle_1);
+                move<U>::lds(tmp[1].x, (addr_2+ 0)^swizzle_2);
+                move<U>::lds(tmp[1].y, (addr_2+ 4)^swizzle_2);
+                move<U>::lds(tmp[3].x, (addr_2+32)^swizzle_2);
+                move<U>::lds(tmp[3].y, (addr_2+36)^swizzle_2);
+                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
+                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
+                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
+                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
+                if(blit) {
+                    #pragma unroll
+                    for(int k = 0; k < 4; k++) {
+                        dst.tiles[i][j].data[k] = T2{dst.tiles[i][j].data[k].y, dst.tiles[i][j].data[k].x};
+                    }
+                }
+            }
+            else {
+                // handle the column-major layout
+                int row = (local_warpid*warp_height + i)*dst.tile_size_row + 2*(warp_laneid % 4);
+                int col = j*dst.tile_size_col + (warp_laneid / 4);
+                U2 tmp[4];
+                move<U>::lds(tmp[0].x, src.idx(shared_addr, {row+0, col+0}));
+                move<U>::lds(tmp[0].y, src.idx(shared_addr, {row+1, col+0}));
+                move<U>::lds(tmp[1].x, src.idx(shared_addr, {row+0, col+8}));
+                move<U>::lds(tmp[1].y, src.idx(shared_addr, {row+1, col+8}));
+                move<U>::lds(tmp[2].x, src.idx(shared_addr, {row+8, col+0}));
+                move<U>::lds(tmp[2].y, src.idx(shared_addr, {row+9, col+0}));
+                move<U>::lds(tmp[3].x, src.idx(shared_addr, {row+8, col+8}));
+                move<U>::lds(tmp[3].y, src.idx(shared_addr, {row+9, col+8}));
+                dst.tiles[i][j].data[0] = base_types::convertor<T2, U2>::convert(tmp[0]);
+                dst.tiles[i][j].data[1] = base_types::convertor<T2, U2>::convert(tmp[1]);
+                dst.tiles[i][j].data[2] = base_types::convertor<T2, U2>::convert(tmp[2]);
+                dst.tiles[i][j].data[3] = base_types::convertor<T2, U2>::convert(tmp[3]);
+            }
+        }
+    }
+}
+
+
+/**
+ * @brief Collaboratively store data into a shared tile from register tiles split across a warpgroup.
+ *
+ * @tparam RT The register tile type
+ * @tparam ST The shared tile type
+ * @param dst[out] The destination shared tile.
+ * @param src[in]  The source register tile.
+ */
+template<ducks::st::all ST, ducks::rt::all RT>
+__device__ inline static void store(ST &dst, const RT &src) {
+    constexpr int height = ST::height;
+    constexpr int warp_height = RT::height;
+    static_assert(height%GROUP_WARPS == 0, "Group load / store requires tile height to be a multiple of GROUP_WARPS.");
+    static_assert(height%warp_height == 0, "Group load / store requires tile height to be a multiple of the RT height.");
+    static_assert(ST::width==RT::width, "Group load / store requires tile widths to match.");
+    int local_warpid;
+    if constexpr(GROUP_WARPS % 4 == 0) local_warpid = (warpid()/4+(warpid()%4)*(GROUP_WARPS/4));
+    else local_warpid = warpid();
+    using T2 = RT::dtype;
+    using U  = ST::dtype;
+    using T  = base_types::packing<T2>::unpacked_type;
+    using U2 = base_types::packing<U>::packed_type;
+    int warp_laneid = ::kittens::laneid();
+
+    // convert to shared state space
+    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+
+    #pragma unroll
+    for(int i = 0; i < warp_height; i++) {
+        #pragma unroll
+        for(int j = 0; j < src.width; j++) {
+            if constexpr (sizeof(typename ST::dtype) == 2) {
+                // handle the row-major layout
+                U2 tmp[4];
+                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
+                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
+                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
+                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
+#ifdef KITTENS_HOPPER
+                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid % 16);
+                int col = j*src.tile_size_col + (warp_laneid / 16) * 8;
+                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
+                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
+                }
+                else {
+                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
+                }
+#else
+                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
+                    int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
+                    int col = j*src.tile_size_col + 2*(warp_laneid % 4);
+                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0]);
+                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[1]);
+                    move<U2>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[2]);
+                    move<U2>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3]);
+                }
+                else {
+                    int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
+                    int col = j*src.tile_size_col + (warp_laneid / 4);
+                    move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
+                    move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
+                    move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
+                    move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
+                    move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
+                    move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
+                    move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
+                    move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
+                }
+#endif
+            }
+            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 1) { 
+                // handle the row-major layout for 8-bit types
+
+                int warp_group_16 = (warp_laneid / 16);  // divide each warp into two groups of 16 threads
+                int lane_in_16 = warp_laneid % 16;       // position in group of 16 threads
+                int row = (local_warpid*warp_height + i)*src.tile_size_row + (lane_in_16 % 16); // find base row for warp in warpgroup and then distribute the 16 threads in the warp across the rows
+                int col = j*src.tile_size_col + warp_group_16 * 16; // find base column and then *16 for second half of the warp
+
+                U2 tmp[4];
+                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
+                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
+                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
+                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
+                if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row>) {
+                    move<U2>::stsm4(dst.idx(shared_addr, {row, col}), tmp[0], tmp[1], tmp[2], tmp[3]);
+                }
+                else {
+                    move<U2>::stsm4t(dst.idx(shared_addr, {row, col}), tmp[0], tmp[2], tmp[1], tmp[3]);
+                }
+            }
+            else if constexpr (std::is_same_v<typename RT::layout, ducks::rt_layout::row> && sizeof(typename ST::dtype) == 4) {
+                // handle the row-major layout for 32-bit types
+                int row = (local_warpid*warp_height + i)*src.tile_size_row + (warp_laneid / 4);
+                int col = j*src.tile_size_col + 2*(warp_laneid % 4);
+                if constexpr (ST::rows != ST::underlying_rows || ST::cols != ST::underlying_cols) { // subtile case
+                    row += dst.row_offset;
+                    col += dst.col_offset;
+                }
+                int blit = sizeof(typename ST::dtype) * ((warp_laneid%4) / 2);
+                T2 reg_tmp[4];
+                if(blit) {
+                    #pragma unroll
+                    for(int k = 0; k < 4; k++) {
+                        reg_tmp[k] = T2{src.tiles[i][j].data[k].y, src.tiles[i][j].data[k].x};
+                    }
+                }
+                else {
+                    #pragma unroll
+                    for(int k = 0; k < 4; k++) {
+                        reg_tmp[k] = src.tiles[i][j].data[k];
+                    }
+                }
+                U2 tmp[4];
+                tmp[0] = base_types::convertor<U2, T2>::convert(reg_tmp[0]);
+                tmp[1] = base_types::convertor<U2, T2>::convert(reg_tmp[1]);
+                tmp[2] = base_types::convertor<U2, T2>::convert(reg_tmp[2]);
+                tmp[3] = base_types::convertor<U2, T2>::convert(reg_tmp[3]);
+                static constexpr int swizzle_repeat = ST::swizzle_bytes * 8;
+                static constexpr int subtile_cols   = ST::swizzle_bytes / sizeof(U);
+                const int outer_idx = col/subtile_cols;
+                const uint32_t addr_1 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+0)*subtile_cols + col%subtile_cols);
+                const uint32_t addr_2 = shared_addr + sizeof(U)*(outer_idx*ST::underlying_rows*subtile_cols + (row+8)*subtile_cols + col%subtile_cols);
+                const int swizzle_1 = blit ^ ((addr_1 % swizzle_repeat) >> 7) << 4;
+                const int swizzle_2 = blit ^ ((addr_2 % swizzle_repeat) >> 7) << 4;
+                move<U>::sts((addr_1+ 0)^swizzle_1, tmp[0].x);
+                move<U>::sts((addr_1+ 4)^swizzle_1, tmp[0].y);
+                move<U>::sts((addr_1+32)^swizzle_1, tmp[2].x);
+                move<U>::sts((addr_1+36)^swizzle_1, tmp[2].y);
+                move<U>::sts((addr_2+ 0)^swizzle_2, tmp[1].x);
+                move<U>::sts((addr_2+ 4)^swizzle_2, tmp[1].y);
+                move<U>::sts((addr_2+32)^swizzle_2, tmp[3].x);
+                move<U>::sts((addr_2+36)^swizzle_2, tmp[3].y);
+            }
+            else {
+                // handle the column-major layout
+                int row = (local_warpid*warp_height + i)*src.tile_size_row + 2*(warp_laneid % 4);
+                int col = j*src.tile_size_col + (warp_laneid / 4);
+                U2 tmp[4];
+                tmp[0] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[0]);
+                tmp[1] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[1]);
+                tmp[2] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[2]);
+                tmp[3] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[3]);
+                move<U>::sts(dst.idx(shared_addr, {row+0, col+0}), tmp[0].x);
+                move<U>::sts(dst.idx(shared_addr, {row+1, col+0}), tmp[0].y);
+                move<U>::sts(dst.idx(shared_addr, {row+0, col+8}), tmp[1].x);
+                move<U>::sts(dst.idx(shared_addr, {row+1, col+8}), tmp[1].y);
+                move<U>::sts(dst.idx(shared_addr, {row+8, col+0}), tmp[2].x);
+                move<U>::sts(dst.idx(shared_addr, {row+9, col+0}), tmp[2].y);
+                move<U>::sts(dst.idx(shared_addr, {row+8, col+8}), tmp[3].x);
+                move<U>::sts(dst.idx(shared_addr, {row+9, col+8}), tmp[3].y);
+            }
+        }
+    }
+}
+
+// Load and store of vectors from/to shared tiles.
+
+template<ducks::rv::naive_layout RV, ducks::st::all ST>
+__device__ inline static auto load(RV &dst, const ST &src, int2 row_col) {
+    KITTENS_CHECK_WARP;
+    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
+    using T = RV::T;
+    using U = ST::T;
+    int warp_laneid = ::kittens::laneid();
+
+    // convert to shared state space
+    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
+
+    #pragma unroll
+    for(int col = warp_laneid; col < dst.length; col+=WARP_THREADS) {
+        U tmp;
+        move<U>::lds(tmp, src.idx(shared_addr, {row_col.x, row_col.y + col}));
+        dst.data[col/WARP_THREADS][0] = base_types::convertor<T, U>::convert(tmp);
+    }
+}
+
+template<ducks::rv::naive_layout RV, ducks::st::all ST>
+__device__ inline static auto store(ST &dst, const RV &src, int2 row_col) {
+    KITTENS_CHECK_WARP;
+    static_assert(ST::cols>=RV::length, "Shared tile must be at least as wide as the vector.");
+    using T = RV::T;
+    using U = ST::T;
+    int warp_laneid = ::kittens::laneid();
+
+    // convert to shared state space
+    uint32_t shared_addr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+
+    #pragma unroll
+    for(int col = warp_laneid; col < src.length; col+=WARP_THREADS) {
+        U tmp = base_types::convertor<U, T>::convert(src.data[col/WARP_THREADS][0]);
+        move<U>::sts(dst.idx(shared_addr, {row_col.x, row_col.y + col}), tmp);
+    }
+}

extra/thunder/cuda/include/ops/group/memory/tile/tensor_to_register.cuh

@@ -0,0 +1,325 @@
+/**
+ * @file
+ * @brief Group (collaborative warp) ops for loading tensor tiles into register tiles.
+ */
+
+/**
+ * @brief Load data from a tensor tile into a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam TM The tensor memory tile type
+ * @param dst[out] The destination register tile.
+ * @param src[in]  The source tensor tile.
+ */
+template<ducks::rt::row_layout RT, ducks::tt::all TM>
+__device__ inline static void load_async(RT &dst, const TM &src) {
+    if constexpr (GROUP_WARPS == 1) {
+        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
+        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
+
+        using T2 = RT::dtype;
+        using U  = typename TM::dtype;
+        using U2 = base_types::packing<typename TM::dtype>::packed_type;
+
+        if constexpr (sizeof(typename TM::dtype) == 1) {
+            #pragma unroll
+            for(int i = 0; i < dst.height; i++) {
+                #pragma unroll
+                for(int j = 0; j < dst.width; j++) {
+                    asm volatile(
+                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
+                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
+                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
+                    );
+                }
+            }
+        } else if constexpr (sizeof(typename TM::dtype) == 2) {
+            #pragma unroll
+            for(int i = 0; i < dst.height; i++) {
+                #pragma unroll
+                for(int j = 0; j < dst.width; j++) {
+                    asm volatile(
+                        "tcgen05.ld.sync.aligned.16x128b.x2.pack::16b.b32 {%0, %1, %2, %3}, [%4];\n"
+                        : "=r"(*(uint32_t*) &dst.tiles[i][j].data[0]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[1]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[2]),
+                            "=r"(*(uint32_t*) &dst.tiles[i][j].data[3])
+                        : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col))
+                    );
+                }
+            }
+        }
+        else if constexpr (sizeof(typename TM::dtype) == 4) {
+            #pragma unroll
+            for(int i = 0; i < dst.height; i++) {
+                if constexpr (dst.width%4 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < dst.width; j+=4) {
+                        U2 data[16];
+                        asm volatile(
+                            "tcgen05.ld.sync.aligned.16x256b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];\n"
+                            : "=f"(data[0].x), "=f"(data[0].y),
+                            "=f"(data[1].x), "=f"(data[1].y),
+                            "=f"(data[2].x), "=f"(data[2].y),
+                            "=f"(data[3].x), "=f"(data[3].y),
+                            "=f"(data[4].x), "=f"(data[4].y),
+                            "=f"(data[5].x), "=f"(data[5].y),
+                            "=f"(data[6].x), "=f"(data[6].y),
+                            "=f"(data[7].x), "=f"(data[7].y),
+                            "=f"(data[8].x), "=f"(data[8].y),
+                            "=f"(data[9].x), "=f"(data[9].y),
+                            "=f"(data[10].x), "=f"(data[10].y),
+                            "=f"(data[11].x), "=f"(data[11].y),
+                            "=f"(data[12].x), "=f"(data[12].y),
+                            "=f"(data[13].x), "=f"(data[13].y),
+                            "=f"(data[14].x), "=f"(data[14].y),
+                            "=f"(data[15].x), "=f"(data[15].y)
+                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
+                        );
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
+                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
+                            dst.tiles[i][j+2].data[k] = base_types::convertor<T2, U2>::convert(data[k+8]);
+                            dst.tiles[i][j+3].data[k] = base_types::convertor<T2, U2>::convert(data[k+12]);
+                        }
+                    }
+                }
+                else if constexpr (dst.width%2 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < dst.width; j+=2) {
+                        U2 data[8];
+                        asm volatile(
+                            "tcgen05.ld.sync.aligned.16x256b.x4.b32 {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];\n"
+                            : "=f"(data[0].x), "=f"(data[0].y),
+                            "=f"(data[1].x), "=f"(data[1].y),
+                            "=f"(data[2].x), "=f"(data[2].y),
+                            "=f"(data[3].x), "=f"(data[3].y),
+                            "=f"(data[4].x), "=f"(data[4].y),
+                            "=f"(data[5].x), "=f"(data[5].y),
+                            "=f"(data[6].x), "=f"(data[6].y),
+                            "=f"(data[7].x), "=f"(data[7].y)
+                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
+                        );
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            dst.tiles[i][j+0].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
+                            dst.tiles[i][j+1].data[k] = base_types::convertor<T2, U2>::convert(data[k+4]);
+                        }
+                    }
+                }
+                else {
+                    #pragma unroll
+                    for(int j = 0; j < dst.width; j++) {
+                        U2 data[4];
+                        asm volatile(
+                            "tcgen05.ld.sync.aligned.16x256b.x2.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];\n"
+                            : "=f"(data[0].x), "=f"(data[0].y),
+                            "=f"(data[1].x), "=f"(data[1].y),
+                            "=f"(data[2].x), "=f"(data[2].y),
+                            "=f"(data[3].x), "=f"(data[3].y)
+                            : "r"(src.addr + ((i * dst.tile_size_row) << 16) + (j * dst.tile_size_col)/(4/(uint32_t)sizeof(U)))
+                        );
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            dst.tiles[i][j].data[k] = base_types::convertor<T2, U2>::convert(data[k]);
+                        }
+                    }
+                }
+            }
+        }
+    }
+    else {
+        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
+        constexpr int warp_rows = TM::rows/GROUP_WARPS;
+        static_assert(TM::cols==RT::cols);
+        static_assert(warp_rows==RT::rows);
+        if constexpr (GROUP_WARPS == 4) {
+            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
+            ::kittens::group<1>::load_async(dst, src_subtile);
+        }
+        else {
+            auto src_subtile = src.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
+            ::kittens::group<1>::load_async(dst, src_subtile);
+        }
+    }
+}
+
+
+/**
+ * @brief Store data into a tensor tile from a register tile.
+ *
+ * @tparam RT The register tile type
+ * @tparam TM The tensor memory tile type
+ * @param dst[out] The destination tensor tile.
+ * @param src[in]  The source register tile.
+ */
+template<ducks::rt::all RT, ducks::tt::all TM>
+__device__ inline static void store_async(TM &dst, const RT &src) {
+    if constexpr (GROUP_WARPS == 1) {
+        static_assert(RT::height == TM::height, "register tile and tensor tile must match height");
+        static_assert(RT::width == TM::width, "register tile and tensor tile must match width");
+
+        using T2 = RT::dtype;
+        using T = base_types::packing<T2>::unpacked_type;
+        using U = TM::dtype;
+        using U2 = base_types::packing<U>::packed_type;
+
+        if constexpr (sizeof(typename TM::dtype) == 2) {
+            #pragma unroll
+            for(int i = 0; i < src.height; i++) {
+                if constexpr (src.width%4 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j+=4) {
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x128b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+2].data[3]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+3].data[3])
+                        );
+                    }
+                }
+                else if constexpr (src.width%2 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j+=2) {
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x128b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+0].data[3]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j+1].data[3])
+                        );
+                    }
+                }
+                else {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j++) {
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x128b.x2.b32 [%0], {%1, %2, %3, %4};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "r"(*(uint32_t*)&src.tiles[i][j].data[0]),
+                            "r"(*(uint32_t*)&src.tiles[i][j].data[1]),
+                            "r"(*(uint32_t*)&src.tiles[i][j].data[2]),
+                            "r"(*(uint32_t*)&src.tiles[i][j].data[3])
+                        );
+                    }
+                }
+            }
+        }
+        else if constexpr (sizeof(typename TM::dtype) == 4) {
+            #pragma unroll
+            for(int i = 0; i < src.height; i++) {
+                if constexpr(src.width%4 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j+=4) {
+                        U2 data[16];
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
+                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
+                            data[k+8] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+2].data[k]);
+                            data[k+12] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+3].data[k]);
+                        }
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x256b.x8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "f"(data[0].x), "f"(data[0].y),
+                            "f"(data[1].x), "f"(data[1].y),
+                            "f"(data[2].x), "f"(data[2].y),
+                            "f"(data[3].x), "f"(data[3].y),
+                            "f"(data[4].x), "f"(data[4].y),
+                            "f"(data[5].x), "f"(data[5].y),
+                            "f"(data[6].x), "f"(data[6].y),
+                            "f"(data[7].x), "f"(data[7].y),
+                            "f"(data[8].x), "f"(data[8].y),
+                            "f"(data[9].x), "f"(data[9].y),
+                            "f"(data[10].x), "f"(data[10].y),
+                            "f"(data[11].x), "f"(data[11].y),
+                            "f"(data[12].x), "f"(data[12].y),
+                            "f"(data[13].x), "f"(data[13].y),
+                            "f"(data[14].x), "f"(data[14].y),
+                            "f"(data[15].x), "f"(data[15].y)
+                        );
+                    }
+                }
+                else if constexpr(src.width%2 == 0) {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j+=2) {
+                        U2 data[8];
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
+                            data[k+4] = base_types::convertor<U2, T2>::convert(src.tiles[i][j+1].data[k]);
+                        }
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x256b.x4.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "f"(data[0].x), "f"(data[0].y),
+                            "f"(data[1].x), "f"(data[1].y),
+                            "f"(data[2].x), "f"(data[2].y),
+                            "f"(data[3].x), "f"(data[3].y),
+                            "f"(data[4].x), "f"(data[4].y),
+                            "f"(data[5].x), "f"(data[5].y),
+                            "f"(data[6].x), "f"(data[6].y),
+                            "f"(data[7].x), "f"(data[7].y)
+                        );
+                    }
+                }
+                else {
+                    #pragma unroll
+                    for(int j = 0; j < src.width; j++) {
+                        U2 data[4];
+                        #pragma unroll
+                        for(int k = 0; k < 4; k++) {
+                            data[k] = base_types::convertor<U2, T2>::convert(src.tiles[i][j].data[k]);
+                        }
+                        asm volatile(
+                            "tcgen05.st.sync.aligned.16x256b.x2.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};\n"
+                            :: "r"(dst.addr + ((i * src.tile_size_row) << 16) + (j * src.tile_size_col)/(4/(uint32_t)sizeof(U))),
+                            "f"(data[0].x), "f"(data[0].y),
+                            "f"(data[1].x), "f"(data[1].y),
+                            "f"(data[2].x), "f"(data[2].y),
+                            "f"(data[3].x), "f"(data[3].y)
+                        );
+                    }
+                }
+            }
+        }
+    }
+    else {
+        static_assert(GROUP_WARPS==4 || GROUP_WARPS==8);
+        constexpr int warp_rows = TM::rows/GROUP_WARPS;
+        static_assert(TM::cols==RT::cols);
+        static_assert(warp_rows==RT::rows);
+        if constexpr (GROUP_WARPS == 4) {
+            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*warpid(), 0);
+            ::kittens::group<1>::store_async(dst_subtile, src);
+        }
+        else {
+            auto dst_subtile = dst.template subtile<tt<typename TM::dtype, warp_rows, TM::cols>>(32*(warpid()%4)+16*(warpid()/4), 0);
+            ::kittens::group<1>::store_async(dst_subtile, src);
+        }
+    }
+}

extra/thunder/cuda/include/ops/group/memory/tile/tile.cuh

@@ -0,0 +1,16 @@
+/**
+ * @file
+ * @brief An aggregate header of group memory operations on tiles.
+ */
+
+#include "shared_to_register.cuh"
+#include "global_to_register.cuh"
+#include "global_to_shared.cuh"
+#ifdef KITTENS_BLACKWELL
+#include "tensor_to_register.cuh"
+#endif
+
+#include "complex/complex_shared_to_register.cuh"
+#include "complex/complex_global_to_register.cuh"
+#include "complex/complex_global_to_shared.cuh"
+

extra/thunder/cuda/include/ops/group/memory/tile/tma.cuh

@@ -0,0 +1,134 @@
+/**
+ * @file
+ * @brief Functions for a group scope to call tile TMA functions.
+ */
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::prefetch<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void prefetch(ST &dst, const GL &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::prefetch<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_add_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_add_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_add_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_add_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_add_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_min_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_min_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_min_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_min_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_min_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_max_async<axis, policy, ST, GL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_max_async(const GL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_max_async<axis, policy, ST, PGL, COORD>(dst, src, idx); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::pgl::all PGL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void store_max_async(const PGL &dst, const ST &src, const COORD &idx) {
+    if(laneid() == 0) {
+        ::kittens::tma::store_max_async<dim::ROW, cache_policy::NORMAL, ST, PGL, COORD>(dst, src, idx);
+    }
+}
+
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
+    if(laneid() == 0) {
+        ::kittens::tma::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar); // Don't do the mask
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar) {
+    if(laneid() == 0) {
+        ::kittens::tma::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar);
+    }
+}

extra/thunder/cuda/include/ops/group/memory/tile/tma_cluster.cuh

@@ -0,0 +1,33 @@
+/**
+ * @file
+ * @brief Functions for a group scope to call tile TMA cluster functions.
+ */
+
+
+#ifdef KITTENS_BLACKWELL
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask, int dst_mbar_cta=-1) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask, dst_mbar_cta);
+    }
+}
+#else
+template<int axis, cache_policy policy, ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::load_async<axis, policy, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
+    }
+}
+template<ducks::st::all ST, ducks::gl::all GL, ducks::coord::tile COORD=coord<ST>>
+__device__ static inline void load_async(ST &dst, const GL &src, const COORD &idx, semaphore& bar, uint16_t cluster_mask) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::load_async<dim::ROW, cache_policy::NORMAL, ST, GL, COORD>(dst, src, idx, bar, cluster_mask);
+    }
+}
+#endif

extra/thunder/cuda/include/ops/group/memory/util/tma.cuh

@@ -0,0 +1,68 @@
+/**
+ * @file
+ * @brief Various utilities for group TMA memory operations.
+ */
+
+/* ----------   Barrier functions for async load  ---------- */
+
+/**
+* @brief Sets the number of bytes expected at the semaphore.
+*
+* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
+* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
+* instruction to set the expected number of bytes.
+*
+* @param semaphore Reference to the semaphore variable.
+* @param bytes The number of bytes expected at the semaphore.
+*/
+__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes) {
+    if(laneid() == 0) {
+        ::kittens::tma::expect_bytes(bar, bytes);
+    }
+}
+/**
+* @brief Sets the number of bytes expected at the semaphore.
+*
+* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
+*/
+template<typename T, typename... args>
+__device__ static inline void expect(semaphore& bar, const T& _1, const args&... _2) {
+    expect_bytes(bar, size_bytes<T, args...>);
+}
+
+/* ----------   Synchronization functions for async store  ---------- */
+
+/**
+ * @brief Commits previous asynchronous TMA stores to a group and performs them.
+*/
+__device__ static inline void store_commit_group() {
+    asm volatile("cp.async.bulk.commit_group;");
+}
+/**
+ * @brief Waits for previous committed TMA store groups to complete.
+ *
+ * @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
+*/
+template <int N=0>
+__device__ static inline void store_async_wait() {
+    asm volatile (
+        "cp.async.bulk.wait_group %0;"
+        :
+        : "n"(N)
+        : "memory"
+    );
+}
+/**
+ * @brief Waits for previous committed TMA store groups to finish reading from shared memory.
+ *
+ * @tparam N The maximum number of remaining TMA store groups. Defaults to 0.
+*/
+template <int N=0>
+__device__ static inline void store_async_read_wait() {
+    asm volatile (
+        "cp.async.bulk.wait_group.read %0;"
+        :
+        : "n"(N)
+        : "memory"
+    );
+}

extra/thunder/cuda/include/ops/group/memory/util/tma_cluster.cuh

@@ -0,0 +1,90 @@
+
+/**
+* @brief Waits for the requested semaphore phase, at cluster scope
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline void wait(semaphore& bar, int kPhaseBit) {
+    void const* const ptr = &bar;
+    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
+
+    asm volatile (
+        "{\n"
+        ".reg .pred                P1;\n"
+        "LAB_WAIT:\n"
+        "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64 P1, [%0], %1;\n"
+        "@P1                       bra.uni DONE;\n"
+        "bra.uni                   LAB_WAIT;\n"
+        "DONE:\n"
+        "}\n"
+        :: "r"(mbar_ptr),
+        "r"(kPhaseBit)
+    );
+}
+
+/**
+* @brief Sets the number of bytes expected at the semaphore, assuming a multicast instruction.
+*
+* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
+* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
+* instruction to set the expected number of bytes.
+* 
+* It's worth being aware that this function is particularly necessary for multicast loads, and
+* distributed shared memory can actually be done with a normal tma::expect followed by wait. See
+* the unit tests of dsmem for an example.
+*
+* @param semaphore Reference to the semaphore variable.
+* @param bytes The number of bytes expected at the semaphore.
+*/
+__device__ static inline void expect_bytes(semaphore& bar, uint32_t bytes, int dst_cta) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::expect_bytes(bar, bytes, dst_cta);
+    }
+}
+/**
+* @brief Sets the number of bytes expected at the semaphore.
+*
+* This function sets the number of bytes expected at the semaphore for the first thread in the warp.
+* It converts the semaphore pointer to a generic shared memory pointer and uses an inline assembly
+* instruction to set the expected number of bytes.
+*
+* @tparam T The type of the data to be stored at the semaphore.
+* @param semaphore Reference to the semaphore variable.
+*/
+/**
+* @brief Sets the number of bytes expected at the semaphore.
+*
+* This function sets the number of bytes expected at the mbarrier before the transaction arrives.
+*/
+template<typename T, typename... args>
+__device__ static inline void expect(semaphore& bar, int dst_cta, const T& _1, const args&... _2) {
+    expect_bytes(bar, size_bytes<T, args...>, dst_cta);
+}
+
+/**
+* @brief Arrives at a semaphore in cluster scope.
+*
+* Marks a thread arrival at an mbarrier
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline void arrive(semaphore& bar, int dst_cta, uint32_t count=1) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::arrive(bar, dst_cta, count);
+    }
+}
+
+// Generic transfer
+__device__ static inline void store_async(void *dst, void *src, int dst_cta, uint32_t size_bytes, semaphore& bar) {
+    if(laneid() == 0) {
+        ::kittens::tma::cluster::store_async(dst, src, dst_cta, size_bytes, bar);
+    }
+}
+
+// Templated transfer for convenience
+template<typename T>
+__device__ static inline void store_async(T &dst_, T &src_, int dst_cta, semaphore& bar) {
+    store_async((void*)&dst_, (void*)&src_, dst_cta, size_bytes<T>, bar);
+}

extra/thunder/cuda/include/ops/group/memory/util/util.cuh

@@ -0,0 +1,168 @@
+/**
+ * @file
+ * @brief Various utilities for group memory operations.
+ */
+
+
+template<int N=0> __device__ static inline void load_async_wait(int bar_id) { // for completing (non-TMA) async loads
+    asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
+    sync(bar_id);
+}
+template<int N=0> __device__ static inline void load_async_wait() { // for completing (non-TMA) async loads
+    KITTENS_CHECK_WARP
+    asm volatile("cp.async.wait_group %0;\n" : : "n"(N) : "memory");
+    __syncwarp();
+}
+
+__device__ static inline void arrive(barrier<GROUP_WARPS> bar) {
+    asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
+}
+__device__ static inline void arrive_and_wait(barrier<GROUP_WARPS> bar) {
+    asm volatile("bar.sync %0, %1;\n" :: "r"(bar.barrier_id), "n"(GROUP_WARPS*WARP_THREADS) : "memory");
+}
+
+/**
+ * @brief Initializes a synchronization semaphore with a transaction count and sets the expected number of bytes.
+ *
+ * This function sets up a semaphore that is used to synchronize threads within a block during asynchronous operations.
+ * It initializes the semaphore with a thread count semaphore.
+ *
+ * Additionally, if it is given a shared tile type, it will also call `set_bytes` to prepare for the memory transaction.
+ *
+ * @param[out] semaphore The semaphore variable to initialize.
+ * @param[in] tc The thread counter for the semaphore.
+ */
+__device__ static inline void init_semaphore(semaphore& bar, int thread_count, int transaction_count=0) {
+    if (laneid() == 0) {
+        void const* const ptr = &bar;
+        uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
+
+        asm volatile (
+            "mbarrier.init.shared::cta.b64 [%0], %1;\n"
+            :: "r"(bar_ptr), "r"(thread_count+transaction_count)
+        );
+    }
+}
+/**
+ * @brief Invalidate an mbarrier
+ *
+ * @param[out] semaphore The semaphore variable to initialize.
+ * @param[in] tc The thread counter for the semaphore.
+ */
+__device__ static inline void invalidate_semaphore(semaphore& bar) {
+    if (laneid() == 0) {
+        void const* const ptr = &bar;
+        uint32_t bar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
+        asm volatile (
+            "mbarrier.inval.shared::cta.b64 [%0];\n"
+            :: "r"(bar_ptr)
+        );
+    }
+}
+
+/**
+* @brief Arrives at a semaphore.
+*
+* Marks a warp arrival at an mbarrier
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline void arrive(semaphore& sem) {
+    if(laneid() == 0) {
+            uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem)); 
+            asm volatile (
+                "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];\n"
+            :
+            : "r"(mbar_ptr)
+            : "memory"
+        );
+    }
+}
+template<int num_warps> __device__ static inline void arrive(barrier<num_warps> bar) {
+    asm volatile("bar.arrive %0, %1;\n" :: "r"(bar.barrier_id), "n"(num_warps*WARP_THREADS) : "memory");
+}
+
+#ifdef KITTENS_HOPPER
+/**
+* @brief Arrives at a semaphore.
+*
+* Marks a warp arrival at an mbarrier
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline void arrive(semaphore& sem, uint32_t count) {
+    if(laneid() == 0) {
+        uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&sem));
+        asm volatile (
+            "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0], %1;\n"
+            :
+            : "r"(mbar_ptr), "r"(count)
+            : "memory"
+        );
+    }
+}
+#endif
+
+/**
+* @brief Waits for the requested semaphore phase.
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline void wait(semaphore& sem, int kPhaseBit) {
+    void const* const ptr = &sem;
+    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr)); 
+
+#ifdef KITTENS_HOPPER
+    asm volatile (
+        "{\n"
+        ".reg .pred                P1;\n"
+        "LAB_WAIT:\n"
+        "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+        "@P1                       bra.uni DONE;\n"
+        "bra.uni                   LAB_WAIT;\n"
+        "DONE:\n"
+        "}\n"
+        :: "r"(mbar_ptr),
+        "r"(kPhaseBit)
+    );
+#else
+    asm volatile (
+        "{\n"
+        ".reg .pred                P1;\n"
+        "LAB_WAIT:\n"
+        "mbarrier.test_wait.parity.shared::cta.b64 P1, [%0], %1;\n"
+        "@P1                       bra.uni DONE;\n"
+        "nanosleep.u32 5;\n" // wait a few nanoseconds on pre-Hopper architectures to save instruction issue slots
+        "bra.uni                   LAB_WAIT;\n"
+        "DONE:\n"
+        "}\n"
+        :: "r"(mbar_ptr),
+        "r"(kPhaseBit)
+    );
+#endif
+}
+
+/**
+* @brief Checks if the requested semaphore phase is ready.
+*
+* @param semaphore Reference to the semaphore variable.
+* @param kPhaseBit The phase bit used for the semaphore.
+*/
+__device__ static inline int test_wait(semaphore& sem, int kPhaseBit) {
+    void const* const ptr = &sem;
+    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+    int result;
+    asm volatile (
+        "{\n"
+        ".reg .pred P1;\n"
+        "mbarrier.test_wait.parity.shared::cta.b64 P1, [%1], %2;\n"
+        "selp.u32 %0,1,0,P1;"
+        "}\n"
+        : "=r"(result)
+        : "r"(mbar_ptr), "r"(kPhaseBit)
+    );
+    return result;
+}

extra/thunder/cuda/include/ops/group/memory/vec/global_to_register.cuh

@@ -0,0 +1,138 @@
+/**
+ * @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<ducks::rv::all RV, ducks::gl::all GL>
+__device__ inline static void load(RV &dst, const GL &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
+    if constexpr (GROUP_WARPS == 1) {
+        using T2 = RV::dtype;
+        using U = typename GL::dtype;
+        using U2 = base_types::packing<U>::packed_type;
+        using T = base_types::packing<T2>::unpacked_type;
+
+        U *src_ptr = (U*)&src[(idx.template unit_coord<-1, 3>())];
+        int laneid = ::kittens::laneid();
+        
+        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
+            #pragma unroll
+            for(auto w = 0; w < (dst.outer_dim+3)/4; w++) {
+                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
+                int o_dim = w*4 + (laneid/4) / 2;
+                int i_dim = (laneid/4) % 2;
+                // this should be a maximally coalesced load.
+                if(idx < dst.outer_dim*16)
+                    dst[o_dim][i_dim] = base_types::convertor<T2, U2>::convert(*(U2*)&src_ptr[idx]);
+            }
+            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                int leader = 8*(w%4) + (laneid%4); // repeats every 64 columns
+                dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader);
+                dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4);
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
+            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
+            // otherwise there will be some pain :/
+            #pragma unroll
+            for(auto w = 0; w < (dst.outer_dim+1)/2; w++) {
+                int idx = w*32 + (laneid%4)*8 + (laneid/4);
+                int o_dim = w*2 + (laneid%4) / 2;
+                // this should be a maximally coalesced load.
+                if(idx < dst.outer_dim*16) {
+                    T tmp = base_types::convertor<T, U>::convert(src_ptr[idx]);
+                    if(laneid%2==0) dst[o_dim][0].x = tmp;
+                    else dst[o_dim][0].y = tmp;
+                }
+            }
+            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                int leader = (laneid/4)*4 + 2*(w%2); // repeats every 64 columns
+                dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader);
+                dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1);
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                if(w < dst.outer_dim-1 || dst.length%32 == 0 || laneid<16) {
+                    dst[w][0] = base_types::convertor<T, U>::convert(src_ptr[w*32 + laneid]);
+                }
+            }
+        }
+    }
+    else {
+        // Call warp level load
+        ::kittens::group<1>::load(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
+    }
+}
+/**
+ * @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<ducks::rv::all RV, ducks::gl::all GL>
+__device__ inline static void store(GL &dst, const RV &src, const coord<rv<typename RV::T, GROUP_WARPS*RV::length, typename RV::layout>> &idx) {
+    if constexpr (GROUP_WARPS == 1) {
+        using T2 = RV::dtype;
+        using U = typename GL::dtype;
+        using U2 = base_types::packing<U>::packed_type;
+        using T = base_types::packing<T2>::unpacked_type;
+        
+        U *dst_ptr = (U*)&dst[(idx.template unit_coord<-1, 3>())];
+        int laneid = ::kittens::laneid();
+        
+        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
+            #pragma unroll
+            for(auto w = 0; w < (src.outer_dim+3)/4; w++) {
+                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
+                int o_dim = w*4 + (laneid/4) / 2;
+                int i_dim = (laneid/4) % 2;
+                // this should be a maximally coalesced store. I hope!
+                if(idx < src.outer_dim*16)
+                    *(U2*)&dst_ptr[idx] = base_types::convertor<U2, T2>::convert(src[o_dim][i_dim]);
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
+            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
+            // otherwise there will be some pain :/
+            #pragma unroll
+            for(auto w = 0; w < (src.outer_dim+1)/2; w++) {
+                int idx = w*32 + (laneid%4)*8 + (laneid/4);
+                int o_dim = w*2 + (laneid%4) / 2;
+                // this should be a maximally coalesced load.
+                if(idx < src.outer_dim*16) {
+                    U tmp;
+                    if(laneid%2==0) tmp = base_types::convertor<U, T>::convert(src[o_dim][0].x);
+                    else tmp = base_types::convertor<U, T>::convert(src[o_dim][0].y);
+                    dst_ptr[idx] = tmp;
+                }
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
+            #pragma unroll
+            for(auto w = 0; w < src.outer_dim; w++) {
+                if(w < src.outer_dim-1 || src.length%32 == 0 || laneid<16) {
+                    dst_ptr[w*32 + laneid] = base_types::convertor<U, T>::convert(src[w][0]);
+                }
+            }
+        }
+    }
+    else {
+        // Call warp level store
+        ::kittens::group<1>::store(dst, src, coord<RV>(idx.b, idx.d, idx.r, idx.c*GROUP_WARPS+warpid()));
+    }
+}

extra/thunder/cuda/include/ops/group/memory/vec/global_to_shared.cuh

@@ -0,0 +1,77 @@
+/**
+ * @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<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void load(SV &dst, const GL &src, const COORD &idx) {
+    constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
+    constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
+    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())];
+    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+    #pragma unroll
+    for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
+        if(i * elem_per_transfer < dst.length) {
+            float4 tmp;
+            move<float4>::ldg(tmp, (float4*)&src_ptr[i*elem_per_transfer]);
+            move<float4>::sts(dst_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer, tmp);
+        }
+    }
+}
+
+/**
+ * @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<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store(GL &dst, const SV &src, const COORD &idx) {
+    constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
+    constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
+    typename GL::dtype *dst_ptr = (typename GL::dtype*)&dst[(idx.template unit_coord<-1, 3>())];
+    uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
+    #pragma unroll
+    for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
+        if(i * elem_per_transfer < src.length) {
+            float4 tmp;
+            move<float4>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*i*elem_per_transfer);
+            move<float4>::stg((float4*)&dst_ptr[i*elem_per_transfer], tmp);
+        }
+    }
+}
+
+template<ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx) {
+    constexpr uint32_t elem_per_transfer = sizeof(float4) / sizeof(typename SV::dtype);
+    constexpr uint32_t total_calls = SV::length / elem_per_transfer; // guaranteed to divide
+    typename GL::dtype *src_ptr = (typename GL::dtype*)&src[(idx.template unit_coord<-1, 3>())];
+    uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+    #pragma unroll
+    for(uint32_t i = threadIdx.x%GROUP_THREADS; i < total_calls; i+=GROUP_THREADS) {
+        if(i * elem_per_transfer < dst.length) {
+            asm volatile(
+                "cp.async.cg.shared.global.L2::128B [%0], [%1], 16;\n"
+                :: "r"(dst_ptr + (uint32_t)sizeof(typename SV::dtype)*i*elem_per_transfer), "l"((uint64_t)&src_ptr[i*elem_per_transfer])
+                : "memory"
+            );
+        }
+    }
+    asm volatile("cp.async.commit_group;\n" ::: "memory");
+}

extra/thunder/cuda/include/ops/group/memory/vec/shared_to_register.cuh

@@ -0,0 +1,159 @@
+/**
+ * @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<ducks::rv::all RV, ducks::sv::all SV>
+__device__ inline static void load(RV &dst, const SV &src) {
+    using T2 = RV::dtype;
+    using U = SV::dtype;
+    using U2 = base_types::packing<U>::packed_type;
+    using T = base_types::packing<T2>::unpacked_type;
+    if constexpr (GROUP_WARPS == 1) {
+        static_assert(SV::length == RV::length);
+        
+        int laneid = ::kittens::laneid();
+        uint32_t src_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&src.data[0]));
+        
+        __syncwarp();
+        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
+            #pragma unroll
+            for(auto w = 0; w < (dst.outer_dim+3)/4; w++) {
+                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
+                int o_dim = w*4 + (laneid/4) / 2;
+                int i_dim = (laneid/4) % 2;
+                // this should be a maximally coalesced load.
+                if(idx < dst.outer_dim*16) {
+                    U2 tmp;
+                    move<U2>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx);
+                    dst[o_dim][i_dim] = base_types::convertor<T2, U2>::convert(tmp);
+                }
+            }
+            __syncwarp();
+            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                int leader = 8*(w%4) + (laneid%4); // repeats every 64 columns
+                dst[w][0] = packed_shfl_sync(MASK_ALL, dst[w][0], leader);
+                dst[w][1] = packed_shfl_sync(MASK_ALL, dst[w][1], leader+4);
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
+            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
+            // otherwise there will be some pain :/
+            #pragma unroll
+            for(auto w = 0; w < (dst.outer_dim+1)/2; w++) {
+                int idx = w*32 + (laneid%4)*8 + (laneid/4);
+                int o_dim = w*2 + (laneid%4) / 2;
+                // this should be a maximally coalesced load.
+                if(idx < dst.outer_dim*16) {
+                    U tmp;
+                    move<U>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*idx);
+                    if(laneid%2==0) dst[o_dim][0].x =  base_types::convertor<T, U>::convert(tmp);
+                    else dst[o_dim][0].y = base_types::convertor<T, U>::convert(tmp);
+                }
+            }
+            __syncwarp();
+            // now we need to do a bunch of shuffle_sync's to make sure everyone has everything they need.
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                int leader = (laneid/4)*4 + 2*(w%2); // repeats every 64 columns
+                dst[w][0].x = __shfl_sync(MASK_ALL, dst[w][0].x, leader);
+                dst[w][0].y = __shfl_sync(MASK_ALL, dst[w][0].y, leader+1);
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
+            #pragma unroll
+            for(auto w = 0; w < dst.outer_dim; w++) {
+                if(w < dst.outer_dim-1 || RV::length%32 == 0 || laneid<16) {
+                    U tmp;
+                    move<U>::lds(tmp, src_ptr + sizeof(typename SV::dtype)*(w*32 + laneid));
+                    dst[w][0] = base_types::convertor<T, U>::convert(tmp);
+                }
+            }
+        }
+    }
+    else {
+        static_assert(SV::length == RV::length*GROUP_WARPS);// confirm size correct
+        auto &_src = src.template subvec<RV::length>(warpid()); // pretend it's smaller and do warp-level load
+
+        ::kittens::group<1>::load(dst, _src); // 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<ducks::sv::all SV, ducks::rv::all RV>
+__device__ inline static void store(SV &dst, const RV &src) {
+    using T2 = RV::dtype;
+    using U = SV::dtype;
+    using U2 = base_types::packing<U>::packed_type;
+    using T = base_types::packing<T2>::unpacked_type;
+
+    if constexpr (GROUP_WARPS == 1) {
+        static_assert(SV::length == RV::length);
+        
+        int laneid = ::kittens::laneid();
+        uint32_t dst_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&dst.data[0]));
+
+        __syncwarp();
+        if constexpr (std::is_same_v<typename RV::layout, align_l>) {
+            #pragma unroll
+            for(auto w = 0; w < (src.outer_dim+3)/4; w++) {
+                int idx = w*64 + (laneid/4)*8 + 2*(laneid%4);
+                int o_dim = w*4 + (laneid/4) / 2;
+                int i_dim = (laneid/4) % 2;
+                // this should be a maximally coalesced store. I hope!
+                if(idx < src.outer_dim*16) {
+                    U2 tmp = base_types::convertor<U2, T2>::convert(src[o_dim][i_dim]);
+                    move<U2>::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp);
+                }
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, ortho_l>) {
+            // really hoping https://stackoverflow.com/questions/15029765/is-coalescing-triggered-for-accessing-memory-in-reverse-order is still true
+            // otherwise there will be some pain :/
+            #pragma unroll
+            for(auto w = 0; w < (src.outer_dim+1)/2; w++) {
+                int idx = w*32 + (laneid%4)*8 + (laneid/4);
+                int o_dim = w*2 + (laneid%4) / 2;
+                // this should be a maximally coalesced load.
+                if(idx < src.outer_dim*16) {
+                    U tmp;
+                    if(laneid%2==0) tmp = base_types::convertor<U, T>::convert(src[o_dim][0].x);
+                    else tmp = base_types::convertor<U, T>::convert(src[o_dim][0].y);
+                    move<U>::sts(dst_ptr + sizeof(typename SV::dtype)*idx, tmp);
+                }
+            }
+        }
+        else if constexpr (std::is_same_v<typename RV::layout, naive_l>) {
+            #pragma unroll
+            for(auto w = 0; w < src.outer_dim; w++) {
+                if(w < src.outer_dim-1 || RV::length%32 == 0 || laneid<16) {
+                    U tmp = base_types::convertor<U, T>::convert(src[w][0]);
+                    move<U>::sts(dst_ptr + sizeof(typename SV::dtype)*(w*32 + laneid), tmp);
+                }
+            }
+        }
+    }
+    else {
+        static_assert(SV::length == RV::length*GROUP_WARPS);// confirm size correct
+        auto &_dst = dst.template subvec<RV::length>(warpid()); // pretend it's smaller and do warp-level load
+
+        ::kittens::group<1>::store(_dst, src); // warp-level
+    }
+}

extra/thunder/cuda/include/ops/group/memory/vec/tma.cuh

@@ -0,0 +1,221 @@
+/**
+ * @file
+ * @brief Functions for a group scope to call vec TMA functions.
+ */
+
+/* ----------   Prefetch Tensor Map  ---------- */
+
+/**
+ * @brief Prefetches data from global memory into a shared memory vector, along with the tensormap.
+ *
+ * @tparam SV A shared vector type with a TMA-compatible layout
+ * @param[out] dst The destination shared memory vector.
+ * @param[in] src_tma_map The source tensormap address in global memory
+ * @param[in] vec_idx The coord of the requested vector.
+ */
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void prefetch(SV &dst, const GL &src, const COORD &idx) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        ::kittens::detail::tma::vec_prefetch_tma_internal<policy>(tma_ptr, tma_coord);
+    }
+}
+__KITTENS_TMA_DEFINE_DEFAULT_LOAD_CACHE_VEC__(prefetch)
+
+
+/* ----------   Async load and store data from gmem/smem  ---------- */
+
+/**
+ * @brief Asynchronously stores data into global memory from a shared memory vector.
+ *
+ * This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
+ *
+ * @tparam SV A shared vector type with a TMA-compatible layout
+ * @param[out] dst_tma_map The destination tensormap address in global memory
+ * @param[in] src The source shared memory vector.
+ * @param[in] vec_idx The coord of the vector destination.
+ */
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_async(const GL &dst, const SV &src, const COORD &idx) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_async)
+
+template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_async(const PGL &dst, const SV &src, const COORD &idx) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_async)
+
+
+/**
+* @brief Asynchronously performs an add reduction and stores the result into global memory.
+*
+* This function performs an asynchronous add reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
+*
+* @tparam SV A shared vector type with a TMA-compatible layout
+* @param[out] dst_tma_map The destination tensormap address in global memory
+* @param[in] src The source shared memory vector.
+* @param[in] vec_idx The coord of the vector destination.
+*/
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_add_async(const GL &dst, const SV &src, const COORD &idx) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_add_async)
+
+template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_add_async(const PGL &dst, const SV &src, const COORD &idx) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_add_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_add_async)
+
+
+/**
+* @brief Asynchronously performs an min reduction and stores the result into global memory.
+*
+* This function performs an asynchronous min reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
+*
+* @tparam SV A shared vector type with a TMA-compatible layout
+* @param[out] dst_tma_map The destination tensormap address in global memory
+* @param[in] src The source shared memory vector.
+* @param[in] vec_idx The coord of the vector destination.
+*/
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_min_async(const GL &dst, const SV &src, const COORD &idx) {
+    static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_min_async)
+
+template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_min_async(const PGL &dst, const SV &src, const COORD &idx) {
+    static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_min_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_min_async)
+
+/**
+* @brief Asynchronously performs an max reduction and stores the result into global memory.
+*
+* This function performs an asynchronous max reduction operation using CUDA's cp.reduce.async.bulk.tensor instruction.
+*
+* @tparam SV A shared vector type with a TMA-compatible layout
+* @param[out] dst_tma_map The destination tensormap address in global memory
+* @param[in] src The source shared memory vector.
+* @param[in] vec_idx The coord of the vector destination.
+*/
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_max_async(const GL &dst, const SV &src, const COORD &idx) {
+    static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_DEFAULT_STORE_CACHE_VEC__(store_max_async)
+
+template<cache_policy policy, ducks::sv::all SV, ducks::pgl::all PGL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void store_max_async(const PGL &dst, const SV &src, const COORD &idx) {
+    static_assert(!std::is_same_v<typename SV::dtype, float>, "TMA does not support async min/max reductions for fp32 types.");
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(dst.template get_tma<SV, -1>());
+    uint32_t src_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&src));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t src_i_ptr = src_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_store_max_async_tma_internal<policy>(tma_ptr, src_i_ptr, tma_coord);
+    }
+    store_commit_group();
+}
+__KITTENS_TMA_DEFINE_PGL_DEFAULT_STORE_CACHE_VEC__(store_max_async)
+
+/**
+ * @brief Asynchronously loads data from global memory into a shared memory vector.
+ *
+ * This function performs an asynchronous copy operation using CUDA's cp.async.bulk.tensor instruction.
+ *
+ * @tparam SV A shared vector type with a TMA-compatible layout
+ * @param[out] dst The destination shared memory vector.
+ * @param[in] src_tma_map The source tensormap address in global memory
+ * @param[in] vec_idx The coord of the requested vector.
+ * @param[in,out] bar The semaphore used for synchronization of the asynchronous copy.
+ */
+template<cache_policy policy, ducks::sv::all SV, ducks::gl::all GL, ducks::coord::vec COORD=coord<SV>>
+__device__ static inline void load_async(SV &dst, const GL &src, const COORD &idx, semaphore& bar) {
+    coord<> unit_coord = idx.template unit_coord<-1, 3>();
+    uint64_t tma_ptr  = reinterpret_cast<uint64_t>(src.template get_tma<SV, -1>());
+    uint32_t mbar_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(&bar));
+    uint32_t dst_ptr  = static_cast<uint32_t>(__cvta_generic_to_shared(&dst));
+    for(int i = ::kittens::laneid(); i < ::kittens::detail::tma::sv_tma_dim2<SV>; i += WARP_THREADS) {
+        coord<> tma_coord = unit_coord;
+        tma_coord.c += i * ::kittens::detail::tma::sv_tma_dim1<SV>;
+        uint32_t dst_i_ptr = dst_ptr + i*::kittens::detail::tma::sv_tma_dim1<SV>*sizeof(typename SV::dtype);
+        ::kittens::detail::tma::vec_load_async_tma_internal<policy>(tma_ptr, dst_i_ptr, mbar_ptr, tma_coord);
+    }
+}
+__KITTENS_TMA_DEFINE_SEMAPHORE_CACHE_VEC__(load_async)