remove AMD_LLVM=0 in mlperf and search ci

tinybox updated to llvm 20

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

@@ -29,10 +29,6 @@ inputs:
     description: "Install CUDA?"
     required: false
     default: 'false'
-  ocelot:
-    description: "Install gpuocelot?"
-    required: false
-    default: 'false'
   webgpu:
     description: "Install webgpu?"
     required: false
@@ -197,14 +193,14 @@ runs:
           sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
         cargo build --release --manifest-path ./extra/remu/Cargo.toml
 
-    # **** gpuocelot ****
+    # **** CUDA ****
 
     - name: Install gpuocelot dependencies (MacOS)
-      if: inputs.ocelot == 'true' && runner.os == 'macOS'
+      if: inputs.cuda == '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.ocelot == 'true'
+      if: inputs.cuda == 'true'
       id: cache-build
       uses: actions/cache@v4
       env:
@@ -213,7 +209,7 @@ runs:
         path: ${{ github.workspace }}/gpuocelot/ocelot
         key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
     - name: Clone/compile gpuocelot
-      if: inputs.ocelot == 'true' && steps.cache-build.outputs.cache-hit != 'true'
+      if: inputs.cuda == 'true' && steps.cache-build.outputs.cache-hit != 'true'
       shell: bash
       run: |
         git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
@@ -224,7 +220,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.ocelot == 'true'
+      if: inputs.cuda == '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_padded TestLinearizer.test_tensor_cores_padded_uops
+      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
     - name: Test AMX tensor cores
       run: |
-        DEBUG=2 CPU=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops TestFloat4.test_float4_multidim_amx TestFloat4.test_float4_multidim_unaligned_load_amx
+        DEBUG=2 CPU=1 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
     - 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 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 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
     - 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_padded TestLinearizer.test_tensor_cores_padded_uops
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
+        NV=1 ALLOW_TF32=1 python3 test/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
     - 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 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 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 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,8 +607,12 @@ 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

@@ -10,7 +10,6 @@ jobs:
   run_script_job:
     runs-on: [self-hosted, Linux, tinybox]
     if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 100
 
     steps:
     - name: Checkout Code
@@ -28,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 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 FUSE_ARANGE=1 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 JITBEAM=4 BS=1024 STEPS=500 python examples/hlb_cifar10.py
+        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

.github/workflows/mlperf.yml

@@ -12,7 +12,6 @@ jobs:
   run_script_job:
     runs-on: [self-hosted, Linux, tinybox]
     if: github.repository_owner == 'tinygrad'
-    timeout-minutes: 240
 
     steps:
     - name: Checkout Code

.github/workflows/test.yml

@@ -239,8 +239,8 @@ jobs:
       run: |
         PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
         PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
         PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 EMULATE_AMD=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded_amd TestLinearizer.test_tensor_cores_padded_uops
     - name: Test emulated AMD MFMA tensor cores
@@ -252,8 +252,8 @@ jobs:
       run: |
         PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
         PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=0 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
-        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 ATOL=1e-3 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=16 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
+        PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 N=64 HALF=1 ACC_HALF=1 python3 ./extra/gemm/simple_matmul.py
         PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores
         PYTHONPATH=. DEBUG=2 EMULATE_AMD_RDNA4=1 FORWARD_ONLY=1 PYTHON=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores_padded TestLinearizer.test_tensor_cores_padded_uops
     - name: Test emulated CUDA tensor cores
@@ -272,6 +272,14 @@ 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
@@ -302,6 +310,8 @@ 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
 
@@ -326,21 +336,16 @@ jobs:
       run: |
         pip3 install --upgrade --force-reinstall ruff==0.11.0
         python3 -m ruff check .
-        python3 -m ruff check extra/onnx.py extra/onnx_parser.py
-        python3 -m ruff check examples/mlperf/ --ignore E501
+        python3 -m ruff check examples/mlperf/model_train.py --ignore E501
     - name: Lint tinygrad with pylint
       run: python -m pylint tinygrad/
     - name: Run mypy
-      run: |
-        python -m mypy --strict-equality --lineprecision-report .
-        cat lineprecision.txt
-        python -m mypy --strict-equality extra/onnx_parser.py
-        python -m mypy --strict-equality extra/onnx.py
+      run: python -m mypy --strict-equality --lineprecision-report . && cat lineprecision.txt
 
   unittest:
     name: Unit Tests
     runs-on: ubuntu-latest
-    timeout-minutes: 15
+    timeout-minutes: 10
 
     steps:
     - name: Checkout Code
@@ -366,15 +371,8 @@ jobs:
       run: PYTHONPATH="." python test/external/external_uop_gc.py
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
-    - name: Regen dataset on test_tiny
-      run: |
-        test/external/process_replay/reset.py
-        CAPTURE_PROCESS_REPLAY=1 python test/test_tiny.py TestTiny.test_plus
-        PYTHONPATH=. python extra/optimization/extract_dataset.py
-        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 < 15500 lines
-      run: MAX_LINE_COUNT=15500 python sz.py
+    - name: Repo line count < 14000 lines
+      run: MAX_LINE_COUNT=14000 python sz.py
 
   fuzzing:
     name: Fuzzing
@@ -427,29 +425,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  testgendataset:
-    name: 'GPU Generate Kernel Dataset'
-    runs-on: ubuntu-22.04
-    timeout-minutes: 10
-    env:
-      IGNORE_OOB: 0
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: gen-dataset
-          deps: testing_minimal
-          opencl: 'true'
-      - name: Generate Dataset
-        run: PYTHONPATH="." extra/optimization/generate_dataset.sh
-      - name: Upload artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: sops.gz
-          path: /tmp/sops.gz
-
   testopenpilot:
     name: 'openpilot Compile Tests'
     runs-on: ubuntu-22.04
@@ -468,17 +443,13 @@ jobs:
           llvm: 'true'
       - name: Test openpilot model kernel count and gate usage
         run: |
-          PYTHONPATH="." ALLOWED_KERNEL_COUNT=208 ALLOWED_READ_IMAGE=2134 ALLOWED_GATED_READ_IMAGE=13 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
+          PYTHONPATH="." ALLOWED_KERNEL_COUNT=209 ALLOWED_READ_IMAGE=2137 ALLOWED_GATED_READ_IMAGE=29 FLOAT16=0 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot alt model correctness (float32)
         run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/3799fe46b3a629e491d4b8498b8ae83e4c88c304/selfdrive/modeld/models/supercombo.onnx
       - name: Test openpilot fastvits model correctness (float32)
         run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      # - name: Test openpilot simple_plan vision model correctness (float32)
-      #   run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/35ff4f4577002f2685e50c8346addae33fe8da27a41dd4d6a0f14d1f4b1af81b
       - 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
 
@@ -503,12 +474,13 @@ 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 Additional ONNX Ops (CPU)
         run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
       - name: Test Quantize ONNX
         run: CPU=1 PYTHONPATH=. python3 test/test_quantize_onnx.py
+      - name: Run REMOTE=1 Test
+        run: |
+          REMOTEDEV=CPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
@@ -532,10 +504,14 @@ jobs:
           opencl: 'true'
       - name: Test ONNX (GPU)
         run: GPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
+      - name: Run REMOTE=1 Test
+        run: |
+          REMOTEDEV=GPU REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
+          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
@@ -545,20 +521,6 @@ jobs:
       - name: Run process replay tests
         uses: ./.github/actions/process-replay
 
-  testllm:
-    name: Test LLM
-    runs-on: ubuntu-24.04
-    timeout-minutes: 15
-    steps:
-      - name: Checkout Code
-        uses: actions/checkout@v4
-      - name: Setup Environment
-        uses: ./.github/actions/setup-tinygrad
-        with:
-          key: apps_llm
-      - name: Test 1B LLM
-        run: echo "What's a male chicken called? Answer with only one word." | MAX_BUFFER_SIZE=0 python3 -m tinygrad.apps.llm | grep -i rooster
-
   testmodels:
     name: Models (llvm+cpu+gpu)
     runs-on: ubuntu-22.04
@@ -598,7 +560,7 @@ jobs:
       with:
         key: dsp-minimal
         deps: testing_minimal
-        pydeps: "onnx==1.18.0 onnxruntime pillow"
+        pydeps: "onnx==1.17.0 onnxruntime pillow"
         llvm: "true"
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
@@ -648,7 +610,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 --durations=20
+          --ignore=test/test_fuzz_shape_ops.py --ignore=test/test_linearizer_failures.py --durations=20
     - name: Run process replay tests
       uses: ./.github/actions/process-replay
 
@@ -686,9 +648,7 @@ jobs:
         if: matrix.backend=='amdllvm'
         run: python test/test_amd_llvm.py
       - name: Run pytest (amd)
-        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py --durations=20
-      - name: Run pytest (amd)
-        run: python -m pytest test/external/external_test_am.py --durations=20
+        run: python -m pytest -n=auto test/test_ops.py test/test_dtype.py test/test_dtype_alu.py test/test_linearizer.py test/test_randomness.py test/test_jit.py test/test_graph.py test/test_multitensor.py test/test_hcq.py test/external/external_test_am.py --durations=20
       - name: Run TRANSCENDENTAL math
         run: TRANSCENDENTAL=2 python -m pytest -n=auto test/test_ops.py::TestOps::test_sin test/test_ops.py::TestOps::test_cos test/test_ops.py::TestOps::test_tan test/test_ops.py::TestOps::test_exp test/test_ops.py::TestOps::test_log --durations=20
       - name: Run TestOps.test_add with SQTT
@@ -717,7 +677,6 @@ 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
@@ -785,7 +744,6 @@ 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
@@ -801,8 +759,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)
@@ -816,6 +774,7 @@ 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
@@ -841,7 +800,7 @@ jobs:
       uses: ./.github/actions/setup-tinygrad
       with:
         key: osx-webgpu
-        deps: testing
+        deps: testing_minimal
         webgpu: 'true'
     - name: Test infinity math in WGSL
       run: WEBGPU=1 python -m pytest -n=auto test/test_renderer_failures.py::TestWGSLFailures::test_multiply_infinity --durations=20
@@ -851,21 +810,17 @@ jobs:
       run: npm cache clean --force
     - name: Install Puppeteer
       run: npm install puppeteer
-    # 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
-    - name: Test ONNX Runner (WEBGPU)
-      run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
+    - 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
 
   osxremote:
     name: MacOS (remote metal)
@@ -892,70 +847,32 @@ jobs:
           python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_tensor_variable.py
 
   amdremote:
-    name: Linux (remote)
+    name: Linux (remote amd)
     runs-on: ubuntu-22.04
     timeout-minutes: 20
     env:
       REMOTE: 1
+      REMOTEDEV: 'AMD'
       PYTHONPATH: ${{ github.workspace }}
+      MOCKGPU: 1
     steps:
       - name: Checkout Code
         uses: actions/checkout@v4
       - name: Setup Environment
         uses: ./.github/actions/setup-tinygrad
         with:
-          key: linux-remote
+          key: linux-remote-amd
           deps: testing_minimal
           amd: 'true'
           llvm: 'true'
-          opencl: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV="$2" \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$3" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-amd-1" "AMD" 6667
-          start_server "remote-server-amd-2" "AMD" 6668
-          start_server "remote-server-gpu" "GPU" 7667
-          start_server "remote-server-cpu" "CPU" 8667
       - name: Check Device.DEFAULT and print some source
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
         run: |
           python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
           python -c "from tinygrad import Device; assert Device.default.properties.real_device == 'AMD', Device.default.properties.real_device"
           DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
-      - name: Run REMOTE=1 Test (AMD)
-        env:
-          HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+      - 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_remote.py test/test_tensor_variable.py
-      - name: Run REMOTE=1 Test (GPU)
-        env:
-          HOST: 127.0.0.1:7667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_image_dtype.py test/test_jit.py
-          IMAGE=2 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
-      - name: Run REMOTE=1 Test (CPU)
-        env:
-          HOST: 127.0.0.1:8667*6
-        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_multitensor.py
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-amd-1 --no-pager
-          journalctl --user -u remote-server-amd-2 --no-pager
-          journalctl --user -u remote-server-gpu --no-pager
-          journalctl --user -u remote-server-cpu --no-pager
+          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
 
   osxtests:
     strategy:
@@ -996,7 +913,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        backend: [llvm, cpu, webgpu]
+        backend: [llvm, cpu]
 
     name: Windows (${{ matrix.backend }})
     runs-on: windows-latest
@@ -1009,12 +926,11 @@ jobs:
       - name: Setup Environment
         uses: ./.github/actions/setup-tinygrad
         with:
-          key: windows-${{ matrix.backend }}-minimal
+          key: windows-minimal
           deps: testing_unit
-          pydeps: ${{ matrix.backend == 'webgpu' && 'dawn-python' || '' }}
       - name: Set env
         shell: bash
-        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1' || matrix.backend == 'webgpu' && 'WEBGPU=1'}}" >> $GITHUB_ENV
+        run:  printf "${{ matrix.backend == 'llvm' && 'LLVM=1' || matrix.backend == 'cpu' && 'CPU=1'}}" >> $GITHUB_ENV
       - name: Run unit tests
         if: matrix.backend=='llvm'
         run: python -m pytest -n=auto test/unit/ --ignore=test/unit/test_disk_tensor.py --ignore=test/unit/test_elf.py --ignore=test/unit/test_tar.py

README.md

@@ -39,8 +39,9 @@ Try a matmul. See how, despite the style, it is fused into one kernel with the p
 
 ```sh
 DEBUG=3 python3 -c "from tinygrad import Tensor;
-N = 1024; a, b = Tensor.empty(N, N), Tensor.empty(N, N);
-(a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2).realize()"
+N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
+c = (a.reshape(N, 1, N) * b.T.reshape(1, N, N)).sum(axis=2);
+print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
 ```
 
 And we can change `DEBUG` to `4` to see the generated code.

autogen_stubs.sh

@@ -118,9 +118,7 @@ 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 \
@@ -149,7 +147,6 @@ generate_nv() {
     $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlc36f.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrlcb33.h \
     $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrla06c.h \
-    $NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl/ctrl90f1.h \
     --clang-args="-include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/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_gpu.py
   fixup $BASE/nv_gpu.py
@@ -167,30 +164,8 @@ generate_nv() {
   sed -n '1i\
 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
+
   python3 -c "import tinygrad.runtime.autogen.nv_gpu"
-
-  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 \
-    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/fsp/fsp_nvdm_format.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.nv"
 }
 
 generate_amd() {
@@ -415,8 +390,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_2.py
-  fixup $BASE/am/smu_v14_0_2.py
+    -o $BASE/am/smu_v14_0_3.py
+  fixup $BASE/am/smu_v14_0_3.py
 }
 
 generate_sqtt() {
@@ -444,7 +419,7 @@ generate_libusb() {
     -o $BASE/libusb.py
 
   fixup $BASE/libusb.py
-  sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libusb.py
+  sed -i "s\import ctypes\import ctypes, os\g" $BASE/libusb.py
   sed -i "s/FIXME_STUB/libusb/g" "$BASE/libusb.py"
   sed -i "s/libusb_le16_to_cpu = libusb_cpu_to_le16//g" "$BASE/libusb.py"
   sed -i "s/FunctionFactoryStub()/None if (lib_path:=os.getenv('LIBUSB_PATH', ctypes.util.find_library('usb-1.0'))) is None else ctypes.CDLL(lib_path)/g" "$BASE/libusb.py"
@@ -461,7 +436,6 @@ 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

@@ -7,30 +7,28 @@
 
 print("******** first, the runtime ***********")
 
-from tinygrad.runtime.ops_cpu import ClangJITCompiler, CPUDevice, CPUProgram
-
-cpu = CPUDevice()
+from tinygrad.runtime.ops_cpu import ClangJITCompiler, MallocAllocator, CPUProgram
 
 # allocate some buffers
-out = cpu.allocator.alloc(4)
-a = cpu.allocator.alloc(4)
-b = cpu.allocator.alloc(4)
+out = MallocAllocator.alloc(4)
+a = MallocAllocator.alloc(4)
+b = MallocAllocator.alloc(4)
 
 # load in some values (little endian)
-cpu.allocator._copyin(a, memoryview(bytearray([2,0,0,0])))
-cpu.allocator._copyin(b, memoryview(bytearray([3,0,0,0])))
+MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
+MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))
 
 # compile a program to a binary
 lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
 
 # create a runtime for the program
-fxn = cpu.runtime("add", lib)
+fxn = CPUProgram("add", lib)
 
 # run the program
 fxn(out, a, b)
 
 # check the data out
-print(val := cpu.allocator._as_buffer(out).cast("I").tolist()[0])
+print(val := MallocAllocator._as_buffer(out).cast("I").tolist()[0])
 assert val == 5
 
 
@@ -48,7 +46,7 @@ from tinygrad.shape.shapetracker import ShapeTracker
 out = Buffer(DEVICE, 1, dtypes.int32).allocate()
 a = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b = Buffer(DEVICE, 1, dtypes.int32).allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
-# NOTE: a._buf is the same as the return from cpu.allocator.alloc
+# NOTE: a._buf is the same as the return from MallocAllocator.alloc
 
 # describe the computation
 buf_1 = UOp(Ops.DEFINE_GLOBAL, dtypes.int32.ptr(), (), 1)
@@ -61,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_program, CompiledRunner
-program = get_program(s, Device[DEVICE].renderer)
+from tinygrad.engine.realize import get_kernel, CompiledRunner
+kernel = get_kernel(Device[DEVICE].renderer, s).linearize()
 
 # compile a program (and print the source)
-fxn = CompiledRunner(program)
+fxn = CompiledRunner(kernel.to_program())
 print(fxn.p.src)
 # NOTE: fxn.clprg is the CPUProgram
 
@@ -80,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.schedule.kernelize import get_kernelize_map
+from tinygrad.engine.grouper import get_kernelize_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)

docs/developer/hcq.md

@@ -52,7 +52,7 @@ Signals are device-dependent structures used for synchronization and timing in H
 The following Python code demonstrates the usage of signals:
 
 ```python
-signal = your_device.new_signal(value=0)
+signal = your_device.signal_t()
 
 HWQueue().timestamp(signal) \
          .signal(signal, value_to_fire) \

docs/developer/layout.md

@@ -1,66 +0,0 @@
-# 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/schedule
-
-Group UOps into kernels.
-
-::: tinygrad.schedule.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.ops import UOp, Ops
+from tinygrad.uop import UOp, Ops
 
 # first, we'll construct some const UOps
 a = UOp(Ops.CONST, dtypes.int, arg=2)

docs/runtime.md

@@ -12,7 +12,7 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
 | [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
 | [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.3.0). |
+| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
 
 ## Interoperability
 

docs/tensor/elementwise.md

@@ -35,7 +35,6 @@ 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

docs/tensor/movement.md

@@ -26,6 +26,5 @@
 ::: tinygrad.Tensor.transpose
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
-::: tinygrad.Tensor.diag
 ::: tinygrad.Tensor.roll
 ::: tinygrad.Tensor.rearrange

examples/benchmark_onnx.py

@@ -1,10 +1,11 @@
-import sys, time
+import sys, time, pickle
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate
 
 def load_onnx_model(onnx_file):
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx_load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)
   run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(None) for k,v in kwargs.items()}).values())), prune=True, optimize=True)
   return run_onnx_jit, run_onnx.graph_inputs
 

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
+from tinygrad import Tensor, nn, Device, dtypes
 from tinygrad.helpers import Timing, colored, getenv, fetch
 from extra.models.llama import Transformer, convert_from_huggingface, fix_bf16
 from sentencepiece import SentencePieceProcessor

examples/compile_tensorflow.py

@@ -10,7 +10,6 @@ import tensorflow as tf
 import tf2onnx
 from tinygrad.frontend.onnx import OnnxRunner
 from tinygrad.tensor import Tensor
-from tinygrad.helpers import to_mv
 from extra.export_model import export_model_clang, compile_net, jit_model
 
 def get_uncompiled_model2(dataset_size=32, output_size=4):
@@ -26,7 +25,7 @@ class TinyOnnx:
   def __init__(self, keras_model):
     input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
     onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-    self.run_onnx = OnnxRunner(Tensor(onnx_model.SerializeToString(), device="PYTHON"))
+    self.run_onnx = OnnxRunner(onnx_model)
 
   def forward(self, x):
     return self.run_onnx({"x": x}, debug=False)['predictions']
@@ -48,8 +47,8 @@ def compile_onnx_model(onnx_model):
   cprog.append("void initialize(float *weights) {")
   weights = bytes()
   for name,cl in bufs_to_save.items():
-    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {cl._buf.size});")
-    weights += bytes(to_mv(cl._buf.va_addr, cl._buf.size))
+    cprog.append(f"memcpy({name}, weights + {len(weights)//4}, {len(cl._buf)*4});")
+    weights += bytes(cl._buf)
   cprog.append("}")
 
   # write the weights to disk

examples/handcode_opt.py

@@ -2,14 +2,13 @@ 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.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from tinygrad.uop.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.engine.search import beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
 from extra.optimization.helpers import time_linearizer
-from tinygrad.engine.realize import get_program
 
 def get_sched_resnet():
   mdl = ResNet50()
@@ -109,7 +108,7 @@ if __name__ == "__main__":
     choices = []
     for lin, nm in lins:
       tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=get_program(lin.get_optimized_ast(), lin.opts)).estimates.ops
+      ops = (prg:=lin.to_program()).estimates.ops
       gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
       choices.append((tm, gflops, lin, prg, nm))
 
@@ -122,7 +121,7 @@ if __name__ == "__main__":
     if getenv("SRC"):
       print(si.ast)
       print(lin.applied_opts)
-      print(get_program(lin.get_optimized_ast(), lin.opts).src)
+      print(lin.to_program().src)
     total_tm += tm
     running_gflops += gflops * tm
     if (key := str([str(m) for m in si.metadata])) not in usage: usage[key] = (0, 0)

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, Variable
-from tinygrad.nn.state import get_state_dict
+from tinygrad import nn, dtypes, Tensor, Device, GlobalCounters, TinyJit
+from tinygrad.nn.state import get_state_dict, get_parameters
 from tinygrad.nn import optim
 from tinygrad.helpers import Context, BEAM, WINO, getenv, colored, prod
 from extra.bench_log import BenchEvent, WallTimeEvent
@@ -145,7 +145,6 @@ hyp = {
   },
 }
 
-@Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1))
 def train_cifar():
 
   def set_seed(seed):
@@ -202,37 +201,24 @@ 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))
 
-  # 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 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)))
+    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))
 
-  def cutmix(X, Y, order, mask_size=3):
+  def cutmix(X:Tensor, Y:Tensor, mask_size=3):
+    # fill the square with randomly selected images from the same batch
     mask = make_square_mask(X.shape, mask_size)
-    X_patch, Y_patch = X[order], Y[order]
+    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_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
@@ -240,16 +226,28 @@ def train_cifar():
       st = time.monotonic()
       X, Y = X_in, Y_in
       if is_train:
-        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
+        # 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()
       et = time.monotonic()
       print(f"shuffling {'training' if is_train else 'test'} dataset in {(et-st)*1e3:.2f} ms ({epoch=})")
-
-      vi = Variable("i", 0, (full_batches := (X.shape[0] // BS) * BS) - BS)
-      for i in range(0, full_batches, BS):
+      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)
         step += 1
-        vib = vi.bind(i)
-        yield X[vib:vib+BS], Y[vib:vib+BS]
+        yield x, y
       epoch += 1
       if not is_train: break
 

examples/llama3.py

@@ -157,11 +157,7 @@ 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
@@ -240,7 +236,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", "405B"], default="1B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
   parser.add_argument("--shard", type=int, default=1, help="Shard the model across multiple devices")
   parser.add_argument("--quantize", choices=["int8", "nf4", "float16"], help="Quantization method")
   parser.add_argument("--no_api", action="store_true", help="Disable the api and run a cli test interface")
@@ -248,7 +244,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=float, default=0.85, help="Temperature")
+  parser.add_argument("--temperature", type=int, 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
+from tinygrad import Device, nn, Tensor, dtypes, Variable
 Device.DEFAULT = "CPU"
 from train_gpt2 import GPT, GPTConfig
-from tinygrad.helpers import dedup, flatten, getenv, GlobalCounters, to_function_name
-from tinygrad.engine.realize import get_kernel
+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.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,5 +1,4 @@
-import math
-from tinygrad import dtypes
+from tinygrad import Tensor, dtypes
 from tinygrad.nn.optim import Optimizer
 
 from extra.lr_scheduler import LR_Scheduler
@@ -21,19 +20,3 @@ 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, string
+import re
+import string
 from collections import Counter
-from tinygrad import Tensor
 
 def levenshtein(a, b):
   n, m = len(a), len(b)
@@ -59,11 +59,3 @@ 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, AdamW
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, Adam
 
 from extra.lr_scheduler import LRSchedulerGroup
 from examples.mlperf.helpers import get_training_state, load_training_state
@@ -914,26 +914,18 @@ def train_rnnt():
   pass
 
 @TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, GPUS, grad_acc:int, **kwargs):
+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])
   optimizer.zero_grad()
 
-  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}"]
+  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 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)
+  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()
@@ -1007,19 +999,16 @@ def train_bert():
     MLLOGGER = None
 
   # ** hyperparameters **
-  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
+  BS                 = config["GLOBAL_BATCH_SIZE"]      = getenv("BS", 11 * len(GPUS) if dtypes.default_float in (dtypes.float16, dtypes.bfloat16) else 8 * len(GPUS))
   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(GBS/96))
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.000175 * math.sqrt(BS/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 // GBS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3600000 // BS)
   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 * GBS + 3000000) / 25000) * 25000) / GBS)) # Round down
+  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
   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")
@@ -1077,7 +1066,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 global batch size {GBS} for one epoch with {train_steps} steps")
+  print(f"training with batch size {BS} for one epoch with {train_steps} steps")
 
   # log mlperf hparams
   if MLLOGGER:
@@ -1126,11 +1115,11 @@ def train_bert():
   # ** train loop **
   wc_start = time.perf_counter()
 
-  i, train_data = start_step, [next(train_it) for _ in range(grad_acc)]
+  i, train_data = start_step, next(train_it)
 
   if RUNMLPERF:
     if MLLOGGER:
-      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*GBS, metadata={"epoch_num": i*GBS})
+      MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})
 
   while train_data is not None and i < train_steps and not achieved:
     if getenv("TRAIN", 1):
@@ -1139,13 +1128,14 @@ def train_bert():
       st = time.perf_counter()
       GlobalCounters.reset()
       with WallTimeEvent(BenchEvent.STEP):
-        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)
+        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)
 
         pt = time.perf_counter()
 
         try:
-          next_data = [next(train_it) for _ in range(grad_acc)]
+          next_data = next(train_it)
         except StopIteration:
           next_data = None
 
@@ -1166,7 +1156,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)*GBS})
+                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
 
       train_data, next_data = next_data, None
       i += 1
@@ -1181,7 +1171,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*GBS, "step_num": i})
+        MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "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 = []
@@ -1231,11 +1221,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)*GBS})
+                    "eval/clsf_accuracy": avg_clsf_acc, "eval/forward_time": avg_fw_time, "epoch": (i+1)*BS})
 
       if MLLOGGER and RUNMLPERF:
-        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})
+        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})
 
       # save model if achieved target
       if not achieved and avg_lm_acc >= target:
@@ -1250,10 +1240,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 * GBS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
+        print(f"Reference Convergence point reached after {i * BS} datasamples and {hours}h{minutes}m{seconds:.2f}s.")
         achieved = True
         if MLLOGGER and RUNMLPERF:
-          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*GBS, metadata={"epoch_num": i*GBS})
+          MLLOGGER.event(key=mllog_constants.EPOCH_STOP, value=i*BS, metadata={"epoch_num": i*BS})
           MLLOGGER.end(key=mllog_constants.RUN_STOP, metadata=dict(status=mllog_constants.SUCCESS))
         # stop once hitting the target
         break
@@ -1281,78 +1271,12 @@ 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 * GBS, "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 * BS, "step_num": i, "first_step_num": i+1})
         previous_step = i
 
-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=}")
+def train_maskrcnn():
+  # TODO: Mask RCNN
+  pass
 
 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
+export BENCHMARK=10 BERT_LAYERS=2 DEBUG=2
 
 python3 examples/mlperf/model_train.py

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

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

@@ -5,26 +5,29 @@ if "IMAGE" not in os.environ: os.environ["IMAGE"] = "2"
 if "NOLOCALS" not in os.environ: os.environ["NOLOCALS"] = "1"
 if "JIT_BATCH_SIZE" not in os.environ: os.environ["JIT_BATCH_SIZE"] = "0"
 
-from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device, dtypes
+from tinygrad import fetch, Tensor, TinyJit, Context, GlobalCounters, Device
 from tinygrad.helpers import DEBUG, getenv
+from tinygrad.tensor import _from_np_dtype
 from tinygrad.engine.realize import CompiledRunner
 
 import onnx
-from tinygrad.frontend.onnx import OnnxRunner
+from onnx.helper import tensor_dtype_to_np_dtype
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx"
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
 
 def compile(onnx_file):
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx_load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)
   print("loaded model")
 
-  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
-  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
+  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
   # Float inputs and outputs to tinyjits for openpilot are always float32
-  input_types = {k:(dtypes.float32 if v is dtypes.float16 else v) for k,v in input_types.items()}
+  input_types = {k:(np.float32 if v==np.float16 else v) for k,v in input_types.items()}
   Tensor.manual_seed(100)
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in sorted(input_shapes.items())}
+  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in sorted(input_shapes.items())}
   new_inputs_numpy = {k:v.numpy() for k,v in new_inputs.items()}
   print("created tensors")
 
@@ -54,11 +57,11 @@ def compile(onnx_file):
       gated_read_image_count += ei.prg.p.src.count("?read_image")
   print(f"{kernel_count=},  {read_image_count=}, {gated_read_image_count=}")
   if (allowed_kernel_count:=getenv("ALLOWED_KERNEL_COUNT", -1)) != -1:
-    assert kernel_count == allowed_kernel_count, f"different kernels! {kernel_count=}, {allowed_kernel_count=}"
+    assert kernel_count <= allowed_kernel_count, f"too many kernels! {kernel_count=}, {allowed_kernel_count=}"
   if (allowed_read_image:=getenv("ALLOWED_READ_IMAGE", -1)) != -1:
     assert read_image_count == allowed_read_image, f"different read_image! {read_image_count=}, {allowed_read_image=}"
   if (allowed_gated_read_image:=getenv("ALLOWED_GATED_READ_IMAGE", -1)) != -1:
-    assert gated_read_image_count == allowed_gated_read_image, f"different gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
+    assert gated_read_image_count <= allowed_gated_read_image, f"too many gated read_image! {gated_read_image_count=}, {allowed_gated_read_image=}"
 
   with open(OUTPUT, "wb") as f:
     pickle.dump(run_onnx_jit, f)

examples/openpilot/compile4.py

@@ -1,8 +1,8 @@
-import sys
-from tinygrad import Tensor, fetch, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp
+import sys, onnx
+from tinygrad import Tensor, fetch, GlobalCounters
+from tinygrad.uop import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.schedule.kernelize import get_kernelize_map
+from tinygrad.engine.grouper import get_kernelize_map
 from tinygrad.engine.schedule import create_schedule_with_vars
 from tinygrad.engine.realize import run_schedule
 
@@ -12,10 +12,12 @@ OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/comm
 OUTPUT = sys.argv[2] if len(sys.argv) > 2 else "/tmp/openpilot.pkl"
 
 if __name__ == "__main__":
+  fn = fetch(OPENPILOT_MODEL)
   onnx_file = fetch(OPENPILOT_MODEL)
-  run_onnx = OnnxRunner(onnx_file)
+  onnx_model = onnx.load(onnx_file)
+  run_onnx = OnnxRunner(onnx_model)
 
-  inputs = run_onnx.get_empty_input_data("npy", dtypes.float32)
+  inputs = run_onnx.get_empty_input_data("npy")
   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()]
@@ -35,12 +37,12 @@ if __name__ == "__main__":
   independent = UOp.sink(*independent_set.keys())
   kernelized = get_kernelize_map(independent)
   independent = independent.substitute(kernelized)
-  schedule, var_vals = create_schedule_with_vars(independent)
+  schedule, var_vals, becomes_map = create_schedule_with_vars(independent)
   run_schedule(schedule)
 
   print("**** real ****")
   GlobalCounters.reset()
-  out.uop = root.substitute(kernelized)
+  out.uop = root.substitute(kernelized).substitute(becomes_map)
   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  # noqa: F401
+    import tinygrad.frontend.torch
     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
+from tinygrad.nn.state import safe_load, load_state_dict, get_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, fetch
+from tinygrad.helpers import getenv
 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,6 +19,10 @@ 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")
@@ -93,7 +97,7 @@ class ContentVec:
     return res, padding_mask
   @classmethod
   def load_from_pretrained(cls, checkpoint_path:str, checkpoint_url:str) -> ContentVec:
-    fetch(checkpoint_url, checkpoint_path)
+    download_if_not_present(checkpoint_path, checkpoint_url)
     cfg = load_fairseq_cfg(checkpoint_path)
     enc = cls(cfg.model)
     _ = load_checkpoint_enc(checkpoint_path, enc, None)
@@ -320,9 +324,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:
-    fetch(config_url, config_path)
+    download_if_not_present(config_path, config_url)
     hps = get_hparams_from_file(config_path)
-    fetch(weights_url, weights_path)
+    download_if_not_present(weights_path, weights_url)
     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}")
@@ -598,7 +602,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: fetch(DEMO_URL, DEMO_PATH)
+  if audio_path == DEMO_PATH: download_if_not_present(DEMO_PATH, DEMO_URL)
   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
+from tinygrad.helpers import trange, DEBUG
 
 # STEPS=70 python3 examples/stunning_mnist.py
 # NOTE: it's broken with STACK=1, why?

examples/torch_cuda_kernel.py

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

examples/vgg7.py

@@ -2,6 +2,8 @@ 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, fetch
+from tinygrad.helpers import getenv, DEBUG, fetch
 
 import numpy as np
 import librosa
@@ -321,7 +321,7 @@ if __name__ == "__main__":
         log_spec = prep_audio(total.reshape(1, -1), model.batch_size, truncate=True)
         encoded_audio = model.encoder.encode(Tensor(log_spec))
       # pass the previously inferred tokens as 'prefix' - https://github.com/openai/whisper/discussions/117#discussioncomment-3727051
-      out = model.decoder(Tensor([lst]), 0, encoded_audio).realize()
+      out = model.decoder(Tensor([lst]), 0, encoded_audio, streaming=True).realize()
       idx = int(out[0,-1].argmax().numpy().item())
       lst.append(idx)
       dec = enc.decode(lst)

examples/yolov8-onnx.py

@@ -2,12 +2,14 @@
 import os
 from ultralytics import YOLO
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+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():
   model = YOLO("yolov8n-seg.pt")
   model.export(format="onnx", imgsz=[480,640])
-run_onnx = OnnxRunner("yolov8n-seg.onnx")
+onnx_model = onnx_load(open("yolov8n-seg.onnx", "rb"))
+run_onnx = OnnxRunner(onnx_model)
 run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)

examples/yolov8.py

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

extra/amdpci/am_smi.py

@@ -8,6 +8,8 @@ from tinygrad.runtime.support.hcq import MMIOInterface
 from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager, AMPageTableEntry
 from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
 
+AM_VERSION = 0xA0000005
+
 def bold(s): return f"\033[1m{s}\033[0m"
 
 def trim(s:str, length:int) -> str:
@@ -71,12 +73,21 @@ class AMSMI(AMDev):
     self._run_discovery()
     self._build_regs()
 
-    if self.reg("regSCRATCH_REG7").read() != AMDev.Version:
+    if self.reg("regSCRATCH_REG7").read() != AM_VERSION:
       raise Exception(f"Unsupported AM version: {self.reg('regSCRATCH_REG7').read():x}")
 
-    self.is_booting = True
-    self.init_sw(smi_dev=True)
+    self.is_booting, self.smi_dev = True, True
     self.partial_boot = True # do not init anything
+    self.mm = AMMemoryManager(self, self.vram_size)
+
+    # Initialize IP blocks
+    self.soc:AM_SOC = AM_SOC(self)
+    self.gmc:AM_GMC = AM_GMC(self)
+    self.ih:AM_IH = AM_IH(self)
+    self.psp:AM_PSP = AM_PSP(self)
+    self.smu:AM_SMU = AM_SMU(self)
+
+    for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]: ip.init_sw()
 
   def read_pci_state(self):
     with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
@@ -125,7 +136,7 @@ class SMICtx:
         if d.pci_state == "D0": d._init_from_d0()
         os.system('clear')
 
-      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AMDev.Version:
+      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AM_VERSION:
         self.devs.remove(d)
         self.opened_pcidevs.remove(d.pcibus)
         os.system('clear')
@@ -284,8 +295,8 @@ if __name__ == "__main__":
           while True:
             try: pid = subprocess.check_output(['sudo', 'lsof', '-t', dev]).decode('utf-8').split('\n')[0]
             except subprocess.CalledProcessError: break
-            if stopped_pids[pid] > 0: time.sleep(0.1)
-            if stopped_pids[pid] == 64:
+            if stopped_pids[pid] > 0: time.sleep(0.5)
+            if stopped_pids[pid] == 10:
               print(f"{dev[8:-5]}: can't stop process {pid}, exitting")
               exit(1)
 

extra/assembly/assembly.py

@@ -1,5 +1,5 @@
 from typing import Tuple, List, NamedTuple, Any, Dict, Optional, Union, DefaultDict, cast
-from tinygrad.opt.kernel import Ops, MemOp, UOp
+from tinygrad.codegen.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_REG:
+    elif uop == Ops.DEFINE_ACC:
       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.opt.kernel import Ops, UOp
+from tinygrad.codegen.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.opt.kernel import Ops, UOp
+from tinygrad.codegen.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.opt.kernel import Ops
+from tinygrad.codegen.kernel import Ops
 from tinygrad.uop.ops import BinaryOps, UnaryOps, TernaryOps
 from tinygrad.runtime.ops_gpu import ROCM_LLVM_PATH
 

extra/backends/rdna.py

@@ -15,8 +15,8 @@ def uops_to_rdna(function_name:str, uops:UOpGraph) -> str:
         u.vin = tuple(n if x == o else x for x in u.vin)
     # pointer indexing
     if u.uop in {UOps.LOAD, UOps.STORE} and u.vin[0].dtype.itemsize > 1:
-      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_at=uops.uops.index(u))
-      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_at=uops.uops.index(u))
+      val = UOp(UOps.CONST, dtypes.int, tuple(), arg=u.vin[0].dtype.itemsize, insert_before=uops.uops.index(u))
+      ptr = UOp(UOps.ALU, dtypes.int, (u.vin[1], val), arg=BinaryOps.MUL, insert_before=uops.uops.index(u))
       u.vin = (u.vin[0], ptr) + u.vin[2:]
   #uops.print()
 

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.opt.kernel import  UOp, Ops
+from tinygrad.codegen.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_REG: kk(f"{ssa(u, 'acc')} = {define_scalar(local_size, dtype, args).replace('//', '/')}")
+    elif uop == Ops.DEFINE_ACC: 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/export_model.py

@@ -4,7 +4,7 @@ from tinygrad.renderer import ProgramSpec
 from tinygrad.tensor import Device, Tensor
 from tinygrad.engine.jit import TinyJit
 from tinygrad.nn.state import get_state_dict
-from tinygrad.helpers import Context, to_mv
+from tinygrad.helpers import Context
 from tinygrad.dtype import dtypes
 from tinygrad.uop.ops import Ops
 import json
@@ -68,7 +68,7 @@ def export_model_clang(functions:Dict[str,str], statements:Dict[str,Tuple[str,in
 
   if not wasm:
     for name,cl in bufs_to_save.items():
-      weight = ''.join(["\\x%02X"%x for x in bytes(to_mv(cl._buf.va_addr, cl._buf.size))])
+      weight = ''.join(["\\x%02X"%x for x in bytes(cl._buf)])
       cprog.append(f"unsigned char {name}_data[] = \"{weight}\";")
     cprog += [f"{dtype_map[dtype]} {name}[{len}];" if name not in bufs_to_save else f"{dtype_map[dtype]} *{name} = ({dtype_map[dtype]} *){name}_data;" for name,(len,dtype,_key) in bufs.items() if name not in input_names+output_names]
     cprog += [f"void net({forward_args}) {{"] + [f"{name}({', '.join(args)});" for (name, args, _global_size, _local_size) in statements] + ["}"]

extra/gemm/amd_matmul.py

@@ -1,39 +1,96 @@
 # kernel8_batched_gmem.s from https://seb-v.github.io/optimization/update/2025/01/20/Fast-GPU-Matrix-multiplication.html
-# sudo PATH=/opt/homebrew/Cellar/llvm/20.1.6/bin:$PATH AMD_LLVM=0 AMD=1 DEBUG=2 python3 extra/gemm/amd_matmul.py
 import pathlib
+import numpy as np
 from dataclasses import replace
-from tinygrad import Tensor, Device, Context, GlobalCounters
+from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, Ops, UOp
+
+# TODO: on METAL for `DEBUG=4 python3 extra/gemm/amd_matmul.py`
+#  * fix load grouping (like float4). idk why it's not working, need new devectorizer (this is a Monday project)
+#  * DONE - remove extra barrier
+#  * DONE (moved Ops.ADD) - fix load order to be in order (the +0 one is last!)
+#  * explore async (fast) global load -> local store
+#  * why is TC=3 broken for 4096x4096?
+#  * write syntactic sugar for these local additions + use it in tensor core kernel.py
 
 N = 4096
+LN = 16
 run_count = 5
 
-if __name__ == "__main__":
-  ast = (Tensor.empty(N, N)@Tensor.empty(N, N)).schedule()[-1].ast
-  prg = get_program(ast, Device.default.renderer)
-
-  if getenv("ASM") == 1:
-    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel8_batched_gmem.s").read_text()
-    prgfast = replace(prg, name="kernel", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
-  elif getenv("ASM") == -1:
-    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel3_registers.cpp").read_text()
-    prgfast = replace(prg, name="kernel3_registers", src=src, global_size=[N//128, N//128, 1], local_size=[256, 1, 1])
+from tinygrad.shape.shapetracker import ShapeTracker, View
+def transform_load(ctx:tuple[Kernel, set[UOp]], x:UOp):
+  if x.src[0].op is not Ops.DEFINE_GLOBAL: return None
+  if x in ctx[1]: return None
+  print(ctx[0].colored_shape())
+  ctx[1].add(x)
+  input_st: ShapeTracker = x.src[1].arg
+  #strides = input_st.real_strides()
+  #strides = (0,0)+strides[2:]
+  if input_st.real_strides()[2] == 0:
+    perm = (0,1,5,3,4,2)
+    strides = (0,0,LN*4,4,0,0,1,0)
+  elif input_st.real_strides()[3] == 0:
+    perm = (0,1,2,5,4,3)
+    strides = (0,0,LN*4,4,0,0,0,1)
   else:
-    src = (pathlib.Path(__file__).parent / "amd_seb" / "kernel5_lds_optim.cpp").read_text()
-    prgfast = replace(prg, name="kernel5_lds_optim", src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
-  runner = CompiledRunner(prgfast)
+    return None
+  if len(input_st.shape) == 8:
+    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1,4,4), strides),))
+    perm = perm + (6,7)
+  else:
+    local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
+  #local_st = ShapeTracker(views=(View.create((1,1,LN,LN,1,1)),))
+  load_st = local_st.permute(perm)
+  input_st = input_st.permute(perm)
+  lcl = UOp(Ops.DEFINE_LOCAL, x.dtype.ptr(local_st.real_size(), local=True), (), f"temp{x.src[0].arg}")
+  global_load = x.replace(src=(x.src[0], input_st.to_uop()))
+  ret = UOp(Ops.STORE, src=(lcl, local_st.to_uop(), global_load))
+  return UOp(Ops.LOAD, x.dtype, src=(lcl, load_st.to_uop(), ret))
 
-  a = Tensor.randn(N, N).realize()
-  b = Tensor.randn(N, N).realize()
-  c = Tensor.zeros(N, N).contiguous().realize()
+local_loads_pm = PatternMatcher([
+  (UPat(Ops.LOAD, name="x"), transform_load),
+])
 
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    for _ in range(run_count): tc = (a@b).realize()
+def ast_transform(k, ast):
+  #return ast
+  ast = graph_rewrite(ast, local_loads_pm, ctx=(k, set()))
+  #ast = ast.replace(arg=replace(ast.arg, upcasted=0))
+  print(ast)
+  return ast
 
-  GlobalCounters.reset()
-  ei = ExecItem(runner, [a.uop.buffer, b.uop.buffer, c.uop.buffer])
+if __name__ == "__main__":
+  rng = np.random.default_rng()
+  a = Tensor(na:=rng.random((4096, 4096), dtype=np.float32)).realize()
+  b = Tensor(nb:=rng.random((4096, 4096), dtype=np.float32)).realize()
+  c = a @ b
+  si = c.schedule()[-1]
+  k = Kernel(si.ast, opts=Device[Device.DEFAULT].renderer)
+  #opts = [Opt(op=OptOps.LOCAL, axis=1, arg=16),
+  #        Opt(op=OptOps.LOCAL, axis=0, arg=8),
+  #        Opt(op=OptOps.UPCAST, axis=2, arg=4),
+  #        Opt(op=OptOps.UPCAST, axis=1, arg=4),
+  #        Opt(op=OptOps.UPCAST, axis=0, arg=2)]
+  #opts = [Opt(op=OptOps.UPCAST, axis=1, arg=4),
+  #        Opt(op=OptOps.UPCAST, axis=0, arg=4),
+  #        Opt(op=OptOps.LOCAL, axis=1, arg=8),
+  #        Opt(op=OptOps.LOCAL, axis=0, arg=4)]
+  opts = [Opt(op=OptOps.UNROLL, axis=0, arg=LN),
+          #Opt(op=OptOps.UPCAST, axis=0, arg=4),
+          #Opt(op=OptOps.UPCAST, axis=1, arg=4),
+          Opt(op=OptOps.LOCAL, axis=1, arg=LN),
+          Opt(op=OptOps.LOCAL, axis=0, arg=LN)]
+  k.apply_opts(opts)
+  prg = k.to_program(ast_transform=ast_transform)
+  if getenv("FAST", 1) and Device.DEFAULT == "AMD":
+    #src = (pathlib.Path(__file__).parent / "fp32_sgemm_amd" / "src" / "kernel8_batched_gmem.s").read_text()
+    src = (pathlib.Path(__file__).parent / "kernel8_batched_gmem.s").read_text()
+    prg = replace(prg, src=src, global_size=[N//128, N//128, 1], local_size=[128, 1, 1])
+  print(prg.global_size, prg.local_size)
+  ei = ExecItem(CompiledRunner(prg), [x.ensure_allocated() for x in si.bufs], si.metadata)
   with Context(DEBUG=2):
     for _ in range(run_count): ei.run(wait=True)
-  print(f"custom  {(c-tc).square().mean().item()}")
+  nc = c.numpy()
+  np.testing.assert_allclose(na@nb, nc, rtol=1e-5)

extra/gemm/amd_seb/kernel3_registers.cpp

@@ -1,142 +0,0 @@
-typedef long unsigned int size_t;
-extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
-extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
-struct Dim3 { size_t x, y, z; };
-#define __shared__ __attribute__((shared, aligned(16)))
-__attribute__((device)) inline void __syncthreads() {
-  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
-  __builtin_amdgcn_s_barrier();
-  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
-}
-
-#define BLOCK_SIZE 256
-extern "C" __attribute__((global)) void kernel3_registers(float *a, float *b, float *c)
-{
-  constexpr int N = 4096;
-  constexpr float alpha = 1.0;
-  constexpr float beta = 0.0;
-
-  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
-  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
-
-  // Block Tile size
-  constexpr int BN = 128;
-  constexpr int BM = 128;
-  // Number of Row or column we read per batch
-  constexpr int BK = 8;
-
-  // Thread Tile size
-  constexpr int TN = 4;
-  constexpr int TM = 4;
-
-  constexpr int nbWaves = BLOCK_SIZE / 32;
-  // Wave Tile size
-  constexpr int WN = 64;
-  constexpr int WM = BN * BM / nbWaves / WN;
-
-  // Number of wave on X & Y axis in the Block tile
-  constexpr int nbWaveX = BN / WN;
-  constexpr int nbWaveY = BM / WM;
-
-  const int waveIndex = threadIdx.x / 32;
-  const int waveIdx = waveIndex % nbWaveX;
-  const int waveIdy = waveIndex / nbWaveX;
-  const int indexInWave = threadIdx.x % 32;
-
-  // A wave is a block of 8x4 of the output matrix
-  constexpr int nbThreadXPerWave = 8;
-  constexpr int nbThreadYPerWave = 4;
-
-  // Thread coordinates in Wave
-  const int idxInWave = indexInWave % nbThreadXPerWave;
-  const int idyInWave = indexInWave / nbThreadXPerWave;
-
-  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
-  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
-
-  // Wave Sub-tile size
-  constexpr int SUBWN = WN / nbIterWaveN;
-  constexpr int SUBWM = WM / nbIterWaveM;
-
-  // Thread mapping to read BKxBN block from A
-  int rAIdx = threadIdx.x % BK;
-  int rAIdy = threadIdx.x / BK;
-  // Thread mapping to read BNxBK block from B
-  int rBIdx = threadIdx.x % BN;
-  int rBIdy = threadIdx.x / BN;
-
-  constexpr int strideReadB = BLOCK_SIZE / BN;
-  constexpr int strideReadA = BLOCK_SIZE / BK;
-  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
-  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
-
-  float A_col[nbIterWaveM * TM];
-  float B_row[nbIterWaveN * TN];
-
-  __shared__ float As[BK][BM];
-  __shared__ float Bs[BK][BN];
-
-  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
-
-  // Iteration over BK blocks.
-  for (int kId = 0; kId < N; kId += BK) {
-    __syncthreads();
-
-    // We populate the Shared Memory with Ks row and columns
-    for (int i = 0; i < nbReadsB; i++) {
-      int index_x = BN * blockIdx.x + rBIdx;
-      int index_y = rBIdy + i * strideReadB + kId;
-      Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
-    }
-
-    for (int i = 0; i < nbReadsA; i++) {
-      int index_x = rAIdx + kId;
-      int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-      As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
-    }
-
-    __syncthreads();
-    for (int k = 0; k < BK; k++) {
-      // we cache A & B for the entire Wave tile
-      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
-        for (int i = 0; i < TN; i++) {
-          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
-          B_row[iterWave * TN + i] = Bs[k][index];
-        }
-      }
-
-      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
-        for (int i = 0; i < TM; i++) {
-          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
-          A_col[iterWave * TM + i] = As[k][index];
-        }
-      }
-
-      // we accumulate to C_regs
-      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-          for (int yt = 0; yt < TM; yt++) {
-            for (int xt = 0; xt < TN; xt++) {
-              const int x = iterWaveN * TN + xt;
-              const int y = iterWaveM * TM + yt;
-              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
-            }
-          }
-        }
-      }
-    }
-  }
-
-  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
-      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
-      for (int yt = 0; yt < TM; yt++) {
-        for (int xt = 0; xt < TN; xt++) {
-          int indexC = N * (yOut + yt) + xOut + xt;
-          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
-        }
-      }
-    }
-  }
-}

extra/gemm/amd_seb/kernel5_lds_optim.cpp

@@ -1,172 +0,0 @@
-typedef long unsigned int size_t;
-extern "C" __attribute__((device, const)) size_t __ockl_get_local_id(unsigned int);
-extern "C" __attribute__((device, const)) size_t __ockl_get_group_id(unsigned int);
-struct Dim3 { size_t x, y, z; };
-#define __shared__ __attribute__((shared, aligned(16)))
-__attribute__((device)) inline void __syncthreads() {
-  __builtin_amdgcn_fence(__ATOMIC_RELEASE, "workgroup");
-  __builtin_amdgcn_s_barrier();
-  __builtin_amdgcn_fence(__ATOMIC_ACQUIRE, "workgroup");
-}
-
-#define BLOCK_SIZE 128
-extern "C" __attribute__((global)) void __attribute__((amdgpu_flat_work_group_size(1, BLOCK_SIZE)))
-kernel5_lds_optim(float *a, float *b, float *c)
-{
-  constexpr int N = 4096;
-  constexpr float alpha = 1.0;
-  constexpr float beta = 0.0;
-
-  const Dim3 blockIdx{ __ockl_get_group_id(0), __ockl_get_group_id(1), __ockl_get_group_id(2) };
-  const Dim3 threadIdx{ __ockl_get_local_id(0), __ockl_get_local_id(1), __ockl_get_local_id(2) };
-
-  // Block Tile size
-  constexpr int BN = 128;
-  constexpr int BM = 128;
-  // Number of Row or column we read per batch
-  constexpr int BK = 8;
-
-  // Thread Tile size . 4x4
-  constexpr int TN = 4;
-  constexpr int TM = 4;
-
-  constexpr int nbWaves = BLOCK_SIZE / 32;
-  // Wave Tile size
-  constexpr int WN = 128;
-  constexpr int WM = BN * BM / nbWaves / WN;
-
-  // Number of wave on X & Y axis in the Block tile
-  constexpr int nbWaveX = BN / WN;
-  constexpr int nbWaveY = BM / WM;
-
-  const int waveIndex = threadIdx.x / 32;
-  const int waveIdx = waveIndex % nbWaveX;
-  const int waveIdy = waveIndex / nbWaveX;
-  const int indexInWave = threadIdx.x % 32;
-
-  // A wave is a block of 8x4 of the output matrix
-  constexpr int nbThreadXPerWave = 8;
-  constexpr int nbThreadYPerWave = 4;
-
-  // Thread coordinates in Wave
-  const int idxInWave = indexInWave % nbThreadXPerWave;
-  const int idyInWave = indexInWave / nbThreadXPerWave;
-
-  constexpr int nbIterWaveN = WN / (nbThreadXPerWave * TN);
-  constexpr int nbIterWaveM = WM / (nbThreadYPerWave * TM);
-
-  // Wave Sub-tile size
-  constexpr int SUBWN = WN / nbIterWaveN;
-  constexpr int SUBWM = WM / nbIterWaveM;
-
-  // Thread mapping to read BKxBN block from A
-  int rAIdx = threadIdx.x % BK;
-  int rAIdy = threadIdx.x / BK;
-  // Thread mapping to read BNxBK block from B
-  int rBIdx = threadIdx.x % BN;
-  int rBIdy = threadIdx.x / BN;
-
-  constexpr int strideReadB = BLOCK_SIZE / BN;
-  constexpr int strideReadA = BLOCK_SIZE / BK;
-  constexpr int nbReadsB = BN * BK / BLOCK_SIZE;
-  constexpr int nbReadsA = BM * BK / BLOCK_SIZE;
-
-  float A_col[nbIterWaveM * TM];
-  float B_row[nbIterWaveN * TN];
-
-  __shared__ float As[BK][BM+4]; // 4 padding to avoid bank conflicts
-  __shared__ float Bs[BK][BN];
-
-  float c_regs[TM * nbIterWaveM * TN * nbIterWaveN] = {0.0f};
-
-  // initial copy into shared memory
-  for (int i = 0; i < nbReadsB; i++) {
-    int index_x = BN * blockIdx.x + rBIdx;
-    int index_y = rBIdy + i * strideReadB;
-    Bs[index_y % BK][index_x % BN] = b[N * index_y + index_x];
-  }
-  for (int i = 0; i < nbReadsA; i++) {
-    int index_x = rAIdx;
-    int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-    As[(index_x % BK)][(index_y % BM)] = a[N * index_y + index_x];
-  }
-
-  __syncthreads();
-  // Iteration over BK blocks.
-  for (int kId = 0; kId < N; kId += BK) {
-    float regA[nbReadsA];
-    float regB[nbReadsB];
-    if (kId < N - BK) {
-      // We populate the Shared Memory with Ks row and columns
-      for (int i = 0; i < nbReadsB; i++) {
-        int index_x = BN * blockIdx.x + rBIdx;
-        int index_y = rBIdy + i * strideReadB + kId + BK;
-        regB[i] = b[N * index_y + index_x];
-      }
-
-      for (int i = 0; i < nbReadsA; i++) {
-        int index_x = rAIdx + kId + BK;
-        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-        regA[i] = a[N * index_y + index_x];
-      }
-    }
-
-    for (int k = 0; k < BK; k++) {
-      // we cache A & B for the entire Wave tile
-      for (int iterWave = 0; iterWave < nbIterWaveN; iterWave++) {
-        for (int i = 0; i < TN; i++) {
-          int index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i;
-          B_row[iterWave * TN + i] = Bs[k][index];
-        }
-      }
-
-      for (int iterWave = 0; iterWave < nbIterWaveM; iterWave++) {
-        for (int i = 0; i < TM; i++) {
-          int index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i;
-          A_col[iterWave * TM + i] = As[k][index];
-        }
-      }
-
-      // we accumulate to C_regs
-      for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-        for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-          for (int yt = 0; yt < TM; yt++) {
-            for (int xt = 0; xt < TN; xt++) {
-              const int x = iterWaveN * TN + xt;
-              const int y = iterWaveM * TM + yt;
-              c_regs[y * TN * nbIterWaveN + x] += A_col[y] * B_row[x];
-            }
-          }
-        }
-      }
-    }
-    __syncthreads();
-    if (kId < N - BK) {
-      for (int i = 0; i < nbReadsB; i++) {
-        int index_x = BN * blockIdx.x + rBIdx;
-        int index_y = rBIdy + i * strideReadB + kId + BK;
-        Bs[index_y % BK][index_x % BN] = regB[i]; // row
-      }
-
-      for (int i = 0; i < nbReadsA; i++) {
-        int index_x = rAIdx + kId + BK;
-        int index_y = BM * blockIdx.y + rAIdy + i * strideReadA;
-        As[(index_x % BK)][(index_y % BM)] = regA[i];
-      }
-      __syncthreads();
-    }
-  }
-
-  for (int iterWaveM = 0; iterWaveM < nbIterWaveM; iterWaveM++) {
-    for (int iterWaveN = 0; iterWaveN < nbIterWaveN; iterWaveN++) {
-      int xOut = blockIdx.x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave;
-      int yOut = blockIdx.y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave;
-      for (int yt = 0; yt < TM; yt++) {
-        for (int xt = 0; xt < TN; xt++) {
-          int indexC = N * (yOut + yt) + xOut + xt;
-          c[indexC] = beta * c[indexC] + alpha * c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)];
-        }
-      }
-    }
-  }
-}

extra/gemm/amd_uop_matmul.py

@@ -1,148 +0,0 @@
-from tinygrad import Tensor, Device, Context, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
-from tinygrad.uop.ops import Ops, UOp
-from tinygrad.dtype import AddrSpace
-
-N = 4096
-run_count = 5
-
-def hand_spec_kernel3():
-  # block tile size
-  BN = 128
-  BM = 128
-  # number of row/column we read per batch
-  BK = 8
-
-  # thread tile size 4x4
-  TN = 4
-  TM = 4
-
-  BLOCK_SIZE = 128
-  nbWaves = BLOCK_SIZE // 32
-  # wave tile size
-  WN = 128
-  WM = BN * BM // nbWaves // WN
-
-  nbWaveX = BN // WN
-  nbWaveY = BM // WM
-
-  threadIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("lidx0", BLOCK_SIZE))
-  waveIndex = threadIdx_x // 32
-  waveIdx = waveIndex % nbWaveX
-  waveIdy = waveIndex // nbWaveX
-  indexInWave = threadIdx_x % 32
-
-  nbThreadXPerWave = 8
-  nbThreadYPerWave = 4
-
-  idxInWave = indexInWave % nbThreadXPerWave
-  idyInWave = indexInWave // nbThreadXPerWave
-
-  nbIterWaveN = WN // (nbThreadXPerWave * TN)
-  nbIterWaveM = WM // (nbThreadYPerWave * TM)
-
-  SUBWN = WN // nbIterWaveN
-  SUBWM = WM // nbIterWaveM
-
-  # Thread mapping to read BKxBN block from A
-  rAIdx = threadIdx_x % BK
-  rAIdy = threadIdx_x // BK
-  # Thread mapping to read BNxBK block from B
-  rBIdx = threadIdx_x % BN
-  rBIdy = threadIdx_x // BN
-
-  strideReadB = BLOCK_SIZE // BN
-  strideReadA = BLOCK_SIZE // BK
-  nbReadsB = BN * BK // BLOCK_SIZE
-  nbReadsA = BM * BK // BLOCK_SIZE
-
-  blockIdx_x = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx0", N//BN))
-  blockIdx_y = UOp(Ops.SPECIAL, dtypes.int, arg=("gidx1", N//BM))
-
-  a = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=0)
-  b = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=1)
-  c = UOp(Ops.DEFINE_GLOBAL, dtypes.float.ptr(N*N), arg=2)
-
-  junk = UOp.const(dtypes.float, 0)
-  A_col = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveM * TM, AddrSpace.REG), src=(junk,), arg=0)
-  B_row = UOp(Ops.DEFINE_REG, dtypes.float.ptr(nbIterWaveN * TN, AddrSpace.REG), src=(junk,), arg=1)
-
-  As = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*(BM+4), AddrSpace.LOCAL), arg=0) # 4 padding to avoid bank conflicts
-  Bs = UOp(Ops.DEFINE_LOCAL, dtypes.float.ptr(BK*BN, AddrSpace.LOCAL), arg=1)
-
-  c_regs = UOp(Ops.DEFINE_REG, dtypes.float.ptr(TM * nbIterWaveM * TN * nbIterWaveN), src=(junk,), arg=2)
-
-  kId_range = UOp.range(dtypes.int, N//BK, 0)
-  kId = kId_range*BK
-
-  # load from globals into locals
-  i = UOp.range(dtypes.int, nbReadsB, 1)
-  index_x = BN * blockIdx_x + rBIdx
-  index_y = rBIdy + i * strideReadB + kId
-  Bs_store = Bs[(index_y % BK)*BN + index_x%BN].store(b[N * index_y + index_x].load(), i)
-
-  i = UOp.range(dtypes.int, nbReadsA, 2)
-  index_x = rAIdx + kId
-  index_y = BM * blockIdx_y + rAIdy + i * strideReadA
-  As_store = As[(index_x % BK)*(BM+4) + index_y%BM].store(a[N * index_y + index_x].load(), i)
-
-  barrier = UOp(Ops.BARRIER, src=(As_store, Bs_store))
-
-  k = UOp.range(dtypes.int, BK, 3)
-
-  # load from locals into registers
-  iterWave = UOp.range(dtypes.int, nbIterWaveN, 4)
-  i = UOp.range(dtypes.int, TN, 5)
-  index = waveIdx * WN + iterWave * SUBWN + TN * idxInWave + i
-  B_row_store = B_row[iterWave*TN + i].store(Bs[k*BN + index].load(barrier), iterWave, i)
-
-  iterWave = UOp.range(dtypes.int, nbIterWaveM, 6)
-  i = UOp.range(dtypes.int, TM, 7)
-  index = waveIdy * WM + iterWave * SUBWM + TM * idyInWave + i
-  A_col_store = A_col[iterWave*TM + i].store(As[k*(BM+4) + index].load(barrier), iterWave, i)
-
-  # do the GEMM math
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 8)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 9)
-  yt = UOp.range(dtypes.int, TM, 10)
-  xt = UOp.range(dtypes.int, TN, 11)
-  x = iterWaveN * TN + xt
-  y = iterWaveM * TM + yt
-  c_regs_idx = c_regs[y * TN * nbIterWaveN + x]
-  sink = c_regs_idx.store(c_regs_idx.load() + A_col[y].load(A_col_store) * B_row[x].load(B_row_store),
-                          iterWaveM, iterWaveN, yt, xt, k, kId_range)
-
-  # store c_regs into c
-  iterWaveM = UOp.range(dtypes.int, nbIterWaveM, 12)
-  iterWaveN = UOp.range(dtypes.int, nbIterWaveN, 13)
-  xOut = blockIdx_x * BN + waveIdx * WN + iterWaveN * SUBWN + TN * idxInWave
-  yOut = blockIdx_y * BM + waveIdy * WM + iterWaveM * SUBWM + TM * idyInWave
-  yt = UOp.range(dtypes.int, TM, 14)
-  xt = UOp.range(dtypes.int, TN, 15)
-  indexC = N * (yOut + yt) + xOut + xt
-  sink = c[indexC].store(c_regs[TN * nbIterWaveN * (iterWaveM * TM + yt) + (iterWaveN * TN + xt)].load(sink), iterWaveM, iterWaveN, yt, xt)
-
-  return sink.sink(arg=KernelInfo(name="tinygemm"))
-
-if __name__ == "__main__":
-  hprg = hand_spec_kernel3()
-  prg = get_program(hprg, Device.default.renderer)
-  print(prg.src)
-  hrunner = CompiledRunner(prg)
-
-  a = Tensor.randn(N, N).realize()
-  b = Tensor.randn(N, N).realize()
-  hc = Tensor.zeros(N, N).contiguous().realize()
-
-  GlobalCounters.reset()
-  with Context(DEBUG=2):
-    for _ in range(run_count): tc = (a@b).realize()
-
-  GlobalCounters.reset()
-  ei = ExecItem(hrunner, [a.uop.buffer, b.uop.buffer, hc.uop.buffer])
-  with Context(DEBUG=2):
-    for _ in range(run_count): ei.run(wait=True)
-  err = (hc-tc).square().mean().item()
-  print(f"hrunner {err}")
-  if err > 1e-06: raise RuntimeError("matmul is wrong!")

extra/gemm/halide_gemm.py

@@ -1,90 +0,0 @@
-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/kernel8_batched_gmem.s

@@ -1,5 +1,6 @@
 	.text
 	.amdgcn_target "amdgcn-amd-amdhsa--gfx1100"
+	;.amdhsa_code_object_version 5
 	.protected	kernel                  ; -- Begin function kernel
 	.globl	kernel
 	.p2align	8
@@ -8,7 +9,7 @@ kernel:                                 ; @kernel
 ; %bb.0:                                ; %.preheader193
 
 	;; Init code for matrix A and B buffer Loads  - START
-	s_load_b128 s[20:23], s[0:1], 0x0 ; Matrix A and B
+	s_load_b128 s[20:23], s[0:1], 0x8 ; Matrix A and B
 	s_waitcnt lgkmcnt(0)
 
 	; Matrix B offsets:
@@ -75,12 +76,14 @@ kernel:                                 ; @kernel
 
 
 	s_clause 0x1
-	; s_load_b128 s[4:7], s[0:1], 0x18
-	; N=4096, alpha=1.0, beta=0.0
-	s_mov_b32 s4, 4096
-	s_mov_b32 s5, 0x3F800000
-	s_mov_b32 s6, 0
-	s_load_b128 s[8:11], s[0:1], 0x0
+	;s_load_b128 s[4:7], s[0:1], 0x18 ; N, alpha, beta, ???
+	s_load_b128 s[8:11], s[0:1], 0x8 ; Matrix A and B
+
+	s_mov_b32 s4, 4096 ; hardcode 4096
+	s_mov_b32 s5, 0x3f800000 ; alpha
+	s_mov_b32 s6, 0    ; beta
+	s_mov_b32 s7, 0
+
 	s_lshl_b32 s2, s14, 7
 	v_lshrrev_b32_e32 v4, 3, v0
 	v_or_b32_e32 v1, s2, v0
@@ -90,7 +93,7 @@ kernel:                                 ; @kernel
 	v_or_b32_e32 v22, s3, v4
 	v_ashrrev_i32_e32 v2, 31, v1
 	s_lshr_b32 s12, s12, 25
-	s_load_b64 s[0:1], s[0:1], 0x10
+	s_load_b64 s[0:1], s[0:1], 0  ; Matrix C
 	v_lshlrev_b32_e32 v135, 2, v118
 	s_delay_alu instid0(VALU_DEP_2) | instskip(SKIP_3) | instid1(VALU_DEP_3)
 	v_lshlrev_b64 v[5:6], 2, v[1:2]
@@ -460,7 +463,7 @@ kernel:                                 ; @kernel
 
 	v_mov_b32_e32 v5, 0
 	v_mov_b32_e32 v3, 0
-	s_add_i32 s7, s4, -8
+	s_add_i32 s7, s4, -1
 	s_add_u32 s8, s8, 32
 	s_addc_u32 s9, s9, 0
 	s_mov_b32 s12, 0
@@ -2395,9 +2398,18 @@ amdhsa.kernels:
         .offset:         16
         .size:           8
         .value_kind:     global_buffer
+      - .offset:         24
+        .size:           4
+        .value_kind:     by_value
+      - .offset:         28
+        .size:           4
+        .value_kind:     by_value
+      - .offset:         32
+        .size:           4
+        .value_kind:     by_value
     .group_segment_fixed_size: 8320
     .kernarg_segment_align: 8
-    .kernarg_segment_size: 24
+    .kernarg_segment_size: 36
     .language:       OpenCL C
     .language_version:
       - 2

extra/gemm/max_matmul.py

@@ -2,12 +2,11 @@ import numpy as np, os
 from tinygrad.helpers import getenv, flat_mv
 from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self
-from tinygrad.engine.realize import get_program
 
 # for copied uops
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.kernel import Kernel, KernelOptError
 from tinygrad.uop.ops import UOp, Ops, BinaryOps, UnaryOps, TernaryOps, KernelInfo
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.dtype import PtrDType, DType, DTYPES_DICT
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -56,7 +55,7 @@ def randoms():
 def ast_to_cuda_prog(compiler, ast, opts):
   k = Kernel(ast)
   k.apply_opts(opts)
-  p = get_program(k.get_optimized_ast(), k.opts)
+  p = k.to_program()
   return CUDAProgram(device, p.function_name, compiler.compile(p.src))
 
 if __name__ == "__main__":

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.opt.kernel import OptOps
+from tinygrad.codegen.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,7 +1,7 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.engine.realize import CompiledRunner, ExecItem, get_program
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.engine.realize import CompiledRunner, ExecItem
 from dataclasses import replace
 
 N = 4096
@@ -29,7 +29,7 @@ if __name__ == "__main__":
       Opt(op=OptOps.LOCAL, axis=0, amt=2),
     ]
   k.apply_opts(opts)
-  prg = get_program(k.get_optimized_ast(), k.opts)
+  prg = k.to_program()
   new_src = prg.src
   # can mod source here
   prg = replace(prg, src=new_src)

extra/gemm/tvm_gemm.py

@@ -37,10 +37,10 @@ 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.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel
 from tinygrad.device import CompilerOptions
 lin = Kernel(sched[-1].ast, CompilerOptions(has_local=False, supports_float4=False))
-lin.to_program()
+lin.linearize()
 from tinygrad.runtime.ops_cpu import renderer
 src = renderer("mmult", lin.uops)
 print(src)

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
-PYTHONPATH=. RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
+RUN_FILES="hcq,allocator" python3 extra/hcqfuzz/fuzzer.py
 ```
 Or skip tests from any file:
 ```bash
-PYTHONPATH=. SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
+SKIP_FILES="allocator" python3 extra/hcqfuzz/fuzzer.py
 ```

extra/hip_gpu_driver/hip_ioctl.py

@@ -5,7 +5,7 @@ start = time.perf_counter()
 # *** ioctl lib ***
 libc = ctypes.CDLL(ctypes.util.find_library("c"))
 # platform.processor calls `uname -p` which can return `unknown` on some systems
-processor = os.getenv("IOCTL_PROCESSOR") or platform.processor() or platform.machine()
+processor = os.getenv("IOCTL_PROCESSOR") or platform.processor()
 IOCTL_SYSCALL = {"aarch64": 0x1d, "x86_64":16}[processor]
 
 def get_struct(argp, stype):

extra/huggingface_onnx/run_models.py

@@ -1,6 +1,6 @@
 import onnx, yaml, tempfile, time, collections, pprint, argparse, json
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx import get_onnx_ops
 from extra.onnx_helpers import validate, get_example_inputs
 
@@ -13,7 +13,8 @@ def get_config(root_path: Path):
   return ret
 
 def run_huggingface_validate(onnx_model_path, config, rtol, atol):
-  onnx_runner = OnnxRunner(onnx_model_path)
+  onnx_model = onnx_load(onnx_model_path)
+  onnx_runner = OnnxRunner(onnx_model)
   inputs = get_example_inputs(onnx_runner.graph_inputs, config)
   validate(onnx_model_path, inputs, rtol=rtol, atol=atol)
 
@@ -45,7 +46,7 @@ def retrieve_op_stats(models:dict[str, tuple[Path, Path]]) -> dict:
   for model_id, (root_path, relative_path) in models.items():
     print(f"examining {model_id}")
     model_path = root_path / relative_path
-    onnx_runner = OnnxRunner(model_path)
+    onnx_runner = OnnxRunner(onnx.load(model_path))
     for node in onnx_runner.graph_nodes:
       op_counter[node.op] += 1
       if node.op not in supported_ops:

extra/mcts_search.py

@@ -4,10 +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.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel
 from tinygrad.device import Buffer, Device, CompileError
-from tinygrad.opt.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
-from tinygrad.engine.realize import get_program
+from tinygrad.engine.search import _ensure_buffer_alloc, get_kernel_actions, _time_program
 
 class MCTSNode:
   def __init__(self, kernel:Kernel, parent=None):
@@ -111,7 +110,7 @@ def mcts_search(lin:Kernel, rawbufs:List[Buffer], amt:int) -> Kernel:
       seen_asts[opt_ast.key] = node
 
       # lowering (50% of the time)
-      p = get_program(node.kernel.get_optimized_ast(name_override="test"), node.kernel.opts)
+      p = node.kernel.to_program(name_override="test")
 
       # rollout
       tm1 = time.perf_counter()

extra/models/llama.py

@@ -1,5 +1,5 @@
 from typing import Union, Optional, Any
-import collections, math
+import collections
 from tinygrad import Tensor, Variable, TinyJit, dtypes, nn, Device
 from tinygrad.helpers import getenv, DEBUG
 
@@ -166,29 +166,27 @@ 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, 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)]
+               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)]
     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().requires_grad_(False)
+    self.freqs_cis = precompute_freqs_cis(dim // n_heads, self.max_context * 2, rope_theta).contiguous()
     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).contiguous()
+    self.freqs_cis = self.freqs_cis.cast(h.dtype).kernelize()
     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) 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).kernelize() 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)
+    return sample(logits.flatten(), temperature, top_k, top_p, alpha_f, alpha_p).kernelize()
 
   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

@@ -1,77 +0,0 @@
-/*
- * 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

@@ -1,455 +0,0 @@
-/*
- * 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

@@ -1,223 +0,0 @@
-/*
- * 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
-{
-    NvU8 version;
-    NvU8 headerSize;
-    NvU8 entrySize;
-    NvU8 entryCount;
-} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_HEADER_V1;
-
-typedef struct
-{
-    NvU32 id;
-    NvU32 dmemOffset;
-} __attribute__((packed)) FALCON_APPLICATION_INTERFACE_ENTRY_V1;
-
-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

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

extra/nvpci/nv_smi.py

@@ -1,65 +0,0 @@
-#!/usr/bin/env python3
-
-from tinygrad.runtime.support.system import System
-import argparse, glob, os, re, time, subprocess, sys
-
-def scan_devs_based_on_lock(prefix:str) -> list[str]:
-  devs = []
-  for dev in glob.glob(f'/tmp/{prefix}_*.lock'):
-    dev_id = dev[8:-5]
-    if os.path.exists(f"/sys/bus/pci/devices/{dev_id}"): devs.append(dev_id)
-  return devs
-
-def _do_reset_device(pci_bus): System.pci_reset(pci_bus)
-def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
-
-def cmd_remove_module(args):
-  to_unload = [m for m in ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if _is_module_loaded(m)]
-  if not to_unload:
-    print("NVIDIA kernel modules are not loaded")
-  else:
-    print("Removing NVIDIA kernel modules:", ", ".join(to_unload))
-    try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
-    except subprocess.CalledProcessError as e:
-      print("Failed to unload all modules — they may be in use.", file=sys.stderr)
-      sys.exit(e.returncode)
-
-def cmd_insert_module(args):
-  cmd_remove_module(args)
-  cmd_reset_devices(args)
-
-  if not os.path.exists("/sys/module/nvidia"):
-    print("Inserting nvidia kernel module")
-    subprocess.run(["nvidia-smi"], check=True)
-  else: print("Nvidia kernel module already loaded")
-
-def cmd_reset_devices(args):
-  devs = scan_devs_based_on_lock("nv")
-  dev_to_reset = args.pci_bus if 'pci_bus' in args.__dir__() else ""
-
-  for dev in devs:
-    if dev.startswith(dev_to_reset):
-      print(f"Resetting device {dev}")
-      _do_reset_device(dev)
-      time.sleep(0.2)
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  subparsers = parser.add_subparsers(required=True, dest="cmd")
-
-  parser_insmod = subparsers.add_parser('insmod', help='Insert a nvidia kernel module')
-  parser_insmod.set_defaults(func=cmd_insert_module)
-
-  parser_rmmod = subparsers.add_parser('rmmod', help='Remove a nvidia kernel module')
-  parser_rmmod.set_defaults(func=cmd_remove_module)
-
-  parser_reset = subparsers.add_parser('reset', help='Reset a nvidia device')
-  parser_reset.add_argument('--pci_bus', type=str, default="", help='PCI bus ID of the device to reset')
-  parser_reset.set_defaults(func=cmd_reset_devices)
-
-  args = parser.parse_args()
-  if args.cmd is None:
-    parser.print_help(sys.stderr)
-    sys.exit(1)
-
-  args.func(args)

extra/onnx.py

@@ -1,67 +1,71 @@
-# mypy: disable-error-code="misc, list-item, assignment, attr-defined, operator, index, arg-type"
 from types import SimpleNamespace
-from typing import Any, Sequence, cast, Literal, Callable, get_args, NamedTuple
-import dataclasses, functools, io, math, types, warnings, pathlib, sys, enum
+from typing import Any, Sequence, cast, Literal, Callable
+import dataclasses, functools, io, math, types
 from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
-from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element
-from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
+from tinygrad.helpers import getenv, DEBUG, all_same, prod, flatten, make_tuple, argsort
+from tinygrad.dtype import DType, ConstType, dtypes, ImageDType
 from tinygrad.device import is_dtype_supported, Device
-from extra.onnx_parser import onnx_load
-
-# 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, TensorProto, TypeProto
+from onnx import AttributeProto, ModelProto, TensorProto, TypeProto, helper
+import numpy as np
 
 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, 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 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 attribute_parse(onnx_attribute: AttributeProto):
-  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)
+  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)
 
 def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
   if onnx_tensor.string_data: raise NotImplementedError("Parsing for buffer with string data is not implemented.")
-  to_dtype, true_dtype = dtype_parse(onnx_tensor.data_type, "buffer parse"), data_types[onnx_tensor.data_type]
-  shape = tuple(onnx_tensor.dims)
-  keys = ['float_data', 'int32_data', 'int64_data', 'double_data', 'uint64_data', "raw_data"]
-  data = next((val for k in keys if (val := getattr(onnx_tensor, k)) is not None), None)
-  if data is None: raise RuntimeError("empty buffer")
-  if not isinstance(data, Tensor): return Tensor(data, dtype=to_dtype).reshape(shape)
-  assert data.dtype is dtypes.uint8, data.dtype
-  data = data.bitcast(true_dtype).reshape(shape)
-  data = data.to(Device.DEFAULT) if true_dtype is to_dtype else data.to("cpu").cast(to_dtype).to(Device.DEFAULT)
-  if shape == ():
-    if data.dtype is dtypes.float16 and sys.version_info < (3, 12): data = data.cast(dtypes.float32)
-    return Tensor(data.item(), dtype=to_dtype).reshape(shape)
-  return data
+  dtype, shape = dtype_parse(onnx_tensor.data_type), 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
+  elif len(onnx_tensor.int64_data): data = onnx_tensor.int64_data
+  elif len(onnx_tensor.double_data): data = onnx_tensor.double_data
+  elif len(onnx_tensor.uint64_data): data = onnx_tensor.uint64_data
+  if isinstance(data, 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(),
+                                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)
+    return ret
+  return Tensor(None)
 
 def type_parse(onnx_type: TypeProto):
   elem_type = onnx_type
@@ -72,7 +76,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 = data_types[elem_type.tensor_type.elem_type]
+    dtype = dtype_parse(elem_type.tensor_type.elem_type)
     return OnnxValue(shape, dtype, is_optional, is_sequence)
   raise RuntimeError(f"TypeProto was not parsed properly: {onnx_type=}")
 
@@ -84,24 +88,10 @@ class OnnxValue:
   is_optional: bool
   is_sequence: bool
 
-class Domain(enum.Enum):
-  ONNX = "ai.onnx"
-  ONNX_ML = "ai.onnx.ml"
-  AI_ONNX_TRAINING = "ai.onnx.training"
-  AI_ONNX_PREVIEW_TRAINING = "ai.onnx.preview.training"
-  MICROSOFT_CONTRIB_OPS = "com.microsoft"
-  @classmethod
-  def from_onnx(cls, domain: str | None) -> "Domain": return cls.ONNX if domain is None or domain == "" else cls(domain)
-
-class OpSetId(NamedTuple):
-  domain: Domain
-  version: int
-
 @dataclasses.dataclass(frozen=True)
 class OnnxNode:
   num: int
   op: str
-  opset_id: OpSetId
   inputs: tuple[str, ...]
   outputs: tuple[str, ...]
   opts: dict[str, Any]
@@ -136,66 +126,49 @@ def to_python_const(t:Any, op:str, idx:int) -> list[ConstType]|ConstType|bytes:
 debug = int(getenv("DEBUGONNX", "0"))
 limit = int(getenv("ONNXLIMIT", "-1"))
 class OnnxRunner:
-  """
-  `OnnxRunner` executes an ONNX model using Tinygrad.
-
-  Args:
-    model_path: The ONNX model, provided as a file path (a string or Path object) or a Tensor.
-  """
-  def __init__(self, model_path: Tensor | str | pathlib.Path):
-    model = onnx_load(model_path)
-    self.is_training = any(n.domain in {Domain.AI_ONNX_TRAINING, Domain.AI_ONNX_PREVIEW_TRAINING} for n in model.graph.node)
+  def __init__(self, model: ModelProto|SimpleNamespace):
+    # parse model protobuf
+    self.is_training = any(n.domain in {"ai.onnx.training", "ai.onnx.preview.training"} for n in model.graph.node)
     self.old_training = Tensor.training
     Tensor.training = True if self.is_training else False
     self.graph_values = {"": None, **{x.name:buffer_parse(x) for x in model.graph.initializer}}
     self.graph_inputs = {x.name:type_parse(x.type) for x in model.graph.input if x.name not in self.graph_values}
     self.graph_outputs = tuple(x.name for x in model.graph.output)
-    opset_imports = {Domain.from_onnx(getattr(x, "domain", "")):x.version for x in model.opset_import}
-    self.graph_nodes = []
-    for num, n in enumerate(model.graph.node):
-      domain = Domain.from_onnx(n.domain)
-      opset_id = OpSetId(domain, opset_imports.get(domain, 1))
-      self.graph_nodes.append(OnnxNode(num, n.op_type, opset_id, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute}))
-    self.graph_nodes = tuple(self.graph_nodes)
+    self.graph_nodes = tuple(OnnxNode(num, n.op_type, tuple(n.input), tuple(n.output), {x.name:attribute_parse(x) for x in n.attribute})
+                       for num,n in enumerate(model.graph.node))
+    self.opset_version = model.opset_import[0].version
     self.variable_dims: dict[str, int] = {}
 
     self.onnx_ops = onnx_ops
 
   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"input {name} received {value}, expected a sequence type")
+      if not isinstance(value, Sequence): raise RuntimeError(f"{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 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}")
+      if not all_same(tuple(t.shape for t in sequence)): raise RuntimeError(f"Shapes for {name} sequence must be homogeneous")
       return sequence
-    dtype = _from_np_dtype(value.dtype) if is_numpy_ndarray(value) 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}.")
+    tensor = Tensor(value, dtype=spec.dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
     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"input {name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
+      if user_dim_input != onnx_dim: raise RuntimeError(f"{name} has mismatch on {dim=}. Expected {onnx_dim}, received {user_dim_input}.")
     return tensor
 
-  def _select_op(self, op:str, required_opset:OpSetId) -> types.FunctionType:
-    if op not in self.onnx_ops: raise NotImplementedError(f"{op=} is not supported")
-    # return default implementation if no opset_id is specified
-    if isinstance(impl := self.onnx_ops[op], types.FunctionType): return impl
-    # match domain and select implementation with latest compatible version
-    eligible_ops = {impl_opset.version:impl_fxn for impl_opset,impl_fxn in impl.items()
-                    if impl_opset.domain == required_opset.domain and impl_opset.version <= required_opset.version}
-    if not eligible_ops: raise NotImplementedError(f"{op=} is not supported for domain {required_opset.domain} and version {required_opset.version}")
-    return eligible_ops[max(eligible_ops.keys())]
+  def _dispatch_op(self, op, inps, opts):
+    if op in self.onnx_ops:
+      fxn = self.onnx_ops[op]
+      if isinstance(fxn, dict):
+        for k in sorted(fxn.keys()):
+          if k <= self.opset_version:
+            real_fxn = fxn[k]
+      else: real_fxn = fxn
+      return real_fxn(*inps, **opts)
+    raise NotImplementedError(f"{op=} not supported")
 
-  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 to(self, device:str|None):
-    self.graph_values = {k:v.to(device) if isinstance(v, Tensor) else v for k,v in self.graph_values.items()}
-    self.graph_nodes = tuple(OnnxNode(n.num, n.op, n.opset_id, tuple(n.inputs), tuple(n.outputs),
-                                      {k:v.to(device) if isinstance(v, Tensor) else v for k,v in n.opts.items()}) for n in self.graph_nodes)
-    return self
+  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 __call__(self, inputs:dict[str, Any], debug=debug):
     for name, input_spec in self.graph_inputs.items():
@@ -212,7 +185,7 @@ class OnnxRunner:
 
       if debug >= 1: print(f"{node.num}: op '{node.op}' opt {opts}")
       if debug >= 2 and node.inputs: print("\tinputs:\n" + "\n".join(f"\t\t{x} - {i!r}" for x,i in zip(node.inputs, inps)))
-      ret = self._select_op(node.op, node.opset_id)(*inps, **opts)
+      ret = self._dispatch_op(node.op, inps, opts)
       ret = ret if isinstance(ret, tuple) else (ret,)
       if debug >= 2: print("\toutputs:\n" + "\n".join(f"\t\t{x} - {o!r}" for x,o in zip(node.outputs, ret)))
 
@@ -227,10 +200,8 @@ class OnnxRunner:
 ####################
 ##### ONNX OPS #####
 ####################
-def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionType]]:
+def get_onnx_ops():
   # ***** helper functions *****
-  def _resolve_const(x: Sequence[ConstType]|ConstType): return x if isinstance(x, get_args(ConstType)) else get_single_element(x)
-
   def _axes(axes, noop_with_empty_axes): return axes or ([] if noop_with_empty_axes else None)
 
   # (padding_top, padding_left, ..., padding_bottom, padding_right, ...) -> (padding_left, padding_right, padding_top, padding_bottom, ...)
@@ -301,8 +272,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     if value_string is not None or value_strings is not None and sparse_value is not None:
       raise NotImplementedError('Constant OP not implemented for value_string, value_strings and sparse_value')
 
-  def Range(start:float|int|list[float|int], limit:float|int|list[float|int], delta:float|int|list[float|int]):
-    return Tensor.arange(start=_resolve_const(start), stop=_resolve_const(limit), step=_resolve_const(delta))
+  def Range(start:float|int, limit:float|int, delta:float|int): return Tensor.arange(start=start, stop=limit, step=delta)
 
   def ImageDecoder(encoded_stream:bytes, pixel_format="RGB"):
     try: import PIL.Image
@@ -314,7 +284,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     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, "EyeLike op") if dtype is not None else x.dtype)
+    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype) 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)
@@ -329,13 +299,13 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   # ***** Unary Ops (math) *****
   def Not(x:Tensor): return x.logical_not()
-  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)  # noqa: A002
+  def Clip(x: Tensor, min:Tensor|None=None, max:Tensor|None=None): return x if min is None and max is None else x.clip(min, max)
   def IsInf(x:Tensor, detect_negative:int=1, detect_positive:int=1): return x.isinf(bool(detect_positive), bool(detect_negative))
 
   # ***** Unary Ops (activation) *****
-  def softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
-  def softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
-  Softmax = {OpSetId(Domain.ONNX, 1):softmax_1, OpSetId(Domain.ONNX, 13):softmax_13}
+  def Softmax_1(x:Tensor, axis:int=1): return x.softmax(axis)
+  def Softmax_13(x:Tensor, axis:int=-1): return x.softmax(axis)
+  Softmax = {1:Softmax_1, 13:Softmax_13}
   def HardSigmoid(x:Tensor, alpha:float=0.2, beta:float=0.5): return (alpha*x + beta).clip(0, 1)
   def Gelu(x:Tensor, approximate:str|None=None): return x.gelu() if approximate == "tanh" else 0.5 * x * (1 + (x/math.sqrt(2)).erf())
   def BiasGelu(x: Tensor, bias: Tensor, approximate: str | None = None) -> Tensor: return Gelu(x + bias, approximate)
@@ -347,7 +317,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   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
+  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 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
@@ -368,7 +338,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   # ***** Casting Ops *****
   # TODO: saturate
-  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to, "Cast op"))
+  def Cast(x:Tensor, to:int, saturate:int=1): return x.cast(dtype_parse(to))
   def CastLike(x:Tensor, target_type:Tensor, saturate:int=1): return x.cast(target_type.dtype)
 
   # ***** Reduce Ops *****
@@ -472,7 +442,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
                         zip(strides, input_shape, output_padding, kernel_shape, dilations, output_shape)], auto_pad)
     if pads is None: # we generate pads
       output_shape = output_shape or [X.shape[i+2] * strides[i] for i in range(len(strides))]
-      pads = [strides[i]*(input_shape[i]-1)+output_padding[i]+((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
+      pads = [strides[i]*(input_shape[i]-1) + output_padding[i] + ((kernel_shape[i]-1)*dilations[i]+1)-output_shape[i] for i in range(len(input_shape))]
       pads = _auto_pad(pads, auto_pad) if auto_pad != "NOTSET" else [0] * len(input_shape) * 2
     pads = _onnx_pads_to_tiny_pads(pads)
     return X.conv_transpose2d(W, B, stride=strides, groups=group, dilation=dilations, padding=pads, output_padding=output_padding)
@@ -490,16 +460,14 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   def Einsum(*Inputs:list[Tensor], equation:str): return Tensor.einsum(equation, *Inputs)
 
-  def CumSum(X:Tensor, axis:int|list[int], exclusive:int=0, reverse:int=0):
-    axis = X._resolve_dim(_resolve_const(axis))
+  def CumSum(X:Tensor, axis:int|list, exclusive:int=0, reverse:int=0):
+    axis = X._resolve_dim(axis[0] if isinstance(axis, list) else axis)
     if reverse: X = X.flip(axis)
     if exclusive: X = X.pad(tuple((1,0) if i == axis else None for i in range(X.ndim)))\
                         .shrink(tuple((0,X.shape[axis]) if i == axis else None for i in range(X.ndim)))
     return X.cumsum(axis).flip(axis) if reverse else X.cumsum(axis)
 
-  def Trilu(x:Tensor, k:int|list[int]=0, upper:int=1):
-    k_ = _resolve_const(k)
-    return x.triu(k_) if upper else x.tril(k_)
+  def Trilu(x:Tensor, k:int=0, upper:int=1): return x.triu(k) if upper else x.tril(k)
 
   def Resize(X:Tensor, roi:list[float]|None=None, scales:list[float]|None=None, sizes:list[int]|None=None, antialias:int=0,
             axes:list[int]|None=None, coordinate_transformation_mode:str='half_pixel', cubic_coeff_a:float=-0.75, exclude_outside:int=0,
@@ -560,8 +528,8 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return X.permute(*argsort(perm)) if perm else X
   def Upsample(X, scales, mode): return Resize(X=X, scales=scales, mode=mode)  # deprecated
 
-  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):  # noqa: A002
-    val, idx = X.topk(_resolve_const(K), axis, largest, sorted)
+  def TopK(X:Tensor, K:int|list[int], axis:int=-1, largest:int=1, sorted:int=1):
+    val, idx = X.topk(K if isinstance(K, int) else K[0], axis, largest, sorted)
     return val, idx.cast(dtypes.int64)
 
   # ***** Neural Network Ops *****
@@ -623,9 +591,9 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def MeanVarianceNormalization(x:Tensor, axis:list[int]=[0,2,3]):
     return (x - x.mean(axis, keepdim=True)) / (x.std(axis, keepdim=True, correction=0) + 1e-9)
 
-  def OneHot(indices:Tensor, depth:float|int|list[int|float], values:Tensor, axis:int=-1):
+  def OneHot(indices:Tensor, depth:float|int|list, values:Tensor, axis:int=-1):
     # Scalar or Rank 1 tensor containing exactly one element
-    depth = int(_resolve_const(depth))
+    depth = int(depth[0] if isinstance(depth, list) else depth)
     indices = indices.int()
     indices = (indices < 0).where(indices+depth, indices)
     return indices.unsqueeze(axis)._one_hot_along_dim(depth, dim=axis).where(values[1], values[0])
@@ -636,18 +604,13 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     return X.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=blocksize, w1=blocksize)
 
   # Reimplemented here because you need legacy RNG for passing ONNX tests.
-  def dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
-    import numpy as np
-    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
+  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
   # 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 = {OpSetId(Domain.ONNX, 6):dropout_6, OpSetId(Domain.ONNX, 7):dropout_7}
+  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}
 
   def LRN(x:Tensor, size:int, alpha:float=1e-4, beta:float=0.75, bias:float=1.0):
     pooled_x = (x**2).rearrange('b c h w -> b 1 c (h w)').pad((0,0,(size-1)//2, size//2)).avg_pool2d((size, 1), 1)
@@ -662,17 +625,16 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
   def AffineGrid(theta:Tensor, size:list[int], align_corners:int=0):
     N, _, *spatial_dims = size
     def generate_grid(steps):
-      if align_corners: return Tensor.linspace(-1, 1, steps, device=theta.device)
-      return Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
+      return Tensor.linspace(-1, 1, steps, device=theta.device) if align_corners else Tensor.linspace(-1+1/steps, 1-1/steps, steps, device=theta.device)
     grids = Tensor.meshgrid(*(generate_grid(d) for d in spatial_dims))
     base_grid = Tensor.stack(*reversed(grids), Tensor.ones_like(grids[0], device=theta.device), dim=-1)
     base_grid = base_grid.reshape(1, prod(spatial_dims), len(grids)+1).expand(N, -1, -1)
     return (base_grid @ theta.transpose(1, 2)).reshape(N, *spatial_dims, -1)
 
-  def attention_contrib(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None,
-                        attention_bias:Tensor|None=None, past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0,
-                        num_heads:int|None=None, past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None,
-                        rotary_embedding_dim:int|None=None, scale:float|None=None, unidirectional:int=0):
+  def Attention(x:Tensor, weights:Tensor, bias:Tensor|None=None, mask_index:Tensor|None=None, past:Tensor|None=None, attention_bias:Tensor|None=None,
+                past_sequence_length:Tensor|None=None,  do_rotary:int=0, mask_filter_value:float=-10000.0, num_heads:int|None=None,
+                past_present_share_buffer:int|None=None, qkv_hidden_sizes:list[int]|None=None, rotary_embedding_dim:int|None=None,
+                scale:float|None=None, unidirectional:int=0):
     assert not do_rotary and not attention_bias, "TODO"
     if qkv_hidden_sizes is None: qkv_hidden_sizes = [weights.shape[1] // 3] * 3
     qkv = x.linear(weights, bias)
@@ -713,86 +675,6 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
     output = output.transpose(1, 2).reshape(batch_size, seq_len, -1)
     return output, present
 
-  def attention_onnx(Q:Tensor, K:Tensor, V:Tensor, attn_mask:Tensor|None=None, past_key:Tensor|None=None, past_value:Tensor|None=None,
-                     is_causal:int=0, kv_num_heads:int|None=None, q_num_heads:int|None=None, qk_matmul_output_mode:int=0, scale:float|None=None,
-                     softcap:float=0.0, softmax_precision:int|None=None):
-    input_shape_len = Q.ndim
-    if input_shape_len == 3:
-      assert q_num_heads is not None and kv_num_heads is not None
-      Q = Q.reshape(Q.shape[0], q_num_heads, Q.shape[1], -1)
-      K = K.reshape(K.shape[0], kv_num_heads, K.shape[1], -1)
-      V = V.reshape(V.shape[0], kv_num_heads, V.shape[1], -1)
-
-    if past_key is not None: K = past_key.cat(K, dim=2)
-    if past_value is not None: V = past_value.cat(V, dim=2)
-    present_key, present_value = K, V
-
-    _q_heads, _kv_heads = q_num_heads or Q.shape[1], kv_num_heads or K.shape[1]
-    if _q_heads != _kv_heads:
-      K = K.repeat((1, _q_heads // _kv_heads, 1, 1))
-      V = V.repeat((1, _q_heads // _kv_heads, 1, 1))
-
-    effective_scale = scale if scale is not None else 1.0 / (Q.shape[-1] ** 0.5)
-    scores = (Q @ K.transpose(-1, -2)) * effective_scale
-    qk_matmul_return_val = scores
-
-    if is_causal:
-      causal_mask = Tensor.ones(Q.shape[-2], K.shape[-2], device=Q.device, dtype=dtypes.bool, requires_grad=False).tril(0)
-      scores = scores.masked_fill(causal_mask.logical_not(), -float("inf"))
-
-    if attn_mask is not None:
-      mask_to_add = attn_mask.where(0, -float("inf")) if attn_mask.dtype == dtypes.bool else attn_mask
-      scores = scores + mask_to_add
-    if qk_matmul_output_mode == 1: qk_matmul_return_val = scores
-
-    if softcap > 0.0: scores = (scores / softcap).tanh() * softcap
-    if qk_matmul_output_mode == 2: qk_matmul_return_val = scores
-
-    if softmax_precision: scores = scores.cast({1: dtypes.float32, 10: dtypes.float16, 16: dtypes.bfloat16}[softmax_precision])
-    qk_softmax = scores.softmax(-1).cast(Q.dtype)
-    if qk_matmul_output_mode == 3: qk_matmul_return_val = qk_softmax
-
-    output = (qk_softmax @ V).cast(Q.dtype)
-    if input_shape_len == 3: output = output.permute(0, 2, 1, 3).reshape(Q.shape[0], Q.shape[2], -1)
-    return output, present_key, present_value, qk_matmul_return_val
-  Attention = {OpSetId(Domain.ONNX, 1): attention_onnx, OpSetId(Domain.MICROSOFT_CONTRIB_OPS, 1): attention_contrib}
-
-  def RMSNormalization(X:Tensor, scale:Tensor, axis:int=-1, epsilon:float=1e-5):
-    norm = X.square().mean(axis=tuple(range(axis + X.ndim if axis < 0 else axis, X.ndim)), keepdim=True).add(epsilon).rsqrt()
-    return X * norm * scale
-
-  def RotaryEmbedding(X:Tensor, cos_cache:Tensor, sin_cache:Tensor, position_ids:Tensor|None=None, interleaved:int=0, num_heads:int|None=None,
-                      rotary_embedding_dim:int=0):
-    original_input_shape = X.shape
-
-    if X.ndim == 4: X = X.permute(0, 2, 1, 3)
-    elif X.ndim == 3:
-      assert num_heads is not None, "num_heads must be provided for 3D input"
-      X = X.reshape(*X.shape[:-1], num_heads, X.shape[-1] // num_heads)
-
-    head_size = X.shape[-1]
-    rot_dim = rotary_embedding_dim or head_size
-    x_rotate, x_pass = X[..., :rot_dim], X[..., rot_dim:]
-
-    cos = cos_cache[position_ids] if position_ids is not None else cos_cache[:X.shape[1]]
-    sin = sin_cache[position_ids] if position_ids is not None else sin_cache[:X.shape[1]]
-    cos = cos[..., :rot_dim//2].unsqueeze(2)
-    sin = sin[..., :rot_dim//2].unsqueeze(2)
-
-    if interleaved:
-      x1, x2 = x_rotate[..., ::2], x_rotate[..., 1::2]
-      real = x1 * cos - x2 * sin
-      imag = x1 * sin + x2 * cos
-      x_rotated = Tensor.stack(real, imag, dim=-1).flatten(start_dim=-2)
-    else:
-      x1, x2 = x_rotate.chunk(2, dim=-1)
-      real = x1 * cos - x2 * sin
-      imag = x1 * sin + x2 * cos
-      x_rotated = real.cat(imag, dim=-1)
-
-    output = x_rotated.cat(x_pass, dim=-1)
-    return output.flatten(start_dim=2) if len(original_input_shape) == 3 else output.permute(0, 2, 1, 3)
-
   # ***** Indexing Ops *****
   def ArrayFeatureExtractor(x:Tensor, indices:Tensor): return x[..., indices]
 
@@ -849,9 +731,7 @@ def get_onnx_ops() -> dict[str, types.FunctionType|dict[OpSetId, types.FunctionT
 
   # ***** 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):
-    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
+    out_dtype = y_zero_point.dtype if isinstance(y_zero_point, Tensor) else dtype_parse(output_dtype) if output_dtype else 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_helpers.py

@@ -1,6 +1,6 @@
 from tinygrad import Tensor
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.frontend.onnx import OnnxRunner
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx import OnnxValue
 import numpy as np
 import onnxruntime as ort
@@ -46,7 +46,7 @@ def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
   return ret
 
 def validate(onnx_file, inputs, rtol=1e-5, atol=1e-5):
-  run_onnx = OnnxRunner(onnx_file)
+  run_onnx = OnnxRunner(onnx_load(onnx_file))
 
   ort_options = ort.SessionOptions()
   ort_options.log_severity_level = 3

extra/onnx_parser.py

@@ -2,9 +2,10 @@
 
 import os, pathlib, struct
 from io import BufferedReader
+from typing import Tuple, Union
 from types import SimpleNamespace
 from tinygrad.nn.state import TensorIO
-from tinygrad.tensor import Tensor
+from tinygrad.tensor import Tensor, dtypes
 
 # Protobuf Wire Types
 WIRETYPE_VARINT = 0; WIRETYPE_FIXED64 = 1; WIRETYPE_LENGTH_DELIMITED = 2; WIRETYPE_START_GROUP = 3; WIRETYPE_END_GROUP = 4; WIRETYPE_FIXED32 = 5 # noqa: E702
@@ -21,13 +22,12 @@ class AttributeType:
 
 class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6; BYTES = 7; SUB = 8 # noqa: E702
 
-PB_INFOS: dict[str, dict] = {
+PB_INFOS = {
   "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
   "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
   "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),
-    10: ("double_data", PBType.FLOATS), 11: ("uint64_data", PBType.INTS)},
+    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},
   "TensorShapeProtoDimension": {1: ("dim_value", PBType.INT), 2: ("dim_param", PBType.STRING)},
   "TensorShapeProto": {1: ("dim", PBType.SUB, True, "TensorShapeProtoDimension")},
   "ModelProto": {1: ("ir_version", PBType.INT), 5: ("model_version", PBType.INT),
@@ -36,16 +36,16 @@ PB_INFOS: dict[str, dict] = {
     8: ("opset_import",PBType.SUB, True, "OperatorSetIdProto")},
   "GraphProto": {2: ("name", PBType.STRING), 10: ("doc_string", PBType.STRING),
     1: ("node", PBType.SUB, True, ("NodeProto", lambda: {"input": [], "output": [], "attribute": [], "domain": None})),
-    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None,
-                                                                  "int64_data": None, "double_data": None, "uint64_data": None, "raw_data": None})),
+    5: ("initializer", PBType.SUB, True, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [],
+                                                                  "int64_data": [], "double_data": [], "uint64_data": []})),
     11: ("input", PBType.SUB, True, "ValueInfoProto"), 12: ("output", PBType.SUB, True, "ValueInfoProto")},
   "NodeProto": { 1: ("input", PBType.STRING, True), 2: ("output", PBType.STRING, True), 3: ("name", PBType.STRING),
     4: ("op_type", PBType.STRING), 6: ("doc_string", PBType.STRING), 7: ("domain", PBType.STRING),
     5: ("attribute", PBType.SUB, True, ("AttributeProto", lambda: {"floats": [], "ints": [], "strings": []}))},
   "AttributeProto": {1: ("name", PBType.STRING), 20: ("type", PBType.INT), 3: ("i", PBType.INT), 8: ("ints", PBType.INT, True),
     2: ("f", PBType.FLOAT), 7: ("floats", PBType.FLOAT, True), 4: ("s", PBType.BYTES), 9: ("strings", PBType.BYTES, True),
-    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": None, "int32_data": None, "string_data": None, "int64_data": None,
-                                                        "double_data": None, "uint64_data": None, "raw_data": None}))},
+    5:("t", PBType.SUB, False, ("TensorProto", lambda: {"dims": [], "float_data": [], "int32_data": [], "string_data": [], "int64_data": [],
+                                                        "double_data": [], "uint64_data": []}))},
   "ValueInfoProto": {1: ("name", PBType.STRING), 2: ("type", PBType.SUB, False, "TypeProto"), 3: ("doc_string", PBType.STRING)},
   "TypeProto": {1: ("tensor_type", PBType.SUB, False, "TypeProtoTensor"), 4: ("sequence_type", PBType.SUB, False, "TypeProtoSequence"),
     9: ("optional_type", PBType.SUB, False, "TypeProtoOptional"), 6: ("denotation", PBType.STRING)},
@@ -54,7 +54,7 @@ PB_INFOS: dict[str, dict] = {
   "TypeProtoTensor": {1: ("elem_type", PBType.INT), 2: ("shape", PBType.SUB, False, ("TensorShapeProto", lambda: {"dim": []}))},
 }
 
-def onnx_load(fn: Tensor|str|pathlib.Path, load_external_data: bool=True):
+def onnx_load(fn: Union[Tensor, str, pathlib.Path], load_external_data: bool=True):
   parser = OnnxParser(fn, load_external_data)
   onnx_model = parser.parse()
   model = dict_to_namespace(onnx_model)
@@ -70,8 +70,8 @@ def dict_to_namespace(d):
   return d
 
 class OnnxParser:
-  def __init__(self, inp: Tensor|str|pathlib.Path, load_external_data: bool=True):
-    self.file_path: pathlib.Path|None = None
+  def __init__(self, inp: Union[Tensor, str, pathlib.Path], load_external_data: bool=True):
+    self.file_path: Union[pathlib.Path, None] = None
     self.load_external_data = load_external_data
     if not isinstance(inp, Tensor):
       self.file_path = pathlib.Path(inp)
@@ -89,6 +89,7 @@ class OnnxParser:
         elif len(config) == 4: name, attr, repeated, parser_fn = config
         handler_fn = self.attr_func_dict[attr]
         def _wrapper_handler(obj, reader, wt, h=handler_fn, n=name, p=parser_fn, r=repeated): return h(obj, n, reader, wt, parser_func=p, repeated=r)
+        _wrapper_handler._debug_info = f"{fid}, {name} => {handler_fn}"
         res[fid] = _wrapper_handler
       self.registered_handles[pb_name] = res
 
@@ -129,19 +130,16 @@ class OnnxParser:
     if message_field_handlers_name == "TensorProto" and self.load_external_data and obj.get("data_location", 0) == 1: self._parse_external_data(obj)
     return obj
 
-  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tensor|bytes:
+  def _handle_delimited(self, reader:BufferedReader, use_tensor=False) -> Tuple[bytes, Tensor]:
     str_len = self.decode_varint(reader)
     if not use_tensor: return reader.read(str_len)
-    raw = reader.raw
-    assert isinstance(raw, TensorIO)
-    res = raw._tensor[reader.tell():(reader.tell()+str_len)]
+    res = reader.raw._tensor[reader.tell():(reader.tell()+str_len)]
     reader.seek(str_len, os.SEEK_CUR)
     return res
 
   def _handle_string(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
     if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for string field '{key_name}'")
     value = self._handle_delimited(reader)
-    assert isinstance(value, bytes)
     gen_result(obj, key_name, value.decode("utf-8"), repeated)
 
   def _handle_bytes(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
@@ -167,17 +165,16 @@ class OnnxParser:
     while reader.tell() < total_bytes_len + old_pos:
       val = self.decode_varint(reader) # need copy here because packed ints are varint
       values.append(val - 2**64 if val & (1 << 63) else val)
-    obj[key_name] = values
+    obj[key_name] = Tensor(values, dtype=dtypes.int64)
 
   def _handle_packed_floats(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
     if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError("Packed floats expected length_delimited")
     value = self._handle_delimited(reader, use_tensor=True)
-    obj[key_name] = value
+    obj[key_name] = value.bitcast(dtypes.float32)
 
   def _handle_sub_message(self, obj, key_name, reader, wire_type, parser_func=None, repeated=False):
     if wire_type != WIRETYPE_LENGTH_DELIMITED: raise ValueError(f"Expected length-delimited for sub-message field '{key_name}'")
     value = self._handle_delimited(reader, use_tensor=True)
-    assert isinstance(value, Tensor)
     if isinstance(parser_func, str): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func)
     elif isinstance(parser_func, tuple): sub_obj = self._parse_message(BufferedReader(TensorIO(value)), parser_func[0], parser_func[1])
     else: sub_obj = parser_func(BufferedReader(TensorIO(value)))
@@ -196,7 +193,7 @@ class OnnxParser:
     if self.file_path is None:
       # get onnx file path from Tensor
       if isinstance(self.tensor.device, str) and self.tensor.device.startswith("DISK:"):
-        self.file_path = pathlib.Path(self.tensor.device[5:])
+        self.file_path = self.tensor.device[5:]
         if not (ext_path := self.file_path.parent.joinpath(location)).exists():
           raise Exception(f"external location not exists: {ext_path}, may caused by symbolic link, try passing onnx file path to onnx_load")
       else: raise Exception("onnx external_data need the origin file path, try passing onnx file path to onnx_load")

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[1]).replace("\n", "").replace(" ", "")+"\n")
+  open(LOGOPS, "a").write(str(args[0]).replace("\n", "").replace(" ", "")+"\n")
   return None
 
 if __name__ == "__main__":
-  _pmap({"get_program":extract_ast})
+  _pmap("kernel", 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.opt.search import actions
+from tinygrad.engine.search import actions
 from extra.optimization.helpers import load_worlds, ast_str_to_lin, lin_to_feats, assert_same_lin
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.kernel import Opt, OptOps
 
 # more stuff
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.search import actions
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.engine.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,12 +6,24 @@ export CAPTURE_PROCESS_REPLAY=1
 rm $LOGOPS
 test/external/process_replay/reset.py
 
-CI=1 python3 -m pytest -n=auto test/test_ops.py test/test_nn.py test/test_winograd.py test/models/test_real_world.py --durations=20
-GPU=1 python3 -m pytest test/test_tiny.py
+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
 
 # extract, sort and uniq
 extra/optimization/extract_dataset.py
 sort -u /tmp/ops > /tmp/sops
-ls -lh /tmp/ops /tmp/sops
-gzip -k /tmp/sops
-# mv /tmp/sops.gz extra/datasets/
+ls -lh /tmp/ops /tmp/sops	

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.opt.search import actions
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.engine.search import actions
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from tinygrad.helpers import tqdm
 
 tactions = set()

extra/optimization/helpers.py

@@ -1,17 +1,15 @@
 # stuff needed to unpack a kernel
 from tinygrad import Variable
-from tinygrad.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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
-from tinygrad.engine.realize import get_program
 inf, nan = float('inf'), float('nan')
 UOps = Ops
 
 # kernel unpacker
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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):
@@ -28,7 +26,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) >= getenv("MIN_ASTS", 1000), f"dataset size = {len(ast_strs)} is too small"
+  assert len(ast_strs) > 5000, 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]
@@ -103,7 +101,7 @@ def lin_to_feats(lin:Kernel, use_sts=True):
   return ret
 
 from tinygrad.device import Device, Buffer
-from tinygrad.opt.search import _ensure_buffer_alloc, _time_program
+from tinygrad.engine.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
@@ -116,7 +114,7 @@ def time_linearizer(lin:Kernel, rawbufs:list[Buffer], allow_test_size=True, max_
 
   rawbufs = _ensure_buffer_alloc(rawbufs)
   var_vals: dict[Variable, int] = {k:int(k.vmax+k.vmin)//2 for k in lin.ast.variables()}
-  p = get_program(lin.get_optimized_ast(), lin.opts)
+  p = lin.to_program()
   tms = _time_program(p, dev.compiler.compile(p.src), var_vals, rawbufs,
                       max_global_size=max_global_size if allow_test_size else None, clear_l2=clear_l2, cnt=cnt, name=to_function_name(lin.name))
 

extra/optimization/pretrain_valuenet.py

@@ -1,4 +1,4 @@
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.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.opt.kernel import Opt, OptOps
+from tinygrad.codegen.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.opt.search import actions, bufs_from_lin, get_kernel_actions
+from tinygrad.engine.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.opt.kernel import Kernel
-from tinygrad.opt.search import get_kernel_actions, actions
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.engine.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.opt.kernel import Kernel
-from tinygrad.opt.search import beam_search, bufs_from_lin
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.engine.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.opt.search import bufs_from_lin, actions, get_kernel_actions
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.engine.search import bufs_from_lin, actions, get_kernel_actions
+from tinygrad.codegen.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.opt.search import bufs_from_lin, get_kernel_actions
+from tinygrad.engine.search import bufs_from_lin, get_kernel_actions
 
 if __name__ == "__main__":
   ast_strs = load_worlds()

extra/perfetto/to_perfetto.py

@@ -1,38 +0,0 @@
-import sys, pickle, decimal, json
-from tinygrad.device import ProfileDeviceEvent, ProfileGraphEvent
-from tinygrad.helpers import tqdm, temp, ProfileEvent, ProfileRangeEvent
-
-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

@@ -0,0 +1,173 @@
+<!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` (default) 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` 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

@@ -4,10 +4,10 @@ from tinygrad import Device, Context, Tensor, GlobalCounters
 from tinygrad.device import Buffer
 from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
-from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item, get_program
+from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 from tinygrad.renderer import ProgramSpec
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 import numpy as np
 
 def move_jit_captured_to_dev(captured, device="DSP"):
@@ -58,7 +58,7 @@ if __name__ == "__main__":
             GlobalCounters.kernel_count -= 1
 
         if not getenv("NOOPT"): k.apply_opts(hand_coded_optimizations(k))
-        p2 = get_program(k.get_optimized_ast(), k.opts)
+        p2 = k.to_program()
         new_ei = replace(ei, prg=CompiledRunner(p2))
         new_ei.run()
         new_jit.append(new_ei)

extra/sqtt/README.md

@@ -26,7 +26,7 @@ extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -o /tmp/gpu0.rgp
 
 Then load gpu0.rgp into Radeon GPU Profiler. It works just fine both in wine (macos, native version available for linux) and via ssh X forwarding
 
-If multiple gpus are used you can select which one to export with `-d` like this:
+If multiplle gpus are used you can select which one to export with `-d` like this:
 
 ```bash
 extra/sqtt/rgptool.py create "/tmp/profile.pkl.$USER" -d 'AMD:5' -o /tmp/gpu5.rgp

extra/torch_backend/backend.py

@@ -59,7 +59,6 @@ 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))
@@ -352,11 +351,6 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
   unwrap(indices).assign(out_indices.cast(dtypes.int64))
   return wrap(out_values), wrap(out_indices)
 
-@torch.library.impl("aten::_linalg_svd", "privateuseone")
-def _linalg_svd(self, full_matrices=False):
-  U, S, Vh = unwrap(self).svd(full_matrices)
-  return wrap(U), wrap(S), wrap(Vh)
-
 # register some decompositions
 from torch._decomp import get_decompositions
 decomps = [
@@ -374,7 +368,6 @@ 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,
@@ -383,7 +376,6 @@ decomps = [
   aten.softshrink,
   aten.hardshrink,
   aten.log_sigmoid_forward,
-  aten.log_sigmoid_backward,
   aten.isneginf,
   aten.isposinf,
   aten.nan_to_num,
@@ -417,7 +409,6 @@ decomps = [
   #aten.lgamma,
   # this needs copy_strided
   #aten.lerp,
-  aten.norm,
 ]
 for k,v in get_decompositions(decomps).items():
   key = str(k._schema).split("(")[0]
@@ -479,7 +470,6 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
   "aten.fmax.out": lambda input,other: Tensor.where(input.isnan() & ~other.isnan(), other, Tensor.where(~input.isnan() & other.isnan(), input, Tensor.maximum(input, other))),
   "aten.fmin.out": lambda input,other: Tensor.where(input.isnan() & ~other.isnan(), other, Tensor.where(~input.isnan() & other.isnan(), input, Tensor.minimum(input, other))),
   "aten.amax.out": lambda self,dim=None: self.max(axis=dim),
-  "aten.amin.out": lambda self,dim=None: self.min(axis=dim),
   # TODO: this gets the shape wrong
   #"aten.arange.start_out": Tensor.arange,
   "aten.lerp.Scalar_out": Tensor.lerp,

extra/torch_backend/test.py

@@ -103,27 +103,6 @@ class TestTorchBackend(unittest.TestCase):
     expected = np.array([[4.7, 12.9, 12.3], [16.9, 24.9, 23.6]], dtype=np.float32)
     np.testing.assert_equal(y3.cpu().numpy(), expected)
 
-
-  def test_amin(self):
-    x = torch.tensor([[[ 1.5,  2.3,  3.1,  4.7],
-                       [ 5.2,  6.8,  7.4,  12.9],
-                       [ 9.0, 12.3, 11.6, 10.1]],
-                      [[13.2, 16.9, 15.5, 14.1],
-                       [17.1, 24.9, 19.8, 20.2],
-                       [21.0, 22.3, 23.6, 18.4]]], device=device)
-
-    y1 = torch.amin(x)
-    expected = np.array([1.5], dtype=np.float32)
-    np.testing.assert_equal(y1.cpu().numpy(), expected)
-
-    y2 = torch.amin(x, dim=(1,2))
-    expected = np.array([1.5, 13.2], dtype=np.float32)
-    np.testing.assert_equal(y2.cpu().numpy(), expected)
-
-    y3 = torch.amin(x, dim=2)
-    expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
-    np.testing.assert_equal(y3.cpu().numpy(), expected)
-
   def test_isfinite(self):
     a = torch.ones(4, device=device)
     np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -227,16 +206,5 @@ 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()

mkdocs.yml

@@ -22,7 +22,6 @@ 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,17 +36,9 @@ 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,10 +25,9 @@ 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.schedule',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv', 'tinygrad.apps'],
-      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
+                  '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/*']},
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"
@@ -55,7 +54,7 @@ setup(name='tinygrad',
         ],
         'testing': testing_minimal + [
             "pillow",
-            "onnx==1.18.0",
+            "onnx==1.17.0",
             "onnx2torch",
             "onnxruntime",
             "opencv-python",

test/external/external_benchmark_hcopt.py

@@ -1,7 +1,7 @@
 import random
 from tinygrad.helpers import getenv
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.codegen.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

@@ -1,35 +0,0 @@
-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_openpilot.py

@@ -1,21 +1,24 @@
 import time, sys, hashlib
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
+from onnx.helper import tensor_dtype_to_np_dtype
+from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from tinygrad import Tensor, dtypes, TinyJit
 from tinygrad.helpers import IMAGE, GlobalCounters, fetch, colored, getenv, trange
+from tinygrad.tensor import _from_np_dtype
 import numpy as np
 from extra.bench_log import BenchEvent, WallTimeEvent
 
 OPENPILOT_MODEL = sys.argv[1] if len(sys.argv) > 1 else "https://github.com/commaai/openpilot/raw/v0.9.4/selfdrive/modeld/models/supercombo.onnx"
 
 if __name__ == "__main__":
-  run_onnx = OnnxRunner(fetch(OPENPILOT_MODEL))
+  onnx_model = onnx_load(onnx_path := fetch(OPENPILOT_MODEL))
+  run_onnx = OnnxRunner(onnx_model)
 
   Tensor.manual_seed(100)
-  input_shapes = {name: spec.shape for name, spec in run_onnx.graph_inputs.items()}
-  input_types = {name: spec.dtype for name, spec in run_onnx.graph_inputs.items()}
-  new_inputs = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
-  new_inputs_junk = {k:Tensor.randn(*shp, dtype=input_types[k]).mul(8).realize() for k,shp in input_shapes.items()}
+  input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input}
+  new_inputs = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in input_shapes.items()}
+  new_inputs_junk = {k:Tensor.randn(*shp, dtype=_from_np_dtype(input_types[k])).mul(8).realize() for k,shp in input_shapes.items()}
   new_inputs_junk_numpy = {k:v.numpy() for k,v in new_inputs_junk.items()}
 
   # benchmark

test/external/external_benchmark_schedule.py

@@ -3,11 +3,10 @@ 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.opt.kernel import Kernel
-from tinygrad.opt.heuristic import hand_coded_optimizations
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.heuristic import hand_coded_optimizations
 from tinygrad.codegen import get_rewrites_for_renderer, apply_rewrites, rewrites_for_linearizer
-from tinygrad.opt.search import beam_search, bufs_from_lin
-from tinygrad.uop.spec import type_verify
+from tinygrad.engine.search import beam_search, bufs_from_lin
 
 if __name__ == "__main__":
   mdl = ResNet50()
@@ -57,6 +56,4 @@ 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

@@ -1,5 +1,4 @@
 from tinygrad import Tensor, dtypes, GlobalCounters
-from tinygrad.engine.realize import get_program
 
 if __name__ == "__main__":
   t = Tensor.empty(81920, 4096, dtype=dtypes.half)
@@ -8,7 +7,7 @@ if __name__ == "__main__":
   GlobalCounters.reset()
   t.softmax(-1, dtype="half", _single_kernel=True).realize()
 
-  from tinygrad.opt.kernel import Kernel, Opt, OptOps
+  from tinygrad.codegen.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())
@@ -24,5 +23,5 @@ if __name__ == "__main__":
   #k.apply_opt(Opt(OptOps.GROUP, 1, 32))
   #k.apply_opt(Opt(OptOps.GROUP, 0, 32))
   from tinygrad.engine.realize import CompiledRunner, ExecItem
-  run = CompiledRunner(prg:=get_program(k.get_optimized_ast(), k.opts))
+  run = CompiledRunner(prg:=k.to_program())
   ExecItem(run, si.bufs).run()

test/external/external_debug_metal_sd_conv.py

@@ -1,8 +1,8 @@
 # ruff: noqa: E501
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
 from tinygrad.dtype import dtypes
-from tinygrad.engine.realize import CompiledRunner, get_program
-from tinygrad.opt.search import bufs_from_lin
+from tinygrad.engine.realize import CompiledRunner
+from tinygrad.engine.search import bufs_from_lin
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
@@ -35,7 +35,7 @@ k = Kernel(ast)
 k.apply_opts(opts)
 bufs = bufs_from_lin(k)
 
-prg = CompiledRunner(get_program(k.get_optimized_ast(), k.opts))
+prg = CompiledRunner(k.to_program())
 
 for i in range(10):
   speed = prg(bufs, var_vals={}, wait=True)