experiments with multi being range

fix assert in test_schedule (#12745 )
* fix assert in test_schedule updated kernel counts and some old tests * fix
2025-10-17 14:11:55 +08:00 · 2025-10-16 15:39:50 -04:00 · 2025-10-16 14:11:33 -04:00 · 2025-10-16 12:43:27 -04:00 · 2025-10-16 11:40:36 -04:00 · 2025-10-16 11:17:47 -04:00
179 changed files with 25354 additions and 3453 deletions
@@ -41,6 +41,10 @@ inputs:
    description: "Install LLVM?"
    required: false
    default: 'false'
+  mesa:
+    description: "Install mesa"
+    required: false
+    default: 'false'
 runs:
  using: "composite"
  steps:
@@ -289,3 +293,13 @@ runs:
      if: inputs.llvm == 'true' && runner.os == 'macOS'
      shell: bash
      run: brew install llvm@20
+
+    # **** mesa ****
+    - name: Install mesa (linux)
+      if: inputs.mesa == 'true' && runner.os == 'Linux'
+      shell: bash
+      run: sudo curl -L https://github.com/sirhcm/tinymesa/releases/download/tinymesa-32dc66c/libtinymesa_cpu-mesa-25.2.4-linux-amd64.so -o /usr/lib/libtinymesa_cpu.so
+    - name: Install mesa (macOS)
+      if: inputs.mesa == 'true' && runner.os == 'macOS'
+      shell: bash
+      run: brew install sirhcm/tinymesa/tinymesa
@@ -36,8 +36,9 @@ jobs:
        cuda: 'true'
        webgpu: 'true'
        llvm: 'true'
+        pydeps: 'pyyaml mako'
    - name: Install autogen support packages
-      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev
+      run: sudo apt-get install -y --no-install-recommends llvm-14-dev libclang-14-dev llvm-20-dev
    - name: Verify OpenCL autogen
      run: |
        cp tinygrad/runtime/autogen/opencl.py /tmp/opencl.py.bak
@@ -89,3 +90,8 @@ jobs:
        cp tinygrad/runtime/autogen/llvm.py /tmp/llvm.py.bak
        ./autogen_stubs.sh llvm
        diff /tmp/llvm.py.bak tinygrad/runtime/autogen/llvm.py
+    - name: Verify mesa autogen
+      run: |
+        cp tinygrad/runtime/autogen/mesa.py /tmp/mesa.py.bak
+        ./autogen_stubs.sh mesa
+        diff /tmp/mesa.py.bak tinygrad/runtime/autogen/mesa.py
@@ -51,17 +51,18 @@ jobs:
        rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
    - name: reset process replay
      run: python3.11 test/external/process_replay/reset.py
+    - name: Print macOS version
+      run: sw_vers
    - name: Run Stable Diffusion
-      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
+      run: BENCHMARK_LOG=stable_diffusion JIT=1 ASSERT_MIN_STEP_TIME=720 python3.11 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
    - name: Run Stable Diffusion without fp16
-      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=900 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
+      run: BENCHMARK_LOG=stable_diffusion_fp32 JIT=1 ASSERT_MIN_STEP_TIME=800 python3.11 examples/stable_diffusion.py --seed 0 --noshow --timing | tee sd_no_fp16.txt
    - name: Run Stable Diffusion v2
      # TODO: very slow step time
-      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=10000 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
+      run: BENCHMARK_LOG=stable_diffusion_v2 JIT=1 ASSERT_MIN_STEP_TIME=4500 python3.11 examples/sdv2.py --fp16 --seed 0 --noshow --timing | tee sdv2.txt
    # process replay can't capture this, the graph is too large
-    # TODO: too slow
-    # - name: Run SDXL
-    #   run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+    - name: Run SDXL
+      run: BENCHMARK_LOG=stable_diffusion_xl ASSERT_MIN_STEP_TIME=5000 CAPTURE_PROCESS_REPLAY=0 JIT=1 python3.11 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
    - name: Run model inference benchmark
      run: METAL=1 python3.11 test/external/external_model_benchmark.py
    - name: Test speed vs torch
@@ -632,7 +633,7 @@ jobs:
      run: PYTHONPATH="." ASSERT_MIN_STEP_TIME=12 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.9/selfdrive/modeld/models/dmonitoring_model.onnx
    - name: openpilot compile3 Space Lab policy + vision
      run: |
-        PYTHONPATH="." ASSERT_MIN_STEP_TIME=4 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
+        PYTHONPATH="." ASSERT_MIN_STEP_TIME=5 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/22aec22a10ce09384d4a4af2a0bbff08d54af7e0c888503508f356fae4ff0e29
        PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 QCOM=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/c824f68646a3b94f117f01c70dc8316fb466e05fbd42ccdba440b8a8dc86914b
    - name: benchmark MobileNetV2 on DSP
      run: |
@@ -641,7 +642,7 @@ jobs:
        ln -s /data/home/tiny/tinygrad/testsig-*.so .
        PYTHONPATH=. CC=clang-19 CPU=1 CPU_LLVM=0 QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
        # benchmark on DSP with NOOPT=1, the devectorizer has issues
-        PYTHONPATH=. CC=clang-19 DSP=1 DONT_REALIZE_EXPAND=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
+        PYTHONPATH=. CC=clang-19 DSP=1 NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
    - name: Run process replay tests
      run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
    - uses: actions/upload-artifact@v4
@@ -12,7 +12,7 @@ jobs:
  run_script_job:
    runs-on: [self-hosted, Linux, tinybox]
    if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 360
+    timeout-minutes: 720

    steps:
    - name: Checkout Code
@@ -89,64 +89,65 @@ jobs:
        clang -O2 recognize.c -lm -o recognize
        cat test/models/efficientnet/Chicken.jpg | ./recognize | grep cock

-  torchbackend:
-    name: Torch Backend Tests
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        pydeps: "pillow torchvision expecttest"
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Lint with ruff
-      run: |
-        pip3 install --upgrade --force-reinstall ruff==0.11.0
-        python3 -m ruff check extra/torch_backend/backend.py
-    - name: Test one op
-      run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
-    - name: Test ResNet-18
-      run: DEBUG=2 python3 extra/torch_backend/example.py
-    - name: My (custom) tests
-      run: python3 extra/torch_backend/test.py
-    - name: Test one op in torch tests
-      run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
-    - name: Test Ops with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
-    - name: Test in-place operations on views
-      run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
-    - name: Test multi-gpu
-      run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py
+  # TODO: fix the torch backend and reenable
+  # torchbackend:
+  #   name: Torch Backend Tests
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       pydeps: "pillow torchvision expecttest"
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Lint with ruff
+  #     run: |
+  #       pip3 install --upgrade --force-reinstall ruff==0.11.0
+  #       python3 -m ruff check extra/torch_backend/backend.py
+  #   - name: Test one op
+  #     run: FORWARD_ONLY=1 TINY_BACKEND=1 python3 test/test_ops.py TestOps.test_add
+  #   - name: Test ResNet-18
+  #     run: DEBUG=2 python3 extra/torch_backend/example.py
+  #   - name: My (custom) tests
+  #     run: python3 extra/torch_backend/test.py
+  #   - name: Test one op in torch tests
+  #     run: DEBUG=2 python3 extra/torch_backend/torch_tests.py TestTinyBackendPRIVATEUSE1.test_unary_log_tiny_float32
+  #   - name: Test Ops with TINY_BACKEND
+  #     run: CPU=1 CPU_LLVM=1 LLVMOPT=0 TINY_BACKEND=1 python3 -m pytest -n auto test/test_ops.py --durations=20
+  #   - name: Test in-place operations on views
+  #     run: TORCH_DEBUG=1 python3 extra/torch_backend/test_inplace.py
+  #   - name: Test multi-gpu
+  #     run: CPU=1 CPU_LLVM=1 GPUS=4 TORCH_DEBUG=1 python3 extra/torch_backend/test_multigpu.py

-  torchbackendmore:
-    name: Torch Backend Tests More
-    runs-on: ubuntu-latest
-    timeout-minutes: 15
-    steps:
-    - name: Checkout Code
-      uses: actions/checkout@v4
-    - name: Setup Environment
-      uses: ./.github/actions/setup-tinygrad
-      with:
-        key: torch-backend-pillow-torchvision-et-pt
-        deps: testing_minimal
-        llvm: 'true'
-    - name: Install ninja
-      run: |
-        sudo apt update || true
-        sudo apt install -y --no-install-recommends ninja-build
-    - name: Test beautiful_mnist in torch with TINY_BACKEND
-      run: CPU=1 CPU_LLVM=1 TARGET_EVAL_ACC_PCT=96.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
-    - name: Test some torch tests (expect failure)
-      run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true
+  # torchbackendmore:
+  #   name: Torch Backend Tests More
+  #   runs-on: ubuntu-latest
+  #   timeout-minutes: 15
+  #   steps:
+  #   - name: Checkout Code
+  #     uses: actions/checkout@v4
+  #   - name: Setup Environment
+  #     uses: ./.github/actions/setup-tinygrad
+  #     with:
+  #       key: torch-backend-pillow-torchvision-et-pt
+  #       deps: testing_minimal
+  #       llvm: 'true'
+  #   - name: Install ninja
+  #     run: |
+  #       sudo apt update || true
+  #       sudo apt install -y --no-install-recommends ninja-build
+  #   - name: Test beautiful_mnist in torch with TINY_BACKEND
+  #     run: STEPS=20 CPU=1 TARGET_EVAL_ACC_PCT=90.0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
+  #   - name: Test some torch tests (expect failure)
+  #     run: python3 -m pytest extra/torch_backend/torch_tests.py -v --tb=no || true

  bepython:
    name: Python Backend
@@ -238,8 +239,6 @@ jobs:
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
        python3 -m ruff check examples/mlperf/ --ignore E501
-    - name: Lint tinygrad with pylint
-      run: python -m pylint tinygrad/
    - name: Run mypy
      run: |
        python -m mypy --strict-equality --lineprecision-report .
@@ -274,6 +273,8 @@ jobs:
    #   run: NULL=1 DEBUG=1 python3 examples/sdxl.py --seed 0 --noshow --timing --fakeweights
    - name: Run Clip tests for SD MLPerf on NULL backend
      run: NULL=1 python -m pytest -n=auto test/external/mlperf_stable_diffusion/external_test_models.py::TestOpenClip --durations=20
+    - name: Run AMD emulated BERT training on NULL backend
+      run: EMULATE=AMD_RDNA4 NULL=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=1 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
    # TODO: support fake weights
    #- name: Run LLaMA 7B on 4 fake devices
    #  run: NULL=1 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 3 --temperature 0 --timing
@@ -309,10 +310,6 @@ jobs:
      run: python test/external/fuzz_symbolic.py
    - name: Fuzz Test fast idiv
      run: python test/external/fuzz_fast_idiv.py
-    - name: Fuzz Test shapetracker
-      run: |
-        python test/external/fuzz_shapetracker.py
-        python test/external/fuzz_shapetracker_math.py
    - name: Fuzz Test shape ops
      run: python test/external/fuzz_shape_ops.py

@@ -377,7 +374,7 @@ jobs:
          llvm: 'true'
      - name: Test openpilot model kernel count and gate usage
        run: |
-          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2041 ALLOWED_GATED_READ_IMAGE=41 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          ALLOWED_KERNEL_COUNT=190 ALLOWED_READ_IMAGE=2081 ALLOWED_GATED_READ_IMAGE=28 FLOAT16=0 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot alt model correctness (float32)
        run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot fastvits model correctness (float32)
@@ -386,8 +383,6 @@ jobs:
      #   run: FLOAT16=0 DEBUGCL=1 CL=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
      - name: Test openpilot LLVM compile
        run: CPU=1 CPU_LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot compile4
-        run: NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -677,7 +672,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [llvm, cpu, opencl]
+        backend: [llvm, cpu, opencl, lvp]

    name: Linux (${{ matrix.backend }})
    runs-on: ubuntu-22.04
@@ -691,9 +686,10 @@ jobs:
          key: ${{ matrix.backend }}-minimal
          deps: testing_minimal
          opencl: ${{ matrix.backend == 'opencl' && 'true' }}
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' }}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'opencl' && 'CL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
          python3 -c "from tinygrad import Device; assert Device.DEFAULT in ['CPU','CL'], Device.DEFAULT"
@@ -895,7 +891,7 @@ jobs:
    strategy:
      fail-fast: false
      matrix:
-        backend: [metal, llvm, cpu]
+        backend: [metal, llvm, cpu, lvp]
    name: MacOS (${{ matrix.backend }})
    runs-on: macos-15
    timeout-minutes: 20
@@ -908,12 +904,13 @@ jobs:
          key: macos-${{ matrix.backend }}-minimal
          deps: testing_minimal
          pydeps: "capstone"
-          llvm: ${{ matrix.backend == 'llvm' && 'true' }}
+          llvm: ${{ matrix.backend == 'llvm' || matrix.backend == 'lvp' }}
+          mesa: ${{ matrix.backend == 'lvp' && 'true' }}
      - name: Set env
-        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1'}}" >> $GITHUB_ENV
+        run: printf "${{ matrix.backend == 'llvm' && 'CPU=1\nCPU_LLVM=1' || matrix.backend == 'cpu' && 'CPU=1\nCPU_LLVM=0\nCPU_COUNT=2' || matrix.backend == 'metal' && 'METAL=1' || matrix.backend == 'lvp' && 'CPU=1\nCPU_LVP=1' }}" >> $GITHUB_ENV
      - name: Check Device.DEFAULT and print some source
        run: |
-          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
+          python -c "from tinygrad import Device; assert Device.DEFAULT == {'LLVM':'CPU','LVP':'CPU'}.get(x:='${{ matrix.backend }}'.upper(), x), Device.DEFAULT"
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run pytest (${{ matrix.backend }})
        run: python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit --durations=20
@@ -38,6 +38,7 @@ extra/huggingface_onnx/models/*
 extra/huggingface_onnx/*.yaml
 extra/weights
 venv
+venv_sd_mlperf
 examples/**/net.*[js,json]
 examples/**/*.safetensors
 node_modules
@@ -20,21 +20,15 @@ repos:
        language: system
        always_run: true
        pass_filenames: false
-      - id: tests
-        name: subset of tests
-        entry: env PYTHONPATH="." python3 -m pytest -n=4 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
-        language: system
-        always_run: true
-        pass_filenames: false
      - id: example
-        name: multi device tests
+        name: test all devices
        entry: python3 test/external/external_test_example.py
        language: system
        always_run: true
        pass_filenames: false
-      - id: pylint
-        name: pylint
-        entry: python3 -m pylint tinygrad/
+      - id: tests
+        name: subset of tests
+        entry: env PYTHONPATH="." python3 -m pytest -n=8 test/test_ops.py test/test_dtype.py test/test_schedule.py test/test_assign.py
        language: system
        always_run: true
        pass_filenames: false
@@ -30,10 +30,6 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=

-# When enabled, pylint would attempt to guess common misconfiguration and emit
-# user-friendly hints instead of false-positive error messages
-suggestion-mode=yes
-
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no
@@ -435,8 +435,8 @@ generate_sqtt() {
    -o extra/sqtt/rocprof/rocprof.py
  fixup extra/sqtt/rocprof/rocprof.py
  sed -i '1s/^/# pylint: skip-file\n/' extra/sqtt/rocprof/rocprof.py
-  sed -i "s/import ctypes/import ctypes\nfrom tinygrad.helpers import fetch/g" extra/sqtt/rocprof/rocprof.py
-  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so')))|g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util/g" extra/sqtt/rocprof/rocprof.py
+  sed -i "s|FunctionFactoryStub()|ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder'))|g" extra/sqtt/rocprof/rocprof.py
 }

 generate_webgpu() {
@@ -461,6 +461,85 @@ generate_libusb() {
  python3 -c "import tinygrad.runtime.autogen.libusb"
 }

+generate_mesa() {
+  MESA_TAG="mesa-25.2.4"
+  MESA_SRC=/tmp/mesa-$MESA_TAG
+  TINYMESA_TAG=tinymesa-32dc66c
+  TINYMESA_DIR=/tmp/tinymesa-$MESA_TAG-$TINYMESA_TAG/
+  TINYMESA_SO=$TINYMESA_DIR/libtinymesa_cpu.so
+  if [ ! -d "$MESA_SRC" ]; then
+    git clone --depth 1 --branch $MESA_TAG https://gitlab.freedesktop.org/mesa/mesa.git $MESA_SRC
+    pushd .
+    cd $MESA_SRC
+    git reset --hard $MESA_COMMIT_HASH
+    # clang 14 doesn't support packed enums
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/headers/nv_device_info.h
+    sed -i "s/enum \w\+ \(\w\+\);$/uint8_t \1;/" $MESA_SRC/src/nouveau/compiler/nak.h
+    sed -i "s/nir_instr_type \(\w\+\);/uint8_t \1;/" $MESA_SRC/src/compiler/nir/nir.h
+    mkdir -p gen/util/format
+    python3 src/util/format/u_format_table.py src/util/format/u_format.yaml --enums > gen/util/format/u_format_gen.h
+    python3 src/compiler/nir/nir_opcodes_h.py > gen/nir_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/nir/nir_intrinsics_indices_h.py --outdir gen
+    python3 src/compiler/nir/nir_builder_opcodes_h.py > gen/nir_builder_opcodes.h
+    python3 src/compiler/nir/nir_intrinsics_h.py --outdir gen
+    python3 src/compiler/builtin_types_h.py gen/builtin_types.h
+    popd
+  fi
+
+  if [ ! -d "$TINYMESA_DIR" ]; then
+    mkdir $TINYMESA_DIR
+    curl -L https://github.com/sirhcm/tinymesa/releases/download/$TINYMESA_TAG/libtinymesa_cpu-$MESA_TAG-linux-amd64.so -o $TINYMESA_SO
+  fi
+
+  clang2py -k cdefstu \
+    $MESA_SRC/src/compiler/nir/nir.h \
+    $MESA_SRC/src/compiler/nir/nir_builder.h \
+    $MESA_SRC/src/compiler/nir/nir_shader_compiler_options.h \
+    $MESA_SRC/src/compiler/nir/nir_serialize.h \
+    $MESA_SRC/gen/nir_intrinsics.h \
+    $MESA_SRC/src/nouveau/headers/nv_device_info.h \
+    $MESA_SRC/src/nouveau/compiler/nak.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_passmgr.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_misc.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_type.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_init.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_nir.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_struct.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_jit_types.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_flow.h \
+    $MESA_SRC/src/gallium/auxiliary/gallivm/lp_bld_const.h \
+    $MESA_SRC/src/compiler/glsl_types.h \
+    $MESA_SRC/src/util/blob.h \
+    $MESA_SRC/src/util/ralloc.h \
+    --clang-args="-DHAVE_ENDIAN_H -DHAVE_STRUCT_TIMESPEC -DHAVE_PTHREAD -I$MESA_SRC/src -I$MESA_SRC/include -I$MESA_SRC/gen -I$MESA_SRC/src/compiler/nir -I$MESA_SRC/src/gallium/auxiliary -I$MESA_SRC/src/gallium/include -I$(llvm-config-20 --includedir)" \
+    -l $TINYMESA_SO \
+    -o $BASE/mesa.py
+
+  LVP_NIR_OPTIONS=$(./extra/mesa/lvp_nir_options.sh $MESA_SRC)
+
+  fixup $BASE/mesa.py
+  patch_dlopen $BASE/mesa.py tinymesa_cpu "(BASE:=os.getenv('MESA_PATH', f\"/usr{'/local/' if helpers.OSX else '/'}lib\"))+'/libtinymesa_cpu'+(EXT:='.dylib' if helpers.OSX else '.so')" "f'{BASE}/libtinymesa{EXT}'" "f'{brew_prefix()}/lib/libtinymesa_cpu.dylib'"
+  echo "lvp_nir_options = gzip.decompress(base64.b64decode('$LVP_NIR_OPTIONS'))" >> $BASE/mesa.py
+  cat <<EOF | sed -i "/import ctypes.*/r /dev/stdin" $BASE/mesa.py
+def brew_prefix():
+  try: return subprocess.check_output(['brew', '--prefix', 'tinymesa']).decode().strip()
+  except Exception: return ''
+EOF
+  sed -i "/in_dll/s/.*/try: &\nexcept AttributeError: pass/" $BASE/mesa.py
+  sed -i "s/import ctypes/import ctypes, ctypes.util, os, gzip, base64, subprocess, tinygrad.helpers as helpers/" $BASE/mesa.py
+  sed -i "s/ctypes.CDLL('.\+')/(dll := _try_dlopen_tinymesa_cpu())/" $BASE/mesa.py
+  echo "def __getattr__(nm): raise AttributeError() if dll else FileNotFoundError(f'libtinymesa not found (MESA_PATH={BASE}). See https://github.com/sirhcm/tinymesa ($TINYMESA_TAG, $MESA_TAG)')" >> $BASE/mesa.py
+  sed -i "s/ctypes.glsl_base_type/glsl_base_type/" $BASE/mesa.py
+  # bitfield bug in clang2py
+  sed -i "s/('fp_fast_math', ctypes.c_bool, 9)/('fp_fast_math', ctypes.c_uint32, 9)/" $BASE/mesa.py
+  sed -i "s/('\(\w\+\)', pipe_shader_type, 8)/('\1', ctypes.c_ubyte)/" $BASE/mesa.py
+  sed -i "s/\([0-9]\+\)()/\1/" $BASE/mesa.py
+  sed -i "s/\(struct_nir_builder._pack_\) = 1/\1 = 0/" $BASE/mesa.py
+  python3 -c "import tinygrad.runtime.autogen.mesa"
+}
+
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -484,6 +563,7 @@ elif [ "$1" == "pci" ]; then generate_pci
 elif [ "$1" == "vfio" ]; then generate_vfio
 elif [ "$1" == "webgpu" ]; then generate_webgpu
 elif [ "$1" == "libusb" ]; then generate_libusb
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+elif [ "$1" == "mesa" ]; then generate_mesa
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu; generate_mesa
 else echo "usage: $0 <type>"
 fi
@@ -232,7 +232,7 @@ if __name__ == "__main__":
  gpt2 = GPT2.build_gguf(args.model_size) if args.model_size.startswith("gpt2_gguf_") else GPT2.build(args.model_size)

  if args.benchmark != -1:
-    gpt2.model(Tensor.rand(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
+    gpt2.model(Tensor.randint(args.batch_size, args.benchmark), Variable("a", 0, MAX_CONTEXT).bind(0)).realize()
  else:
    texts = gpt2.generate(args.prompt, args.count, args.temperature, timing=args.timing, batch_size=args.batch_size)
    if not args.noshow:
@@ -1188,7 +1188,9 @@ def train_bert():
      if MLLOGGER and RUNMLPERF:
        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*GBS, "step_num": i})
      if getenv("RESET_STEP"): train_step_bert.reset()
-      elif getenv("FREE_INTERMEDIATE", 1) and train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
+      elif getenv("FREE_INTERMEDIATE", 0) and train_step_bert.captured is not None:
+        # TODO: FREE_INTERMEDIATE nan'ed after jit step 2
+        train_step_bert.captured.free_intermediates()
      eval_lm_losses = []
      eval_clsf_losses = []
      eval_lm_accs = []
@@ -1222,7 +1224,7 @@ def train_bert():
          return

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

      del eval_data
      avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=1 BS=128 EVAL_BS=128
+
+export IGNORE_OOB=1
+
+export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+# export BEAM_LOG_SURPASS_MAX=1
+# export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_8xMI300X"
+export DEFAULT_FLOAT="HALF" GPUS=8 BS=1024 EVAL_BS=1024
+
+# similar to https://github.com/mlcommons/training_results_v3.1/blob/d06288b2bd675a9d88e0e6181f5bb5626b71ec19/Quanta_Cloud_Technology/results/D54U-3U/bert/result_1.txt#L54
+export OPT_BASE_LEARNING_RATE=0.0011 OPT_LAMB_BETA_1=0.60466 OPT_LAMB_BETA_2=0.85437 DECAY=0.1
+export TRAIN_STEPS=3900
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
+
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,17 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,28 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=8 BEAM_UOPS_MAX=10000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,69 @@
+# 1. Problem
+
+This problem uses BERT for NLP.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+Also install gdown (for dataset), numpy, tqdm and tensorflow.
+```
+pip install gdown numpy tqdm tensorflow
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download and verify data
+
+### 1. Download raw data
+
+```
+BASEDIR="/raid/datasets/wiki" WIKI_TRAIN=1 VERIFY_CHECKSUM=1 python3 extra/datasets/wikipedia_download.py
+```
+
+### 2. Preprocess train and validation data
+
+Note: The number of threads used for preprocessing is limited by available memory. With 128GB of RAM, a maximum of 16 threads is recommended. 
+
+#### Training:
+```
+BASEDIR="/raid/datasets/wiki" NUM_WORKERS=16 python3 extra/datasets/wikipedia.py pre-train all
+```
+
+Generating a specific topic (Between 0 and 499)
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-train 42
+```
+
+#### Validation:
+```
+BASEDIR="/raid/datasets/wiki" python3 extra/datasets/wikipedia.py pre-eval
+```
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_green/run_and_time.sh
+```
+
+### tinybox_red
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh
+```
+### tinybox_8xMI300X
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh
+```
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BEAM_LOG_SURPASS_MAX=1
+export BASEDIR="/raid/datasets/wiki"
+
+export RESET_STEP=1
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+export WANDB=1 PARALLEL=0
+
+RUNMLPERF=1 python3 examples/mlperf/model_train.py
@@ -0,0 +1,31 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="bert"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export IGNORE_OOB=1
+export REWRITE_STACK_LIMIT=500000
+
+export BEAM=5 BEAM_UOPS_MAX=8000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/wiki"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="bert_red_${DATETIME}_${SEED}.log"
+
+export HCQDEV_WAIT_TIMEOUT_MS=100000  # prevents hang?
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=2
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=10 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,50 @@
+# 1. Problem
+
+This problem uses the ResNet-50 CNN to do image classification.
+
+## Requirements
+
+Install tinygrad and mlperf-logging from master.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+### tinybox_red
+Disable cwsr
+This is the default on production tinybox red.
+```
+sudo vi /etc/modprobe.d/amdgpu.conf
+cat <<EOF > /etc/modprobe.d/amdgpu.conf
+options amdgpu cwsr_enable=0
+EOF
+sudo update-initramfs -u
+sudo reboot
+
+# validate
+sudo cat /sys/module/amdgpu/parameters/cwsr_enable #= 0
+```
+
+# 2. Directions
+
+## Steps to download and verify data
+
+```
+IMGNET_TRAIN=1 python3 extra/datasets/imagenet_download.py
+```
+
+## Steps for one time setup
+
+### tinybox_red
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/setup.sh
+```
+
+## Steps to run benchmark
+```
+examples/mlperf/training_submission_v4.0/tinycorp/benchmarks/resnet/implementations/tinybox_red/run_and_time.sh
+```
@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=10 DEBUG=${DEBUG:-2}
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export EVAL_START_EPOCH=3 EVAL_FREQ=4
+
+export WANDB=1 PARALLEL=0
+
+python3 examples/mlperf/model_train.py
@@ -0,0 +1,26 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." AMD=1
+export MODEL="resnet"
+export SUBMISSION_PLATFORM="tinybox_red"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=1536 EVAL_BS=192
+
+export RESET_STEP=0
+
+export TRAIN_BEAM=4 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=96 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=${LOGMLPERF:-1}
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="resnet_red_${DATETIME}_${SEED}.log"
+
+# init
+sleep 5 && sudo rmmod amdgpu || true
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 EVAL_START_EPOCH=3 EVAL_FREQ=4 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,8 @@
+#!/bin/bash
+
+rocm-smi --setprofile compute
+rocm-smi --setmclk 3
+rocm-smi --setperflevel high
+
+# power cap to 350W
+echo "350000000" | sudo tee /sys/class/drm/card{1..6}/device/hwmon/hwmon*/power1_cap
@@ -0,0 +1,38 @@
+# 1. Problem
+
+This problem uses RetinaNet for SSD.
+
+## Requirements
+
+Install tinygrad and mlperf-logging (uncomment mlperf from setup.py) from branch mlperf_training_v5.0.
+```
+git clone https://github.com/tinygrad/tinygrad.git
+python3 -m pip install -e ".[mlperf]"
+```
+
+Also install the following dependencies:
+```
+pip install tqdm numpy pycocotools boto3 pandas torch torchvision
+```
+
+### tinybox_green
+Install the p2p driver per [README](https://github.com/tinygrad/open-gpu-kernel-modules/blob/550.54.15-p2p/README.md)
+This is the default on production tinybox green.
+
+# 2. Directions
+
+## Steps to download data
+
+Run the following:
+```
+BASEDIR=/raid/datasets/openimages python3 extra/datasets/openimages.py
+```
+
+## Running
+
+### tinybox_green
+
+#### Steps to run benchmark
+```
+examples/mlperf/training_submission_v5.0/tinycorp/benchmarks/retinanet/implementations/tinybox_green/run_and_time.sh
+```
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,25 @@
+#!/bin/bash
+set -e  # Exit on any error
+set -o pipefail  # Make pipeline fail if any command fails
+
+export PYTHONPATH="." NV=1
+export MODEL="retinanet"
+export SUBMISSION_PLATFORM="tinybox_green"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+
+export TRAIN_BEAM=2 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+export IGNORE_JIT_FIRST_BEAM=1
+export BASEDIR="/raid/datasets/openimages"
+
+# pip install -e ".[mlperf]"
+export LOGMLPERF=1
+
+export SEED=$RANDOM
+DATETIME=$(date "+%m%d%H%M")
+LOGFILE="retinanet_green_${DATETIME}_${SEED}.log"
+
+# init
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+
+# run
+PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
@@ -0,0 +1,14 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export BENCHMARK=5 DEBUG=2
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+export PYTHONPATH="." AMD=1
+export MODEL="retinanet"
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=96 EVAL_BS=96
+export BASEDIR="/raid/datasets/openimages"
+
+# export RESET_STEP=0
+
+export TRAIN_BEAM=2 IGNORE_JIT_FIRST_BEAM=1 BEAM_UOPS_MAX=1500 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
+
+export WANDB=1 PARALLEL=0
+export RUNMLPERF=1
+
+python examples/mlperf/model_train.py
@@ -0,0 +1,38 @@
+{
+    "submitter": "tinycorp",
+    "division": "closed",
+    "status": "Available on-premise",
+    "system_name": "tinybox 8xMI300X",
+    "number_of_nodes": "1",
+    "host_processors_per_node": "2",
+    "host_processor_model_name": "AMD EPYC 9354",
+    "host_processor_core_count": "32",
+    "host_processor_vcpu_count": "64",
+    "host_processor_frequency": "",
+    "host_processor_caches": "",
+    "host_processor_interconnect": "",
+    "host_memory_capacity": "2304GB",
+    "host_storage_type": "NVMe SSD",
+    "host_storage_capacity": "3x 4TB raid array",
+    "host_networking": "",
+    "host_networking_topology": "",
+    "host_memory_configuration": "24x 96GB DDR5",
+    "accelerators_per_node": "8",
+    "accelerator_model_name": "AMD Instinct MI300X 192GB HBM3",
+    "accelerator_host_interconnect": "PCIe 5.0 x16",
+    "accelerator_frequency": "",
+    "accelerator_on-chip_memories": "",
+    "accelerator_memory_configuration": "HBM3",
+    "accelerator_memory_capacity": "192GB",
+    "accelerator_interconnect": "",
+    "accelerator_interconnect_topology": "",
+    "cooling": "air",
+    "hw_notes": "",
+    "framework": "tinygrad, branch mlperf_training_v5.0",
+    "other_software_stack": {
+        "python": "3.10.16",
+        "ROCm": "3.0.0+94441cb"
+    },
+    "operating_system": "Ubuntu 24.04.1 LTS",
+    "sw_notes": ""
+  }
@@ -0,0 +1,38 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox green",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "NVIDIA GeForce RTX 4090",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6X",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12",
+    "CUDA": "12.4"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -0,0 +1,37 @@
+{
+  "submitter": "tinycorp",
+  "division": "closed",
+  "status": "Available on-premise",
+  "system_name": "tinybox red",
+  "number_of_nodes": "1",
+  "host_processors_per_node": "1",
+  "host_processor_model_name": "AMD EPYC 7532",
+  "host_processor_core_count": "32",
+  "host_processor_vcpu_count": "64",
+  "host_processor_frequency": "",
+  "host_processor_caches": "",
+  "host_processor_interconnect": "",
+  "host_memory_capacity": "128GB",
+  "host_storage_type": "NVMe SSD",
+  "host_storage_capacity": "4 TB raid array + 1 TB boot",
+  "host_networking": "",
+  "host_networking_topology": "",
+  "host_memory_configuration": "8x 16GB DDR4",
+  "accelerators_per_node": "6",
+  "accelerator_model_name": "AMD Radeon RX 7900 XTX",
+  "accelerator_host_interconnect": "PCIe 4.0 x16",
+  "accelerator_frequency": "",
+  "accelerator_on-chip_memories": "",
+  "accelerator_memory_configuration": "GDDR6",
+  "accelerator_memory_capacity": "24GB",
+  "accelerator_interconnect": "",
+  "accelerator_interconnect_topology": "",
+  "cooling": "air",
+  "hw_notes": "",
+  "framework": "tinygrad, branch mlperf_training_v5.0",
+  "other_software_stack": {
+    "python": "3.10.12"
+  },
+  "operating_system": "Ubuntu 22.04.4",
+  "sw_notes": ""
+}
@@ -1,47 +0,0 @@
-import sys
-from tinygrad import Tensor, fetch, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp
-from tinygrad.nn.onnx import OnnxRunner
-from tinygrad.schedule.rangeify import get_rangeify_map
-from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.realize import run_schedule
-
-# NOLOCALS=1 CL=1 IMAGE=2 FLOAT16=1 VIZ=1 DEBUG=2 python3 examples/openpilot/compile4.py
-
-OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
-OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
-
-if __name__ == "__main__":
-  onnx_file = fetch(OPENPILOT_MODEL)
-  run_onnx = OnnxRunner(onnx_file)
-
-  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
-  out: Tensor = next(iter(run_onnx({k:v.to(None) for k,v in inputs.items()}).values())).to('cpu')
-  root = out.uop
-  targets = [x.uop for x in inputs.values()]
-  print(targets)
-
-  # TODO: abstract this from gradient?
-
-  # compute the target path (top down)
-  in_target_path: dict[UOp, bool] = {}
-  for u in root.toposort(): in_target_path[u] = any(x in targets or in_target_path[x] for x in u.src)
-  independent_set = {}
-  for u in root.toposort():
-    if in_target_path[u]:
-      for s in u.src:
-        if not in_target_path[s]:
-          independent_set[s] = None
-  independent = UOp.sink(*independent_set.keys())
-  kernelized = get_rangeify_map(independent)
-  independent = independent.substitute(kernelized)
-  schedule, var_vals = create_schedule_with_vars(independent)
-  run_schedule(schedule)
-
-  print("**** real ****")
-  GlobalCounters.reset()
-  out.uop = root.substitute(kernelized)
-  out.kernelize()
-
-  # realize
-  out.realize()
@@ -269,7 +269,8 @@ if __name__ == "__main__":

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

    if args.fp16:
      for k,v in get_state_dict(model).items():
@@ -19,8 +19,8 @@ from tinygrad.helpers import fetch, getenv

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

 def imagenet_dataloader(cnt=0):
  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
@@ -0,0 +1,23 @@
+#!/bin/sh
+
+if [ "$#" -ne 1 ] || ! [ -d $1 ]; then
+  echo "usage: $0 MESA_PREFIX"
+  exit 1
+fi
+
+TMP=$(mktemp)
+trap 'rm -f "$TMP"' EXIT
+
+(
+  cat <<EOF
+#define HAVE_ENDIAN_H
+#define HAVE_STRUCT_TIMESPEC
+#define HAVE_PTHREAD
+#include <unistd.h>
+#include "nir_shader_compiler_options.h"
+#include "compiler/shader_enums.h"
+EOF
+  sed -n '/struct nir_shader_compiler_options/,/^}/{p;/^}/q}' $1/src/gallium/drivers/llvmpipe/lp_screen.c
+  echo "int main(void) { write(1, &gallivm_nir_options, sizeof(gallivm_nir_options)); }"
+) | cc -x c -o $TMP - -I$1/src/compiler/nir -I$1/src -I$1/include && $TMP | gzip | base64 -w0
+
@@ -3,8 +3,21 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
+ort_options = ort.SessionOptions()
+ort_options.log_severity_level = 3

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
+  """
+  Generate example input tensors based on the provided ONNX graph input specifications.
+
+  NOTE: This is not guaranteed to be reliable. It's a best-effort helper
+  that uses heuristics to guess input shapes and values.
+
+  Example:
+    from tinygrad.nn.onnx import OnnxRunner
+    from extra.onnx_helpers import get_example_inputs
+    inputs = get_example_inputs(OnnxRunner(model_path).graph_inputs)
+  """
  def _get_shape(onnx_shape: tuple[str|int]):
    shape = []
    for onnx_dim in onnx_shape:
@@ -44,11 +57,9 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
    ret.update({name:value})
  return ret

-def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+def _get_tinygrad_and_ort_np_outputs(onnx_file, inputs):
  run_onnx = OnnxRunner(onnx_file)

-  ort_options = ort.SessionOptions()
-  ort_options.log_severity_level = 3
  ort_sess = ort.InferenceSession(onnx_file, ort_options, ["CPUExecutionProvider"])
  np_inputs = {k:v.numpy() if isinstance(v, Tensor) else v for k,v in inputs.items()}
  out_names = list(run_onnx.graph_outputs)
@@ -56,9 +67,101 @@ def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
  ort_out = dict(zip(out_names, out_values))

  tinygrad_out = run_onnx(inputs)
+  Tensor.realize(*(x for x in tinygrad_out.values() if x is not None))
+  tinygrad_out = {k:v.numpy() if v is not None else None for k,v in tinygrad_out.items()}
+  return tinygrad_out, ort_out
+
+def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares the final output tensors of an onnx model run in tinygrad and onnxruntime.
+  """
+  tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(onnx_file, inputs)

  assert tinygrad_out.keys() == ort_out.keys()
  for k in tinygrad_out.keys():
    tiny_v, onnx_v = tinygrad_out[k], ort_out[k]
    if tiny_v is None: assert onnx_v is None, f"{k}: {tiny_v=}, {onnx_v=}"
-    else: np.testing.assert_allclose(tiny_v.numpy(), onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+    else: np.testing.assert_allclose(tiny_v, onnx_v, rtol=rtol, atol=atol, err_msg=f"For tensor '{k}' in {tinygrad_out.keys()}")
+
+def validate_all_intermediates(onnx_file, inputs, rtol=1e-5, atol=1e-5):
+  """
+  Compares all intermediate node output of an onnx model run in tinygrad and onnxruntime.
+  """
+  report = generate_node_output_report(onnx_file, inputs)
+  for i, node in enumerate(report):
+    node_name = node["node"]
+    op = node["op"]
+    outputs = node["outputs"]
+    for output in outputs:
+      output_name = output["name"]
+      tinygrad_out = output["tinygrad"]
+      ort_out = output["onnxruntime"]
+      try:
+        if tinygrad_out is None: assert ort_out is None, f"None outputs are not equal {tinygrad_out=} {ort_out=}"
+        else: np.testing.assert_allclose(tinygrad_out, ort_out, rtol=rtol, atol=atol)
+        print(f"Validated {i}: {op=} {node_name=} {output_name=}")
+      except AssertionError as e:
+        print(f"FAILED {i}: {op=} {node_name=} {output_name=}")
+        print(str(e).strip() + "\n")
+
+def generate_node_output_report(onnx_file, inputs):
+  """
+  Build a report of all ONNX node outputs from tinygrad and onnxruntime
+
+  Returns:
+    A list of dictionaries, where each entry corresponds to one
+    node in the ONNX graph. The structure is as follows:
+    [
+      {
+        "node": str,  # The name of the ONNX node.
+        "op": str,    # The operation type of the ONNX node.
+        "outputs": [
+          {
+            "name": str,                       # The name of the output tensor.
+            "tinygrad": np.ndarray | None,     # The output value from tinygrad.
+            "onnxruntime": np.ndarray | None,  # The output value from onnxruntime.
+          },
+          ...
+        ]
+      },
+      ...
+    ]
+  """
+  import onnx_graphsurgeon as gs
+  import onnx
+  import tempfile
+
+  # rewrite the model to output all the node outputs
+  # `infer_shapes` here tries to fill the shapes and dtypes of intermediate values which graphsurgeon requires when assigning them as outputs
+  inferred_model = onnx.shape_inference.infer_shapes(onnx.load(onnx_file))
+  model = gs.import_onnx(inferred_model)
+  model_nodes = model.nodes
+  node_outputs = [n.outputs for n in model.nodes]
+  model.outputs = [
+      each_output for outputs in node_outputs for each_output in outputs
+      if not (each_output.dtype is None and each_output.shape is None)  # output with None dtype and None shape is likely a `None` value
+  ]
+  rewritten_model = gs.export_onnx(model)
+
+  # TODO: remove this once ORT supports 1.18.0
+  if getattr(rewritten_model, "ir_version", 0) > 10:
+    rewritten_model.ir_version = 10
+
+  with tempfile.NamedTemporaryFile(suffix=".onnx") as f:
+    onnx.save(rewritten_model, f.name)
+    rewritten_model_path = f.name
+    tinygrad_out, ort_out = _get_tinygrad_and_ort_np_outputs(rewritten_model_path, inputs)
+
+  report = []
+  for node in model_nodes:
+    outputs = []
+    for each_output in node.outputs:
+      if each_output.dtype is None and each_output.shape is None:
+        continue
+      name = each_output.name
+      tinygrad_output = tinygrad_out[name]
+      ort_output = ort_out[name]
+      outputs.append({"name": name, "tinygrad": tinygrad_output, "onnxruntime": ort_output})
+    report.append({"node": node.name, "op": node.op, "outputs": outputs})
+
+  return report
@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import os, shutil
+from pathlib import Path
+from tinygrad.helpers import fetch, OSX
+
+DEST = Path("/usr/local/lib")
+DEST.mkdir(exist_ok=True)
+
+if __name__ == "__main__":
+  if OSX:
+    fp = fetch("https://github.com/ROCm/rocprof-trace-decoder/releases/download/0.1.4/rocprof-trace-decoder-macos-arm64-0.1.4-Darwin.sh")
+    lib = fp.parent/"rocprof-trace-decoder-macos-arm64-0.1.4-Darwin"/"lib"/"librocprof-trace-decoder.dylib"
+    os.chmod(fp, 0o755)
+    os.system(f"sudo {fp} --prefix={fp.parent} --include-subdir")
+  else:
+    lib = fetch("https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so", name="librocprof-trace-decoder.so")
+  shutil.copy2(lib, DEST)
+  print(f"Installed {lib.name} to", DEST)
@@ -7,8 +7,7 @@
 # POINTER_SIZE is: 8
 # LONGDOUBLE_SIZE is: 16
 #
-import ctypes
-from tinygrad.helpers import fetch
+import ctypes, ctypes.util


 class AsDictMixin:
@@ -156,7 +155,7 @@ class FunctionFactoryStub:
 # You can either re-run clan2py with -l /path/to/library.so
 # Or manually fix this by comment the ctypes.CDLL loading
 _libraries = {}
-_libraries['FIXME_STUB'] = ctypes.CDLL(str(fetch('https://github.com/ROCm/rocprof-trace-decoder/raw/5420409ad0963b2d76450add067b9058493ccbd0/releases/linux_glibc_2_28_x86_64/librocprof-trace-decoder.so'))) #  ctypes.CDLL('FIXME_STUB')
+_libraries['FIXME_STUB'] = ctypes.CDLL(ctypes.util.find_library('rocprof-trace-decoder')) #  ctypes.CDLL('FIXME_STUB')



@@ -0,0 +1,11 @@
+from tinygrad.tensor import Tensor
+import argparse
+
+if __name__ == "__main__":
+  parser = argparse.ArgumentParser()
+  parser.add_argument("hash", type=str, required=True, help="file hash to fetch")
+  parser.add_argument("len", type=int, required=True, help="file length to fetch")
+  parser.add_argument("dest", type=str, required=True, help="destination path to save the file")
+  args = parser.parse_args()
+
+  Tensor(bytes.fromhex(args.hash), device="CPU").load(args.len).to(f"disk:{args.dest}").realize()
@@ -0,0 +1,39 @@
+import json, multiprocessing
+from pathlib import Path
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm, getenv
+
+raid_root = Path(getenv("RAID_ROOT", "/raid"))
+
+def fetch_file(item):
+  path, info = item
+  h, size = info["hash"], info["size"]
+
+  path = raid_root / Path(path)
+  path.parent.mkdir(parents=True, exist_ok=True)
+
+  try:
+    pt = Tensor(bytes.fromhex(h), device="CPU").load(size).to(f"disk:{path.as_posix()}").realize()
+  except Exception as e:
+    print(f"error fetching {path}, {h}, {size}: {e}")
+    raise
+
+  pt.uop.buffer.deallocate()
+
+def fetch_mapping():
+  mapping_tensor = Tensor(bytes.fromhex("d734f5e3be9f1e9d863bfaa4fc6c1ef2")).load(175866113).realize()
+  mapping = mapping_tensor.data().tobytes().decode()
+  mapping = json.loads(mapping)
+  mapped_files = mapping.items()
+  return list(mapped_files)
+
+if __name__ == "__main__":
+  with multiprocessing.Pool(processes=1) as pool:
+    mapped_files = pool.apply(fetch_mapping)
+
+  print(f"fetched mapping for {len(mapped_files)} files")
+
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
+      pass
@@ -0,0 +1,31 @@
+from pathlib import Path
+import multiprocessing, json
+
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import tqdm
+
+raid_root = Path("/raid")
+
+def upload_file(path: Path):
+  pt = Tensor(path).realize()
+  h = pt.store().realize()
+  pt.uop.realized.deallocate()
+  return h.data().hex(), path, pt.nbytes()
+
+if __name__ == "__main__":
+  raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
+  print(f"found {len(raid_files)} files in /raid")
+
+  mapping = {}
+  with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
+    for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
+      mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
+
+  # sort the mapping by key
+  mapping = dict(sorted(mapping.items()))
+
+  mapping = json.dumps(mapping).encode()
+  mapping_tensor = Tensor(mapping, device="CPU")
+  h = mapping_tensor.store().realize()
+
+  print(f"final hash: {h.data().hex()}, size: {len(mapping)}")
@@ -155,16 +155,14 @@ def index_tensor(x, y):
 def zero_(x):
  if TORCH_DEBUG: print(f"zero_ {x.shape}")
  tt = unwrap(x)
-  # NOTE: unconditional contiguous covers if x is contiguous (match it) or if x is view (realize for inplace)
-  # TODO: consolidate
-  tt.assign(tt.zeros_like().contiguous())
+  tt.assign(tt.zeros_like())

@torch.library.impl("aten::fill_.Scalar", "privateuseone")
@inplace_fn("x")
 def fill_scalar(x, y):
  if TORCH_DEBUG: print(f"fill_.Scalar {x.shape} {y}")
  tt = unwrap(x)
-  tt.assign(tt.full_like(y).contiguous())
+  tt.assign(tt.full_like(y))

@torch.library.impl("aten::_local_scalar_dense", "privateuseone")
 def _local_scalar_dense(tensor): return unwrap(tensor).item()
@@ -644,10 +642,11 @@ def get_real_tinygrad_buffers():
 torch.nn.modules.module.register_module_buffer_registration_hook(register_torch_buffer)

 from torch.nn.modules import Module
-def backward_hook(model:Module, _grad_input, _grad_out):
-  grads_to_realize = [unwrap(p.grad) for p in model.parameters() if p.grad is not None]
-  if len(grads_to_realize): Tensor.realize(*grads_to_realize)
-def module_hook(module:Module, _name, _submodule): module.register_backward_hook(backward_hook)
+def param_hook(_grad):
+  if _grad is not None and _grad.is_tiny: Tensor.realize(unwrap(_grad))
+def module_hook(module:Module, _name, _submodule):
+  for param in _submodule.parameters(recurse=False):
+    if param.requires_grad: param.register_hook(param_hook)
 torch.nn.modules.module.register_module_module_registration_hook(module_hook)

 def realize_optimizer_step(optimizer: torch.optim.Optimizer, *args, **kwargs):
@@ -0,0 +1,13 @@
+xcuserdata/
+
+**/*.xcodeproj/project.xcworkspace/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata
+
+**/*.xcodeproj/project.xcworkspace/xcshareddata/*
+!**/*.xcodeproj/project.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings
+
+**/*.playground/playground.xcworkspace/*
+!**/*.playground/playground.xcworkspace/xcshareddata
+
+**/*.playground/playground.xcworkspace/xcshareddata/*
+!**/*.playground/playground.xcworkspace/xcshareddata/WorkspaceSettings.xcsettings
@@ -0,0 +1,11 @@
+{
+  "colors" : [
+    {
+      "idiom" : "universal"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,148 @@
+{
+  "images" : [
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "2x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "iphone",
+      "scale" : "3x",
+      "size" : "60x60"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "20x20"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "29x29"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "40x40"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "1x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "76x76"
+    },
+    {
+      "idiom" : "ipad",
+      "scale" : "2x",
+      "size" : "83.5x83.5"
+    },
+    {
+      "idiom" : "ios-marketing",
+      "scale" : "1x",
+      "size" : "1024x1024"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "16x16"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "32x32"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "128x128"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "256x256"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "1x",
+      "size" : "512x512"
+    },
+    {
+      "idiom" : "mac",
+      "scale" : "2x",
+      "size" : "512x512"
+    }
+  ],
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,6 @@
+{
+  "info" : {
+    "author" : "xcode",
+    "version" : 1
+  }
+}
@@ -0,0 +1,19 @@
+import AppKit
+import SwiftUI
+
+final class AppDelegate: NSObject, NSApplicationDelegate {
+	func applicationShouldTerminateAfterLastWindowClosed(_ sender: NSApplication) -> Bool {
+		true
+	}
+}
+
+@main
+struct TinyGPUApp: App {
+	@NSApplicationDelegateAdaptor(AppDelegate.self) var appDelegate
+
+    var body: some Scene {
+        WindowGroup {
+			TinyGPUView()
+        }
+    }
+}
@@ -0,0 +1,33 @@
+import SwiftUI
+
+struct TinyGPUView: View {
+	@ObservedObject var viewModel = TinyGPUViewModel()
+
+	var body: some View {
+#if os(macOS)
+		VStack(alignment: .center) {
+			Text("TinyGPU Intsaller")
+				.padding()
+				.font(.title)
+			Text(self.viewModel.dextLoadingState)
+				.multilineTextAlignment(.center)
+			HStack {
+				Button(
+					action: {
+						self.viewModel.activateMyDext()
+					}, label: {
+						Text("Install extension")
+					}
+				)
+			}
+		}
+		.frame(width: 500, height: 200, alignment: .center)
+#endif
+	}
+}
+
+struct TinyGPUView_Previews: PreviewProvider {
+    static var previews: some View {
+		TinyGPUView()
+    }
+}
@@ -0,0 +1,149 @@
+import Foundation
+import os.log
+import SystemExtensions
+
+class TinyGPUDriverLoadingStateMachine {
+  enum State { case unloaded, activating, needsApproval, activated, activationError }
+}
+
+class TinyGPUViewModel: NSObject {
+
+	@Published private var state: TinyGPUDriverLoadingStateMachine.State = .unloaded
+	
+	override init() {
+		super.init()
+		refreshInitialDextState()
+	}
+	
+	private func refreshInitialDextState() {
+#if os(macOS)
+		Task.detached { [dextIdentifier] in
+		  let newState = Self.queryDextState(bundleID: dextIdentifier)
+		  await MainActor.run { self.state = newState }
+		}
+#endif
+	}
+
+#if os(macOS)
+	private static func queryDextState(bundleID: String) -> TinyGPUDriverLoadingStateMachine.State {
+		let tool = "/usr/bin/systemextensionsctl"
+		let p = Process()
+		p.executableURL = URL(fileURLWithPath: tool)
+		p.arguments = ["list"]
+
+		let pipe = Pipe()
+		p.standardOutput = pipe
+		p.standardError = Pipe()
+
+		do {
+			try p.run()
+			p.waitUntilExit()
+			let data = pipe.fileHandleForReading.readDataToEndOfFile()
+			guard let output = String(data: data, encoding: .utf8) else { return .unloaded }
+
+			// Look for our bundle id line
+			if let line = output.split(separator: "\n").first(where: { $0.contains(bundleID) }) {
+				if line.contains("[activated enabled]") { return .activated }
+				if line.contains("[activated waiting for user]") { return .needsApproval }
+				if line.contains("terminated waiting to uninstall") { return .unloaded }
+				return .activating
+			} else {
+				return .unloaded
+			}
+		} catch {
+			return .unloaded
+		}
+	}
+#endif
+
+	private let dextIdentifier: String = "org.tinygrad.tinygpu.edriver"
+
+	public var dextLoadingState: String {
+		switch state {
+		case .unloaded:
+			return "TinyGPUDriver isn't loaded."
+		case .activating:
+			return "Activating TinyGPUDriver, please wait."
+		case .needsApproval:
+			return "Please follow the prompt to approve TinyGPUDriver."
+		case .activated:
+			return "TinyGPUDriver has been activated and is ready to use. You can close the installer."
+		case .activationError:
+			return "TinyGPUDriver has experienced an error during activation.\nPlease check the logs to find the error."
+		}
+	}
+}
+
+extension TinyGPUViewModel: ObservableObject {
+
+#if os(macOS)
+	func activateMyDext() {
+		activateExtension(dextIdentifier)
+	}
+	
+	func deactivateMyDext() {
+		deactivateExtension(dextIdentifier)
+	}
+
+	func activateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest
+			.activationRequest(forExtensionWithIdentifier: dextIdentifier,
+							   queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .activating
+	}
+	
+	func deactivateExtension(_ dextIdentifier: String) {
+
+		let request = OSSystemExtensionRequest.deactivationRequest(forExtensionWithIdentifier: dextIdentifier, queue: .main)
+		request.delegate = self
+		OSSystemExtensionManager.shared.submitRequest(request)
+
+		self.state = .unloaded
+	}
+#endif
+}
+
+#if os(macOS)
+extension TinyGPUViewModel: OSSystemExtensionRequestDelegate {
+
+	func request(
+		_ request: OSSystemExtensionRequest,
+		actionForReplacingExtension existing: OSSystemExtensionProperties,
+		withExtension ext: OSSystemExtensionProperties) -> OSSystemExtensionRequest.ReplacementAction {
+
+		var replacementAction: OSSystemExtensionRequest.ReplacementAction
+
+		os_log("sysex actionForReplacingExtension: %@ %@", existing, ext)
+
+		// Add appropriate logic here to determine whether to replace the extension
+		// with the new extension. Common things to check for include
+		// testing whether the new extension's version number is newer than
+		// the current version number, or whether the bundleIdentifier is different.
+		// For simplicity, this sample always replaces the current extension
+		// with the new one.
+		replacementAction = .replace
+
+		self.state = .activating
+		return replacementAction
+	}
+
+	func requestNeedsUserApproval(_ request: OSSystemExtensionRequest) {
+		os_log("sysex requestNeedsUserApproval")
+		self.state = .needsApproval
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFinishWithResult result: OSSystemExtensionRequest.Result) {
+		os_log("sysex didFinishWithResult: %d", result.rawValue)
+		self.state = .activated
+	}
+
+	func request(_ request: OSSystemExtensionRequest, didFailWithError error: Error) {
+		os_log("sysex didFailWithError: %@", error.localizedDescription)
+		self.state = .activationError
+	}
+}
+#endif
@@ -0,0 +1,589 @@
+// !$*UTF8*$!
+{
+	archiveVersion = 1;
+	classes = {
+	};
+	objectVersion = 54;
+	objects = {
+
+/* Begin PBXBuildFile section */
+		0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 0ACB55382E9CB880007029EF /* PCIDriverKit.framework */; };
+		54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 54798268286A3512009785F6 /* CoreAudio.framework */; };
+		549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */; };
+		549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */ = {isa = PBXBuildFile; fileRef = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */; settings = {ATTRIBUTES = (RemoveHeadersOnCopy, ); }; };
+		549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 549EB130286A2B98009D38AB /* IOKit.framework */; };
+		54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */; };
+		54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 54E42BB9286A1696000E1E9A /* TinyGPUView.swift */; };
+		54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 54E42BBA286A1697000E1E9A /* Assets.xcassets */; };
+		C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */; };
+		C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */; };
+		C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */; };
+		C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */; };
+		C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */; };
+		C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */ = {isa = PBXBuildFile; fileRef = C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */; };
+/* End PBXBuildFile section */
+
+/* Begin PBXContainerItemProxy section */
+		549EB126286A1D66009D38AB /* PBXContainerItemProxy */ = {
+			isa = PBXContainerItemProxy;
+			containerPortal = C5B7D9B326128AC50089B4C3 /* Project object */;
+			proxyType = 1;
+			remoteGlobalIDString = C5B7D9BB26128AC50089B4C3;
+			remoteInfo = SimpleAudioDriver;
+		};
+/* End PBXContainerItemProxy section */
+
+/* Begin PBXCopyFilesBuildPhase section */
+		549EB122286A1D3A009D38AB /* Embed System Extensions */ = {
+			isa = PBXCopyFilesBuildPhase;
+			buildActionMask = 2147483647;
+			dstPath = "$(SYSTEM_EXTENSIONS_FOLDER_PATH)";
+			dstSubfolderSpec = 16;
+			files = (
+				549EB123286A1D48009D38AB /* org.tinygrad.tinygpu.edriver.dext in Embed System Extensions */,
+			);
+			name = "Embed System Extensions";
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXCopyFilesBuildPhase section */
+
+/* Begin PBXFileReference section */
+		0ACB55382E9CB880007029EF /* PCIDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = PCIDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/PCIDriverKit.framework; sourceTree = SDKROOT; };
+		54798268286A3512009785F6 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUViewModel.swift; sourceTree = "<group>"; usesTabs = 1; };
+		549EB130286A2B98009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX13.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		549EB132286A2B9D009D38AB /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/iPhoneOS.platform/Developer/SDKs/iPhoneOS16.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUApp.swift; sourceTree = "<group>"; };
+		54E42BB9286A1696000E1E9A /* TinyGPUView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TinyGPUView.swift; sourceTree = "<group>"; };
+		54E42BBA286A1697000E1E9A /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
+		54E42BC4286A1697000E1E9A /* TinyGPU.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TinyGPU.app; sourceTree = BUILT_PRODUCTS_DIR; };
+		54E42BC6286A1697000E1E9A /* macOS.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = macOS.entitlements; sourceTree = "<group>"; };
+		C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */ = {isa = PBXFileReference; explicitFileType = "wrapper.driver-extension"; includeInIndex = 0; path = org.tinygrad.tinygpu.edriver.dext; sourceTree = BUILT_PRODUCTS_DIR; };
+		C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriver.cpp; sourceTree = "<group>"; usesTabs = 1; };
+		C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriver.iig; sourceTree = "<group>"; };
+		C5B7D9C626128AC50089B4C3 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
+		C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = SDKROOT; };
+		C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TinyGPUDriver.entitlements; sourceTree = "<group>"; };
+		C5C0063326178F98003345D8 /* AppKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AppKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/AppKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C006352617ACB8003345D8 /* CoreAudio.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = CoreAudio.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/CoreAudio.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = AudioDriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/AudioDriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = DriverKit.framework; path = Platforms/DriverKit.platform/Developer/SDKs/DriverKit.MacOSX21.0.Internal.sdk/System/DriverKit/System/Library/Frameworks/DriverKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.iig; path = TinyGPUDriverUserClient.iig; sourceTree = "<group>"; };
+		C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.cpp; path = TinyGPUDriverUserClient.cpp; sourceTree = "<group>"; };
+		C5D787B026169723006047E5 /* IOKit.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = IOKit.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/IOKit.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B22616973F006047E5 /* SystemExtensions.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = SystemExtensions.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/SystemExtensions.framework; sourceTree = DEVELOPER_DIR; };
+		C5D787B426169747006047E5 /* Foundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Foundation.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX12.0.sdk/System/Library/Frameworks/Foundation.framework; sourceTree = DEVELOPER_DIR; };
+/* End PBXFileReference section */
+
+/* Begin PBXFrameworksBuildPhase section */
+		54E42BC1286A1697000E1E9A /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54798269286A3512009785F6 /* CoreAudio.framework in Frameworks */,
+				549EB131286A2B98009D38AB /* IOKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B926128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXFrameworksBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5C3BBB32612ACDC003C7BFE /* AudioDriverKit.framework in Frameworks */,
+				C5C3BBB52612ACEF003C7BFE /* DriverKit.framework in Frameworks */,
+				0ACB55392E9CB880007029EF /* PCIDriverKit.framework in Frameworks */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXFrameworksBuildPhase section */
+
+/* Begin PBXGroup section */
+		54E42BB7286A1696000E1E9A /* Shared */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BB8286A1696000E1E9A /* TinyGPUApp.swift */,
+				54E42BB9286A1696000E1E9A /* TinyGPUView.swift */,
+				549EB11F286A1A37009D38AB /* TinyGPUViewModel.swift */,
+				54E42BBA286A1697000E1E9A /* Assets.xcassets */,
+			);
+			path = Shared;
+			sourceTree = "<group>";
+		};
+		54E42BC5286A1697000E1E9A /* macOS */ = {
+			isa = PBXGroup;
+			children = (
+				54E42BC6286A1697000E1E9A /* macOS.entitlements */,
+			);
+			path = macOS;
+			sourceTree = "<group>";
+		};
+		C5B7D9B226128AC50089B4C3 = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */,
+				54E42BB7286A1696000E1E9A /* Shared */,
+				54E42BC5286A1697000E1E9A /* macOS */,
+				C5B7D9BE26128AC50089B4C3 /* Frameworks */,
+				C5B7D9BD26128AC50089B4C3 /* Products */,
+			);
+			sourceTree = "<group>";
+			usesTabs = 1;
+		};
+		C5B7D9BD26128AC50089B4C3 /* Products */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */,
+				54E42BC4286A1697000E1E9A /* TinyGPU.app */,
+			);
+			name = Products;
+			sourceTree = "<group>";
+		};
+		C5B7D9BE26128AC50089B4C3 /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				0ACB55382E9CB880007029EF /* PCIDriverKit.framework */,
+				54798268286A3512009785F6 /* CoreAudio.framework */,
+				549EB130286A2B98009D38AB /* IOKit.framework */,
+				549EB132286A2B9D009D38AB /* IOKit.framework */,
+				C5C006352617ACB8003345D8 /* CoreAudio.framework */,
+				C5C0063326178F98003345D8 /* AppKit.framework */,
+				C5D787B426169747006047E5 /* Foundation.framework */,
+				C5D787B22616973F006047E5 /* SystemExtensions.framework */,
+				C5D787B026169723006047E5 /* IOKit.framework */,
+				C5C3BBB42612ACEF003C7BFE /* DriverKit.framework */,
+				C5B7D9CC26128ADA0089B4C3 /* AudioDriverKit.framework */,
+				C5C3BBB12612ACD3003C7BFE /* AudioDriverKit.framework */,
+				C5B7D9BF26128AC50089B4C3 /* DriverKit.framework */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
+		C5B7D9C126128AC50089B4C3 /* TinyGPUDriverExtension */ = {
+			isa = PBXGroup;
+			children = (
+				C5B7D9C226128AC50089B4C3 /* TinyGPUDriver.cpp */,
+				C5B7D9C426128AC50089B4C3 /* TinyGPUDriver.iig */,
+				C5D787AD26168D1E006047E5 /* TinyGPUDriverUserClient.cpp */,
+				C5D787AB261667FC006047E5 /* TinyGPUDriverUserClient.iig */,
+				C5B7D9C626128AC50089B4C3 /* Info.plist */,
+				C5B7D9CE26128B150089B4C3 /* TinyGPUDriver.entitlements */,
+			);
+			path = TinyGPUDriverExtension;
+			sourceTree = "<group>";
+		};
+/* End PBXGroup section */
+
+/* Begin PBXHeadersBuildPhase section */
+		C5B7D9B726128AC50089B4C3 /* Headers */ = {
+			isa = PBXHeadersBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXHeadersBuildPhase section */
+
+/* Begin PBXNativeTarget section */
+		54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = 54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */;
+			buildPhases = (
+				54E42BC0286A1697000E1E9A /* Sources */,
+				54E42BC1286A1697000E1E9A /* Frameworks */,
+				54E42BC2286A1697000E1E9A /* Resources */,
+				549EB122286A1D3A009D38AB /* Embed System Extensions */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+				549EB127286A1D66009D38AB /* PBXTargetDependency */,
+			);
+			name = "TinyGPU Installer (macOS)";
+			productName = "SimpleAudioDriverExtension2 (macOS)";
+			productReference = 54E42BC4286A1697000E1E9A /* TinyGPU.app */;
+			productType = "com.apple.product-type.application";
+		};
+		C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */ = {
+			isa = PBXNativeTarget;
+			buildConfigurationList = C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */;
+			buildPhases = (
+				C5B7D9B726128AC50089B4C3 /* Headers */,
+				C5B7D9B826128AC50089B4C3 /* Sources */,
+				C5B7D9B926128AC50089B4C3 /* Frameworks */,
+				C5B7D9BA26128AC50089B4C3 /* Resources */,
+			);
+			buildRules = (
+			);
+			dependencies = (
+			);
+			name = TinyGPUDriver;
+			productName = SimpleAudioDriverExtension;
+			productReference = C5B7D9BC26128AC50089B4C3 /* org.tinygrad.tinygpu.edriver.dext */;
+			productType = "com.apple.product-type.driver-extension";
+		};
+/* End PBXNativeTarget section */
+
+/* Begin PBXProject section */
+		C5B7D9B326128AC50089B4C3 /* Project object */ = {
+			isa = PBXProject;
+			attributes = {
+				BuildIndependentTargetsInParallel = YES;
+				DefaultBuildSystemTypeForWorkspace = Latest;
+				LastSwiftUpdateCheck = 1400;
+				LastUpgradeCheck = 1600;
+				ORGANIZATIONNAME = Apple;
+				TargetAttributes = {
+					54E42BC3286A1697000E1E9A = {
+						CreatedOnToolsVersion = 14.0;
+						LastSwiftMigration = 1400;
+					};
+					C5B7D9BB26128AC50089B4C3 = {
+						CreatedOnToolsVersion = 13.0;
+					};
+				};
+			};
+			buildConfigurationList = C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */;
+			compatibilityVersion = "Xcode 9.3";
+			developmentRegion = en;
+			hasScannedForEncodings = 0;
+			knownRegions = (
+				en,
+				Base,
+			);
+			mainGroup = C5B7D9B226128AC50089B4C3;
+			productRefGroup = C5B7D9BD26128AC50089B4C3 /* Products */;
+			projectDirPath = "";
+			projectRoot = "";
+			targets = (
+				C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */,
+				54E42BC3286A1697000E1E9A /* TinyGPU Installer (macOS) */,
+			);
+		};
+/* End PBXProject section */
+
+/* Begin PBXResourcesBuildPhase section */
+		54E42BC2286A1697000E1E9A /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				54E42BCC286A1697000E1E9A /* Assets.xcassets in Resources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9BA26128AC50089B4C3 /* Resources */ = {
+			isa = PBXResourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXResourcesBuildPhase section */
+
+/* Begin PBXSourcesBuildPhase section */
+		54E42BC0286A1697000E1E9A /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				549EB121286A1A37009D38AB /* TinyGPUViewModel.swift in Sources */,
+				54E42BCA286A1697000E1E9A /* TinyGPUView.swift in Sources */,
+				54E42BC8286A1697000E1E9A /* TinyGPUApp.swift in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+		C5B7D9B826128AC50089B4C3 /* Sources */ = {
+			isa = PBXSourcesBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+				C5B7D9C526128AC50089B4C3 /* TinyGPUDriver.iig in Sources */,
+				C5D787AE26168E59006047E5 /* TinyGPUDriverUserClient.cpp in Sources */,
+				C5D787AC261667FC006047E5 /* TinyGPUDriverUserClient.iig in Sources */,
+				C5B7D9C326128AC50089B4C3 /* TinyGPUDriver.cpp in Sources */,
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+		};
+/* End PBXSourcesBuildPhase section */
+
+/* Begin PBXTargetDependency section */
+		549EB127286A1D66009D38AB /* PBXTargetDependency */ = {
+			isa = PBXTargetDependency;
+			target = C5B7D9BB26128AC50089B4C3 /* TinyGPUDriver */;
+			targetProxy = 549EB126286A1D66009D38AB /* PBXContainerItemProxy */;
+		};
+/* End PBXTargetDependency section */
+
+/* Begin XCBuildConfiguration section */
+		54E42BCF286A1697000E1E9A /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Debug;
+		};
+		54E42BD0286A1697000E1E9A /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
+				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
+				CLANG_ENABLE_MODULES = YES;
+				CODE_SIGN_ENTITLEMENTS = macOS/macOS.entitlements;
+				CODE_SIGN_IDENTITY = "-";
+				"CODE_SIGN_IDENTITY[sdk=macosx*]" = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				COMBINE_HIDPI_IMAGES = YES;
+				CURRENT_PROJECT_VERSION = 1;
+				DEAD_CODE_STRIPPING = YES;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				ENABLE_HARDENED_RUNTIME = YES;
+				ENABLE_PREVIEWS = YES;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_KEY_NSHumanReadableCopyright = "";
+				LD_RUNPATH_SEARCH_PATHS = (
+					"$(inherited)",
+					"@executable_path/../Frameworks",
+				);
+				MACOSX_DEPLOYMENT_TARGET = 12.1;
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.installer;
+				PRODUCT_NAME = TinyGPU;
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				SDKROOT = macosx;
+				SWIFT_COMPILATION_MODE = wholemodule;
+				SWIFT_EMIT_LOC_STRINGS = YES;
+				SWIFT_OPTIMIZATION_LEVEL = "-O";
+				SWIFT_VERSION = 5.0;
+			};
+			name = Release;
+		};
+		C5B7D9C726128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = dwarf;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				ENABLE_TESTABILITY = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_DYNAMIC_NO_PIC = NO;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_OPTIMIZATION_LEVEL = 0;
+				GCC_PREPROCESSOR_DEFINITIONS = (
+					"DEBUG=1",
+					"$(inherited)",
+				);
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE;
+				MTL_FAST_MATH = YES;
+				ONLY_ACTIVE_ARCH = YES;
+				SDKROOT = driverkit;
+			};
+			name = Debug;
+		};
+		C5B7D9C826128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				ALWAYS_SEARCH_USER_PATHS = NO;
+				CLANG_ANALYZER_NONNULL = YES;
+				CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE;
+				CLANG_CXX_LANGUAGE_STANDARD = "gnu++17";
+				CLANG_CXX_LIBRARY = "libc++";
+				CLANG_ENABLE_MODULES = YES;
+				CLANG_ENABLE_OBJC_ARC = YES;
+				CLANG_ENABLE_OBJC_WEAK = YES;
+				CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES;
+				CLANG_WARN_BOOL_CONVERSION = YES;
+				CLANG_WARN_COMMA = YES;
+				CLANG_WARN_CONSTANT_CONVERSION = YES;
+				CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES;
+				CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR;
+				CLANG_WARN_DOCUMENTATION_COMMENTS = YES;
+				CLANG_WARN_EMPTY_BODY = YES;
+				CLANG_WARN_ENUM_CONVERSION = YES;
+				CLANG_WARN_INFINITE_RECURSION = YES;
+				CLANG_WARN_INT_CONVERSION = YES;
+				CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES;
+				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
+				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
+				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
+				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
+				CLANG_WARN_STRICT_PROTOTYPES = YES;
+				CLANG_WARN_SUSPICIOUS_MOVE = YES;
+				CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE;
+				CLANG_WARN_UNREACHABLE_CODE = YES;
+				CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
+				COPY_PHASE_STRIP = NO;
+				DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_NS_ASSERTIONS = NO;
+				ENABLE_STRICT_OBJC_MSGSEND = YES;
+				GCC_C_LANGUAGE_STANDARD = gnu11;
+				GCC_NO_COMMON_BLOCKS = YES;
+				GCC_WARN_64_TO_32_BIT_CONVERSION = YES;
+				GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR;
+				GCC_WARN_UNDECLARED_SELECTOR = YES;
+				GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE;
+				GCC_WARN_UNUSED_FUNCTION = YES;
+				GCC_WARN_UNUSED_VARIABLE = YES;
+				MTL_ENABLE_DEBUG_INFO = NO;
+				MTL_FAST_MATH = YES;
+				SDKROOT = driverkit;
+				SWIFT_COMPILATION_MODE = wholemodule;
+			};
+			name = Release;
+		};
+		C5B7D9CA26128AC50089B4C3 /* Debug */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Debug;
+		};
+		C5B7D9CB26128AC50089B4C3 /* Release */ = {
+			isa = XCBuildConfiguration;
+			buildSettings = {
+				AD_HOC_CODE_SIGNING_ALLOWED = YES;
+				CODE_SIGN_ENTITLEMENTS = TinyGPUDriverExtension/TinyGPUDriver.entitlements;
+				CODE_SIGN_IDENTITY = "Apple Development";
+				CODE_SIGN_STYLE = Automatic;
+				CURRENT_PROJECT_VERSION = 1;
+				DEVELOPMENT_TEAM = 9YG3G8543N;
+				DRIVERKIT_DEPLOYMENT_TARGET = 21.0;
+				ENABLE_USER_SCRIPT_SANDBOXING = YES;
+				FRAMEWORK_SEARCH_PATHS = (
+					"$(inherited)",
+					"$(SDKROOT)/System/DriverKit/System/Library/Frameworks",
+				);
+				GENERATE_INFOPLIST_FILE = YES;
+				INFOPLIST_FILE = TinyGPUDriverExtension/Info.plist;
+				INFOPLIST_KEY_OSBundleUsageDescription = "Sample Code Audio Driver Kit Extension";
+				MARKETING_VERSION = 1.0;
+				PRODUCT_BUNDLE_IDENTIFIER = org.tinygrad.tinygpu.edriver;
+				PRODUCT_NAME = "$(inherited)";
+				PROVISIONING_PROFILE_SPECIFIER = "";
+				RUN_CLANG_STATIC_ANALYZER = YES;
+				SDKROOT = driverkit;
+				SKIP_INSTALL = YES;
+			};
+			name = Release;
+		};
+/* End XCBuildConfiguration section */
+
+/* Begin XCConfigurationList section */
+		54E42BD2286A1697000E1E9A /* Build configuration list for PBXNativeTarget "TinyGPU Installer (macOS)" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				54E42BCF286A1697000E1E9A /* Debug */,
+				54E42BD0286A1697000E1E9A /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9B626128AC50089B4C3 /* Build configuration list for PBXProject "TinyGPUDriverExtension" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9C726128AC50089B4C3 /* Debug */,
+				C5B7D9C826128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+		C5B7D9C926128AC50089B4C3 /* Build configuration list for PBXNativeTarget "TinyGPUDriver" */ = {
+			isa = XCConfigurationList;
+			buildConfigurations = (
+				C5B7D9CA26128AC50089B4C3 /* Debug */,
+				C5B7D9CB26128AC50089B4C3 /* Release */,
+			);
+			defaultConfigurationIsVisible = 0;
+			defaultConfigurationName = Release;
+		};
+/* End XCConfigurationList section */
+	};
+	rootObject = C5B7D9B326128AC50089B4C3 /* Project object */;
+}
@@ -0,0 +1,10 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>BuildSystemType</key>
+	<string>Latest</string>
+	<key>DerivedDataLocationStyle</key>
+	<string>Default</string>
+</dict>
+</plist>
@@ -0,0 +1,37 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>IOKitPersonalities</key>
+	<dict>
+		<key>TinyGPUDriver</key>
+		<dict>
+			<key>CFBundleIdentifier</key>
+			<string>$(PRODUCT_BUNDLE_IDENTIFIER)</string>
+			<key>IOClass</key>
+			<string>IOUserService</string>
+			<key>IOMatchCategory</key>
+			<string>TinyGPUDriver</string>
+			<key>IOPCIClassMatch</key>
+			<string>0x03000000</string>
+			<key>IOPCITunnelCompatible</key>
+			<true/>
+			<key>IOProviderClass</key>
+			<string>IOPCIDevice</string>
+			<key>IOResourceMatch</key>
+			<string>IOKit</string>
+			<key>IOUserClass</key>
+			<string>TinyGPUDriver</string>
+			<key>IOUserServerName</key>
+			<string>org.tinygrad.tinygpu.Driver</string>
+			<key>TinyGPUDriverUserClientProperties</key>
+			<dict>
+				<key>IOClass</key>
+				<string>IOUserUserClient</string>
+				<key>IOUserClass</key>
+				<string>TinyGPUDriverUserClient</string>
+			</dict>
+		</dict>
+	</dict>
+</dict>
+</plist>
@@ -0,0 +1,259 @@
+#include "TinyGPUDriver.h"
+#include "TinyGPUDriverUserClient.h"
+#include <AudioDriverKit/AudioDriverKit.h>
+#include <DriverKit/IOUserServer.h>
+#include <DriverKit/IOLib.h>
+#include <DriverKit/OSString.h>
+#include <DriverKit/IOMemoryMap.h>
+#include <DriverKit/IODMACommand.h>
+#include <DriverKit/IODispatchQueue.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+#include <DriverKit/OSAction.h>
+
+struct TinyGPUDriver_IVars
+{
+	IOPCIDevice *pci = nullptr;
+};
+
+bool TinyGPUDriver::init()
+{
+	os_log(OS_LOG_DEFAULT, "tinygpu: init");
+
+	auto answer = super::init();
+	if (!answer) {
+		return false;
+	}
+
+	ivars = new TinyGPUDriver_IVars();
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriver::free()
+{
+	if (ivars != nullptr) {
+		
+	}
+	IOSafeDeleteNULL(ivars, TinyGPUDriver_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriver::Start_Impl(IOService* in_provider)
+{
+	IOServiceName service_name;
+	os_log(OS_LOG_DEFAULT, "tinygpu: on gpu detected");
+
+	kern_return_t err = Start(in_provider, SUPERDISPATCH);
+	if (err) return err;
+
+	ivars->pci = OSDynamicCast(IOPCIDevice, in_provider);
+	if (!ivars->pci) return kIOReturnNoDevice;
+
+	err = ivars->pci->Open(this, 0);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: Open() failed 0x%08x", err);
+		ivars->pci = nullptr;
+		return err;
+	}
+
+	uint16_t ven = 0, dev = 0;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetVendorID, &ven);
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetDeviceID, &dev);
+	os_log(OS_LOG_DEFAULT, "tinygpu: opened device ven=0x%04x dev=0x%04x", ven, dev);
+
+#if 0
+	uint32_t off = 0x100;
+	while (off) {
+		uint32_t hdr = 0, next = 0, cap_id = 0;
+		ivars->pci->ConfigurationRead32(off, &hdr);
+		cap_id = hdr & 0xFFFFu;
+		next = (hdr >> 20) & 0xFFCu;
+		os_log(OS_LOG_DEFAULT, "tinygpu: cap: %u", cap_id);
+		if (cap_id == 0x15) {
+			uint32_t cap = 0, ctrl = 0;
+			ivars->pci->ConfigurationRead32(off+0x4, &cap);
+			ivars->pci->ConfigurationRead32(off+0x8, &ctrl);
+
+			uint32_t new_bar_size = 31 - __builtin_clz(cap >> 4);
+			uint32_t new_ctrl = (ctrl & ~0x1f00) | (new_bar_size << 8);
+			ivars->pci->ConfigurationWrite32(off+0x8, new_ctrl);
+
+			os_log(OS_LOG_DEFAULT, "tinygpu: rebar: cap=%u ctrl=%u new_bar_size=%u new_ctrl=%u", cap, ctrl, new_bar_size, new_ctrl);
+			ivars->pci->Reset(0);
+			break;
+		}
+		off = next;
+	}
+	ivars->pci->Reset(kIOPCIDeviceResetTypeHotReset);
+#endif
+
+	uint16_t commandRegister;
+	ivars->pci->ConfigurationRead16(kIOPCIConfigurationOffsetCommand, &commandRegister);
+	commandRegister |= (kIOPCICommandIOSpace | kIOPCICommandBusMaster | kIOPCICommandMemorySpace);
+	ivars->pci->ConfigurationWrite16(kIOPCIConfigurationOffsetCommand, commandRegister);
+
+	memcpy((void*)service_name, (void*)"tinygpu\0", 8);
+	SetName(service_name);
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: will register service %s", service_name);
+	RegisterService();
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: service started %s", service_name);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::Stop_Impl(IOService* in_provider)
+{
+	ivars->pci->Close(this, 0);
+	return 0;
+}
+
+kern_return_t TinyGPUDriver::NewUserClient_Impl(uint32_t in_type, IOUserClient** out_user_client)
+{
+	kern_return_t err = 0;
+
+	IOService* user_client_service = nullptr;
+	err = Create(this, "TinyGPUDriverUserClientProperties", &user_client_service);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create NewUserClient");
+		goto error;
+	}
+	*out_user_client = OSDynamicCast(IOUserClient, user_client_service);
+	os_log(OS_LOG_DEFAULT, "tinygpu: NewUserClient created");
+
+error:
+	return err;
+}
+
+kern_return_t TinyGPUDriver::MapBar(uint32_t bar, IOMemoryDescriptor** memory)
+{
+	kern_return_t err = 0;
+	uint8_t barMemoryIndex, barMemoryType;
+	uint64_t barMemorySize;
+	err = ivars->pci->GetBARInfo(bar, &barMemoryIndex, &barMemorySize, &barMemoryType);
+	if (err) return err;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: requested bar mapping %d, %d", bar, (uint32_t)barMemoryIndex);
+	err = ivars->pci->_CopyDeviceMemoryWithIndex(barMemoryIndex, memory, this);
+	return err;
+}
+
+kern_return_t TinyGPUDriver::CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc)
+{
+	kern_return_t err = 0;
+	IOMemoryMap* memoryMap = nullptr;
+	IOBufferMemoryDescriptor* sharedBuf = nullptr;
+	IODMACommand* dmaCmd = nullptr;
+	uint64_t flags = kIOMemoryDirectionInOut;
+	uint32_t segCount = 32;
+	IOAddressSegment segments[32];
+	IODMACommandSpecification dmaSpec = {
+		.options = 0,
+		.maxAddressBits = 40,
+	};
+
+	err = IOBufferMemoryDescriptor::Create(kIOMemoryDirectionInOut, size, IOVMPageSize, &sharedBuf);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to alloc user buffer, err=%d", err);
+		goto error;
+	}
+	
+	err = IODMACommand::Create(ivars->pci, kIODMACommandCreateNoOptions, &dmaSpec, &dmaCmd);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to create dma command, err=%d", err);
+		goto error;
+	}
+
+	err = dmaCmd->PrepareForDMA(kIODMACommandPrepareForDMANoOptions, sharedBuf, 0, size,
+								&flags, &segCount, segments);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to prepare for dma, err=%d", err);
+		goto error;
+	}
+
+	// pass addresses to userland
+	{
+		// debug
+		for (int i = 0; i < segCount; i++) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: new dma mapping (sz=0x%zx) %d 0x%llx 0x%llx", size, i, segments[i].address, segments[i].length);
+		}
+
+		err = sharedBuf->CreateMapping(0, 0, 0, IOVMPageSize, IOVMPageSize, &memoryMap); // one page should be fine
+		if (err) {
+			os_log(OS_LOG_DEFAULT, "tinygpu: failed to map memory, err=%d", err);
+			goto error;
+		}
+
+		// Send back gpu addresses
+		uint64_t* addr = (uint64_t*)memoryMap->GetAddress();
+		for (int i = 0; i < segCount; i++) {
+			addr[i * 2] = segments[i].address;
+			addr[i * 2 + 1] = segments[i].length;
+		}
+		addr[segCount * 2] = 0;
+		addr[segCount * 2 + 1] = 0;
+
+		// free memoryMap
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+
+	dmaDesc->sharedBuf = sharedBuf;
+	dmaDesc->dmaCmd = dmaCmd;
+	return 0;
+
+error:
+	if (memoryMap) {
+		memoryMap->release();
+		memoryMap = nullptr;
+	}
+	if (dmaCmd) {
+		dmaCmd->CompleteDMA(kIODMACommandCompleteDMANoOptions);
+		dmaCmd->release();
+		dmaCmd = nullptr;
+	}
+	if (sharedBuf) {
+		sharedBuf->release();
+		sharedBuf = nullptr;
+	}
+	return err;
+}
+
+kern_return_t TinyGPUDriver::CfgRead(uint32_t off, uint32_t size, uint32_t* outVal)
+{
+  if (!ivars->pci || !outVal) return kIOReturnNotReady;
+
+  if (size == 1) {
+	uint8_t v8 = 0;
+	ivars->pci->ConfigurationRead8(off, &v8);
+	*outVal = v8;
+  } else if (size == 2) {
+	uint16_t v16 = 0;
+	ivars->pci->ConfigurationRead16(off, &v16);
+	*outVal = v16;
+  } else if (size == 4) {
+	uint32_t v32 = 0;
+	ivars->pci->ConfigurationRead32(off, &v32);
+	*outVal = v32;
+  }
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::CfgWrite(uint32_t off, uint32_t size, uint32_t val)
+{
+  if (!ivars->pci) return kIOReturnNotReady;
+  if (size == 1) ivars->pci->ConfigurationWrite8 (off, (uint8_t)val);
+  else if (size == 2) ivars->pci->ConfigurationWrite16(off, (uint16_t)val);
+  else if (size == 4) ivars->pci->ConfigurationWrite32(off, (uint32_t)val);
+  return 0;
+}
+
+kern_return_t TinyGPUDriver::ResetDevice()
+{
+	if (!ivars->pci) return kIOReturnNotReady;
+	ivars->pci->Reset(kIOPCIDeviceResetTypeFunctionReset);
+	return 0;
+}
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.developer.driverkit.transport.pci</key>
+	<array>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x00001002&amp;0x0000FFFF</string>
+		</dict>
+		<dict>
+			<key>IOPCIMatch</key>
+			<string>0x000010de&amp;0x0000FFFF</string>
+		</dict>
+	</array>
+	<key>com.apple.developer.driverkit.allow-any-userclient-access</key>
+	<true/>
+	<key>com.apple.developer.driverkit</key>
+	<true/>
+</dict>
+</plist>
@@ -0,0 +1,38 @@
+#ifndef TinyGPUDriver_h
+#define TinyGPUDriver_h
+
+#include <Availability.h>
+#include <DriverKit/IOService.iig>
+#include <PCIDriverKit/IOPCIDevice.iig>
+#include <DriverKit/IOMemoryMap.iig>
+#include <DriverKit/IODMACommand.iig>
+
+struct TinyGPUCreateDMAResp
+{
+	IOBufferMemoryDescriptor* sharedBuf;
+	IODMACommand* dmaCmd;
+};
+
+class TinyGPUDriver: public IOService
+{
+public:
+	virtual bool init() override;
+
+	virtual void free() override;
+
+	virtual kern_return_t Start(IOService * provider) override;
+
+	virtual kern_return_t Stop(IOService * provider) override;
+
+	virtual kern_return_t NewUserClient(uint32_t in_type, IOUserClient** out_user_client) override;
+
+	kern_return_t MapBar(uint32_t bar, IOMemoryDescriptor** memory) LOCALONLY;
+	kern_return_t CreateDMA(size_t size, TinyGPUCreateDMAResp* dmaDesc) LOCALONLY;
+
+	kern_return_t CfgRead(uint32_t off, uint32_t size, uint32_t* val) LOCALONLY;
+	kern_return_t CfgWrite(uint32_t off, uint32_t size, uint32_t val) LOCALONLY;
+	kern_return_t ResetDevice() LOCALONLY;
+	kern_return_t BarInfo() LOCALONLY;
+};
+
+#endif /* TinyGPUDriver_h */
@@ -0,0 +1,127 @@
+#include "TinyGPUDriverUserClient.h"
+#include "TinyGPUDriver.h"
+#include <DriverKit/DriverKit.h>
+#include <DriverKit/OSSharedPtr.h>
+#include <PCIDriverKit/PCIDriverKit.h>
+
+struct TinyGPUDriverUserClient_IVars
+{
+	OSSharedPtr<TinyGPUDriver> provider = nullptr;
+};
+
+bool TinyGPUDriverUserClient::init()
+{
+	auto theAnswer = super::init();
+	if (!theAnswer) {
+		return false;
+	}
+
+	ivars = IONewZero(TinyGPUDriverUserClient_IVars, 1);
+	if (ivars == nullptr) {
+		return false;
+	}
+
+	return true;
+}
+
+void TinyGPUDriverUserClient::free()
+{
+	if (ivars != nullptr) {
+		ivars->provider.reset();
+	}
+
+	IOSafeDeleteNULL(ivars, TinyGPUDriverUserClient_IVars, 1);
+	super::free();
+}
+
+kern_return_t TinyGPUDriverUserClient::Start_Impl(IOService* in_provider)
+{
+	kern_return_t err = kIOReturnSuccess;
+	if (!in_provider) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: provider is null");
+		err = kIOReturnBadArgument;
+		goto error;
+	}
+
+	err = Start(in_provider, SUPERDISPATCH);
+	if (err) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: failed to start super (%d)", err);
+		goto error;
+	}
+
+	ivars->provider = OSSharedPtr(OSDynamicCast(TinyGPUDriver, in_provider), OSRetain);
+	return 0;
+
+error:
+	ivars->provider.reset();
+	return err;
+}
+
+kern_return_t TinyGPUDriverUserClient::Stop_Impl(IOService* in_provider)
+{
+	return Stop(in_provider, SUPERDISPATCH);
+}
+
+kern_return_t TinyGPUDriverUserClient::ExternalMethod(uint64_t selector, IOUserClientMethodArguments* args, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference)
+{
+	kern_return_t err = 0;
+
+	os_log(OS_LOG_DEFAULT, "tinygpu: rpc (%llu) in:%d, out:%d", selector, args->scalarInputCount, args->scalarOutputCount);
+
+	if (selector == TinyGPURPC::ReadCfg) {
+		if (args->scalarInputCount != 2 or args->scalarOutputCount < 1) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+
+		uint32_t val = 0;
+		err = ivars->provider->CfgRead(off, size, &val);
+		os_log(OS_LOG_DEFAULT, "tinygpu: read cfg off:%x sz:%d, val:%x", off, size, val);
+
+		if (!err) {
+			args->scalarOutput[0] = val;
+			args->scalarOutputCount = 1;
+		}
+		return err;
+	} else if (selector == TinyGPURPC::WriteCfg) {
+		if (args->scalarInputCount != 3) return kIOReturnBadArgument;
+
+		uint32_t off = uint32_t(args->scalarInput[0]);
+		uint32_t size = uint32_t(args->scalarInput[1]);
+		uint32_t val = uint32_t(args->scalarInput[2]);
+
+		os_log(OS_LOG_DEFAULT, "tinygpu: wr cfg off:%x sz:%d, val:%x", off, size, val);
+		return ivars->provider->CfgWrite(off, size, val);
+	} else if (selector == TinyGPURPC::Reset) {
+		os_log(OS_LOG_DEFAULT, "tinygpu: reset");
+		return ivars->provider->ResetDevice();
+	}
+
+	return kIOReturnUnsupported;
+}
+
+kern_return_t IMPL(TinyGPUDriverUserClient, CopyClientMemoryForType)
+{
+	if (!memory) {
+		return kIOReturnBadArgument;
+	}
+
+	if (ivars->provider.get() == nullptr) {
+		return kIOReturnNotAttached;
+	}
+
+	if (type < 6) {
+		uint32_t bar = (uint32_t)type;
+		return ivars->provider->MapBar(bar, memory);
+	}
+
+	// dma page buffer
+	TinyGPUCreateDMAResp buf;
+	kern_return_t err = ivars->provider->CreateDMA(type, &buf);
+	if (err) {
+		return err;
+	}
+
+	*memory = buf.sharedBuf;
+	return 0;
+}
@@ -0,0 +1,28 @@
+#ifndef TinyGPUDriverUserClient_h
+#define TinyGPUDriverUserClient_h
+
+#include <DriverKit/IOUserClient.iig>
+
+enum TinyGPURPC
+{
+	ReadCfg,
+	WriteCfg,
+	Reset
+};
+
+class TinyGPUDriverUserClient : public IOUserClient
+{
+public:
+	virtual bool init() final;
+	virtual void free() final;
+
+	virtual kern_return_t Start(IOService* in_provider) final;
+	virtual kern_return_t Stop(IOService* in_provider) final;
+
+	virtual kern_return_t ExternalMethod(uint64_t in_selector, IOUserClientMethodArguments* in_arguments, const IOUserClientMethodDispatch* in_dispatch, OSObject* in_target, void* in_reference) final;
+
+	virtual kern_return_t CopyClientMemoryForType(
+		uint64_t type, uint64_t *options, IOMemoryDescriptor **memory) final;
+};
+
+#endif /* TinyGPUDriverUserClient_h */
@@ -0,0 +1,12 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN" "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0">
+<dict>
+	<key>com.apple.security.app-sandbox</key>
+	<false/>
+	<key>com.apple.security.files.user-selected.read-only</key>
+	<true/>
+	<key>com.apple.developer.system-extension.install</key>
+	<true/>
+</dict>
+</plist>
@@ -0,0 +1,39 @@
+#include <CoreFoundation/CoreFoundation.h>
+#include <IOKit/IOKitLib.h>
+#include <mach/mach.h>
+#include <stdio.h>
+#include <inttypes.h>
+
+static io_connect_t open_uc_by_name(const char *svc_name) {
+  io_connect_t conn = IO_OBJECT_NULL;
+  io_service_t service = IOServiceGetMatchingService(kIOMasterPortDefault, IOServiceNameMatching(svc_name));
+  if (!service) { fprintf(stderr, "service not found: %s\n", svc_name); return IO_OBJECT_NULL; }
+  kern_return_t kr = IOServiceOpen(service, mach_task_self(), /*type*/0, &conn);
+  IOObjectRelease(service);
+  if (kr) { fprintf(stderr, "IOServiceOpen 0x%x\n", kr); return IO_OBJECT_NULL; }
+  return conn;
+}
+
+int main(int argc, char **argv) {
+  uint32_t bar = (argc > 1) ? (uint32_t)strtoul(argv[1], NULL, 0) : 0;  // pick BAR index
+  io_connect_t conn = open_uc_by_name("tinygpu");
+  if (!conn) return 2;
+
+  mach_vm_address_t addr = 0;
+  mach_vm_size_t    size = 0;
+  kern_return_t kr = IOConnectMapMemory64(conn, bar, mach_task_self(), &addr, &size, kIOMapAnywhere);
+  if (kr) { fprintf(stderr, "Map BAR%u failed 0x%x\n", bar, kr); IOServiceClose(conn); return 3; }
+
+  printf("BAR%u mapped at 0x%llx, size 0x%llx\n", bar, (unsigned long long)addr, (unsigned long long)size);
+
+  // example: read a 32-bit register at offset 0x0 (make sure it’s safe!)
+  volatile uint32_t *mmio = (volatile uint32_t*)(uintptr_t)addr;
+  uint32_t v = mmio[0];
+  printf("mmio[0]=0x%08x\n", v);
+
+  kr = IOConnectUnmapMemory64(conn, bar, mach_task_self(), addr);
+  if (kr) fprintf(stderr, "Unmap failed 0x%x\n", kr);
+
+  IOServiceClose(conn);
+  return 0;
+}
@@ -0,0 +1,82 @@
+import ctypes, ctypes.util, sys
+
+cf = ctypes.CDLL(ctypes.util.find_library("CoreFoundation"))
+iokit = ctypes.CDLL(ctypes.util.find_library("IOKit"))
+libsys = ctypes.CDLL(ctypes.util.find_library("System"))
+
+kern_return_t = ctypes.c_int
+mach_port_t = ctypes.c_uint
+io_object_t = mach_port_t
+io_service_t = io_object_t
+io_connect_t = mach_port_t
+CFMutableDictionaryRef = ctypes.c_void_p
+CFStringRef = ctypes.c_void_p
+
+kIOMasterPortDefault = mach_port_t(0)
+
+libsys.mach_task_self_.restype = mach_port_t
+
+iokit.IOServiceNameMatching.argtypes = [ctypes.c_char_p]
+iokit.IOServiceNameMatching.restype  = CFMutableDictionaryRef
+
+iokit.IOServiceGetMatchingService.argtypes = [mach_port_t, CFMutableDictionaryRef]
+iokit.IOServiceGetMatchingService.restype  = io_service_t
+
+iokit.IOObjectRelease.argtypes = [io_object_t]
+iokit.IOObjectRelease.restype  = kern_return_t
+
+iokit.IOServiceOpen.argtypes = [io_service_t, mach_port_t, ctypes.c_uint32, ctypes.POINTER(io_connect_t)]
+iokit.IOServiceOpen.restype  = kern_return_t
+
+iokit.IOConnectCallMethod.argtypes = [io_connect_t, ctypes.c_uint32, ctypes.POINTER(ctypes.c_uint64), ctypes.c_uint32, ctypes.c_void_p,
+  ctypes.c_size_t, ctypes.POINTER(ctypes.c_uint64), ctypes.POINTER(ctypes.c_uint32), ctypes.c_void_p, ctypes.POINTER(ctypes.c_size_t)]
+iokit.IOConnectCallMethod.restype = kern_return_t
+
+def open_userclient_by_name(name: str, uc_type: int = 0) -> io_connect_t:
+  mdict = iokit.IOServiceNameMatching(name.encode("utf-8"))
+  if not mdict: raise RuntimeError("IOServiceNameMatching returned NULL")
+
+  # Grab the first matching service
+  service = iokit.IOServiceGetMatchingService(kIOMasterPortDefault, mdict)
+  if not service: raise RuntimeError(f'service "{name}" not found')
+
+  # print("lol", service)
+  # print(libsys.mach_task_self_)
+  # cast libsys.mach_task_self_ to uint and print
+  # print("lol", ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value)
+
+  try:
+    # Open user client (type -> passed to NewUserClient_Impl)
+    conn = io_connect_t(0)
+    # print("lol", libsys.mach_task_self_)
+    kr = iokit.IOServiceOpen(service, ctypes.cast(libsys.mach_task_self_, ctypes.POINTER(ctypes.c_uint)).contents.value,
+                             ctypes.c_uint32(uc_type), ctypes.byref(conn))
+    if kr != 0: raise OSError(kr, f"IOServiceOpen failed (0x{kr:08x})")
+    return conn
+  finally: iokit.IOObjectRelease(service)
+
+def external_method(conn: io_connect_t, selector: int = 0) -> int:
+  # no scalars in/out, no struct in/out — just ping selector 0
+  in_scalars = ctypes.POINTER(ctypes.c_uint64)()  # NULL
+  out_scalars = (ctypes.c_uint64 * 1)()       # space if driver returns something
+  out_scalars_cnt = ctypes.c_uint32(0)      # driver can set this
+
+  return iokit.IOConnectCallMethod(conn, ctypes.c_uint32(selector), in_scalars, ctypes.c_uint32(0), None, ctypes.c_size_t(0),
+    out_scalars, ctypes.byref(out_scalars_cnt), None, ctypes.byref(ctypes.c_size_t(0)))
+
+def close_userclient(conn: io_connect_t) -> None:
+  # IOServiceClose is a macro; exported symbol is IOServiceClose in IOKit
+  iokit.IOServiceClose.argtypes = [io_connect_t]
+  iokit.IOServiceClose.restype  = kern_return_t
+  iokit.IOServiceClose(conn)
+
+if __name__ == "__main__":
+  try:
+    conn = open_userclient_by_name("tinygpu", uc_type=0)
+    kr = external_method(conn, selector=0)
+    print(f"ExternalMethod(0) -> 0x{kr:08x}")
+  except Exception as e:
+    print(e)
+    sys.exit(1)
+  finally:
+    if 'conn' in locals() and conn.value: close_userclient(conn)
@@ -50,4 +50,7 @@ exclude = [
  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
  "FURB110", "RUF018", "F541", "F841"
-]
+]
+
+[format]
+exclude = ["*"]
@@ -42,7 +42,6 @@ setup(name='tinygrad',
        'tinygrad.runtime.support.am',
        'tinygrad.runtime.support.nv',
        'tinygrad.schedule',
-        'tinygrad.shape',
        'tinygrad.uop',
        'tinygrad.viz',
      ],
@@ -1,4 +1,4 @@
-from tinygrad import Tensor, dtypes, Context, GlobalCounters
+from tinygrad import Tensor, dtypes, GlobalCounters
 dtypes.default_float = dtypes.float16
 from tinygrad.dtype import to_dtype
 from tinygrad.helpers import getenv
@@ -13,6 +13,5 @@ if __name__ == "__main__":

  # test single kernel softmax
  GlobalCounters.reset()
-  with Context(DONT_GROUP_REDUCES=1):
-    single_kernel_softmax(t, -1, acc_dtype).realize()
+  single_kernel_softmax(t, -1, acc_dtype).realize()

@@ -1,46 +0,0 @@
-# ruff: noqa: E501
-from tinygrad.codegen.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.codegen.opt.search import bufs_from_lin
-from tinygrad.uop.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, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=2, arg=0), Opt(op=OptOps.UNROLL, axis=1, arg=0), Opt(op=OptOps.LOCAL, axis=1, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=8), Opt(op=OptOps.LOCAL, axis=2, arg=2)]
-
-k = Kernel(ast)
-k.apply_opts(opts)
-bufs = bufs_from_lin(k)
-
-prg = CompiledRunner(get_program(k.ast, k.opts, k.applied_opts))
-
-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
@@ -0,0 +1,39 @@
+import subprocess, unittest, os, sys
+from tinygrad.device import Device
+
+class TestTinygradSlow(unittest.TestCase):
+  def test_env_overwrite_default_device(self):
+    subprocess.run([f'{Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+    subprocess.run([f'DISK=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+    subprocess.run([f'NPY=1 {Device.DEFAULT}=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                    shell=True, check=True)
+
+    if Device.DEFAULT != "CPU":
+      # setting multiple devices fail
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'{Device.DEFAULT}=1 CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+      # setting device via DEV
+      subprocess.run([f'DEV={Device.DEFAULT.capitalize()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.lower()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+      subprocess.run([f'DEV={Device.DEFAULT.upper()} python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                      shell=True, check=True)
+
+      with self.assertRaises(subprocess.CalledProcessError):
+        subprocess.run([f'DEV={Device.DEFAULT} CPU=1 python3 -c "from tinygrad import Device; assert Device.DEFAULT == \\"{Device.DEFAULT}\\""'],
+                        shell=True, check=True)
+
+class TestRunAsModule(unittest.TestCase):
+  def test_module_runs(self):
+    p = subprocess.run([sys.executable, "-m", "tinygrad.device"],stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+      env={**os.environ, "DEBUG": "1"}, timeout=40,)
+    out = (p.stdout + p.stderr).decode()
+    self.assertEqual(p.returncode, 0, msg=out)
+
+if __name__ == '__main__':
+  unittest.main()
@@ -58,8 +58,8 @@ class TestExample(unittest.TestCase):
      print(f"WARNING: {device} test isn't running")
      return

-    x = Tensor.eye(64, device=device, requires_grad=True)
-    y = Tensor.eye(64, device=device, requires_grad=True)
+    x = Tensor.eye(8, device=device, requires_grad=True)
+    y = Tensor.eye(8, device=device, requires_grad=True)
    z = y.matmul(x).sum()
    z.backward()

@@ -1,41 +1,50 @@
+import functools, multiprocessing
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
+from tinygrad.apps.llm import SimpleTokenizer
 from tinygrad.helpers import tqdm, getenv, partition

+@functools.cache
+def get_tokenizers():
+  print("getting tokenizers")
+  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids)
+  simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+  return base_tokenizer, simple_tokenizer
+
+def test_tokenize(samp) -> bool:
+  base_tokenizer, simple_tokenizer = get_tokenizers()
+  idx, txt = samp
+  try: simple_tokens = tuple(simple_tokenizer.encode(txt))
+  except RuntimeError: simple_tokens = ()
+  base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False))
+  if simple_tokens != base_tokens:
+    print(f"tokens mismatch at index: {idx}.\n")
+    color_codes = [91, 92, 94, 93, 95]
+    def color_tokens(tids):
+      return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
+    print("simple:  ", color_tokens(simple_tokens))
+    print("official:", color_tokens(base_tokens) + "\n")
+    return False
+  if simple_tokenizer.decode(simple_tokens) != txt:
+    print(f"decode mismatch at {idx}")
+    return False
+  return True
+
 # use ALLOW_FAILED=-1 to go over the entire dataset without printing.
 if __name__ == "__main__":
-  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
-    lambda e: e[1] in base_tokenizer.all_special_ids)
-  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
-  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
-
-  color_codes = [ 91, 92, 94, 93, 95 ]
-  def color_tokens(tids):
-    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
-
+  print("loading datasets")
  ds = load_dataset("OpenAssistant/oasst1")
+  loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])]
+  print(f"loaded {len(loaded_ds)}")
+
  allow_failed = getenv("ALLOW_FAILED", 10)
-
  fail_count, total = 0, 0
-
-  for idx, el in enumerate(tqdm(ds["train"])):
-    total += 1
-
-    try: simple_tokens = tuple(simple_tokenizer.encode(el["text"]))
-    except RuntimeError: simple_tokens = ()
-    base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False))
-
-    if simple_tokens != base_tokens:
-      fail_count += 1
-      allow_failed -= 1
-
-      if allow_failed >= 0:
-        print(f"tokens mismatch at index: {idx}.\n")
-
-        print("simple:  ", color_tokens(simple_tokens))
-        print("official:", color_tokens(base_tokens) + "\n")
-
-      if allow_failed == 0: break
-  print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
+  with multiprocessing.Pool(16) as pool:
+    for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)):
+      total += 1
+      if not good:
+        fail_count += 1
+        allow_failed -= 1
+        if allow_failed == 0: break
+    print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
@@ -1,55 +0,0 @@
-# ruff: noqa: E501
-import unittest
-
-from tinygrad.uop.ops import UOp, Ops
-from .search import Opt, OptOps
-from tinygrad.dtype import dtypes
-from tinygrad.shape.shapetracker import ShapeTracker
-from tinygrad.shape.view import View
-from tinygrad.codegen.opt.kernel import Kernel
-
-from test.external.fuzz_linearizer import run_linearizer
-
-class TestTrainGpt2Kernel(unittest.TestCase):
-  def test_1(self):
-    # kernel 244
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(206045184), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 1), strides=(51511296, 50304, 1, 0), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=16), Opt(op=OptOps.UPCAST, axis=1, arg=3), Opt(op=OptOps.LOCAL, axis=0, arg=2)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-  def test_2(self):
-    # kernel 254
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(3145728), arg=0, src=()),
-        UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 1, 768), strides=(786432, 768, 0, 1), offset=0, mask=None, contiguous=True),)), src=()),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (2,)), src=(
-        UOp(Ops.MUL, dtypes.float, arg=None, src=(
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(38633472), arg=1, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(0, 0, 768, 1), offset=0, mask=None, contiguous=False),)), src=()),)),
-          UOp(Ops.LOAD, dtypes.float, arg=None, src=(
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(205852672), arg=2, src=()),
-          UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4, 1024, 50304, 768), strides=(51463168, 50257, 1, 0), offset=0, mask=((0, 4), (0, 1024), (0, 50257), (0, 768)), contiguous=False),)), src=()),)),)),)),)),))
-
-    opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16), Opt(op=OptOps.LOCAL, axis=0, arg=8), Opt(op=OptOps.UPCAST, axis=2, arg=4), Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.LOCAL, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=3, arg=4)]
-    kernel = Kernel(ast)
-    kernel.apply_opts(opts)
-    run_linearizer(kernel)
-
-if __name__ == "__main__":
-  unittest.main()
@@ -1,6 +1,5 @@
 import gc
 from tinygrad import Tensor, UOp, Device, nn
-from tinygrad.shape.shapetracker import views_to_valid_uop
 from tinygrad.engine.realize import method_cache, get_program
 from tinygrad.schedule.indexing import apply_movement_op
 from test.test_tiny import TestTiny
@@ -69,7 +68,6 @@ if __name__ == "__main__":

    # these caches will keep uops alive
    method_cache.clear()
-    views_to_valid_uop.cache_clear()
    apply_movement_op.cache_clear()
    Tensor._device_seeds.clear()
    Tensor._device_rng_counters.clear()
@@ -42,13 +42,13 @@ class ProcessReplayWarning(Warning): pass

 # *** replay the function and convert return values to string

-def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
+def replay_get_rangeify_map(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[Any, ...]]:
  UOp.unique_num = itertools.count(max([u.arg for u in big_sink.toposort() if u.op is Ops.UNIQUE], default=0)+1)
  new_sink = big_sink.substitute(get_rangeify_map(big_sink))
  def to_str(ret:UOp) -> str:
    asts = [repr(u.arg.ast) for u in ret.toposort() if u.op is Ops.KERNEL]
    return "\n".join([f"{len(asts)} kernels", *asts])
-  return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)
+  return to_str(new_sink), to_str(big_sink.substitute(ret)), (big_sink,)

 def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) -> tuple[str, str, tuple[Any, ...]]:
  # NOTE: this always uses the opts_to_apply path
@@ -65,7 +65,7 @@ def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer|None=None, opts
  ast_repr = codecs.decode(str(input_ast), "unicode_escape")
  return to_str(p2), to_str(p), (ast_repr, renderer)

-replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_kernelize_map":replay_kernelize, "get_program":replay_get_program}
+replayers: dict[str, Callable[..., tuple[str, str, tuple[Any, ...]]]] = {"get_rangeify_map":replay_get_rangeify_map, "get_program":replay_get_program}

 # *** run replayers on captured rows and print diffs

@@ -2,7 +2,6 @@ import unittest
 from tinygrad import Device, Tensor, dtypes
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.codegen.opt import Opt, OptOps
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.engine.realize import get_program
 from tinygrad.helpers import AMX

@@ -149,33 +148,5 @@ class TestFloat4(unittest.TestCase):

    assert TestFloat4.count_float4(uops) == (1, 1)

-  @unittest.skip("Ops.VIEW no longer exists")
-  def test_half4_load_unrolled(self):
-    # from llama 7B shard 4 gpus
-    ast = UOp(Ops.SINK, dtypes.void, arg=None, src=(
-      UOp(Ops.STORE, dtypes.void, arg=None, src=(
-        UOp(Ops.VIEW, dtypes.float.ptr(96000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1), strides=(0, 32000, 1, 0), offset=0, mask=None, contiguous=True),)), src=( # noqa: E501
-          UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(96000), arg=0, src=()),)),
-        UOp(Ops.REDUCE_AXIS, dtypes.float, arg=(Ops.ADD, (3,)), 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.VIEW, dtypes.half.ptr(9216), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 4096, 0, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(9216), arg=1, src=()),)),)),
-              UOp(Ops.LOAD, dtypes.half, arg=None, src=(
-                UOp(Ops.VIEW, dtypes.half.ptr(32768000), arg=ShapeTracker(views=(View(shape=(1, 3, 32000, 1024), strides=(0, 0, 1024, 1), offset=0, mask=None, contiguous=False),)), src=( # noqa: E501
-                  UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(32768000), arg=2, src=()),)),)),)),)),)),)),))
-
-    # TODO: fix this, expected might change but should be positive
-    for expected, opts in [
-      ((7, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UPCAST, axis=0, arg=3), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((5, 0), [Opt(op=OptOps.UPCAST, axis=1, arg=4), Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-      ((2, 0), [Opt(op=OptOps.UNROLL, axis=0, arg=4)]),
-    ]:
-      program = get_program(ast, Device[Device.DEFAULT].renderer, opts=opts)
-
-      count = TestFloat4.count_half4(program.uops)
-      assert count == expected, f"{count=}, {expected=}"
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,6 +1,5 @@
 import unittest
 from tinygrad import Device, Tensor, dtypes
-from tinygrad.helpers import CI
 from tinygrad.codegen.opt import Opt, OptOps, KernelOptError

 # TODO: write a clean version of this
@@ -177,9 +176,7 @@ class TestKernelOpts(unittest.TestCase):
    ], apply_tc=True, atol=atol, rtol=rtol)

  def test_padto_matmul(self):
-    if (CI and Device.DEFAULT in ["AMD", "NV", "CUDA"]):
-      self.skipTest("super slow on CUDA and AMD because of the big grid dims")
-    N = 17 * 17
+    N = 17
    Tensor.manual_seed(289)
    a = Tensor.rand(N, N)
    b = Tensor.rand(N, N)
@@ -213,7 +210,7 @@ class TestKernelOpts(unittest.TestCase):
      helper_linearizer_opt(a@b, [[Opt(OptOps.UNROLL, 0, 0), Opt(OptOps.PADTO, 2, 8)]])

  def test_padto_sum_ok(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
    a = Tensor.rand(N, N).realize().shrink(((0, 17), (0, 17))) * 100
    b = (Tensor.rand(N, N) < 0.5).realize().shrink(((0, 17), (0, 17)))
@@ -244,7 +241,7 @@ class TestKernelOpts(unittest.TestCase):
    helper_linearizer_opt(a.sum(0).exp(), [[Opt(OptOps.PADTO, 1, 32)],])

  def test_padto_sum_not_ok(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one dimension
    a = Tensor.rand(N, N).shrink(((0, 17), (0, 17))).exp()
    # exp is not safe to pad
@@ -261,7 +258,7 @@ class TestKernelOpts(unittest.TestCase):
      helper_linearizer_opt(b.sum(0), [[Opt(OptOps.PADTO, 1, 32)],])

  def test_padto_max(self):
-    N = 18 * 18
+    N = 18
    # NOTE: this setup prevents 17 * 17 contiguous merged into one axis
    a = -Tensor.rand(N, N).shrink(((0, 17), (0, 17))) * 100

@@ -282,7 +279,7 @@ class TestKernelOpts(unittest.TestCase):

  def test_padto_where(self):
    Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
    a = (Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1).where(1, 0)
    helper_linearizer_opt(a.max(0), [
      [Opt(OptOps.PADTO, 0, 32)],
@@ -291,7 +288,7 @@ class TestKernelOpts(unittest.TestCase):

  def test_padto_where_multioutput(self):
    Tensor.manual_seed(0)
-    N = 17 * 17
+    N = 17
    r = Tensor.randn(N, N).realize().max(axis=0, keepdim=True) > 1
    a0 = r.where(1, 0)
    a1 = r.where(2, 0)
@@ -131,8 +131,7 @@ class TestIndexing(unittest.TestCase):
    # llama3 is 128256
    vocab_size, embed_size = (10, 3) if CI else (32000, 4096)
    emb = nn.Embedding(vocab_size, embed_size)
-    # TODO: why is a new realize needed here
-    emb_w = emb.weight.realize().numpy()
+    emb_w = emb.weight.numpy()
    x = Tensor([1,2,3,4])
    with Context(NOOPT=noopt):
      GlobalCounters.reset()
@@ -129,6 +129,7 @@ class TestAssign(unittest.TestCase):
  @unittest.expectedFailure
  def test_assign_changes_realized_alt(self): return self.test_assign_changes_alt(realize=True)

+  @unittest.skip("assign to contiguous shouldn't change the base buffer")
  def test_assign_changes_buffer_alt(self):
    a, b = [Tensor(Tensor(0).contiguous().realize().uop.as_buf()) for _ in range(2)]
    Tensor.realize(a.contiguous().assign(1), b.contiguous().assign(2))
@@ -1,7 +1,7 @@
 import unittest, io
 from contextlib import redirect_stdout
 from tinygrad import Tensor, dtypes, Device
-from tinygrad.helpers import OSX, CPU_LLVM
+from tinygrad.helpers import OSX, CPU_LLVM, CPU_LVP
 from tinygrad.engine.realize import lower_schedule
 from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import get_program
@@ -19,7 +19,7 @@ class TestCompileFailures(unittest.TestCase):

 class TestDisassembly(unittest.TestCase):
  # TODO: fails on llvm. llvm.LLVMGetHostCPUName() returns "generic"
-  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not CPU_LLVM and OSX, "m series cpus support fp16 arithmetic")
+  @unittest.skipUnless(Device.DEFAULT in ("CPU",) and not (CPU_LLVM or CPU_LVP) and OSX, "m series cpus support fp16 arithmetic")
  def test_float16_alu(self):
    c = Tensor([1], dtype=dtypes.float16) + Tensor([1], dtype=dtypes.float16)
    s = c.schedule()[-1]
@@ -6,6 +6,7 @@ from tinygrad.device import is_dtype_supported
 from tinygrad.helpers import getenv, DEBUG, CI
 from tinygrad.dtype import DType, DTYPES_DICT, least_upper_dtype, fp8_to_float, float_to_fp8, _to_np_dtype, _to_torch_dtype, truncate
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from tinygrad import Device, Tensor, dtypes
 from hypothesis import given, settings, strategies as strat
 from test.helpers import rand_for_dtype
@@ -102,7 +103,7 @@ class TestDType(unittest.TestCase):
    ))

  @unittest.skipIf(Device.DEFAULT == "PYTHON", "skip for now")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "skip for now")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "skip for now")
  def test_uint_overflow(self):
    if not dtypes.is_unsigned(self.DTYPE): raise unittest.SkipTest("only for unsigned")
    v = dtypes.max(self.DTYPE)
@@ -261,7 +262,7 @@ class TestFloatDType(TestDType):
 class TestDoubleDType(TestDType):
  DTYPE = dtypes.double
  @unittest.skipIf((CI and Device.DEFAULT in {"CUDA", "NV"}) or \
-                    isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "conversion not supported on CI CUDA and PTX")  # TODO: why not?
+   isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, NIRRenderer)), "conversion not supported on CI CUDA, PTX, and NIR")  # TODO: why not?
  def test_float64_increased_precision(self):
    for func in [
      lambda t: t.exp(),
@@ -6,6 +6,7 @@ from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
 from tinygrad.runtime.ops_python import from_storage_scalar
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 import numpy as np
 import pytest
 from hypothesis import assume, given, strategies as strat, settings, HealthCheck
@@ -29,8 +30,8 @@ unary_operations = [(Tensor.exp, np.exp), (Tensor.log, np.log), (Tensor.sin, np.
 # TODO: enable this (this is a dtype issue)
 #binary_operations.append(operator.truediv)

-# TODO: CI CUDA segfaults on sin, WEBGPU sin is not precise enough for large numbers
-if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU":
+# TODO: CI CUDA segfaults on sin, WEBGPU and NIR sines are not precise enough for large numbers
+if (getenv("MOCKGPU") and Device.DEFAULT in {"NV", "CUDA"}) or Device.DEFAULT == "WEBGPU" or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer):
  unary_operations.remove((Tensor.sin, np.sin))
  unary_operations.remove((Tensor.cos, np.cos))

@@ -184,8 +185,8 @@ class TestDTypeALU(unittest.TestCase):
  @given(ht.int32, ht.int32, ht.float32, strat.sampled_from(integer_binary_operations), strat.sampled_from(binary_operations))
  def test_int32_midcast_float(self, a, b, c, op1, op2): universal_test_midcast(a, b, c, op1, op2, dtypes.int32, dtypes.float32)

-  # Metal and CUDA and HIP behave differently than numpy in CI for overflows
-  skip_overflow = CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}
+  # Metal and CUDA and HIP and NIR behave differently than numpy in CI for overflows
+  skip_overflow = (CI and Device.DEFAULT in {"AMD", "NV", "CUDA"}) or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer)
  @given(strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         strat.floats(width=32, min_value=0, max_value=10.0) if skip_overflow else ht.float32,
         ht.int32, strat.sampled_from(binary_operations), strat.sampled_from(integer_binary_operations))
@@ -26,8 +26,9 @@ import unittest
 import numpy as np
 import torch
 from tinygrad import Tensor, dtypes, nn
-from tinygrad.device import is_dtype_supported
+from tinygrad.device import Device, is_dtype_supported
 from tinygrad.helpers import getenv
+from tinygrad.renderer.nir import NIRRenderer

 MOCKGPU = getenv("MOCKGPU")

@@ -206,7 +207,8 @@ class TestUOpValidationIssue(unittest.TestCase):
  # these fail with UOp verification error.
  # we want more of these with diverse errors!

-  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU, "hangs gpuocelot")
+  @unittest.skipIf((not is_dtype_supported(dtypes.long)) or MOCKGPU or isinstance(Device[Device.DEFAULT].renderer, NIRRenderer),
+                   "hangs gpuocelot, NIR cannot render")
  def test_tensor_index_overflow(self):
    val = Tensor([1])
    big = val.expand(2**31 + 3)
@@ -8,9 +8,11 @@ from tinygrad.uop.ops import UOp, Ops, GroupOp
 from tinygrad.device import Device, Buffer, is_dtype_supported
 from tinygrad.tensor import Tensor, _to_np_dtype
 from tinygrad.engine.realize import run_schedule, lower_schedule, CompiledRunner, get_program
-from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT
+from tinygrad.helpers import Context, flatten, dedup, TC_SELECT, TC_OPT, getenv
 from tinygrad.dtype import DType, dtypes, PtrDType, AddrSpace
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.cstyle import CUDARenderer
+MOCKGPU = getenv("MOCKGPU")

 class TestLinearizer(unittest.TestCase):
  def test_arg_dedup(self):
@@ -68,7 +70,8 @@ class TestLinearizer(unittest.TestCase):
    ranges = [i for i,u in enumerate(uops) if u.op is Ops.RANGE]
    # RANGE -> ALU -> RANGE -> ALU + LOAD -> STORE
    assert any(x.op in GroupOp.ALU for x in uops[ranges[0]:ranges[1]])
-    assert not any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
+    # the index of the load doesnt depend on the second range
+    assert any(x.op is Ops.LOAD for x in uops[ranges[0]:ranges[1]])
    assert any(x.op in {*GroupOp.ALU, Ops.LOAD} for x in uops[ranges[1]:])

  def test_range_outer_op_before_phi(self):
@@ -314,7 +317,7 @@ class TestLinearizer(unittest.TestCase):
    a.realize()
    np.testing.assert_equal(a.flatten().numpy(), [1.,1.,1.,1.,2.,2.,2.,2.,1.,1.,1.,1.,1.,1.,1.,1.])

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "PTX indexes differently. might be ok?")
+  @unittest.skipIf(MOCKGPU and isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, CUDARenderer)), "PTX indexes differently. might be ok?")
  def test_where_fold(self):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.shrink(((1, 2), None)).pad(((1, 2), None))
@@ -4,12 +4,9 @@

 import unittest
 from tinygrad import Device, dtypes
-from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops, AxisType, KernelInfo
-from tinygrad.shape.shapetracker import ShapeTracker, View
 from tinygrad.codegen.opt.search import Opt, OptOps
 from tinygrad.engine.realize import get_program
-from tinygrad.renderer.ptx import PTXRenderer

 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
@@ -30,91 +27,5 @@ class TestLinearizerFailure(unittest.TestCase):
    ast = c12.sink(arg=KernelInfo(name='test', axis_types=(), dont_use_locals=False, applied_opts=(Opt(op=OptOps.GROUP, axis=1, arg=16),), opts_to_apply=None))
    _ = get_program(ast, Device["METAL"].renderer)

-class TestLinearizerDumb(unittest.TestCase):
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.has_local, "need local")
-  @unittest.skip("Ops.VALID no longer exists")
-  def test_max_simplify_and_cancel(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1000), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1000), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(1, 0), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c9 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1001, 1999), strides=(0, 0), offset=0, mask=((0, 1001), (999, 1999)), contiguous=False), View(shape=(1000, 1000), strides=(1, 2000), offset=0, mask=None, contiguous=False))), src=())
-    c10 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(1000, 1000), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c11 = c1.store((c4.alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c8)).cast(dtypes.int)*(c9.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, -1, src=c10), UOp.const(dtypes.int, 0, src=c10)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,)))+UOp.const(dtypes.int, 1000, src=c8))))
-    ast = c11.sink()
-    #opts = [Opt(op=OptOps.UNROLL, axis=0, arg=4), Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    opts = [Opt(op=OptOps.LOCAL, axis=0, arg=8)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    assert prg.uops is not None and not any(uop.op is Ops.MAX for uop in prg.uops), "leftover MAX"
-
-  # this was a bug in embedding, someday we should fold this anyway
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), f"half dtype not supported on {Device.DEFAULT}")
-  @unittest.skip("UOp.view is no longer supported")
-  def test_llama_embedding(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(4096), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4096, 1, 1), strides=(1, 0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(32001, 63999), strides=(0, 0), offset=0, mask=((0, 32001), (31999, 63999)), contiguous=False), View(shape=(4096, 32000, 32000), strides=(0, 1, 64000), offset=0, mask=None, contiguous=False))), src=())
-    c3 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 32000), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c4 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.int.ptr(1), arg=1, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = UOp(Ops.DEFINE_GLOBAL, dtypes.half.ptr(131072000), arg=2, src=())
-    c9 = c8.view(ShapeTracker(views=(View(shape=(4096, 32000, 1), strides=(1, 4096, 0), offset=0, mask=None, contiguous=False),)))
-    c10 = c9.load()
-    c11 = c1.store(((c2.f(Ops.VALID, dtype=dtypes.bool).where(UOp.const(dtypes.int, 1, src=c3), UOp.const(dtypes.int, 0, src=c3)).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (2,)))+UOp.const(dtypes.int, -1, src=c4)).alu(Ops.CMPNE, c7).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c4)).cast(dtypes.half)*c10).cast(dtypes.float).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (1,))).cast(dtypes.half))
-    ast = c11.sink()
-    prg = get_program(ast, Device[Device.DEFAULT].renderer)
-    print(prg.src)
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  def test_unrolled_float4_align(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(1, 1), strides=(0, 0), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.long.ptr(18), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.VIEW, dtypes.void, arg=ShapeTracker(views=(View(shape=(3, 6), strides=(0, 0), offset=0, mask=None, contiguous=False),)), src=())
-    c6 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(18), arg=2, src=())
-    c7 = c6.view(ShapeTracker(views=(View(shape=(3, 6), strides=(6, 1), offset=0, mask=None, contiguous=True),)))
-    c8 = c7.load()
-    c9 = c1.store(c4.alu(Ops.CMPNE, UOp.const(dtypes.long, -1, src=c5)).alu(Ops.CMPNE, UOp.const(dtypes.bool, True, src=c5)).where(UOp.const(dtypes.float, 0.0, src=c5), c8).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (0, 1))))
-    ast = c9.sink()
-    opts = [Opt(op=OptOps.UNROLL, axis=0, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    load_idxs = [x.src[1] for x in prg.uops if x.op is Ops.LOAD and x.src[0].arg == 2]
-    assert load_idxs[0] < load_idxs[1], f"first loaded idx {load_idxs[0].arg} then {load_idxs[1].arg}!"
-
-  @unittest.expectedFailure
-  @unittest.skipUnless(Device[Device.DEFAULT].renderer.supports_float4, "need float4")
-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "this is somehow correct in PTX")
-  def test_upcasted_stores_out_of_order(self):
-    c0 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(9360), arg=0, src=())
-    c1 = c0.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 1, 1, 4, 3, 3), strides=(2340, 468, 36, 0, 0, 0, 0, 0, 9, 3, 1), offset=0, mask=None, contiguous=True),)))
-    c2 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(144), arg=1, src=())
-    c3 = c2.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(0, 0, 0, 0, 0, 0, 1, 0, 4, 48, 16), offset=0, mask=None, contiguous=False),)))
-    c4 = c3.load()
-    c5 = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(1040), arg=2, src=())
-    c6 = c5.view(ShapeTracker(views=(View(shape=(4, 5, 13, 1, 1, 1, 4, 1, 4, 3, 3), strides=(260, 13, 1, 0, 0, 0, 65, 0, 0, 0, 0), offset=0, mask=None, contiguous=False),)))
-    c7 = c6.load()
-    c8 = c1.store((c4*c7).f(Ops.REDUCE_AXIS, arg=(Ops.ADD, (6,))))
-    ast = c8.sink()
-    opts = [Opt(op=OptOps.UPCAST, axis=3, arg=0), Opt(op=OptOps.UPCAST, axis=2, arg=0)]
-    prg = get_program(ast, Device[Device.DEFAULT].renderer, opts)
-    print(prg.src)
-    store_idxs = [x.src[1] for x in prg.uops if x.op is Ops.STORE]
-    for i in range(len(store_idxs) - 1):
-      first_bounds = store_idxs[i].vmin+store_idxs[i].vmax
-      next_bounds = store_idxs[i+1].vmin+store_idxs[i+1].vmax
-      assert first_bounds < next_bounds, f"first stored (max) idx {first_bounds} then {next_bounds}!"
-
 if __name__ == '__main__':
  unittest.main()
@@ -390,7 +390,6 @@ class TestMultiTensor(unittest.TestCase):

  # NOTE: this is failing on LLVM CI, no idea why. Works locally.
  @unittest.skipIf(CI and REAL_DEV in ("CUDA", "NV", "CPU", "AMD"), "slow, and flaky on CPU")
-  @unittest.skip("TODO: pm_rangeify hangs")
  def test_data_parallel_resnet(self):
    from extra.models.resnet import ResNet18

@@ -659,7 +658,7 @@ class TestMultiTensor(unittest.TestCase):

  # it doesn't work like this anymore
  # NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
-  @unittest.expectedFailure
+  @unittest.skip("this test is broken")
  def test_mlb_assign_change_axis(self):
    t_none = Tensor.zeros((16, 16)).shard(devices_2).contiguous().realize()
    t_zero = Tensor.ones((16, 16)).shard(devices_2, axis=0)
@@ -2,7 +2,7 @@ import time, math, unittest, functools, platform, warnings
 import numpy as np
 from typing import List, Callable
 import torch
-from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, AMD_LLVM
+from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, CPU_LLVM, CPU_LVP, AMD_LLVM
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.device import is_dtype_supported
@@ -698,8 +698,8 @@ class TestOps(unittest.TestCase):

  def test_pow_zero_tensor(self):
    helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.0]])
-    # TODO: fix WEBGPU
-    if Device.DEFAULT != "WEBGPU":
+    # TODO: fix WEBGPU and LVP
+    if Device.DEFAULT != "WEBGPU" and not CPU_LVP:
      helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [0.3]])
      helper_test_op(None, lambda x,y: x**y, vals=[[0.0], [-0.3]])
  def test_pow_zero_const(self):
@@ -830,6 +830,7 @@ class TestOps(unittest.TestCase):
    self.assertEqual(a, b)
    self.assertEqual(Tensor(-1).contiguous().idiv(4).item(), 0)  # NOTE this is trunc-div behaviour

+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_sin(self):
    helper_test_op([(45,65)], lambda x: x.sin())
    helper_test_op([()], lambda x: x.sin())
@@ -839,6 +840,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.sin(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_cos(self):
    helper_test_op([(45,65)], lambda x: x.cos())
    helper_test_op([()], lambda x: x.cos())
@@ -847,6 +849,7 @@ class TestOps(unittest.TestCase):
      helper_test_op(None, lambda x: x.cos(), vals=[[1e1, 1e2, 1e3, 1e4, 1e5, 1e6, -1e1, -1e2, -1e3, -1e4, -1e5, -1e6]],
                    atol=3e-3, rtol=3e-3, grad_atol=3e-3, grad_rtol=3e-3)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU" and platform.system() == "Windows", "Not accurate enough with DirectX backend")
+  @unittest.skipIf(getenv("NV_NAK"), "MUFU.SIN is not accurate enough")
  def test_tan(self):
    # NOTE: backward has much higher diff with input close to pi/2 and -pi/2
    helper_test_op([(45,65)], lambda x: x.tan(), low=-1.5, high=1.5)
@@ -3177,6 +3180,7 @@ class TestOps(unittest.TestCase):
  def test_bitcast(self):
    helper_test_op([(3, 3)], lambda x: x.view(torch.int32), lambda x: x.bitcast(dtypes.int32), forward_only=True)

+  @unittest.skip("we have test_linalg, no need to test here. TODO: should be in torch backend tests")
  def test_svd(self):
    # test for tiny backend. real svd tests are in test_linalg
    A = torch.randn(5, 5)
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, Device
-from tinygrad.helpers import CPU_LLVM
+from tinygrad.helpers import CPU_LLVM, CPU_LVP
 from tinygrad.codegen.opt import Opt, OptOps
 from tinygrad.engine.realize import get_program

@@ -12,7 +12,7 @@ class TestOpts(unittest.TestCase):
    out = (a+b).contiguous(arg=opts)
    s = out.schedule()
    self.assertEqual(s[-1].ast.arg.opts_to_apply, opts)
-    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM:
+    if Device.DEFAULT in {"CPU", "CL", "METAL"} and not CPU_LLVM and not CPU_LVP:
      prg = get_program(s[-1].ast)
      self.assertIn('float4', prg.src)

@@ -217,7 +217,7 @@ class TestProfiler(unittest.TestCase):
        Tensor.realize(a, b)
    profile, _ = helper_profile_filter_device(profile, TestProfiler.d0.device)
    exec_points = [e for e in profile if isinstance(e, ProfilePointEvent) and e.name == "exec"]
-    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent)]
+    range_events = [e for e in profile if isinstance(e, ProfileRangeEvent) and not e.is_copy]
    self.assertEqual(len(exec_points), len(range_events), 2)
    self.assertEqual(len(dedup(e.key for e in exec_points)), 1)
    self.assertEqual(len(dedup(e.arg['metadata'] for e in exec_points)), 1)
@@ -72,7 +72,7 @@ class TestQuantizeOnnxCPU(unittest.TestCase):
    out_file = get_quantized_model(sz)
    run_onnx = OnnxRunner(out_file)
    inp = Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32))
-    with Context(DONT_REALIZE_EXPAND=1, QUANTIZE=1):
+    with Context(QUANTIZE=1):
      sched = run_onnx({"input":inp})["output"].schedule()
      ei = lower_schedule_item(sched[-2])
      daccs = [u for u in ei.prg.p.uops if u.op is Ops.DEFINE_REG]
@@ -86,8 +86,7 @@ class TestQuantizeOnnx(unittest.TestCase):
    # divide is ~1500-2000 without reduce_range, 750-900 with it
    out_file = get_quantized_model(sz)
    run_onnx_jit, _ = load_onnx_model(out_file)
-    with Context(DONT_REALIZE_EXPAND=1):
-      run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))
+    run_onnx_jit(input=Tensor(np.random.uniform(size=(sz, sz)).astype(np.float32)))

  def test_prequant_conv2d_1x1(self):
    X = Tensor(np.random.uniform(0, 255, size=(1, 32, 128, 128)).astype(np.uint8))
@@ -109,11 +108,10 @@ class TestQuantizeOnnx(unittest.TestCase):
    N = 512
    X = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(xi))
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(wi))
-    with Context(DONT_REALIZE_EXPAND=1):
-      # this divide is interesting and forces the accumulator to actually be an int
-      out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
-      sexec(out, opts)
+    # this divide is interesting and forces the accumulator to actually be an int
+    out = (X.cast("int").matmul(W.cast("int"))//1000).cast("int8")
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)]
+    sexec(out, opts)

  def test_prequant_gemm_handcode(self):
    src = """typedef int int128 __attribute__((aligned(512),vector_size(512)));
@@ -203,14 +201,12 @@ class TestQuantizeOnnx(unittest.TestCase):
  def test_prequant_gemm_intacc(self, xi=np.uint8, wi=np.uint8, replace_src=None, N=512, clip=True, opts=None):
    X = Tensor(m1:=(np.random.uniform(0, 255, size=(N,N)).astype(xi))).realize()
    W = Tensor(m2:=(np.random.uniform(0, 255, size=(N,N)).astype(wi))).realize()
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    tg_dtype = dtypes.int8 if xi == np.int8 else dtypes.uint8
-    with Context(DONT_REALIZE_EXPAND=1):
-      out = (X.int().matmul(W.int())//1000)
-      if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
-      out = out.cast(tg_dtype)
-      opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
-      sexec(out, opts, replace_src, run_count=1)
+    out = (X.int().matmul(W.int())//1000)
+    if clip: out = out.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
+    out = out.cast(tg_dtype)
+    opts = [Opt(op=OptOps.UPCAST, axis=1, arg=128), Opt(op=OptOps.UNROLL, axis=0, arg=4)] if opts is None else opts
+    sexec(out, opts, replace_src, run_count=1)
    tout = out.numpy()
    mout = ((m1.astype(np.int32) @ m2.astype(np.int32)) // 1000)
    if clip: mout = mout.clip(dtypes.min(tg_dtype),dtypes.max(tg_dtype))
@@ -225,7 +221,6 @@ class TestQuantizeOnnx(unittest.TestCase):

  def test_prequant_gemv(self):
    N = 2048
-    # ugh, it's so broken with those casts. need DONT_REALIZE_EXPAND=1 python3 test/test_quantize_onnx.py TestQuantizeOnnx.test_prequant
    X = Tensor(np.random.uniform(0, 255, size=(1,N)).astype(np.uint8)).realize()
    W = Tensor(np.random.uniform(0, 255, size=(N,N)).astype(np.uint8)).realize()
    #out = X.cast(dtypes.int) @ W.cast(dtypes.int)
@@ -6,6 +6,7 @@ from tinygrad.helpers import getenv, CI, OSX
 from tinygrad.device import is_dtype_supported
 from tinygrad.engine.realize import lower_schedule, CompiledRunner
 from tinygrad.renderer.ptx import PTXRenderer
+from tinygrad.renderer.nir import NIRRenderer
 from test.helpers import not_support_multi_device

 import numpy as np
@@ -100,7 +101,7 @@ class TestRandomness(unittest.TestCase):

    np.testing.assert_allclose(jr, r)

-  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "fails with PTX")
+  @unittest.skipIf(isinstance(Device[Device.DEFAULT].renderer, (NIRRenderer, PTXRenderer)), "PTX and NIR use pointer arithmetic")
  def test_threefry_doesnt_use_long(self):
    for (_,ei) in lower_schedule(Tensor.rand(20).schedule()):
      if isinstance(ei.prg, CompiledRunner):
@@ -1,6 +1,6 @@
 import unittest
 from tinygrad import Tensor, nn
-from tinygrad.helpers import Context, GlobalCounters
+from tinygrad.helpers import Context, GlobalCounters, CI
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops

 class TestRangeifyAssign(unittest.TestCase):
@@ -17,8 +17,22 @@ class TestRangeifyAssign(unittest.TestCase):
    self.assertListEqual(lst, lst3)
    self.assertListEqual(lst2, B.permute(1, 0).tolist())

+class TestRangeifyEdgeCase(unittest.TestCase):
+  def test_matmul_relu_cat(self):
+    a = Tensor.ones(100, 512).contiguous().realize()
+    c = Tensor.ones(1, 512).contiguous().realize()
+    cm = Tensor.ones(512, 512)
+    c = c @ cm
+    c = c.relu()
+
+    res = Tensor.cat(a, c, dim=0)
+    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
+
+# *** non CI rangeify tests below this line ***
+
 N = 256

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeifyOpt(unittest.TestCase):
  def test_randperm(self):
    Tensor.randperm(10000).realize()
@@ -54,6 +68,7 @@ class TestRangeifyOpt(unittest.TestCase):
    A = Tensor.empty(8,8,8,8).permute(1,0,3,2).flatten()
    A.sum().realize()

+@unittest.skipIf(CI, "useless in CI, doesn't test anything")
 class TestRangeify(unittest.TestCase):
  def test_groupnorm(self):
    # ranges 1 and 3 are merging
@@ -280,16 +295,5 @@ class TestRangeifyPM(unittest.TestCase):
    b = self.base.pad(((0,1),(0,0))).pad(((0,0),(0,1)))
    self.assert_same(a, b)

-class TestRangeifyEdgeCase(unittest.TestCase):
-  def test_matmul_relu_cat(self):
-    a = Tensor.ones(100, 512).contiguous().realize()
-    c = Tensor.ones(1, 512).contiguous().realize()
-    cm = Tensor.ones(512, 512)
-    c = c @ cm
-    c = c.relu()
-
-    res = Tensor.cat(a, c, dim=0)
-    self.assertEqual(res.numpy()[-1, :16].tolist(), [512] * 16)
-
 if __name__ == '__main__':
  unittest.main()
@@ -2,16 +2,14 @@
 # schedule confirms the right things are capable of fusing
 # NOTE: this has overlap with external_test_opt.py

-import unittest
+import unittest, functools
 import numpy as np
-import functools
 from typing import cast
 from hypothesis import assume, given, settings, strategies as strat

 from tinygrad import nn, dtypes, Device, Tensor, Variable
 from tinygrad.device import is_dtype_supported
 from tinygrad.dtype import DType, ImageDType
-from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.uop.ops import UOp, Ops, GroupOp, UPat
 from tinygrad.helpers import CI, DEBUG, SPLIT_REDUCEOP, GlobalCounters, Context, getenv, all_same, temp
 from tinygrad.schedule.rangeify import get_rangeify_map, Kernel
@@ -32,7 +30,6 @@ def check_schedule(t:Tensor|list[Tensor]|UOp, allowed:int, to_prerealize:list[Te
  # test lowering all the ScheduleItems to ExecItems
  kernel_cnt = len([si for si,ei in lower_schedule(sched.copy()) if isinstance(ei.prg, CompiledRunner) or not filter_sink])
  if kernel_cnt != allowed:
-    return sched # allow different kernel count, TODO: fix the asserts
    print(f"SCHEDULE ISSUE, expecting {allowed} got {len(sched)}")
    if DEBUG >= 3:
      for i,s in enumerate(sched):
@@ -118,8 +115,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    with self.assertRaisesRegex(RuntimeError, "all buffers must be on the same device"): check_schedule(c, 2)

-  @unittest.skipUnless(is_dtype_supported(dtypes.half) and getenv("CAST_AFTER_EXPAND"), "need half and CAST_AFTER_EXPAND=1")
-  @unittest.skip("CAST_AFTER_EXPAND is not supported")
+  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
  def test_expand_buffer_before_cast(self):
    a = Tensor.randn(4, 2, 1).realize().permute((1, 0, 2))
    b = a.cast(dtypes.half).expand((2, 4, 4))+2
@@ -129,7 +125,7 @@ class TestSchedule(unittest.TestCase):
  def test_indexing_scalars_simple(self):
    X = Tensor.randn(2, 2).realize()
    xt = X[Tensor(1)][Tensor(0)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[1][0])

  @unittest.skipIf(CI and Device.DEFAULT == "NV", "crashes on NV CI")
@@ -149,31 +145,30 @@ class TestSchedule(unittest.TestCase):
    assume(a<x and b<y)
    X = Tensor.randn(x, y).realize()
    xt = X[Tensor(a)][Tensor(b)]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[a][b])

  def test_push_pads_elementwise(self):
    x = Tensor.full((4,4), 2.).contiguous().realize()
    y = Tensor.full((4,4), 4.).contiguous().realize()
    z = (x.reciprocal()*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2))
+    run_schedule(check_schedule(z, 1))
    self.assertEqual(z.item(), 32)

  def test_push_pads_contiguous(self):
    x = Tensor.full((4,1), 2.).contiguous()
    y = Tensor.full((4,4), 4.).contiguous()
    z = (x.reciprocal().expand(4,4)*y).pad((None, (0,1),)).sum()
-    run_schedule(check_schedule(z, 2, [x,y]))
+    run_schedule(check_schedule(z, 1, [x,y]))
    self.assertEqual(z.item(), 32)

  def test_rand(self):
    x = Tensor.rand(32)
-    check_schedule(x, 4, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_rand_recompute_arange(self):
    x = Tensor.rand(32)
-    with Context(DONT_GROUP_REDUCES=1):
-      check_schedule(x, 3, [Tensor._device_rng_counters[x.device]])
+    check_schedule(x, 1, [Tensor._device_rng_counters[x.device]])

  def test_empty_is_not_realized(self):
    a = Tensor.empty(10)
@@ -190,10 +185,7 @@ class TestSchedule(unittest.TestCase):

  def test_simplify_padded_const(self):
    a = Tensor.empty(1022).cummax(axis=0)
-    check_schedule(a, 5)
-    # TODO: what is this testing?
-    #ast = sched[0].ast
-    #self.assertLessEqual(len([u for u in ast.toposort() if u.op is Ops.WHERE]), 6)
+    check_schedule(a, 3)

  def test_basic_binop_fusion(self):
    a = Tensor.empty(10)
@@ -266,18 +258,17 @@ class TestSchedule(unittest.TestCase):
    c = a.sum(axis=0) + b
    check_schedule(c, 1)

-  # not pushing permutes through reduces
  def test_reduce_permute_binop_fusion(self):
    a = Tensor.empty(10,10,10)
    b = Tensor.empty(10,10,1)
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    check_schedule(c, 2)
+    check_schedule(c, 1)

  def test_allow_push_permutes(self):
    a = Tensor.randn(10,10,10).realize()
    b = Tensor.randn(10,10,1).realize()
    c = a.sum(axis=0, keepdim=True).permute(2,1,0) + b
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(c, 1))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_allclose(c.numpy(), np.sum(a.numpy(), axis=0, keepdims=True).transpose(2,1,0)+b.numpy())

  def test_binop_early_reshape_reduce_fusion(self):
@@ -342,7 +333,7 @@ class TestSchedule(unittest.TestCase):
    r1 = (x - r0).sum(axis=0).div(2)
    out0 = r0 + y
    out1 = r1 + y
-    schedule = check_schedule([out0, out1], 2)
+    schedule = check_schedule([out0, out1], 4)
    reduceops = [x for si in schedule for x in si.ast.toposort() if x.op in {Ops.REDUCE_AXIS, Ops.REDUCE}]
    self.assertEqual(len(reduceops), 2) # why is RANGEIFY different?

@@ -375,7 +366,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4,), 2.).contiguous()
    first = a.assign(b)
    second = a.assign(b)
-    check_schedule([first, second], 1)
+    check_schedule([first, second], 2) # TODO: 1?

  # NOTE: this is causing "LAZYCACHE=1 incorrectly reuses contiguous const" #4562
  # should contiguous dedup?
@@ -455,7 +446,7 @@ class TestSchedule(unittest.TestCase):
  @unittest.skipUnless(is_dtype_supported(dtypes.ulong), "Needs ulong")
  def test_fold_conv_batchnorm_optim(self):
    # this is too high
-    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 11)]:
+    for optim, cnt in [(nn.optim.Adam, 30), (nn.optim.SGD, 13)]:
      with self.subTest(optim=optim.__name__):
        with Tensor.train():
          img = Tensor.ones(1,3,4,4)
@@ -476,7 +467,7 @@ class TestSchedule(unittest.TestCase):
      fw = bn(x).contiguous_backward().relu().contiguous()
      fw.sum().backward()
      # TODO: this is too many
-      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 10)
+      check_schedule([x.grad, bn.weight.grad, bn.bias.grad, fw], 9)

  def test_fold_conv_relu(self):
    c1 = nn.Conv2d(3,16,3)
@@ -519,9 +510,8 @@ class TestSchedule(unittest.TestCase):
    img = Tensor.empty(64,64)
    x = (img.sum(0) + img.sum(1))
    out = x.relu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

-  #@unittest.skip("failing in old lazy")
  def test_push_permute_through_reshape(self):
    a = Tensor.empty(16,16)
    b = Tensor.empty(16,16)
@@ -555,7 +545,7 @@ class TestSchedule(unittest.TestCase):
    c = a+b
    d = a.reshape(10,1)+b.reshape(10,1)
    out = c.sum() + d.sum()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_children_dont_push(self):
    a = Tensor.empty(10, 10, 1)
@@ -563,7 +553,7 @@ class TestSchedule(unittest.TestCase):
    d = (a+b).expand(10, 10, 10)
    e = (a+b).permute(2,1,0)
    f = d+e
-    check_schedule(f, 2)
+    check_schedule(f, 1)

  # failing in new lazy
  @unittest.skip("always fusing elementwise")
@@ -602,13 +592,13 @@ class TestSchedule(unittest.TestCase):
    e = c[0] * d
    check_schedule(e, 1)

-  def test_expand_nofuse(self):
+  def test_expand_fuse(self):
    a = Tensor.empty(1, 16)
    b = Tensor.empty(1, 16)
    c = a * b
    d = Tensor.empty(8192, 16)
    e = c * d
-    check_schedule(e, 2)
+    check_schedule(e, 1)

  # this is the failing case in openpilot...it's very simple like this
  def test_image_conv_fusion(self):
@@ -626,7 +616,7 @@ class TestSchedule(unittest.TestCase):

    # NOOP, 3 convs, contiguous
    #check_schedule(x, 5)
-    check_schedule(x, 8)
+    check_schedule(x, 7)

  def test_image_conv_fusion_minimal(self):
    b1 = Tensor.empty(16)
@@ -809,13 +799,13 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.empty(32, 32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(axis=2).T+y
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_two_elus_sum(self):
    x = Tensor.empty(32, 32)
    y = Tensor.empty(32, 32)
    out = x.sum(1).relu().elu() + y.sum(1).relu().elu()
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  @unittest.skipUnless(SPLIT_REDUCEOP, "Testing split reducop requires SPLIT_REDUCEOP")
  def test_preserve_multistage_reduce(self):
@@ -828,7 +818,7 @@ class TestSchedule(unittest.TestCase):
  def test_multistage_reduce(self):
    x = Tensor.empty(32, 32, 32)
    out = x.sum(2).relu().sum(1)
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multistage_reduce_fork(self):
    x = Tensor.empty(32, 32, 32)
@@ -844,7 +834,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_contig(self):
@@ -853,7 +843,7 @@ class TestSchedule(unittest.TestCase):
    z = y.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.ones((64,64)))

  def test_example_matmul_same(self):
@@ -861,7 +851,7 @@ class TestSchedule(unittest.TestCase):
    z = x.matmul(x).sum()
    z.backward()
    out = x.grad.contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    # NOTE: the gradient flows twice
    np.testing.assert_allclose(out.numpy(), 2*np.ones((64,64)))

@@ -884,8 +874,7 @@ class TestSchedule(unittest.TestCase):
    x = x.sum(1)
    x = x[:16]
    out = x + y
-    # NOTE: this could be 1 kernel if we mask the store?
-    check_schedule(out, 2)
+    check_schedule(out, 1)

  def test_multireduce_shrink(self):
    Tensor.manual_seed(0)
@@ -897,8 +886,7 @@ class TestSchedule(unittest.TestCase):
    b_out = b.sum(1)
    b_out = b_out[:16]
    out = a_out + b_out + c
-    # run_schedule(check_schedule(out, 2))  # TODO: this should be 1 (can we make it 1 with the new linearizer?)
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy().sum(axis=1)[:16] + b.numpy().sum(axis=1)[:16] + c.numpy(), atol=1e-4, rtol=1e-4)

  # broken due to const folding and two contiguous are different kernels
@@ -915,7 +903,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + 2
    out1 = a.sum() + 4
    out2 = out0 * out1
-    run_schedule(check_schedule([out0, out1, out2], 1))
+    run_schedule(check_schedule([out0, out1, out2], 3)) # TODO: 1?
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), out0_np*out1_np, atol=1e-4, rtol=1e-6)
@@ -926,7 +914,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum().exp2()
    # out1 has two paths to a.sum()
    out1 = a.sum() + out0
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
    np.testing.assert_allclose(out0.numpy(), out0_np:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+out0_np, atol=1e-4, rtol=1e-6)

@@ -939,7 +927,7 @@ class TestSchedule(unittest.TestCase):
    out2 = b.sum().exp2()
    out3 = b.sum() + out2
    # run_schedule(check_schedule([out0, out1, out2, out3], 1))
-    run_schedule(check_schedule([out0, out1, out2, out3], 6))
+    run_schedule(check_schedule([out0, out1, out2, out3], 4))
    np.testing.assert_allclose(out0.numpy(), np_out0:=np.exp2(a.numpy().sum()), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), np_out1:=a.numpy().sum()+np_out0, atol=1e-4, rtol=1e-4)
    np_b = (a.numpy() + np_out0 + np_out1)
@@ -954,7 +942,7 @@ class TestSchedule(unittest.TestCase):
    out0 = a.sum() + b.sum() + 2
    out1 = a.sum() + b.sum() + 4
    # run_schedule(check_schedule([out0, out1], 1))
-    run_schedule(check_schedule([out0, out1], 4))
+    run_schedule(check_schedule([out0, out1], 2))
    np.testing.assert_allclose(out0.numpy(), a.numpy().sum()+b.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out1.numpy(), a.numpy().sum()+b.numpy().sum()+4, atol=1e-4, rtol=1e-4)

@@ -981,7 +969,7 @@ class TestSchedule(unittest.TestCase):
    out1 = b.max() + out0*2
    out2 = a.sum() + out1
    # run_schedule(check_schedule([out0, out1, out2], 1))
-    run_schedule(check_schedule([out0, out1, out2], 4))
+    run_schedule(check_schedule([out0, out1, out2], 3))
    np.testing.assert_allclose(out0.numpy(), out0_np:=a.numpy().sum()+4, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out1.numpy(), out1_np:=b.numpy().max() + out0_np*2, atol=1e-4, rtol=1e-6)
    np.testing.assert_allclose(out2.numpy(), a.numpy().sum() + out1_np, atol=1e-4, rtol=1e-6)
@@ -1018,7 +1006,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.empty(10,)
    c = a.sum() + b[0]
    d = a.sum() + 2
-    check_schedule([c, d], 1)
+    check_schedule([c, d], 2) # TODO: 1?

  def test_reduce_multiple_paths_midshrink(self):
    a = Tensor.empty(4, 4)
@@ -1047,7 +1035,7 @@ class TestSchedule(unittest.TestCase):
    k = Tensor.randn(32,8,16,8).realize()
    v = Tensor.randn(32,8,16,8).realize()
    out = Tensor.scaled_dot_product_attention(q,k,v)
-    run_schedule(check_schedule(out, 5))
+    run_schedule(check_schedule(out, 4))
    if getenv("CHECK", 1):
      import torch
      compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1055,7 +1043,7 @@ class TestSchedule(unittest.TestCase):

    with Context(FUSE_ATTENTION=1):
      out = Tensor.scaled_dot_product_attention(q,k,v)
-      run_schedule(check_schedule(out, 1))
+      run_schedule(check_schedule(out, 4)) # TODO: should be 1?
      if getenv("CHECK", 1):
        import torch
        compare = torch.nn.functional.scaled_dot_product_attention(torch.tensor(q.numpy()),torch.tensor(k.numpy()),torch.tensor(v.numpy()))
@@ -1068,7 +1056,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.randn(4, 32).realize()
    out = (c * a.sum(-1, keepdim=True)).sum(-1) + (b * a.sum(-1, keepdim=True)).sum(-1) # a.sum has >1 children but should still fuse
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), \
      (c.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1) + (b.numpy()*a.numpy().sum(axis=-1,keepdims=True)).sum(-1), atol=1e-4, rtol=1e-4)

@@ -1113,8 +1101,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = y.sum(axis=-1) + x.sum(axis=-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), y.numpy().sum(axis=-1) + x.numpy().sum(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential(self):
@@ -1131,7 +1118,7 @@ class TestSchedule(unittest.TestCase):
    y = Tensor.randn(4, 32).realize()
    out = x.std(-1) + y.std(-1)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 3))
    np.testing.assert_allclose(out.numpy(), x.numpy().std(axis=-1, ddof=1) + y.numpy().std(axis=-1, ddof=1), atol=1e-4, rtol=1e-4)

  def test_multireduce_diffops_sequential(self):
@@ -1147,8 +1134,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(4, 32).realize()
    y = Tensor.randn(4, 32).realize()
    out = x.sum(-1) + y.max(-1)
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), x.numpy().sum(axis=-1) + y.numpy().max(axis=-1), atol=1e-4, rtol=1e-4)

  def test_multireduce_fusion_sequential_and_parallel(self):
@@ -1160,7 +1146,7 @@ class TestSchedule(unittest.TestCase):
    np_mu = (x.numpy() - x.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True) + \
      (y.numpy() - y.numpy().max(axis=-1, keepdims=True)).mean(axis=-1, keepdims=True)
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 6))
+    run_schedule(check_schedule(out, 5))
    np.testing.assert_allclose(out[0].numpy(), np.sqrt(np.square(x.numpy() - np_mu).sum(-1)/x.shape[-1]), atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(out[1].numpy(), np.sqrt(np.square(y.numpy() - np_mu).sum(-1)/y.shape[-1]), atol=1e-4, rtol=1e-4)

@@ -1169,8 +1155,7 @@ class TestSchedule(unittest.TestCase):
    a,b = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    c,d = Tensor.randn(4, 64).realize(), Tensor.rand(64,8).realize()
    out = a@b + c@d
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), a.numpy()@b.numpy() + c.numpy()@d.numpy(), atol=1e-4, rtol=1e-4)

  def test_softmax_fusion(self):
@@ -1181,17 +1166,15 @@ class TestSchedule(unittest.TestCase):
    expected = (x_exp:=np.exp(x.numpy()-x.numpy().max(-1, keepdims=True)))/x_exp.sum(-1, keepdims=True)
    np.testing.assert_allclose(out.numpy(), expected, atol=1e-4, rtol=1e-4)

-  # TODO: rangeify stores the output in float32
  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
  def test_softmax_upcast(self):
    # input half, softmax in float
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 12, 64, 64, dtype=dtypes.half).realize()
    out = x.softmax(dtype=dtypes.float)
    sched = out.schedule()
-    self.assertEqual(len(sched), 2)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
+    self.assertEqual(len(sched), 3)
+    self.assertEqual(sched[0].bufs[0].dtype, dtypes.float)

    # input float, softmax in float
    Tensor.manual_seed(0)
@@ -1223,12 +1206,12 @@ class TestSchedule(unittest.TestCase):
  def test_scaled_dot_product_attention_fusion(self):
    x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4))
    out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_scaled_dot_product_attention_causal_fusion(self):
    x, y, z = (Tensor.empty(32, 8, 16, 16) for _ in range(3))
    out = Tensor.scaled_dot_product_attention(x, y, z, is_causal=True)
-    check_schedule(out, 5)
+    check_schedule(out, 4)

  def test_adam_step_fusion(self):
    with Tensor.train():
@@ -1258,7 +1241,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.Adam(nn.state.get_parameters([c1, c2]), lr=1e-4)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 20)
+      check_schedule(opt.schedule_step(), 18)

  def test_sgd_conv_fuse(self):
    with Tensor.train():
@@ -1268,7 +1251,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters(c1))
      opt.zero_grad()
      c1(img).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 3)
+      check_schedule(opt.schedule_step(), 5) # TODO: 3?

  def test_sgd_2convs_fuse(self):
    with Tensor.train():
@@ -1291,7 +1274,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2]), nesterov=True, momentum=0.9, weight_decay=0.1)
      opt.zero_grad()
      c2(c1(img).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 13)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse(self):
    with Tensor.train():
@@ -1304,7 +1287,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 17)
+      check_schedule(opt.schedule_step(), 15)

  def test_sgd_4convs_fuse_conv_bw(self):
    with Tensor.train():
@@ -1317,50 +1300,7 @@ class TestSchedule(unittest.TestCase):
      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
      opt.zero_grad()
      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-      check_schedule(opt.schedule_step(), 14)
-
-  @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
-  @unittest.expectedFailure
-  def test_prefer_half_buffer(self):
-    x = Tensor.ones(4).contiguous().realize()
-    # y = Tensor.ones(4).contiguous().realize()
-    z = Tensor.ones(4, 4).contiguous().realize()
-
-    # should not create extra kernel if output will be realized anyways
-    dummy = x.sum().half().float()
-    check_schedule(dummy, 1)
-    dummy = x.sum().half().float().contiguous() + 1
-    check_schedule(dummy, 2)
-
-    # shared between two outputs
-    shared = x.sum().half().float()
-    a = shared * 2
-    b = shared * 3
-    sched = check_schedule([a, b], 3)
-    # store reduceop in half
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    # fuse cast with the child kernel
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-    self.assertEqual(sched[2].bufs[0].dtype, dtypes.float)
-
-    # reduce
-    a = z.sum(axis=0).half().float().sum(axis=0)
-    sched = check_schedule(a, 2)
-    self.assertEqual(sched[0].bufs[0].dtype, dtypes.half)
-    self.assertEqual(sched[1].bufs[0].dtype, dtypes.float)
-
-    # expand
-    # expand will realize just after the .float(), so requires change to realize-before-expand
-    # normal = (x.sum().half().float().reshape(1) * y).sum()
-    # sched = check_schedule(normal, 2)
-    # for si in sched[:-1]: assert all(out.dtype == dtypes.half for out in si.outputs[:-1])
-
-    # parallel reduce
-    # a = x.sum().half().float() * y.sum().half().float()
-    # b = a + 1
-    # c = a + 2
-    # sched = check_schedule([b, c], 4)
-    # doesn't store either in half because it doesn't chase
+      check_schedule(opt.schedule_step(), 15)

  def test_reduce_simple_chase(self):
    a = Tensor.empty(4, 4, 4)
@@ -1409,7 +1349,7 @@ class TestSchedule(unittest.TestCase):
    c = Tensor.empty(16, )
    r = a.sum(1) + c
    d = r[:4] * b
-    check_schedule(d, 2)
+    check_schedule(d, 1)

  def test_multireduce_push_shrink_chase(self):
    Tensor.manual_seed(0)
@@ -1419,22 +1359,20 @@ class TestSchedule(unittest.TestCase):
    d = Tensor.randn(16, 16).realize()
    r = a.sum(1) + c
    out = r[:4] * b + d.sum(1)[:4]
-    # schedule = check_schedule(out, 2)
-    schedule = check_schedule(out, 3)
+    schedule = check_schedule(out, 1)
    run_schedule(schedule)
    np.testing.assert_allclose(out.numpy(), (a.numpy().sum(1) + c.numpy())[:4] * b.numpy() + d.numpy().sum(1)[:4], atol=1e-4, rtol=1e-4)

  def test_midreduce_nochase(self):
    a = Tensor.empty(16, 16)
    b = (a.sum(0) + a.max(1)) + 2
-    check_schedule(b, 2)
+    check_schedule(b, 1)

  def test_multireduce_midreduce_nochase(self):
    Tensor.manual_seed(0)
    a = Tensor.randn(16, 16).realize()
    b = (a.sum(0)+a.max(0) + a.max(1)+a.sum(1)) + 2
-    # schedule = check_schedule(b, 2)
-    schedule = check_schedule(b, 4)
+    schedule = check_schedule(b, 1)
    run_schedule(schedule)
    np.testing.assert_allclose(b.numpy(), a.numpy().sum(0)+a.numpy().max(0) + a.numpy().max(1)+a.numpy().sum(1)+2, atol=1e-4, rtol=1e-4)

@@ -1446,7 +1384,7 @@ class TestSchedule(unittest.TestCase):
    c = a.sum() + 2
    d = (a.sum() - b.sum()) * 4
    # run_schedule(check_schedule([c, d], 1))
-    run_schedule(check_schedule([c, d], 3))
+    run_schedule(check_schedule([c, d], 2))
    np.testing.assert_allclose(c.numpy(), a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), (a.numpy().sum() - b.numpy().sum()) * 4, atol=1e-4, rtol=1e-4)

@@ -1472,7 +1410,7 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = b.sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 2))
+    run_schedule(check_schedule([c, d, e, f], 4))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1487,7 +1425,7 @@ class TestSchedule(unittest.TestCase):
    e = c * d
    f = (b - d).sum() - e
    # run_schedule(check_schedule([c, d, e, f], 1))
-    run_schedule(check_schedule([c, d, e, f], 5))
+    run_schedule(check_schedule([c, d, e, f], 4))
    np.testing.assert_allclose(c.numpy(), c_np:=a.numpy().sum()+2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(d.numpy(), d_np:=a.numpy().sum()*2, atol=1e-4, rtol=1e-4)
    np.testing.assert_allclose(e.numpy(), e_np:=c_np*d_np, atol=1e-4, rtol=1e-4)
@@ -1506,8 +1444,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.randn(3, 4, 5).realize()
    b = Tensor.randn(3, 4, 5).realize()
    out = (a.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum(keepdim=True)+b.pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()).contiguous()
-    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(a.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(keepdims=True) + \
                                                   np.pad(b.numpy(), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-4, rtol=1e-4)

@@ -1515,7 +1452,7 @@ class TestSchedule(unittest.TestCase):
    Tensor.manual_seed(0)
    a = Tensor.rand(3, 4, 5).realize()
    out = a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=1e-5, rtol=1e-6)

  def test_multireduce_pad_reduce_unsafe(self):
@@ -1524,9 +1461,9 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.randn(3, 4, 5).abs().realize()
    out = (a.log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum()+b).abs().log2().pad(((0, 1), (0, 1), (0, 1)), value=1.0).sum().contiguous()
    # run_schedule(check_schedule(out, 1))
-    run_schedule(check_schedule(out, 4))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), np.pad(np.log2(np.abs(np.pad(np.log2(a.numpy()), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum() + \
-                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-6)
+                                                   b.numpy())), ((0, 1), (0, 1), (0, 1)), constant_values=1.0).sum(), atol=3e-4, rtol=1e-5)

  def test_shrink_pad_safe(self):
    a = Tensor.ones((3, )).contiguous().realize()
@@ -1538,7 +1475,7 @@ class TestSchedule(unittest.TestCase):
  def test_shrink_pad_unsafe(self):
    a = Tensor.ones((3, )).contiguous().realize()
    out = a.exp2().shrink(((0, 1),)).pad(((0, 1),)).contiguous()
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    np.testing.assert_equal(out.numpy(), [2, 0])

  def test_base_change_shrink_pad(self):
@@ -1546,7 +1483,7 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:-1, :-1]
    d = c.pad(((0, 1), (0, 1))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:-1, :-1], ((0, 1), (0, 1)))*2)

  def test_base_change_expand_pad(self):
@@ -1554,14 +1491,14 @@ class TestSchedule(unittest.TestCase):
    b = a.exp2()
    c = b[:, None, :]
    d = c.pad(((0, 0), (1, 1), (0, 0))) * 2
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    np.testing.assert_equal(d.numpy(), np.pad(np.exp2(a.numpy())[:, None, :], ((0, 0), (1, 1), (0, 0)))*2)

  def test_fuse_arange_pad_replicate_mode(self):
    x = Tensor.empty(3,3,3,3, requires_grad=True)
    y = x.pad((-1,2,2,-1), mode="replicate")
    dx = y.sum().gradient(x)[0]
-    sched = check_schedule(dx, 3)
+    sched = check_schedule(dx, 1)
    run_schedule(sched)
    np.testing.assert_allclose(dx.numpy(), [[[[0.,3.,9.],[0,1.,3.],[0.,0.,0.]]]*3]*3)

@@ -1571,7 +1508,7 @@ class TestSchedule(unittest.TestCase):
    a = Tensor.ones(4, 4).contiguous().realize()
    b = a.cast(dtypes.half).expand(2, 4, 4)
    c = b.cast(dtypes.int).expand(2, 2, 4, 4)
-    run_schedule(check_schedule(c, 2))
+    run_schedule(check_schedule(c, 1))
    np.testing.assert_equal(c.numpy(), np.ones(((2, 2, 4, 4)), dtype=np.int32))

  def test_base_change_pad_expand(self):
@@ -1579,7 +1516,7 @@ class TestSchedule(unittest.TestCase):
    b = Tensor.full((4, 4), 2.).contiguous().realize()
    c = (a + b).pad(((1, 1), (1, 1)))
    d = c.cast(dtypes.int).expand((2, 6, 6)) * 4
-    run_schedule(check_schedule(d, 2))
+    run_schedule(check_schedule(d, 1))
    c_np = np.pad((np.full((4, 4), 2., dtype=np.float32) + np.full((4, 4), 1., dtype=np.float32)), ((1, 1), (1, 1)), constant_values=0.0)
    np.testing.assert_equal(d.numpy(), np.broadcast_to(c_np.astype(np.half), (2, *c_np.shape)) * 4)

@@ -1678,7 +1615,7 @@ class TestSchedule(unittest.TestCase):
    self._test_fusion([(4, 4), (1, 4)], lambda a,b:a.sum(1).reshape(b.shape)+b, 1)

  def test_late_fusion_post_permute(self):
-    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 2)
+    self._test_fusion([(4, 6, 4), (4, 4, 1)], lambda a,b:a.sum(1, keepdim=True).permute((2, 0, 1))+b, 1)

  def test_late_fusion_double_transpose(self):
    self._test_fusion([(32, 16, 1)],
@@ -1716,6 +1653,7 @@ class TestSchedule(unittest.TestCase):
    self.assertListEqual(realized_const_view.tolist(), [[1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1]])

  @given(strat.sampled_from(dtypes.all), strat.sampled_from(dtypes.all))
+  @unittest.skip("kernel count depends on input")
  def test_cast_padded_const(self, dt1, dt2):
    assume(is_dtype_supported(dt1) and is_dtype_supported(dt2))
    a = Tensor(1, dtype=dt1).reshape(1, 1).pad(((1, 1), None))
@@ -1729,7 +1667,7 @@ class TestSchedule(unittest.TestCase):
    X = Tensor.randn(10, 10).realize()
    idxs = Tensor([0, 2]).realize()
    xt = X[idxs]
-    run_schedule(check_schedule(xt, 2))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), X.numpy()[idxs.numpy()])

  def test_simple_indexing_alt(self):
@@ -1747,7 +1685,7 @@ class TestSchedule(unittest.TestCase):
  def test_advanced_indexing_alt(self):
    X = Tensor.arange(6).reshape(3, 2)+1
    xt = X[[Tensor([2]), Tensor([1])]]
-    run_schedule(check_schedule(xt, 3))
+    run_schedule(check_schedule(xt, 1))
    np.testing.assert_equal(xt.numpy(), 6)

  def test_advanced_simple_indexing_combined(self):
@@ -1795,7 +1733,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.full((2,2), 16)
    y = x.idiv(Tensor.linspace(2, 8, steps=4, dtype=dtypes.int).reshape(2,2)).pad(((1,1), (1,1)))
    out = y.sum(axis=1)
-    run_schedule(check_schedule(out, 2))
+    run_schedule(check_schedule(out, 1))
    self.assertListEqual(out.tolist(), [0, 12, 4, 0])

  def test_arange_transposed_descendants(self):
@@ -1828,7 +1766,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = Tensor.arange(10).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+np.arange(10)[2]).sum(), atol=1e-5, rtol=1e-6)

  def test_arange_index_child(self):
@@ -1844,7 +1782,7 @@ class TestSchedule(unittest.TestCase):
    x = Tensor.randn(5, 2).realize()
    a = (Tensor.arange(10)+1).contiguous()
    out = (x + a[2]).sum()
-    run_schedule(check_schedule(out, 3))
+    run_schedule(check_schedule(out, 2))
    np.testing.assert_allclose(out.numpy(), (x.numpy()+(np.arange(10)+1)[2]).sum(), atol=1e-5, rtol=1e-6)

  @unittest.skip("BUFFER_VIEW no longer supported on non-disk devices")
@@ -1859,10 +1797,10 @@ class TestSchedule(unittest.TestCase):
    from extra.models.llama import precompute_freqs_cis
    args = {"dim":32 if CI else 128, "end":2048 if CI else 8192, "theta":10000}
    fused = precompute_freqs_cis(**args)
-    run_schedule(check_schedule(fused, 3))
+    run_schedule(check_schedule(fused, 1))
    if getenv("CHECK", 1):
      ref = precompute_freqs_cis(**args)
-      run_schedule(check_schedule(ref, 3))
+      run_schedule(check_schedule(ref, 1))
      np.testing.assert_equal(fused.numpy(), ref.numpy())

  def test_fuse_assign_contiguous(self):
@@ -1904,7 +1842,7 @@ class TestSchedule(unittest.TestCase):
    X = Tensor([[0, 2, 3], [1, 2, 3]]).realize()
    Y = Tensor([1, 2]).realize()
    loss = X.sparse_categorical_crossentropy(Y)
-    run_schedule(check_schedule(loss, 4))
+    run_schedule(check_schedule(loss, 3))
    np.testing.assert_allclose(loss.item(), 0.878309, atol=1e-5, rtol=1e-6)

  def test_const_folding_alt(self):
@@ -1925,7 +1863,7 @@ class TestSchedule(unittest.TestCase):
    yt = Tensor.randn(BS, 10).realize()
    with Context(SPLIT_REDUCEOP=0):
      loss = yt.sparse_categorical_crossentropy(Y_train[samples])
-      run_schedule(check_schedule(loss, 6))
+      run_schedule(check_schedule(loss, 5))
      loss_fused = loss.numpy()
    loss_ref = torch.nn.CrossEntropyLoss()(torch.tensor(yt.numpy()), torch.tensor(Y_train.numpy())[torch.tensor(samples.numpy())])
    np.testing.assert_allclose(loss_fused, loss_ref.numpy(), atol=1e-6, rtol=1e-6)
@@ -1935,7 +1873,7 @@ class TestSchedule(unittest.TestCase):
    r = (X+Tensor.arange(16).reshape(4, 4)).sum()
    out0 = r+2
    out1 = r+3
-    run_schedule(check_schedule([out0, out1], 1))
+    run_schedule(check_schedule([out0, out1], 2)) # TODO: 1?
    r_ref = (X.numpy()+np.arange(16).reshape(4, 4)).sum()
    np.testing.assert_allclose(out0.numpy(), r_ref+2, rtol=2e-7)
    np.testing.assert_allclose(out1.numpy(), r_ref+3, rtol=2e-7)
@@ -1977,8 +1915,7 @@ class TestSwizzle(unittest.TestCase):
      a = Tensor.randint(32, 32).realize()
    r = (a+a).sum(1).sum(0)
    # double reduce collapses to a single reduce
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(r, 1))
+    run_schedule(check_schedule(r, 1))
    self.assertEqual(r.numpy(), (a.numpy()+a.numpy()).sum(1).sum(0))

  def test_single_swizzle(self):
@@ -1998,33 +1935,29 @@ class TestSwizzle(unittest.TestCase):
      b = Tensor.randint(4,).realize()
    # parallel reduce!
    add = a.sum(0)+b.sum(0)
-    with Context(DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(add, 1))
+    run_schedule(check_schedule(add, 1))
    self.assertEqual(add.numpy(), a.numpy().sum(0)+b.numpy().sum(0))

-  @unittest.skip("TODO: how do we express the norm")
  def test_softmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(32, 32).realize()
    t = a.softmax()
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1):
-      check_schedule(t, 1)
+    check_schedule(t, 3) # TODO: 1?

  def test_argmax_one_kernel(self):
    Tensor.manual_seed(0)
    with Context(DEBUG=0, TRACK_MATCH_STATS=0):
      a = Tensor.randn(10, 20).realize()
    t = a.argmax(0)
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): t.realize()
+    check_schedule(t, 2) # TODO: 1?

  def test_swizzle_reduceop(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4,4).realize()
    y = Tensor.randn(4,4,4).realize()
    out = x.reshape(4,4,1).expand(4,4,4).sum(axis=(1,))+y
-    with Context(DONT_REALIZE_EXPAND=1, DONT_GROUP_REDUCES=1):
-      run_schedule(check_schedule(out, 1))
+    run_schedule(check_schedule(out, 2)) # TODO: 1?
    np.testing.assert_allclose(out.numpy(), np.tile(x.numpy().reshape(4,4,1), (1,1,4)).sum(axis=1)+y.numpy())

  def test_permute_rewrite(self):
@@ -2032,7 +1965,7 @@ class TestSwizzle(unittest.TestCase):
    y = Tensor.randn(4, 1, 16).realize()
    z = Tensor.randn(4, 4, 1).realize()
    t = (x*y).sum(axis=(0, 2)).reshape(1, 4, 1).permute(0, 2, 1)+z
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 2)) # TODO: 1?
    t_np = (x.numpy()*y.numpy()).sum(axis=(0, 2)).reshape(1, 4, 1).transpose(0, 2, 1)+z.numpy()
    np.testing.assert_allclose(t.numpy(), t_np, atol=1e-6, rtol=1e-3)

@@ -2043,14 +1976,14 @@ class TestSwizzle(unittest.TestCase):
    a_reduce = a.sum(axis=(2,), keepdim=True).sum(axis=(1,))
    b_reduce = b.sum(axis=(0,))
    t = a_reduce+b_reduce
-    with Context(DONT_GROUP_REDUCES=1, DONT_REALIZE_EXPAND=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))

  def test_parallel_reduce_possible(self):
    Tensor.manual_seed(0)
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 2, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  # kernels can only have 1 or n in each dim
@@ -2059,7 +1992,7 @@ class TestSwizzle(unittest.TestCase):
    x = Tensor.randn(4, 2, 2).realize()
    y = Tensor.randn(4, 3, 2).realize()
    t = x.sum(axis=1)+y.sum(axis=1)
-    with Context(DONT_GROUP_REDUCES=1): run_schedule(check_schedule(t, 1))
+    run_schedule(check_schedule(t, 1))
    np.testing.assert_allclose(t.numpy(), x.numpy().sum(axis=1)+y.numpy().sum(axis=1), atol=1e-6, rtol=1e-3)

  def test_unsafe_pad(self):
@@ -2152,7 +2085,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(3).realize()
    zeros = Tensor.zeros(3).realize()
    b = (a*zeros).to("CPU")
-    run_schedule(check_schedule(b, 0, filter_sink=False))
+    run_schedule(check_schedule(b, 2, filter_sink=False)) # TODO: 0?
    self.assertListEqual(b.tolist(), [0, 0, 0])
    self.assertEqual(b.device, "CPU")

@@ -2160,8 +2093,8 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.ones((4,)).to("CPU")
    b = Tensor.empty(4, device="CPU")
    add = a+b
-    add.kernelize()
    assert all_same([x.device for x in add.uop.src]), f"ALU has different devices! {[x.device for x in add.src]}"
+    add.kernelize()

  def test_alu_before_copy(self):
    buf = Tensor.ones(1).contiguous().realize()
@@ -2172,12 +2105,12 @@ class TestCopyFolding(unittest.TestCase):
  def test_copy_to_same_device(self):
    a = Tensor.empty(4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_alt(self):
    a = Tensor.empty(4, 4).uop
    b = a.copy_to_device(a.device)
-    check_schedule(b, 0, filter_sink=False)
+    check_schedule(b, 1, filter_sink=False) # TODO: 0?

  def test_copy_to_same_device_sched(self):
    a = Tensor.ones(4).contiguous().realize().uop.as_buf()
@@ -2192,13 +2125,11 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.empty(4)
    check_schedule(a.clone(), 1, filter_sink=False)

-  # NOTE: moving copy before view might change this
  def test_shrink_copy(self):
    a = Tensor.arange(4)
    view = a.shrink(((0, 2),))
    b = view.clone()
-    # NOTE: this was sort of a bug making this 2
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 2)
    self.assertEqual(b.uop.size, 2)
    self.assertListEqual(b.tolist(), [0, 1])
@@ -2207,7 +2138,7 @@ class TestCopyFolding(unittest.TestCase):
    a = Tensor.arange(2)
    view = a.reshape(2, 1).expand(2, 2)
    b = view.clone()
-    run_schedule(check_schedule(b, 2, filter_sink=False))
+    run_schedule(check_schedule(b, 1, filter_sink=False))
    self.assertEqual(b.uop.base.buffer.size, 4)
    self.assertEqual(b.uop.size, 4)
    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])
@@ -2251,7 +2182,7 @@ class TestBufferUOp(unittest.TestCase):
  def test_buffer_has_buffer(self):
    buf = Tensor.empty(10)
    self.assertIsNotNone(buf.uop.buffer)
-    self.assertEqual(buf.uop.st, ShapeTracker.from_shape((10,)))
+    self.assertEqual(buf.uop.shape, (10,))
    # the device Buffer remains unallocated until it's we run the schedule
    self.assertFalse(buf.uop.buffer.is_allocated())
    add = buf+1
@@ -2330,7 +2261,7 @@ class TestContiguous(unittest.TestCase):
  def test_double_contiguous_realizes_once(self):
    a = Tensor.empty(4, 1)
    b = a.expand((4, 4)).contiguous().contiguous()
-    check_schedule(b, 1)
+    check_schedule(b, 2) # TODO: should be 1?

  def test_view_does_not_realize(self):
    a = Tensor.empty(4)
@@ -2466,10 +2397,6 @@ class TestUOpBecome(unittest.TestCase):
    c = (a.reshape(1, 1, 4, 4)+0).shrink(((0, 1), (0, 1), (0, 3), (0, 3)))+0
    check_schedule([b, c], 0)
    assert all_same([x.uop.base.realized for x in [a,b,c]])
-    # these movement ops result in the same ShapeTracker
-    assert b.uop.st == c.uop.st
-    assert b.uop is c.uop
-    assert UPat(Ops.VIEW, src=(UPat(Ops.BUFFER),)).match(c.uop, {})

  def test_setitem_becomes_subbuffer(self):
    a = Tensor.full((4,), 2.).contiguous().realize()
@@ -52,7 +52,6 @@ class TestSetitem(unittest.TestCase):

  def test_setitem_into_noncontiguous(self):
    t = Tensor.ones(4)
-    self.assertFalse(t.uop.st.contiguous)
    with self.assertRaises(RuntimeError): t[1] = 5

  @unittest.skip("TODO: flaky")
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
geohot	978502be46	experiments with multi being range	2025-10-17 14:11:55 +08:00
chenyuandGitHub	9561803cb0	fix assert in test_schedule (#12745 ) * fix assert in test_schedule updated kernel counts and some old tests * fix	2025-10-16 15:39:50 -04:00
chenyuandGitHub	285534ce64	delete DONT_REALIZE_EXPAND and DONT_GROUP_REDUCES (#12744 ) does nothing now	2025-10-16 14:11:33 -04:00
chenyuandGitHub	98239f1156	few shapetracker cleanups (#12741 )	2025-10-16 12:43:27 -04:00
chenyuandGitHub	53478c741d	relax ASSERT_MIN_STEP_TIME for space lab policy (#12742 )	2025-10-16 11:40:36 -04:00
geohotstanandGitHub	5d209ee7ec	onnx helper intermediate node output validation (#12740 ) * start * update comments * good * add comments and better printing * done	2025-10-16 11:17:47 -04:00
sirhcmandGitHub	bce2bc0465	Revert "use RTLD_GLOBAL on macos" (#12738 ) This reverts commit `89fe3e574d`.	2025-10-16 10:07:21 -04:00
chenyuandGitHub	f34f26bca0	fix gpt2 with benchmark (#12736 ) `CPU=1 python3 examples/gpt2.py --benchmark 128` works now	2025-10-16 09:55:20 -04:00
Sieds LyklesandGitHub	55db1b0e0e	reduce where that is cut from two sides (#12733 ) * better rule * correct pattern * shorten line	2025-10-16 15:25:15 +02:00
nimlgenandGitHub	cf9baeea61	Revert "nv: check if jitlink is avail (#12731 )" (#12735 ) This reverts commit `a069a45d14`.	2025-10-16 20:41:49 +08:00
George HotzandGitHub	8be7844b2e	use apply uop for assign to fix assign metadata (#12732 ) * use apply uop for assign * fix metadata for assign * fix backward metadata * those aren't real tests	2025-10-16 20:34:12 +08:00
nimlgenandGitHub	3aa2277b8f	nv: usb4 (#12696 ) * hackish * prog * match * l * simpler * refactor * not osx * apple things * tiny changes * fix mask * match fix * nn	2025-10-16 20:11:19 +08:00
nimlgenandGitHub	a069a45d14	nv: check if jitlink is avail (#12731 )	2025-10-16 19:58:50 +08:00
George HotzandGitHub	a498ec9c18	cleanup names of postrange + fast FUSE_OPTIM (#12730 ) * cleanup names of postrange * make FUSE_OPTIM not slow * delete junk in def r	2025-10-16 19:38:31 +08:00
Sieds LyklesandGitHub	8f740e07ff	no broadcasting/vectors in reduce collapse (#12729 )	2025-10-16 13:22:57 +02:00
qazalandGitHub	533f18b22c	viz: add trace data for inflight buffers (#12728 ) * viz: add trace data for inflight buffers * add test_inflight_buf * temp stores the keys * update tests / use Tensor.ones	2025-10-16 19:15:03 +08:00
George HotzandGitHub	af4479c169	faster stable diffusion load (#12725 ) * faster stable diffusion load * failing tests	2025-10-16 18:31:59 +08:00
nimlgenandGitHub	e7c057d5dc	system: alloc_sysmem return view (#12724 ) * system: alloc_sysmem return view * e	2025-10-16 17:55:01 +08:00
nimlgenandGitHub	b86a33a312	ptx: support bw (#12722 )	2025-10-16 15:38:08 +08:00
nimlgenandGitHub	b8cd66c7a2	nv: support all gb20x and small bar (#12721 )	2025-10-16 15:37:54 +08:00
George HotzandGitHub	1d1e1d9d88	delete the ShapeTracker (#12720 ) * delete the ShapeTracker * fix tests * fix more * fix gc test	2025-10-16 15:36:22 +08:00
George HotzandGitHub	592e86f6f5	remove UOp.st (#12716 ) * remove UOp.st * fix tests * torch backend disable	2025-10-16 14:44:09 +08:00
wozeparrotandGitHub	cc2dfe22f5	tinyfs: fetch file utility (#12719 )	2025-10-15 23:38:56 -07:00
nimlgenandGitHub	3ed543f956	system: reorder funcs + barrier on macos (#12714 )	2025-10-16 14:38:01 +08:00
qazalandGitHub	b77bdbbc62	viz: count unpickle in server startup time (#12715 ) * viz: count unpickle in server startup time * type checking	2025-10-16 13:07:46 +08:00
George HotzandGitHub	7c19db00f1	remove st from jit/split_reduceop (#12713 ) * remove st from jit * fix by merging reshapes * no st usage in rangeify * hmm, stop early works * fix speed regressions	2025-10-16 12:50:58 +08:00
qazalandGitHub	069177c1be	trace buffer producer and consumers (#12639 ) * trace buffer producer and consumers * work * generic colored util * fix batched * basic clicking works * generic javascript that works for producer and consumers * keep focused shape * idle time * timings for producer and consumers dedup * from sd test * tiny cleanups * timeline * work * up to here * assert * list it * work	2025-10-16 11:11:31 +08:00
George HotzandGitHub	4a151e7533	make xcode signing happy, waiting for entitlement (#12712 )	2025-10-16 10:20:34 +08:00
chenyuandGitHub	c3278e5622	clean up old tests (#12708 )	2025-10-15 17:53:17 -04:00
chenyuandGitHub	b8cf35fb77	print macOS version in CI (#12705 )	2025-10-15 15:05:33 -04:00
Daniel GitHub chenyu	d65bd669f8	update tiny torch backend hook (#12575 ) * update the backend to fix torch deprecation warning * use param_hook to avoid full backward hook needlessly firing on inputs which do not require gradients * fix indentation --------- Co-authored-by: chenyu <[email protected]>	2025-10-15 14:02:33 -04:00
nimlgenandGitHub	db5ae846aa	nv: do not use va_addr for cpu accesses (#12697 ) * nv: do not use va_addr for cpu accesses * mypy	2025-10-15 22:48:12 +08:00
nimlgenandGitHub	3ab23af829	nv: copy prog with copyin (#12701 ) * nv: copy prog with copyin * to bytes * fix test	2025-10-15 22:48:01 +08:00
nimlgenandGitHub	fafbf3daea	memory: reserve ptable (#12702 )	2025-10-15 22:47:50 +08:00
geohot	85a907605c	hotfix: only 20 steps of beautiful_mnist_torch, some CI machines are slow	2025-10-15 22:29:34 +08:00
sirhcmandGitHub	e1996d358c	use RTLD_GLOBAL on macos (#12699 )	2025-10-15 22:24:50 +08:00
chenyuandGitHub	312c622d35	support None in pad_to and shrink_to (#12700 )	2025-10-15 09:25:31 -04:00
George HotzandGitHub	612e3d6143	replace mop arg with vectorized index (#12695 ) * replace mop arg with vectorized index * tests passing * better viz * no compile4	2025-10-15 20:50:06 +08:00
wozeparrotandGitHub	9ec4c06d7d	feat: one request per device (#12698 )	2025-10-15 05:22:07 -07:00
Sieds LyklesandGitHub	99aa3bd5f9	reduce collapse reduce only the cut range (#12687 )	2025-10-15 13:57:41 +02:00
Sieds LyklesandGitHub	91ac4f1f92	late merging of where and load (#12694 )	2025-10-15 13:33:06 +02:00
qazalandGitHub	768dc952de	viz ui cleanups / renaming (#12691 ) * better viz names * delete unused * don't use opacity, it's multiplicative * keep styles * scrollbar coloring * pyrender doesn't work here beautiful_mnist r_64_16_32_36@lower all index dtypes	2025-10-15 18:40:22 +08:00
chenyuandGitHub	2e50ed0767	increase timeout of resnet cron (#12693 ) does not finish in 6 hours now	2025-10-15 06:08:58 -04:00
sirhcm GitHub George Hotz nimlgen	0aabc1e938	Mesa NIR backend (NAK/LLVMpipe) (#12089 ) * nak works * TestOps::test_add works * testop has no crashes * fix bool casts * fix typo * add disassemble * RANGE and locals/regs * simplify NAKCompiler * disass cleanup * cleanup nir codegen * almost all tests passing * cleanup notes in extra/ * old notes * only import nak if NIR=1 * fix new SPECIAL syntax * fix local/shared memory * more tests passing * add DEFINE_VAR support * llvmpipe kinda works * diskcache * some mypy stuff * lvp passing test_ops.py * fix imports * actually fix imports * remove 'stdout' * fix llvm import * fix mypy issues * nicer errors * simpler test_dtype skips * test lvp in CI * fix github action syntax * fix more actions typos * switch to mesa 25.1.0 * diskcache_put * better generation for lvp nir_options * b64encode shader blobs * Revert diskcache changes This reverts commits `930fa3de8a` and `8428c694b3`. * general cleanup * better error messages * fix llvm import * fix windows tests * link with libm and libgcc_s * fix some errors * dont check for 'float4' * NIR uses pointer arithmetic * use tinymesa * bump tinymesa * bump tinymesa again * update lvp nir_options * print nir shader with DEBUG * simplify LVPCompiler * more tests * "gated" STORE * NAK is cacheable * more tests * all tests pass locally for NAK * test autogen in CI * autogen deps * more deps * fix uop_gc * fix macos * mypy * save 2 lines * save two more lines * save 1 line * save 4 lines * save more lines * Revert "save more lines" This reverts commit `dd3a720c5a`. * save more lines * fix LVP on windows * refactor * reorganize some code * refactor lib_gpu * move LVP check * out of order loads * remove support.mesa * bump tinymesa version * simplify LVP jit * macos * macos ci * shell: bash * testing * more testing * compute brew prefix * stupid typo * actually fix * lib * stdout on macos * inline gallivm_compile_module * Revert "inline gallivm_compile_module" This reverts commit `b65983b151`. * elf macos * semicolon * inherit from CPULLVMCompiler * ruff * disas test * fix libm linking * default is fine actually * arm works * add elf loader link test * fix NAK beam * pylint is too smart by half --------- Co-authored-by: George Hotz <[email protected]> Co-authored-by: nimlgen <[email protected]>	2025-10-15 17:38:33 +08:00
qazalandGitHub	f0268d13f6	cleanup viz server (#12688 )	2025-10-15 15:58:36 +08:00
nimlgenandGitHub	aa81bde150	amd: usb4/thunderbolt on macs (#12641 ) * tbgpu * works * cleaner * this * zero size * h * fix * simpler * prio over usb * c * not needed * linter * this way * mappings * mypy * mypy * mypy 2 * nn	2025-10-15 13:02:01 +08:00
George HotzandGitHub	236c4590c3	use margs as intermediate for new style mops (#12686 ) * use marg to prepare for movement op change * clean up forced reshape * move marg * more marg * more	2025-10-15 12:43:00 +08:00
qazalandGitHub	7597e1dcac	pyrender in viz (#12682 ) * pyrender in viz * keep profile still print_tree * keep special in render	2025-10-15 11:53:30 +08:00
qazalandGitHub	60e03eec37	viz: add View Program option (#12683 )	2025-10-15 11:37:51 +08:00
George HotzandGitHub	a59439d013	use UOp.shape property instead of UOp.st (#12664 ) * work on shape property * reshape causing issues * more mops * all mops * need to cache it * _shape is like _device * mostly works * shape is good * const uses _shape * fix tests * size doesn't use st * close * test is broken * one less st * hack for 3 op assign * oops, i didn't mean to change that * support emulate in the NullDevice * reproed failure in emulation * fix wmma	2025-10-15 10:01:34 +08:00
chenyuandGitHub	89df6f611d	reenable sdxl mac benchmark (#12680 ) also updated faster sd step times	2025-10-14 17:36:17 -04:00
chenyuandGitHub	d25ceffe8d	update padto opts tests (#12679 )	2025-10-14 17:00:42 -04:00
chenyuandGitHub	e8380968f2	add venv_sd_mlperf to gitignore (#12676 ) training stable diffusion stuff	2025-10-14 12:51:36 -04:00
wozeparrotandGitHub	f228c03f9f	fetch raid from cloud (#10799 ) * feat: initial tinyfs device * feat: don't allow compute on tinyfs device * feat: tensor helpers to load and store * feat: bufferview for tinyfs * fix: keep copy sizes correct * fix: recv large * clean: unneeded * feat: comment * clean: unneeded * clean: remove * clean: remove * feat: get request tag * feat: rename to cloud * feat: send request_id * feat: start computing tree * feat: compute store tree on this side * feat: jank chunked load * feat: more debugging * feat: rename to just load and store * feat: correct chunk count * fix: fix load for < 1mb * feat: comments * feat: don't truncate on block devices * feat: better way of testing block device * feat: don't need to pad that much * feat: connect to nodes directly on load * feat: cache connections * feat: don't hard code chunk size * feat: close mmap when closing file handle * feat: don't overwrite stuff on disk if storing from disk * clean: debug print * fix: close mmap * feat: await workers * feat: fast copy from tinyfs to disk * feat: don't copy to device on last * feat: use single socket per device * feat: raid in tinyfs * clean: remove import * clean: type * feat: maintain single event loop * feat: lower worker count * feat: use connection pool * feat: fetch mapping in its own process * fix: release lock * feat: don't fetch if exists * feat: req id only on stores * feat: always fetch * fix: rangeify * feat: allow specifying raid root * fix: dealloc buffer * feat: start support non 0 offset * clean: use cleaner * feat: don't pass to threadpool * clean: typing	2025-10-14 07:53:55 -07:00
chenyuandGitHub	70dd297a05	BS=96 for bert (#12675 ) 96 trains fine now	2025-10-14 09:07:43 -04:00
Sieds LyklesandGitHub	852d80dff9	better where on load folding (#12651 ) * move where clauses to load * shorten line * drop clauses if they are duplicated * add rule for swapped where branch * where on ungated load * dont move clause if load is in the clause * parse_valid returns None * no data dependent branches * fix rule * enable swapped rule * remove those	2025-10-14 13:30:47 +02:00
nimlgenandGitHub	c7e63601fd	gfx1200 tc for AMD_LLVM (#12673 )	2025-10-14 19:17:48 +08:00
George HotzandGitHub	db4a359374	fix up some slow tests that launch python (#12672 ) * fix up some slow tests that launch python * svd nonfull in parallel * split test_advancedindex	2025-10-14 19:13:55 +08:00
nimlgenandGitHub	4918c827c2	amd: lib_gpu does not need cpu_access (#12670 )	2025-10-14 18:34:34 +08:00
nimlgenandGitHub	0c9d47deab	hcq: add alignment to kernargs (#12669 )	2025-10-14 18:33:12 +08:00
qazalandGitHub	d3bfcd3277	minor patches for SQTT over usb on gfx12 (#12627 ) * disable cpu_access in the sqtt buffer allocation not sure if this is required, it results in a very slow call to pcie_mem_write over USB GPU, removing it worked fine. * fix itrace_se_mask on gfx12 on gfx11 it gave 6 se, on gfx11 this value is 2 so no instructions were traced. * Revert "fix itrace_se_mask on gfx12" This reverts commit `0644adbcd1`.	2025-10-14 18:07:46 +08:00
Sieds LyklesandGitHub	1e6e5a0efd	`parse_valid` returns None instead of raising (#12663 ) * parse_valid returns None * change there too	2025-10-14 11:57:38 +02:00
qazalandGitHub	471bd30d16	cleanup viz/serve.py (#12665 ) * use load_pickle * update comment	2025-10-14 17:50:39 +08:00
George HotzandGitHub	fb61f3519f	remove assign contiguous hack (#12659 ) * remove assign contiguous hack * remove bad contiguous usage in torch backend * assign	2025-10-14 16:42:14 +08:00
George HotzandGitHub	30ee7c4c26	cleanup Device usage in Tensor (#12662 )	2025-10-14 16:22:22 +08:00
Sieds LyklesandGitHub	e06cbfcb8a	combine `pm_drop_and_clauses` (#12660 ) * combine those * wino kernels decreased	2025-10-14 10:09:41 +02:00
George HotzandGitHub	84d4589ed4	remove pylint from pre-commit and CI (#12658 ) * remove pylint from pre-commit and CI * multidevice test is fast * faster pre-commit * 8 is faster than 4 * better name * how did that typecheck?	2025-10-14 15:39:59 +08:00
qazalandGitHub	8ecaf839e2	cleanup UOp tracing [pr] (#12657 )	2025-10-14 14:50:59 +08:00
George HotzandGitHub	b9eb5b5d49	clean up the LLM tokenizer (#12653 ) * clean up the LLM tokenizer * simple tokenizer is actually simple * ugh write good code	2025-10-14 14:22:01 +08:00
qazalandGitHub	a9ef93176f	viz: add colored text helper (#12654 )	2025-10-14 13:05:26 +08:00
George HotzandGitHub	ecdc7539a2	add typing to MathTraits (#12650 ) * add typing to MathTraits * fix assign	2025-10-14 12:35:20 +08:00
qazalandGitHub	9bf032de69	viz: keep focused shape in view (#12648 )	2025-10-14 10:49:08 +08:00
chenyuandGitHub	77b5e6774e	fix bert training config (#12647 ) FREE_INTERMEDIATE=0 REWRITE_STACK_LIMIT=500000	2025-10-13 15:03:47 -04:00
nimlgenandGitHub	f1041dc0ac	pylint 4.0.0 (#12642 ) * cpu: fix spacing * fix pylint * fix pylint * pylint 4.0.0 * lambda * keep eval for now * im so sorry	2025-10-13 23:28:36 +08:00
wozeparrotandGitHub	47e0c43976	feat: Tensor.{load, store} (#12629 )	2025-10-13 08:04:41 -07:00
chenyuandGitHub	0f776c6e46	examples/mlperf/training_submission_v6.0 (#12644 ) copied from v5.1	2025-10-13 09:58:25 -04:00
Sieds LyklesandGitHub	e0139fafc1	UOp symbolic tests use eval to check against string (#12643 )	2025-10-13 14:19:42 +02:00
b1tg GitHub wozeparrot	218225e8d0	pylint error (#12630 ) Co-authored-by: wozeparrot <[email protected]>	2025-10-13 05:05:12 -07:00
nimlgenandGitHub	9096d7cc2e	amd: support for rx9060 (#12640 )	2025-10-13 19:44:15 +08:00
qazalandGitHub	066d25f5fb	refactor to trace_num property in buffers (#12638 )	2025-10-13 18:06:55 +08:00
qazalandGitHub	cd6aeebfee	sqtt: osx decoder installer (#12637 )	2025-10-13 17:26:12 +08:00
Sieds LyklesandGitHub	e537e895b1	drop unused invalid conditions (#12635 ) * drop where conditions if the ranges are not used inside the index * remove allow_any_len	2025-10-13 10:52:21 +02:00
wozeparrotandGitHub	9ab06dffad	hotfix: block from env (#12628 )	2025-10-12 08:07:32 -07:00
wozeparrotandGitHub	12435a2dab	actual tinyfs device (#12620 )	2025-10-12 07:51:17 -07:00
chenyuandGitHub	8f5f57c7d9	smaller CNT fuzz shapetracker (#12626 )	2025-10-12 08:52:30 -04:00
George HotzandGitHub	1ecf403294	cleanup long lines [pr] (#12623 ) * cleanup long lines * more * a few more * all noqa fixed * fix amd + cuda * clean that up	2025-10-12 20:18:05 +08:00
qazalandGitHub	fd51ecf983	process_replay for get_rangeify_map (#12624 )	2025-10-12 15:14:40 +03:00
qazalandGitHub	b5afa3848e	viz: fix memory graph total nbytes (#12622 ) * viz: fix memory graph total nbytes * post increment * simple regression test * loop with markers + slightly off text baseline * cpu events clear	2025-10-12 14:32:46 +03:00
nimlgenandGitHub	822eab057f	cpu: respect taskset + allow all cores (#12619 ) * cpu: account taskset + allow all cores * spaces	2025-10-12 14:31:40 +08:00
chenyuandGitHub	7ac74d1550	remove unused type ignore [pr] (#12618 )	2025-10-11 21:24:04 -04:00
Sieds LyklesandGitHub	772a8dfe31	reshape uses valid when simplifying (#12597 ) * reshape uses valid when simplifying * try with IGNORE_OOB=0 * is it this test? * skipif gpuocelot	2025-10-11 17:02:54 +02:00
nimlgenandGitHub	08e62454b6	amd: use cpu_view() in sqtt (#12610 )	2025-10-11 18:11:25 +08:00
Sieds LyklesandGitHub	a2ae56674a	`uop_given_valid` try multiple clauses (#12615 ) * uop_given_valid uses less simplify * enable test * try all expressions together * enable test	2025-10-11 11:53:42 +02:00
Sieds LyklesandGitHub	dccdd190aa	uop_given_valid uses less simplify (#12612 ) * uop_given_valid uses less simplify * enable test	2025-10-11 10:57:39 +02:00
qazalandGitHub	9205527db0	viz: draw highlights above shapes (#12613 )	2025-10-11 11:39:13 +03:00
George HotzandGitHub	cab034b863	improve typing (#12611 ) * improve typing and bump to 3.11 * no need for Self yet * improve typing * binop also	2025-10-11 16:20:23 +08:00