can't allocate this at once, need parts back

looking for better ways to write this. first adding some tests

.github/workflows/benchmark.yml

@@ -49,9 +49,11 @@ jobs:
     - name: reset process replay
       run: python3.11 test/external/process_replay/reset.py
     - name: Run Stable Diffusion
-      run: JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd.txt
-    - name: Run Stable Diffusion with fp16
-      run: JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd_fp16.txt
+      run: JIT=1 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+    - name: Run Stable Diffusion without fp16
+      run: JIT=1 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+    - name: Run Stable Diffusion v2
+      run: JIT=1 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
     - name: Run SDXL
       run: JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run model inference benchmark
@@ -60,10 +62,12 @@ jobs:
       run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
       run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-    - name: Run Tensor Core GEMM
-      run: |
-        DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
-        DEBUG=2 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+    - name: Run Tensor Core GEMM (float)
+      run: DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
+    - name: Run Tensor Core GEMM (half)
+      run: DEBUG=2 HALF=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_half.txt
+    - name: Run Tensor Core GEMM (bfloat16)
+      run: DEBUG=2 BFLOAT16=1 python3.11 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
     - name: Fuzz Padded Tensor Core GEMM
       run: METAL=1 M_START=6 M_STOP=10 M_STEP=1 N_START=6 N_STOP=10 N_STEP=1 K_START=6 K_STOP=24 K_STEP=1 TC_OPT=2 DEBUG=2 python3.11 ./extra/gemm/fuzz_matmul.py
     - name: Run LLaMA
@@ -91,7 +95,7 @@ jobs:
     - name: Run 10 CIFAR training steps
       run: JIT=1 STEPS=10 python3.11 examples/hlb_cifar10.py | tee train_cifar.txt
     - name: Run 10 CIFAR training steps w HALF
-      run: JIT=1 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
+      run: JIT=2 STEPS=10 DEFAULT_FLOAT=HALF python3.11 examples/hlb_cifar10.py | tee train_cifar_half.txt
     #- name: Run 10 CIFAR training steps w BF16
     #  run: STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3.11 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     - name: Run 10 CIFAR training steps w winograd
@@ -114,8 +118,10 @@ jobs:
           gpt2_half_beam.txt
           matmul.txt
           matmul_half.txt
+          matmul_bfloat16.txt
           sd.txt
-          sd_fp16.txt
+          sd_no_fp16.txt
+          sdv2.txt
           sdxl.txt
           beautiful_mnist.txt
           train_cifar.txt
@@ -159,7 +165,7 @@ jobs:
     - name: Test speed vs torch
       run: NV=1 RUN_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
-      run: NV=1 IGNORE_BEAM_CACHE=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
         NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
@@ -172,10 +178,12 @@ jobs:
       run: NV=1 PTX=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
     - name: Run Tensor Core GEMM (NV)
       run: NV=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_nv.txt
-    # - name: Run Tensor Core GEMM (NV) with BEAM
-    #   run: BEAM=4 NV=1 HALF=1 IGNORE_BEAM_CACHE=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
+    - name: Test NV=1
+      run: DEBUG=2 NV=1 python -m pytest -rA test/test_tiny.py
+    - name: Test CUDA=1
+      run: DEBUG=2 CUDA=1 python -m pytest -rA test/test_tiny.py
     - name: Run Stable Diffusion
-      run: NV=1 python3 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd.txt
+      run: NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run SDXL
       run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run LLaMA
@@ -189,11 +197,11 @@ jobs:
     # - name: Run LLaMA 7B on 6 GPUs
     #   run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
     - name: Run LLaMA-3 8B BEAM
-      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
     - name: Run LLaMA-3 8B on 4 GPUs
-      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
     - name: Run LLaMA-3 8B on 6 GPUs
-      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
     - name: Run LLaMA-2 70B
       run: NV=1 RUN_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
     - name: Run Mixtral 8x7B
@@ -274,7 +282,7 @@ jobs:
     - name: Run 10 CIFAR training steps w BF16
       run: 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: NV=1 RUN_PROCESS_REPLAY=0 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: NV=1 RUN_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 NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
     - name: Run full CIFAR training steps w 6 GPUS
@@ -345,17 +353,20 @@ jobs:
     #    python3 -c "import torch; print(torch.__version__)"
     #    LD_PRELOAD="/opt/rocm/lib/libhsa-runtime64.so" HSA=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
-      run: AMD=1 IGNORE_BEAM_CACHE=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
+      run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
-      run: |
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+      run: AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
+    - name: Test AMD=1
+      run: DEBUG=2 AMD=1 python -m pytest -rA test/test_tiny.py
+    - name: Test HIP=1
+      run: DEBUG=2 HIP=1 python -m pytest -rA test/test_tiny.py
     # TODO: AMD compiler bug causes this to fail
     #- name: Fuzz Padded Tensor Core GEMM
     #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
     - name: Run Stable Diffusion
-      run: AMD=1 python3 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd.txt
+      run: AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run SDXL
       run: AMD=1 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run LLaMA 7B
@@ -369,11 +380,11 @@ jobs:
     # - name: Run LLaMA 7B on 6 GPUs
     #   run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
     - name: Run LLaMA-3 8B BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
     - name: Run LLaMA-3 8B on 4 GPUs
-      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
     - name: Run LLaMA-3 8B on 6 GPUs
-      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
     - name: Run LLaMA-2 70B
       run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
     - name: Run Mixtral 8x7B
@@ -452,7 +463,7 @@ jobs:
     - name: Run 10 CIFAR training steps w BF16
       run: 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: AMD=1 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: 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 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
@@ -497,10 +508,6 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
-    - name: openpilot compile 0.9.4
-      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py | tee openpilot_compile_0_9_4.txt
-    - name: openpilot compile 0.9.7
-      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_compile_0_9_7.txt
     - name: validate openpilot 0.9.7
       run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
     - name: benchmark openpilot 0.9.4

.github/workflows/test.yml

@@ -1,7 +1,7 @@
 name: Unit Tests
 env:
   # increment this when downloads substantially change to avoid the internet
-  DOWNLOAD_CACHE_VERSION: '7'
+  DOWNLOAD_CACHE_VERSION: '8'
   RUN_PROCESS_REPLAY: 1
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
   PYTHONPATH: .
@@ -37,10 +37,10 @@ jobs:
         python -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
         pip install mypy
         mypy -c "from tinygrad.tensor import Tensor; print(Tensor([1,2,3,4,5]))"
-    - name: Run beautiful_mnist without numpy
+    - name: Run beautiful_mnist with tinygrad only
       run: |
-        mkdir $HOME/test_no_numpy_dir
-        cd $HOME/test_no_numpy_dir
+        mkdir $GITHUB_WORKSPACE/test_dir
+        cd $GITHUB_WORKSPACE/test_dir
         python -m venv venv
         source venv/bin/activate
         pip install $GITHUB_WORKSPACE
@@ -229,7 +229,7 @@ jobs:
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
     - name: Run mypy
-      run: python -m mypy --strict-equality
+      run: python -m mypy --strict-equality --lineprecision-report . && cat lineprecision.txt
     - name: Test README
       run: awk '/```python/{flag=1;next}/```/{flag=0}flag' README.md > README.py &&  PYTHONPATH=. python README.py
     - name: Run unit tests
@@ -293,19 +293,15 @@ jobs:
           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
       - if: ${{ matrix.task == 'optimage' }}
-        name: Test openpilot model compile and size
+        name: Test openpilot model kernel count and gate usage
         run: |
-          PYTHONPATH="." DEBUG=2 ALLOWED_KERNEL_COUNT=208 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=1 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile2.py
-          python -c 'import os; assert os.path.getsize("/tmp/output.thneed") < 100_000_000'
-      - if: ${{ matrix.task == 'optimage' }}
-        name: Test openpilot model correctness (float32)
-        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile2.py
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2131 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
       - if: ${{ matrix.task == 'optimage' }}
         name: Test openpilot alt model correctness (float32)
-        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile2.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
+        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
       - if: ${{ matrix.task == 'optimage' }}
         name: Test openpilot fastvits model correctness (float32)
-        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile2.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+        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
       - if: ${{ matrix.task == 'onnx' }}
         name: Test ONNX (GPU)
         run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
@@ -348,46 +344,47 @@ jobs:
           export COMMIT_MESSAGE=$(git show -s --format=%B ${{ github.event.pull_request.head.sha }})
           cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
-  #testwebgpu:
-  #  name: WebGPU Tests
-  #  runs-on: macos-13
-  #  timeout-minutes: 20
-  #  steps:
-  #  - name: Checkout Code
-  #    uses: actions/checkout@v4
-  #  - name: Set up Python 3.11
-  #    uses: actions/setup-python@v5
-  #    with:
-  #      python-version: 3.11
-  #  - name: Cache python packages
-  #    uses: actions/cache@v4
-  #    with:
-  #      path: /Users/runner/Library/Python/3.11/lib/python/site-packages
-  #      key: webgpu-testing-user3-packages-${{ hashFiles('**/setup.py') }}
-  #  - name: Install Dependencies
-  #    run: pip install --user -e '.[webgpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
-  #  - name: Cache downloads
-  #    uses: actions/cache@v4
-  #    with:
-  #      path: ~/Library/Caches/tinygrad/downloads/
-  #      key: downloads-cache-webgpu-${{ env.DOWNLOAD_CACHE_VERSION }}
-  #  - name: Check Device.DEFAULT (WEBGPU) and print some source
-  #    run: |
-  #      WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
-  #      WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
-    #- name: Run webgpu pytest
-    #  run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m pytest -n=auto
-  #  - name: Run selected webgpu tests
-  #    run: |
-  #      WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m pytest -n=auto test/test_ops.py test/test_dtype.py \
-  #      test/test_jit.py test/test_symbolic_ops.py test/test_symbolic_jit.py test/test_linearizer.py \
-  #      test/test_linearizer_failures.py test/test_nn.py
-  #  - name: Build WEBGPU Efficientnet
-  #    run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python -m examples.compile_efficientnet
-  #  - name: Install Puppeteer
-  #    run: npm install puppeteer
-  #  - name: Run WEBGPU Efficientnet
-  #    run: node test/web/test_webgpu.js
+  testwebgpu:
+    name: WebGPU Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v4
+    - name: Set up Python 3.11
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.11
+    - name: Cache python packages
+      uses: actions/cache@v4
+      with:
+        path: ~/.local/lib/python3.11/site-packages
+        key: webgpu-testing-user3-packages-${{ hashFiles('**/setup.py') }}
+    - name: Install Dependencies
+      run: pip install --user -e '.[webgpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+    - name: Install dependencies for software-based vulkan
+      run: |
+        sudo apt update -y || true
+        sudo apt install -y libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers
+    - name: Cache downloads
+      uses: actions/cache@v4
+      with:
+        path: ~/.cache/tinygrad/downloads/
+        key: downloads-cache-webgpu-${{ env.DOWNLOAD_CACHE_VERSION }}
+    - name: Check Device.DEFAULT (WEBGPU) and print some source
+      run: |
+        WEBGPU=1 python -c "from tinygrad import Device; assert Device.DEFAULT == 'WEBGPU', Device.DEFAULT"
+        WEBGPU=1 DEBUG=4 FORWARD_ONLY=1 python3 test/test_ops.py TestOps.test_add
+    - name: Run selected webgpu tests
+      run: |
+          WEBGPU=1 WGPU_BACKEND_TYPE=Vulkan python3 -m pytest -n=auto test/ --ignore=test/external --ignore=test/models --ignore=test/unit \
+          --ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py --ignore=test/test_speed_v_torch.py --ignore=test/test_transcendental.py \
+          --ignore=test/test_fuzz_shape_ops.py --ignore=test/test_linearizer_failures.py --durations=20
+    - name: Run process replay tests
+      run: |
+        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
+        export COMMIT_MESSAGE=$(git show -s --format=%B ${{ github.event.pull_request.head.sha }})
+        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
   testmetal:
     name: Metal Tests
@@ -409,7 +406,7 @@ jobs:
         path: /Users/runner/Library/Python/3.11/lib/python/site-packages
         key: metal-m1-testing-user3-packages-${{ hashFiles('**/setup.py') }}
     - name: Install Dependencies
-      run: pip install --user -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+      run: pip install --user -e '.[webgpu,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
     - name: Cache downloads
       uses: actions/cache@v4
       with:
@@ -434,50 +431,26 @@ jobs:
     - name: Test Beam Search
       run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
     - name: Fuzz Test linearizer
-      run: PYTHONPATH="." METAL=1 FUZZ_ALL_ACTIONS=1 DEPTH=2 FUZZ_N=24 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+      run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     # - name: Fuzz Test models schedule
     # run: FUZZ_SCHEDULE=1 FUZZ_SCHEDULE_MAX_PATHS=5 python -m pytest test/models/test_train.py test/models/test_end2end.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
+    # WebGPU e2e tests
+    - name: Build WEBGPU Efficientnet
+      run: WEBGPU=1 WGPU_BACKEND_TYPE=Metal python3 -m examples.compile_efficientnet
+    - name: Clean npm cache
+      run: npm cache clean --force
+    - name: Install Puppeteer
+      run: npm install puppeteer
+    - name: Run WEBGPU Efficientnet
+      run: node test/web/test_webgpu.js
     - name: Run process replay tests
       run: |
         export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
         export COMMIT_MESSAGE=$(git show -s --format=%B ${{ github.event.pull_request.head.sha }})
         cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
-#  testwebgl:
-#    name: WebGL Tests
-#    runs-on: ubuntu-latest
-#    timeout-minutes: 20
-#
-#    steps:
-#    - name: Checkout Code
-#      uses: actions/checkout@v3
-#    - name: Set up Python 3.11
-#      uses: actions/setup-python@v4
-#      with:
-#        python-version: 3.11
-#    - name: Cache python packages
-#      uses: actions/cache@v4
-#      with:
-#        path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
-#        key: webgl-testing-packages-${{ hashFiles('**/setup.py') }}
-#    - name: Install Dependencies
-#      run: pip install -e '.[webgl,testing]' --extra-index-url https://download.pytorch.org/whl/cpu
-#    - name: Cache downloads
-#      uses: actions/cache@v4
-#      with:
-#        path: ~/Library/Caches/tinygrad/downloads/
-#        key: downloads-cache-webgl-${{ env.DOWNLOAD_CACHE_VERSION }}
-#    - name: Prepare
-#      run: |
-#        sudo apt-get -y install xvfb
-#        sudo /usr/bin/Xvfb :0 -screen 0 4096x4096x24+32 &
-#    - name: Run selected webgl tests
-#      run: WEBGL=1 python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_jit.py
-#    - name: Build WebGL Efficientnet
-#      run: WEBGL=1 python -m examples.compile_efficientnet
-
   tests:
     strategy:
       fail-fast: false
@@ -493,14 +466,14 @@ jobs:
         uses: actions/checkout@v4
         with:
           fetch-depth: 2 # NOTE: this fetches the HEAD commit of the PR
-      - name: Set up Python 3.11
+      - name: Set up Python 3.12
         uses: actions/setup-python@v5
         with:
-          python-version: 3.11
+          python-version: 3.12
       - name: Cache python packages
         uses: actions/cache@v4
         with:
-          path: ${{ env.Python3_ROOT_DIR }}/lib/python3.11/site-packages
+          path: ${{ env.Python3_ROOT_DIR }}/lib/python3.12/site-packages
           key: ${{ matrix.backend }}-packages-${{ hashFiles('**/setup.py') }}
       - name: Cache downloads
         uses: actions/cache@v4
@@ -549,7 +522,7 @@ jobs:
         if: matrix.backend == 'ptx' || matrix.backend == 'triton' || matrix.backend == 'nv'
         run: |
           cd ${{ github.workspace }}/gpuocelot/ocelot/build
-          sudo ninja install -d explain
+          sudo cp libgpuocelot.so /usr/lib/libgpuocelot.so
       - name: Install packages (amd)
         if: matrix.backend == 'amd'
         run: |
@@ -578,9 +551,6 @@ jobs:
       - name: Run pytest (not cuda or amd)
         if: matrix.backend!='ptx' && matrix.backend!='triton' && matrix.backend != 'amd' && matrix.backend != 'nv'
         run: python -m pytest -n=auto test/ --ignore=test/unit --durations=20
-      # - name: Run test_ops with FUZZ_UOPS=1
-      #   if: matrix.backend!='cuda' && matrix.backend!='ptx' && matrix.backend!='triton' && matrix.backend != 'amd' && matrix.backend != 'nv'
-      #   run: FUZZ_UOPS=1 python -m pytest -n=auto test/test_ops.py --durations=20
       - name: Run ONNX (only LLVM)
         if: matrix.backend == 'llvm'
         run: python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20

.gitignore

@@ -35,8 +35,9 @@ extra/datasets/COCO/
 extra/datasets/audio*
 extra/weights
 venv
-examples/**/net.*[js,json]
+examples/**/**/net*.*[js,json]
 examples/**/*.safetensors
+examples/webgpu/stable_diffusion/*.*[js,mjs]
 node_modules
 package.json
 package-lock.json

.pylintrc

@@ -54,7 +54,7 @@ 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
+disable=C,R,W0613,W0511,W0212,W0201,W0106,W0603,W0621,W0703,W1201,W1203,E1136,W1514,E1101,W0221,W0105,E0401,abstract-method
 # E1101 for function binding
 # W0221 for Function class
 # W0105 for comment strings

README.md

@@ -88,6 +88,7 @@ tinygrad already supports numerous accelerators, including:
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
 - [x] [NV](tinygrad/runtime/ops_nv.py)
 - [x] [QCOM](tinygrad/runtime/ops_qcom.py)
+- [x] [WEBGPU](tinygrad/runtime/ops_webgpu.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 
@@ -125,8 +126,8 @@ y = Tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()
 
-print(x.grad.numpy())  # dz/dx
-print(y.grad.numpy())  # dz/dy
+print(x.grad.tolist())  # dz/dx
+print(y.grad.tolist())  # dz/dy
 ```
 
 The same thing but in PyTorch:
@@ -138,8 +139,8 @@ y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
 z = y.matmul(x).sum()
 z.backward()
 
-print(x.grad.numpy())  # dz/dx
-print(y.grad.numpy())  # dz/dy
+print(x.grad.tolist())  # dz/dx
+print(y.grad.tolist())  # dz/dy
 ```
 
 ## Contributing

autogen_stubs.sh

@@ -162,6 +162,8 @@ generate_amd() {
   clang2py -k cdefstum \
     extra/hip_gpu_driver/sdma_registers.h \
     extra/hip_gpu_driver/nvd.h \
+    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
+    extra/hip_gpu_driver/soc21_enum.h \
     extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
     extra/hip_gpu_driver/gc_11_0_0_offset.h \
     extra/hip_gpu_driver/gc_10_3_0_offset.h \

docs/abstractions2.py

@@ -15,8 +15,8 @@ a = MallocAllocator.alloc(4)
 b = MallocAllocator.alloc(4)
 
 # load in some values (little endian)
-MallocAllocator.copyin(a, memoryview(bytearray([2,0,0,0])))
-MallocAllocator.copyin(b, memoryview(bytearray([3,0,0,0])))
+MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))
 
 # compile a program to a binary
 lib = ClangCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
@@ -28,7 +28,7 @@ fxn = ClangProgram("add", lib)
 fxn(out, a, b)
 
 # check the data out
-print(val := MallocAllocator.as_buffer(out).cast("I").tolist()[0])
+print(val := MallocAllocator._as_buffer(out).cast("I").tolist()[0])
 assert val == 5
 
 
@@ -76,13 +76,12 @@ assert out.as_buffer().cast('I')[0] == 5
 
 print("******** third, the LazyBuffer ***********")
 
-from tinygrad.engine.lazy import LazyBuffer
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule
 
 # allocate some values + load in values
-a = LazyBuffer.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
-b = LazyBuffer.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+a = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+b = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
 del a.srcs

docs/developer/developer.md

@@ -9,17 +9,13 @@ There is a good [bunch of tutorials](https://mesozoic-egg.github.io/tinygrad-not
 
 ## Frontend
 
-Everything in [Tensor](../tensor/index.md) is syntactic sugar around [function.py](function.md), where the forwards and backwards passes are implemented for the different functions. There's about 25 of them, implemented using about 20 basic ops. Those basic ops go on to construct a graph of:
+Everything in [Tensor](../tensor/index.md) is syntactic sugar around [function.py](function.md), where the forwards and backwards passes are implemented for the different functions. There's about 25 of them, implemented using about 20 basic ops. Those basic ops go on to construct a graph of [UOps](../developer/uop.md).
 
-::: tinygrad.engine.lazy.LazyBuffer
-    options:
-        show_source: false
-
-The `LazyBuffer` graph specifies the compute in terms of low level tinygrad ops. Not all LazyBuffers will actually become realized. There's two types of LazyBuffers, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
+The `UOp` graph specifies the compute in terms of low level tinygrad ops. Not all UOps will actually become realized. There's two types of UOps, base and view. base contains compute into a contiguous buffer, and view is a view (specified by a ShapeTracker). Inputs to a base can be either base or view, inputs to a view can only be a single base.
 
 ## Scheduling
 
-The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of LazyBuffers into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
+The [scheduler](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/engine/schedule.py) converts the graph of UOps into a list of `ScheduleItem`. One `ScheduleItem` is one kernel on the GPU, and the scheduler is responsible for breaking the large compute graph into subgraphs that can fit in a kernel. `ast` specifies what compute to run, and `bufs` specifies what buffers to run it on.
 
 ::: tinygrad.engine.schedule.ScheduleItem
 

docs/developer/hcq.md

@@ -6,57 +6,29 @@ The main aspect of HCQ-compatible runtimes is how they interact with devices. In
 
 ### Command Queues
 
-To interact with devices, there are 2 types of queues: `HWComputeQueue` and `HWCopyQueue`. Commands which are defined in a base `HWCommandQueue` class should be supported by both queues. These methods are timestamp and synchronization methods like [signal](#tinygrad.runtime.support.hcq.HWCommandQueue.signal) and [wait](#tinygrad.runtime.support.hcq.HWCommandQueue.wait).
+To interact with devices you create a `HWQueue`. Some methods are required, like timestamp and synchronization methods like [signal](#tinygrad.runtime.support.hcq.HWQueue.signal) and [wait](#tinygrad.runtime.support.hcq.HWQueue.wait), while others are dependent on it being a compute or copy queue.
 
 For example, the following Python code enqueues a wait, execute, and signal command on the HCQ-compatible device:
 ```python
-HWComputeQueue().wait(signal_to_wait, value_to_wait) \
-                .exec(program, args_state, global_dims, local_dims) \
-                .signal(signal_to_fire, value_to_fire) \
-                .submit(your_device)
+HWQueue().wait(signal_to_wait, value_to_wait) \
+         .exec(program, args_state, global_dims, local_dims) \
+         .signal(signal_to_fire, value_to_fire) \
+         .submit(your_device)
 ```
 
-Each runtime should implement the required functions that are defined in the `HWCommandQueue`, `HWComputeQueue`, and `HWCopyQueue` classes.
+Each runtime should implement the required functions that are defined in the `HWQueue` classes.
 
-::: tinygrad.runtime.support.hcq.HWCommandQueue
+::: tinygrad.runtime.support.hcq.HWQueue
     options:
         members: [
             "signal",
             "wait",
             "timestamp",
-            "update_signal",
-            "update_wait",
             "bind",
             "submit",
-        ]
-        show_source: false
-
-::: tinygrad.runtime.support.hcq.HWComputeQueue
-    options:
-        members: [
             "memory_barrier",
             "exec",
-            "update_exec",
-        ]
-        show_source: false
-
-::: tinygrad.runtime.support.hcq.HWCopyQueue
-    options:
-        members: [
             "copy",
-            "update_copy",
-        ]
-        show_source: false
-
-#### Implementing custom commands
-
-To implement custom commands in the queue, use the @hcq_command decorator for your command implementations.
-
-::: tinygrad.runtime.support.hcq.hcq_command
-    options:
-        members: [
-            "copy",
-            "update_copy",
         ]
         show_source: false
 
@@ -82,9 +54,9 @@ The following Python code demonstrates the usage of signals:
 ```python
 signal = your_device.signal_t()
 
-HWComputeQueue().timestamp(signal) \
-                .signal(signal, value_to_fire) \
-                .submit(your_device)
+HWQueue().timestamp(signal) \
+         .signal(signal, value_to_fire) \
+         .submit(your_device)
 
 signal.wait(value_to_fire)
 signaled_value = signal.value # should be the same as `value_to_fire`
@@ -134,17 +106,17 @@ Backends must adhere to the `HCQBuffer` protocol when returning allocation resul
         members: true
         show_source: false
 
-**Lifetime**: The `HCQArgsState` is passed to `HWComputeQueue.exec` and is guaranteed not to be freed until `HWComputeQueue.submit` for the same queue is called.
+**Lifetime**: The `HCQArgsState` is passed to `HWQueue.exec` and is guaranteed not to be freed until `HWQueue.submit` for the same queue is called.
 
 ### Synchronization
 
 HCQ-compatible devices use a global timeline signal for synchronizing all operations. This mechanism ensures proper ordering and completion of tasks across the device. By convention, `self.timeline_value` points to the next value to signal. So, to wait for all previous operations on the device to complete, wait for `self.timeline_value - 1` value. The following Python code demonstrates the typical usage of signals to synchronize execution to other operations on the device:
 
 ```python
-HWComputeQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
-                .exec(...)
-                .signal(your_device.timeline_signal, your_device.timeline_value) \
-                .submit(your_device)
+HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
+         .exec(...)
+         .signal(your_device.timeline_signal, your_device.timeline_value) \
+         .submit(your_device)
 your_device.timeline_value += 1
 
 # Optionally wait for execution
@@ -153,5 +125,5 @@ your_device.timeline_signal.wait(your_device.timeline_value - 1)
 
 ## HCQGraph
 
-[HCQGraph](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/graph/hcq.py) is a core feature that implements `GraphRunner` for HCQ-compatible devices. `HCQGraph` builds a static `HWComputeQueue` and `HWCopyQueue` for all operations per device. To optimize enqueue time, only the necessary parts of the queues are updated for each run using the update APIs of the queues, avoiding a complete rebuild.
+[HCQGraph](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/graph/hcq.py) is a core feature that implements `GraphRunner` for HCQ-compatible devices. `HCQGraph` builds static `HWQueue` for all operations per device. To optimize enqueue time, only the necessary parts of the queues are updated for each run using the symbolic variables, avoiding a complete rebuild.
 Optionally, queues can implement a `bind` API, which allows further optimization by eliminating the need to copy the queues into the device ring.

docs/tensor/elementwise.md

@@ -13,6 +13,9 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sin
 ::: tinygrad.Tensor.cos
 ::: tinygrad.Tensor.tan
+::: tinygrad.Tensor.asin
+::: tinygrad.Tensor.acos
+::: tinygrad.Tensor.atan
 ::: tinygrad.Tensor.trunc
 ::: tinygrad.Tensor.ceil
 ::: tinygrad.Tensor.floor
@@ -34,6 +37,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.hardsigmoid
 ::: tinygrad.Tensor.elu
 ::: tinygrad.Tensor.celu
+::: tinygrad.Tensor.selu
 ::: tinygrad.Tensor.swish
 ::: tinygrad.Tensor.silu
 ::: tinygrad.Tensor.relu6

docs/tensor/ops.md

@@ -27,9 +27,11 @@
 ::: tinygrad.Tensor.matmul
 ::: tinygrad.Tensor.einsum
 ::: tinygrad.Tensor.cumsum
+::: tinygrad.Tensor.cummax
 ::: tinygrad.Tensor.triu
 ::: tinygrad.Tensor.tril
 ::: tinygrad.Tensor.interpolate
+::: tinygrad.Tensor.scatter
 
 ## Neural Network (functional)
 

examples/compile_efficientnet.py

@@ -1,7 +1,7 @@
 from pathlib import Path
 from extra.models.efficientnet import EfficientNet
 from tinygrad.tensor import Tensor
-from tinygrad.nn.state import safe_save
+from tinygrad.nn.state import get_state_dict, safe_save, safe_load, load_state_dict
 from extra.export_model import export_model
 from tinygrad.helpers import getenv, fetch
 import ast
@@ -9,12 +9,16 @@ import ast
 if __name__ == "__main__":
   model = EfficientNet(0)
   model.load_from_pretrained()
-  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else "webgl" if getenv("WEBGL", "") != "" else ""
-  prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
   dirname = Path(__file__).parent
+  # exporting a model that's loaded from safetensors doesn't work without loading in from safetensors first
+  # loading the state dict from a safetensor file changes the generated kernels
+  if getenv("WEBGPU"):
+    safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
+    load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
+  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
+  prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
   if getenv("CLANG", "") == "":
-    safe_save(state, (dirname / "net.safetensors").as_posix())
-    ext = "js" if getenv("WEBGPU", "") != "" or getenv("WEBGL", "") != "" else "json"
+    ext = "js" if getenv("WEBGPU", "") != "" else "json"
     with open(dirname / f"net.{ext}", "w") as text_file:
       text_file.write(prg)
   else:

examples/gpt2.py

@@ -1,11 +1,8 @@
 #!/usr/bin/env python3
-import os
+import os, argparse
 from typing import Optional, Union
-import argparse
-import numpy as np
 import tiktoken
-from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable
-from tinygrad.dtype import dtypes
+from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
 from tinygrad.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange
 from tinygrad.nn import Embedding, Linear, LayerNorm
@@ -192,7 +189,7 @@ class GPT2:
           tokens = Variable("tokens", 0, VOCAB_SIZE).bind(toks[0][start_pos])
         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).bind(start_pos), temperature).numpy().tolist()
+        tok = self.model(tokens, Variable("start_pos", 1 if start_pos else 0, MAX_CONTEXT).bind(start_pos), temperature).tolist()
       start_pos = len(toks[0])
       for i,t in enumerate(tok): toks[i].append(t)
     return [self.tokenizer.decode(x) for x in toks]
@@ -218,7 +215,6 @@ if __name__ == "__main__":
 
   if args.seed is not None:
     Tensor.manual_seed(args.seed)
-    np.random.seed(args.seed)
 
   print(f"using {args.model_size}")
   gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

examples/llama.py

@@ -6,8 +6,6 @@
 from pathlib import Path
 from typing import List, Optional
 import argparse, json
-import numpy as np
-np.set_printoptions(linewidth=200)
 from tinygrad import Tensor, Device, GlobalCounters, nn
 from tinygrad.helpers import Context, Timing, Profiling, DEBUG, JIT, getenv, colored
 from tinygrad.nn.state import safe_load, torch_load, load_state_dict, get_parameters
@@ -250,6 +248,8 @@ class LLaMa:
     self.tokenizer = tokenizer
 
   def greedy_until(self, prompt:str, until, max_length, temperature):
+    # only used in old eval script
+    import numpy as np
     toks = [self.tokenizer.bos_id()] + self.tokenizer.encode(prompt)
     start_pos = 0
     for i in range(max_length):
@@ -498,14 +498,16 @@ After you are done speaking, output [EOS]. You are not Chad.
     if not chatbot: break
 
   # validate output!
-  if args.temperature == 0 and args.count == 10 and args.prompt == "Hello." and not args.quantize:
+  if args.temperature == 0 and args.count == 10 and args.prompt == "Hello.":
     text = llama.tokenizer.decode(toks)
-    key = (args.gen, args.size)
+    key = (args.gen, args.size, args.quantize)
     expected = {
-      ("1", "7B"): "Hello. I'm a 20 year old male",
-      ("2", "7B"): "Hello. I'm a 20 year old girl",
-      ("2", "70B"): "Hello. I am a 20 year old female.",
-      ("3", "8B"): "Hello. I am a 20 year old female. I",
+      ("1", "7B", None): "Hello. I'm a 20 year old male",
+      ("1", "7B", "int8"): "Hello. I'm a 20 year old male",
+      ("1", "7B", "nf4"): "Hello. I'm a 20 year old male",
+      ("2", "7B", None): "Hello. I'm a 20 year old girl",
+      ("2", "70B", None): "Hello. I am a 20 year old female.",
+      ("3", "8B", None): "Hello. I am a 20 year old female. I",
     }
     try:
       assert text == expected[key], f"invalid output: `{colored(text, 'red')}` != `{expected[key]}`"

examples/llama3.py

@@ -166,7 +166,8 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, device=N
 
   with Context(BEAM=0):
     # quantize
-    if quantize is not None:
+    if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
+    elif quantize is not None:
       weights = linear.quantize(weights, device)
       for _,v in weights.items(): v.realize()
 
@@ -219,9 +220,9 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--download_model", action="store_true", help="Download a model")
   parser.add_argument("--model", type=Path, help="Model path")
-  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="8B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
   parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
-  parser.add_argument("--quantize", choices=["int8", "nf4"], help="Quantization method")
+  parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
   parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
   parser.add_argument("--host", type=str, default="0.0.0.0", help="Web server bind address")
   parser.add_argument("--port", type=int, default=7776, help="Web server port")
@@ -233,8 +234,8 @@ if __name__ == "__main__":
   parser.add_argument("--profile", action="store_true", help="Output profile data")
   args = parser.parse_args()
 
-  assert (args.model and not args.download_model) or (not args.model and args.download_model), "either download or provide model"
-  if args.download_model:
+  # download_model is the default without a model passed in
+  if args.download_model or not args.model:
     if args.size == "1B":
       fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
       args.model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf", "Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
@@ -245,6 +246,12 @@ if __name__ == "__main__":
       fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00003-of-00004.safetensors", "model-00003-of-00004.safetensors", subdir="llama3-8b-sfr")
       fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00004-of-00004.safetensors", "model-00004-of-00004.safetensors", subdir="llama3-8b-sfr")
       args.model = fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/raw/main/model.safetensors.index.json", "model.safetensors.index.json", subdir="llama3-8b-sfr")
+    elif args.size == "70B":
+      subdir = "Llama-3.1-Nemotron-70B-Instruct-HF"
+      args.model = fetch("https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
+      fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir=subdir)
+      for i in range(30):
+        fetch(f"https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model-{i+1:05d}-of-00030.safetensors?download=true", f"model-{i+1:05d}-of-00030.safetensors", subdir=subdir)
 
   assert args.model is not None, "please provide --model option"
 
@@ -281,8 +288,9 @@ if __name__ == "__main__":
       for key, value in cors_headers.items(): response.set_header(key, value)
 
     @app.route("/<filename>")
-    def server_static(filename):
-      return static_file(filename, root=(Path(__file__).parent / "tinychat").as_posix())
+    def server_static(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat").as_posix())
+    @app.route("/assets/<filename:path>")
+    def server_assets(filename): return static_file(filename, root=(Path(__file__).parent / "tinychat" / "assets").as_posix())
     @app.route("/")
     def index():
       return static_file("index.html", root=(Path(__file__).parent / "tinychat").as_posix())

examples/mlperf/dataloader.py

@@ -67,11 +67,11 @@ def loader_process(q_in, q_out, X:Tensor, seed):
 
       # broken out
       #img_tensor = Tensor(img.tobytes(), device='CPU')
-      #storage_tensor = X[idx].contiguous().realize().lazydata.realized
+      #storage_tensor = X[idx].contiguous().realize().lazydata.base.realized
       #storage_tensor._copyin(img_tensor.numpy())
 
       # faster
-      X[idx].contiguous().realize().lazydata.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
+      X[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
 
       # ideal
       #X[idx].assign(img.tobytes())   # NOTE: this is slow!
@@ -267,8 +267,8 @@ def load_unet3d_data(preprocessed_dataset_dir, seed, queue_in, queue_out, X:Tens
       x = random_brightness_augmentation(x)
       x = gaussian_noise(x)
 
-    X[idx].contiguous().realize().lazydata.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
-    Y[idx].contiguous().realize().lazydata.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
+    X[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = x.tobytes()
+    Y[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = y.tobytes()
 
     queue_out.put(idx)
   queue_out.put(None)

examples/mlperf/model_train.py

@@ -3,12 +3,14 @@ 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, tqdm
+from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
 from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
 
 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
+# TODO: fix benchmark logging and use tinygrad tqdm
+from tqdm import tqdm
 
 def train_resnet():
   from extra.models import resnet

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_beam.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="."
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/dev_run.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="."
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_green/run_and_time.sh

@@ -5,7 +5,7 @@ export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_green"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
 

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_beam.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="."
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/dev_run.sh

@@ -4,7 +4,7 @@ export PYTHONPATH="."
 export MODEL="resnet"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 

examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh

@@ -5,7 +5,7 @@ export MODEL="resnet"
 export SUBMISSION_PLATFORM="tinybox_red"
 export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
 
-export LAZYCACHE=0 RESET_STEP=0
+export RESET_STEP=0
 
 export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
 

examples/openpilot/compile2.py

@@ -1,211 +0,0 @@
-#!/usr/bin/env python3
-import os, sys, io, pathlib, json, struct
-import numpy as np
-sys.path.insert(0, str(pathlib.Path(__file__).parents[1]))
-
-if "FLOAT16" not in os.environ: os.environ["FLOAT16"] = "1"
-if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
-if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
-
-OPENPILOT_MODEL = "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
-
-import onnx
-from typing import Tuple, List, Optional, Dict, cast
-from extra.onnx import get_run_onnx
-from tinygrad import Tensor, Device, GlobalCounters, dtypes
-from tinygrad.dtype import ImageDType
-from tinygrad.device import Buffer
-from tinygrad.helpers import partition, Context, fetch, getenv, DEBUG, tqdm
-from tinygrad.engine.realize import run_schedule, lower_schedule, ExecItem, CompiledRunner
-from tinygrad.engine.memory import memory_planner
-from tinygrad.engine.schedule import ScheduleItem, create_schedule
-from tinygrad.ops import Ops
-from tinygrad.tensor import _to_np_dtype
-Device.DEFAULT = "GPU"
-
-def get_schedule(onnx_data) -> Tuple[List[ScheduleItem], List[ScheduleItem]]:
-  Tensor.no_grad = True
-  Tensor.training = False
-
-  # load the model
-  onnx_model = onnx.load(io.BytesIO(onnx_data))
-  run_onnx = get_run_onnx(onnx_model)
-  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-
-  # run the model
-  inputs = {k:Tensor.empty(*shp) for k,shp in input_shapes.items()}
-  ret: Tensor = next(iter(run_onnx(inputs).values())).cast(dtypes.float32).contiguous()
-  schedule = create_schedule([ret.lazydata])
-
-  # filter schedule that don't depend on the inputs
-  input_lb = [x.lazydata.base.buffer for x in inputs.values()]
-  depends = set(input_lb)
-  for si in schedule:
-    if any(b in depends for b in si.inputs):
-      for out in si.outputs: depends.add(out)
-
-  # run all kernels that don't depend on the inputs
-  # NOTE: there's two extra kernels due to fusions that now happen since the weights aren't realized
-  schedule, schedule_independent = partition(schedule, lambda si: any(out in depends for out in si.outputs))
-  print(f"{len(schedule)} schedule items depend on the input, {len(schedule_independent)} don't")
-
-  # confirm no non-sink metaop in the (non independent) schedule except for the ones that load the input buffers
-  assert all(si.ast.op is Ops.SINK or out in input_lb for si in schedule for out in si.outputs), "has non SINK ops, can't compile to Thneed"
-  return schedule, schedule_independent, inputs
-
-def test_vs_onnx(onnx_data, eis:Optional[List[ExecItem]], inputs:Dict[str, Tensor]):
-  import onnx
-  #import pyopencl as cl
-  #from extra.thneed import Thneed
-  import numpy as np
-  onnx_model = onnx.load(io.BytesIO(onnx_data))
-
-  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
-  Tensor.manual_seed(1337)
-  new_inputs = {k:Tensor.randn(*shp, requires_grad=False)*8 for k,shp in input_shapes.items()}
-  new_np_inputs = {k:v.realize().numpy() for k,v in new_inputs.items()}
-
-  if getenv("ORT"):
-    # test with onnxruntime
-    import onnxruntime as ort
-    onnx_session = ort.InferenceSession(onnx_data)
-    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_np_inputs.items()})
-    new_torch_out = onnx_output[0]
-    print("got ort outputs")
-  else:
-    # test with torch
-    from test.models.test_onnx import run_onnx_torch
-    new_torch_out = run_onnx_torch(onnx_model, new_np_inputs).numpy()
-    print("got torch outputs")
-
-  # if you don't have a schedule
-  if eis is None:
-    run_onnx = get_run_onnx(onnx_model)
-    new_tinygrad_out = next(iter(run_onnx(new_inputs).values())).cast(dtypes.float32).numpy()
-    np.testing.assert_allclose(new_torch_out, new_tinygrad_out, atol=1e-4, rtol=1e-2)
-    print("classic self-test passed!")
-    return
-
-  # set inputs
-  for k,v in inputs.items(): v.lazydata.base.realized.copyin(new_np_inputs[k].data)
-
-  # run code (all buffers have been allocated)
-  GlobalCounters.reset()
-  output = eis[-1].bufs[0]
-  for ei in eis: ei.run()
-
-  new_tinygrad_out = np.frombuffer(output.as_buffer(), dtype=_to_np_dtype(output.dtype))
-  np.testing.assert_allclose(new_torch_out.reshape(new_tinygrad_out.shape), new_tinygrad_out, atol=1e-4, rtol=1e-2)
-  print("semi-thneed self-test passed!")
-
-if __name__ == "__main__":
-  onnx_data = fetch(sys.argv[1] if len(sys.argv) > 1 else OPENPILOT_MODEL).read_bytes()
-
-  # quick test for ONNX issues
-  #thneed_test_onnx(onnx_data, None)
-  #exit(0)
-
-  schedule, schedule_independent, inputs = get_schedule(onnx_data)
-  schedule, schedule_input = partition(schedule, lambda x: x.ast.op is Ops.SINK)
-  print(f"{len(schedule_input)} inputs")
-
-  run_schedule(schedule_independent)
-  run_schedule(schedule_input)
-  with Context(DEBUG=max(DEBUG.value, 2), BEAM=getenv("LATEBEAM")):
-    schedule = memory_planner(schedule)
-    for si in schedule:
-      for b in si.outputs:
-        assert not b.is_allocated(), "output should not be allocated"
-    image_count = sum(isinstance(out.dtype, ImageDType) for si in schedule for out in si.outputs)
-    print(f"**** compiling real kernels {image_count}/{len(schedule)} images ****")
-    eis = list(tqdm(lower_schedule(schedule), total=len(schedule)))
-
-  print("kernel count:", len(eis))
-  assert len(eis) <= getenv("ALLOWED_KERNEL_COUNT", 0) or getenv("ALLOWED_KERNEL_COUNT", 0) == 0, "too many kernels!"
-
-  # new simple thneed
-  def to_ref(b:Buffer): return struct.pack("Q", id(b)).decode("latin_1")
-
-  seen_buffers = set()
-  input_buffers = [x.lazydata.buffer for x in inputs.values()]
-  jdat = {"binaries": [], "programs": {}, "kernels": [], "objects": []}
-  jdat["inputs"] = {k:to_ref(v.lazydata.buffer) for k,v in inputs.items()}
-  jdat["outputs"] = [to_ref(eis[-1].bufs[0])]
-  weights = []
-  for i,ei in enumerate(eis):
-    #print("***", i)
-    for b in ei.bufs:
-      needs_load = b.is_allocated() and b not in input_buffers
-      #print(b, needs_load)
-      if b in seen_buffers: continue
-      seen_buffers.add(b)
-      if isinstance(b.dtype, ImageDType):
-        base_dtype = dtypes.float16 if b.dtype.fmt == 'e' else dtypes.float32
-        row_pitch = (b.dtype.shape[0]*4*base_dtype.itemsize + 63)//64 * 64
-        size = row_pitch * b.dtype.shape[1]
-        jdat['objects'].append({
-          "id": to_ref(b), "needs_load": needs_load, "size": size, "arg_type": "image2d_t",
-          "width": b.dtype.shape[0], "height": b.dtype.shape[1], "row_pitch": row_pitch, "float32": b.dtype.base == dtypes.float32,
-        })
-        if needs_load:
-          t = Tensor.empty(b.dtype.shape, dtype=b.dtype)
-          t.lazydata.buffer = b
-          data = t.cast(dtypes.float32).pad(((0, row_pitch//(4*base_dtype.itemsize)-b.dtype.shape[0]), (0,0), (0,0))).contiguous().numpy()
-          # NOTE: this cast must be done in numpy for platforms that don't support half
-          if base_dtype == dtypes.float16: data = data.astype(np.float16)
-          weights.append(data.tobytes())
-          assert len(weights[-1]) == size, "wrong size buffer"
-      else:
-        jdat['objects'].append({
-          "id": to_ref(b), "arg_type": b.dtype.name + "*", "needs_load": needs_load, "size": b.nbytes,
-        })
-        if needs_load:
-          weights.append(b.as_buffer())
-          assert len(weights[-1]) == b.nbytes, "wrong size buffer"
-
-  saved_binaries = set()
-  binaries = []
-  gated_read_image_count = 0
-  GlobalCounters.reset()
-  with Context(DEBUG=max(DEBUG.value, 2)):
-    for ei in eis:
-      prg = cast(CompiledRunner, ei.prg)
-      assert len(prg.p.vars) == 0
-      if prg.p.function_name not in saved_binaries:
-        jdat['binaries'].append({"name":prg.p.function_name, "length":len(prg.lib)})
-        binaries.append(prg.lib)
-        saved_binaries.add(prg.p.function_name)
-      gated_read_image_count += prg.p.src.count("?read_image")
-      ei.run()
-      jdat['kernels'].append({
-        "name": prg.p.function_name,
-        "work_dim": len(prg.p.global_size),
-        "global_work_size": prg.p.global_size,
-        "local_work_size": prg.p.local_size,
-        "num_args": len(ei.bufs),
-        "args": [to_ref(b) for b in ei.bufs],
-        "arg_size": [8]*len(ei.bufs),
-      })
-
-  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
-    assert gated_read_image_count <= allowed_gated_read_image, \
-      f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
-
-  output_fn = sys.argv[2] if len(sys.argv) >= 3 else "/tmp/output.thneed"
-  print(f"saving thneed to {output_fn} with {len(weights)} buffers and {len(binaries)} binaries")
-  with open(output_fn, "wb") as f:
-    j = json.dumps(jdat, ensure_ascii=False).encode('latin_1')
-    f.write(struct.pack("I", len(j)))
-    f.write(j)
-    for w in weights: f.write(w)
-    for b in binaries: f.write(b)
-    print("saved", f.tell(), "bytes")
-
-  FLOAT16 = getenv("FLOAT16", 0)
-  if FLOAT16 == 0:
-    try:
-      test_vs_onnx(onnx_data, eis, inputs)
-    except ModuleNotFoundError as e:
-      print(f"TEST NOT HAPPENING {e}")
-
-

examples/openpilot/compile3.py

@@ -5,9 +5,10 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
-from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
 from tinygrad.helpers import DEBUG, getenv
 from tinygrad.tensor import _from_np_dtype
+from tinygrad.engine.realize import CompiledRunner
 
 import onnx
 from onnx.helper import tensor_dtype_to_np_dtype
@@ -16,64 +17,120 @@ from extra.onnx import get_run_onnx   # TODO: port to main tinygrad
 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 = "/tmp/openpilot.pkl"
 
-def compile():
+def compile(onnx_file):
+  onnx_model = onnx.load(onnx_file)
   Tensor.no_grad = True
   Tensor.training = False
 
-  onnx_bytes = fetch(OPENPILOT_MODEL)
-  onnx_model = onnx.load(onnx_bytes)
   run_onnx = get_run_onnx(onnx_model)
   print("loaded model")
 
   input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
   input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
+  if getenv("FLOAT16", 0) == 0: input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
   Tensor.manual_seed(100)
   new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
   print("created tensors")
 
-  run_onnx_jit = TinyJit(lambda **kwargs: run_onnx(kwargs), prune=True)
+  run_onnx_jit = TinyJit(lambda **kwargs:
+                         next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())).cast('float32'), prune=True)
   for i in range(3):
     GlobalCounters.reset()
     print(f"run {i}")
+    inputs = {**{k:v.clone() for k,v in new_inputs.items() if 'img' in k},
+              **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
     with Context(DEBUG=max(DEBUG.value, 2 if i == 2 else 1)):
-      ret = next(iter(run_onnx_jit(**new_inputs).values())).cast('float32').numpy()
+      ret = run_onnx_jit(**inputs).numpy()
     # copy i == 1 so use of JITBEAM is okay
     if i == 1: test_val = np.copy(ret)
   print(f"captured {len(run_onnx_jit.captured.jit_cache)} kernels")
-  np.testing.assert_equal(test_val, ret)
+  np.testing.assert_equal(test_val, ret, "JIT run failed")
   print("jit run validated")
 
+  # checks from compile2
+  kernel_count = 0
+  read_image_count = 0
+  gated_read_image_count = 0
+  for ei in run_onnx_jit.captured.jit_cache:
+    if isinstance(ei.prg, CompiledRunner):
+      kernel_count += 1
+      read_image_count += ei.prg.p.src.count("read_image")
+      gated_read_image_count += ei.prg.p.src.count("?read_image")
+  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"too many kernels! {kernel_count=}, {allowed_kernel_count=}"
+  if (allowed_read_image:=getenv("ALLOWED_READ_IMAGE", -1)) != -1:
+    assert read_image_count == allowed_read_image, f"different read_image! {read_image_count=}, {allowed_read_image=}"
+  if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
+    assert gated_read_image_count <= allowed_gated_read_image, f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+
   with open(OUTPUT, "wb") as f:
     pickle.dump(run_onnx_jit, f)
-  mdl_sz = os.path.getsize(onnx_bytes)
+  mdl_sz = os.path.getsize(onnx_file)
   pkl_sz = os.path.getsize(OUTPUT)
   print(f"mdl size is {mdl_sz/1e6:.2f}M")
   print(f"pkl size is {pkl_sz/1e6:.2f}M")
   print("**** compile done ****")
   return test_val
 
-def test(test_val=None):
-  with open(OUTPUT, "rb") as f:
-    run = pickle.load(f)
-  Tensor.manual_seed(100)
-  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
-                sorted(zip(run.captured.expected_names, run.captured.expected_st_vars_dtype_device))}
+def test_vs_compile(run, new_inputs, test_val=None):
   new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+
+  # create fake "from_blob" tensors for the inputs, and wrapped NPY tensors for the numpy inputs (these have the same underlying memory)
+  inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
+            **{k:Tensor(v, device="NPY").realize() for k,v in new_inputs_numpy.items() if 'img' not in k}}
+
+  # run 20 times
   for _ in range(20):
     st = time.perf_counter()
-    # Need to cast non-image inputs from numpy, this is only realistic way to run it
-    inputs = {**{k:v for k,v in new_inputs.items() if 'img' in k},
-              **{k:Tensor(v) for k,v in new_inputs_numpy.items() if 'img' not in k}}
     out = run(**inputs)
     mt = time.perf_counter()
-    val = out['outputs'].numpy()
+    val = out.numpy()
     et = time.perf_counter()
     print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
   print(out, val.shape, val.dtype)
   if test_val is not None: np.testing.assert_equal(test_val, val)
   print("**** test done ****")
 
-if __name__ == "__main__":
-  test_val = compile() if not getenv("RUN") else None
-  test(test_val)
+  # test that changing the numpy changes the model outputs
+  for v in new_inputs_numpy.values(): v *= 2
+  out = run(**inputs)
+  changed_val = out.numpy()
+  np.testing.assert_raises(AssertionError, np.testing.assert_array_equal, val, changed_val)
+  return val
 
+def test_vs_onnx(new_inputs, test_val, onnx_file):
+  new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
+  onnx_model = onnx.load(onnx_file)
+
+  if getenv("ORT"):
+    # test with onnxruntime
+    import onnxruntime as ort
+    onnx_session = ort.InferenceSession(onnx_file)
+    onnx_output = onnx_session.run([onnx_model.graph.output[0].name], {k:v.astype(np.float16) for k,v in new_inputs_numpy.items()})
+    new_torch_out = onnx_output[0]
+    print("got ort outputs")
+  else:
+    # test with torch
+    from test.models.test_onnx import run_onnx_torch
+    # NOTE: we have to correct the order here
+    new_torch_out = run_onnx_torch(onnx_model, {k.name:new_inputs_numpy[k.name] for k in onnx_model.graph.input}).numpy()
+    print("got torch outputs")
+
+  np.testing.assert_allclose(new_torch_out.reshape(test_val.shape), test_val, atol=1e-4, rtol=1e-2)
+  print("test vs onnx passed")
+
+if __name__ == "__main__":
+  onnx_file = fetch(OPENPILOT_MODEL)
+  test_val = compile(onnx_file) if not getenv("RUN") else None
+
+  with open(OUTPUT, "rb") as f: pickle_loaded = pickle.load(f)
+
+  # same randomness as compile
+  Tensor.manual_seed(100)
+  new_inputs = {nm:Tensor.randn(*st.shape, dtype=dtype).mul(8).realize() for nm, (st, _, dtype, _) in
+                sorted(zip(pickle_loaded.captured.expected_names, pickle_loaded.captured.expected_st_vars_dtype_device))}
+
+  test_val = test_vs_compile(pickle_loaded, new_inputs, test_val)
+  if not getenv("FLOAT16"): test_vs_onnx(new_inputs, test_val, onnx_file)

examples/openpilot/go.sh

@@ -1,2 +0,0 @@
-#!/bin/bash
-NOLOCALS=1 FLOAT16=1 DEBUGCL=1 IMAGE=2 GPU=1 python3 examples/openpilot/compile2.py

examples/qwq.py

@@ -0,0 +1,117 @@
+import argparse
+import os
+import sys
+
+from transformers import AutoTokenizer
+from pathlib import Path
+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.helpers import fetch, colored, GlobalCounters, Timing, DEBUG
+from tinygrad.nn.state import load_state_dict, get_parameters
+
+MODELS = {
+  "32B": {
+    "model_params": {"dim": 5120, "n_heads": 40, "n_kv_heads": 8, "n_layers": 64, "norm_eps": 1e-5, "rope_theta": 1000000, "vocab_size": 152064, "hidden_dim": 27648},
+    "total_num_weights": 17,
+    "tokenizer": "Qwen/QwQ-32B-Preview"
+  }
+}
+
+def download_weights(total_num_weights:int) -> Path:
+  model = fetch("https://huggingface.co/Qwen/QwQ-32B-Preview/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=(subdir:="qwq_32b_preview"))
+
+  for i in range(1, total_num_weights + 1):
+    filename = f"model-{i:05d}-of-{total_num_weights:05d}.safetensors"
+    fetch(f"https://huggingface.co/Qwen/QwQ-32B-Preview/resolve/main/{filename}?download=true", filename, subdir=subdir)
+
+  return Path(os.path.dirname(model))
+
+def load_model(model_path:Path, model_params:Dict[str, Union[int, float]]) -> Transformer:
+  # build model
+  model = Transformer(**model_params, linear=nn.Linear)
+
+  # update layers to add bias
+  updated_layers = []
+  for layer in model.layers:
+    head_dim = model_params["dim"] // model_params["n_heads"]
+    layer.attention.wq = nn.Linear(model_params["dim"], model_params["n_heads"] * head_dim, bias=True)
+    layer.attention.wk = nn.Linear(model_params["dim"], model_params["n_kv_heads"] * head_dim, bias=True)
+    layer.attention.wv = nn.Linear(model_params["dim"], model_params["n_kv_heads"] * head_dim, bias=True)
+    updated_layers.append(layer)
+  model.layers = updated_layers
+
+  # load weights
+  weights = fix_bf16(convert_from_huggingface(load(str(model_path / "model.safetensors.index.json")), model, model_params["n_heads"], model_params["n_kv_heads"], permute_layers=False))
+
+  # replace weights in model
+  load_state_dict(model, weights, strict=False, consume=True)
+  return model
+
+
+if __name__ == "__main__":
+  Tensor.no_grad = True
+
+  parser = argparse.ArgumentParser(description="Run QwQ in tinygrad", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+  parser.add_argument("--size", choices=["32B"], default="32B", help="Model size")
+  parser.add_argument("--count", type=int, default=30, help="Max number of tokens to generate")
+  parser.add_argument("--temperature", type=float, default=0.7, help="Temperature in the softmax")
+  parser.add_argument("--prompt", type=str, default="Hello.", help="Phrase to start with")
+  parser.add_argument("--weights", type=str, default=None, help="Path to the downloaded weights")
+  parser.add_argument("--timing", action="store_true", help="Print timing per token")
+  args = parser.parse_args()
+
+  model_info = MODELS[args.size]
+
+  model_path = Path(args.weights) if args.weights else download_weights(model_info["total_num_weights"])
+  transformer = load_model(model_path, model_info["model_params"])
+  tokenizer = AutoTokenizer.from_pretrained(model_info["tokenizer"])
+  param_bytes = sum(x.lazydata.size * x.dtype.itemsize for x in get_parameters(transformer))
+
+  outputted = args.prompt
+  start_pos, toks = 0, tokenizer(outputted)["input_ids"]
+  print(outputted, end="", flush=True)
+
+  tok_tensor = None
+  for i in range(args.count):
+    GlobalCounters.reset()
+
+    if args.timing: print("")
+    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 "") +
+                  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)
+      tok = tok_tensor.item()
+
+    # use the kv cache
+    start_pos = len(toks)
+
+    # add the new token
+    toks.append(tok)
+
+    cur = tokenizer.decode(toks, skip_special_tokens=True)
+    sys.stdout.write(cur[len(outputted):])
+    sys.stdout.flush()
+    outputted = cur
+
+  if args.temperature == 0:
+    text = tokenizer.decode(toks)
+    key = (args.size, args.count, args.prompt)
+    expected = {
+      ("32B", 10, "Hello."): "Hello. I'm trying to make a program that will read",
+      (
+        "32B",
+        50,
+        "Can you tell me more about machine learning?"
+      ): "Can you tell me more about machine learning? Sure, I'd be happy to help! Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data and make predictions or decisions without being explicitly programmed to do so. It's a fascinating field with a lot of real"
+    }
+    try:
+      assert text == expected[key], f"invalid output: `{colored(text, 'red')}` != `{expected[key]}`"
+      print("\n" + colored("output validated", "green"))
+    except KeyError:
+      pass

examples/self_tokenize.py

@@ -0,0 +1,35 @@
+import os, pathlib
+from examples.llama3 import Tokenizer
+from tabulate import tabulate
+from tinygrad import fetch
+from tinygrad.helpers import flatten
+
+# llama 3 tokenizer
+tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
+
+def read_code(base_path):
+  ret = []
+  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
+    for name in files:
+      if not name.endswith(".py"): continue
+      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
+      fullpath = os.path.join(path, name)
+      code = pathlib.Path(fullpath).read_text()
+      ret += [(fullpath.split("tinygrad/", 1)[1], code)]
+  return ret
+
+if __name__ == "__main__":
+  ret = read_code(".")
+
+  table = []
+  for name,code in ret:
+    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
+  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))
+
+  code_str = '\x00'.join(flatten(ret))
+  print(f"code has {len(code_str)} chars")
+  newline_count = code_str.count('\n')
+  print(f"code has {newline_count} newlines")
+
+  encoded = tokenizer.encode(code_str)
+  print(f"code has {len(encoded)} tokens")

examples/stable_diffusion.py

@@ -189,7 +189,7 @@ class StableDiffusion:
     # make image correct size and scale
     x = (x + 1.0) / 2.0
     x = x.reshape(3,512,512).permute(1,2,0).clip(0,1)*255
-    return x.cast(dtypes.uint8) if Device.DEFAULT != "WEBGPU" else x
+    return x.cast(dtypes.uint8)
 
   def __call__(self, unconditional_context, context, latent, timestep, alphas, alphas_prev, guidance):
     e_t = self.get_model_output(unconditional_context, context, latent, timestep, guidance)
@@ -280,7 +280,7 @@ if __name__ == "__main__":
   print(x.shape)
 
   # save image
-  im = Image.fromarray(x.numpy().astype(np.uint8, copy=False))
+  im = Image.fromarray(x.numpy())
   print(f"saving {args.out}")
   im.save(args.out)
   # Open image.
@@ -290,5 +290,5 @@ if __name__ == "__main__":
   if args.prompt == default_prompt and args.steps == 6 and args.seed == 0 and args.guidance == 7.5:
     ref_image = Tensor(np.array(Image.open(Path(__file__).parent / "stable_diffusion_seed0.png")))
     distance = (((x.cast(dtypes.float) - ref_image.cast(dtypes.float)) / ref_image.max())**2).mean().item()
-    assert distance < 70e-5, colored(f"validation failed with {distance=}", "red")
+    assert distance < 3e-3, colored(f"validation failed with {distance=}", "red")  # higher distance with WINO
     print(colored(f"output validated with {distance=}", "green"))

examples/tinychat/assets/cdn.jsdelivr.net/npm/@alpinejs/intersect@3.x.x/dist/cdn.min.js

@@ -0,0 +1 @@
+(()=>{function o(e){e.directive("intersect",e.skipDuringClone((t,{value:i,expression:l,modifiers:n},{evaluateLater:r,cleanup:c})=>{let s=r(l),a={rootMargin:x(n),threshold:f(n)},u=new IntersectionObserver(d=>{d.forEach(h=>{h.isIntersecting!==(i==="leave")&&(s(),n.includes("once")&&u.disconnect())})},a);u.observe(t),c(()=>{u.disconnect()})}))}function f(e){if(e.includes("full"))return .99;if(e.includes("half"))return .5;if(!e.includes("threshold"))return 0;let t=e[e.indexOf("threshold")+1];return t==="100"?1:t==="0"?0:Number(`.${t}`)}function p(e){let t=e.match(/^(-?[0-9]+)(px|%)?$/);return t?t[1]+(t[2]||"px"):void 0}function x(e){let t="margin",i="0px 0px 0px 0px",l=e.indexOf(t);if(l===-1)return i;let n=[];for(let r=1;r<5;r++)n.push(p(e[l+r]||""));return n=n.filter(r=>r!==void 0),n.length?n.join(" ").trim():i}document.addEventListener("alpine:init",()=>{window.Alpine.plugin(o)});})();

examples/tinychat/assets/unpkg.com/@highlightjs/cdn-assets@11.9.0/styles/vs2015.min.css

@@ -0,0 +1 @@
+pre code.hljs{display:block;overflow-x:auto;padding:1em}code.hljs{padding:3px 5px}.hljs{background:#1e1e1e;color:#dcdcdc}.hljs-keyword,.hljs-literal,.hljs-name,.hljs-symbol{color:#569cd6}.hljs-link{color:#569cd6;text-decoration:underline}.hljs-built_in,.hljs-type{color:#4ec9b0}.hljs-class,.hljs-number{color:#b8d7a3}.hljs-meta .hljs-string,.hljs-string{color:#d69d85}.hljs-regexp,.hljs-template-tag{color:#9a5334}.hljs-formula,.hljs-function,.hljs-params,.hljs-subst,.hljs-title{color:#dcdcdc}.hljs-comment,.hljs-quote{color:#57a64a;font-style:italic}.hljs-doctag{color:#608b4e}.hljs-meta,.hljs-meta .hljs-keyword,.hljs-tag{color:#9b9b9b}.hljs-template-variable,.hljs-variable{color:#bd63c5}.hljs-attr,.hljs-attribute{color:#9cdcfe}.hljs-section{color:gold}.hljs-emphasis{font-style:italic}.hljs-strong{font-weight:700}.hljs-bullet,.hljs-selector-attr,.hljs-selector-class,.hljs-selector-id,.hljs-selector-pseudo,.hljs-selector-tag{color:#d7ba7d}.hljs-addition{background-color:#144212;display:inline-block;width:100%}.hljs-deletion{background-color:#600;display:inline-block;width:100%}

examples/tinychat/assets/unpkg.com/@marcreichel/alpine-autosize@1.3.x/dist/alpine-autosize.min.js

@@ -0,0 +1,2 @@
+!function(e){"function"==typeof define&&define.amd?define(e):e()}((function(){"use strict";var e=new Map;function t(t){var o=e.get(t);o&&o.destroy()}function o(t){var o=e.get(t);o&&o.update()}var r=null;"undefined"==typeof window?((r=function(e){return e}).destroy=function(e){return e},r.update=function(e){return e}):((r=function(t,o){return t&&Array.prototype.forEach.call(t.length?t:[t],(function(t){return function(t){if(t&&t.nodeName&&"TEXTAREA"===t.nodeName&&!e.has(t)){var o,r=null,n=window.getComputedStyle(t),i=(o=t.value,function(){s({testForHeightReduction:""===o||!t.value.startsWith(o),restoreTextAlign:null}),o=t.value}),l=function(o){t.removeEventListener("autosize:destroy",l),t.removeEventListener("autosize:update",a),t.removeEventListener("input",i),window.removeEventListener("resize",a),Object.keys(o).forEach((function(e){return t.style[e]=o[e]})),e.delete(t)}.bind(t,{height:t.style.height,resize:t.style.resize,textAlign:t.style.textAlign,overflowY:t.style.overflowY,overflowX:t.style.overflowX,wordWrap:t.style.wordWrap});t.addEventListener("autosize:destroy",l),t.addEventListener("autosize:update",a),t.addEventListener("input",i),window.addEventListener("resize",a),t.style.overflowX="hidden",t.style.wordWrap="break-word",e.set(t,{destroy:l,update:a}),a()}function s(e){var o,i,l=e.restoreTextAlign,a=void 0===l?null:l,d=e.testForHeightReduction,u=void 0===d||d,c=n.overflowY;if(0!==t.scrollHeight&&("vertical"===n.resize?t.style.resize="none":"both"===n.resize&&(t.style.resize="horizontal"),u&&(o=function(e){for(var t=[];e&&e.parentNode&&e.parentNode instanceof Element;)e.parentNode.scrollTop&&t.push([e.parentNode,e.parentNode.scrollTop]),e=e.parentNode;return function(){return t.forEach((function(e){var t=e[0],o=e[1];t.style.scrollBehavior="auto",t.scrollTop=o,t.style.scrollBehavior=null}))}}(t),t.style.height=""),i="content-box"===n.boxSizing?t.scrollHeight-(parseFloat(n.paddingTop)+parseFloat(n.paddingBottom)):t.scrollHeight+parseFloat(n.borderTopWidth)+parseFloat(n.borderBottomWidth),"none"!==n.maxHeight&&i>parseFloat(n.maxHeight)?("hidden"===n.overflowY&&(t.style.overflow="scroll"),i=parseFloat(n.maxHeight)):"hidden"!==n.overflowY&&(t.style.overflow="hidden"),t.style.height=i+"px",a&&(t.style.textAlign=a),o&&o(),r!==i&&(t.dispatchEvent(new Event("autosize:resized",{bubbles:!0})),r=i),c!==n.overflow&&!a)){var f=n.textAlign;"hidden"===n.overflow&&(t.style.textAlign="start"===f?"end":"start"),s({restoreTextAlign:f,testForHeightReduction:!0})}}function a(){s({testForHeightReduction:!0,restoreTextAlign:null})}}(t)})),t}).destroy=function(e){return e&&Array.prototype.forEach.call(e.length?e:[e],t),e},r.update=function(e){return e&&Array.prototype.forEach.call(e.length?e:[e],o),e});var n=r;document.addEventListener("alpine:init",(()=>{var e;(e=window.Alpine).directive("autosize",((e,{modifiers:t},{cleanup:o})=>{n(e);const r=Array.from(e.attributes);let i=!1;for(let{nodeName:e}of r)if("wire:model"===e||e.startsWith("wire:model.")){i=!0;break}!e.hasAttribute("wire:ignore")&&i&&e.setAttribute("wire:ignore","");const l=()=>{n.update(e)};e.addEventListener("autosize",l),o((()=>{n.destroy(e),e.removeEventListener("autosize",l)}))})),e.magic("autosize",(e=>t=>{(t||e).dispatchEvent(new Event("autosize"))}))}))}));
+//# sourceMappingURL=alpine-autosize.min.js.map

examples/tinychat/assets/unpkg.com/marked-highlight@2.1.2/lib/index.umd.js

@@ -0,0 +1,97 @@
+(function (global, factory) {
+  typeof exports === 'object' && typeof module !== 'undefined' ? factory(exports) :
+  typeof define === 'function' && define.amd ? define(['exports'], factory) :
+  (global = typeof globalThis !== 'undefined' ? globalThis : global || self, factory(global.markedHighlight = {}));
+})(this, (function (exports) { 'use strict';
+
+  function markedHighlight(options) {
+    if (typeof options === 'function') {
+      options = {
+        highlight: options
+      };
+    }
+
+    if (!options || typeof options.highlight !== 'function') {
+      throw new Error('Must provide highlight function');
+    }
+
+    if (typeof options.langPrefix !== 'string') {
+      options.langPrefix = 'language-';
+    }
+
+    return {
+      async: !!options.async,
+      walkTokens(token) {
+        if (token.type !== 'code') {
+          return;
+        }
+
+        const lang = getLang(token.lang);
+
+        if (options.async) {
+          return Promise.resolve(options.highlight(token.text, lang, token.lang || '')).then(updateToken(token));
+        }
+
+        const code = options.highlight(token.text, lang, token.lang || '');
+        if (code instanceof Promise) {
+          throw new Error('markedHighlight is not set to async but the highlight function is async. Set the async option to true on markedHighlight to await the async highlight function.');
+        }
+        updateToken(token)(code);
+      },
+      useNewRenderer: true,
+      renderer: {
+        code({ text, lang, escaped }) {
+          const language = getLang(lang);
+          const classAttr = language
+            ? ` class="${options.langPrefix}${escape(language)}"`
+            : '';
+          text = text.replace(/\n$/, '');
+          return `<pre><code${classAttr}>${escaped ? text : escape(text, true)}\n</code></pre>`;
+        }
+      }
+    };
+  }
+
+  function getLang(lang) {
+    return (lang || '').match(/\S*/)[0];
+  }
+
+  function updateToken(token) {
+    return (code) => {
+      if (typeof code === 'string' && code !== token.text) {
+        token.escaped = true;
+        token.text = code;
+      }
+    };
+  }
+
+  // copied from marked helpers
+  const escapeTest = /[&<>"']/;
+  const escapeReplace = new RegExp(escapeTest.source, 'g');
+  const escapeTestNoEncode = /[<>"']|&(?!(#\d{1,7}|#[Xx][a-fA-F0-9]{1,6}|\w+);)/;
+  const escapeReplaceNoEncode = new RegExp(escapeTestNoEncode.source, 'g');
+  const escapeReplacements = {
+    '&': '&amp;',
+    '<': '&lt;',
+    '>': '&gt;',
+    '"': '&quot;',
+    "'": '&#39;'
+  };
+  const getEscapeReplacement = (ch) => escapeReplacements[ch];
+  function escape(html, encode) {
+    if (encode) {
+      if (escapeTest.test(html)) {
+        return html.replace(escapeReplace, getEscapeReplacement);
+      }
+    } else {
+      if (escapeTestNoEncode.test(html)) {
+        return html.replace(escapeReplaceNoEncode, getEscapeReplacement);
+      }
+    }
+
+    return html;
+  }
+
+  exports.markedHighlight = markedHighlight;
+
+}));

examples/tinychat/fetch_assets.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+fetch() {
+  echo "fetch $1"
+  mkdir -p assets/$1
+  rmdir assets/$1
+  curl -L -o assets/$1 https://$1
+}
+fetch "cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"
+fetch "cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"
+fetch "cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"
+fetch "unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"
+fetch "unpkg.com/[email protected]/dist/cdn.min.js"
+fetch "unpkg.com/[email protected]/dist/purify.min.js"
+fetch "unpkg.com/[email protected]/marked.min.js"
+fetch "unpkg.com/[email protected]/lib/index.umd.js"
+fetch "unpkg.com/@highlightjs/[email protected]/highlight.min.js"
+fetch "cdn.jsdelivr.net/npm/[email protected]/build/base-min.css"
+fetch "cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
+fetch "unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css"

examples/tinychat/index.css

@@ -271,7 +271,7 @@ p {
 
 /* fonts */
 .megrim-regular {
-  font-family: "Megrim", system-ui;
+  font-family: monospace;
   font-weight: 400;
   font-style: normal;
 }

examples/tinychat/index.html

@@ -5,28 +5,24 @@
   <meta name="viewport" content="width=device-width, initial-scale=1">
   <link rel="icon" href="favicon.svg" type="image/svg+xml">
 
-  <script defer src="https://cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="https://cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="https://cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="https://unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"></script>
-  <script defer src="https://unpkg.com/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="assets/cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
+  <script defer src="assets/unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"></script>
+  <script defer src="assets/unpkg.com/[email protected]/dist/cdn.min.js"></script>
 
-  <script src="https://unpkg.com/[email protected]/dist/purify.min.js"></script>
-  <script src="https://unpkg.com/[email protected]/marked.min.js"></script>
-  <script src="https://unpkg.com/[email protected]/lib/index.umd.js"></script>
-  <script src="https://unpkg.com/@highlightjs/[email protected]/highlight.min.js"></script>
+  <script src="assets/unpkg.com/[email protected]/dist/purify.min.js"></script>
+  <script src="assets/unpkg.com/[email protected]/marked.min.js"></script>
+  <script src="assets/unpkg.com/[email protected]/lib/index.umd.js"></script>
+  <script src="assets/unpkg.com/@highlightjs/[email protected]/highlight.min.js"></script>
 
   <script src="index.js"></script>
 
-  <link rel="preconnect" href="https://fonts.googleapis.com">
-  <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
-  <link href="https://fonts.googleapis.com/css2?family=Megrim&display=swap" rel="stylesheet">
-
-  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/build/base-min.css">
-  <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
+  <link rel="stylesheet" href="assets/cdn.jsdelivr.net/npm/[email protected]/build/base-min.css">
+  <link rel="stylesheet" href="assets/cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
     integrity="sha512-SnH5WK+bZxgPHs44uWIX+LLJAJ9/2PkPKZ5QiAj6Ta86w+fsb2TkcmfRyVX3pBnMFcV7oQPJkl9QevSCWr3W6A=="
     crossorigin="anonymous" referrerpolicy="no-referrer" />
-  <link rel="stylesheet" href="https://unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css">
+  <link rel="stylesheet" href="assets/unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css">
 
   <link rel="stylesheet" href="index.css">
   <link rel="stylesheet" href="common.css">

examples/webgpu/efficientnet/index.html

@@ -17,8 +17,11 @@ canvas {  display: none; }
 * {  text-align: center;  font-family: monospace; }
 </style>
 <title>tinygrad has WebGPU</title>
-<script src="./net.js"></script>
 <link rel="icon" type="image/x-icon" href="https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png">
+<script type="module">
+	import model from "../../net.js";
+	window.model = model;
+</script>
 </head>
 <body>
 <h1>WebGPU <a href="https://github.com/geohot/tinygrad">tinygrad</a> EfficientNet!</h1>
@@ -61,8 +64,6 @@ canvas {  display: none; }
 
 	const getLabels = async () => (await fetch("https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json")).json();
 
-	const getSavetensorBuffer = async () => new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer());
-
 	const reorderChannelsAndRemoveAlpha = (data) => {
 		const out = [];
 		let i = 0;
@@ -97,9 +98,8 @@ canvas {  display: none; }
 		try {
 			resultText.innerHTML = "loading..."
 			labels = await getLabels();
-			const safetensor = await getSavetensorBuffer();
 			const device = await getDevice();
-			net = await timer(() => setupNet(device, safetensor), "(compilation)");
+			net = await timer(() => model.load(device, '../../net.safetensors'), "(compilation)");
 			resultText.innerHTML = "ready"
 		} catch (e) {
 			error(e)

examples/webgpu/stable_diffusion/compile.py

@@ -1,231 +1,59 @@
 import os
-from extra.export_model import compile_net, jit_model
+from extra.export_model import export_model
+from extra.f16_decompress import u32_to_f16
 from examples.stable_diffusion import StableDiffusion
-from tinygrad.nn.state import get_state_dict, safe_save, safe_load_metadata, torch_load, load_state_dict
+from tinygrad.nn.state import safe_save, torch_load, load_state_dict
 from tinygrad.tensor import Tensor
-from tinygrad import Device
+from tinygrad import Device, dtypes
 from tinygrad.helpers import fetch
-from typing import NamedTuple, Any, List
-from pathlib import Path
-import argparse
+import requests
 import numpy as np
+from pathlib import Path
 
 def convert_f32_to_f16(input_file, output_file):
   with open(input_file, 'rb') as f:
     metadata_length_bytes = f.read(8)
     metadata_length = int.from_bytes(metadata_length_bytes, byteorder='little', signed=False)
     metadata_json_bytes = f.read(metadata_length)
-    float32_values = np.fromfile(f, dtype=np.float32)
+    values = np.fromfile(f, dtype=np.float32)
 
-  first_text_model_offset = 3772703308
-  num_elements = int((first_text_model_offset)/4)
-  front_float16_values = float32_values[:num_elements].astype(np.float16)
-  rest_float32_values = float32_values[num_elements:]
+  f16_values = values.astype(np.float16)
 
   with open(output_file, 'wb') as f:
     f.write(metadata_length_bytes)
     f.write(metadata_json_bytes)
-    front_float16_values.tofile(f)
-    rest_float32_values.tofile(f)
+    f16_values.tofile(f)
 
-def split_safetensor(fn):
-  _, json_len, metadata = safe_load_metadata(fn)
-  text_model_offset = 3772703308
-  chunk_size = 536870912
-
-  for k in metadata:
-    # safetensor is in fp16, except for text moel
-    if (metadata[k]["data_offsets"][0] < text_model_offset):
-      metadata[k]["data_offsets"][0] = int(metadata[k]["data_offsets"][0]/2)
-      metadata[k]["data_offsets"][1] = int(metadata[k]["data_offsets"][1]/2)
-
-  last_offset = 0
-  part_end_offsets = []
-
-  for k in metadata:
-    offset = metadata[k]['data_offsets'][0]
-
-    if offset == text_model_offset:
-      break
-
-    part_offset = offset - last_offset
-
-    if (part_offset >= chunk_size):
-      part_end_offsets.append(8+json_len+offset)
-      last_offset = offset
-
-  text_model_start = int(text_model_offset/2)
-  net_bytes = bytes(open(fn, 'rb').read())
-  part_end_offsets.append(text_model_start+8+json_len)
-  cur_pos = 0
-
-  for i, end_pos in enumerate(part_end_offsets):
-    with open(f'./net_part{i}.safetensors', "wb+") as f:
-      f.write(net_bytes[cur_pos:end_pos])
-      cur_pos = end_pos
-
-  with open(f'./net_textmodel.safetensors', "wb+") as f:
-    f.write(net_bytes[text_model_start+8+json_len:])
-
-  return part_end_offsets
+def fetch_dep(file, url):
+  with open(file, "w", encoding="utf-8") as f:
+    f.write(requests.get(url).text.replace("https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/bpe_simple_vocab_16e6.mjs", "./bpe_simple_vocab_16e6.mjs"))
 
 if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description='Run Stable Diffusion', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument('--remoteweights', action='store_true', help="Use safetensors from Huggingface, or from local")
-  args = parser.parse_args()
-  Device.DEFAULT = "WEBGPU"
+  fetch_dep(os.path.join(os.path.dirname(__file__), "clip_tokenizer.js"), "https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/clip_tokenizer.js")
+  fetch_dep(os.path.join(os.path.dirname(__file__), "bpe_simple_vocab_16e6.mjs"), "https://huggingface.co/wpmed/tinygrad-sd-f16/raw/main/bpe_simple_vocab_16e6.mjs")
 
+  Device.DEFAULT = "WEBGPU"
   Tensor.no_grad = True
   model = StableDiffusion()
-
-  # load in 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)
 
-  class Step(NamedTuple):
-    name: str = ""
-    input: List[Tensor] = []
-    forward: Any = None
-
-  sub_steps = [
-    Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
-    Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
-    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode)
+  model_parts = [
+    ("textModel", [Tensor.randn(1, 77)], model.cond_stage_model.transformer.text_model),
+    ("diffusor", [
+      Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), 
+      Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)
+    ], model),
+    ("decoder", [Tensor.randn(1,4,64,64)], model.decode),
+    ("f16tof32", [Tensor.randn(2097120, dtype=dtypes.uint32)], u32_to_f16)
   ]
 
-  prg = ""
+  for model in model_parts:
+    prg, inp_sizes, out_sizes, state = export_model(model[2], Device.DEFAULT.lower(), *model[1], model_name=model[0])
+    dirname = Path(__file__).parent
+    weight_loc =  (dirname / f"net_{model[0]}.safetensors").as_posix()
+    safe_save(state, weight_loc)
+    if model[0] == "diffusor":
+      convert_f32_to_f16(weight_loc, (dirname / f"net_diffusor_f16.safetensors").as_posix())
 
-  def compile_step(model, step: Step):
-    run, special_names = jit_model(step, *step.input)
-    functions, statements, bufs, _ = compile_net(run, special_names)
-    state = get_state_dict(model)
-    weights = {id(x.lazydata.base.realized): name for name, x in state.items()}
-    kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
-    kernel_names = ', '.join([name for (name, _, _, _) in statements])
-    kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
-    bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
-    gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
-    input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,(_,value) in enumerate(special_names.items()) if value != "output0"])
-    return f"""\n    var {step.name} = function() {{
-
-    {kernel_code}
-
-    return {{
-      "setup": async (device, safetensor) => {{
-        const metadata = getTensorMetadata(safetensor[0]);
-
-        {bufs}
-
-        {gpu_write_bufs}
-        const gpuReadBuffer = device.createBuffer({{ size: output0.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});
-
-        const kernels = [{kernel_names}];
-        const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
-
-        return async ({",".join([f'data{i}' for i,(k,v) in enumerate(special_names.items()) if v != "output0"])}) => {{
-            const commandEncoder = device.createCommandEncoder();
-
-            {input_writer}
-
-            {kernel_calls}
-            commandEncoder.copyBufferToBuffer(output0, 0, gpuReadBuffer, 0, output0.size);
-            const gpuCommands = commandEncoder.finish();
-            device.queue.submit([gpuCommands]);
-
-            await gpuReadBuffer.mapAsync(GPUMapMode.READ);
-            const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
-            resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
-            gpuReadBuffer.unmap();
-            return resultBuffer;
-        }}
-      }}
-    }}
-  }}
-  """
-
-  for step in sub_steps:
-    print(f'Executing step={step.name}')
-    prg += compile_step(model, step)
-
-    if step.name == "diffusor":
-      if args.remoteweights:
-        base_url = "https://huggingface.co/wpmed/tinygrad-sd-f16/resolve/main"
-      else:
-        state = get_state_dict(model)
-        safe_save(state, os.path.join(os.path.dirname(__file__), "net.safetensors"))
-        convert_f32_to_f16("./net.safetensors", "./net_conv.safetensors")
-        split_safetensor("./net_conv.safetensors")
-        os.remove("net.safetensors")
-        os.remove("net_conv.safetensors")
-        base_url = "."
-
-  prekernel = f"""
-    window.MODEL_BASE_URL= "{base_url}";
-    const getTensorMetadata = (safetensorBuffer) => {{
-      const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
-      const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
-      return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {{...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}}]));
-    }};
-
-  const getTensorBuffer = (safetensorParts, tensorMetadata, key) => {{
-    let selectedPart = 0;
-    let counter = 0;
-    let partStartOffsets = [1131408336, 2227518416, 3308987856, 4265298864];
-    let correctedOffsets = tensorMetadata.data_offsets;
-    let prev_offset = 0;
-
-    for (let start of partStartOffsets) {{
-      prev_offset = (counter == 0) ? 0 : partStartOffsets[counter-1];
-
-      if (tensorMetadata.data_offsets[0] < start) {{
-        selectedPart = counter;
-        correctedOffsets = [correctedOffsets[0]-prev_offset, correctedOffsets[1]-prev_offset];
-        break;
-      }}
-
-      counter++;
-    }}
-
-    let allZero = true;
-    let out = safetensorParts[selectedPart].subarray(...correctedOffsets);
-
-    for (let i = 0; i < out.length; i++) {{
-        if (out[i] !== 0) {{
-            allZero = false;
-            break;
-        }}
-    }}
-
-    if (allZero) {{
-        console.log("Error: weight '" + key + "' is all zero.");
-    }}
-
-    return safetensorParts[selectedPart].subarray(...correctedOffsets);
-  }}
-
-  const getWeight = (safetensors, key) => {{
-    let uint8Data = getTensorBuffer(safetensors, getTensorMetadata(safetensors[0])[key], key);
-    return new Float32Array(uint8Data.buffer, uint8Data.byteOffset, uint8Data.byteLength / Float32Array.BYTES_PER_ELEMENT);
-  }}
-
-  const createEmptyBuf = (device, size) => {{
-      return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
-  }};
-
-  const createWeightBuf = (device, size, data) => {{
-    const buf = device.createBuffer({{ mappedAtCreation: true, size, usage: GPUBufferUsage.STORAGE }});
-    new Uint8Array(buf.getMappedRange()).set(data);
-    buf.unmap();
-    return buf;
-  }};
-
-  const addComputePass = (device, commandEncoder, pipeline, bufs, workgroup) => {{
-    const bindGroup = device.createBindGroup({{layout: pipeline.getBindGroupLayout(0), entries: bufs.map((buffer, index) => ({{ binding: index, resource: {{ buffer }} }}))}});
-    const passEncoder = commandEncoder.beginComputePass();
-    passEncoder.setPipeline(pipeline);
-    passEncoder.setBindGroup(0, bindGroup);
-    passEncoder.dispatchWorkgroups(...workgroup);
-    passEncoder.end();
-  }};"""
-
-  with open(os.path.join(os.path.dirname(__file__), "net.js"), "w") as text_file:
-    text_file.write(prekernel + prg)
+    with open(dirname / f"net_{model[0]}.js", "w") as text_file:
+      text_file.write(prg)

examples/webgpu/stable_diffusion/f16_to_f32.js

@@ -1,64 +0,0 @@
-const f16tof32 = `
-fn u16_to_f16(x: u32) -> f32 {
-    let sign = f32((x >> 15) & 0x1);
-    let exponent = f32((x >> 10) & 0x1F);
-    let fraction = f32(x & 0x3FF);
-
-    let sign_multiplier = select(1.0, -1.0, sign == 1.0);
-    if (exponent == 0.0) {
-        return sign_multiplier * 6.103515625e-5 * (fraction / 1024.0);
-    } else {
-        return sign_multiplier * exp2(exponent - 15.0) * (1.0 + fraction / 1024.0);
-    }
-}
-
-@group(0) @binding(0) var<storage,read_write> data0: array<u32>;
-@group(0) @binding(1) var<storage,read_write> data1: array<f32>;
-@compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let gidx = gid.x;
-    let outgidx = gidx*2;
-
-    if (gidx >= arrayLength(&data0)) {
-        return;
-    }
-
-    let oo = data0[gidx];
-    let oo1 = (oo >> 16);
-    let oo2 = oo & 0xFFFFu;
-
-    let f1 = u16_to_f16(oo2);
-    let f2 = u16_to_f16(oo1);
-    
-    data1[outgidx] = f1;
-    data1[outgidx + 1] = f2;
-}`;
-
-window.f16tof32GPU = async(device, inf16) => {
-    const input = device.createBuffer({size: inf16.length, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
-    const output = device.createBuffer({size: inf16.length*2, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST });
-
-    const gpuWriteBuffer = device.createBuffer({size: input.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE });
-    const gpuReadBuffer = device.createBuffer({ size: output.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
-    const commandEncoder = device.createCommandEncoder();
-    await gpuWriteBuffer.mapAsync(GPUMapMode.WRITE);
-
-    const alignedUint32View = new Uint32Array(inf16.buffer, inf16.byteOffset, inf16.length / 4);
-    new Uint32Array(gpuWriteBuffer.getMappedRange()).set(alignedUint32View);
-
-    gpuWriteBuffer.unmap();
-    commandEncoder.copyBufferToBuffer(gpuWriteBuffer, 0, input, 0, gpuWriteBuffer.size);
-    const pipeline = await device.createComputePipelineAsync({layout: "auto", compute: { module: device.createShaderModule({ code: f16tof32 }), entryPoint: "main" }});
-
-    addComputePass(device, commandEncoder, pipeline, [input, output], [Math.ceil(inf16.length/(4*256)), 1, 1]);
-
-    commandEncoder.copyBufferToBuffer(output, 0, gpuReadBuffer, 0, output.size);
-    const gpuCommands = commandEncoder.finish();
-    device.queue.submit([gpuCommands]);
-
-    await gpuReadBuffer.mapAsync(GPUMapMode.READ);
-    const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
-    resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
-    gpuReadBuffer.unmap();
-
-    return resultBuffer;
-}

examples/webgpu/stable_diffusion/index.html

@@ -5,105 +5,222 @@
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>tinygrad has WebGPU</title>
-
     <style>
-        body {
-            font-family: 'Arial', sans-serif;
-            text-align: center;
-            padding: 30px;
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
         }
 
-        a {
-            text-decoration: none;
-            color: #4A90E2;
+        body {
+            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+            background: #f4f7fb;
+            color: #333;
+            display: flex;
+            justify-content: center;
+            align-items: center;
+            min-height: 100vh;
+            padding: 20px;
+            flex-direction: column;
+            text-align: center;
         }
 
         h1 {
-            font-size: 36px;
-            font-weight: normal;
+            font-size: 2.2rem;
             margin-bottom: 20px;
+            color: #4A90E2;
+        }
+
+        #wgpuError {
+            color: red;
+            font-size: 1.2rem;
+            margin-top: 20px;
+            display: none;
+        }
+
+        #sdTitle {
+            font-size: 1.5rem;
+            margin-bottom: 30px;
         }
 
         #mybox {
-            display: flex;
-            flex-direction: column;
-            align-items: center;
-            gap: 20px;
-            width: 50%;
-            margin: 0 auto;
+            background: #ffffff;
+            padding: 15px;
+            border-radius: 10px;
+            box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
+            width: 120%;
+            max-width: 550px;
+            margin-bottom: 10px;
         }
 
-        #promptText, #stepRange, #btnRunNet, #guidanceRange {
-            font-size: 18px;
+        input[type="text"] {
             width: 100%;
+            padding: 12px;
+            margin-bottom: 20px;
+            font-size: 1rem;
+            border: 1px solid #ccc;
+            border-radius: 8px;
+            outline: none;
+            transition: all 0.3s ease;
         }
 
-        #result {
-            font-size: 48px;
-        }
-
-        #time {
-            font-size: 16px;
-            color: grey;
-        }
-
-        canvas {
-            margin-top: 20px;
-            border: 1px solid #000; 
+        input[type="text"]:focus {
+            border-color: #4A90E2;
         }
 
         label {
             display: flex;
+            justify-content: space-between;
+            font-size: 1rem;
+            margin-bottom: 15px;
             align-items: center;
-            gap: 10px;
-            width: 100%;
         }
-        
-        #sliderValue {
-            margin-right: 10px;/
+
+        input[type="range"] {
+            width: 100%;
+            margin-left: 10px;
+            -webkit-appearance: none;
+            appearance: none;
+            height: 8px;
+            border-radius: 4px;
+            background: #ddd;
+            outline: none;
+            transition: background 0.3s ease;
+        }
+
+        input[type="range"]:focus {
+            background: #4A90E2;
+        }
+
+        #stepRange,
+        #guidanceRange {
+            width: 80%;
+        }
+
+        span {
+            font-size: 1.1rem;
+            font-weight: 600;
+            color: #333;
+        }
+
+        input[type="button"] {
+            padding: 12px 25px;
+            background-color: #4A90E2;
+            color: #fff;
+            font-size: 1.2rem;
+            border: none;
+            border-radius: 8px;
+            cursor: pointer;
+            transition: background-color 0.3s ease;
+            width: 100%;
+            margin-bottom: 20px;
+        }
+
+        input[type="button"]:disabled {
+            background-color: #ccc;
+            cursor: not-allowed;
+        }
+
+        input[type="button"]:hover {
+            background-color: #357ABD;
+        }
+
+        #divModelDl, #divStepProgress {
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+            margin-bottom: 20px;
+        }
+
+        #modelDlProgressBar,
+        #progressBar {
+            width: 80%;
+            height: 12px;
+            border-radius: 6px;
+            background-color: #e0e0e0;
+        }
+
+        #modelDlProgressBar::-webkit-progress-bar,
+        #progressBar::-webkit-progress-bar {
+            border-radius: 6px;
+        }
+
+        #modelDlProgressValue, #progressFraction {
+            font-size: 1rem;
+            font-weight: 600;
+            color: #333;
+        }
+
+        canvas {
+            max-width: 100%;
+            max-height: 450px;
+            margin-top: 10px;
+            border-radius: 8px;
+            border: 1px solid #ddd;
         }
     </style>
 
     <script type="module">
-        import ClipTokenizer from 'https://softwired.nyc3.cdn.digitaloceanspaces.com/sd/clip_tokenizer.js';
+        import ClipTokenizer from './clip_tokenizer.js';
+        import textModel from './net_textModel.js';
+        import diffusor from './net_diffusor.js';
+        import decoder from './net_decoder.js';
+        import f16tof32 from './net_f16tof32.js';
+
         window.clipTokenizer = new ClipTokenizer();
+        window.textModel = textModel;
+        window.diffusor = diffusor;
+        window.decoder = decoder;
+        window.f16tof32 = f16tof32;
     </script>
-    <script src="./f16_to_f32.js"></script>
-    <script src="./net.js"></script>
 </head>
-
 <body>
-<h1 id="wgpuError" style="display: none; color: red;">WebGPU is not supported in this browser</h1>
-<h1 id="sdTitle">StableDiffusion by <a href="https://github.com/tinygrad/tinygrad" target="_blank">tinygrad</a> WebGPU</h1>
-<div id="mybox">
-    <input id="promptText" type="text" placeholder="Enter your prompt here" value="a horse sized cat eating a bagel">
+    <h1 id="wgpuError" style="display: none;">WebGPU is not supported in this browser</h1>
+    <h1 id="sdTitle">StableDiffusion powered by <a href="https://github.com/tinygrad/tinygrad" target="_blank" style="color: #4A90E2;">tinygrad</a></h1>
+    <a href="https://github.com/tinygrad/tinygrad" target="_blank" style="position: absolute; top: 20px; right: 20px;">
+        <img src="https://upload.wikimedia.org/wikipedia/commons/9/91/Octicons-mark-github.svg"
+             alt="GitHub Logo"
+             style="width: 32px; height: 32px;">
+    </a>
     
-    <label>
-        Steps: <span id="stepValue">8</span>
-        <input id="stepRange" type="range" min="5" max="20" value="8" step="1">
-    </label>
+    <div id="mybox">
+        <form id="promptForm">
+            <input id="promptText" type="text" placeholder="Enter your prompt here" value="a human standing on the surface of mars">
 
-    <label>
-        Guidance: <span id="guidanceValue">7.5</span>
-        <input id="guidanceRange" type="range" min="3" max="15" value="7.5" step="0.1">
-    </label>
-    
-    <input id="btnRunNet" type="button" value="Run" disabled>
+            <label>
+                Steps: <span id="stepValue">9</span>
+                <input id="stepRange" type="range" min="5" max="20" value="9" step="1">
+            </label>
 
-    <div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
-        <span id="modelDlTitle">Downloading model</span>
-        <progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
-        <span id="modelDlProgressValue"></span>
+            <label>
+                Guidance: <span id="guidanceValue">8.0</span>
+                <input id="guidanceRange" type="range" min="3" max="15" value="8.0" step="0.1">
+            </label>
+
+            <input id="btnRunNet" type="button" value="Run" disabled>
+        </form>
+
+        <div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
+            <span id="modelDlTitle">Downloading model</span>
+            <progress id="modelDlProgressBar" value="0" max="100" style="flex-grow: 1;"></progress>
+            <span id="modelDlProgressValue"></span>
+        </div>
+
+        <div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
+            <progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
+            <span id="progressFraction"></span>
+        </div>
+
+        <div id="divStepTime" style="display: none; align-items: center; width: 100%; gap: 10px;">
+            <span id="stepTimeValue">0 ms</span>
+        </div>
     </div>
 
-    <div id="divStepProgress" style="display: none; align-items: center; width: 100%; gap: 10px;">
-        <progress id="progressBar" value="0" max="100" style="flex-grow: 1;"></progress>
-        <span id="progressFraction"></span>
-    </div>
-</div>
-<canvas id="canvas" width="512" height="512"></canvas>
+    <canvas id="canvas" width="512" height="512"></canvas>
 
 <script>
+    let f16decomp = null;
+
     function initDb() {
         return new Promise((resolve, reject) => {
             let db;
@@ -129,30 +246,36 @@
     }
 
     function saveTensorToDb(db, id, tensor) {
-        return new Promise((resolve, reject) => {
-            if (db == null) {
-                resolve(null);
+        return readTensorFromDb(db, id).then((result) => {
+            if (!result) {
+                new Promise((resolve, reject) => {
+                    if (db == null) {
+                        resolve(null);
+                    }
+
+                    const transaction = db.transaction(['tensors'], 'readwrite');
+                    const store = transaction.objectStore('tensors');
+                    const request = store.put({ id: id, content: tensor });
+
+                    transaction.onabort = (event) => {
+                        console.log("Transaction error while saving tensor: " + event.target.error);
+                        resolve(null);
+                    };
+
+                    request.onsuccess = () => {
+                        console.log('Tensor saved successfully.');
+                        resolve();
+                    };
+
+                    request.onerror = (event) => {
+                        console.error('Tensor save failed:', event.target.error);
+                        resolve(null);
+                    };
+                });
+            } else {
+                return null;
             }
-
-            const transaction = db.transaction(['tensors'], 'readwrite');
-            const store = transaction.objectStore('tensors');
-            const request = store.put({ id: id, content: tensor });
-
-            transaction.onabort = (event) => {
-                console.log("Transaction error while saving tensor: " + event.target.error);
-                resolve(null);
-            };
-
-            request.onsuccess = () => {
-                console.log('Tensor saved successfully.');
-                resolve();
-            };
-
-            request.onerror = (event) => {
-                console.error('Tensor save failed:', event.target.error);
-                resolve(null);
-            };
-        });
+        }).catch(()=> null);
     }
 
     function readTensorFromDb(db, id) {
@@ -173,10 +296,8 @@
             request.onsuccess = (event) => {
                 const result = event.target.result;
                 if (result) {
-                    console.log("Cache hit: " + id);
                     resolve(result);
                 } else {
-                    console.log("Cache miss: " + id);
                     resolve(null);
                 }
             };
@@ -190,7 +311,7 @@
 
     window.addEventListener('load', async function() {
         if (!navigator.gpu) {
-            document.getElementById("wgpuError").style.display = "";
+            document.getElementById("wgpuError").style.display = "block";
             document.getElementById("sdTitle").style.display = "none";
             return;
         }
@@ -242,151 +363,51 @@
             return res.arrayBuffer();
         };
   
-        const getAndDecompressF16Safetensors = async (device, progress) => {
+        const decompressf16Safetensor = async (device, progress, f16safeTensor) => {
             let totalLoaded = 0;
             let totalSize = 0;
             let partSize = {};
 
-            const progressCallback = (part, loaded, total) => {
-                totalLoaded += loaded;
-
-                if (!partSize[part]) {
-                    totalSize += total;
-                    partSize[part] = true;
-                }
-                
-                progress(totalLoaded, totalSize);
-            };
-
-            let combinedBuffer = await readTensorFromDb(db, "net.f16");
-            let textModelU8 = await readTensorFromDb(db, "net.text");
-            let textModelFetched = false;
-
-            if (combinedBuffer == null) {
-                let dlParts = [
-                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part0.safetensors', progressCallback),
-                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part1.safetensors', progressCallback),
-                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part2.safetensors', progressCallback),
-                    getProgressDlForPart(window.MODEL_BASE_URL + '/net_part3.safetensors', progressCallback)
-                ];
-
-                if (textModelU8 == null) {
-                    dlParts.push(getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
-                }
-
-                let buffers = await Promise.all(dlParts);
-
-                // Combine everything except for text model, since that's alreafy f32
-                const totalLength = buffers.reduce((acc, buffer, index, array) => {
-                    if (index < 4) {
-                        return acc + buffer.byteLength;
-                    } else {
-                        return acc;
-                    }
-                    }, 0
-                );
-
-                combinedBuffer = new Uint8Array(totalLength);
-
-                let offset = 0;
-                buffers.forEach((buffer, index) => {
-                    if (index < 4) {
-                        combinedBuffer.set(new Uint8Array(buffer), offset);
-                        offset += buffer.byteLength;
-                        buffer = null;
-                    }
-                });
-
-                await saveTensorToDb(db, "net.f16", combinedBuffer);
-
-                if (textModelU8 == null) {
-                    textModelFetched = true;
-                    textModelU8 = new Uint8Array(buffers[4]);
-                    await saveTensorToDb(db, "net.text", textModelU8);
-                }
-            } else {
-                combinedBuffer = combinedBuffer.content;
-            }
-
-            if (textModelU8 == null) {
-                textModelU8 = new Uint8Array(await getProgressDlForPart(window.MODEL_BASE_URL + '/net_textmodel.safetensors', progressCallback));
-                await saveTensorToDb(db, "net.text", textModelU8);
-            } else if (!textModelFetched) {
-                textModelU8 = textModelU8.content;
-            }
-
             document.getElementById("modelDlTitle").innerHTML = "Decompressing model";
 
-            const textModelOffset = 3772703308;
-            const metadataLength = Number(new DataView(combinedBuffer.buffer).getBigUint64(0, true));
-            const metadata = JSON.parse(new TextDecoder("utf8").decode(combinedBuffer.subarray(8, 8 + metadataLength)));
+            const metadataLength = Number(new DataView(f16safeTensor.buffer).getBigUint64(0, true));
+            const metadata = JSON.parse(new TextDecoder("utf8").decode(f16safeTensor.subarray(8, 8 + metadataLength)));
 
-            const allToDecomp = combinedBuffer.byteLength - (8 + metadataLength);
-            const decodeChunkSize = 67107840;
+            const allToDecomp = f16safeTensor.byteLength - (8 + metadataLength);
+            const decodeChunkSize = 8388480;
             const numChunks = Math.ceil(allToDecomp/decodeChunkSize);
 
             console.log(allToDecomp + " bytes to decompress");
             console.log("Will be decompressed in " + numChunks+ " chunks");
-
-            let partOffsets = [{start: 0, end: 1131408336}, {start: 1131408336, end: 2227518416}, {start: 2227518416, end: 3308987856}, {start: 3308987856, end: 4265298864}];
-            let parts = [];
-
-            for (let offsets of partOffsets) {
-                parts.push(new Uint8Array(offsets.end-offsets.start));
-            }
-            parts[0].set(new Uint8Array(new BigUint64Array([BigInt(metadataLength)]).buffer), 0);
-            parts[0].set(combinedBuffer.subarray(8, 8 + metadataLength), 8);
-            parts[3].set(textModelU8, textModelOffset+8+metadataLength - partOffsets[3].start);
+            
+            f32safeTensor = new Uint8Array(allToDecomp*2);
+            f32safeTensor.set(new Uint8Array(new BigUint64Array([BigInt(metadataLength)]).buffer), 0);
+            f32safeTensor.set(f16safeTensor.subarray(8, 8 + metadataLength), 8);
 
             let start = Date.now();
-            let cursor = 0;
 
             for (let i = 0; i < numChunks; i++) {
                 progress(i, numChunks);
-                let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
+                let metaOffset = 8 + metadataLength;
+                let chunkStartF16 = metaOffset + (decodeChunkSize * i);
                 let chunkEndF16 = chunkStartF16 + decodeChunkSize;
-                let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
-
-                if (chunk.byteLength %4 != 0) {
-                    const paddingBytes = 4 - (chunk.byteLength % 4);
-                    const alignedBuffer = new ArrayBuffer(chunk.byteLength + paddingBytes);
-                    const alignedView = new Uint8Array(alignedBuffer);
-                    alignedView.set(new Uint8Array(chunk));
-                    chunk = alignedView;
-                }
-
-                let result = await f16tof32GPU(device, chunk);
+                let chunk = f16safeTensor.subarray(chunkStartF16, chunkEndF16);
+                let uint32Chunk = new Uint32Array(chunk.buffer, chunk.byteOffset, chunk.byteLength / 4);
+                let result = await f16decomp(uint32Chunk);
                 let resultUint8 = new Uint8Array(result.buffer);
-                let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
-                let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;
-                let offsetInPart = chunkStartF32 - partOffsets[cursor].start;
-
-                if (chunkEndF32 < partOffsets[cursor].end || cursor === parts.length - 1) {
-                    parts[cursor].set(resultUint8, offsetInPart);
-                } else {
-                    let spaceLeftInCurrentPart = partOffsets[cursor].end - chunkStartF32;
-                    parts[cursor].set(resultUint8.subarray(0, spaceLeftInCurrentPart), offsetInPart);
-
-                    cursor++;
-
-                    if (cursor < parts.length) {
-                        let nextPartOffset = spaceLeftInCurrentPart;
-                        let nextPartLength = resultUint8.length - nextPartOffset;
-                        parts[cursor].set(resultUint8.subarray(nextPartOffset, nextPartOffset + nextPartLength), 0);
-                    }
-                }
-
+                let f32offset = metaOffset + (decodeChunkSize * i * 2);
+                f32safeTensor.set(resultUint8, f32offset);
                 resultUint8 = null;
                 result = null;
             }
 
-            combinedBuffer = null;
+            f16safeTensor = null;
 
             let end = Date.now();
             console.log("Decoding took: " + ((end - start) / 1000) + " s");
             console.log("Avarage " + ((end - start) / numChunks) + " ms per chunk");
 
-            return parts;
+            return f32safeTensor;
         };
 
         const loadNet = async () => {
@@ -398,17 +419,15 @@
             }
 
             const device = await getDevice();
-            safetensorParts = await getAndDecompressF16Safetensors(device, progress);
 
             modelDlTitle.innerHTML = "Compiling model"
 
-            let models = ["textModel", "diffusor", "decoder"];
+            let netText = await textModel.load(device, "./net_textModel.safetensors");
+            let netDiffusor = await diffusor.load(device, "./net_diffusor_f16.safetensors");
+            let netDecoder = await decoder.load(device, "./net_decoder.safetensors");
+            let funcF16Decomp = await f16tof32.load(device);
 
-            nets = await timer(() => Promise.all([
-                textModel().setup(device, safetensorParts),
-                diffusor().setup(device, safetensorParts),
-                decoder().setup(device, safetensorParts)
-            ]).then((loadedModels) => loadedModels.reduce((acc, model, index) => { acc[models[index]] = model; return acc; }, {})), "(compilation)")
+            decompressf16Safetensor(device, progress, diffusor.getWeights());
 
             progress(1, 1);
 
@@ -421,7 +440,7 @@
             document.getElementById("btnRunNet").disabled = false;
         }
 
-        function runStableDiffusion(prompt, steps, guidance) {
+        function runStableDiffusion(prompt, steps, guidance, showStep) {
             return new Promise(async (resolve, reject) => {
                 let context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP(prompt)));
                 let unconditional_context = await timer(() => nets["textModel"](clipTokenizer.encodeForCLIP("")));
@@ -456,8 +475,16 @@
                 
                 for (let i = timesteps.length - 1; i >= 0; i--) {
                     let timestep = new Float32Array([timesteps[i]]);
-                    let x_prev = await timer(() => nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance])));
+                    let start = performance.now()
+                    let x_prev = await nets["diffusor"](unconditional_context, context, latent, timestep, new Float32Array([alphas[i]]), new Float32Array([alphas_prev[i]]), new Float32Array([guidance]));
+                    document.getElementById("divStepTime").style.display = "block";
+                    document.getElementById("stepTimeValue").innerText = `${(performance.now() - start).toFixed(1)} ms / step`
                     latent = x_prev;
+
+                    if (showStep != null) {
+                        showStep(await nets["decoder"](latent));
+                    }
+
                     document.getElementById("progressBar").value = ((steps - i) / steps) * 100
                     document.getElementById("progressFraction").innerHTML = (steps - i) + "/" + steps
                 }
@@ -466,29 +493,53 @@
             });
         }
 
-        document.getElementById("btnRunNet").addEventListener("click", function(e) {
-            e.target.disabled = true;
+        function renderImage(image) {
+            let pixels = []
+            let pixelCounter = 0
 
-            runStableDiffusion(document.getElementById("promptText").value, document.getElementById("stepRange").value, document.getElementById("guidanceRange").value).then((image) => {
-                let pixels = []
-                let pixelCounter = 0
-
-                for (let j = 0; j < 512; j++) {
-                    for (let k = 0; k < 512; k++) {
-                        pixels.push(image[pixelCounter])
-                        pixels.push(image[pixelCounter+1])
-                        pixels.push(image[pixelCounter+2])
-                        pixels.push(255)
-                        pixelCounter += 3
-                    }
+            for (let j = 0; j < 512; j++) {
+                for (let k = 0; k < 512; k++) {
+                    pixels.push(image[pixelCounter])
+                    pixels.push(image[pixelCounter+1])
+                    pixels.push(image[pixelCounter+2])
+                    pixels.push(255)
+                    pixelCounter += 3
                 }
-                
-                ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
-                console.log(image);
-                console.log("Success");
-                e.target.disabled = false;
+            }
+
+            ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
+        }
+
+        const handleRunNetAndRenderResult = () => {
+            document.getElementById("btnRunNet").disabled = true;
+            const canvas = document.getElementById("canvas");
+            ctx.clearRect(0, 0, canvas.width, canvas.height);
+
+            const prevTitleValue = document.getElementById("modelDlTitle").innerHTML;
+            document.getElementById("modelDlTitle").innerHTML = "Running model";
+
+            runStableDiffusion(
+                document.getElementById("promptText").value,
+                document.getElementById("stepRange").value,
+                document.getElementById("guidanceRange").value,
+                // Decode at each step
+                null
+            ).then((image) => {
+                renderImage(image);
+            }).finally(() => {
+                document.getElementById("modelDlTitle").innerHTML = prevTitleValue;
+                document.getElementById("btnRunNet").disabled = false;
             });
-        }, false);
+        };
+
+        document.getElementById("btnRunNet").addEventListener("click", handleRunNetAndRenderResult, false);
+
+        document.getElementById("promptForm").addEventListener("submit", function (event) {
+            event.preventDefault();
+            if (document.getElementById("btnRunNet").disabled) return;
+
+            handleRunNetAndRenderResult();
+        })
 
         const stepSlider = document.getElementById('stepRange');
         const stepValue = document.getElementById('stepValue');

examples/webgpu/yolov8/compile.py

@@ -1,22 +1,18 @@
 from pathlib import Path
-from examples.yolov8 import YOLOv8
+from examples.yolov8 import YOLOv8, get_weights_location
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import safe_save
 from extra.export_model import export_model
-from tinygrad.helpers import fetch
-from tinygrad.helpers import getenv
 from tinygrad.device import Device
 from tinygrad.nn.state import safe_load, load_state_dict
 
 if __name__ == "__main__":
-    Device.DEFAULT = "WEBGL"
+    Device.DEFAULT = "WEBGPU"
     yolo_variant = 'n'
     yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
-    weights_location = Path(__file__).parents[1] / "weights" / f'yolov8{yolo_variant}.safetensors'
-    fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
-    state_dict = safe_load(weights_location)
+    state_dict = safe_load(get_weights_location(yolo_variant))
     load_state_dict(yolo_infer, state_dict)
-    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,640,640))
+    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,416,416), model_name="yolov8")
     dirname = Path(__file__).parent
     safe_save(state, (dirname / "net.safetensors").as_posix())
     with open(dirname / f"net.js", "w") as text_file:

examples/webgpu/yolov8/index.html

@@ -1,12 +1,15 @@
-<!-- Pre and post-processing functions from: https://github.com/AndreyGermanov/yolov8_onnx_javascript -->
 <!DOCTYPE html>
 <html lang="en">
 <head>
     <meta charset="UTF-8">
-    <title>YOLOv8 tinygrad WebGL</title>
-    <script src="./net.js"></script>
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>YOLOv8 tinygrad WebGPU</title>
+    <script type="module">
+        import yolov8 from "./net.js"
+        window.yolov8 = yolov8;
+    </script>
     <style>
-       body {
+        body {
             text-align: center;
             font-family: Arial, sans-serif;
             margin: 0;
@@ -17,7 +20,11 @@
         .video-container {
             position: relative;
             width: 100%;
+            height: 100vh;
             margin: 0 auto;
+            display: flex;
+            align-items: center;
+            justify-content: center;
         }
 
         #video, #canvas {
@@ -28,35 +35,100 @@
             height: auto;
         }
 
+        .loader {
+            width: 48px;
+            height: 48px;
+            border: 5px solid #FFF;
+            border-bottom-color: transparent;
+            border-radius: 50%;
+            display: inline-block;
+            box-sizing: border-box;
+            animation: rotation 1s linear infinite;
+        }
+
+        @keyframes rotation {
+            0% {
+                transform: rotate(0deg);
+            }
+            100% {
+                transform: rotate(360deg);
+            }
+        }
+
         #canvas {
             background: transparent;
         }
 
+        #fps-meter {
+            position: absolute;
+            top: 20px;
+            right: 20px;
+            background-color: rgba(0, 0, 0, 0.7);
+            color: white;
+            padding: 10px;
+            font-size: 18px;
+            border-radius: 5px;
+            z-index: 10;
+        }
+
         h1 {
             margin-top: 20px;
         }
+
+        .loading-container {
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            position: fixed;
+            top: 0;
+            left: 0;
+            width: 100%;
+            height: 100%;
+            background-color: rgba(0, 0, 0, 0.6);
+            z-index: 10;
+        }
+
+        .loading-text {
+            font-size: 24px;
+            color: white;
+            margin-bottom: 20px;
+        }
     </style>
 </head>
 <body>
-    <h1>YOLOv8 tinygrad WebGL</h1>
+    <h2>YOLOv8 tinygrad WebGPU</h2>
+    <h2 id="wgpu-error" style="display: none; color: red;">Error: WebGPU is not supported in this browser</h2>
     <div class="video-container">
         <video id="video" muted autoplay playsinline></video>
         <canvas id="canvas"></canvas>
+        <div id="fps-meter"></div>
+
+        <div id="div-loading" class="loading-container">
+            <p class="loading-text">Loading model</p>
+            <span class="loader"></span>
+        </div>
     </div>
-    
     <script>
         let net = null;
+        const modelInputSize = 416;
+        let lastCalledTime;
+        let fps = 0, accumFps = 0, frameCounter = 0;
 
         const video = document.getElementById('video');
         const canvas = document.getElementById('canvas');
         const context = canvas.getContext('2d');
         const offscreenCanvas = document.createElement('canvas');
-        offscreenCanvas.width = 640;
-        offscreenCanvas.height = 640;
+        const fpsMeter = document.getElementById('fps-meter');
+        const loadingContainer = document.getElementById('div-loading');
+        const wgpuError = document.getElementById('wgpu-error');
+        offscreenCanvas.width = modelInputSize;
+        offscreenCanvas.height = modelInputSize;
         const offscreenContext = offscreenCanvas.getContext('2d');
 
+
         if (navigator.mediaDevices && navigator.mediaDevices.getUserMedia) {
-            navigator.mediaDevices.getUserMedia({ audio: false, video: true }).then(function (stream) {
+            navigator.mediaDevices.getUserMedia({ audio: false,  video: { facingMode: { ideal: "environment" }}}).then(function (stream) {
                 video.srcObject = stream;
                 video.onloadedmetadata = function() {
                     canvas.width = video.clientWidth;
@@ -66,62 +138,103 @@
         }
 
         async function processFrame() {
-            offscreenContext.drawImage(video, 0, 0, 640, 640);
+            if (video.videoWidth == 0 || video.videoHeight == 0) {
+                requestAnimationFrame(processFrame);
+                return;
+            }
+
+            if (!lastCalledTime) {
+                lastCalledTime = performance.now();
+                fps = 0;
+            } else {
+                const now = performance.now();
+                delta = (now - lastCalledTime)/1000.0;
+                lastCalledTime = now;
+                accumFps += 1/delta;
+
+                if (frameCounter++ >= 10) {
+                    fps = accumFps/frameCounter;
+                    frameCounter = 0;
+                    accumFps = 0;
+                    fpsMeter.innerText = `FPS: ${fps.toFixed(1)}`
+                }
+            }
+
+            const videoAspectRatio = video.videoWidth / video.videoHeight;
+            let targetWidth, targetHeight;
+
+            if (videoAspectRatio > 1) {
+                targetWidth = modelInputSize;
+                targetHeight = modelInputSize / videoAspectRatio;
+            } else {
+                targetHeight = modelInputSize;
+                targetWidth = modelInputSize * videoAspectRatio;
+            }
+
+            const offsetX = (modelInputSize - targetWidth) / 2;
+            const offsetY = (modelInputSize - targetHeight) / 2;
+            offscreenContext.clearRect(0, 0, modelInputSize, modelInputSize);
+            offscreenContext.drawImage(video, offsetX, offsetY, targetWidth, targetHeight);
             const boxes = await detectObjectsOnFrame(offscreenContext);
-            drawBoxes(offscreenCanvas, boxes);
+            drawBoxes(offscreenCanvas, boxes, targetWidth, targetHeight, offsetX, offsetY);
             requestAnimationFrame(processFrame);
         }
 
         requestAnimationFrame(processFrame);
 
-        function drawBoxes(offscreenCanvas, boxes) {
-            const canvas = document.querySelector("canvas");
-            const ctx = canvas.getContext("2d");
+        function drawBoxes(offscreenCanvas, boxes, targetWidth, targetHeight, offsetX, offsetY) {
+            const ctx = document.querySelector("canvas").getContext("2d");
             ctx.clearRect(0, 0, canvas.width, canvas.height);
             ctx.lineWidth = 3;
-            ctx.font = "20px serif";
-            const scaleX = canvas.width / 640;
-            const scaleY = canvas.height / 640;
+            ctx.font = "30px serif";
+            const scaleX = canvas.width / targetWidth;
+            const scaleY = canvas.height / targetHeight;
 
             boxes.forEach(([x1, y1, x2, y2, label]) => {
                 const classIndex = yolo_classes.indexOf(label);
                 const color = classColors[classIndex];
-                const textWidth = ctx.measureText(label).width;
                 ctx.strokeStyle = color;
                 ctx.fillStyle = color;
-
-                let adjustedX1 = x1 * scaleX;
-                let adjustedY1 = y1 * scaleY;
-                let adjustedX2 = x2 * scaleX;
-                let adjustedY2 = y2 * scaleY;
-                let boxWidth = adjustedX2 - adjustedX1;
-                let boxHeight = adjustedY2 - adjustedY1;
-
+                const adjustedX1 = (x1 - offsetX) * scaleX;
+                const adjustedY1 = (y1 - offsetY) * scaleY;
+                const adjustedX2 = (x2 - offsetX) * scaleX;
+                const adjustedY2 = (y2 - offsetY) * scaleY;
+                const boxWidth = adjustedX2 - adjustedX1;
+                const boxHeight = adjustedY2 - adjustedY1;
                 ctx.strokeRect(adjustedX1, adjustedY1, boxWidth, boxHeight);
+                const textWidth = ctx.measureText(label).width;
                 ctx.fillRect(adjustedX1, adjustedY1 - 25, textWidth + 10, 25);
-                ctx.fillStyle = "#000000";
-                ctx.fillText(label, adjustedX1, adjustedY1 - 7);
+                ctx.fillStyle = "#FFFFFF";
+                ctx.fillText(label, adjustedX1 + 5, adjustedY1 - 7);
             });
         }
 
         async function detectObjectsOnFrame(offscreenContext) {
-            if (!net) net = await loadNet();
+            if (!net) {
+                let device = await getDevice();
+                if (!device) {
+                    wgpuError.style.display = "block";
+                    loadingContainer.style.display = "none";
+                }
+                net = await yolov8.load(device, "./net.safetensors");
+                loadingContainer.style.display = "none";
+            }
             let start = performance.now();
             const [input,img_width,img_height] = await prepareInput(offscreenContext);
             console.log("Preprocess took: " + (performance.now() - start) + " ms");
             start = performance.now();
-            const output = net(new Float32Array(input));
+            const output = await net(new Float32Array(input));
             console.log("Inference took: " + (performance.now() - start) + " ms");
             start = performance.now();
-            let out = processOutput(output,img_width,img_height);
+            let out = processOutput(output[0],img_width,img_height);
             console.log("Postprocess took: " + (performance.now() - start) + " ms");
             return out;
         }
 
         async function prepareInput(offscreenContext) {
             return new Promise(resolve => {
-                const [img_width,img_height] = [640, 640]
-                const imgData = offscreenContext.getImageData(0,0,640,640);
+                const [img_width,img_height] = [modelInputSize, modelInputSize]
+                const imgData = offscreenContext.getImageData(0,0,modelInputSize,modelInputSize);
                 const pixels = imgData.data;
                 const red = [], green = [], blue = [];
 
@@ -134,36 +247,31 @@
                 resolve([input, img_width, img_height])
             })
         }
-        
-        const loadNet = async () => {
-            try {
-                const safetensor = await (new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer()));
-                const gl = document.createElement("canvas").getContext("webgl2");
-                return setupNet(gl, safetensor);
-            } catch (e) {
-                console.log(e);
-                return null;
-            }
-        }
 
+        const getDevice = async () => {
+            if (!navigator.gpu) return false;
+            const adapter = await navigator.gpu.requestAdapter();
+            return await adapter.requestDevice();
+        };
+        
         function processOutput(output, img_width, img_height) {
             let boxes = [];
-            for (let index=0;index<8400;index++) {
+            const numPredictions = Math.pow(modelInputSize/32, 2) * 21;
+            for (let index=0;index<numPredictions;index++) {
                 const [class_id,prob] = [...Array(80).keys()]
-                    .map(col => [col, output[8400*(col+4)+index]])
+                    .map(col => [col, output[numPredictions*(col+4)+index]])
                     .reduce((accum, item) => item[1]>accum[1] ? item : accum,[0,0]);
-                if (prob < 0.25) {
-                    continue;
-                }
+
+                if (prob < 0.25) continue;
                 const label = yolo_classes[class_id];
                 const xc = output[index];
-                const yc = output[8400+index];
-                const w = output[2*8400+index];
-                const h = output[3*8400+index];
-                const x1 = (xc-w/2)/640*img_width;
-                const y1 = (yc-h/2)/640*img_height;
-                const x2 = (xc+w/2)/640*img_width;
-                const y2 = (yc+h/2)/640*img_height;
+                const yc = output[numPredictions+index];
+                const w = output[2*numPredictions+index];
+                const h = output[3*numPredictions+index];
+                const x1 = (xc-w/2)/modelInputSize*img_width;
+                const y1 = (yc-h/2)/modelInputSize*img_height;
+                const x2 = (xc+w/2)/modelInputSize*img_width;
+                const y2 = (yc+h/2)/modelInputSize*img_height;
                 boxes.push([x1,y1,x2,y2,label,prob]);
             }
 

examples/yolov8.py

@@ -1,5 +1,7 @@
 from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
+from tinygrad.device import is_dtype_supported
+from tinygrad import dtypes
 import numpy as np
 from itertools import chain
 from pathlib import Path
@@ -8,6 +10,7 @@ from collections import defaultdict
 import time, sys
 from tinygrad.helpers import fetch
 from tinygrad.nn.state import safe_load, load_state_dict
+import json
 
 #Model architecture from https://github.com/ultralytics/ultralytics/issues/189
 #The upsampling class has been taken from this pull request https://github.com/tinygrad/tinygrad/pull/784 by dc-dc-dc. Now 2(?) models use upsampling. (retinet and this)
@@ -385,6 +388,32 @@ class YOLOv8:
     yolov8_head_weights = [(22, self.head)]
     return [*zip(backbone_modules, self.net.return_modules()), *zip(yolov8neck_modules, self.fpn.return_modules()), *yolov8_head_weights]
 
+def convert_f16_safetensor_to_f32(input_file: Path, output_file: Path):
+  with open(input_file, 'rb') as f:
+    metadata_length = int.from_bytes(f.read(8), 'little')
+    metadata = json.loads(f.read(metadata_length).decode())
+    float32_values = np.fromfile(f, dtype=np.float16).astype(np.float32)
+
+  for v in metadata.values():
+    if v["dtype"] == "F16": v.update({"dtype": "F32", "data_offsets": [offset * 2 for offset in v["data_offsets"]]})
+
+  with open(output_file, 'wb') as f:
+    new_metadata_bytes = json.dumps(metadata).encode()
+    f.write(len(new_metadata_bytes).to_bytes(8, 'little'))
+    f.write(new_metadata_bytes)
+    float32_values.tofile(f)
+
+def get_weights_location(yolo_variant: str) -> Path:
+  weights_location = Path(__file__).parents[1] / "weights" / f'yolov8{yolo_variant}.safetensors'
+  fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors', weights_location)
+
+  if not is_dtype_supported(dtypes.half):
+    f32_weights = weights_location.with_name(f"{weights_location.stem}_f32.safetensors")
+    if not f32_weights.exists(): convert_f16_safetensor_to_f32(weights_location, f32_weights)
+    weights_location = f32_weights
+
+  return weights_location
+
 if __name__ == '__main__':
 
   # usage : python3 yolov8.py "image_URL OR image_path" "v8 variant" (optional, n is default)
@@ -410,8 +439,7 @@ if __name__ == '__main__':
   # Different YOLOv8 variants use different w , r, and d multiples. For a list , refer to this yaml file (the scales section) https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/models/v8/yolov8.yaml
   depth, width, ratio = get_variant_multiples(yolo_variant)
   yolo_infer = YOLOv8(w=width, r=ratio, d=depth, num_classes=80)
-
-  state_dict = safe_load(fetch(f'https://gitlab.com/r3sist/yolov8_weights/-/raw/master/yolov8{yolo_variant}.safetensors'))
+  state_dict = safe_load(get_weights_location(yolo_variant))
   load_state_dict(yolo_infer, state_dict)
 
   st = time.time()

extra/backends/clang_graph.py

@@ -32,8 +32,8 @@ class ClangGraph(GraphRunner):
     if DEBUG >= 4: print("\n".join(code))
     compiler = Device["CLANG"].compiler
     assert compiler is not None
-    self.clprg = ClangProgram("batched", compiler.compile(prgs+"\n"+"\n".join(code))) # no point in caching the pointers
+    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):
     return cpu_time_execution(
-    lambda: self.clprg(*[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].expr)]), enable=wait)

extra/backends/cstyle.py

@@ -1,70 +0,0 @@
-# TODO: how much of this can be merged with above?
-class WGSLLanguage(CStyleLanguage):
-  code_for_workitem = {"g": lambda x: f"i32(gindex.{'xyz'[x]})", "l": lambda x: f"i32(lindex.{'xyz'[x]})"}
-  size_prefix = "let"
-  barrier="workgroupBarrier();"
-  generic_var_prefix = "var "
-  external_local_bufs = True
-  code_for_op = { **CStyleLanguage().code_for_op,
-                 BinaryOps.CMPLT: lambda x,y,dtype: f"f32({x}<{y})", BinaryOps.CMPEQ: lambda x,y,dtype: f"f32({x}=={y})",
-                 TernaryOps.MULACC: lambda x,y,z,dtype: f"fma({x},{y},{z})", TernaryOps.WHERE: lambda a,b,c,dtype: f"select({c},{b},bool({a}))" }
-  # HACK: write bool as f32
-  type_map = {dtypes.float: "f32", dtypes.half: "f16", dtypes.int32: "i32", dtypes.uint32: "u32", dtypes.bool: "f32"}
-
-  def render_local(self, name: str, dtype:DType, size: int): return f"var<workgroup> {name}: array<{self.type_map[dtype]},{size}>;"
-
-  def render_const(self, x:Union[float,int], var_dtype) -> str:
-    if math.isnan(x): return "nan()"
-    elif math.isinf(x): return ("-" if x < 0 else "") + "inf(1.0)"
-    return f"({super().render_const(x, var_dtype)})"
-
-  def render_if(self, cond: str): return f"if (bool({cond})) {{"
-
-  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], local_size:List[int], prekernel:List[str]) -> str:
-    local_size = local_size[::-1] if local_size else [1]
-    bind_it = iter(range(len(bufs)))
-    prg = "fn nan() -> f32 { let bits = 0xffffffffu; return bitcast<f32>(bits); }\nfn inf(a: f32) -> f32 { return a/0.0; }\n"
-    prg += "\n".join(prekernel+[f"@group(0) @binding({next(bind_it)}) {'var<storage,read_write>' if isinstance(dtype, PtrDType) else 'var<uniform>'} {name}: {f'array<{self.type_map[dtype]}>' if isinstance(dtype, PtrDType) else 'i32'};" for name,dtype in bufs])  # noqa: E501
-    prg += f"\n@compute @workgroup_size({','.join([str(x) for x in local_size])}) fn {function_name}(@builtin(workgroup_id) gindex: vec3<u32>, @builtin(local_invocation_id) lindex: vec3<u32>) {{\n" + "\n".join(kernel) + "\n}"  # noqa: E501
-    return prg
-
-  def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
-    if self.type_map[var_dtype]: return f"bitcast<{self.type_map[var_dtype]}>({x[0]})" if bitcast else f"{self.type_map[var_dtype]}({x[0]})"
-    raise NotImplementedError(f"no cast for {var_dtype}")
-WGSLRenderer = functools.partial(uops_to_cstyle, WGSLLanguage())
-
-
-class GLSLLanguage(CStyleLanguage):
-  type_map = {dtypes.float: "float", dtypes.half: "float", dtypes.int32: "int", dtypes.uint32: "uint", dtypes.bool: "bool"}
-  sampler_prefix = {dtypes.float64: "d", dtypes.float: "", dtypes.half: "", dtypes.int32: "i", dtypes.uint32: "u", dtypes.bool: "i"}
-  fragment_center_offset = 0.5
-  code_for_workitem = {"i": lambda x, offset=fragment_center_offset:f"int(gl_FragCoord.y-{offset}) * width + int(gl_FragCoord.x-{offset})"}
-  code_for_op = {**CStyleLanguage().code_for_op, **{op: lambda a,b,dtype,charforop=charforop: f"bool(int({a}){charforop}int({b}))" \
-    if dtype == dtypes.bool else f"({a}{charforop}{b})" for op,charforop in [(BinaryOps.MUL,"*"),(BinaryOps.ADD,"+"),(BinaryOps.DIV,"/")]},
-    BinaryOps.CMPLT: lambda a,b,dtype: f"(float({a})<float({b}))" if dtype == dtypes.bool else f"({a}<{b})",
-    BinaryOps.MOD: lambda a,b,dtype: f"(int({a})%int({b}))", TernaryOps.WHERE: lambda a,b,c,dtype: f"(float({a})!=0.0?{b}:{c})"}
-
-  def render_const(self, x:Union[float,int], var_dtype) -> str:
-    if math.isnan(x): return "(0.0 / 0.0)"
-    elif math.isinf(x): return ("-" if x < 0 else "") + "(1./0.)"
-    return self.render_cast(["({:.1f})".format(x) if x == int(x) and dtypes.is_float(var_dtype) else f"({x})"]*var_dtype.sz, var_dtype)
-
-  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,DType]], local_size:List[int], prekernel:List[str]) -> str:
-    prg = "#version 330\nprecision highp float;\nprecision highp int;\nin vec2 uv;\nuniform int width;\n"
-    prg += "\n".join([f"uniform {self.sampler_prefix[dtype]}sampler2D {name};" for name,dtype in bufs if name != "data0"])
-    prg += f"\nout {'int' if bufs[0][1] == dtypes.bool else self.type_map[bufs[0][1]]} out_data;\n"
-    return prg + "\nvoid main() {\n" + "\n".join(kernel) + "\n}"
-
-  def render_cast(self, x:List[str], var_dtype:DType, bitcast=False) -> str:
-    if self.type_map[var_dtype]: return f"{self.type_map[var_dtype]}({x[0]})"
-    raise NotImplementedError(f"no cast for {var_dtype}")
-
-  def render_load(self, output_dtype, buf_name, buf_dtype, idx, local=False) -> str:
-    x_calc = f"float(int({idx})%textureSize({buf_name}, 0).x)"
-    y_calc = f"float(int({idx})/textureSize({buf_name}, 0).x)"
-    out_val = f"texture({buf_name}, vec2(float({x_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).x),\
-    float({y_calc} + {self.fragment_center_offset}f)/float(textureSize({buf_name}, 0).y))).r"
-    return f"{self.render_cast([out_val], output_dtype)}"
-
-  def render_store(self, buf_name:str, buf_dtype:DType, var_name:str, var_dtype:DType, idx, local=False) -> str:
-    return f"out_data = {'int' if buf_dtype == dtypes.bool else self.type_map[buf_dtype]}({var_name});"

extra/backends/graph_hip.py

@@ -12,7 +12,7 @@ class HIPGraph(CUDAGraph):
   def __del__(self):
     if hasattr(self, 'graph'): check(hip.hipGraphDestroy(self.graph))
     if hasattr(self, 'instance'): check(hip.hipGraphExecDestroy(self.instance))
-  def set_device(self): hip_set_device(self.device)
+  def set_device(self): hip_set_device(self.dev)
   def encode_args_info(self): return (hip.hipDeviceptr_t, (1,2,3))
   def graph_create(self): return init_c_var(hip.hipGraph_t(), lambda x: check(hip.hipGraphCreate(ctypes.byref(x), 0)))
   def graph_instantiate(self, graph):

extra/backends/hsa_graph.py

@@ -1,7 +1,7 @@
 import ctypes, collections, time, itertools
 from typing import List, Any, Dict, cast, Optional, Tuple
 from tinygrad.helpers import init_c_var, round_up
-from tinygrad.device import Buffer, BufferOptions
+from tinygrad.device import Buffer, BufferSpec
 from tinygrad.device import Compiled, Device
 from tinygrad.ops import Variable
 from tinygrad.runtime.ops_hsa import HSADevice, PROFILE, Profiler
@@ -32,7 +32,7 @@ class HSAGraph(MultiGraphRunner):
     # Check all jit items are compatible.
     compiled_devices = set()
     for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): compiled_devices.add(ji.prg.device)
+      if isinstance(ji.prg, CompiledRunner): compiled_devices.add(ji.prg.dev)
       elif isinstance(ji.prg, BufferXfer):
         for x in ji.bufs[0:2]: compiled_devices.add(Device[cast(Buffer, x).device])
       else: raise GraphException
@@ -43,15 +43,15 @@ class HSAGraph(MultiGraphRunner):
     # Allocate kernel args.
     kernargs_size: Dict[Compiled, int] = collections.defaultdict(int)
     for ji in self.jit_cache:
-      if isinstance(ji.prg, CompiledRunner): kernargs_size[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
-    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferOptions()) for dev,sz in kernargs_size.items()}
+      if isinstance(ji.prg, CompiledRunner): kernargs_size[ji.prg.dev] += round_up(ctypes.sizeof(ji.prg._prg.args_struct_t), 16)
+    kernargs_ptrs: Dict[Compiled, int] = {dev:dev.allocator._alloc(sz, BufferSpec()) for dev,sz in kernargs_size.items()}
 
     # Fill initial arguments.
     self.ji_kargs_structs: Dict[int, ctypes.Structure] = {}
     for j,ji in enumerate(self.jit_cache):
       if not isinstance(ji.prg, CompiledRunner): continue
-      self.ji_kargs_structs[j] = ji.prg.clprg.args_struct_t.from_address(kernargs_ptrs[ji.prg.device])
-      kernargs_ptrs[ji.prg.device] += round_up(ctypes.sizeof(ji.prg.clprg.args_struct_t), 16)
+      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]])
 
@@ -72,13 +72,13 @@ class HSAGraph(MultiGraphRunner):
       if isinstance(ji.prg, CompiledRunner):
         wait_signals = self.access_resources(ji.bufs, ji.prg.p.outs, new_dependency=j, sync_with_aql_packets=False)
         for i in range(0, len(wait_signals), 5):
-          self.virt_aql_queues[ji.prg.device].submit_barrier(wait_signals[i:i+5])
-        self.packets[j] = hsa.hsa_kernel_dispatch_packet_t.from_address(self.virt_aql_queues[ji.prg.device].write_addr)
+          self.virt_aql_queues[ji.prg.dev].submit_barrier(wait_signals[i:i+5])
+        self.packets[j] = hsa.hsa_kernel_dispatch_packet_t.from_address(self.virt_aql_queues[ji.prg.dev].write_addr)
 
         sync_signal = self.alloc_signal(reset_on_start=True) if PROFILE else None
-        self.virt_aql_queues[ji.prg.device].submit_kernel(ji.prg.clprg, *ji.prg.p.launch_dims(var_vals), #type:ignore
+        self.virt_aql_queues[ji.prg.dev].submit_kernel(ji.prg._prg, *ji.prg.p.launch_dims(var_vals), #type:ignore
                                                           ctypes.addressof(self.ji_kargs_structs[j]), completion_signal=sync_signal)
-        if PROFILE: self.profile_info[ji.prg.device].append((sync_signal, ji.prg.clprg.name, False))
+        if PROFILE: self.profile_info[ji.prg.dev].append((sync_signal, ji.prg._prg.name, False))
       elif isinstance(ji.prg, BufferXfer):
         dest, src = [cast(Buffer, x) for x in ji.bufs[0:2]]
         dest_dev, src_dev = cast(HSADevice, Device[dest.device]), cast(HSADevice, Device[src.device])

extra/backends/ops_dsp.py

@@ -1,14 +1,46 @@
 from __future__ import annotations
-from typing import Tuple, Any
+from typing import Tuple, Any, List
 import ctypes, os, mmap, tempfile, pathlib, array, functools, threading, contextlib, sys
 assert sys.platform != 'win32'
-from tinygrad.device import BufferOptions, Compiled, Allocator
+from tinygrad.device import BufferSpec, Compiled, Allocator
+from tinygrad.dtype import dtypes, DType, PtrDType
+from tinygrad.ops import Ops, UOp
 from tinygrad.helpers import from_mv, getenv, round_up, mv_address, to_mv
 from tinygrad.runtime.ops_clang import ClangCompiler
-from tinygrad.renderer.cstyle import DSPRenderer
+from tinygrad.renderer.cstyle import ClangRenderer
 from tinygrad.runtime.autogen import libc, qcom_dsp
 if getenv("IOCTL"): import extra.dsp.run # noqa: F401 # pylint: disable=unused-import
 
+class DSPRenderer(ClangRenderer):
+  device = "DSP"
+  supports_float4 = False
+  buffer_suffix = " restrict __attribute__((align_value(128)))"
+  kernel_prefix = "__attribute__((noinline)) "
+  type_map = { **ClangRenderer.type_map, dtypes.uint64: "unsigned long long", dtypes.int64: "long long" }
+  code_for_op = {**ClangRenderer.code_for_op, Ops.SIN: lambda x,dtype: f"__builtin_sin({x})",
+                 Ops.LOG2: lambda x,dtype: f"__builtin_log2l({x})" if dtype == dtypes.float64 else f"__builtin_log2f({x})",
+                 Ops.EXP2: lambda x,dtype: f"__builtin_exp2l({x})" if dtype == dtypes.float64 else f"__builtin_exp2f({x})"}
+
+  def render_kernel(self, function_name:str, kernel:List[str], bufs:List[Tuple[str,Tuple[DType,bool]]], uops:List[UOp], prefix=None) -> str:
+    ret = super().render_kernel(function_name, kernel, bufs, uops, prefix)
+    msrc = ['''struct dcvs_v2_req { int type; int _pad; _Bool dcvs_enable; char dcvs_option; _Bool set_latency; int latency; _Bool set_dcvs_params;
+                 short _pad2; char target_corner; char min_corner; char max_corner; int _pad3[3]; };''', 'int HAP_power_set(void*, void*);',
+            'typedef union { struct { void *pv; unsigned int len; } buf; struct { int fd; unsigned int offset; } dma; } remote_arg;',
+            'void* HAP_mmap(void *addr, int len, int prot, int flags, int fd, long offset);', 'int HAP_munmap(void *addr, int len);',
+            'unsigned long long HAP_perf_get_time_us(void);', 'int entry(unsigned long long handle, unsigned int sc, remote_arg* pra) {',
+            'struct dcvs_v2_req req = {.type=7, .dcvs_enable=0, .set_latency=1, .latency=100, .set_dcvs_params=1, .target_corner = 6 /* TURBO */};',
+            'HAP_power_set((void*)handle, (void*)&req);']
+    msrc += ['if ((sc>>24) != 2) return 0;']
+    msrc += [f'int sz_or_val_{i} = ((int*)pra[0].buf.pv)[{i}];' for i,b in enumerate(bufs)]
+    msrc += [f'int off{i} = ((int*)pra[1].buf.pv)[{i}];' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += [f'void *buf_{i} = HAP_mmap(0,sz_or_val_{i},3,0,pra[{i+3}].dma.fd,0)+off{i};' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += ["unsigned long long start = HAP_perf_get_time_us();"]
+    msrc += [f"{function_name}({', '.join([(f'buf_{i}' if isinstance(b[1][0], PtrDType) else f'sz_or_val_{i}') for i,b in enumerate(bufs)])});"]
+    msrc += ["*(unsigned long long *)(pra[2].buf.pv) = HAP_perf_get_time_us() - start;"]
+    msrc += [f'HAP_munmap(buf_{i}, sz_or_val_{i});' for i,b in enumerate(bufs) if isinstance(b[1][0], PtrDType)]
+    msrc += ["return 0; }"]
+    return ret + '\n' + '\n'.join(msrc)
+
 def rpc_sc(method=0, ins=0, outs=0, fds=0): return (method << 24) | (ins << 16) | (outs << 8) | fds
 def rpc_prep_args(ins=None, outs=None, in_fds=None):
   ins, outs, in_fds = ins or list(), outs or list(), in_fds or list()
@@ -21,8 +53,8 @@ def rpc_prep_args(ins=None, outs=None, in_fds=None):
   return pra, fds, attrs, (ins, outs)
 
 class DSPProgram:
-  def __init__(self, device:DSPDevice, name:str, lib:bytes):
-    self.device, self.lib = device, lib
+  def __init__(self, dev:DSPDevice, name:str, lib:bytes):
+    self.dev, self.lib = dev, lib
 
   def __call__(self, *bufs, vals:Tuple[int, ...]=(), wait=False):
     if len(bufs) >= 16: raise RuntimeError(f"Too many buffers to execute: {len(bufs)}")
@@ -31,7 +63,7 @@ class DSPProgram:
                                        outs=[timer:=memoryview(bytearray(8)).cast('Q')], in_fds=[b.share_info.fd for b in bufs])
     var_vals_mv.cast('i')[:] = array.array('i', tuple(b.size for b in bufs) + vals)
     off_mv.cast('I')[:] = array.array('I', tuple(b.offset for b in bufs))
-    self.device.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs)
+    self.dev.exec_lib(self.lib, rpc_sc(method=2, ins=2, outs=1, fds=len(bufs)), pra, fds, attrs)
     return timer[0] / 1e6
 
 class DSPBuffer:
@@ -39,25 +71,25 @@ class DSPBuffer:
     self.va_addr, self.size, self.share_info, self.offset = va_addr, size, share_info, offset
 
 class DSPAllocator(Allocator):
-  def __init__(self, device:DSPDevice):
-    self.device = device
+  def __init__(self, dev:DSPDevice):
+    self.dev = dev
     super().__init__()
 
-  def _alloc(self, size:int, options:BufferOptions):
-    b = qcom_dsp.ION_IOC_ALLOC(self.device.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
-    share_info = qcom_dsp.ION_IOC_SHARE(self.device.ion_fd, handle=b.handle)
+  def _alloc(self, size:int, options:BufferSpec):
+    b = qcom_dsp.ION_IOC_ALLOC(self.dev.ion_fd, len=size, align=0x200, heap_id_mask=1<<qcom_dsp.ION_SYSTEM_HEAP_ID, flags=qcom_dsp.ION_FLAG_CACHED)
+    share_info = qcom_dsp.ION_IOC_SHARE(self.dev.ion_fd, handle=b.handle)
     va_addr = libc.mmap(0, size, mmap.PROT_READ|mmap.PROT_WRITE, mmap.MAP_SHARED, share_info.fd, 0)
     return DSPBuffer(va_addr, size, share_info, offset=0)
 
-  def _free(self, opaque:DSPBuffer, options:BufferOptions):
+  def _free(self, opaque:DSPBuffer, options:BufferSpec):
     libc.munmap(opaque.va_addr, opaque.size)
     os.close(opaque.share_info.fd)
-    qcom_dsp.ION_IOC_FREE(self.device.ion_fd, handle=opaque.share_info.handle)
+    qcom_dsp.ION_IOC_FREE(self.dev.ion_fd, handle=opaque.share_info.handle)
 
-  def as_buffer(self, src:DSPBuffer) -> memoryview: return to_mv(src.va_addr, src.size)
-  def copyin(self, dest:DSPBuffer, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), src.nbytes)
-  def copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(from_mv(dest), src.va_addr, dest.nbytes)
-  def offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)
+  def _as_buffer(self, src:DSPBuffer) -> memoryview: return to_mv(src.va_addr, src.size)
+  def _copyin(self, dest:DSPBuffer, src:memoryview): ctypes.memmove(dest.va_addr, from_mv(src), src.nbytes)
+  def _copyout(self, dest:memoryview, src:DSPBuffer): ctypes.memmove(from_mv(dest), src.va_addr, dest.nbytes)
+  def _offset(self, buf, size:int, offset:int): return DSPBuffer(buf.va_addr+offset, size, buf.share_info, buf.offset+offset)
 
 class DSPDevice(Compiled):
   def __init__(self, device:str=""):
@@ -75,7 +107,7 @@ class DSPDevice(Compiled):
                      ClangCompiler("compile_dsp", args=compiler_args, objdump_tool='llvm-objdump'), functools.partial(DSPProgram, self))
 
     fastrpc_shell = memoryview(bytearray(pathlib.Path('/dsp/cdsp/fastrpc_shell_3').read_bytes()))
-    self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferOptions(nolru=True))
+    self.shell_buf = self.allocator.alloc(round_up(fastrpc_shell.nbytes, 0x1000), BufferSpec(nolru=True))
     ctypes.memmove(self.shell_buf.va_addr, mv_address(fastrpc_shell), fastrpc_shell.nbytes)
 
     self.init_dsp()

extra/backends/ops_hip.py

@@ -1,197 +0,0 @@
-from __future__ import annotations
-import ctypes, functools, subprocess, io
-from typing import Tuple, TypeVar, List, Any, cast, Set
-import tinygrad.runtime.autogen.hip as hip
-from tinygrad.helpers import DEBUG, getenv, init_c_var
-from tinygrad.helpers import from_mv, round_up, to_mv, colored, init_c_struct_t
-from tinygrad.device import Compiled, LRUAllocator, BufferOptions, Runner, Device, Buffer, MallocAllocator, update_stats, Compiler, CompilerOptions
-from tinygrad.renderer.cstyle import HIPRenderer
-from tinygrad.runtime.support.hip_comgr import compile_hip
-from tinygrad.renderer.rdna import uops_to_rdna
-
-class RDNACompiler(Compiler):
-  linearizer_opts = LinearizerOptions("HIP", has_tensor_cores=True)
-  def __init__(self, arch:str):
-    self.arch = arch
-    super().__init__(f"compile_rdna_{self.arch}")
-  def render(self, name:str, uops) -> str: return uops_to_rdna(name, uops)
-  def compile(self, src:str) -> bytes:
-    ret = compile_hip(src, self.arch, True)
-    #with open("/tmp/out.so", "wb") as f: f.write(ret)
-    return ret
-
-class HIPCompiler(Compiler):
-  compiler_opts = CompilerOptions("HIP", has_tensor_cores=True, shared_max=65536)
-  def __init__(self, arch:str):
-    self.arch = arch
-    super().__init__(f"compile_hip_{self.arch}")
-  def render(self, name:str, uops) -> str: return HIPRenderer(name, uops)
-  def compile(self, src:str) -> bytes: return compile_hip(src, self.arch)
-
-hip_current_device = None
-def hip_set_device(d:int):
-  global hip_current_device
-  if d == hip_current_device: return
-  check(hip.hipSetDevice(d))
-  hip_current_device = d
-
-def check(status):
-  if status != 0: raise RuntimeError(f"HIP Error {status}, {ctypes.string_at(hip.hipGetErrorString(status)).decode()}")
-
-class HIPProgram:
-  def __init__(self, device:int, name:str, lib:bytes):
-    self.device, self.name, self.lib = device, name, lib
-
-    if DEBUG >= 6:
-      asm = subprocess.check_output(["/opt/rocm/llvm/bin/llvm-objdump", '-d', '-'], input=lib)
-      print('\n'.join([x for x in asm.decode('utf-8').split("\n") if 's_code_end' not in x]))
-
-    hip_set_device(self.device)
-    self.module = init_c_var(hip.hipModule_t(), lambda x: check(hip.hipModuleLoadData(ctypes.byref(x), lib)))
-    self.prg = init_c_var(hip.hipFunction_t(), lambda x: check(hip.hipModuleGetFunction(ctypes.byref(x), self.module, name.encode("utf-8"))))
-
-  def __del__(self):
-    if hasattr(self, 'module'): check(hip.hipModuleUnload(self.module))
-
-  def __call__(self, *args, global_size:Tuple[int,int,int]=(1,1,1), local_size:Tuple[int,int,int]=(1,1,1), vals:Tuple[int, ...]=(), wait=False):
-    hip_set_device(self.device)
-    if not hasattr(self, "vargs"):
-      self.c_args = init_c_struct_t(tuple([(f'f{i}', hip.hipDeviceptr_t) for i in range(len(args))] +
-                                          [(f'v{i}', ctypes.c_int) for i in range(len(vals))]))(*args, *vals)
-      self.vargs = (ctypes.c_void_p * 5)(ctypes.c_void_p(1), ctypes.cast(ctypes.byref(self.c_args), ctypes.c_void_p),
-                                         ctypes.c_void_p(2), ctypes.cast(ctypes.byref(ctypes.c_size_t(ctypes.sizeof(self.c_args))), ctypes.c_void_p),
-                                         ctypes.c_void_p(3))
-    else:
-      for i in range(len(args)): self.c_args.__setattr__(f'f{i}', args[i])
-      for i in range(len(vals)): self.c_args.__setattr__(f'v{i}', vals[i])
-    if wait:
-      evs = [init_c_var(hip.hipEvent_t(), lambda x: hip.hipEventCreate(ctypes.byref(x), 0)) for _ in range(2)]
-      check(hip.hipEventRecord(evs[0], None))
-    check(hip.hipModuleLaunchKernel(self.prg, *global_size, *local_size, 0, None, None, self.vargs))
-    if wait:
-      check(hip.hipEventRecord(evs[1], None))
-      check(hip.hipEventSynchronize(evs[1]))
-      check(hip.hipEventElapsedTime(ctypes.byref(ret := ctypes.c_float()), evs[0], evs[1]))
-      for ev in evs: check(hip.hipEventDestroy(ev))
-      return ret.value * 1e-3
-    return None
-
-T = TypeVar("T")
-CHUNK_SIZE, PAGE_SIZE = 256*1024*1024, 0x1000
-class HIPAllocator(LRUAllocator):
-  def __init__(self, device:HIPDevice):
-    self.device = device
-    self.track_cross_device: Set[HIPDevice] = set()
-    super().__init__()
-  def full_synchronize(self):
-    self.device.synchronize()
-    for x in self.track_cross_device: x.synchronize()
-    self.track_cross_device.clear()
-  def free_cache(self):
-    self.full_synchronize()
-    return super().free_cache()
-  def _alloc(self, size:int):
-    hip_set_device(self.device.device)
-    return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipMalloc(ctypes.byref(x), size)))
-  def _alloc_with_options(self, size:int, options:BufferOptions):
-    hip_set_device(self.device.device)
-    if options.uncached:
-      return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipExtMallocWithFlags(ctypes.byref(x), size, 3)))  # hipDeviceMallocUncached = 3
-    elif options.host:
-      return init_c_var(hip.hipDeviceptr_t(), lambda x: check(hip.hipHostMalloc(ctypes.byref(x), size, 2 if options.signal else 0)))
-    else:
-      raise Exception("no options")
-  def _free(self, opaque:T): check(hip.hipFree(opaque))
-  def copy_from_fd(self, dest, fd, offset, size):
-    hip_set_device(self.device.device)
-    if not hasattr(self, 'hb'):
-      self.hb = [self._alloc_with_options(CHUNK_SIZE, BufferOptions(host=True)) for _ in range(2)]
-      self.hb_events = [None, None]
-      self.hb_polarity = 0
-    fo = io.FileIO(fd, "a+b", closefd=False)
-    fo.seek(offset - (minor_offset:=offset % PAGE_SIZE))
-    copied_in = 0
-    for local_offset in range(0, size+minor_offset, CHUNK_SIZE):
-      local_size = min(round_up(size+minor_offset, PAGE_SIZE)-local_offset, CHUNK_SIZE)
-      if self.hb_events[self.hb_polarity] is not None:
-        # NOTE: block doesn't work here because we modify the CPU memory
-        check(hip.hipEventSynchronize(self.hb_events[self.hb_polarity]))
-        check(hip.hipEventDestroy(self.hb_events[self.hb_polarity]))
-        self.hb_events[self.hb_polarity] = None
-      fo.readinto(to_mv(self.hb[self.hb_polarity], local_size))
-      check(hip.hipMemcpyAsync(ctypes.c_void_p(dest.value + copied_in), ctypes.c_void_p(self.hb[self.hb_polarity].value + minor_offset),
-                               copy_size:=min(local_size-minor_offset, size-copied_in), hip.hipMemcpyHostToDevice, None))
-      self.hb_events[self.hb_polarity] = init_c_var(hip.hipEvent_t(), lambda x: check(hip.hipEventCreate(ctypes.byref(x))))
-      check(hip.hipEventRecord(self.hb_events[self.hb_polarity], None))
-      copied_in += copy_size
-      self.hb_polarity = (self.hb_polarity+1) % len(self.hb)
-      minor_offset = 0 # only on the first
-  def copyin(self, dest:T, src: memoryview):
-    hip_set_device(self.device.device)
-    host_mem = self._alloc_with_options(len(src), BufferOptions(host=True))
-    self.device.pending_copyin.append(host_mem)
-    ctypes.memmove(host_mem, from_mv(src), len(src))
-    check(hip.hipMemcpyAsync(dest, host_mem, len(src), hip.hipMemcpyHostToDevice, None))
-  def copyout(self, dest:memoryview, src:T):
-    self.full_synchronize()
-    hip_set_device(self.device.device)
-    check(hip.hipMemcpy(from_mv(dest), src, len(dest), hip.hipMemcpyDeviceToHost))
-  def transfer(self, dest:T, src:T, sz:int, **kwargs):
-    hip_set_device(self.device.device)
-    check(hip.hipMemcpyAsync(dest, src, sz, hip.hipMemcpyDeviceToDevice, None))
-
-class HIPSyncEvent(Runner):
-  def __init__(self, lb):
-    self.lb, self.device, self.dname = lb, cast(HIPDevice, Device[lb.device]), lb.device
-    super().__init__()
-  def __call__(self, rawbufs:List[Buffer], var_vals, wait=False, jit=False):
-    to_mv(rawbufs[0]._buf, 4).cast("I")[0] = 0
-    hip_set_device(self.device.device)
-    check(hip.hipStreamWriteValue32(None, rawbufs[0]._buf, 1, 0))
-    update_stats(colored("sync", "red"), 0, 0, {}, None, 1, jit, device=self.dname)
-
-class HIPWaitEvent(Runner):
-  def __init__(self, device):
-    self.device, self.dname = cast(HIPDevice, Device[device]), device
-    super().__init__()
-  def __call__(self, rawbufs:List[Buffer], var_vals, wait=False, jit=False):
-    hip_set_device(self.device.device)
-    check(hip.hipStreamWaitValue32(None, rawbufs[0]._buf, 1, 1, 0xFFFFFFFF))
-    update_stats(colored("wait", "RED"), 0, 0, {}, None, 1, jit, device=self.dname)
-
-if getenv("HIPCPU"):
-  rhip = ctypes.CDLL("/usr/local/lib/libremu.so")
-  class RHIPProgram:
-    def __init__(self, name:str, lib:bytes):
-      self.name, self.lib = name, lib
-    def __call__(self, *args, global_size, local_size, vals=(), wait=False):
-      args = (*args, *vals)
-      rhip.hipModuleLaunchKernel(self.lib, len(self.lib), *global_size, *local_size, 0, None, None,
-                                len(args), (ctypes.c_void_p * len(args))(*[ctypes.cast(x, ctypes.c_void_p) for x in args]))
-
-class HIPDevice(Compiled):
-  def __init__(self, device:str=""):
-    self.device = int(device.split(":")[1]) if ":" in device else 0
-    self.pending_copyin: List[ctypes.c_void_p] = []
-    self.track_cross_buffer: List[Any] = []
-    self.peers: Set[int] = set()
-
-    if getenv("HIPCPU"):
-      super().__init__(device, MallocAllocator, HIPCompiler("gfx1100"), RHIPProgram)
-    else:
-      self.arch = init_c_var(hip.hipDeviceProp_t(), lambda x: check(hip.hipGetDeviceProperties(x, self.device))).gcnArchName.decode()
-      from tinygrad.runtime.graph.hip import HIPGraph
-      super().__init__(device, HIPAllocator(self), RDNACompiler(self.arch) if getenv("RDNA") else HIPCompiler(self.arch),
-                       functools.partial(HIPProgram, self.device), HIPGraph)
-  def synchronize(self):
-    if getenv("HIPCPU"): return
-    hip_set_device(self.device)
-    check(hip.hipDeviceSynchronize())
-    for opaque in self.pending_copyin: check(hip.hipFree(opaque))
-    self.track_cross_buffer.clear()
-    self.pending_copyin.clear()
-  def enable_peer(self, dnum):
-    if self.device == dnum or dnum in self.peers: return
-    hip_set_device(self.device)
-    check(hip.hipDeviceEnablePeerAccess(dnum, 0))
-    self.peers.add(dnum)

extra/backends/ops_hsa.py

@@ -3,7 +3,7 @@ import ctypes, functools, subprocess, io, atexit, collections, json
 from typing import Tuple, TypeVar, List, Dict, Any
 import tinygrad.runtime.autogen.hsa as hsa
 from tinygrad.helpers import DEBUG, init_c_var, from_mv, round_up, to_mv, init_c_struct_t, getenv, PROFILE
-from tinygrad.device import Compiled, Compiler, CompileError, BufferOptions, LRUAllocator
+from tinygrad.device import Compiled, Compiler, CompileError, BufferSpec, LRUAllocator
 from tinygrad.renderer.cstyle import HIPRenderer
 from tinygrad.runtime.support.hsa import check, scan_agents, find_memory_pool, AQLQueue
 from tinygrad.runtime.support.hip_comgr import compile_hip
@@ -102,7 +102,7 @@ class HSAAllocator(LRUAllocator):
     self.device = device
     super().__init__()
 
-  def _alloc(self, size:int, options:BufferOptions):
+  def _alloc(self, size:int, options:BufferSpec):
     if options.host:
       check(hsa.hsa_amd_memory_pool_allocate(HSADevice.cpu_mempool, size, 0, ctypes.byref(mem := ctypes.c_void_p())))
       check(hsa.hsa_amd_agents_allow_access(2, (hsa.hsa_agent_t*2)(HSADevice.cpu_agent, self.device.agent), None, mem))
@@ -112,14 +112,14 @@ class HSAAllocator(LRUAllocator):
     check(hsa.hsa_amd_agents_allow_access(len(HSADevice.agents[hsa.HSA_DEVICE_TYPE_GPU]), c_agents, None, buf))
     return buf.value
 
-  def _free(self, opaque:T, options:BufferOptions):
+  def _free(self, opaque:T, options:BufferSpec):
     HSADevice.synchronize_system()
     check(hsa.hsa_amd_memory_pool_free(opaque))
 
-  def copyin(self, dest:T, src: memoryview):
+  def _copyin(self, dest:T, src: memoryview):
     # Async copyin sync model uses barriers on the main hw queue, since barriers are guaranteed to execute in order with all other packets.
     self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
-    mem = self._alloc(src.nbytes, BufferOptions(host=True))
+    mem = self._alloc(src.nbytes, BufferSpec(host=True))
     ctypes.memmove(mem, from_mv(src), src.nbytes)
     check(hsa.hsa_amd_memory_async_copy_on_engine(dest, self.device.agent, mem, HSADevice.cpu_agent, src.nbytes, 1, ctypes.byref(sync_signal),
                                                   copy_signal := self.device.alloc_signal(reusable=True), hsa.HSA_AMD_SDMA_ENGINE_0, True))
@@ -131,7 +131,7 @@ class HSAAllocator(LRUAllocator):
     self.device.hw_queue.submit_barrier([], sync_signal := self.device.alloc_signal(reusable=True))
 
     if not hasattr(self, 'hb'):
-      self.hb = [self._alloc(CHUNK_SIZE, BufferOptions(host=True)) for _ in range(2)]
+      self.hb = [self._alloc(CHUNK_SIZE, BufferSpec(host=True)) for _ in range(2)]
       self.hb_signals = [self.device.alloc_signal(reusable=False) for _ in range(2)]
       self.hb_polarity = 0
       self.sdma = [hsa.HSA_AMD_SDMA_ENGINE_0, hsa.HSA_AMD_SDMA_ENGINE_1]
@@ -164,7 +164,7 @@ class HSAAllocator(LRUAllocator):
     if copies_called > 1: wait_signals.append(self.hb_signals[self.hb_polarity])
     self.device.hw_queue.submit_barrier(wait_signals)
 
-  def copyout(self, dest:memoryview, src:T):
+  def _copyout(self, dest:memoryview, src:T):
     HSADevice.synchronize_system()
     copy_signal = self.device.alloc_signal(reusable=True)
     c_agents = (hsa.hsa_agent_t*2)(self.device.agent, HSADevice.cpu_agent)
@@ -256,7 +256,7 @@ class HSADevice(Compiled):
 
   def _new_kernargs_region(self, sz:int):
     if hasattr(self, 'kernarg_start_addr'): self.delayed_free.append(self.kernarg_start_addr)
-    self.kernarg_start_addr: int = self.allocator._alloc(sz, BufferOptions())
+    self.kernarg_start_addr: int = self.allocator._alloc(sz, BufferSpec())
     self.kernarg_next_addr = self.kernarg_start_addr
     self.kernarg_pool_sz: int = sz
 

extra/backends/ops_rhip.py

@@ -1,18 +0,0 @@
-import ctypes
-from tinygrad.device import Compiled, MallocAllocator
-from tinygrad.renderer.cstyle import HIPRenderer
-from tinygrad.runtime.ops_hsa import HSACompiler
-
-rhip = ctypes.CDLL("/usr/local/lib/libremu.so")
-class RHIPProgram:
-  def __init__(self, name:str, lib:bytes):
-    self.name, self.lib = name, lib
-  def __call__(self, *args, global_size, local_size, vals=(), wait=False):
-    args = (*args, *vals)
-    rhip.hipModuleLaunchKernel(self.lib, len(self.lib), *global_size, *local_size, 0, None, None,
-                              len(args), (ctypes.c_void_p * len(args))(*[ctypes.cast(x, ctypes.c_void_p) for x in args]))
-
-class RHIPDevice(Compiled):
-  def __init__(self, device:str=""):
-    self.device = int(device.split(":")[1]) if ":" in device else 0
-    super().__init__(device, MallocAllocator, HIPRenderer(), HSACompiler("gfx1100"), RHIPProgram)

extra/backends/ops_webgl.py

@@ -1,52 +0,0 @@
-import numpy as np
-import functools
-from tinygrad.dtype import dtypes, ImageDType
-from tinygrad.device import Compiled, Allocator, CompilerOptions
-from tinygrad.codegen.kernel import OptOps
-from tinygrad.renderer.cstyle import uops_to_cstyle
-from tinygrad.renderer.cstyle import GLSLLanguage
-import moderngl
-
-ctx = moderngl.create_standalone_context()
-max_dims = 4096
-dtype_map = { dtypes.float64: "f8", dtypes.float: "f4", dtypes.half: "f2", dtypes.int32: "i4", dtypes.uint32: "u4", dtypes.bool: "i1"}
-vertex_shader="#version 330\nprecision highp float;\nin vec2 in_position;in vec2 in_uv;out vec2 uv;void main(){\
-gl_Position=vec4(in_position,0.0,1.0);uv=in_uv;}"
-class WebGLProgram:
-  def __init__(self, name: str, prg: str, bufs:int=0, vars:int=0):
-    self.name, self.prg = name, ctx.program(vertex_shader=vertex_shader, fragment_shader=prg)
-  def __call__(self, *bufs, global_size, local_size=None, vals=(), wait=False):
-    vert = ctx.buffer(np.asarray([-1, 1, -1, -1, 1, 1, 1, -1], dtype='f4').tobytes())
-    uv = ctx.buffer(np.asarray([0, 1, 0, 0, 1, 1, 1, 0], dtype='f4').tobytes())
-    self.vao = ctx.vertex_array(self.prg, [])
-    self.vao.bind(self.prg["in_position"].location if "in_position" in self.prg else 0, buffer=vert, cls='f', fmt='2f4')
-    self.vao.bind(self.prg["in_uv"].location if "in_uv" in self.prg else 1, buffer=uv, cls='f', fmt='2f4')
-    self.vao.vertices = vert.size//4//2
-    self.fbo = ctx.framebuffer(color_attachments=[bufs[0]])
-
-    for i, x in enumerate(bufs[1:], start=1):
-      if f"data{i}" in self.prg:
-        self.prg[f"data{i}"] = i
-        x.use(i)
-
-    if ("width" in self.prg): self.prg["width"].value = self.fbo.size[0]
-    ctx.viewport = (0, 0, self.fbo.size[0], self.fbo.size[1])
-    self.fbo.use()
-    self.vao.render(mode=moderngl.TRIANGLE_STRIP)
-
-class RawWebGLAllocator(Allocator):
-  def _alloc_image(self, dtype:ImageDType):
-    tex = ctx.texture(dtype.shape, 1, dtype=dtype_map[dtype.base])
-    tex.filter = (moderngl.NEAREST, moderngl.NEAREST)
-    return tex
-  def copyin(self, dest:moderngl.Texture, src: memoryview): dest.write(src)
-  def copyout(self, dest:memoryview, src: moderngl.Texture):
-    src.read_into(dest)
-    return dest
-
-class WebGlDevice(Compiled):
-  def __init__(self, device:str):
-    super().__init__(RawWebGLAllocator(),
-      CompilerOptions(device="WEBGL", global_max=[4096*4096,1,1], unsupported_opts=[OptOps.UPCAST, OptOps.UPCASTMID],
-      supports_float4=False, supports_float4_alu=False, has_local=False, has_shared=False, dont_use_locals=True),
-      functools.partial(uops_to_cstyle, GLSLLanguage()), lambda x: x, WebGLProgram)

extra/backends/ops_webgpu.py

@@ -1,40 +0,0 @@
-from wgpu.utils.device import get_default_device
-from tinygrad.device import Compiled, Allocator, CompilerOptions
-from tinygrad.renderer.cstyle import WGSLRenderer
-import wgpu
-
-wgpu_device = get_default_device()
-def create_uniform(val: int) -> wgpu.GPUBuffer:
-  buf = wgpu_device.create_buffer(size=4, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
-  wgpu_device.queue.write_buffer(buf, 0, val.to_bytes(4, "little"))
-  return buf
-
-class WebGPUProgram:
-  def __init__(self, name:str, lib:bytes):
-    self.name, self.lib, self.prg = name, lib, wgpu_device.create_shader_module(code=lib)   # NOTE: this is the compiler
-  def __call__(self, *bufs, global_size, local_size, vals=(), wait=False):
-    assert len(bufs) <= 8, "WEBGPU only supports 8 buffers"
-    binding_layouts = [{"binding": i, "visibility": wgpu.ShaderStage.COMPUTE, "buffer": {"type": wgpu.BufferBindingType.uniform if i >= len(bufs) else wgpu.BufferBindingType.storage }} for i in range(len(bufs)+len(vals))]  # noqa: E501
-    bindings = [{"binding": i, "resource": {"buffer": create_uniform(x) if i >= len(bufs) else x, "offset": 0, "size": 4 if i >= len(bufs) else x.size}} for i,x in enumerate(bufs+vals)]  # noqa: E501
-    bind_group_layout = wgpu_device.create_bind_group_layout(entries=binding_layouts)
-    pipeline_layout = wgpu_device.create_pipeline_layout(bind_group_layouts=[bind_group_layout])
-    bind_group = wgpu_device.create_bind_group(layout=bind_group_layout, entries=bindings)
-    compute_pipeline = wgpu_device.create_compute_pipeline(layout=pipeline_layout,compute={"module": self.prg, "entry_point": self.name},)
-    command_encoder = wgpu_device.create_command_encoder()
-    compute_pass = command_encoder.begin_compute_pass()
-    compute_pass.set_pipeline(compute_pipeline)
-    compute_pass.set_bind_group(0, bind_group, [], 0, 999999) # last 2 not used
-    compute_pass.dispatch_workgroups(*global_size)  # x y z
-    compute_pass.end()
-    wgpu_device.queue.submit([command_encoder.finish()])
-
-class WebGpuAllocator(Allocator):
-  def _alloc(self, size: int):
-    return wgpu_device.create_buffer(size=size, usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.COPY_SRC)
-  def copyin(self, dest, src: memoryview): wgpu_device.queue.write_buffer(dest, 0, src)
-  def copyout(self, dest, src: memoryview): dest[:] = wgpu_device.queue.read_buffer(src, 0)    # TODO: remove this copy
-
-class WebGpuDevice(Compiled):
-  def __init__(self, device:str):
-    super().__init__(WebGpuAllocator(), CompilerOptions(device="WEBGPU", supports_float4=False, local_max=[256, 256, 64],
-                                                          global_max=[65535, 65535, 65535]), WGSLRenderer, lambda x: x, WebGPUProgram)

extra/debug_sd_speed.py

@@ -1,99 +0,0 @@
-# NOTE: this is written in a way that checkout back to old commit still works
-# fast SD 297ms step on M1 Max, 4444e6d  https://github.com/tinygrad/tinygrad/pull/2129
-# lazy rewrite, 1765849  https://github.com/tinygrad/tinygrad/pull/2878
-# SD 415ms step on M1 Max on master around 11/15/2024
-
-import time
-from typing import Optional
-try: from tinygrad.jit import TinyJit
-except ImportError: from tinygrad import TinyJit
-from tinygrad.tensor import Tensor, Device
-from tinygrad.helpers import GlobalCounters
-from tinygrad.nn import Linear, LayerNorm
-from tinygrad.nn.state import get_parameters
-
-class CrossAttention:
-  def __init__(self, query_dim:int, ctx_dim:int, n_heads:int, d_head:int):
-    self.to_q = Linear(query_dim, n_heads*d_head, bias=False)
-    self.to_k = Linear(ctx_dim,   n_heads*d_head, bias=False)
-    self.to_v = Linear(ctx_dim,   n_heads*d_head, bias=False)
-    self.num_heads = n_heads
-    self.head_size = d_head
-    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)
-    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.dim_out = dim_out
-
-  def __call__(self, x:Tensor) -> Tensor:
-    x, gate = self.proj(x).chunk(2, dim=-1)
-    return x * gate.gelu()
-
-class FeedForward:
-  def __init__(self, dim:int, mult:int=4):
-    self.net = [
-      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)
-
-class BasicTransformerBlock:
-  def __init__(self, dim:int, ctx_dim:int, n_heads:int, d_head:int):
-    self.attn1 = CrossAttention(dim, dim, n_heads, d_head)
-    self.ff    = FeedForward(dim)
-    self.attn2 = CrossAttention(dim, ctx_dim, n_heads, d_head)
-    self.norm1 = LayerNorm(dim)
-    self.norm2 = LayerNorm(dim)
-    self.norm3 = LayerNorm(dim)
-
-  def __call__(self, x:Tensor, ctx:Optional[Tensor]=None) -> Tensor:
-    x = x + self.attn1(self.norm1(x))  # 5.4 before, # 6.8 master
-    x = x + self.attn2(self.norm2(x), ctx=ctx)  # 12 before, 12 master
-    x = x + self.ff(self.norm3(x))  # 23 before, # 27 master
-    return x
-
-def helper_test(gen, model):
-  tms = []
-  for _ in range(5):
-    early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
-    GlobalCounters.reset()
-    Device[Device.DEFAULT].synchronize()
-    st = time.perf_counter_ns()
-    model(*early_gen)
-    Device[Device.DEFAULT].synchronize()
-    tms.append(time.perf_counter_ns() - st)
-  print(f"{min(tms)/1e6=:.2f} ms")
-
-def derandomize_model(model):
-  for p in get_parameters(model):
-    p.lazydata = Tensor.empty(*p.shape, device=p.device, dtype=p.dtype).lazydata
-    p.realize()
-
-def test_transformer_block():
-  # dim, d_head, x = 320, 40, (4096, 320)  # 137ms 4444e6d 115ms master
-  # dim, d_head, x = 640, 80, (1024, 640)  #  36ms 4444e6d, 31ms master
-  dim, d_head, x = 1280, 160, (256, 1280)  # 23ms 4444e6d, 28ms master, 31ms on 176584993
-
-  model = [BasicTransformerBlock(dim, 768, 8, d_head) for _ in range(4)]
-
-  derandomize_model(model)
-  @TinyJit
-  def test(t, t2):
-    for l in model: t = l(t, t2)
-    return t.realize()
-  helper_test(lambda: (Tensor.empty(2, *x), Tensor.empty(2, 77, 768)), test)
-
-if __name__ == "__main__":
-  test_transformer_block()

extra/export_model.py

@@ -1,6 +1,6 @@
-from typing import Tuple, Dict, List
+from typing import Tuple, Dict, List, Optional
 from tinygrad.dtype import DType
-from tinygrad.renderer import Program
+from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
@@ -8,23 +8,12 @@ from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes
 import json
 
-EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "WEBGL", "CLANG", "CUDA", "GPU"]
-web_utils = {
-  "getTensorBuffer":
-  """const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {
-    return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
-  }""",
-  "getTensorMetadata": """const getTensorMetadata = (safetensorBuffer) => {
-    const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
-    const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
-    return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}]));
-  };"""
-}
+EXPORT_SUPPORTED_DEVICE = ["WEBGPU", "CLANG", "CUDA", "GPU"]
 
 def compile_net(run:TinyJit, special_names:Dict[int,str]) -> Tuple[Dict[str,str],List[Tuple[str,List[str],List[int]]],Dict[str,Tuple[int,DType,int]],Dict[str,Tensor]]:
   functions, bufs, bufs_to_save, statements, bufnum = {}, {}, {}, [], 0
   for ji in run.jit_cache:
-    fxn: Program = ji.prg.p
+    fxn: ProgramSpec = ji.prg.p
     functions[fxn.function_name] = fxn.src   # NOTE: this assumes all with the same name are the same
     cargs = []
     for i,arg in enumerate(ji.bufs):
@@ -79,193 +68,59 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
   cprog += [f"void net({inputs}, {outputs}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]
   return '\n'.join(cprog)
 
-def export_model_webgl(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names) -> str:
-  header = f"""
-  function setupNet(gl, safetensor) {{
-    function createShaderProgram(gl, code) {{
-      const vertexShader = loadShader(gl, gl.VERTEX_SHADER, '#version 300 es\\nin vec2 in_position;in vec2 in_uv;out vec2 uv;void main(){{gl_Position=vec4(in_position,0.0,1.0);uv=in_uv;}}');
-      const fragmentShader = loadShader(gl, gl.FRAGMENT_SHADER, code);
-      const shaderProgram = gl.createProgram();
-      gl.attachShader(shaderProgram, vertexShader);
-      gl.attachShader(shaderProgram, fragmentShader);
-      gl.linkProgram(shaderProgram);
+def dtype_to_js_type(dtype: DType) -> str:
+  return f"{'Uint' if dtype in dtypes.uints else 'Int' if (dtype in dtypes.sints or dtype == dtypes.bool) else 'Float'}{8*dtype.itemsize}Array"
 
-      if (!gl.getProgramParameter(shaderProgram, gl.LINK_STATUS)) {{
-        console.log(`Unable to initialize the shader program: ${{gl.getProgramInfoLog(shaderProgram)}}`);
-        return null;
-      }}
-
-      return shaderProgram;
-    }}
-
-    function loadShader(gl, type, source) {{
-      const shader = gl.createShader(type);
-      gl.shaderSource(shader, source);
-      gl.compileShader(shader);
-
-      if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {{
-        console.log(`An error occurred compiling the shaders: ${{gl.getShaderInfoLog(shader)}}`);
-        gl.deleteShader(shader);
-        return null;
-      }}
-
-      return shader;
-    }}
-
-    function setupVertexData(gl, program, vertices) {{
-      let vao = gl.createVertexArray();
-      gl.bindVertexArray(vao);
-      let vertexBuffer = gl.createBuffer();
-      gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
-      gl.bufferData(gl.ARRAY_BUFFER, new Float32Array(vertices), gl.STATIC_DRAW);
-      const positionLocation = gl.getAttribLocation(program, 'in_position');
-      const uvLocation = gl.getAttribLocation(program, 'in_uv');
-      gl.enableVertexAttribArray(positionLocation);
-      gl.vertexAttribPointer(positionLocation, 2, gl.FLOAT, false, 4 * 4, 0);
-      gl.enableVertexAttribArray(uvLocation);
-      gl.vertexAttribPointer(uvLocation, 2, gl.FLOAT, false, 4 * 4, 2 * 4);
-      gl.bindVertexArray(null);
-
-      return vao;
-    }}
-
-    function runProgram(gl, kernelName, program, textures) {{
-      let framebuffer = gl.createFramebuffer();
-      gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
-      gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, textures[0].tex, 0);
-      gl.useProgram(program);
-      gl.uniform1i(gl.getUniformLocation(program, "width"), textures[0].width);
-
-      const vao = setupVertexData(gl, program, [-1, 1, 0, 1, -1, -1, 0, 0, 1, 1, 1, 1, 1, -1, 1, 0]);
-      gl.bindVertexArray(vao);
-      // Texture 0 is the framebuffer texture, so we skip that
-      for (let i = 1; i < textures.length; i++) {{
-        gl.activeTexture(gl.TEXTURE0 + i-1);
-        gl.bindTexture(gl.TEXTURE_2D, textures[i].tex);
-        gl.uniform1i(gl.getUniformLocation(program, 'data' + i), i-1);
-      }}
-
-      gl.viewport(0, 0, textures[0].width, textures[0].height);
-      gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
-      gl.bindFramebuffer(gl.FRAMEBUFFER, null);
-
-      for (let i = 1; i < textures.length; i++) {{
-        gl.activeTexture(gl.TEXTURE0 + i-1);
-        gl.bindTexture(gl.TEXTURE_2D, null);
-      }}
-
-      console.log("Finished running: " + kernelName);
-    }}
-
-    function limitTextureDims(size, threshold) {{
-      if (size <= threshold) {{ return [size, 1] }};
-
-      for (let i = 2; i < threshold + 1; i++) {{
-        if ((size % i == 0) && (Math.floor(size / i) <= threshold)) {{
-          return [Math.floor(size / i), i];
-        }}
-      }}
-
-      return [size, 1];
-    }}
-
-    function updateTextureData(gl, texture, data, isHalf) {{
-      gl.bindTexture(gl.TEXTURE_2D, texture.tex);
-      gl.texSubImage2D(gl.TEXTURE_2D, 0, 0, 0, texture.width, texture.height, gl.RED, (isHalf) ? gl.HALF_FLOAT : gl.FLOAT, data);
-      gl.bindTexture(gl.TEXTURE_2D, null);
-    }}
-
-    function readTextureData(gl, texture) {{
-      const framebuffer = gl.createFramebuffer();
-      gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
-      gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture.tex, 0);
-
-      if (gl.checkFramebufferStatus(gl.FRAMEBUFFER) !== gl.FRAMEBUFFER_COMPLETE) {{
-        throw new Error('Framebuffer not complete');
-      }}
-
-      let data = new Float32Array(texture.width * texture.height);
-      gl.readPixels(0, 0, texture.width, texture.height, gl.RED, gl.FLOAT, data);
-      gl.bindFramebuffer(gl.FRAMEBUFFER, null);
-      gl.deleteFramebuffer(framebuffer);
-
-      return data;
-    }}
-
-    function createTexture(gl, size, isHalf, tensorBuffer) {{
-      const texture = gl.createTexture();
-      gl.bindTexture(gl.TEXTURE_2D, texture);
-      const internalFormat = gl.RGBA;
-      const texSize = limitTextureDims(size, gl.getParameter(gl.MAX_TEXTURE_SIZE));
-      let weights;
-
-      if (tensorBuffer != null) {{
-        if (!isHalf)
-          weights = new Float32Array(tensorBuffer.buffer, tensorBuffer.byteOffset, tensorBuffer.byteLength / Float32Array.BYTES_PER_ELEMENT);
-        else
-          weights = new Uint16Array(tensorBuffer.buffer, tensorBuffer.byteOffset, tensorBuffer.byteLength / Uint16Array.BYTES_PER_ELEMENT);
-      }} else {{
-        if (!isHalf)
-          weights = new Float32Array(size).fill(0.0);
-        else
-          weights = new Uint16Array(size).fill(0.0);
-      }}
-
-      if (size != weights.length)
-        console.log("Weights length: " + weights.length + ", texsize: " + texSize[0]*texSize[1]);
-
-      gl.texImage2D(gl.TEXTURE_2D, 0, (isHalf) ? gl.R16F : gl.R32F, texSize[0], texSize[1], 0, gl.RED, (isHalf) ? gl.HALF_FLOAT : gl.FLOAT, weights);
-      gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
-      gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
-      gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
-      gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
-      gl.bindTexture(gl.TEXTURE_2D, null);
-      return {{ tex: texture, width: texSize[0], height: texSize[1] }};
-    }}
-
-    {web_utils["getTensorBuffer"]}
-    {web_utils["getTensorMetadata"]}
-
-    const metadata = getTensorMetadata(safetensor);
-  """
-
-  textures = '\n    '.join([f"const {name} = " + (f"createTexture(gl, {size/(2 if dtype == dtypes.half else 4)}, {'true' if dtype == dtypes.half else 'false'});" if _key not in weight_names else f"createTexture(gl, {size/(2 if dtype == dtypes.half else 4)}, {'true' if dtype == dtypes.half else 'false'}, getTensorBuffer(safetensor, metadata['{weight_names[_key]}']))") + ";"  for name,(size,dtype,_key) in bufs.items()])
-  kernels = '\n\n'.join([f"const {key} = `{code.replace(key, 'main').replace('version 330', 'version 300 es')}`;" for key, code in functions.items()])
-  kernel_names = ', '.join([name for (name, _args, _global_size, _local_size) in statements])
-  kernel_calls = '\n        '.join([f"runProgram(gl, '{name}', programs[{i}], [{', '.join(args)}]);" for i, (name, args, _global_size, _local_size) in enumerate(statements) ])
-  copy_inputs = "\n".join([f'updateTextureData(gl, {name}, _{name}, {"true" if dtype == dtypes.half else "false"});' for name,(size,dtype,_key) in bufs.items() if "input" in name])
-  entry_point = f"""
-    return function({",".join([f"_{name}" for name,(size,dtype,_key) in bufs.items() if "input" in name])}) {{
-      const ext = gl.getExtension('EXT_color_buffer_float');
-      {copy_inputs}
-      {kernel_calls}
-
-      return readTextureData(gl, output0);
-    }}
-  """
-  programs = f"let programs = [{kernel_names}].map((code) => createShaderProgram(gl, code));"
-  return f"{header}\n{kernels}\n{textures}\n{programs}\n{entry_point}}}"
-
-def export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names) -> Tuple[str,int,int]:
+def export_model_webgpu(functions, statements, bufs, weight_names, input_names, output_names, model_name) -> Tuple[str,int,int]:
+  exported_name = "model" if model_name == None else model_name
   kernel_code = '\n\n'.join([f"const {key} = `{code.replace(key, 'main')}`;" for key, code in functions.items()])
-  kernel_names = ', '.join([name for (name, _args, _global_size, _local_size) in statements])
-  kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
+  kernel_names = ', '.join([name for (name, _, _, _) in statements])
+  create_bind_group_layouts = ",".join([
+    "device.createBindGroupLayout({{entries: [{{binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: {{ type: 'uniform' }}}}, {}]}})".format(
+        ",".join([f"{{binding: {argIdx+1}, visibility: GPUShaderStage.COMPUTE, buffer: {{ type: 'storage' }} }}" for argIdx, _ in enumerate(args)])
+    )
+    for _, (_, args, _, _) in enumerate(statements)
+  ])
+  layouts = f"const layouts=[{create_bind_group_layouts}]"
+  kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, pipelines[{i}], layouts[{i}], infinityBuf, [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
   _bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weight_names else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weight_names[_key]}']))") + ";"  for name,(size,dtype,_key) in bufs.items()])
   gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:{input_name}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,input_name in enumerate(input_names)])
-  input_writers = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n        new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set(" + f'_{inp_name});' + f"\n        gpuWriteBuffer{i}.unmap();\n        commandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, {inp_name}, 0, gpuWriteBuffer{i}.size);"  for i,inp_name in enumerate(input_names)])
+  input_buffer_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
+  output_buffer_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
+  input_writers = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n        new {input_buffer_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'_{inp_name});' + f"\n        gpuWriteBuffer{i}.unmap();\n        commandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, {inp_name}, 0, gpuWriteBuffer{i}.size);"  for i,inp_name in enumerate(input_names)])
   gpu_read_bufs = '\n    '.join([f"const gpuReadBuffer{i} = device.createBuffer({{size:{output_name}.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});" for i,output_name in enumerate(output_names)])
   outbuf_copies = '\n        '.join([f"commandEncoder.copyBufferToBuffer({output_name}, 0, gpuReadBuffer{i}, 0, output{i}.size);" for i,output_name in enumerate(output_names)])
-  output_readers = '\n        '.join([f"await gpuReadBuffer{i}.mapAsync(GPUMapMode.READ);\n        const resultBuffer{i} = new Float32Array(gpuReadBuffer{i}.size);\n        resultBuffer{i}.set(new Float32Array(gpuReadBuffer{i}.getMappedRange()));\n        gpuReadBuffer{i}.unmap();" for i in range(len(output_names))])
+  output_readers = '\n        '.join([f"await gpuReadBuffer{i}.mapAsync(GPUMapMode.READ);\n        const resultBuffer{i} = new {output_buffer_types[i]}(gpuReadBuffer{i}.size/{bufs[output_names[i]][1].itemsize});\n        resultBuffer{i}.set(new {output_buffer_types[i]}(gpuReadBuffer{i}.getMappedRange()));\n        gpuReadBuffer{i}.unmap();" for i in range(len(output_names))])
   output_return = '[{}]'.format(",".join([f'resultBuffer{i}' for i in range(len(output_names))]))
   return f"""
-{web_utils["getTensorBuffer"]}
+const {exported_name} = (() => {{
+let weights = null;
+const getTensorBuffer = (safetensorBuffer, tensorMetadata) => {{
+  return safetensorBuffer.subarray(...tensorMetadata.data_offsets);
+}};
 
-{web_utils["getTensorMetadata"]}
+const getTensorMetadata = (safetensorBuffer) => {{
+    const metadataLength = Number(new DataView(safetensorBuffer.buffer).getBigUint64(0, true));
+    const metadata = JSON.parse(new TextDecoder("utf8").decode(safetensorBuffer.subarray(8, 8 + metadataLength)));
+    return Object.fromEntries(Object.entries(metadata).filter(([k, v]) => k !== "__metadata__").map(([k, v]) => [k, {{...v, data_offsets: v.data_offsets.map(x => 8 + metadataLength + x)}}]));
+}};
 
 const createEmptyBuf = (device, size) => {{
     return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
 }};
 
+const createInfinityUniformBuf = (device) => {{
+  const size = 4;
+  const buf = device.createBuffer({{
+    mappedAtCreation: true,
+    size,
+    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST
+  }});
+  new Float32Array(buf.getMappedRange())[0] = Infinity;
+  buf.unmap();
+  return buf;
+}};
+
 const createWeightBuf = (device, size, data) => {{
   const buf = device.createBuffer({{ mappedAtCreation: true, size, usage: GPUBufferUsage.STORAGE }});
   new Uint8Array(buf.getMappedRange()).set(data);
@@ -273,8 +128,15 @@ const createWeightBuf = (device, size, data) => {{
   return buf;
 }};
 
-const addComputePass = (device, commandEncoder, pipeline, bufs, workgroup) => {{
-  const bindGroup = device.createBindGroup({{layout: pipeline.getBindGroupLayout(0), entries: bufs.map((buffer, index) => ({{ binding: index, resource: {{ buffer }} }}))}});
+const addComputePass = (device, commandEncoder, pipeline, layout, infinityUniformBuf, bufs, workgroup) => {{
+  const bindGroup = device.createBindGroup({{
+    layout: layout,
+    entries: [
+      {{ binding: 0, resource: {{ buffer: infinityUniformBuf }} }},
+      ...bufs.map((buffer, index) => ({{ binding: index + 1, resource: {{ buffer }} }}))
+    ]
+  }});
+
   const passEncoder = commandEncoder.beginComputePass();
   passEncoder.setPipeline(pipeline);
   passEncoder.setBindGroup(0, bindGroup);
@@ -285,7 +147,11 @@ const addComputePass = (device, commandEncoder, pipeline, bufs, workgroup) => {{
 {kernel_code}
 
 const setupNet = async (device, safetensor) => {{
-    const metadata = getTensorMetadata(safetensor);
+    weights = safetensor;
+    const metadata = safetensor ? getTensorMetadata(safetensor) : null;
+    const infinityBuf = createInfinityUniformBuf(device);
+
+    {layouts}
 
     {_bufs}
 
@@ -294,7 +160,19 @@ const setupNet = async (device, safetensor) => {{
     {gpu_read_bufs}
 
     const kernels = [{kernel_names}];
-    const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
+    const pipelines = await Promise.all(kernels.map(async (name, i) => {{
+      return await device.createComputePipelineAsync({{
+          layout: device.createPipelineLayout({{
+              bindGroupLayouts: [layouts[i]],
+          }}),
+          compute: {{
+              module: device.createShaderModule({{
+                  code: name,
+              }}),
+              entryPoint: "main",
+          }},
+      }});
+  }}))
 
     return async ({",".join([f"_{input_name}" for input_name in input_names])}) => {{
         const commandEncoder = device.createCommandEncoder();
@@ -308,10 +186,18 @@ const setupNet = async (device, safetensor) => {{
         return {output_return};
     }}
 }}
-  """ + f"\n\nconst loadNet = async (device) => {{ return await fetch('net.safetensors').then(x => x.arrayBuffer()).then(x => setupNet(device, new Uint8Array(x))); }}"
+const load = async (device, weight_path) =>
+{{
+  const buffer = weight_path ? await fetch(weight_path).then(x => x.arrayBuffer()) : null;
+  return setupNet(device, buffer ? new Uint8Array(buffer) : null);
+}}
+return {{ load, getWeights: () => weights }};
+}})();
+export default {exported_name};
+"""
 
-def export_model(model, target:str, *inputs):
-  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, "only WEBGPU, WEBGL, CLANG, CUDA, GPU, METAL are supported"
+def export_model(model, target:str, *inputs, model_name: Optional[str] = None):
+  assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, "only WEBGPU, CLANG, CUDA, GPU, METAL are supported"
   with Context(JIT=2): run,special_names = jit_model(model, *inputs)
   functions, statements, bufs, bufs_to_save = compile_net(run, special_names)
   state = get_state_dict(model)
@@ -322,9 +208,7 @@ def export_model(model, target:str, *inputs):
   if target == "clang":
     prg = export_model_clang(functions, statements, bufs, bufs_to_save, input_names, output_names)
   elif target == "webgpu":
-    prg = export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names)
-  elif target == "webgl":
-    prg = export_model_webgl(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names)
+    prg = export_model_webgpu(functions, statements, bufs, weight_names, input_names, output_names, model_name)
   else:
     prg = json.dumps({
       "backend": Device.DEFAULT,

extra/f16_decompress.py

@@ -0,0 +1,16 @@
+from tinygrad import Tensor
+
+def bit_extract(x: Tensor, e: int, s: int) -> Tensor:
+  mask = (1 << (e - s + 1)) - 1
+  return (x >> s) & mask
+
+def u16_to_f16(x: Tensor) -> Tensor:
+  sign = bit_extract(x, 15, 15).float()
+  exponent = bit_extract(x, 14, 10).float()
+  fraction = bit_extract(x, 9, 0).float()
+  return sign.where(-1, 1) * exponent.where((exponent - 15.0).exp2() * (1 + fraction / 1024.0), 6.103515625e-5 * (fraction / 1024.0))
+
+def u32_to_f16(oo: Tensor) -> Tensor:
+  f1 = u16_to_f16(oo>>16)
+  f2 = u16_to_f16(oo&0xFFFF)
+  return Tensor.cat(f2.reshape(-1, 1), f1.reshape(-1, 1), dim=1).flatten()

extra/f16_w_uint32.py

@@ -1,40 +0,0 @@
-import numpy as np
-from tinygrad import Device, dtypes, Tensor
-
-# TODO: will be better when tinygrad does math in the target dtype, can remove the floor and use a mul
-def bit_extract(x, s, e) -> Tensor:
-  # extract the top bits we don't want
-  top_bits = (x / (1<<(s+1))).floor() * (1<<(s+1))
-  x = (x - top_bits) / (1<<e)
-  return x.contiguous()
-
-def u16_to_f16(x):
-  sign = bit_extract(x, 15, 15).float()
-  exponent = bit_extract(x, 14, 10).float()
-  fraction = bit_extract(x, 9, 0).float()
-  return sign.where(-1, 1) * exponent.where((exponent - 15).exp2() * (1 + fraction / 0x400), 6.103515625e-5 * (fraction / 0x400))
-
-def u32_to_f16(oo):
-  oo1 = (oo/0x10000).floor().contiguous()
-  # TODO: this is wrong and unextractable until we do this math in u32
-  oo2 = (oo-(oo1*0x10000)).floor().contiguous()
-  f1 = u16_to_f16(oo1)
-  f2 = u16_to_f16(oo2)
-  return Tensor.cat(f2.reshape(-1, 1), f1.reshape(-1, 1), dim=1).flatten()
-
-if __name__ == "__main__":
-  # random float16
-  Tensor.manual_seed(2)
-  a = Tensor.randn(100, dtype=dtypes.float16)
-
-  # this converts it to u32 on disk
-  oo = a.to("disk:/tmp/f16").cast(dtypes.uint32)[:50].to(Device.DEFAULT).realize()
-
-  # convert to 2xf16 using tinygrad math ops
-  f16 = u32_to_f16(oo)
-
-  ref = a.numpy()
-  out = f16.numpy().astype(np.float16)
-  print(ref-out)
-
-  np.testing.assert_allclose(ref, out)

extra/gemm/amx.py

@@ -48,8 +48,8 @@ a = MallocAllocator.alloc(na.size * np.dtype(np.float32).itemsize)
 b = MallocAllocator.alloc(nb.size * np.dtype(np.float32).itemsize)
 c = MallocAllocator.alloc(nc.size * np.dtype(np.float32).itemsize)
 
-MallocAllocator.copyin(b, flat_mv(nb.data))
-MallocAllocator.copyin(c, flat_mv(nc.data))
+MallocAllocator._copyin(b, flat_mv(nb.data))
+MallocAllocator._copyin(c, flat_mv(nc.data))
 
 module = ir.Module(name=__file__)
 func = ir.Function(module, ir.FunctionType(ir.IntType(64), [ir.FloatType().as_pointer()]*3), name='exec')
@@ -171,7 +171,7 @@ def timeit(fxn):
   return time.perf_counter() - st
 
 tm = min([timeit(lambda: prog(a, b, c, N**2)) for _ in range(20)])
-MallocAllocator.copyout(flat_mv(na.data), a)
+MallocAllocator._copyout(flat_mv(na.data), a)
 print(f"{N*N:10d} {tm*1e6:9.2f} us, {BW*1e-9/tm:.2f} GB/s")
 
 np.testing.assert_allclose(na[:ns.shape[0]], ns, atol=1e-4, rtol=1e-4)

extra/gemm/cuda_matmul.py

@@ -23,8 +23,8 @@ a = cudaalloc.alloc(N*N*2 if FLOAT16 else N*N*4)
 b = cudaalloc.alloc(N*N*2 if FLOAT16 else N*N*4)
 c = cudaalloc.alloc(N*N*4)
 
-cudaalloc.copyin(a, bytearray(na))
-cudaalloc.copyin(b, bytearray(nb))
+cudaalloc._copyin(a, bytearray(na))
+cudaalloc._copyin(b, bytearray(nb))
 
 FLOPS = N*N*N*2
 BW = N*N*3*4
@@ -103,5 +103,5 @@ extern "C" __global__ void wmma_example({'half' if FLOAT16 else 'float'} *a, {'h
 global_size, local_size = [(N//16)//4, (N//16)//4, 1], [32, 1, 1]
 tm = min([prog(a, b, c, global_size=global_size, local_size=local_size, wait=True) for _ in range(20)])
 print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")
-cudaalloc.copyout(flat_mv(nc.data), c)
+cudaalloc._copyout(flat_mv(nc.data), c)
 np.testing.assert_allclose(na.T.astype(np.float32) @ nb.T.astype(np.float32), nc.reshape(N,N).T, atol=1e-2)

extra/gemm/hip_matmul.py

@@ -40,8 +40,8 @@ c = hipallocator.alloc(N*N*2)
 na = np.empty(N*N, np.float32)
 nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
 nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
-hipallocator.copyin(b, memoryview(bytearray(nb)))
-hipallocator.copyin(c, memoryview(bytearray(nc)))
+hipallocator._copyin(b, memoryview(bytearray(nb)))
+hipallocator._copyin(c, memoryview(bytearray(nc)))
 
 prog_str = f"""
 #define F32
@@ -126,13 +126,13 @@ def timeit(fxn):
   if RAND:
     nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
     nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32).astype(np.float16)
-    hipallocator.copyin(b, memoryview(bytearray(nb)))
-    hipallocator.copyin(c, memoryview(bytearray(nc)))
+    hipallocator._copyin(b, memoryview(bytearray(nb)))
+    hipallocator._copyin(c, memoryview(bytearray(nc)))
   return et
 
 print("global/local size", global_size, local_size, f"local_size:{prod(local_size)} total_size:{prod(global_size+local_size)}")
 tm = min([timeit(lambda: prog(a, b, c, global_size=global_size, local_size=local_size, wait=True)) for _ in range(CNT)])
-hipallocator.copyout(flat_mv(na.data),a)
+hipallocator._copyout(flat_mv(na.data),a)
 na = na.reshape(N,N)
 comp = nb.astype(np.float32) @ nc.astype(np.float32)
 print(f"{N*N:10d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matmul, {BW*1e-9/tm:.2f} GB/s")

extra/gemm/metal_matmul.py

@@ -20,8 +20,8 @@ na = np.zeros((N,N),dtype=np.float32)
 nb = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)N
 nc = np.random.default_rng().standard_normal(size=(N,N), dtype=np.float32) #.astype(np.int32).astype(np.float32)
 
-metalalloc.copyin(b,nb.tobytes())
-metalalloc.copyin(c,nc.tobytes())
+metalalloc._copyin(b,nb.tobytes())
+metalalloc._copyin(c,nc.tobytes())
 
 FLOPS = N*N*N*2
 BW = N*N*3*4
@@ -96,7 +96,7 @@ def timeit(fxn):
   return time.perf_counter() - st
 tm = min([timeit(lambda: prog(a, b, c, global_size=[N//(8*4), N//(8*4*LID), 1], local_size=[32, LID, 1], wait=True)) for _ in range(20)])
 comp = nb@nc
-metalalloc.copyout(flat_mv(na.data), a)
+metalalloc._copyout(flat_mv(na.data), a)
 if N <= 32:
   print(na)
   print(comp)

extra/gemm/metal_matvec.py

@@ -80,8 +80,8 @@ kernel void test(device float* data0, const device float* data1, const device fl
 a = metalalloc.alloc(M*4)
 b = metalalloc.alloc(N*4)
 c = metalalloc.alloc(N*M*4)
-metalalloc.copyin(b,nb.tobytes())
-metalalloc.copyin(c,nc.tobytes())
+metalalloc._copyin(b,nb.tobytes())
+metalalloc._copyin(c,nc.tobytes())
 def metalrun():
   prog(a, b, c, global_size=GLOBAL_SIZE, local_size=LOCAL_SIZE, wait=True)
   return a
@@ -93,7 +93,7 @@ def timeit(fxn):
 tm = min([timeit(metalrun) for _ in range(200)])
 print(f"{N:d}x{M:d} {tm*1e6:9.2f} us, would be {FLOPS*1e-9/tm:9.2f} GFLOPS matvec in metal")
 metal_a = np.zeros(M, dtype=np.float32)
-metalalloc.copyout(flat_mv(metal_a.data), a)
+metalalloc._copyout(flat_mv(metal_a.data), a)
 np.testing.assert_allclose(metal_a, torch_a, atol=5e-3)
 
 b = Tensor(nb)

extra/gemm/triton_nv_matmul.py

@@ -4,7 +4,7 @@ import triton.language as tl
 from triton.compiler import AttrsDescriptor, ASTSource, compile as triton_compile
 import numpy as np
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.engine.realize import CompiledRunner, ExecItem, Program
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ProgramSpec
 from tinygrad.helpers import getenv
 np.set_printoptions(suppress=True)
 
@@ -85,7 +85,7 @@ if __name__ == "__main__":
   # remove debug sections
   src = src.split("\t.file")[0]
   assert '.extern .shared' not in src
-  prg = Program("matmul_kernel", src, dname=Device.DEFAULT,
+  prg = ProgramSpec("matmul_kernel", src, device=Device.DEFAULT,
                 global_size=[M//BLOCK_SIZE_M, N//BLOCK_SIZE_N, 1], local_size=[32*compiled.metadata.num_warps, 1, 1],
                 mem_estimate=A.nbytes() + B.nbytes() + C.nbytes())
   ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)

extra/hip_gpu_driver/kfd_pm4_headers_ai.h

@@ -0,0 +1,247 @@
+/* SPDX-License-Identifier: GPL-2.0 OR MIT */
+/*
+ * Copyright 2016-2022 Advanced Micro Devices, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
+ * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
+ * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+ * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+ * OTHER DEALINGS IN THE SOFTWARE.
+ *
+ */
+
+#ifndef F32_MES_PM4_PACKETS_H
+#define F32_MES_PM4_PACKETS_H
+
+#define uint32_t unsigned int
+#define int32_t int
+
+#ifndef PM4_MES_HEADER_DEFINED
+#define PM4_MES_HEADER_DEFINED
+union PM4_MES_TYPE_3_HEADER {
+	struct {
+		uint32_t reserved1 : 8; /* < reserved */
+		uint32_t opcode    : 8; /* < IT opcode */
+		uint32_t count     : 14;/* < number of DWORDs - 1 in the
+					 *   information body.
+					 */
+		uint32_t type      : 2; /* < packet identifier.
+					 *   It should be 3 for type 3 packets
+					 */
+	};
+	uint32_t u32All;
+};
+#endif /* PM4_MES_HEADER_DEFINED */
+
+#ifndef PM4_MEC_RELEASE_MEM_DEFINED
+#define PM4_MEC_RELEASE_MEM_DEFINED
+
+enum mec_release_mem_event_index_enum {
+	event_index__mec_release_mem__end_of_pipe = 5,
+	event_index__mec_release_mem__shader_done = 6
+};
+
+enum mec_release_mem_cache_policy_enum {
+	cache_policy__mec_release_mem__lru = 0,
+	cache_policy__mec_release_mem__stream = 1
+};
+
+enum mec_release_mem_pq_exe_status_enum {
+	pq_exe_status__mec_release_mem__default = 0,
+	pq_exe_status__mec_release_mem__phase_update = 1
+};
+
+enum mec_release_mem_dst_sel_enum {
+	dst_sel__mec_release_mem__memory_controller = 0,
+	dst_sel__mec_release_mem__tc_l2 = 1,
+	dst_sel__mec_release_mem__queue_write_pointer_register = 2,
+	dst_sel__mec_release_mem__queue_write_pointer_poll_mask_bit = 3
+};
+
+enum mec_release_mem_int_sel_enum {
+	int_sel__mec_release_mem__none = 0,
+	int_sel__mec_release_mem__send_interrupt_only = 1,
+	int_sel__mec_release_mem__send_interrupt_after_write_confirm = 2,
+	int_sel__mec_release_mem__send_data_after_write_confirm = 3,
+	int_sel__mec_release_mem__unconditionally_send_int_ctxid = 4,
+	int_sel__mec_release_mem__conditionally_send_int_ctxid_based_on_32_bit_compare = 5,
+	int_sel__mec_release_mem__conditionally_send_int_ctxid_based_on_64_bit_compare = 6
+};
+
+enum mec_release_mem_data_sel_enum {
+	data_sel__mec_release_mem__none = 0,
+	data_sel__mec_release_mem__send_32_bit_low = 1,
+	data_sel__mec_release_mem__send_64_bit_data = 2,
+	data_sel__mec_release_mem__send_gpu_clock_counter = 3,
+	data_sel__mec_release_mem__send_cp_perfcounter_hi_lo = 4,
+	data_sel__mec_release_mem__store_gds_data_to_memory = 5
+};
+
+struct pm4_mec_release_mem {
+	union {
+		union PM4_MES_TYPE_3_HEADER header;     /*header */
+		unsigned int ordinal1;
+	};
+
+	union {
+		struct {
+			unsigned int event_type:6;
+			unsigned int reserved1:2;
+			enum mec_release_mem_event_index_enum event_index:4;
+			unsigned int tcl1_vol_action_ena:1;
+			unsigned int tc_vol_action_ena:1;
+			unsigned int reserved2:1;
+			unsigned int tc_wb_action_ena:1;
+			unsigned int tcl1_action_ena:1;
+			unsigned int tc_action_ena:1;
+			uint32_t reserved3:1;
+			uint32_t tc_nc_action_ena:1;
+			uint32_t tc_wc_action_ena:1;
+			uint32_t tc_md_action_ena:1;
+			uint32_t reserved4:3;
+			enum mec_release_mem_cache_policy_enum cache_policy:2;
+			uint32_t reserved5:2;
+			enum mec_release_mem_pq_exe_status_enum pq_exe_status:1;
+			uint32_t reserved6:2;
+		} bitfields2;
+		unsigned int ordinal2;
+	};
+
+	union {
+		struct {
+			uint32_t reserved7:16;
+			enum mec_release_mem_dst_sel_enum dst_sel:2;
+			uint32_t reserved8:6;
+			enum mec_release_mem_int_sel_enum int_sel:3;
+			uint32_t reserved9:2;
+			enum mec_release_mem_data_sel_enum data_sel:3;
+		} bitfields3;
+		unsigned int ordinal3;
+	};
+
+	union {
+		struct {
+			uint32_t reserved10:2;
+			unsigned int address_lo_32b:30;
+		} bitfields4;
+		struct {
+			uint32_t reserved11:3;
+			uint32_t address_lo_64b:29;
+		} bitfields4b;
+		uint32_t reserved12;
+		unsigned int ordinal4;
+	};
+
+	union {
+		uint32_t address_hi;
+		uint32_t reserved13;
+		uint32_t ordinal5;
+	};
+
+	union {
+		uint32_t data_lo;
+		uint32_t cmp_data_lo;
+		struct {
+			uint32_t dw_offset:16;
+			uint32_t num_dwords:16;
+		} bitfields6c;
+		uint32_t reserved14;
+		uint32_t ordinal6;
+	};
+
+	union {
+		uint32_t data_hi;
+		uint32_t cmp_data_hi;
+		uint32_t reserved15;
+		uint32_t reserved16;
+		uint32_t ordinal7;
+	};
+
+	uint32_t int_ctxid;
+
+};
+
+#endif
+
+#ifndef PM4_MEC_WRITE_DATA_DEFINED
+#define PM4_MEC_WRITE_DATA_DEFINED
+
+enum WRITE_DATA_dst_sel_enum {
+	dst_sel___write_data__mem_mapped_register = 0,
+	dst_sel___write_data__tc_l2 = 2,
+	dst_sel___write_data__gds = 3,
+	dst_sel___write_data__memory = 5,
+	dst_sel___write_data__memory_mapped_adc_persistent_state = 6,
+};
+
+enum WRITE_DATA_addr_incr_enum {
+	addr_incr___write_data__increment_address = 0,
+	addr_incr___write_data__do_not_increment_address = 1
+};
+
+enum WRITE_DATA_wr_confirm_enum {
+	wr_confirm___write_data__do_not_wait_for_write_confirmation = 0,
+	wr_confirm___write_data__wait_for_write_confirmation = 1
+};
+
+enum WRITE_DATA_cache_policy_enum {
+	cache_policy___write_data__lru = 0,
+	cache_policy___write_data__stream = 1
+};
+
+
+struct pm4_mec_write_data_mmio {
+	union {
+		union PM4_MES_TYPE_3_HEADER header;     /*header */
+		unsigned int ordinal1;
+	};
+
+	union {
+		struct {
+			unsigned int reserved1:8;
+			unsigned int dst_sel:4;
+			unsigned int reserved2:4;
+			unsigned int addr_incr:1;
+			unsigned int reserved3:2;
+			unsigned int resume_vf:1;
+			unsigned int wr_confirm:1;
+			unsigned int reserved4:4;
+			unsigned int cache_policy:2;
+			unsigned int reserved5:5;
+		} bitfields2;
+		unsigned int ordinal2;
+	};
+
+	union {
+		struct {
+			unsigned int dst_mmreg_addr:18;
+			unsigned int reserved6:14;
+		} bitfields3;
+		unsigned int ordinal3;
+	};
+
+	uint32_t reserved7;
+
+	uint32_t data;
+
+};
+
+#endif
+
+enum {
+	CACHE_FLUSH_AND_INV_TS_EVENT = 0x00000014
+};
+#endif
+

extra/introspection.py

@@ -1,6 +1,6 @@
 import gc
 from tinygrad.helpers import prod
-from tinygrad.engine.lazy import LazyBuffer
+from tinygrad.ops import UOp
 from tinygrad.device import Buffer
 from tinygrad import Tensor, GlobalCounters
 
@@ -8,7 +8,7 @@ def print_objects():
   #gc.collect()
   tensors = [x for x in gc.get_objects() if isinstance(x, Tensor)]
   tensor_ram_used = sum([prod(x.shape)*4 for x in tensors])
-  lazybuffers = [x for x in gc.get_objects() if isinstance(x, LazyBuffer)]
+  lazybuffers = [x for x in gc.get_objects() if isinstance(x, UOp)]
   gpubuffers = [x for x in gc.get_objects() if isinstance(x, Buffer) and hasattr(x, "_buf")]
   realized_buffers = [x.realized for x in lazybuffers if x.base == x and x.realized]
   gpubuffers_orphaned = [x for x in gpubuffers if x not in realized_buffers]

extra/models/llama.py

@@ -3,11 +3,10 @@ from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
 from tinygrad.helpers import getenv
 
 # https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0, dtype=dtypes.half) -> Tensor:
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
   freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
   freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
-  # TODO: move dtype outside this
-  return Tensor.stack(freqs.cos().cast(dtype), freqs.sin().cast(dtype), dim=-1).reshape(1, end, 1, dim//2, 2)
+  return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)
 
 # (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
 def complex_mult(A, c, d):
@@ -135,7 +134,7 @@ def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
 
     # approximate top p
     # because we are already limited to top k elements we can do top p "without sorting"
-    output_cumsum = output[::-1]._cumsum()[::-1] + t.sum()
+    output_cumsum = output[::-1].cumsum()[::-1] + t.sum()
     output = (output_cumsum >= (1 - p)) * output
     output_indices = (output_cumsum >= (1 - p)) * output_indices
 
@@ -163,9 +162,11 @@ class Transformer:
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
     _bsz, seqlen = tokens.shape
+    h = self.tok_embeddings(tokens)
+
+    self.freqs_cis = self.freqs_cis.cast(h.dtype).realize()
     freqs_cis = self.freqs_cis.shrink((None, (start_pos, start_pos+seqlen),None,None,None))
 
-    h = self.tok_embeddings(tokens)
     mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1).realize() if seqlen > 1 else None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
     logits = self.output(self.norm(h)).float()[:, -1, :]
@@ -180,7 +181,7 @@ class Transformer:
 
 # *** helpers ***
 
-def convert_from_huggingface(weights:Dict[str, Tensor], model: Transformer, n_heads: int, n_kv_heads: int):
+def convert_from_huggingface(weights:Dict[str, Tensor], model: Transformer, n_heads: int, n_kv_heads: int, permute_layers: bool = True):
   def permute(v: Tensor, n_heads: int):
     return v.reshape(n_heads, 2, v.shape[0] // n_heads // 2, v.shape[1]).transpose(1, 2).reshape(*v.shape[:2])
 
@@ -188,6 +189,7 @@ def convert_from_huggingface(weights:Dict[str, Tensor], model: Transformer, n_he
     "model.embed_tokens.weight": "tok_embeddings.weight",
     **{f"model.layers.{l}.input_layernorm.weight": f"layers.{l}.attention_norm.weight" for l in range(len(model.layers))},
     **{f"model.layers.{l}.self_attn.{x}_proj.weight": f"layers.{l}.attention.w{x}.weight" for x in ["q", "k", "v", "o"] for l in range(len(model.layers))},
+    **{f"model.layers.{l}.self_attn.{x}_proj.bias": f"layers.{l}.attention.w{x}.bias" for x in ["q", "k", "v", "o"] for l in range(len(model.layers))},
     **{f"model.layers.{l}.post_attention_layernorm.weight": f"layers.{l}.ffn_norm.weight" for l in range(len(model.layers))},
     **{f"model.layers.{l}.mlp.{x}_proj.weight": f"layers.{l}.feed_forward.w{y}.weight" for x, y in {"gate": "1", "down": "2", "up": "3"}.items() for l in range(len(model.layers))},
     "model.norm.weight": "norm.weight",
@@ -198,9 +200,9 @@ def convert_from_huggingface(weights:Dict[str, Tensor], model: Transformer, n_he
     if ".rotary_emb." in k: continue
     v = v.to(Device.DEFAULT)
     if "model.layers" in k:
-      if "q_proj" in k:
+      if "q_proj" in k and permute_layers:
         v = permute(v, n_heads)
-      elif "k_proj" in k:
+      elif "k_proj" in k and permute_layers:
         v = permute(v, n_kv_heads)
     sd[keymap[k]] = v
   return sd

extra/onnx.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
-from typing import List, Dict, Union
-import importlib
-from functools import lru_cache
+from typing import List, Dict, Union, Callable, Any, Sequence
+import importlib, functools
 import numpy as np
-from tinygrad import Tensor, dtypes, Device
-from tinygrad.tensor import _to_np_dtype
-from tinygrad.helpers import getenv, DEBUG, CI, OSX
-from tinygrad.dtype import ConstType, DType
+from tinygrad import Tensor, dtypes
+from tinygrad.helpers import getenv, DEBUG, all_same
+from tinygrad.dtype import DType, ConstType
 from tinygrad.device import is_dtype_supported
-from onnx import AttributeProto, ModelProto, TensorProto, TypeProto
+from onnx import AttributeProto, ModelProto, TensorProto, ValueInfoProto
 try:
   from onnx.helper import tensor_dtype_to_np_dtype
 except ImportError:
@@ -17,177 +15,134 @@ except ImportError:
   def tensor_dtype_to_np_dtype(tensor_dtype:int) -> np.dtype: return TENSOR_TYPE_TO_NP_TYPE[tensor_dtype]
 
 cache_misses = 0
-@lru_cache(None)
-def _cached_to_python_const(t:Tensor, tobytes): return t.data().tobytes() if tobytes else t.tolist()
+@functools.lru_cache(None)
+def _cached_to_python_const(t:Tensor):
+  if t.dtype is dtypes.uint8: return t.data().tobytes()
+  if 0 in t.shape: return []
+  return t.tolist()
 
 # Tensor -> python value cache for parameters
-def to_python_const(t, tobytes=False) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]:
+def to_python_const(t) -> Union[List[ConstType], List[bytes], Union[ConstType, bytes]]:
   if not isinstance(t, Tensor): return t
   global cache_misses
-  ret = _cached_to_python_const(t, tobytes)
+  ret = _cached_to_python_const(t)
   if (info := _cached_to_python_const.cache_info()).misses > cache_misses and DEBUG >= 3:
-    print(f"Cache miss for {t}, {tobytes=}")
+    print(f"Cache miss for {t}")
     cache_misses = info.misses
   return ret
 
-# src: onnx/mapping.py  https://onnx.ai/onnx/api/mapping.html#l-mod-onnx-mapping
-# not supported: STRING = 8 COMPLEX64 = 14, COMPLEX128 = 15, UINT4 = 21, INT4 = 22
-# TODO: use dtypes.float16 for FLOAT16
-DTYPE_MAP: Dict[TensorProto.DataType, DType] = {
-  TensorProto.FLOAT:dtypes.float, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8, TensorProto.UINT16:dtypes.uint16,
-  TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64, TensorProto.BOOL:dtypes.bool,
-  TensorProto.FLOAT16:dtypes.float, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32, TensorProto.UINT64:dtypes.uint64,
-  TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float, TensorProto.FLOAT8E4M3FNUZ:dtypes.float,
-  TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float
+# TODO: use real float16
+# src: onnx/mapping.py
+DTYPE_MAP: Dict[TensorProto.DataType | int, DType] = {
+  TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8,
+  TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64,
+  TensorProto.BOOL:dtypes.bool, TensorProto.FLOAT16:dtypes.float32, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32,
+  TensorProto.UINT64:dtypes.uint64, TensorProto.BFLOAT16:dtypes.bfloat16, TensorProto.FLOAT8E4M3FN:dtypes.float,
+  TensorProto.FLOAT8E4M3FNUZ:dtypes.float, TensorProto.FLOAT8E5M2:dtypes.float, TensorProto.FLOAT8E5M2FNUZ:dtypes.float
 }
+def dtype_parse(onnx_dtype: TensorProto.DataType | int) -> DType:
+  if onnx_dtype not in DTYPE_MAP: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported")
+  return DTYPE_MAP[onnx_dtype] if is_dtype_supported(DTYPE_MAP[onnx_dtype]) else dtypes.float
+
+# src: onnx/onnx_ml_pb2.pyi
+ATTRIBUTE_MAP: Dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = {
+  AttributeProto.FLOAT: lambda a: float(a.f), AttributeProto.INT: lambda a: int(a.i),
+  AttributeProto.STRING: lambda a: a.s.decode("utf-8"), AttributeProto.TENSOR: lambda a: buffer_parse(a.t),
+  AttributeProto.FLOATS: lambda a: tuple(float(x) for x in a.floats), AttributeProto.INTS: lambda a: tuple(int(x) for x in a.ints),
+  AttributeProto.STRINGS: lambda a: tuple(x.decode("utf-8") for x in a.strings)
+}
+def attribute_parse(onnx_attribute: AttributeProto):
+  if onnx_attribute.type not in ATTRIBUTE_MAP:
+    raise NotImplementedError(f"attribute with type {AttributeProto.AttributeType.Name(onnx_attribute.type)} is not supported")
+  return ATTRIBUTE_MAP[onnx_attribute.type](onnx_attribute)
+
+def buffer_parse(inp: TensorProto) -> Tensor:
+  if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data):
+    return Tensor(dat, dtype=dtype_parse(inp.data_type), requires_grad=False).reshape(tuple(inp.dims))
+  if len(inp.raw_data) > 0:
+    return Tensor(np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).copy().reshape(tuple(inp.dims)),
+                  dtype=dtype_parse(inp.data_type), requires_grad=False)
+  raise NotImplementedError(f"buffer with data type {TensorProto.DataType.Name(inp.data_type)} is not supported")
 
 onnx_ops = importlib.import_module('extra.onnx_ops')
-
 ONNXLIMIT = getenv("ONNXLIMIT", -1)
 
 def get_run_onnx(onnx_model: ModelProto):
-  def type_parse(type_proto: TypeProto):
-    ret = []
-    while True:
-      attr = type_proto.WhichOneof('value')
-      if attr == 'tensor_type':
-        if "dim_value" not in type_proto.tensor_type.shape.dim.__dir__(): return () # variable type, unable to determine shape
-        elif not ret:
-          return tuple([x.dim_value for x in type_proto.tensor_type.shape.dim])
-        else:
-          ret.extend([(x.dim_value,) for x in type_proto.tensor_type.shape.dim])
-          return tuple(ret)
-      elif attr == 'sequence_type':
-        type_proto = getattr(type_proto, attr).elem_type
-        ret.append(1)
-      elif attr == 'optional_type': type_proto = getattr(type_proto, attr).elem_type
-      elif attr == 'map_type': raise NotImplementedError(f"map_type is not implemented: {type_proto}")
-      elif attr == 'opaque_type': raise NotImplementedError(f"opaque_type is not implemented: {type_proto}")
-      elif attr == 'sparse_tensor_type': raise NotImplementedError(f"sparse_tensor_type is not implemented: {type_proto}")
-      else: raise AttributeError(f"unknown attr: {attr}, {type_proto}")
-
-  def buffer_parse(inp: TensorProto) -> Tensor:
-    if inp.data_type not in DTYPE_MAP:
-      raise NotImplementedError(f"data type not supported {inp.name} {inp.dims} {inp.data_type}")
-    dtype = DTYPE_MAP[inp.data_type] if is_dtype_supported(DTYPE_MAP[inp.data_type]) else dtypes.float32
-    if dat := list(inp.float_data) or list(inp.int32_data) or list(inp.int64_data):
-      return Tensor(dat, dtype=dtype, requires_grad=False).reshape(tuple(inp.dims))
-    if len(inp.raw_data) > 0:
-      data = np.frombuffer(inp.raw_data, dtype=tensor_dtype_to_np_dtype(inp.data_type)).astype(_to_np_dtype(dtype)).copy()
-      return Tensor(data.reshape(tuple(inp.dims)), requires_grad=False)
-    return Tensor(None, requires_grad=False)
-
-  def attribute_parse(a: AttributeProto) -> float | int | str | Tensor | tuple[float] | tuple[int]:
-    # TODO: this is not complete, see onnx/onnx_ml_pb2.pyi for a complete list
-    if a.type == AttributeProto.FLOAT: return float(a.f)
-    elif a.type == AttributeProto.INT: return int(a.i)
-    elif a.type == AttributeProto.STRING: return a.s.decode("utf-8")
-    elif a.type == AttributeProto.TENSOR: return buffer_parse(a.t) # TENSOR
-    elif a.type == AttributeProto.FLOATS: return tuple(float(x) for x in a.floats)
-    elif a.type == AttributeProto.INTS: return tuple(int(x) for x in a.ints)
-    elif a.type == AttributeProto.STRINGS: return tuple(x.decode("utf-8") for x in a.strings)
-    elif a.type == AttributeProto.GRAPH: raise NotImplementedError(f"graph not implemented: {a.g}\n likely an OP requiring control flow")
-    else: raise RuntimeError(f"can't parse {a.type} {a}")
-
-  tensors: Dict[str, Tensor] = {}
-
-  # get weights and biases
-  for inp in onnx_model.graph.initializer:
-    tensors[inp.name] = buffer_parse(inp)
-
-  # preparse the attributes
-  attribute_dict = {}
-  domain = ""
-  for num,n in enumerate(onnx_model.graph.node):
-    attribute_dict[num] = {x.name:attribute_parse(x) for x in n.attribute}
-    if n.domain: domain = n.domain
+  # model initialization data
+  model_parameters = {inp.name:buffer_parse(inp) for inp in onnx_model.graph.initializer}
+  model_attributes = {num:{x.name:attribute_parse(x) for x in n.attribute} for num,n in enumerate(onnx_model.graph.node)}
 
+  # model descriptions
+  # TODO: need a better way of controlling training vs non-training
+  is_onnx_preview_training = any(n.HasField("domain") and n.domain == "ai.onnx.preview.training" for n in onnx_model.graph.node)
   onnx_model_version = onnx_model.opset_import[0].version
 
+  # mapping from onnx ops to tensor.py ops
+  tensor_methods = {
+    op:op.lower() for op in ("Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Asin", "Acos", "Atan",
+    "Relu", "Sigmoid", "MatMul", "Floor", "Ceil", "IsInf", "IsNaN", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh",
+    "Softsign", "Asinh", "Acosh", "Atanh",  "Elu", "Celu", "Selu", "Xor", "Round", "Erf")
+  }
+
+  # src: https://onnx.ai/onnx/repo-docs/IR.html#input-output-data-types
+  # parses and validates inputs based on their shape and dtype specified by model
+  def prepare_input(user_input:Any, model_input:ValueInfoProto):
+    type_proto = model_input.type
+    if type_proto.HasField("optional_type"):
+      if user_input is None: return Tensor(None)
+      type_proto = type_proto.optional_type.elem_type
+    if type_proto.HasField("sequence_type"):
+      if not isinstance(user_input, Sequence): raise RuntimeError(f"{model_input.name} received {user_input}, expected sequence type")
+      dtype = dtype_parse(type_proto.sequence_type.elem_type.tensor_type.elem_type)
+      sequence = [Tensor(i, dtype=dtype, requires_grad=is_onnx_preview_training) if not isinstance(i, Tensor) else i for i in user_input]
+      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"shapes for {model_input.name} must be homogeneous")
+      # TODO: need true float16 for dtype checking
+      # if not all(t.dtype is dtype for t in sequence): raise RuntimeError(f"{model_input.name} received wrong dtype, expected {dtype}")
+      return sequence
+    if type_proto.HasField("tensor_type"):
+      dtype = dtype_parse(type_proto.tensor_type.elem_type)
+      tensor = Tensor(user_input, dtype=dtype, requires_grad=is_onnx_preview_training) if not isinstance(user_input, Tensor) else user_input
+      # TODO: need true float16 for dtype checking
+      # if dtype is not tensor.dtype: raise RuntimeError(f"{model_input.name} received dtype {inp.dtype}, expected {dtype}")
+      for d,onnx_dim in enumerate(type_proto.tensor_type.shape.dim):
+        # NOTE: dim is a variable dimension when `dim_param` is specified, e.g. dim {dim_param: "N"} is a variable dim
+        if onnx_dim.dim_param is None and onnx_dim.dim_value != user_input.shape[d]:
+          raise RuntimeError(f"{model_input.name} received value {user_input.shape[d]} on dim {d}, expected {onnx_dim.dim_value}")
+      return tensor
+    type_field_names = [field.name for field,_ in type_proto.ListFields()]
+    raise NotImplementedError(f"{model_input.name} with {type_field_names=} is not supported")
+
   def run_onnx(inputs={}, debug=0):
     debug = getenv("DEBUGONNX") or debug
-    input_tensors: Dict[str,Tensor|List[Tensor]] = {}
     intermediate_tensors: Dict[str,Tensor] = {}
-    output_tensor_names = [x.name for x in onnx_model.graph.output]
 
-    # get inputs
+    input_tensors: Dict[str, Tensor | List[Tensor]] = {}
     for model_input in onnx_model.graph.input:
-      name = model_input.name
-      if name in tensors: continue
-      shape = type_parse(model_input.type)
-      if name in inputs:
-        if isinstance(inputs[name], Tensor):
-          input_tensors[name] = inputs[name]
-        elif isinstance(inputs[name], list):
-          input_tensors[name] = [Tensor(i, requires_grad=False) for i in inputs[name]]
-        elif domain == "ai.onnx.preview.training": # not sure if in real use the domain is "ai.onnx.preview.training"
-          input_tensors[name] = Tensor(inputs[name], requires_grad=True) # TODO there isn't a good way to parse which inp requires_grad, some are manually turned off in optimizer ops
-        else:
-          input_tensors[name] = Tensor(inputs[name], requires_grad=False)
-        if shape: # if only input_tensor is not variable type
-          ts = input_tensors[name]
-          input_shape = ts.shape if isinstance(ts, Tensor) else (1, *[i.shape for i in ts])
-          assert input_shape == shape, f"wrong shape for input {name}, {input_shape} isn't {shape}"
-      else:
-        raise RuntimeError(f"no data for {name} with shape {shape}")
+      if model_input.name in inputs: input_tensors[model_input.name] = prepare_input(inputs[model_input.name], model_input)
+      elif model_input.name not in model_parameters: raise RuntimeError(f"Please provide input data for {model_input.name}")
 
     def fetch_tensor(x: str):
-      if x in tensors: return tensors[x]
+      if x in model_parameters: return model_parameters[x]
       if x in intermediate_tensors: return intermediate_tensors[x]
       if x != "": return input_tensors[x]
       return None
 
     for num,n in enumerate(onnx_model.graph.node):
-      inp: List[Tensor] = []
-      if debug >= 3: print("inputs:")
-      for x in n.input:
-        t = fetch_tensor(x)
-        if debug >= 3: print(f"\t{x} - {t}")
-        inp.append(t)
-      opt: Dict = attribute_dict[num]
-      if debug >= 1: print(f"{num}: op {n.op_type} shape {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}")
+      inp = [fetch_tensor(x) for x in n.input]
+      opt = model_attributes[num]
+
+      if debug >= 1: print(f"{num}: op \"{n.op_type}\" input shapes {[x.shape if isinstance(x, Tensor) else x for x in inp]} opt {opt}")
+      if debug >= 3: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {t}" for i,(x,t) in enumerate(zip(n.input, inp))))
+
+      if n.op_type in tensor_methods:
+        ret = getattr(Tensor, tensor_methods[n.op_type])(*inp, **opt)
 
       # NOTE some ops live here because they require access to some local variables
-      # have to use n.output for cases when num_outputs is absent
-      if n.op_type in onnx_ops.tensor_methods:
-        ret = getattr(Tensor, n.op_type.lower())(*inp, **opt)
       elif n.op_type == "Split":
-        axis = opt.get("axis", 0)
-        split = None if len(inp) == 1 else to_python_const(inp[1])
-        if split is None:
-          split = [inp[0].shape[axis] // len(n.output)] * len(n.output)
-          for i in range(inp[0].shape[axis] % len(n.output)):
-            split[i] += 1
-        i, ret = 0, []
-        arg = [None] * inp[0].ndim
-        for s in split:
-          arg[axis] = (i,i+s)
-          ret.append(inp[0].shrink(arg=tuple(arg)))
-          i = i+s
-        ret = tuple(ret)
-
-      # need to check onnx_model_version
-      elif n.op_type == "Slice":
-        if onnx_model_version < 10:
-          axes, ends, starts, steps = list(opt.get("axes", range(inp[0].ndim))), list(opt["ends"]), list(opt["starts"]), [1]*inp[0].ndim
-        else:
-          starts, ends = inp[1:3]
-          axes = list(range(inp[0].ndim)) if len(inp) <= 3 else to_python_const(inp[3].cast(dtypes.int32))
-          steps = inp[4].cast(dtypes.int32).tolist() if len(inp) > 4 else [1]*inp[0].ndim
-          starts, ends = to_python_const(starts), to_python_const(ends)
-        arg = [(0,x,1) for x in inp[0].shape]
-        for i, axis in enumerate(axes):
-          axis = int(axis) + inp[0].ndim if axis < 0 else int(axis)
-          if starts[i] < 0: starts[i] += inp[0].shape[axis]
-          if ends[i] < 0: ends[i] += inp[0].shape[axis]
-          starts[i], ends[i] = max(0, min(starts[i], inp[0].shape[axis])), max(0, min(ends[i], inp[0].shape[axis]))
-          if starts[i] > ends[i] and steps[i] >= 0: steps[i] = -steps[i]
-          arg[axis] = (starts[i], ends[i], steps[i])
-        new_shape = tuple((s, e) if st > 0 else (e+1, s+1) for s, e, st in arg)
-        if any(s==e for s,e in new_shape): ret = inp[0].shrink(new_shape)
-        else: ret = inp[0][tuple([slice(s,e,st) for s,e,st in arg])]
-
-      # need to call backward on intermediate_tensors
+        axis, n_outputs  = opt.get('axis', 0), opt.get('num_outputs') or len(n.output)
+        sz = inp[0].shape[axis]
+        sizes = to_python_const(inp[1]) if len(inp) == 2 else [sz // n_outputs + (1 if i < sz % n_outputs else 0) for i in range(n_outputs)]
+        ret = inp[0].split(sizes, axis)
       elif n.op_type == "Gradient":
         assert len(opt["xs"]) == len(inp), f"len(opt['xs']):{len(opt['xs'])}, len(inp):{len(inp)} output and input has to match"
         y = opt["y"]
@@ -208,16 +163,12 @@ def get_run_onnx(onnx_model: ModelProto):
         print("UNSUPPORTED", n.op_type, n.input, n.output)
         raise NotImplementedError(f"op_type {n.op_type} not supported")
 
+      # finalization after running the op
       if not isinstance(ret, tuple): ret = (ret, )
-      assert len(n.output) <= len(ret), f"expected output size must be less than {len(ret)}, it's {n.output}"
-      if debug >= 2: print([x.shape if isinstance(x, Tensor) else None for x in ret])
-      if debug >= 2: print("outputs:")
-      for i in range(len(n.output)):
-        if debug >= 2: print(f"\t{n.output[i]} - {ret[i]}")
-        intermediate_tensors[n.output[i]] = ret[i]
-      if num == ONNXLIMIT:
-        output_tensor_names = n.output
-        break
+      if len(n.output) > len(ret): raise RuntimeError(f"expected output size must be less than {len(ret)}, it's {n.output}")
+      for i in range(len(n.output)): intermediate_tensors[n.output[i]] = ret[i]
+      if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{n.output[i]} - {ret[i]}" for i in range(len(n.output))))
 
-    return {outp:intermediate_tensors[outp] for outp in output_tensor_names}
+      if num == ONNXLIMIT: return {name:intermediate_tensors[name] for name in n.output}
+    return {x.name:intermediate_tensors[x.name] for x in onnx_model.graph.output}
   return run_onnx

extra/onnx_ops.py

@@ -1,15 +1,11 @@
 import functools, io, math
-from typing import Union, Tuple, Optional, List, Any
+from typing import Union, Tuple, Optional, List, Any, cast
 from tinygrad.tensor import Tensor, _broadcast_shape
 from tinygrad.dtype import ImageDType, dtypes
 from tinygrad.helpers import prod, flatten
-from extra.onnx import DTYPE_MAP, to_python_const
+from extra.onnx import dtype_parse, to_python_const
 import numpy as np
 
-tensor_methods = {"Neg", "Reciprocal", "Pow", "Sqrt", "Sign", "Abs", "Exp", "Log", "Mish", "Sin", "Cos", "Tan", "Relu", "Sigmoid", "MatMul",
-                  "Floor", "Ceil", "Softplus", "HardSwish", "Where", "Mul", "Sinh", "Cosh", "Tanh", "Softsign", "Asinh", "Acosh", "Atanh",
-                  "Elu", "Celu", "Xor", "Round", "Erf"}
-
 # **************** Free Ops ****************
 
 def Identity(x: Tensor): return x
@@ -21,12 +17,16 @@ def LessOrEqual(x:Tensor,y:Tensor): return x <= y
 def Greater(x:Tensor,y:Tensor): return x > y
 def GreaterOrEqual(x:Tensor,y:Tensor): return x >= y
 def Equal(x:Tensor,y:Tensor): return x == y
+def BitwiseNot(x:Tensor): return ~x
+def BitwiseOr(x:Tensor, y:Tensor): return x | y
+def BitwiseAnd(x:Tensor, y:Tensor): return x & y
+def BitwiseXor(x:Tensor, y:Tensor): return x ^ y
 def Max(*data_0): return functools.reduce(Tensor.maximum, data_0)
 def Min(*data_0): return functools.reduce(Tensor.minimum, data_0)
 def Sum(*data_0): return functools.reduce(Tensor.add, data_0)
 def Mean(*data_0): return Sum(*data_0) / len(data_0)
 # NOTE: does not support saturate
-def Cast(x: Tensor, to: int, saturate=1): return x.cast(DTYPE_MAP[to])
+def Cast(x: Tensor, to: int, saturate=1): return x.cast(dtype_parse(to))
 def CastLike(x: Tensor, target_type: Tensor, saturate=1): return x.cast(target_type.dtype)
 
 # **************** Simple Ops ****************
@@ -44,7 +44,6 @@ def Constant(value:Optional[Tensor]=None, value_float=None, value_floats=None, v
 
 def HardSigmoid(x: Tensor, alpha=0.2, beta=0.5): return (alpha*x + beta).clip(0, 1)
 def Gelu(x:Tensor, approximate=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
-def Selu(X: Tensor, alpha=1.67326319217681884765625, gamma=1.05070102214813232421875): return gamma * (X.relu() - (-alpha*X.exp()+alpha).relu())
 def PRelu(X:Tensor, slope:Tensor):
   slope = slope[0] if slope.shape[-1] != X.shape[-1] else slope # HACK OnnxBackendPyTorchConvertedModelTest HAS WEIRD SLOPE WHERE IT'S [0.25, 0.25, 0.25] FOR ANY X.SHAPE
   return (X > 0).where(X, X * slope)
@@ -88,26 +87,18 @@ def And(x:Tensor, y:Tensor): return (x==y).where(x, False)
 def Or(x:Tensor, y:Tensor): return (x==y).where(x, True)
 def Not(x:Tensor): return x.logical_not()
 
-def Asin(x): return Atan(x / (1 - x * x).sqrt())
-def Acos(x: Tensor):
-  negate = (x < 0)
-  x = x.abs()
-  ret = ((((-0.0187293 * x) + 0.0742610)*x - 0.2121144) * x + 1.5707288) * (1.0 - x).sqrt()
-  ret = ret - 2 * negate * ret
-  return negate * math.pi + ret
-def Atan(y: Tensor):
-  t1 = y.abs()
-  t3 = (1 > t1).where(t1, t1.reciprocal())
-  t4 = t3 * t3
-  t0 = ((((-0.013480470 * t4 + 0.057477314) * t4 - 0.121239071) * t4 + 0.195635925) * t4 - 0.332994597) * t4 + 0.999995630
-  t3 = t0 * t3
-  t3 = (t1 > 1).where(1.570796327 - t3, t3)
-  return y.sign() * t3
-
 def Trilu(x: Tensor, k: Union[Tensor, int]=0, upper=1):
   k = to_python_const(k) if isinstance(k, Tensor) else 0 # onnx passes k as a tensor int64 with one element, default is 0
   return x.triu(k) if upper else x.tril(k)
 
+def Slice(data: Tensor, starts:Tensor, ends:Tensor, axes:Optional[Tensor]=None, steps:Optional[Tensor]=None):
+  if axes is None: axes = list(range(data.ndim))
+  if steps is None: steps = [1] * data.ndim
+  starts, ends, axes, steps = (to_python_const(x) for x in (starts, ends, axes, steps))
+  slices = [slice(0,x,1) for x in data.shape]
+  for i, axis in enumerate(axes): slices[axis] = slice(starts[i], ends[i], steps[i])
+  return data[tuple(slices)]
+
 def Squeeze(data: Tensor, axes):
   if isinstance(axes, Tensor): axes = to_python_const(axes)
   axes = [data._resolve_dim(x) for x in axes]
@@ -143,16 +134,11 @@ def Gemm(A: Tensor, B: Tensor, C: Tensor=None, alpha=1.0, beta=1.0, transA=0, tr
 def Einsum(*Inputs: List[Tensor], equation): return Tensor.einsum(equation, Inputs)
 
 def CumSum(X:Tensor, axis:Tensor, exclusive=0, reverse=0):
-  axis = to_python_const(axis)
-  if axis < 0: axis += X.ndim
+  if (axis := to_python_const(axis)) < 0: axis += X.ndim
   if reverse: X = X.flip(axis)
-  if exclusive:
-    pad_arg, shrink_arg = [None] * X.ndim, [None] * X.ndim
-    pad_arg[axis] = (1, 0)
-    shrink_arg[axis] = (0, X.shape[axis])
-    X = X.pad(tuple(pad_arg)).shrink(tuple(shrink_arg))
-  if reverse: return X.cumsum(axis).flip(axis)
-  return X.cumsum(axis)
+  if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
+                      .shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
+  return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)
 
 # TODO: this is copied from tinygrad/nn/__init__.py
 # spatial is from opset 7 and has since been removed
@@ -225,34 +211,13 @@ def _auto_pad(X: Tensor, auto_pad, strides, kernel_shape, dilations):
     return pad_shape[::2] + pad_shape[1::2] if auto_pad == "SAME_UPPER" else pad_shape[1::2] + pad_shape[::2]
   raise NotImplementedError(f"auto_pad={auto_pad} not implemented")
 
-def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Tensor=None, axes: Tensor=None, mode="constant", value: float=0.):
-  constant_value = value if constant_value is None else float(to_python_const(constant_value))
-  seq_pads = list(pads) if isinstance(pads, tuple) else to_python_const(pads)
-  seq_pads = [math.ceil(i) for i in seq_pads]
-  seq_axes = to_python_const(axes) if axes is not None else None
-  base_shape = x.shape
-  pads = _format_padding(seq_pads, ndims=len(x.shape), axes=seq_axes)
-  if mode == "wrap":
-    repeat_args = [math.ceil(dim[0]/sh) + math.ceil(dim[1]/sh) + 1 for dim, sh in zip(pads, base_shape)]
-    new_shape = [s*r for s,r in zip(base_shape, repeat_args)]
-    shrink_args = [(sh-dim[0]%sh if dim[0]%sh != 0 else 0, nsh-(sh-dim[1]%sh if dim[1]%sh != 0 else 0)) for dim, sh, nsh in zip(pads, base_shape, new_shape)]
-    return x.repeat(tuple(repeat_args)).shrink(tuple(shrink_args))
-  if mode == "reflect":
-    for i,s in enumerate(x.shape):
-      if pads[i] != (0,0):
-        xL = x.flip(i).shrink(tuple((s-pads[i][0]-1, s_-1) if i_ == i else None for i_,s_ in enumerate(x.shape)))
-        xR = x.flip(i).shrink(tuple((1, pads[i][1]+1) if i_ == i else None for i_ in range(x.ndim)))
-        x = xL.cat(x, xR, dim=i)
-    return x
-  if mode == "edge":
-    for i,s in enumerate(x.shape):
-      if pads[i] != (0,0):
-        xL = x.shrink(tuple((0,1) if i_ == i else None for i_ in range(x.ndim))).expand([pads[i][0] if i_ == i else None for i_ in range(x.ndim)])
-        xR = x.shrink(tuple((s_-1, s_) if i_ == i else None for i_,s_ in enumerate(x.shape))).expand([pads[i][1] if i_ == i else None for i_ in range(x.ndim)])
-        x = xL.cat(x, xR, dim=i)
-    return x
-  if mode == "constant":
-    return _padded(x, seq_pads, axes=seq_axes, constant_value=constant_value)
+# (x1_begin, x2_begin, ..., x1_end, x2_end, ...) -> (..., x2_start, x2_end, x1_start, x1_end)
+def _onnx_pads_to_pad2d_pads(pads): return flatten(reversed(list((pB, pE) for pB, pE in zip(pads, pads[len(pads)//2:]))))
+def Pad(x: Tensor, pads: Union[Tensor, Tuple[int, ...]], constant_value: Optional[Tensor]=None, axes: Optional[Tensor]=None, mode="constant", value=0):
+  pads, value, axes = to_python_const(pads), to_python_const(constant_value) or value or 0, to_python_const(axes) or list(range(x.ndim))
+  real_pads = [0] * (x.ndim*2)
+  for i,axis in enumerate(axes): real_pads[axis%x.ndim], real_pads[axis%x.ndim+x.ndim] = pads[i], pads[i+len(axes)]
+  return x.pad(padding=_onnx_pads_to_pad2d_pads(to_python_const(real_pads)), mode={"edge":"replicate", "wrap":"circular"}.get(mode, mode), value=value)
 
 def AveragePool(X: Tensor, kernel_shape, auto_pad="NOTSET", ceil_mode=0, count_include_pad=0, dilations=1, pads=None, strides=1):
   pixel_axes = tuple(range(2, X.ndim))
@@ -312,15 +277,9 @@ def ConvTranspose(X: Tensor, W: Tensor, B:Optional[Tensor]=None, auto_pad="NOTSE
   return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads if pads is not None else 0, output_padding=output_padding)
 
 def DepthToSpace(X:Tensor, blocksize:int, mode:str="DCR"):
-  b, c, h, w = X.shape
-  if mode == "DCR":
-    return X.reshape(b, blocksize, blocksize, c // (blocksize**2), h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, c // (blocksize**2), h * blocksize, w * blocksize)
-  elif mode == "CRD":
-    return X.reshape(b, c // (blocksize ** 2), blocksize, blocksize, h, w).permute(0, 1, 4, 2, 5, 3).reshape(b, c // (blocksize ** 2), h * blocksize, w * blocksize)
-
+  return X.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)" if mode=="CRD" else "b (h1 w1 c) h w -> b c (h h1) (w w1)", h1=blocksize, w1=blocksize)
 def SpaceToDepth(X:Tensor, blocksize:int):
-  b, c, h, w = X.shape
-  return X.reshape(b, c, h // blocksize, blocksize, w // blocksize, blocksize).permute(0, 3, 5, 1, 2, 4).reshape(b, c * (blocksize**2), h // blocksize, w // blocksize)
+  return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)
 
 # Reimplemented here because you need legacy RNG for passing ONNX tests.
 def Dropout(data: Tensor, ratio=0.5, training_mode=False, seed=None):
@@ -333,44 +292,17 @@ def Dropout(data: Tensor, ratio=0.5, training_mode=False, seed=None):
   return data * mask * (1/(1.0 - ratio)), mask
 
 def LRN(x: Tensor, size, alpha=1e-4, beta=0.75, bias=1.0):
-  bs, c, iy, ix = x.shape
-  return x / x.mul(x).reshape(bs,1,c,iy*ix).pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1).reshape(bs,c,iy,ix).mul(alpha).add(bias).pow(beta)
+  pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
+  return x / (pooled_x.reshape(x.shape) * alpha + bias).pow(beta)
 
-def MeanVarianceNormalization(x: Tensor, axis=(0, 2, 3)):
-  mean = x.mean(axis, keepdim=True)
-  std = x.std(axis, keepdim=True, correction=0)
-  return (x - mean) / (std + 1e-9)
+def MeanVarianceNormalization(x: Tensor, axis=(0, 2, 3)): return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
 
 def NegativeLogLikelihoodLoss(x: Tensor, target: Tensor, weight=None, ignore_index=None, reduction="mean"):
-  N, C, i_shape = x.shape[0], x.shape[1], x.shape
-  t_shape = target.shape
-  if len(x.shape) != 3:
-    x = x.reshape((N, C, -1))
-    target = target.reshape((N, -1))
-  if weight is not None:
-    mask = target.unsqueeze(-1) == Tensor.arange(C).repeat((N, 1, 1))
-    weight = (mask * weight).sum(axis=-1)
-  if ignore_index is not None:
-    cond = target == ignore_index
-    weight = cond.where(0, weight) if weight is not None else cond.where(0, 1)
-  mask = target[:, None, :] == Tensor.arange(C).reshape([1, C] + [1]*(x.ndim -2))
-  loss = -(mask * x).sum(axis=1) * (1 if weight is None else weight)
-  if reduction == "mean": return loss.mean() if weight is None else loss.sum() / weight.sum()
-  if reduction == "sum": return loss.sum()
-  return loss.reshape(t_shape) if len(i_shape) != 3 else loss
+  return x.nll_loss(target, weight, ignore_index, reduction)
 
 def SoftmaxCrossEntropyLoss(scores: Tensor, labels: Tensor, weights=None, ignore_index=None, reduction="mean"):
-  _N, C, *s_dimensions = scores.shape
-  if ignore_index is not None: labels = (labels == ignore_index).where(C+1, labels)
-  mask = labels.unsqueeze(1) == Tensor.arange(C).reshape(1, C, *[1]*len(s_dimensions))
-  y = scores.log_softmax(axis=1)
-  loss = (mask * -y).sum(1)
-  if weights is not None:
-    weights = weights[labels, ...]
-    loss = loss * weights
-  if reduction == "mean": loss = loss.sum() / ((loss != 0).sum() if weights is None else weights.sum())
-  elif reduction == "sum": loss = loss.sum()
-  return loss, y
+  log_probs = scores.log_softmax(1)
+  return log_probs.nll_loss(labels, weights, ignore_index, reduction), log_probs
 
 def ArrayFeatureExtractor(x: Tensor, indices: Tensor): return x[..., indices]
 
@@ -384,114 +316,81 @@ def Gather(x: Tensor, indices: Tensor, axis=0):
     return x.shrink(arg=tuple(args[0])).cat(*[x.shrink(arg=tuple(arg)) for arg in args[1:]], dim=axis).reshape(ret_shape)
   # NOTE faster gather, fixed number of kernels, but exceeds limited kernels for openpilot
   return x[tuple([slice(None) if i != axis else indices for i in range(x.ndim)])]
+def Scatter(*args, **kwargs): return ScatterElements(*args, **kwargs) # deprecated
 
+def ScatterElements(x: Tensor, indices: Tensor, updates: Tensor, axis=0, reduction:Optional[str]=None):
+  if reduction in {"min", "max"}: raise NotImplementedError("min and max reduction not supported")
+  indices = (indices < 0).where(x.shape[axis], 0) + indices
+  return x.scatter(axis, indices, updates, reduction)
 def GatherElements(x: Tensor, indices: Tensor, axis):
   indices = (indices < 0).where(x.shape[axis], 0) + indices
   return x.gather(axis, indices)
 
-# TODO clean this up, it's taking the longest in CI
 def Resize(X:Tensor, roi=None, scales=None, sizes=None, antialias=0, axes=None, coordinate_transformation_mode='half_pixel',
-           cubic_coeff_a=-0.75, exclude_outside=0, extrapolation_value=0.0, keep_aspect_ratio_policy='stretch',
-           mode='nearest', nearest_mode='round_prefer_floor'):
-  def _nearest_gather(X: Tensor, x_out, y_out): return X[:,:,y_out,:][:,:,:,x_out]
-  def _nearest_mode(x_resized: Tensor, nearest_mode: str, x_len):
-    if nearest_mode == "round_prefer_floor": ret = (x_resized - 0.5).ceil()
-    elif nearest_mode == "round_prefer_ceil": ret = (x_resized + 0.5).floor()
-    elif nearest_mode == "floor": ret = x_resized.floor()
-    elif nearest_mode == "ceil": ret = x_resized.ceil()
-    return ret.cast(dtypes.int32).clip(0, x_len-1)
-  def _coordinate_transformation(x_out, y_out, output_shape, scales_, roi=None):
-    if coordinate_transformation_mode == "half_pixel":
-      x_out = (x_out + 0.5) / scales_[-1] - 0.5
-      y_out = (y_out + 0.5) / scales_[-2] - 0.5
-    elif coordinate_transformation_mode == "align_corners":
-      x_out = x_out * (X.shape[-1] - 1) / (output_shape[-1] - 1)
-      y_out = y_out * (X.shape[-2] - 1) / (output_shape[-2] - 1)
-    elif coordinate_transformation_mode == "asymmetric":
-      x_out = x_out / scales_[-1]
-      y_out = y_out / scales_[-2]
-    elif coordinate_transformation_mode == "half_pixel_symmetric":
-      x_out = X.shape[-1] / 2 * (1 - int(output_shape[-1]) / output_shape[-1]) + (x_out + 0.5) / scales_[-1] - 0.5
-      y_out = X.shape[-2] / 2 * (1 - int(output_shape[-2]) / output_shape[-2]) + (y_out + 0.5) / scales_[-2] - 0.5
-    elif coordinate_transformation_mode == "pytorch_half_pixel":
-      x_out = (x_out + 0.5) / scales_[-1] - 0.5 if output_shape[-1] > 1 else Tensor([0])
-      y_out = (y_out + 0.5) / scales_[-2] - 0.5 if output_shape[-2] > 1 else Tensor([0])
-    elif coordinate_transformation_mode == "tf_crop_and_resize":
-      x_out = roi[-1][0] * (X.shape[-1] - 1) + x_out * ((roi[-1][1] - roi[-1][0]) * (X.shape[-1] - 1) / (output_shape[-1] - 1)) if output_shape[-1] > 1 else Tensor([0.5 * (roi[-1][0] + roi[-1][1]) * (X.shape[-1] - 1)])
-      y_out = roi[-2][0] * (X.shape[-2] - 1) + y_out * ((roi[-2][1] - roi[-2][0]) * (X.shape[-2] - 1) / (output_shape[-2] - 1)) if output_shape[-2] > 1 else Tensor([0.5 * (roi[-2][0] + roi[-2][1]) * (X.shape[-2] - 1)])
-    return x_out.clip(0, X.shape[-1]-1), y_out.clip(0, X.shape[-2]-1)
-  if roi is not None:
-    roi = to_python_const(roi)
-    roi = [(st,ed) for st, ed in zip(roi[:len(roi)//2], roi[len(roi)//2:])]
-    roi_ = [(1,1)] * 4
-    if axes is not None:
-      for a,r in zip(axes, roi):
-        roi_[a] = r
-      roi = roi_
-  if scales is not None:
-    scales = to_python_const(scales)
-    if axes is not None:
-      scales_ = [1]*X.ndim
-      for a,s in zip(axes, scales):
-        scales_[a] = s
-      scales = scales_
-  elif sizes is not None:
-    sizes = to_python_const(sizes)
-    scales = []
-    if axes is not None:
-      sizes_ = [1]*X.ndim
-      for a,s in zip(axes, sizes):
-        sizes_[a] = s
-        scales.append(s/X.shape[a])
-      sizes = sizes_
-    else: scales = [si/xs for xs, si in zip(X.shape, sizes)]
-    if keep_aspect_ratio_policy == "not_larger":
-      scale = min(scales)
-      sizes = list(X.shape[:-2]) + [math.ceil(sh*scale) for sh in X.shape[-2:]]
-    elif keep_aspect_ratio_policy == "not_smaller":
-      scale = max(scales)
-      sizes = list(X.shape[:-2]) + [math.ceil(sh*scale) for sh in X.shape[-2:]]
-  output_shape = sizes if sizes else [math.floor(x*s) for x,s in zip(X.shape, scales)]
-  output_shape_ = sizes if sizes else [x*s for x,s in zip(X.shape, scales)]
-  scales_ = [os/xs for xs, os in zip(X.shape, output_shape)]
-  x_out = Tensor.arange(output_shape[-1], dtype=dtypes.default_float)
-  y_out = Tensor.arange(output_shape[-2], dtype=dtypes.default_float)
-  if mode == "nearest":
-    x_out, y_out = _coordinate_transformation(x_out, y_out, output_shape, scales_, roi)
-    x_out = _nearest_mode(x_out, nearest_mode, X.shape[-1])
-    y_out = _nearest_mode(y_out, nearest_mode, X.shape[-1])
-    return _nearest_gather(X, x_out, y_out)
-  if mode == "linear":
-    x_out, y_out = _coordinate_transformation(x_out, y_out, output_shape_, scales, roi)
-    ret = []
-    for y in to_python_const(y_out):
-      for x in to_python_const(x_out):
-        x_floor, y_floor = int(x), int(y)
-        y_shrink = (y_floor, math.ceil(y)+1)
-        x_shrink = (x_floor, math.ceil(x)+1)
-        corners = to_python_const(X.shrink((None, None, y_shrink, x_shrink)))[0][0]
+          cubic_coeff_a=-0.75, exclude_outside=0, extrapolation_value=0.0, keep_aspect_ratio_policy='stretch',
+          mode='nearest', nearest_mode='round_prefer_floor'):
+  def _apply_nearest_mode(index: Tensor, input_dim, mode: str):
+    if mode == "round_prefer_floor": index = (index - 0.5).ceil()
+    elif mode == "round_prefer_ceil": index = (index + 0.5).floor()
+    elif mode in ["floor", "ceil"]: index = getattr(index, mode)()
+    else: raise ValueError(f"invalid {nearest_mode=}")
+    return index.cast(dtypes.int32).clip(0, input_dim-1)
+  def _apply_transformation(index: Tensor, input_dim, scale_dim, roi_dim, sizes_frac, mode):
+    # TODO: needs more testing, not confident in this
+    # NOTE: their reference implementation differ from the implementation in their reference docs
+    # https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_resize.py
+    # https://github.com/onnx/onnx/blob/main/docs/Operators.md#Resize
+    output_dim = scale_dim * input_dim
+    if mode == "half_pixel": index = (index + 0.5) / scale_dim - 0.5
+    elif mode == "align_corners": index = index * (input_dim - 1) / (output_dim - 1) if output_dim != 1 else Tensor([0])
+    elif mode == "asymmetric": index = index / scale_dim
+    elif mode == "pytorch_half_pixel": index = (index + 0.5) / scale_dim - 0.5 if output_dim != 1 else Tensor([-0.5])
+    elif mode == "half_pixel_symmetric": index = input_dim / 2 * (1 - int(output_dim) / sizes_frac) + (index + 0.5) / scale_dim - 0.5
+    elif mode == "tf_crop_and_resize": index = roi_dim[0] * (input_dim - 1) + index * ((roi_dim[1] - roi_dim[0]) * (input_dim - 1) / (output_dim - 1)) # noqa: E501
+    else: raise ValueError(f"invalid {coordinate_transformation_mode=}")
+    return index.clip(0, input_dim-1)
 
-        wx, wy = math.ceil(x) - x, math.ceil(y) - y
-        if x == x_floor and y == y_floor:
-          weighted = corners[0][0]
-        elif x == x_floor:
-          weighted = corners[0][0] * wy + corners[1][0] * (1-wy)
-        elif y == y_floor:
-          weighted = corners[0][0] * wx + corners[0][1] * (1-wx)
-        else:
-          weighted = (corners[0][0] * wx + corners[0][1] * (1-wx)) * wy + \
-                     (corners[1][0] * (wx) + corners[1][1] * (1-wx)) * (1-wy)
-        ret.append(weighted)
-    return Tensor(ret).reshape(output_shape)
-  if mode == "cubic":
-    raise NotImplementedError("cubic interpolation is not implemented")
+  roi, scales, sizes = (to_python_const(a) for a in (roi, scales, sizes))
+  scales, sizes = (None if scales is None else scales[-2:]), (None if sizes is None else sizes[-2:])
+  # we pre permute the axes and permute back after resize
+  axes, input_shape, = (axes or list(range(X.ndim))), X.shape[2:],
+  perm = [a for a in range(len(X.shape)) if a not in axes] + list(axes)
+  X = X.permute(*perm)
+
+  if sizes is not None:
+    if keep_aspect_ratio_policy in ["not_larger", "not_smaller"]:
+      scale_fxn = min if keep_aspect_ratio_policy == "not_larger" else max
+      scales = scale_fxn([sizes[i] / input_shape[i] for i in range(X.ndim-2) if i+2 in axes])
+      sizes = [int((scales * input_shape[i]) + 0.5) if i+2 in axes else input_shape[i] for i in range(X.ndim-2)]
+    else: scales = [sizes[-2] / X.size(-2), sizes[-1] / X.size(-1)]
+  else: sizes = [int(sc*sh) for sc, sh in zip(scales, input_shape)]
+  scales = [scales] * 2 if not isinstance(scales, list) else scales
+  roi = [[st, ed] for st, ed in zip(roi, roi[len(roi)//2:])] if isinstance(roi, list) else [None] * (X.ndim-2)
+
+  # NOTE: this transformation makes it so that we can't just call Tensor.interpolate
+  # in Tensor.interpolate, we use indexes without any transformation
+  indexes = []
+  for shape, size, scale, region in zip(input_shape, sizes, scales, roi):
+    indexes.append(_apply_transformation(Tensor.arange(size), shape, scale, region, shape * scale, coordinate_transformation_mode))
+
+  if mode == "nearest":
+    indexes = [_apply_nearest_mode(index, shape, nearest_mode) for (index, shape) in zip(indexes, input_shape)]
+    X = X[(..., *Tensor.meshgrid(*indexes))]
+  if mode == "linear":
+    expand = list(X.shape)
+    for i in range(-len(sizes), 0):
+      reshape, index = [1] * X.ndim, indexes[i]
+      reshape[i] = expand[i] = sizes[i]
+      low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor(), index.ceil(), index - index.floor())]
+      X = X.gather(i, low).lerp(X.gather(i, high), perc)
+  if mode == "cubic": raise NotImplementedError("cubic interpolation is not implemented")
+  return X.permute(*[perm.index(i) for i in range(len(perm))]) if perm else X
 
 def CenterCropPad(t: Tensor, shape: Tensor, axes=None):
-  if not axes: axes = list(range(t.ndim))
+  shape = to_python_const(shape)
   shrink_arg = [None] * t.ndim
   pad_arg = [None] * t.ndim
-  shape = to_python_const(shape)
-  for s, x in zip(shape, axes):
+  for s, x in zip(shape, axes or range(t.ndim)):
     tx = t.shape[x]
     if s < tx: shrink_arg[x] = (tx//2 - (s+1)//2, tx//2 + s//2)
     elif s > tx: pad_arg[x] = ((s-tx)//2, (s-tx+1)//2)
@@ -499,67 +398,43 @@ def CenterCropPad(t: Tensor, shape: Tensor, axes=None):
 
 def OneHot(indices: Tensor, depth: Tensor, values: Tensor, axis=-1):
   depth = int(to_python_const(depth))
-  indices, rank = (indices < 0).where(indices+depth, indices), indices.ndim
-  if axis < 0: axis += rank + 1
-  ls, rs = indices.shape[0:axis], indices.shape[axis: rank]
-  cond = indices[:,None] == Tensor.arange(depth).reshape((1,) * len(ls) + (depth,) + (1,) * len(rs))
-  return cond.where(values[1], values[0])
+  # Scalar or Rank 1 tensor containing exactly one element
+  depth, indices = depth[0] if isinstance(depth, list) else depth, (indices < 0).where(indices+depth, indices),
+  return indices[:, None]._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
 
 def Compress(inp: Tensor, condition: Tensor, axis=None):
   if axis is None:
     inp = inp.flatten()
     axis = 0
-
   if axis < 0: axis += inp.ndim
-
   con_np = to_python_const(condition)
   con = Tensor(np.arange(condition.shape[0])[con_np]) # no boolean indexing in Tensor
   return inp[tuple(con if i == axis else slice(None) for i in range(inp.ndim))]
 
 def EyeLike(x: Tensor, dtype=None, k=0):
-  if dtype is None: dtype = x.dtype
-  else: dtype = DTYPE_MAP[int(dtype)]
-  dim = min(x.shape)
-  if x.shape[0] == x.shape[1]:
-    return Tensor.eye(dim, dtype=dtype)
-  padarg = tuple(None if d == dim else (k, d-dim-k) for d in x.shape)
-  return Tensor.eye(dim, dtype=dtype).pad(padarg)
+  ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) if dtype else x.dtype)
+  return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.size(0)-k) for d in x.shape))
 
 def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)
 
-def IsInf(x: Tensor, detect_negative=1, detect_positive=1):
-  return (x == float("inf")) * bool(detect_positive) + (x == float("-inf")) * bool(detect_negative)
-
 def DequantizeLinear(x: Tensor, x_scale: Tensor, x_zero_point: Union[Tensor, int] = 0, axis=1, block_size=0):
-  def numpy_repeat(t: Tensor, axis, repeats, out_shape):
-    t = t.reshape(tuple(-1 if i == axis-1 else 1 if i == axis else sh for i,sh in enumerate(t.shape)))
-    return t.repeat([repeats if i == axis else 1 for i in range(t.ndim)]).reshape(out_shape)
   if axis < 0: axis += x.ndim
-  if block_size:
-    x_zer, x_sc = numpy_repeat(x_zero_point, axis, block_size, x.shape), numpy_repeat(x_scale, axis, block_size, x.shape)
+  if not isinstance(x_zero_point, Tensor): x_zero_point = Tensor(x_zero_point)
+  if block_size: x_zer, x_sc = x_zero_point.repeat_interleave(block_size, axis), x_scale.repeat_interleave(block_size, axis)
   else:
-    x_sc = x_scale.reshape(*[1]*axis, *x_scale.shape, *[1]*(x.ndim - axis - x_scale.ndim))
-    x_zer = x_zero_point.reshape(*[1]*axis, *x_scale.shape, *[1]*(x.ndim - axis - x_scale.ndim)) if isinstance(x_zero_point, Tensor) else x_zero_point
+    shape = (*[1]*axis, *x_scale.shape, *[1]*(x.ndim - axis - x_scale.ndim))
+    x_sc, x_zer = x_scale.reshape(shape), x_zero_point.reshape(shape)
   return ((x.float() - x_zer) * x_sc).cast(x_scale.dtype)
 
-def IsNaN(x: Tensor): return x != x
-
 # copied from https://github.com/onnx/onnx/blob/main/onnx/reference/ops/op_image_decoder.py
 # without importing PIL we'll have to manually decode a bunch of image formats like PNG, JPEG, WebP, etc
 def ImageDecoder(encoded_stream: Tensor, pixel_format="RGB"):
-  try:
-    import PIL.Image
-  except ImportError as e:
-    raise ImportError("Pillow must be installed to use the reference implementation of the ImageDecoder operator") from e
-  img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream, tobytes=True)))
-  if pixel_format == "BGR":
-    return Tensor(np.array(img))[:, :, ::-1]
-  if pixel_format == "RGB":
-    return Tensor(np.array(img))
-  if pixel_format == "Grayscale":
-    img = img.convert("L")
-    decoded = Tensor(np.array(img))
-    return decoded.unsqueeze(-1) # (H, W) to (H, W, 1)
+  try: import PIL.Image
+  except ImportError as e: raise ImportError("Pillow must be installed to use the reference implementation of the ImageDecoder operator") from e
+  img = PIL.Image.open(io.BytesIO(to_python_const(encoded_stream)))
+  if pixel_format == "BGR": return Tensor(np.array(img))[:, :, ::-1]
+  if pixel_format == "RGB": return Tensor(np.array(img))
+  if pixel_format == "Grayscale": return Tensor(np.array(img.convert("L"))).unsqueeze(-1) # (H, W) to (H, W, 1)
   raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
 
 def AffineGrid(theta: Tensor, size: Tensor, align_corners=0):
@@ -592,9 +467,8 @@ def EmbedLayerNormalization(input_ids: Tensor, segment_ids:Optional[Tensor]=None
   compute_seg_emb = (segment_embedding is not None and segment_ids is not None)
   vocab_size, max_position_embeddings, type_vocab_size = word_embedding.shape[0], position_embedding.shape[0], (segment_embedding.shape[0] if compute_seg_emb else None)
 
-  def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:  # TODO from nn.Embedding. Could probably upstream this to Tensor
-    vocab_counter = Tensor.arange(vocab_size, dtype=x.dtype, requires_grad=False).reshape(1, 1, vocab_size).expand(*x.shape, vocab_size)
-    return (vocab_counter == x.unsqueeze(2).expand(*x.shape, vocab_size)) @ weight
+  def embedding(x:Tensor, vocab_size, weight:Tensor) -> Tensor:
+    return x.unsqueeze(-1).expand(*x.shape, vocab_size)._one_hot_along_dim(vocab_size) @ weight
 
   # bert embedding layer
   if epsilon is None: epsilon = 1e-12
@@ -644,53 +518,58 @@ def Attention(x:Tensor, weights, bias:Optional[Tensor]=None, mask_index:Optional
   return out, present
 
 # **************** ai.onnx.preview.training Ops ****************
+# NOTE: onnx test coverage only covers `T==0` cases, so for all `T>0` this isn't tested
+# NOTE: onnx training ops actually don't need the state for optim, all the ops work in a functional way, but we still can reuse optim.py code
 
-# TODO not entirely sure these optimizers are correct
+from tinygrad.nn.optim import Adam as TinyAdam
+from tinygrad.nn.optim import SGD
+
+def onnx_training(input_group_size):
+  def _decorator(func):
+    def __wrapper(R, T, *inputs, **kwargs):
+      old_training = Tensor.training
+      Tensor.training = True
+      T, R = to_python_const(T), R.detach()
+      groups = len(inputs) // input_group_size
+      ret = [func(R, T, *inps, **kwargs) for inps in (inputs[i::groups] for i in range(groups))]
+      Tensor.training = old_training
+      return tuple(flatten(zip(*ret)))
+    return __wrapper
+  return _decorator
+
+@onnx_training(3)
 def Adagrad(R, T, *inputs, decay_factor=0.0, epsilon=0.0, norm_coefficient=0.0):
-  groups = len(inputs) // 3
-  grouped_inputs = [inputs[i::groups] for i in range(groups)]
-  r = to_python_const(R / (1 + T * decay_factor))
-  ret = []
-  for X, G, H in grouped_inputs:
-    X.grad = norm_coefficient * X + G
-    X.grad.requires_grad, H.requires_grad = False, False # TODO manually turning off requires_grad, see TODO under (domain == "ai.onnx.preview.training") in onnx.py
-    H.assign(H.detach() + X.grad * X.grad).realize()
-    H_adaptive = H.sqrt() + epsilon
-    X.assign(X.detach() - r * X.grad / H_adaptive)
-    ret.extend([X, H])
-  ret = ret[::2] + ret[1::2]
-  return tuple(ret)
+  X, G, H = (i.detach() for i in inputs)
+  grad = norm_coefficient * X + G
+  H.assign(H + grad.square())
+  up = grad / (H.sqrt() + epsilon)
+  r = R / (1 + T * decay_factor)
+  X.assign(X.detach() - r * up)
+  return [X, H]
 
-def Momentum(R, T, *inputs, alpha, beta, mode, norm_coefficient):
-  groups = len(inputs) // 3
-  grouped_inputs = [inputs[i::groups] for i in range(groups)]
-  T, R.requires_grad = to_python_const(T), False
-  beta_adjusted = beta if T > 0 else 1
-  ret = []
-  for X, G, V in grouped_inputs:
-    X.grad = (norm_coefficient * X + G).realize()
-    X.grad.requires_grad, V.requires_grad = False, False
-    V.assign(alpha * V + beta_adjusted * X.grad).realize()
-    if mode == "standard": X.assign(X.detach() - R * V).realize()
-    elif mode == "nesterov": X.assign(X.detach() - R * (X.grad + alpha + V)).realize()
-    ret.extend([X, V])
-  ret = ret[::2] + ret[1::2]
-  return tuple(ret)
-
-# copied from tinygrad/nn/optim.py: LAMB with some edits
+@onnx_training(4)
 def Adam(R, T, *inputs, alpha=0.9, beta=0.999, epsilon=0.0, norm_coefficient=0.0, norm_coefficient_post=0.0):
-  groups = len(inputs) // 4
-  grouped_inputs = [inputs[i::groups] for i in range(groups)]
-  T, R.requires_grad = to_python_const(T), False
-  ret = []
-  for X, G, V, H in grouped_inputs:
-    X.grad = (norm_coefficient * X + G).realize()
-    V.requires_grad, H.requires_grad, X.grad.requires_grad = False, False, False
-    V.assign(alpha * V + (1.0 - alpha) * X.grad).realize()
-    H.assign(beta * H + (1.0 - beta) * (X.grad * X.grad)).realize()
-    up = (V / (1.0 - alpha**T)) / ((H / (1.0 - beta**T)).sqrt() + epsilon) if T > 0 else V / (H.sqrt() + epsilon)
-    X.assign(X.detach() - R * up).realize()
-    X = (1 - norm_coefficient_post) * X
-    ret.extend([X, V, H])
-  ret = ret[::3] + ret[1::3] + ret[2::3]
-  return tuple(ret)
+  X, G, V, H = inputs
+  G, V, H = G.detach(), V.detach(), H.detach()  # TODO we shouldn't need these detaches
+  X.grad = norm_coefficient * X.detach() + G
+  opt = TinyAdam([X], b1=alpha, b2=beta, eps=epsilon)
+  opt.m, opt.v, opt.lr = [V], [H], R
+  # need no-op for m_hat and v_hat if T == 0
+  if T == 0: opt.b1_t, opt.b2_t = opt.b1_t.zeros_like(), opt.b2_t.zeros_like()
+  else:
+    # `T-1` since it's applied again at the start of `_step`
+    opt.b1_t = Tensor([alpha**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
+    opt.b2_t = Tensor([beta**(T-1)], dtype=dtypes.float32, device=X.device, requires_grad=False)
+  opt.step()
+  X = (1 - norm_coefficient_post) * X
+  return [X, V, H]
+
+@onnx_training(3)
+def Momentum(R, T, *inputs, alpha, beta, mode, norm_coefficient):
+  X, G, V = inputs
+  G, V = G.detach(), V.detach()
+  X.grad = (norm_coefficient * X.detach() + G) * (beta if T > 0 else 1)
+  opt = SGD([X], momentum=alpha, nesterov=(mode=="nesterov"))
+  opt.b, opt.lr = [V], R
+  opt.step()
+  return [X, V]

extra/optimization/extract_dataset.py

@@ -5,9 +5,9 @@ from test.external.process_replay.process_replay import _pmap
 
 LOGOPS = os.getenv("LOGOPS", "/tmp/sops")
 
-def extract_ast(*args) -> bool:
+def extract_ast(*args) -> None:
   open(LOGOPS, "a").write(str(args[0]).replace("\n", "").replace(" ", "")+"\n")
-  return args[-1]
+  return None
 
 if __name__ == "__main__":
   _pmap("kernel", extract_ast)

extra/optimization/generate_dataset.sh

@@ -1,5 +1,6 @@
 #!/bin/bash
 export PAGE_SIZE=1
+export PYTHONPATH=.
 export LOGOPS=/tmp/ops
 export RUN_PROCESS_REPLAY=1
 rm $LOGOPS
@@ -24,5 +25,5 @@ JIT=2 BIG=1 MPS=1 python -m pytest test/test_speed_v_torch.py
 
 # extract, sort and uniq
 extra/optimization/extract_dataset.py
-sort -u /tmp/ops > /tmp/sops	
+sort -u /tmp/ops > /tmp/sops
 ls -lh /tmp/ops /tmp/sops	

ruff.toml

@@ -1,6 +1,6 @@
 indent-width = 2
 preview = true
-target-version = "py38"
+target-version = "py310"
 
 lint.select = [
   "F",  # Pyflakes

setup.py

@@ -29,7 +29,7 @@ setup(name='tinygrad',
         'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
         'linting': [
             "pylint",
-            "mypy==1.11.2",
+            "mypy==1.13.0",
             "typing-extensions",
             "pre-commit",
             "ruff",
@@ -59,6 +59,7 @@ setup(name='tinygrad',
             "bottle",
             "ggml-python"
         ],
+        'webgpu': ["wgpu"],
         'docs': [
             "mkdocs",
             "mkdocs-material",

test/external/external_benchmark_multitensor_allreduce.py

@@ -1,14 +1,13 @@
 import time
 from tinygrad import Tensor, Device, GlobalCounters, TinyJit
-from tinygrad.engine.lazy import LazyBuffer
-from tinygrad.ops import Ops
+from tinygrad.ops import Ops, UOp
 from tinygrad.multi import MultiLazyBuffer, all_reduce
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import run_schedule
 from tinygrad.helpers import getenv, Context, RING, DEBUG
 from typing import List, Union
 
-def realize(x: Union[LazyBuffer, List[LazyBuffer]]):
+def realize(x: Union[UOp, List[UOp]]):
   x = x if isinstance(x, list) else [x]
   run_schedule(create_schedule(x))
   for lb in x: Device[lb.device].synchronize()

test/external/external_benchmark_schedule.py

@@ -31,13 +31,14 @@ if __name__ == "__main__":
         if (restrict_kernel := getenv("RESTRICT_KERNEL", -1)) != -1: asts = asts[restrict_kernel:restrict_kernel+1]
         kernels: List[Kernel] = []
         with Timing(f"***** model opts({len(asts):2d}) in  "):
-          for ast in asts:
-            k = Kernel(ast)
-            if BEAM:
-              with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value)
-            elif NOOPT: pass
-            else: k.hand_coded_optimizations()
-            kernels.append(k)
+          with Profiling(PROFILE >= 3):
+            for ast in asts:
+              k = Kernel(ast)
+              if BEAM:
+                with Context(DEBUG=max(2, DEBUG.value)): k = beam_search(k, bufs_from_lin(k), BEAM.value)
+              elif NOOPT: pass
+              else: k.hand_coded_optimizations()
+              kernels.append(k)
 
         with Timing("***** model lower in     "): uops = [rewrite_shapetracker_with_index(k.get_optimized_ast(), k.opts) for k in kernels]
         with Profiling(PROFILE, fn="/tmp/rewrite.prof"):
@@ -48,7 +49,8 @@ if __name__ == "__main__":
                 rewritten_uops.append(full_graph_rewrite(u, k.opts))
             uops = rewritten_uops
         if getenv("LINEARIZE", 1):
-          with Timing("***** model linearize in "): uops = [linearize_uop(u) for u in uops]
+          with Profiling(PROFILE >= 2):
+            with Timing("***** model linearize in "): uops = [linearize_uop(u) for u in uops]
           print(sum(len(u) for u in uops))
           if getenv("SRC", 0):
             renderer = Device[Device.DEFAULT].renderer

test/external/external_debug_metal_sd_conv.py

@@ -0,0 +1,46 @@
+# ruff: noqa: E501
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.dtype import dtypes
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.search import bufs_from_lin
+from tinygrad.ops import UOp, Ops
+from tinygrad.shape.shapetracker import ShapeTracker
+from tinygrad.shape.view import View
+
+ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
+  UOp(Ops.STORE, dtypes.void, arg=None, src=(
+    UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=0, src=()),
+    UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(81920, 0, 64, 8, 1, 0, 0, 0), offset=0, mask=None, contiguous=True),)), src=()),
+    UOp(Ops.ADD, dtypes.half, arg=None, src=(
+      UOp(Ops.ADD, dtypes.half, arg=None, src=(
+        UOp(Ops.CAST, dtypes.half, arg=None, src=(
+          UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (5, 6, 7)), src=(
+            UOp(Ops.CAST, dtypes.float, arg=None, src=(
+              UOp(Ops.MUL, dtypes.half, arg=None, src=(
+                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
+                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=1, src=()),
+                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1, 2, 1, 2560, 4, 10, 4, 10), strides=(0, 163840, 0, 64, 0, 8, 0, 1), offset=-9, mask=((0, 1), (0, 2), (0, 1), (0, 2560), (0, 4), (1, 9), (0, 4), (1, 9)), contiguous=False), View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(4096000, 0, 0, 40, 1, 1600, 440, 11), offset=0, mask=None, contiguous=False))), src=()),)),
+                UOp(Ops.LOAD, dtypes.half, arg=None, src=(
+                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=2, src=()),
+                  UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 2560, 3, 3), strides=(0, 0, 23040, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),)),
+        UOp(Ops.LOAD, dtypes.half, arg=None, src=(
+          UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=3, src=()),
+          x17:=UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(2, 1, 1280, 8, 8, 1, 1, 1), strides=(0, 0, 1, 0, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)), src=()),)),)),
+      UOp(Ops.LOAD, dtypes.half, arg=None, src=(
+        UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(), arg=4, src=()),
+         x17,)),)),)),))
+opts = [Opt(op=OptOps.UPCAST, axis=3, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.UNROLL, axis=2, amt=0), Opt(op=OptOps.UNROLL, axis=1, amt=0), Opt(op=OptOps.LOCAL, axis=1, amt=8), Opt(op=OptOps.LOCAL, axis=2, amt=8), Opt(op=OptOps.LOCAL, axis=2, amt=2)]
+
+k = Kernel(ast)
+for opt in opts: k.apply_opt(opt)
+bufs = bufs_from_lin(k)
+
+prg = CompiledRunner(k.to_program())
+
+for i in range(10):
+  speed = prg(bufs, var_vals={}, wait=True)
+  print(f"kernel time: {speed*1e3:.2f} ms")
+
+# on M1 Max
+# 11ms before block 9b0859d71780fef5cf3831e317f74e53f2483229
+# 15ms after block cbcc1c20eb09a1342f6581cfbb99632bade982a8

test/external/external_metal_compile_fail.py

@@ -225,6 +225,6 @@ from tinygrad.runtime.ops_metal import MetalDevice, MetalCompiler, MetalProgram
 
 if __name__ == "__main__":
   dev = MetalDevice("METAL")
-  lib = MetalCompiler(dev).compile(src)
+  lib = MetalCompiler().compile(src)
   prg = MetalProgram(dev, "r_64_32_8_16_4_6_6_4", lib)
 

test/external/external_test_hcq.py

@@ -1,7 +1,7 @@
 import unittest, ctypes, struct, time, array
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.helpers import to_mv, CI
-from tinygrad.device import Buffer, BufferOptions
+from tinygrad.device import Buffer, BufferSpec
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import get_runner
 
@@ -22,27 +22,27 @@ class TestHCQ(unittest.TestCase):
     TestHCQ.a = Tensor([0.,1.], device=Device.DEFAULT).realize()
     TestHCQ.b = self.a + 1
     si = create_schedule([self.b.lazydata])[-1]
-    TestHCQ.runner = get_runner(TestHCQ.d0.dname, si.ast)
+    TestHCQ.runner = get_runner(TestHCQ.d0.device, si.ast)
     TestHCQ.b.lazydata.buffer.allocate()
     # wow that's a lot of abstraction layers
     TestHCQ.addr = struct.pack("QQ", TestHCQ.b.lazydata.buffer._buf.va_addr, TestHCQ.a.lazydata.buffer._buf.va_addr)
     TestHCQ.addr2 = struct.pack("QQ", TestHCQ.a.lazydata.buffer._buf.va_addr, TestHCQ.b.lazydata.buffer._buf.va_addr)
-    TestHCQ.kernargs_off = TestHCQ.runner.clprg.kernargs_offset
-    TestHCQ.kernargs_size = TestHCQ.runner.clprg.kernargs_alloc_size
+    TestHCQ.kernargs_off = TestHCQ.runner._prg.kernargs_offset
+    TestHCQ.kernargs_size = TestHCQ.runner._prg.kernargs_alloc_size
     ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_off, TestHCQ.addr, len(TestHCQ.addr))
     ctypes.memmove(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size+TestHCQ.kernargs_off, TestHCQ.addr2, len(TestHCQ.addr2))
 
     if Device.DEFAULT == "AMD":
-      from tinygrad.runtime.ops_amd import HWCopyQueue, HWPM4Queue
+      from tinygrad.runtime.ops_amd import HWQueue, HWPM4Queue
       TestHCQ.compute_queue = HWPM4Queue
-      TestHCQ.copy_queue = HWCopyQueue
+      TestHCQ.copy_queue = HWQueue
     elif Device.DEFAULT == "NV":
-      from tinygrad.runtime.ops_nv import HWCopyQueue, HWComputeQueue
+      from tinygrad.runtime.ops_nv import HWQueue, HWQueue
       # nv need to copy constbuffer there as well
-      to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner.clprg.constbuffer_0)
-      to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner.clprg.constbuffer_0)
-      TestHCQ.compute_queue = HWComputeQueue
-      TestHCQ.copy_queue = HWCopyQueue
+      to_mv(TestHCQ.d0.kernargs_ptr, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner._prg.constbuffer_0)
+      to_mv(TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, 0x160).cast('I')[:] = array.array('I', TestHCQ.runner._prg.constbuffer_0)
+      TestHCQ.compute_queue = HWQueue
+      TestHCQ.copy_queue = HWQueue
 
   def setUp(self):
     TestHCQ.d0.synchronize()
@@ -54,11 +54,11 @@ class TestHCQ(unittest.TestCase):
     temp_signal, temp_value = TestHCQ.d0._alloc_signal(value=0), 0
     q = TestHCQ.compute_queue()
     for _ in range(1000):
-      q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
       q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
       temp_value += 1
 
-      q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
       q.signal(temp_signal, temp_value + 1).wait(temp_signal, temp_value + 1)
       temp_value += 1
 
@@ -72,9 +72,9 @@ class TestHCQ(unittest.TestCase):
   def test_run_1000_times(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
            TestHCQ.runner.p.local_size)
     for _ in range(1000):
       TestHCQ.d0._set_signal(temp_signal, 1)
@@ -88,11 +88,11 @@ class TestHCQ(unittest.TestCase):
   def test_run_to_3(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(temp_signal, 1).wait(temp_signal, 1)
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
@@ -102,7 +102,7 @@ class TestHCQ(unittest.TestCase):
   def test_update_exec(self):
     q = TestHCQ.compute_queue()
     exec_cmd_idx = len(q)
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.update_exec(exec_cmd_idx, (1,1,1), (1,1,1))
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
@@ -116,9 +116,9 @@ class TestHCQ(unittest.TestCase):
   def test_bind_run(self):
     temp_signal = TestHCQ.d0._alloc_signal(value=0)
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(temp_signal, 2).wait(temp_signal, 2)
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr+TestHCQ.kernargs_size, TestHCQ.runner.p.global_size,
            TestHCQ.runner.p.local_size)
     q.bind(TestHCQ.d0)
     for _ in range(1000):
@@ -134,7 +134,7 @@ class TestHCQ(unittest.TestCase):
   def test_update_exec_binded(self):
     q = TestHCQ.compute_queue()
     exec_ptr = q.ptr()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     q.bind(TestHCQ.d0)
 
@@ -171,7 +171,7 @@ class TestHCQ(unittest.TestCase):
 
   def test_run_normal(self):
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     TestHCQ.d0.timeline_value += 1
@@ -202,7 +202,7 @@ class TestHCQ(unittest.TestCase):
 
   def test_run_signal(self):
     q = TestHCQ.compute_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
     q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
     q.submit(TestHCQ.d0)
     TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
@@ -255,8 +255,8 @@ class TestHCQ(unittest.TestCase):
   def test_copy_bandwidth(self):
     # THEORY: the bandwidth is low here because it's only using one SDMA queue. I suspect it's more stable like this at least.
     SZ = 2_000_000_000
-    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate()
-    b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate()
+    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    b = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
     q = TestHCQ.copy_queue()
     q.copy(a._buf.va_addr, b._buf.va_addr, SZ)
     et = _time_queue(q, TestHCQ.d0)
@@ -266,8 +266,8 @@ class TestHCQ(unittest.TestCase):
 
   def test_cross_device_copy_bandwidth(self):
     SZ = 2_000_000_000
-    b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate()
-    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferOptions(nolru=True)).allocate()
+    b = Buffer(f"{Device.DEFAULT}:1", SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
+    a = Buffer(Device.DEFAULT, SZ, dtypes.uint8, options=BufferSpec(nolru=True)).allocate()
     TestHCQ.d0._gpu_map(b._buf)
     q = TestHCQ.copy_queue()
     q.copy(a._buf.va_addr, b._buf.va_addr, SZ)
@@ -279,7 +279,7 @@ class TestHCQ(unittest.TestCase):
   def test_interleave_compute_and_copy(self):
     q = TestHCQ.compute_queue()
     qc = TestHCQ.copy_queue()
-    q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)  # b = [1, 2]
+    q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)  # b = [1, 2]
     q.signal(sig:=TestHCQ.d0._alloc_signal(value=0), value=1)
     qc.wait(sig, value=1)
     qc.copy(TestHCQ.a.lazydata.buffer._buf.va_addr, TestHCQ.b.lazydata.buffer._buf.va_addr, 8)
@@ -316,7 +316,7 @@ class TestHCQ(unittest.TestCase):
     for _ in range(40):
       q = TestHCQ.compute_queue()
       q.wait(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value - 1)
-      q.exec(TestHCQ.runner.clprg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
+      q.exec(TestHCQ.runner._prg, TestHCQ.d0.kernargs_ptr, TestHCQ.runner.p.global_size, TestHCQ.runner.p.local_size)
       q.signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value).submit(TestHCQ.d0)
       TestHCQ.d0._wait_signal(TestHCQ.d0.timeline_signal, TestHCQ.d0.timeline_value)
       TestHCQ.d0.timeline_value += 1

test/external/external_test_nv.py

@@ -2,7 +2,7 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.runtime.ops_nv import NVDevice, HWComputeQueue
+from tinygrad.runtime.ops_nv import NVDevice, HWQueue
 from tinygrad.engine.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 from tinygrad.engine.realize import get_runner, CompiledRunner
@@ -21,7 +21,7 @@ class TestNV(unittest.TestCase):
     TestNV.a = Tensor([0.,1.], device="NV").realize()
     TestNV.b = self.a + 1
     si = create_schedule([self.b.lazydata])[-1]
-    TestNV.d0_runner = get_runner(TestNV.d0.dname, si.ast)
+    TestNV.d0_runner = get_runner(TestNV.d0.device, si.ast)
     TestNV.b.lazydata.buffer.allocate()
     TestNV.addr = struct.pack("QQ", TestNV.b.lazydata.buffer._buf.va_addr, TestNV.a.lazydata.buffer._buf.va_addr)
 
@@ -44,19 +44,19 @@ class TestNV(unittest.TestCase):
   def test_buf4_usage(self):
     TestNV.along = Tensor([105615], device="NV").realize()
     ast = LazyOp(op=BufferOps.STORE, src=(LazyOp(op=Ops.SIN, src=(LazyOp(op=Ops.CAST, src=(LazyOp(op=BufferOps.LOAD, src=(), arg=MemBuffer(idx=1, dtype=dtypes.ulong, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))),), arg=dtypes.float),), arg=None),), arg=MemBuffer(idx=0, dtype=dtypes.float, st=ShapeTracker(views=(View(shape=(3,), strides=(1,), offset=0, mask=None, contiguous=True),)))) # noqa: E501
-    temp_runner = get_runner(TestNV.d0.dname, (ast,))
+    temp_runner = get_runner(TestNV.d0.device, (ast,))
     temp_runner([TestNV.b.lazydata.buffer, TestNV.along.lazydata.buffer], var_vals={})
     val = TestNV.b.lazydata.buffer.as_buffer().cast("f")[0]
     assert abs(val - 0.80647) < 0.001, f"got val {val}"
 
   def test_kernargs_no_oob_access(self):
     kernargs_start = TestNV.d0._gpu_alloc((2 << 20), map_to_cpu=True).va_addr
-    kernargs = kernargs_start + ((2 << 20) - TestNV.d0_runner.clprg.kernargs_alloc_size)
-    to_mv(kernargs, 0x160).cast('I')[:] = array.array('I', TestNV.d0_runner.clprg.constbuffer_0)
-    ctypes.memmove(kernargs + TestNV.d0_runner.clprg.kernargs_offset, TestNV.addr, len(TestNV.addr))
+    kernargs = kernargs_start + ((2 << 20) - TestNV.d0_runner._prg.kernargs_alloc_size)
+    to_mv(kernargs, 0x160).cast('I')[:] = array.array('I', TestNV.d0_runner._prg.constbuffer_0)
+    ctypes.memmove(kernargs + TestNV.d0_runner._prg.kernargs_offset, TestNV.addr, len(TestNV.addr))
 
-    q = HWComputeQueue()
-    q.exec(TestNV.d0_runner.clprg, kernargs, TestNV.d0_runner.global_size, TestNV.d0_runner.local_size)
+    q = HWQueue()
+    q.exec(TestNV.d0_runner._prg, kernargs, TestNV.d0_runner.global_size, TestNV.d0_runner.local_size)
     q.signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value).submit(TestNV.d0)
     TestNV.d0._wait_signal(TestNV.d0.timeline_signal, TestNV.d0.timeline_value)
     TestNV.d0.timeline_value += 1

test/external/external_test_onnx_backend.py

@@ -40,10 +40,11 @@ class TinygradBackend(Backend):
 
 backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
 
-# TODO figure out why it's returning wrong values, geohotstan's uneducated guess is it's due to imprecision from float64 (double) -> float32
-# see Type Constraints: https://onnx.ai/onnx/operators/onnx_aionnxpreviewtraining_Adam.html#type-constraints
+# TODO: there isn't an AttributeProto for `epsilon` in the NodeProto for 'test_adam_multiple_cpu'
+# [x.name for x in n.attribute] -> ['alpha', 'beta', 'norm_coefficient']
+# but in their documentation https://github.com/onnx/onnx/blob/main/docs/Operators.md#examples-176, it states there being an epsilon of 1e-2
+# test passes with epsilon = 1e-2
 backend_test.exclude('test_adam_multiple_cpu')
-backend_test.exclude('test_nesterov_momentum_cpu')
 
 # about different dtypes
 if not is_dtype_supported(dtypes.float64):
@@ -69,11 +70,16 @@ backend_test.exclude('BFLOAT16')  # not supported in numpy
 # TODO: fix these with true onnx float16
 backend_test.exclude('to_FLOAT16')
 backend_test.exclude('cast_no_saturate')
+backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
+backend_test.exclude('test_max_float16_cpu')
+backend_test.exclude('test_min_float16_cpu')
 
-backend_test.exclude('test_pow_types_int*')
 backend_test.exclude('test_convinteger_*')
 backend_test.exclude('test_matmulinteger_*')
 
+backend_test.exclude('test_dequantizelinear_int4_cpu')
+backend_test.exclude('test_dequantizelinear_uint4_cpu')
+
 # we don't support indexes
 backend_test.exclude('test_nonzero_*')
 
@@ -116,7 +122,6 @@ backend_test.exclude('test_affine_grid_3d_expanded_cpu')
 backend_test.exclude('test_range_int32_type_negative_delta_expanded_cpu')
 
 # unsupported (strange) ops
-backend_test.exclude('test_bitwise_*')
 backend_test.exclude('test_blackmanwindow_*')
 backend_test.exclude('test_bernoulli_*')
 backend_test.exclude('test_det_*')
@@ -142,7 +147,7 @@ backend_test.exclude('test_melweightmatrix_*')
 backend_test.exclude('test_basic_deform_conv_*')
 backend_test.exclude('test_deform_conv_*')
 backend_test.exclude('test_lppool_*')
-backend_test.exclude('test_scan*')
+backend_test.exclude('test_scan_*')
 backend_test.exclude('test_split_to_sequence_*')
 backend_test.exclude('test_resize_downsample_scales_cubic_*') # unsure how to implement cubic
 backend_test.exclude('test_resize_downsample_sizes_cubic_*') # unsure how to implement cubic
@@ -157,6 +162,8 @@ backend_test.exclude('test_resize_tf_crop_and_resize_cpu') # unsure about fill v
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_value_only_mapping_cpu') # bad data type string
 backend_test.exclude('test_ai_onnx_ml_label_encoder_tensor_mapping_cpu') # bad data type string
 backend_test.exclude('test_group_normalization_*') # numerical inaccuracy problem. Current Group Normalization OP fails test
+backend_test.exclude('test_scatter_elements_with_reduction_min_cpu') # min not yet supported
+backend_test.exclude('test_scatter_elements_with_reduction_max_cpu') # max not yet supported
 
 if Device.DEFAULT in ['GPU', 'METAL']:
   backend_test.exclude('test_resize_upsample_sizes_nearest_axes_2_3_cpu')

test/external/external_test_speed_llama.py

@@ -14,7 +14,7 @@ class FakeProgram:
 
 class FakeAllocator(Allocator):
   def _alloc(self, sz, options): return None
-  def copyin(self, dest, src:memoryview): pass
+  def _copyin(self, dest, src:memoryview): pass
 
 class TestLLaMASpeed(unittest.TestCase):
   def test_llama_compile(self):

test/external/fuzz_graph.py

@@ -29,7 +29,7 @@ def alloc_rawbuffer(device, fill=False):
   if fill:
     with Context(DEBUG=0):
       data = np.random.randint(-10000, 10000, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data).realize().lazydata.realized.as_buffer())
+      rawbuf.copyin(Tensor(data).realize().lazydata.base.realized.as_buffer())
   return rawbuf
 
 def gen_kernel_ji(device, deps):

test/external/fuzz_kfd.py

@@ -3,7 +3,7 @@ import random
 from tqdm import trange
 from typing import List
 from tinygrad import Device
-from tinygrad.runtime.ops_amd import AMDDevice, HWCopyQueue, HWComputeQueue
+from tinygrad.runtime.ops_amd import AMDDevice, HWQueue
 
 if __name__ == "__main__":
   dev: List[AMDDevice] = [Device[f"KFD:{i}"] for i in range(6)]
@@ -15,9 +15,9 @@ if __name__ == "__main__":
     d1, b1 = random.choice(buffers)
     d2, b2 = random.choice(buffers)
     d1._gpu_map(b2)
-    q = HWComputeQueue()
+    q = HWQueue()
     q.signal(sig:=AMDDevice._alloc_signal(10))
-    qc = HWCopyQueue()
+    qc = HWQueue()
     qc.wait(sig)
     qc.copy(b1.va_addr, b2.va_addr, min(b1.size, b2.size))
     d1.completion_signal.value = 1

test/external/fuzz_linearizer.py

@@ -75,7 +75,7 @@ def get_fuzz_rawbufs(lin):
         data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
       else:
         data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
-      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().lazydata.realized.as_buffer())
+      rawbuf.copyin(Tensor(data, device=lin.opts.device).realize().lazydata.base.realized.as_buffer())
   return rawbufs
 
 def get_fuzz_rawbuf_like(old_rawbuf, zero=False, copy=False, size=None, force_device=None):
@@ -228,7 +228,7 @@ def fuzz_linearizer(lin: Kernel, rtol=1e-2, atol=1e-2, opts_list=None):
           validate_lin = test_lin.copy()
           validate_lin.opts = validate_device.renderer
           if validate_rawbufs is None:
-            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.dname) for x in rawbufs]
+            validate_rawbufs = [get_fuzz_rawbuf_like(x, copy=True, force_device=validate_device.device) for x in rawbufs]
           (_msg, _, _, _, state2) = compare_linearizer(validate_lin, validate_rawbufs, var_vals, ground_truth, rtol=rtol, atol=atol)
 
           if _msg != "PASS": failures[f"VALIDATE_DEV_{_msg}"].append((validate_lin.ast, validate_lin.applied_opts))

test/external/fuzz_schedule.py

@@ -4,15 +4,14 @@ from typing import DefaultDict, Dict, List, Set, Tuple, TypeVar, Union
 from tinygrad.device import Buffer
 from tinygrad.engine.realize import capturing, lower_schedule_item
 from tinygrad.helpers import DEBUG, MULTIOUTPUT, colored, getenv
-from tinygrad.engine.lazy import LazyBuffer
 from tinygrad.engine.schedule import LBScheduleItem, _graph_schedule, ScheduleItem
-from tinygrad.ops import Ops
+from tinygrad.ops import Ops, UOp
 from tinygrad.tensor import Tensor, _to_np_dtype
 
 ctx_vars = { MULTIOUTPUT: (0, 1) }
 FUZZ_SCHEDULE_MAX_PATHS = getenv("FUZZ_SCHEDULE_MAX_PATHS", 10)
 
-def fuzz_schedule(outs:List[LazyBuffer]):
+def fuzz_schedule(outs:List[UOp]):
   # find toposorts across all tunable params
   unique_ts: Dict[Tuple[LBScheduleItem, ...], Dict[str, int]] = {}
   for combination in itertools.product(*ctx_vars.values()):
@@ -24,10 +23,10 @@ def fuzz_schedule(outs:List[LazyBuffer]):
   if DEBUG >= 1: print(colored(f"fuzzing {len(toposorts)} schedule permutations", "yellow"))
 
   # setup ground truth
-  ground_truth: Dict[LazyBuffer, memoryview] = {}
-  assign_targets: Dict[LazyBuffer, LazyBuffer] = {}
+  ground_truth: Dict[UOp, memoryview] = {}
+  assign_targets: Dict[UOp, UOp] = {}
   # IMPORTANT: freeze prerealized bufs before ScheduleItem exec
-  prerealized: Dict[LazyBuffer, memoryview] = {}
+  prerealized: Dict[UOp, memoryview] = {}
   seed = Tensor._seed
   ts,_ = toposorts[0]
   for lsi in ts:
@@ -47,7 +46,7 @@ def fuzz_schedule(outs:List[LazyBuffer]):
   # exec and validate each permutation with new Buffers
   for i, (ts, ctx) in enumerate(toposorts[1:]):
     if DEBUG >= 1: print(colored(f"testing permutation {i} {ctx}", "yellow"))
-    rawbufs: Dict[LazyBuffer, Buffer] = {}
+    rawbufs: Dict[UOp, Buffer] = {}
     for lsi in ts:
       for out in lsi.outputs:
         base = rawbufs[lsi.inputs[0]].base if out.op is Ops.BUFFER_VIEW else None

test/external/fuzz_uops.py

@@ -3,7 +3,7 @@ from collections import defaultdict
 import numpy as np
 from dataclasses import replace
 from typing import DefaultDict, Dict, List, Tuple
-from tinygrad.ops import END_FOR_UOP, UOp, print_uops
+from tinygrad.ops import UOp, print_uops, Ops
 from tinygrad.device import Buffer, Device
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import DEBUG, colored
@@ -11,6 +11,8 @@ from tinygrad.ops import Variable
 from tinygrad.tensor import _to_np_dtype
 from test.external.fuzz_schedule import FUZZ_SCHEDULE_MAX_PATHS, find_all_toposorts
 
+END_FOR_UOP = {Ops.IF:(Ops.STORE, Ops.ENDIF), Ops.RANGE:(Ops.ASSIGN, Ops.ENDRANGE)}
+
 def fuzz_uops(uops:List[UOp]) -> List[Tuple[UOp, ...]]:
   blocks: List[List[UOp]] = [[]]
   for u in uops:
@@ -50,10 +52,10 @@ class UOpsFuzzerRunner(CompiledRunner):
       # setup prg
       uops = list(path)
       if DEBUG >= 5: print_uops(uops)
-      self.p = replace(self.p, name=(name:=f"{init_name}fuzz{i}"), src=Device[self.p.dname].renderer.render(name, uops), uops=uops)
+      self.p = replace(self.p, name=(name:=f"{init_name}fuzz{i}"), src=Device[self.p.device].renderer.render(name, uops), uops=uops)
       if DEBUG >= 4: print(self.p.src)
-      self.lib = Device[self.p.dname].compiler.compile_cached(self.p.src)
-      self.clprg = Device[self.p.dname].runtime(name, self.lib)
+      self.lib = Device[self.p.device].compiler.compile_cached(self.p.src)
+      self.clprg = Device[self.p.device].runtime(name, self.lib)
       for x in (rawbufs:=[init_globals[i] for i in self.p.globals]): x.copyin(init_rawbufs[x])
       # verify
       super().__call__(rawbufs, var_vals, wait)

test/external/process_replay/process_replay.py

@@ -1,9 +1,9 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, functools
-from typing import Callable, List, Tuple, Union, cast
+import os, multiprocessing, logging, pickle, sqlite3, difflib, functools, warnings
+from typing import Callable, List, Set, Tuple, Union, cast
 from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
-from tinygrad.engine.schedule import full_ast_rewrite
+from tinygrad.engine.schedule import ScheduleContext, full_ast_rewrite
 from tinygrad.codegen.kernel import Kernel, Opt
 from tinygrad.renderer import Renderer
 from tinygrad.ops import UOp
@@ -25,10 +25,12 @@ ASSERT_DIFF = int((flag:="[pr]") in os.getenv("COMMIT_MESSAGE", flag) or flag in
 if not getenv("ASSERT_PROCESS_REPLAY", 1): ASSERT_DIFF = 0
 SKIP_PROCESS_REPLAY = (k:="[skip_process_replay]") in os.getenv("COMMIT_MESSAGE", "") or k in os.getenv("PR_TITLE", "")
 if REF == "master": SKIP_PROCESS_REPLAY = True
+class ProcessReplayWarning(Warning): pass
 
 # *** recreators
 
-def recreate_sched(*args) -> UOp: return full_ast_rewrite(*args[0])[0]
+def recreate_sched(ast:UOp, assigns:Set[UOp]) -> UOp:
+  return full_ast_rewrite(ast, ScheduleContext(assigns=assigns))[0]
 def recreate_kernel(ast:UOp, opts:Renderer, applied_opts:List[Opt], name:str, _) -> str:
   k = Kernel(ast, opts=opts)
   for opt in applied_opts: k.apply_opt(opt)
@@ -53,10 +55,10 @@ def diff(offset:int, name:str, fxn:Callable) -> Union[Tuple[int, int], bool]:
     # try recreate
     try:
       with Context(**{k:v for k,v in args[-2].items() if k in ContextVar._cache and k != "DEBUG"}): good = fxn(*args[:-2])
+      if good is None: continue
     except Exception as e:
-      logging.warning(f"FAILED TO RECREATE KERNEL {e}")
+      warnings.warn(f"FAILED TO RECREATE KERNEL {e}", ProcessReplayWarning)
       for x in args[:-1]: logging.info(x)
-      if ASSERT_DIFF: return True
       continue
     # diff kernels
     try: assert args[-1] == good
@@ -83,7 +85,7 @@ def _pmap(name:str, fxn:Callable, maxtasksperchild:int=16) -> None:
   cur = conn.cursor()
   try: row_count = cur.execute(f"select count(*) from '{name}_{TABLE_NAME}'").fetchone()[0]
   except sqlite3.OperationalError:
-    logging.warning(f"{name}_{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
+    warnings.warn(f"{name}_{TABLE_NAME} isn't accessible in master, did DB_VERSION change?", ProcessReplayWarning)
     return None
   conn.commit()
   cur.close()
@@ -98,7 +100,7 @@ def _pmap(name:str, fxn:Callable, maxtasksperchild:int=16) -> None:
     logging.info(f"{sum(changed)} kernels changed")
     if sum(insertion) != 0: logging.info(colored(f"{sum(insertion)} insertions(+)", "green"))
     if sum(deletions) != 0: logging.info(colored(f"{sum(deletions)} deletions(-)", "red"))
-    if any(changed) and ASSERT_DIFF: raise AssertionError("process replay detected changes")
+    if any(changed): warnings.warn("process replay detected changes", ProcessReplayWarning)
 
 # *** main loop
 
@@ -107,6 +109,7 @@ if __name__ == "__main__":
     logging.info("skipping process replay.")
     exit(0)
 
+  if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
   for name,fxn in [("schedule", recreate_sched), ("kernel", recreate_kernel)]:
     logging.info(f"***** {name} diff")
     try: _pmap(name, fxn)

test/external/speed_v_theoretical.py

@@ -1,7 +1,8 @@
-import unittest, time
+import unittest
 from tinygrad import Tensor, TinyJit, Device
-from tinygrad.helpers import Context, DEBUG
+from tinygrad.helpers import Context, DEBUG, GlobalCounters
 from tinygrad.nn import Conv2d
+from tinygrad.nn.state import get_parameters
 
 class TestKernelSpeed(unittest.TestCase):
   def _get_tensor(self, *shape:int):
@@ -40,14 +41,15 @@ class TestKernelSpeed(unittest.TestCase):
     if K is None: K = M
     tms = []
     with Context(BEAM=3):
-      for _ in range(10):
+      for i in range(10):
         a = self._get_tensor(M, K)
         b = self._get_tensor(K, N)
-        Device.default.synchronize()
-        st = time.perf_counter()
-        c = f(a, b)
-        Device.default.synchronize()
-        tms.append(time.perf_counter() - st)
+        if i >= 3:
+          GlobalCounters.time_sum_s = 0
+          with Context(DEBUG=max(DEBUG, 2)): c = f(a, b)
+          tms.append(GlobalCounters.time_sum_s)
+        else:
+          c = f(a, b)
 
     ops = 2 * M * N * K
     mems = a.dtype.itemsize * M * K + b.dtype.itemsize * K * N + c.dtype.itemsize * M * N
@@ -61,35 +63,39 @@ class TestKernelSpeed(unittest.TestCase):
     def f(conv, x) -> Tensor: return conv(x).realize()
     tms = []
     K = 3
-    with Context(BEAM=2):
+    with Context(BEAM=0, DEBUG=0):
       conv = Conv2d(CIN, COUT, K, padding=1)
-      for _ in range(10):
+      Tensor.realize(*get_parameters(conv))
+
+    with Context(BEAM=2):
+      for i in range(10):
         x = self._get_tensor(BS, CIN, H, W)
-        Device.default.synchronize()
-        st = time.perf_counter()
-        _c = f(conv, x)
-        Device.default.synchronize()
-        tms.append(time.perf_counter() - st)
+        if i >= 3:
+          GlobalCounters.time_sum_s = 0
+          with Context(DEBUG=max(DEBUG, 2)): _c = f(conv, x)
+          tms.append(GlobalCounters.time_sum_s)
+        else:
+          _c = f(conv, x)
 
     # naive algo
     ops = 2 * BS * CIN * COUT * K * K * H * W
-    # TODO: what should this be?
-    mems = 0
+    mems = x.nbytes() + conv.weight.nbytes() + conv.bias.nbytes() + _c.nbytes()
     tm = min(tms)
     tflops = ops / tm / 1e12
     gbs = mems / tm / 1e9
     self._compare(tm, tflops, gbs, nv_tflops, nv_gbs, amd_tflops, amd_gbs)
 
-  # TODO: smaller ones has other overhead in synchronize
-  # def test_gemm_1024(self): self._test_matmul(1024, nv_tflops=8, amd_tflops=7)
-  # def test_gemm_2048(self): self._test_matmul(2048, nv_tflops=50, amd_tflops=30)
-  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=100, amd_tflops=70)
-  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=130, amd_tflops=70)
+  # NOTE: tiny7 was slower than tiny12
+  # TODO: why are convs so slow?!?
+  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=27, amd_tflops=24)
 
-  def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=430, amd_gbs=400)
-  def test_gemv_4096_16384(self): self._test_matmul(4096, 16384, 1, nv_gbs=430, amd_gbs=400)
+  # theoretical is nv_tflops=165, amd_tflops=123
+  def test_gemm_4096(self): self._test_matmul(4096, nv_tflops=115, amd_tflops=80)
+  def test_gemm_8192(self): self._test_matmul(8192, nv_tflops=130, amd_tflops=75)
 
-  def test_conv_3x3_256_32_32_256_256(self): self._test_conv_3x3(256, 32, 32, 256, 256, nv_tflops=30, amd_tflops=22)
+  # theoretical is nv_gbs=1008, amd_gbs=960
+  def test_gemv_16384_4096(self): self._test_matmul(16384, 4096, 1, nv_gbs=840, amd_gbs=750)
+  def test_gemv_4096_16384(self): self._test_matmul(4096, 16384, 1, nv_gbs=830, amd_gbs=760)
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/imported/test_indexing.py

@@ -4,6 +4,7 @@ import unittest, random, copy, warnings
 import numpy as np
 
 from tinygrad import Tensor, dtypes, Device, TinyJit
+from tinygrad.device import is_dtype_supported
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.helpers import CI, all_same, prod
@@ -22,7 +23,7 @@ def consec(shape, start=1):
 def set_(reference: Tensor, shape, strides, offset):
   if reference.lazydata.base.realized is None: reference.realize()
   assert reference.lazydata.base.realized, "base has to be realized before setting it to strided's base"
-  strided = Tensor(reference.lazydata._view(ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),))))
+  strided = Tensor(reference.lazydata.view(ShapeTracker((View.create(shape=shape, strides=strides, offset=offset),))))
   assert strided.lazydata.st.real_strides() == strides, "real_strides should equal strides for strided"
   return strided
 
@@ -810,6 +811,7 @@ class TestIndexing(unittest.TestCase):
       numpy_testing_assert_equal_helper(output, input_list)
   '''
 
+  @unittest.skipUnless(is_dtype_supported(dtypes.long), f"long dtype not supported on {Device.DEFAULT}")
   def test_index_ind_dtype(self):
     x = Tensor.randn(4, 4)
     # ind_long = torch.randint(4, (4,), dtype=torch.long)
@@ -1062,9 +1064,9 @@ class TestIndexing(unittest.TestCase):
     numpy_testing_assert_equal_helper(a[0, one], a[zero, 1])
 
     # indexing by a scalar should slice (not copy)
-    self.assertEqual(data_ptr(a[0, 1]), data_ptr(a[zero, one]))
-    self.assertEqual(data_ptr(a[1]), data_ptr(a[one.cast(dtypes.int32)]))
-    self.assertEqual(data_ptr(a[1]), data_ptr(a[one.cast(dtypes.int16)]))
+    numpy_testing_assert_equal_helper(a[0, 1], a[zero, one])
+    numpy_testing_assert_equal_helper(a[1], a[one.cast(dtypes.int32)])
+    numpy_testing_assert_equal_helper(a[1], a[one.cast(dtypes.int16)])
 
     # scalar indexed with scalar
     r = Tensor.randn()
@@ -1105,6 +1107,21 @@ class TestIndexing(unittest.TestCase):
     np.testing.assert_allclose(9.9, r, rtol=1e-7)
   '''
 
+  @unittest.skip("getitem expects elementwise ops folding to be instant, but it happens late in the scheduler")
+  def test_getitem_casted_scalars_folding(self):
+    Tensor.manual_seed(0)
+    # cast of const is just another const, don't need extra kernels for this
+    a = Tensor.randn(2, 3)
+    one = Tensor(1, dtype=dtypes.int64)
+    self.assertEqual(data_ptr(a[1]), data_ptr(a[one.cast(dtypes.int32)]))
+    self.assertEqual(data_ptr(a[1]), data_ptr(a[one.cast(dtypes.int16)]))
+
+  def test_getitem_scalars_simple_folding(self):
+    a = Tensor.randn(2, 3)
+    zero = Tensor(0, dtype=dtypes.int64)
+    one = Tensor(1, dtype=dtypes.int64)
+    self.assertEqual(data_ptr(a[0, 1]), data_ptr(a[zero, one]))
+
   def test_basic_advanced_combined(self):
     # From the NumPy indexing example
     x = Tensor.arange(0, 12).reshape(4, 3)
@@ -1320,8 +1337,9 @@ class TestNumpy(unittest.TestCase):
     self.assertRaises(IndexError, lambda: a[0, 0, -1.4])
     self.assertRaises(IndexError, lambda: a[-1.4, 0, 0])
     self.assertRaises(IndexError, lambda: a[0, -1.4, 0])
-    self.assertRaises(IndexError, lambda: a[0.0:, 0.0])
-    self.assertRaises(IndexError, lambda: a[0.0:, 0.0,:])
+    # these two trigger slice internal type verification first
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0])
+    self.assertRaises(TypeError, lambda: a[0.0:, 0.0,:])
 
   def test_none_index(self):
     # `None` index adds newaxis
@@ -1554,4 +1572,4 @@ class TestNumpy(unittest.TestCase):
   '''
 
 if __name__ == '__main__':
-  unittest.main()
+  unittest.main()

test/models/test_real_world.py

@@ -17,24 +17,25 @@ from extra.models.unet import ResBlock
 
 global_mem_used = 0
 def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jitted=False):
-  tms = []
-  for _ in range(4):
-    early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
-    GlobalCounters.reset()
-    Device[Device.DEFAULT].synchronize()
-    st = time.perf_counter_ns()
-    model(*early_gen)
-    Device[Device.DEFAULT].synchronize()
-    tms.append(time.perf_counter_ns() - st)
-  mem_used = GlobalCounters.mem_used - global_mem_used
+  with Context(JIT=2):
+    tms = []
+    for _ in range(4):
+      early_gen = [x.realize() if isinstance(x, Tensor) else x for x in gen()]
+      GlobalCounters.reset()
+      Device[Device.DEFAULT].synchronize()
+      st = time.perf_counter_ns()
+      model(*early_gen)
+      Device[Device.DEFAULT].synchronize()
+      tms.append(time.perf_counter_ns() - st)
+    mem_used = GlobalCounters.mem_used - global_mem_used
 
-  # TODO: jit should expose this correctly with graph
-  kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
-  print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
-  assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB"
-  assert not kernels_used or kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels"
-  if all_jitted:
-    assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501
+    # TODO: jit should expose this correctly with graph
+    kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
+    print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
+    assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB - {mem_used/1e9:.2} GB used"
+    assert not kernels_used or kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels"
+    if all_jitted:
+      assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted"  # noqa: E501
 
 class TestRealWorld(unittest.TestCase):
   def setUp(self):
@@ -68,7 +69,7 @@ class TestRealWorld(unittest.TestCase):
     def test(t, t2):
       for l in model: t = l(t, t2)
       return t.realize()
-    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.01, 43)
+    helper_test("test_unet_resblock", lambda: (Tensor.empty(4, 16, 8, 8), Tensor.empty(1, 24)), test, 0.01, 37)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_llama(self):
@@ -80,7 +81,7 @@ class TestRealWorld(unittest.TestCase):
     @TinyJit
     def test(t): return model(t, 0).realize()
     # TODO: test first token vs rest properly
-    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.27, 192, all_jitted=True)
+    helper_test("test_llama", lambda: (Tensor([[1,2,3,4]]),), test, 0.27, 168, all_jitted=True)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_gpt2(self):
@@ -92,7 +93,7 @@ class TestRealWorld(unittest.TestCase):
     @TinyJit
     def test(t, v):
       with Context(JIT=0): return model(t, v).realize()
-    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 164 if CI else 468, all_jitted=True)
+    helper_test("test_gpt2", lambda: (Tensor([[1,]]),Variable("pos", 1, 100).bind(1)), test, 0.23 if CI else 0.9, 137 if CI else 396, all_jitted=True)
 
   @unittest.skipIf(CI and Device.DEFAULT == "CLANG", "slow")
   def test_train_mnist(self):
@@ -110,7 +111,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 127)
+      helper_test("train_mnist", lambda: (Tensor.randn(BS, 1, 28, 28),), train, 0.07, 63)
 
   @unittest.skipIf(CI and Device.DEFAULT in {"CLANG", "GPU", "LLVM"}, "slow")
   def test_train_cifar(self):
@@ -127,7 +128,7 @@ class TestRealWorld(unittest.TestCase):
         loss.backward()
         optimizer.step()
 
-      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, (1.0/48)*BS, 142 if CI else 154)   # it's 154 on metal
+      helper_test("train_cifar", lambda: (Tensor.randn(BS, 3, 32, 32),), train, (1.0/48)*BS, 123)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.float16), "need dtypes.float16")
   def test_train_cifar_hyp(self):

test/test_arange.py

@@ -1,11 +1,12 @@
 import unittest, contextlib
 import numpy as np
-from tinygrad import Tensor, GlobalCounters, dtypes, nn
+from tinygrad import Tensor, GlobalCounters, dtypes, nn, Device
 from tinygrad.helpers import CI, Context, getenv
 from tinygrad.engine.realize import run_schedule
 from tinygrad.codegen.kernel import Opt, OptOps, Kernel, KernelOptError
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from tinygrad.engine.search import get_kernel_actions
+from tinygrad.ops import Ops
 
 class TestArange(unittest.TestCase):
   def _get_flops(self, N, opts=None):
@@ -65,6 +66,7 @@ class TestArange(unittest.TestCase):
     return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, amt=0)])
 
 class TestIndexing(unittest.TestCase):
+  @unittest.expectedFailure
   def test_arange_2_reduce(self):
     needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
     needle[1337] = 1
@@ -86,7 +88,7 @@ class TestIndexing(unittest.TestCase):
     print("*** indexing ***")
     with Context(NOOPT=1, FUSE_ARANGE=1):
       GlobalCounters.reset()
-      rng = Tensor.ones(4, 256, 16384, dtype=dtypes.int)._cumsum(axis=-1, _first_zero=True).reshape(4, 256, 16384, 1)
+      rng = Tensor.ones(4, 256, 16384, dtype=dtypes.int)._cumalu(axis=-1, op=Ops.ADD, _include_initial=True).reshape(4, 256, 16384, 1)
       idxs = idxs.reshape(4,1,1,1).expand(4, 256, 16384, 1)
       reshape_dataset = dataset.T.reshape(1, 256, 16384, 1).expand(4, 256, 16384, 1)
       full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, 256, 16384, 1))
@@ -138,7 +140,7 @@ class TestIndexing(unittest.TestCase):
       np.testing.assert_equal(X.numpy(), 0)
 
   @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
-  def test_index_mnist(self, noopt=1, op_limit=512*784*5):
+  def test_index_mnist(self, noopt=1, op_limit=512*784*13):
     from tinygrad.nn.datasets import mnist
     X_train, Y_train, _, _ = mnist()
     with Context(NOOPT=noopt, FUSE_ARANGE=1, SPLIT_REDUCEOP=0):
@@ -152,7 +154,7 @@ class TestIndexing(unittest.TestCase):
   @unittest.skip("not ready")
   def test_index_mnist_opt(self): self.test_index_mnist(0)
 
-  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
+  @unittest.skipIf(getenv("PTX") or Device.DEFAULT == "WEBGPU", "broken on ptx and WebGPU for some reason")
   def test_llama_embedding(self, noopt=1, op_limit=65536):
     # llama3 is 128256
     vocab_size, embed_size = (10, 3) if CI else (32000, 4096)

test/test_const_folding.py

@@ -10,7 +10,7 @@ def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
   schedule = create_schedule(t.lazydata.lbs)
   asts = [s for s in schedule if s.ast.op is Ops.SINK]
-  assert len(asts) == desired_count
+  assert len(asts) == desired_count, f"{len(asts)} != {desired_count}"
 
 class TestUnaryOpsConstFolding(unittest.TestCase):
   def test_all_consts_ops(self):
@@ -103,7 +103,8 @@ class TestIndexingConstFolding(unittest.TestCase):
   def test_scalar_index(self):
     t = Tensor.arange(16).float().reshape(1,1,4,4).realize()
     _check_ast_count(0, t[:,:,Tensor(1),:])
-    _check_ast_count(0, t[:,:,Tensor(1)+2,:])
+    # NOTE: this is no longer supported because the 1+2 isn't folding early.
+    #_check_ast_count(0, t[:,:,Tensor(1)+2,:])
     _check_ast_count(0, t[:,:,Tensor(1),Tensor(0)])
 
   @unittest.expectedFailure
@@ -158,6 +159,37 @@ class TestReduceOpsConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum())
     np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2)
 
+  def test_bool_zero_max(self):
+    _check_ast_count(0, Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)))
+    np.testing.assert_equal(Tensor.full((1, 2), True).shrink(((0, 1), (0, 0))).max((1, 0)).numpy(), False)
+
+  def test_zero_size_ops(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]: # lambda x:x.max() NOTE: numpy gives "reduction operation maximum which has no identity"
+      _check_ast_count(0, reduceop(Tensor.empty(1, 0)))
+      np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty(shape)))
+
+  def test_zero_size_ops_view(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
+      _check_ast_count(0, reduceop(Tensor.empty(1, 0, 4).permute((1, 2, 0)).contiguous()))
+      np.testing.assert_equal(reduceop(Tensor.empty(shape:=(1, 0))).numpy(), reduceop(np.empty((shape))))
+
+  def test_zero_size_ops_realized(self):
+    for reduceop in [lambda x:x.prod(), lambda x:x.sum()]:
+      _check_ast_count(0, reduceop((Tensor.randn(0, 1)+1).realize()))
+      np.testing.assert_equal(reduceop((Tensor.randn(shape:=(0, 1))+1).realize()).numpy(), reduceop(np.empty(shape)))
+
+  def test_zero_size_realize_folded(self):
+    # non contiguous folded output doesn't realize
+    _check_ast_count(0, Tensor.empty(1, 0).sum())
+    # contiguous folded const can still schedule
+    a = Tensor.empty(1, 0).sum().contiguous()
+    _check_ast_count(2, a+2)
+    self.assertIsNotNone(a.lazydata.base.realized)
+    np.testing.assert_equal((Tensor.empty(1, 0).sum().contiguous()+2).numpy(), 2)
+    # otherwise we just fuse it
+    _check_ast_count(1, (Tensor.empty(1, 0).sum()+2).contiguous())
+    np.testing.assert_equal((Tensor.empty(1, 0).sum()+2).numpy(), 2)
+
   def test_const_prod(self):
     _check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod())
     np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3))