Progress with pack/unpack wgsl built-ins

.github/actions/process-replay/action.yml

@@ -1,15 +0,0 @@
-name: Run process replay tests
-description: Verify process replay compared to master
-runs:
-  using: "composite"
-  steps:
-    - name: Run process replay tests
-      shell: bash
-      run: |
-        export PR_TITLE=$(jq -r .pull_request.title "$GITHUB_EVENT_PATH")
-        export CURRENT_SHA=${{ github.event.pull_request && github.event.pull_request.head.sha || github.sha }}
-        git fetch origin $CURRENT_SHA
-        export COMMIT_MESSAGE=$(git show -s --format=%B "$CURRENT_SHA")
-        export CURRENT_HEAD=$(git rev-parse HEAD)
-        cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
-        git checkout $CURRENT_HEAD  # restore to branch

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

@@ -1,224 +0,0 @@
-name: Setup Python & Install
-description: Sets up Python and installs project dependencies.
-inputs:
-  python-version:
-    description: 'Python version to use'
-    required: false
-    default: '3.12'
-  key:
-    description: 'Key for the python cache'
-    required: false
-    default: ''  # if you don't set a key, it doesn't cache
-  deps:
-    description: 'Extra dependency groups (comma separated)'
-    required: false
-    default: ''
-  pydeps:
-    description: 'Extra Python dependency groups (space separated)'
-    required: false
-    default: ''
-  opencl:
-    description: "Install OpenCL?"
-    required: false
-    default: 'false'
-  amd:
-    description: "Install AMD?"
-    required: false
-    default: 'false'
-  cuda:
-    description: "Install CUDA?"
-    required: false
-    default: 'false'
-  webgpu:
-    description: "Install webgpu?"
-    required: false
-    default: 'false'
-  llvm:
-    description: "Install LLVM?"
-    required: false
-    default: 'false'
-runs:
-  using: "composite"
-  steps:
-    - name: Set up Python ${{ inputs.python-version }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ inputs.python-version }}
-
-    # **** Caching packages ****
-    # TODO: key should include input.deps, but it can't since it can't contain commas
-
-    - name: Cache Python packages (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
-      with:
-        path: ${{ env.Python3_ROOT_DIR }}/lib/python${{ inputs.python-version }}/site-packages
-        key: python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
-    - name: Cache Python packages (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
-      uses: actions/cache@v4
-      with:
-        path: /Users/runner/Library/Python/${{ inputs.python-version }}/lib/python/site-packages
-        key: osx-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
-    - name: Cache Python packages (Windows)
-      if: inputs.key != '' && runner.os == 'Windows'
-      uses: actions/cache@v4
-      with:
-        path: ${{ env.Python3_ROOT_DIR }}\Lib\site-packages
-        key: windows-python-package-${{ inputs.key }}-${{ hashFiles('**/setup.py') }}
-
-    # **** Caching downloads ****
-
-    - name: Cache downloads (Linux)
-      if: inputs.key != '' && runner.os == 'Linux'
-      uses: actions/cache@v4
-      with:
-        path: ~/.cache/tinygrad/downloads/
-        key: downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
-    - name: Cache downloads (macOS)
-      if: inputs.key != '' && runner.os == 'macOS'
-      uses: actions/cache@v4
-      with:
-        path: ~/Library/Caches/tinygrad/downloads/
-        key: osx-downloads-cache-${{ inputs.key }}-${{ env.DOWNLOAD_CACHE_VERSION }}
-
-    # **** Python deps ****
-
-    - name: Install dependencies (with extra)
-      if: inputs.deps != ''
-      shell: bash
-      run: pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e ".[${{ inputs.deps }}]" ${{ inputs.pydeps }} --extra-index-url https://download.pytorch.org/whl/cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
-    - name: Install dependencies (without extra)
-      if: inputs.deps == ''
-      shell: bash
-      run: pip install ${{ (runner.os == 'macOS' && '--user') || (runner.os != 'macOS' && '') }} -e . ${{ inputs.pydeps }}
-
-    # **** OpenCL ****
-
-    - name: Install OpenCL
-      if: inputs.opencl == 'true'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        echo "deb [ allow-insecure=yes ] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list
-        sudo apt update || true
-        sudo apt install --allow-unauthenticated -y --no-install-recommends opencl-headers \
-          intel-oneapi-runtime-openmp=2023.2.1-16 intel-oneapi-runtime-compilers-common=2023.2.1-16 intel-oneapi-runtime-compilers=2023.2.1-16 \
-          intel-oneapi-runtime-dpcpp-sycl-opencl-cpu=2023.2.1-16 intel-oneapi-runtime-tbb-common=2021.10.0-49541 \
-          intel-oneapi-runtime-tbb=2021.10.0-49541 intel-oneapi-runtime-opencl=2023.2.1-16
-
-    # **** AMD ****
-
-    - name: Install AMD (Linux)
-      if: inputs.amd == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
-        sudo tee /etc/apt/sources.list.d/rocm.list <<'EOF'
-        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.1.2 jammy main
-        EOF
-        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
-        sudo apt update || true
-        sudo apt install --no-install-recommends --allow-unauthenticated -y hsa-rocr comgr hsa-rocr-dev liburing-dev libc6-dev
-        curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-        jq -r '.assets[] | select(.name == "libremu.so").browser_download_url' | \
-        sudo xargs curl -L -o /usr/local/lib/libremu.so
-        sudo tee --append /etc/ld.so.conf.d/rocm.conf <<'EOF'
-          /opt/rocm/lib
-          /opt/rocm/lib64
-        EOF
-        sudo ldconfig
-    - name: Install AMD comgr+remu (macOS)
-      if: inputs.amd == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        sudo mkdir -p /usr/local/lib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/nimlgen/amdcomgr_dylib/releases/latest | \
-          jq -r '.assets[] | select(.name == "libamd_comgr.dylib").browser_download_url' | \
-          sudo xargs curl -L -o /usr/local/lib/libamd_comgr.dylib
-        curl -s -H "Authorization: token $GH_TOKEN" curl -s https://api.github.com/repos/Qazalin/remu/releases/latest | \
-          jq -r '.assets[] | select(.name == "libremu.dylib").browser_download_url' | \
-          sudo xargs curl -L -o /usr/local/lib/libremu.dylib
-
-    # **** CUDA ****
-
-    - name: Install cuda packages (Linux)
-      if: inputs.cuda == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        sudo apt update -y || true
-        sudo apt install -y --no-install-recommends git g++ cmake ninja-build llvm-15-dev zlib1g-dev libglew-dev \
-          flex bison libfl-dev libboost-thread-dev libboost-filesystem-dev nvidia-cuda-toolkit-gcc libzstd-dev
-    - name: Install gpuocelot dependencies (MacOS)
-      if: inputs.cuda == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        brew update
-        brew install cmake ninja llvm@15 zlib glew flex bison boost zstd ncurses
-    - name: Cache gpuocelot
-      if: inputs.cuda == 'true'
-      id: cache-build
-      uses: actions/cache@v4
-      env:
-        cache-name: cache-gpuocelot-build
-      with:
-        path: ${{ github.workspace }}/gpuocelot/ocelot
-        key: ${{ runner.os }}-gpuocelot-b16039dc940dc6bc4ea0a98380495769ff35ed99-rebuild-0
-    - name: Clone/compile gpuocelot
-      if: inputs.cuda == 'true' && steps.cache-build.outputs.cache-hit != 'true'
-      shell: bash
-      run: |
-        git clone --recurse-submodules https://github.com/gpuocelot/gpuocelot.git ${{ github.workspace }}/gpuocelot
-        cd ${{ github.workspace }}/gpuocelot/ocelot
-        git checkout b16039dc940dc6bc4ea0a98380495769ff35ed99
-        mkdir build
-        cd build
-        cmake .. -Wno-dev -G Ninja -DOCELOT_BUILD_TOOLS=OFF -DCMAKE_BUILD_ALWAYS=0 -DBUILD_TESTS_CUDA=OFF
-        ninja
-    - name: Install gpuocelot
-      if: inputs.cuda == 'true'
-      shell: bash
-      run: |
-        cd ${{ github.workspace }}/gpuocelot/ocelot/build
-        sudo cp libgpuocelot.${{ runner.os == 'macOS' && 'dylib' || 'so' }} /usr/${{ runner.os == 'macOS' && 'local/' || ''}}lib/
-
-    # **** WebGPU ****
-
-    - name: Install WebGPU dawn (Linux)
-      if: inputs.webgpu == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        sudo curl -L https://github.com/wpmed92/pydawn/releases/download/v0.1.6/libwebgpu_dawn.so -o /usr/local/lib/libwebgpu_dawn.so
-    - name: Install dependencies for software-based vulkan
-      if: inputs.webgpu == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        sudo apt update -y || true
-        sudo apt install -y libegl1-mesa libgl1-mesa-dri libxcb-xfixes0-dev mesa-vulkan-drivers
-
-    - name: Install WebGPU dawn (macOS)
-      if: inputs.webgpu == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        brew tap wpmed92/dawn
-        brew install dawn
-
-    # **** LLVM ****
-
-    - name: Install LLVM (Linux)
-      if: inputs.llvm == 'true' && runner.os == 'Linux'
-      shell: bash
-      run: |
-        echo 'Acquire::http::Pipeline-Depth "5";' | sudo tee -a /etc/apt/apt.conf.d/99parallel
-        wget -qO- https://apt.llvm.org/llvm-snapshot.gpg.key | sudo tee /etc/apt/trusted.gpg.d/apt.llvm.org.asc
-        echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-19 main" | sudo tee /etc/apt/sources.list.d/llvm.list
-        sudo apt update -y || true
-        sudo apt install -y --no-install-recommends libllvm19 clang-19 lld-19
-
-    - name: Install LLVM (macOS)
-      if: inputs.llvm == 'true' && runner.os == 'macOS'
-      shell: bash
-      run: |
-        brew install llvm

.github/workflows/benchmark.yml

@@ -2,7 +2,7 @@ name: Benchmarks
 env:
   # TODO: this rescheduling makes gpt2, mixtral and llama unjitted slower
   # TODO: very slow for llama 70B and resnet training 6 GPU
-  CAPTURE_PROCESS_REPLAY: "1"
+  RUN_PROCESS_REPLAY: "1"
   ASSERT_PROCESS_REPLAY: "0"
   PYTHONPATH: .
   GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@@ -62,10 +62,6 @@ jobs:
       run: BIG=2 MPS=1 python3.11 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test tensor cores
       run: METAL=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
-    - name: 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
-        DEBUG=2 LLVM=1 AMX=1 python3.11 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
     - name: Run Tensor Core GEMM (float)
       run: DEBUG=2 python3.11 extra/gemm/simple_matmul.py | tee matmul.txt
     - name: Run Tensor Core GEMM (half)
@@ -84,8 +80,8 @@ jobs:
       run: |
         python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize int8 | tee llama_int8.txt
         python3.11 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing --quantize nf4 | tee llama_nf4.txt
-    #- name: Run LLaMA 7B on 4 (virtual) GPUs
-    #  run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    - name: Run LLaMA 7B on 4 (virtual) GPUs
+      run: python3.11 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
     - name: Run GPT2
       run: |
         JIT=0 python3.11 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
@@ -93,7 +89,7 @@ jobs:
     - name: Run GPT2 w HALF
       run: HALF=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
     - name: Run GPT2 w HALF/BEAM
-      run: HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3.11 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
     - name: Train MNIST
       run: time PYTHONPATH=. TARGET_EVAL_ACC_PCT=96.0 python3.11 examples/beautiful_mnist.py | tee beautiful_mnist.txt
     - name: Run 10 CIFAR training steps
@@ -138,7 +134,7 @@ jobs:
   testnvidiabenchmark:
     name: tinybox green Benchmark
     runs-on: [self-hosted, Linux, tinyboxgreen]
-    timeout-minutes: 30
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -o pipefail {0}
@@ -165,22 +161,19 @@ jobs:
     - name: reset process replay
       run: test/external/process_replay/reset.py
     - name: Run model inference benchmark
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
+      run: NV=1 RUN_PROCESS_REPLAY=0 NOCLANG=1 python3 test/external/external_model_benchmark.py
     - name: Test speed vs torch
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 HALF=1 BIG=2 TORCHCUDA=1 python3 test/test_speed_v_torch.py | tee torch_speed.txt
     - name: Test speed vs theoretical
       run: NV=1 IGNORE_BEAM_CACHE=1 BEAM_DEBUG=1 DEBUG=1 python -m pytest -rA test/external/speed_v_theoretical.py --durations=20
-    - name: Test benchmark allreduce
-      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
-        PTX=1 ALLOW_TF32=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
+        PTX=1 NV=1 python3 test/test_linearizer.py TestLinearizer.test_tensor_cores TestLinearizer.test_tensor_cores_padded
     - name: Run Tensor Core GEMM (CUDA)
       run: |
         CUDA=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul.txt
         CUDA=1 BFLOAT16=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_bfloat16.txt
-        CUDA=1 ALLOW_TF32=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_tf32.txt
     - name: Run Tensor Core GEMM (PTX)
       run: NV=1 PTX=1 HALF=1 DEBUG=2 python3 extra/gemm/simple_matmul.py | tee matmul_ptx.txt
     - name: Run Tensor Core GEMM (NV)
@@ -192,7 +185,7 @@ jobs:
     - name: Run Stable Diffusion
       run: NV=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run SDXL
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/sdxl.py --seed 0 --noshow --timing | tee sdxl.txt
     - name: Run LLaMA
       run: |
         NV=1 JIT=0 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_unjitted.txt
@@ -200,19 +193,19 @@ jobs:
     - name: Run LLaMA with BEAM
       run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
     # - name: Run LLaMA 7B on 4 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
     # - name: Run LLaMA 7B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
     - name: Run LLaMA-3 8B BEAM
-      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
-    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
-    # - name: Run LLaMA-2 70B
-    #   run: NV=1 CAPTURE_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+      run: NV=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+    - name: Run LLaMA-3 8B on 4 GPUs
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+    - name: Run LLaMA-3 8B on 6 GPUs
+      run: NV=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    - name: Run LLaMA-2 70B
+      run: NV=1 RUN_PROCESS_REPLAY=0 MAX_CONTEXT=256 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
     - name: Run Mixtral 8x7B
-      run: time NV=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
+      run: time NV=1 RUN_PROCESS_REPLAY=0 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
     - name: Run GPT2
       run: |
         NV=1 JIT=0 python3 examples/gpt2.py --prompt "Hello." --count 10 --temperature 0 --timing | tee gpt2_unjitted.txt
@@ -220,7 +213,7 @@ jobs:
     - name: Run GPT2 w HALF
       run: NV=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
     - name: Run GPT2 w HALF/BEAM
-      run: NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: NV=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (NVIDIA)
@@ -229,7 +222,6 @@ jobs:
           torch_speed.txt
           matmul.txt
           matmul_bfloat16.txt
-          matmul_tf32.txt
           matmul_ptx.txt
           matmul_nv.txt
           sd.txt
@@ -290,20 +282,17 @@ jobs:
     - name: Run 10 CIFAR training steps w BF16
       run: NV=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     - name: Run 10 CIFAR training steps w winograd
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
     - name: Run full CIFAR training w 1 GPU
       run: time NV=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
     - name: Run full CIFAR training steps w 6 GPUS
-      run: time CAPTURE_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
+      run: time RUN_PROCESS_REPLAY=0 NV=1 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_six_gpu.txt
     - name: Run MLPerf resnet eval on training data
       run: time NV=1 MODEL=resnet python3 examples/mlperf/model_eval.py
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       run: NV=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
     - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: NV=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: NV=1 RUN_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (NVIDIA Training)
@@ -314,10 +303,9 @@ jobs:
           train_cifar_bf16.txt
           train_cifar_wino.txt
           train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
           train_resnet.txt
           train_resnet_one_gpu.txt
-          train_bert.txt
+          train_cifar_six_gpu.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -332,8 +320,6 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
-    - name: Insert amdgpu
-      run: sudo modprobe amdgpu
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -379,12 +365,6 @@ jobs:
     # TODO: AMD compiler bug causes this to fail
     #- name: Fuzz Padded Tensor Core GEMM
     #  run: HSA=1 M_START=12 M_STOP=20 M_STEP=1 N_START=12 N_STOP=20 N_STEP=1 K_START=28 K_STOP=36 K_STEP=1 HALF=1 TC_OPT=2 DEBUG=2 python3 ./extra/gemm/fuzz_matmul.py
-    - name: Remove amdgpu
-      run: sleep 5 && sudo rmmod amdgpu # sleep a bit to let the driver unload the prev pid.
-    - name: Test AM cold start time
-      run: time AMD=1 AM_RESET=1 python3 test/test_tiny.py TestTiny.test_plus
-    - name: Test AM warm start time
-      run: time AMD=1 python3 test/test_tiny.py TestTiny.test_plus
     - name: Run Stable Diffusion
       run: AMD=1 python3 examples/stable_diffusion.py --fp16 --seed 0 --noshow --timing | tee sd.txt
     - name: Run SDXL
@@ -396,19 +376,17 @@ jobs:
     - name: Run LLaMA 7B with BEAM
       run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama.py --gen 1 --prompt "Hello." --count 10 --temperature 0 --timing | tee llama_beam.txt
     # - name: Run LLaMA 7B on 4 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
+    #   run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 4 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_four_gpu.txt
     # - name: Run LLaMA 7B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
+    #   run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 1 --size 7B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_six_gpu.txt
     - name: Run LLaMA-3 8B BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --size 8B --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
-    - name: Run LLaMA-3 8B on 4 GPUs with BEAM
-      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
-    # - name: Run LLaMA-3 8B on 6 GPUs
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama3.py --size 8B --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
-    - name: Restore amdgpu
-      run: sudo modprobe amdgpu
-    # - name: Run LLaMA-2 70B
-    #   run: AMD=1 CAPTURE_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
+      run: AMD=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/llama3.py --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_beam.txt
+    - name: Run LLaMA-3 8B on 4 GPUs
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 4 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_four_gpu.txt
+    - name: Run LLaMA-3 8B on 6 GPUs
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama3.py --shard 6 --model weights/LLaMA-3/8B-SF-DPO/ --benchmark --temperature 0 | tee llama3_six_gpu.txt
+    - name: Run LLaMA-2 70B
+      run: AMD=1 RUN_PROCESS_REPLAY=0 python3 examples/llama.py --gen 2 --size 70B --shard 6 --prompt "Hello." --count 10 --temperature 0  --timing | tee llama_2_70B.txt
     - name: Run Mixtral 8x7B
       run: time AMD=1 python3 examples/mixtral.py --temperature 0 --count 10 --timing | tee mixtral.txt
     - name: Run GPT2
@@ -418,7 +396,7 @@ jobs:
     - name: Run GPT2 w HALF
       run: AMD=1 HALF=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half.txt
     - name: Run GPT2 w HALF/BEAM
-      run: AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
+      run: AMD=1 HALF=1 JITBEAM=2 IGNORE_BEAM_CACHE=1 CAST_BEFORE_VIEW=0 python3 examples/gpt2.py --count 10 --temperature 0 --timing | tee gpt2_half_beam.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD)
@@ -447,7 +425,7 @@ jobs:
   testmoreamdbenchmark:
     name: tinybox red Training Benchmark
     runs-on: [self-hosted, Linux, tinybox]
-    timeout-minutes: 30
+    timeout-minutes: 20
     defaults:
       run:
         shell: bash -o pipefail {0}
@@ -455,8 +433,6 @@ jobs:
     steps:
     - name: Checkout Code
       uses: actions/checkout@v4
-    - name: Remove amdgpu
-      run: sudo rmmod amdgpu || true
     - name: Symlink models and datasets
       run: |
         mkdir -p weights
@@ -474,6 +450,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
+    - name: setup perflevel
+      run: |
+        examples/mlperf/training_submission_v4.1/tinycorp/benchmarks/bert/implementations/tinybox_red/setup.sh
+        rocm-smi
     - name: Train MNIST
       run: time PYTHONPATH=. AMD=1 TARGET_EVAL_ACC_PCT=96.0 python3 examples/beautiful_mnist.py | tee beautiful_mnist.txt
     - name: Run 10 CIFAR training steps
@@ -483,7 +463,7 @@ jobs:
     - name: Run 10 CIFAR training steps w BF16
       run: AMD=1 STEPS=10 DEFAULT_FLOAT=BFLOAT16 python3 examples/hlb_cifar10.py | tee train_cifar_bf16.txt
     - name: Run 10 CIFAR training steps w winograd
-      run: AMD=1 WINO=1 STEPS=10 DEFAULT_FLOAT=HALF python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
+      run: AMD=1 WINO=1 STEPS=10 python3 examples/hlb_cifar10.py | tee train_cifar_wino.txt
     - name: Run full CIFAR training w 1 GPU
       run: time AMD=1 DEFAULT_FLOAT=HALF LATEWINO=1 STEPS=1000 TARGET_EVAL_ACC_PCT=93.2 python3 examples/hlb_cifar10.py | tee train_cifar_one_gpu.txt
     - name: Run full CIFAR training steps w 6 GPUS
@@ -493,10 +473,7 @@ jobs:
     - name: Run 10 MLPerf ResNet50 training steps (1 gpu)
       run: AMD=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet_one_gpu.txt
     - name: Run 10 MLPerf ResNet50 training steps (6 gpu)
-      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
-    - name: Run 10 MLPerf Bert training steps (6 gpu)
-      # TODO: remove BERT_LAYERS once scheduler is fast
-      run: AMD=1 CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=66 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py | tee train_bert.txt
+      run: AMD=1 RUN_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py | tee train_resnet.txt
     - uses: actions/upload-artifact@v4
       with:
         name: Speed (AMD Training)
@@ -507,10 +484,9 @@ jobs:
           train_cifar_bf16.txt
           train_cifar_wino.txt
           train_cifar_one_gpu.txt
-          train_cifar_six_gpu.txt
           train_resnet.txt
           train_resnet_one_gpu.txt
-          train_bert.txt
+          train_cifar_six_gpu.txt
     - name: Run process replay tests
       run: cp test/external/process_replay/process_replay.py ./process_replay.py && git fetch origin master && git -c advice.detachedHead=false checkout origin/master && PYTHONPATH=. python3 process_replay.py
 
@@ -532,6 +508,10 @@ jobs:
         rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
     - name: reset process replay
       run: test/external/process_replay/reset.py
+    - name: openpilot compile 0.9.4
+      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py | tee openpilot_compile_0_9_4.txt
+    - name: openpilot compile 0.9.7
+      run: PYTHONPATH=. NOLOCALS=1 FLOAT16=1 IMAGE=2 QCOM=1 taskset -c 4-7 python examples/openpilot/compile2.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_compile_0_9_7.txt
     - name: validate openpilot 0.9.7
       run: PYTHONPATH=. FLOAT16=0 IMAGE=2 QCOM=1 taskset -c 4-7 python3 test/external/external_benchmark_openpilot.py https://github.com/commaai/openpilot/raw/v0.9.7/selfdrive/modeld/models/supercombo.onnx | tee openpilot_image_0_9_7.txt
     - name: benchmark openpilot 0.9.4

.gitignore

@@ -10,7 +10,6 @@ notebooks
 *.so
 *.txt
 build
-!examples/tinychat/assets/cdn.jsdelivr.net/npm/[email protected]/build/
 /dist
 *.egg-info
 /env
@@ -34,8 +33,6 @@ extra/datasets/open-images-v6-mlperf
 extra/datasets/kits/
 extra/datasets/COCO/
 extra/datasets/audio*
-extra/huggingface_onnx/models/*
-extra/huggingface_onnx/*.yaml
 extra/weights
 venv
 examples/**/net.*[js,json]
@@ -58,5 +55,4 @@ weights
 comgr_*
 *.pkl
 site/
-profile_stats
-*.log
+master_schedule.py

.pylintrc

@@ -7,7 +7,7 @@ extension-pkg-whitelist=scipy,cereal.messaging.messaging_pyx,PyQt5,av
 
 # Add files or directories to the blacklist. They should be base names, not
 # paths.
-ignore=CVS,autogen,msm_kgsl.py,runtime
+ignore=CVS,autogen,msm_kgsl.py
 
 # Add files or directories matching the regex patterns to the blacklist. The
 # regex matches against base names, not paths.

README.md

@@ -81,7 +81,7 @@ See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full vers
 tinygrad already supports numerous accelerators, including:
 
 - [x] [GPU (OpenCL)](tinygrad/runtime/ops_gpu.py)
-- [x] [CPU (C Code)](tinygrad/runtime/ops_cpu.py)
+- [x] [CLANG (C Code)](tinygrad/runtime/ops_clang.py)
 - [x] [LLVM](tinygrad/runtime/ops_llvm.py)
 - [x] [METAL](tinygrad/runtime/ops_metal.py)
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
@@ -151,7 +151,7 @@ We'll start with what will get your PR closed with a pointer to this section:
 
 - No code golf! While low line count is a guiding light of this project, anything that remotely looks like code golf will be closed. The true goal is reducing complexity and increasing readability, and deleting `\n`s does nothing to help with that.
 - All docs and whitespace changes will be closed unless you are a well-known contributor. The people writing the docs should be those who know the codebase the absolute best. People who have not demonstrated that shouldn't be messing with docs. Whitespace changes are both useless *and* carry a risk of introducing bugs.
-- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainability and readability.
+- Anything you claim is a "speedup" must be benchmarked. In general, the goal is simplicity, so even if your PR makes things marginally faster, you have to consider the tradeoff with maintainablity and readablity.
 - In general, the code outside the core `tinygrad/` folder is not well tested, so unless the current code there is broken, you shouldn't be changing it.
 - If your PR looks "complex", is a big diff, or adds lots of lines, it won't be reviewed or merged. Consider breaking it up into smaller PRs that are individually clear wins. A common pattern I see is prerequisite refactors before adding new functionality. If you can (cleanly) refactor to the point that the feature is a 3 line change, this is great, and something easy for us to review.
 

autogen_stubs.sh

@@ -69,7 +69,7 @@ generate_comgr() {
   --clang-args="-D__HIP_PLATFORM_AMD__ -I/opt/rocm/include -x c++" -o $BASE/comgr.py -l /opt/rocm/lib/libamd_comgr.so
   fixup $BASE/comgr.py
   sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/comgr.py
-  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'" "'/usr/local/lib/libamd_comgr.dylib'" "'/opt/homebrew/lib/libamd_comgr.dylib'"
+  patch_dlopen $BASE/comgr.py amd_comgr "'/opt/rocm/lib/libamd_comgr.so'" "os.getenv('ROCM_PATH', '')+'/lib/libamd_comgr.so'"
   sed -i "s\ctypes.CDLL('/opt/rocm/lib/libamd_comgr.so')\_try_dlopen_amd_comgr()\g" $BASE/comgr.py
   python3 -c "import tinygrad.runtime.autogen.comgr"
 }
@@ -79,10 +79,6 @@ generate_kfd() {
 
   fixup $BASE/kfd.py
   sed -i "s\import ctypes\import ctypes, os\g" $BASE/kfd.py
-  sed -i "s\import fcntl, functools\import functools" $BASE/kfd.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/kfd.py
-  sed -i "s\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd, **kwargs):\def _do_ioctl(__idir, __base, __nr, __user_struct, __fd:HWInterface, **kwargs):\g" $BASE/kfd.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kfd.py
   python3 -c "import tinygrad.runtime.autogen.kfd"
 }
 
@@ -171,7 +167,6 @@ generate_amd() {
     extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
     extra/hip_gpu_driver/gc_11_0_0_offset.h \
     extra/hip_gpu_driver/gc_10_3_0_offset.h \
-    extra/hip_gpu_driver/sienna_cichlid_ip_offset.h \
     --clang-args="-I/opt/rocm/include -x c++" \
     -o $BASE/amd_gpu.py
 
@@ -212,10 +207,8 @@ generate_libc() {
   clang2py -k cdefstum \
     $(dpkg -L libc6-dev | grep sys/mman.h) \
     $(dpkg -L libc6-dev | grep sys/syscall.h) \
-    /usr/include/string.h \
     /usr/include/elf.h \
     /usr/include/unistd.h \
-    /usr/include/asm-generic/mman-common.h \
     -o $BASE/libc.py
 
   sed -i "s\import ctypes\import ctypes, ctypes.util, os\g" $BASE/libc.py
@@ -225,30 +218,11 @@ generate_libc() {
   fixup $BASE/libc.py
 }
 
-generate_llvm() {
-  INC="$(llvm-config-14 --includedir)"
-  clang2py -k cdefstum \
-    $(find "$INC/llvm-c/" -type f -name '*.h' | sort) \
-    "$INC/llvm/Config/Targets.def" \
-    "$INC/llvm/Config/AsmPrinters.def" \
-    "$INC/llvm/Config/AsmParsers.def" \
-    "$INC/llvm/Config/Disassemblers.def" \
-    --clang-args="$(llvm-config-14 --cflags)" \
-    -o "$BASE/llvm.py"
-
-  sed -i "s\import ctypes\import ctypes, tinygrad.runtime.support.llvm as llvm_support\g" "$BASE/llvm.py"
-  sed -i "s\FIXME_STUB\llvm\g" "$BASE/llvm.py"
-  sed -i "s\FunctionFactoryStub()\ctypes.CDLL(llvm_support.LLVM_PATH)\g" "$BASE/llvm.py"
-
-  fixup "$BASE/llvm.py"
-}
-
 generate_kgsl() {
   clang2py extra/qcom_gpu_driver/msm_kgsl.h -o $BASE/kgsl.py -k cdefstum
   fixup $BASE/kgsl.py
   sed -i "s\import ctypes\import ctypes, os\g" $BASE/kgsl.py
   sed -nE 's/#define ([A-Za-z0-9_]+)_SHIFT\s*[^\S\r\n]*[0-9]*$/def \1(val): return (val << \1_SHIFT) \& \1_MASK/p' extra/qcom_gpu_driver/msm_kgsl.h >> $BASE/kgsl.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\__fd.ioctl((__idir<<30)\g" $BASE/kgsl.py
   python3 -c "import tinygrad.runtime.autogen.kgsl"
 }
 
@@ -273,248 +247,6 @@ generate_qcom() {
   python3 -c "import tinygrad.runtime.autogen.qcom_dsp"
 }
 
-generate_pci() {
-  clang2py -k cdefstum \
-    /usr/include/linux/pci_regs.h \
-    -o $BASE/pci.py
-  fixup $BASE/pci.py
-}
-
-generate_vfio() {
-  clang2py -k cdefstum \
-    /usr/include/linux/vfio.h \
-    -o $BASE/vfio.py
-  fixup $BASE/vfio.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/vfio.py
-  sed -i "s\import fcntl, functools\import functools" $BASE/vfio.py
-  sed -i "s\import ctypes,os\a from tinygrad.runtime.support import HWInterface\g" $BASE/vfio.py
-  sed -i "s\fcntl.ioctl(__fd, (__idir<<30)\return __fd.ioctl((__idir<<30)\g" $BASE/vfio.py
-}
-
-generate_am() {
-  clang2py -k cdefstum \
-    extra/amdpci/headers/v11_structs.h \
-    extra/amdpci/headers/v12_structs.h \
-    extra/amdpci/headers/amdgpu_vm.h \
-    extra/amdpci/headers/discovery.h \
-    extra/amdpci/headers/amdgpu_ucode.h \
-    extra/amdpci/headers/psp_gfx_if.h \
-    extra/amdpci/headers/amdgpu_psp.h \
-    extra/amdpci/headers/amdgpu_irq.h \
-    extra/amdpci/headers/amdgpu_doorbell.h \
-    extra/amdpci/headers/soc15_ih_clientid.h \
-    --clang-args="-include stdint.h" \
-    -o $BASE/am/am.py
-  fixup $BASE/am/am.py
-  sed -i "s\(int64_t)\ \g" $BASE/am/am.py
-  sed -i "s\AMDGPU_PTE_MTYPE_VG10(2)\AMDGPU_PTE_MTYPE_VG10(0, 2)\g" $BASE/am/am.py # incorrect parsing (TODO: remove when clang2py is gone).
-
-  clang2py -k cdefstum \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/soc15d.h \
-    -o $BASE/am/pm4_soc15.py
-  fixup $BASE/am/pm4_soc15.py
-
-  clang2py -k cdefstum \
-    extra/hip_gpu_driver/kfd_pm4_headers_ai.h \
-    extra/hip_gpu_driver/nvd.h \
-    -o $BASE/am/pm4_nv.py
-  fixup $BASE/am/pm4_nv.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/vega10_enum.h \
-    -o $BASE/am/vega10.py
-  fixup $BASE/am/vega10.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/navi10_enum.h \
-    -o $BASE/am/navi10.py
-  fixup $BASE/am/navi10.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/soc21_enum.h \
-    -o $BASE/am/soc21.py
-  fixup $BASE/am/soc21.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/soc24_enum.h \
-    -o $BASE/am/soc24.py
-  fixup $BASE/am/soc24.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_13_0_0_offset.h \
-    extra/amdpci/headers/mp_13_0_0_sh_mask.h \
-    -o $BASE/am/mp_13_0_0.py
-  fixup $BASE/am/mp_13_0_0.py
-
-  # 14_0_3 reuses 14_0_2
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_14_0_2_offset.h \
-    extra/amdpci/headers/mp_14_0_2_sh_mask.h \
-    -o $BASE/am/mp_14_0_3.py
-  fixup $BASE/am/mp_14_0_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mp_11_0_offset.h \
-    extra/amdpci/headers/mp_11_0_sh_mask.h \
-    -o $BASE/am/mp_11_0.py
-  fixup $BASE/am/mp_11_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_9_4_3_offset.h \
-    extra/amdpci/headers/gc_9_4_3_sh_mask.h \
-    extra/amdpci/overlay/gc_9_4_3.h \
-    -o $BASE/am/gc_9_4_3.py
-  fixup $BASE/am/gc_9_4_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_10_3_0_offset.h \
-    extra/amdpci/headers/gc_10_3_0_sh_mask.h \
-    -o $BASE/am/gc_10_3_0.py
-  fixup $BASE/am/gc_10_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_11_0_0_offset.h \
-    extra/amdpci/headers/gc_11_0_0_sh_mask.h \
-    -o $BASE/am/gc_11_0_0.py
-  fixup $BASE/am/gc_11_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/gc_12_0_0_offset.h \
-    extra/amdpci/headers/gc_12_0_0_sh_mask.h \
-    -o $BASE/am/gc_12_0_0.py
-  fixup $BASE/am/gc_12_0_0.py
-
-  clang2py -k cdefstum \
-    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/vega10_sdma_pkt_open.h \
-    --clang-args="-I/opt/rocm/include -x c++" \
-    -o $BASE/am/sdma_4_0_0.py
-  fixup $BASE/am/sdma_4_0_0.py
-
-  clang2py -k cdefstum \
-    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/navi10_sdma_pkt_open.h \
-    --clang-args="-I/opt/rocm/include -x c++" \
-    -o $BASE/am/sdma_5_0_0.py
-  fixup $BASE/am/sdma_5_0_0.py
-
-  clang2py -k cdefstum \
-    extra/hip_gpu_driver/sdma_registers.h \
-    extra/hip_gpu_driver/sdma_v6_0_0_pkt_open.h \
-    --clang-args="-I/opt/rocm/include -x c++" \
-    -o $BASE/am/sdma_6_0_0.py
-  fixup $BASE/am/sdma_6_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_0_offset.h \
-    extra/amdpci/headers/mmhub_3_0_0_sh_mask.h \
-    -o $BASE/am/mmhub_3_0_0.py
-  fixup $BASE/am/mmhub_3_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_3_0_2_offset.h \
-    extra/amdpci/headers/mmhub_3_0_2_sh_mask.h \
-    -o $BASE/am/mmhub_3_0_2.py
-  fixup $BASE/am/mmhub_3_0_2.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_2_3_offset.h \
-    extra/amdpci/headers/nbio_2_3_sh_mask.h \
-    -o $BASE/am/nbio_2_3_0.py
-  fixup $BASE/am/nbio_2_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/mmhub_4_1_0_offset.h \
-    extra/amdpci/headers/mmhub_4_1_0_sh_mask.h \
-    -o $BASE/am/mmhub_4_1_0.py
-  fixup $BASE/am/mmhub_4_1_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_4_3_0_offset.h \
-    extra/amdpci/headers/nbio_4_3_0_sh_mask.h \
-    -o $BASE/am/nbio_4_3_0.py
-  fixup $BASE/am/nbio_4_3_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbif_6_3_1_offset.h \
-    extra/amdpci/headers/nbif_6_3_1_sh_mask.h \
-    -o $BASE/am/nbif_6_3_1.py
-  fixup $BASE/am/nbif_6_3_1.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_7_9_0_offset.h \
-    extra/amdpci/headers/nbio_7_9_0_sh_mask.h \
-    -o $BASE/am/nbio_7_9_0.py
-  fixup $BASE/am/nbio_7_9_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/nbio_7_11_0_offset.h \
-    extra/amdpci/headers/nbio_7_11_0_sh_mask.h \
-    -o $BASE/am/nbio_7_11_0.py
-  fixup $BASE/am/nbio_7_11_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/osssys_6_0_0_offset.h \
-    extra/amdpci/headers/osssys_6_0_0_sh_mask.h \
-    -o $BASE/am/osssys_6_0_0.py
-  fixup $BASE/am/osssys_6_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/osssys_7_0_0_offset.h \
-    extra/amdpci/headers/osssys_7_0_0_sh_mask.h \
-    -o $BASE/am/osssys_7_0_0.py
-  fixup $BASE/am/osssys_7_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/smu_v13_0_0_ppsmc.h \
-    extra/amdpci/headers/smu13_driver_if_v13_0_0.h \
-    extra/amdpci/headers/amdgpu_smu.h \
-    -o $BASE/am/smu_v13_0_0.py
-  fixup $BASE/am/smu_v13_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/smu_v14_0_0_pmfw.h \
-    extra/amdpci/headers/smu_v14_0_2_ppsmc.h \
-    extra/amdpci/headers/smu14_driver_if_v14_0_0.h \
-    extra/amdpci/headers/smu14_driver_if_v14_0.h \
-    extra/amdpci/headers/amdgpu_smu.h \
-    --clang-args="-include stdint.h" \
-    -o $BASE/am/smu_v14_0_3.py
-  fixup $BASE/am/smu_v14_0_3.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/hdp_6_0_0_offset.h \
-    extra/amdpci/headers/hdp_6_0_0_sh_mask.h \
-    -o $BASE/am/hdp_6_0_0.py
-  fixup $BASE/am/hdp_6_0_0.py
-
-  clang2py -k cdefstum \
-    extra/amdpci/headers/hdp_7_0_0_offset.h \
-    extra/amdpci/headers/hdp_7_0_0_sh_mask.h \
-    -o $BASE/am/hdp_7_0_0.py
-  fixup $BASE/am/hdp_7_0_0.py
-}
-
-generate_sqtt() {
-  clang2py -k cdefstum \
-    extra/sqtt/sqtt.h \
-    -o $BASE/sqtt.py
-
-  fixup $BASE/sqtt.py
-  sed -i "s\import ctypes\import ctypes, os\g" $BASE/sqtt.py
-  python3 -c "import tinygrad.runtime.autogen.sqtt"
-}
-
-generate_webgpu() {
-  clang2py extra/webgpu/webgpu.h -o $BASE/webgpu.py
-  fixup $BASE/webgpu.py
-  sed -i "s/FIXME_STUB/webgpu/g" "$BASE/webgpu.py"
-  sed -i "s/FunctionFactoryStub()/ctypes.CDLL(webgpu_support.WEBGPU_PATH)/g" "$BASE/webgpu.py"
-  sed -i "s/import ctypes/import ctypes, tinygrad.runtime.support.webgpu as webgpu_support/g" "$BASE/webgpu.py"
-  python3 -c "import tinygrad.runtime.autogen.webgpu"
-}
-
 if [ "$1" == "opencl" ]; then generate_opencl
 elif [ "$1" == "hip" ]; then generate_hip
 elif [ "$1" == "comgr" ]; then generate_comgr
@@ -524,17 +256,11 @@ elif [ "$1" == "hsa" ]; then generate_hsa
 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" == "sqtt" ]; then generate_sqtt
 elif [ "$1" == "qcom" ]; then generate_qcom
 elif [ "$1" == "io_uring" ]; then generate_io_uring
 elif [ "$1" == "libc" ]; then generate_libc
-elif [ "$1" == "llvm" ]; then generate_llvm
 elif [ "$1" == "kgsl" ]; then generate_kgsl
 elif [ "$1" == "adreno" ]; then generate_adreno
-elif [ "$1" == "pci" ]; then generate_pci
-elif [ "$1" == "vfio" ]; then generate_vfio
-elif [ "$1" == "webgpu" ]; then generate_webgpu
-elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc; generate_am; generate_webgpu
+elif [ "$1" == "all" ]; then generate_opencl; generate_hip; generate_comgr; generate_cuda; generate_nvrtc; generate_hsa; generate_kfd; generate_nv; generate_amd; generate_io_uring; generate_libc
 else echo "usage: $0 <type>"
 fi

docs/abstractions2.py

@@ -7,7 +7,7 @@
 
 print("******** first, the runtime ***********")
 
-from tinygrad.runtime.ops_cpu import ClangJITCompiler, MallocAllocator, CPUProgram
+from tinygrad.runtime.ops_clang import ClangProgram, ClangCompiler, MallocAllocator
 
 # allocate some buffers
 out = MallocAllocator.alloc(4)
@@ -19,10 +19,10 @@ MallocAllocator._copyin(a, memoryview(bytearray([2,0,0,0])))
 MallocAllocator._copyin(b, memoryview(bytearray([3,0,0,0])))
 
 # compile a program to a binary
-lib = ClangJITCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
+lib = ClangCompiler().compile("void add(int *out, int *a, int *b) { out[0] = a[0] + b[0]; }")
 
-# create a runtime for the program
-fxn = CPUProgram("add", lib)
+# create a runtime for the program (ctypes.CDLL)
+fxn = ClangProgram("add", lib)
 
 # run the program
 fxn(out, a, b)
@@ -34,7 +34,7 @@ assert val == 5
 
 print("******** second, the Device ***********")
 
-DEVICE = "CPU"   # NOTE: you can change this!
+DEVICE = "CLANG"   # NOTE: you can change this!
 
 import struct
 from tinygrad.dtype import dtypes
@@ -65,7 +65,7 @@ kernel = get_kernel(Device[DEVICE].renderer, s).linearize()
 # compile a program (and print the source)
 fxn = CompiledRunner(kernel.to_program())
 print(fxn.p.src)
-# NOTE: fxn.clprg is the CPUProgram
+# NOTE: fxn.clprg is the ClangProgram
 
 # run the program
 fxn.exec([out, a, b])
@@ -76,23 +76,24 @@ assert out.as_buffer().cast('I')[0] == 5
 
 print("******** third, the LazyBuffer ***********")
 
+from tinygrad.engine.lazy import LazyBuffer
 from tinygrad.engine.realize import run_schedule
-from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.engine.schedule import create_schedule
 
 # allocate some values + load in values
-a = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
-b = UOp.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+a = LazyBuffer.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
+b = LazyBuffer.metaop(Ops.EMPTY, (1,), dtypes.int32, DEVICE)
 a.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 2))))
 b.buffer.allocate().copyin(memoryview(bytearray(struct.pack("I", 3))))
+del a.srcs
+del b.srcs
 
 # describe the computation
 out = a.alu(Ops.ADD, b)
 
 # schedule the computation as a list of kernels
-sched, _, becomes_map = create_schedule_with_vars(out.sink())
-for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CPU
-# NOTE: UOps are no longer mutable, the scheduler gives you a map to lookup which BUFFER the result was written to
-out = becomes_map[out]
+sched = create_schedule([out])
+for si in sched: print(si.ast.op)  # NOTE: the first two convert it to CLANG
 
 # DEBUGGING: print the compute ast
 print(sched[-1].ast)
@@ -102,7 +103,7 @@ print(sched[-1].ast)
 run_schedule(sched)
 
 # check the data out
-assert out.is_realized and out.buffer.as_buffer().cast('I')[0] == 5
+assert out.realized is not None and out.realized.as_buffer().cast('I')[0] == 5
 
 
 print("******** fourth, the Tensor ***********")

docs/abstractions3.py

@@ -26,11 +26,10 @@ l1n, l2n = l1.numpy(), l2.numpy()
 from tinygrad.nn.optim import SGD
 optim = SGD([l1, l2])
 
-Tensor.training = True
 X, Y = X_train[(samples:=Tensor.randint(128, high=X_train.shape[0]))], Y_train[samples]
 optim.zero_grad()
 model(X).sparse_categorical_crossentropy(Y).backward()
-optim.schedule_step()   # this will step the optimizer without running realize
+optim._step()   # this will step the optimizer without running realize
 
 # *****
 # 3. Create a schedule.
@@ -48,7 +47,7 @@ for si in schedule: print(str(si)[:80])
 # 4. Lower a schedule.
 
 from tinygrad.engine.realize import lower_schedule_item, ExecItem
-lowered: List[ExecItem] = [lower_schedule_item(si) for si in tqdm(schedule)]
+lowered: List[ExecItem] = [ExecItem(lower_schedule_item(si).prg, list(si.bufs)) for si in tqdm(schedule)]
 
 # *****
 # 5. Run the schedule

docs/developer/am.md

@@ -1,39 +0,0 @@
-# AM Driver
-
-AM driver is a userspace driver targeting AMD's 7900XTX. You only need tinygrad to send compute tasks to your GPU!
-
-## How to run?
-Make sure that amdgpu module is unloaded and just run tinygrad with `AMD=1`!
-
-Optional requirements:
-
-* System without IOMMU for P2P / SDMA support
-* vfio-pci module for IRQ handling
-
-## Environment Variables
-
-| Variable | Possible Value(s) | Description |
-|----------|------------------|-------------|
-| AM_RESET | [1] | Performs a full GPU reset (reloading all firmware and IP blocks) |
-| AM_DEBUG | [0-4] | Sets the level of additional debugging information |
-
-## AM Driver Details
-
-### Compute & SDMA Queues
-
-AM binds compute queues directly to MEC (bypassing MES). Tinygrad uses only one compute queue, which is bound at `pipe=0 queue=0`. Similarly, the single SDMA queue is bound at `engine=0 queue=0`.
-
-### Boot
-
-The GPU being passed can be in one of several states:
-1. Not initialized
-2. Initialized by amdgpu
-3. Initialized by AM
-
-The first and second states require a full GPU setup since their states are unknown. The second state also requires a mode1 reset to reinitialize all components.
-
-The third state can be set up partially to optimize boot time. In this case, only the GFX and SDMA IPs need to be initialized. To enable this, AM uses a separate boot memory that is guaranteed not to be overwritten. This physical memory is utilized for all blocks that are initialized only during the initial AM boot. To determine if the GPU is in the third state, AM uses `regSCRATCH_REG7` as a flag.
-
-### VM Management
-
-Each AM device sets up only a single `VMID=0` and one page directory. The page directory used is 3-level and thus supports up to 512GB of virtual addresses. All AM devices are located in one virtual address space.

docs/developer/developer.md

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

docs/developer/function.md

@@ -0,0 +1,33 @@
+::: tinygrad.function
+    options:
+        members: [
+            "Contiguous",
+            "ContiguousBackward",
+            "Cast",
+            "Neg",
+            "Reciprocal",
+            "Sin",
+            "Relu",
+            "Log",
+            "Exp",
+            "Sqrt",
+            "Sigmoid",
+            "Sign",
+            "Less",
+            "Eq",
+            "Xor",
+            "Add",
+            "Sub",
+            "Mul",
+            "Div",
+            "Where",
+            "Sum",
+            "Max",
+            "Expand",
+            "Reshape",
+            "Permute",
+            "Pad",
+            "Shrink",
+            "Flip",
+        ]
+        show_source: false

docs/developer/hcq.md

@@ -115,8 +115,9 @@ HCQ-compatible devices use a global timeline signal for synchronizing all operat
 ```python
 HWQueue().wait(your_device.timeline_signal, your_device.timeline_value - 1) \
          .exec(...)
-         .signal(your_device.timeline_signal, your_device.next_timeline()) \
+         .signal(your_device.timeline_signal, your_device.timeline_value) \
          .submit(your_device)
+your_device.timeline_value += 1
 
 # Optionally wait for execution
 your_device.timeline_signal.wait(your_device.timeline_value - 1)

docs/developer/runtime.md

@@ -36,9 +36,9 @@ The `Allocator` class is responsible for managing memory on the device. There is
 
 ### Program
 
-The `Program` class is created for each loaded program. It is responsible for executing the program on the device. As an example, here is a `CPUProgram` implementation which loads program and runs it.
+The `Program` class is created for each loaded program. It is responsible for compiling and executing the program on the device. As an example, here is a `ClangProgram` implementation which loads program and runs it.
 
-::: tinygrad.runtime.ops_cpu.CPUProgram
+::: tinygrad.runtime.ops_clang.ClangProgram
     options:
         members: true
 

docs/developer/speed.md

@@ -1,71 +0,0 @@
-# speed in tinygrad
-
-## Overview
-
-Speed refers to many different things. To break it down to four, there's:
-
-- Compile Speed (Python)
-- Execution Speed (driver)
-- Model Speed (scheduler)
-- Kernel Speed (codegen)
-
-## Compile Speed (Python)
-
-This is how long the first run of your model takes. It's limited largely by the runtime of the Python doing UOp rewrites. Currently it's a bit slow, but on par with torch.compile. It gets even slower if you are using BEAM, since that's compiling many variants of each kernel.
-
-This will be improved by writing faster graph_rewrite, doing less graph_rewrite, and better parallelization.
-
-## Execution Speed (driver)
-
-After your model is compiled, you are often using the `TinyJIT`. tinygrad has the best execution speed of any framework because it usually bypasses the GPU driver and prebuilds the command queue. It's tons faster than normal CUDA, and often even faster than CUDA Graph.
-
-There's very little to improve here, as this is almost never the bottleneck.
-
-## Model Speed (scheduler)
-
-The scheduler determines how operations are grouped into kernels and which Tensors are written to memory. This is currently a big bottleneck of training speed.
-
-The decisions are often not obvious. For example, when is it worth recomputing an arithmetic operation instead of storing and loading from memory? Example:
-
-```python
-from tinygrad import Tensor
-a = Tensor.rand(100)
-b = Tensor.rand(100)
-c = Tensor.rand(100)
-d = Tensor.rand(100)
-out1 = a+b+c
-out2 = a+b+d
-Tensor.realize(out1, out2)
-```
-
-The real answer is obvious, compute both `out1` and `out2` in the same kernel. But you can't always do that. If you can't, should `a+b` first be saved to a subbuffer? Or should both the `out1` and `out2` kernels recompute `a+b`?
-
-In this case: with recompute (6 reads + 2 writes), no recompute (6 reads + 3 writes), so we should probably recompute. However, once you add movement ops and casts this is even harder to figure out. tinygrad doesn't yet have a systematic way to do it.
-
-## Kernel Speed (codegen)
-
-Given that you have decided how the model ops will be grouped and what will be written to memory, kernel speed determines how fast that operation is done. This is what BEAM changes, it searches over a set of equivalent kernels which all perform the same operation and finds the one which performs the task the fastest.
-
-In `kernel.py` we have a set of `OptOps`, these control the parameters of the speed optimizations applied to the kernel.
-
-### Memory
-
-The main bottleneck in most kernels is accessing memory. In a freshman algorithms class, you'll learn about cache aware matrix multiplication, and this is all forms of that. While the same math is run, the order in which you run it can have large impacts on the speed depending on if the data you are loading. OptOps will change this order.
-
-Memory, even cache, is often much slower than accessing the register file. The amount of times data is used in math is called the "arithmetic intensity". For operations like BS=1 GEMV, the arithmetic intensity is 1, but for GEMMs and convs it can be much higher. OptOps like UPCAST and UNROLL can increase this, but be careful of making them too large, as if there's too much register pressure on the GPU the warp scheduler may not be able to fit many warps, or even worse, it could be spilling to local memory.
-
-4090s have 1 TB/s of ram bandwidth and ~160 TFLOPS of compute, so you need to use each loaded value ~100 times. The L1 cache has around 40 TB/s of bandwidth, so in order to get full compute utilization you need to use each value ~4 times.
-
-A lot of work can still be done here. For example, we never copy the inputs to on chip SRAM, but this is often quite helpful for kernel speed. Also, we aren't doing a good job with L2 cache awareness (the locals handle L1 quite well)
-
-### Tensor Cores
-
-Many accelerators have Tensor Cores / MAC arrays / systolic arrays. The main value of these is that, since they are 2-D, they create an n^2 ratio between the compute and the input data.
-
-GPUs use Tensor Cores instead of MAC arrays to fit better in the GPU warp paradigm. This is because the output of Tensor Cores is O(n) wrt the input, while the output of MAC arrays like the AMX is O(n^2)
-
-We have a simple framework in tinygrad for adding these ALU blocks and achieving good performance from them.
-
-### Indexing
-
-Indexing determines the address of the memory we need to load. GPUs often have less integer math resources than floating point math, so this can sometimes be the bottleneck. We have a symbolic math engine in our rewrite rules to simplifiy indexing before it's emitted to the kernel. Newer NVIDIA GPUs have a "Tensor Memory Accelerator" to assist with fast indexing, however, this is not supported in tinygrad yet.

docs/env_vars.md

@@ -31,22 +31,20 @@ These control the behavior of core tinygrad even when used as a library.
 Variable | Possible Value(s) | Description
 ---|---|---
 DEBUG               | [1-6]      | enable debugging output, with 4 you get operations, timings, speed, generated code and more
-GPU                 | [1]        | enable the GPU (OpenCL) backend
+GPU                 | [1]        | enable the GPU backend
 CUDA                | [1]        | enable CUDA backend
 AMD                 | [1]        | enable AMD backend
 NV                  | [1]        | enable NV backend
 METAL               | [1]        | enable Metal backend (for Mac M1 and after)
 METAL_XCODE         | [1]        | enable Metal using macOS Xcode SDK
-CPU                 | [1]        | enable CPU (Clang) backend
+CLANG               | [1]        | enable Clang backend
 LLVM                | [1]        | enable LLVM backend
 BEAM                | [#]        | number of beams in kernel beam search
 DEFAULT_FLOAT       | [HALF, ...]| specify the default float dtype (FLOAT32, HALF, BFLOAT16, FLOAT64, ...), default to FLOAT32
 IMAGE               | [1-2]      | enable 2d specific optimizations
 FLOAT16             | [1]        | use float16 for images instead of float32
 PTX                 | [1]        | enable the specialized [PTX](https://docs.nvidia.com/cuda/parallel-thread-execution/) assembler for Nvidia GPUs. If not set, defaults to generic CUDA codegen backend.
-PROFILE             | [1]        | enable profiling. This feature is supported in NV, AMD, QCOM and METAL backends.
+PROFILE             | [1]        | enable output of [perfetto](https://ui.perfetto.dev/) compatible profile. This feature is supported in NV and AMD backends.
 VISIBLE_DEVICES     | [list[int]]| restricts the NV/AMD devices that are available. The format is a comma-separated list of identifiers (indexing starts with 0).
 JIT                 | [0-2]      | 0=disabled, 1=[jit enabled](quickstart.md#jit) (default), 2=jit enabled, but graphs are disabled
-VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)
-ALLOW_TF32          | [1]        | enable TensorFloat-32 tensor cores on Ampere or newer GPUs.
-WEBGPU_BACKEND      | [WGPUBackendType_Metal, ...]          | Force select a backend for WebGPU (Metal, DirectX, OpenGL, Vulkan...)
+VIZ                 | [1]        | 0=disabled, 1=[viz enabled](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/viz)

docs/mnist.md

@@ -17,7 +17,7 @@ from tinygrad import Device
 print(Device.DEFAULT)
 ```
 
-You will see `CUDA` here on a GPU instance, or `CPU` here on a CPU instance.
+You will see `CUDA` here on a GPU instance, or `CLANG` here on a CPU instance.
 
 ## A simple model
 

docs/nn.md

@@ -29,12 +29,4 @@
 ::: tinygrad.nn.state.get_state_dict
 ::: tinygrad.nn.state.get_parameters
 ::: tinygrad.nn.state.load_state_dict
-::: tinygrad.nn.state.tar_extract
-    options:
-        show_signature: false
-        separate_signature: false
 ::: tinygrad.nn.state.torch_load
-    options:
-        show_signature: false
-        separate_signature: false
-::: tinygrad.nn.state.gguf_load

docs/quickstart.md

@@ -110,7 +110,7 @@ class TinyNet:
 
   def __call__(self, x):
     x = self.l1(x)
-    x = x.leaky_relu()
+    x = x.leakyrelu()
     x = self.l2(x)
     return x
 
@@ -118,7 +118,7 @@ net = TinyNet()
 ```
 
 We can see that the forward pass of our neural network is just the sequence of operations performed on the input tensor `x`.
-We can also see that functional operations like `leaky_relu` are not defined as classes and instead are just methods we can just call.
+We can also see that functional operations like `leakyrelu` are not defined as classes and instead are just methods we can just call.
 Finally, we just initialize an instance of our neural network, and we are ready to start training it.
 
 ## Training

docs/runtime.md

@@ -1,6 +1,6 @@
 # Runtimes
 
-tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CPU=1`).
+tinygrad supports various runtimes, enabling your code to scale across a wide range of devices. The default runtime can be automatically selected based on the available hardware, or you can force a specific runtime to be default using environment variables (e.g., `CLANG=1`).
 
 | Runtime | Description | Requirements |
 |---------|-------------|--------------|
@@ -10,57 +10,5 @@ tinygrad supports various runtimes, enabling your code to scale across a wide ra
 | [METAL](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_metal.py) | Utilizes Metal for acceleration on Apple devices | M1+ Macs; Metal 3.0+ for `bfloat` support |
 | [CUDA](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cuda.py) | Utilizes CUDA for acceleration on NVIDIA GPUs | NVIDIA GPU with CUDA support |
 | [GPU (OpenCL)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_gpu.py) | Accelerates computations using OpenCL on GPUs | OpenCL 2.0 compatible device |
-| [CPU (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_cpu.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
-| [LLVM (LLVM IR)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | llvm libraries installed and findable |
-| [WEBGPU](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_webgpu.py) | Runs on GPU using the Dawn WebGPU engine (used in Google Chrome) | Dawn library installed and findable. Download binaries [here](https://github.com/wpmed92/pydawn/releases/tag/v0.1.6). |
-
-## Interoperability
-
-tinygrad provides interoperability with OpenCL and PyTorch, allowing efficient tensor data sharing between frameworks through the `Tensor.from_blob` API. This enables zero-copy operations by working directly with external memory pointers.
-
-**Important**: When using external memory pointers with tinygrad tensors, you must ensure these pointers remain valid throughout the entire lifetime of the tinygrad tensor to prevent memory corruption.
-
-### `CUDA`/`METAL` PyTorch Interoperability
-
-You can seamlessly work with CUDA/MPS tensors between PyTorch and tinygrad without data copying:
-```python
-from tinygrad.dtype import _from_torch_dtype
-tensor1 = torch.tensor([1.0, 2.0, 3.0], device=torch.device("cuda"))
-tiny_tensor1 = Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype=_from_torch_dtype(tensor1.dtype), device='CUDA')
-
-# Before tinygrad calculations, mps needs to be synchronized to make sure data is valid.
-if data.device.type == "mps": torch.mps.synchronize()
-else: torch.cuda.synchronize()
-
-x = (tiny_tensor1 + 1).realize()
-```
-
-### `QCOM` OpenCL Interoperability
-
-tinygrad supports OpenCL interoperability on `QCOM` backend.
-
-Buffer interop allows direct access to OpenCL memory buffers:
-```python
-# create raw opencl buffer.
-cl_buf = cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status := ctypes.c_int32())
-
-# extract pointers
-cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_buf), 8).cast('Q')[0]
-rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
-
-# create tiny tensor
-tiny = Tensor.from_blob(rawbuf_ptr, (8, 8), dtype=dtypes.int, device='QCOM')
-```
-
-And the same for the images:
-```python
-# create cl image.
-cl_img = cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status := ctypes.c_int32())
-
-# extract pointers
-cl_buf_desc_ptr = to_mv(ctypes.addressof(cl_img), 8).cast('Q')[0]
-rawbuf_ptr = to_mv(cl_buf_desc_ptr, 0x100).cast('Q')[20] # offset 0xA0 is a raw gpu pointer.
-
-# create tiny tensor
-tiny = Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtype=dtypes.imagef((h,w)), device='QCOM')
-```
+| [CLANG (C Code)](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_clang.py) | Runs on CPU using the clang compiler | `clang` compiler in system `PATH` |
+| [LLVM](https://github.com/tinygrad/tinygrad/tree/master/tinygrad/runtime/ops_llvm.py) | Runs on CPU using the LLVM compiler infrastructure | `llvmlite` package installed |

docs/tensor/creation.md

@@ -20,7 +20,6 @@
 
 ::: tinygrad.Tensor.manual_seed
 ::: tinygrad.Tensor.rand
-::: tinygrad.Tensor.rand_like
 ::: tinygrad.Tensor.randn
 ::: tinygrad.Tensor.randint
 ::: tinygrad.Tensor.normal

docs/tensor/elementwise.md

@@ -22,7 +22,6 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.round
 ::: tinygrad.Tensor.isinf
 ::: tinygrad.Tensor.isnan
-::: tinygrad.Tensor.isfinite
 ::: tinygrad.Tensor.lerp
 ::: tinygrad.Tensor.square
 ::: tinygrad.Tensor.clamp
@@ -53,7 +52,7 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.erf
 ::: tinygrad.Tensor.gelu
 ::: tinygrad.Tensor.quick_gelu
-::: tinygrad.Tensor.leaky_relu
+::: tinygrad.Tensor.leakyrelu
 ::: tinygrad.Tensor.mish
 ::: tinygrad.Tensor.softplus
 ::: tinygrad.Tensor.softsign
@@ -64,19 +63,13 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
 ::: tinygrad.Tensor.sub
 ::: tinygrad.Tensor.mul
 ::: tinygrad.Tensor.div
-::: tinygrad.Tensor.idiv
-::: tinygrad.Tensor.mod
-::: tinygrad.Tensor.bitwise_xor
-::: tinygrad.Tensor.bitwise_and
-::: tinygrad.Tensor.bitwise_or
-::: tinygrad.Tensor.bitwise_not
+::: tinygrad.Tensor.xor
 ::: tinygrad.Tensor.lshift
 ::: tinygrad.Tensor.rshift
 ::: tinygrad.Tensor.pow
 ::: tinygrad.Tensor.maximum
 ::: tinygrad.Tensor.minimum
 ::: tinygrad.Tensor.where
-::: tinygrad.Tensor.copysign
 
 ## Casting Ops
 

docs/tensor/movement.md

@@ -10,7 +10,6 @@
 
 ## Movement (high level)
 
-::: tinygrad.Tensor.__getitem__
 ::: tinygrad.Tensor.gather
 ::: tinygrad.Tensor.cat
 ::: tinygrad.Tensor.stack
@@ -25,5 +24,3 @@
 ::: tinygrad.Tensor.transpose
 ::: tinygrad.Tensor.flatten
 ::: tinygrad.Tensor.unflatten
-::: tinygrad.Tensor.roll
-::: tinygrad.Tensor.rearrange

docs/tensor/ops.md

@@ -6,10 +6,8 @@
 ::: tinygrad.Tensor.min
 ::: tinygrad.Tensor.any
 ::: tinygrad.Tensor.all
-::: tinygrad.Tensor.isclose
 ::: tinygrad.Tensor.mean
 ::: tinygrad.Tensor.var
-::: tinygrad.Tensor.var_mean
 ::: tinygrad.Tensor.std
 ::: tinygrad.Tensor.std_mean
 ::: tinygrad.Tensor.softmax
@@ -23,7 +21,6 @@
 
 ::: tinygrad.Tensor.avg_pool2d
 ::: tinygrad.Tensor.max_pool2d
-::: tinygrad.Tensor.max_unpool2d
 ::: tinygrad.Tensor.conv2d
 ::: tinygrad.Tensor.conv_transpose2d
 ::: tinygrad.Tensor.dot
@@ -35,10 +32,6 @@
 ::: tinygrad.Tensor.tril
 ::: tinygrad.Tensor.interpolate
 ::: tinygrad.Tensor.scatter
-::: tinygrad.Tensor.scatter_reduce
-::: tinygrad.Tensor.masked_select
-::: tinygrad.Tensor.sort
-::: tinygrad.Tensor.topk
 
 ## Neural Network (functional)
 

docs/tensor/properties.md

@@ -25,15 +25,10 @@
 ::: tinygrad.Tensor.replace
 ::: tinygrad.Tensor.assign
 ::: tinygrad.Tensor.detach
-::: tinygrad.Tensor.clone
 ::: tinygrad.Tensor.to
 ::: tinygrad.Tensor.to_
 ::: tinygrad.Tensor.shard
 ::: tinygrad.Tensor.shard_
 ::: tinygrad.Tensor.contiguous
 ::: tinygrad.Tensor.contiguous_backward
-
-## Gradient
-
-::: tinygrad.Tensor.gradient
 ::: tinygrad.Tensor.backward

examples/benchmark_onnx.py

@@ -1,41 +0,0 @@
-import sys, onnx, time, pickle
-from tinygrad import TinyJit, Device, GlobalCounters, fetch, getenv
-from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx_helpers import get_example_inputs, validate
-
-def load_onnx_model(onnx_file):
-  onnx_model = onnx.load(onnx_file)
-  run_onnx = OnnxRunner(onnx_model)
-  run_onnx_jit = TinyJit(lambda **kwargs: next(iter(run_onnx({k:v.to(Device.DEFAULT) for k,v in kwargs.items()}).values())), prune=True)
-  return run_onnx_jit, run_onnx.graph_inputs
-
-if __name__ == "__main__":
-  onnx_file = fetch(sys.argv[1])
-  run_onnx_jit, input_specs = load_onnx_model(onnx_file)
-  print("loaded model")
-
-  for i in range(3):
-    new_inputs = get_example_inputs(input_specs)
-    GlobalCounters.reset()
-    print(f"run {i}")
-    run_onnx_jit(**new_inputs)
-
-  # run 20 times
-  for _ in range(20):
-    new_inputs = get_example_inputs(input_specs)
-    GlobalCounters.reset()
-    st = time.perf_counter()
-    out = run_onnx_jit(**new_inputs)
-    mt = time.perf_counter()
-    val = out.numpy()
-    et = time.perf_counter()
-    print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {(et-st)*1e3:6.2f} ms")
-
-  if getenv("ORT"):
-    validate(onnx_file, new_inputs, rtol=1e-3, atol=1e-3)
-    print("model validated")
-
-  if (fn:=getenv("SAVE_PKL", "")) != "":
-    with open(fn, "wb") as f:
-      pickle.dump(run_onnx_jit, f)
-    print(f"pkl saved to {fn}")

examples/compile_efficientnet.py

@@ -15,9 +15,9 @@ if __name__ == "__main__":
   if getenv("WEBGPU"):
     safe_save(get_state_dict(model), (dirname / "net.safetensors").as_posix())
     load_state_dict(model, safe_load(str(dirname / "net.safetensors")))
-  mode = "clang" if getenv("CPU", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
+  mode = "clang" if getenv("CLANG", "") != "" else "webgpu" if getenv("WEBGPU", "") != "" else ""
   prg, inp_sizes, out_sizes, state = export_model(model, mode, Tensor.randn(1,3,224,224))
-  if getenv("CPU", "") == "":
+  if getenv("CLANG", "") == "":
     ext = "js" if getenv("WEBGPU", "") != "" else "json"
     with open(dirname / f"net.{ext}", "w") as text_file:
       text_file.write(prg)
@@ -68,6 +68,6 @@ if __name__ == "__main__":
     else printf("%s\\n", lbls[best_idx]);
   }""")
 
-    # CPU=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
+    # CLANG=1 python3 examples/compile_efficientnet.py | clang -O2 -lm -x c - -o recognize && DEBUG=1 time ./recognize docs/showcase/stable_diffusion_by_tinygrad.jpg
     # category : 281 (tabby, tabby cat) with 9.452788
     print('\n'.join(cprog))

examples/compile_tensorflow.py

@@ -1,14 +1,14 @@
 # An example to compile a small Tensorflow model to extremely portable C code
 
 import os, sys
-os.environ["CPU"] = '1'
+os.environ["CLANG"] = '1'
 os.environ["JIT"] = '2'
 
 import numpy as np
 import subprocess
 import tensorflow as tf
 import tf2onnx
-from tinygrad.frontend.onnx import OnnxRunner
+from extra.onnx import get_run_onnx
 from tinygrad.tensor import Tensor
 from extra.export_model import export_model_clang, compile_net, jit_model
 
@@ -25,7 +25,7 @@ class TinyOnnx:
   def __init__(self, keras_model):
     input_signature = [tf.TensorSpec([1,32], tf.float32, name='x')]
     onnx_model, _ = tf2onnx.convert.from_keras(keras_model, input_signature, opset=13)
-    self.run_onnx = OnnxRunner(onnx_model)
+    self.run_onnx = get_run_onnx(onnx_model)
 
   def forward(self, x):
     return self.run_onnx({"x": x}, debug=False)['predictions']

examples/conversation.py

@@ -117,7 +117,7 @@ def tts(
   stn_tst = text_mapper.get_text(text_to_synthesize, hps.data.add_blank, hps.data.text_cleaners)
   init_shape = stn_tst.shape
   assert init_shape[0] < pad_length, "text is too long"
-  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), value=1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
+  x_tst, x_tst_lengths = stn_tst.pad(((0, pad_length - init_shape[0]),), 1).unsqueeze(0), Tensor([init_shape[0]], dtype=dtypes.int64)
   sid = Tensor([speaker_id], dtype=dtypes.int64) if model_has_multiple_speakers else None
 
   # Perform inference.

examples/gpt2.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
-import os, argparse, contextlib
+import os, argparse
 from typing import Optional, Union
-with contextlib.suppress(ImportError): import tiktoken
+import tiktoken
 from tinygrad import Tensor, TinyJit, Device, GlobalCounters, Variable, dtypes
 from tinygrad.ops import UOp
 from tinygrad.helpers import Timing, DEBUG, JIT, getenv, fetch, colored, trange

examples/handcode_opt.py

@@ -1,3 +1,4 @@
+from typing import List, Tuple
 from extra.models.resnet import ResNet50
 from extra.mcts_search import mcts_search
 from examples.mlperf.helpers import get_mlperf_bert_model
@@ -5,9 +6,9 @@ from tinygrad import Tensor, Device, dtypes, nn
 from tinygrad.codegen.kernel import Kernel
 from tinygrad.ops import Ops, sym_infer
 from tinygrad.device import Compiled
-from tinygrad.engine.search import beam_search, bufs_from_lin
+from tinygrad.engine.schedule import create_schedule
+from tinygrad.engine.search import time_linearizer, beam_search, bufs_from_lin
 from tinygrad.helpers import DEBUG, ansilen, getenv, colored, TRACEMETA
-from extra.optimization.helpers import time_linearizer
 
 def get_sched_resnet():
   mdl = ResNet50()
@@ -17,12 +18,12 @@ def get_sched_resnet():
   # run model twice to get only what changes, these are the kernels of the model
   for _ in range(2):
     out = mdl(Tensor.empty(BS, 3, 224, 224))
-    targets = [out]
+    targets = [out.lazydata]
     if getenv("BACKWARD"):
       optim.zero_grad()
       out.sparse_categorical_crossentropy(Tensor.empty(BS, dtype=dtypes.int)).backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
+      targets += [x.lazydata for x in optim.schedule_step()]
+    sched = create_schedule(targets)
     print(f"schedule length {len(sched)}")
   return sched
 
@@ -41,16 +42,17 @@ def get_sched_bert():
   next_sentence_labels = Tensor.empty((BS, 1), dtype=dtypes.float32)
 
   # run model twice to get only what changes, these are the kernels of the model
+  seen = set()
   for _ in range(2):
     lm_logits, seq_relationship_logits = mdl(input_ids, attention_mask, masked_positions, segment_ids)
-    targets = [lm_logits, seq_relationship_logits]
+    targets = [lm_logits.lazydata, seq_relationship_logits.lazydata]
     if getenv("BACKWARD"):
       optim.zero_grad()
       loss = mdl.loss(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
       # ignore grad norm and loss scaler for now
       loss.backward()
-      targets += [x for x in optim.schedule_step()]
-    sched = Tensor.schedule(*targets)
+      targets += [x.lazydata for x in optim.schedule_step()]
+    sched = create_schedule(targets)
     print(f"schedule length {len(sched)}")
   return sched
 
@@ -79,7 +81,7 @@ if __name__ == "__main__":
     rawbufs = bufs_from_lin(Kernel(si.ast))
 
     # "linearize" the op into uops in different ways
-    lins: list[tuple[Kernel, str]] = []
+    lins: List[Tuple[Kernel, str]] = []
 
     # always try hand coded opt
     lin = Kernel(si.ast, opts=device.renderer)
@@ -107,7 +109,7 @@ if __name__ == "__main__":
     choices = []
     for lin, nm in lins:
       tm = time_linearizer(lin, rawbufs, allow_test_size=False, cnt=10, disable_cache=True)
-      ops = (prg:=lin.to_program()).estimates.ops
+      ops = (prg:=lin.to_program()).op_estimate
       gflops = sym_infer(ops, {k:k.min for k in lin.ast.variables()})*1e-9/tm
       choices.append((tm, gflops, lin, prg, nm))
 

examples/hlb_cifar10.py

@@ -11,6 +11,7 @@ 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 tinygrad.multi import MultiLazyBuffer
 
 cifar_mean = [0.4913997551666284, 0.48215855929893703, 0.4465309133731618]
 cifar_std = [0.24703225141799082, 0.24348516474564, 0.26158783926049628]
@@ -34,6 +35,8 @@ class UnsyncedBatchNorm:
     self.num_batches_tracked = Tensor.zeros(1, dtype=dtypes.int, requires_grad=False)
 
   def __call__(self, x:Tensor):
+    if isinstance(x.lazydata, MultiLazyBuffer): assert x.lazydata.axis is None or x.lazydata.axis == 0 and len(x.lazydata.lbs) == self.num_devices
+
     xr = x.reshape(self.num_devices, -1, *x.shape[1:]).cast(dtypes.float32)
     batch_mean, batch_invstd = self.calc_stats(xr)
     ret = xr.batchnorm(

examples/index.html

@@ -17,6 +17,7 @@ canvas {  display: none; }
 * {  text-align: center;  font-family: monospace; }
 </style>
 <title>tinygrad has WebGPU</title>
+<script src="./net.js"></script>
 <link rel="icon" type="image/x-icon" href="https://raw.githubusercontent.com/tinygrad/tinygrad/master/docs/logo.png">
 </head>
 <body>
@@ -45,10 +46,7 @@ canvas {  display: none; }
 	const getDevice = async () => {
 		if (!navigator.gpu) error("WebGPU not supported.");
 		const adapter = await navigator.gpu.requestAdapter();
-		return await adapter.requestDevice({
-			requiredFeatures: ["shader-f16"],
-			powerPreference: "high-performance"
-		});
+		return await adapter.requestDevice();
 	};
 
 	const timer = async (func, label = "") => {
@@ -63,6 +61,8 @@ canvas {  display: none; }
 
 	const getLabels = async () => (await fetch("https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json")).json();
 
+	const getSavetensorBuffer = async () => new Uint8Array(await (await fetch("./net.safetensors")).arrayBuffer());
+
 	const reorderChannelsAndRemoveAlpha = (data) => {
 		const out = [];
 		let i = 0;
@@ -97,9 +97,9 @@ canvas {  display: none; }
 		try {
 			resultText.innerHTML = "loading..."
 			labels = await getLabels();
+			const safetensor = await getSavetensorBuffer();
 			const device = await getDevice();
-			const model = (await import("../../net.js")).default;
-			net = await timer(() => model.load(device, '../../net.safetensors'), "(compilation)");
+			net = await timer(() => setupNet(device, safetensor), "(compilation)");
 			resultText.innerHTML = "ready"
 		} catch (e) {
 			error(e)

examples/llama3.py

@@ -47,7 +47,7 @@ def concat_weights(models, device=None):
     disk_tensors: List[Tensor] = [model[name] for model in models]
     if len(disk_tensors) == 1 or len(disk_tensors[0].shape) == 1:
       return disk_tensors[0].to(device=device)
-    axis = 1 if name.endswith((".attention.wo.weight", ".feed_forward.w2.weight")) else 0
+    axis = 1 if name.endswith(".attention.wo.weight") or name.endswith(".feed_forward.w2.weight") else 0
     lazy_tensors = [data.to(device=device) for data in disk_tensors]
     return lazy_tensors[0].cat(*lazy_tensors[1:], dim=axis)
   return {name: convert(name) for name in {name: None for model in models for name in model}}
@@ -73,17 +73,16 @@ class Int8Linear:
     self.scale = Tensor.ones(out_features, dtype=dtypes.half)
 
   def __call__(self, x):
-    return x.dot(self.weight.cast(self.scale.dtype).T*self.scale)
+    return x.dot(self.weight.cast(dtype=dtypes.half).T*self.scale)
 
   @staticmethod
-  def quantize(tensors, device, scale_dtype=dtypes.float16, quantize_embeds=False):
+  def quantize(tensors, device):
     new_tensors = {}
     for name,v in tensors.items():
-      if "feed_forward" in name or "attention.w" in name or (quantize_embeds and "tok_embeddings.weight" in name):
+      if "feed_forward" in name or "attention.w" in name:
         assert "weight" in name, name
-        v = v.cast(scale_dtype)
         scale = v.abs().max(axis=1) / 127.0
-        int8_weight = (v.T/scale).T.round().cast(dtype=dtypes.int8) # without round(), cast truncates -34.9 to -34
+        int8_weight = (v.T/scale).T.cast(dtype=dtypes.int8)
         new_tensors[name] = int8_weight
         new_tensors[name.replace('weight', 'scale')] = scale
         if isinstance(device, tuple):
@@ -91,20 +90,8 @@ class Int8Linear:
           new_tensors[name.replace('weight', 'scale')].shard_(device, axis=None)
       else:
         new_tensors[name] = v
-    if quantize_embeds: new_tensors.update({"output.weight": new_tensors["tok_embeddings.weight"], "output.scale": new_tensors["tok_embeddings.scale"]})
     return new_tensors
 
-class Int8Embedding:
-  def __init__(self, vocab_size:int, embed_size:int):
-    self.vocab_sz, self.embed_sz = vocab_size, embed_size
-    self.weight, self.scale = Tensor.ones(vocab_size, embed_size, dtype=dtypes.int8), Tensor.ones(vocab_size, dtype=dtypes.half)
-
-  def __call__(self, idx:Tensor) -> Tensor:
-    if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).unsqueeze(-1)
-    big_shp = idx.shape+(self.vocab_sz, self.embed_sz)
-    arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1)).expand(big_shp), (self.weight.cast(self.scale.dtype).T*self.scale).T
-    return (arange == idx).mul(vals).sum(-2, dtype=vals.dtype)
-
 def NF4Linear(block_size):
   _CODE = [
     -1.0, -0.6961928009986877, -0.5250730514526367, -0.39491748809814453, -0.28444138169288635, -0.18477343022823334, -0.09105003625154495, 0.0,
@@ -126,7 +113,7 @@ def NF4Linear(block_size):
       return x.linear(unscaled.reshape(self.out_features, self.in_features).T)
 
     @staticmethod
-    def quantize(state_dict: dict[str, Tensor], device, scale_dtype=dtypes.float16) -> dict[str, Tensor]:
+    def quantize(state_dict: dict[str, Tensor], device) -> dict[str, Tensor]:
       new_state_dict = {}
       for k, v in state_dict.items():
         if "feed_forward" in k or "attention.w" in k:
@@ -134,7 +121,7 @@ def NF4Linear(block_size):
           scale = (grouped.abs().max(axis=1, keepdim=True))
           coded = ((grouped / scale).unsqueeze(-1) - CODE.to(v.device)).abs().argmin(axis=-1).cast(dtypes.uint8).flatten()
           new_state_dict[k] = coded[::2] * 2 ** 4 + coded[1::2]
-          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(scale_dtype)
+          new_state_dict[k.replace(".weight", ".scale")] = scale.cast(dtypes.float16)
           if isinstance(device, tuple):
             new_state_dict[k].shard_(device, axis=-1)
             new_state_dict[k.replace('weight', 'scale')].shard_(device, axis=None)
@@ -157,14 +144,13 @@ MODEL_PARAMS = {
     "files": 8
   }
 }
-def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dtype=dtypes.float16, device=None, max_context=8192, load_weights=True):
+def build_transformer(model_path: Path, model_size="8B", quantize=None, device=None):
   # build model
-  if quantize == "int8": linear, embedding, quantize_embeds = Int8Linear, Int8Embedding, True
-  elif quantize == "nf4": linear, embedding, quantize_embeds = NF4Linear(64), nn.Embedding, False
-  else: linear, embedding, quantize_embeds = nn.Linear, nn.Embedding, False
-  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, embedding=embedding, max_context=max_context, jit=True)
+  if quantize == "int8": linear = Int8Linear
+  elif quantize == "nf4": linear = NF4Linear(64)
+  else: linear = nn.Linear
+  model = Transformer(**MODEL_PARAMS[model_size]["args"], linear=linear, max_context=8192, jit=True)
 
-  if not load_weights: return model
   # load weights
   if model_path.is_dir():
     if (model_path / "model.safetensors.index.json").exists(): weights = load(str(model_path / "model.safetensors.index.json"))
@@ -182,7 +168,7 @@ def build_transformer(model_path: Path, model_size="8B", quantize=None, scale_dt
     # quantize
     if quantize == "float16": weights = {k:v.cast(quantize).contiguous() for k,v in weights.items()}
     elif quantize is not None:
-      weights = linear.quantize(weights, device, scale_dtype, quantize_embeds)
+      weights = linear.quantize(weights, device)
       for _,v in weights.items(): v.realize()
 
     # shard
@@ -234,7 +220,7 @@ if __name__ == "__main__":
   parser = argparse.ArgumentParser()
   parser.add_argument("--download_model", action="store_true", help="Download a model")
   parser.add_argument("--model", type=Path, help="Model path")
-  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="1B", help="Model size")
+  parser.add_argument("--size", choices=["1B", "8B", "70B"], default="8B", 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,8 +234,8 @@ if __name__ == "__main__":
   parser.add_argument("--profile", action="store_true", help="Output profile data")
   args = parser.parse_args()
 
-  # download_model is the default without a model passed in
-  if args.download_model or not args.model:
+  assert (args.model and not args.download_model) or (not args.model and args.download_model), "either download or provide model"
+  if args.download_model:
     if args.size == "1B":
       fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
       args.model = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-Q6_K.gguf", "Llama-3.2-1B-Instruct-Q6_K.gguf", subdir="llama3-1b-instruct")
@@ -261,11 +247,11 @@ if __name__ == "__main__":
       fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/resolve/main/model-00004-of-00004.safetensors", "model-00004-of-00004.safetensors", subdir="llama3-8b-sfr")
       args.model = fetch("https://huggingface.co/TriAiExperiments/SFR-Iterative-DPO-LLaMA-3-8B-R/raw/main/model.safetensors.index.json", "model.safetensors.index.json", subdir="llama3-8b-sfr")
     elif args.size == "70B":
-      subdir = "DeepSeek-R1-Distill-Llama-70B"
-      args.model = fetch("https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
+      subdir = "Llama-3.1-Nemotron-70B-Instruct-HF"
+      args.model = fetch("https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model.safetensors.index.json?download=true", "model.safetensors.index.json", subdir=subdir)
       fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir=subdir)
-      for i in range(17):
-        fetch(f"https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Llama-70B/resolve/main/model-{i+1:05d}-of-000017.safetensors?download=true", f"model-{i+1:05d}-of-000017.safetensors", subdir=subdir)
+      for i in range(30):
+        fetch(f"https://huggingface.co/nvidia/Llama-3.1-Nemotron-70B-Instruct-HF/resolve/main/model-{i+1:05d}-of-00030.safetensors?download=true", f"model-{i+1:05d}-of-00030.safetensors", subdir=subdir)
 
   assert args.model is not None, "please provide --model option"
 

examples/llm.c/export.py

@@ -2,9 +2,10 @@
 import os
 if "NOOPT" not in os.environ: os.environ["NOOPT"] = "1"
 from tinygrad import Device, nn, Tensor, dtypes, Variable
-Device.DEFAULT = "CPU"
+Device.DEFAULT = "CLANG"
 from train_gpt2 import GPT, GPTConfig
 from tinygrad.helpers import dedup, to_function_name, flatten, getenv, GlobalCounters, ansilen, to_function_name
+from tinygrad.engine.schedule import create_schedule
 from tinygrad.engine.realize import get_kernel, run_schedule
 from tinygrad.engine.memory import memory_planner
 from tinygrad.ops import Ops
@@ -36,16 +37,16 @@ if __name__ == "__main__":
       tensors = optimizer.schedule_step()
     else:
       tensors = []
-    sched = loss.schedule(*tensors)
+    sched = create_schedule([loss.lazydata] + [x.lazydata for x in tensors])
     print(f"calls {i}:", len(sched))
     #run_schedule(sched[:])
   sched = memory_planner(sched)
   ast_dedup = dedup([si.ast for si in sched if si.ast.op is Ops.SINK])
   srcs = {}
   for ast in ast_dedup:
-    k = get_kernel(Device["CPU"].renderer, ast)
+    k = get_kernel(Device["CLANG"].renderer, ast)
     k.linearize()
-    src = Device["CPU"].renderer.render(to_function_name(k.name), k.uops)
+    src = Device["CLANG"].renderer.render(to_function_name(k.name), k.uops)
     srcs[ast] = (k.name, src)
   print("functions:", len(srcs))
   used_buffers = dedup(flatten([si.bufs for si in sched]))

examples/mlperf/dataloader.py

@@ -170,13 +170,13 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
 
 def process_batch_bert(data: List[dict]) -> dict[str, Tensor]:
   return {
-    "input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32, device="CPU"),
-    "next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.int32, device="CPU"),
+    "input_ids": Tensor(np.concatenate([s["input_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "input_mask": Tensor(np.concatenate([s["input_mask"] for s in data], axis=0), dtype=dtypes.default_float),
+    "segment_ids": Tensor(np.concatenate([s["segment_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_positions": Tensor(np.concatenate([s["masked_lm_positions"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_ids": Tensor(np.concatenate([s["masked_lm_ids"] for s in data], axis=0), dtype=dtypes.float32),
+    "masked_lm_weights": Tensor(np.concatenate([s["masked_lm_weights"] for s in data], axis=0), dtype=dtypes.float32),
+    "next_sentence_labels": Tensor(np.concatenate([s["next_sentence_labels"] for s in data], axis=0), dtype=dtypes.float32),
   }
 
 def load_file(file: str):
@@ -223,8 +223,14 @@ def batch_load_train_bert(BS:int):
   assert cycle_length > 0, "cycle_length must be greater than 0"
 
   dataset = InterleavedDataset(train_files, cycle_length)
+  buffer = [dataset.get() for _ in range(1000)]
   while True:
-    yield process_batch_bert([dataset.get() for _ in range(BS)])
+    batch = []
+    for _ in range(BS):
+      index = random.randint(0, 999)
+      batch.append(buffer[index])
+      buffer[index] = dataset.get()
+    yield process_batch_bert(batch)
 
 # Reference: https://github.com/mlcommons/training/blob/1c8a098ae3e70962a4f7422c0b0bd35ae639e357/language_model/tensorflow/bert/run_pretraining.py, Line 416
 def batch_load_val_bert(BS:int):
@@ -348,167 +354,6 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
       # happens with BENCHMARK set
       pass
 
-### RetinaNet
-
-def load_retinanet_data(base_dir:Path, val:bool, queue_in:Queue, queue_out:Queue,
-                        imgs:Tensor, boxes:Tensor, labels:Tensor, matches:Tensor|None=None,
-                        anchors:Tensor|None=None, seed:int|None=None):
-  from extra.datasets.openimages import image_load, random_horizontal_flip, resize
-  from examples.mlperf.helpers import box_iou, find_matches, generate_anchors
-  import torch
-
-  while (data:=queue_in.get()) is not None:
-    idx, img, tgt = data
-    img = image_load(base_dir, img["subset"], img["file_name"])
-
-    if val:
-      img = resize(img)[0]
-    else:
-      if seed is not None:
-        np.random.seed(seed)
-        random.seed(seed)
-        torch.manual_seed(seed)
-
-      img, tgt = random_horizontal_flip(img, tgt)
-      img, tgt, _ = resize(img, tgt=tgt)
-      match_quality_matrix = box_iou(tgt["boxes"], (anchor := np.concatenate(generate_anchors((800, 800)))))
-      match_idxs = find_matches(match_quality_matrix, allow_low_quality_matches=True)
-      clipped_match_idxs = np.clip(match_idxs, 0, None)
-      clipped_boxes, clipped_labels = tgt["boxes"][clipped_match_idxs], tgt["labels"][clipped_match_idxs]
-
-      boxes[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_boxes.tobytes()
-      labels[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = clipped_labels.tobytes()
-      matches[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = match_idxs.tobytes()
-      anchors[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = anchor.tobytes()
-
-    imgs[idx].contiguous().realize().lazydata.base.realized.as_buffer(force_zero_copy=True)[:] = img.tobytes()
-
-    queue_out.put(idx)
-  queue_out.put(None)
-
-def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, shuffle:bool=True, seed:int|None=None):
-  def _enqueue_batch(bc):
-    from extra.datasets.openimages import prepare_target
-    for idx in range(bc * batch_size, (bc+1) * batch_size):
-      img = dataset.loadImgs(next(dataset_iter))[0]
-      ann = dataset.loadAnns(dataset.getAnnIds(img_id:=img["id"]))
-      tgt = prepare_target(ann, img_id, (img["height"], img["width"]))
-
-      if img_ids is not None:
-        img_ids[idx] = img_id
-
-      if img_sizes is not None:
-        img_sizes[idx] = tgt["image_size"]
-
-      queue_in.put((idx, img, tgt))
-
-  def _setup_shared_mem(shm_name:str, size:tuple[int, ...], dtype:dtypes) -> tuple[shared_memory.SharedMemory, Tensor]:
-    if os.path.exists(f"/dev/shm/{shm_name}"): os.unlink(f"/dev/shm/{shm_name}")
-    shm = shared_memory.SharedMemory(name=shm_name, create=True, size=prod(size))
-    shm_tensor = Tensor.empty(*size, dtype=dtype, device=f"disk:/dev/shm/{shm_name}")
-    return shm, shm_tensor
-
-  image_ids = sorted(dataset.imgs.keys())
-  batch_count = min(32, len(image_ids) // batch_size)
-
-  queue_in, queue_out = Queue(), Queue()
-  procs, data_out_count = [], [0] * batch_count
-
-  shm_imgs, imgs = _setup_shared_mem("retinanet_imgs", (batch_size * batch_count, 800, 800, 3), dtypes.uint8)
-
-  if val:
-    boxes, labels, matches, anchors = None, None, None, None
-    img_ids, img_sizes = [None] * (batch_size * batch_count), [None] * (batch_size * batch_count)
-  else:
-    img_ids, img_sizes = None, None
-    shm_boxes, boxes = _setup_shared_mem("retinanet_boxes", (batch_size * batch_count, 120087, 4), dtypes.float32)
-    shm_labels, labels = _setup_shared_mem("retinanet_labels", (batch_size * batch_count, 120087), dtypes.int64)
-    shm_matches, matches = _setup_shared_mem("retinanet_matches", (batch_size * batch_count, 120087), dtypes.int64)
-    shm_anchors, anchors = _setup_shared_mem("retinanet_anchors", (batch_size * batch_count, 120087, 4), dtypes.float64)
-
-  shutdown = False
-  class Cookie:
-    def __init__(self, bc):
-      self.bc = bc
-    def __del__(self):
-      if not shutdown:
-        try: _enqueue_batch(self.bc)
-        except StopIteration: pass
-
-  def shuffle_indices(indices, seed):
-    rng = random.Random(seed)
-    rng.shuffle(indices)
-
-  if shuffle: shuffle_indices(image_ids, seed=seed)
-  dataset_iter = iter(image_ids)
-
-  try:
-    for _ in range(cpu_count()):
-      proc = Process(
-        target=load_retinanet_data,
-        args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
-        kwargs={"matches": matches, "anchors": anchors, "seed": seed}
-      )
-      proc.daemon = True
-      proc.start()
-      procs.append(proc)
-
-    for bc in range(batch_count):
-      _enqueue_batch(bc)
-
-    for _ in range(len(image_ids) // batch_size):
-      while True:
-        bc = queue_out.get() // batch_size
-        data_out_count[bc] += 1
-        if data_out_count[bc] == batch_size: break
-
-      data_out_count[bc] = 0
-
-      if val:
-        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
-               img_ids[bc * batch_size:(bc + 1) * batch_size],
-               img_sizes[bc * batch_size:(bc + 1) * batch_size],
-               Cookie(bc))
-      else:
-        yield (imgs[bc * batch_size:(bc + 1) * batch_size],
-               boxes[bc * batch_size:(bc + 1) * batch_size],
-               labels[bc * batch_size:(bc + 1) * batch_size],
-               matches[bc * batch_size:(bc + 1) * batch_size],
-               anchors[bc * batch_size:(bc + 1) * batch_size],
-               Cookie(bc))
-  finally:
-    shutdown = True
-
-    for _ in procs: queue_in.put(None)
-    queue_in.close()
-
-    for _ in procs:
-      while queue_out.get() is not None: pass
-    queue_out.close()
-
-    # shutdown processes
-    for proc in procs: proc.join()
-
-    shm_imgs.close()
-
-    if not val:
-      shm_boxes.close()
-      shm_labels.close()
-      shm_matches.close()
-      shm_anchors.close()
-
-    try:
-      shm_imgs.unlink()
-
-      if not val:
-        shm_boxes.unlink()
-        shm_labels.unlink()
-        shm_matches.unlink()
-        shm_anchors.unlink()
-    except FileNotFoundError:
-      # happens with BENCHMARK set
-      pass
-
 if __name__ == "__main__":
   def load_unet3d(val):
     assert not val, "validation set is not supported due to different sizes on inputs"
@@ -529,14 +374,6 @@ if __name__ == "__main__":
       for x,y,c in batch_load_resnet(val=val):
         pbar.update(x.shape[0])
 
-  def load_retinanet(val):
-    from extra.datasets.openimages import BASEDIR, download_dataset
-    from pycocotools.coco import COCO
-    dataset = COCO(download_dataset(base_dir:=getenv("BASE_DIR", BASEDIR), "validation" if val else "train"))
-    with tqdm(total=len(dataset.imgs.keys())) as pbar:
-      for x in batch_load_retinanet(dataset, val, base_dir):
-        pbar.update(x[0].shape[0])
-
   load_fn_name = f"load_{getenv('MODEL', 'resnet')}"
   if load_fn_name in globals():
     globals()[load_fn_name](getenv("VAL", 1))

examples/mlperf/helpers.py

@@ -1,7 +1,6 @@
 from collections import OrderedDict
 import unicodedata
 from typing import Optional
-import math
 import numpy as np
 from tinygrad.nn import state
 from tinygrad.tensor import Tensor, dtypes
@@ -196,18 +195,20 @@ def get_bert_qa_prediction(features, example, start_end_logits):
   return "empty"
 
 def get_mlperf_bert_config():
-  """benchmark is BERT-large"""
-  ret = {"attention_probs_dropout_prob": 0.1, "hidden_dropout_prob": 0.1, "vocab_size": 30522, "type_vocab_size": 2, "max_position_embeddings": 512}
+  """Config is BERT-large"""
+  return {
+    "attention_probs_dropout_prob": 0.1,
+    "hidden_dropout_prob": 0.1,
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "max_position_embeddings": 512,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "type_vocab_size": 2,
+    "vocab_size": 30522
+  }
 
-  match (bert_size:=getenv("BERT_SIZE", "large")):
-    case "large": ret.update({"hidden_size": 1024, "intermediate_size": 4096, "num_attention_heads": 16, "num_hidden_layers": 24})
-    case "tiny": ret.update({"hidden_size": 128, "intermediate_size": 512, "num_attention_heads": 2, "num_hidden_layers": 2})
-    case _: raise RuntimeError(f"unhandled {bert_size=}")
-
-  if (bert_layers:=getenv("BERT_LAYERS")): ret["num_hidden_layers"] = bert_layers
-  return ret
-
-def get_mlperf_bert_model():
+def get_mlperf_bert_model(checkpoint_path:Optional[str]=None):
   from extra.models import bert
   from examples.mlperf.initializers import LinearBert, EmbeddingBert, LayerNormBert
 
@@ -219,79 +220,21 @@ def get_mlperf_bert_model():
   config = get_mlperf_bert_config()
   if getenv("DISABLE_DROPOUT", 0):
     config["hidden_dropout_prob"] = config["attention_probs_dropout_prob"] = 0.0
-  return BertForPretraining(**config)
+  model = BertForPretraining(**config)
+  return model.load_from_pretrained(checkpoint_path) if checkpoint_path else model
 
-def get_fake_data_bert(BS:int):
+def get_data_bert(GPUS:list[str], it):
+  data: dict[str, Tensor] = next(it)
+  for key in data.keys(): data[key].shard_(GPUS, axis=0)
+  return data
+
+def get_fake_data_bert(GPUS:list[str], BS:int):
   return {
-    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.int32, device="CPU"),
-    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32, device="CPU"),
-    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.int32, device="CPU"),
+    "input_ids": Tensor.empty((BS, 512), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "input_mask": Tensor.empty((BS, 512), dtype=dtypes.default_float).contiguous().shard_(GPUS, axis=0),
+    "segment_ids": Tensor.empty((BS, 512), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_positions": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_ids": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "masked_lm_weights": Tensor.empty((BS, 76), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
+    "next_sentence_labels": Tensor.empty((BS, 1), dtype=dtypes.float32).contiguous().shard_(GPUS, axis=0),
   }
-
-def find_matches(match_quality_matrix:np.ndarray, high_threshold:float=0.5, low_threshold:float=0.4, allow_low_quality_matches:bool=False) -> np.ndarray:
-  BELOW_LOW_THRESHOLD, BETWEEN_THRESHOLDS = -1, -2
-
-  def _set_low_quality_matches_(matches:np.ndarray, all_matches:np.ndarray, match_quality_matrix:np.ndarray):
-    highest_quality_foreach_gt = np.max(match_quality_matrix, axis=1)
-    pred_inds_to_update = np.nonzero(match_quality_matrix == highest_quality_foreach_gt[:, None])[1]
-    matches[pred_inds_to_update] = all_matches[pred_inds_to_update]
-
-  assert low_threshold <= high_threshold
-
-  matched_vals, matches = match_quality_matrix.max(axis=0), match_quality_matrix.argmax(axis=0)
-  all_matches = np.copy(matches) if allow_low_quality_matches else None
-  below_low_threshold = matched_vals < low_threshold
-  between_thresholds = (matched_vals >= low_threshold) & (matched_vals < high_threshold)
-  matches[below_low_threshold] = BELOW_LOW_THRESHOLD
-  matches[between_thresholds] = BETWEEN_THRESHOLDS
-
-  if allow_low_quality_matches:
-    assert all_matches is not None
-    _set_low_quality_matches_(matches, all_matches, match_quality_matrix)
-
-  return matches
-
-def box_iou(boxes1:np.ndarray, boxes2:np.ndarray) -> np.ndarray:
-  def _box_area(boxes:np.ndarray) -> np.ndarray: return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
-
-  def _box_inter_union(boxes1:np.ndarray, boxes2:np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    area1, area2 = _box_area(boxes1), _box_area(boxes2)
-    lt, rb = np.maximum(boxes1[:, None, :2], boxes2[:, :2]), np.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
-    wh = np.clip(rb - lt, a_min=0, a_max=None)
-    inter = wh[:, :, 0] * wh[:, :, 1]
-    union = area1[:, None] + area2 - inter
-    return inter, union
-
-  inter, union = _box_inter_union(boxes1, boxes2)
-  return inter / union
-
-def generate_anchors(input_size:tuple[int, int], scales:Optional[tuple[Tensor, ...]]=None, aspect_ratios:Optional[tuple[Tensor, ...]]=None) -> list[np.ndarray]:
-  def _compute_grid_sizes(input_size:tuple[int, int]) -> np.ndarray:
-    return np.ceil(np.array(input_size)[None, :] / 2 ** np.arange(3, 8)[:, None])
-
-  scales = tuple((i, int(i * 2 ** (1/3)), int(i * 2 ** (2/3))) for i in 2 ** np.arange(5, 10)) if scales is None else scales
-  aspect_ratios = ((0.5, 1.0, 2.0),) * len(scales) if aspect_ratios is None else aspect_ratios
-  aspect_ratios = tuple(ar for ar in aspect_ratios)
-  grid_sizes = _compute_grid_sizes(input_size)
-
-  assert len(scales) == len(aspect_ratios) == len(grid_sizes), "scales, aspect_ratios, and grid_sizes must have the same length"
-
-  anchors = []
-  for s, ar, gs in zip(scales, aspect_ratios, grid_sizes):
-    s, ar = np.array(s), np.array(ar)
-    h_ratios = np.sqrt(ar)
-    w_ratios = 1 / h_ratios
-    ws = (w_ratios[:, None] * s[None, :]).reshape(-1)
-    hs = (h_ratios[:, None] * s[None, :]).reshape(-1)
-    base_anchors = (np.stack([-ws, -hs, ws, hs], axis=1) / 2).round()
-    stride_h, stride_w = input_size[0] // gs[0], input_size[1] // gs[1]
-    shifts_x, shifts_y = np.meshgrid(np.arange(gs[1]) * stride_w, np.arange(gs[0]) * stride_h)
-    shifts_x, shifts_y = shifts_x.reshape(-1), shifts_y.reshape(-1)
-    shifts = np.stack([shifts_x, shifts_y, shifts_x, shifts_y], axis=1, dtype=np.float32)
-    anchors.append((shifts[:, None] + base_anchors[None, :]).reshape(-1, 4))
-
-  return anchors

examples/mlperf/initializers.py

@@ -1,5 +1,5 @@
 import math
-from typing import Union
+from typing import Union, Tuple
 
 from tinygrad import Tensor, nn, dtypes
 from tinygrad.helpers import prod, argfix
@@ -53,10 +53,10 @@ class EmbeddingBert(nn.Embedding):
     arange_shp, weight_shp, big_shp = (1, 1, self.vocab_sz, 1), (1, 1, self.vocab_sz, self.embed_sz), idx.shape+(self.vocab_sz, self.embed_sz,)
     if not hasattr(self, 'arange'): self.arange = Tensor.arange(self.vocab_sz, requires_grad=False, device=self.weight.device).reshape(arange_shp)
     arange, idx, vals = self.arange.expand(big_shp), idx.reshape(idx.shape+(1, 1,)).expand(big_shp), self.weight.cast(dtypes.default_float).reshape(weight_shp).expand(big_shp)
-    return (arange == idx).mul(vals).sum(2, dtype=vals.dtype)
+    return (arange == idx).mul(vals).sum(2, acc_dtype=vals.dtype)
 
 class LayerNormBert:
-  def __init__(self, normalized_shape:Union[int, tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
+  def __init__(self, normalized_shape:Union[int, Tuple[int, ...]], eps:float=1e-12, elementwise_affine:bool=True):
     self.normalized_shape = (normalized_shape,) if isinstance(normalized_shape, int) else tuple(normalized_shape)
     self.axis, self.eps, self.elementwise_affine = tuple(-1-i for i in range(len(self.normalized_shape))), eps, elementwise_affine
     self.weight, self.bias = (Tensor.ones(*self.normalized_shape, dtype=dtypes.float32), Tensor.zeros(*self.normalized_shape, dtype=dtypes.float32)) if elementwise_affine else (None, None)

examples/mlperf/losses.py

@@ -1,29 +1,6 @@
 from examples.mlperf.metrics import dice_score
-from tinygrad import Tensor
 
 def dice_ce_loss(pred, tgt):
   ce = pred.permute(0, 2, 3, 4, 1).sparse_categorical_crossentropy(tgt.squeeze(1))
   dice = (1.0 - dice_score(pred, tgt, argmax=False, to_one_hot_x=False)).mean()
   return (dice + ce) / 2
-
-def sigmoid_focal_loss(pred:Tensor, tgt:Tensor, alpha:float=0.25, gamma:float=2.0, reduction:str="none") -> Tensor:
-  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
-  p, ce_loss = pred.sigmoid(), pred.binary_crossentropy_logits(tgt, reduction="none")
-  p_t = p * tgt + (1 - p) * (1 - tgt)
-  loss = ce_loss * ((1 - p_t) ** gamma)
-
-  if alpha >= 0:
-    alpha_t = alpha * tgt + (1 - alpha) * (1 - tgt)
-    loss = loss * alpha_t
-
-  if reduction == "mean": loss = loss.mean()
-  elif reduction == "sum": loss = loss.sum()
-  return loss
-
-def l1_loss(pred:Tensor, tgt:Tensor, reduction:str="none") -> Tensor:
-  assert reduction in ["mean", "sum", "none"], f"unsupported reduction {reduction}"
-  loss = (pred - tgt).abs()
-
-  if reduction == "mean": loss = loss.mean()
-  elif reduction == "sum": loss = loss.sum()
-  return loss

examples/mlperf/model_train.py

@@ -79,7 +79,7 @@ def train_resnet():
   lr_warmup_epochs  = config["lr_warmup_epochs"]  = getenv("WARMUP_EPOCHS", 2)
   decay             = config["decay"]             = getenv("DECAY", 2e-4)
 
-  loss_scaler       = config["LOSS_SCALER"]       = getenv("LOSS_SCALER", 256.0 if dtypes.default_float == dtypes.float16 else 1.0)
+  loss_scaler       = config["LOSS_SCALER"]       = getenv("LOSS_SCALER", 128.0 if dtypes.default_float == dtypes.float16 else 1.0)
 
   target, achieved  = getenv("TARGET", 0.759), False
   eval_start_epoch  = getenv("EVAL_START_EPOCH", 0)
@@ -273,7 +273,7 @@ def train_resnet():
       else:
         it = iter(tqdm(batch_load_resnet(batch_size=EVAL_BS, val=True, shuffle=False, pad_first_batch=True), total=steps_in_val_epoch))
         i, proc = 0, data_get(it)
-
+        
       prev_cookies = []
       while proc is not None:
         GlobalCounters.reset()
@@ -446,7 +446,7 @@ def train_unet3d():
     loss.backward()
     optim.step()
     return loss.realize()
-
+  
   @Tensor.train(mode=False)
   @Tensor.test()
   def eval_step(model, x, y):
@@ -455,7 +455,7 @@ def train_unet3d():
     loss = dice_ce_loss(y_hat, y)
     score = dice_score(y_hat, y)
     return loss.realize(), score.realize()
-
+  
   if WANDB: wandb.init(config=config, project=PROJ_NAME)
 
   step_times, start_epoch = [], 1
@@ -464,7 +464,7 @@ def train_unet3d():
   next_eval_at = start_eval_at
 
   print(f"Training on {GPUS}")
-
+  
   if BENCHMARK: print("Benchmarking UNet3D")
   else: print(f"Start evaluation at epoch {start_eval_at} and every {evaluate_every} epoch(s) afterwards")
 
@@ -572,11 +572,7 @@ def train_rnnt():
   pass
 
 @TinyJit
-def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor,
-                    masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
   optimizer.zero_grad()
 
   lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
@@ -584,7 +580,7 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Te
   (loss * loss_scaler).backward()
 
   global_norm = Tensor([0.0], dtype=dtypes.float32, device=optimizer[0].device).realize()
-  for p in optimizer.params:
+  for p in optimizer.params: 
     p.grad = p.grad / loss_scaler
     global_norm += p.grad.float().square().sum()
   global_norm = global_norm.sqrt()
@@ -592,28 +588,23 @@ def train_step_bert(model, optimizer, scheduler, loss_scaler:float, input_ids:Te
 
   optimizer.step()
   scheduler.step()
-  # TODO: no to("CPU") here because it blocks and messes the python time
-  Tensor.realize(loss, global_norm, optimizer.optimizers[0].lr)
-  return loss, global_norm, optimizer.optimizers[0].lr
+  return loss.realize()
 
 @TinyJit
-def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor,
-                   masked_lm_weights:Tensor, next_sentence_labels:Tensor, GPUS):
-  for t in [input_ids, segment_ids, attention_mask, masked_positions, masked_lm_ids, masked_lm_weights, next_sentence_labels]:
-    if len(GPUS) > 1: t.shard_(GPUS, axis=0)
-    else: t.to_(GPUS[0])
+def eval_step_bert(model, input_ids:Tensor, segment_ids:Tensor, attention_mask:Tensor, masked_positions:Tensor, masked_lm_ids:Tensor, masked_lm_weights:Tensor, next_sentence_labels:Tensor):
   lm_logits, seq_relationship_logits = model(input_ids, attention_mask, masked_positions, segment_ids)
-  masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = \
-    model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
-  for t in [masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss]:
-    t.to_("CPU")
-  Tensor.realize(masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss)
-  return masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss
+  masked_lm_accuracy, seq_relationship_accuracy, masked_lm_loss, next_sentence_loss = model.accuracy(lm_logits, seq_relationship_logits, masked_lm_ids, masked_lm_weights, next_sentence_labels)
+  return {
+    "masked_lm_accuracy": masked_lm_accuracy.realize(),
+    "next_sentence_accuracy": seq_relationship_accuracy.realize(),
+    "masked_lm_loss": masked_lm_loss.realize(),
+    "next_sentence_loss": next_sentence_loss.realize()
+  }
 
 def train_bert():
   # NOTE: pip install tensorflow, wandb required
   from examples.mlperf.dataloader import batch_load_train_bert, batch_load_val_bert
-  from examples.mlperf.helpers import get_mlperf_bert_model, get_fake_data_bert
+  from examples.mlperf.helpers import get_mlperf_bert_model, get_data_bert, get_fake_data_bert
   from examples.mlperf.lr_schedulers import PolynomialDecayWithWarmup
 
   config = {}
@@ -658,9 +649,9 @@ def train_bert():
   # ** hyperparameters **
   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(BS/96))
+  max_lr             = config["OPT_BASE_LEARNING_RATE"] = getenv("OPT_BASE_LEARNING_RATE", 0.0001 * math.sqrt(BS/66))
 
-  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3300000 // BS)
+  train_steps        = config["TRAIN_STEPS"]            = getenv("TRAIN_STEPS", 3630000 // 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 * BS + 3000000) / 25000) * 25000) / BS)) # Round down
@@ -669,7 +660,7 @@ def train_bert():
   save_ckpt_dir      = config["SAVE_CKPT_DIR"]          = getenv("SAVE_CKPT_DIR", "./ckpts")
   init_ckpt          = config["INIT_CKPT_DIR"]          = getenv("INIT_CKPT_DIR", BASEDIR)
 
-  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**11 if dtypes.default_float == dtypes.float16 else 1.0)
+  loss_scaler        = config["LOSS_SCALER"]            = getenv("LOSS_SCALER", 2.0**10 if dtypes.default_float == dtypes.float16 else 1.0)
   decay              = config["DECAY"]                  = getenv("DECAY", 0.01)
   epsilon            = config["EPSILON"]                = getenv("EPSILON", 1e-6)
   poly_power         = config["POLY_POWER"]             = getenv("POLY_POWER", 1.0)
@@ -694,18 +685,11 @@ def train_bert():
 
   # ** init model **
 
-  model = get_mlperf_bert_model()
-  if RUNMLPERF:
-    model.load_from_pretrained(init_ckpt)
-  else:
-    # for init, zero out all weights
-    for p in get_parameters(model):
-      p = p.assign(Tensor.zeros_like(p).contiguous()).realize()
-
+  model = get_mlperf_bert_model(init_ckpt if RUNMLPERF else None)
+  
+  for _, x in get_state_dict(model).items():
+    x.realize().to_(GPUS)
   parameters = get_parameters(model)
-  if len(GPUS) > 1:
-    for p in parameters:
-      p.to_(GPUS)
 
   # ** Log run config **
   for key, value in config.items(): print(f'HParam: "{key}": {value}')
@@ -751,7 +735,7 @@ def train_bert():
   previous_step = None
   if ckpt:=getenv("RESUME", ""):
     load_training_state(model, optimizer_group, scheduler_group, safe_load(ckpt))
-    start_step = int(scheduler_wd.epoch_counter.item())
+    start_step = int(scheduler_wd.epoch_counter.numpy().item())
     print(f"resuming from {ckpt} at step {start_step}")
 
   if RUNMLPERF:
@@ -759,74 +743,70 @@ def train_bert():
     eval_it = iter(batch_load_val_bert(EVAL_BS))
     train_it = iter(tqdm(batch_load_train_bert(BS), total=train_steps, disable=BENCHMARK))
     for _ in range(start_step): next(train_it) # Fast forward
-  else:
-    # repeat fake data
-    def repeat_fake(bs):
-      while True: yield get_fake_data_bert(bs)
-    eval_it = iter(repeat_fake(EVAL_BS))
-    train_it = iter(repeat_fake(BS))
+
 
   step_times = []
   # ** train loop **
   wc_start = time.perf_counter()
-
-  i, train_data = start_step, next(train_it)
-
   if RUNMLPERF:
+    # only load real data with RUNMLPERF
+    i, train_data = start_step, get_data_bert(GPUS, train_it)
     if MLLOGGER:
       MLLOGGER.start(key=mllog_constants.EPOCH_START, value=i*BS, metadata={"epoch_num": i*BS})
+  else:
+    i, train_data = start_step, get_fake_data_bert(GPUS, BS)
 
   while train_data is not None and i < train_steps and not achieved:
-    if getenv("TRAIN", 1):
-      Tensor.training = True
-      BEAM.value = TRAIN_BEAM
-      st = time.perf_counter()
-      GlobalCounters.reset()
-      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)
+    Tensor.training = True
+    BEAM.value = TRAIN_BEAM
+    st = time.perf_counter()
+    GlobalCounters.reset()
+    loss = 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"])
 
-      pt = time.perf_counter()
+    pt = time.perf_counter()
 
-      try:
-        next_data = next(train_it)
-      except StopIteration:
-        next_data = None
+    try:
+      if RUNMLPERF:
+        next_data = get_data_bert(GPUS, train_it)
+      else:
+        next_data = get_fake_data_bert(GPUS, BS)
+    except StopIteration:
+      next_data = None
 
-      dt = time.perf_counter()
+    dt = time.perf_counter()
 
-      device_str = parameters[0].device if isinstance(parameters[0].device, str) else f"{parameters[0].device[0]} * {len(parameters[0].device)}"
-      loss = loss.item()
-      lr = lr.item()
+    device_str = loss.device if isinstance(loss.device, str) else f"{loss.device[0]} * {len(loss.device)}"
+    loss = loss.numpy().item()
 
-      cl = time.perf_counter()
-      if BENCHMARK: step_times.append(cl - st)
+    cl = time.perf_counter()
+    if BENCHMARK: step_times.append(cl - st)
 
-      tqdm.write(
-        f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
-        f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {lr:.6f} LR, "
-        f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
-      if WANDB:
-        wandb.log({"lr": lr, "train/loss": loss, "train/global_norm": global_norm.item(), "train/step_time": cl - st,
-                    "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
-                    "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
+    tqdm.write(
+      f"{i:5} {((cl - st)) * 1000.0:7.2f} ms run, {(pt - st) * 1000.0:7.2f} ms python, {(dt - pt) * 1000.0:6.2f} ms fetch data, "
+      f"{(cl - dt) * 1000.0:7.2f} ms {device_str}, {loss:5.2f} loss, {optimizer_wd.lr.numpy()[0]:.6f} LR, "
+      f"{GlobalCounters.mem_used / 1e9:.2f} GB used, {GlobalCounters.global_ops * 1e-9 / (cl - st):9.2f} GFLOPS")
+    if WANDB:
+      wandb.log({"lr": optimizer_wd.lr.numpy(), "train/loss": loss, "train/step_time": cl - st,
+                  "train/python_time": pt - st, "train/data_time": dt - pt, "train/cl_time": cl - dt,
+                  "train/GFLOPS": GlobalCounters.global_ops * 1e-9 / (cl - st), "epoch": (i+1)*BS})
 
-      train_data, next_data = next_data, None
-      i += 1
+    train_data, next_data = next_data, None
+    i += 1
 
-      if i == BENCHMARK:
-        median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
-        estimated_total_minutes = int(median_step_time * train_steps / 60)
-        print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
-        print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
-              f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")
+    if i == BENCHMARK:
+      median_step_time = sorted(step_times)[(BENCHMARK + 1) // 2]  # in seconds
+      estimated_total_minutes = int(median_step_time * train_steps / 60)
+      print(f"Estimated training time: {estimated_total_minutes // 60}h{estimated_total_minutes % 60}m")
+      print(f"epoch global_ops: {train_steps * GlobalCounters.global_ops:_}, "
+            f"epoch global_mem: {train_steps * GlobalCounters.global_mem:_}")
 
     # ** eval loop **
-    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK) or i == train_steps:
+    if i % eval_step_freq == 0 or (BENCHMARK and i == BENCHMARK):
       if MLLOGGER and RUNMLPERF:
         MLLOGGER.start(key=mllog_constants.EVAL_START, value=None, metadata={"epoch_num": i*BS, "step_num": i})
-      if getenv("RESET_STEP", 0): train_step_bert.reset()
-      elif train_step_bert.captured is not None: train_step_bert.captured.free_intermediates()
+      if getenv("RESET_STEP", 1): train_step_bert.reset()
       eval_lm_losses = []
       eval_clsf_losses = []
       eval_lm_accs = []
@@ -836,14 +816,19 @@ def train_bert():
       BEAM.value = EVAL_BEAM
 
       for j in tqdm(range(max_eval_steps), desc="Evaluating", total=max_eval_steps, disable=BENCHMARK):
-        eval_data = next(eval_it)
+        if RUNMLPERF:
+          eval_data = get_data_bert(GPUS, eval_it)
+        else:
+          eval_data = get_fake_data_bert(GPUS, EVAL_BS)
         GlobalCounters.reset()
         st = time.time()
 
-        lm_acc, clsf_acc, lm_loss, clsf_loss = eval_step_bert(model,
+        eval_result: dict[str, Tensor] = eval_step_bert(model,
           eval_data["input_ids"], eval_data["segment_ids"], eval_data["input_mask"], eval_data["masked_lm_positions"],
-          eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"], GPUS)
-        lm_acc, clsf_acc, lm_loss, clsf_loss = lm_acc.item(), clsf_acc.item(), lm_loss.item(), clsf_loss.item()
+          eval_data["masked_lm_ids"], eval_data["masked_lm_weights"], eval_data["next_sentence_labels"])
+
+        lm_loss, clsf_loss  = eval_result["masked_lm_loss"].item(), eval_result["next_sentence_loss"].item()
+        lm_acc, clsf_acc = eval_result["masked_lm_accuracy"].item(), eval_result["next_sentence_accuracy"].item()
 
         eval_lm_losses.append(lm_loss)
         eval_clsf_losses.append(clsf_loss)
@@ -859,9 +844,8 @@ def train_bert():
             MLLOGGER.event(key=mllog_constants.INIT_STOP, value=None)
           return
 
-      if getenv("RESET_STEP", 0): eval_step_bert.reset()
-      elif eval_step_bert.captured is not None: eval_step_bert.captured.free_intermediates()
-      del eval_data
+      if getenv("RESET_STEP", 1): eval_step_bert.reset()
+      del eval_data, eval_result
       avg_lm_loss = sum(eval_lm_losses) / len(eval_lm_losses)
       avg_clsf_loss = sum(eval_clsf_losses) / len(eval_clsf_losses)
       avg_lm_acc = sum(eval_lm_accs) / len(eval_lm_accs)

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

@@ -2,11 +2,10 @@
 
 export PYTHONPATH="."
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"
 
 export BENCHMARK=10 DEBUG=2

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

@@ -2,9 +2,9 @@
 
 export PYTHONPATH="."
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 

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

@@ -3,9 +3,9 @@
 export PYTHONPATH="."
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_green"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=24
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=4 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=4 BEAM_UOPS_MAX=2000 BEAM_UPCAST_MAX=64 BEAM_LOCAL_MAX=512
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 
@@ -17,7 +17,7 @@ DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_green_${DATETIME}_${SEED}.log"
 
 # init
-BENCHMARK=10 INITMLPERF=1 BEAM_LOG_SURPASS_MAX=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
 
 # run
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE

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

@@ -2,14 +2,12 @@
 
 export PYTHONPATH="."
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=3
 export IGNORE_JIT_FIRST_BEAM=1
-export BEAM_LOG_SURPASS_MAX=1
 export BASEDIR="/raid/datasets/wiki"
 
-export RESET_STEP=1
 export BENCHMARK=10 DEBUG=2
 
 python3 examples/mlperf/model_train.py

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

@@ -2,9 +2,9 @@
 
 export PYTHONPATH="."
 export MODEL="bert"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=3
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 

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

@@ -3,9 +3,9 @@
 export PYTHONPATH="."
 export MODEL="bert"
 export SUBMISSION_PLATFORM="tinybox_red"
-export DEFAULT_FLOAT="HALF" SUM_DTYPE="HALF" GPUS=6 BS=96 EVAL_BS=96
+export DEFAULT_FLOAT="HALF" GPUS=6 BS=72 EVAL_BS=6
 
-export BEAM=3 BEAM_UOPS_MAX=4000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5
+export BEAM=3
 export IGNORE_JIT_FIRST_BEAM=1
 export BASEDIR="/raid/datasets/wiki"
 
@@ -17,11 +17,7 @@ DATETIME=$(date "+%m%d%H%M")
 LOGFILE="bert_red_${DATETIME}_${SEED}.log"
 
 # init
-sudo rmmod amdgpu || true
-BENCHMARK=10 INITMLPERF=1 BEAM_LOG_SURPASS_MAX=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
+BENCHMARK=10 INITMLPERF=1 python3 examples/mlperf/model_train.py | tee $LOGFILE
 
 # run
-# TODO: AMD driver hangs during init, but is 5% faster per step in real run.
-sudo modprobe amdgpu
 PARALLEL=0 RUNMLPERF=1 python3 examples/mlperf/model_train.py | tee -a $LOGFILE
-sudo rmmod amdgpu || true

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

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

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

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

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

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

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

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

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

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

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

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

examples/mnist_gan.py

@@ -16,9 +16,9 @@ class LinearGen:
     self.l4 = Tensor.scaled_uniform(1024, 784)
 
   def forward(self, x):
-    x = x.dot(self.l1).leaky_relu(0.2)
-    x = x.dot(self.l2).leaky_relu(0.2)
-    x = x.dot(self.l3).leaky_relu(0.2)
+    x = x.dot(self.l1).leakyrelu(0.2)
+    x = x.dot(self.l2).leakyrelu(0.2)
+    x = x.dot(self.l3).leakyrelu(0.2)
     x = x.dot(self.l4).tanh()
     return x
 
@@ -31,9 +31,9 @@ class LinearDisc:
 
   def forward(self, x):
     # balance the discriminator inputs with const bias (.add(1))
-    x = x.dot(self.l1).add(1).leaky_relu(0.2).dropout(0.3)
-    x = x.dot(self.l2).leaky_relu(0.2).dropout(0.3)
-    x = x.dot(self.l3).leaky_relu(0.2).dropout(0.3)
+    x = x.dot(self.l1).add(1).leakyrelu(0.2).dropout(0.3)
+    x = x.dot(self.l2).leakyrelu(0.2).dropout(0.3)
+    x = x.dot(self.l3).leakyrelu(0.2).dropout(0.3)
     x = x.dot(self.l4).log_softmax()
     return x
 

examples/olmoe.py

@@ -1,94 +0,0 @@
-# https://arxiv.org/pdf/2409.02060
-import time
-import numpy as np
-np.set_printoptions(suppress=True, linewidth=1000)
-import functools
-from tinygrad import Tensor, nn, Device, GlobalCounters
-from tinygrad.helpers import Timing, getenv
-from extra.models.llama import Transformer, convert_from_huggingface
-
-class MixtureFeedForward:
-  def __init__(self, num_experts:int, activated_experts:int, dim:int, hidden_dim:int, linear=nn.Linear):
-    self.activated_experts = activated_experts
-    self.gate = nn.Linear(dim, num_experts, bias=False)
-    self.up_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
-    self.down_proj = Tensor.zeros(num_experts, dim, hidden_dim, dtype='bfloat16')
-    self.gate_proj = Tensor.zeros(num_experts, hidden_dim, dim, dtype='bfloat16')
-  def __call__(self, x:Tensor) -> Tensor:
-    assert x.shape[0] == 1, "only BS=1"
-    assert x.shape[1] == 1, "only length=1"
-    g = self.gate(x).float().softmax(-1)
-
-    g = g.squeeze() # (BS, length, num_experts) -> (num_experts,)
-    probs, sel = g.topk(self.activated_experts)
-
-    # run MoE
-    x_up_gate = x.dot(self.gate_proj[sel].permute(0,2,1)).silu() * x.dot(self.up_proj[sel].permute(0,2,1))
-    x_down = x_up_gate.dot(self.down_proj[sel].permute(0,2,1))
-    return (x_down.float() * probs.reshape(self.activated_experts, 1, 1)).sum(axis=0)
-
-# model is bf16, 1.3B active, 6.9B total
-# M3 Max is 400 GB/s, so 400/2.6 = ~154 tok/s
-
-def fetch_weights() -> dict[str, Tensor]:
-  # TODO: make this lazy so the 3 fetches can happen in parallel
-  m1 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00001-of-00003.safetensors").to(Device.DEFAULT)
-  m2 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00002-of-00003.safetensors").to(Device.DEFAULT)
-  m3 = Tensor.from_url("https://huggingface.co/allenai/OLMoE-1B-7B-0924/resolve/main/model-00003-of-00003.safetensors").to(Device.DEFAULT)
-  return {**nn.state.safe_load(m1), **nn.state.safe_load(m2), **nn.state.safe_load(m3)}
-
-if __name__ == "__main__":
-  if getenv("TORCH"):
-    from transformers import OlmoeForCausalLM, AutoTokenizer
-    model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924")
-    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
-    inputs = tokenizer("Hello", return_tensors="pt")
-    generate_ids = model.generate(inputs.input_ids, max_length=30)
-    out = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    print(out)
-    exit(0)
-
-  with Timing("create model: "):
-    model = Transformer(n_layers=16, dim=2048, hidden_dim=1024, n_heads=16, norm_eps=1e-5, qk_norm=1e-5, max_context=1024,
-                        vocab_size=50304, feed_forward=functools.partial(MixtureFeedForward, 64, 8))
-    model_state_dict = nn.state.get_state_dict(model)
-    del model_state_dict['freqs_cis']
-
-  with Timing("load weights to GPU: "):
-    nhf_state = convert_from_huggingface(fetch_weights(), model, 16, 16)
-    # NOTE: i'm not sure this actually needs float32, it may just change the type of things downstream from it. but doesn't match torch w/o this
-    for needs_float32 in ['tok_embeddings.weight']: nhf_state[needs_float32] = nhf_state[needs_float32].float()
-  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
-
-  with Timing("unpack weights: "):
-    nn.state.load_state_dict(model, nhf_state, verbose=False, strict=False, consume=True, realize=False)
-    assert len(nhf_state) == 0
-    Tensor.realize(*list(nn.state.get_state_dict(model).values()))
-  print(f"ram used: {GlobalCounters.mem_used/1e9:.2f} GB")
-
-  count = 30
-  temperature = 0
-
-  with Timing("load tokenizer: "):
-    from transformers import AutoTokenizer
-    tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924")
-
-  toks = [12092]
-  start_pos = 0
-  timings = []
-  for i in range(count):
-    GlobalCounters.reset()
-    st = time.perf_counter()
-    tok = model(Tensor([toks[start_pos:]]), start_pos, temperature).item()
-    timings.append(time.perf_counter()-st)
-    toks.append(tok)
-    start_pos += 1
-    print(toks)
-    print(tokenizer.decode(toks))
-  print(f"fastest token {min(timings)*1e3:.2f} ms, {1/min(timings):.1f} tok/s")
-
-  if temperature == 0:
-    # Hello, I am a newbie to this forum and I am trying to get a better understanding of the different types of data that can be stored in a
-    assert toks == [12092, 13, 309, 717, 247, 747, 17782, 281, 436, 12209, 285, 309, 717, 2820, 281, 755,
-                    247, 1805, 4685, 273, 253, 1027, 3510, 273, 941, 326, 476, 320, 7141, 275, 247], "BAD OUTPUT!"
-

examples/openpilot/compile2.py

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

examples/openpilot/compile3.py

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

examples/openpilot/go.sh

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

examples/other_mnist/beautiful_mnist_torch.py

@@ -1,5 +1,5 @@
-from tinygrad import dtypes, getenv, Device
-from tinygrad.helpers import trange, colored, DEBUG, temp
+from tinygrad import dtypes
+from tinygrad.helpers import trange
 from tinygrad.nn.datasets import mnist
 import torch
 from torch import nn, optim
@@ -26,20 +26,14 @@ class Model(nn.Module):
     return self.lin(torch.flatten(x, 1))
 
 if __name__ == "__main__":
-  if getenv("TINY_BACKEND"):
-    import tinygrad.frontend.torch
-    device = torch.device("tiny")
-  else:
-    device = torch.device({"METAL":"mps","NV":"cuda"}.get(Device.DEFAULT, "cpu"))
-  if DEBUG >= 1: print(f"using torch backend {device}")
+  mps_device = torch.device("mps")
   X_train, Y_train, X_test, Y_test = mnist()
-  X_train = torch.tensor(X_train.float().numpy(), device=device)
-  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=device)
-  X_test = torch.tensor(X_test.float().numpy(), device=device)
-  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=device)
+  X_train = torch.tensor(X_train.float().numpy(), device=mps_device)
+  Y_train = torch.tensor(Y_train.cast(dtypes.int64).numpy(), device=mps_device)
+  X_test = torch.tensor(X_test.float().numpy(), device=mps_device)
+  Y_test = torch.tensor(Y_test.cast(dtypes.int64).numpy(), device=mps_device)
 
-  if getenv("TORCHVIZ"): torch.cuda.memory._record_memory_history()
-  model = Model().to(device)
+  model = Model().to(mps_device)
   optimizer = optim.Adam(model.parameters(), 1e-3)
 
   loss_fn = nn.CrossEntropyLoss()
@@ -59,11 +53,3 @@ if __name__ == "__main__":
     loss = step(samples)
     if i%10 == 9: test_acc = ((model(X_test).argmax(axis=-1) == Y_test).sum() * 100 / X_test.shape[0]).item()
     t.set_description(f"loss: {loss.item():6.2f} test_accuracy: {test_acc:5.2f}%")
-
-  # verify eval acc
-  if target := getenv("TARGET_EVAL_ACC_PCT", 0.0):
-    if test_acc >= target and test_acc != 100.0: print(colored(f"{test_acc=} >= {target}", "green"))
-    else: raise ValueError(colored(f"{test_acc=} < {target}", "red"))
-  if getenv("TORCHVIZ"):
-    torch.cuda.memory._dump_snapshot(fp:=temp("torchviz.pkl", append_user=True))
-    print(f"saved torch memory snapshot to {fp}, view in https://pytorch.org/memory_viz")

examples/sdxl.py

@@ -3,10 +3,10 @@
 # Stability-AI/generative-models | MIT     | https://github.com/Stability-AI/generative-models/blob/fbdc58cab9f4ee2be7a5e1f2e2787ecd9311942f/LICENSE-CODE
 # mlfoundations/open_clip        | MIT     | https://github.com/mlfoundations/open_clip/blob/58e4e39aaabc6040839b0d2a7e8bf20979e4558a/LICENSE
 
-from tinygrad import Tensor, TinyJit, dtypes, GlobalCounters
+from tinygrad import Tensor, TinyJit, dtypes
 from tinygrad.nn import Conv2d, GroupNorm
 from tinygrad.nn.state import safe_load, load_state_dict
-from tinygrad.helpers import fetch, trange, colored, Timing
+from tinygrad.helpers import fetch, trange, colored, Timing, GlobalCounters
 from extra.models.clip import Embedder, FrozenClosedClipEmbedder, FrozenOpenClipEmbedder
 from extra.models.unet import UNetModel, Upsample, Downsample, timestep_embedding
 from examples.stable_diffusion import ResnetBlock, Mid
@@ -345,7 +345,6 @@ class DPMPP2MSampler:
     old_denoised = None
     for i in trange(num_sigmas - 1):
       with Timing("step in ", enabled=timing, on_exit=lambda _: f", using {GlobalCounters.mem_used/1e9:.2f} GB"):
-        GlobalCounters.reset()
         x, old_denoised = self.sampler_step(
           old_denoised=old_denoised,
           prev_sigma=(None if i==0 else sigmas[i-1].expand(x.shape[0])),

examples/self_tokenize.py

@@ -1,48 +0,0 @@
-import os, pathlib, argparse
-from examples.llama3 import Tokenizer
-from tabulate import tabulate
-from tinygrad import fetch
-from tinygrad.helpers import flatten
-
-# llama 3 tokenizer
-tokenizer = Tokenizer(fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model").as_posix())
-
-def read_code(base_path):
-  ret = []
-  for path, _, files in os.walk(os.path.join(base_path, "tinygrad")):
-    for name in files:
-      if not name.endswith(".py"): continue
-      if 'tinygrad/runtime/autogen' in path.replace('\\', '/'): continue
-      fullpath = os.path.join(path, name)
-      code = pathlib.Path(fullpath).read_text()
-      ret.append(("### " + fullpath.split("tinygrad/", 1)[1], code))
-  return ret
-
-def write_code_to_file(filename, code_list):
-  """Writes the combined code to a specified file."""
-  with open(filename, 'w') as f:
-    f.write('\n'.join(flatten(code_list)))
-
-if __name__ == "__main__":
-  parser = argparse.ArgumentParser(description="Analyze and optionally save tinygrad code.")
-  parser.add_argument("--output", help="Output file to write the combined code to.")
-  args = parser.parse_args()
-
-  ret = read_code(".")
-
-  table = []
-  for name,code in ret:
-    table.append([name, len(tokenizer.encode(name+"\x00"+code))])
-  print(tabulate([["name", "llm tokens"]]+sorted(table, key=lambda x: -x[1]), headers="firstrow"))
-
-  code_str = '\x00'.join(flatten(ret))
-  print(f"code has {len(code_str)} chars")
-  newline_count = code_str.count('\n')
-  print(f"code has {newline_count} newlines")
-
-  encoded = tokenizer.encode(code_str)
-  print(f"code has {len(encoded)} tokens")
-
-  if args.output:
-    write_code_to_file(args.output, ret)
-    print(f"Combined code written to {args.output}")

examples/so_vits_svc.py

@@ -437,14 +437,14 @@ class Generator:
     x = self.conv_pre(x)
     if g is not None:  x = x + self.cond(g)
     for i in range(self.num_upsamples):
-      x, xs = self.ups[i](x.leaky_relu(LRELU_SLOPE)), None
+      x, xs = self.ups[i](x.leakyrelu(LRELU_SLOPE)), None
       x_source = self.noise_convs[i](har_source)
       x = x + x_source
       for j in range(self.num_kernels):
         if xs is None: xs = self.resblocks[i * self.num_kernels + j].forward(x)
         else: xs += self.resblocks[i * self.num_kernels + j].forward(x)
       x = xs / self.num_kernels
-    return self.conv_post(x.leaky_relu()).tanh()
+    return self.conv_post(x.leakyrelu()).tanh()
 
 # **** helpers ****
 
@@ -504,7 +504,7 @@ def load_checkpoint_enc(checkpoint_path, model: ContentVec, optimizer=None, skip
         obj, v = getattr(parent, "weight"), weight_norm(weight_v, weight_g, 0)
         weight_g, weight_v, parent, skip = None, None, None, False
       if not skip and obj.shape == v.shape:
-        if "feature_extractor" in key and (isinstance(parent, (nn.GroupNorm, nn.LayerNorm))):  # cast
+        if "feature_extractor" in key and (isinstance(parent, nn.GroupNorm) or isinstance(parent, nn.LayerNorm)):  # cast
           obj.assign(v.to(obj.device).float())
         else:
           obj.assign(v.to(obj.device))

examples/test_f16.py

@@ -0,0 +1,5 @@
+from tinygrad import Tensor, dtypes
+
+if __name__ == "__main__":
+    a = Tensor([1.0,2.0,3.0], dtype=dtypes.half)
+    print((a*2.0).numpy())

examples/test_onnx_imagenet.py

@@ -1,81 +0,0 @@
-import random, sys
-import numpy as np
-from extra.datasets.imagenet import get_imagenet_categories, get_val_files, center_crop
-from examples.benchmark_onnx import load_onnx_model
-from PIL import Image
-from tinygrad import Tensor, dtypes, GlobalCounters
-from tinygrad.helpers import fetch, getenv
-
-# works:
-#  ~70% - https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/resnet/model/resnet50-v2-7.onnx
-#  ~43% - https://github.com/onnx/models/raw/refs/heads/main/Computer_Vision/alexnet_Opset16_torch_hub/alexnet_Opset16.onnx
-#  ~72% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-#  ~71% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenetv2_1.0.opt.onnx
-#  ~67% - https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7-quantized.onnx
-# broken:
-#  https://github.com/MTlab/onnx2caffe/raw/refs/heads/master/model/MobileNetV2.onnx
-#  https://huggingface.co/qualcomm/MobileNet-v2-Quantized/resolve/main/MobileNet-v2-Quantized.onnx
-#  ~35% - https://github.com/axinc-ai/onnx-quantization/raw/refs/heads/main/models/mobilenev2_quantized.onnx
-
-# QUANT=1 python3 examples/test_onnx_imagenet.py
-#   https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx
-# DONT_REALIZE_EXPAND=1 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
-# VIZ=1 DONT_REALIZE_EXPAND=1 python3 examples/benchmark_onnx.py /tmp/model.quant.onnx
-
-def imagenet_dataloader(cnt=0):
-  input_mean = Tensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)
-  input_std = Tensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)
-  files = get_val_files()
-  random.shuffle(files)
-  if cnt != 0: files = files[:cnt]
-  cir = get_imagenet_categories()
-  for fn in files:
-    img = Image.open(fn)
-    img = img.convert('RGB') if img.mode != "RGB" else img
-    img = center_crop(img)
-    img = np.array(img)
-    img = Tensor(img).permute(2,0,1).reshape(1,3,224,224)
-    img = ((img.cast(dtypes.float32)/255.0) - input_mean) / input_std
-    y = cir[fn.split("/")[-2]]
-    yield img,y
-
-if __name__ == "__main__":
-  fn = sys.argv[1]
-  if getenv("QUANT"):
-    from onnxruntime.quantization import quantize_dynamic, quantize_static, QuantFormat, QuantType, CalibrationDataReader
-    model_fp32 = fetch(fn)
-    fn = '/tmp/model.quant.onnx'
-    if getenv("DYNAMIC"):
-      quantize_dynamic(model_fp32, fn)
-    else:
-      class ImagenetReader(CalibrationDataReader):
-        def __init__(self):
-          self.iter = imagenet_dataloader(cnt=1000)
-        def get_next(self) -> dict:
-          try:
-            img,y = next(self.iter)
-          except StopIteration:
-            return None
-          return {"input": img.numpy()}
-      quantize_static(model_fp32, fn, ImagenetReader(), quant_format=QuantFormat.QDQ, per_channel=False,
-                      activation_type=QuantType.QUInt8, weight_type=QuantType.QUInt8,
-                      extra_options={"ActivationSymmetric": False})
-
-  run_onnx_jit, input_specs = load_onnx_model(fetch(fn))
-  t_name, t_spec = list(input_specs.items())[0]
-  assert t_spec.shape[1:] == (3,224,224), f"shape is {t_spec.shape}"
-
-  hit = 0
-  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
-    GlobalCounters.reset()
-    p = run_onnx_jit(**{t_name:img})
-    assert p.shape == (1,1000)
-    t = p.to('cpu').argmax().item()
-    hit += y==t
-    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")
-
-  MS_TARGET = 13.4
-  print(f"need {GlobalCounters.global_ops/1e9*(1000/MS_TARGET):.2f} GFLOPS for {MS_TARGET:.2f} ms")
-
-  import pickle
-  with open("/tmp/im.pkl", "wb") as f: pickle.dump(run_onnx_jit, f)

examples/test_pkl_imagenet.py

@@ -1,19 +0,0 @@
-import sys, pickle
-from tinygrad import GlobalCounters
-from tinygrad.helpers import fetch, getenv
-from examples.test_onnx_imagenet import imagenet_dataloader
-
-if __name__ == "__main__":
-  with open(fetch(sys.argv[1]), "rb") as f:
-    run_onnx_jit = pickle.load(f)
-  input_name = run_onnx_jit.captured.expected_names[0]
-  device = run_onnx_jit.captured.expected_st_vars_dtype_device[0][-1]
-  print(f"input goes into {input_name=} on {device=}")
-  hit = 0
-  for i,(img,y) in enumerate(imagenet_dataloader(cnt=getenv("CNT", 100))):
-    GlobalCounters.reset()
-    p = run_onnx_jit(**{input_name:img.to(device)})
-    assert p.shape == (1,1000)
-    t = p.argmax().item()
-    hit += y==t
-    print(f"target: {y:3d}  pred: {t:3d}  acc: {hit/(i+1)*100:.2f}%")

examples/tinychat/assets/cdn.jsdelivr.net/npm/purecss@3.0.0/build/base-min.css

@@ -1,11 +0,0 @@
-/*!
-Pure v3.0.0
-Copyright 2013 Yahoo!
-Licensed under the BSD License.
-https://github.com/pure-css/pure/blob/master/LICENSE
-*/
-/*!
-normalize.css v | MIT License | https://necolas.github.io/normalize.css/
-Copyright (c) Nicolas Gallagher and Jonathan Neal
-*/
-/*! normalize.css v8.0.1 | MIT License | github.com/necolas/normalize.css */html{line-height:1.15;-webkit-text-size-adjust:100%}body{margin:0}main{display:block}h1{font-size:2em;margin:.67em 0}hr{box-sizing:content-box;height:0;overflow:visible}pre{font-family:monospace,monospace;font-size:1em}a{background-color:transparent}abbr[title]{border-bottom:none;text-decoration:underline;-webkit-text-decoration:underline dotted;text-decoration:underline dotted}b,strong{font-weight:bolder}code,kbd,samp{font-family:monospace,monospace;font-size:1em}small{font-size:80%}sub,sup{font-size:75%;line-height:0;position:relative;vertical-align:baseline}sub{bottom:-.25em}sup{top:-.5em}img{border-style:none}button,input,optgroup,select,textarea{font-family:inherit;font-size:100%;line-height:1.15;margin:0}button,input{overflow:visible}button,select{text-transform:none}[type=button],[type=reset],[type=submit],button{-webkit-appearance:button}[type=button]::-moz-focus-inner,[type=reset]::-moz-focus-inner,[type=submit]::-moz-focus-inner,button::-moz-focus-inner{border-style:none;padding:0}[type=button]:-moz-focusring,[type=reset]:-moz-focusring,[type=submit]:-moz-focusring,button:-moz-focusring{outline:1px dotted ButtonText}fieldset{padding:.35em .75em .625em}legend{box-sizing:border-box;color:inherit;display:table;max-width:100%;padding:0;white-space:normal}progress{vertical-align:baseline}textarea{overflow:auto}[type=checkbox],[type=radio]{box-sizing:border-box;padding:0}[type=number]::-webkit-inner-spin-button,[type=number]::-webkit-outer-spin-button{height:auto}[type=search]{-webkit-appearance:textfield;outline-offset:-2px}[type=search]::-webkit-search-decoration{-webkit-appearance:none}::-webkit-file-upload-button{-webkit-appearance:button;font:inherit}details{display:block}summary{display:list-item}template{display:none}[hidden]{display:none}html{font-family:sans-serif}.hidden,[hidden]{display:none!important}.pure-img{max-width:100%;height:auto;display:block}

examples/tinychat/fetch_assets.sh

@@ -9,15 +9,11 @@ fetch "cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"
 fetch "cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"
 fetch "cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"
 fetch "unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"
-fetch "unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js.map"
 fetch "unpkg.com/[email protected]/dist/cdn.min.js"
 fetch "unpkg.com/[email protected]/dist/purify.min.js"
-fetch "unpkg.com/[email protected]/dist/purify.min.js.map"
 fetch "unpkg.com/[email protected]/marked.min.js"
 fetch "unpkg.com/[email protected]/lib/index.umd.js"
 fetch "unpkg.com/@highlightjs/[email protected]/highlight.min.js"
 fetch "cdn.jsdelivr.net/npm/[email protected]/build/base-min.css"
 fetch "cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
 fetch "unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css"
-fetch "cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/webfonts/fa-solid-900.ttf"
-fetch "cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/webfonts/fa-solid-900.woff2"

examples/tinychat/tinychat-browser/.gitignore

@@ -1,5 +0,0 @@
-net_*
-llama3-2.tiktoken
-tiktoken.js
-tiktoken_bg.wasm
-transformer*

examples/tinychat/tinychat-browser/README.md

@@ -1,8 +0,0 @@
-# How to build and run tinychat in browser (WebGPU and WASM)
-- `PYTHONPATH=. python examples/tinychat/tinychat-browser/compile.py`
-- `./examples/tinychat/tinychat-browser/compile_wasm.sh`
-    - Prerequisite: [install emscripten](https://emscripten.org/docs/getting_started/downloads.html). This script looks for `~/emsdk/emsdk_env.sh`, adjust this based on your installation.
-- `./examples/tinychat/tinychat-browser/make_tiktoken_js.sh`
-    - Prerequisite: install `npm`, `webpack`.
-- `cd examples/tinychat && python -m http.server 7776`
-- In browser: open either `localhost:7776/tinychat-browser` (WebGPU), or `localhost:7776/tinychat-browser/?backend=wasm` (WASM)

examples/tinychat/tinychat-browser/compile.py

@@ -1,149 +0,0 @@
-import os, json, hashlib, math
-from extra.export_model import export_model
-from examples.llama3 import build_transformer, Tokenizer
-from tinygrad.nn.state import get_state_dict, load_state_dict
-from tinygrad import Device, Variable, Tensor, dtypes, TinyJit
-from tinygrad.helpers import fetch, Context
-from tiktoken.load import load_tiktoken_bpe, dump_tiktoken_bpe
-
-def prepare_browser_chunks(model):
-  # split weights into browser-friendly chunks
-  state_dict = get_state_dict(model)
-  del state_dict['output.weight'], state_dict['output.scale'] # same as tok_embeddings; ensures consistency with model export
-  chunk_size = 16 * 1024 * 1024 # small chunks based on iphone browser constraints
-  metadata = {}
-  # We won't export cache_kv bytes (because we start inference on client at start_pos=0), but we will tell the client how big cache_kv needs to be
-  t_infos = [(v.lazydata.base.realized.nbytes, k, v.dtype) for k,v in state_dict.items() if "cache_kv" not in k]
-  empty_t_infos = [(v.lazydata.base.realized.nbytes, k, v.dtype) for k,v in state_dict.items() if "cache_kv" in k]
-
-  split_t_infos = []
-  for size, name, dtype in t_infos:
-    if size <= chunk_size:
-      split_t_infos.append((size, name, dtype, ()))
-    else: # split large weights into multiple parts
-      for i in range(0, size, chunk_size):
-        split_t_infos.append((min(chunk_size, size-i), f"{name}_part{math.ceil(i/chunk_size)}", dtype, (i, min(i+chunk_size, size))))
-
-  files = []
-  # pack weights into files with FFD bin packing
-  split_t_infos = sorted(split_t_infos, reverse=True)
-  for info in split_t_infos:
-    placed = False
-    for file in files:
-      if sum(i[0] for i in file) + info[0] <= chunk_size:
-        if info[3] and any(i[3] for i in file): continue # no two split tensors can touch the same file, due to wasm loading constraints
-        file.append(info)
-        placed = True
-        break
-    if not placed:
-      files.append([info])
-
-  tinygrad_dtypes = {dtypes.float32: "float32", dtypes.float16: "float16", dtypes.int8: "int8", dtypes.int32: "int32"}
-  for i, file in enumerate(files):
-    cursor = 0
-    with open(os.path.join(os.path.dirname(__file__), f'./net_part{i}.chunk'), "wb+") as writer:
-      for size, name, dtype, offsets in file:
-        name, part_num = (name, 0) if "_part" not in name else (name.split("_part")[0], int(name.split("_part")[1]))
-        default = {"parts": {}, "dtype": tinygrad_dtypes[dtype]}
-        weight_metadata = metadata.get(name, default)
-        weight_metadata["parts"][part_num] = {"file": i, "file_start_pos": cursor, "size": size}
-        metadata[name] = weight_metadata
-        data = bytes(state_dict[name].lazydata.base.realized.as_buffer())
-        data = data if not offsets else data[offsets[0]:offsets[1]]
-        writer.write(data)
-        cursor += size
-
-  metadata.update({name: {"parts": {0: {"empty": True, "size": size}}, "dtype": tinygrad_dtypes[dtype]} for size, name, dtype in empty_t_infos})
-
-  for k in metadata:
-    metadata[k]["parts"] = [part for part_num, part in sorted(metadata[k]["parts"].items(), key = lambda x: x[0])]
-    cursor = 0
-    for i, part in enumerate(metadata[k]["parts"]):
-      metadata[k]["parts"][i]["target_start_pos"] = cursor
-      cursor += part["size"]
-    metadata[k]["size"] = cursor
-
-  # compute hashes, which client app will check to determine whether to update with new weights and/or detect integrity issues
-  state_dict_hash = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode("utf-8")).hexdigest()
-  metadata = {"state_dict": metadata, "state_dict_hash": state_dict_hash, "files": []}
-  hashes = set()
-  for i in range(len(files)):
-    with open(os.path.join(os.path.dirname(__file__), f'./net_part{i}.chunk'), "rb") as reader:
-      hash = hashlib.sha256(reader.read()).hexdigest()
-      hashes.add(hash)
-      metadata["files"].append({"name": f'net_part{i}.chunk', "hash": hash})
-  if len(hashes) != len(files): print(f"WARNING: {len(files)} files were exported, but only {len(hashes)} are unique: something may have gone wrong")
-  metadata_hash = hashlib.sha256(json.dumps(metadata, sort_keys=True).encode("utf-8")).hexdigest()
-  metadata = {"metadata": metadata, "metadata_hash": metadata_hash}
-
-  with open(os.path.join(os.path.dirname(__file__), f'./net_metadata.json'), "w") as writer: json.dump(metadata, writer, indent=4)
-  return metadata
-
-def validate_model(model, tokenizer):
-  prompt = "yo"
-  toks = [tokenizer.bos_id]
-  toks += [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode("user") + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
-  toks += tokenizer.encode(prompt) + [tokenizer.special_tokens["<|eot_id|>"]]
-  toks += [tokenizer.special_tokens["<|start_header_id|>"]] + tokenizer.encode("assistant") + [tokenizer.special_tokens["<|end_header_id|>"]] + tokenizer.encode("\n\n")
-  start_pos = 0
-  run = TinyJit(model.forward)
-  for tok in toks[:-1]:
-    run(Tensor([[tok]]), Variable("start_pos", 0, model.max_context).bind(start_pos), 0.0, 0, 0.0, 0.0, 0.0).realize()
-    start_pos += 1
-  tok = toks[-1]
-  result = ""
-  expected = "How's it going?"
-  while True:
-    tok = run(Tensor([[tok]]), Variable("start_pos", 0, model.max_context).bind(start_pos), 0.0, 0, 0.0, 0.0, 0.0).item()
-    start_pos += 1
-    if tok in tokenizer.stop_tokens or len(result) > len(expected): break
-    result += tokenizer.decode([tok])
-  assert result == expected, f"Model validation failed, expected output: {expected}, actual output: {result}"
-
-if __name__=="__main__":
-  # Export BPE data for use with tiktoken.js
-  tokenizer_path = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model", "tokenizer.model", subdir="llama3-1b-instruct")
-  mergeable_ranks = load_tiktoken_bpe(str(tokenizer_path))
-  bpe_path = os.path.join(os.path.dirname(__file__), "llama3-2.tiktoken")
-  dump_tiktoken_bpe(mergeable_ranks, bpe_path)
-  tokenizer = Tokenizer(str(tokenizer_path))
-
-  model_path = fetch("https://huggingface.co/bartowski/Llama-3.2-1B-Instruct-GGUF/resolve/main/Llama-3.2-1B-Instruct-f16.gguf", "Llama-3.2-1B-Instruct-f16.gguf", subdir="llama3-1b-instruct")
-  Tensor.no_grad = True
-  max_context=1024
-  tok = 128000
-  TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P = 0.95, 0, 0.0, 0.0, 0.0
-  start_pos = Variable("start_pos", 0, max_context).bind(0)
-  model_input = lambda: [Tensor([[tok]]), start_pos, TEMPERATURE, TOP_K, TOP_P, ALPHA_F, ALPHA_P]
-
-  Device.DEFAULT="CPU"
-  model = build_transformer(model_path, model_size="1B", quantize="int8", scale_dtype=dtypes.float32, device=Device.DEFAULT, max_context=max_context)
-  state_dict = get_state_dict(model)
-  validate_model(model, tokenizer)
-  model_name = "transformer"
-
-  with Context(BEAM=3):
-    cprog, js_wrapper = export_model(model, "wasm", *model_input(), model_name=model_name)
-    # ensure consistency with exported weights
-    js_wrapper = js_wrapper.replace("output.weight", "tok_embeddings.weight").replace("output.scale", "tok_embeddings.scale")
-
-  with open(os.path.join(os.path.dirname(__file__), f"{model_name}.c"), "w") as f: f.write(cprog)
-  with open(os.path.join(os.path.dirname(__file__), "net_clang.js"), "w") as f: f.write(js_wrapper)
-
-  Device.DEFAULT="WEBGPU"
-  # float16 is not yet supported for dawn/Vulkan/NVIDIA stack, see: https://issues.chromium.org/issues/42251215
-  # therefore for now, we used CLANG to quantize the float16 llama to int8 with float32 scales, then load to WEBGPU
-  model = build_transformer(model_path, model_size="1B", quantize="int8", max_context=max_context, load_weights=False)
-  load_state_dict(model, state_dict)
-  # these were the same before load_state_dict
-  model.output.weight, model.output.scale = model.tok_embeddings.weight, model.tok_embeddings.scale
-
-  validate_model(model, tokenizer)
-  metadata = prepare_browser_chunks(model) # export weights to disk
-
-  with Context(BEAM=3):
-    prg, input_sizes, output_sizes, state = export_model(model, "webgpu", *model_input(), model_name=model_name, stream_weights=True)
-    # ensure consistency with exported weights
-    prg = prg.replace("output.weight", "tok_embeddings.weight").replace("output.scale", "tok_embeddings.scale")
-
-  with open(os.path.join(os.path.dirname(__file__), "net.js"), "w") as f: f.write(prg)

examples/tinychat/tinychat-browser/compile_wasm.sh

@@ -1,23 +0,0 @@
-#!/usr/bin/env bash
-cd "$(dirname "$0")"
-
-# prereq: install emscripten: https://emscripten.org/docs/getting_started/downloads.html
-EMSCRIPTEN_PATH=~/emsdk/emsdk_env.sh
-source $EMSCRIPTEN_PATH
-step="transformer"
-initial_memory=6553600
-max_memory=1500053504
-exported_functions='["_net", "_malloc", "_free", "_set_buf"]'
-
-emcc "${step}.c" \
-  -O3 -msimd128 -ffast-math -flto \
-  -o "${step}.js" \
-  -s MODULARIZE=1 \
-  -s EXPORT_ES6=1 \
-  -s EXPORTED_FUNCTIONS="${exported_functions}" \
-  -s ENVIRONMENT='worker' \
-  -s FILESYSTEM=0 \
-  -s EVAL_CTORS \
-  -s ALLOW_MEMORY_GROWTH=1 \
-  -s INITIAL_MEMORY="$initial_memory" \
-  -s MAXIMUM_MEMORY="$max_memory"

examples/tinychat/tinychat-browser/index.css

@@ -1,322 +0,0 @@
-/* define colors */
-:root {
-  --primary-color: #fff;
-  --secondary-color: #2a2a2a;
-  --secondary-color-transparent: #ffffff66;
-  --primary-bg-color: #1a1a1a;
-  --foreground-color: #f0f0f0;
-}
-
-main {
-  width: 100%;
-  height: 100%;
-
-  display: flex;
-  flex-direction: column;
-
-  place-items: center;
-}
-
-.home {
-  width: 100%;
-  height: 90%;
-
-  margin-bottom: 10rem;
-}
-
-.title {
-  font-size: 3rem;
-  margin: 1rem 0;
-  margin-top: 3rem;
-}
-
-.histories-container-container {
-  width: 100%;
-  max-height: 75%;
-
-  position: relative;
-}
-
-.histories-container {
-  overflow-y: auto;
-  overflow-x: hidden;
-  width: 100%;
-  height: 100%;
-
-  display: flex;
-  flex-direction: column;
-  gap: 1rem;
-  align-items: center;
-
-  margin: 0;
-  padding: 3rem 1rem;
-}
-
-.histories-start {
-  height: 3rem;
-  width: 100%;
-
-  z-index: 999;
-  top: 0;
-  position: absolute;
-
-  background: linear-gradient(
-    180deg,
-    var(--primary-bg-color) 0%,
-    transparent 100%
-  );
-}
-.histories-end {
-  height: 3rem;
-  width: 100%;
-
-  z-index: 999;
-  bottom: 0;
-  position: absolute;
-
-  background: linear-gradient(
-    0deg,
-    var(--primary-bg-color) 0%,
-    transparent 100%
-  );
-}
-
-.history {
-  padding: 1rem;
-  width: 100%;
-  max-width: 40rem;
-
-  background-color: var(--secondary-color);
-  border-radius: 10px;
-  border-left: 2px solid var(--primary-color);
-
-  cursor: pointer;
-
-  transform: translateX(calc(1px * var(--tx, 0)));
-  opacity: var(--opacity, 1);
-}
-.history:hover {
-  background-color: var(--secondary-color);
-}
-
-.history-delete-button {
-  position: absolute;
-  top: 0;
-  right: 0;
-  padding: 0.5rem;
-  margin: 0;
-  outline: none;
-  border: none;
-  background-color: var(--secondary-color);
-  color: var(--foreground-color);
-  border-radius: 0 0 0 10px;
-  cursor: pointer;
-  transition: 0.2s;
-}
-.history-delete-button:hover {
-  background-color: var(--secondary-color);
-  padding: 0.75rem;
-}
-
-.messages {
-  overflow-y: auto;
-  height: 100%;
-  width: 100%;
-  max-width: 1200px;
-
-  display: flex;
-  flex-direction: column;
-  gap: 1rem;
-  align-items: center;
-  padding-top: 1rem;
-  padding-bottom: 11rem;
-}
-
-.message {
-  max-width: 75%;
-  padding: 0.5rem 1rem;
-  border-radius: 20px;
-}
-.message-role-assistant {
-  background-color: var(--secondary-color);
-  margin-right: auto;
-  color: #fff;
-}
-.message-role-user {
-  margin-left: auto;
-  background-color: var(--primary-color);
-  color: #000;
-}
-
-.message > pre {
-  white-space: pre-wrap;
-}
-
-.hljs {
-  width: 100%;
-  position: relative;
-  border-radius: 10px;
-  /* wrap code blocks */
-  white-space: pre-wrap;
-}
-/* put clipboard button in the top right corner of the code block */
-.clipboard-button {
-  position: absolute;
-  top: 0;
-  right: 0;
-  padding: 0.5rem;
-  margin: 0;
-  outline: none;
-  border: none;
-  background-color: var(--secondary-color);
-  color: var(--foreground-color);
-  border-radius: 0 0 0 10px;
-  cursor: pointer;
-  transition: 0.2s;
-}
-.clipboard-button:hover {
-  background-color: var(--secondary-color);
-  padding: 0.75rem;
-}
-
-.input-container {
-  position: absolute;
-  bottom: 0;
-
-  /* linear gradient from background-color to transparent on the top */
-  background: linear-gradient(
-    0deg,
-    var(--primary-bg-color) 55%,
-    transparent 100%
-  );
-
-  width: 100%;
-  max-width: 1200px;
-  display: flex;
-  flex-direction: column;
-  justify-content: center;
-  align-items: center;
-  z-index: 999;
-}
-
-.input-performance {
-  margin-top: 4rem;
-
-  display: flex;
-  flex-direction: row;
-  gap: 1rem;
-}
-
-.input-performance-point {
-  display: flex;
-  flex-direction: row;
-  place-items: center;
-  gap: 0.5rem;
-}
-.input-performance-point > p {
-  height: 1rem;
-  line-height: normal;
-}
-
-.input {
-  width: 90%;
-  min-height: 3rem;
-  flex-shrink: 0;
-
-  display: flex;
-  flex-direction: row;
-  justify-content: center;
-  gap: 0.5rem;
-
-  align-items: flex-end;
-  margin-bottom: 2rem;
-}
-
-.input-form {
-  width: 100%;
-  padding: 1rem;
-  min-height: 3rem;
-  max-height: 8rem;
-
-  background-color: var(--secondary-color);
-  color: var(--foreground-color);
-  border-radius: 10px;
-  border: none;
-  resize: none;
-  outline: none;
-}
-.mobile .input-form { /* prevent auto-zoom on touching prompt box */
-  font-size: 16px;
-}
-
-.input-button {
-  height: 3rem;
-  width: 4rem;
-
-  background-color: var(--primary-color);
-  color: var(--secondary-color);
-  border-radius: 10px;
-  padding: 0.5rem;
-  cursor: pointer;
-}
-.input-button:hover {
-  background-color: var(--secondary-color-transparent);
-}
-.input-button:disabled {
-  background-color: var(--secondary-color);
-  cursor: not-allowed;
-}
-
-/* wrap text */
-p {
-  white-space: pre-wrap;
-}
-
-/* fonts */
-.megrim-regular {
-  font-family: monospace;
-  font-weight: 400;
-  font-style: normal;
-}
-
-.monospace {
-  font-family: monospace;
-}
-
-.loading-bar {
-  display: flex;
-  flex-direction: row;
-  align-items: center;
-  gap: 0.5rem;
-  width: 100%;
-  min-height: 3rem;
-  margin-bottom: 2rem;
-}
-
-.loading-text {
-  color: var(--foreground-color);
-  font-size: 1rem;
-  white-space: nowrap;
-}
-
-#progress-percentage {
-  color: var(--foreground-color);
-  font-size: 1rem;
-  white-space: nowrap;
-}
-
-.progress-bar {
-  flex-grow: 1;
-  height: 0.5rem;
-  background-color: var(--secondary-color);
-  border-radius: 5px;
-  overflow: hidden;
-  position: relative;
-}
-
-.progress {
-  width: 0%;
-  height: 100%;
-  background-color: var(--primary-color);
-  transition: width 0.2s ease-in-out;
-}

examples/tinychat/tinychat-browser/index.html

@@ -1,182 +0,0 @@
-<!DOCTYPE html>
-
-<head>
-  <title>tinychat</title>
-  <meta name="viewport" content="width=device-width, initial-scale=1">
-  <link rel="icon" href="../favicon.svg" type="image/svg+xml">
-
-  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpine-collective/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="../assets/cdn.jsdelivr.net/npm/@alpinejs/[email protected]/dist/cdn.min.js"></script>
-  <script defer src="../assets/unpkg.com/@marcreichel/[email protected]/dist/alpine-autosize.min.js"></script>
-  <script defer src="../assets/unpkg.com/[email protected]/dist/cdn.min.js"></script>
-
-  <script src="../assets/unpkg.com/[email protected]/dist/purify.min.js"></script>
-  <script src="../assets/unpkg.com/[email protected]/marked.min.js"></script>
-  <script src="../assets/unpkg.com/[email protected]/lib/index.umd.js"></script>
-  <script src="../assets/unpkg.com/@highlightjs/[email protected]/highlight.min.js"></script>
-
-  <script src="index.js"></script>
-
-  <link rel="stylesheet" href="../assets/cdn.jsdelivr.net/npm/[email protected]/build/base-min.css">
-  <link rel="stylesheet" href="../assets/cdnjs.cloudflare.com/ajax/libs/font-awesome/6.5.2/css/all.min.css"
-    integrity="sha512-SnH5WK+bZxgPHs44uWIX+LLJAJ9/2PkPKZ5QiAj6Ta86w+fsb2TkcmfRyVX3pBnMFcV7oQPJkl9QevSCWr3W6A=="
-    crossorigin="anonymous" referrerpolicy="no-referrer" />
-  <link rel="stylesheet" href="../assets/unpkg.com/@highlightjs/[email protected]/styles/vs2015.min.css">
-
-  <link rel="stylesheet" href="index.css">
-  <link rel="stylesheet" href="../common.css">
-</head>
-
-<body>
-  <main x-data="state" x-init="console.log(endpoint)">
-    <div class="home centered" x-show="home === 0" x-transition x-effect="
-      $refs.inputForm.focus();
-      if (home === 1) setTimeout(() => home = 2, 100);
-      if (home === -1) setTimeout(() => home = 0, 100);
-    " @popstate.window="
-      if (home === 2) {
-        cancelGeneration = true;
-        if (maxContextReached) generating = false;
-        if (!generating) cstate = { time: null, messages: [] };
-        home = -1;
-        time_till_first = 0;
-        tokens_per_second = 0;
-        total_tokens = 0;
-      }
-    ">
-      <h1 class="title megrim-regular">tinychat</h1>
-      <div class="histories-container-container">
-        <template x-if="histories.length">
-          <div class="histories-start"></div>
-        </template>
-        <div class="histories-container" x-intersect="
-          $el.scrollTo({ top: 0, behavior: 'smooth' });
-        ">
-          <template x-for="_state in histories.toSorted((a, b) => b.time - a.time)">
-            <div x-data="{ otx: 0, trigger: 75 }" class="history" @click="
-            cstate = _state;
-            updateTotalTokens(cstate.messages);
-            home = 1;
-            // ensure that going back in history will go back to home
-            window.history.pushState({}, '', window.TINYCHAT_ROOT || '/');
-          " @touchstart="
-            otx = $event.changedTouches[0].clientX;
-          " @touchmove="
-            $el.style.setProperty('--tx', $event.changedTouches[0].clientX - otx);
-            $el.style.setProperty('--opacity', 1 - (Math.abs($event.changedTouches[0].clientX - otx) / trigger));
-          " @touchend="
-            if (Math.abs($event.changedTouches[0].clientX - otx) > trigger) removeHistory(_state);
-            $el.style.setProperty('--tx', 0);
-            $el.style.setProperty('--opacity', 1);
-          ">
-              <h3 x-text="new Date(_state.time).toLocaleString()"></h3>
-              <p x-text="$truncate(_state.messages[0].content, 80)"></p>
-              <!-- delete button -->
-              <button class="history-delete-button" @click.stop="removeHistory(_state);">
-                <i class=" fas fa-trash"></i>
-              </button>
-            </div>
-          </template>
-        </div>
-        <template x-if="histories.length">
-          <div class="histories-end"></div>
-        </template>
-      </div>
-    </div>
-    <div x-ref="messages" class="messages" x-init="
-      $watch('cstate', value => {
-        $el.innerHTML = '';
-        value.messages.forEach(({ role, content }) => {
-          const div = document.createElement('div');
-          div.className = `message message-role-${role}`;
-          try {
-            div.innerHTML = DOMPurify.sanitize(marked.parse(content));
-          } catch (e) {
-            console.log(content);
-            console.error(e);
-          }
-
-          // add a clipboard button to all code blocks
-          const codeBlocks = div.querySelectorAll('.hljs');
-          codeBlocks.forEach(codeBlock => {
-            const button = document.createElement('button');
-            button.className = 'clipboard-button';
-            button.innerHTML = '<i class=\'fas fa-clipboard\'></i>';
-            button.onclick = () => {
-              // navigator.clipboard.writeText(codeBlock.textContent);
-              const range = document.createRange();
-              range.setStartBefore(codeBlock);
-              range.setEndAfter(codeBlock);
-              window.getSelection()?.removeAllRanges();
-              window.getSelection()?.addRange(range);
-              document.execCommand('copy');
-              window.getSelection()?.removeAllRanges();
-
-              button.innerHTML = '<i class=\'fas fa-check\'></i>';
-              setTimeout(() => button.innerHTML = '<i class=\'fas fa-clipboard\'></i>', 1000);
-            };
-            codeBlock.appendChild(button);
-          });
-
-          $el.appendChild(div);
-        });
-
-        $el.scrollTo({ top: $el.scrollHeight, behavior: 'smooth' });
-      });
-    " x-intersect="
-      $el.scrollTo({ top: $el.scrollHeight, behavior: 'smooth' });
-    " x-show="home === 2" x-transition>
-    </div>
-    <div class="input-container">
-      <div class="input-performance">
-        <span class="input-performance-point">
-          <p class="monospace" x-text="(time_till_first / 1000).toFixed(2)"></p>
-          <p class="megrim-regular">SEC TO FIRST TOKEN</p>
-        </span>
-        <span class="input-performance-point">
-          <p class="monospace" x-text="tokens_per_second.toFixed(1)"></p>
-          <p class="megrim-regular">TOKENS/SEC</p>
-        </span>
-        <span class="input-performance-point">
-          <p class="monospace" x-text="total_tokens"></p>
-          <p class="megrim-regular">TOKENS</p>
-        </span>
-      </div>
-      <div class="loading-bar" x-show="loadingMessage !== ''">
-        <p class="loading-text" id="loading-message">Loading:</p>
-        <span id="progress-percentage">0%</span>
-        <div class="progress-bar">
-          <div class="progress"></div>
-        </div>
-      </div>
-      <div class="input" x-show="loadingMessage === ''">
-        <textarea x-ref="inputForm" id="input-form" class="input-form" autofocus rows=1 x-autosize
-          :placeholder="generating ? placeholderText : 'Say something'" :disabled="generating" @input="
-            home = (home === 0) ? 1 : home
-            if (cstate.messages.length === 0 && $el.value === '') home = -1;
-
-            if ($el.value !== '') {
-              const messages = [...cstate.messages];
-              messages.push({ role: 'user', content: $el.value });
-              updateTotalTokens(messages);
-            } else {
-              if (cstate.messages.length === 0) total_tokens = 0;
-              else updateTotalTokens(cstate.messages);
-            }
-          " x-effect="
-            console.log(generating);
-            if (!generating) $nextTick(() => {
-              $el.focus();
-              setTimeout(() => $refs.messages.scrollTo({ top: $refs.messages.scrollHeight, behavior: 'smooth' }), 100);
-            });
-          " @keydown.enter="await handleEnter($event)" @keydown.escape.window="$focus.focus($el)"></textarea>
-        <button class="input-button" :disabled="generating" @click="await handleSend()">
-          <i class="fas" :class="generating ? 'fa-spinner fa-spin' : 'fa-paper-plane'"></i>
-        </button>
-      </div>
-    </div>
-  </main>
-</body>
-
-</html>

examples/tinychat/tinychat-browser/index.js

@@ -1,927 +0,0 @@
-window.TINYCHAT_ROOT = "/tinychat-browser/";
-const queryParams = new URLSearchParams(window.location.search);
-const normalizedParams = Object.fromEntries([...queryParams].map(([key, value]) => [key.toUpperCase(), value.toUpperCase()]));
-window.BACKEND = (normalizedParams["BACKEND"] === "WASM") ? "WASM" : "WebGPU";
-const isMobileAgent = /Mobi|Android|iPhone|iPad|iPod/i.test(navigator.userAgent);
-const hasTouchScreen = 'ontouchstart' in window || navigator.maxTouchPoints > 0;
-window.isMobile = isMobileAgent || hasTouchScreen;
-if (window.isMobile) document.documentElement.classList.add('mobile'); // prevent annoying auto-zoom when entering prompt on mobile
-// MODEL_BASE_URL is where the weights are hosted, WEBGPU_EXPORT is the JS-wrapped WebGPU code exported from tinygrad
-window.PC_MODEL_BASE_URL = ".";
-window.PC_WEBGPU_EXPORT = './net.js'
-window.PC_MAX_CONTEXT = 1024;
-window.MOBILE_MODEL_BASE_URL = ".";
-window.MOBILE_WEBGPU_EXPORT = './net.js'
-window.MOBILE_MAX_CONTEXT = 1024;
-
-const tiktokenReady = (async () => {
-  const { init, get_encoding, Tiktoken, load } = await import('./tiktoken.js');
-  window.Tiktoken = Tiktoken;
-  window.tiktokenInit = init;
-  window.tiktokenGetEncoding = get_encoding;
-  window.tiktokenLoad = load;
-})();
-
-async function getDevice() {
-  let adapter;
-  try {
-    adapter = await navigator.gpu.requestAdapter();
-    if (!adapter) {
-      this.loadingMessage = "Loading WASM (WebGPU not enabled):";
-      throw new Error("No WebGPU adapter found");
-    }
-  } catch(error) {
-      this.loadingMessage = "Loading WASM (WebGPU not enabled):";
-    throw error;
-  }
-  const requiredLimits = {};
-  const maxBufferSize = 322122544;
-  requiredLimits.maxStorageBufferBindingSize = maxBufferSize;
-  requiredLimits.maxBufferSize = maxBufferSize;
-  requiredLimits.maxComputeInvocationsPerWorkgroup = 512; // may need to vary based on what the WEBGPU backend produces
-            
-  try {
-    return await adapter.requestDevice({ requiredLimits });
-  } catch(error) {
-    this.loadingMessage = "Loading WASM (WebGPU error):";
-    throw error;
-  }
-};
-
-// copied from examples/webgpu/stable_diffusion/index.html 
-function initDb() {
-  return new Promise((resolve, reject) => {
-    let db;
-    const request = indexedDB.open('tinydb', 1);
-    request.onerror = (event) => {
-      console.error('Database error:', event.target.error);
-      resolve(null);
-    };
-
-    request.onsuccess = (event) => {
-      db = event.target.result;
-      console.log("Db initialized.");
-      resolve(db);
-    };
-
-    request.onupgradeneeded = (event) => {
-      db = event.target.result;
-      if (!db.objectStoreNames.contains('tensors')) {
-        db.createObjectStore('tensors', { keyPath: 'id' });
-      }
-    };
-  });
-}
-
-// copied from examples/webgpu/stable_diffusion/index.html 
-function readTensorFromDb(db, id) {
-  return new Promise((resolve, reject) => {
-    if (db == null) {
-      resolve(null);
-    }
-            
-      const transaction = db.transaction(['tensors'], 'readonly');
-      const store = transaction.objectStore('tensors');
-      const request = store.get(id);
-
-      transaction.onabort = (event) => {
-        console.log("Transaction error while reading tensor: " + event.target.error);
-        resolve(null);
-      };
-
-      request.onsuccess = (event) => {
-        const result = event.target.result;
-        if (result) {
-          resolve(result);
-        } else {
-          resolve(null);
-        }
-      };
-
-      request.onerror = (event) => {
-        console.error('Tensor retrieve failed: ', event.target.error);
-        resolve(null);
-      };
-  });
-}
-
-function getAllKeysFromDb(db) {
-  return new Promise((resolve, reject) => {
-    if (db == null) {resolve([]);}
-    const transaction = db.transaction(['tensors'], 'readonly');
-    const store = transaction.objectStore('tensors');
-    const request = store.getAllKeys();
-    transaction.onabort = (event) => {
-      console.log("Transaction error while reading IndexedDB keys: " + event.target.error);
-      resolve([]);
-    };
-    request.onsuccess = function (event) {resolve(event.target.result);};
-    request.onerror = (event) => {
-      console.error('Retrieval of IndexedDB keys failed: ', event.target.error);
-      resolve([]);
-    };
-  });
-}
-
-// modified from examples/webgpu/stable_diffusion/index.html 
-function saveTensorToDb(db, id, tensor) {
-  return readTensorFromDb(db, id).then((result) => {
-    if (!result) {
-      new Promise((resolve, reject) => {
-        if (db == null) {
-          resolve(null);
-        }
-
-        const transaction = db.transaction(['tensors'], 'readwrite');
-        const store = transaction.objectStore('tensors');
-        const request = store.put({ id: id, content: tensor });
-
-        transaction.onabort = (event) => {
-          console.log("Transaction error while saving tensor: " + event.target.error);
-          resolve(null);
-        };
-
-        request.onsuccess = () => {
-          console.log('Tensor saved successfully.');
-          resolve();
-        };
-
-        request.onerror = (event) => {
-          console.error('Tensor save failed:', event.target.error);
-          resolve(null);
-        };
-      });
-    } else {
-      return null;
-    }
-  }).catch(()=> null);
-}
-
-function deleteTensorFromDb(db, id) {
-  return new Promise((resolve, reject) => {
-    if (db == null) {
-      console.error("Database is not initialized.");
-      resolve(null);
-      return;
-    }
-
-    const transaction = db.transaction(['tensors'], 'readwrite');
-    const store = transaction.objectStore('tensors');
-    const request = store.delete(id);
-
-    transaction.oncomplete = () => {
-      console.log(`Tensor with ID '${id}' deleted successfully.`);
-      resolve();
-    };
-
-    transaction.onerror = (event) => {
-      console.error("Transaction error while deleting tensor:", event.target.error);
-      resolve(null);
-    };
-
-    request.onerror = (event) => {
-      console.error('Tensor deletion failed:', event.target.error);
-      resolve(null);
-    };
-
-    request.onsuccess = () => {
-      console.log(`Delete request for tensor with ID '${id}' succeeded.`);
-    };
-  });
-}
-
-function makeProgress(total) {
-  let acc = 0;
-  const ret = function progress(amount, message) {
-    if (amount >= 0) { // allow updating message only
-      acc += amount;
-      const percentage = total ? Math.trunc((acc / total) * 100) : 0;
-      document.querySelector('.progress').style.width = `${percentage}%`;
-      document.getElementById('progress-percentage').textContent = `${percentage}%`;
-    }
-    if (message) {
-      this.loadingMessage = message;
-      document.getElementById('loading-message').textContent = this.loadingMessage;
-    }
-  }.bind(this);
-  ret.total = total;
-  return ret;
-}
-
-function sendMessageToWorker(worker, message) {
-  return new Promise((resolve, reject) => {
-    const onMessage = (event) => {
-      resolve(event.data);
-      worker.removeEventListener('message', onMessage);
-      worker.removeEventListener('error', onError);
-    };
-
-    const onError = (error) => {
-      reject(error);
-      worker.removeEventListener('message', onMessage);
-      worker.removeEventListener('error', onError);
-    };
-
-    worker.addEventListener('message', onMessage);
-    worker.addEventListener('error', onError);
-
-    if (message.header === "token") worker.postMessage(message.data);
-    else if (message.header === "load_state_dict") {
-      if (message.data === "done") worker.postMessage(message.data);
-      else worker.postMessage(message.data, message.data.map(file => file.bytes.buffer));
-    }
-    else if (message.header === "init") worker.postMessage("init");
-  });
-}
-
-async function load_state_dict (data, device, progress) {
-  let state_dict = data.metadata.state_dict;
-  let completed = 0;
-
-  // modified from examples/webgpu/stable_diffusion/index.html getProgressDlForPart
-  const loadPart = async (part) => {
-      const response = await fetch(part);
-      const res = new Response(new ReadableStream({
-          async start(controller) {
-              const reader = response.body.getReader();
-              for (;;) {
-                  const { done, value } = await reader.read();
-                  if (done) break;
-                  progress(value.byteLength);
-                  controller.enqueue(value);
-              }
-              controller.close();
-          },
-      }));
-        
-      return res.arrayBuffer();
-  };
-
-  let db = await initDb();
-
-  const getPart = async(filename, hash) => {
-    let part = await readTensorFromDb(db, hash);
-
-    if (part) {
-      console.log(`Cache hit: ${filename}, hash: ${hash}`);
-      progress(part.content.byteLength);
-      return Promise.resolve(part.content);
-    } else {
-      console.log(`Cache miss: ${filename}, hash: ${hash}`);
-      return loadPart(`${window.MODEL_BASE_URL}/${filename}`);
-    }
-  }
-
-  const correctHashes = data.metadata.files.map(file => file.hash)
-  // delete unused cached buffers to free disk space -- if we update weights, user will otherwise have obsolete cached buffers
-  const dbKeys = await getAllKeysFromDb(db);
-  const correctHashesSet = new Set(correctHashes);
-  const notInCorrectHashes = dbKeys.filter(key => !correctHashesSet.has(key));
-  // await these right before starting to save new stuff
-  const deletionPromises = notInCorrectHashes.map(async (hash) => deleteTensorFromDb(db, hash));
-
-  // instantiates empty weight buffers on WebGPU, attaches buffers to state_dict
-  let model;
-  if (window.BACKEND === "WebGPU") {
-    //model = await transformer().setup(device, state_dict, progress);
-    model = await transformer.setupNet(device, state_dict);
-    progress(0.15 * progress.total);
-
-  }
-  else if (window.BACKEND === "WASM") {
-    progress(0.02 * progress.total);
-    model = new Worker(`./worker.js?version=${Date.now()}`);
-    await sendMessageToWorker(model, {header: "init"});
-    progress(0.02 * progress.total);
-    progress(0.11 * progress.total);
-  }
-
-  const downloaded = [];
-  const triggerChainDownload = async (toDownload) => {
-    const numDownloaders = window.isMobile ? 4 : toDownload.length; // TODO: dynamically base this on DL file size? current assumption is 16 MiB chunks
-
-    const chainDownload = async() => {
-      const file = toDownload.shift();
-      loadPart(`${window.MODEL_BASE_URL}/${file.name}`) // triggers download
-      .then(async (arraybuf) => { 
-        downloaded.push({ ...file, bytes: new Uint8Array(arraybuf)});
-        // pause downloads if further processing is a bottleneck
-        while (toDownload.length && downloaded.length >= numDownloaders) await new Promise(resolve => setTimeout(resolve, 5));
-        if (toDownload.length && downloaded.length < numDownloaders) chainDownload(); // start next download
-      })
-    }
-    for (let i=0; i<numDownloaders; i++) if (toDownload.length) chainDownload();
-  }
-
-  const loadFileToStateDict = async(file) => {
-    if (window.BACKEND === "WebGPU") {
-      for (const part of file.parts) {
-        if (part.empty) continue;
-        part.bytes = (part.size === file.bytes.length) ? file.bytes : file.bytes.slice(part.file_start_pos, part.file_start_pos + part.size);
-        device.queue.writeBuffer(state_dict[part.key].bytes, part.target_start_pos, part.bytes); // improves stability over mappedAtCreation writing
-        part.bytes = null;
-      }
-    }
-    else if (window.BACKEND === "WASM") {
-      await sendMessageToWorker(model, {header: "load_state_dict", data: [file]});
-    }
-    file.bytes = null;
-  }
-
-  if (window.BACKEND === "WebGPU") { // contiguous loading not needed for WebGPU stability
-    const files = data.tensor_file_groups.flatMap(obj => obj.files);
-    data.tensor_file_groups = [{contiguous: false, files: files}];
-  }
-
-  for (const group of data.tensor_file_groups) {
-    const contiguous = group.contiguous;
-    const files = group.files;
-    const tensor_file_indices = files.map(file => file.index);
-    const contiguousFiles = [];
-    const fileHashes = new Set(files.map(file => file.hash));
-    const cachedFileHashes = new Set(dbKeys.filter(key => fileHashes.has(key)));
-    const cachedFiles = files.filter(file => cachedFileHashes.has(file.hash));
-    const toDownload = files.filter(file => !cachedFileHashes.has(file.hash));
-    triggerChainDownload(toDownload);
-
-    const loadDelay = 5;
-    await Promise.all(deletionPromises);
-
-    while (completed < files.length) {
-      const start = performance.now();
-      // prioritize files from downloaded queue, so we can continue downloading more files
-      if (downloaded.length) {
-        const file = downloaded.shift();
-        await saveTensorToDb(db, file.hash, file.bytes); // for wasm, must await to prevent race between indexedDB and transfer to worker
-        if (!contiguous) await loadFileToStateDict(file);
-        else contiguousFiles.push(file);
-        completed += 1;
-      }
-      else if (!downloaded.length && cachedFiles.length) {
-        const file = cachedFiles.shift();
-        file.bytes = await getPart(file.name, file.hash); // reads data from IndexedDB
-        if (!contiguous) await loadFileToStateDict(file);
-        else contiguousFiles.push(file);
-        completed += 1;
-      }
-      const end = performance.now();
-      const elapsed = end - start;
-      if (elapsed < loadDelay) await new Promise(resolve => setTimeout(resolve, loadDelay - elapsed));
-    }
-    if (contiguous) {
-      const orderMap = tensor_file_indices.reduce((acc, id, index) => {acc[id] = index; return acc;}, {});
-      contiguousFiles.sort((a, b) => orderMap[a.index] - orderMap[b.index]); // glue files together in the right order
-      await sendMessageToWorker(model, {header: "load_state_dict", data: contiguousFiles});
-    }
-    completed = 0;
-  }
-
-  // initialize empty kv_caches, which were part of exported model's state_dict, but which we didn't want to package/download
-  if (window.BACKEND === "WASM") {
-    for (const [k, v] of Object.entries(state_dict).filter(([_, v]) => v.empty === true)) {
-      v.parts[0].file_start_pos = 0;
-      const file = { parts: v.parts, size: v.size, bytes: new Uint8Array(v.size).fill(0) };
-      await loadFileToStateDict(file);
-    }
-  }
-
-  return model;
-};
-
-document.addEventListener("alpine:init", () => {
-  Alpine.data("state", () => ({
-    // loadingMessage updates the user on page load progress, including weights download and decompression
-    // if loadingMessage is not '', then prompt box will be hidden: this is default behavior on page load
-    placeholderText: "Generating...",
-    loadingMessage: `Loading ${window.BACKEND} model:`,
-    // model
-    nets: {},
-    tokenizer: null,
-    max_context: 1024,
-    lastSeenToks: [],
-
-    progress: null,
-
-    async init() {
-      var device = null;
-      var webgpuErrorMessage = null;
-      if (window.BACKEND === "WebGPU") {
-        try {
-          device = await getDevice.call(this);
-          console.log("WebGPU device initialized");
-        } catch (error) {
-          window.BACKEND = "WASM";
-          console.log(`error: ${error}\nFailed to launch WebGPU. Loading WASM model instead...`); // return;
-          webgpuErrorMessage = this.loadingMessage;
-        }
-      }
-
-      window.MODEL_BASE_URL = (window.BACKEND === "WebGPU" && !window.isMobile) ? window.PC_MODEL_BASE_URL : window.MOBILE_MODEL_BASE_URL;
-      this.max_context = (window.BACKEND === "WebGPU" && !window.isMobile) ? window.PC_MAX_CONTEXT : window.MOBILE_MAX_CONTEXT;
-
-      const kernelsReady = (async () => {
-        if (window.BACKEND === "WASM") {var exports = await import(`./net_clang.js?version=${Date.now()}`);}
-        else if (window.BACKEND === "WebGPU" && !window.isMobile) {var exports = await import(`${PC_WEBGPU_EXPORT}?version=${Date.now()}`);}
-        else if (window.BACKEND === "WebGPU" && window.isMobile) {var exports = await import(`${MOBILE_WEBGPU_EXPORT}?version=${Date.now()}`);}
-        self.transformer = exports.default;
-      })();
-
-      const response = await fetch(`${window.MODEL_BASE_URL}/net_metadata.json`);
-      // TODO: cache metadata (and everything else, including tokenizer)
-      // TODO: use service worker to reload page when offline
-      const data = await response.json();
-      data.metadata.files = data.metadata.files.map((file, index) => ({...file, index}));
-      const state_dict = data.metadata.state_dict;
-
-      /*
-      - allocating memory to WASM on mobile has longstanding issues: https://github.com/WebAssembly/design/issues/1397
-
-      - the below pattern, while yielding a succesfully-functioning model when it doesn't crash, causes regular crashes on iOS Safari (iphone 15 iOS 18.3):
-        - call WASM malloc (to fit all tensors, or one per tensor) for all tensors up front, then load tensor byte chunks into the buffers in random order
-
-      - the below pattern has been stable on iOS Safari (iphone 15 iOS 18.3):
-        - call only one WASM malloc at a time before filling the allocated bytes, as small as possible (malloc up to 256 MiB has been tested)
-        - fill the malloc'd memory in linear order from start to end (what has been tested is calling wasm.HEAPU8.set on 16 MiB chunks from start to end)
-        - use ALLOW_MEMORY_GROWTH=1 in wasm compilation, minimize initial memory
-
-      - additional considerations affecting loading design, for WASM:
-        - it seems that copying bytes into wasm memory cannot be zero-copy without sharedarraybuffer, which isn't currently used due to increased hosting complexity
-        - non-zero copies create memory pressure, which is not reliably capped because of lack of control over garbage collection
-        - to minimize peak memory pressure if GC is delayed, we process (i.e. download + copy into WASM) large tensors (> 16 MiB) one at a time, in descending size order
-      */
-      data.tensor_file_groups = []; // see above: for WASM, limit processing of multi-file Tensors to one at a time, in descending order based on Tensor size
-      const unsplit_tensors = [];
-      const sortedEntries = Object.entries(state_dict).sort(([, objA], [, objB]) => objB.size - objA.size);
-
-      let totalSize = 0;
-      const seen = new Set();
-      for (const [k,v] of sortedEntries) {
-        const files_in_tensor = [];
-        for (const part of v.parts) {
-          part.key = k;
-          if (part.empty) state_dict[k].empty = true; // assumes no other parts of this weight exist and are non-empty
-          else {
-            const file = data.metadata.files[part.file];
-            if (!seen.has(file.index)) {
-              seen.add(file.index);
-              files_in_tensor.push(file);
-            }
-            totalSize += part.size;
-            part.dtype = v.dtype;
-            if (!data.metadata.files[part.file].parts) data.metadata.files[part.file].parts = [];
-            data.metadata.files[part.file].size ??= 0;
-            data.metadata.files[part.file].size += part.size;
-            data.metadata.files[part.file].parts.push(part);
-          }
-        }
-        if (files_in_tensor.length > 1) data.tensor_file_groups.push({contiguous: true, files: files_in_tensor}); // [tensorN_file0, tensorN_file1, ...]
-        else if (files_in_tensor.length > 0) unsplit_tensors.push(files_in_tensor[0]);
-      }
-      data.tensor_file_groups.push({contiguous: false, files: unsplit_tensors});
-
-      data.totalSize = totalSize;
-      totalSize = totalSize / 0.8; // give space in progress bar for initializing model bufs, and tokenizer
-      this.progress = makeProgress.call(this, totalSize); // creates closure with totalSize
-
-      try {
-        this.progress(0.01 * totalSize, "Loading tokenizer:");
-        const wasmResponse = await fetch(`${window.MODEL_BASE_URL}/tiktoken_bg.wasm`);
-        this.progress(0.01 * totalSize);
-        const wasmBytes = await wasmResponse.arrayBuffer();
-        await tiktokenReady;
-        await window.tiktokenInit((imports) => WebAssembly.instantiate(wasmBytes, imports));
-        this.progress(0.01 * totalSize);
-
-        this.tokenizer = await createTokenizer(`${window.MODEL_BASE_URL}/llama3-2.tiktoken`);
-        const tokenizer_works = (new TextDecoder().decode(this.tokenizer.decode(this.tokenizer.encode("hello world"))) === "hello world");
-        console.log("tokenizer works:", tokenizer_works)
-        this.progress(0.01 * totalSize);
-      } catch (error) {this.progress(-1, `Error launching tokenizer: ${error}`); console.log(error); return;}
-
-      try {
-        const loadModelMessage = (webgpuErrorMessage) ? webgpuErrorMessage : `Loading ${window.BACKEND} model:`
-        this.progress(0, loadModelMessage);
-        await kernelsReady;
-        const model = await load_state_dict(data, device, this.progress);
-
-        if (window.BACKEND === "WebGPU") {
-          this.nets = {"transformer": model};
-        }
-        else if (window.BACKEND === "WASM") {
-          const msg = await sendMessageToWorker(model, {header: "load_state_dict", data: "done"});
-          this.nets = {"transformer": async (tok, start_pos) => sendMessageToWorker(model, {header: "token", data: [tok, start_pos]})};
-        }
-        this.progress(0.01 * totalSize, `Launching ${window.BACKEND} model:`);
-        this.loadingMessage = ""; // Triggers removal of loading bar, display of prompt box
-      } catch (error) {this.progress(-1, `Error launching model: ${error}`); console.log(error); return;}
-    },
-
-    // current state
-    cstate: {
-      time: null,
-      messages: [],
-    },
-
-    // historical state
-    histories: JSON.parse(localStorage.getItem("histories")) || [],
-
-    home: 0,
-    generating: false,
-    maxContextReached: false,
-    cancelGeneration: false,
-    endpoint: `${window.location.origin}/v1`,
-
-    // performance tracking
-    time_till_first: 0,
-    tokens_per_second: 0,
-    total_tokens: 0,
-    max_context: 0,
-
-    removeHistory(cstate) {
-      const index = this.histories.findIndex((state) => {
-        return state.time === cstate.time;
-      });
-      if (index !== -1) {
-        this.histories.splice(index, 1);
-        localStorage.setItem("histories", JSON.stringify(this.histories));
-      }
-    },
-
-    async handleSend() {
-      const el = document.getElementById("input-form");
-      const value = el.value.trim();
-      if (!value) return;
-
-      if (this.generating) return;
-      this.maxContextReached = false;
-      this.placeholderText = "Generating...";
-      this.generating = true;
-      this.cancelGeneration = false;
-      if (this.home === 0) this.home = 1;
-
-      // ensure that going back in history will go back to home
-      window.history.pushState({}, "", window.TINYCHAT_ROOT || "/");
-
-      // add message to list
-      this.cstate.messages.push({ role: "user", content: value });
-
-      // clear textarea
-      el.value = "";
-      el.style.height = "auto";
-      el.style.height = el.scrollHeight + "px";
-
-      // reset performance tracking
-      const prefill_start = Date.now();
-      let start_time = 0;
-      let tokens = 0;
-      this.tokens_per_second = 0;
-
-      let gottenFirstChunk = false;
-      try {
-        for await (
-          const chunk of this.openaiChatCompletion(this.cstate.messages)
-        ) {
-          if (!gottenFirstChunk) {
-            this.cstate.messages.push({ role: "assistant", content: "" });
-            gottenFirstChunk = true;
-          }
-
-          // add chunk to the last message
-          // TODO: handle localStorage overflow
-          //   possible example: this.cstate.messages[...] was undefined when trying to prompt within an old cstate (chat session)
-          this.cstate.messages[this.cstate.messages.length - 1].content += chunk;
-
-          // calculate performance tracking
-          tokens += 1;
-          this.total_tokens += 1;
-          if (start_time === 0) {
-            start_time = Date.now();
-            this.time_till_first = start_time - prefill_start;
-          } else {
-            const diff = Date.now() - start_time;
-            if (diff > 0) {
-              this.tokens_per_second = tokens / (diff / 1000);
-            }
-          }
-          this.checkMaxContext(this.total_tokens);
-          if (this.cancelGeneration) break;
-        }
-      } finally {
-        // update the state in histories or add it if it doesn't exist
-        const index = this.histories.findIndex((cstate) => {
-          return cstate.time === this.cstate.time;
-        });
-        this.cstate.time = Date.now();
-        if (index !== -1) {
-          // update the time
-          this.histories[index] = this.cstate;
-        } else {
-          this.histories.push(this.cstate);
-        }
-        // update in local storage
-        localStorage.setItem("histories", JSON.stringify(this.histories));
-
-        if (!this.maxContextReached) this.generating = false;
-        if (this.cancelGeneration && !this.maxContextReached) this.cstate = { time: null, messages: [] };
-      }
-    },
-
-    async handleEnter(event) {
-      // if shift is not pressed
-      if (!event.shiftKey) {
-        event.preventDefault();
-        await this.handleSend();
-      }
-    },
-
-    updateTotalTokens(messages) {
-      try {
-        let toks = [this.tokenizer.bos_id];
-        messages.forEach((message) => {
-          if (!message.role || !message.content) {
-            throw new Error("Each message must have a 'role' and 'content' property.");
-          }
-          toks = toks.concat(this.tokenizer.encodeMessage(message.role, message.content));
-
-          if (messages.length > 0 && messages[messages.length - 1].role === "user") {
-            toks = toks.concat(this.tokenizer.encodeRole("assistant"));
-          }
-          this.total_tokens = toks.length;
-        });
-      } catch (error) {
-        console.error("Error updating total tokens:", error);
-      }
-    },
-
-    checkMaxContext(num_tokens) {
-      if (num_tokens >= this.max_context) {
-        this.cancelGeneration = true;
-        this.maxContextReached = true;
-        this.placeholderText = `Max context reached: ${this.max_context} tokens`;
-      }
-    },
-
-    async *openaiChatCompletion(messages) {
-      let tokens = [this.tokenizer.bos_id];
-      for (const message of messages) {
-        tokens = tokens.concat(this.tokenizer.encodeMessage(message.role, message.content));
-      }
-      tokens = tokens.concat(this.tokenizer.encodeRole("assistant"));
-      this.checkMaxContext(tokens.length); // don't waste time prefilling if we know we're over the token limit
-      let startPos = 0
-      const prefillToks = tokens.slice(0, -1);
-
-      // Skip the largest possible sequence of tokens already represented at the beginning of the model's kv caches
-      for (let i=0; i <= prefillToks.length; i++) {
-        startPos = i;
-        if (i == prefillToks.length) break;
-        if (i == this.lastSeenToks.length) break;
-        if (prefillToks[i] !== this.lastSeenToks[i]) break;
-      }
-      //this.lastSeenToks = prefillToks;
-      //prefillToks = prefillToks.slice(startPos);
-      const unprocessedPrefillToks = prefillToks.slice(startPos);
-      this.lastSeenToks = prefillToks.slice(0, startPos);
-
-      this.progress = makeProgress(unprocessedPrefillToks.length);
-      this.loadingMessage = (window.BACKEND === "WebGPU") ? "Reading input:" : "Loading (enable WebGPU for speed):";
-      this.progress(0, this.loadingMessage);
-      for (const tok of unprocessedPrefillToks) {
-        if (this.cancelGeneration) {this.loadingMessage=""; return;}
-        if (window.BACKEND === "WebGPU") {await this.nets["transformer"](new Int32Array([tok]), new Int32Array([startPos]));}
-        else {await this.nets["transformer"](tok, startPos);}
-        this.lastSeenToks.push(tok)
-        startPos += 1;
-        this.progress(1);
-      }
-      this.loadingMessage = ""; // hides progress bar
-
-      let lastTok = tokens[tokens.length - 1];
-      while (true) {
-        if (window.BACKEND === "WebGPU") {var tok = await this.nets["transformer"](new Int32Array([lastTok]), new Int32Array([startPos])); tok = tok[0][0];}
-        else {var tok = await this.nets["transformer"](lastTok, startPos);}
-        this.lastSeenToks.push(lastTok); // lets us skip prefilling with these tokens at the next prompt in this chain
-        startPos += 1;
-        lastTok = tok;
-        if (this.tokenizer.stop_tokens.has(lastTok)) break;
-        yield new TextDecoder().decode(this.tokenizer.decode([lastTok]));
-      }
-    },
-  }));
-});
-
-const { markedHighlight } = globalThis.markedHighlight;
-marked.use(markedHighlight({
-  langPrefix: "hljs language-",
-  highlight(code, lang, _info) {
-    const language = hljs.getLanguage(lang) ? lang : "plaintext";
-    return hljs.highlight(code, { language }).value;
-  },
-}));
-
-// **** eventsource-parser ****
-class EventSourceParserStream extends TransformStream {
-  constructor() {
-    let parser;
-
-    super({
-      start(controller) {
-        parser = createParser((event) => {
-          if (event.type === "event") {
-            controller.enqueue(event);
-          }
-        });
-      },
-
-      transform(chunk) {
-        parser.feed(chunk);
-      },
-    });
-  }
-}
-
-function createParser(onParse) {
-  let isFirstChunk;
-  let buffer;
-  let startingPosition;
-  let startingFieldLength;
-  let eventId;
-  let eventName;
-  let data;
-  reset();
-  return {
-    feed,
-    reset,
-  };
-  function reset() {
-    isFirstChunk = true;
-    buffer = "";
-    startingPosition = 0;
-    startingFieldLength = -1;
-    eventId = void 0;
-    eventName = void 0;
-    data = "";
-  }
-  function feed(chunk) {
-    buffer = buffer ? buffer + chunk : chunk;
-    if (isFirstChunk && hasBom(buffer)) {
-      buffer = buffer.slice(BOM.length);
-    }
-    isFirstChunk = false;
-    const length = buffer.length;
-    let position = 0;
-    let discardTrailingNewline = false;
-    while (position < length) {
-      if (discardTrailingNewline) {
-        if (buffer[position] === "\n") {
-          ++position;
-        }
-        discardTrailingNewline = false;
-      }
-      let lineLength = -1;
-      let fieldLength = startingFieldLength;
-      let character;
-      for (
-        let index = startingPosition;
-        lineLength < 0 && index < length;
-        ++index
-      ) {
-        character = buffer[index];
-        if (character === ":" && fieldLength < 0) {
-          fieldLength = index - position;
-        } else if (character === "\r") {
-          discardTrailingNewline = true;
-          lineLength = index - position;
-        } else if (character === "\n") {
-          lineLength = index - position;
-        }
-      }
-      if (lineLength < 0) {
-        startingPosition = length - position;
-        startingFieldLength = fieldLength;
-        break;
-      } else {
-        startingPosition = 0;
-        startingFieldLength = -1;
-      }
-      parseEventStreamLine(buffer, position, fieldLength, lineLength);
-      position += lineLength + 1;
-    }
-    if (position === length) {
-      buffer = "";
-    } else if (position > 0) {
-      buffer = buffer.slice(position);
-    }
-  }
-  function parseEventStreamLine(lineBuffer, index, fieldLength, lineLength) {
-    if (lineLength === 0) {
-      if (data.length > 0) {
-        onParse({
-          type: "event",
-          id: eventId,
-          event: eventName || void 0,
-          data: data.slice(0, -1),
-          // remove trailing newline
-        });
-
-        data = "";
-        eventId = void 0;
-      }
-      eventName = void 0;
-      return;
-    }
-    const noValue = fieldLength < 0;
-    const field = lineBuffer.slice(
-      index,
-      index + (noValue ? lineLength : fieldLength),
-    );
-    let step = 0;
-    if (noValue) {
-      step = lineLength;
-    } else if (lineBuffer[index + fieldLength + 1] === " ") {
-      step = fieldLength + 2;
-    } else {
-      step = fieldLength + 1;
-    }
-    const position = index + step;
-    const valueLength = lineLength - step;
-    const value = lineBuffer.slice(position, position + valueLength).toString();
-    if (field === "data") {
-      data += value ? "".concat(value, "\n") : "\n";
-    } else if (field === "event") {
-      eventName = value;
-    } else if (field === "id" && !value.includes("\0")) {
-      eventId = value;
-    } else if (field === "retry") {
-      const retry = parseInt(value, 10);
-      if (!Number.isNaN(retry)) {
-        onParse({
-          type: "reconnect-interval",
-          value: retry,
-        });
-      }
-    }
-  }
-}
-const BOM = [239, 187, 191];
-function hasBom(buffer) {
-  return BOM.every((charCode, index) => buffer.charCodeAt(index) === charCode);
-}
-
-const PAT_STR = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+";
-
-async function createTokenizer(bpeUrl) {
-  const num_base_tokens = 128000;
-  const special_tokens = {
-    "<|begin_of_text|>": 128000,
-    "<|end_of_text|>": 128001,
-    "<|start_header_id|>": 128006,
-    "<|end_header_id|>": 128007,
-    "<|eot_id|>": 128009
-  };
-  const model = await window.tiktokenLoad({
-        "load_tiktoken_bpe": bpeUrl,
-        "special_tokens": special_tokens,
-        "pat_str": PAT_STR
-    });
-  const tokenizer = new window.Tiktoken(model.bpe_ranks, model.special_tokens, model.pat_str)
-
-  return {
-    get bos_id() {
-      return special_tokens["<|begin_of_text|>"];
-    },
-
-    get stop_tokens() {
-      return new Set([
-        special_tokens["<|end_of_text|>"],
-        special_tokens["<|eot_id|>"],
-      ]);
-    },
-
-    decode(toks) {
-      const filtered = toks.filter((t) => t < num_base_tokens);
-      return tokenizer.decode(filtered);
-    },
-
-    encode(text, allow_special = false) {
-      const allowedSpecial = allow_special ? "all" : new Set();
-      const disallowedSpecial = new Set();
-      return tokenizer.encode(text, allowedSpecial, disallowedSpecial);
-    },
-
-    encodeRole(role) {
-      const tokens = [];
-      tokens.push(special_tokens["<|start_header_id|>"]);
-      tokens.push(...this.encode(role));
-      tokens.push(special_tokens["<|end_header_id|>"]);
-      tokens.push(...this.encode("\n\n"));
-      return tokens;
-    },
-
-    encodeMessage(role, content) {
-      const roleTokens = this.encodeRole(role);
-      const contentTokens = this.encode(content.trim());
-      return [...roleTokens, ...contentTokens, special_tokens["<|eot_id|>"]];
-    },
-  };
-}

examples/tinychat/tinychat-browser/make_tiktoken_js.sh

@@ -1,11 +0,0 @@
-#!/usr/bin/env bash
-cd "$(dirname "$0")"
-npm init -y && \
-npm install --save-dev webpack webpack-cli && \
-npm install tiktoken && \
-jq '.scripts.build = "webpack"' package.json > package.tmp.json && \
-mv package.tmp.json package.json && \
-npm run build && \
-mv dist/*.wasm ./tiktoken_bg.wasm && \
-mv dist/* ./ && \
-rm -rf dist node_modules package-lock.json package.json

examples/tinychat/tinychat-browser/tiktoken-export.js

@@ -1,5 +0,0 @@
-// Force Webpack to copy the WASM
-import 'tiktoken/tiktoken_bg.wasm'; 
-import { init, get_encoding, encoding_for_model, Tiktoken } from 'tiktoken/init';
-import { load } from 'tiktoken/load';
-export { init, get_encoding, encoding_for_model, Tiktoken, load };

examples/tinychat/tinychat-browser/webpack.config.js

@@ -1,25 +0,0 @@
-const path = require("path");
-
-module.exports = {
-  mode: "production",
-  entry: "./tiktoken-export.js",
-  output: {
-    filename: "tiktoken.js",
-    path: path.resolve(__dirname, "dist"),
-    library: {
-      type: "module"
-    }
-  },
-  experiments: {
-    outputModule: true,
-    asyncWebAssembly: true
-  },
-  module: {
-    rules: [
-      {
-        test: /\.wasm$/,
-        type: "asset/resource",
-      }
-    ]
-  }
-};

examples/tinychat/tinychat-browser/worker.js

@@ -1,62 +0,0 @@
-const kernelsReady = (async () => {
-  // can't get browser to use updated versions except with cache-busting query string
-  const exports = await import(`./net_clang.js?version=${Date.now()}`);
-  Object.assign(self, exports);
-})();
-
-async function init(event) {
-  await kernelsReady;
-  self.model = await self.transformer();
-  self.addEventListener("message", loadStateDict);
-  self.removeEventListener("message", init);
-  self.postMessage("success");
-}
-
-function loadStateDict(event) {
-  if (event.data === "done") {
-    self.addEventListener("message", inference);
-    self.removeEventListener("message", loadStateDict);
-  }
-  else {
-    if (event.data.length > 1) {
-      // the bytes from files are set contiguously in WASM memory
-      const malloc_size = event.data.reduce((sum, file) => sum + file.bytes.length, 0);
-      const malloc_ptr = self.model.wasm._malloc(malloc_size);
-      let cursor = 0;
-      for (const file of event.data) {
-        self.model.wasm.HEAPU8.set(file.bytes, malloc_ptr + cursor);
-        for (const part of file.parts) {
-          if (part.target_start_pos === 0) {
-            // tell WASM code where the tensor is in memory
-            self.model.wasm._set_buf(self.transformer_name_to_id[part.key], malloc_ptr + cursor);
-          }
-          cursor += part.size;
-        }
-        file.bytes = null;
-      }
-    }
-    else {
-      // the bytes from files are not guaranteed to be set contiguously in WASM memory
-      const file = event.data[0];
-      const malloc_ptr = self.model.wasm._malloc(file.size);
-      self.model.wasm.HEAPU8.set(file.bytes, malloc_ptr);
-      for (const part of file.parts) {
-        if (part.target_start_pos === 0) {
-          self.model.wasm._set_buf(self.transformer_name_to_id[part.key], malloc_ptr + part.file_start_pos);
-        }
-      }
-      file.bytes = null;
-    }
-  }
-  self.postMessage("success");
-}
-
-function inference(event) {
-  const [tok, start_pos] = event.data;
-  const int32tok = new Int32Array([tok]);
-  const model_out = self.model.run(new Uint8Array(int32tok.buffer), start_pos);
-  const int32nextTok = new Int32Array(model_out[0].buffer);
-  self.postMessage(int32nextTok[0]);
-}
-
-self.addEventListener("message", init);

examples/torch_cuda_kernel.py

@@ -1,38 +0,0 @@
-#!POPCORN leaderboard grayscale
-#!POPCORN gpu A100
-# not a stable API, but works
-
-import torch, functools
-from tinygrad import Tensor, TinyJit, Device
-from tinygrad.engine.realize import CompiledRunner
-from tinygrad.helpers import get_single_element, Context, OSX
-from tinygrad.dtype import _from_torch_dtype
-
-@TinyJit
-def f(tg_out, tg_data): return tg_out.assign(tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140).realize()
-
-def custom_kernel(data: torch.Tensor, device="CUDA") -> torch.Tensor:
-  assert data.dtype == torch.float32
-  tg_data = Tensor.from_blob(data.data_ptr(), data.shape, dtype=_from_torch_dtype(data.dtype), device=device)
-
-  out = torch.empty((data.shape[0], data.shape[1]), dtype=data.dtype, device=data.device)
-  tg_out = Tensor.from_blob(out.data_ptr(), out.shape, dtype=_from_torch_dtype(out.dtype), device=device)
-
-  # Need to sync torch to make sure the data is valid.
-  if data.device.type == "mps": torch.mps.synchronize()
-  else: torch.cuda.synchronize()
-
-  with Context(BEAM=2): f(tg_out, tg_data)
-
-  # Wait for computation to finish and the data is valid.
-  Device[device].synchronize()
-
-  return out
-
-if __name__ == "__main__":
-  for i in range(3):
-    if OSX:
-      out = custom_kernel(inp:=torch.rand(16, 16, 3, device=torch.device("mps")), device="METAL")
-    else:
-      out = custom_kernel(inp:=torch.rand(16, 16, 3, device=torch.device("cuda")), device="CUDA")
-    assert torch.allclose(out, inp[:, :, 0] * 0.2989 + inp[:, :, 1] * 0.5870 + inp[:, :, 2] * 0.1140)

examples/vgg7_helpers/waifu2x.py

@@ -105,12 +105,12 @@ class Vgg7:
     Output format: (1, 3, Y - 14, X - 14)
     (the - 14 represents the 7-pixel context border that is lost)
     """
-    x = self.conv1.forward(x).leaky_relu(0.1)
-    x = self.conv2.forward(x).leaky_relu(0.1)
-    x = self.conv3.forward(x).leaky_relu(0.1)
-    x = self.conv4.forward(x).leaky_relu(0.1)
-    x = self.conv5.forward(x).leaky_relu(0.1)
-    x = self.conv6.forward(x).leaky_relu(0.1)
+    x = self.conv1.forward(x).leakyrelu(0.1)
+    x = self.conv2.forward(x).leakyrelu(0.1)
+    x = self.conv3.forward(x).leakyrelu(0.1)
+    x = self.conv4.forward(x).leakyrelu(0.1)
+    x = self.conv5.forward(x).leakyrelu(0.1)
+    x = self.conv6.forward(x).leakyrelu(0.1)
     x = self.conv7.forward(x)
     return x
 

examples/vits.py

@@ -42,10 +42,10 @@ class Synthesizer:
     if pad_length > -1:
       # Pad flow forward inputs to enable JIT
       assert pad_length > row_len, "pad length is too small"
-      y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len))).cast(z_p.dtype)
+      y_mask = y_mask.pad(((0, 0), (0, 0), (0, pad_length - row_len)), 0).cast(z_p.dtype)
       # New y_mask tensor to remove sts mask
       y_mask = Tensor(y_mask.numpy(), device=y_mask.device, dtype=y_mask.dtype, requires_grad=y_mask.requires_grad)
-      z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), value=1).unsqueeze(0)
+      z_p = z_p.squeeze(0).pad(((0, 0), (0, pad_length - z_p.shape[2])), 1).unsqueeze(0)
     z = self.flow.forward(z_p.realize(), y_mask.realize(), g=g.realize(), reverse=True)
     result_length = reduce(lambda x, y: x * y, self.dec.upsample_rates, row_len)
     o = self.dec.forward((z * y_mask)[:, :, :max_len], g=g)[:, :, :result_length]
@@ -114,7 +114,7 @@ class StochasticDurationPredictor:
     flows = flows[:-2] + [flows[-1]] # remove a useless vflow
     z = Tensor.randn(x.shape[0], 2, x.shape[2], dtype=x.dtype).to(device=x.device) * noise_scale
     for flow in flows: z = flow.forward(z, x_mask, g=x, reverse=reverse)
-    z0, z1 = z.split([1, 1], 1)
+    z0, z1 = split(z, [1, 1], 1)
     return z0.realize()
 
 class DurationPredictor:
@@ -147,7 +147,7 @@ class TextEncoder:
     x = x.transpose(1, -1)  # [b, t, h] -transpose-> [b, h, t]
     x_mask = sequence_mask(x_lengths, x.shape[2]).unsqueeze(1).cast(x.dtype)
     x = self.encoder.forward(x * x_mask, x_mask)
-    m, logs = (self.proj(x) * x_mask).split(self.out_channels, dim=1)
+    m, logs = split(self.proj(x) * x_mask, self.out_channels, dim=1)
     return x.realize(), m.realize(), logs.realize(), x_mask.realize()
 
 class ResidualCouplingBlock:
@@ -193,10 +193,10 @@ class Generator:
     x = self.conv_pre(x)
     if g is not None:  x = x + self.cond(g)
     for i in range(self.num_upsamples):
-      x = self.ups[i](x.leaky_relu(LRELU_SLOPE))
+      x = self.ups[i](x.leakyrelu(LRELU_SLOPE))
       xs = sum(self.resblocks[i * self.num_kernels + j].forward(x) for j in range(self.num_kernels))
       x = (xs / self.num_kernels).realize()
-    res = self.conv_post(x.leaky_relu()).tanh().realize()
+    res = self.conv_post(x.leakyrelu()).tanh().realize()
     return res
 
 class LayerNorm(nn.LayerNorm):
@@ -238,8 +238,8 @@ class ResBlock1:
     self.convs2 = [nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)) for _ in range(3)]
   def forward(self, x: Tensor, x_mask=None):
     for c1, c2 in zip(self.convs1, self.convs2):
-      xt = x.leaky_relu(LRELU_SLOPE)
-      xt = c1(xt if x_mask is None else xt * x_mask).leaky_relu(LRELU_SLOPE)
+      xt = x.leakyrelu(LRELU_SLOPE)
+      xt = c1(xt if x_mask is None else xt * x_mask).leakyrelu(LRELU_SLOPE)
       x = c2(xt if x_mask is None else xt * x_mask) + x
     return x if x_mask is None else x * x_mask
 
@@ -282,7 +282,7 @@ class ConvFlow:
     self.convs = DDSConv(filter_channels, kernel_size, n_layers, p_dropout=0.)
     self.proj = nn.Conv1d(filter_channels, self.half_channels * (num_bins * 3 - 1), 1)
   def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = x.split([self.half_channels] * 2, 1)
+    x0, x1 = split(x, [self.half_channels] * 2, 1)
     h = self.proj(self.convs.forward(self.pre(x0), x_mask, g=g)) * x_mask
     b, c, t = x0.shape
     h = h.reshape(b, c, -1, t).permute(0, 1, 3, 2) # [b, cx?, t] -> [b, c, t, ?]
@@ -302,10 +302,10 @@ class ResidualCouplingLayer:
     self.enc = WN(hidden_channels, kernel_size, dilation_rate, n_layers, p_dropout=p_dropout, gin_channels=gin_channels)
     self.post = nn.Conv1d(hidden_channels, self.half_channels * (2 - mean_only), 1)
   def forward(self, x, x_mask, g=None, reverse=False):
-    x0, x1 = x.split([self.half_channels] * 2, 1)
+    x0, x1 = split(x, [self.half_channels] * 2, 1)
     stats = self.post(self.enc.forward(self.pre(x0) * x_mask, x_mask, g=g)) * x_mask
     if not self.mean_only:
-      m, logs = stats.split([self.half_channels] * 2, 1)
+      m, logs = split(stats, [self.half_channels] * 2, 1)
     else:
       m = stats
       logs = Tensor.zeros_like(m)
@@ -420,7 +420,7 @@ def piecewise_rational_quadratic_transform(inputs, un_normalized_widths, un_norm
   return spline_fn(inputs=inputs, un_normalized_widths=un_normalized_widths, un_normalized_heights=un_normalized_heights, un_normalized_derivatives=un_normalized_derivatives, inverse=inverse, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative, **spline_kwargs)
 def unconstrained_rational_quadratic_spline(inputs, un_normalized_widths, un_normalized_heights, un_normalized_derivatives, inverse=False, tails='linear', tail_bound=1., min_bin_width=DEFAULT_MIN_BIN_WIDTH, min_bin_height=DEFAULT_MIN_BIN_HEIGHT, min_derivative=DEFAULT_MIN_DERIVATIVE):
   if not tails == 'linear': raise RuntimeError('{} tails are not implemented.'.format(tails))
-  constant = np.log(np.exp(1 - min_derivative) - 1).item()
+  constant = np.log(np.exp(1 - min_derivative) - 1)
   un_normalized_derivatives = cat_lr(un_normalized_derivatives, constant, constant)
   output, log_abs_det = rational_quadratic_spline(inputs=inputs.squeeze(dim=0).squeeze(dim=0), unnormalized_widths=un_normalized_widths.squeeze(dim=0).squeeze(dim=0), unnormalized_heights=un_normalized_heights.squeeze(dim=0).squeeze(dim=0), unnormalized_derivatives=un_normalized_derivatives.squeeze(dim=0).squeeze(dim=0), inverse=inverse, left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound, min_bin_width=min_bin_width, min_bin_height=min_bin_height, min_derivative=min_derivative)
   return output.unsqueeze(dim=0).unsqueeze(dim=0), log_abs_det.unsqueeze(dim=0).unsqueeze(dim=0)
@@ -478,7 +478,16 @@ def get_shape(tensor):
   return tuple(shape)
 def convert_pad_shape(pad_shape): return tuple(tuple(x) for x in pad_shape)
 def get_padding(kernel_size, dilation=1): return int((kernel_size*dilation - dilation)/2)
-
+def split(tensor, split_sizes, dim=0):  # if split_sizes is an integer, convert it to a tuple of size split_sizes elements
+  if isinstance(split_sizes, int): split_sizes = (split_sizes,) * (tensor.shape[dim] // split_sizes)
+  assert sum(split_sizes) == tensor.shape[
+    dim], "Sum of split_sizes must equal the dimension size of tensor along the given dimension."
+  start, slices = 0, []
+  for size in split_sizes:
+    slice_range = [(start, start + size) if j == dim else None for j in range(len(tensor.shape))]
+    slices.append(slice_range)
+    start += size
+  return [tensor._slice(s) for s in slices]
 def gather(x, indices, axis):
   indices = (indices < 0).where(indices + x.shape[axis], indices).transpose(0, axis)
   permute_args = list(range(x.ndim))

examples/webgpu/stable_diffusion/compile.py

@@ -1,10 +1,9 @@
 import os
-from extra.export_model import compile_net, jit_model, dtype_to_js_type
-from extra.f16_decompress import u32_to_f16
+from extra.export_model import compile_net, jit_model
 from examples.stable_diffusion import StableDiffusion
 from tinygrad.nn.state import get_state_dict, safe_save, safe_load_metadata, torch_load, load_state_dict
 from tinygrad.tensor import Tensor
-from tinygrad import Device, dtypes
+from tinygrad import Device
 from tinygrad.helpers import fetch
 from typing import NamedTuple, Any, List
 import requests
@@ -30,7 +29,7 @@ def convert_f32_to_f16(input_file, output_file):
     rest_float32_values.tofile(f)
 
 def split_safetensor(fn):
-  _, data_start, metadata = safe_load_metadata(fn)
+  _, json_len, metadata = safe_load_metadata(fn)
   text_model_offset = 3772703308
   chunk_size = 536870912
 
@@ -52,12 +51,12 @@ def split_safetensor(fn):
     part_offset = offset - last_offset
 
     if (part_offset >= chunk_size):
-      part_end_offsets.append(data_start+offset)
+      part_end_offsets.append(8+json_len+offset)
       last_offset = offset
 
   text_model_start = int(text_model_offset/2)
   net_bytes = bytes(open(fn, 'rb').read())
-  part_end_offsets.append(text_model_start+data_start)
+  part_end_offsets.append(text_model_start+8+json_len)
   cur_pos = 0
 
   for i, end_pos in enumerate(part_end_offsets):
@@ -66,7 +65,7 @@ def split_safetensor(fn):
       cur_pos = end_pos
 
   with open(os.path.join(os.path.dirname(__file__), f'./net_textmodel.safetensors'), "wb+") as f:
-    f.write(net_bytes[text_model_start+data_start:])
+    f.write(net_bytes[text_model_start+8+json_len:])
 
   return part_end_offsets
 
@@ -96,8 +95,7 @@ if __name__ == "__main__":
   sub_steps = [
     Step(name = "textModel", input = [Tensor.randn(1, 77)], forward = model.cond_stage_model.transformer.text_model),
     Step(name = "diffusor", input = [Tensor.randn(1, 77, 768), Tensor.randn(1, 77, 768), Tensor.randn(1,4,64,64), Tensor.rand(1), Tensor.randn(1), Tensor.randn(1), Tensor.randn(1)], forward = model),
-    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode),
-    Step(name = "f16tof32", input = [Tensor.randn(2097120, dtype=dtypes.uint32)], forward = u32_to_f16)
+    Step(name = "decoder", input = [Tensor.randn(1,4,64,64)], forward = model.decode)
   ]
 
   prg = ""
@@ -118,23 +116,19 @@ if __name__ == "__main__":
     weights = {id(x.lazydata.base.realized): name for name, x in state.items()}
     kernel_code = '\n\n'.join([f"const {key} = `{fixup_code(code, key)}`;" for key, code in functions.items()])
     kernel_names = ', '.join([name for (name, _, _, _) in statements])
-    input_names = [name for _,name in special_names.items() if "input" in name]
-    output_names = [name for _,name in special_names.items() if "output" in name]
-    input_buf_types = [dtype_to_js_type(bufs[inp_name][1]) for inp_name in input_names]
-    output_buf_types = [dtype_to_js_type(bufs[out_name][1]) for out_name in output_names]
     kernel_calls = '\n        '.join([f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});" for i, (_name, args, global_size, _local_size) in enumerate(statements) ])
-    exported_bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
+    bufs =  '\n    '.join([f"const {name} = " + (f"createEmptyBuf(device, {size});" if _key not in weights else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weights[_key]}'], '{weights[_key]}'))") + ";"  for name,(size,dtype,_key) in bufs.items()])
     gpu_write_bufs =  '\n    '.join([f"const gpuWriteBuffer{i} = device.createBuffer({{size:input{i}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});" for i,(_,value) in enumerate(special_names.items()) if "output" not in value])
-    input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new {input_buf_types[i]}(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,_ in enumerate(input_names)])
+    input_writer = '\n    '.join([f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n    new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set(" + f'data{i});' + f"\n    gpuWriteBuffer{i}.unmap();\ncommandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, input{i}, 0, gpuWriteBuffer{i}.size);"  for i,(_,value) in enumerate(special_names.items()) if value != "output0"])
     return f"""\n    var {step.name} = function() {{
 
     {kernel_code}
 
     return {{
       "setup": async (device, safetensor) => {{
-        const metadata = safetensor ? getTensorMetadata(safetensor[0]) : null;
+        const metadata = getTensorMetadata(safetensor[0]);
 
-        {exported_bufs}
+        {bufs}
 
         {gpu_write_bufs}
         const gpuReadBuffer = device.createBuffer({{ size: output0.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});
@@ -153,8 +147,8 @@ if __name__ == "__main__":
             device.queue.submit([gpuCommands]);
 
             await gpuReadBuffer.mapAsync(GPUMapMode.READ);
-            const resultBuffer = new {output_buf_types[0]}(gpuReadBuffer.size/{bufs[output_names[0]][1].itemsize});
-            resultBuffer.set(new {output_buf_types[0]}(gpuReadBuffer.getMappedRange()));
+            const resultBuffer = new Float32Array(gpuReadBuffer.size/4);
+            resultBuffer.set(new Float32Array(gpuReadBuffer.getMappedRange()));
             gpuReadBuffer.unmap();
             return resultBuffer;
         }}

examples/webgpu/stable_diffusion/index.html

@@ -165,6 +165,10 @@
         import ClipTokenizer from './clip_tokenizer.js';
         window.clipTokenizer = new ClipTokenizer();
     </script>
+    <script type="module">
+        import { f16tof32GPU } from 'https://unpkg.com/[email protected]/src/index.js';
+        window.f16tof32GPU = f16tof32GPU;
+    </script>
     <script src="./net.js"></script>
 </head>
 <body>
@@ -177,21 +181,19 @@
     </a>
     
     <div id="mybox">
-        <form id="promptForm">
-            <input id="promptText" type="text" placeholder="Enter your prompt here" value="a human standing on the surface of mars">
+        <input id="promptText" type="text" placeholder="Enter your prompt here" value="a human standing on the surface of mars">
 
-            <label>
-                Steps: <span id="stepValue">9</span>
-                <input id="stepRange" type="range" min="5" max="20" value="9" step="1">
-            </label>
+        <label>
+            Steps: <span id="stepValue">9</span>
+            <input id="stepRange" type="range" min="5" max="20" value="9" step="1">
+        </label>
 
-            <label>
-                Guidance: <span id="guidanceValue">8.0</span>
-                <input id="guidanceRange" type="range" min="3" max="15" value="8.0" step="0.1">
-            </label>
+        <label>
+            Guidance: <span id="guidanceValue">8.0</span>
+            <input id="guidanceRange" type="range" min="3" max="15" value="8.0" step="0.1">
+        </label>
 
-            <input id="btnRunNet" type="button" value="Run" disabled>
-        </form>
+        <input id="btnRunNet" type="button" value="Run" disabled>
 
         <div id="divModelDl" style="display: flex; align-items: center; width: 100%; gap: 10px;">
             <span id="modelDlTitle">Downloading model</span>
@@ -212,8 +214,6 @@
     <canvas id="canvas" width="512" height="512"></canvas>
 
 <script>
-    let f16decomp = null;
-
     function initDb() {
         return new Promise((resolve, reject) => {
             let db;
@@ -322,9 +322,7 @@
             requiredLimits.maxBufferSize = maxBufferSizeInSDModel;
             
             return await adapter.requestDevice({
-                requiredLimits,
-                requiredFeatures: ["shader-f16"],
-                powerPreference: "high-performance"
+                requiredLimits
             });
         };
 
@@ -418,7 +416,7 @@
             const metadata = JSON.parse(new TextDecoder("utf8").decode(combinedBuffer.subarray(8, 8 + metadataLength)));
 
             const allToDecomp = combinedBuffer.byteLength - (8 + metadataLength);
-            const decodeChunkSize = 8388480;
+            const decodeChunkSize = 67107840;
             const numChunks = Math.ceil(allToDecomp/decodeChunkSize);
 
             console.log(allToDecomp + " bytes to decompress");
@@ -442,8 +440,7 @@
                 let chunkStartF16 = 8 + metadataLength + (decodeChunkSize * i);
                 let chunkEndF16 = chunkStartF16 + decodeChunkSize;
                 let chunk = combinedBuffer.subarray(chunkStartF16, chunkEndF16);
-                let uint32Chunk = new Uint32Array(chunk.buffer, chunk.byteOffset, chunk.byteLength / 4);
-                let result = await f16decomp(uint32Chunk);
+                let result = await f16tof32GPU(chunk);
                 let resultUint8 = new Uint8Array(result.buffer);
                 let chunkStartF32 = 8 + metadataLength + (decodeChunkSize * i * 2);
                 let chunkEndF32 = chunkStartF32 + resultUint8.byteLength;
@@ -486,7 +483,6 @@
             }
 
             const device = await getDevice();
-            f16decomp =  await f16tof32().setup(device, safetensorParts),
             safetensorParts = await getAndDecompressF16Safetensors(device, progress);
 
             modelDlTitle.innerHTML = "Compiling model"
@@ -563,7 +559,7 @@
             });
         }
 
-        function renderImage(image) {
+        function renderImage(e, image) {
             let pixels = []
             let pixelCounter = 0
 
@@ -578,16 +574,14 @@
             }
 
             ctx.putImageData(new ImageData(new Uint8ClampedArray(pixels), 512, 512), 0, 0);
+            e.target.disabled = false;
         }
 
-        const handleRunNetAndRenderResult = () => {
-            document.getElementById("btnRunNet").disabled = true;
+        document.getElementById("btnRunNet").addEventListener("click", function(e) {
+            e.target.disabled = true;
             const canvas = document.getElementById("canvas");
             ctx.clearRect(0, 0, canvas.width, canvas.height);
 
-            const prevTitleValue = document.getElementById("modelDlTitle").innerHTML;
-            document.getElementById("modelDlTitle").innerHTML = "Running model";
-
             runStableDiffusion(
                 document.getElementById("promptText").value,
                 document.getElementById("stepRange").value,
@@ -595,21 +589,9 @@
                 // Decode at each step
                 null
             ).then((image) => {
-                renderImage(image);
-            }).finally(() => {
-                document.getElementById("modelDlTitle").innerHTML = prevTitleValue;
-                document.getElementById("btnRunNet").disabled = false;
+                renderImage(e, image);
             });
-        };
-
-        document.getElementById("btnRunNet").addEventListener("click", handleRunNetAndRenderResult, false);
-
-        document.getElementById("promptForm").addEventListener("submit", function (event) {
-            event.preventDefault();
-            if (document.getElementById("btnRunNet").disabled) return;
-
-            handleRunNetAndRenderResult();
-        })
+        }, false);
 
         const stepSlider = document.getElementById('stepRange');
         const stepValue = document.getElementById('stepValue');

examples/webgpu/yolov8/compile.py

@@ -12,7 +12,7 @@ if __name__ == "__main__":
     yolo_infer = YOLOv8(w=0.25, r=2.0, d=0.33, num_classes=80)
     state_dict = safe_load(get_weights_location(yolo_variant))
     load_state_dict(yolo_infer, state_dict)
-    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,416,416), model_name="yolov8")
+    prg, inp_sizes, out_sizes, state = export_model(yolo_infer, Device.DEFAULT.lower(), Tensor.randn(1,3,256,256))
     dirname = Path(__file__).parent
     safe_save(state, (dirname / "net.safetensors").as_posix())
     with open(dirname / f"net.js", "w") as text_file:

examples/webgpu/yolov8/index.html

@@ -4,10 +4,7 @@
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
     <title>YOLOv8 tinygrad WebGPU</title>
-    <script type="module">
-        import yolov8 from "./net.js"
-        window.yolov8 = yolov8;
-    </script>
+    <script src="./net.js"></script>
     <style>
         body {
             text-align: center;
@@ -98,7 +95,6 @@
 </head>
 <body>
     <h2>YOLOv8 tinygrad WebGPU</h2>
-    <h2 id="wgpu-error" style="display: none; color: red;">Error: WebGPU is not supported in this browser</h2>
     <div class="video-container">
         <video id="video" muted autoplay playsinline></video>
         <canvas id="canvas"></canvas>
@@ -111,7 +107,7 @@
     </div>
     <script>
         let net = null;
-        const modelInputSize = 416;
+        const modelInputSize = 256
         let lastCalledTime;
         let fps = 0, accumFps = 0, frameCounter = 0;
 
@@ -121,7 +117,6 @@
         const offscreenCanvas = document.createElement('canvas');
         const fpsMeter = document.getElementById('fps-meter');
         const loadingContainer = document.getElementById('div-loading');
-        const wgpuError = document.getElementById('wgpu-error');
         offscreenCanvas.width = modelInputSize;
         offscreenCanvas.height = modelInputSize;
         const offscreenContext = offscreenCanvas.getContext('2d');
@@ -152,7 +147,7 @@
                 lastCalledTime = now;
                 accumFps += 1/delta;
 
-                if (frameCounter++ >= 10) {
+                if (frameCounter++ >= 30) {
                     fps = accumFps/frameCounter;
                     frameCounter = 0;
                     accumFps = 0;
@@ -211,12 +206,7 @@
 
         async function detectObjectsOnFrame(offscreenContext) {
             if (!net) {
-                let device = await getDevice();
-                if (!device) {
-                    wgpuError.style.display = "block";
-                    loadingContainer.style.display = "none";
-                }
-                net = await yolov8.load(device, "./net.safetensors");
+                net = await loadNet(await getDevice());
                 loadingContainer.style.display = "none";
             }
             let start = performance.now();
@@ -249,12 +239,9 @@
         }
 
         const getDevice = async () => {
-            if (!navigator.gpu) return false;
+            if (!navigator.gpu) error("WebGPU not supported.");
             const adapter = await navigator.gpu.requestAdapter();
-            return await adapter.requestDevice({
-                requiredFeatures: ["shader-f16"],
-                powerPreference: "high-performance"
-		    });
+            return await adapter.requestDevice();
         };
         
         function processOutput(output, img_width, img_height) {

examples/yolov3.py

@@ -228,7 +228,7 @@ class Darknet:
           module.append(BatchNorm2d(filters, eps=1e-05, track_running_stats=True))
         # LeakyReLU activation
         if activation == "leaky":
-          module.append(lambda x: x.leaky_relu(0.1))
+          module.append(lambda x: x.leakyrelu(0.1))
       elif module_type == "maxpool":
         size, stride = int(x["size"]), int(x["stride"])
         module.append(lambda x: x.max_pool2d(kernel_size=(size, size), stride=stride))

examples/yolov8-onnx.py

@@ -3,8 +3,7 @@ import os
 from ultralytics import YOLO
 import onnx
 from pathlib import Path
-from tinygrad.frontend.onnx import OnnxRunner
-from extra.onnx_helpers import get_example_inputs
+from extra.onnx import get_run_onnx
 from tinygrad.tensor import Tensor
 
 os.chdir("/tmp")
@@ -12,5 +11,8 @@ if not Path("yolov8n-seg.onnx").is_file():
   model = YOLO("yolov8n-seg.pt")
   model.export(format="onnx", imgsz=[480,640])
 onnx_model = onnx.load(open("yolov8n-seg.onnx", "rb"))
-run_onnx = OnnxRunner(onnx_model)
-run_onnx(get_example_inputs(run_onnx.graph_inputs), debug=True)
+# TODO: move get example inputs to onnx
+input_shapes = {inp.name:tuple(x.dim_value for x in inp.type.tensor_type.shape.dim) for inp in onnx_model.graph.input}
+print(input_shapes)
+run_onnx = get_run_onnx(onnx_model)
+run_onnx({"images": Tensor.zeros(1,3,480,640)}, debug=True)

extra/amdpci/am_smi.py

@@ -1,245 +0,0 @@
-#!/usr/bin/env python3
-
-import time, mmap, sys, shutil, os, glob, subprocess
-from tinygrad.helpers import to_mv, DEBUG, colored, ansilen
-from tinygrad.runtime.autogen import libc
-from tinygrad.runtime.autogen.am import smu_v13_0_0
-from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
-from tinygrad.runtime.support.am.ip import AM_SOC, AM_GMC, AM_IH, AM_PSP, AM_SMU, AM_GFX, AM_SDMA
-
-AM_VERSION = 0xA0000002
-SMU_11_0_GFX_BUSY_THRESHOLD = 15
-
-def bold(s): return f"\033[1m{s}\033[0m"
-
-def trim(s:str, length:int) -> str:
-  if len(s) > length: return s[:length-3] + "..."
-  return s
-
-def color_temp(temp):
-  if temp >= 87: return colored(f"{temp:>3}", "red")
-  elif temp >= 80: return colored(f"{temp:>3}", "yellow")
-  return f"{temp:>3}"
-
-def color_voltage(voltage): return colored(f"{voltage/1000:>5.3f}V", "cyan")
-
-def draw_bar(percentage, width=40, fill='█', empty='░'):
-  filled_width = int(width * percentage)
-  bar = fill * filled_width + empty * (width - filled_width)
-  return f'[{bar}] {percentage*100:5.1f}%'
-
-def same_line(strs:list[list[str]|None], split=8) -> list[str]:
-  strs = [s for s in strs if s is not None]
-  if len(strs) == 0: return []
-
-  ret = []
-  max_width_in_block = [max(ansilen(line) for line in block) for block in strs]
-  max_height = max(len(block) for block in strs)
-  for i in range(max_height):
-    line = []
-    for bid, block in enumerate(strs):
-      if i < len(block): line.append(block[i] + ' ' * (split + max_width_in_block[bid] - ansilen(block[i])))
-      else: line.append(' ' * (split + max_width_in_block[bid]))
-    ret.append(' '.join(line))
-  return ret
-
-def get_bar0_size(pcibus):
-  resource_file = f"/sys/bus/pci/devices/{pcibus}/resource"
-  if not os.path.exists(resource_file): raise FileNotFoundError(f"Resource file not found: {resource_file}")
-
-  with open(resource_file, "r") as f: lines = f.readlines()
-  bar0_info = lines[0].split()
-  if len(bar0_info) < 3: raise ValueError("Unexpected resource file format for BAR0.")
-
-  start_hex, end_hex, _flags = bar0_info
-  return int(end_hex, 16) - int(start_hex, 16) + 1
-
-class AMSMI(AMDev):
-  def __init__(self, pcibus, vram_bar:memoryview, doorbell_bar:memoryview, mmio_bar:memoryview):
-    self.pcibus = pcibus
-    self.vram, self.doorbell64, self.mmio = vram_bar, doorbell_bar, mmio_bar
-    self.pci_state = self.read_pci_state()
-    if self.pci_state == "D0": self._init_from_d0()
-
-  def _init_from_d0(self):
-    self._run_discovery()
-    self._build_regs()
-
-    if self.reg("regSCRATCH_REG7").read() != AM_VERSION:
-      raise Exception(f"Unsupported AM version: {self.reg('regSCRATCH_REG7').read():x}")
-
-    self.is_booting, self.smi_dev = True, True
-    self.partial_boot = True # do not init anything
-    self.mm = AMMemoryManager(self, self.vram_size)
-
-    # Initialize IP blocks
-    self.soc:AM_SOC = AM_SOC(self)
-    self.gmc:AM_GMC = AM_GMC(self)
-    self.ih:AM_IH = AM_IH(self)
-    self.psp:AM_PSP = AM_PSP(self)
-    self.smu:AM_SMU = AM_SMU(self)
-
-    for ip in [self.soc, self.gmc, self.ih, self.psp, self.smu]: ip.init_sw()
-
-  def read_pci_state(self):
-    with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
-
-class SMICtx:
-  def __init__(self):
-    self.devs = []
-    self.opened_pcidevs = []
-    self.opened_pci_resources = {}
-    self.prev_lines_cnt = 0
-    self.prev_terminal_width = 0
-
-    remove_parts = ["Advanced Micro Devices, Inc. [AMD/ATI]", "VGA compatible controller:"]
-    lspci = subprocess.check_output(["lspci"]).decode("utf-8").splitlines()
-    self.lspci = {l.split()[0]: l.split(" ", 1)[1] for l in lspci}
-    for k,v in self.lspci.items():
-      for part in remove_parts: self.lspci[k] = self.lspci[k].replace(part, "").strip().rstrip()
-
-  def _open_am_device(self, pcibus):
-    if pcibus not in self.opened_pci_resources:
-      bar_fds = {bar: os.open(f"/sys/bus/pci/devices/{pcibus}/resource{bar}", os.O_RDWR | os.O_SYNC) for bar in [0, 2, 5]}
-      bar_size = {0: get_bar0_size(pcibus), 2: os.fstat(bar_fds[2]).st_size, 5: os.fstat(bar_fds[5]).st_size}
-
-      def map_pci_range(bar):
-        return to_mv(libc.mmap(0, bar_size[bar], mmap.PROT_READ | mmap.PROT_WRITE, mmap.MAP_SHARED, bar_fds[bar], 0), bar_size[bar])
-      self.opened_pci_resources[pcibus] = (map_pci_range(0), None, map_pci_range(5).cast('I'))
-
-    try:
-      self.devs.append(AMSMI(pcibus, *self.opened_pci_resources[pcibus]))
-    except Exception as e:
-      if DEBUG >= 2: print(f"Failed to open AM device {pcibus}: {e}")
-      return
-
-    self.opened_pcidevs.append(pcibus)
-    if DEBUG >= 2: print(f"Opened AM device {pcibus}")
-
-  def rescan_devs(self):
-    pattern = os.path.join('/tmp', 'am_*.lock')
-    for d in [f[8:-5] for f in glob.glob(pattern)]:
-      if d not in self.opened_pcidevs:
-        self._open_am_device(d)
-
-    for d in self.devs:
-      if d.read_pci_state() != d.pci_state:
-        d.pci_state = d.read_pci_state()
-        if d.pci_state == "D0": d._init_from_d0()
-        os.system('clear')
-
-      if d.pci_state == "D0" and d.reg("regSCRATCH_REG7").read() != AM_VERSION:
-        self.devs.remove(d)
-        self.opened_pcidevs.remove(d.pcibus)
-        os.system('clear')
-        if DEBUG >= 2: print(f"Removed AM device {d.pcibus}")
-
-  def collect(self): return {d: d.smu.read_metrics() if d.pci_state == "D0" else None for d in self.devs}
-
-  def draw(self):
-    terminal_width, terminal_height = shutil.get_terminal_size()
-    if self.prev_terminal_width != terminal_width or self.prev_terminal_height != terminal_height:
-      os.system('clear')
-    self.prev_terminal_width, self.prev_terminal_height = terminal_width, terminal_height
-
-    activity_line_width = 50 if self.prev_terminal_width > 170 else \
-                         (30 if self.prev_terminal_width > 130 else \
-                         (16 if self.prev_terminal_width > 92 else \
-                         max(0, self.prev_terminal_width - 77)))
-    max_col_size = terminal_width // 2
-
-    dev_metrics = self.collect()
-    dev_content = []
-    for dev, metrics in dev_metrics.items():
-      if dev.pci_state != "D0":
-        dev_content.append([f"{colored('(sleep)', 'yellow')} {bold(dev.pcibus)}: {self.lspci[dev.pcibus[5:]]}"] +
-                           [f"PCI State: {dev.pci_state}"] + [" "*107])
-        continue
-
-      device_line = [f"{bold(dev.pcibus)} {trim(self.lspci[dev.pcibus[5:]], max_col_size - 24)}"] + [""]
-      activity_line = [f"GFX Activity {draw_bar(metrics.SmuMetrics.AverageGfxActivity / 100, activity_line_width)}"] \
-                    + [f"MEM Activity {draw_bar(metrics.SmuMetrics.AverageUclkActivity / 100, activity_line_width)}"]
-
-      # draw_metrics_table(metrics, dev)
-      temps_keys = [(k, name) for k, name in smu_v13_0_0.c__EA_TEMP_e__enumvalues.items()
-                                if k < smu_v13_0_0.TEMP_COUNT and metrics.SmuMetrics.AvgTemperature[k] != 0]
-      temps_table = ["=== Temps (°C) ==="] + [f"{name:<15}: {color_temp(metrics.SmuMetrics.AvgTemperature[k])}" for k, name in temps_keys]
-      temps_table_compact = [f"Temps (°C): {color_temp(metrics.SmuMetrics.AvgTemperature[smu_v13_0_0.TEMP_HOTSPOT])} hotspot /" \
-                          + f"{color_temp(metrics.SmuMetrics.AvgTemperature[smu_v13_0_0.TEMP_MEM])} mem"]
-
-      voltage_keys = [(k, name) for k, name in smu_v13_0_0.c__EA_SVI_PLANE_e__enumvalues.items() if k < smu_v13_0_0.SVI_PLANE_COUNT]
-      power_table = ["=== Power ==="] \
-                  + [f"Fan Speed: {metrics.SmuMetrics.AvgFanRpm} RPM"] \
-                  + [f"Fan Power: {metrics.SmuMetrics.AvgFanPwm}%"]
-      power_line = [f"Power: {metrics.SmuMetrics.AverageSocketPower:>3}W " +
-                       draw_bar(metrics.SmuMetrics.AverageSocketPower / metrics.SmuMetrics.dGPU_W_MAX, 16)]
-      power_line_compact = [f"Power:  {metrics.SmuMetrics.AverageSocketPower:>3}W " +
-                       draw_bar(metrics.SmuMetrics.AverageSocketPower / metrics.SmuMetrics.dGPU_W_MAX, activity_line_width)]
-
-      voltage_table = ["=== Voltages ==="] + [f"{name:<20}: {color_voltage(metrics.SmuMetrics.AvgVoltage[k])}" for k, name in voltage_keys]
-
-      gfx_freq = (metrics.SmuMetrics.AverageGfxclkFrequencyPostDs if metrics.SmuMetrics.AverageGfxActivity <= SMU_11_0_GFX_BUSY_THRESHOLD else \
-                  metrics.SmuMetrics.AverageGfxclkFrequencyPreDs)
-      fclk_freq = (metrics.SmuMetrics.AverageFclkFrequencyPostDs if metrics.SmuMetrics.AverageUclkActivity <= SMU_11_0_GFX_BUSY_THRESHOLD else \
-                   metrics.SmuMetrics.AverageFclkFrequencyPreDs)
-      mclk_freq = (metrics.SmuMetrics.AverageMemclkFrequencyPostDs if metrics.SmuMetrics.AverageUclkActivity <= SMU_11_0_GFX_BUSY_THRESHOLD else \
-                   metrics.SmuMetrics.AverageMemclkFrequencyPreDs)
-
-      frequency_table = ["=== Frequencies ===",
-        f"GFXCLK Target : {metrics.SmuMetrics.AverageGfxclkFrequencyTarget:>4} MHz",
-        f"GFXCLK PreDs  : {metrics.SmuMetrics.AverageGfxclkFrequencyPreDs:>4} MHz",
-        f"GFXCLK PostDs : {metrics.SmuMetrics.AverageGfxclkFrequencyPostDs:>4} MHz",
-        f"FCLK PreDs    : {metrics.SmuMetrics.AverageFclkFrequencyPreDs:>4} MHz",
-        f"FCLK PostDs   : {metrics.SmuMetrics.AverageFclkFrequencyPostDs:>4} MHz",
-        f"MCLK PreDs    : {metrics.SmuMetrics.AverageMemclkFrequencyPreDs:>4} MHz",
-        f"MCLK PostDs   : {metrics.SmuMetrics.AverageMemclkFrequencyPostDs:>4} MHz",
-        f"VCLK0         : {metrics.SmuMetrics.AverageVclk0Frequency:>4} MHz",
-        f"DCLK0         : {metrics.SmuMetrics.AverageDclk0Frequency:>4} MHz",
-        f"VCLK1         : {metrics.SmuMetrics.AverageVclk1Frequency:>4} MHz",
-        f"DCLK1         : {metrics.SmuMetrics.AverageDclk1Frequency:>4} MHz"]
-
-      frequency_table_compact = ["=== Frequencies ===",
-        f"GFXCLK: {gfx_freq:>4} MHz",
-        f"FCLK  : {fclk_freq:>4} MHz",
-        f"MCLK  : {mclk_freq:>4} MHz"]
-
-      if self.prev_terminal_width >= 231:
-        power_table += power_line + [""] + voltage_table
-        activity_line += [""]
-      elif self.prev_terminal_width >= 171:
-        power_table += power_line + [""] + frequency_table_compact
-        activity_line += [""]
-        frequency_table = None
-      elif self.prev_terminal_width >= 121:
-        temps_table = None
-        frequency_table = frequency_table_compact
-        activity_line += power_line_compact
-      else:
-        temps_table = None
-        power_table = None
-        frequency_table = None
-        activity_line += power_line_compact
-
-      dev_content.append(device_line + activity_line + same_line([temps_table, power_table, frequency_table]))
-
-    raw_text = 'AM Monitor'.center(terminal_width) + "\n" + "=" * terminal_width + "\n\n"
-    for i in range(0, len(dev_content), 2):
-      if i + 1 < len(dev_content): raw_text += '\n'.join(same_line([dev_content[i], dev_content[i+1]]))
-      else: raw_text += '\n'.join(dev_content[i])
-      if i + 2 < len(dev_content): raw_text += "\n" + "=" * terminal_width + "\n\n"
-
-    sys.stdout.write(f'\033[{self.prev_lines_cnt}A')
-    sys.stdout.flush()
-    print(raw_text)
-
-    self.prev_lines_cnt = len(raw_text.splitlines()) + 2
-
-if __name__ == "__main__":
-  try:
-    os.system('clear')
-    smi_ctx = SMICtx()
-    while True:
-      smi_ctx.rescan_devs()
-      smi_ctx.draw()
-      time.sleep(1)
-  except KeyboardInterrupt: print("Exiting...")

extra/amdpci/headers/amdgpu_doorbell.h

@@ -1,279 +0,0 @@
-/*
- * Copyright 2018 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- */
-
-#ifndef AMDGPU_DOORBELL_H
-#define AMDGPU_DOORBELL_H
-
-enum AMDGPU_DOORBELL_ASSIGNMENT {
-	AMDGPU_DOORBELL_KIQ                     = 0x000,
-	AMDGPU_DOORBELL_HIQ                     = 0x001,
-	AMDGPU_DOORBELL_DIQ                     = 0x002,
-	AMDGPU_DOORBELL_MEC_RING0               = 0x010,
-	AMDGPU_DOORBELL_MEC_RING1               = 0x011,
-	AMDGPU_DOORBELL_MEC_RING2               = 0x012,
-	AMDGPU_DOORBELL_MEC_RING3               = 0x013,
-	AMDGPU_DOORBELL_MEC_RING4               = 0x014,
-	AMDGPU_DOORBELL_MEC_RING5               = 0x015,
-	AMDGPU_DOORBELL_MEC_RING6               = 0x016,
-	AMDGPU_DOORBELL_MEC_RING7               = 0x017,
-	AMDGPU_DOORBELL_GFX_RING0               = 0x020,
-	AMDGPU_DOORBELL_sDMA_ENGINE0            = 0x1E0,
-	AMDGPU_DOORBELL_sDMA_ENGINE1            = 0x1E1,
-	AMDGPU_DOORBELL_IH                      = 0x1E8,
-	AMDGPU_DOORBELL_MAX_ASSIGNMENT          = 0x3FF,
-	AMDGPU_DOORBELL_INVALID                 = 0xFFFF
-};
-
-enum AMDGPU_VEGA20_DOORBELL_ASSIGNMENT {
-
-	/* Compute + GFX: 0~255 */
-	AMDGPU_VEGA20_DOORBELL_KIQ                     = 0x000,
-	AMDGPU_VEGA20_DOORBELL_HIQ                     = 0x001,
-	AMDGPU_VEGA20_DOORBELL_DIQ                     = 0x002,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING0               = 0x003,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING1               = 0x004,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING2               = 0x005,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING3               = 0x006,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING4               = 0x007,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING5               = 0x008,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING6               = 0x009,
-	AMDGPU_VEGA20_DOORBELL_MEC_RING7               = 0x00A,
-	AMDGPU_VEGA20_DOORBELL_USERQUEUE_START	       = 0x00B,
-	AMDGPU_VEGA20_DOORBELL_USERQUEUE_END	       = 0x08A,
-	AMDGPU_VEGA20_DOORBELL_GFX_RING0               = 0x08B,
-	/* SDMA:256~335*/
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE0            = 0x100,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE1            = 0x10A,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE2            = 0x114,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE3            = 0x11E,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE4            = 0x128,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE5            = 0x132,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE6            = 0x13C,
-	AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE7            = 0x146,
-	/* IH: 376~391 */
-	AMDGPU_VEGA20_DOORBELL_IH                      = 0x178,
-	/* MMSCH: 392~407
-	 * overlap the doorbell assignment with VCN as they are  mutually exclusive
-	 * VCN engine's doorbell is 32 bit and two VCN ring share one QWORD
-	 */
-	AMDGPU_VEGA20_DOORBELL64_VCN0_1                  = 0x188, /* VNC0 */
-	AMDGPU_VEGA20_DOORBELL64_VCN2_3                  = 0x189,
-	AMDGPU_VEGA20_DOORBELL64_VCN4_5                  = 0x18A,
-	AMDGPU_VEGA20_DOORBELL64_VCN6_7                  = 0x18B,
-
-	AMDGPU_VEGA20_DOORBELL64_VCN8_9                  = 0x18C, /* VNC1 */
-	AMDGPU_VEGA20_DOORBELL64_VCNa_b                  = 0x18D,
-	AMDGPU_VEGA20_DOORBELL64_VCNc_d                  = 0x18E,
-	AMDGPU_VEGA20_DOORBELL64_VCNe_f                  = 0x18F,
-
-	AMDGPU_VEGA20_DOORBELL64_UVD_RING0_1             = 0x188,
-	AMDGPU_VEGA20_DOORBELL64_UVD_RING2_3             = 0x189,
-	AMDGPU_VEGA20_DOORBELL64_UVD_RING4_5             = 0x18A,
-	AMDGPU_VEGA20_DOORBELL64_UVD_RING6_7             = 0x18B,
-
-	AMDGPU_VEGA20_DOORBELL64_VCE_RING0_1             = 0x18C,
-	AMDGPU_VEGA20_DOORBELL64_VCE_RING2_3             = 0x18D,
-	AMDGPU_VEGA20_DOORBELL64_VCE_RING4_5             = 0x18E,
-	AMDGPU_VEGA20_DOORBELL64_VCE_RING6_7             = 0x18F,
-
-	AMDGPU_VEGA20_DOORBELL64_FIRST_NON_CP            = AMDGPU_VEGA20_DOORBELL_sDMA_ENGINE0,
-	AMDGPU_VEGA20_DOORBELL64_LAST_NON_CP             = AMDGPU_VEGA20_DOORBELL64_VCE_RING6_7,
-
-	/* kiq/kcq from second XCD. Max 8 XCDs */
-	AMDGPU_VEGA20_DOORBELL_XCC1_KIQ_START             = 0x190,
-	/* 8 compute rings per GC. Max to 0x1CE */
-	AMDGPU_VEGA20_DOORBELL_XCC1_MEC_RING0_START       = 0x197,
-
-	/* AID1 SDMA: 0x1D0 ~ 0x1F7 */
-	AMDGPU_VEGA20_DOORBELL_AID1_sDMA_START           = 0x1D0,
-
-	AMDGPU_VEGA20_DOORBELL_MAX_ASSIGNMENT            = 0x1F7,
-	AMDGPU_VEGA20_DOORBELL_INVALID                   = 0xFFFF
-};
-
-enum AMDGPU_NAVI10_DOORBELL_ASSIGNMENT {
-
-	/* Compute + GFX: 0~255 */
-	AMDGPU_NAVI10_DOORBELL_KIQ			= 0x000,
-	AMDGPU_NAVI10_DOORBELL_HIQ			= 0x001,
-	AMDGPU_NAVI10_DOORBELL_DIQ			= 0x002,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING0		= 0x003,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING1		= 0x004,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING2		= 0x005,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING3		= 0x006,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING4		= 0x007,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING5		= 0x008,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING6		= 0x009,
-	AMDGPU_NAVI10_DOORBELL_MEC_RING7		= 0x00A,
-	AMDGPU_NAVI10_DOORBELL_MES_RING0	        = 0x00B,
-	AMDGPU_NAVI10_DOORBELL_MES_RING1		= 0x00C,
-	AMDGPU_NAVI10_DOORBELL_USERQUEUE_START		= 0x00D,
-	AMDGPU_NAVI10_DOORBELL_USERQUEUE_END		= 0x08A,
-	AMDGPU_NAVI10_DOORBELL_GFX_RING0		= 0x08B,
-	AMDGPU_NAVI10_DOORBELL_GFX_RING1		= 0x08C,
-	AMDGPU_NAVI10_DOORBELL_GFX_USERQUEUE_START	= 0x08D,
-	AMDGPU_NAVI10_DOORBELL_GFX_USERQUEUE_END	= 0x0FF,
-
-	/* SDMA:256~335*/
-	AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0		= 0x100,
-	AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE1		= 0x10A,
-	AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE2		= 0x114,
-	AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE3		= 0x11E,
-	/* IH: 376~391 */
-	AMDGPU_NAVI10_DOORBELL_IH			= 0x178,
-	/* MMSCH: 392~407
-	 * overlap the doorbell assignment with VCN as they are  mutually exclusive
-	 * VCE engine's doorbell is 32 bit and two VCE ring share one QWORD
-	 */
-	AMDGPU_NAVI10_DOORBELL64_VCN0_1			= 0x188, /* lower 32 bits for VNC0 and upper 32 bits for VNC1 */
-	AMDGPU_NAVI10_DOORBELL64_VCN2_3			= 0x189,
-	AMDGPU_NAVI10_DOORBELL64_VCN4_5			= 0x18A,
-	AMDGPU_NAVI10_DOORBELL64_VCN6_7			= 0x18B,
-
-	AMDGPU_NAVI10_DOORBELL64_VCN8_9			= 0x18C,
-	AMDGPU_NAVI10_DOORBELL64_VCNa_b			= 0x18D,
-	AMDGPU_NAVI10_DOORBELL64_VCNc_d			= 0x18E,
-	AMDGPU_NAVI10_DOORBELL64_VCNe_f			= 0x18F,
-
-	AMDGPU_NAVI10_DOORBELL64_VPE			= 0x190,
-
-	AMDGPU_NAVI10_DOORBELL64_FIRST_NON_CP		= AMDGPU_NAVI10_DOORBELL_sDMA_ENGINE0,
-	AMDGPU_NAVI10_DOORBELL64_LAST_NON_CP		= AMDGPU_NAVI10_DOORBELL64_VPE,
-
-	AMDGPU_NAVI10_DOORBELL_MAX_ASSIGNMENT		= AMDGPU_NAVI10_DOORBELL64_VPE,
-	AMDGPU_NAVI10_DOORBELL_INVALID			= 0xFFFF
-};
-
-/*
- * 64bit doorbell, offset are in QWORD, occupy 2KB doorbell space
- */
-enum AMDGPU_DOORBELL64_ASSIGNMENT {
-	/*
-	 * All compute related doorbells: kiq, hiq, diq, traditional compute queue, user queue, should locate in
-	 * a continues range so that programming CP_MEC_DOORBELL_RANGE_LOWER/UPPER can cover this range.
-	 *  Compute related doorbells are allocated from 0x00 to 0x8a
-	 */
-
-
-	/* kernel scheduling */
-	AMDGPU_DOORBELL64_KIQ                     = 0x00,
-
-	/* HSA interface queue and debug queue */
-	AMDGPU_DOORBELL64_HIQ                     = 0x01,
-	AMDGPU_DOORBELL64_DIQ                     = 0x02,
-
-	/* Compute engines */
-	AMDGPU_DOORBELL64_MEC_RING0               = 0x03,
-	AMDGPU_DOORBELL64_MEC_RING1               = 0x04,
-	AMDGPU_DOORBELL64_MEC_RING2               = 0x05,
-	AMDGPU_DOORBELL64_MEC_RING3               = 0x06,
-	AMDGPU_DOORBELL64_MEC_RING4               = 0x07,
-	AMDGPU_DOORBELL64_MEC_RING5               = 0x08,
-	AMDGPU_DOORBELL64_MEC_RING6               = 0x09,
-	AMDGPU_DOORBELL64_MEC_RING7               = 0x0a,
-
-	/* User queue doorbell range (128 doorbells) */
-	AMDGPU_DOORBELL64_USERQUEUE_START         = 0x0b,
-	AMDGPU_DOORBELL64_USERQUEUE_END           = 0x8a,
-
-	/* Graphics engine */
-	AMDGPU_DOORBELL64_GFX_RING0               = 0x8b,
-
-	/*
-	 * Other graphics doorbells can be allocated here: from 0x8c to 0xdf
-	 * Graphics voltage island aperture 1
-	 * default non-graphics QWORD index is 0xe0 - 0xFF inclusive
-	 */
-
-	/* For vega10 sriov, the sdma doorbell must be fixed as follow
-	 * to keep the same setting with host driver, or it will
-	 * happen conflicts
-	 */
-	AMDGPU_DOORBELL64_sDMA_ENGINE0            = 0xF0,
-	AMDGPU_DOORBELL64_sDMA_HI_PRI_ENGINE0     = 0xF1,
-	AMDGPU_DOORBELL64_sDMA_ENGINE1            = 0xF2,
-	AMDGPU_DOORBELL64_sDMA_HI_PRI_ENGINE1     = 0xF3,
-
-	/* Interrupt handler */
-	AMDGPU_DOORBELL64_IH                      = 0xF4,  /* For legacy interrupt ring buffer */
-	AMDGPU_DOORBELL64_IH_RING1                = 0xF5,  /* For page migration request log */
-	AMDGPU_DOORBELL64_IH_RING2                = 0xF6,  /* For page migration translation/invalidation log */
-
-	/* VCN engine use 32 bits doorbell  */
-	AMDGPU_DOORBELL64_VCN0_1                  = 0xF8, /* lower 32 bits for VNC0 and upper 32 bits for VNC1 */
-	AMDGPU_DOORBELL64_VCN2_3                  = 0xF9,
-	AMDGPU_DOORBELL64_VCN4_5                  = 0xFA,
-	AMDGPU_DOORBELL64_VCN6_7                  = 0xFB,
-
-	/* overlap the doorbell assignment with VCN as they are  mutually exclusive
-	 * VCE engine's doorbell is 32 bit and two VCE ring share one QWORD
-	 */
-	AMDGPU_DOORBELL64_UVD_RING0_1             = 0xF8,
-	AMDGPU_DOORBELL64_UVD_RING2_3             = 0xF9,
-	AMDGPU_DOORBELL64_UVD_RING4_5             = 0xFA,
-	AMDGPU_DOORBELL64_UVD_RING6_7             = 0xFB,
-
-	AMDGPU_DOORBELL64_VCE_RING0_1             = 0xFC,
-	AMDGPU_DOORBELL64_VCE_RING2_3             = 0xFD,
-	AMDGPU_DOORBELL64_VCE_RING4_5             = 0xFE,
-	AMDGPU_DOORBELL64_VCE_RING6_7             = 0xFF,
-
-	AMDGPU_DOORBELL64_FIRST_NON_CP            = AMDGPU_DOORBELL64_sDMA_ENGINE0,
-	AMDGPU_DOORBELL64_LAST_NON_CP             = AMDGPU_DOORBELL64_VCE_RING6_7,
-
-	AMDGPU_DOORBELL64_MAX_ASSIGNMENT          = 0xFF,
-	AMDGPU_DOORBELL64_INVALID                 = 0xFFFF
-};
-
-enum AMDGPU_DOORBELL_ASSIGNMENT_LAYOUT1 {
-
-	/* XCC0: 0x00 ~20, XCC1: 20 ~ 2F ... */
-
-	/* KIQ/HIQ/DIQ */
-	AMDGPU_DOORBELL_LAYOUT1_KIQ_START		= 0x000,
-	AMDGPU_DOORBELL_LAYOUT1_HIQ			= 0x001,
-	AMDGPU_DOORBELL_LAYOUT1_DIQ			= 0x002,
-	/* Compute: 0x08 ~ 0x20  */
-	AMDGPU_DOORBELL_LAYOUT1_MEC_RING_START		= 0x008,
-	AMDGPU_DOORBELL_LAYOUT1_MEC_RING_END		= 0x00F,
-	AMDGPU_DOORBELL_LAYOUT1_USERQUEUE_START		= 0x010,
-	AMDGPU_DOORBELL_LAYOUT1_USERQUEUE_END		= 0x01F,
-	AMDGPU_DOORBELL_LAYOUT1_XCC_RANGE		= 0x020,
-
-	/* SDMA: 0x100 ~ 0x19F */
-	AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_START	= 0x100,
-	AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_END		= 0x19F,
-	/* IH: 0x1A0 ~ 0x1AF */
-	AMDGPU_DOORBELL_LAYOUT1_IH                      = 0x1A0,
-	/* VCN: 0x1B0 ~ 0x1E8 */
-	AMDGPU_DOORBELL_LAYOUT1_VCN_START               = 0x1B0,
-	AMDGPU_DOORBELL_LAYOUT1_VCN_END                 = 0x1E8,
-
-	AMDGPU_DOORBELL_LAYOUT1_FIRST_NON_CP		= AMDGPU_DOORBELL_LAYOUT1_sDMA_ENGINE_START,
-	AMDGPU_DOORBELL_LAYOUT1_LAST_NON_CP		= AMDGPU_DOORBELL_LAYOUT1_VCN_END,
-
-	AMDGPU_DOORBELL_LAYOUT1_MAX_ASSIGNMENT          = 0x1E8,
-	AMDGPU_DOORBELL_LAYOUT1_INVALID                 = 0xFFFF
-};
-
-#endif

extra/amdpci/headers/amdgpu_irq.h

@@ -1,87 +0,0 @@
-/*
- * Copyright 2014 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- */
-
-#ifndef __AMDGPU_IRQ_H__
-#define __AMDGPU_IRQ_H__
-
-// #include <linux/irqdomain.h>
-// #include "soc15_ih_clientid.h"
-// #include "amdgpu_ih.h"
-
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
-#define bool _Bool
-#define u32 unsigned int
-
-#define AMDGPU_MAX_IRQ_SRC_ID		0x100
-#define AMDGPU_MAX_IRQ_CLIENT_ID	0x100
-
-#define AMDGPU_IRQ_CLIENTID_LEGACY	0
-#define AMDGPU_IRQ_CLIENTID_MAX		SOC15_IH_CLIENTID_MAX
-
-#define AMDGPU_IRQ_SRC_DATA_MAX_SIZE_DW	4
-
-struct amdgpu_device;
-
-enum amdgpu_interrupt_state {
-	AMDGPU_IRQ_STATE_DISABLE,
-	AMDGPU_IRQ_STATE_ENABLE,
-};
-
-struct amdgpu_iv_entry {
-	// struct amdgpu_ih_ring *ih;
-	unsigned client_id;
-	unsigned src_id;
-	unsigned ring_id;
-	unsigned vmid;
-	unsigned vmid_src;
-	uint64_t timestamp;
-	unsigned timestamp_src;
-	unsigned pasid;
-	unsigned node_id;
-	unsigned src_data[AMDGPU_IRQ_SRC_DATA_MAX_SIZE_DW];
-	const uint32_t *iv_entry;
-};
-
-enum interrupt_node_id_per_aid {
-	AID0_NODEID = 0,
-	XCD0_NODEID = 1,
-	XCD1_NODEID = 2,
-	AID1_NODEID = 4,
-	XCD2_NODEID = 5,
-	XCD3_NODEID = 6,
-	AID2_NODEID = 8,
-	XCD4_NODEID = 9,
-	XCD5_NODEID = 10,
-	AID3_NODEID = 12,
-	XCD6_NODEID = 13,
-	XCD7_NODEID = 14,
-	NODEID_MAX,
-};
-
-#endif

extra/amdpci/headers/amdgpu_psp.h

@@ -1,559 +0,0 @@
-/*
- * Copyright 2016 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- * Author: Huang Rui
- *
- */
-#ifndef __AMDGPU_PSP_H__
-#define __AMDGPU_PSP_H__
-
-// #include "amdgpu.h"
-// #include "psp_gfx_if.h"
-// #include "ta_xgmi_if.h"
-// #include "ta_ras_if.h"
-// #include "ta_rap_if.h"
-// #include "ta_secureDisplay_if.h"
-
-#define PSP_FENCE_BUFFER_SIZE	0x1000
-#define PSP_CMD_BUFFER_SIZE	0x1000
-#define PSP_1_MEG		0x100000
-#define PSP_TMR_SIZE(adev)	((adev)->asic_type == CHIP_ALDEBARAN ? 0x800000 : 0x400000)
-#define PSP_TMR_ALIGNMENT	0x100000
-#define PSP_FW_NAME_LEN		0x24
-
-// extern const struct attribute_group amdgpu_flash_attr_group;
-
-enum psp_shared_mem_size {
-	PSP_ASD_SHARED_MEM_SIZE				= 0x0,
-	PSP_XGMI_SHARED_MEM_SIZE			= 0x4000,
-	PSP_RAS_SHARED_MEM_SIZE				= 0x4000,
-	PSP_HDCP_SHARED_MEM_SIZE			= 0x4000,
-	PSP_DTM_SHARED_MEM_SIZE				= 0x4000,
-	PSP_RAP_SHARED_MEM_SIZE				= 0x4000,
-	PSP_SECUREDISPLAY_SHARED_MEM_SIZE	= 0x4000,
-};
-
-enum ta_type_id {
-	TA_TYPE_XGMI = 1,
-	TA_TYPE_RAS,
-	TA_TYPE_HDCP,
-	TA_TYPE_DTM,
-	TA_TYPE_RAP,
-	TA_TYPE_SECUREDISPLAY,
-
-	TA_TYPE_MAX_INDEX,
-};
-
-struct psp_context;
-struct psp_xgmi_node_info;
-struct psp_xgmi_topology_info;
-struct psp_bin_desc;
-
-enum psp_bootloader_cmd {
-	PSP_BL__LOAD_SYSDRV		= 0x10000,
-	PSP_BL__LOAD_SOSDRV		= 0x20000,
-	PSP_BL__LOAD_KEY_DATABASE	= 0x80000,
-	PSP_BL__LOAD_SOCDRV             = 0xB0000,
-	PSP_BL__LOAD_DBGDRV             = 0xC0000,
-	PSP_BL__LOAD_HADDRV		= PSP_BL__LOAD_DBGDRV,
-	PSP_BL__LOAD_INTFDRV		= 0xD0000,
-	PSP_BL__LOAD_RASDRV		= 0xE0000,
-	PSP_BL__LOAD_IPKEYMGRDRV	= 0xF0000,
-	PSP_BL__DRAM_LONG_TRAIN		= 0x100000,
-	PSP_BL__DRAM_SHORT_TRAIN	= 0x200000,
-	PSP_BL__LOAD_TOS_SPL_TABLE	= 0x10000000,
-};
-
-enum psp_ring_type {
-	PSP_RING_TYPE__INVALID = 0,
-	/*
-	 * These values map to the way the PSP kernel identifies the
-	 * rings.
-	 */
-	PSP_RING_TYPE__UM = 1, /* User mode ring (formerly called RBI) */
-	PSP_RING_TYPE__KM = 2  /* Kernel mode ring (formerly called GPCOM) */
-};
-
-// struct psp_ring {
-// 	enum psp_ring_type		ring_type;
-// 	struct psp_gfx_rb_frame		*ring_mem;
-// 	uint64_t			ring_mem_mc_addr;
-// 	void				*ring_mem_handle;
-// 	uint32_t			ring_size;
-// 	uint32_t			ring_wptr;
-// };
-
-/* More registers may will be supported */
-enum psp_reg_prog_id {
-	PSP_REG_IH_RB_CNTL        = 0,  /* register IH_RB_CNTL */
-	PSP_REG_IH_RB_CNTL_RING1  = 1,  /* register IH_RB_CNTL_RING1 */
-	PSP_REG_IH_RB_CNTL_RING2  = 2,  /* register IH_RB_CNTL_RING2 */
-	PSP_REG_LAST
-};
-
-// struct psp_funcs {
-// 	int (*init_microcode)(struct psp_context *psp);
-// 	int (*wait_for_bootloader)(struct psp_context *psp);
-// 	int (*bootloader_load_kdb)(struct psp_context *psp);
-// 	int (*bootloader_load_spl)(struct psp_context *psp);
-// 	int (*bootloader_load_sysdrv)(struct psp_context *psp);
-// 	int (*bootloader_load_soc_drv)(struct psp_context *psp);
-// 	int (*bootloader_load_intf_drv)(struct psp_context *psp);
-// 	int (*bootloader_load_dbg_drv)(struct psp_context *psp);
-// 	int (*bootloader_load_ras_drv)(struct psp_context *psp);
-// 	int (*bootloader_load_ipkeymgr_drv)(struct psp_context *psp);
-// 	int (*bootloader_load_sos)(struct psp_context *psp);
-// 	int (*ring_create)(struct psp_context *psp,
-// 			   enum psp_ring_type ring_type);
-// 	int (*ring_stop)(struct psp_context *psp,
-// 			    enum psp_ring_type ring_type);
-// 	int (*ring_destroy)(struct psp_context *psp,
-// 			    enum psp_ring_type ring_type);
-// 	bool (*smu_reload_quirk)(struct psp_context *psp);
-// 	int (*mode1_reset)(struct psp_context *psp);
-// 	int (*mem_training)(struct psp_context *psp, uint32_t ops);
-// 	uint32_t (*ring_get_wptr)(struct psp_context *psp);
-// 	void (*ring_set_wptr)(struct psp_context *psp, uint32_t value);
-// 	int (*load_usbc_pd_fw)(struct psp_context *psp, uint64_t fw_pri_mc_addr);
-// 	int (*read_usbc_pd_fw)(struct psp_context *psp, uint32_t *fw_ver);
-// 	int (*update_spirom)(struct psp_context *psp, uint64_t fw_pri_mc_addr);
-// 	int (*vbflash_stat)(struct psp_context *psp);
-// 	int (*fatal_error_recovery_quirk)(struct psp_context *psp);
-// 	bool (*get_ras_capability)(struct psp_context *psp);
-// 	bool (*is_aux_sos_load_required)(struct psp_context *psp);
-// };
-
-// struct ta_funcs {
-// 	int (*fn_ta_initialize)(struct psp_context *psp);
-// 	int (*fn_ta_invoke)(struct psp_context *psp, uint32_t ta_cmd_id);
-// 	int (*fn_ta_terminate)(struct psp_context *psp);
-// };
-
-#define AMDGPU_XGMI_MAX_CONNECTED_NODES		64
-// struct psp_xgmi_node_info {
-// 	uint64_t				node_id;
-// 	uint8_t					num_hops;
-// 	uint8_t					is_sharing_enabled;
-// 	enum ta_xgmi_assigned_sdma_engine	sdma_engine;
-// 	uint8_t					num_links;
-// 	struct xgmi_connected_port_num		port_num[TA_XGMI__MAX_PORT_NUM];
-// };
-
-// struct psp_xgmi_topology_info {
-// 	uint32_t			num_nodes;
-// 	struct psp_xgmi_node_info	nodes[AMDGPU_XGMI_MAX_CONNECTED_NODES];
-// };
-
-// struct psp_bin_desc {
-// 	uint32_t fw_version;
-// 	uint32_t feature_version;
-// 	uint32_t size_bytes;
-// 	uint8_t *start_addr;
-// };
-
-// struct ta_mem_context {
-// 	struct amdgpu_bo		*shared_bo;
-// 	uint64_t		shared_mc_addr;
-// 	void			*shared_buf;
-// 	enum psp_shared_mem_size	shared_mem_size;
-// };
-
-// struct ta_context {
-// 	bool			initialized;
-// 	uint32_t		session_id;
-// 	uint32_t		resp_status;
-// 	struct ta_mem_context	mem_context;
-// 	struct psp_bin_desc		bin_desc;
-// 	enum psp_gfx_cmd_id		ta_load_type;
-// 	enum ta_type_id		ta_type;
-// };
-
-// struct ta_cp_context {
-// 	struct ta_context		context;
-// 	struct mutex			mutex;
-// };
-
-// struct psp_xgmi_context {
-// 	struct ta_context		context;
-// 	struct psp_xgmi_topology_info	top_info;
-// 	bool				supports_extended_data;
-// 	uint8_t				xgmi_ta_caps;
-// };
-
-// struct psp_ras_context {
-// 	struct ta_context		context;
-// 	struct amdgpu_ras		*ras;
-// };
-
-#define MEM_TRAIN_SYSTEM_SIGNATURE		0x54534942
-#define GDDR6_MEM_TRAINING_DATA_SIZE_IN_BYTES	0x1000
-#define GDDR6_MEM_TRAINING_OFFSET		0x8000
-/*Define the VRAM size that will be encroached by BIST training.*/
-#define BIST_MEM_TRAINING_ENCROACHED_SIZE	0x2000000
-
-enum psp_memory_training_init_flag {
-	PSP_MEM_TRAIN_NOT_SUPPORT	= 0x0,
-	PSP_MEM_TRAIN_SUPPORT		= 0x1,
-	PSP_MEM_TRAIN_INIT_FAILED	= 0x2,
-	PSP_MEM_TRAIN_RESERVE_SUCCESS	= 0x4,
-	PSP_MEM_TRAIN_INIT_SUCCESS	= 0x8,
-};
-
-enum psp_memory_training_ops {
-	PSP_MEM_TRAIN_SEND_LONG_MSG	= 0x1,
-	PSP_MEM_TRAIN_SAVE		= 0x2,
-	PSP_MEM_TRAIN_RESTORE		= 0x4,
-	PSP_MEM_TRAIN_SEND_SHORT_MSG	= 0x8,
-	PSP_MEM_TRAIN_COLD_BOOT		= PSP_MEM_TRAIN_SEND_LONG_MSG,
-	PSP_MEM_TRAIN_RESUME		= PSP_MEM_TRAIN_SEND_SHORT_MSG,
-};
-
-// struct psp_memory_training_context {
-// 	/*training data size*/
-// 	u64 train_data_size;
-// 	/*
-// 	 * sys_cache
-// 	 * cpu virtual address
-// 	 * system memory buffer that used to store the training data.
-// 	 */
-// 	void *sys_cache;
-
-// 	/*vram offset of the p2c training data*/
-// 	u64 p2c_train_data_offset;
-
-// 	/*vram offset of the c2p training data*/
-// 	u64 c2p_train_data_offset;
-// 	struct amdgpu_bo *c2p_bo;
-
-// 	enum psp_memory_training_init_flag init;
-// 	u32 training_cnt;
-// 	bool enable_mem_training;
-// };
-
-/** PSP runtime DB **/
-#define PSP_RUNTIME_DB_SIZE_IN_BYTES		0x10000
-#define PSP_RUNTIME_DB_OFFSET			0x100000
-#define PSP_RUNTIME_DB_COOKIE_ID		0x0ed5
-#define PSP_RUNTIME_DB_VER_1			0x0100
-#define PSP_RUNTIME_DB_DIAG_ENTRY_MAX_COUNT	0x40
-
-enum psp_runtime_entry_type {
-	PSP_RUNTIME_ENTRY_TYPE_INVALID		= 0x0,
-	PSP_RUNTIME_ENTRY_TYPE_TEST		= 0x1,
-	PSP_RUNTIME_ENTRY_TYPE_MGPU_COMMON	= 0x2,  /* Common mGPU runtime data */
-	PSP_RUNTIME_ENTRY_TYPE_MGPU_WAFL	= 0x3,  /* WAFL runtime data */
-	PSP_RUNTIME_ENTRY_TYPE_MGPU_XGMI	= 0x4,  /* XGMI runtime data */
-	PSP_RUNTIME_ENTRY_TYPE_BOOT_CONFIG	= 0x5,  /* Boot Config runtime data */
-	PSP_RUNTIME_ENTRY_TYPE_PPTABLE_ERR_STATUS = 0x6, /* SCPM validation data */
-};
-
-/* PSP runtime DB header */
-// struct psp_runtime_data_header {
-// 	/* determine the existence of runtime db */
-// 	uint16_t cookie;
-// 	/* version of runtime db */
-// 	uint16_t version;
-// };
-
-// /* PSP runtime DB entry */
-// struct psp_runtime_entry {
-// 	/* type of runtime db entry */
-// 	uint32_t entry_type;
-// 	/* offset of entry in bytes */
-// 	uint16_t offset;
-// 	/* size of entry in bytes */
-// 	uint16_t size;
-// };
-
-// /* PSP runtime DB directory */
-// struct psp_runtime_data_directory {
-// 	/* number of valid entries */
-// 	uint16_t			entry_count;
-// 	/* db entries*/
-// 	struct psp_runtime_entry	entry_list[PSP_RUNTIME_DB_DIAG_ENTRY_MAX_COUNT];
-// };
-
-/* PSP runtime DB boot config feature bitmask */
-enum psp_runtime_boot_cfg_feature {
-	BOOT_CFG_FEATURE_GECC                       = 0x1,
-	BOOT_CFG_FEATURE_TWO_STAGE_DRAM_TRAINING    = 0x2,
-};
-
-/* PSP run time DB SCPM authentication defines */
-enum psp_runtime_scpm_authentication {
-	SCPM_DISABLE                     = 0x0,
-	SCPM_ENABLE                      = 0x1,
-	SCPM_ENABLE_WITH_SCPM_ERR        = 0x2,
-};
-
-/* PSP runtime DB boot config entry */
-// struct psp_runtime_boot_cfg_entry {
-// 	uint32_t boot_cfg_bitmask;
-// 	uint32_t reserved;
-// };
-
-// /* PSP runtime DB SCPM entry */
-// struct psp_runtime_scpm_entry {
-// 	enum psp_runtime_scpm_authentication scpm_status;
-// };
-
-// struct psp_context {
-// 	struct amdgpu_device		*adev;
-// 	struct psp_ring			km_ring;
-// 	struct psp_gfx_cmd_resp		*cmd;
-
-// 	const struct psp_funcs		*funcs;
-// 	const struct ta_funcs		*ta_funcs;
-
-// 	/* firmware buffer */
-// 	struct amdgpu_bo		*fw_pri_bo;
-// 	uint64_t			fw_pri_mc_addr;
-// 	void				*fw_pri_buf;
-
-// 	/* sos firmware */
-// 	const struct firmware		*sos_fw;
-// 	struct psp_bin_desc		sys;
-// 	struct psp_bin_desc		sos;
-// 	struct psp_bin_desc		toc;
-// 	struct psp_bin_desc		kdb;
-// 	struct psp_bin_desc		spl;
-// 	struct psp_bin_desc		rl;
-// 	struct psp_bin_desc		soc_drv;
-// 	struct psp_bin_desc		intf_drv;
-// 	struct psp_bin_desc		dbg_drv;
-// 	struct psp_bin_desc		ras_drv;
-// 	struct psp_bin_desc		ipkeymgr_drv;
-
-// 	/* tmr buffer */
-// 	struct amdgpu_bo		*tmr_bo;
-// 	uint64_t			tmr_mc_addr;
-
-// 	/* asd firmware */
-// 	const struct firmware		*asd_fw;
-
-// 	/* toc firmware */
-// 	const struct firmware		*toc_fw;
-
-// 	/* cap firmware */
-// 	const struct firmware		*cap_fw;
-
-// 	/* fence buffer */
-// 	struct amdgpu_bo		*fence_buf_bo;
-// 	uint64_t			fence_buf_mc_addr;
-// 	void				*fence_buf;
-
-// 	/* cmd buffer */
-// 	struct amdgpu_bo		*cmd_buf_bo;
-// 	uint64_t			cmd_buf_mc_addr;
-// 	struct psp_gfx_cmd_resp		*cmd_buf_mem;
-
-// 	/* fence value associated with cmd buffer */
-// 	atomic_t			fence_value;
-// 	/* flag to mark whether gfx fw autoload is supported or not */
-// 	bool				autoload_supported;
-// 	/* flag to mark whether psp use runtime TMR or boottime TMR */
-// 	bool				boot_time_tmr;
-// 	/* flag to mark whether df cstate management centralized to PMFW */
-// 	bool				pmfw_centralized_cstate_management;
-
-// 	/* xgmi ta firmware and buffer */
-// 	const struct firmware		*ta_fw;
-// 	uint32_t			ta_fw_version;
-
-// 	uint32_t			cap_fw_version;
-// 	uint32_t			cap_feature_version;
-// 	uint32_t			cap_ucode_size;
-
-// 	struct ta_context		asd_context;
-// 	struct psp_xgmi_context		xgmi_context;
-// 	struct psp_ras_context		ras_context;
-// 	struct ta_cp_context		hdcp_context;
-// 	struct ta_cp_context		dtm_context;
-// 	struct ta_cp_context		rap_context;
-// 	struct ta_cp_context		securedisplay_context;
-// 	struct mutex			mutex;
-// 	struct psp_memory_training_context mem_train_ctx;
-
-// 	uint32_t			boot_cfg_bitmask;
-
-// 	/* firmware upgrades supported */
-// 	bool				sup_pd_fw_up;
-// 	bool				sup_ifwi_up;
-
-// 	char				*vbflash_tmp_buf;
-// 	size_t				vbflash_image_size;
-// 	bool				vbflash_done;
-// };
-
-// struct amdgpu_psp_funcs {
-// 	bool (*check_fw_loading_status)(struct amdgpu_device *adev,
-// 					enum AMDGPU_UCODE_ID);
-// };
-
-
-// #define psp_ring_create(psp, type) (psp)->funcs->ring_create((psp), (type))
-// #define psp_ring_stop(psp, type) (psp)->funcs->ring_stop((psp), (type))
-// #define psp_ring_destroy(psp, type) ((psp)->funcs->ring_destroy((psp), (type)))
-// #define psp_init_microcode(psp) \
-// 		((psp)->funcs->init_microcode ? (psp)->funcs->init_microcode((psp)) : 0)
-// #define psp_bootloader_load_kdb(psp) \
-// 		((psp)->funcs->bootloader_load_kdb ? (psp)->funcs->bootloader_load_kdb((psp)) : 0)
-// #define psp_bootloader_load_spl(psp) \
-// 		((psp)->funcs->bootloader_load_spl ? (psp)->funcs->bootloader_load_spl((psp)) : 0)
-// #define psp_bootloader_load_sysdrv(psp) \
-// 		((psp)->funcs->bootloader_load_sysdrv ? (psp)->funcs->bootloader_load_sysdrv((psp)) : 0)
-// #define psp_bootloader_load_soc_drv(psp) \
-// 		((psp)->funcs->bootloader_load_soc_drv ? (psp)->funcs->bootloader_load_soc_drv((psp)) : 0)
-// #define psp_bootloader_load_intf_drv(psp) \
-// 		((psp)->funcs->bootloader_load_intf_drv ? (psp)->funcs->bootloader_load_intf_drv((psp)) : 0)
-// #define psp_bootloader_load_dbg_drv(psp) \
-// 		((psp)->funcs->bootloader_load_dbg_drv ? (psp)->funcs->bootloader_load_dbg_drv((psp)) : 0)
-// #define psp_bootloader_load_ras_drv(psp) \
-// 		((psp)->funcs->bootloader_load_ras_drv ? \
-// 		(psp)->funcs->bootloader_load_ras_drv((psp)) : 0)
-// #define psp_bootloader_load_ipkeymgr_drv(psp) \
-// 		((psp)->funcs->bootloader_load_ipkeymgr_drv ? \
-// 		 (psp)->funcs->bootloader_load_ipkeymgr_drv((psp)) : 0)
-// #define psp_bootloader_load_sos(psp) \
-// 		((psp)->funcs->bootloader_load_sos ? (psp)->funcs->bootloader_load_sos((psp)) : 0)
-// #define psp_smu_reload_quirk(psp) \
-// 		((psp)->funcs->smu_reload_quirk ? (psp)->funcs->smu_reload_quirk((psp)) : false)
-// #define psp_mode1_reset(psp) \
-// 		((psp)->funcs->mode1_reset ? (psp)->funcs->mode1_reset((psp)) : false)
-// #define psp_mem_training(psp, ops) \
-// 	((psp)->funcs->mem_training ? (psp)->funcs->mem_training((psp), (ops)) : 0)
-
-// #define psp_ring_get_wptr(psp) (psp)->funcs->ring_get_wptr((psp))
-// #define psp_ring_set_wptr(psp, value) (psp)->funcs->ring_set_wptr((psp), (value))
-
-// #define psp_load_usbc_pd_fw(psp, fw_pri_mc_addr) \
-// 	((psp)->funcs->load_usbc_pd_fw ? \
-// 	(psp)->funcs->load_usbc_pd_fw((psp), (fw_pri_mc_addr)) : -EINVAL)
-
-// #define psp_read_usbc_pd_fw(psp, fw_ver) \
-// 	((psp)->funcs->read_usbc_pd_fw ? \
-// 	(psp)->funcs->read_usbc_pd_fw((psp), fw_ver) : -EINVAL)
-
-// #define psp_update_spirom(psp, fw_pri_mc_addr) \
-// 	((psp)->funcs->update_spirom ? \
-// 	(psp)->funcs->update_spirom((psp), fw_pri_mc_addr) : -EINVAL)
-
-// #define psp_vbflash_status(psp) \
-// 	((psp)->funcs->vbflash_stat ? \
-// 	(psp)->funcs->vbflash_stat((psp)) : -EINVAL)
-
-// #define psp_fatal_error_recovery_quirk(psp) \
-// 	((psp)->funcs->fatal_error_recovery_quirk ? \
-// 	(psp)->funcs->fatal_error_recovery_quirk((psp)) : 0)
-
-// #define psp_is_aux_sos_load_required(psp) \
-// 	((psp)->funcs->is_aux_sos_load_required ? (psp)->funcs->is_aux_sos_load_required((psp)) : 0)
-
-// extern const struct amd_ip_funcs psp_ip_funcs;
-
-// extern const struct amdgpu_ip_block_version psp_v3_1_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v10_0_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v11_0_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v11_0_8_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v12_0_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v13_0_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v13_0_4_ip_block;
-// extern const struct amdgpu_ip_block_version psp_v14_0_ip_block;
-
-// extern int psp_wait_for(struct psp_context *psp, uint32_t reg_index,
-// 			uint32_t field_val, uint32_t mask, bool check_changed);
-// extern int psp_wait_for_spirom_update(struct psp_context *psp, uint32_t reg_index,
-// 			uint32_t field_val, uint32_t mask, uint32_t msec_timeout);
-
-// int psp_execute_ip_fw_load(struct psp_context *psp,
-// 			   struct amdgpu_firmware_info *ucode);
-
-// int psp_gpu_reset(struct amdgpu_device *adev);
-
-// int psp_ta_init_shared_buf(struct psp_context *psp,
-// 				  struct ta_mem_context *mem_ctx);
-// void psp_ta_free_shared_buf(struct ta_mem_context *mem_ctx);
-// int psp_ta_unload(struct psp_context *psp, struct ta_context *context);
-// int psp_ta_load(struct psp_context *psp, struct ta_context *context);
-// int psp_ta_invoke(struct psp_context *psp,
-// 			uint32_t ta_cmd_id,
-// 			struct ta_context *context);
-
-// int psp_xgmi_initialize(struct psp_context *psp, bool set_extended_data, bool load_ta);
-// int psp_xgmi_terminate(struct psp_context *psp);
-// int psp_xgmi_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
-// int psp_xgmi_get_hive_id(struct psp_context *psp, uint64_t *hive_id);
-// int psp_xgmi_get_node_id(struct psp_context *psp, uint64_t *node_id);
-// int psp_xgmi_get_topology_info(struct psp_context *psp,
-// 			       int number_devices,
-// 			       struct psp_xgmi_topology_info *topology,
-// 			       bool get_extended_data);
-// int psp_xgmi_set_topology_info(struct psp_context *psp,
-// 			       int number_devices,
-// 			       struct psp_xgmi_topology_info *topology);
-// int psp_ras_initialize(struct psp_context *psp);
-// int psp_ras_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
-// int psp_ras_enable_features(struct psp_context *psp,
-// 		union ta_ras_cmd_input *info, bool enable);
-// int psp_ras_trigger_error(struct psp_context *psp,
-// 			  struct ta_ras_trigger_error_input *info, uint32_t instance_mask);
-// int psp_ras_terminate(struct psp_context *psp);
-// int psp_ras_query_address(struct psp_context *psp,
-// 			  struct ta_ras_query_address_input *addr_in,
-// 			  struct ta_ras_query_address_output *addr_out);
-
-// int psp_hdcp_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
-// int psp_dtm_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
-// int psp_rap_invoke(struct psp_context *psp, uint32_t ta_cmd_id, enum ta_rap_status *status);
-// int psp_securedisplay_invoke(struct psp_context *psp, uint32_t ta_cmd_id);
-
-// int psp_rlc_autoload_start(struct psp_context *psp);
-
-// int psp_reg_program(struct psp_context *psp, enum psp_reg_prog_id reg,
-// 		uint32_t value);
-// int psp_ring_cmd_submit(struct psp_context *psp,
-// 			uint64_t cmd_buf_mc_addr,
-// 			uint64_t fence_mc_addr,
-// 			int index);
-// int psp_init_asd_microcode(struct psp_context *psp,
-// 			   const char *chip_name);
-// int psp_init_toc_microcode(struct psp_context *psp,
-// 			   const char *chip_name);
-// int psp_init_sos_microcode(struct psp_context *psp,
-// 			   const char *chip_name);
-// int psp_init_ta_microcode(struct psp_context *psp,
-// 			  const char *chip_name);
-// int psp_init_cap_microcode(struct psp_context *psp,
-// 			  const char *chip_name);
-// int psp_get_fw_attestation_records_addr(struct psp_context *psp,
-// 					uint64_t *output_ptr);
-
-// int psp_load_fw_list(struct psp_context *psp,
-// 		     struct amdgpu_firmware_info **ucode_list, int ucode_count);
-// void psp_copy_fw(struct psp_context *psp, uint8_t *start_addr, uint32_t bin_size);
-
-// int psp_spatial_partition(struct psp_context *psp, int mode);
-
-// int is_psp_fw_valid(struct psp_bin_desc bin);
-
-// int amdgpu_psp_wait_for_bootloader(struct amdgpu_device *adev);
-// bool amdgpu_psp_get_ras_capability(struct psp_context *psp);
-#endif

extra/amdpci/headers/amdgpu_smu.h

@@ -1,347 +0,0 @@
-/*
- * Copyright 2019 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- */
-#ifndef __AMDGPU_SMU_H__
-#define __AMDGPU_SMU_H__
-
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
-#define bool _Bool
-#define u32 unsigned int
-
-#define SMU_THERMAL_MINIMUM_ALERT_TEMP		0
-#define SMU_THERMAL_MAXIMUM_ALERT_TEMP		255
-#define SMU_TEMPERATURE_UNITS_PER_CENTIGRADES	1000
-#define SMU_FW_NAME_LEN			0x24
-
-#define SMU_DPM_USER_PROFILE_RESTORE (1 << 0)
-#define SMU_CUSTOM_FAN_SPEED_RPM     (1 << 1)
-#define SMU_CUSTOM_FAN_SPEED_PWM     (1 << 2)
-
-// Power Throttlers
-#define SMU_THROTTLER_PPT0_BIT			0
-#define SMU_THROTTLER_PPT1_BIT			1
-#define SMU_THROTTLER_PPT2_BIT			2
-#define SMU_THROTTLER_PPT3_BIT			3
-#define SMU_THROTTLER_SPL_BIT			4
-#define SMU_THROTTLER_FPPT_BIT			5
-#define SMU_THROTTLER_SPPT_BIT			6
-#define SMU_THROTTLER_SPPT_APU_BIT		7
-
-// Current Throttlers
-#define SMU_THROTTLER_TDC_GFX_BIT		16
-#define SMU_THROTTLER_TDC_SOC_BIT		17
-#define SMU_THROTTLER_TDC_MEM_BIT		18
-#define SMU_THROTTLER_TDC_VDD_BIT		19
-#define SMU_THROTTLER_TDC_CVIP_BIT		20
-#define SMU_THROTTLER_EDC_CPU_BIT		21
-#define SMU_THROTTLER_EDC_GFX_BIT		22
-#define SMU_THROTTLER_APCC_BIT			23
-
-// Temperature
-#define SMU_THROTTLER_TEMP_GPU_BIT		32
-#define SMU_THROTTLER_TEMP_CORE_BIT		33
-#define SMU_THROTTLER_TEMP_MEM_BIT		34
-#define SMU_THROTTLER_TEMP_EDGE_BIT		35
-#define SMU_THROTTLER_TEMP_HOTSPOT_BIT		36
-#define SMU_THROTTLER_TEMP_SOC_BIT		37
-#define SMU_THROTTLER_TEMP_VR_GFX_BIT		38
-#define SMU_THROTTLER_TEMP_VR_SOC_BIT		39
-#define SMU_THROTTLER_TEMP_VR_MEM0_BIT		40
-#define SMU_THROTTLER_TEMP_VR_MEM1_BIT		41
-#define SMU_THROTTLER_TEMP_LIQUID0_BIT		42
-#define SMU_THROTTLER_TEMP_LIQUID1_BIT		43
-#define SMU_THROTTLER_VRHOT0_BIT		44
-#define SMU_THROTTLER_VRHOT1_BIT		45
-#define SMU_THROTTLER_PROCHOT_CPU_BIT		46
-#define SMU_THROTTLER_PROCHOT_GFX_BIT		47
-
-// Other
-#define SMU_THROTTLER_PPM_BIT			56
-#define SMU_THROTTLER_FIT_BIT			57
-
-struct smu_hw_power_state {
-	unsigned int magic;
-};
-
-struct smu_power_state;
-
-enum smu_state_ui_label {
-	SMU_STATE_UI_LABEL_NONE,
-	SMU_STATE_UI_LABEL_BATTERY,
-	SMU_STATE_UI_TABEL_MIDDLE_LOW,
-	SMU_STATE_UI_LABEL_BALLANCED,
-	SMU_STATE_UI_LABEL_MIDDLE_HIGHT,
-	SMU_STATE_UI_LABEL_PERFORMANCE,
-	SMU_STATE_UI_LABEL_BACO,
-};
-
-enum smu_state_classification_flag {
-	SMU_STATE_CLASSIFICATION_FLAG_BOOT                     = 0x0001,
-	SMU_STATE_CLASSIFICATION_FLAG_THERMAL                  = 0x0002,
-	SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE      = 0x0004,
-	SMU_STATE_CLASSIFICATION_FLAG_RESET                    = 0x0008,
-	SMU_STATE_CLASSIFICATION_FLAG_FORCED                   = 0x0010,
-	SMU_STATE_CLASSIFICATION_FLAG_USER_3D_PERFORMANCE      = 0x0020,
-	SMU_STATE_CLASSIFICATION_FLAG_USER_2D_PERFORMANCE      = 0x0040,
-	SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE           = 0x0080,
-	SMU_STATE_CLASSIFICATION_FLAG_AC_OVERDIRVER_TEMPLATE   = 0x0100,
-	SMU_STATE_CLASSIFICATION_FLAG_UVD                      = 0x0200,
-	SMU_STATE_CLASSIFICATION_FLAG_3D_PERFORMANCE_LOW       = 0x0400,
-	SMU_STATE_CLASSIFICATION_FLAG_ACPI                     = 0x0800,
-	SMU_STATE_CLASSIFICATION_FLAG_HD2                      = 0x1000,
-	SMU_STATE_CLASSIFICATION_FLAG_UVD_HD                   = 0x2000,
-	SMU_STATE_CLASSIFICATION_FLAG_UVD_SD                   = 0x4000,
-	SMU_STATE_CLASSIFICATION_FLAG_USER_DC_PERFORMANCE      = 0x8000,
-	SMU_STATE_CLASSIFICATION_FLAG_DC_OVERDIRVER_TEMPLATE   = 0x10000,
-	SMU_STATE_CLASSIFICATION_FLAG_BACO                     = 0x20000,
-	SMU_STATE_CLASSIFICATIN_FLAG_LIMITED_POWER_SOURCE2      = 0x40000,
-	SMU_STATE_CLASSIFICATION_FLAG_ULV                      = 0x80000,
-	SMU_STATE_CLASSIFICATION_FLAG_UVD_MVC                  = 0x100000,
-};
-
-struct smu_state_classification_block {
-	enum smu_state_ui_label         ui_label;
-	enum smu_state_classification_flag  flags;
-	int                          bios_index;
-	bool                      temporary_state;
-	bool                      to_be_deleted;
-};
-
-struct smu_state_pcie_block {
-	unsigned int lanes;
-};
-
-enum smu_refreshrate_source {
-	SMU_REFRESHRATE_SOURCE_EDID,
-	SMU_REFRESHRATE_SOURCE_EXPLICIT
-};
-
-struct smu_state_display_block {
-	bool              disable_frame_modulation;
-	bool              limit_refreshrate;
-	enum smu_refreshrate_source refreshrate_source;
-	int                  explicit_refreshrate;
-	int                  edid_refreshrate_index;
-	bool              enable_vari_bright;
-};
-
-struct smu_state_memory_block {
-	bool              dll_off;
-	uint8_t                 m3arb;
-	uint8_t                 unused[3];
-};
-
-struct smu_state_software_algorithm_block {
-	bool disable_load_balancing;
-	bool enable_sleep_for_timestamps;
-};
-
-struct smu_temperature_range {
-	int min;
-	int max;
-	int edge_emergency_max;
-	int hotspot_min;
-	int hotspot_crit_max;
-	int hotspot_emergency_max;
-	int mem_min;
-	int mem_crit_max;
-	int mem_emergency_max;
-	int software_shutdown_temp;
-	int software_shutdown_temp_offset;
-};
-
-struct smu_state_validation_block {
-	bool single_display_only;
-	bool disallow_on_dc;
-	uint8_t supported_power_levels;
-};
-
-struct smu_uvd_clocks {
-	uint32_t vclk;
-	uint32_t dclk;
-};
-
-/**
-* Structure to hold a SMU Power State.
-*/
-
-enum smu_power_src_type {
-	SMU_POWER_SOURCE_AC,
-	SMU_POWER_SOURCE_DC,
-	SMU_POWER_SOURCE_COUNT,
-};
-
-enum smu_ppt_limit_type {
-	SMU_DEFAULT_PPT_LIMIT = 0,
-	SMU_FAST_PPT_LIMIT,
-};
-
-enum smu_ppt_limit_level {
-	SMU_PPT_LIMIT_MIN = -1,
-	SMU_PPT_LIMIT_CURRENT,
-	SMU_PPT_LIMIT_DEFAULT,
-	SMU_PPT_LIMIT_MAX,
-};
-
-enum smu_memory_pool_size {
-    SMU_MEMORY_POOL_SIZE_ZERO   = 0,
-    SMU_MEMORY_POOL_SIZE_256_MB = 0x10000000,
-    SMU_MEMORY_POOL_SIZE_512_MB = 0x20000000,
-    SMU_MEMORY_POOL_SIZE_1_GB   = 0x40000000,
-    SMU_MEMORY_POOL_SIZE_2_GB   = 0x80000000,
-};
-
-enum smu_clk_type {
-	SMU_GFXCLK,
-	SMU_VCLK,
-	SMU_DCLK,
-	SMU_VCLK1,
-	SMU_DCLK1,
-	SMU_ECLK,
-	SMU_SOCCLK,
-	SMU_UCLK,
-	SMU_DCEFCLK,
-	SMU_DISPCLK,
-	SMU_PIXCLK,
-	SMU_PHYCLK,
-	SMU_FCLK,
-	SMU_SCLK,
-	SMU_MCLK,
-	SMU_PCIE,
-	SMU_LCLK,
-	SMU_OD_CCLK,
-	SMU_OD_SCLK,
-	SMU_OD_MCLK,
-	SMU_OD_VDDC_CURVE,
-	SMU_OD_RANGE,
-	SMU_OD_VDDGFX_OFFSET,
-	SMU_OD_FAN_CURVE,
-	SMU_OD_ACOUSTIC_LIMIT,
-	SMU_OD_ACOUSTIC_TARGET,
-	SMU_OD_FAN_TARGET_TEMPERATURE,
-	SMU_OD_FAN_MINIMUM_PWM,
-	SMU_CLK_COUNT,
-};
-
-struct smu_user_dpm_profile {
-	uint32_t fan_mode;
-	uint32_t power_limit;
-	uint32_t fan_speed_pwm;
-	uint32_t fan_speed_rpm;
-	uint32_t flags;
-	uint32_t user_od;
-
-	/* user clock state information */
-	uint32_t clk_mask[SMU_CLK_COUNT];
-	uint32_t clk_dependency;
-};
-
-#define SMU_TABLE_INIT(tables, table_id, s, a, d)	\
-	do {						\
-		tables[table_id].size = s;		\
-		tables[table_id].align = a;		\
-		tables[table_id].domain = d;		\
-	} while (0)
-
-struct smu_table {
-	uint64_t size;
-	uint32_t align;
-	uint8_t domain;
-	uint64_t mc_address;
-	void *cpu_addr;
-	struct amdgpu_bo *bo;
-	uint32_t version;
-};
-
-enum smu_perf_level_designation {
-	PERF_LEVEL_ACTIVITY,
-	PERF_LEVEL_POWER_CONTAINMENT,
-};
-
-struct smu_performance_level {
-	uint32_t core_clock;
-	uint32_t memory_clock;
-	uint32_t vddc;
-	uint32_t vddci;
-	uint32_t non_local_mem_freq;
-	uint32_t non_local_mem_width;
-};
-
-struct smu_clock_info {
-	uint32_t min_mem_clk;
-	uint32_t max_mem_clk;
-	uint32_t min_eng_clk;
-	uint32_t max_eng_clk;
-	uint32_t min_bus_bandwidth;
-	uint32_t max_bus_bandwidth;
-};
-
-struct smu_bios_boot_up_values {
-	uint32_t			revision;
-	uint32_t			gfxclk;
-	uint32_t			uclk;
-	uint32_t			socclk;
-	uint32_t			dcefclk;
-	uint32_t			eclk;
-	uint32_t			vclk;
-	uint32_t			dclk;
-	uint16_t			vddc;
-	uint16_t			vddci;
-	uint16_t			mvddc;
-	uint16_t			vdd_gfx;
-	uint8_t				cooling_id;
-	uint32_t			pp_table_id;
-	uint32_t			format_revision;
-	uint32_t			content_revision;
-	uint32_t			fclk;
-	uint32_t			lclk;
-	uint32_t			firmware_caps;
-};
-
-enum smu_table_id {
-	SMU_TABLE_PPTABLE = 0,
-	SMU_TABLE_WATERMARKS,
-	SMU_TABLE_CUSTOM_DPM,
-	SMU_TABLE_DPMCLOCKS,
-	SMU_TABLE_AVFS,
-	SMU_TABLE_AVFS_PSM_DEBUG,
-	SMU_TABLE_AVFS_FUSE_OVERRIDE,
-	SMU_TABLE_PMSTATUSLOG,
-	SMU_TABLE_SMU_METRICS,
-	SMU_TABLE_DRIVER_SMU_CONFIG,
-	SMU_TABLE_ACTIVITY_MONITOR_COEFF,
-	SMU_TABLE_OVERDRIVE,
-	SMU_TABLE_I2C_COMMANDS,
-	SMU_TABLE_PACE,
-	SMU_TABLE_ECCINFO,
-	SMU_TABLE_COMBO_PPTABLE,
-	SMU_TABLE_WIFIBAND,
-	SMU_TABLE_COUNT,
-};
-
-
-#endif

extra/amdpci/headers/amdgpu_ucode.h

@@ -1,634 +0,0 @@
-/*
- * Copyright 2012 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- */
-#ifndef __AMDGPU_UCODE_H__
-#define __AMDGPU_UCODE_H__
-
-// #include "amdgpu_socbb.h"
-#define int32_t int
-#define uint32_t unsigned int
-#define int8_t signed char
-#define uint8_t unsigned char
-#define uint16_t unsigned short
-#define int16_t short
-#define uint64_t unsigned long long
-#define bool _Bool
-#define u32 unsigned int
-
-struct common_firmware_header {
-	uint32_t size_bytes; /* size of the entire header+image(s) in bytes */
-	uint32_t header_size_bytes; /* size of just the header in bytes */
-	uint16_t header_version_major; /* header version */
-	uint16_t header_version_minor; /* header version */
-	uint16_t ip_version_major; /* IP version */
-	uint16_t ip_version_minor; /* IP version */
-	uint32_t ucode_version;
-	uint32_t ucode_size_bytes; /* size of ucode in bytes */
-	uint32_t ucode_array_offset_bytes; /* payload offset from the start of the header */
-	uint32_t crc32;  /* crc32 checksum of the payload */
-};
-
-/* version_major=1, version_minor=0 */
-struct mc_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t io_debug_size_bytes; /* size of debug array in dwords */
-	uint32_t io_debug_array_offset_bytes; /* payload offset from the start of the header */
-};
-
-/* version_major=1, version_minor=0 */
-struct smc_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_start_addr;
-};
-
-/* version_major=2, version_minor=0 */
-struct smc_firmware_header_v2_0 {
-	struct smc_firmware_header_v1_0 v1_0;
-	uint32_t ppt_offset_bytes; /* soft pptable offset */
-	uint32_t ppt_size_bytes; /* soft pptable size */
-};
-
-struct smc_soft_pptable_entry {
-        uint32_t id;
-        uint32_t ppt_offset_bytes;
-        uint32_t ppt_size_bytes;
-};
-
-/* version_major=2, version_minor=1 */
-struct smc_firmware_header_v2_1 {
-        struct smc_firmware_header_v1_0 v1_0;
-        uint32_t pptable_count;
-        uint32_t pptable_entry_offset;
-};
-
-struct psp_fw_legacy_bin_desc {
-	uint32_t fw_version;
-	uint32_t offset_bytes;
-	uint32_t size_bytes;
-};
-
-/* version_major=1, version_minor=0 */
-struct psp_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	struct psp_fw_legacy_bin_desc sos;
-};
-
-/* version_major=1, version_minor=1 */
-struct psp_firmware_header_v1_1 {
-	struct psp_firmware_header_v1_0 v1_0;
-	struct psp_fw_legacy_bin_desc toc;
-	struct psp_fw_legacy_bin_desc kdb;
-};
-
-/* version_major=1, version_minor=2 */
-struct psp_firmware_header_v1_2 {
-	struct psp_firmware_header_v1_0 v1_0;
-	struct psp_fw_legacy_bin_desc res;
-	struct psp_fw_legacy_bin_desc kdb;
-};
-
-/* version_major=1, version_minor=3 */
-struct psp_firmware_header_v1_3 {
-	struct psp_firmware_header_v1_1 v1_1;
-	struct psp_fw_legacy_bin_desc spl;
-	struct psp_fw_legacy_bin_desc rl;
-	struct psp_fw_legacy_bin_desc sys_drv_aux;
-	struct psp_fw_legacy_bin_desc sos_aux;
-};
-
-struct psp_fw_bin_desc {
-	uint32_t fw_type;
-	uint32_t fw_version;
-	uint32_t offset_bytes;
-	uint32_t size_bytes;
-};
-
-enum psp_fw_type {
-	PSP_FW_TYPE_UNKOWN,
-	PSP_FW_TYPE_PSP_SOS,
-	PSP_FW_TYPE_PSP_SYS_DRV,
-	PSP_FW_TYPE_PSP_KDB,
-	PSP_FW_TYPE_PSP_TOC,
-	PSP_FW_TYPE_PSP_SPL,
-	PSP_FW_TYPE_PSP_RL,
-	PSP_FW_TYPE_PSP_SOC_DRV,
-	PSP_FW_TYPE_PSP_INTF_DRV,
-	PSP_FW_TYPE_PSP_DBG_DRV,
-	PSP_FW_TYPE_PSP_RAS_DRV,
-	PSP_FW_TYPE_PSP_IPKEYMGR_DRV,
-	PSP_FW_TYPE_MAX_INDEX,
-};
-
-/* version_major=2, version_minor=0 */
-struct psp_firmware_header_v2_0 {
-	struct common_firmware_header header;
-	uint32_t psp_fw_bin_count;
-	struct psp_fw_bin_desc psp_fw_bin[1];
-};
-
-/* version_major=2, version_minor=1 */
-struct psp_firmware_header_v2_1 {
-	struct common_firmware_header header;
-	uint32_t psp_fw_bin_count;
-	uint32_t psp_aux_fw_bin_index;
-	struct psp_fw_bin_desc psp_fw_bin[1];
-};
-
-/* version_major=1, version_minor=0 */
-struct ta_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	struct psp_fw_legacy_bin_desc xgmi;
-	struct psp_fw_legacy_bin_desc ras;
-	struct psp_fw_legacy_bin_desc hdcp;
-	struct psp_fw_legacy_bin_desc dtm;
-	struct psp_fw_legacy_bin_desc securedisplay;
-};
-
-enum ta_fw_type {
-	TA_FW_TYPE_UNKOWN,
-	TA_FW_TYPE_PSP_ASD,
-	TA_FW_TYPE_PSP_XGMI,
-	TA_FW_TYPE_PSP_RAS,
-	TA_FW_TYPE_PSP_HDCP,
-	TA_FW_TYPE_PSP_DTM,
-	TA_FW_TYPE_PSP_RAP,
-	TA_FW_TYPE_PSP_SECUREDISPLAY,
-	TA_FW_TYPE_MAX_INDEX,
-};
-
-/* version_major=2, version_minor=0 */
-struct ta_firmware_header_v2_0 {
-	struct common_firmware_header header;
-	uint32_t ta_fw_bin_count;
-	struct psp_fw_bin_desc ta_fw_bin[1];
-};
-
-/* version_major=1, version_minor=0 */
-struct gfx_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t jt_offset; /* jt location */
-	uint32_t jt_size;  /* size of jt */
-};
-
-/* version_major=2, version_minor=0 */
-struct gfx_firmware_header_v2_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t ucode_size_bytes;
-	uint32_t ucode_offset_bytes;
-	uint32_t data_size_bytes;
-	uint32_t data_offset_bytes;
-	uint32_t ucode_start_addr_lo;
-	uint32_t ucode_start_addr_hi;
-};
-
-/* version_major=1, version_minor=0 */
-struct mes_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t mes_ucode_version;
-	uint32_t mes_ucode_size_bytes;
-	uint32_t mes_ucode_offset_bytes;
-	uint32_t mes_ucode_data_version;
-	uint32_t mes_ucode_data_size_bytes;
-	uint32_t mes_ucode_data_offset_bytes;
-	uint32_t mes_uc_start_addr_lo;
-	uint32_t mes_uc_start_addr_hi;
-	uint32_t mes_data_start_addr_lo;
-	uint32_t mes_data_start_addr_hi;
-};
-
-/* version_major=1, version_minor=0 */
-struct rlc_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t save_and_restore_offset;
-	uint32_t clear_state_descriptor_offset;
-	uint32_t avail_scratch_ram_locations;
-	uint32_t master_pkt_description_offset;
-};
-
-/* version_major=2, version_minor=0 */
-struct rlc_firmware_header_v2_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t jt_offset; /* jt location */
-	uint32_t jt_size;  /* size of jt */
-	uint32_t save_and_restore_offset;
-	uint32_t clear_state_descriptor_offset;
-	uint32_t avail_scratch_ram_locations;
-	uint32_t reg_restore_list_size;
-	uint32_t reg_list_format_start;
-	uint32_t reg_list_format_separate_start;
-	uint32_t starting_offsets_start;
-	uint32_t reg_list_format_size_bytes; /* size of reg list format array in bytes */
-	uint32_t reg_list_format_array_offset_bytes; /* payload offset from the start of the header */
-	uint32_t reg_list_size_bytes; /* size of reg list array in bytes */
-	uint32_t reg_list_array_offset_bytes; /* payload offset from the start of the header */
-	uint32_t reg_list_format_separate_size_bytes; /* size of reg list format array in bytes */
-	uint32_t reg_list_format_separate_array_offset_bytes; /* payload offset from the start of the header */
-	uint32_t reg_list_separate_size_bytes; /* size of reg list array in bytes */
-	uint32_t reg_list_separate_array_offset_bytes; /* payload offset from the start of the header */
-};
-
-/* version_major=2, version_minor=1 */
-struct rlc_firmware_header_v2_1 {
-	struct rlc_firmware_header_v2_0 v2_0;
-	uint32_t reg_list_format_direct_reg_list_length; /* length of direct reg list format array */
-	uint32_t save_restore_list_cntl_ucode_ver;
-	uint32_t save_restore_list_cntl_feature_ver;
-	uint32_t save_restore_list_cntl_size_bytes;
-	uint32_t save_restore_list_cntl_offset_bytes;
-	uint32_t save_restore_list_gpm_ucode_ver;
-	uint32_t save_restore_list_gpm_feature_ver;
-	uint32_t save_restore_list_gpm_size_bytes;
-	uint32_t save_restore_list_gpm_offset_bytes;
-	uint32_t save_restore_list_srm_ucode_ver;
-	uint32_t save_restore_list_srm_feature_ver;
-	uint32_t save_restore_list_srm_size_bytes;
-	uint32_t save_restore_list_srm_offset_bytes;
-};
-
-/* version_major=2, version_minor=2 */
-struct rlc_firmware_header_v2_2 {
-	struct rlc_firmware_header_v2_1 v2_1;
-	uint32_t rlc_iram_ucode_size_bytes;
-	uint32_t rlc_iram_ucode_offset_bytes;
-	uint32_t rlc_dram_ucode_size_bytes;
-	uint32_t rlc_dram_ucode_offset_bytes;
-};
-
-/* version_major=2, version_minor=3 */
-struct rlc_firmware_header_v2_3 {
-    struct rlc_firmware_header_v2_2 v2_2;
-    uint32_t rlcp_ucode_version;
-    uint32_t rlcp_ucode_feature_version;
-    uint32_t rlcp_ucode_size_bytes;
-    uint32_t rlcp_ucode_offset_bytes;
-    uint32_t rlcv_ucode_version;
-    uint32_t rlcv_ucode_feature_version;
-    uint32_t rlcv_ucode_size_bytes;
-    uint32_t rlcv_ucode_offset_bytes;
-};
-
-/* version_major=2, version_minor=4 */
-struct rlc_firmware_header_v2_4 {
-    struct rlc_firmware_header_v2_3 v2_3;
-    uint32_t global_tap_delays_ucode_size_bytes;
-    uint32_t global_tap_delays_ucode_offset_bytes;
-    uint32_t se0_tap_delays_ucode_size_bytes;
-    uint32_t se0_tap_delays_ucode_offset_bytes;
-    uint32_t se1_tap_delays_ucode_size_bytes;
-    uint32_t se1_tap_delays_ucode_offset_bytes;
-    uint32_t se2_tap_delays_ucode_size_bytes;
-    uint32_t se2_tap_delays_ucode_offset_bytes;
-    uint32_t se3_tap_delays_ucode_size_bytes;
-    uint32_t se3_tap_delays_ucode_offset_bytes;
-};
-
-/* version_major=1, version_minor=0 */
-struct sdma_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t ucode_change_version;
-	uint32_t jt_offset; /* jt location */
-	uint32_t jt_size; /* size of jt */
-};
-
-/* version_major=1, version_minor=1 */
-struct sdma_firmware_header_v1_1 {
-	struct sdma_firmware_header_v1_0 v1_0;
-	uint32_t digest_size;
-};
-
-/* version_major=2, version_minor=0 */
-struct sdma_firmware_header_v2_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t ctx_ucode_size_bytes; /* context thread ucode size */
-	uint32_t ctx_jt_offset; /* context thread jt location */
-	uint32_t ctx_jt_size; /* context thread size of jt */
-	uint32_t ctl_ucode_offset;
-	uint32_t ctl_ucode_size_bytes; /* control thread ucode size */
-	uint32_t ctl_jt_offset; /* control thread jt location */
-	uint32_t ctl_jt_size; /* control thread size of jt */
-};
-
-/* version_major=1, version_minor=0 */
-struct vpe_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t ucode_feature_version;
-	uint32_t ctx_ucode_size_bytes; /* context thread ucode size */
-	uint32_t ctx_jt_offset; /* context thread jt location */
-	uint32_t ctx_jt_size; /* context thread size of jt */
-	uint32_t ctl_ucode_offset;
-	uint32_t ctl_ucode_size_bytes; /* control thread ucode size */
-	uint32_t ctl_jt_offset; /* control thread jt location */
-	uint32_t ctl_jt_size; /* control thread size of jt */
-};
-
-/* version_major=1, version_minor=0 */
-struct umsch_mm_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t umsch_mm_ucode_version;
-	uint32_t umsch_mm_ucode_size_bytes;
-	uint32_t umsch_mm_ucode_offset_bytes;
-	uint32_t umsch_mm_ucode_data_version;
-	uint32_t umsch_mm_ucode_data_size_bytes;
-	uint32_t umsch_mm_ucode_data_offset_bytes;
-	uint32_t umsch_mm_irq_start_addr_lo;
-	uint32_t umsch_mm_irq_start_addr_hi;
-	uint32_t umsch_mm_uc_start_addr_lo;
-	uint32_t umsch_mm_uc_start_addr_hi;
-	uint32_t umsch_mm_data_start_addr_lo;
-	uint32_t umsch_mm_data_start_addr_hi;
-};
-
-/* version_major=3, version_minor=0 */
-struct sdma_firmware_header_v3_0 {
-    struct common_firmware_header header;
-    uint32_t ucode_feature_version;
-    uint32_t ucode_offset_bytes;
-    uint32_t ucode_size_bytes;
-};
-
-/* gpu info payload */
-struct gpu_info_firmware_v1_0 {
-	uint32_t gc_num_se;
-	uint32_t gc_num_cu_per_sh;
-	uint32_t gc_num_sh_per_se;
-	uint32_t gc_num_rb_per_se;
-	uint32_t gc_num_tccs;
-	uint32_t gc_num_gprs;
-	uint32_t gc_num_max_gs_thds;
-	uint32_t gc_gs_table_depth;
-	uint32_t gc_gsprim_buff_depth;
-	uint32_t gc_parameter_cache_depth;
-	uint32_t gc_double_offchip_lds_buffer;
-	uint32_t gc_wave_size;
-	uint32_t gc_max_waves_per_simd;
-	uint32_t gc_max_scratch_slots_per_cu;
-	uint32_t gc_lds_size;
-};
-
-struct gpu_info_firmware_v1_1 {
-	struct gpu_info_firmware_v1_0 v1_0;
-	uint32_t num_sc_per_sh;
-	uint32_t num_packer_per_sc;
-};
-
-/* gpu info payload
- * version_major=1, version_minor=1 */
-// struct gpu_info_firmware_v1_2 {
-// 	struct gpu_info_firmware_v1_1 v1_1;
-// 	struct gpu_info_soc_bounding_box_v1_0 soc_bounding_box;
-// };
-
-/* version_major=1, version_minor=0 */
-struct gpu_info_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint16_t version_major; /* version */
-	uint16_t version_minor; /* version */
-};
-
-/* version_major=1, version_minor=0 */
-struct dmcu_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t intv_offset_bytes; /* interrupt vectors offset from end of header, in bytes */
-	uint32_t intv_size_bytes;  /* size of interrupt vectors, in bytes */
-};
-
-/* version_major=1, version_minor=0 */
-struct dmcub_firmware_header_v1_0 {
-	struct common_firmware_header header;
-	uint32_t inst_const_bytes; /* size of instruction region, in bytes */
-	uint32_t bss_data_bytes; /* size of bss/data region, in bytes */
-};
-
-/* version_major=1, version_minor=0 */
-struct imu_firmware_header_v1_0 {
-    struct common_firmware_header header;
-    uint32_t imu_iram_ucode_size_bytes;
-    uint32_t imu_iram_ucode_offset_bytes;
-    uint32_t imu_dram_ucode_size_bytes;
-    uint32_t imu_dram_ucode_offset_bytes;
-};
-
-/* header is fixed size */
-union amdgpu_firmware_header {
-	struct common_firmware_header common;
-	struct mc_firmware_header_v1_0 mc;
-	struct smc_firmware_header_v1_0 smc;
-	struct smc_firmware_header_v2_0 smc_v2_0;
-	struct psp_firmware_header_v1_0 psp;
-	struct psp_firmware_header_v1_1 psp_v1_1;
-	struct psp_firmware_header_v1_3 psp_v1_3;
-	struct psp_firmware_header_v2_0 psp_v2_0;
-	struct psp_firmware_header_v2_0 psp_v2_1;
-	struct ta_firmware_header_v1_0 ta;
-	struct ta_firmware_header_v2_0 ta_v2_0;
-	struct gfx_firmware_header_v1_0 gfx;
-	struct gfx_firmware_header_v2_0 gfx_v2_0;
-	struct rlc_firmware_header_v1_0 rlc;
-	struct rlc_firmware_header_v2_0 rlc_v2_0;
-	struct rlc_firmware_header_v2_1 rlc_v2_1;
-	struct rlc_firmware_header_v2_2 rlc_v2_2;
-	struct rlc_firmware_header_v2_3 rlc_v2_3;
-	struct rlc_firmware_header_v2_4 rlc_v2_4;
-	struct sdma_firmware_header_v1_0 sdma;
-	struct sdma_firmware_header_v1_1 sdma_v1_1;
-	struct sdma_firmware_header_v2_0 sdma_v2_0;
-	struct sdma_firmware_header_v3_0 sdma_v3_0;
-	struct gpu_info_firmware_header_v1_0 gpu_info;
-	struct dmcu_firmware_header_v1_0 dmcu;
-	struct dmcub_firmware_header_v1_0 dmcub;
-	struct imu_firmware_header_v1_0 imu;
-	uint8_t raw[0x100];
-};
-
-#define UCODE_MAX_PSP_PACKAGING (((sizeof(union amdgpu_firmware_header) - sizeof(struct common_firmware_header) - 4) / sizeof(struct psp_fw_bin_desc)) * 2)
-
-/*
- * fw loading support
- */
-enum AMDGPU_UCODE_ID {
-	AMDGPU_UCODE_ID_CAP = 0,
-	AMDGPU_UCODE_ID_SDMA0,
-	AMDGPU_UCODE_ID_SDMA1,
-	AMDGPU_UCODE_ID_SDMA2,
-	AMDGPU_UCODE_ID_SDMA3,
-	AMDGPU_UCODE_ID_SDMA4,
-	AMDGPU_UCODE_ID_SDMA5,
-	AMDGPU_UCODE_ID_SDMA6,
-	AMDGPU_UCODE_ID_SDMA7,
-	AMDGPU_UCODE_ID_SDMA_UCODE_TH0,
-	AMDGPU_UCODE_ID_SDMA_UCODE_TH1,
-	AMDGPU_UCODE_ID_SDMA_RS64,
-	AMDGPU_UCODE_ID_CP_CE,
-	AMDGPU_UCODE_ID_CP_PFP,
-	AMDGPU_UCODE_ID_CP_ME,
-	AMDGPU_UCODE_ID_CP_RS64_PFP,
-	AMDGPU_UCODE_ID_CP_RS64_ME,
-	AMDGPU_UCODE_ID_CP_RS64_MEC,
-	AMDGPU_UCODE_ID_CP_RS64_PFP_P0_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_PFP_P1_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_ME_P0_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_ME_P1_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_MEC_P0_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_MEC_P1_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_MEC_P2_STACK,
-	AMDGPU_UCODE_ID_CP_RS64_MEC_P3_STACK,
-	AMDGPU_UCODE_ID_CP_MEC1,
-	AMDGPU_UCODE_ID_CP_MEC1_JT,
-	AMDGPU_UCODE_ID_CP_MEC2,
-	AMDGPU_UCODE_ID_CP_MEC2_JT,
-	AMDGPU_UCODE_ID_CP_MES,
-	AMDGPU_UCODE_ID_CP_MES_DATA,
-	AMDGPU_UCODE_ID_CP_MES1,
-	AMDGPU_UCODE_ID_CP_MES1_DATA,
-	AMDGPU_UCODE_ID_IMU_I,
-	AMDGPU_UCODE_ID_IMU_D,
-	AMDGPU_UCODE_ID_GLOBAL_TAP_DELAYS,
-	AMDGPU_UCODE_ID_SE0_TAP_DELAYS,
-	AMDGPU_UCODE_ID_SE1_TAP_DELAYS,
-	AMDGPU_UCODE_ID_SE2_TAP_DELAYS,
-	AMDGPU_UCODE_ID_SE3_TAP_DELAYS,
-	AMDGPU_UCODE_ID_RLC_RESTORE_LIST_CNTL,
-	AMDGPU_UCODE_ID_RLC_RESTORE_LIST_GPM_MEM,
-	AMDGPU_UCODE_ID_RLC_RESTORE_LIST_SRM_MEM,
-	AMDGPU_UCODE_ID_RLC_IRAM,
-	AMDGPU_UCODE_ID_RLC_DRAM,
-	AMDGPU_UCODE_ID_RLC_P,
-	AMDGPU_UCODE_ID_RLC_V,
-	AMDGPU_UCODE_ID_RLC_G,
-	AMDGPU_UCODE_ID_STORAGE,
-	AMDGPU_UCODE_ID_SMC,
-	AMDGPU_UCODE_ID_PPTABLE,
-	AMDGPU_UCODE_ID_UVD,
-	AMDGPU_UCODE_ID_UVD1,
-	AMDGPU_UCODE_ID_VCE,
-	AMDGPU_UCODE_ID_VCN,
-	AMDGPU_UCODE_ID_VCN1,
-	AMDGPU_UCODE_ID_DMCU_ERAM,
-	AMDGPU_UCODE_ID_DMCU_INTV,
-	AMDGPU_UCODE_ID_VCN0_RAM,
-	AMDGPU_UCODE_ID_VCN1_RAM,
-	AMDGPU_UCODE_ID_DMCUB,
-	AMDGPU_UCODE_ID_VPE_CTX,
-	AMDGPU_UCODE_ID_VPE_CTL,
-	AMDGPU_UCODE_ID_VPE,
-	AMDGPU_UCODE_ID_UMSCH_MM_UCODE,
-	AMDGPU_UCODE_ID_UMSCH_MM_DATA,
-	AMDGPU_UCODE_ID_UMSCH_MM_CMD_BUFFER,
-	AMDGPU_UCODE_ID_P2S_TABLE,
-	AMDGPU_UCODE_ID_JPEG_RAM,
-	AMDGPU_UCODE_ID_ISP,
-	AMDGPU_UCODE_ID_MAXIMUM,
-};
-
-/* engine firmware status */
-enum AMDGPU_UCODE_STATUS {
-	AMDGPU_UCODE_STATUS_INVALID,
-	AMDGPU_UCODE_STATUS_NOT_LOADED,
-	AMDGPU_UCODE_STATUS_LOADED,
-};
-
-enum amdgpu_firmware_load_type {
-	AMDGPU_FW_LOAD_DIRECT = 0,
-	AMDGPU_FW_LOAD_PSP,
-	AMDGPU_FW_LOAD_SMU,
-	AMDGPU_FW_LOAD_RLC_BACKDOOR_AUTO,
-};
-
-/* conform to smu_ucode_xfer_cz.h */
-#define AMDGPU_SDMA0_UCODE_LOADED	0x00000001
-#define AMDGPU_SDMA1_UCODE_LOADED	0x00000002
-#define AMDGPU_CPCE_UCODE_LOADED	0x00000004
-#define AMDGPU_CPPFP_UCODE_LOADED	0x00000008
-#define AMDGPU_CPME_UCODE_LOADED	0x00000010
-#define AMDGPU_CPMEC1_UCODE_LOADED	0x00000020
-#define AMDGPU_CPMEC2_UCODE_LOADED	0x00000040
-#define AMDGPU_CPRLC_UCODE_LOADED	0x00000100
-
-/* amdgpu firmware info */
-struct amdgpu_firmware_info {
-	/* ucode ID */
-	enum AMDGPU_UCODE_ID ucode_id;
-	/* request_firmware */
-	const struct firmware *fw;
-	/* starting mc address */
-	uint64_t mc_addr;
-	/* kernel linear address */
-	void *kaddr;
-	/* ucode_size_bytes */
-	uint32_t ucode_size;
-	/* starting tmr mc address */
-	uint32_t tmr_mc_addr_lo;
-	uint32_t tmr_mc_addr_hi;
-};
-
-// struct amdgpu_firmware {
-// 	struct amdgpu_firmware_info ucode[AMDGPU_UCODE_ID_MAXIMUM];
-// 	enum amdgpu_firmware_load_type load_type;
-// 	struct amdgpu_bo *fw_buf;
-// 	unsigned int fw_size;
-// 	unsigned int max_ucodes;
-// 	/* firmwares are loaded by psp instead of smu from vega10 */
-// 	const struct amdgpu_psp_funcs *funcs;
-// 	struct amdgpu_bo *rbuf;
-// 	struct mutex mutex;
-
-// 	/* gpu info firmware data pointer */
-// 	const struct firmware *gpu_info_fw;
-
-// 	void *fw_buf_ptr;
-// 	uint64_t fw_buf_mc;
-// };
-
-// void amdgpu_ucode_print_mc_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_smc_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_imu_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_gfx_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_rlc_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_sdma_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_psp_hdr(const struct common_firmware_header *hdr);
-// void amdgpu_ucode_print_gpu_info_hdr(const struct common_firmware_header *hdr);
-// int amdgpu_ucode_request(struct amdgpu_device *adev, const struct firmware **fw,
-// 			 const char *fw_name);
-// void amdgpu_ucode_release(const struct firmware **fw);
-// bool amdgpu_ucode_hdr_version(union amdgpu_firmware_header *hdr,
-// 				uint16_t hdr_major, uint16_t hdr_minor);
-
-// int amdgpu_ucode_init_bo(struct amdgpu_device *adev);
-// int amdgpu_ucode_create_bo(struct amdgpu_device *adev);
-// int amdgpu_ucode_sysfs_init(struct amdgpu_device *adev);
-// void amdgpu_ucode_free_bo(struct amdgpu_device *adev);
-// void amdgpu_ucode_sysfs_fini(struct amdgpu_device *adev);
-
-// enum amdgpu_firmware_load_type
-// amdgpu_ucode_get_load_type(struct amdgpu_device *adev, int load_type);
-
-// const char *amdgpu_ucode_name(enum AMDGPU_UCODE_ID ucode_id);
-
-// void amdgpu_ucode_ip_version_decode(struct amdgpu_device *adev, int block_type, char *ucode_prefix, int len);
-
-#endif

extra/amdpci/headers/amdgpu_vm.h

@@ -1,665 +0,0 @@
-/*
- * Copyright 2016 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- * Authors: Christian König
- */
-#ifndef __AMDGPU_VM_H__
-#define __AMDGPU_VM_H__
-
-// #include <linux/idr.h>
-// #include <linux/kfifo.h>
-// #include <linux/rbtree.h>
-// #include <drm/gpu_scheduler.h>
-// #include <drm/drm_file.h>
-// #include <drm/ttm/ttm_bo.h>
-// #include <linux/sched/mm.h>
-
-// #include "amdgpu_sync.h"
-// #include "amdgpu_ring.h"
-// #include "amdgpu_ids.h"
-
-// struct drm_exec;
-
-// struct amdgpu_bo_va;
-// struct amdgpu_job;
-// struct amdgpu_bo_list_entry;
-// struct amdgpu_bo_vm;
-// struct amdgpu_mem_stats;
-
-/*
- * GPUVM handling
- */
-
-/* Maximum number of PTEs the hardware can write with one command */
-#define AMDGPU_VM_MAX_UPDATE_SIZE	0x3FFFF
-
-/* number of entries in page table */
-#define AMDGPU_VM_PTE_COUNT(adev) (1 << (adev)->vm_manager.block_size)
-
-#define AMDGPU_PTE_VALID	(1ULL << 0)
-#define AMDGPU_PTE_SYSTEM	(1ULL << 1)
-#define AMDGPU_PTE_SNOOPED	(1ULL << 2)
-
-/* RV+ */
-#define AMDGPU_PTE_TMZ		(1ULL << 3)
-
-/* VI only */
-#define AMDGPU_PTE_EXECUTABLE	(1ULL << 4)
-
-#define AMDGPU_PTE_READABLE	(1ULL << 5)
-#define AMDGPU_PTE_WRITEABLE	(1ULL << 6)
-
-#define AMDGPU_PTE_FRAG(x)	((x & 0x1fULL) << 7)
-
-/* TILED for VEGA10, reserved for older ASICs  */
-#define AMDGPU_PTE_PRT		(1ULL << 51)
-
-/* PDE is handled as PTE for VEGA10 */
-#define AMDGPU_PDE_PTE		(1ULL << 54)
-
-#define AMDGPU_PTE_LOG          (1ULL << 55)
-
-/* PTE is handled as PDE for VEGA10 (Translate Further) */
-#define AMDGPU_PTE_TF		(1ULL << 56)
-
-/* MALL noalloc for sienna_cichlid, reserved for older ASICs  */
-#define AMDGPU_PTE_NOALLOC	(1ULL << 58)
-
-/* PDE Block Fragment Size for VEGA10 */
-#define AMDGPU_PDE_BFS(a)	((uint64_t)a << 59)
-
-/* Flag combination to set no-retry with TF disabled */
-#define AMDGPU_VM_NORETRY_FLAGS	(AMDGPU_PTE_EXECUTABLE | AMDGPU_PDE_PTE | \
-				AMDGPU_PTE_TF)
-
-/* Flag combination to set no-retry with TF enabled */
-#define AMDGPU_VM_NORETRY_FLAGS_TF (AMDGPU_PTE_VALID | AMDGPU_PTE_SYSTEM | \
-				   AMDGPU_PTE_PRT)
-/* For GFX9 */
-#define AMDGPU_PTE_MTYPE_VG10_SHIFT(mtype)	((uint64_t)(mtype) << 57)
-#define AMDGPU_PTE_MTYPE_VG10_MASK	AMDGPU_PTE_MTYPE_VG10_SHIFT(3ULL)
-#define AMDGPU_PTE_MTYPE_VG10(flags, mtype)			\
-	(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_VG10_MASK)) |	\
-	  AMDGPU_PTE_MTYPE_VG10_SHIFT(mtype))
-
-#define AMDGPU_MTYPE_NC 0
-#define AMDGPU_MTYPE_CC 2
-
-#define AMDGPU_PTE_DEFAULT_ATC  (AMDGPU_PTE_SYSTEM      \
-                                | AMDGPU_PTE_SNOOPED    \
-                                | AMDGPU_PTE_EXECUTABLE \
-                                | AMDGPU_PTE_READABLE   \
-                                | AMDGPU_PTE_WRITEABLE  \
-                                | AMDGPU_PTE_MTYPE_VG10(AMDGPU_MTYPE_CC))
-
-/* gfx10 */
-#define AMDGPU_PTE_MTYPE_NV10_SHIFT(mtype)	((uint64_t)(mtype) << 48)
-#define AMDGPU_PTE_MTYPE_NV10_MASK     AMDGPU_PTE_MTYPE_NV10_SHIFT(7ULL)
-#define AMDGPU_PTE_MTYPE_NV10(flags, mtype)			\
-	(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_NV10_MASK)) |	\
-	  AMDGPU_PTE_MTYPE_NV10_SHIFT(mtype))
-
-/* gfx12 */
-#define AMDGPU_PTE_PRT_GFX12		(1ULL << 56)
-#define AMDGPU_PTE_PRT_FLAG(adev)	\
-	((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PTE_PRT_GFX12 : AMDGPU_PTE_PRT)
-
-#define AMDGPU_PTE_MTYPE_GFX12_SHIFT(mtype)	((uint64_t)(mtype) << 54)
-#define AMDGPU_PTE_MTYPE_GFX12_MASK	AMDGPU_PTE_MTYPE_GFX12_SHIFT(3ULL)
-#define AMDGPU_PTE_MTYPE_GFX12(flags, mtype)				\
-	(((uint64_t)(flags) & (~AMDGPU_PTE_MTYPE_GFX12_MASK)) |	\
-	  AMDGPU_PTE_MTYPE_GFX12_SHIFT(mtype))
-
-#define AMDGPU_PTE_IS_PTE		(1ULL << 63)
-
-/* PDE Block Fragment Size for gfx v12 */
-#define AMDGPU_PDE_BFS_GFX12(a)		((uint64_t)((a) & 0x1fULL) << 58)
-#define AMDGPU_PDE_BFS_FLAG(adev, a)	\
-	((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PDE_BFS_GFX12(a) : AMDGPU_PDE_BFS(a))
-/* PDE is handled as PTE for gfx v12 */
-#define AMDGPU_PDE_PTE_GFX12		(1ULL << 63)
-#define AMDGPU_PDE_PTE_FLAG(adev)	\
-	((amdgpu_ip_version((adev), GC_HWIP, 0) >= IP_VERSION(12, 0, 0)) ? AMDGPU_PDE_PTE_GFX12 : AMDGPU_PDE_PTE)
-
-/* How to program VM fault handling */
-#define AMDGPU_VM_FAULT_STOP_NEVER	0
-#define AMDGPU_VM_FAULT_STOP_FIRST	1
-#define AMDGPU_VM_FAULT_STOP_ALWAYS	2
-
-/* How much VRAM be reserved for page tables */
-#define AMDGPU_VM_RESERVED_VRAM		(8ULL << 20)
-
-/*
- * max number of VMHUB
- * layout: max 8 GFXHUB + 4 MMHUB0 + 1 MMHUB1
- */
-#define AMDGPU_MAX_VMHUBS			13
-#define AMDGPU_GFXHUB_START			0
-#define AMDGPU_MMHUB0_START			8
-#define AMDGPU_MMHUB1_START			12
-#define AMDGPU_GFXHUB(x)			(AMDGPU_GFXHUB_START + (x))
-#define AMDGPU_MMHUB0(x)			(AMDGPU_MMHUB0_START + (x))
-#define AMDGPU_MMHUB1(x)			(AMDGPU_MMHUB1_START + (x))
-
-#define AMDGPU_IS_GFXHUB(x) ((x) >= AMDGPU_GFXHUB_START && (x) < AMDGPU_MMHUB0_START)
-#define AMDGPU_IS_MMHUB0(x) ((x) >= AMDGPU_MMHUB0_START && (x) < AMDGPU_MMHUB1_START)
-#define AMDGPU_IS_MMHUB1(x) ((x) >= AMDGPU_MMHUB1_START && (x) < AMDGPU_MAX_VMHUBS)
-
-/* Reserve space at top/bottom of address space for kernel use */
-#define AMDGPU_VA_RESERVED_CSA_SIZE		(2ULL << 20)
-#define AMDGPU_VA_RESERVED_CSA_START(adev)	(((adev)->vm_manager.max_pfn \
-						  << AMDGPU_GPU_PAGE_SHIFT)  \
-						 - AMDGPU_VA_RESERVED_CSA_SIZE)
-#define AMDGPU_VA_RESERVED_SEQ64_SIZE		(2ULL << 20)
-#define AMDGPU_VA_RESERVED_SEQ64_START(adev)	(AMDGPU_VA_RESERVED_CSA_START(adev) \
-						 - AMDGPU_VA_RESERVED_SEQ64_SIZE)
-#define AMDGPU_VA_RESERVED_TRAP_SIZE		(2ULL << 12)
-#define AMDGPU_VA_RESERVED_TRAP_START(adev)	(AMDGPU_VA_RESERVED_SEQ64_START(adev) \
-						 - AMDGPU_VA_RESERVED_TRAP_SIZE)
-#define AMDGPU_VA_RESERVED_BOTTOM		(1ULL << 16)
-#define AMDGPU_VA_RESERVED_TOP			(AMDGPU_VA_RESERVED_TRAP_SIZE + \
-						 AMDGPU_VA_RESERVED_SEQ64_SIZE + \
-						 AMDGPU_VA_RESERVED_CSA_SIZE)
-
-/* See vm_update_mode */
-#define AMDGPU_VM_USE_CPU_FOR_GFX (1 << 0)
-#define AMDGPU_VM_USE_CPU_FOR_COMPUTE (1 << 1)
-
-/* VMPT level enumerate, and the hiberachy is:
- * PDB2->PDB1->PDB0->PTB
- */
-enum amdgpu_vm_level {
-	AMDGPU_VM_PDB2,
-	AMDGPU_VM_PDB1,
-	AMDGPU_VM_PDB0,
-	AMDGPU_VM_PTB
-};
-
-// /* base structure for tracking BO usage in a VM */
-// struct amdgpu_vm_bo_base {
-// 	/* constant after initialization */
-// 	struct amdgpu_vm		*vm;
-// 	struct amdgpu_bo		*bo;
-
-// 	/* protected by bo being reserved */
-// 	struct amdgpu_vm_bo_base	*next;
-
-// 	/* protected by spinlock */
-// 	struct list_head		vm_status;
-
-// 	/* protected by the BO being reserved */
-// 	bool				moved;
-// };
-
-// /* provided by hw blocks that can write ptes, e.g., sdma */
-// struct amdgpu_vm_pte_funcs {
-// 	/* number of dw to reserve per operation */
-// 	unsigned	copy_pte_num_dw;
-
-// 	/* copy pte entries from GART */
-// 	void (*copy_pte)(struct amdgpu_ib *ib,
-// 			 uint64_t pe, uint64_t src,
-// 			 unsigned count);
-
-// 	/* write pte one entry at a time with addr mapping */
-// 	void (*write_pte)(struct amdgpu_ib *ib, uint64_t pe,
-// 			  uint64_t value, unsigned count,
-// 			  uint32_t incr);
-// 	/* for linear pte/pde updates without addr mapping */
-// 	void (*set_pte_pde)(struct amdgpu_ib *ib,
-// 			    uint64_t pe,
-// 			    uint64_t addr, unsigned count,
-// 			    uint32_t incr, uint64_t flags);
-// };
-
-// struct amdgpu_task_info {
-// 	char		process_name[TASK_COMM_LEN];
-// 	char		task_name[TASK_COMM_LEN];
-// 	pid_t		pid;
-// 	pid_t		tgid;
-// 	struct kref	refcount;
-// };
-
-// /**
-//  * struct amdgpu_vm_update_params
-//  *
-//  * Encapsulate some VM table update parameters to reduce
-//  * the number of function parameters
-//  *
-//  */
-// struct amdgpu_vm_update_params {
-
-// 	/**
-// 	 * @adev: amdgpu device we do this update for
-// 	 */
-// 	struct amdgpu_device *adev;
-
-// 	/**
-// 	 * @vm: optional amdgpu_vm we do this update for
-// 	 */
-// 	struct amdgpu_vm *vm;
-
-// 	/**
-// 	 * @immediate: if changes should be made immediately
-// 	 */
-// 	bool immediate;
-
-// 	/**
-// 	 * @unlocked: true if the root BO is not locked
-// 	 */
-// 	bool unlocked;
-
-// 	/**
-// 	 * @pages_addr:
-// 	 *
-// 	 * DMA addresses to use for mapping
-// 	 */
-// 	dma_addr_t *pages_addr;
-
-// 	/**
-// 	 * @job: job to used for hw submission
-// 	 */
-// 	struct amdgpu_job *job;
-
-// 	/**
-// 	 * @num_dw_left: number of dw left for the IB
-// 	 */
-// 	unsigned int num_dw_left;
-
-// 	/**
-// 	 * @needs_flush: true whenever we need to invalidate the TLB
-// 	 */
-// 	bool needs_flush;
-
-// 	/**
-// 	 * @allow_override: true for memory that is not uncached: allows MTYPE
-// 	 * to be overridden for NUMA local memory.
-// 	 */
-// 	bool allow_override;
-
-// 	/**
-// 	 * @tlb_flush_waitlist: temporary storage for BOs until tlb_flush
-// 	 */
-// 	struct list_head tlb_flush_waitlist;
-// };
-
-// struct amdgpu_vm_update_funcs {
-// 	int (*map_table)(struct amdgpu_bo_vm *bo);
-// 	int (*prepare)(struct amdgpu_vm_update_params *p, struct dma_resv *resv,
-// 		       enum amdgpu_sync_mode sync_mode);
-// 	int (*update)(struct amdgpu_vm_update_params *p,
-// 		      struct amdgpu_bo_vm *bo, uint64_t pe, uint64_t addr,
-// 		      unsigned count, uint32_t incr, uint64_t flags);
-// 	int (*commit)(struct amdgpu_vm_update_params *p,
-// 		      struct dma_fence **fence);
-// };
-
-// struct amdgpu_vm_fault_info {
-// 	/* fault address */
-// 	uint64_t	addr;
-// 	/* fault status register */
-// 	uint32_t	status;
-// 	/* which vmhub? gfxhub, mmhub, etc. */
-// 	unsigned int	vmhub;
-// };
-
-// struct amdgpu_vm {
-// 	/* tree of virtual addresses mapped */
-// #ifndef HAVE_TREE_INSERT_HAVE_RB_ROOT_CACHED
-// 	struct rb_root	va;
-// #else
-// 	struct rb_root_cached	va;
-// #endif
-
-// 	/* Lock to prevent eviction while we are updating page tables
-// 	 * use vm_eviction_lock/unlock(vm)
-// 	 */
-// 	struct mutex		eviction_lock;
-// 	bool			evicting;
-// 	unsigned int		saved_flags;
-
-// 	/* Lock to protect vm_bo add/del/move on all lists of vm */
-// 	spinlock_t		status_lock;
-
-// 	/* Per-VM and PT BOs who needs a validation */
-// 	struct list_head	evicted;
-
-// 	/* BOs for user mode queues that need a validation */
-// 	struct list_head	evicted_user;
-
-// 	/* PT BOs which relocated and their parent need an update */
-// 	struct list_head	relocated;
-
-// 	/* per VM BOs moved, but not yet updated in the PT */
-// 	struct list_head	moved;
-
-// 	/* All BOs of this VM not currently in the state machine */
-// 	struct list_head	idle;
-
-// 	/* regular invalidated BOs, but not yet updated in the PT */
-// 	struct list_head	invalidated;
-
-// 	/* BO mappings freed, but not yet updated in the PT */
-// 	struct list_head	freed;
-
-// 	/* BOs which are invalidated, has been updated in the PTs */
-// 	struct list_head        done;
-
-// 	/* PT BOs scheduled to free and fill with zero if vm_resv is not hold */
-// 	struct list_head	pt_freed;
-// 	struct work_struct	pt_free_work;
-
-// 	/* contains the page directory */
-// 	struct amdgpu_vm_bo_base     root;
-// 	struct dma_fence	*last_update;
-
-// 	/* Scheduler entities for page table updates */
-// 	struct drm_sched_entity	immediate;
-// 	struct drm_sched_entity	delayed;
-
-// 	/* Last finished delayed update */
-// 	atomic64_t		tlb_seq;
-// 	struct dma_fence	*last_tlb_flush;
-// 	atomic64_t		kfd_last_flushed_seq;
-// 	uint64_t		tlb_fence_context;
-
-// 	/* How many times we had to re-generate the page tables */
-// 	uint64_t		generation;
-
-// 	/* Last unlocked submission to the scheduler entities */
-// 	struct dma_fence	*last_unlocked;
-
-// 	unsigned int		pasid;
-// 	bool			reserved_vmid[AMDGPU_MAX_VMHUBS];
-
-// 	/* Flag to indicate if VM tables are updated by CPU or GPU (SDMA) */
-// 	bool					use_cpu_for_update;
-
-// 	/* Functions to use for VM table updates */
-// 	const struct amdgpu_vm_update_funcs	*update_funcs;
-
-// 	/* Up to 128 pending retry page faults */
-// 	DECLARE_KFIFO(faults, u64, 128);
-
-// 	/* Points to the KFD process VM info */
-// 	struct amdkfd_process_info *process_info;
-
-// 	/* List node in amdkfd_process_info.vm_list_head */
-// 	struct list_head	vm_list_node;
-
-// 	/* Valid while the PD is reserved or fenced */
-// 	uint64_t		pd_phys_addr;
-
-// 	/* Some basic info about the task */
-// 	struct amdgpu_task_info *task_info;
-
-// 	/* Store positions of group of BOs */
-// 	struct ttm_lru_bulk_move lru_bulk_move;
-// 	/* Flag to indicate if VM is used for compute */
-// 	bool			is_compute_context;
-
-// 	/* Memory partition number, -1 means any partition */
-// 	int8_t			mem_id;
-
-// 	/* cached fault info */
-// 	struct amdgpu_vm_fault_info fault_info;
-// };
-
-// struct amdgpu_vm_manager {
-// 	/* Handling of VMIDs */
-// 	struct amdgpu_vmid_mgr			id_mgr[AMDGPU_MAX_VMHUBS];
-// 	unsigned int				first_kfd_vmid;
-// 	bool					concurrent_flush;
-
-// 	/* Handling of VM fences */
-// 	u64					fence_context;
-// 	unsigned				seqno[AMDGPU_MAX_RINGS];
-
-// 	uint64_t				max_pfn;
-// 	uint32_t				num_level;
-// 	uint32_t				block_size;
-// 	uint32_t				fragment_size;
-// 	enum amdgpu_vm_level			root_level;
-// 	/* vram base address for page table entry  */
-// 	u64					vram_base_offset;
-// 	/* vm pte handling */
-// 	const struct amdgpu_vm_pte_funcs	*vm_pte_funcs;
-// 	struct drm_gpu_scheduler		*vm_pte_scheds[AMDGPU_MAX_RINGS];
-// 	unsigned				vm_pte_num_scheds;
-// 	struct amdgpu_ring			*page_fault;
-
-// 	/* partial resident texture handling */
-// 	spinlock_t				prt_lock;
-// 	atomic_t				num_prt_users;
-
-// 	/* controls how VM page tables are updated for Graphics and Compute.
-// 	 * BIT0[= 0] Graphics updated by SDMA [= 1] by CPU
-// 	 * BIT1[= 0] Compute updated by SDMA [= 1] by CPU
-// 	 */
-// 	int					vm_update_mode;
-
-// 	/* PASID to VM mapping, will be used in interrupt context to
-// 	 * look up VM of a page fault
-// 	 */
-// #ifdef HAVE_STRUCT_XARRAY
-// 	struct xarray				pasids;
-// #else
-// 	struct idr				pasid_idr;
-// 	spinlock_t				pasid_lock;
-// #endif
-// 	/* Global registration of recent page fault information */
-// 	struct amdgpu_vm_fault_info	fault_info;
-// };
-
-// struct amdgpu_bo_va_mapping;
-
-// #define amdgpu_vm_copy_pte(adev, ib, pe, src, count) ((adev)->vm_manager.vm_pte_funcs->copy_pte((ib), (pe), (src), (count)))
-// #define amdgpu_vm_write_pte(adev, ib, pe, value, count, incr) ((adev)->vm_manager.vm_pte_funcs->write_pte((ib), (pe), (value), (count), (incr)))
-// #define amdgpu_vm_set_pte_pde(adev, ib, pe, addr, count, incr, flags) ((adev)->vm_manager.vm_pte_funcs->set_pte_pde((ib), (pe), (addr), (count), (incr), (flags)))
-
-// extern const struct amdgpu_vm_update_funcs amdgpu_vm_cpu_funcs;
-// extern const struct amdgpu_vm_update_funcs amdgpu_vm_sdma_funcs;
-
-// void amdgpu_vm_manager_init(struct amdgpu_device *adev);
-// void amdgpu_vm_manager_fini(struct amdgpu_device *adev);
-
-// int amdgpu_vm_set_pasid(struct amdgpu_device *adev, struct amdgpu_vm *vm,
-// 			u32 pasid);
-
-// long amdgpu_vm_wait_idle(struct amdgpu_vm *vm, long timeout);
-// int amdgpu_vm_init(struct amdgpu_device *adev, struct amdgpu_vm *vm, int32_t xcp_id);
-// int amdgpu_vm_make_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-// void amdgpu_vm_release_compute(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-// void amdgpu_vm_fini(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-// int amdgpu_vm_lock_pd(struct amdgpu_vm *vm, struct drm_exec *exec,
-// 		      unsigned int num_fences);
-// bool amdgpu_vm_ready(struct amdgpu_vm *vm);
-// uint64_t amdgpu_vm_generation(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-// int amdgpu_vm_validate(struct amdgpu_device *adev, struct amdgpu_vm *vm,
-// 		       struct ww_acquire_ctx *ticket,
-// 		       int (*callback)(void *p, struct amdgpu_bo *bo),
-// 		       void *param);
-// int amdgpu_vm_flush(struct amdgpu_ring *ring, struct amdgpu_job *job, bool need_pipe_sync);
-// int amdgpu_vm_update_pdes(struct amdgpu_device *adev,
-// 			  struct amdgpu_vm *vm, bool immediate);
-// int amdgpu_vm_clear_freed(struct amdgpu_device *adev,
-// 			  struct amdgpu_vm *vm,
-// 			  struct dma_fence **fence);
-// int amdgpu_vm_handle_moved(struct amdgpu_device *adev,
-// 			   struct amdgpu_vm *vm,
-// 			   struct ww_acquire_ctx *ticket);
-// int amdgpu_vm_flush_compute_tlb(struct amdgpu_device *adev,
-// 				struct amdgpu_vm *vm,
-// 				uint32_t flush_type,
-// 				uint32_t xcc_mask);
-// void amdgpu_vm_bo_base_init(struct amdgpu_vm_bo_base *base,
-// 			    struct amdgpu_vm *vm, struct amdgpu_bo *bo);
-// int amdgpu_vm_update_range(struct amdgpu_device *adev, struct amdgpu_vm *vm,
-// 			   bool immediate, bool unlocked, bool flush_tlb, bool allow_override,
-// 			   struct dma_resv *resv, uint64_t start, uint64_t last,
-// 			   uint64_t flags, uint64_t offset, uint64_t vram_base,
-// 			   struct ttm_resource *res, dma_addr_t *pages_addr,
-// 			   struct dma_fence **fence);
-// int amdgpu_vm_bo_update(struct amdgpu_device *adev,
-// 			struct amdgpu_bo_va *bo_va,
-// 			bool clear);
-// bool amdgpu_vm_evictable(struct amdgpu_bo *bo);
-// void amdgpu_vm_bo_invalidate(struct amdgpu_device *adev,
-// 			     struct amdgpu_bo *bo, bool evicted);
-// uint64_t amdgpu_vm_map_gart(const dma_addr_t *pages_addr, uint64_t addr);
-// struct amdgpu_bo_va *amdgpu_vm_bo_find(struct amdgpu_vm *vm,
-// 				       struct amdgpu_bo *bo);
-// struct amdgpu_bo_va *amdgpu_vm_bo_add(struct amdgpu_device *adev,
-// 				      struct amdgpu_vm *vm,
-// 				      struct amdgpu_bo *bo);
-// int amdgpu_vm_bo_map(struct amdgpu_device *adev,
-// 		     struct amdgpu_bo_va *bo_va,
-// 		     uint64_t addr, uint64_t offset,
-// 		     uint64_t size, uint64_t flags);
-// int amdgpu_vm_bo_replace_map(struct amdgpu_device *adev,
-// 			     struct amdgpu_bo_va *bo_va,
-// 			     uint64_t addr, uint64_t offset,
-// 			     uint64_t size, uint64_t flags);
-// int amdgpu_vm_bo_unmap(struct amdgpu_device *adev,
-// 		       struct amdgpu_bo_va *bo_va,
-// 		       uint64_t addr);
-// int amdgpu_vm_bo_clear_mappings(struct amdgpu_device *adev,
-// 				struct amdgpu_vm *vm,
-// 				uint64_t saddr, uint64_t size);
-// struct amdgpu_bo_va_mapping *amdgpu_vm_bo_lookup_mapping(struct amdgpu_vm *vm,
-// 							 uint64_t addr);
-// void amdgpu_vm_bo_trace_cs(struct amdgpu_vm *vm, struct ww_acquire_ctx *ticket);
-// void amdgpu_vm_bo_del(struct amdgpu_device *adev,
-// 		      struct amdgpu_bo_va *bo_va);
-// void amdgpu_vm_adjust_size(struct amdgpu_device *adev, uint32_t min_vm_size,
-// 			   uint32_t fragment_size_default, unsigned max_level,
-// 			   unsigned max_bits);
-// int amdgpu_vm_ioctl(struct drm_device *dev, void *data, struct drm_file *filp);
-// bool amdgpu_vm_need_pipeline_sync(struct amdgpu_ring *ring,
-// 				  struct amdgpu_job *job);
-// void amdgpu_vm_check_compute_bug(struct amdgpu_device *adev);
-
-// struct amdgpu_task_info *
-// amdgpu_vm_get_task_info_pasid(struct amdgpu_device *adev, u32 pasid);
-
-// struct amdgpu_task_info *
-// amdgpu_vm_get_task_info_vm(struct amdgpu_vm *vm);
-
-// void amdgpu_vm_put_task_info(struct amdgpu_task_info *task_info);
-
-// bool amdgpu_vm_handle_fault(struct amdgpu_device *adev, u32 pasid,
-// 			    u32 vmid, u32 node_id, uint64_t addr,
-// 			    bool write_fault);
-
-// void amdgpu_vm_set_task_info(struct amdgpu_vm *vm);
-
-// void amdgpu_vm_move_to_lru_tail(struct amdgpu_device *adev,
-// 				struct amdgpu_vm *vm);
-// void amdgpu_vm_get_memory(struct amdgpu_vm *vm,
-// 			  struct amdgpu_mem_stats *stats);
-
-// int amdgpu_vm_pt_clear(struct amdgpu_device *adev, struct amdgpu_vm *vm,
-// 		       struct amdgpu_bo_vm *vmbo, bool immediate);
-// int amdgpu_vm_pt_create(struct amdgpu_device *adev, struct amdgpu_vm *vm,
-// 			int level, bool immediate, struct amdgpu_bo_vm **vmbo,
-// 			int32_t xcp_id);
-// void amdgpu_vm_pt_free_root(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-
-// int amdgpu_vm_pde_update(struct amdgpu_vm_update_params *params,
-// 			 struct amdgpu_vm_bo_base *entry);
-// int amdgpu_vm_ptes_update(struct amdgpu_vm_update_params *params,
-// 			  uint64_t start, uint64_t end,
-// 			  uint64_t dst, uint64_t flags);
-// void amdgpu_vm_pt_free_work(struct work_struct *work);
-// void amdgpu_vm_pt_free_list(struct amdgpu_device *adev,
-// 			    struct amdgpu_vm_update_params *params);
-
-// #if defined(CONFIG_DEBUG_FS)
-// void amdgpu_debugfs_vm_bo_info(struct amdgpu_vm *vm, struct seq_file *m);
-// #endif
-
-// int amdgpu_vm_pt_map_tables(struct amdgpu_device *adev, struct amdgpu_vm *vm);
-
-// bool amdgpu_vm_is_bo_always_valid(struct amdgpu_vm *vm, struct amdgpu_bo *bo);
-
-// /**
-//  * amdgpu_vm_tlb_seq - return tlb flush sequence number
-//  * @vm: the amdgpu_vm structure to query
-//  *
-//  * Returns the tlb flush sequence number which indicates that the VM TLBs needs
-//  * to be invalidated whenever the sequence number change.
-//  */
-// static inline uint64_t amdgpu_vm_tlb_seq(struct amdgpu_vm *vm)
-// {
-// 	unsigned long flags;
-// 	spinlock_t *lock;
-
-// 	/*
-// 	 * Workaround to stop racing between the fence signaling and handling
-// 	 * the cb. The lock is static after initially setting it up, just make
-// 	 * sure that the dma_fence structure isn't freed up.
-// 	 */
-// 	rcu_read_lock();
-// 	lock = vm->last_tlb_flush->lock;
-// 	rcu_read_unlock();
-
-// 	spin_lock_irqsave(lock, flags);
-// 	spin_unlock_irqrestore(lock, flags);
-
-// 	return atomic64_read(&vm->tlb_seq);
-// }
-
-// /*
-//  * vm eviction_lock can be taken in MMU notifiers. Make sure no reclaim-FS
-//  * happens while holding this lock anywhere to prevent deadlocks when
-//  * an MMU notifier runs in reclaim-FS context.
-//  */
-// static inline void amdgpu_vm_eviction_lock(struct amdgpu_vm *vm)
-// {
-// 	mutex_lock(&vm->eviction_lock);
-// 	vm->saved_flags = memalloc_noreclaim_save();
-// }
-
-// static inline bool amdgpu_vm_eviction_trylock(struct amdgpu_vm *vm)
-// {
-// 	if (mutex_trylock(&vm->eviction_lock)) {
-// 		vm->saved_flags = memalloc_noreclaim_save();
-// 		return true;
-// 	}
-// 	return false;
-// }
-
-// static inline void amdgpu_vm_eviction_unlock(struct amdgpu_vm *vm)
-// {
-// 	memalloc_noreclaim_restore(vm->saved_flags);
-// 	mutex_unlock(&vm->eviction_lock);
-// }
-
-// void amdgpu_vm_update_fault_cache(struct amdgpu_device *adev,
-// 				  unsigned int pasid,
-// 				  uint64_t addr,
-// 				  uint32_t status,
-// 				  unsigned int vmhub);
-// void amdgpu_vm_tlb_fence_create(struct amdgpu_device *adev,
-// 				 struct amdgpu_vm *vm,
-// 				 struct dma_fence **fence);
-
-#endif

extra/amdpci/headers/discovery.h

@@ -1,600 +0,0 @@
-/*
- * Copyright 2018 Advanced Micro Devices, Inc.
- *
- * Permission is hereby granted, free of charge, to any person obtaining a
- * copy of this software and associated documentation files (the "Software"),
- * to deal in the Software without restriction, including without limitation
- * the rights to use, copy, modify, merge, publish, distribute, sublicense,
- * and/or sell copies of the Software, and to permit persons to whom the
- * Software is furnished to do so, subject to the following conditions:
- *
- * The above copyright notice and this permission notice shall be included in
- * all copies or substantial portions of the Software.
- *
- * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
- * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
- * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
- * THE COPYRIGHT HOLDER(S) OR AUTHOR(S) BE LIABLE FOR ANY CLAIM, DAMAGES OR
- * OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
- * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
- * OTHER DEALINGS IN THE SOFTWARE.
- *
- */
-
- #ifndef _DISCOVERY_H_
- #define _DISCOVERY_H_
- 
- #define PSP_HEADER_SIZE                 256
- #define BINARY_SIGNATURE                0x28211407
- #define DISCOVERY_TABLE_SIGNATURE       0x53445049
- #define GC_TABLE_ID                     0x4347
- #define HARVEST_TABLE_SIGNATURE         0x56524148
- #define VCN_INFO_TABLE_ID               0x004E4356
- #define MALL_INFO_TABLE_ID              0x4C4C414D
- #define NPS_INFO_TABLE_ID 0x0053504E
- 
- typedef enum {
-	 IP_DISCOVERY = 0,
-	 GC,
-	 HARVEST_INFO,
-	 VCN_INFO,
-	 MALL_INFO,
-	 NPS_INFO,
-	 TOTAL_TABLES = 6
- } table;
- 
- #pragma pack(1)
- 
- typedef struct table_info
- {
-	 uint16_t offset;   /* Byte offset */
-	 uint16_t checksum; /* Byte sum of the table */
-	 uint16_t size;     /* Table size */
-	 uint16_t padding;
- } table_info;
- 
- typedef struct binary_header
- {
-	 /* psp structure should go at the top of this structure */
-	 uint32_t binary_signature; /* 0x7, 0x14, 0x21, 0x28 */
-	 uint16_t version_major;
-	 uint16_t version_minor;
-	 uint16_t binary_checksum;  /* Byte sum of the binary after this field */
-	 uint16_t binary_size;      /* Binary Size*/
-	 table_info table_list[TOTAL_TABLES];
- } binary_header;
- 
- typedef struct die_info
- {
-	 uint16_t die_id;
-	 uint16_t die_offset; /* Points to the corresponding die_header structure */
- } die_info;
- 
- 
- typedef struct ip_discovery_header
- {
-	 uint32_t signature;    /* Table Signature */
-	 uint16_t version;      /* Table Version */
-	 uint16_t size;         /* Table Size */
-	 uint32_t id;           /* Table ID */
-	 uint16_t num_dies;     /* Number of Dies */
-	 die_info die_info[16]; /* list die information for up to 16 dies */
-	 union {
-		 uint16_t padding[1];	/* version <= 3 */
-		 struct {		/* version == 4 */
-			 uint8_t base_addr_64_bit : 1; /* ip structures are using 64 bit base address */
-			 uint8_t reserved : 7;
-			 uint8_t reserved2;
-		 };
-	 };
- } ip_discovery_header;
- 
- typedef struct ip
- {
-	 uint16_t hw_id;           /* Hardware ID */
-	 uint8_t number_instance;  /* instance of the IP */
-	 uint8_t num_base_address; /* Number of Base Addresses */
-	 uint8_t major;            /* HCID Major */
-	 uint8_t minor;            /* HCID Minor */
-	 uint8_t revision;         /* HCID Revision */
- #if defined(__BIG_ENDIAN)
-	 uint8_t reserved : 4;     /* Placeholder field */
-	 uint8_t harvest : 4;      /* Harvest */
- #else
-	 uint8_t harvest : 4;      /* Harvest */
-	 uint8_t reserved : 4;     /* Placeholder field */
- #endif
-	 uint32_t base_address[]; /* variable number of Addresses */
- } ip;
- 
- typedef struct ip_v3
- {
-	 uint16_t hw_id;                         /* Hardware ID */
-	 uint8_t instance_number;                /* Instance number for the IP */
-	 uint8_t num_base_address;               /* Number of base addresses*/
-	 uint8_t major;                          /* Hardware ID.major version */
-	 uint8_t minor;                          /* Hardware ID.minor version */
-	 uint8_t revision;                       /* Hardware ID.revision version */
- #if defined(__BIG_ENDIAN)
-	 uint8_t variant : 4;                    /* HW variant */
-	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
- #else
-	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-	 uint8_t variant : 4;                    /* HW variant */
- #endif
-	 uint32_t base_address[];		/* Base Address list. Corresponds to the num_base_address field*/
- } ip_v3;
- 
- typedef struct ip_v4 {
-	 uint16_t hw_id;                         /* Hardware ID */
-	 uint8_t instance_number;                /* Instance number for the IP */
-	 uint8_t num_base_address;               /* Number of base addresses*/
-	 uint8_t major;                          /* Hardware ID.major version */
-	 uint8_t minor;                          /* Hardware ID.minor version */
-	 uint8_t revision;                       /* Hardware ID.revision version */
- #if defined(LITTLEENDIAN_CPU)
-	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
-	 uint8_t variant : 4;                    /* HW variant */
- #elif defined(BIGENDIAN_CPU)
-	 uint8_t variant : 4;                    /* HW variant */
-	 uint8_t sub_revision : 4;               /* HCID Sub-Revision */
- #endif
- } ip_v4;
- 
- typedef struct die_header
- {
-	 uint16_t die_id;
-	 uint16_t num_ips;
- } die_header;
- 
- typedef struct ip_structure
- {
-	 ip_discovery_header* header;
-	 struct die
-	 {
-		 die_header *die_header;
-		 union
-		 {
-			 ip *ip_list;
-			 ip_v3 *ip_v3_list;
-			 ip_v4 *ip_v4_list;
-		 };                                  /* IP list. Variable size*/
-	 } die;
- } ip_structure;
- 
- struct gpu_info_header {
-	 uint32_t table_id;      /* table ID */
-	 uint16_t version_major; /* table version */
-	 uint16_t version_minor; /* table version */
-	 uint32_t size;          /* size of the entire header+data in bytes */
- };
- 
- struct gc_info_v1_0 {
-	 struct gpu_info_header header;
- 
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_wgp0_per_sa;
-	 uint32_t gc_num_wgp1_per_sa;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_gl2c;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_sa_per_se;
-	 uint32_t gc_num_packer_per_sc;
-	 uint32_t gc_num_gl2a;
- };
- 
- struct gc_info_v1_1 {
-	 struct gpu_info_header header;
- 
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_wgp0_per_sa;
-	 uint32_t gc_num_wgp1_per_sa;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_gl2c;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_sa_per_se;
-	 uint32_t gc_num_packer_per_sc;
-	 uint32_t gc_num_gl2a;
-	 uint32_t gc_num_tcp_per_sa;
-	 uint32_t gc_num_sdp_interface;
-	 uint32_t gc_num_tcps;
- };
- 
- struct gc_info_v1_2 {
-	 struct gpu_info_header header;
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_wgp0_per_sa;
-	 uint32_t gc_num_wgp1_per_sa;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_gl2c;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_sa_per_se;
-	 uint32_t gc_num_packer_per_sc;
-	 uint32_t gc_num_gl2a;
-	 uint32_t gc_num_tcp_per_sa;
-	 uint32_t gc_num_sdp_interface;
-	 uint32_t gc_num_tcps;
-	 uint32_t gc_num_tcp_per_wpg;
-	 uint32_t gc_tcp_l1_size;
-	 uint32_t gc_num_sqc_per_wgp;
-	 uint32_t gc_l1_instruction_cache_size_per_sqc;
-	 uint32_t gc_l1_data_cache_size_per_sqc;
-	 uint32_t gc_gl1c_per_sa;
-	 uint32_t gc_gl1c_size_per_instance;
-	 uint32_t gc_gl2c_per_gpu;
- };
- 
- struct gc_info_v1_3 {
-	 struct gpu_info_header header;
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_wgp0_per_sa;
-	 uint32_t gc_num_wgp1_per_sa;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_gl2c;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_sa_per_se;
-	 uint32_t gc_num_packer_per_sc;
-	 uint32_t gc_num_gl2a;
-	 uint32_t gc_num_tcp_per_sa;
-	 uint32_t gc_num_sdp_interface;
-	 uint32_t gc_num_tcps;
-	 uint32_t gc_num_tcp_per_wpg;
-	 uint32_t gc_tcp_l1_size;
-	 uint32_t gc_num_sqc_per_wgp;
-	 uint32_t gc_l1_instruction_cache_size_per_sqc;
-	 uint32_t gc_l1_data_cache_size_per_sqc;
-	 uint32_t gc_gl1c_per_sa;
-	 uint32_t gc_gl1c_size_per_instance;
-	 uint32_t gc_gl2c_per_gpu;
-	 uint32_t gc_tcp_size_per_cu;
-	 uint32_t gc_tcp_cache_line_size;
-	 uint32_t gc_instruction_cache_size_per_sqc;
-	 uint32_t gc_instruction_cache_line_size;
-	 uint32_t gc_scalar_data_cache_size_per_sqc;
-	 uint32_t gc_scalar_data_cache_line_size;
-	 uint32_t gc_tcc_size;
-	 uint32_t gc_tcc_cache_line_size;
- };
- 
- struct gc_info_v2_0 {
-	 struct gpu_info_header header;
- 
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_cu_per_sh;
-	 uint32_t gc_num_sh_per_se;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_tccs;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_packer_per_sc;
- };
- 
- struct gc_info_v2_1 {
-	 struct gpu_info_header header;
- 
-	 uint32_t gc_num_se;
-	 uint32_t gc_num_cu_per_sh;
-	 uint32_t gc_num_sh_per_se;
-	 uint32_t gc_num_rb_per_se;
-	 uint32_t gc_num_tccs;
-	 uint32_t gc_num_gprs;
-	 uint32_t gc_num_max_gs_thds;
-	 uint32_t gc_gs_table_depth;
-	 uint32_t gc_gsprim_buff_depth;
-	 uint32_t gc_parameter_cache_depth;
-	 uint32_t gc_double_offchip_lds_buffer;
-	 uint32_t gc_wave_size;
-	 uint32_t gc_max_waves_per_simd;
-	 uint32_t gc_max_scratch_slots_per_cu;
-	 uint32_t gc_lds_size;
-	 uint32_t gc_num_sc_per_se;
-	 uint32_t gc_num_packer_per_sc;
-	 /* new for v2_1 */
-	 uint32_t gc_num_tcp_per_sh;
-	 uint32_t gc_tcp_size_per_cu;
-	 uint32_t gc_num_sdp_interface;
-	 uint32_t gc_num_cu_per_sqc;
-	 uint32_t gc_instruction_cache_size_per_sqc;
-	 uint32_t gc_scalar_data_cache_size_per_sqc;
-	 uint32_t gc_tcc_size;
- };
- 
- typedef struct harvest_info_header {
-	 uint32_t signature; /* Table Signature */
-	 uint32_t version;   /* Table Version */
- } harvest_info_header;
- 
- typedef struct harvest_info {
-	 uint16_t hw_id;          /* Hardware ID */
-	 uint8_t number_instance; /* Instance of the IP */
-	 uint8_t reserved;        /* Reserved for alignment */
- } harvest_info;
- 
- typedef struct harvest_table {
-	 harvest_info_header header;
-	 harvest_info list[32];
- } harvest_table;
- 
- struct mall_info_header {
-	 uint32_t table_id; /* table ID */
-	 uint16_t version_major; /* table version */
-	 uint16_t version_minor; /* table version */
-	 uint32_t size_bytes; /* size of the entire header+data in bytes */
- };
- 
- struct mall_info_v1_0 {
-	 struct mall_info_header header;
-	 uint32_t mall_size_per_m;
-	 uint32_t m_s_present;
-	 uint32_t m_half_use;
-	 uint32_t m_mall_config;
-	 uint32_t reserved[5];
- };
- 
- struct mall_info_v2_0 {
-	 struct mall_info_header header;
-	 uint32_t mall_size_per_umc;
-	 uint32_t reserved[8];
- };
- 
- #define VCN_INFO_TABLE_MAX_NUM_INSTANCES 4
- 
- struct vcn_info_header {
-	 uint32_t table_id; /* table ID */
-	 uint16_t version_major; /* table version */
-	 uint16_t version_minor; /* table version */
-	 uint32_t size_bytes; /* size of the entire header+data in bytes */
- };
- 
- struct vcn_instance_info_v1_0
- {
-	 uint32_t instance_num; /* VCN IP instance number. 0 - VCN0; 1 - VCN1 etc*/
-	 union _fuse_data {
-		 struct {
-			 uint32_t av1_disabled : 1;
-			 uint32_t vp9_disabled : 1;
-			 uint32_t hevc_disabled : 1;
-			 uint32_t h264_disabled : 1;
-			 uint32_t reserved : 28;
-		 } bits;
-		 uint32_t all_bits;
-	 } fuse_data;
-	 uint32_t reserved[2];
- };
- 
- struct vcn_info_v1_0 {
-	 struct vcn_info_header header;
-	 uint32_t num_of_instances; /* number of entries used in instance_info below*/
-	 struct vcn_instance_info_v1_0 instance_info[VCN_INFO_TABLE_MAX_NUM_INSTANCES];
-	 uint32_t reserved[4];
- };
- 
- #define NPS_INFO_TABLE_MAX_NUM_INSTANCES 12
- 
- struct nps_info_header {
-	 uint32_t table_id; /* table ID */
-	 uint16_t version_major; /* table version */
-	 uint16_t version_minor; /* table version */
-	 uint32_t size_bytes; /* size of the entire header+data in bytes = 0x000000D4 (212) */
- };
- 
- struct nps_instance_info_v1_0 {
-	 uint64_t base_address;
-	 uint64_t limit_address;
- };
- 
- struct nps_info_v1_0 {
-	 struct nps_info_header header;
-	 uint32_t nps_type;
-	 uint32_t count;
-	 struct nps_instance_info_v1_0
-		 instance_info[NPS_INFO_TABLE_MAX_NUM_INSTANCES];
- };
- 
- enum amd_hw_ip_block_type {
-	GC_HWIP = 1,
-	HDP_HWIP,
-	SDMA0_HWIP,
-	SDMA1_HWIP,
-	SDMA2_HWIP,
-	SDMA3_HWIP,
-	SDMA4_HWIP,
-	SDMA5_HWIP,
-	SDMA6_HWIP,
-	SDMA7_HWIP,
-	LSDMA_HWIP,
-	MMHUB_HWIP,
-	ATHUB_HWIP,
-	NBIO_HWIP,
-	MP0_HWIP,
-	MP1_HWIP,
-	UVD_HWIP,
-	VCN_HWIP = UVD_HWIP,
-	JPEG_HWIP = VCN_HWIP,
-	VCN1_HWIP,
-	VCE_HWIP,
-	VPE_HWIP,
-	DF_HWIP,
-	DCE_HWIP,
-	OSSSYS_HWIP,
-	SMUIO_HWIP,
-	PWR_HWIP,
-	NBIF_HWIP,
-	THM_HWIP,
-	CLK_HWIP,
-	UMC_HWIP,
-	RSMU_HWIP,
-	XGMI_HWIP,
-	DCI_HWIP,
-	PCIE_HWIP,
-	ISP_HWIP,
-	MAX_HWIP
-};
-
-#define HWIP_MAX_INSTANCE	44
-
-#define HW_ID_MAX		300
-
-// HW ID
-#define MP1_HWID                                           1
-#define MP2_HWID                                           2
-#define THM_HWID                                           3
-#define SMUIO_HWID                                         4
-#define FUSE_HWID                                          5
-#define CLKA_HWID                                          6
-#define PWR_HWID                                          10
-#define GC_HWID                                           11
-#define UVD_HWID                                          12
-#define VCN_HWID                                          UVD_HWID
-#define AUDIO_AZ_HWID                                     13
-#define ACP_HWID                                          14
-#define DCI_HWID                                          15
-#define DMU_HWID                                         271
-#define DCO_HWID                                          16
-#define DIO_HWID                                         272
-#define XDMA_HWID                                         17
-#define DCEAZ_HWID                                        18
-#define DAZ_HWID                                         274
-#define SDPMUX_HWID                                       19
-#define NTB_HWID                                          20
-#define VPE_HWID                                          21
-#define IOHC_HWID                                         24
-#define L2IMU_HWID                                        28
-#define VCE_HWID                                          32
-#define MMHUB_HWID                                        34
-#define ATHUB_HWID                                        35
-#define DBGU_NBIO_HWID                                    36
-#define DFX_HWID                                          37
-#define DBGU0_HWID                                        38
-#define DBGU1_HWID                                        39
-#define OSSSYS_HWID                                       40
-#define HDP_HWID                                          41
-#define SDMA0_HWID                                        42
-#define SDMA1_HWID                                        43
-#define ISP_HWID                                          44
-#define DBGU_IO_HWID                                      45
-#define DF_HWID                                           46
-#define CLKB_HWID                                         47
-#define FCH_HWID                                          48
-#define DFX_DAP_HWID                                      49
-#define L1IMU_PCIE_HWID                                   50
-#define L1IMU_NBIF_HWID                                   51
-#define L1IMU_IOAGR_HWID                                  52
-#define L1IMU3_HWID                                       53
-#define L1IMU4_HWID                                       54
-#define L1IMU5_HWID                                       55
-#define L1IMU6_HWID                                       56
-#define L1IMU7_HWID                                       57
-#define L1IMU8_HWID                                       58
-#define L1IMU9_HWID                                       59
-#define L1IMU10_HWID                                      60
-#define L1IMU11_HWID                                      61
-#define L1IMU12_HWID                                      62
-#define L1IMU13_HWID                                      63
-#define L1IMU14_HWID                                      64
-#define L1IMU15_HWID                                      65
-#define WAFLC_HWID                                        66
-#define FCH_USB_PD_HWID                                   67
-#define SDMA2_HWID                                        68
-#define SDMA3_HWID                                        69
-#define PCIE_HWID                                         70
-#define PCS_HWID                                          80
-#define DDCL_HWID                                         89
-#define SST_HWID                                          90
-#define LSDMA_HWID                                        91
-#define IOAGR_HWID                                       100
-#define NBIF_HWID                                        108
-#define IOAPIC_HWID                                      124
-#define SYSTEMHUB_HWID                                   128
-#define NTBCCP_HWID                                      144
-#define UMC_HWID                                         150
-#define SATA_HWID                                        168
-#define USB_HWID                                         170
-#define CCXSEC_HWID                                      176
-#define XGMI_HWID                                        200
-#define XGBE_HWID                                        216
-#define MP0_HWID                                         255
-
-static int hw_id_map[MAX_HWIP] = {
-	[GC_HWIP]	= GC_HWID,
-	[HDP_HWIP]	= HDP_HWID,
-	[SDMA0_HWIP]	= SDMA0_HWID,
-	[SDMA1_HWIP]	= SDMA1_HWID,
-	[SDMA2_HWIP]    = SDMA2_HWID,
-	[SDMA3_HWIP]    = SDMA3_HWID,
-	[LSDMA_HWIP]    = LSDMA_HWID,
-	[MMHUB_HWIP]	= MMHUB_HWID,
-	[ATHUB_HWIP]	= ATHUB_HWID,
-	[NBIO_HWIP]	= NBIF_HWID,
-	[MP0_HWIP]	= MP0_HWID,
-	[MP1_HWIP]	= MP1_HWID,
-	[UVD_HWIP]	= UVD_HWID,
-	[VCE_HWIP]	= VCE_HWID,
-	[DF_HWIP]	= DF_HWID,
-	[DCE_HWIP]	= DMU_HWID,
-	[OSSSYS_HWIP]	= OSSSYS_HWID,
-	[SMUIO_HWIP]	= SMUIO_HWID,
-	[PWR_HWIP]	= PWR_HWID,
-	[NBIF_HWIP]	= NBIF_HWID,
-	[THM_HWIP]	= THM_HWID,
-	[CLK_HWIP]	= CLKA_HWID,
-	[UMC_HWIP]	= UMC_HWID,
-	[XGMI_HWIP]	= XGMI_HWID,
-	[DCI_HWIP]	= DCI_HWID,
-	[PCIE_HWIP]	= PCIE_HWID,
-	[VPE_HWIP]	= VPE_HWID,
-	[ISP_HWIP]	= ISP_HWID,
-};
-
-#endif
-
-