remove AMD_LLVM=0 in mlperf and search ci

tinybox updated to llvm 20
2025-06-11 16:16:30 -07:00
216 changed files with 6096 additions and 14866 deletions
@@ -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
@@ -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
-        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
+        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
@@ -27,7 +27,7 @@ jobs:
        BENCHMARK_LOG=search_sdxl_cached PYTHONPATH=. AMD=1 JITBEAM=2 python examples/sdxl.py --noshow --timing --seed 0
    - name: Run winograd cifar with new search
      run: |
-        BENCHMARK_LOG=search_wino_cifar WINO=1 DEFAULT_FLOAT=HALF 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
@@ -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,7 +336,7 @@ jobs:
      run: |
        pip3 install --upgrade --force-reinstall ruff==0.11.0
        python3 -m ruff check .
-        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
@@ -361,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 < 14600 lines
-      run: MAX_LINE_COUNT=14600 python sz.py
+    - name: Repo line count < 14000 lines
+      run: MAX_LINE_COUNT=14000 python sz.py

  fuzzing:
    name: Fuzzing
@@ -447,8 +450,6 @@ jobs:
        run: PYTHONPATH="." FLOAT16=0 DEBUGCL=1 GPU=1 IMAGE=2 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
      - name: Test openpilot LLVM compile
        run: PYTHONPATH="." LLVM=1 LLVMOPT=1 JIT=2 BEAM=0 IMAGE=0 python examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
-      - name: Test openpilot compile4
-        run: PYTHONPATH="." NOLOCALS=1 GPU=1 IMAGE=2 FLOAT16=1 DEBUG=2 python3 examples/openpilot/compile4.py
      - name: Run process replay tests
        uses: ./.github/actions/process-replay

@@ -473,10 +474,6 @@ jobs:
        run: CPU=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
      - name: Test ONNX (LLVM)
        run: LLVM=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      - name: Test ONNX Runner (CPU)
-        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
-      - name: Test ONNX Runner (WEBGPU)
-        run: WEBGPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_runner.py
      - name: Test Additional ONNX Ops (CPU)
        run: CPU=1 PYTHONPATH=. python3 test/external/external_test_onnx_ops.py
      - name: Test Quantize ONNX
@@ -513,8 +510,8 @@ jobs:
          REMOTEDEV=GPU IMAGE=2 REMOTE=1 python3 -m pytest test/test_tiny.py test/test_image_dtype.py
      - name: Test Optimization Helpers
        run: PYTHONPATH="." DEBUG=1 python3 extra/optimization/test_helpers.py
-      #- name: Test Action Space
-      #  run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
+      - name: Test Action Space
+        run: PYTHONPATH="." DEBUG=1 GPU=1 python3 extra/optimization/get_action_space.py
      - name: Test Beam Search
        run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
      - name: Test MLPerf stuff
@@ -613,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

@@ -680,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
@@ -748,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
@@ -764,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)
@@ -779,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
@@ -814,19 +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: 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)
@@ -858,8 +852,9 @@ jobs:
    timeout-minutes: 20
    env:
      REMOTE: 1
-      HOST: 127.0.0.1:6667*6,127.0.0.1:6668*6
+      REMOTEDEV: 'AMD'
      PYTHONPATH: ${{ github.workspace }}
+      MOCKGPU: 1
    steps:
      - name: Checkout Code
        uses: actions/checkout@v4
@@ -870,21 +865,6 @@ jobs:
          deps: testing_minimal
          amd: 'true'
          llvm: 'true'
-      - name: Start remote server
-        run: |
-          start_server() {
-            systemd-run --user \
-              --unit="$1" \
-              --setenv=REMOTEDEV=AMD \
-              --setenv=MOCKGPU=1 \
-              --setenv=PYTHONPATH=. \
-              --setenv=PORT="$2" \
-              --working-directory="$(pwd)" \
-              python tinygrad/runtime/ops_remote.py
-          }
-
-          start_server "remote-server-1" 6667
-          start_server "remote-server-2" 6668
      - name: Check Device.DEFAULT and print some source
        run: |
          python -c "from tinygrad import Device; assert Device.DEFAULT == 'REMOTE', Device.DEFAULT"
@@ -892,12 +872,7 @@ jobs:
          DEBUG=4 python3 test/test_tiny.py TestTiny.test_plus
      - name: Run REMOTE=1 Test
        run: |
-          python3 -m pytest test/test_tiny.py test/test_jit.py test/test_subbuffer.py test/test_graph.py test/test_multitensor.py test/test_remote.py test/test_tensor_variable.py
-      - name: Show remote server logs
-        if: always()
-        run: |
-          journalctl --user -u remote-server-1 --no-pager
-          journalctl --user -u remote-server-2 --no-pager
+          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:
@@ -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 \
@@ -167,26 +165,6 @@ generate_nv() {
 nv_status_codes = {}
 /^NV_STATUS_CODE/ { s/^NV_STATUS_CODE(\([^,]*\), *\([^,]*\), *"\([^"]*\)") *.*$/\1 = \2\nnv_status_codes[\1] = "\3"/; p }' $NVKERN_SRC/src/common/sdk/nvidia/inc/nvstatuscodes.h >> $BASE/nv_gpu.py

-  clang2py -k cdefstum \
-    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/fsp/kern_fsp_cot_payload.h \
-    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gspifpub.h \
-    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_wpr_meta.h \
-    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/gsp/gsp_fw_sr_meta.h \
-    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
-    $NVKERN_SRC/src/nvidia/inc/kernel/gpu/gsp/gsp_init_args.h \
-    $NVKERN_SRC/src/common/uproc/os/common/include/libos_init_args.h \
-    $NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc/rmRiscvUcode.h \
-    $NVKERN_SRC/src/common/shared/msgq/inc/msgq/msgq_priv.h \
-    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_headers.h \
-    $NVKERN_SRC/src/nvidia/inc/kernel/vgpu/rpc_global_enums.h \
-    $NVKERN_SRC/src/nvidia/generated/g_rpc-structures.h \
-    extra/nv_gpu_driver/g_rpc-message-header.h \
-    extra/nv_gpu_driver/gsp_static_config.h \
-    extra/nv_gpu_driver/vbios.h \
-    --clang-args="-DRPC_MESSAGE_STRUCTURES -DRPC_STRUCTURES -include $NVKERN_SRC/src/common/sdk/nvidia/inc/nvtypes.h -I$NVKERN_SRC/src/nvidia/generated -I$NVKERN_SRC/src/common/inc -I$NVKERN_SRC/src/nvidia/inc -I$NVKERN_SRC/src/nvidia/interface/ -I$NVKERN_SRC/src/nvidia/inc/kernel -I$NVKERN_SRC/src/nvidia/inc/libraries -I$NVKERN_SRC/src/nvidia/arch/nvalloc/common/inc -I$NVKERN_SRC/kernel-open/nvidia-uvm -I$NVKERN_SRC/kernel-open/common/inc -I$NVKERN_SRC/src/common/sdk/nvidia/inc -I$NVKERN_SRC/src/nvidia/arch/nvalloc/unix/include -I$NVKERN_SRC/src/common/sdk/nvidia/inc/ctrl" \
-    -o $BASE/nv/nv.py
-
-  fixup $BASE/nv/nv.py
  python3 -c "import tinygrad.runtime.autogen.nv_gpu"
 }

@@ -412,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() {
@@ -458,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
@@ -59,11 +59,11 @@ st_0 = UOp(Ops.STORE, dtypes.void, (output_buf.view(ShapeTracker.from_shape((1,)
 s = UOp(Ops.SINK, dtypes.void, (st_0,))

 # convert the computation to a "linearized" format (print the format)
-from tinygrad.engine.realize import get_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

@@ -78,7 +78,7 @@ print("******** third, the UOp ***********")

 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
-from tinygrad.kernelize.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)
@@ -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/kernelize
-
-Group UOps into kernels.
-
-::: tinygrad.kernelize.kernelize.get_kernelize_map
-    options:
-        members: false
-        show_labels: false
-        show_source: false
-
---
-
-## tinygrad/opt
-
-Transforms the ast into an optimized ast. This is where BEAM search and heuristics live.
-
-::: tinygrad.opt.get_optimized_ast
-    options:
-        members: false
-        show_labels: false
-        show_source: false
-
---
-
-## tinygrad/codegen
-
-Transform the optimized ast into a linearized list of UOps.
-
-::: tinygrad.codegen.full_rewrite
-    options:
-        members: false
-        show_labels: false
-        show_source: false
-
---
-
-## tinygrad/renderer
-
-Transform the linearized list of UOps into a program, represented as a string.
-
-::: tinygrad.renderer.Renderer
-    options:
-        members:
-            - render
-        show_labels: false
-        show_source: false
-
---
-
-## tinygrad/engine
-
-Abstracted high level interface to the runtimes.
-
-::: tinygrad.engine.realize.get_program
-    options:
-        members: false
-        show_labels: false
-        show_source: false
@@ -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)
@@ -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
@@ -1,4 +1,4 @@
-import sys, time
+import sys, time, pickle
 from tinygrad import TinyJit, GlobalCounters, fetch, getenv
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs, validate
@@ -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
@@ -2,11 +2,11 @@ from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
 from tinygrad import Tensor, Device, dtypes, nn
-from tinygrad.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

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

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

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

@@ -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)
@@ -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)
@@ -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')
@@ -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
@@ -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
@@ -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
@@ -1,8 +1,8 @@
 import sys, onnx
-from tinygrad import Tensor, fetch, GlobalCounters, dtypes
-from tinygrad.uop.ops import UOp
+from tinygrad import Tensor, fetch, GlobalCounters
+from tinygrad.uop import UOp
 from tinygrad.frontend.onnx import OnnxRunner
-from tinygrad.kernelize.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

@@ -17,7 +17,7 @@ if __name__ == "__main__":
  onnx_model = onnx.load(onnx_file)
  run_onnx = OnnxRunner(onnx_model)

-  inputs = run_onnx.get_empty_input_data("npy", 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()]
@@ -37,12 +37,12 @@ if __name__ == "__main__":
  independent = UOp.sink(*independent_set.keys())
  kernelized = get_kernelize_map(independent)
  independent = independent.substitute(kernelized)
-  schedule, var_vals = 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
@@ -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"))
@@ -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
@@ -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)
@@ -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?
@@ -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
@@ -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
@@ -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
@@ -4,6 +4,7 @@ from ultralytics import YOLO
 from pathlib import Path
 from tinygrad.frontend.onnx import OnnxRunner, onnx_load
 from extra.onnx_helpers import get_example_inputs
+from tinygrad.tensor import Tensor

 os.chdir("/tmp")
 if not Path("yolov8n-seg.onnx").is_file():
@@ -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
@@ -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:
@@ -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

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

@@ -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('//', '/')}")
@@ -4,7 +4,7 @@ import numpy as np
 from dataclasses import replace
 from tinygrad import Tensor, Device, Context
 from tinygrad.helpers import getenv
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+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

@@ -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.")
@@ -4,9 +4,9 @@ from tinygrad import dtypes
 from typing import Optional, List, Tuple, cast, Dict, Final, DefaultDict, Self

 # 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
@@ -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
@@ -1,6 +1,6 @@
 from tinygrad import Tensor, dtypes, Device
 from tinygrad.helpers import getenv, DEBUG
-from tinygrad.opt.kernel import Kernel, Opt, OptOps
+from tinygrad.codegen.kernel import Kernel, Opt, OptOps
 from tinygrad.engine.realize import CompiledRunner, ExecItem
 from dataclasses import replace

@@ -37,7 +37,7 @@ B = Tensor.rand(K, N, device="CPU")
 C = (A.reshape(M, 1, K) * B.permute(1,0).reshape(1, N, K)).sum(axis=2)

 sched = C.schedule()
-from tinygrad.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.linearize()
@@ -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
 ```
@@ -4,9 +4,9 @@ import numpy as np
 np.set_printoptions(suppress=True)
 import math, functools, time, random, statistics
 from tinygrad.helpers import DEBUG, getenv, CACHELEVEL, diskcache_get, diskcache_put, colored, Profiling
-from tinygrad.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.search import _ensure_buffer_alloc, get_kernel_actions, _time_program

 class MCTSNode:
  def __init__(self, kernel:Kernel, parent=None):
@@ -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?
@@ -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
@@ -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 */
@@ -1,209 +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
-{
-    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 */
@@ -1,2 +0,0 @@
-GPU="$1"
-echo 1 | sudo tee /sys/bus/pci/devices/$GPU/reset 2>/dev/null
@@ -1,28 +1,11 @@
 from types import SimpleNamespace
 from typing import Any, Sequence, cast, Literal, Callable
-import dataclasses, functools, io, math, types, warnings, sys
+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
-from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype
+from tinygrad.dtype import DType, ConstType, dtypes, ImageDType
 from tinygrad.device import is_dtype_supported, Device

-# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L500-L544
-data_types: dict[int, DType] = {
-  1:dtypes.float32, 2:dtypes.uint8, 3:dtypes.int8, 4:dtypes.uint16, 5:dtypes.int16, 6:dtypes.int32, 7:dtypes.int64,
-  9:dtypes.bool, 10:dtypes.float16, 11:dtypes.double, 12:dtypes.uint32, 13:dtypes.uint64, 16:dtypes.bfloat16,
-}
-
-# https://github.com/onnx/onnx/blob/rel-1.17.0/onnx/onnx.proto3#L128-L145
-attribute_types: dict[int, Callable] = {
-  1: lambda a: float(a.f),
-  2: lambda a: int(a.i),
-  3: lambda a: a.s.data().tobytes().decode("utf8") if isinstance(a.s, Tensor) else a.s.decode("utf8"),
-  4: lambda a: buffer_parse(a.t),
-  6: lambda a: tuple(float(x) for x in a.floats),
-  7: lambda a: tuple(int(x) for x in a.ints),
-  8: lambda a: tuple(x.data().tobytes().decode("utf8") for x in a.strings)
-}
-
 # ***** protobuf parsing ******
 from onnx import AttributeProto, ModelProto, TensorProto, TypeProto, helper
 import numpy as np
@@ -31,24 +14,39 @@ 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.")
-  dtype, shape = dtype_parse(onnx_tensor.data_type, "buffer parse"), tuple(onnx_tensor.dims)
+  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
@@ -59,17 +57,13 @@ def buffer_parse(onnx_tensor: TensorProto) -> Tensor:
    if len(data) == 1: return Tensor(data.tolist()[0], dtype=dtype).reshape(shape)
    return data.cast(dtype).reshape(shape).to(Device.DEFAULT)
  if has_field(onnx_tensor, "raw_data"):
-    raw_data = onnx_tensor.raw_data
-    if not isinstance(raw_data, Tensor): raw_data = Tensor(raw_data)
-    if not is_dtype_supported(data_types[onnx_tensor.data_type]):
-      np_buffer = np.frombuffer(raw_data.data().tobytes(),
+    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 = raw_data.bitcast(dtype).reshape(shape).to(Device.DEFAULT)
-    if shape == ():
-      if ret.dtype is dtypes.float16 and sys.version_info < (3, 12): ret = ret.cast(dtypes.float32)
-      ret = Tensor(ret.item(), dtype=dtype).reshape(shape)
+    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)

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

@@ -149,19 +143,17 @@ class OnnxRunner:

  def _parse_input(self, name: str, value: Any, spec: OnnxValue):
    if spec.is_optional and value is None: return None
+    # TODO: need true float16 for dtype checking
    if spec.is_sequence:
-      if not isinstance(value, Sequence): raise RuntimeError(f"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 str(type(value)) == "<class 'numpy.ndarray'>" else spec.dtype
-    tensor = Tensor(value, dtype=dtype, requires_grad=self.is_training) if not isinstance(value, Tensor) else value
-    if tensor.dtype is not spec.dtype: warnings.warn(f"input {name} has mismatch on dtype. Expected {spec.dtype}, received {tensor.dtype}.")
+    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 _dispatch_op(self, op, inps, opts):
@@ -175,8 +167,8 @@ class OnnxRunner:
      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 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():
@@ -292,7 +284,7 @@ def get_onnx_ops():
    raise ValueError(f"pixel_format={pixel_format!r} is not supported.")

  def EyeLike(x:Tensor, dtype:int|None=None, k:int=0):
-    ret = Tensor.eye(cast(int, min(x.shape)), dtype=dtype_parse(dtype, "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)
@@ -325,7 +317,7 @@ def get_onnx_ops():
  def Binarizer(x:Tensor, threshold:float=0.0): return (x > threshold).float()

  # ***** Unary Ops (broadcasted) *****
-  def Add(x:Tensor,y:Tensor, broadcast=None, axis=None): return x + y
+  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
@@ -346,7 +338,7 @@ def get_onnx_ops():

  # ***** 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 *****
@@ -613,13 +605,9 @@ def get_onnx_ops():

  # Reimplemented here because you need legacy RNG for passing ONNX tests.
  def Dropout_7(data:Tensor, ratio:float=0.5, training_mode:bool=False, seed:int|None=None):
-    if not training_mode: return data, 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
+    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 = {6:Dropout_6, 7:Dropout_7}
@@ -743,9 +731,7 @@ def get_onnx_ops():

  # ***** Quantization Ops *****
  def QuantizeLinear(x:Tensor, y_scale:Tensor, y_zero_point:Tensor|int=0, axis:int=1, block_size:int=0, output_dtype:int=0, saturate=1):
-    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
@@ -25,7 +25,6 @@ class PBType: FLOAT = 1; INT = 2; STRING = 3; FLOATS = 4; INTS = 5; STRINGS = 6;
 PB_INFOS = {
  "OperatorSetIdProto": {1: ("domain", PBType.STRING), 2: ("version", PBType.INT)},
  "StringStringEntryProto": {1: ("key", PBType.STRING), 2: ("value", PBType.STRING)},
-  # TODO: support uint64 parsing (11: "uint64_data") and double parsing (10: "double_data")
  "TensorProto": {1: ("dims", PBType.INT, True), 2: ("data_type", PBType.INT), 4: ("float_data", PBType.FLOATS),
    13: ("external_data", PBType.SUB, True, "StringStringEntryProto"), 14: ("data_location", PBType.INT),
    5: ("int32_data", PBType.INTS), 7: ("int64_data", PBType.INTS), 8: ("name", PBType.STRING), 9: ("raw_data", PBType.BYTES)},
@@ -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)
@@ -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:
@@ -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
@@ -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	
@@ -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()
@@ -1,16 +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
 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):
@@ -27,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]
@@ -102,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
@@ -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

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

@@ -6,8 +6,8 @@ from tinygrad.helpers import getenv, BEAM
 from tinygrad.engine.jit import TinyJit
 from tinygrad.engine.realize import CompiledRunner, ExecItem, ScheduleItem, lower_schedule_item
 from tinygrad.renderer import ProgramSpec
-from tinygrad.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"):
@@ -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))
@@ -369,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,
@@ -378,7 +376,6 @@ decomps = [
  aten.softshrink,
  aten.hardshrink,
  aten.log_sigmoid_forward,
-  aten.log_sigmoid_backward,
  aten.isneginf,
  aten.isposinf,
  aten.nan_to_num,
@@ -206,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()
@@ -1,73 +0,0 @@
-# play with upcasted warps
-from tinygrad import Tensor, Device
-from tinygrad.uop.ops import KernelInfo
-from tinygrad.opt import get_optimized_ast
-from tinygrad.opt.kernel import OptOps, Opt
-from tinygrad.engine.realize import get_program
-
-if __name__ == "__main__":
-  renderer = Device.default.renderer
-  N = 64
-
-  """
-  a = Tensor.empty(N,N)
-
-  out = (a + 1) #.sum(axis=2)
-  ast = out.schedule()[-1].ast
-  opts = tuple()
-  opts += (Opt(OptOps.UPCAST, 0, 32),)
-  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
-  ast = get_optimized_ast(ast, renderer)
-  prg = get_program(ast, renderer)
-  print(prg.src)
-  """
-
-  # how you split the store determines everything if you don't allow cross warp comms.
-  # actually not everything, there's also the split before the horizontal (unrolled) reduces
-
-  # new flow
-  #  - pull out any dimensions from the store that you want to upcast.
-  #  - decide how you want to assign them to registers. GPUs have a 512-byte memory LOAD/STORE which loads into 4 regs. see BUFFER_LOAD_B128
-  #    - the loads and stores can be shuffled, but only in restrictive ways. in kernels without reduces, the store determines everything
-  #    - it loads 16 bytes from up 32 different places = 512 bytes
-  #  - in kernels with reduces, you now have more flexibility. the final target of the reduce must be what is stored
-  #    - warp dimensions can be in the reduce (this is GROUP)
-
-  # every dimension can be assigned to <global, local, loop, upcast, warp>
-
-  """
-  out = a.sum(axis=1)
-  ast = out.schedule()[-1].ast
-  opts = tuple()
-  opts += (Opt(OptOps.UPCAST, 0, 8),)
-  opts += (Opt(OptOps.UNROLL, 0, 8),)
-  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
-  ast = get_optimized_ast(ast, renderer)
-  prg = get_program(ast, renderer)
-  print(prg.src)
-
-  out = a.sum(axis=1)
-  ast = out.schedule()[-1].ast
-  opts = tuple()
-  opts += (Opt(OptOps.UNROLL, 0, 8),)
-  opts += (Opt(OptOps.UPCAST, 0, 8),)
-  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
-  ast = get_optimized_ast(ast, renderer)
-  prg = get_program(ast, renderer)
-  print(prg.src)
-  """
-
-  # gemm
-  b = Tensor.empty(N,N)
-  # metal TC
-  #opts = (Opt(OptOps.UPCAST, 0, 2), # not the warp
-  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2), Opt(OptOps.UPCAST, 1, 2),
-  #        Opt(OptOps.UPCAST, 0, 2), Opt(OptOps.UPCAST, 1, 2))
-  # new TC should just be able to extract from this and swizzle as needed
-  opts = (Opt(OptOps.UPCAST, 0, 8), Opt(OptOps.UPCAST, 1, 8), Opt(OptOps.UNROLL, 0, 8))
-  c = (a@b)
-  ast = c.schedule()[-1].ast
-  ast = ast.replace(arg=KernelInfo(opts_to_apply=opts))
-  ast = get_optimized_ast(ast, renderer)
-  prg = get_program(ast, renderer)
-  print(prg.src)
@@ -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:
@@ -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"
-]
@@ -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.kernelize',
-                  'tinygrad.runtime.support.am', 'tinygrad.runtime.graph', 'tinygrad.shape', 'tinygrad.uop', 'tinygrad.opt',
-                  'tinygrad.runtime.support.nv'],
-      package_data = {'tinygrad': ['py.typed'], 'tinygrad.viz': ['index.html', 'assets/**/*', 'js/*']},
+                  '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"
@@ -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):
@@ -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()
@@ -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))
@@ -7,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())
@@ -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
-from tinygrad.opt.search import bufs_from_lin
+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
@@ -1,5 +1,4 @@
 import random
-from typing import Optional
 from tinygrad.helpers import round_up
 from tinygrad.runtime.support.am.amdev import AMPageTableTraverseContext
 from test.external.external_test_am import helper_read_entry_components, FakeAM
@@ -30,7 +29,7 @@ class AMPTFuzzer:
        self.d.vram[pte['paddr']] = pattern # Mark this page
        assert pte['valid'] == 1

-        # If page has contiguous fragment, all range should be this valid memory
+        # If page has contigous fragment, all range should be this valid memory
        frags_cnt = pte['fragment']
        contig_range = (1 << (frags_cnt + 12))
        start_vaddr = _vaddr & ~(contig_range - 1)
@@ -60,7 +59,7 @@ class AMPTFuzzer:

    return True

-  def random_alloc(self) -> Optional[int]:
+  def random_alloc(self):
    if self.total_size - self.alloc_payload < self.min_alloc_size: return None

    size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))
@@ -1,7 +1,7 @@
 import random
 from typing import Dict, Optional
 from tinygrad.helpers import getenv
-from tinygrad.runtime.support.memory import TLSFAllocator
+from tinygrad.runtime.support.allocator import TLSFAllocator

 class AllocatorFuzzer:
  def __init__(self, total_size):
@@ -30,8 +30,6 @@ class AllocatorFuzzer:
    return True

  def random_alloc(self) -> Optional[int]:
-    if self.total_size - self.alloc_payload < self.min_alloc_size: return None
-
    size = random.randint(self.min_alloc_size, min(self.max_alloc_size, self.total_size - self.alloc_payload))

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

 MODELS = {
  "resnet50": "https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet50-caffe2-v1-9.onnx",
@@ -27,7 +27,6 @@ MODELS = {
  # really slow
  # "resnet18": "https://github.com/onnx/models/raw/main/archive/vision/classification/resnet/model/resnet18-v2-7.onnx",
 }
-half_models = ["openpilot", "commavq"]

 CSV = {}
 open_csv = None
@@ -55,6 +54,7 @@ def benchmark_model(m, devices, validate_outs=False):
  excluded = {inp.name for inp in onnx_model.graph.initializer}
  input_shapes = {inp.name:tuple(x.dim_value if hasattr(x, "dim_value") and x.dim_value != 0 else 1 for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input if inp.name not in excluded}  # noqa: E501
  input_types = {inp.name: tensor_dtype_to_np_dtype(inp.type.tensor_type.elem_type) for inp in onnx_model.graph.input if inp.name not in excluded}
+  #input_types = {k:v if v!=np.float16 else np.float32 for k,v in input_types.items()}  # cast
  np_inputs = {k:torch.randn(shp).numpy().astype(input_types[k]) for k,shp in input_shapes.items()}
  assert len(input_shapes) < 30, f"too many input shapes {len(input_shapes)}"

@@ -106,12 +106,7 @@ def benchmark_model(m, devices, validate_outs=False):
    for device in devices:
      rtol, atol = 2e-3, 2e-3  # tolerance for fp16 models
      Device.DEFAULT = device
-      # force half inputs to float for numerical stability when validating
-      # this will reply on automatic dtype promotion for converting half weights inside the graph
-      if m in half_models:
-        inputs = {k:Tensor(inp, dtype=dtypes.float32) if inp.dtype == np.float16 else Tensor(inp) for k,inp in np_inputs.items()}
-      else:
-        inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
+      inputs = {k:Tensor(inp) for k,inp in np_inputs.items()}
      tinygrad_model = OnnxRunner(onnx_model)
      tinygrad_out = tinygrad_model(inputs)

@@ -1,8 +1,7 @@
 import unittest
-from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableEntry
+from tinygrad.runtime.support.am.amdev import AMMemoryManager, AMPageTableTraverseContext
 from tinygrad.runtime.support.am.ip import AM_GMC
 from tinygrad.runtime.support.hcq import MMIOInterface
-from tinygrad.runtime.support.memory import PageTableTraverseContext
 from tinygrad.runtime.autogen.am import am
 from tinygrad.helpers import mv_address

@@ -24,9 +23,7 @@ class FakeAM:
    self.vram_mv = memoryview(bytearray(4 << 30))
    self.vram = MMIOInterface(mv_address(self.vram_mv), self.vram_mv.nbytes)
    self.gmc = FakeGMC(self)
-    self.mm = AMMemoryManager(self, 4 << 30, boot_size=(32 << 20), pt_t=AMPageTableEntry, pte_cnt=[512, 512, 512, 512],
-      pte_covers=[(1 << ((9 * (3-lv)) + 12)) for lv in range(4)], first_lv=am.AMDGPU_VM_PDB1, first_page_lv=am.AMDGPU_VM_PDB2,
-      va_base=AMMemoryManager.va_allocator.base)
+    self.mm = AMMemoryManager(self, vram_size=4 << 30)
    self.is_booting = False
    self.ip_ver = {am.GC_HWIP: (11, 0, 0)}
  def paddr2cpu(self, paddr:int) -> int: return paddr + mv_address(self.vram)
@@ -68,7 +65,7 @@ class TestAMPageTable(unittest.TestCase):
      exteranl_va = va + AMMemoryManager.va_allocator.base
      mm.map_range(vaddr=exteranl_va, size=sz, paddrs=[(va, sz)])

-      ctx = PageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
+      ctx = AMPageTableTraverseContext(self.d[0], mm.root_page_table, exteranl_va)
      results = list(ctx.next(sz))

      total_covered = 0
@@ -94,17 +91,6 @@ class TestAMPageTable(unittest.TestCase):
          assert pte['paddr'] == 0
          assert pte['valid'] == 0

-  def test_map_notaligned(self):
-    mm0 = self.d[0].mm
-
-    for (va1,sz1),(va2,sz2) in [((0x10000, (0x1000)), (0x11000, (2 << 20)))]:
-      exteranl_va1 = va1 + AMMemoryManager.va_allocator.base
-      exteranl_va2 = va2 + AMMemoryManager.va_allocator.base
-      mm0.map_range(vaddr=exteranl_va1, size=sz1, paddrs=[(va1, sz1)])
-      mm0.map_range(vaddr=exteranl_va2, size=sz2, paddrs=[(va2, sz2)])
-      mm0.unmap_range(va2, sz2)
-      mm0.unmap_range(va1, sz1)
-
  def test_double_map(self):
    mm0 = self.d[0].mm

@@ -129,7 +115,7 @@ class TestAMPageTable(unittest.TestCase):
      # Finally can map and check paddrs
      mm0.map_range(vaddr=exteranl_va + 0x2000, size=0x100000, paddrs=[(0xdead0000, 0x1000), (0xdead1000, 0xff000)])

-      ctx = PageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
+      ctx = AMPageTableTraverseContext(self.d[0], mm0.root_page_table, exteranl_va + 0x2000)
      for tup in ctx.next(0x100000):
        _offset, _pt, _pte_idx, _n_ptes, _pte_covers = tup
        for i in range(_n_ptes):
@@ -64,7 +64,6 @@ class TestKiTS19Dataset(ExternalTestDatasets):

    return iter(dataset)

-  @unittest.skip("flaky")
  def test_training_set(self):
    preproc_pth, preproc_img_pths, preproc_lbl_pths = self._create_samples(False)
    ref_dataset = self._create_ref_dataloader(preproc_img_pths, preproc_lbl_pths, False)
@@ -4,10 +4,10 @@ os.environ["VALIDATE_HCQ"]="1"

 import unittest, random
 import numpy as np
-from tinygrad.opt.kernel import Kernel, KernelOptError
+from tinygrad.codegen.kernel import Kernel, KernelOptError
 from tinygrad.device import is_dtype_supported
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from tinygrad import Device, dtypes, Tensor
 from test.external.fuzz_linearizer import compare_linearizer, compare_states, get_fuzz_rawbuf_like

@@ -3,7 +3,7 @@ from tinygrad.runtime.support.hip_comgr import compile_hip
 from tinygrad import Tensor
 from tinygrad.device import Device
 from tinygrad.engine.schedule import create_schedule
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel

 class TestHIPCompileSpeed(unittest.TestCase):
  @unittest.skipIf(Device.DEFAULT != "HIP", "only run on HIP")
@@ -1,10 +1,10 @@
 from tinygrad import Tensor
 from test.external.mlperf_unet3d.dice import DiceScore
-from examples.mlperf.metrics import dice_score, log_perplexity
+from examples.mlperf.metrics import dice_score

 import numpy as np
 import torch
-import unittest, math
+import unittest

 class ExternalTestMetrics(unittest.TestCase):
  def _test_metrics(self, tinygrad_metrics, orig_metrics, pred, label, atol=1e-8, rtol=1e-7):
@@ -16,27 +16,5 @@ class ExternalTestMetrics(unittest.TestCase):
    pred, label = np.random.rand(1, 3, 128, 128, 128).astype(np.float32), np.ones((1, 1, 128, 128, 128)).astype(np.uint8)
    self._test_metrics(dice_score, DiceScore(), pred, label)

-  def test_log_perplexity(self):
-    # equally likely
-    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0, 1, 1, 1]]]), Tensor([[2]])).numpy(), math.log(4))
-    np.testing.assert_allclose(log_perplexity(Tensor([[[1.0]*256]*32]), Tensor([[2]*32])).numpy(), math.log(256), rtol=1e-6)
-    # pretty correct and incorrect
-    np.testing.assert_allclose(log_perplexity(Tensor([[[10000., 0, 0, 0]]]), Tensor([[0]])).numpy(), 0)
-    np.testing.assert_allclose(log_perplexity(Tensor([[[0.0, 10000, 10000, 10000]]]), Tensor([[0]])).numpy(), 10000, rtol=1e-3)
-    # higher logit -> lower loss
-    x = Tensor([[[4.0, 3, 2, 1]]])
-    for i in range(x.numel()-1): self.assertLess(log_perplexity(x, Tensor([[i]])).item(), log_perplexity(x, Tensor([[i+1]])).item())
-    # torch eval examples
-    np.testing.assert_allclose(
-      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]])).exp().numpy(),
-      2.7593, rtol=1e-5)
-    np.testing.assert_allclose(
-      log_perplexity(Tensor([[[0.3, 0.7, 0.3, 0.1], [0.5, 0.4, 0.1, 0.4],[0.1, 0.1, 0.2, 0.5]],
-                             [[0.1, 0.6, 0.1, 0.5], [0.3, 0.7, 0.3, 0.4], [0.3, 0.7, 0.3, 0.4]]]), Tensor([[2, 1, 3],  [1, 0, 1]])).exp().numpy(),
-      3.6216, rtol=1e-5)
-    np.testing.assert_allclose(
-      log_perplexity(Tensor([[[0.3659, 0.7025, 0.3104], [0.0097, 0.6577, 0.1947]]]), Tensor([[2, 1]]), ignore_index=1).exp().numpy(),
-      3.5372, rtol=1e-4)
-
 if __name__ == '__main__':
  unittest.main()
@@ -2,12 +2,12 @@ import unittest, struct, array, ctypes
 from tinygrad import Device, dtypes, Tensor
 from tinygrad.helpers import to_mv
 from tinygrad.runtime.ops_nv import NVDevice, HWQueue
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from test.test_linearizer_failures import helper_test_lin
 from tinygrad.engine.realize import get_runner, CompiledRunner
 from test.external.fuzz_linearizer import get_fuzz_rawbufs

-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel
 from tinygrad.uop.ops import LazyOp, Ops, ReduceOps, BufferOps, MemBuffer
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
@@ -44,9 +44,6 @@ class TinygradBackend(Backend):

 backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__)

-# BUG: segfaults
-backend_test.exclude('test_MaxPool1d_stride_padding_dilation_cpu')
-
 # BUG: buggy onnx tests
 backend_test.exclude('test_adam_multiple_cpu')

@@ -94,6 +91,13 @@ backend_test.exclude('FLOAT8')
 backend_test.exclude('INT4')
 backend_test.exclude('UINT4')
 backend_test.exclude('BFLOAT16')  # not supported in numpy
+# TODO: fix these with true onnx float16
+backend_test.exclude('to_FLOAT16')
+backend_test.exclude('cast_no_saturate')
+backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
+backend_test.exclude('test_max_float16_cpu')
+backend_test.exclude('test_min_float16_cpu')
+backend_test.exclude('test_mod_mixed_sign_float16_cpu')

 backend_test.exclude('test_dequantizelinear_int4_cpu')
 backend_test.exclude('test_dequantizelinear_uint4_cpu')
@@ -107,7 +111,6 @@ backend_test.exclude('test_quantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_quantizelinear_e5m2_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_cpu')
 backend_test.exclude('test_dequantizelinear_e4m3fn_zero_point_cpu')
-backend_test.exclude('test_dequantizelinear_e4m3fn_float16_cpu')
 backend_test.exclude('test_dequantizelinear_e5m2_cpu')

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

 from tinygrad.tensor import Tensor
-from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup, AdamW
+from tinygrad.nn.optim import LAMB, LARS, SGD, OptimizerGroup

 from test.external.mlperf_resnet.lars_optimizer import LARSOptimizer

-from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup, CosineAnnealingLRWithWarmup
+from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
 from test.external.mlperf_resnet.lars_util import PolynomialDecayWithWarmup as PolynomialDecayWithWarmup_tf

 np.random.seed(1337)
@@ -171,26 +171,5 @@ class ExternalTestOptim(unittest.TestCase):
      'warmup': steps_per_epoch * warmup_epochs,
    }, 1e-5, 1e-5, do_optim=False)

-
-class TestCosineAnnealingLRWithWarmup(unittest.TestCase):
-  # only tests the lr
-  def _test_lr(self, base_lr, end_lr, warmup_steps, decay_steps):
-    net = TinyNet()
-    optim = AdamW([net.W], lr=0.0)
-    tiny_lr = CosineAnnealingLRWithWarmup(optim, base_lr, end_lr, warmup_steps, decay_steps)
-    lr = []
-    for _ in range(warmup_steps+decay_steps):
-      lr.append(optim.lr.item())
-      tiny_lr.step()
-    # reimplemented in python
-    expected = []
-    for i in range(warmup_steps): expected.append((i+1)/warmup_steps*base_lr)
-    for i in range(decay_steps): expected.append(end_lr+(base_lr-end_lr)*(1+math.cos((i+1)/decay_steps*math.pi))/2)
-    np.testing.assert_allclose(lr, expected, rtol=1e-5)
-
-  def test_lr_0(self): self._test_lr(3e-4, 8e-5, 3, 5)
-  def test_lr_1(self): self._test_lr(3e-4, 8e-5, 10, 20)
-  def test_lr_llama3(self): self._test_lr(8e-5, 8e-7, 20, 100)
-
 if __name__ == '__main__':
  unittest.main()
@@ -1,5 +1,5 @@
 import unittest
-from tinygrad.runtime.support.memory import TLSFAllocator
+from tinygrad.runtime.support.allocator import TLSFAllocator

 class TestTLSFAllocator(unittest.TestCase):
  def setUp(self):
@@ -2,11 +2,11 @@
 import unittest

 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel

 from test.external.fuzz_linearizer import run_linearizer

@@ -3,11 +3,11 @@ import unittest

 from tinygrad import Device
 from tinygrad.uop.ops import UOp, Ops
-from tinygrad.opt.search import Opt, OptOps
+from tinygrad.engine.search import Opt, OptOps
 from tinygrad.dtype import dtypes
 from tinygrad.shape.shapetracker import ShapeTracker
 from tinygrad.shape.view import View
-from tinygrad.opt.kernel import Kernel
+from tinygrad.codegen.kernel import Kernel

 class TestOpenpilotValidhack(unittest.TestCase):
  def test_valid_removal(self):
@@ -1,7 +1,7 @@
 import gc
 from tinygrad import Tensor, UOp, Device
 from tinygrad.shape.shapetracker import views_to_indexed_uops
-from tinygrad.engine.realize import method_cache, get_program
+from tinygrad.engine.realize import method_cache, get_kernel

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

 from tinygrad import Tensor, Device, dtypes
 from tinygrad.tensor import _to_np_dtype
-from tinygrad.opt.kernel import Kernel
-from tinygrad.opt.kernel import Opt, OptOps
-from tinygrad.opt.search import get_kernel_actions, bufs_from_lin
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.codegen.kernel import Opt, OptOps
+from tinygrad.engine.search import get_kernel_actions, bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import getenv, from_mv, prod, colored, Context, DEBUG, Timing
 from tinygrad.uop.ops import UOp, Ops
@@ -1,12 +1,11 @@
 #!/usr/bin/env python3
 # compare kernels created by HEAD against master
-import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools, functools
+import os, multiprocessing, logging, pickle, sqlite3, difflib, warnings, itertools
 from typing import Callable, Any
-from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm
-from tinygrad.kernelize.kernelize import get_kernelize_map
-from tinygrad.renderer import Renderer, ProgramSpec
-from tinygrad.engine.realize import get_program
-from tinygrad.uop.ops import UOp, Ops, KernelInfo
+from tinygrad.helpers import VERSION, Context, ContextVar, colored, db_connection, getenv, tqdm, to_function_name
+from tinygrad.engine.grouper import get_kernelize_map
+from tinygrad.codegen.kernel import Kernel
+from tinygrad.uop.ops import UOp, Ops

 # *** process replay settings

@@ -40,16 +39,21 @@ def replay_kernelize(ret:dict[UOp, UOp], big_sink:UOp) -> tuple[str, str, tuple[
    return "\n".join([f"{len(asts)} kernels", *asts])
  return to_str(new_sink), to_str(ret[big_sink]), (big_sink,)

-def replay_get_program(p:ProgramSpec, ast:UOp, renderer:Renderer) -> tuple[str, str, tuple[Any, ...]]:
-  p2 = get_program(ast.replace(arg=KernelInfo(opts_to_apply=p.applied_opts, name=p.name)) if ast.arg is None else ast, renderer)
-  def to_str(ret:ProgramSpec) -> str: return ret.src
-  return to_str(p2), to_str(p), (p.ast, renderer, p.applied_opts)
+def replay_linearize(k:Kernel, _:Kernel, name_override=None, ast_transform=None) -> tuple[str, str, tuple[Any, ...]]:
+  # create a copy because the Kernel class contains optimization parameters (other than applied_opts) in its state
+  # this should be made fully functional. It's fine for process replay since copy returns a fresh instance
+  k2 = k.copy()
+  k2.linearize(name_override=name_override or to_function_name(k.name), ast_transform=ast_transform)
+  def to_str(ret:Kernel) -> str:
+    try: return ret.opts.render(ret.uops)
+    except NotImplementedError: return "" # NULL backend doesn't have a renderer, this is okay
+  return to_str(k2), to_str(k), (k.ast, k.opts, k.applied_opts)

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

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

-def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
+def diff(offset:int) -> None:
  if ASSERT_DIFF: warnings.filterwarnings("error", category=ProcessReplayWarning)
  if early_stop.is_set(): return None
  conn = db_connection()
@@ -64,10 +68,8 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
    try:
      name, args, kwargs, ctx_vals, loc, ret = pickle.loads(row[0])
      ctx_vars = {k:v.value for k,v in ctx_vals.items() if k != "DEBUG" and (var:=ContextVar._cache.get(k)) is not None and var.value != v.value}
-      if (replayer:=fxns.get(name)) is None: continue
-      with Context(**ctx_vars):
-        if (ret:=replayer(ret, *args, **kwargs)) is None: continue
-        good, compare, metadata = ret
+      if (replayer:=replayers.get(name)) is None: continue
+      with Context(**ctx_vars): good, compare, metadata = replayer(ret, *args, **kwargs)
      if good != compare:
        for m in metadata: trunc_log(m)
        logging.info(loc)
@@ -81,25 +83,6 @@ def diff(offset:int, fxns:dict[str, Callable[..., tuple|None]]) -> None:
  conn.commit()
  cur.close()

-# *** generic runner to map rows of a table to a function in parallel
-
-def _pmap(fxns:dict[str, Callable]) -> None:
-  conn = db_connection()
-  cur = conn.cursor()
-  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
-  except sqlite3.OperationalError:
-    raise RuntimeError(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?")
-  finally:
-    conn.commit()
-    cur.close()
-
-  with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
-    inputs = list(range(0, row_count, PAGE_SIZE))
-    list(tqdm(pool.imap_unordered(functools.partial(diff, fxns=fxns), inputs), total=len(inputs)))
-    pool.close()
-    pool.join()
-    pool.terminate()
-
 # *** main loop

 if __name__ == "__main__":
@@ -107,8 +90,20 @@ if __name__ == "__main__":
    logging.info("skipping process replay.")
    exit(0)

-  logging.info(f"running process replay with {ASSERT_DIFF=}")
-  try: _pmap(replayers)
-  except Exception as e:
-    logging.info("process replay err", e)
+  conn = db_connection()
+  cur = conn.cursor()
+  try: row_count = cur.execute(f"select count(*) from '{TABLE_NAME}'").fetchone()[0]
+  except sqlite3.OperationalError:
+    warnings.warn(f"{TABLE_NAME} isn't accessible in master, did DB_VERSION change?", ProcessReplayWarning)
    exit(int(ASSERT_DIFF))
+  finally:
+    conn.commit()
+    cur.close()
+
+  logging.info(f"running process replay with {ASSERT_DIFF=}")
+  with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
+    inputs = list(range(0, row_count, PAGE_SIZE))
+    list(tqdm(pool.imap_unordered(diff, inputs), total=len(inputs)))
+    pool.close()
+    pool.join()
+    pool.terminate()
--- a/Show More
+++ b/Show More