notes

* tc3 shape expand [pr]

* remove unused stuff in lowerer

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

@@ -29,6 +29,10 @@ inputs:
     description: "Install CUDA?"
     required: false
     default: 'false'
+  ocelot:
+    description: "Install gpuocelot?"
+    required: false
+    default: 'false'
   webgpu:
     description: "Install webgpu?"
     required: false
@@ -193,14 +197,14 @@ runs:
           sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
 
-    # **** CUDA ****
+    # **** gpuocelot ****
 
     - name: Install gpuocelot dependencies (MacOS)
-      if: inputs.cuda == 'true' && runner.os == 'macOS'
+      if: inputs.ocelot == 'true' && runner.os == 'macOS'
       shell: bash
       run: brew install --quiet cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
     - name: Cache gpuocelot
-      if: inputs.cuda == 'true'
+      if: inputs.ocelot == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
@@ -209,7 +213,7 @@ runs:
         path: ${{ github.workspace }}/gpuocelot/ocelot
         key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
     - name: Clone/compile gpuocelot
-      if: inputs.cuda == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
         git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -220,7 +224,7 @@ runs:
         cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF -DCMAKE_POLICY_VERSION_MINIMUM=3.5
         ninja
     - name: Install gpuocelot
-      if: inputs.cuda == 'true'
+      if: inputs.ocelot == 'true'
       shell: bash
       run: |
         cd ${{ github.workspace }}/gpuocelot/ocelot/build

.github/workflows/benchmark.yml

@@ -67,11 +67,11 @@ jobs:
     - name: Test speed vs torch
       run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
-      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+      run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Test AMX tensor cores
       run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Run Tensor Core GEMM (float)
       run: DEBUG=2 SHOULD_USE_TC=1 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
     - name: Run Tensor Core GEMM (half)
@@ -123,7 +123,7 @@ jobs:
     - name: UsbGPU copy speeds
       run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
     - name: UsbGPU openpilot test
-      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB AMD_LLVM=1 NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      run: sudo -E PYTHONPATH=. AMD=1 AMD_IFACE=USB NOLOCALS=0 IMAGE=0 GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (Mac)
@@ -196,8 +196,8 @@ jobs:
       run: NV=1 python test/external/external_benchmark_multitensor_allreduce.py
     - name: Test tensor cores
       run: |
-        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        NV=1 ALLOW_TF32=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Run Tensor Core GEMM (CUDA)
       run: |
         CUDA=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
@@ -396,8 +396,8 @@ jobs:
       run: AMD=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
     - name: Test tensor cores
       run: |
-        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
-        AMD=1 AMD_LLVM=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_emulation TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 AMD_LLVM=0 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
+        AMD=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
         AMD=1 SHOULD_USE_TC=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py
     - name: Run Tensor Core GEMM (AMD)
       run: AMD=1 SHOULD_USE_TC=1 HALF=1 DEBUG=2 ATOL=2e-2 python3 extra/gemm/simple_matmul.py | tee matmul_amd.txt
@@ -607,12 +607,8 @@ jobs:
       run: test/external/process_replay/reset.py
     - name: validate openpilot 0.9.7
       run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
-    - name: benchmark openpilot 0.9.4
-      run: BENCHMARK_LOG=openpilot_0_9_4 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_4.txt
     - name: benchmark openpilot 0.9.7
       run: BENCHMARK_LOG=openpilot_0_9_7 PYTHONPATH=. QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_0_9_7.txt
-    - name: benchmark openpilot w IMAGE=2 0.9.4
-      run: BENCHMARK_LOG=openpilot_0_9_4_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_4.txt
     - name: benchmark openpilot w IMAGE=2 0.9.7
       run: BENCHMARK_LOG=openpilot_0_9_7_image PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
     - name: openpilot compile3 0.9.7

.github/workflows/benchmark_search.yml

@@ -27,7 +27,7 @@ jobs:
         BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
     - name: Run winograd cifar with new search
       run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 IGNORE_BEAM_CACHE=1 DISABLE_COMPILER_CACHE=1 BS=1024 STEPS=500 python examples/hlb_cifar10.py
     - name: Run winograd cifar with cached search
       run: |
-        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF FUSE_ARANGE=1 JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        BENCHMARK_LOG=search_wino_cifar_cached WINO=1 DEFAULT_FLOAT=HALF JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py

.github/workflows/test.yml

@@ -272,14 +272,6 @@ jobs:
         PYTHONPATH=. DEBUG=2 EMULATE_CUDA=1 ALLOW_TF32=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 EMULATE_INTEL=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 AMX=1 EMULATE_AMX=1 FORWARD_ONLY=1 PYTHON=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores
-    - name: Test tensor cores (TC=3)
-      run: |
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_METAL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMD_MFMA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_CUDA=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_INTEL=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
-        PYTHONPATH=. DEBUG=2 PYTHON=1 EMULATE_AMX=1 AMX=1 python3 ./test/test_linearizer.py TestLinearizer.test_tensor_cores_emulation
     - name: Test device flop counts
       run: |
         PYTHONPATH=. DEBUG=2 EMULATE_METAL=1 PYTHON=1 python3 ./test/test_uops_stats.py TestUOpsStatsMatmulHalf
@@ -310,8 +302,6 @@ jobs:
       run: PYTHON=1 python3 -m pytest test/test_uops.py --durations=20
     - name: Test symbolic with Python emulator
       run: PYTHONPATH=. PYTHON=1 python3 test/test_symbolic_ops.py
-    - name: test_linearizer_failures with Python emulator
-      run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_linearizer_failures.py::TestLinearizerFailures::test_failure_1
     - name: test_renderer_failures with Python emulator
       run: PYTHONPATH=. PYTHON=1 python3 -m pytest -rA test/test_renderer_failures.py::TestRendererFailures
 
@@ -336,7 +326,7 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
-        python3 -m ruff check examples/mlperf/model_train.py --ignore E501
+        python3 -m ruff check examples/mlperf/ --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
     - name: Run mypy
@@ -371,8 +361,15 @@ jobs:
       run: PYTHONPATH="." python test/external/external_uop_gc.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-    - name: Repo line count < 14000 lines
-      run: MAX_LINE_COUNT=14000 python sz.py
+    - name: Regen dataset on test_tiny
+      run: |
+        test/external/process_replay/reset.py
+        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
+        PYTHONPATH=. python extra/optimization/extract_dataset.py
+        gzip -c /tmp/sops > extra/datasets/sops.gz
+        DEBUG=1 MIN_ASTS=1 PYTHONPATH=. python extra/optimization/get_action_space.py
+    - name: Repo line count < 14600 lines
+      run: MAX_LINE_COUNT=14600 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -450,6 +447,8 @@ jobs:
         run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot LLVM compile
         run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
+      - name: Test openpilot compile4
+        run: PYTHONPATH="." NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -474,6 +473,10 @@ jobs:
         run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
       - name: Test ONNX (LLVM)
         run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Test ONNX Runner (CPU)
+        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+      - name: Test ONNX Runner (WEBGPU)
+        run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
       - name: Test Additional ONNX Ops (CPU)
         run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
@@ -510,8 +513,8 @@ jobs:
           REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
       - name: Test Optimization Helpers
         run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
-      - name: Test Action Space
-        run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
+      #- name: Test Action Space
+      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
       - name: Test Beam Search
         run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
       - name: Test MLPerf stuff
@@ -610,7 +613,7 @@ jobs:
       run: |
           WEBGPU=1 WEBGPU_BACKEND="WGPUBackendType_Vulkan" python3 -m pytest -n=auto test/ --ignore=test/models --ignore=test/unit \
           --ignore=test/test_copy_speed.py --ignore=test/test_rearrange_einops.py \
-          --ignore=test/test_fuzz_shape_ops.py --ignore=test/test_linearizer_failures.py --durations=20
+          --ignore=test/test_fuzz_shape_ops.py --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -677,6 +680,7 @@ jobs:
           key: ${{ matrix.backend }}-minimal
           deps: testing_minimal
           cuda: 'true'
+          ocelot: 'true'
       - name: Set env
         run: printf "${{ matrix.backend == 'PTX' && 'FORWARD_ONLY=1\nJIT=1\nOPT=2\nCUDA=1\nPTX=1\nMOCKGPU=1' || matrix.backend == 'nv' && 'NV=1\nMOCKGPU=1\nFORWARD_ONLY=1' }}" >> $GITHUB_ENV
       - name: Check Device.DEFAULT and print some source
@@ -744,6 +748,7 @@ jobs:
         python-version: '3.11'
         amd: 'true'
         cuda: 'true'
+        ocelot: 'true'
         llvm: 'true'
     - name: Run real world test
       run: METAL=1 python -m pytest -n=auto test/models/test_real_world.py --durations=20
@@ -759,8 +764,8 @@ jobs:
       run: PYTHONPATH="." METAL=1 python test/external/external_test_speed_llama.py
     - name: Test Beam Search
       run: PYTHONPATH="." METAL=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
-    - name: Fuzz Test linearizer
-      run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
+    #- name: Fuzz Test linearizer
+    #  run: PYTHONPATH="." METAL=1 DEPTH=4 FUZZ_N=50 FUZZ_MAX_SIZE=1000000 python test/external/fuzz_linearizer.py
     - name: Run TRANSCENDENTAL math
       run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
     - name: Run pytest (amd)
@@ -774,7 +779,6 @@ jobs:
       env:
         MOCKGPU: 1
         AMD: 1
-        AMD_LLVM: 1
         FORWARD_ONLY: 1
       run: |
         python -m pytest -n=auto test/test_hcq.py test/test_tiny.py test/test_amd_llvm.py --durations=20
@@ -810,17 +814,19 @@ jobs:
       run: npm cache clean --force
     - name: Install Puppeteer
       run: npm install puppeteer
-    - name: Run WEBGPU Efficientnet
-      run: node test/web/test_webgpu.js
-    - name: Run VIZ tests as external package
-      run: |
-        mkdir $GITHUB_WORKSPACE/test_dir
-        cd $GITHUB_WORKSPACE/test_dir
-        python -m venv venv
-        source venv/bin/activate
-        pip install $GITHUB_WORKSPACE
-        cp $GITHUB_WORKSPACE/test/web/test_viz.js .
-        node test_viz.js
+    # this is also flaky
+    #- name: Run WEBGPU Efficientnet
+    #  run: node test/web/test_webgpu.js
+    # this is flaky
+    #- name: Run VIZ tests as external package
+    #  run: |
+    #    mkdir $GITHUB_WORKSPACE/test_dir
+    #    cd $GITHUB_WORKSPACE/test_dir
+    #    python -m venv venv
+    #    source venv/bin/activate
+    #    pip install $GITHUB_WORKSPACE
+    #    cp $GITHUB_WORKSPACE/test/web/test_viz.js .
+    #    node test_viz.js
 
   osxremote:
     name: MacOS (remote metal)
@@ -852,9 +858,8 @@ jobs:
     timeout-minutes: 20
     env:
       REMOTE: 1
-      REMOTEDEV: 'AMD'
+      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
       PYTHONPATH: ${{ github.workspace }}
-      MOCKGPU: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
@@ -865,6 +870,21 @@ jobs:
           deps: testing_minimal
           amd: 'true'
           llvm: 'true'
+      - name: Start remote server
+        run: |
+          start_server() {
+            systemd-run --user \
+              --unit="$1" \
+              --setenv=REMOTEDEV=AMD \
+              --setenv=MOCKGPU=1 \
+              --setenv=PYTHONPATH=. \
+              --setenv=PORT="$2" \
+              --working-directory="$(pwd)" \
+              python tinygrad/runtime/ops_remote.py
+          }
+
+          start_server "remote-server-1" 6667
+          start_server "remote-server-2" 6668
       - name: Check Device.DEFAULT and print some source
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
@@ -872,7 +892,12 @@ jobs:
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
       - name: Run REMOTE=1 Test
         run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
+          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
+      - name: Show remote server logs
+        if: always()
+        run: |
+          journalctl --user -u remote-server-1 --no-pager
+          journalctl --user -u remote-server-2 --no-pager
 
   osxtests:
     strategy:

autogen_stubs.sh

@@ -118,7 +118,9 @@ generate_nv() {
   clang2py -k cdefstum \
     extra/nv_gpu_driver/clc6c0qmd.h \
     extra/nv_gpu_driver/clcec0qmd.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0000.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl0080.h \
+    $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/cl2080_notification.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc56f.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/class/clc86f.h \
@@ -165,6 +167,26 @@ generate_nv() {
 nv_status_codes = {}
 /^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py
 
+  clang2py -k cdefstum \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
+    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
+    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
+    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
+    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
+    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
+    extra/nv_gpu_driver/g_rpc-message-header.h \
+    extra/nv_gpu_driver/gsp_static_config.h \
+    extra/nv_gpu_driver/vbios.h \
+    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
+    -o $BASE/nv/nv.py
+
+  fixup $BASE/nv/nv.py
   python3 -c "import tinygrad.runtime.autogen.nv_gpu"
 }
 
@@ -390,8 +412,8 @@ generate_am() {
     $AMKERN_AMD/pm/swsmu/inc/pmfw_if/smu14_driver_if_v14_0.h \
     extra/amdpci/headers/amdgpu_smu.h \
     --clang-args="-include stdint.h" \
-    -o $BASE/am/smu_v14_0_3.py
-  fixup $BASE/am/smu_v14_0_3.py
+    -o $BASE/am/smu_v14_0_2.py
+  fixup $BASE/am/smu_v14_0_2.py
 }
 
 generate_sqtt() {
@@ -436,6 +458,7 @@ elif [ "$1" == "kfd" ]; then generate_kfd
 elif [ "$1" == "nv" ]; then generate_nv
 elif [ "$1" == "amd" ]; then generate_amd
 elif [ "$1" == "am" ]; then generate_am
+elif [ "$1" == "nvdrv" ]; then generate_nvdrv
 elif [ "$1" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring

docs/abstractions2.py

@@ -59,11 +59,11 @@ st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,)
 s = UOp(Ops.SINK, dtypes.void, (st_0,))
 
 # convert the computation to a "linearized" format (print the format)
-from tinygrad.engine.realize import get_kernel, CompiledRunner
-kernel = get_kernel(Device[DEVICE].renderer, s).linearize()
+from tinygrad.engine.realize import get_program, CompiledRunner
+program = get_program(s, Device[DEVICE].renderer)
 
 # compile a program (and print the source)
-fxn = CompiledRunner(kernel.to_program())
+fxn = CompiledRunner(program)
 print(fxn.p.src)
 # NOTE: fxn.clprg is the CPUProgram
 
@@ -78,7 +78,7 @@ print("******** third, the UOp ***********")
 
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.engine.grouper import get_kernelize_map
+from tinygrad.kernelize.kernelize import get_kernelize_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)

docs/developer/layout.md

@@ -0,0 +1,66 @@
+# tinygrad directory layout
+
+This explains the flow of a big graph down to programs.
+
+Directories are listed in order of how they are processed.
+
+---
+
+## tinygrad/kernelize
+
+Group UOps into kernels.
+
+::: tinygrad.kernelize.kernelize.get_kernelize_map
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/opt
+
+Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
+
+::: tinygrad.opt.get_optimized_ast
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/codegen
+
+Transform the optimized ast into a linearized list of UOps.
+
+::: tinygrad.codegen.full_rewrite
+    options:
+        members: false
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/renderer
+
+Transform the linearized list of UOps into a program, represented as a string.
+
+::: tinygrad.renderer.Renderer
+    options:
+        members:
+            - render
+        show_labels: false
+        show_source: false
+
+---
+
+## tinygrad/engine
+
+Abstracted high level interface to the runtimes.
+
+::: tinygrad.engine.realize.get_program
+    options:
+        members: false
+        show_labels: false
+        show_source: false

docs/ramp.py

@@ -239,7 +239,7 @@ print("******* PART 3 *******")
 # it's much simpler than what's in LLVM or MLIR
 
 from tinygrad import dtypes
-from tinygrad.uop import UOp, Ops
+from tinygrad.uop.ops import UOp, Ops
 
 # first, we'll construct some const UOps
 a = UOp(Ops.CONST, dtypes.int, arg=2)

docs/tensor/elementwise.md

@@ -35,6 +35,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 
 ::: tinygrad.Tensor.relu
 ::: tinygrad.Tensor.sigmoid
+::: tinygrad.Tensor.logsigmoid
 ::: tinygrad.Tensor.hardsigmoid
 ::: tinygrad.Tensor.elu
 ::: tinygrad.Tensor.celu

examples/benchmark_onnx.py

@@ -1,4 +1,4 @@
-import sys, time, pickle
+import sys, time
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate

examples/coder.py

@@ -4,7 +4,7 @@ sys.path.append(os.getcwd())
 
 from io import StringIO
 from contextlib import redirect_stdout
-from tinygrad import Tensor, nn, Device, dtypes
+from tinygrad import Tensor, nn
 from tinygrad.helpers import Timing, colored, getenv, fetch
 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from sentencepiece import SentencePieceProcessor

examples/handcode_opt.py

@@ -2,11 +2,11 @@ from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
 

examples/hlb_cifar10.py

@@ -7,8 +7,8 @@ import random, time
 import numpy as np
 from typing import Optional
 from extra.lr_scheduler import OneCycleLR
-from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit
-from tinygrad.nn.state import get_state_dict, get_parameters
+from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit, Variable
+from tinygrad.nn.state import get_state_dict
 from tinygrad.nn import optim
 from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -145,6 +145,7 @@ hyp = {
   },
 }
 
+@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():
 
   def set_seed(seed):
@@ -201,24 +202,37 @@ def train_cifar():
     idx_y = Tensor.arange(H, dtype=dtypes.int32).reshape((1,1,H,1))
     return (idx_x >= low_x) * (idx_x < (low_x + mask_size)) * (idx_y >= low_y) * (idx_y < (low_y + mask_size))
 
-  def random_crop(X:Tensor, crop_size=32):
-    mask = make_square_mask(X.shape, crop_size)
-    mask = mask.expand((-1,3,-1,-1))
-    X_cropped = Tensor(X.numpy()[mask.numpy()])
-    return X_cropped.reshape((-1, 3, crop_size, crop_size))
+  # Similar, but different enough.
+  def make_random_crop_indices(shape, mask_size) -> Tensor:
+    BS, _, H, W = shape
+    low_x = Tensor.randint(BS, low=0, high=W-mask_size).reshape(BS,1,1,1)
+    low_y = Tensor.randint(BS, low=0, high=H-mask_size).reshape(BS,1,1,1)
+    idx_x = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,1,mask_size))
+    idx_y = Tensor.arange(mask_size, dtype=dtypes.int32).reshape((1,1,mask_size,1))
+    return low_x, low_y, idx_x, idx_y
 
-  def cutmix(X:Tensor, Y:Tensor, mask_size=3):
-    # fill the square with randomly selected images from the same batch
+  def random_crop(X:Tensor, crop_size=32):
+    Xs, Ys, Xi, Yi = make_random_crop_indices(X.shape, crop_size)
+    return X.gather(-1, (Xs + Xi).expand(-1, 3, X.shape[2], -1)).gather(-2, ((Ys+Yi).expand(-1, 3, crop_size, crop_size)))
+
+  def cutmix(X, Y, order, mask_size=3):
     mask = make_square_mask(X.shape, mask_size)
-    order = list(range(0, X.shape[0]))
-    random.shuffle(order)
-    X_patch = Tensor(X.numpy()[order], device=X.device, dtype=X.dtype)
-    Y_patch = Tensor(Y.numpy()[order], device=Y.device, dtype=Y.dtype)
+    X_patch, Y_patch = X[order], Y[order]
     X_cutmix = mask.where(X_patch, X)
     mix_portion = float(mask_size**2)/(X.shape[-2]*X.shape[-1])
     Y_cutmix = mix_portion * Y_patch + (1. - mix_portion) * Y
     return X_cutmix, Y_cutmix
 
+  @TinyJit
+  def augmentations(X:Tensor, Y:Tensor):
+    perms = Tensor.randperm(X.shape[0], device=X.device) # We reuse perms for cutmix, because they are expensivne to generate
+    if getenv("RANDOM_CROP", 1):
+      X = random_crop(X, crop_size=32)
+    if getenv("RANDOM_FLIP", 1):
+      X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
+    X, Y = X[perms], Y[perms]
+    return X, Y, *cutmix(X, Y, perms, mask_size=hyp['net']['cutmix_size'])
+
   # the operations that remain inside batch fetcher is the ones that involves random operations
   def fetch_batches(X_in:Tensor, Y_in:Tensor, BS:int, is_train:bool):
     step, epoch = 0, 0
@@ -226,28 +240,16 @@ def train_cifar():
       st = time.monotonic()
       X, Y = X_in, Y_in
       if is_train:
-        # TODO: these are not jitted
-        if getenv("RANDOM_CROP", 1):
-          X = random_crop(X, crop_size=32)
-        if getenv("RANDOM_FLIP", 1):
-          X = (Tensor.rand(X.shape[0],1,1,1) < 0.5).where(X.flip(-1), X) # flip LR
-        if getenv("CUTMIX", 1):
-          if step >= hyp['net']['cutmix_steps']:
-            X, Y = cutmix(X, Y, mask_size=hyp['net']['cutmix_size'])
-        order = list(range(0, X.shape[0]))
-        random.shuffle(order)
-        X, Y = X.numpy()[order], Y.numpy()[order]
-      else:
-        X, Y = X.numpy(), Y.numpy()
+        X, Y, X_cm, Y_cm = augmentations(X, Y)
+        if getenv("CUTMIX", 1) and step >= hyp['net']['cutmix_steps']: X, Y = X_cm, Y_cm
       et = time.monotonic()
       print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
-      for i in range(0, X.shape[0], BS):
-        # pad the last batch  # TODO: not correct for test
-        batch_end = min(i+BS, Y.shape[0])
-        x = Tensor(X[batch_end-BS:batch_end], device=X_in.device, dtype=X_in.dtype)
-        y = Tensor(Y[batch_end-BS:batch_end], device=Y_in.device, dtype=Y_in.dtype)
+
+      vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
+      for i in range(0, full_batches, BS):
         step += 1
-        yield x, y
+        vib = vi.bind(i)
+        yield X[vib:vib+BS], Y[vib:vib+BS]
       epoch += 1
       if not is_train: break
 

examples/llama3.py

@@ -157,7 +157,11 @@ MODEL_PARAMS = {
   "70B": {
     "args": {"dim": 8192, "n_heads": 64, "n_kv_heads": 8, "n_layers": 80, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256,  "hidden_dim": 28672},
     "files": 8
-  }
+  },
+  "405B": {
+    "args": {"dim": 16384, "n_heads": 128, "n_kv_heads": 8, "n_layers": 126, "norm_eps": 1e-5, "rope_theta": 500000, "vocab_size": 128256,  "hidden_dim": 53248},
+    "files": 191
+  },
 }
 def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
   # build model
@@ -236,7 +240,7 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--download_model", action="store_true", help="Download a model")
   parser.add_argument("--model", type=Path, help="Model path")
-  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B", "405B"], default="1B", help="Model size")
   parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
   parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
   parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
@@ -244,7 +248,7 @@ if __name__ == "__main__":
   parser.add_argument("--port", type=int, default=7776, help="Web server port")
   parser.add_argument("--debug", action="store_true", help="Enable debug mode")
   parser.add_argument("--seed", type=int, help="Random seed")
-  parser.add_argument("--temperature", type=int, default=0.85, help="Temperature")
+  parser.add_argument("--temperature", type=float, default=0.85, help="Temperature")
   parser.add_argument("--benchmark", action="store_true", help="Run a benchmark")
   parser.add_argument("--timing", action="store_true", help="Print timing per token")
   parser.add_argument("--profile", action="store_true", help="Output profile data")

examples/llm.c/export.py

@@ -1,11 +1,11 @@
 #!/usr/bin/env python3
 import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
-from tinygrad import Device, nn, Tensor, dtypes, Variable
+from tinygrad import Device, nn, Tensor, dtypes
 Device.DEFAULT = "CPU"
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
-from tinygrad.engine.realize import get_kernel, run_schedule
+from tinygrad.helpers import dedup, flatten, getenv, GlobalCounters, to_function_name
+from tinygrad.engine.realize import get_kernel
 from tinygrad.engine.memory import memory_planner
 from tinygrad.uop.ops import Ops
 

examples/mlperf/helpers.py

@@ -212,7 +212,7 @@ def get_mlperf_bert_model():
   from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert
 
   bert.Linear = LinearBert
-  bert.Embedding = EmbeddingBert 
+  bert.Embedding = EmbeddingBert
   bert.LayerNorm = LayerNormBert
 
   from extra.models.bert import BertForPretraining

examples/mlperf/initializers.py

@@ -39,7 +39,7 @@ class LinearBert(nn.Linear):
   def __init__(self, in_features, out_features, bias=True, std=0.02):
     self.weight = std * rand_truncn(out_features, in_features, dtype=dtypes.float32)
     self.bias = Tensor.zeros(out_features, dtype=dtypes.float32) if bias else None
-  
+
   def __call__(self, x:Tensor):
     return x.cast(dtypes.default_float).linear(self.weight.cast(dtypes.default_float).transpose(), self.bias.cast(dtypes.default_float) if self.bias is not None else None)
 

examples/mlperf/lr_schedulers.py

@@ -1,4 +1,5 @@
-from tinygrad import Tensor, dtypes
+import math
+from tinygrad import dtypes
 from tinygrad.nn.optim import Optimizer
 
 from extra.lr_scheduler import LR_Scheduler
@@ -20,3 +21,19 @@ class PolynomialDecayWithWarmup(LR_Scheduler):
     warmup_lr = (self.epoch_counter * (1.0 / self.warmup)) * self.initial_lr
     x = (1 - (self.epoch_counter - self.warmup) / (self.epochs - self.warmup + 1))
     return (self.epoch_counter <= self.warmup).where(warmup_lr, (self.initial_lr - self.end_lr) * x ** self.power + self.end_lr).cast(self.optimizer.lr.dtype)
+
+class CosineAnnealingLRWithWarmup(LR_Scheduler):
+  def __init__(self, optimizer:Optimizer, base_lr, end_lr, warmup_steps:int, decay_steps:int):
+    assert warmup_steps > 0 and decay_steps > 0
+    super().__init__(optimizer)
+    self.base_lr = base_lr
+    self.end_lr = end_lr
+    self.warmup_steps = warmup_steps
+    self.decay_steps = decay_steps
+    # set lr for first warmup step
+    self.optimizer.lr.assign(self.get_lr()).realize()
+
+  def get_lr(self):
+    warmup_lr = ((self.epoch_counter+1) / self.warmup_steps) * self.base_lr
+    decay_lr = self.end_lr + 0.5 * (self.base_lr-self.end_lr) * (1 + (((self.epoch_counter+1-self.warmup_steps)/self.decay_steps) * math.pi).cos())
+    return (self.epoch_counter < self.warmup_steps).where(warmup_lr, decay_lr).cast(self.optimizer.lr.dtype)

examples/mlperf/metrics.py

@@ -1,6 +1,6 @@
-import re
-import string
+import re, string
 from collections import Counter
+from tinygrad import Tensor
 
 def levenshtein(a, b):
   n, m = len(a), len(b)
@@ -59,3 +59,11 @@ def f1_score(x, y):
   p = ns / len(xt)
   r = ns / len(yt)
   return 2 * p * r / (p + r)
+
+def log_perplexity(logit:Tensor, target:Tensor, ignore_index:int|None=None):
+  # logit has shape (n_samples, seq_len, vocab_size), target has shape (n_samples, seq_len)
+  assert logit.ndim == 3, logit.ndim
+  assert target.ndim == 2, target.ndim
+  assert logit.shape[:2] == target.shape, f"{logit.shape[:2]=}, {target.shape=}"
+  log_prob = logit.log_softmax(axis=-1)
+  return log_prob.transpose(1, 2).nll_loss(target, ignore_index=ignore_index)

examples/mlperf/model_train.py

@@ -5,7 +5,7 @@ import multiprocessing
 from tinygrad import Device, GlobalCounters, Tensor, TinyJit, dtypes
 from tinygrad.helpers import getenv, BEAM, WINO, round_up, diskcache_clear, FUSE_CONV_BW, Profiling
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_load, safe_save
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam, AdamW
 
 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
@@ -914,18 +914,26 @@ def train_rnnt():
   pass
 
 @TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
-                    masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
   optimizer.zero_grad()
 
-  lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-  loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-  (loss * loss_scaler).backward()
+  for i in range(grad_acc):
+    input_ids, segment_ids = kwargs[f"input_ids{i}"], kwargs[f"segment_ids{i}"]
+    # NOTE: these two have different names
+    attention_mask, masked_positions = kwargs[f"input_mask{i}"], kwargs[f"masked_lm_positions{i}"]
+    masked_lm_ids, masked_lm_weights, next_sentence_labels = kwargs[f"masked_lm_ids{i}"], kwargs[f"masked_lm_weights{i}"], kwargs[f"next_sentence_labels{i}"]
 
-  global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device)
+    for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
+      if len(GPUS) > 1: t.shard_(GPUS, axis=0)
+      else: t.to_(GPUS[0])
+
+    lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
+    loss = model.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+    (loss * loss_scaler).backward()
+    # TODO: OOM without this realize with large grad_acc
+    Tensor.realize(*[p.grad for p in optimizer.params])
+
+  global_norm = Tensor(0.0, dtype=dtypes.float32, device=optimizer[0].device)
   for p in optimizer.params:
     p.grad = p.grad / loss_scaler
     global_norm += p.grad.float().square().sum()
@@ -999,16 +1007,19 @@ def train_bert():
     MLLOGGER = None
 
   # ** hyperparameters **
-  BS                 = config["GLOBAL_BATCH_SIZE"]      = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
+  BS                 = config["BS"]                     = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  # TODO: mlperf logging
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
   EVAL_BS            = config["EVAL_BS"]                = getenv("EVAL_BS", 1 * len(GPUS))
-  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(BS/96))
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(GBS/96))
   opt_lamb_beta_1    = config["OPT_LAMB_BETA_1"]        = getenv("OPT_LAMB_BETA_1", 0.9)
   opt_lamb_beta_2    = config["OPT_LAMB_BETA_2"]        = getenv("OPT_LAMB_BETA_2", 0.999)
 
-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // GBS)
   warmup_steps       = config["NUM_WARMUP_STEPS"]       = getenv("NUM_WARMUP_STEPS", 1)
   max_eval_steps     = config["MAX_EVAL_STEPS"]         = getenv("MAX_EVAL_STEPS", (10000 + EVAL_BS - 1) // EVAL_BS) # EVAL_BS * MAX_EVAL_STEPS >= 10000
-  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * BS + 3000000) / 25000) * 25000) / BS)) # Round down
+  eval_step_freq     = config["EVAL_STEP_FREQ"]         = getenv("EVAL_STEP_FREQ", int((math.floor(0.05 * (230.23 * GBS + 3000000) / 25000) * 25000) / GBS)) # Round down
   save_ckpt_freq     = config["SAVE_CKPT_FREQ"]         = getenv("SAVE_CKPT_FREQ", 1000)
   keep_ckpt_amount   = config["KEEP_CKPT_AMOUNT"]       = getenv("KEEP_CKPT_AMOUNT", 5)
   save_ckpt_dir      = config["SAVE_CKPT_DIR"]          = getenv("SAVE_CKPT_DIR", "./ckpts")
@@ -1066,7 +1077,7 @@ def train_bert():
   scheduler_wd = PolynomialDecayWithWarmup(optimizer_wd, max_lr, 0, train_steps, warmup_steps, power=poly_power)
   scheduler_no_wd = PolynomialDecayWithWarmup(optimizer_no_wd, max_lr, 0, train_steps, warmup_steps, power=poly_power)
   scheduler_group = LRSchedulerGroup(scheduler_wd, scheduler_no_wd)
-  print(f"training with batch size {BS} for one epoch with {train_steps} steps")
+  print(f"training with global batch size {GBS} for one epoch with {train_steps} steps")
 
   # log mlperf hparams
   if MLLOGGER:
@@ -1115,11 +1126,11 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, next(train_it)
+  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
 
   if RUNMLPERF:
     if MLLOGGER:
-      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})
+      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
 
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
@@ -1128,14 +1139,13 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler,
-          train_data["input_ids"], train_data["segment_ids"], train_data["input_mask"], train_data["masked_lm_positions"], \
-          train_data["masked_lm_ids"], train_data["masked_lm_weights"], train_data["next_sentence_labels"], GPUS)
+        data = {f"{k}{i}":v for i,d in enumerate(train_data) for k,v in d.items()}
+        loss, global_norm, lr = train_step_bert(model, optimizer_group, scheduler_group, loss_scaler, GPUS, grad_acc, **data)
 
         pt = time.perf_counter()
 
         try:
-          next_data = next(train_it)
+          next_data = [next(train_it) for _ in range(grad_acc)]
         except StopIteration:
           next_data = None
 
@@ -1156,7 +1166,7 @@ def train_bert():
       if WANDB:
         wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
                     "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
+                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*GBS})
 
       train_data, next_data = next_data, None
       i += 1
@@ -1171,7 +1181,7 @@ def train_bert():
     # ** eval loop **
     if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
       if MLLOGGER and RUNMLPERF:
-        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
+        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()
       eval_lm_losses = []
@@ -1221,11 +1231,11 @@ def train_bert():
 
       if WANDB:
         wandb.log({"eval/lm_loss": avg_lm_loss, "eval/clsf_loss": avg_clsf_loss, "eval/lm_accuracy": avg_lm_acc, \
-                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*BS})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*GBS})
 
       if MLLOGGER and RUNMLPERF:
-        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*BS, metadata={"epoch_count": i*BS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
-        MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=avg_lm_acc, metadata={"epoch_num": i*BS, "masked_lm_accuracy": avg_lm_acc})
+        MLLOGGER.end(key=mllog_constants.EVAL_STOP, value=i*GBS, metadata={"epoch_count": i*GBS, "step_num": i, "samples_count": config["EVAL_BS"] * config["MAX_EVAL_STEPS"]})
+        MLLOGGER.event(key=mllog_constants.EVAL_ACCURACY, value=avg_lm_acc, metadata={"epoch_num": i*GBS, "masked_lm_accuracy": avg_lm_acc})
 
       # save model if achieved target
       if not achieved and avg_lm_acc >= target:
@@ -1240,10 +1250,10 @@ def train_bert():
         hours = int(total_seconds // 3600)
         minutes = int((total_seconds % 3600) // 60)
         seconds = total_seconds % 60
-        print(f"Reference Convergence point reached after {i * BS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
+        print(f"Reference Convergence point reached after {i * GBS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
         achieved = True
         if MLLOGGER and RUNMLPERF:
-          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*BS, metadata={"epoch_num": i*BS})
+          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*GBS, metadata={"epoch_num": i*GBS})
           MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata=dict(status=mllog_constants.SUCCESS))
         # stop once hitting the target
         break
@@ -1271,12 +1281,78 @@ def train_bert():
         os.remove(os.path.join(ckpt_dir, last))
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.end(key="checkpoint_stop", value=None, metadata={"step_num": i})
-        MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * BS, "step_num": i, "first_step_num": i+1})
+        MLLOGGER.start(key=mllog_constants.BLOCK_START, value=None, metadata={"first_epoch_num": 1, "epoch_num": 1, "epoch_count": 1, "samples_count": i * GBS, "step_num": i, "first_step_num": i+1})
         previous_step = i
 
-def train_maskrcnn():
-  # TODO: Mask RCNN
-  pass
+def train_llama3():
+  from extra.models.llama import Transformer
+  from examples.llama3 import MODEL_PARAMS
+  from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
+
+  config = {}
+  BS                 = config["BS"]                     = getenv("BS", 4)
+  grad_acc           = config["GRADIENT_ACC_STEPS"]     = getenv("GRADIENT_ACC_STEPS", 1)
+  GBS                = config["GLOBAL_BATCH_SIZE"]      = BS * grad_acc
+
+  opt_adamw_beta_1 = 0.9
+  opt_adamw_beta_2 = 0.95
+  opt_adamw_epsilon = 1e-5
+  opt_adamw_weight_decay = 0.1
+
+  opt_gradient_clip_norm = 1.0
+  sequence_length = 8192
+  opt_learning_rate_warmup_steps = getenv("WARMUP_STEPS", math.ceil(8000 * 1152 / GBS))
+  opt_learning_rate_decay_steps = getenv("DECAY_STEPS", math.ceil(1_200_000 * 1152 / GBS) - opt_learning_rate_warmup_steps)
+  opt_base_learning_rate = getenv("LR", 8e-5 * GBS / 1152)  # NOTE: cannot change for benchmark
+  opt_end_learning_rate = 8e-7
+
+  # TODO: confirm weights are in bf16
+  # vocab_size from the mixtral tokenizer
+  model = Transformer(**(MODEL_PARAMS[getenv("LLAMA3_SIZE", "8B")]["args"]|{"vocab_size": 32000}), max_context=sequence_length, jit=False, disable_kv_cache=True)
+
+  optim = AdamW(get_parameters(model), lr=0.0,
+                b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay)
+  scheduler = CosineAnnealingLRWithWarmup(optim, opt_base_learning_rate, opt_end_learning_rate, opt_learning_rate_warmup_steps, opt_learning_rate_decay_steps)
+
+  @TinyJit
+  @Tensor.train()
+  def train_step(model, x, y):
+    optim.zero_grad()
+    logits:Tensor = model(x, start_pos=0, temperature=math.nan)
+    loss = logits.cross_entropy(y)
+    loss.backward()
+
+    # L2 norm grad clip
+    # https://github.com/NVIDIA/NeMo/blob/3368c3fc0b4a186ab33a1d68a504315100c0b2a6/nemo/collections/nlp/modules/common/megatron/clip_grads.py#L57
+    # https://docs.pytorch.org/docs/stable/generated/torch.nn.utils.clip_grad_norm_.html
+    if not getenv("DISABLE_GRAD_CLIP_NORM"):
+      total_norm = Tensor(0.0, dtype=dtypes.float32, device=optim.params[0].device)
+      for p in optim.params:
+        total_norm += p.grad.float().square().sum()
+      total_norm = total_norm.sqrt().contiguous()
+      for p in optim.params:
+        p.grad = p.grad * opt_gradient_clip_norm / (total_norm + 1e-6)
+
+    optim.step()
+    scheduler.step()
+
+    lr = optim.lr
+    loss.realize(lr)
+    return loss, lr
+
+  # overfitting this example should give cross_entropy log(BS)
+  fake_input = Tensor([list(range(getenv("SEQLEN", 10)))], dtype="int16").expand(BS, -1)
+  fake_label = Tensor(list(range(BS)), dtype="int16")
+
+  for _ in range(100):
+    GlobalCounters.reset()
+    loss, lr = train_step(model, fake_input, fake_label)
+    # BS=2 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    # uses 43% ~= 83GB
+    # 8B bf16 = 16GB. model + grad + optim m and v = 64GB
+    # TODO: this OOM
+    # BS=1 SEQLEN=4000 OPTIM_DTYPE=bfloat16 LLAMA3_SIZE=8B WARMUP_STEPS=2 DECAY_STEPS=300 PYTHONPATH=. AMD=1 MODEL=llama3 python3 examples/mlperf/model_train.py
+    print(loss.item(), lr.item(), f"{GlobalCounters.global_mem//10**9=}")
 
 if __name__ == "__main__":
   multiprocessing.set_start_method('spawn')

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/dev_beam.sh

@@ -9,6 +9,6 @@ export BEAM=3 BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MI
 export IGNORE_JIT_FIRST_BEAM=1 FREE_INTERMEDIATE=0
 export BASEDIR="/raid/datasets/wiki"
 
-export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
+export BENCHMARK=10 BERT_LAYERS=2
 
 python3 examples/mlperf/model_train.py

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_8xMI300X/run_and_time.sh

@@ -22,8 +22,7 @@ export SEED=$RANDOM
 DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_8xMI300x_${DATETIME}_${SEED}.log"
 
-# init  # TODO: without DEBUG=2 it hangs
-BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 DEBUG=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
+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

examples/mlperf/training_submission_v5.1/tinycorp/benchmarks/bert/implementations/tinybox_red/run_and_time.sh

@@ -27,6 +27,4 @@ sleep 5 && sudo rmmod amdgpu || true
 BENCHMARK=10 INITMLPERF=1 BERT_LAYERS=2 python3 examples/mlperf/model_train.py | tee $LOGFILE
 
 # run
-# TODO: AM driver resulted in nan
-sudo modprobe amdgpu
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

examples/openpilot/compile4.py

@@ -1,8 +1,8 @@
 import sys, onnx
-from tinygrad import Tensor, fetch, GlobalCounters
-from tinygrad.uop import UOp
+from tinygrad import Tensor, fetch, GlobalCounters, dtypes
+from tinygrad.uop.ops import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.engine.grouper import get_kernelize_map
+from tinygrad.kernelize.kernelize import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule
 
@@ -17,7 +17,7 @@ if __name__ == "__main__":
   onnx_model = onnx.load(onnx_file)
   run_onnx = OnnxRunner(onnx_model)
 
-  inputs = run_onnx.get_empty_input_data("npy")
+  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()]
@@ -37,12 +37,12 @@ if __name__ == "__main__":
   independent = UOp.sink(*independent_set.keys())
   kernelized = get_kernelize_map(independent)
   independent = independent.substitute(kernelized)
-  schedule, var_vals, becomes_map = create_schedule_with_vars(independent)
+  schedule, var_vals = create_schedule_with_vars(independent)
   run_schedule(schedule)
 
   print("**** real ****")
   GlobalCounters.reset()
-  out.uop = root.substitute(kernelized).substitute(becomes_map)
+  out.uop = root.substitute(kernelized)
   out.kernelize()
 
   # realize

examples/other_mnist/beautiful_mnist_torch.py

@@ -27,7 +27,7 @@ class Model(nn.Module):
 
 if __name__ == "__main__":
   if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch
+    import tinygrad.frontend.torch  # noqa: F401
     device = torch.device("tiny")
   else:
     device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))

examples/sdxl.py

@@ -5,7 +5,7 @@
 
 from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
 from tinygrad.nn import Conv2d, GroupNorm
-from tinygrad.nn.state import safe_load, load_state_dict, get_state_dict
+from tinygrad.nn.state import safe_load, load_state_dict
 from tinygrad.helpers import fetch, trange, colored, Timing
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding

examples/so_vits_svc.py

@@ -5,7 +5,7 @@ from functools import partial, reduce
 from pathlib import Path
 from typing import Tuple, Optional, Type
 from tinygrad import nn, dtypes, Tensor
-from tinygrad.helpers import getenv
+from tinygrad.helpers import getenv, fetch
 from tinygrad.nn.state import torch_load
 from examples.vits import ResidualCouplingBlock, PosteriorEncoder, Encoder, ResBlock1, ResBlock2, LRELU_SLOPE, sequence_mask, split, get_hparams_from_file, load_checkpoint, weight_norm, HParams
 from examples.sovits_helpers import preprocess
@@ -19,10 +19,6 @@ F0_MIN = 50.0
 F0_MEL_MIN = 1127 * np.log(1 + F0_MIN / 700)
 F0_MEL_MAX = 1127 * np.log(1 + F0_MAX / 700)
 
-def download_if_not_present(file_path: Path, url: str):
-  if not os.path.isfile(file_path): download_file(url, file_path)
-  return file_path
-
 class SpeechEncoder:
   def __init__(self, hidden_dim, model:ContentVec): self.hidden_dim, self.model = hidden_dim, model
   def encode(self, ): raise NotImplementedError("implement me")
@@ -97,7 +93,7 @@ class ContentVec:
     return res, padding_mask
   @classmethod
   def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
-    download_if_not_present(checkpoint_path, checkpoint_url)
+    fetch(checkpoint_url, checkpoint_path)
     cfg = load_fairseq_cfg(checkpoint_path)
     enc = cls(cfg.model)
     _ = load_checkpoint_enc(checkpoint_path, enc, None)
@@ -324,9 +320,9 @@ class Synthesizer:
     return f0_coarse
   @classmethod
   def load_from_pretrained(cls, config_path:str, config_url:str, weights_path:str, weights_url:str) -> Synthesizer:
-    download_if_not_present(config_path, config_url)
+    fetch(config_url, config_path)
     hps = get_hparams_from_file(config_path)
-    download_if_not_present(weights_path, weights_url)
+    fetch(weights_url, weights_path)
     net_g = cls(hps.data.filter_length // 2 + 1, hps.train.segment_size // hps.data.hop_length, **hps.model)
     _ = load_checkpoint(weights_path, net_g, None, skip_list=["f0_decoder"])
     logging.debug(f"{cls.__name__}:Loaded model with hps: {hps}")
@@ -602,7 +598,7 @@ if __name__=="__main__":
   speaker = args.speaker if args.speaker is not None else list(hps.spk.__dict__.keys())[0]
 
   ### Loading audio and slicing ###
-  if audio_path == DEMO_PATH: download_if_not_present(DEMO_PATH, DEMO_URL)
+  if audio_path == DEMO_PATH: fetch(DEMO_URL, DEMO_PATH)
   assert Path(audio_path).is_file() and Path(audio_path).suffix == ".wav"
   chunks = preprocess.cut(audio_path, db_thresh=slice_db)
   audio_data, audio_sr = preprocess.chunks2audio(audio_path, chunks)

examples/stunning_mnist.py

@@ -7,7 +7,7 @@
 from examples.beautiful_mnist import Model
 from tinygrad import Tensor, nn, getenv, GlobalCounters, Variable
 from tinygrad.nn.datasets import mnist
-from tinygrad.helpers import trange, DEBUG
+from tinygrad.helpers import trange
 
 # STEPS=70 python3 examples/stunning_mnist.py
 # NOTE: it's broken with STACK=1, why?

examples/torch_cuda_kernel.py

@@ -2,10 +2,9 @@
 #!POPCORN gpu A100
 # not a stable API, but works
 
-import torch, functools
+import torch
 from tinygrad import Tensor, TinyJit, Device
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import get_single_element, Context, OSX
+from tinygrad.helpers import Context, OSX
 from tinygrad.dtype import _from_torch_dtype
 
 @TinyJit

examples/vgg7.py

@@ -2,8 +2,6 @@ import sys
 import random
 import json
 import numpy
-from pathlib import Path
-from PIL import Image
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import SGD
 from tinygrad.nn.state import safe_save, safe_load, get_state_dict, load_state_dict

examples/whisper.py

@@ -5,7 +5,7 @@ from typing import Optional, Union, Literal, List
 
 from tinygrad import Tensor, TinyJit, Variable, nn
 from tinygrad.nn.state import torch_load, load_state_dict
-from tinygrad.helpers import getenv, DEBUG, fetch
+from tinygrad.helpers import getenv, fetch
 
 import numpy as np
 import librosa

examples/yolov8-onnx.py

@@ -4,7 +4,6 @@ from ultralytics import YOLO
 from pathlib import Path
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs
-from tinygrad.tensor import Tensor
 
 os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():

examples/yolov8.py

@@ -1,7 +1,5 @@
 from tinygrad.nn import Conv2d, BatchNorm2d
 from tinygrad.tensor import Tensor
-from tinygrad.device import is_dtype_supported
-from tinygrad import dtypes
 import numpy as np
 from itertools import chain
 from pathlib import Path

extra/assembly/assembly.py

@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.codegen.kernel import Ops, MemOp, UOp
+from tinygrad.opt.kernel import Ops, MemOp, UOp
 from tinygrad.uop.ops import BinaryOps, UnaryOps
 from tinygrad.dtype import DType, dtypes
 from tinygrad.helpers import DEBUG
@@ -156,7 +156,7 @@ def uops_to_asmstyle(lang, function_name:str, uops:List[UOp]):
         lang.ins.append(AssemblyInstruction(Ops.ALU, out, [tmp], args))
       else:
         lang.ins.append(AssemblyInstruction(Ops.ALU, out, [lang.tor[x] for x in vin], args))
-    elif uop == Ops.DEFINE_ACC:
+    elif uop == Ops.DEFINE_REG:
       reg = lang.newreg(u, dtype=dtype)
       lang.ins.append(AssemblyInstruction(Ops.LOAD, reg, [], args))
     elif uop == Ops.SPECIAL:

extra/assembly/assembly_arm64.py

@@ -3,7 +3,7 @@ from platform import system
 from typing import Tuple, Dict, List, Optional
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad.helpers import CI
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
 

extra/assembly/assembly_ptx.py

@@ -1,7 +1,7 @@
 from typing import List
 import struct
 from tinygrad.codegen.assembly import uops_to_asmstyle, AssemblyLanguage
-from tinygrad.codegen.kernel import Ops, UOp
+from tinygrad.opt.kernel import Ops, UOp
 from tinygrad import dtypes
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_cuda import arch

extra/assembly/assembly_rdna.py

@@ -2,7 +2,7 @@ import yaml
 from typing import Tuple, Set, Dict
 from tinygrad import dtypes
 from tinygrad.codegen.assembly import AssemblyCodegen, Register
-from tinygrad.codegen.kernel import Ops
+from tinygrad.opt.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
 

extra/backends/triton.py

@@ -2,7 +2,7 @@ from typing import Dict, List, Final, Callable, DefaultDict
 from collections import defaultdict
 from tinygrad.uop.ops import UnaryOps, BinaryOps, TernaryOps, Op
 from tinygrad.helpers import DType, PtrDType, dtypes, ImageDType, DEBUG, getenv
-from tinygrad.codegen.kernel import  UOp, Ops
+from tinygrad.opt.kernel import  UOp, Ops
 from triton.compiler import compile as triton_compile
 import linecache
 import math
@@ -88,7 +88,7 @@ def uops_to_triton(function_name:str, uops:List[UOp]):
       assert dtype is not None
       if len(vin) == 2: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.load({r[vin[0]]} + { fill_dims_for_idx(r[vin[1]], dims)}, mask = {render_valid(valid)})', dtype)}")
       else: kk(f"{ssa(u, 'val')} = {render_cast(f'tl.where({r[vin[2]]}, tl.load({r[vin[0]]}+{fill_dims_for_idx(r[vin[1]],dims)} , mask={render_valid(valid+[r[vin[2]]])}), 0.0)', dtype)}")
-    elif uop == Ops.DEFINE_ACC: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
+    elif uop == Ops.DEFINE_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
     elif uop == Ops.CONST: r[u] = define_scalar([], dtype, args)
     elif uop == Ops.ASSIGN:
       kk(f"{r[vin[0]]} = {r[vin[1]].replace('//', '/')}")

extra/gemm/amd_matmul.py

@@ -4,7 +4,7 @@ import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
 

extra/gemm/halide_gemm.py

@@ -0,0 +1,90 @@
+import numpy as np
+import halide as hl
+from tinygrad.helpers import Timing, getenv
+
+# HL_DEBUG_CODEGEN=1
+N = getenv("N", 1024)
+
+def gemm_pipeline(gpu=False):
+  # ---------------- Vars & Parameters ----------------
+  i, j = hl.Var("i"), hl.Var("j")  # output tile coordinates
+
+  A = hl.InputBuffer(hl.Float(32), 2)  # [M, K]
+  B = hl.InputBuffer(hl.Float(32), 2)  # [K, N]
+
+  A.dim(0).set_bounds(0, N)
+  A.dim(1).set_bounds(0, N)
+  B.dim(0).set_bounds(0, N)
+  B.dim(1).set_bounds(0, N)
+
+  # ---------------- Definition ----------------
+
+  k = hl.RDom([(0, N)])
+
+  partial = hl.Func("partial")
+  partial[i, j] = 0.0
+  partial[i, j] += A[i, k] * B[k, j]
+
+  C = hl.Func("C")
+  C[i, j] = partial[i, j]
+
+  if not gpu:
+    # ---------------- Schedule ----------------
+    VEC = 16
+    TILE_I = 64
+    TILE_J = 64
+
+    io, jo, ii, ji = hl.Var("io"), hl.Var("jo"), hl.Var("ii"), hl.Var("ji")
+    C.update().tile(i, j, io, jo, ii, ji, TILE_I, TILE_J).fuse(io, jo, io).parallel(io).vectorize(ji, VEC)
+  else:
+    # ---------------- Schedule ----------------
+    GRP_I    = 8     # output tile size
+    GRP_J    = 16
+
+    #partial.store_in(hl.MemoryType.Register)
+    #partial.update().unroll(k, 4)
+
+    io, jo, ii, ji = hl.Var(), hl.Var(), hl.Var(), hl.Var()
+    C.gpu_tile(i, j, io, jo, ii, ji, GRP_I, GRP_J, hl.TailStrategy.RoundUp)
+
+  return C, A, B
+
+if __name__ == "__main__":
+  pipe, A, B = gemm_pipeline(gpu=True)
+
+  # NOTE: meteal does nothing
+  target = hl.get_host_target().with_feature(hl.TargetFeature.Metal)
+
+  a_np = np.random.randn(N, N).astype(np.float32)
+  b_np = np.random.randn(N, N).astype(np.float32)
+
+  # reverse order is correct!
+  a_hal = hl.Buffer(b_np)
+  b_hal = hl.Buffer(a_np)
+  A.set(a_hal)
+  B.set(b_hal)
+
+  pipe.compile_to_lowered_stmt("/tmp/my_function.html", [A, B], hl.StmtOutputFormat.HTML, target=target)
+  #exit(0)
+
+  c_hal = hl.Buffer(hl.Float(32), [N,N])
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+    c_out = np.array(c_hal)
+  print(c_out)
+
+  # tinygrad gets 60 ms with no BEAM, 20 ms with BEAM on CPU
+  with Timing("halide gemm "):
+    pipe.realize(c_hal, target)
+    c_hal.copy_to_host()
+
+  # Check correctness
+  with Timing("numpy gemm "):
+    ref = a_np @ b_np
+  max_err = np.abs(ref - c_out).max()
+  print("Max absolute error:", max_err)
+  assert max_err < 1e-4, "GEMM result incorrect!"
+
+  print("Pipeline ran on", target)
+  print("Success - GEMM Halide-Python output matches NumPy.")

extra/gemm/max_matmul.py

@@ -4,9 +4,9 @@ from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
 
 # for copied uops
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker

extra/gemm/simple_matmul.py

@@ -2,7 +2,7 @@ import numpy as np
 from tinygrad import dtypes, Tensor
 from tinygrad.helpers import getenv, get_single_element
 from tinygrad.dtype import _to_np_dtype
-from tinygrad.codegen.kernel import OptOps
+from tinygrad.opt.kernel import OptOps
 from tinygrad.engine.realize import lower_schedule
 
 dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float

extra/gemm/tinygrad_nv_matmul.py

@@ -1,6 +1,6 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from dataclasses import replace
 

extra/gemm/tvm_gemm.py

@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)
 
 sched = C.schedule()
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
 lin.linearize()

extra/hcqfuzz/readme

@@ -4,9 +4,9 @@ To add a new test, define a `TestSpec`-based class in a file in the `tests/` fol
 
 You can choose which tests to load from which file:
 ```bash
-RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
 ```
 Or skip tests from any file:
 ```bash
-SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
+PYTHONPATH=. SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
 ```

extra/mcts_search.py

@@ -4,9 +4,9 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.engine.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 
 class MCTSNode:
   def __init__(self, kernel:Kernel, parent=None):

extra/models/llama.py

@@ -1,5 +1,5 @@
 from typing import Union, Optional, Any
-import collections
+import collections, math
 from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
 from tinygrad.helpers import getenv, DEBUG
 
@@ -166,27 +166,29 @@ def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):
 
 class Transformer:
   def __init__(self, dim:int, hidden_dim:int, n_heads:int, n_layers:int, norm_eps:float, vocab_size, linear=nn.Linear, embedding=nn.Embedding,
-               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None):
-    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, max_context, linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
+               n_kv_heads=None, rope_theta=10000, max_context=1024, jit=True, feed_forward=FeedForward, qk_norm=None, disable_kv_cache=False):
+    self.layers = [TransformerBlock(dim, hidden_dim, n_heads, n_kv_heads, norm_eps, 0 if disable_kv_cache else max_context,
+                                    linear, feed_forward=feed_forward, qk_norm=qk_norm) for _ in range(n_layers)]
     self.norm = nn.RMSNorm(dim, norm_eps)
     self.tok_embeddings = embedding(vocab_size, dim)
     self.output = nn.Linear(dim, vocab_size, bias=False) if embedding == nn.Embedding else linear(dim, vocab_size, bias=False)
     self.max_context = max_context
-    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous()
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous().requires_grad_(False)
     self.forward_jit = TinyJit(self.forward) if jit else None
 
   def forward(self, tokens:Tensor, start_pos:Union[Variable,int], temperature:float, top_k:int, top_p:float, alpha_f:float, alpha_p:float):
     _bsz, seqlen = tokens.shape
     h = self.tok_embeddings(tokens)
 
-    self.freqs_cis = self.freqs_cis.cast(h.dtype).kernelize()
+    self.freqs_cis = self.freqs_cis.cast(h.dtype).contiguous()
     freqs_cis = self.freqs_cis[:, start_pos:start_pos+seqlen, :, :, :]
 
-    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1).kernelize() if seqlen > 1 else None
+    mask = Tensor.full((1, 1, seqlen, start_pos+seqlen), float("-inf"), dtype=h.dtype, device=h.device).triu(start_pos+1) if seqlen > 1 else None
     for layer in self.layers: h = layer(h, start_pos, freqs_cis, mask)
     logits = self.output(self.norm(h)).float()[:, -1, :]
+    if math.isnan(temperature): return logits
 
-    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p).kernelize()
+    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p)
 
   def __call__(self, tokens:Tensor, start_pos:int, temperature:float=0.0, top_k:int=0, top_p:float=0.8, alpha_f:float=0.0, alpha_p:float=0.0):
     # TODO: better way to handle the first call v.s. the rest?

extra/nv_gpu_driver/g_rpc-message-header.h

@@ -0,0 +1,77 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2008-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+/*
+ * WARNING: This is an autogenerated file. DO NOT EDIT.
+ * This file is generated using below files:
+ * template file: inc/kernel/vgpu/gt_rpc-message.h
+ * definition file: inc/kernel/vgpu/rpc-message-header.def
+ */
+
+typedef struct GSP_MSG_QUEUE_ELEMENT
+{
+    NvU8  authTagBuffer[16];         // Authentication tag buffer.
+    NvU8  aadBuffer[16];             // AAD buffer.
+    NvU32 checkSum;                  // Set to value needed to make checksum always zero.
+    NvU32 seqNum;                    // Sequence number maintained by the message queue.
+    NvU32 elemCount;                 // Number of message queue elements this message has.
+    NvU32 padding;                  // Reserved for future use.
+} GSP_MSG_QUEUE_ELEMENT;
+
+#ifdef RPC_MESSAGE_STRUCTURES
+typedef union rpc_message_rpc_union_field_v03_00
+{
+    NvU32      spare;
+    NvU32      cpuRmGfid;
+} rpc_message_rpc_union_field_v03_00;
+
+typedef rpc_message_rpc_union_field_v03_00 rpc_message_rpc_union_field_v;
+
+typedef struct rpc_message_header_v03_00
+{
+    NvU32      header_version;
+    NvU32      signature;
+    NvU32      length;
+    NvU32      function;
+    NvU32      rpc_result;
+    NvU32      rpc_result_private;
+    NvU32      sequence;
+    rpc_message_rpc_union_field_v u;
+    // rpc_generic_union rpc_message_data[];
+} rpc_message_header_v03_00;
+
+typedef rpc_message_header_v03_00 rpc_message_header_v;
+
+
+#endif
+
+#ifdef RPC_MESSAGE_GENERIC_UNION
+// This is a generic union, that will be used for the communication between the vmioplugin & guest RM.
+typedef union rpc_message_generic_union {
+    rpc_message_rpc_union_field_v03_00 rpc_union_field_v03_00;
+    rpc_message_rpc_union_field_v rpc_union_field_v;
+    rpc_message_header_v03_00 header_v03_00;
+    rpc_message_header_v header_v;
+} rpc_message_generic_union;
+
+#endif

extra/nv_gpu_driver/gsp_static_config.h

@@ -0,0 +1,455 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef GSP_STATIC_CONFIG_H
+#define GSP_STATIC_CONFIG_H
+
+//
+// This header describes the set of static GPU configuration information
+// that is collected during GSP RM init and made available to the
+// CPU RM (aka GSP client) via NV_RM_RPC_GET_GSP_STATIC_INFO() call.
+
+#include "ctrl/ctrl0080/ctrl0080gpu.h"
+#include "ctrl/ctrl2080/ctrl2080bios.h"
+#include "ctrl/ctrl2080/ctrl2080fb.h"
+#include "ctrl/ctrl2080/ctrl2080gpu.h"
+
+#include "vgpu/rpc_headers.h"
+#include "nvacpitypes.h"
+
+#include "ctrl/ctrl0073/ctrl0073system.h"
+
+#define MAX_DSM_SUPPORTED_FUNCS_RTN_LEN 8 // # bytes to store supported functions
+#define NV_ACPI_GENERIC_FUNC_COUNT                  8
+
+#define REGISTRY_TABLE_ENTRY_TYPE_UNKNOWN  0
+#define REGISTRY_TABLE_ENTRY_TYPE_DWORD    1
+#define REGISTRY_TABLE_ENTRY_TYPE_BINARY   2
+#define REGISTRY_TABLE_ENTRY_TYPE_STRING   3
+typedef struct PACKED_REGISTRY_ENTRY
+{
+    NvU32                   nameOffset;
+    NvU8                    type;
+    NvU32                   data;
+    NvU32                   length;
+} PACKED_REGISTRY_ENTRY;
+
+typedef struct PACKED_REGISTRY_TABLE
+{
+    NvU32                   size;
+    NvU32                   numEntries;
+} PACKED_REGISTRY_TABLE;
+
+/* Indicates the current state of mux */
+typedef enum
+{
+    dispMuxState_None = 0,
+    dispMuxState_IntegratedGPU,
+    dispMuxState_DiscreteGPU,
+} DISPMUXSTATE;
+
+typedef struct {
+    // supported function status and cache
+    NvU32  suppFuncStatus;
+    NvU8   suppFuncs[MAX_DSM_SUPPORTED_FUNCS_RTN_LEN];
+    NvU32  suppFuncsLen;
+    NvBool bArg3isInteger;
+    // callback status and cache
+    NvU32  callbackStatus;
+    NvU32  callback;
+} ACPI_DSM_CACHE;
+
+typedef struct {
+
+    ACPI_DSM_CACHE                   dsm[ACPI_DSM_FUNCTION_COUNT];
+    ACPI_DSM_FUNCTION                dispStatusHotplugFunc;
+    ACPI_DSM_FUNCTION                dispStatusConfigFunc;
+    ACPI_DSM_FUNCTION                perfPostPowerStateFunc;
+    ACPI_DSM_FUNCTION                stereo3dStateActiveFunc;
+    NvU32                            dsmPlatCapsCache[ACPI_DSM_FUNCTION_COUNT];
+    NvU32                            MDTLFeatureSupport;
+
+    // cache of generic func/subfunction remappings.
+    ACPI_DSM_FUNCTION                dsmCurrentFunc[NV_ACPI_GENERIC_FUNC_COUNT];
+    NvU32                            dsmCurrentSubFunc[NV_ACPI_GENERIC_FUNC_COUNT];
+    NvU32                            dsmCurrentFuncSupport;
+
+} ACPI_DATA;
+
+typedef struct DOD_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     acpiIdListLen;
+    NvU32     acpiIdList[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+} DOD_METHOD_DATA;
+
+typedef struct JT_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     jtCaps;
+    NvU16     jtRevId;
+    NvBool    bSBIOSCaps;
+} JT_METHOD_DATA;
+
+typedef struct MUX_METHOD_DATA_ELEMENT
+{
+    NvU32       acpiId;
+    NvU32       mode;
+    NV_STATUS   status;
+} MUX_METHOD_DATA_ELEMENT;
+
+typedef struct MUX_METHOD_DATA
+{
+    NvU32                       tableLen;
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxModeTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxPartTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];
+    MUX_METHOD_DATA_ELEMENT     acpiIdMuxStateTable[NV0073_CTRL_SYSTEM_ACPI_ID_MAP_MAX_DISPLAYS];    
+} MUX_METHOD_DATA;
+
+typedef struct CAPS_METHOD_DATA
+{
+    NV_STATUS status;
+    NvU32     optimusCaps;
+} CAPS_METHOD_DATA;
+
+typedef struct ACPI_METHOD_DATA
+{
+    NvBool                                               bValid;
+    DOD_METHOD_DATA                                      dodMethodData;
+    JT_METHOD_DATA                                       jtMethodData;
+    MUX_METHOD_DATA                                      muxMethodData;
+    CAPS_METHOD_DATA                                     capsMethodData;
+} ACPI_METHOD_DATA;
+
+#define MAX_GROUP_COUNT 2
+
+// #include "gpu/nvbitmask.h"
+typedef enum
+{
+    RM_ENGINE_TYPE_NULL                            =       (0x00000000),
+    RM_ENGINE_TYPE_GR0                             =       (0x00000001),
+    RM_ENGINE_TYPE_GR1                             =       (0x00000002),
+    RM_ENGINE_TYPE_GR2                             =       (0x00000003),
+    RM_ENGINE_TYPE_GR3                             =       (0x00000004),
+    RM_ENGINE_TYPE_GR4                             =       (0x00000005),
+    RM_ENGINE_TYPE_GR5                             =       (0x00000006),
+    RM_ENGINE_TYPE_GR6                             =       (0x00000007),
+    RM_ENGINE_TYPE_GR7                             =       (0x00000008),
+    RM_ENGINE_TYPE_COPY0                           =       (0x00000009),
+    RM_ENGINE_TYPE_COPY1                           =       (0x0000000a),
+    RM_ENGINE_TYPE_COPY2                           =       (0x0000000b),
+    RM_ENGINE_TYPE_COPY3                           =       (0x0000000c),
+    RM_ENGINE_TYPE_COPY4                           =       (0x0000000d),
+    RM_ENGINE_TYPE_COPY5                           =       (0x0000000e),
+    RM_ENGINE_TYPE_COPY6                           =       (0x0000000f),
+    RM_ENGINE_TYPE_COPY7                           =       (0x00000010),
+    RM_ENGINE_TYPE_COPY8                           =       (0x00000011),
+    RM_ENGINE_TYPE_COPY9                           =       (0x00000012),
+    RM_ENGINE_TYPE_COPY10                          =       (0x00000013),
+    RM_ENGINE_TYPE_COPY11                          =       (0x00000014),
+    RM_ENGINE_TYPE_COPY12                          =       (0x00000015),
+    RM_ENGINE_TYPE_COPY13                          =       (0x00000016),
+    RM_ENGINE_TYPE_COPY14                          =       (0x00000017),
+    RM_ENGINE_TYPE_COPY15                          =       (0x00000018),
+    RM_ENGINE_TYPE_COPY16                          =       (0x00000019),
+    RM_ENGINE_TYPE_COPY17                          =       (0x0000001a),
+    RM_ENGINE_TYPE_COPY18                          =       (0x0000001b),
+    RM_ENGINE_TYPE_COPY19                          =       (0x0000001c),
+    RM_ENGINE_TYPE_NVDEC0                          =       (0x0000001d),
+    RM_ENGINE_TYPE_NVDEC1                          =       (0x0000001e),
+    RM_ENGINE_TYPE_NVDEC2                          =       (0x0000001f),
+    RM_ENGINE_TYPE_NVDEC3                          =       (0x00000020),
+    RM_ENGINE_TYPE_NVDEC4                          =       (0x00000021),
+    RM_ENGINE_TYPE_NVDEC5                          =       (0x00000022),
+    RM_ENGINE_TYPE_NVDEC6                          =       (0x00000023),
+    RM_ENGINE_TYPE_NVDEC7                          =       (0x00000024),
+    RM_ENGINE_TYPE_NVENC0                          =       (0x00000025),
+    RM_ENGINE_TYPE_NVENC1                          =       (0x00000026),
+    RM_ENGINE_TYPE_NVENC2                          =       (0x00000027),
+    // Bug 4175886 - Use this new value for all chips once GB20X is released
+    RM_ENGINE_TYPE_NVENC3                          =       (0x00000028),
+    RM_ENGINE_TYPE_VP                              =       (0x00000029),
+    RM_ENGINE_TYPE_ME                              =       (0x0000002a),
+    RM_ENGINE_TYPE_PPP                             =       (0x0000002b),
+    RM_ENGINE_TYPE_MPEG                            =       (0x0000002c),
+    RM_ENGINE_TYPE_SW                              =       (0x0000002d),
+    RM_ENGINE_TYPE_TSEC                            =       (0x0000002e),
+    RM_ENGINE_TYPE_VIC                             =       (0x0000002f),
+    RM_ENGINE_TYPE_MP                              =       (0x00000030),
+    RM_ENGINE_TYPE_SEC2                            =       (0x00000031),
+    RM_ENGINE_TYPE_HOST                            =       (0x00000032),
+    RM_ENGINE_TYPE_DPU                             =       (0x00000033),
+    RM_ENGINE_TYPE_PMU                             =       (0x00000034),
+    RM_ENGINE_TYPE_FBFLCN                          =       (0x00000035),
+    RM_ENGINE_TYPE_NVJPEG0                         =       (0x00000036),
+    RM_ENGINE_TYPE_NVJPEG1                         =       (0x00000037),
+    RM_ENGINE_TYPE_NVJPEG2                         =       (0x00000038),
+    RM_ENGINE_TYPE_NVJPEG3                         =       (0x00000039),
+    RM_ENGINE_TYPE_NVJPEG4                         =       (0x0000003a),
+    RM_ENGINE_TYPE_NVJPEG5                         =       (0x0000003b),
+    RM_ENGINE_TYPE_NVJPEG6                         =       (0x0000003c),
+    RM_ENGINE_TYPE_NVJPEG7                         =       (0x0000003d),
+    RM_ENGINE_TYPE_OFA0                            =       (0x0000003e),
+    RM_ENGINE_TYPE_OFA1                            =       (0x0000003f),
+    RM_ENGINE_TYPE_RESERVED40                      =       (0x00000040),
+    RM_ENGINE_TYPE_RESERVED41                      =       (0x00000041),
+    RM_ENGINE_TYPE_RESERVED42                      =       (0x00000042),
+    RM_ENGINE_TYPE_RESERVED43                      =       (0x00000043),
+    RM_ENGINE_TYPE_RESERVED44                      =       (0x00000044),
+    RM_ENGINE_TYPE_RESERVED45                      =       (0x00000045),
+    RM_ENGINE_TYPE_RESERVED46                      =       (0x00000046),
+    RM_ENGINE_TYPE_RESERVED47                      =       (0x00000047),
+    RM_ENGINE_TYPE_RESERVED48                      =       (0x00000048),
+    RM_ENGINE_TYPE_RESERVED49                      =       (0x00000049),
+    RM_ENGINE_TYPE_RESERVED4a                      =       (0x0000004a),
+    RM_ENGINE_TYPE_RESERVED4b                      =       (0x0000004b),
+    RM_ENGINE_TYPE_RESERVED4c                      =       (0x0000004c),
+    RM_ENGINE_TYPE_RESERVED4d                      =       (0x0000004d),
+    RM_ENGINE_TYPE_RESERVED4e                      =       (0x0000004e),
+    RM_ENGINE_TYPE_RESERVED4f                      =       (0x0000004f),
+    RM_ENGINE_TYPE_RESERVED50                      =       (0x00000050),
+    RM_ENGINE_TYPE_RESERVED51                      =       (0x00000051),
+    RM_ENGINE_TYPE_RESERVED52                      =       (0x00000052),
+    RM_ENGINE_TYPE_RESERVED53                      =       (0x00000053),
+    RM_ENGINE_TYPE_LAST                            =       (0x00000054),
+} RM_ENGINE_TYPE;
+
+//
+// The duplicates in the RM_ENGINE_TYPE. Using define instead of putting them
+// in the enum to make sure that each item in the enum has a unique number.
+//
+#define RM_ENGINE_TYPE_GRAPHICS                 RM_ENGINE_TYPE_GR0
+#define RM_ENGINE_TYPE_BSP                      RM_ENGINE_TYPE_NVDEC0
+#define RM_ENGINE_TYPE_MSENC                    RM_ENGINE_TYPE_NVENC0
+#define RM_ENGINE_TYPE_CIPHER                   RM_ENGINE_TYPE_TSEC
+#define RM_ENGINE_TYPE_NVJPG                    RM_ENGINE_TYPE_NVJPEG0
+
+#define RM_ENGINE_TYPE_COPY_SIZE 20
+// Bug 4175886 - Use this new value for all chips once GB20X is released
+#define RM_ENGINE_TYPE_NVENC_SIZE 4
+#define RM_ENGINE_TYPE_NVJPEG_SIZE 8
+#define RM_ENGINE_TYPE_NVDEC_SIZE 8
+#define RM_ENGINE_TYPE_OFA_SIZE 2
+#define RM_ENGINE_TYPE_GR_SIZE 8
+
+#define NVGPU_ENGINE_CAPS_MASK_BITS                32
+#define NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX           ((RM_ENGINE_TYPE_LAST-1)/NVGPU_ENGINE_CAPS_MASK_BITS + 1)
+#define NVGPU_GET_ENGINE_CAPS_MASK(caps, id)       (caps[(id)/NVGPU_ENGINE_CAPS_MASK_BITS] & NVBIT((id) % NVGPU_ENGINE_CAPS_MASK_BITS))
+#define NVGPU_SET_ENGINE_CAPS_MASK(caps, id)       (caps[(id)/NVGPU_ENGINE_CAPS_MASK_BITS] |= NVBIT((id) % NVGPU_ENGINE_CAPS_MASK_BITS))
+
+
+// #include "gpu/gpu.h" // COMPUTE_BRANDING_TYPE
+// #include "gpu/gpu_acpi_data.h" // ACPI_METHOD_DATA
+// #include "vgpu/rpc_headers.h" // MAX_GPC_COUNT
+// #include "platform/chipset/chipset.h" // BUSINFO
+// #include "gpu/nvbitmask.h" // NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX
+
+typedef struct
+{
+    NvU16               deviceID;           // deviceID
+    NvU16               vendorID;           // vendorID
+    NvU16               subdeviceID;        // subsystem deviceID
+    NvU16               subvendorID;        // subsystem vendorID
+    NvU8                revisionID;         // revision ID
+} BUSINFO;
+
+// VF related info for GSP-RM
+typedef struct GSP_VF_INFO
+{
+    NvU32  totalVFs;
+    NvU32  firstVFOffset;
+    NvU64  FirstVFBar0Address;
+    NvU64  FirstVFBar1Address;
+    NvU64  FirstVFBar2Address;
+    NvBool b64bitBar0;
+    NvBool b64bitBar1;
+    NvBool b64bitBar2;
+} GSP_VF_INFO;
+
+// Cache config registers from pcie space
+typedef struct
+{
+    // Link capabilities
+    NvU32 linkCap;
+} GSP_PCIE_CONFIG_REG;
+
+typedef struct
+{
+    NvU32 ecidLow;
+    NvU32 ecidHigh;
+    NvU32 ecidExtended;
+} EcidManufacturingInfo;
+
+typedef struct
+{
+    NvU64 nonWprHeapOffset;
+    NvU64 frtsOffset;
+} FW_WPR_LAYOUT_OFFSET;
+
+// Fetched from GSP-RM into CPU-RM
+typedef struct GspStaticConfigInfo_t
+{
+    NvU8 grCapsBits[NV0080_CTRL_GR_CAPS_TBL_SIZE];
+    NV2080_CTRL_GPU_GET_GID_INFO_PARAMS gidInfo;
+    NV2080_CTRL_BIOS_GET_SKU_INFO_PARAMS SKUInfo;
+    NV2080_CTRL_CMD_FB_GET_FB_REGION_INFO_PARAMS fbRegionInfoParams;
+
+    NV0080_CTRL_GPU_GET_SRIOV_CAPS_PARAMS sriovCaps;
+    NvU32 sriovMaxGfid;
+
+    NvU32 engineCaps[NVGPU_ENGINE_CAPS_MASK_ARRAY_MAX];
+
+    NvBool poisonFuseEnabled;
+
+    NvU64 fb_length;
+    NvU64 fbio_mask;
+    NvU32 fb_bus_width;
+    NvU32 fb_ram_type;
+    NvU64 fbp_mask;
+    NvU32 l2_cache_size;
+
+    NvU8 gpuNameString[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvU8 gpuShortNameString[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvU16 gpuNameString_Unicode[NV2080_GPU_MAX_NAME_STRING_LENGTH];
+    NvBool bGpuInternalSku;
+    NvBool bIsQuadroGeneric;
+    NvBool bIsQuadroAd;
+    NvBool bIsNvidiaNvs;
+    NvBool bIsVgx;
+    NvBool bGeforceSmb;
+    NvBool bIsTitan;
+    NvBool bIsTesla;
+    NvBool bIsMobile;
+    NvBool bIsGc6Rtd3Allowed;
+    NvBool bIsGc8Rtd3Allowed;
+    NvBool bIsGcOffRtd3Allowed;
+    NvBool bIsGcoffLegacyAllowed;
+    NvBool bIsMigSupported;
+
+    /* "Total Board Power" refers to power requirement of GPU,
+     * while in GC6 state. Majority of this power will be used
+     * to keep V-RAM active to preserve its content.
+     * Some energy maybe consumed by Always-on components on GPU chip.
+     * This power will be provided by 3.3v voltage rail.
+     */
+    NvU16  RTD3GC6TotalBoardPower;
+
+    /* PERST# (i.e. PCI Express Reset) is a sideband signal
+     * generated by the PCIe Host to indicate the PCIe devices,
+     * that the power-rails and the reference-clock are stable.
+     * The endpoint device typically uses this signal as a global reset.
+     */
+    NvU16  RTD3GC6PerstDelay;
+
+    NvU64 bar1PdeBase;
+    NvU64 bar2PdeBase;
+
+    NvBool bVbiosValid;
+    NvU32 vbiosSubVendor;
+    NvU32 vbiosSubDevice;
+
+    NvBool bPageRetirementSupported;
+
+    NvBool bSplitVasBetweenServerClientRm;
+
+    NvBool bClRootportNeedsNosnoopWAR;
+
+    VIRTUAL_DISPLAY_GET_NUM_HEADS_PARAMS displaylessMaxHeads;
+    VIRTUAL_DISPLAY_GET_MAX_RESOLUTION_PARAMS displaylessMaxResolution;
+    NvU64 displaylessMaxPixels;
+
+    // Client handle for internal RMAPI control.
+    NvHandle hInternalClient;
+
+    // Device handle for internal RMAPI control.
+    NvHandle hInternalDevice;
+
+    // Subdevice handle for internal RMAPI control.
+    NvHandle hInternalSubdevice;
+
+    NvBool bSelfHostedMode;
+    NvBool bAtsSupported;
+
+    NvBool bIsGpuUefi;
+    NvBool bIsEfiInit;
+
+    EcidManufacturingInfo ecidInfo[MAX_GROUP_COUNT];
+
+    FW_WPR_LAYOUT_OFFSET fwWprLayoutOffset;
+} GspStaticConfigInfo;
+
+// Pushed from CPU-RM to GSP-RM
+typedef struct GspSystemInfo
+{
+    NvU64 gpuPhysAddr;
+    NvU64 gpuPhysFbAddr;
+    NvU64 gpuPhysInstAddr;
+    NvU64 gpuPhysIoAddr;
+    NvU64 nvDomainBusDeviceFunc;
+    NvU64 simAccessBufPhysAddr;
+    NvU64 notifyOpSharedSurfacePhysAddr;
+    NvU64 pcieAtomicsOpMask;
+    NvU64 consoleMemSize;
+    NvU64 maxUserVa;
+    NvU32 pciConfigMirrorBase;
+    NvU32 pciConfigMirrorSize;
+    NvU32 PCIDeviceID;
+    NvU32 PCISubDeviceID;
+    NvU32 PCIRevisionID;
+    NvU32 pcieAtomicsCplDeviceCapMask;
+    NvU8 oorArch;
+    NvU64 clPdbProperties;
+    NvU32 Chipset;
+    NvBool bGpuBehindBridge;
+    NvBool bFlrSupported;
+    NvBool b64bBar0Supported;
+    NvBool bMnocAvailable;
+    NvU32  chipsetL1ssEnable;
+    NvBool bUpstreamL0sUnsupported;
+    NvBool bUpstreamL1Unsupported;
+    NvBool bUpstreamL1PorSupported;
+    NvBool bUpstreamL1PorMobileOnly;
+    NvBool bSystemHasMux;
+    NvU8   upstreamAddressValid;
+    BUSINFO FHBBusInfo;
+    BUSINFO chipsetIDInfo;
+    ACPI_METHOD_DATA acpiMethodData;
+    NvU32 hypervisorType;
+    NvBool bIsPassthru;
+    NvU64 sysTimerOffsetNs;
+    GSP_VF_INFO gspVFInfo;
+    NvBool bIsPrimary;
+    NvBool isGridBuild;
+    GSP_PCIE_CONFIG_REG pcieConfigReg;
+    NvU32 gridBuildCsp;
+    NvBool bPreserveVideoMemoryAllocations;
+    NvBool bTdrEventSupported;
+    NvBool bFeatureStretchVblankCapable;
+    NvBool bEnableDynamicGranularityPageArrays;
+    NvBool bClockBoostSupported;
+    NvBool bRouteDispIntrsToCPU;
+    NvU64  hostPageSize;
+} GspSystemInfo;
+
+
+#endif /* GSP_STATIC_CONFIG_H */

extra/nv_gpu_driver/vbios.h

@@ -0,0 +1,209 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: MIT
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a
+ * copy of this software and associated documentation files (the "Software"),
+ * to deal in the Software without restriction, including without limitation
+ * the rights to use, copy, modify, merge, publish, distribute, sublicense,
+ * and/or sell copies of the Software, and to permit persons to whom the
+ * Software is furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+ * THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+ * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+ * DEALINGS IN THE SOFTWARE.
+ */
+
+#ifndef VBIOS_H
+#define VBIOS_H
+
+#include "gpu/vbios/bios_types.h"
+#define FALCON_APPLICATION_INTERFACE_ENTRY_ID_DMEMMAPPER     (0x4)
+
+typedef struct
+{
+    NvU32 signature;
+    NvU16 version;
+    NvU16 size;
+    NvU32 cmd_in_buffer_offset;
+    NvU32 cmd_in_buffer_size;
+    NvU32 cmd_out_buffer_offset;
+    NvU32 cmd_out_buffer_size;
+    NvU32 nvf_img_data_buffer_offset;
+    NvU32 nvf_img_data_buffer_size;
+    NvU32 printfBufferHdr;
+    NvU32 ucode_build_time_stamp;
+    NvU32 ucode_signature;
+    NvU32 init_cmd;
+    NvU32 ucode_feature;
+    NvU32 ucode_cmd_mask0;
+    NvU32 ucode_cmd_mask1;
+    NvU32 multiTgtTbl;
+} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3;
+
+#define FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3_CMD_FRTS (0x15)
+#define FALCON_APPLICATION_INTERFACE_DMEM_MAPPER_V3_CMD_SB   (0x19)
+
+#define BIT_HEADER_ID                     0xB8FF
+#define BIT_HEADER_SIGNATURE              0x00544942  // "BIT\0"
+#define BIT_HEADER_SIZE_OFFSET            8
+
+struct __attribute__((packed)) BIT_HEADER_V1_00
+{
+    unsigned short Id;
+    unsigned int Signature;
+    unsigned short BCD_Version;
+    unsigned char HeaderSize;
+    unsigned char TokenSize;
+    unsigned char TokenEntries;
+    unsigned char HeaderChksum;
+};
+#define BIT_HEADER_V1_00_FMT "1w1d1w4b"
+typedef struct BIT_HEADER_V1_00 BIT_HEADER_V1_00;
+
+struct __attribute__((packed)) BIT_TOKEN_V1_00
+{
+    unsigned char TokenId;
+    unsigned char DataVersion;
+    unsigned short DataSize;
+    unsigned int DataPtr;
+};
+
+#define BIT_TOKEN_V1_00_SIZE_6     6U
+#define BIT_TOKEN_V1_00_SIZE_8     8U
+
+#define BIT_TOKEN_V1_00_FMT_SIZE_6 "2b2w"
+#define BIT_TOKEN_V1_00_FMT_SIZE_8 "2b1w1d"
+typedef struct BIT_TOKEN_V1_00 BIT_TOKEN_V1_00;
+
+#define BIT_TOKEN_BIOSDATA          0x42
+
+// structure for only version info from BIT_DATA_BIOSDATA_V1 and BIT_DATA_BIOSDATA_V2
+typedef struct
+{
+    unsigned int Version;     // BIOS Binary Version Ex. 5.40.00.01.12 = 0x05400001
+    unsigned char OemVersion;  // OEM Version Number  Ex. 5.40.00.01.12 = 0x12
+} __attribute__((packed)) BIT_DATA_BIOSDATA_BINVER;
+
+#define BIT_DATA_BIOSDATA_VERSION_1         0x1
+#define BIT_DATA_BIOSDATA_VERSION_2         0x2
+
+#define BIT_DATA_BIOSDATA_BINVER_FMT "1d1b"
+#define BIT_DATA_BIOSDATA_BINVER_SIZE_5    5
+
+#define BIT_TOKEN_FALCON_DATA       0x70
+
+typedef struct
+{
+    unsigned int FalconUcodeTablePtr;
+} __attribute__((packed)) BIT_DATA_FALCON_DATA_V2;
+
+#define BIT_DATA_FALCON_DATA_V2_4_FMT       "1d"
+#define BIT_DATA_FALCON_DATA_V2_SIZE_4      4
+
+typedef struct
+{
+    unsigned char Version;
+    unsigned char HeaderSize;
+    unsigned char EntrySize;
+    unsigned char EntryCount;
+    unsigned char DescVersion;
+    unsigned char DescSize;
+} __attribute__((packed)) FALCON_UCODE_TABLE_HDR_V1;
+
+#define FALCON_UCODE_TABLE_HDR_V1_VERSION   1
+#define FALCON_UCODE_TABLE_HDR_V1_SIZE_6    6
+#define FALCON_UCODE_TABLE_HDR_V1_6_FMT     "6b"
+
+typedef struct
+{
+    unsigned char ApplicationID;
+    unsigned char TargetID;
+    unsigned int DescPtr;
+} __attribute__((packed)) FALCON_UCODE_TABLE_ENTRY_V1;
+
+#define FALCON_UCODE_TABLE_ENTRY_V1_VERSION             1
+#define FALCON_UCODE_TABLE_ENTRY_V1_SIZE_6              6
+#define FALCON_UCODE_TABLE_ENTRY_V1_6_FMT               "2b1d"
+
+#define FALCON_UCODE_ENTRY_APPID_FIRMWARE_SEC_LIC       0x05
+#define FALCON_UCODE_ENTRY_APPID_FWSEC_DBG              0x45
+#define FALCON_UCODE_ENTRY_APPID_FWSEC_PROD             0x85
+
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION                0:0
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_UNAVAILABLE    0x00
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_VERSION_AVAILABLE      0x01
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_RESERVED               1:1
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_FLAGS_ENCRYPTED              2:2
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_RESERVED                     7:3
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION                      15:8
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V1                   0x01
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V2                   0x02
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V3                   0x03
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_VERSION_V4                   0x04
+#define NV_BIT_FALCON_UCODE_DESC_HEADER_VDESC_SIZE                         31:16
+
+typedef struct
+{
+    unsigned int vDesc;
+} __attribute__((packed)) FALCON_UCODE_DESC_HEADER;
+#define FALCON_UCODE_DESC_HEADER_FORMAT   "1d"
+
+typedef struct {
+    FALCON_UCODE_DESC_HEADER Hdr;
+    unsigned int StoredSize;
+    unsigned int PKCDataOffset;
+    unsigned int InterfaceOffset;
+    unsigned int IMEMPhysBase;
+    unsigned int IMEMLoadSize;
+    unsigned int IMEMVirtBase;
+    unsigned int DMEMPhysBase;
+    unsigned int DMEMLoadSize;
+    unsigned short EngineIdMask;
+    unsigned char UcodeId;
+    unsigned char SignatureCount;
+    unsigned short SignatureVersions;
+    unsigned short Reserved;
+} FALCON_UCODE_DESC_V3;
+
+#define FALCON_UCODE_DESC_V3_SIZE_44    44
+#define FALCON_UCODE_DESC_V3_44_FMT     "9d1w2b2w"
+#define BCRT30_RSA3K_SIG_SIZE 384
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU64 gfwImageOffset;
+    NvU32 gfwImageSize;
+    NvU32 flags;
+} __attribute__((packed)) FWSECLIC_READ_VBIOS_DESC;
+
+#define FWSECLIC_READ_VBIOS_STRUCT_FLAGS (2)
+
+typedef struct
+{
+    NvU32 version;
+    NvU32 size;
+    NvU32 frtsRegionOffset4K;
+    NvU32 frtsRegionSize;
+    NvU32 frtsRegionMediaType;
+} __attribute__((packed)) FWSECLIC_FRTS_REGION_DESC;
+
+#define FWSECLIC_FRTS_REGION_MEDIA_FB (2)
+#define FWSECLIC_FRTS_REGION_SIZE_1MB_IN_4K (0x100)
+
+typedef struct
+{
+    FWSECLIC_READ_VBIOS_DESC readVbiosDesc;
+    FWSECLIC_FRTS_REGION_DESC frtsRegionDesc;
+} __attribute__((packed)) FWSECLIC_FRTS_CMD;
+
+#endif /* VBIOS_H */

extra/nvpci/gpu_reset.sh

@@ -0,0 +1,2 @@
+GPU="$1"
+echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null

extra/onnx.py

@@ -1,11 +1,28 @@
 from types import SimpleNamespace
 from typing import Any, Sequence, cast, Literal, Callable
-import dataclasses, functools, io, math, types
+import dataclasses, functools, io, math, types, warnings, sys
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
 from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort
-from tinygrad.dtype import DType, ConstType, dtypes, ImageDType
+from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
 from tinygrad.device import is_dtype_supported, Device
 
+# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
+data_types: dict[int, DType] = {
+  1:dtypes.float32, 2:dtypes.uint8, 3:dtypes.int8, 4:dtypes.uint16, 5:dtypes.int16, 6:dtypes.int32, 7:dtypes.int64,
+  9:dtypes.bool, 10:dtypes.float16, 11:dtypes.double, 12:dtypes.uint32, 13:dtypes.uint64, 16:dtypes.bfloat16,
+}
+
+# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L128-L145
+attribute_types: dict[int, Callable] = {
+  1: lambda a: float(a.f),
+  2: lambda a: int(a.i),
+  3: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
+  4: lambda a: buffer_parse(a.t),
+  6: lambda a: tuple(float(x) for x in a.floats),
+  7: lambda a: tuple(int(x) for x in a.ints),
+  8: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
+}
+
 # ***** protobuf parsing ******
 from onnx import AttributeProto, ModelProto, TensorProto, TypeProto, helper
 import numpy as np
@@ -14,39 +31,24 @@ def has_field(onnx_type: TypeProto|SimpleNamespace, field):
   if isinstance(onnx_type, TypeProto): return onnx_type.HasField(field)
   return hasattr(onnx_type, field)
 
-def dtype_parse(onnx_dtype: int) -> DType:
-  supported: dict[int, DType] = {
-    TensorProto.FLOAT:dtypes.float32, TensorProto.UINT8:dtypes.uint8, TensorProto.INT8:dtypes.int8,
-    TensorProto.UINT16:dtypes.uint16, TensorProto.INT16:dtypes.int16, TensorProto.INT32:dtypes.int32, TensorProto.INT64:dtypes.int64,
-    TensorProto.BOOL:dtypes.bool, TensorProto.FLOAT16:dtypes.float32, TensorProto.DOUBLE:dtypes.double, TensorProto.UINT32:dtypes.uint32,
-    TensorProto.UINT64:dtypes.uint64, TensorProto.BFLOAT16:dtypes.bfloat16,
-  }
-  unsupported = {
-    TensorProto.UNDEFINED, TensorProto.STRING, TensorProto.COMPLEX64, TensorProto.COMPLEX128, TensorProto.FLOAT8E4M3FN, TensorProto.FLOAT8E4M3FNUZ,
-    TensorProto.FLOAT8E5M2, TensorProto.FLOAT8E5M2FNUZ, TensorProto.UINT4, TensorProto.INT4
-  }
-  if onnx_dtype in unsupported: raise NotImplementedError(f"onnx dtype {TensorProto.DataType.Name(onnx_dtype)} is not supported")
-  return supported[onnx_dtype] if is_dtype_supported(supported[onnx_dtype]) else dtypes.float
+def dtype_parse(onnx_dtype: int, fallback_context: str | None = None) -> DType:
+  if onnx_dtype not in data_types: raise NotImplementedError(f"onnx dtype id {onnx_dtype} is not supported")
+  if is_dtype_supported(dtype := data_types[onnx_dtype]): return dtype
+  # if fallback_context is provided, we can fall back to a default dtype
+  if fallback_context is not None:
+    default_dtype = dtypes.default_int if dtypes.is_int(dtype) else dtypes.default_float
+    warnings.warn(f"dtype {dtype} on {Device.DEFAULT} from {fallback_context} is not supported, falling back to {default_dtype}")
+    assert is_dtype_supported(default_dtype), f"dtype {default_dtype} must be supported on {Device.DEFAULT}"
+    return default_dtype
+  raise RuntimeError(f"dtype {dtype} on device {Device.DEFAULT} is not supported")
 
 def attribute_parse(onnx_attribute: AttributeProto):
-  supported: dict[AttributeProto.AttributeType, Callable[[AttributeProto], Any]] = {
-    AttributeProto.FLOAT: lambda a: float(a.f), AttributeProto.INT: lambda a: int(a.i),
-    AttributeProto.STRING: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
-    AttributeProto.TENSOR: lambda a: buffer_parse(a.t),
-    AttributeProto.FLOATS: lambda a: tuple(float(x) for x in a.floats), AttributeProto.INTS: lambda a: tuple(int(x) for x in a.ints),
-    AttributeProto.STRINGS: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
-  }
-  unsupported = {
-    AttributeProto.UNDEFINED, AttributeProto.GRAPH, AttributeProto.SPARSE_TENSOR, AttributeProto.TYPE_PROTO, AttributeProto.TENSORS,
-    AttributeProto.GRAPHS, AttributeProto.SPARSE_TENSORS, AttributeProto.TYPE_PROTOS
-  }
-  if onnx_attribute.type in unsupported:
-    raise NotImplementedError(f"attribute with type {AttributeProto.AttributeType.Name(onnx_attribute.type)} is not supported")
-  return supported[onnx_attribute.type](onnx_attribute)
+  if onnx_attribute.type not in attribute_types: raise NotImplementedError(f"attribute type {onnx_attribute.type} is not supported")
+  return attribute_types[onnx_attribute.type](onnx_attribute)
 
 def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
   if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
-  dtype, shape = dtype_parse(onnx_tensor.data_type), tuple(onnx_tensor.dims)
+  dtype, shape = dtype_parse(onnx_tensor.data_type, "buffer parse"), tuple(onnx_tensor.dims)
   data = None
   if len(onnx_tensor.float_data): data = onnx_tensor.float_data
   elif len(onnx_tensor.int32_data): data = onnx_tensor.int32_data
@@ -57,13 +59,17 @@ def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
     if len(data) == 1: return Tensor(data.tolist()[0], dtype=dtype).reshape(shape)
     return data.cast(dtype).reshape(shape).to(Device.DEFAULT)
   if has_field(onnx_tensor, "raw_data"):
-    if onnx_tensor.data_type == TensorProto.FLOAT16:
-      np_buffer = np.frombuffer(onnx_tensor.raw_data.data().tobytes(),
+    raw_data = onnx_tensor.raw_data
+    if not isinstance(raw_data, Tensor): raw_data = Tensor(raw_data)
+    if not is_dtype_supported(data_types[onnx_tensor.data_type]):
+      np_buffer = np.frombuffer(raw_data.data().tobytes(),
                                 dtype=helper.tensor_dtype_to_np_dtype(onnx_tensor.data_type)).copy().reshape(shape)
       if np_buffer.size == 1: return Tensor(np_buffer.item(), dtype=dtype).reshape(shape)
       return Tensor(np_buffer, dtype=dtype)
-    ret = onnx_tensor.raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
-    if shape == (): ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
+    ret = raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
+    if shape == ():
+      if ret.dtype is dtypes.float16 and sys.version_info < (3, 12): ret = ret.cast(dtypes.float32)
+      ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
     return ret
   return Tensor(None)
 
@@ -76,7 +82,7 @@ def type_parse(onnx_type: TypeProto):
   if has_field(elem_type, "tensor_type"):
     shape = tuple(getattr(d, "dim_param", None) or getattr(d, "dim_value") for d in elem_type.tensor_type.shape.dim) \
       if has_field(elem_type.tensor_type, "shape") else None # test_identity_sequence_cpu
-    dtype = dtype_parse(elem_type.tensor_type.elem_type)
+    dtype = data_types[elem_type.tensor_type.elem_type]
     return OnnxValue(shape, dtype, is_optional, is_sequence)
   raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
 
@@ -143,17 +149,19 @@ class OnnxRunner:
 
   def _parse_input(self, name: str, value: Any, spec: OnnxValue):
     if spec.is_optional and value is None: return None
-    # TODO: need true float16 for dtype checking
     if spec.is_sequence:
-      if not isinstance(value, Sequence): raise RuntimeError(f"{name} received {value}, expected a sequence type")
+      if not isinstance(value, Sequence): raise RuntimeError(f"input {name} received {value}, expected a sequence type")
       sequence = [Tensor(v, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(v, Tensor) else v for v in value]
-      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for {name} sequence must be homogeneous")
+      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for input {name} sequence must be homogeneous")
+      if not all(t.dtype is spec.dtype for t in sequence): warnings.warn(f"Dtypes for input {name} sequence aren't all {spec.dtype}")
       return sequence
-    tensor = Tensor(value, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
+    dtype = _from_np_dtype(value.dtype) if str(type(value)) == "<class 'numpy.ndarray'>" else spec.dtype
+    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
+    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
     for dim, (onnx_dim, user_dim_input) in enumerate(zip(spec.shape, tensor.shape, strict=True)):
       if isinstance(onnx_dim, str):
         onnx_dim = self.variable_dims[onnx_dim] if onnx_dim in self.variable_dims else self.variable_dims.setdefault(onnx_dim, int(user_dim_input))
-      if user_dim_input != onnx_dim: raise RuntimeError(f"{name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
+      if user_dim_input != onnx_dim: raise RuntimeError(f"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
     return tensor
 
   def _dispatch_op(self, op, inps, opts):
@@ -167,8 +175,8 @@ class OnnxRunner:
       return real_fxn(*inps, **opts)
     raise NotImplementedError(f"{op=} not supported")
 
-  def get_empty_input_data(self, device:str|None=None) -> dict[str, Tensor]:
-    return {name:Tensor.empty(*spec.shape, device=device, dtype=spec.dtype) for name, spec in self.graph_inputs.items()}
+  def get_empty_input_data(self, device:str|None=None, dtype:DType|None=None) -> dict[str, Tensor]:
+    return {name:Tensor.empty(*spec.shape, device=device, dtype=dtype or spec.dtype) for name, spec in self.graph_inputs.items()}
 
   def __call__(self, inputs:dict[str, Any], debug=debug):
     for name, input_spec in self.graph_inputs.items():
@@ -284,7 +292,7 @@ def get_onnx_ops():
     raise ValueError(f"pixel_format={pixel_format!r} is not supported.")
 
   def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
-    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) if dtype is not None else x.dtype)
+    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype, "EyeLike op") if dtype is not None else x.dtype)
     return ret if x.size(0) == x.size(1) else ret.pad(tuple(None if d == ret.size(0) else (k, d-ret.shape[0]-k) for d in x.shape))
 
   def OptionalHasElement(x:Tensor|None=None): return Tensor(x is not None and x.numel() > 0)
@@ -317,7 +325,7 @@ def get_onnx_ops():
   def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()
 
   # ***** Unary Ops (broadcasted) *****
-  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y if x.dtype == dtypes.float or isinstance(x.dtype, ImageDType) else (x + y).cast(x.dtype)
+  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
   def Sub(x:Tensor|int,y:Tensor): return x - y # some test has input as int
   def Div(x:Tensor,y:Tensor): return x.div(y, rounding_mode='trunc' if dtypes.is_int(x.dtype) else None)
   def Less(x:Tensor,y:Tensor): return x < y
@@ -338,7 +346,7 @@ def get_onnx_ops():
 
   # ***** Casting Ops *****
   # TODO: saturate
-  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to))
+  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to, "Cast op"))
   def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
 
   # ***** Reduce Ops *****
@@ -605,9 +613,13 @@ def get_onnx_ops():
 
   # Reimplemented here because you need legacy RNG for passing ONNX tests.
   def Dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
-    if not training_mode: return data, Tensor.ones(data.shape, dtype=dtypes.bool)  # if mask is requested as output it will contain all True's.
-    mask = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)) >= ratio, requires_grad=False, device=data.device)
-    return data * mask * (1/(1.0 - ratio)), mask
+    if not training_mode: return data, data.full_like(True, dtype=dtypes.bool)
+    if seed is not None:
+      rand = Tensor(np.random.RandomState(seed).random(cast(tuple[int,...], data.shape)), requires_grad=False, dtype=data.dtype, device=data.device)
+    else:
+      rand = data.rand_like(requires_grad=False)
+    mask = rand >= ratio
+    return data * mask / (1.0 - ratio), mask
   # 6 with 'is_test' needed for https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
   def Dropout_6(data:Tensor, ratio:float=0.5, is_test=0): return Dropout_7(data, ratio, training_mode=not is_test)
   Dropout = {6:Dropout_6, 7:Dropout_7}
@@ -731,7 +743,9 @@ def get_onnx_ops():
 
   # ***** Quantization Ops *****
   def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
-    out_dtype = y_zero_point.dtype if isinstance(y_zero_point, Tensor) else dtype_parse(output_dtype) if output_dtype else dtypes.uint8
+    if isinstance(y_zero_point, Tensor): out_dtype = y_zero_point.dtype
+    elif output_dtype != 0: out_dtype = dtype_parse(output_dtype, "QuantizeLinear op")
+    else: out_dtype = dtypes.uint8
     y_scale, y_zero_point = _prepare_quantize(x, y_scale, y_zero_point, axis, block_size)
     if out_dtype == dtypes.uchar:
       # this appears to work in practice, at least for uchar out_dtype. it folds with the quantize stuff

extra/onnx_parser.py

@@ -25,6 +25,7 @@ class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6;
 PB_INFOS = {
   "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
   "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
+  # TODO: support uint64 parsing (11: "uint64_data") and double parsing (10: "double_data")
   "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
     13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
     5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},

extra/optimization/extract_dataset.py

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

extra/optimization/extract_policynet.py

@@ -5,9 +5,9 @@ from tinygrad.nn import Linear
 from tinygrad.tensor import Tensor
 from tinygrad.nn.optim import Adam
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.engine.search import actions
+from tinygrad.opt.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.helpers import getenv
 
 # stuff needed to unpack a kernel
@@ -17,7 +17,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 INNER = 256
 class PolicyNet:

extra/optimization/extract_sa_pairs.py

@@ -10,11 +10,11 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 # more stuff
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import actions
 from extra.optimization.helpers import lin_to_feats
 from extra.optimization.pretrain_valuenet import ValueNet
 from tinygrad.nn.optim import Adam

extra/optimization/generate_dataset.sh

@@ -6,24 +6,12 @@ export CAPTURE_PROCESS_REPLAY=1
 rm $LOGOPS
 test/external/process_replay/reset.py
 
-python3 -m pytest -n=auto test/ --ignore=test/unit --durations=20
-STEPS=3 python3 examples/hlb_cifar10.py
-WINO=1 STEPS=3 python3 examples/hlb_cifar10.py
-python3 examples/stable_diffusion.py --noshow
-python3 examples/llama.py --prompt "hello" --count 5
-python3 examples/gpt2.py --count 5
-HALF=1 python3 examples/gpt2.py --count 5
-python3 examples/beautiful_mnist.py
-python3 examples/beautiful_cartpole.py
-python3 examples/mlperf/model_spec.py
-python3 examples/yolov8.py ./test/models/efficientnet/Chicken.jpg
-examples/openpilot/go.sh
-JIT=2 BIG=1 MPS=1 pytest -n=auto test/ --ignore=test/test_fusion_op.py --ignore=test/test_linearizer_failures.py --ignore=test/test_gc.py --ignore=test/test_speed_v_torch.py --ignore=test/test_jit.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_gc.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_jit.py
-JIT=2 BIG=1 MPS=1 python -m pytest test/test_speed_v_torch.py
+CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
+GPU=1 python3 -m pytest test/test_tiny.py
 
 # extract, sort and uniq
 extra/optimization/extract_dataset.py
 sort -u /tmp/ops > /tmp/sops
-ls -lh /tmp/ops /tmp/sops	
+ls -lh /tmp/ops /tmp/sops
+# gzip -k /tmp/sops
+# mv /tmp/sops.gz extra/datasets/

extra/optimization/get_action_space.py

@@ -1,8 +1,8 @@
 import random
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
-from tinygrad.engine.search import actions
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm
 
 tactions = set()

extra/optimization/helpers.py

@@ -1,15 +1,16 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 from tinygrad.uop.ops import UOp, Ops, KernelInfo
 from tinygrad.dtype import dtypes, PtrDType
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
+from tinygrad.helpers import getenv
 inf, nan = float('inf'), float('nan')
 UOps = Ops
 
 # kernel unpacker
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 def ast_str_to_ast(ast_str:str) -> UOp: return eval(ast_str)
 def ast_str_to_lin(ast_str:str, opts=None): return Kernel(ast_str_to_ast(ast_str), opts=opts)
 def kern_str_to_lin(kern_str:str, opts=None):
@@ -26,7 +27,7 @@ from tinygrad.helpers import dedup, DEBUG
 def load_worlds(filter_reduce=True, filter_noimage=True, filter_novariable=True):
   fn = Path(__file__).parent.parent / "datasets/sops.gz"
   ast_strs = dedup(gzip.open(fn).read().decode('utf-8').strip().split("\n"))
-  assert len(ast_strs) > 5000, f"dataset size = {len(ast_strs)} is too small"
+  assert len(ast_strs) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
   if DEBUG >= 1: print(f"loaded {len(ast_strs)=} before filters")
   if filter_reduce: ast_strs = [x for x in ast_strs if "REDUCE_AXIS" in x]
   if filter_noimage: ast_strs = [x for x in ast_strs if "dtypes.image" not in x]
@@ -101,7 +102,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   return ret
 
 from tinygrad.device import Device, Buffer
-from tinygrad.engine.search import _ensure_buffer_alloc, _time_program
+from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
 from tinygrad.helpers import to_function_name, CACHELEVEL, diskcache_get, diskcache_put
 
 def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_global_size=65536, cnt=3, disable_cache=False, clear_l2=False) -> float:  # noqa: E501

extra/optimization/pretrain_valuenet.py

@@ -1,4 +1,4 @@
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tqdm import tqdm, trange
 import math
 import random
@@ -14,7 +14,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
 from tinygrad.uop.ops import Variable
 inf, nan = float('inf'), float('nan')
-from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.opt.kernel import Opt, OptOps
 
 from extra.optimization.helpers import lin_to_feats, MAX_DIMS
 

extra/optimization/rl.py

@@ -3,7 +3,7 @@ import numpy as np
 import math, random
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.engine.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.opt.search import actions, bufs_from_lin, get_kernel_actions
 from tinygrad.nn.optim import Adam
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer

extra/optimization/run_qnet.py

@@ -1,6 +1,6 @@
 from typing import List, Tuple
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import get_kernel_actions, actions
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import get_kernel_actions, actions
 
 _net = None
 def beam_q_estimate(beam:List[Tuple[Kernel, float]]) -> List[Tuple[Kernel, float]]:

extra/optimization/search.py

@@ -4,8 +4,8 @@ from extra.optimization.helpers import ast_str_to_lin, time_linearizer
 from tinygrad import dtypes
 from tinygrad.helpers import BEAM, getenv
 from tinygrad.device import Device, Compiled
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.search import beam_search, bufs_from_lin
 
 
 if __name__ == '__main__':

extra/optimization/test_net.py

@@ -6,8 +6,8 @@ from copy import deepcopy
 from tinygrad.helpers import getenv, colored
 from tinygrad.tensor import Tensor
 from tinygrad.nn.state import get_parameters, get_state_dict, safe_save, safe_load, load_state_dict
-from tinygrad.engine.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, time_linearizer
 from extra.optimization.extract_policynet import PolicyNet
 from extra.optimization.pretrain_valuenet import ValueNet

extra/optimization/test_time_linearizer.py

@@ -1,5 +1,5 @@
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
-from tinygrad.engine.search import bufs_from_lin, get_kernel_actions
+from tinygrad.opt.search import bufs_from_lin, get_kernel_actions
 
 if __name__ == "__main__":
   ast_strs = load_worlds()

extra/perfetto/to_perfetto.py

@@ -0,0 +1,38 @@
+import sys, pickle, decimal, json
+from tinygrad.device import ProfileEvent, ProfileDeviceEvent, ProfileRangeEvent, ProfileGraphEvent
+from tinygrad.helpers import tqdm, temp
+
+devices:dict[str, tuple[decimal.Decimal, decimal.Decimal, int]] = {}
+def prep_ts(device:str, ts:decimal.Decimal, is_copy): return int(decimal.Decimal(ts) + devices[device][is_copy])
+def dev_to_pid(device:str, is_copy=False): return {"pid": devices[device][2], "tid": int(is_copy)}
+def dev_ev_to_perfetto_json(ev:ProfileDeviceEvent):
+  devices[ev.device] = (ev.comp_tdiff, ev.copy_tdiff if ev.copy_tdiff is not None else ev.comp_tdiff, len(devices))
+  return [{"name": "process_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "args": {"name": ev.device}},
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 0, "args": {"name": "COMPUTE"}},
+          {"name": "thread_name", "ph": "M", "pid": dev_to_pid(ev.device)['pid'], "tid": 1, "args": {"name": "COPY"}}]
+def range_ev_to_perfetto_json(ev:ProfileRangeEvent):
+  return [{"name": ev.name, "ph": "X", "ts": prep_ts(ev.device, ev.st, ev.is_copy), "dur": float(ev.en-ev.st), **dev_to_pid(ev.device, ev.is_copy)}]
+def graph_ev_to_perfetto_json(ev:ProfileGraphEvent, reccnt):
+  ret = []
+  for i,e in enumerate(ev.ents):
+    st, en = ev.sigs[e.st_id], ev.sigs[e.en_id]
+    ret += [{"name": e.name, "ph": "X", "ts": prep_ts(e.device, st, e.is_copy), "dur": float(en-st), **dev_to_pid(e.device, e.is_copy)}]
+    for dep in ev.deps[i]:
+      d = ev.ents[dep]
+      ret += [{"ph": "s", **dev_to_pid(d.device, d.is_copy), "id": reccnt+len(ret), "ts": prep_ts(d.device, ev.sigs[d.en_id], d.is_copy), "bp": "e"}]
+      ret += [{"ph": "f", **dev_to_pid(e.device, e.is_copy), "id": reccnt+len(ret)-1, "ts": prep_ts(e.device, st, e.is_copy), "bp": "e"}]
+  return ret
+def to_perfetto(profile:list[ProfileEvent]):
+  # Start json with devices.
+  prof_json = [x for ev in profile if isinstance(ev, ProfileDeviceEvent) for x in dev_ev_to_perfetto_json(ev)]
+  for ev in tqdm(profile, desc="preparing profile"):
+    if isinstance(ev, ProfileRangeEvent): prof_json += range_ev_to_perfetto_json(ev)
+    elif isinstance(ev, ProfileGraphEvent): prof_json += graph_ev_to_perfetto_json(ev, reccnt=len(prof_json))
+  return {"traceEvents": prof_json}
+
+if __name__ == "__main__":
+  fp = sys.argv[1]
+  with open(fp, "rb") as f: profile = pickle.load(f)
+  ret = to_perfetto(profile)
+  with open(fp:=temp("perfetto.json", append_user=True), "w") as f: json.dump(ret, f)
+  print(f"Saved perfetto output to {fp}. You can use upload this to the perfetto UI or Chrome devtools.")

extra/profiler/index.html

@@ -1,173 +0,0 @@
-<!DOCTYPE html>
-<html>
-<head>
-  <title>tinygrad profiler</title>
-  <meta name="viewport" content="width=device-width, initial-scale=1.0">
-  <meta charset="UTF-8">
-  <link rel="icon" href="data:;base64,iVBORw0KGgo=">
-  <script src="assets/d3js.org/d3.v7.min.js" charset="utf-8"></script>
-  <style>
-  * {
-    box-sizing: border-box;
-    margin: 0;
-    padding: 0;
-  }
-  html, body {
-    color: #f0f0f5;
-    margin: 0;
-    padding: 0;
-    width: 100%;
-    height: 100%;
-    font-family: sans-serif;
-    font-optical-sizing: auto;
-    font-weight: 400;
-    font-style: normal;
-    font-variation-settings: "wdth" 100;
-    font-size: 14px;
-    overflow: hidden;
-    background-color: #08090e;
-  }
-  #root {
-    display:flex;
-    width:100%;
-    height: 100%;
-    padding: 20px;
-  }
-  #process-name {
-    background: #0f1018;
-    padding: 2px;
-    border-radius: 2px;
-  }
-  [id^="thread"] {
-    padding: 2px;
-  }
-  #table-root {
-    position: absolute;
-    width: 100%;
-    height: 300px;
-    background: #0f1018;
-    bottom: 0;
-    left: 0;
-    overflow: auto;
-  }
-  table {
-    border-collapse: collapse;
-    width: 100%;
-  }
-  table thead th {
-    position: sticky;
-    top: 0;
-    z-index: 1
-  }
-  th {
-    background: #1D1F2A;
-    cursor: pointer;
-  }
-  th, td {
-    padding: 8px 16px;
-    text-align: left;
-  }
-  th.sorted-asc::after { content: " ↑"; }
-  th.sorted-desc::after { content: " 	↓"; }
-  </style>
-</head>
-<body>
-<script>
-const colors = ["7aa2f7", "ff9e64", "f7768e", "2ac3de", "7dcfff", "1abc9c", "9ece6a", "e0af68", "bb9af7", "9d7cd8", "ff007c"];
-
-const formatTime = (ms) => {
-  if (ms<=1e3) return `${ms}us`;
-  if (ms<=1e6) return `${(ms*1e-3).toFixed(2)}ms`;
-  return `${(ms*1e-6).toFixed(2)}s`;
-}
-
-async function main() {
-  const { traceEvents } = await (await fetch("/get_profile")).json();
-  const root = createChild("div.root", document.querySelector("body"));
-  const list = createChild("div.list", root);
-  const data = [];
-  const nameColors = {}; // event names get a unique color
-  const procNames = {};
-  for (const e of traceEvents) {
-    if (e.name === "process_name") {
-      const proc = createChild(`div.proc-${e.pid}`, list);
-      createChild("p.process-name", proc).textContent = e.args.name;
-      procNames[e.pid] = e.args.name;
-    }
-    else if (e.name === "thread_name") {
-      const thread = createChild(`div.thread-${e.pid}-${e.tid}`, `proc-${e.pid}`);
-      createChild("p.thread-name", thread).textContent = e.args.name;
-    }
-    else if (e.ph === "X") {
-      const thread = document.getElementById(`thread-${e.pid}-${e.tid}`);
-      if (!(e.name in nameColors)) nameColors[e.name] = colors[data.length%(colors.length-1)];
-      data.push({ ...e, y:rect(thread).y, color:`#${nameColors[e.name]}`, proc:procNames[e.pid] });
-    }
-  }
-  // render graph
-  const svg = d3.select(root).append("svg").attr("width", "100%");
-  const { y, width } = rect(svg.node()); // global coordinates
-  const render = svg.append("g").attr("transform", `translate(0, ${y})`);
-  const timestamps = data.map(t => t.ts);
-  const st = Math.min(...timestamps);
-  const timeScale = d3.scaleLinear().domain([0, Math.max(...timestamps)-st]).range([y, width]);
-  const timeAxis = render.append("g").call(d3.axisTop(timeScale).tickFormat(formatTime));
-  list.style = `margin-top: ${rect(timeAxis.node()).bottom}px;`;
-  // rescale time based coordinates to fit screen
-  for (e of data) {
-    e.st = e.ts-st;
-    e.x = timeScale(e.st);
-    e.width = timeScale(e.dur);
-  }
-  render.selectAll("rect").data(data).join("rect").attr("fill", d => d.color).attr("x", d => d.x).attr("y", d => d.y).attr("width", d => d.width)
-    .attr("height", 20);
-  render.call(d3.brush().on("end", (e) => {
-    if (!e.selection) return renderTable({ data });
-    const [[x0, y0], [x1, y1]] = e.selection;
-    const newData = data.filter(d => d.x>=x0 && d.x<=x1 && d.y>=y0 && d.y<=y1);
-    renderTable({ data: newData });
-  }));
-  createChild("div.table-root", root);
-  renderTable({ data });
-}
-
-const rect = (e) => e.getBoundingClientRect();
-
-const createChild = (es, p) => {
-  const parts = es.split(".", 2);
-  if (typeof p === "string") p = document.getElementById(p);
-  const ret = p.appendChild(document.createElement(parts[0]));
-  if (parts.length !== 1) ret.id = parts[1];
-  return ret;
-}
-
-const columnNames = {"name":"Name", "st":"Start Time", "dur":"Duration", "proc":"Process"};
-const tableState = {data:null, sortBy:null, asc:true};
-function renderTable(newState) {
-  const { data, sortBy, asc } = Object.assign(tableState, newState);
-  const root = document.getElementById("table-root");
-  root.innerHTML = "";
-  const table = createChild("table", root);
-  const thead = createChild("tr", createChild("thead", table));
-  for (const [k,v] of Object.entries(columnNames)) {
-    const th = createChild(`th.${k}`, thead);
-    th.innerText = v;
-    th.onclick = (e) => renderTable(k === sortBy ? { asc:!asc } : { sortBy:k, asc:true });
-  }
-  if (sortBy != null) {
-    data.sort((a, b) => asc ? a[sortBy]-b[sortBy] : b[sortBy]-a[sortBy]); // inplace sort
-    document.getElementById(sortBy).className = asc ? "sorted-asc" : "sorted-desc";
-  }
-  const tbody = createChild("tbody", table);
-  for (const d of data) {
-    const row = createChild("tr", tbody);
-    for (const k of Object.keys(columnNames)) {
-      let formatted = typeof d[k] === "string" ? d[k] : formatTime(d[k]);
-      createChild("td", row).innerText = formatted;
-    }
-  }
-}
-main()
-</script>
-</body>
-</html>

extra/remu/README.md

@@ -22,7 +22,7 @@ This will produce a binary in the `extra/remu/target/release` directory.
 
 The latest binaries are released in https://github.com/Qazalin/remu/releases. Alternatively, you can [build locally](#build-locally).
 
-Tinygrad does not yet output RDNA3 kernels directly. You can either install comgr or use `AMD_LLVM=1` if you have [LLVM@19](https://github.com/tinygrad/tinygrad/blob/e2ed673c946c8f1774d816c75e52a994c2dd8a88/.github/actions/setup-tinygrad/action.yml#L208).
+Tinygrad does not yet output RDNA3 kernels directly. You can either install comgr or use `AMD_LLVM=1` (default) if you have [LLVM@19](https://github.com/tinygrad/tinygrad/blob/e2ed673c946c8f1774d816c75e52a994c2dd8a88/.github/actions/setup-tinygrad/action.yml#L208).
 
 `PYTHONPATH="." MOCKGPU=1 AMD=1 python test/test_tiny.py TestTiny.test_plus` runs an emulated RDNA3 kernel with Remu.
 

extra/replay_pkl.py

@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 from tinygrad.renderer import ProgramSpec
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.heuristic import hand_coded_optimizations
 import numpy as np
 
 def move_jit_captured_to_dev(captured, device="DSP"):

extra/torch_backend/backend.py

@@ -59,6 +59,7 @@ view_ops = {
   "aten.squeeze.dim": Tensor.squeeze,
   "aten.unsqueeze": Tensor.unsqueeze,
   "aten.detach": Tensor.detach,
+  "aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
 }
 
 for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
@@ -368,6 +369,7 @@ decomps = [
   aten.threshold,
   aten.nll_loss_forward,
   aten.nll_loss_backward,
+  aten.nll_loss2d_backward,
   # AttributeError: 'int' object has no attribute '_broadcasted'
   aten.sigmoid_backward,
   aten.tanh_backward,
@@ -376,6 +378,7 @@ decomps = [
   aten.softshrink,
   aten.hardshrink,
   aten.log_sigmoid_forward,
+  aten.log_sigmoid_backward,
   aten.isneginf,
   aten.isposinf,
   aten.nan_to_num,

extra/torch_backend/test.py

@@ -206,5 +206,16 @@ class TestTorchBackend(unittest.TestCase):
       X.cpu(), Y.cpu()
       self.assertLessEqual(GlobalCounters.global_ops, 10_000_000)
 
+  def _test_diagonal(self, *shape):
+    a = torch.randn(*shape, dtype=torch.float32, device=device)
+    ref = np.diagonal(a.cpu().numpy(), axis1=-2, axis2=-1)
+    diag = torch.linalg.diagonal(a)
+    np.testing.assert_equal(diag.cpu().numpy(), ref)
+    np.testing.assert_equal(diag[-1].cpu().numpy(), ref[-1])
+
+  def test_diagonal_cube(self): self._test_diagonal(3, 3, 3)
+  def test_diagonal_rectangular(self): self._test_diagonal(4, 5, 6)
+  def test_diagonal_4d(self): self._test_diagonal(2, 3, 4, 5)
+
 if __name__ == "__main__":
   unittest.main()

extra/upcasted_warps.py

@@ -0,0 +1,73 @@
+# play with upcasted warps
+from tinygrad import Tensor, Device
+from tinygrad.uop.ops import KernelInfo
+from tinygrad.opt import get_optimized_ast
+from tinygrad.opt.kernel import OptOps, Opt
+from tinygrad.engine.realize import get_program
+
+if __name__ == "__main__":
+  renderer = Device.default.renderer
+  N = 64
+
+  """
+  a = Tensor.empty(N,N)
+
+  out = (a + 1) #.sum(axis=2)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 32),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # how you split the store determines everything if you don't allow cross warp comms.
+  # actually not everything, there's also the split before the horizontal (unrolled) reduces
+
+  # new flow
+  #  - pull out any dimensions from the store that you want to upcast.
+  #  - decide how you want to assign them to registers. GPUs have a 512-byte memory LOAD/STORE which loads into 4 regs. see BUFFER_LOAD_B128
+  #    - the loads and stores can be shuffled, but only in restrictive ways. in kernels without reduces, the store determines everything
+  #    - it loads 16 bytes from up 32 different places = 512 bytes
+  #  - in kernels with reduces, you now have more flexibility. the final target of the reduce must be what is stored
+  #    - warp dimensions can be in the reduce (this is GROUP)
+
+  # every dimension can be assigned to <global, local, loop, upcast, warp>
+
+  """
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+
+  out = a.sum(axis=1)
+  ast = out.schedule()[-1].ast
+  opts = tuple()
+  opts += (Opt(OptOps.UNROLL, 0, 8),)
+  opts += (Opt(OptOps.UPCAST, 0, 8),)
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)
+  """
+
+  # gemm
+  b = Tensor.empty(N,N)
+  # metal TC
+  #opts = (Opt(OptOps.UPCAST, 0, 2), # not the warp
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UPCAST, 1, 2),
+  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2))
+  # new TC should just be able to extract from this and swizzle as needed
+  opts = (Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 1, 8), Opt(OptOps.UNROLL, 0, 8))
+  c = (a@b)
+  ast = c.schedule()[-1].ast
+  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
+  ast = get_optimized_ast(ast, renderer)
+  prg = get_program(ast, renderer)
+  print(prg.src)

mkdocs.yml

@@ -22,6 +22,7 @@ nav:
     - Runtime: runtime.md
   - Developer:
     - Intro: developer/developer.md
+    - Layout: developer/layout.md
     - Speed: developer/speed.md
     - UOp: developer/uop.md
     - Grouper:

ruff.toml

@@ -36,9 +36,17 @@ line-length = 150
 
 exclude = [
   "docs/",
-  "examples/",
   "extra/",
   "tinygrad/runtime/autogen",
   "test/external/mlperf_resnet",
   "test/external/mlperf_unet3d",
 ]
+
+# detect unused imports in examples
+[lint.per-file-ignores]
+"examples/**/*.py" = [
+  "W6", "E71", "E72", "E112", "E113", "E203", "E272", "E275",
+  "E303", "E304", "E501", "E702", "E703", "E731", "W191",
+  "W291", "W293", "UP039", "C416", "RET506", "RET507", "A",
+  "FURB110", "RUF018", "F541", "F841"
+]

setup.py

@@ -25,9 +25,10 @@ setup(name='tinygrad',
       long_description=long_description,
       long_description_content_type='text/markdown',
       packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.runtime.autogen.am', 'tinygrad.codegen', 'tinygrad.nn',
-                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop'],
-      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'perfetto.html', 'assets/**/*', 'js/*']},
+                  'tinygrad.renderer', 'tinygrad.engine', 'tinygrad.viz', 'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.kernelize',
+                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
+                  'tinygrad.runtime.support.nv'],
+      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"

test/external/external_benchmark_hcopt.py

@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.engine.search import beam_search, bufs_from_lin
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, time_linearizer
 
 def optimize_kernel(k):

test/external/external_benchmark_llama_schedule.py

@@ -0,0 +1,35 @@
+from tinygrad import nn, Tensor, Device, dtypes
+from tinygrad.helpers import Timing
+
+from extra.models.llama import Transformer
+from examples.llama3 import MODEL_PARAMS
+
+if __name__ == "__main__":
+  Device.DEFAULT = "NULL"
+  Tensor.training = True
+  #model_size = "8B"
+  model_size = "405B"
+
+  with Timing("total "):
+    with Timing("***** create model in    "):
+      model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=nn.Linear, embedding=nn.Embedding,
+                          max_context=1024, jit=True, disable_kv_cache=True)
+
+    with Timing("***** fake state in      "):
+      Tensor.realize(*[p.assign(Tensor.empty(*p.shape, device=p.device, dtype=p.dtype)) for p in nn.state.get_parameters(model)])
+
+    with Timing("***** create optim in    "):
+      opt = nn.optim.AdamW(nn.state.get_parameters(model))
+
+    with Timing("***** run model in       "):
+      toks = Tensor.empty(1, 1024, dtype=dtypes.int)
+      out = model(toks, 0, temperature=float('nan'))
+
+    with Timing("***** backward in        "):
+      out.mean().backward()
+
+    with Timing("***** realize in        "):
+      out.realize()
+
+    with Timing("***** step in           "):
+      opt.step()

test/external/external_benchmark_schedule.py

@@ -3,10 +3,11 @@ from extra.models.resnet import ResNet50
 from tinygrad import Tensor, nn
 from tinygrad.helpers import Profiling, Timing, getenv, BEAM, NOOPT, DEBUG, Context, ansilen
 from tinygrad.uop.ops import Ops
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.heuristic import hand_coded_optimizations
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.uop.spec import type_verify
 
 if __name__ == "__main__":
   mdl = ResNet50()
@@ -56,4 +57,6 @@ if __name__ == "__main__":
             uops_line = []
             for u in rewritten_uops:
               uops_line.append(apply_rewrites(u, rewrites_for_linearizer))
+          with Timing("***** model verify in    "):
+            for u in uops_line: type_verify(u.arg.lst)
           print(sum(len(u.arg.lst) for u in uops_line))

test/external/external_benchmark_sdxl_softmax.py

@@ -7,7 +7,7 @@ if __name__ == "__main__":
   GlobalCounters.reset()
   t.softmax(-1, dtype="half", _single_kernel=True).realize()
 
-  from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+  from tinygrad.opt.kernel import Kernel, Opt, OptOps
   from tinygrad.helpers import get_single_element
   GlobalCounters.reset()
   si = get_single_element(t.softmax(-1, dtype="half", _single_kernel=True).schedule())

test/external/external_debug_metal_sd_conv.py

@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.opt.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.engine.realize import CompiledRunner
-from tinygrad.engine.search import bufs_from_lin
+from tinygrad.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

test/external/external_fuzz_ampt.py

@@ -1,4 +1,5 @@
 import random
+from typing import Optional
 from tinygrad.helpers import round_up
 from tinygrad.runtime.support.am.amdev import AMPageTableTraverseContext
 from test.external.external_test_am import helper_read_entry_components, FakeAM
@@ -29,7 +30,7 @@ class AMPTFuzzer:
         self.d.vram[pte['paddr']] = pattern # Mark this page
         assert pte['valid'] == 1
 
-        # If page has contigous fragment, all range should be this valid memory
+        # If page has contiguous fragment, all range should be this valid memory
         frags_cnt = pte['fragment']
         contig_range = (1 << (frags_cnt + 12))
         start_vaddr = _vaddr & ~(contig_range - 1)
@@ -59,7 +60,7 @@ class AMPTFuzzer:
 
     return True
 
-  def random_alloc(self):
+  def random_alloc(self) -> Optional[int]:
     if self.total_size - self.alloc_payload < self.min_alloc_size: return None
 
     size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))

test/external/external_fuzz_tlsf.py

@@ -1,7 +1,7 @@
 import random
 from typing import Dict, Optional
 from tinygrad.helpers import getenv
-from tinygrad.runtime.support.allocator import TLSFAllocator
+from tinygrad.runtime.support.memory import TLSFAllocator
 
 class AllocatorFuzzer:
   def __init__(self, total_size):
@@ -30,6 +30,8 @@ class AllocatorFuzzer:
     return True
 
   def random_alloc(self) -> Optional[int]:
+    if self.total_size - self.alloc_payload < self.min_alloc_size: return None
+
     size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
 
     try:

test/external/external_llama_eval.py

@@ -1,102 +0,0 @@
-from lm_eval.base import BaseLM
-from lm_eval import evaluator, tasks
-import torch, json, argparse
-
-from examples.llama import LLaMa
-from tinygrad.tensor import Tensor
-from tinygrad import Device
-
-class LLaMaAdaptor(BaseLM):
-  def __init__(
-    self,
-    model_size="7B",
-    model_gen=1,
-    device="",
-    quantize=False,
-    batch_size=1,
-    max_batch_size=1,
-    do_sample=False,
-    temperature=1.0,
-    checkpoint_path="",
-    tokenizer_path="",
-  ):
-    super().__init__()
-
-    if batch_size is None:
-      batch_size = 1
-    self.do_sample = do_sample
-    self.temperature = temperature
-    self._device = device
-
-    assert isinstance(model_gen, int)
-    assert isinstance(model_size, str)
-    assert isinstance(batch_size, int)
-    assert isinstance(checkpoint_path, str)
-    assert isinstance(tokenizer_path, str)
-
-    self.llama = LLaMa.build(checkpoint_path, tokenizer_path, model_gen, model_size, quantize)
-
-  @classmethod
-  def create_from_arg_string(cls, arg_string, additional_config=None):
-    kwargs = {el.split("=")[0]: el.split("=")[1] for el in arg_string.split(",")}
-    return cls(**kwargs, **additional_config)
-
-  @property
-  def eot_token_id(self):
-    # we use EOT because end of *text* is more accurate for what we're doing than end of *sentence*
-    return self.llama.tokenizer.eos_id()
-
-  @property
-  def max_length(self):
-    return 1024
-
-  @property
-  def max_gen_toks(self):
-    return 256
-
-  @property
-  def batch_size(self):
-    return 1
-
-  @property
-  def device(self):
-    return self._device
-
-  def tok_encode(self, string: str):
-    return [self.llama.tokenizer.bos_id()] + self.llama.tokenizer.encode(string)
-
-  def tok_decode(self, tokens):
-    return self.llama.tokenizer.decode(tokens)
-
-  def _model_call(self, inps):
-    return torch.Tensor(self.llama.model(Tensor(inps.numpy()), 0).numpy())
-
-  def greedy_until(self, requests):
-    continuations = []
-    for request in requests:
-      prompt, until = request[0], request[1]['until']
-      output = self.llama.greedy_until(prompt, until, max_length=128, temperature=0.0)
-      continuations.append(output[len(prompt):])
-    return continuations
-
-  def _model_generate(self, context, max_length, eos_token_id):
-    raise NotImplementedError()
-
-if __name__ == '__main__':
-  print(f"using {Device.DEFAULT} backend")
-
-  parser = argparse.ArgumentParser(description='Run LLaMA evals in tinygrad', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
-  parser.add_argument('--size', type=str, default="7B", help="Size of model to use [7B, 13B, 30B, 65B] for Gen 1, [7B, 13B] for Gen 2")
-  parser.add_argument('--gen', type=int, default="1", help="Generation of the model to use [1, 2]")
-  parser.add_argument('--quantize', action='store_true', help="Quantize the weights to int8 in memory")
-  parser.add_argument('--eval', type=str, default="arc_easy", help="Run in evaluation mode")
-  parser.add_argument('--limit', type=int, default=None, help="Limit tests in eval")
-  parser.add_argument('--weights', type=str, default="./weights/LLaMa/", help="Location of the weights")
-  parser.add_argument('--tokenizer', type=str, default="./weights/LLaMa/tokenizer.model", help="Location of the tokenizer")
-  args = parser.parse_args()
-
-  # run eval and exit
-  adaptor = LLaMaAdaptor(model_gen=args.gen, model_size=args.size, quantize=args.quantize,
-                         checkpoint_path=args.weights, tokenizer_path=args.tokenizer, device="cpu")
-  results = evaluator.evaluate(adaptor, tasks.get_task_dict(args.eval.split(",")), False, 0, args.limit)
-  print(json.dumps(results, indent=2))

test/external/external_model_benchmark.py

@@ -7,7 +7,7 @@ import onnxruntime as ort
 from onnx2torch import convert
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from tinygrad.helpers import OSX, DEBUG, fetch, getenv
-from tinygrad import Tensor, Device
+from tinygrad import Tensor, Device, dtypes
 
 MODELS = {
   "resnet50": "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
@@ -27,6 +27,7 @@ MODELS = {
   # really slow
   # "resnet18": "https://github.com/onnx/models/raw/main/archive/vision/classification/resnet/model/resnet18-v2-7.onnx",
 }
+half_models = ["openpilot", "commavq"]
 
 CSV = {}
 open_csv = None
@@ -54,7 +55,6 @@ def benchmark_model(m, devices, validate_outs=False):
   excluded = {inp.name for inp in onnx_model.graph.initializer}
   input_shapes = {inp.name:tuple(x.dim_value if hasattr(x, "dim_value") and x.dim_value != 0 else 1 for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input if inp.name not in excluded}  # noqa: E501
   input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input if inp.name not in excluded}
-  #input_types = {k:v if v!=np.float16 else np.float32 for k,v in input_types.items()}  # cast
   np_inputs = {k:torch.randn(shp).numpy().astype(input_types[k]) for k,shp in input_shapes.items()}
   assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"
 
@@ -106,7 +106,12 @@ def benchmark_model(m, devices, validate_outs=False):
     for device in devices:
       rtol, atol = 2e-3, 2e-3  # tolerance for fp16 models
       Device.DEFAULT = device
-      inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
+      # force half inputs to float for numerical stability when validating
+      # this will reply on automatic dtype promotion for converting half weights inside the graph
+      if m in half_models:
+        inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
+      else:
+        inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
       tinygrad_model = OnnxRunner(onnx_model)
       tinygrad_out = tinygrad_model(inputs)
 

test/external/external_test_am.py

@@ -1,7 +1,8 @@
 import unittest
-from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableTraverseContext
+from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
+from tinygrad.runtime.support.memory import PageTableTraverseContext
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address
 
@@ -23,7 +24,9 @@ class FakeAM:
     self.vram_mv = memoryview(bytearray(4 << 30))
     self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
     self.gmc = FakeGMC(self)
-    self.mm = AMMemoryManager(self, vram_size=4 << 30)
+    self.mm = AMMemoryManager(self, 4 << 30, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
+      pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, first_page_lv=am.AMDGPU_VM_PDB2,
+      va_base=AMMemoryManager.va_allocator.base)
     self.is_booting = False
     self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
   def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
@@ -65,7 +68,7 @@ class TestAMPageTable(unittest.TestCase):
       exteranl_va = va + AMMemoryManager.va_allocator.base
       mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])
 
-      ctx = AMPageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
+      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
       results = list(ctx.next(sz))
 
       total_covered = 0
@@ -91,6 +94,17 @@ class TestAMPageTable(unittest.TestCase):
           assert pte['paddr'] == 0
           assert pte['valid'] == 0
 
+  def test_map_notaligned(self):
+    mm0 = self.d[0].mm
+
+    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
+      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
+      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
+      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
+      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
+      mm0.unmap_range(va2, sz2)
+      mm0.unmap_range(va1, sz1)
+
   def test_double_map(self):
     mm0 = self.d[0].mm
 
@@ -115,7 +129,7 @@ class TestAMPageTable(unittest.TestCase):
       # Finally can map and check paddrs
       mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])
 
-      ctx = AMPageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
+      ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
       for tup in ctx.next(0x100000):
         _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
         for i in range(_n_ptes):

test/external/external_test_datasets.py

@@ -64,6 +64,7 @@ class TestKiTS19Dataset(ExternalTestDatasets):
 
     return iter(dataset)
 
+  @unittest.skip("flaky")
   def test_training_set(self):
     preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(False)
     ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, False)

test/external/external_test_hcq_fuzz_failures.py

@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"
 
 import unittest, random
 import numpy as np
-from tinygrad.codegen.kernel import Kernel, KernelOptError
+from tinygrad.opt.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like
 

test/external/external_test_hip_compile.py

@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestHIPCompileSpeed(unittest.TestCase):
   @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")

test/external/external_test_metrics.py

@@ -1,10 +1,10 @@
 from tinygrad import Tensor
 from test.external.mlperf_unet3d.dice import DiceScore
-from examples.mlperf.metrics import dice_score
+from examples.mlperf.metrics import dice_score, log_perplexity
 
 import numpy as np
 import torch
-import unittest
+import unittest, math
 
 class ExternalTestMetrics(unittest.TestCase):
   def _test_metrics(self, tinygrad_metrics, orig_metrics, pred, label, atol=1e-8, rtol=1e-7):
@@ -16,5 +16,27 @@ class ExternalTestMetrics(unittest.TestCase):
     pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
     self._test_metrics(dice_score, DiceScore(), pred, label)
 
+  def test_log_perplexity(self):
+    # equally likely
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0, 1, 1, 1]]]), Tensor([[2]])).numpy(), math.log(4))
+    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0]*256]*32]), Tensor([[2]*32])).numpy(), math.log(256), rtol=1e-6)
+    # pretty correct and incorrect
+    np.testing.assert_allclose(log_perplexity(Tensor([[[10000., 0, 0, 0]]]), Tensor([[0]])).numpy(), 0)
+    np.testing.assert_allclose(log_perplexity(Tensor([[[0.0, 10000, 10000, 10000]]]), Tensor([[0]])).numpy(), 10000, rtol=1e-3)
+    # higher logit -> lower loss
+    x = Tensor([[[4.0, 3, 2, 1]]])
+    for i in range(x.numel()-1): self.assertLess(log_perplexity(x, Tensor([[i]])).item(), log_perplexity(x, Tensor([[i+1]])).item())
+    # torch eval examples
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]])).exp().numpy(),
+      2.7593, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3, 0.7, 0.3, 0.1], [0.5, 0.4, 0.1, 0.4],[0.1, 0.1, 0.2, 0.5]],
+                             [[0.1, 0.6, 0.1, 0.5], [0.3, 0.7, 0.3, 0.4], [0.3, 0.7, 0.3, 0.4]]]), Tensor([[2, 1, 3],  [1, 0, 1]])).exp().numpy(),
+      3.6216, rtol=1e-5)
+    np.testing.assert_allclose(
+      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]]), ignore_index=1).exp().numpy(),
+      3.5372, rtol=1e-4)
+
 if __name__ == '__main__':
   unittest.main()

test/external/external_test_nv.py

@@ -2,12 +2,12 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 from tinygrad.engine.realize import get_runner, CompiledRunner
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View

test/external/external_test_onnx_backend.py

@@ -44,6 +44,9 @@ class TinygradBackend(Backend):
 
 backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)
 
+# BUG: segfaults
+backend_test.exclude('test_MaxPool1d_stride_padding_dilation_cpu')
+
 # BUG: buggy onnx tests
 backend_test.exclude('test_adam_multiple_cpu')
 
@@ -91,13 +94,6 @@ backend_test.exclude('FLOAT8')
 backend_test.exclude('INT4')
 backend_test.exclude('UINT4')
 backend_test.exclude('BFLOAT16')  # not supported in numpy
-# TODO: fix these with true onnx float16
-backend_test.exclude('to_FLOAT16')
-backend_test.exclude('cast_no_saturate')
-backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
-backend_test.exclude('test_max_float16_cpu')
-backend_test.exclude('test_min_float16_cpu')
-backend_test.exclude('test_mod_mixed_sign_float16_cpu')
 
 backend_test.exclude('test_dequantizelinear_int4_cpu')
 backend_test.exclude('test_dequantizelinear_uint4_cpu')
@@ -111,6 +107,7 @@ backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_quantizelinear_e5m2_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
+backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
 backend_test.exclude('test_dequantizelinear_e5m2_cpu')
 
 # we don't support indexes

test/external/external_test_onnx_runner.py

@@ -0,0 +1,77 @@
+import unittest, onnx, tempfile
+from tinygrad import dtypes
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
+from tinygrad.device import is_dtype_supported
+from extra.onnx import data_types
+from hypothesis import given, settings, strategies as st
+import numpy as np
+
+data_types.pop(16) # TODO: this is bf16, need to support double parsing first.
+device_supported_dtypes = [odt for odt, dtype in data_types.items() if is_dtype_supported(dtype)]
+device_unsupported_dtypes = [odt for odt, dtype in data_types.items() if not is_dtype_supported(dtype)]
+
+class TestOnnxRunnerDtypes(unittest.TestCase):
+  def _test_input_spec_dtype(self, onnx_data_type, tinygrad_dtype):
+    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
+    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
+    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    onnx.save(model, tmp.name)
+    tmp.flush()
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(len(runner.graph_inputs), 1)
+    self.assertEqual(runner.graph_inputs['input'].dtype, tinygrad_dtype)
+
+  def _test_initializer_dtype(self, onnx_data_type, tinygrad_dtype):
+    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
+    initializer = onnx.helper.make_tensor('initializer', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
+    input_tensor = onnx.helper.make_tensor_value_info('input', onnx_data_type, ())
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, ())
+    node = onnx.helper.make_node('Identity', inputs=['input'], outputs=['output'])
+    graph = onnx.helper.make_graph([node], 'identity_test', [input_tensor], [output_tensor], [initializer])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    onnx.save(model, tmp.name)
+    tmp.flush()
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(len(runner.graph_inputs), 1)
+    self.assertEqual(runner.graph_values['initializer'].dtype, tinygrad_dtype)
+
+  def _test_node_attribute_dtype(self, onnx_data_type, tinygrad_dtype):
+    arr = np.array([0, 1], dtype=onnx.helper.tensor_dtype_to_np_dtype(onnx_data_type))
+    output_tensor = onnx.helper.make_tensor_value_info('output', onnx_data_type, arr.shape)
+    value_tensor = onnx.helper.make_tensor('value', onnx_data_type, arr.shape, arr.tobytes(), raw=True)
+    node = onnx.helper.make_node('Constant', inputs=[], outputs=['output'], value=value_tensor)
+    graph = onnx.helper.make_graph([node], 'attribute_test', [], [output_tensor])
+    model = onnx.helper.make_model(graph)
+    tmp = tempfile.NamedTemporaryFile(suffix='.onnx')
+    tmp.flush()
+    onnx.save(model, tmp.name)
+    model = onnx_load(tmp.name)
+    runner = OnnxRunner(model)
+    self.assertEqual(runner.graph_nodes[0].opts['value'].dtype, tinygrad_dtype)
+
+  @settings(deadline=1000) # TODO investigate unreliable timing
+  @given(onnx_data_type=st.sampled_from(device_supported_dtypes))
+  def test_supported_dtype_spec(self, onnx_data_type):
+    tinygrad_dtype = data_types[onnx_data_type]
+    self._test_input_spec_dtype(onnx_data_type, tinygrad_dtype)
+    self._test_initializer_dtype(onnx_data_type, tinygrad_dtype)
+    self._test_node_attribute_dtype(onnx_data_type, tinygrad_dtype)
+
+  @unittest.skipUnless(device_unsupported_dtypes, "No unsupported dtypes for this device to test.")
+  @settings(deadline=1000) # TODO investigate unreliable timing
+  @given(onnx_data_type=st.sampled_from(device_unsupported_dtypes))
+  def test_unsupported_dtype_spec(self, onnx_data_type):
+    true_dtype = data_types[onnx_data_type]
+    default_dtype = dtypes.default_int if dtypes.is_int(true_dtype) else dtypes.default_float
+    self._test_input_spec_dtype(onnx_data_type, true_dtype)
+    self._test_initializer_dtype(onnx_data_type, default_dtype)
+    self._test_node_attribute_dtype(onnx_data_type, default_dtype)
+
+if __name__ == '__main__':
+  unittest.main()

test/external/external_test_optim.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-import unittest
+import unittest, math
 import numpy as np
 import tensorflow as tf
 import tensorflow_addons as tfa
@@ -7,11 +7,11 @@ from tensorflow.python.ops import math_ops
 from extra.lr_scheduler import LRSchedulerGroup
 
 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW
 
 from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer
 
-from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
+from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup
 from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf
 
 np.random.seed(1337)
@@ -171,5 +171,26 @@ class ExternalTestOptim(unittest.TestCase):
       'warmup': steps_per_epoch * warmup_epochs,
     }, 1e-5, 1e-5, do_optim=False)
 
+
+class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
+  # only tests the lr
+  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
+    net = TinyNet()
+    optim = AdamW([net.W], lr=0.0)
+    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
+    lr = []
+    for _ in range(warmup_steps+decay_steps):
+      lr.append(optim.lr.item())
+      tiny_lr.step()
+    # reimplemented in python
+    expected = []
+    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
+    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
+    np.testing.assert_allclose(lr, expected, rtol=1e-5)
+
+  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
+  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
+  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
+
 if __name__ == '__main__':
   unittest.main()

test/external/external_test_tlsf.py

@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.runtime.support.allocator import TLSFAllocator
+from tinygrad.runtime.support.memory import TLSFAllocator
 
 class TestTLSFAllocator(unittest.TestCase):
   def setUp(self):

test/external/external_test_train_gpt2.py

@@ -2,11 +2,11 @@
 import unittest
 
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 from test.external.fuzz_linearizer import run_linearizer
 

test/external/external_test_valid_remove.py

@@ -3,11 +3,11 @@ import unittest
 
 from tinygrad import Device
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.engine.search import Opt, OptOps
+from tinygrad.opt.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.codegen.kernel import Kernel
+from tinygrad.opt.kernel import Kernel
 
 class TestOpenpilotValidhack(unittest.TestCase):
   def test_valid_removal(self):

test/external/external_uop_gc.py

@@ -1,7 +1,7 @@
 import gc
 from tinygrad import Tensor, UOp, Device
 from tinygrad.shape.shapetracker import views_to_indexed_uops
-from tinygrad.engine.realize import method_cache, get_kernel
+from tinygrad.engine.realize import method_cache, get_program
 
 def uops_allocated(): return sum([isinstance(x, UOp) for x in gc.get_objects()])
 def print_uops():
@@ -14,12 +14,10 @@ def two_plus_two(): Tensor([2])+Tensor([2])
 def two_plus_two_schedule(): (Tensor([2])+Tensor([2])).schedule()
 def two_plus_two_kernel():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  get_kernel(Device.default.renderer, si.ast)
+  get_program(si.ast, Device.default.renderer)
 def two_plus_two_linearize():
   si = (Tensor([2])+Tensor([2])).schedule()[-1]
-  k = get_kernel(Device.default.renderer, si.ast)
-  k.get_optimized_ast()
-  #k.linearize()
+  get_program(si.ast, Device.default.renderer)
 def two_plus_two_realize(): (Tensor([2])+Tensor([2])).realize()
 def two_plus_two_item(): (Tensor([2])+Tensor([2])).item()
 def gradient_test():
@@ -36,7 +34,7 @@ def kernel_matmul():
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)
   z = y.matmul(x)
   si = z.schedule()[-1]
-  get_kernel(Device.default.renderer, si.ast)
+  get_program(si.ast, Device.default.renderer)
 def realized_matmul():
   x = Tensor.eye(3, requires_grad=True)
   y = Tensor([[2.0,0,-2.0]], requires_grad=True)

test/external/fuzz_linearizer.py

@@ -20,9 +20,9 @@ if os.getenv("VALIDATE_HCQ", 0) != 0:
 
 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.codegen.kernel import Kernel
-from tinygrad.codegen.kernel import Opt, OptOps
-from tinygrad.engine.search import get_kernel_actions, bufs_from_lin
+from tinygrad.opt.kernel import Kernel
+from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops