mawei/StarPilot
1
0
Fork 0
mirror of https://github.com/firestar5683/StarPilot.git synced 2026-09-03 14:43:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

Pin

Browse Source
This commit is contained in:
firestar5683
2026-09-02 09:59:09 -05:00
parent 3d4c625deb
commit 449cb11938
301 changed files with 14621 additions and 7377 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/MODEL_REBUILD.md
+19 -5
View File
@@ -159,7 +159,21 @@ All four files must be updated together.
## Manifest
Generate the base manifest after compilation, then namespace the release artifacts as v23:
The current test branch uses manifest v25 and requests v25 only. Seed the new
manifest from the previous catalog, then replace entries as artifacts are
rebuilt with the pinned runtime:
```bash
cp /path/to/model_names_v24.json /path/to/model_names_v25.json
```
The current tinygrad pin is `f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae`, from
`openpilot` `origin/master` (`bump tg + TC_MIN_GLOBALS`). StarPilot's
multi-model `modeld` remains in place; do not replace it with upstream's
single-model `modeld`.
For the older namespace migration workflow, generate the base manifest after
compilation and namespace the release artifacts as v23:
```bash
python3 scripts/model_rebuild_pipeline.py manifest \
@@ -175,9 +189,9 @@ python3 scripts/namespace_model_artifacts.py \
The namespace command changes IDs such as `tr1422` to `tr14223`, renames the
compiled and upload-ready files, and writes an ID map. It preserves display
names and behavioral versions. The current model manager requests v23 only;
the manifest is fetched from `Models/model_names_v23.json`, while v22 remains
available for devices that have not updated yet.
names and behavioral versions. The current model manager requests v25 only; the
manifest is fetched from `Models/model_names_v25.json`. Devices still running
the prior branch continue to request their existing manifest version.
After importing newly compiled sources, normalize the release namespace before
copying files into either resource repository:
@@ -205,4 +219,4 @@ Compilation validates JIT capture/replay, pickle round-trip, finite outputs, met
4. Confirm `driverStateV2` on both supported camera resolutions.
5. Test download, selection, deletion, randomization, migration, and fallback in both device UIs and Galaxy.
The built-in RDF artifact is `selfdrive/modeld/models/driving_tinygrad.pkl`. If migration cannot download the selected v23 artifact, StarPilot switches to that built-in model.
The built-in RDF artifact is `selfdrive/modeld/models/driving_tinygrad.pkl`. If migration cannot download the selected v25 artifact, StarPilot switches to that built-in model.
scripts/model_compiler.py
+2 -3
View File
@@ -56,14 +56,13 @@ def build_compile_env(*, supercombo: bool = False) -> dict[str, str]:
existing_pythonpath = env.get("PYTHONPATH", "")
env["PYTHONPATH"] = f"{REPO_ROOT}{os.pathsep}{existing_pythonpath}" if existing_pythonpath else str(REPO_ROOT)
defaults = {
"DEBUG": "0",
"FLOAT16": "1",
"IMAGE": "1" if supercombo else "2",
"JIT_BATCH_SIZE": "0",
"NOLOCALS": "1",
"OPENPILOT_HACKS": "1",
} | ({} if supercombo else {
"DEBUG": "0",
})
}
for key, default in defaults.items():
try:
int(str(env.get(key)), 0)
scripts/model_release.py
+1 -1
View File
@@ -31,7 +31,7 @@ OPENPILOT_REPO = "commaai/openpilot"
RESOURCES_REPO = os.environ.get("STARPILOT_RESOURCES_REPO", "firestar5683/StarPilot-Resources")
HF_BUCKET = os.environ.get("STARPILOT_HF_BUCKET", "StarPilot-Driving/StarPilot-Resources")
RESOURCE_BRANCH = "Models"
MANIFEST_VERSION = "v24"
MANIFEST_VERSION = "v25"
DEFAULT_BEHAVIOR_VERSION = "v16"
DEVICE_ROOT = "/data/openpilot"
REPOSITORY_FILE_LIMIT = 100_000_000
scripts/tests/test_model_release.py
+2 -2
View File
@@ -79,14 +79,14 @@ def test_runtime_scan_excludes_model_weights_but_flags_runtime_code():
def test_update_manifest_replaces_one_entry(tmp_path: Path):
manifest = tmp_path / "model_names_v24.json"
manifest = tmp_path / "model_names_v25.json"
manifest.write_text(json.dumps({"models": [{"id": "old"}]}) + "\n")
info = parse_pasted_release(RELEASE_TEXT, "bmrlnapv4", "v16")
path = update_manifest(
tmp_path,
info,
{"size": 123, "sha256": "a" * 64},
"v24",
"v25",
)
payload = json.loads(path.read_text())
assert len(payload["models"]) == 2
selfdrive/modeld/modeld.py
+4
View File
@@ -7,6 +7,10 @@ import struct
from openpilot.system.hardware import HARDWARE, TICI
os.environ['GMMU'] = '0'
os.environ['DEV'] = 'QCOM' if TICI else 'LLVM'
try:
int(os.getenv('DEBUG', '0'), 0)
except ValueError:
os.environ['DEBUG'] = '0'
from tinygrad.device import Device
from tinygrad.tensor import Tensor
import time
starpilot/assets/model_manager.py
+1 -1
View File
@@ -24,7 +24,7 @@ from openpilot.starpilot.common.starpilot_utilities import delete_file
from openpilot.starpilot.common.starpilot_variables import MODELS_PATH
from openpilot.system.hardware.usb import chestnut_firmware_ready
MANIFEST_CANDIDATES = ("v24",)
MANIFEST_CANDIDATES = ("v25",)
MODEL_NAMESPACE_SUFFIX = "3"
DEFAULT_MODEL_KEY = "rdf43"
LOCAL_MODEL_PREFIX = "local-"
starpilot/assets/tests/test_model_pipeline.py
+7 -7
View File
@@ -15,12 +15,12 @@ from openpilot.starpilot.assets.model_manager import MANIFEST_CANDIDATES, ModelM
from openpilot.starpilot.common.model_versions import UNIFIED_ARTIFACT_FORMAT
def test_v24_is_the_only_manifest_candidate():
assert MANIFEST_CANDIDATES == ("v24",)
def test_v25_is_the_only_manifest_candidate():
assert MANIFEST_CANDIDATES == ("v25",)
def test_v24_manifest_is_loaded_from_models_checkout():
assert ModelManager._manifest_paths("v24") == ("Models/model_names_v24.json",)
def test_v25_manifest_is_loaded_from_models_checkout():
assert ModelManager._manifest_paths("v25") == ("Models/model_names_v25.json",)
def test_resource_sources_prefer_huggingface_then_github(monkeypatch):
@@ -33,9 +33,9 @@ def test_resource_sources_prefer_huggingface_then_github(monkeypatch):
def test_huggingface_manifest_has_root_and_manifests_fallbacks():
assert ModelManager._hf_manifest_paths("v24") == (
"model_names_v24.json",
"manifests/model_names_v24.json",
assert ModelManager._hf_manifest_paths("v25") == (
"model_names_v25.json",
"manifests/model_names_v25.json",
)
tinygrad_repo/.github/actions/setup-tinygrad/action.yml
Vendored
+34 -29
View File
@@ -4,7 +4,7 @@ inputs:
python-version:
description: 'Python version to use'
required: false
default: '' # if you don't set a version, the native python version will be used
default: '3.14'
key:
description: 'Key for the python cache'
required: false
@@ -41,12 +41,12 @@ inputs:
description: "Install LLVM?"
required: false
default: 'false'
tinydreno:
description: "Install tinydreno"
qemu:
description: "Install qemu?"
required: false
default: 'false'
qemu:
description: "Install qemu"
ninja:
description: "Install ninja?"
required: false
default: 'false'
runs:
@@ -59,18 +59,18 @@ runs:
echo "OMP_NUM_THREADS=1" >> "$GITHUB_ENV"
# no buffers should be over 300MB in CI
echo "MAX_BUFFER_SIZE=300000000" >> "$GITHUB_ENV"
if [[ "$RUNNER_OS" == "Linux" ]]; then
echo "VIRTUAL_ENV=/opt/venv/${{ inputs.python-version }}" >> "$GITHUB_ENV"
echo "UV_PYTHON_INSTALL_DIR=/opt/python" >> "$GITHUB_ENV"
else
echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
fi
- name: Set up uv
uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b
with:
enable-cache: 'false' # see below for manual caching
- name: Set up Python ${{ inputs.python-version }}
uses: actions/setup-python@v6
if: inputs.python-version != ''
with:
python-version: ${{ inputs.python-version }}
# **** Caching packages ****
- name: Cache Python packages (PR)
@@ -109,15 +109,15 @@ runs:
if: inputs.deps != ''
shell: bash
run: |
uv venv .venv
uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
DEPS="${{ inputs.deps }}"
uv pip install --python .venv -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
uv pip install --python "$VIRTUAL_ENV" -e ".[${DEPS// /,}]" ${{ inputs.pydeps }} --torch-backend cpu --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/
- name: Install dependencies in venv (without extra)
if: inputs.deps == ''
shell: bash
run: |
uv venv .venv
uv pip install --python .venv -e . ${{ inputs.pydeps }}
uv venv --allow-existing --python ${{ inputs.python-version }} "$VIRTUAL_ENV"
uv pip install --python "$VIRTUAL_ENV" -e . ${{ inputs.pydeps }}
- name: Prune uv cache
if: github.event_name != 'pull_request'
shell: bash
@@ -125,16 +125,15 @@ runs:
- name: Configure venv
shell: bash
run: |
echo "VIRTUAL_ENV=${{ github.workspace }}/.venv" >> "$GITHUB_ENV"
if [[ "$RUNNER_OS" == "Windows" ]]; then
echo "${{ github.workspace }}/.venv/Scripts" >> "$GITHUB_PATH"
echo "$VIRTUAL_ENV/Scripts" >> "$GITHUB_PATH"
else
echo "${{ github.workspace }}/.venv/bin" >> "$GITHUB_PATH"
echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH"
fi
# ******************* apt *******************
- name: Setup apt
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
shell: bash
run: |
sudo mkdir -p /var/cache/apt/archives
@@ -162,7 +161,7 @@ runs:
echo "deb http://apt.llvm.org/$(lsb_release -cs)/ llvm-toolchain-$(lsb_release -cs)-20 main" | sudo tee /etc/apt/sources.list.d/llvm.list
- name: Compute Package List + Hash
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
id: apt-pkgs
shell: bash
run: |
@@ -187,25 +186,37 @@ runs:
if [[ "${{ inputs.qemu }}" == "true" ]]; then
pkgs+=" qemu-user-static"
fi
# **** ninja ****
if [[ "${{ inputs.ninja }}" == "true" ]]; then
pkgs+=" ninja-build"
fi
echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"
installed=true
for pkg in $pkgs; do
info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
echo "${pkg}: ${info:-not in dpkg database}"
[[ "$info" == ii* ]] || installed=false
done
echo "installed=$installed" >> "$GITHUB_OUTPUT"
- name: Cache apt (PR)
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name == 'pull_request'
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
uses: actions/cache/restore@v5
with:
path: /var/cache/apt/archives/
key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
- name: Cache apt
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true') && github.event_name != 'pull_request'
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
uses: actions/cache@v5
with:
path: /var/cache/apt/archives/
key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
- name: Run apt Update + Install
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true')
if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
shell: bash
run: |
sudo apt -qq update || true
@@ -277,12 +288,6 @@ runs:
shell: bash
run: brew install llvm@20
# *** tinydreno ***
- name: Install tinydreno (linux)
if: inputs.tinydreno == 'true' && runner.os == 'Linux'
shell: bash
run: sudo curl -fL https://github.com/sirhcm/tinydreno/raw/refs/heads/master/libllvm-qcom.so -o /usr/lib/libllvm-qcom.so
# *** OpenCL ***
- name: Install rusticl
if: inputs.opencl == 'true'
tinygrad_repo/.github/workflows/autogen.yml
Vendored
+4 -43
View File
@@ -35,15 +35,15 @@ jobs:
key: 'autogen'
amd: 'true'
llvm: 'true'
pydeps: 'pyyaml mako'
deps: 'autogen'
- name: Install autogen support packages
run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
- name: Regenerate autogen files
run: |
find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
python3 -c "from tinygrad.runtime.autogen import opencl"
python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
python3 -c "from tinygrad.runtime.autogen.am import *"
python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
@@ -54,7 +54,7 @@ jobs:
python3 -c "from tinygrad.runtime.autogen import mesa"
python3 -c "from tinygrad.runtime.autogen import avcodec"
python3 -c "from tinygrad.runtime.autogen import llvm_qcom"
python3 -c "from tinygrad.runtime.autogen import mlx5"
python3 -c "from tinygrad.runtime.autogen import mlx5, bnxt"
python3 -c "from tinygrad.runtime.autogen import ggml_common"
REGEN=1 python3 -c "from tinygrad.runtime.autogen import libclang"
- name: Check for differences
@@ -102,42 +102,3 @@ jobs:
with:
name: autogen-macos-patch
path: autogen-macos.patch
autogen-comgr-2:
name: In-tree Autogen (comgr 2)
runs-on: ubuntu-24.04
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: 'autogen-comgr'
- name: Install autogen support packages
run: |
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.2 $(lsb_release -cs) 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 -qq update || true
sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
- name: Regenerate autogen files
run: |
rm tinygrad/runtime/autogen/comgr.py
python3 -c "from tinygrad.runtime.autogen import comgr"
- name: Check for differences
run: |
if ! git diff --quiet; then
git diff
git diff > autogen-comgr2.patch
echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
exit 1
fi
- name: Upload patch artifact
if: failure()
uses: actions/upload-artifact@v7
with:
name: autogen-comgr2-patch
path: autogen-comgr2.patch
tinygrad_repo/.github/workflows/benchmark.yml
Vendored
+148 -127
View File
@@ -88,13 +88,13 @@ jobs:
fail-fast: false
matrix:
dev: ['METAL', 'AMD', 'NV']
timeout-minutes: 60
timeout-minutes: 30
defaults:
run:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
HCQ2: '0'
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -102,16 +102,11 @@ jobs:
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py amd rmmod
./extra/hcq/hcq_smi.py amd kill_pids
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
- name: Symlink models and datasets
run: |
mkdir -p weights
ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -121,18 +116,14 @@ jobs:
run: python3 test/external/process_replay/reset.py
- name: Run llama3.2
run: BENCHMARK_LOG=llama32_3b-f16 JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m llama3.2:3b-f16 --benchmark --warmup
- name: Run qwen3.6
# qwen3.6:35b-a3b doesn't fit on mac
- name: Run qwen3.8
# qwen3.8:27b doesn't fit on mac
if: ${{ matrix.dev != 'METAL' }}
run: BENCHMARK_LOG=qwen36_35b-a3b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.6:35b-a3b --benchmark --warmup
run: BENCHMARK_LOG=qwen38_27b JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m qwen3.8:27b --benchmark --warmup
- name: Run olmoe
# just metal for now
if: ${{ matrix.dev == 'METAL' }}
run: BENCHMARK_LOG=olmoe JITBEAM=2 IGNORE_BEAM_CACHE=1 python3 -m tinygrad.llm -m olmoe --benchmark --warmup
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
# only run on machines with multiple gpus
if: ${{ matrix.dev != 'METAL' }}
run: BENCHMARK_LOG=llama3_beam_4gpu 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
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -143,13 +134,13 @@ jobs:
fail-fast: false
matrix:
dev: ['METAL', 'AMD', 'NV']
timeout-minutes: 60
timeout-minutes: 10
defaults:
run:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
HCQ2: '0'
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -157,12 +148,11 @@ jobs:
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py amd rmmod
./extra/hcq/hcq_smi.py amd kill_pids
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -182,10 +172,6 @@ jobs:
# slow on metal
if: ${{ matrix.dev != 'METAL' }}
run: time BENCHMARK_LOG=cifar DEFAULT_FLOAT=HALF STEPS=1000 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run full CIFAR training steps w 6 GPUS
# only run on machines with multiple gpus
if: ${{ matrix.dev != 'METAL' }}
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -196,13 +182,13 @@ jobs:
fail-fast: false
matrix:
dev: ['AMD', 'NV']
timeout-minutes: 60
timeout-minutes: 5
defaults:
run:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
HCQ2: '0'
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -210,12 +196,11 @@ jobs:
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py amd rmmod
./extra/hcq/hcq_smi.py amd kill_pids
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: Symlink models and datasets
run: |
mkdir -p extra/datasets
@@ -227,15 +212,8 @@ 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: Run MLPerf resnet eval on training data
run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
run: BENCHMARK_LOG=resnet_10steps DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -246,13 +224,13 @@ jobs:
fail-fast: false
matrix:
dev: ['METAL', 'AMD', 'NV']
timeout-minutes: 60
timeout-minutes: 15
defaults:
run:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
HCQ2: '0'
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
@@ -260,12 +238,11 @@ jobs:
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py amd rmmod
./extra/hcq/hcq_smi.py amd kill_pids
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -285,6 +262,59 @@ jobs:
- name: Run process replay tests
uses: ./.github/actions/process-replay
multigpubenchmark:
name: Multi-GPU Benchmarks (DEV=${{ matrix.dev }})
runs-on: [self-hosted, "${{ matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
strategy:
fail-fast: false
matrix:
dev: ['AMD', 'NV']
timeout-minutes: 20
defaults:
run:
shell: bash -e -o pipefail {0}
env:
DEV: ${{ matrix.dev }}
HCQ2: '0'
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: Symlink models and datasets
run: |
mkdir -p weights
mkdir -p extra/datasets
ln -s /raid/weights/LLaMA-3 weights/LLaMA-3
ln -s /raid/datasets/imagenet extra/datasets/imagenet
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: reset process replay
run: python3 test/external/process_replay/reset.py
- name: Run LLaMA-3 8B on 4 GPUs with BEAM
run: BENCHMARK_LOG=llama3_beam_4gpu 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
- name: Run full CIFAR training steps w 6 GPUS
run: time BENCHMARK_LOG=cifar_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF STEPS=350 BS=1536 GPUS=6 TARGET_EVAL_ACC_PCT=93.0 python3 examples/hlb_cifar10.py
- name: Run MLPerf resnet eval on training data
run: time BENCHMARK_LOG=resnet_eval MODEL=resnet python3 examples/mlperf/model_eval.py
- name: Run 10 MLPerf ResNet50 training steps (6 gpu)
run: BENCHMARK_LOG=resnet_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=1536 GPUS=6 MODEL=resnet python3 examples/mlperf/model_train.py
- name: Run 10 MLPerf Bert training steps (6 gpu)
# TODO: remove BERT_LAYERS once scheduler is fast
run: BENCHMARK_LOG=bert_10steps_6gpu CAPTURE_PROCESS_REPLAY=0 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=72 GPUS=6 BERT_LAYERS=2 MODEL=bert python3 examples/mlperf/model_train.py
- name: Run process replay tests
uses: ./.github/actions/process-replay
tests:
name: Tests (DEV=${{ matrix.dev }})
runs-on: [self-hosted, "${{ matrix.dev == 'METAL' && 'macOS' || matrix.dev == 'AMD' && 'tinybox' || 'tinyboxgreen' }}"]
@@ -292,7 +322,7 @@ jobs:
fail-fast: false
matrix:
dev: ['METAL', 'AMD', 'NV']
timeout-minutes: 60
timeout-minutes: 10
defaults:
run:
shell: bash -e -o pipefail {0}
@@ -305,12 +335,11 @@ jobs:
- name: Setup (AMD)
if: ${{ matrix.dev == 'AMD' }}
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py amd rmmod
./extra/hcq/hcq_smi.py amd kill_pids
./extra/hcq/hcq_smi.py amd rmmod --expect
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
- name: Setup (NV)
if: ${{ matrix.dev == 'NV' }}
run: sudo lsof -tQ /dev/nvidia* | { xargs -r sudo kill -9 || true; }
run: lsof -tQ /dev/nvidia* | { xargs -r kill -9 || true; }
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
@@ -374,7 +403,7 @@ jobs:
run: python test/external/external_benchmark_multitensor_allreduce.py
- name: HEVC Decode Benchmark
if: ${{ matrix.dev == 'NV' }}
run: VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
- uses: actions/upload-artifact@v7
if: ${{ matrix.dev != 'AMD' }}
with:
@@ -387,7 +416,7 @@ jobs:
testusbgpu:
name: UsbGPU Benchmark
runs-on: [self-hosted, macOS]
timeout-minutes: 10
timeout-minutes: 3
defaults:
run:
shell: bash -e -o pipefail {0}
@@ -402,32 +431,66 @@ jobs:
rm -f /tmp/staging.db /tmp/staging.db-shm /tmp/staging.db-wal
- name: Kill stale pids
run: |
PYTHONPATH=. ./extra/hcq/hcq_smi.py amd kill_pids
PYTHONPATH=. ./extra/hcq/hcq_smi.py nv kill_pids
# since sudo is required for usbgpu on macos, do not write bytecode, as some of the files are owned by root
./extra/hcq/hcq_smi.py amd kill_pids --sudoless
./extra/hcq/hcq_smi.py nv kill_pids --sudoless
- name: UsbGPU boot time
run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
run: GMMU=0 DEBUG=2 AM_RESET=1 DEV=USB+AMD time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU tiny tests
run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
run: GMMU=0 DEV=USB+AMD python3.11 test/test_tiny.py
- name: UsbGPU copy speeds
run: sudo -E PYTHONDONTWRITEBYTECODE=1 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
#- name: UsbGPU openpilot test
# run: sudo -E PYTHONPATH=. GMMU=0 DEV=USB+AMD GRAPH_ONE_KERNEL=1 python3.11 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/9118973ed03c1ae1d40cf69a29507ec2cc78efd7/selfdrive/modeld/models/supercombo.onnx
run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3.11 test/external/external_test_usb_asm24.py TestDevCopySpeeds
- name: UsbGPU (USB4/TB) install script
run: PYTHONPATH=. sh extra/setup_tinygpu_osx.sh
run: sh extra/setup_tinygpu_osx.sh
- name: UsbGPU (USB4/TB) boot time
run: PYTHONPATH=. DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
run: DEBUG=3 DEV=PCI+NV:NAK time python3.11 test/test_tiny.py TestTiny.test_plus
- name: UsbGPU (USB4/TB) tiny tests
run: PYTHONPATH=. DEV=PCI+NV:NAK python3.11 test/test_tiny.py
run: DEV=PCI+NV:NAK python3.11 test/test_tiny.py
testcommalatest:
name: comma Benchmark (0.11.2)
runs-on: [self-hosted, Linux, comma]
timeout-minutes: 12
testcomma:
strategy:
matrix:
dev: ['QCOM', 'QCOM:IR3']
version: ['0.11.0', '0.11.2']
model: ['vision', 'policy', 'supercombo', 'dmonitoring']
# exclude non-existent models
exclude: [{ version: '0.11.0', model: supercombo }, { version: '0.11.2', model: vision }, { version: '0.11.2', model: policy }]
include:
- version: '0.11.0'
model: vision
url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
timing: 18
- version: '0.11.0'
model: policy
url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
timing: 3.4
- version: '0.11.0'
model: dmonitoring
url: https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
timing: 13
- version: '0.11.2'
model: supercombo
url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
timing: 28
- dev: QCOM:IR3
version: '0.11.2'
model: supercombo
timing: 29
- version: '0.11.2'
model: dmonitoring
url: https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
timing: 12.5
fail-fast: false
name: openpilot ${{ matrix.version }} compile3 ${{ matrix.model }} (DEV=${{ matrix.dev }})
runs-on: [self-hosted, Linux, comma4]
timeout-minutes: 5
defaults:
run:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
env:
DEV: ${{ matrix.dev }}
ASSERT_MIN_STEP_TIME: ${{ matrix.timing }}
BENCHMARK_LOG: ${{ matrix.dev == 'QCOM:IR3' && 'ir3_' || '' }}openpilot_${{ matrix.version }}_${{ matrix.model }}
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -438,45 +501,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 compile3 0.11.2 supercombo
run: BENCHMARK_LOG=openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
- name: openpilot compile3 0.11.2 supercombo (from pickle)
run: BENCHMARK_LOG=openpilot_0_11_2_supercombo_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=26 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
- name: IR3 openpilot compile3 0.11.2 supercombo
run: BENCHMARK_LOG=ir3_openpilot_0_11_2_supercombo PYTHONPATH="." ASSERT_MIN_STEP_TIME=41 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/433f85f956837606ad1f1cbee4aa7e2158ad23c768dea914b20436c97232741b
- name: openpilot compile3 0.11.2 dmonitoring
run: BENCHMARK_LOG=openpilot_0_11_2_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://gitlab.com/commaai/openpilot-lfs.git/gitlab-lfs/objects/3e7b31dfbc0a5234f1baf196513b77fc6af12204b8a8ffe8ee0417e48352f316
- name: Run process replay tests
uses: ./.github/actions/process-replay
testcommaold:
name: comma Benchmark (0.11.0)
runs-on: [self-hosted, Linux, comma]
timeout-minutes: 12
defaults:
run:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
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 compile3 0.11.0 driving_vision
run: BENCHMARK_LOG=openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.11.0 driving_vision (from pickle)
run: BENCHMARK_LOG=openpilot_0_11_0_vision_run_pickle RUN_PICKLE=1 PYTHONPATH="." ASSERT_MIN_STEP_TIME=17 DEV=QCOM taskset -c 4-7 python3 examples/openpilot/compile3.py
- name: IR3 openpilot compile3 0.11.0 driving_vision
run: BENCHMARK_LOG=ir3_openpilot_0_11_0_vision PYTHONPATH="." ASSERT_MIN_STEP_TIME=18 DEV=QCOM:IR3 FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_vision.onnx
- name: openpilot compile3 0.11.0 driving_policy
run: BENCHMARK_LOG=openpilot_0_11_0_policy PYTHONPATH="." ASSERT_MIN_STEP_TIME=3.2 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/driving_policy.onnx
- name: openpilot compile3 0.11.0 dmonitoring
run: BENCHMARK_LOG=openpilot_0_11_0_dmonitoring PYTHONPATH="." ASSERT_MIN_STEP_TIME=11 DEV=QCOM FLOAT16=1 IMAGE=1 NOLOCALS=1 taskset -c 4-7 python3 examples/openpilot/compile3.py https://github.com/commaai/openpilot/raw/v0.11.0/selfdrive/modeld/models/dmonitoring_model.onnx
- name: compile
run: FLOAT16=1 IMAGE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py ${{ matrix.url }} openpilot.pkl
- name: run pickle
run: BENCHMARK_LOG="${BENCHMARK_LOG}_run_pickle" RUN_PICKLE=1 taskset -c 4-7 python3 examples/openpilot/compile3.py - openpilot.pkl
- name: Run process replay tests
uses: ./.github/actions/process-replay
@@ -489,15 +517,6 @@ jobs:
shell: bash -e -o pipefail {0}
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: setup staging db
if: github.ref == 'refs/heads/update_benchmark_staging'
run: |
echo "CACHEDB=/tmp/staging.db" >> $GITHUB_ENV
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: Checkout Code
uses: actions/checkout@v6
- name: setup staging db
@@ -510,8 +529,8 @@ jobs:
- name: benchmark MobileNetV2 on DSP
run: |
# generate quantized weights
ln -s /data/home/tiny/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
ln -s /data/home/tiny/tinygrad/testsig-*.so .
ln -s ~/tinygrad/extra/datasets/imagenet extra/datasets/imagenet
ln -s ~/tinygrad/testsig-*.so .
PYTHONPATH=. DEV=CPU QUANT=1 CNT=0 python3 examples/test_onnx_imagenet.py https://github.com/xamcat/mobcat-samples/raw/refs/heads/master/onnx_runtime/InferencingSample/InferencingSample/mobilenetv2-7.onnx /tmp/model.quant.onnx
# benchmark on DSP with NOOPT=1, the devectorizer has issues
PYTHONPATH=. DEV=DSP NOOPT=1 CNT=2 DEBUG=2 python3 examples/test_onnx_imagenet.py /tmp/model.quant.onnx
@@ -521,7 +540,7 @@ jobs:
testcommausbgpubenchmark:
name: UsbGPU Benchmark (comma)
runs-on: [self-hosted, Linux, comma4]
timeout-minutes: 20
timeout-minutes: 10
defaults:
run:
shell: bash -e -o pipefail {0}
@@ -541,7 +560,7 @@ jobs:
- name: openpilot run_pickle big_driving_supercombo
run: BENCHMARK_LOG=usbgpu_openpilot_big_driving_supercombo_run_pickle RUN_PICKLE=1 PICKLE_OOB=1 PYTHONPATH="." GMMU=0 DEV=USB+AMD ASSERT_MIN_STEP_TIME=50 python3 examples/openpilot/compile3.py - openpilot.pkl
- name: Test copy speeds
run: SIZE=64e6 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
run: SIZE=64000000 PYTHONPATH=. GMMU=0 DEV=USB+AMD python3 test/external/external_test_usb_asm24.py TestDevCopySpeeds
driverbenchmarks:
name: PCI Driver Benchmark (DEV=${{ matrix.dev }})
@@ -550,7 +569,7 @@ jobs:
fail-fast: false
matrix:
dev: ['AMD', 'NV']
timeout-minutes: 20
timeout-minutes: 5
defaults:
run:
shell: bash -e -o pipefail {0}
@@ -562,9 +581,8 @@ jobs:
uses: actions/checkout@v6
- name: Setup
run: |
./extra/amdpci/setup_python_cap.sh
./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} rmmod
./extra/hcq/hcq_smi.py ${{ matrix.dev == 'AMD' && 'amd' || 'nv' }} kill_pids
./extra/hcq/hcq_smi.py ${{ matrix.dev }} rmmod --expect
./extra/hcq/hcq_smi.py ${{ matrix.dev }} kill_pids --sudoless
mkdir -p extra/datasets
ln -s /raid/datasets/imagenet extra/datasets/imagenet
- name: setup staging db
@@ -599,6 +617,9 @@ jobs:
run: |
GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyDefaulttoCPUJit
GRAPH_ONE_KERNEL=1 NSZ=8192 python3 test/speed/external_test_copy_speed.py TestCopySpeed.testCopyCPUtoDefaultJit
- name: HEVC Decode Benchmark
if: ${{ matrix.dev == 'NV' }}
run: IGNORE_BEAM_CACHE=1 VALIDATE=1 MAX_FRAMES=100 ASSERT_FPS=1400 JITBEAM=1 PYTHONPATH=. python3 extra/hevc/decode.py
- name: Run 10 MLPerf ResNet50 training steps (1 gpu)
if: ${{ matrix.dev == 'NV' }}
run: BENCHMARK_LOG=resnet_10steps MNISTMOCK=1 DEFAULT_FLOAT=HALF BENCHMARK=10 BS=256 GPUS=1 MODEL=resnet python3 examples/mlperf/model_train.py
@@ -621,12 +642,12 @@ jobs:
llvmspeed:
name: LLVM Speed
runs-on: [self-hosted, Linux, tinyboxrandom]
timeout-minutes: 20
timeout-minutes: 10
if: github.repository_owner == 'tinygrad'
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Speed Test
run: DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
run: DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
- name: Speed Test (BEAM=2)
run: BEAM=2 DEV=CPU:LLVM THREADS=0 python3 test/speed/external_test_speed_v_torch.py
run: IGNORE_BEAM_CACHE=1 BEAM=2 DEV=CPU:LLVM python3 test/speed/external_test_speed_v_torch.py
tinygrad_repo/.github/workflows/docs.yml
Vendored
+1 -1
View File
@@ -8,7 +8,7 @@ permissions:
contents: write
jobs:
deploy:
runs-on: ubuntu-latest
runs-on: ubuntu-24.04
steps:
- uses: actions/checkout@v6
- name: Configure Git Credentials
tinygrad_repo/.github/workflows/platform.yml
Vendored
+1 -33
View File
@@ -166,7 +166,7 @@ jobs:
uses: ./.github/actions/setup-tinygrad
with:
key: windows-${{ matrix.dev }}-minimal
deps: testing_unit
deps: testing_minimal
pydeps: ${{ matrix.dev == 'WEBGPU' && 'dawn-python' || '' }}
- name: Set env
shell: bash
@@ -179,35 +179,3 @@ jobs:
- name: Run test_tiny
shell: bash
run: python -m pytest -n=auto test/test_tiny.py --durations=20
qcomclcompiletests:
name: Compile-only (QCOM CL)
runs-on: ubuntu-24.04-arm
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: compile-qcomcl
deps: testing_unit
tinydreno: 'true'
- name: Set env
shell: bash
run: printf "DEV=NULL:QCOMCL:a630\nNULL_ALLOW_COPYOUT=1" >> $GITHUB_ENV
- name: Run test_ops
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
python -m pytest -n=auto test/backend/test_ops.py --durations=20
- name: Run test_ops (IMAGE)
shell: bash
env:
IMAGE: 1
DEV: "NULL:QCOMCL:a630,IMAGE_PITCH_ALIGNMENT=64"
run: |
DEBUG=4 python test/backend/test_ops.py TestOps.test_gemm | grep read_imagef
python -m pytest -n=auto test/backend/test_ops.py --durations=20
tinygrad_repo/.github/workflows/python-publish.yml
Vendored
+1 -1
View File
@@ -10,7 +10,7 @@ on:
jobs:
deploy:
runs-on: ubuntu-latest
runs-on: ubuntu-24.04
steps:
- uses: actions/checkout@v6
- name: Set up Python
tinygrad_repo/.github/workflows/szdiff.yml
Vendored
+3 -3
View File
@@ -10,7 +10,7 @@ concurrency:
jobs:
checkbranch:
name: Check PR Branch status
runs-on: ubuntu-latest
runs-on: ubuntu-24.04
outputs:
branchstat: ${{ steps.brstat.outputs.stat}}
steps:
@@ -44,7 +44,7 @@ jobs:
permissions:
contents: read
pull-requests: write
runs-on: ubuntu-latest
runs-on: ubuntu-24.04
needs: checkbranch
if: needs.checkbranch.outputs.branchstat == 'false'
steps:
@@ -87,7 +87,7 @@ jobs:
name: Core Library Line Difference
permissions:
pull-requests: write
runs-on: ubuntu-latest
runs-on: ubuntu-24.04
needs: checkbranch
if: needs.checkbranch.outputs.branchstat == 'true'
steps:
tinygrad_repo/.github/workflows/test.yml
Vendored
+66 -53
View File
@@ -21,7 +21,7 @@ concurrency:
jobs:
docs:
name: Docs
runs-on: &linux ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 10
env:
CHECK_OOB: 0
@@ -31,8 +31,7 @@ jobs:
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
deps: docs
pydeps: "capstone torch"
deps: "docs testing_minimal"
- name: Build wheel and show size
run: |
uv build --wheel
@@ -61,7 +60,7 @@ jobs:
torchbackend:
name: Torch Backend Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -73,10 +72,7 @@ jobs:
deps: testing_unit
pydeps: "pillow torchvision expecttest"
llvm: 'true'
- name: Install ninja
run: |
sudo apt update || true
sudo apt install -y --no-install-recommends ninja-build
ninja: 'true'
- name: Test ResNet-18
run: DEBUG=2 python3 extra/torch_backend/example.py
- name: Test one op in torch tests
@@ -86,9 +82,26 @@ jobs:
- name: Custom tests
run: DEV=CPU:LLVM GPUS=4 TINY_BACKEND=1 python3 -m pytest -nauto extra/torch_backend/test.py extra/torch_backend/test_inplace.py extra/torch_backend/test_multigpu.py extra/torch_backend/test_kernel_fusion.py --durations=20
torchbackendtrain:
name: Torch Backend Training
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
uses: actions/checkout@v6
- name: Setup Environment
uses: ./.github/actions/setup-tinygrad
with:
key: torch-backend-pillow-torchvision-et-pt
deps: testing_unit
llvm: 'true'
ninja: 'true'
- name: Test beautiful_mnist in torch with TINY_BACKEND
run: STEPS=20 DEV=CPU TARGET_EVAL_ACC_PCT=90.0 MAX_BUFFER_SIZE=0 TINY_BACKEND=1 python3 examples/other_mnist/beautiful_mnist_torch.py
bepython:
name: Python Backend
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -126,7 +139,7 @@ jobs:
linter:
name: Linters
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 10
steps:
@@ -157,7 +170,7 @@ jobs:
nulltest:
name: Null Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
@@ -191,7 +204,7 @@ jobs:
unittest:
name: Unit Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
@@ -220,7 +233,7 @@ jobs:
- name: Run process replay tests
uses: ./.github/actions/process-replay
- name: Repo line count <= 26000 lines
run: MAX_LINE_COUNT=26000 python sz.py
run: MAX_LINE_COUNT=26500 python sz.py
spec:
strategy:
@@ -228,7 +241,7 @@ jobs:
matrix:
group: [1, 2]
name: SPEC=2 (${{ matrix.group }})
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -244,7 +257,7 @@ jobs:
fuzzing:
name: Fuzzing
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 10
steps:
- name: Checkout Code
@@ -260,7 +273,7 @@ jobs:
testopenclimage:
name: CL IMAGE Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -280,7 +293,7 @@ jobs:
testopenpilot:
name: openpilot Compile Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -309,7 +322,7 @@ jobs:
testonnxcpu:
name: ONNX (CPU) Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 20
steps:
@@ -328,7 +341,7 @@ jobs:
testoptim:
name: Optimization Tests
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 20
steps:
- name: Checkout Code
@@ -360,7 +373,7 @@ jobs:
testllm:
name: Test LLM
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
env:
CHECK_OOB: 0
@@ -377,17 +390,17 @@ jobs:
run: |
parallel --link --tagstring '[{1}]' '{2}' \
::: llama 'llama q4' qwen3.5 qwen \
::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
$'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
# NOTE: qwen is dumb and only knows about female chickens
# ****** Models Tests ******
testmodels:
name: Models
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -407,7 +420,7 @@ jobs:
testdsp:
name: Linux (DSP)
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
steps:
- name: Checkout Code
@@ -435,7 +448,7 @@ jobs:
- 'WEBGPU'
name: Linux (DEV=${{ matrix.dev }})
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 20
steps:
- name: Checkout Code
@@ -461,7 +474,7 @@ jobs:
testamdasm:
name: AMD ASM IDE
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 20
env:
DEV: MOCKKFD+AMD
@@ -491,7 +504,7 @@ jobs:
- name: Run AMD renderer tests (AMD:LLVM)
run: DEV=MOCKKFD+AMD:LLVM python -m pytest -n=auto test/amd/ --durations 20
- name: Run SQTT profiling tests
run: PROFILE=1 SQTT=1 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
run: VIZ=-2 python3 -m pytest -n=auto test/amd/test_sqtt_profiler.py
- name: Run AMD emulated tests on NULL backend
env:
AMD: 0
@@ -507,7 +520,7 @@ jobs:
hcq2:
name: hcq2
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 5
steps:
- name: Checkout Code
@@ -521,16 +534,15 @@ jobs:
- name: Run HCQ2 tests
run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/test_tiny.py
- name: Run HCQ2 multi-device tests
run: |
HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_multitensor.py \
TestMultiTensor.test_simple_add TestMultiTensor.test_shard_reduce \
TestMultiTensor.test_backward_sum TestMultiTensor.test_matmul_shard_0_0
run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest -n=auto test/backend/test_multitensor.py
- name: Run HCQ2 JIT tests
run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python test/unit/test_jit.py
- name: Run HCQ2 unit tests
run: HCQ_RUNTIME_DEV=PYTHON HCQ2=1 DEV=MOCKKFD+AMD FORWARD_ONLY=1 PYTHONPATH=. python -m pytest test/device/test_hcq2.py
testmockam:
name: Linux (am)
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
env:
DEV: MOCKPCI+AMD
@@ -566,7 +578,7 @@ jobs:
arch: [gfx1100, gfx1201, gfx950]
name: Linux (${{ matrix.backend }} ${{ matrix.arch }})
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
env:
DEV: MOCKKFD+AMD:${{ matrix.backend == 'amdllvm' && 'LLVM' || '' }}:${{ matrix.arch }}
@@ -589,7 +601,7 @@ jobs:
if: ${{ matrix.backend == 'amd' && matrix.arch == 'gfx950' }}
run: PYTHONPATH=. DEV=NULL:HIP:gfx950 MXFP4=1 LLAMA_LAYERS=2 BENCHMARK=3 NULL_ALLOW_COPYOUT=1 NO_HIPCC=1 ROCM_PATH=/opt/rocm JITBEAM=0 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/profile.sh
- name: Run pytest (amd)
run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM test/opt/test_tensor_cores.py --durations=20
run: python -m pytest -n=auto test/backend/test_ops.py test/backend/test_dtype.py test/backend/test_dtype_alu.py test/backend/test_linearizer.py test/backend/test_randomness.py test/backend/test_jit.py test/backend/test_graph.py test/backend/test_multitensor.py test/device/test_hcq.py test/external/external_test_am.py test/backend/test_asm_gemm.py::TestAsmGEMM --durations=20
- name: Run disk copy tests
run: python -m pytest test/unit/test_disk_tensor.py -k test_copy_from_disk
- name: Run TRANSCENDENTAL math
@@ -604,7 +616,7 @@ jobs:
backend: [ptx, nv]
name: Linux (${{ matrix.backend }})
runs-on: *linux
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 20
env:
FORWARD_ONLY: 1
@@ -638,10 +650,17 @@ jobs:
strategy:
fail-fast: false
matrix:
backend: [ir3, nak]
name: Compile-only (${{ matrix.backend }})
runs-on: *linux
dev:
- 'NULL:IR3:a630'
- 'NULL:QCOMCL:a630'
- 'NULL:NAK:sm_120'
name: Compile-only (DEV=${{ matrix.dev }})
runs-on: ${{ github.repository == 'tinygrad/tinygrad' && github.event_name == 'pull_request' && github.event.pull_request.author_association == 'COLLABORATOR' && 'namespace-profile-tinygrad' || 'ubuntu-24.04' }}
timeout-minutes: 15
env:
NULL_ALLOW_COPYOUT: 1
DEV: ${{ matrix.dev }}${{ contains(matrix.dev, 'a630') && ',IMAGE_PITCH_ALIGNMENT=64' || '' }}
IMAGE: ${{ contains(matrix.dev, 'a630') && '1' || '0' }}
steps:
- name: Checkout Code
uses: actions/checkout@v6
@@ -650,21 +669,15 @@ jobs:
with:
key: compile-${{ matrix.backend }}
deps: "testing_unit mesa"
- name: Set env
qemu: ${{ contains(matrix.dev, 'QCOMCL') }}
- name: Test IMAGE
shell: bash
run: printf "NULL_ALLOW_COPYOUT=1\n${{ matrix.backend == 'ir3' && 'DEV=NULL:IR3:a630' || matrix.backend == 'nak' && 'DEV=NULL:NAK:sm_120' }}" >> $GITHUB_ENV
if: contains(matrix.dev, 'a630')
run: DEBUG=7 python3 test/backend/test_ops.py TestOps.test_gemm | grep isam
- name: Run test_ops
shell: bash
run: |
python -c "from tinygrad import Device; assert Device.DEFAULT == 'NULL'"
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_add
python -m pytest -n=auto test/backend/test_ops.py --durations=20
- name: Run test_ops (IMAGE)
if: matrix.backend == 'ir3'
shell: bash
env:
IMAGE: 1
DEV: "NULL:IR3:a630,IMAGE_PITCH_ALIGNMENT=64"
run: |
DEBUG=4 python3 test/backend/test_ops.py TestOps.test_gemm | grep image_load
python -m pytest -n=auto test/backend/test_ops.py --durations=20
# QCOMCL compiles in qemu, too slow for parallel workers
${{ contains(matrix.dev, 'QCOMCL') && 'PARALLEL=0' || '' }} python -m pytest -n=auto test/backend/test_ops.py --durations=20
tinygrad_repo/.gitignore
+1
View File
@@ -69,3 +69,4 @@ mutants
dagre/
graphlib/
uv.lock
pi_session_window0.jsonl
tinygrad_repo/.pre-commit-config.yaml
+1 -1
View File
@@ -28,7 +28,7 @@ repos:
pass_filenames: false
- id: tests
name: comprehensive test suite
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/unit/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
entry: env OMP_NUM_THREADS=1 SKIP_SLOW_TEST=1 PYTHONPATH="." python3 -m pytest -n=6 test/backend/test_ops.py test/backend/test_schedule.py test/backend/test_assign.py test/backend/test_tensor.py test/backend/test_jit.py test/unit/test_schedule_cache.py test/null/test_pattern_matcher.py test/null/test_uop_symbolic.py test/unit/test_helpers.py
language: system
always_run: true
pass_filenames: false
tinygrad_repo/AGENTS.md
+1
View File
@@ -4,3 +4,4 @@
- Run `python -m mypy tinygrad/` to typecheck
- Run `python -m ruff check .` to lint
- Read `./tinygrad/viz/README.md` for profiling and debugging rewrite rules
- Do not do amend commits. Always do a new commit if a force push to origin would be required.
tinygrad_repo/README.md
+1 -1
View File
@@ -140,7 +140,7 @@ Documentation along with a quick start guide can be found on the [docs website](
```python
from tinygrad import Tensor
x = Tensor.eye(3)
x = Tensor.eye(3).clone() # clone to make it a buffer
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
z.backward()
tinygrad_repo/TINYGRAD_COMMIT
+1 -1
View File
@@ -1 +1 @@
8611fe22a7fcc7d1928bbde19ded66277cb12f3e
f6fc4e3f2c3db5fae1e19cbfbc3ad9fc579a12ae
tinygrad_repo/conftest.py
+1 -1
View File
@@ -2,7 +2,7 @@ import os, pytest, signal, threading
@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item):
t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
t.start()
try: yield
finally:
tinygrad_repo/docs/abstractions4.py
+1 -1
View File
@@ -122,7 +122,7 @@ def example_5_custom_assembly(a:Tensor, correct):
offset_dwords = (self.labels[inst._target] - inst._pos - inst.size()) // 4
if not -32768 <= offset_dwords <= 32767: raise ValueError(f"branch to '{inst._target}' offset {offset_dwords} exceeds simm16 range")
inst.simm16 = offset_dwords
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in self.instructions]))))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in self.instructions]))))
CU_COUNT = 32
LANES = 64
tinygrad_repo/docs/developer/speed.md
+1 -1
View File
@@ -52,7 +52,7 @@ In `kernel.py` we have a set of `OptOps`, these control the parameters of the sp
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.
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. Splitting an axis into UPCAST 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.
tinygrad_repo/examples/mlperf/dataloader.py
+5 -5
View File
@@ -5,7 +5,7 @@ from multiprocessing import Queue, Process, shared_memory, connection, Lock
import numpy as np
from tinygrad import dtypes, Tensor
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, NUM_CPU_THREADS
from tinygrad.helpers import getenv, prod, Context, round_up, tqdm, OSX, CPU_COUNT
from tinygrad.nn.state import TensorIO
### ResNet
@@ -131,7 +131,7 @@ def batch_load_resnet(batch_size=64, val=False, shuffle=True, seed=None, pad_fir
else: X = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name}")
Y = [None] * (batch_size*BATCH_COUNT)
for _ in range(NUM_CPU_THREADS.value):
for _ in range(CPU_COUNT):
p = Process(target=loader_process, args=(q_in, q_out, X, seed))
p.daemon = True
p.start()
@@ -212,7 +212,7 @@ def batch_load_train_bert(BS:int, seed:int|None=None):
rng.shuffle(fs)
train_files.append(fs.pop(0))
cycle_length = min(NUM_CPU_THREADS.value, len(train_files))
cycle_length = min(CPU_COUNT, len(train_files))
assert cycle_length > 0, "cycle_length must be greater than 0"
dataset = InterleavedDataset(train_files, cycle_length)
@@ -301,7 +301,7 @@ def batch_load_unet3d(preprocessed_dataset_dir:Path, batch_size:int=6, val:bool=
X = Tensor.empty(*sz, dtype=dtypes.float32, device=f"disk:/dev/shm/{shm_name_x}")
Y = Tensor.empty(*sz, dtype=dtypes.uint8, device=f"disk:/dev/shm/{shm_name_y}")
for _ in range(NUM_CPU_THREADS.value):
for _ in range(CPU_COUNT):
proc = Process(target=load_unet3d_data, args=(preprocessed_dataset_dir, seed, queue_in, queue_out, X, Y))
proc.daemon = True
proc.start()
@@ -437,7 +437,7 @@ def batch_load_retinanet(dataset, val:bool, base_dir:Path, batch_size:int=32, sh
dataset_iter = iter(image_ids)
try:
for _ in range(NUM_CPU_THREADS.value):
for _ in range(CPU_COUNT):
proc = Process(
target=load_retinanet_data,
args=(base_dir, val, queue_in, queue_out, imgs, boxes, labels),
tinygrad_repo/examples/mlperf/model_train.py
+30 -42
View File
@@ -1667,15 +1667,14 @@ def train_llama3():
def train_gptoss():
from examples.mlperf.models.gpt_oss import GPTOSS, GPT_OSS_20B, apply_grad, FP8_DTYPE
from examples.mlperf.lr_schedulers import CosineAnnealingLRWithWarmup
from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, clip_grads
from examples.mlperf.optim import GradAccClipAdamW, GradAccClipAdamWGroup, fclip_grads
BENCHMARK = getenv("BENCHMARK")
config = {}
BASEDIR = config["BASEDIR"] = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
BS = config["BS"] = getenv("BS", 16)
grad_acc = config["GRADIENT_ACC_STEPS"] = getenv("GRADIENT_ACC_STEPS", 1)
GBS = config["GLOBAL_BATCH_SIZE"] = BS * grad_acc
GBS = config["GLOBAL_BATCH_SIZE"] = BS
SEED = config["SEED"] = getenv("SEED", 5760)
DATA_SEED = config["DATA_SEED"] = getenv("DATA_SEED", SEED)
SEQLEN = config["SEQLEN"] = getenv("SEQLEN", 8192)
@@ -1737,13 +1736,13 @@ def train_gptoss():
params_wd = [p for p in params if p.ndim >= 3]
params_no_wd = [p for p in params if p.ndim < 3]
optim = GradAccClipAdamWGroup(
GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=grad_acc, device=optim_device),
GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=grad_acc, device=optim_device),
GradAccClipAdamW(params_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=opt_adamw_weight_decay, grad_acc=1, device=optim_device),
GradAccClipAdamW(params_no_wd, lr=0.0, b1=opt_adamw_beta_1, b2=opt_adamw_beta_2, eps=opt_adamw_epsilon, weight_decay=0.0, grad_acc=1, device=optim_device),
)
for p in optim.params:
grad_dtype = dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype
p.grad = p.zeros_like(dtype=grad_dtype).contiguous()
p.grad = p.zeros_like(dtype=dtypes.bfloat16 if p.dtype == FP8_DTYPE else p.dtype).contiguous()
if getattr(p, "_zero2", False): p.grad = optim.optimizers[0]._zero_shard(p.grad)
grads = [p.grad for p in optim.params]
from extra.gemm.cdna_asm_gemm import _mx_block_scale
@@ -1770,31 +1769,32 @@ def train_gptoss():
@TinyJit
@Context(TRAINING=1)
def minibatch(tokens:Tensor):
def step(tokens:Tensor):
if is_dp: tokens = tokens.to(None).shard(device, 0)
if not is_sharding: tokens = tokens.to(None)
logits:Tensor = model(tokens[:, :-1], save=True)
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
if getenv("FUSED_CE", 0):
from extra.llama_kernels.fused_ce import fused_ce_loss
loss = fused_ce_loss(logits.cast(dtypes.bfloat16), tokens[:, 1:], label_smoothing=0.0)
else:
loss = logits.sparse_categorical_crossentropy(tokens[:, 1:])
for g, new_g in zip(grads, loss.gradient(*optim.params)):
apply_grad(g, new_g.uop)
loss_cpu = loss.flatten().float().to("CPU")
return loss_cpu.realize(*grads)
Tensor.realize(loss, *grads)
@TinyJit
def optim_step():
grad_norm = clip_grads(grads, grad_acc, 1.0)
optim.fstep(grads, grad_norm)
clipped_grads, grad_norm = fclip_grads(grads, 1.0)
optim.fstep(clipped_grads, grad_norm)
scheduler.step()
for g in grads: g.assign(0)
loss_cpu = loss.flatten().float().to("CPU")
lr_cpu = optim.lr.float().to("CPU")
grad_norm_cpu = grad_norm.float().to("CPU")
Tensor.realize(lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
Tensor.realize(loss_cpu, lr_cpu, grad_norm_cpu, *grads, *fp8_inv_scales)
return lr_cpu, grad_norm_cpu
return loss_cpu, lr_cpu, grad_norm_cpu
@TinyJit
@Context(TRAINING=0)
@@ -1843,30 +1843,20 @@ def train_gptoss():
profile_marker(f"train @ {i}")
st = time.perf_counter()
stopped = False
losses, data_time, dev_time = [], 0, 0
for _ in range(grad_acc if i >= 2 else 1):
ist = time.perf_counter()
try: tokens = next(train_iter)
except StopIteration:
stopped = True
break
mst = time.perf_counter()
data_time += mst - ist
losses.append(minibatch(tokens).item())
dev_time += time.perf_counter() - mst
if stopped: break
ist = time.perf_counter()
gt = time.perf_counter()
ret = optim_step()
lr, grad_norm = ret[0].item(), ret[1].item()
try: tokens = next(train_iter)
except StopIteration: break
mst = time.perf_counter()
data_time = mst - ist
ret = step(tokens)
dev_time = time.perf_counter() - mst
loss, lr, grad_norm = ret[0].item(), ret[1].item(), ret[2].item()
et = time.perf_counter()
loss = sum(losses) / len(losses)
optim_time = et - gt
dev_time += optim_time
step_time = et - st
gbs_time = gt - st
if BENCHMARK: step_times.append(step_time)
i += 1
@@ -1876,7 +1866,7 @@ def train_gptoss():
gflops = GlobalCounters.global_ops / 1e9 / dev_time
mfu = ((6 * num_params * SEQLEN * GBS) / (dev_time * device_count * 4.6e15)) * 100
tqdm.write(
f"{i:5} {step_time:.3f} s step, {gbs_time:.3f} s gbs, {optim_time:.3f} s optim, {data_time:.3f} s data, {loss:.4f} loss, " \
f"{i:5} {step_time:.3f} s step, {dev_time:.3f} s dev, {data_time:.3f} s data, {loss:.4f} loss, " \
f"{lr:.12f} LR, {grad_norm:.6f} grad_norm, {mem_gb:.2f} GB used, {gflops:9.2f} GFLOPS, {mfu:5.2f}% MFU")
if DEBUG >= 1: tqdm.write(" mem per device: " + ', '.join(f"{dev}: {mem/1e9:.2f} GB" for dev, mem in sorted(GlobalCounters.mem_used_per_device.items())))
@@ -1886,8 +1876,6 @@ def train_gptoss():
"train/lr": lr,
"train/grad_norm": grad_norm,
"train/step_time": step_time,
"train/gbs_time": gbs_time,
"train/optim_time": optim_time,
"train/dev_time": dev_time,
"train/data_time": data_time,
"train/mem": mem_gb,
tinygrad_repo/examples/mlperf/models/gpt_oss.py
+55 -15
View File
@@ -12,7 +12,7 @@ from tinygrad.helpers import Timing, colored, GlobalCounters, profile_marker
from tinygrad.uop.ops import Ops, UOp
from extra.models.llama import apply_rotary_emb
from extra.llama_kernels.rmsnorm import rmsnorm
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8
from extra.gemm.cdna_asm_gemm import _mx_block_scale, _mx_block_scale_3d, quantize_mxfp8, asm_gemm, can_use_asm_gemm
from extra.gemm.moe_gemm import grouped_mx_gemm
from extra.gemm.moe_routing import route, dispatch, combine
@@ -61,7 +61,25 @@ def dequant_weight(w_q:Tensor, w_scale:Tensor) -> Tensor:
call = UOp.maketuple(fxn.uop).call(w_q.uop, w_scale.uop, grad_fxn=_dequant_bwd)
return Tensor(call.gettuple(0))
def matmul_mx(x:Tensor, w_q:Tensor, w_scale:Tensor) -> Tensor:
def matmul_mx(x:Tensor|tuple[Tensor, Tensor], w_q:Tensor, w_scale:Tensor) -> Tensor:
if isinstance(x, tuple):
assert ASM_GEMM, "pre-quantized MXFP8 input requires ASM_GEMM"
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
x_q, x_e8 = x
l_shape, padded = x_q.shape[:-1], x_q.shape[-1]
x_q, x_e8 = x_q.reshape(-1, padded), x_e8.reshape(-1, padded // 32)
K, N = w_q.shape[1], w_q.shape[0]
assert padded >= K and (padded - K) % 32 == 0 and x_e8.shape[-1] == padded // 32
wq, ws = w_q, w_scale
if (pad := padded - K):
wq = wq.pad(((0, 0), (0, pad)))
ws = ws.pad(((0, 0), (0, pad // 32)), value=127).cast(dtypes.uint8)
if (npad := (-N) % 256):
wq = wq.pad(((0, npad), (0, 0)))
ws = ws.pad(((0, npad), (0, 0)), value=127).cast(dtypes.uint8)
assert can_use_asm_gemm(x_q, wq.T)
out = asm_gemm(x_q, wq.T, mx=True, mx_scales=(mx_pack(x_e8), x_e8, mx_pack(ws), ws), mx_w_stored=True)
return (out[:, :N] if npad else out).reshape(*l_shape, N).cast(dtypes.bfloat16)
l_shape = x.shape[:-1]
if ASM_GEMM:
from extra.gemm.cdna_asm_gemm import asm_gemm, can_use_asm_gemm, mx_pack
@@ -146,6 +164,7 @@ class GPTOSS:
return w_q, w_e8.is_param_(False)
if moe:
qs = [_one(*shape[1:]) for _ in range(shape[0])]
for q in qs: q[0]._zero2 = True # grad arrives sharded on the expert axis under ZeRO-2 (moe_gemm)
return [q[0] for q in qs], [q[1] for q in qs]
return _one(*shape)
@@ -174,20 +193,30 @@ class GPTOSS:
def attention(self, x:Tensor, freqs_cis:Tensor, mask:Tensor, sliding:bool, *, attention_norm:Tensor, wqkv:Tensor,
wqkv_scale:Tensor, wqkv_bias:Tensor, wo:Tensor, wo_scale:Tensor, wo_bias:Tensor, sinks:Tensor):
bsz, seqlen, _ = x.shape
x_normed, rrms = rmsnorm(x, self.norm_eps)
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
if getenv("FUSED_RMSNORM_MUL", 0):
from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
x_normed, rrms = rmsnorm_mul(x, attention_norm, self.norm_eps)
norm_saves = [x_normed, rrms]
else:
x_normed, rrms = rmsnorm(x, self.norm_eps)
qkv = matmul_mx(x_normed * attention_norm, wqkv, wqkv_scale) + wqkv_bias
norm_saves = [x_normed, rrms]
qkv = qkv.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep + 2, self.head_dim)
xq = qkv[:, :, :, :self.n_rep].reshape(bsz, seqlen, self.n_heads, self.head_dim)
xk, xv = qkv[:, :, :, self.n_rep], qkv[:, :, :, self.n_rep + 1]
xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
xq, xk, xv = xq.cast(dtypes.bfloat16), xk.cast(dtypes.bfloat16), xv.cast(dtypes.bfloat16) # (B,N,H,D)/(B,N,KV,D)
if sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
elif getenv("HK_FLASH_ATTENTION"):
fa_saves = []
if getenv("HK_FLASH_ATTENTION"):
from extra.thunder.amd.fa import flash_attention
attn, *_ = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks)
attn, _, l_vec = flash_attention(xq, xk, xv, is_causal=True, write_flat=True, sinks=sinks, window=self.sliding_window if sliding else 0)
attn = attn.reshape(bsz, seqlen, self.n_heads * self.head_dim)
fa_saves = [xq, xk, xv, l_vec]
elif sliding:
attn = self._sliding_attention(xq, xk, xv, sinks)
else:
xqm = xq.reshape(bsz, seqlen, self.n_kv_heads, self.n_rep, self.head_dim).permute(0, 2, 3, 1, 4)
xkm, xvm = xk.permute(0, 2, 1, 3).unsqueeze(2), xv.permute(0, 2, 1, 3).unsqueeze(2)
@@ -199,13 +228,19 @@ class GPTOSS:
attn = (w @ xvm).permute(0, 3, 1, 2, 4).reshape(bsz, seqlen, self.n_heads * self.head_dim)
out = matmul_mx(attn, wo, wo_scale) + wo_bias
return out, [x_normed, rrms, attn]
return out, [attn] + norm_saves + fa_saves
def feed_forward(self, x:Tensor, *, ffn_norm:Tensor, gate:Tensor, gate_bias:Tensor,
w_gate_up:Tensor, w_gate_up_scale:Tensor, w_gate_up_bias:Tensor,
w_down:Tensor, w_down_scale:Tensor, w_down_bias:Tensor):
x_normed, rrms = rmsnorm(x, self.norm_eps)
inp = x_normed * ffn_norm
if getenv("FUSED_RMSNORM_MUL", 0):
from extra.gptoss_kernels.rmsnorm import rmsnorm_mul
x_normed, rrms = rmsnorm_mul(x, ffn_norm, self.norm_eps)
inp = x_normed
else:
x_normed, rrms = rmsnorm(x, self.norm_eps)
inp = x_normed * ffn_norm
logits = inp.float() @ gate.float().T + gate_bias.float()
dim, inter = self.dim, self.intermediate_size
@@ -220,6 +255,7 @@ class GPTOSS:
z = grouped_mx_gemm(_pad_cols(y.cast(dtypes.bfloat16)), (w_down, w_down_scale), r.off)[:, :dim] \
+ (onehot @ w_down_bias.float()).cast(dtypes.bfloat16)
out = combine(z, r, inp.shape[0], self.experts_per_tok).reshape(bsz, seqlen, dim)
return out, [x_normed, rrms, xg, h, y, z, r.weights, r.dest_row, r.off]
else:
thresh = logits.topk(self.experts_per_tok)[0][..., -1:]
weights = (logits >= thresh).where(logits, -float("inf")).softmax(-1)
@@ -263,7 +299,11 @@ class GPTOSS:
w_down=self.w_down[i], w_down_scale=self.w_down_scale[i], w_down_bias=self.w_down_bias[i])
h, *_ = self.run_layer(h, freqs_cis, mask_full, i % 2 == 0, attn_kwargs, ffn_kwargs, save=save)
logits = self.norm(h) @ self.output.T
h_normed = self.norm(h)
pad = (-self.dim) % 256
h_padded, w_padded = h_normed.pad((None, None, (0, pad))), self.output.pad(((0, 0), (0, pad)))
if ASM_GEMM and can_use_asm_gemm(h_padded, w_padded.T): logits = asm_gemm(h_padded, w_padded.T)
else: logits = h_normed @ self.output.T
return logits
def _get_pads(uop:UOp) -> list[UOp]:
@@ -274,14 +314,14 @@ def apply_grad(grad_buf:Tensor, new_grad:UOp):
pads = _get_pads(new_grad)
if len(pads) <= 1:
new_grad = new_grad.cast(grad_buf.dtype)
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(grad_buf.uop + new_grad))
grad_buf.uop = grad_buf.uop.after(grad_buf.uop.store(new_grad))
return
cur = grad_buf.uop
for pad in sorted(pads, key=lambda p: p.marg[0][0] if p.op == Ops.PAD else 0, reverse=True):
if pad.op == Ops.PAD:
grad_shrink = tuple([(p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg)])
grad_shrink = tuple((p[0], s+p[0]) for s,p in zip(pad.src[0].shape, pad.marg))
buf_slice = cur.shrink(grad_shrink)
cur = cur.after(buf_slice.store(buf_slice + pad.src[0].cast(cur.dtype)))
cur = cur.after(buf_slice.store(pad.src[0].cast(cur.dtype)))
else:
cur = cur.after(cur.store(cur + pad.cast(cur.dtype)))
grad_buf.uop = cur
tinygrad_repo/examples/mlperf/optim.py
+6 -1
View File
@@ -15,7 +15,7 @@ def stochastic_round_bf16(x:Tensor) -> Tensor:
bits = x.bitcast(dtypes.uint32)
if isinstance(x.device, tuple):
shape = x.uop.shard_shape if x.uop.axis is not None else x.shape
noise = Tensor(UOp(Ops.MSTACK, dtypes.default_float, tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
noise = Tensor(UOp(Ops.MSTACK, src=tuple(Tensor.rand(*shape, device=d).uop for d in x.device)))
else:
noise = x.rand_like()
noise = (noise * 0xFFFF).cast(dtypes.uint32)
@@ -27,6 +27,11 @@ def clip_grads(grads:list[Tensor], grad_acc, clip_norm) -> Tensor:
for g in grads: g.assign((g * (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)).cast(g.dtype))
return total_norm
def fclip_grads(grads:list[Tensor], clip_norm) -> Tensor:
total_norm = Tensor.stack(*[g.float().square().sum() for g in grads]).sum().sqrt().contiguous()
scale = (clip_norm / (total_norm + 1e-6)).clamp(max_=1.0)
return [(g * scale).cast(g.dtype) for g in grads], total_norm
class GradAccClipAdamW(Optimizer):
def __init__(self, params:list[Tensor], lr=0.001, b1=0.9, b2=0.999, eps=1e-6, weight_decay=0.0, grad_acc=1, clip_norm=1.0, device=None, fused=FUSE_OPTIM):
super().__init__(params, lr, device, fused)
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py
+13
View File
@@ -0,0 +1,13 @@
from pathlib import Path
from examples.mlperf.dataloader import get_llama3_dataset
from tinygrad.helpers import getenv
BASEDIR = Path(getenv("BASEDIR", "/raid/datasets/c4-8b/"))
SAMPLES = getenv("SAMPLES", 1_200_000 * 32)
EVAL_SAMPLES = getenv("EVAL_SAMPLES", 1024)
SEQLEN = getenv("SEQLEN", 8192)
DATA_SEED = getenv("DATA_SEED", 5760)
get_llama3_dataset(SAMPLES, SEQLEN, BASEDIR, seed=DATA_SEED, val=False, small=True)
get_llama3_dataset(EVAL_SAMPLES, SEQLEN, BASEDIR, seed=0, val=True, small=True)
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh
+3 -3
View File
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
export GBS=$((BS * GRADIENT_ACC_STEPS))
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
@@ -44,7 +44,7 @@ export SEED=${SEED:-5760}
export DATA_SEED=${DATA_SEED:-5760}
export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export FAKEDATA=${FAKEDATA:-1} BENCHMARK=${BENCHMARK:-10}
if [ -z "$FULL_LAYERS" ]; then
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh
+5 -5
View File
@@ -1,8 +1,8 @@
#!/usr/bin/env bash
export PYTHONPATH="."
export PATH="/opt/rocm-7.1.1/bin:$PATH"
export ROCM_PATH="/opt/rocm-7.1.1"
export ROCM_PATH=${ROCM_PATH:-/opt/rocm-7.1.1}
export PATH="$ROCM_PATH/bin:$PATH"
export DEV=${DEV:-AMD}
export CHECK_OOB=0
export REWRITE_STACK_LIMIT=5000000 HCQDEV_WAIT_TIMEOUT_MS=240000
@@ -26,7 +26,7 @@ export FUSED_SILU_W13=${FUSED_SILU_W13:-1}
export SPLIT_W13=${SPLIT_W13:-0}
export OFFLOAD_OPTIM=${OFFLOAD_OPTIM:-0}
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
export DP=${DP:-8} MP=${MP:-1} BS=${BS:-16} EVAL_BS=${EVAL_BS:-8} GRADIENT_ACC_STEPS=${GRADIENT_ACC_STEPS:-2}
export GBS=$((BS * GRADIENT_ACC_STEPS))
@@ -35,7 +35,7 @@ export BASEDIR="/raid/datasets/c4-8b/"
export SMALL=1
export LLAMA3_SIZE=${LLAMA3_SIZE:-"8B"}
export EVAL_TARGET=3.3 EVAL_FREQ=12288
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=4096 MAX_STEPS=1200000
export LR="1e-3" END_LR="1e-4" WARMUP_SAMPLES=2048 MAX_STEPS=1200000
export WARMUP_STEPS=$((WARMUP_SAMPLES / GBS))
export SAMPLES=$((MAX_STEPS * GBS))
export SEQLEN=${SEQLEN:-8192}
@@ -44,6 +44,6 @@ export SEED=${SEED:-$RANDOM}
export DATA_SEED=${DATA_SEED:-5760}
export JITBEAM=${JITBEAM:-3}
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
python3 examples/mlperf/model_train.py
tinygrad_repo/examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/run_and_time.sh
+3 -3
View File
@@ -17,7 +17,7 @@ export USE_ATOMICS=1
export ASM_GEMM=1
export WQKV=1
export MASTER_WEIGHTS=1
export FP8=1
export MXFP4=1
export ALLREDUCE_CAST=1
export FAST_CE=1
export FUSED_INPUT_QUANTIZE=1
@@ -26,7 +26,7 @@ export FUSED_ADD_NORM_MUL_QUANTIZE=1
export FUSED_SILU_W13=1
export SPLIT_W13=0
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="bfloat16"
export DEFAULT_FLOAT="bfloat16" OPTIM_DTYPE="float32"
export DP=8 MP=1 BS=16 EVAL_BS=8 GRADIENT_ACC_STEPS=2
export GBS=$((BS * GRADIENT_ACC_STEPS))
@@ -44,7 +44,7 @@ export SEED=$RANDOM
export DATA_SEED=$SEED
export JITBEAM=3
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=1
export BEAM_UOPS_MAX=6000 BEAM_UPCAST_MAX=256 BEAM_LOCAL_MAX=1024 BEAM_MIN_PROGRESS=5 BEAM_PADTO=0
export LOGMLPERF=1
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_beam.sh
+1
View File
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
export DEBUG=${DEBUG:-2}
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export GROUPED_MOE=${GROUPED_MOE:-1}
export ALL2ALL=${ALL2ALL:-1}
export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
tinygrad_repo/examples/mlperf/training_submission_v6.1/tinycorp/benchmarks/gpt_oss/implementations/tinybox_8xMI350X/dev_run.sh
+1
View File
@@ -11,6 +11,7 @@ export DEVICE_IN_FUNCTION_BUG=1
export DEBUG=${DEBUG:-0}
export HK_FLASH_ATTENTION=${HK_FLASH_ATTENTION:-1}
export ASM_GEMM=${ASM_GEMM:-1}
export GROUPED_MOE=${GROUPED_MOE:-1}
export ALL2ALL=${ALL2ALL:-1}
export LATE_ALLREDUCE=${LATE_ALLREDUCE:-0}
export ALLREDUCE_CAST=${ALLREDUCE_CAST:-1}
tinygrad_repo/examples/openpilot/compile3.py
+10 -12
View File
@@ -107,14 +107,21 @@ def compile(onnx_file):
return inputs, test_val
def test_vs_compile(run, inputs, test_val=None):
if (log:=bool(getenv("BENCHMARK_LOG", ""))): from extra.bench_log import WallTimeEvent, BenchEvent
# run 20 times
step_times = []
for _ in range(20):
st = time.perf_counter()
out = run(**inputs)
mt = time.perf_counter()
val = out.numpy()
if log:
with WallTimeEvent(BenchEvent.STEP):
out = run(**inputs)
mt = time.perf_counter()
val = out.numpy()
else:
out = run(**inputs)
mt = time.perf_counter()
val = out.numpy()
et = time.perf_counter()
step_times.append((et-st)*1e3)
print(f"enqueue {(mt-st)*1e3:6.2f} ms -- total run {step_times[-1]:6.2f} ms")
@@ -160,12 +167,6 @@ def test_vs_onnx(new_inputs, test_val, onnx_file, tol):
print("test vs onnx passed")
return timings
def bench(run, inputs):
from extra.bench_log import WallTimeEvent, BenchEvent
for _ in range(10):
with WallTimeEvent(BenchEvent.STEP):
run(**inputs).numpy()
if __name__ == "__main__":
if getenv("RUN_PICKLE"):
with open(OUTPUT, "rb") as f: pickle_loaded = load_pickle(f)
@@ -181,6 +182,3 @@ if __name__ == "__main__":
test_vs_compile(pickle_loaded, inputs, outputs)
if getenv("SELFTEST"):
test_vs_onnx(inputs, outputs, onnx_file, 1e-4)
if getenv("BENCHMARK_LOG", ""):
bench(pickle_loaded, inputs)
tinygrad_repo/extra/amdpci/am_smi.py
+4 -2
View File
@@ -84,7 +84,8 @@ class AMSMI(AMDev):
with open(f"/sys/bus/pci/devices/{self.pcibus}/power_state", "r") as f: return f.read().strip().rstrip()
class SMICtx:
def __init__(self):
def __init__(self, dev_filter=None):
self.dev_filter = dev_filter
self.devs = []
self.opened_pcidevs = []
self.opened_pci_resources = {}
@@ -135,6 +136,7 @@ class SMICtx:
pattern = os.path.join('/tmp', 'am_*.lock')
for d in [f[8:-5] for f in glob.glob(pattern)]:
if d.startswith("usb"): continue
if self.dev_filter is not None and d != self.dev_filter: continue
if d not in self.opened_pcidevs:
self._open_am_device(d)
@@ -406,7 +408,7 @@ if __name__ == "__main__":
try:
if not args.list: os.system('clear')
smi_ctx = SMICtx()
smi_ctx = SMICtx(args.dev)
while True:
smi_ctx.rescan_devs()
smi_ctx.draw(args.list)
tinygrad_repo/extra/bench_log.py
+9 -9
View File
@@ -35,7 +35,7 @@ class WallTimeEvent:
return self
def __exit__(self, *_):
self.time = time.monotonic() - self.start
_events[self.event]["wall"].append(self.time)
_events[self.event]["wall"].append((self.time, BENCHMARK_LOG.value))
return False
class KernelTimeEvent:
@@ -47,19 +47,19 @@ class KernelTimeEvent:
self.start = GlobalCounters.time_sum_s
return self
def __exit__(self, *_):
_events[self.event]["kernel"].append(GlobalCounters.time_sum_s - self.start)
_events[self.event]["kernel"].append((GlobalCounters.time_sum_s - self.start, BENCHMARK_LOG.value))
return False
def log_event_instant(event:InstantBenchEvent, value:float):
_events[event].append(value)
_events[event].append((value, BENCHMARK_LOG.value))
if BENCHMARK_LOG:
INFLUXDB_HOST = getenv("INFLUXDB_HOST", "")
INFLUXDB_ORG = getenv("INFLUXDB_ORG", "tiny")
INFLUXDB_TOKEN = getenv("INFLUXDB_TOKEN", "")
def _create_point(run_id, i, attempt, ref, commit, name, value, run):
point = Point(BENCHMARK_LOG.value).tag("id", run_id).tag("index", i)
def _create_point(run_id, i, attempt, ref, commit, name, value, log_name, run):
point = Point(log_name.replace(':', '_').replace('.', '_')).tag("id", run_id).tag("index", i)
point = point.tag("device", Device.DEFAULT)
point = point.tag("attempt", attempt).tag("ref", ref).tag("commit", commit)
point = point.field(name, value).field("x", run)
@@ -91,12 +91,12 @@ if BENCHMARK_LOG:
run_id = str(uuid.uuid4())
if isinstance(event, BenchEvent):
for event_type, values in _events[event].items():
for i, value in enumerate(values):
point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, run)
for i, (value, log_name) in enumerate(values):
point = _create_point(run_id, i, attempt, ref, commit, f"{event.value}_{event_type}", value, log_name, run)
points.append(point)
else:
for i, value in enumerate(_events[event]):
point = _create_point(run_id, i, attempt, ref, commit, event.value, value, run)
for i, (value, log_name) in enumerate(_events[event]):
point = _create_point(run_id, i, attempt, ref, commit, event.value, value, log_name, run)
points.append(point)
write_options = WriteOptions(write_type=WriteType.synchronous, retry_interval=5000, max_retries=5, max_retry_delay=30000, exponential_base=2)
tinygrad_repo/extra/bnxt_driver/bnxtdev.py
+238
View File
@@ -0,0 +1,238 @@
import ctypes, struct
from tinygrad.helpers import ceildiv, getenv, wait_cond, DEBUG
from tinygrad.runtime.autogen import bnxt, pci
from tinygrad.runtime.support.system import PCIDevice, System, ipv4_to_gid
BNXT_DEBUG = getenv("BNXT_DEBUG", 0)
BNXT_ACCESS, BNXT_INIT_MASK, BNXT_RTR_MASK, BNXT_RTS_MASK = 3, 0xd, 0x41515ad, 0xae005
BNXT_CHIMP_COMM, BNXT_CHIMP_COMM_TRIGGER = 0x0, 0x100
BNXT_BACKING_STORE = ((0, 2), (1, 0), (2, 2), (3, 0), (4, 2), (5, 0), (6, 0), (14, 2), (15, 0))
def db_value(xid, typ, index, epoch):
return (xid & bnxt.DBC_DBC_XID_MASK | bnxt.DBC_DBC_PATH_ROCE | typ | bnxt.BNXT_QPLIB_DBR_VALID) << 32 | \
index & bnxt.DBC_DBC_INDEX_MASK | epoch << bnxt.BNXT_QPLIB_DBR_EPOCH_SHIFT
def _pbl(dev, paddrs, queue=False):
if len(paddrs) == 1: return 0, paddrs[0]
values = [p | bnxt.PTU_PTE_VALID for p in paddrs]
if queue:
values[-1] |= bnxt.PTU_PTE_LAST
if len(values) > 1: values[-2] |= bnxt.PTU_PTE_NEXT_TO_LAST
table, table_paddrs = dev.pci_dev.alloc_sysmem(ceildiv(len(values), 512) * 0x1000)
table[:len(values) * 8] = struct.pack(f"<{len(values)}Q", *values)
if len(table_paddrs) == 1: return 1, table_paddrs[0]
top, top_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
top[:len(table_paddrs) * 8] = struct.pack(f"<{len(table_paddrs)}Q", *(p | bnxt.PTU_PTE_VALID for p in table_paddrs))
return 2, top_paddrs[0]
def _queue(dev, stride:int=16, aux=False):
mem, paddrs = dev.pci_dev.alloc_sysmem(0x1000 + aux * 0x400)
level, base = _pbl(dev, paddrs, queue=True)
return {"mem":mem, "paddrs":paddrs, "stride":stride, "prod":0, "cons":0, "level":level, "base":base}
def _qread(q, i):
off = (i & 15) * q["stride"]
return q["mem"][off:off + q["stride"]]
def _qwrite(q, i, data, aux=False):
off = 0x1000 + i % 128 * 8 if aux else (i & 15) * q["stride"]
q["mem"][off:off + len(data)] = data
class BNXTDev:
def __init__(self, pci_dev:PCIDevice, ip:str=getenv("BNXT_IP", "10.0.0.1")):
self.pci_dev, self.devfmt = pci_dev, pci_dev.pcibus
self.bar0, self.db = pci_dev.map_bar(0, fmt='I'), pci_dev.map_bar(2, fmt='Q')
pci_dev.write_config(pci.PCI_COMMAND, pci_dev.read_config(pci.PCI_COMMAND, 2) | pci.PCI_COMMAND_MASTER, 2)
self.resp, self.resp_pa = pci_dev.alloc_sysmem(0x1000)
self.seq = 0
ver = self.hwrm("ver_get")
if DEBUG >= 2: print(f"bnxt {self.devfmt}: firmware {ver.hwrm_fw_maj_8b}.{ver.hwrm_fw_min_8b}.{ver.hwrm_fw_bld_8b}")
self.hwrm("func_reset", timeout_ms=40000)
caps = self.hwrm("func_qcaps", fid=0xffff)
self.mac, self.port_id = int.from_bytes(bytes(caps.mac_address), 'big'), caps.port_id
self.hwrm("func_drv_rgtr")
self.db_off = self.hwrm("func_qcfg", fid=0xffff).legacy_l2_db_size_kb * 1024
self.setup_backing_store()
self._open_rcfw()
self._open_l2()
self.local_gid = ipv4_to_gid(ip)
gids, mac = (ctypes.c_uint32 * 4)(*(int.from_bytes(self.local_gid[i:i + 4], 'big') for i in (12, 8, 4, 0))), self.mac.to_bytes(6, 'big')
smac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac[i:i + 2], 'big') for i in (0, 2, 4)))
self.gid_id = self.rcfw("add_gid", gid=gids, src_mac=smac).xid
if DEBUG >= 2: print(f"bnxt {self.devfmt}: booted mac={self.mac.to_bytes(6, 'big').hex(':')} gid={self.local_gid.hex()}")
def hwrm(self, name, timeout_ms=10000, **fields):
inp, out = getattr(bnxt, f"struct_hwrm_{name}_input"), getattr(bnxt, f"struct_hwrm_{name}_output")
opcode = getattr(bnxt, f"HWRM_{name.upper()}")
self.seq = (self.seq + 1) & 0xffff
data = bytes(inp(req_type=opcode, cmpl_ring=bnxt.BNXT_HWRM_NO_CMPL_RING, seq_id=self.seq, target_id=bnxt.BNXT_HWRM_TARGET,
resp_addr=self.resp_pa[0], **fields))
self.resp[:] = bytes(len(self.resp))
System.memory_barrier()
for i, w in enumerate(memoryview(bytearray(data.ljust(bnxt.HWRM_MAX_REQ_LEN, b'\0'))).cast('I')):
self.bar0[BNXT_CHIMP_COMM // 4 + i] = w
self.bar0[BNXT_CHIMP_COMM_TRIGGER // 4] = 1
def hdr(): return bnxt.struct_hwrm_resp_hdr.from_buffer_copy(bytes(self.resp[:8]))
wait_cond(lambda: (n := hdr().resp_len) and hdr().seq_id == self.seq and self.resp[n - 1], timeout_ms=timeout_ms, msg=f"HWRM {name}")
ret = out.from_buffer_copy(bytes(self.resp[:ctypes.sizeof(out)]))
assert ret.error_code == 0, f"HWRM {name}: {ret.error_code}"
return ret
def setup_backing_store(self):
counts: dict[int, int] = {}
for typ, extra in BNXT_BACKING_STORE:
caps = self.hwrm("func_backing_store_qcaps_v2", type=typ)
size, splits = caps.entry_size, tuple(getattr(caps, f"split_entry_{j}") for j in range(caps.subtype_valid_cnt))
counts[typ] = n = counts[0] if typ == 15 else max(caps.min_num_entries, sum(splits) + extra)
# a zero bitmap means the type has a single instance 0
for instance in [i for i in range(8) if caps.instance_bit_map >> i & 1] or [0]:
mem, paddrs = self.pci_dev.alloc_sysmem(ceildiv(n * size, 0x1000) * 0x1000)
if caps.ctx_init_value:
for off in range(caps.ctx_init_offset, len(mem), size): mem[off] = caps.ctx_init_value
lvl, base = _pbl(self, paddrs)
self.hwrm("func_backing_store_cfg_v2", type=typ, instance=instance, entry_size=size, num_entries=n, page_dir=base,
page_size_pbl_level=lvl, subtype_valid_cnt=len(splits),
flags=bnxt.FUNC_BACKING_STORE_CFG_V2_REQ_FLAGS_BS_CFG_ALL_DONE if typ == 15 else 0,
**{f"split_entry_{j}": v for j, v in enumerate(splits)})
def _open_rcfw(self):
self.rcfw_first = True
self.creq = _queue(self)
self.creq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=self.creq["base"],
page_size=12, page_tbl_depth=self.creq["level"], length=16, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
self.cmdq = _queue(self)
self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, 0, 0)
init = bnxt.struct_cmdq_init(cmdq_pbl=self.cmdq["base"], creq_ring_id=self.creq_id,
cmdq_size_cmdq_lvl=16 << bnxt.CMDQ_INIT_CMDQ_SIZE_SFT)
System.memory_barrier()
for i, w in enumerate(memoryview(bytearray(bytes(init))).cast('I')): self.bar0[bnxt.RCFW_COMM_BASE_OFFSET // 4 + i] = w
_, p = self.pci_dev.alloc_sysmem(0x1000)
self.rcfw("initialize_fw", stat_ctx_id=self.hwrm("stat_ctx_alloc", stats_dma_addr=p[0], stats_dma_length=176).stat_ctx_id,
flags=bnxt.CMDQ_INITIALIZE_FW_FLAGS_HW_REQUESTER_RETX_SUPPORTED)
# RoCE notification ring: never armed or serviced, but CQ and L2 ring allocation require one
nq = _queue(self)
self.nq_id = self.hwrm("ring_alloc", ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_NQ, page_tbl_addr=nq["base"],
page_size=12, page_tbl_depth=nq["level"], length=16, logical_id=1, int_mode=bnxt.RING_ALLOC_REQ_INT_MODE_MSIX).ring_id
def rcfw(self, name, timeout_ms=20000, **fields):
req_t, resp_t = getattr(bnxt, f"struct_cmdq_{name}"), getattr(bnxt, f"struct_creq_{name}_resp")
op = getattr(bnxt, f"CMDQ_BASE_OPCODE_{name.upper()}")
data = bytes(req_t(opcode=op, cmd_size=(slots := ceildiv(ctypes.sizeof(req_t), 16)), **fields)).ljust(slots * 16, b'\0')
for i in range(slots): _qwrite(self.cmdq, self.cmdq["prod"] + i, data[i * 16:(i + 1) * 16])
self.cmdq["prod"] += slots
prod = self.cmdq["prod"] & 0xffff
if self.rcfw_first: prod, self.rcfw_first = prod | 1 << bnxt.FIRMWARE_FIRST_FLAG, False
System.memory_barrier()
self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_PF_VF_COMM_PROD_OFFSET) // 4] = prod
self.bar0[(bnxt.RCFW_COMM_BASE_OFFSET + bnxt.RCFW_COMM_TRIG_OFFSET) // 4] = bnxt.RCFW_CMDQ_TRIG_VAL
def poll():
h = bnxt.struct_creq_base.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
return bool(h.v & bnxt.CREQ_BASE_V) != bool((self.creq["cons"] // 16) & 1)
wait_cond(poll, timeout_ms=timeout_ms, msg=f"RCFW {name}")
ret = resp_t.from_buffer_copy(bytes(_qread(self.creq, self.creq["cons"])))
self.creq["cons"] += 1
# NQ_ARM also publishes the CREQ consumer index, which is what frees ring space for the next command
self.doorbell(self.creq_id, bnxt.DBC_DBC_TYPE_NQ_ARM, self.creq["cons"] & 15, (self.creq["cons"] // 16) & 1)
assert ret.status == 0, f"RCFW {name}: {ret.status}"
if BNXT_DEBUG >= 1: print(f"bnxt {self.devfmt}: rcfw {name} xid={getattr(ret, 'xid', 0):#x}")
return ret
def doorbell(self, xid, typ, index, epoch):
System.memory_barrier()
self.db[self.db_off // 8] = db_value(xid, typ, index, epoch)
# L2 receive path, required for RoCE ingress even though no ethernet receive buffers are posted
def _open_l2(self):
cq = _queue(self)
ci = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_L2_CMPL,
page_tbl_addr=cq["base"], page_size=12, page_tbl_depth=cq["level"], length=16, nq_ring_id=self.nq_id).ring_id
rx = _queue(self)
ri = self.hwrm("ring_alloc", enables=bnxt.RING_ALLOC_REQ_ENABLES_NQ_RING_ID_VALID |
bnxt.RING_ALLOC_REQ_ENABLES_RX_BUF_SIZE_VALID, ring_type=bnxt.RING_ALLOC_REQ_RING_TYPE_RX, page_tbl_addr=rx["base"],
page_size=12, page_tbl_depth=rx["level"], length=16, rx_buf_size=640, nq_ring_id=self.nq_id).ring_id
vi = self.hwrm("vnic_alloc").vnic_id
self.hwrm("vnic_cfg", enables=bnxt.VNIC_CFG_REQ_ENABLES_MRU | bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_RX_RING_ID |
bnxt.VNIC_CFG_REQ_ENABLES_DEFAULT_CMPL_RING_ID, vnic_id=vi, mru=9018,
default_rx_ring_id=ri, default_cmpl_ring_id=ci)
self.hwrm("cfa_l2_filter_alloc", flags=bnxt.CFA_L2_FILTER_ALLOC_REQ_FLAGS_PATH_RX,
enables=bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR | bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_L2_ADDR_MASK |
bnxt.CFA_L2_FILTER_ALLOC_REQ_ENABLES_DST_ID, l2_addr=tuple(self.mac.to_bytes(6, 'big')), l2_addr_mask=(0xff,) * 6, dst_id=vi)
def register_mem(self, paddrs:list[int], size:int, log_page_size:int=12) -> int:
level, base = _pbl(self, paddrs[:ceildiv(size, 1 << log_page_size)])
return self.rcfw("register_mr", flags=bnxt.CMDQ_REGISTER_MR_FLAGS_ALLOC_MR,
log2_pg_size_lvl=level << bnxt.CMDQ_REGISTER_MR_LVL_SFT | log_page_size << bnxt.CMDQ_REGISTER_MR_LOG2_PG_SIZE_SFT,
access=bnxt.CMDQ_REGISTER_MR_ACCESS_LOCAL_WRITE | bnxt.CMDQ_REGISTER_MR_ACCESS_REMOTE_WRITE,
log2_pbl_pg_size=12, pbl=base, va=paddrs[0], mr_size=size).xid
class BNXTQP:
def __init__(self, dev:BNXTDev):
self.dev, self.sq_psn, self.msn = dev, 0, 0
self.cqq = _queue(dev, ctypes.sizeof(bnxt.struct_cq_base))
self.cq_id = dev.rcfw("create_cq", cq_size=16, pbl=self.cqq["base"],
pg_size_lvl=self.cqq["level"], cq_fco_cnq_id=dev.nq_id).xid
self.sq = _queue(dev, aux=True)
self.qpn = dev.rcfw("create_qp", type=bnxt.CMDQ_CREATE_QP_TYPE_RC,
sq_size=16, sq_fwo_sq_sge=1, scq_cid=self.cq_id, rcq_cid=self.cq_id,
sq_pbl=self.sq["base"], sq_pg_size_sq_lvl=self.sq["level"]).xid
self.qp_op(1, BNXT_INIT_MASK, access=BNXT_ACCESS, pkey=0xffff)
def qp_op(self, state, mask, network_type=0, **fields):
self.dev.rcfw("modify_qp", qp_cid=self.qpn, modify_mask=mask,
network_type_en_sqd_async_notify_new_state=state | network_type, **fields)
def connect(self, qpn:int, gid:bytes, mac:int):
network_type = bnxt.CMDQ_MODIFY_QP_NETWORK_TYPE_ROCEV2_IPV4
dgid = (ctypes.c_uint32 * 4)(*(int.from_bytes(gid[i:i + 4], 'little') for i in (0, 4, 8, 12)))
dmac = (ctypes.c_uint16 * 3)(*(int.from_bytes(mac.to_bytes(6, 'big')[i:i + 2], 'little') for i in (0, 2, 4)))
self.qp_op(2, BNXT_RTR_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
pkey=0xffff, dgid=dgid, sgid_index=self.dev.gid_id, hop_limit=64, dest_mac=dmac,
path_mtu_pingpong_push_enable=bnxt.CMDQ_MODIFY_QP_PATH_MTU_MTU_1024, max_dest_rd_atomic=4,
dest_qp_id=qpn)
self.qp_op(3, BNXT_RTS_MASK, network_type=network_type, qp_type=bnxt.CMDQ_MODIFY_QP_QP_TYPE_RC, access=BNXT_ACCESS,
max_rd_atomic=1)
if BNXT_DEBUG >= 1: print(f"bnxt: QP {self.qpn:#x} connected (remote={qpn:#x})")
def _poll(self, timeout):
def poll():
base = bnxt.struct_cq_base.from_buffer_copy(bytes(_qread(self.cqq, self.cqq["cons"])))
return bool(base.cqe_type_toggle & bnxt.CQ_BASE_TOGGLE) == (not bool((self.cqq["cons"] // 16) & 1))
wait_cond(poll, timeout_ms=timeout, msg="BNXT CQ")
raw = bytes(_qread(self.cqq, self.cqq["cons"]))
self.cqq["cons"] += 1
self.dev.doorbell(self.cq_id, bnxt.DBC_DBC_TYPE_CQ, self.cqq["cons"] & 15, (self.cqq["cons"] // 16) & 1)
return raw
def rdma_write(self, rva, rkey, lva, lkey, size, timeout_ms=20000):
start = self.sq["prod"] & 15
hdr = bytes(bnxt.struct_sq_rdma_hdr(wqe_type=bnxt.SQ_RDMA_HDR_WQE_TYPE_WRITE_WQE,
flags=bnxt.SQ_SEND_FLAGS_SIGNAL_COMP, wqe_size=3, length=size, remote_va=rva, remote_key=rkey))
for i, data in enumerate((hdr[:16], hdr[16:32], bytes(bnxt.struct_sq_sge(va_or_pa=lva, l_key=lkey, size=size)))):
_qwrite(self.sq, start + i, data)
nxt = (self.sq_psn + max(1, ceildiv(size, 1024))) & 0xffffff
value = start << bnxt.SQ_MSN_SEARCH_START_IDX_SFT | nxt << bnxt.SQ_MSN_SEARCH_NEXT_PSN_SFT | self.sq_psn
_qwrite(self.sq, self.msn, struct.pack("<Q", value), aux=True)
self.msn, self.sq_psn, self.sq["prod"] = (self.msn + 1) % 128, nxt, self.sq["prod"] + 3
self.dev.doorbell(self.qpn, bnxt.DBC_DBC_TYPE_SQ, self.sq["prod"] & 15, (self.sq["prod"] // 16) & 1)
cqe = bnxt.struct_cq_req.from_buffer_copy(self._poll(timeout_ms))
assert cqe.status == 0
tinygrad_repo/extra/bnxt_driver/connect.py
+118
View File
@@ -0,0 +1,118 @@
#!/usr/bin/env python3
"""Send and validate one RDMA WRITE between two Broadcom BNXT hosts.
This follows ``extra/mlx_driver/connect.py``: sync the driver, start the remote
endpoint over SSH, exchange QP/GID/MAC/MR metadata, move both RC QPs to RTS,
write bytes into the remote MR, and verify the bytes on the remote host.
Both PCI functions must be unbound from bnxt_en/bnxt_re first.
"""
import json
import os
import subprocess
import sys
from typing import Any, IO
TINYGRAD = os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
sys.path.insert(0, TINYGRAD)
from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
from tinygrad.runtime.support.system import PCIDevice
REMOTE_HOST = os.getenv("REMOTE_HOST", "192.168.52.213")
REMOTE_USER = os.getenv("REMOTE_USER", "nimlgen")
LOCAL_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
REMOTE_PCI = os.getenv("REMOTE_PCI", "0000:41:00.0")
LOCAL_IP = os.getenv("LOCAL_IP", "10.0.200.5")
REMOTE_IP = os.getenv("REMOTE_IP", "10.0.200.6")
MESSAGE = os.getenv("RDMA_MESSAGE", "Test message, rdma works!").encode()
REMOTE = f"{REMOTE_USER}@{REMOTE_HOST}"
SSH = ["ssh", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-o", "StrictHostKeyChecking=accept-new", REMOTE]
SYNC_FILES = ("tinygrad/runtime/autogen/bnxt.py", "tinygrad/runtime/support/system.py",
"extra/bnxt_driver/bnxtdev.py", "extra/bnxt_driver/connect.py")
def read_json(stream:IO[str], what:str) -> dict[str, Any]:
for line in iter(stream.readline, ""):
print(f" [remote] {line}", end="")
try: value = json.loads(line)
except json.JSONDecodeError: continue
if isinstance(value, dict): return value
raise RuntimeError(f"remote exited before publishing {what}")
def wait_line(stream:IO[str], text:str) -> str:
for line in iter(stream.readline, ""):
print(f" [remote] {line}", end="")
if text in line: return line
raise RuntimeError(f"remote exited before reporting {text!r}")
def send_line(stream:IO[str], value:str|dict[str, Any]):
stream.write((json.dumps(value) if isinstance(value, dict) else value) + "\n")
stream.flush()
def qp_info(dev:BNXTDev, qp:BNXTQP) -> dict[str, Any]:
return {"qpn":qp.qpn, "mac":dev.mac.to_bytes(6, "big").hex(), "gid":dev.local_gid.hex()}
def server():
dev = BNXTDev(PCIDevice("bnxt", os.getenv("BNXT_PCI", "0000:41:00.0")), ip=os.getenv("BNXT_IP", REMOTE_IP))
qp = BNXTQP(dev)
print(json.dumps(qp_info(dev, qp)), flush=True)
peer = json.loads(sys.stdin.readline())
qp.connect(peer["qpn"], bytes.fromhex(peer["gid"]), int(peer["mac"], 16))
print("connected", flush=True)
target, target_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
target[:0x1000] = bytes(0x1000)
rkey = dev.register_mem(target_paddrs, 0x1000)
print(json.dumps({"target_addr":target_paddrs[0], "rkey":rkey}), flush=True)
assert sys.stdin.readline().strip() == "done"
received = bytes(target).rstrip(b"\0")
print(f"AS TEXT: {received.decode(errors='replace')!r}", flush=True)
print(json.dumps({"data":received.hex()}), flush=True)
def sync_remote():
if os.getenv("SYNC", "1") == "0": return
print("syncing BNXT driver to remote")
subprocess.run(["rsync", "-azR", *SYNC_FILES, f"{REMOTE}:~/tinygrad/"], cwd=TINYGRAD, check=True)
def start_remote() -> subprocess.Popen[str]:
print("booting remote")
command = (f"cd ~/tinygrad && sudo env PYTHONPATH=. PYTHONUNBUFFERED=1 BNXT_DEBUG={os.getenv('BNXT_DEBUG', '0')} "
f"BNXT_PCI={REMOTE_PCI} BNXT_IP={REMOTE_IP} python3 extra/bnxt_driver/connect.py --server")
return subprocess.Popen(SSH + [command], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=sys.stderr, text=True)
def client():
assert 0 < len(MESSAGE) <= 0x1000
sync_remote()
remote = start_remote()
assert remote.stdin is not None and remote.stdout is not None
remote_info = read_json(remote.stdout, "QP information")
print("booting local")
dev = BNXTDev(PCIDevice("bnxt", LOCAL_PCI), ip=LOCAL_IP)
qp = BNXTQP(dev)
send_line(remote.stdin, qp_info(dev, qp))
wait_line(remote.stdout, "connected")
qp.connect(remote_info["qpn"], bytes.fromhex(remote_info["gid"]), int(remote_info["mac"], 16))
print("both QPs in RTS")
remote_target = read_json(remote.stdout, "MR information")
source, source_paddrs = dev.pci_dev.alloc_sysmem(0x1000)
source[:len(MESSAGE)] = MESSAGE
lkey = dev.register_mem(source_paddrs, 0x1000)
print(f"RDMA WRITE {len(MESSAGE)}B to remote phys 0x{remote_target['target_addr']:x}")
qp.rdma_write(remote_target["target_addr"], remote_target["rkey"], source_paddrs[0], lkey, len(MESSAGE))
send_line(remote.stdin, "done")
wait_line(remote.stdout, "AS TEXT")
result = read_json(remote.stdout, "RDMA result")
assert bytes.fromhex(result["data"]) == MESSAGE
print("RDMA WRITE data verified")
remote.stdin.close()
assert remote.wait() == 0
print("RDMA WRITE test complete")
if __name__ == "__main__":
server() if "--server" in sys.argv else client()
tinygrad_repo/extra/bnxt_driver/loopback.py
+46
View File
@@ -0,0 +1,46 @@
#!/usr/bin/env python3
"""Local BNXT RoCEv2 RDMA WRITE loopback using the firmware's PHY loopback mode.
The kernel bnxt_en/bnxt_re modules must be unloaded first.
sudo PYTHONPATH=. BNXT_PCI=0000:41:00.0 BNXT_IP=10.0.200.5 python3 extra/bnxt_driver/loopback.py
"""
import os
import sys
import time
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
from extra.bnxt_driver.bnxtdev import BNXTDev, BNXTQP
from tinygrad.runtime.autogen import bnxt
from tinygrad.runtime.support.system import PCIDevice
BUF_SIZE = 0x1000
BNXT_PCI = os.getenv("BNXT_PCI", "0000:41:00.0")
BNXT_IP = os.getenv("BNXT_IP", "10.0.200.5")
if __name__ == "__main__":
print(f"[init] BNXT at {BNXT_PCI}")
dev = BNXTDev(PCIDevice("bnxt", BNXT_PCI), ip=BNXT_IP)
tx_qp, rx_qp = BNXTQP(dev), BNXTQP(dev)
print(f"[init] loopback-connect TX QP 0x{tx_qp.qpn:x} <-> RX QP 0x{rx_qp.qpn:x}")
tx_qp.connect(rx_qp.qpn, dev.local_gid, dev.mac)
rx_qp.connect(tx_qp.qpn, dev.local_gid, dev.mac)
src, src_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
dst, dst_paddrs = dev.pci_dev.alloc_sysmem(BUF_SIZE)
message = b"Hello from BNXT RoCE PHY loopback!"
src[:BUF_SIZE], dst[:BUF_SIZE] = bytes(BUF_SIZE), bytes(BUF_SIZE)
src[:len(message)] = message
lkey = dev.register_mem(src_paddrs, BUF_SIZE)
rkey = dev.register_mem(dst_paddrs, BUF_SIZE)
print("[loopback] enabling local PHY loopback")
dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_LOCAL)
time.sleep(1)
tx_qp.rdma_write(dst_paddrs[0], rkey, src_paddrs[0], lkey, len(message))
got = bytes(dst[:len(message)])
print(f"[result] {got!r}")
assert got == message
print("BNXT RoCE PHY loopback RDMA WRITE passed")
dev.hwrm("port_phy_cfg", port_id=dev.port_id, enables=bnxt.PORT_PHY_CFG_REQ_ENABLES_LPBK, lpbk=bnxt.PORT_PHY_CFG_REQ_LPBK_NONE)
tinygrad_repo/extra/export_model.py
+2 -3
View File
@@ -35,7 +35,7 @@ def compile_net(linear:UOp, output_bufs:List[Buffer]) -> Tuple[Dict[str,str], Li
return name
for call in iter_kernel_calls(linear):
arg_uops = [b for b in call.src[1:] if b.op is not Ops.BIND]
arg_uops = [b for b in call.src[1:] if not b.is_bound_var]
prg = to_program(call.src[0], Device[arg_uops[0].device].renderer)
info = prg.arg
functions[info.function_name] = prg.src[2].arg
@@ -241,8 +241,7 @@ export default {model_name};
def export_model(model, target:str, *inputs, model_name: Optional[str] = "model", stream_weights=False):
assert Device.DEFAULT in EXPORT_SUPPORTED_DEVICE, f"only {', '.join(EXPORT_SUPPORTED_DEVICE)} are supported"
# NOTE: NUM_CPU_THREADS=1, since export does not support threading
with Context(JIT=2, NUM_CPU_THREADS=1): linear, output_bufs = jit_model(model, *inputs)
with Context(JIT=2): linear, output_bufs = jit_model(model, *inputs)
functions, statements, bufs, bufs_to_save = compile_net(linear, output_bufs)
state = get_state_dict(model)
weight_names = {(id(b), b.offset, b.size, b.dtype): name for name, x in state.items() if (b:=x.uop.base.realized) is not None}
tinygrad_repo/extra/gemm/amd_asm_matmul.py
+1 -1
View File
@@ -462,7 +462,7 @@ def test_matmul():
lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs, arg=KernelInfo(name=colored("kernel", "cyan"),
estimates=Estimates(ops=N*N*N*2, mem=N*N*4*3)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
linear = c.schedule_linear()
tinygrad_repo/extra/gemm/cdna_asm_gemm.py
+4 -3
View File
@@ -122,9 +122,10 @@ def custom_mxfp4_gemm(C:UOp, A:UOp, B:UOp, scale_a:UOp, scale_b:UOp, *extra:UOp,
groups_x, groups_y = UOp.special(ceildiv(N, tile_n), "gidx0"), UOp.special(ceildiv(M, tile_m), "gidx1")
lds = UOp.placeholder((163840,), dtypes.uint8, 0, AddrSpace.LOCAL)
sink = UOp.sink(C.base, A.base, B.base, scale_a.base, scale_b.base, *(x.base for x in extra), lds, threads, groups_x, groups_y,
arg=KernelInfo(f"custom_mxfp4_gemm_{M}_{N}_{K}", estimates=Estimates(ops=2*M*N*K)))
arg=KernelInfo(f"mxfp4_gemm_{M}_{N}_{K}",
estimates=Estimates(ops=2*M*N*K, mem=(M*half_k+N*half_k)*A.dtype.itemsize+M*N*C.dtype.itemsize)))
insts = build_kernel(M, N, K, tile_m, tile_n)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=x) for x in insts))))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts))))
def _mxfp4_gemm_quantized(a_q:Tensor, b_q:Tensor, scale_a:Tensor, scale_b:Tensor) -> Tensor:
M, half_k = a_q.shape
@@ -214,7 +215,7 @@ def custom_uop_gemm(C:UOp, A:UOp, B:UOp) -> UOp:
k = UOp.range(K, 0, AxisType.REDUCE)
mul = (A.flatten().index((m*UOp.const(K)+k))*
B.flatten().index((k*UOp.const(N)+n))).cast(dtypes.float32)
red = mul.reduce(k, arg=Ops.ADD, dtype=dtypes.float32).cast(C.dtype)
red = mul.reduce(k, arg=Ops.ADD).cast(C.dtype)
store = C.flatten().index((m*UOp.const(N)+n)).store(red).end(m, n)
return store.sink(arg=KernelInfo(name=f'uop_gemm_{M}_{N}_{K}'))
tinygrad_repo/extra/gemm/gemm_mxfp4.py
+68 -794
View File
File diff suppressed because it is too large Load Diff
tinygrad_repo/extra/gemm/max_matmul.py
-1
View File
@@ -113,7 +113,6 @@ if __name__ == "__main__":
}
elif GEMM_VARIATION == "hcopt" and M == N == K == 4096 and DTYPE_IN == dtypes.half and DTYPE_OUT == dtypes.half and DTYPE_ACC == dtypes.float:
print("Using CUDA and generated hcopt")
# [Opt(op=OptOps.TC, axis=0, amt=0), Opt(op=OptOps.UPCAST, axis=0, amt=4), Opt(op=OptOps.UPCAST, axis=1, amt=4), Opt(op=OptOps.LOCAL, axis=1, amt=4)]
prog = CUDAProgram(device, "wmma_example", compiler.compile(open(os.path.join(script_dir, 'max_kernels/nv.fp16_fp32_fp16.hcopt.cu')).read()))
args = (c, a, b)
kwargs = {
tinygrad_repo/extra/gemm/moe_gemm.py
+25 -2
View File
@@ -1,10 +1,32 @@
import functools, pathlib
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.uop.ops import UOp, Ops, KernelInfo, AxisType
from tinygrad.helpers import getenv
from tinygrad.renderer import Estimates
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
from extra.gemm.cdna_asm_gemm import quantize_mxfp8, _mx_block_scale, _mx_block_scale_3d
ZERO_OPTIM = getenv("ZERO_OPTIM", 0)
def reduce_scatter_devaxis(out:Tensor, shard_axis:int=0) -> Tensor:
# out: sharded on the device axis, shape (ndev, *rest); return the device-axis sum left sharded on shard_axis.
u = out.uop
devs, rest = u.device, u.shape[1:]
assert rest[shard_axis] % len(devs) == 0, f"reduce_scatter needs even shards: {rest[shard_axis]} % {len(devs)}"
# reach the raw per-device buffer below the UNSHARD, keeping the AFTERs so reads stay ordered after the kernel writes
node, barriers = u, []
while node.op is not Ops.UNSHARD:
if node.op is Ops.AFTER: barriers += node.src[1:]
node = node.src[0]
mbuf = node.src[0].after(*barriers) if barriers else node.src[0]
sz = rest[shard_axis] // len(devs)
shards = []
for i in range(len(devs)):
bounds = tuple((0,s) if a != shard_axis else (i*sz,(i+1)*sz) for a,s in enumerate(rest))
contribs = [mbuf.mselect(j).reshape(rest).shrink(bounds).copy_to_device(devs[i]) for j in range(len(devs))]
shards.append(functools.reduce(lambda a,b: a.alu(Ops.ADD, b), contribs))
return Tensor(UOp.mstack(*shards).unshard(shard_axis, UOp.range(len(devs), -1, AxisType.DEVICE)), device=devs)
@functools.cache
def custom_hk_grouped_mxfp8_gemm(C:UOp, A:UOp, B:UOp, scale_A:UOp, scale_B:UOp, *extra:UOp, dname:str, n_experts:int) -> UOp:
M, K = A.shape
@@ -58,7 +80,8 @@ def grouped_mx_wgrad(g:Tensor, xg:Tensor, expert_off:Tensor, n_experts:int) -> T
out = Tensor(inv.uop.unshard(0), device=g.device) if is_multi else inv
out = Tensor.custom_kernel(out, gT, xT, g_si, x_si, expert_off,
fxn=functools.partial(custom_hk_grouped_mxfp8_wgrad, dname=dname, n_experts=n_experts))[0]
out = out.sum(0) if is_multi else out.squeeze(0)
if is_multi and ZERO_OPTIM: out = reduce_scatter_devaxis(out, 0)
else: out = out.sum(0) if is_multi else out.squeeze(0)
return out.reshape(n_experts, N, K)
def mx_pack_3d(e8:Tensor) -> Tensor:
tinygrad_repo/extra/gemm/moe_routing.py
+1 -1
View File
@@ -53,7 +53,7 @@ def _ggather_bwd(gradient:UOp, kernel:UOp) -> tuple:
g, m, j, jo, ji = _kv_ranges(Gk, M, Dk, _blk_for(Dk))
row = idx.index(g, m).cast(dtypes.weakint)
val = gout.index(g, m, j).load().cast(dtypes.float32)
atomic = UOp(Ops.CUSTOM, dtypes.void, (gtab.index(g, row, j), val), arg=atomic_str)
atomic = UOp(Ops.CUSTOM, src=(gtab.index(g, row, j), val), arg=(atomic_str, dtypes.void))
return atomic.end(g, m, jo, ji).sink(arg=KernelInfo(name=f"ggather_bwd_{M}_{Dk}", opts_to_apply=()))
grad_table = Tensor.custom_kernel(gt, go, Tensor(idx_u, device=dev), fxn=_bwd_kernel)[0]
return (None, grad_table.cast(table_u.dtype).uop, None)
tinygrad_repo/extra/gemm/rdna4_asm_matmul.py
+1 -1
View File
@@ -223,7 +223,7 @@ def test_matmul():
lds = UOp.placeholder((lds_size,), dtypes.uint8, 0, AddrSpace.LOCAL)
sink = UOp.sink(A.base, B.base, C.base, lds, *gidxs, *lidxs,
arg=KernelInfo(name=colored("kernel","cyan"), estimates=Estimates(ops=N*N*N*2, mem=N*N*2*3)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
c = Tensor.custom_kernel(a, b, c, fxn=asm_kernel)[2]
linear = c.schedule_linear()
tinygrad_repo/extra/gemm/tinygrad_nv_matmul.py
+8 -7
View File
@@ -1,6 +1,7 @@
from tinygrad import Tensor, dtypes, Context
from tinygrad.helpers import getenv
from tinygrad.codegen.opt import Opt, OptOps
from tinygrad.uop.ops import AxisType
from tinygrad.engine.realize import run_linear
from dataclasses import replace
@@ -13,17 +14,17 @@ if __name__ == "__main__":
C = A.matmul(B)
if getenv("GEMV"):
opts = [
Opt(op=OptOps.UNROLL, axis=0, amt=8),
Opt(op=OptOps.GROUP, axis=0, amt=32),
Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UNROLL)),
Opt(op=OptOps.SPLIT, axis=1, arg=(32, AxisType.GROUP_REDUCE)),
]
else:
opts = [
Opt(op=OptOps.TC, axis=0, amt=0),
Opt(op=OptOps.UPCAST, axis=0, amt=4),
Opt(op=OptOps.UPCAST, axis=1, amt=8),
Opt(op=OptOps.LOCAL, axis=0, amt=2),
Opt(op=OptOps.LOCAL, axis=1, amt=2),
Opt(op=OptOps.LOCAL, axis=0, amt=2),
Opt(op=OptOps.SPLIT, axis=0, arg=(4, AxisType.UPCAST)),
Opt(op=OptOps.SPLIT, axis=1, arg=(8, AxisType.UPCAST)),
Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
Opt(op=OptOps.SPLIT, axis=1, arg=(2, AxisType.LOCAL)),
Opt(op=OptOps.SPLIT, axis=0, arg=(2, AxisType.LOCAL)),
]
linear = C.schedule_linear()
call = linear.src[-1]
tinygrad_repo/extra/gemm/triton_nv_matmul.py
+1 -1
View File
@@ -79,7 +79,7 @@ if __name__ == "__main__":
linear, var_vals = C.linear_with_vars()
last_call = linear.src[-1]
ast = last_call.src[0]
bufs = [s.buffer for s in last_call.src[1:] if s.op is not Ops.BIND]
bufs = [s.buffer for s in last_call.src[1:] if not s.is_bound_var]
src = compiled.asm["ptx"]
# specify the shared memory here so we don't need to do it dynamically
tinygrad_repo/extra/gptoss_kernels/rmsnorm/__init__.py
+29
View File
@@ -0,0 +1,29 @@
from __future__ import annotations
import functools
from tinygrad import Tensor
from tinygrad.uop.ops import UOp
def rmsnorm_mul_fwd(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
x = x_in.float()
rrms = (x.square().mean(-1, keepdim=True) + eps).rsqrt()
return ((x * rrms) * weight.float()).cast(x_in.dtype), rrms
@functools.cache
def _rmsnorm_mul_fwd_fxn(x_in_p, w_p, eps, device):
return rmsnorm_mul_fwd(Tensor(x_in_p, device=device), Tensor(w_p, device=device), eps)
def _rmsnorm_mul_bwd(grad:UOp, call:UOp) -> tuple:
x = Tensor(call.src[1]).float(); weight = Tensor(call.src[2]).float()
rrms = Tensor(call.gettuple(1))
x_normed = x * rrms # recompute unweighted normed (x is call.src[1])
d_y = Tensor(grad).float()
dxn = d_y * weight # d/d(x_normed)
d_x = rrms * (dxn - x_normed * (dxn * x_normed).mean(-1, keepdim=True))
dw = d_y * x_normed
d_weight = dw.sum(axis=tuple(range(dw.ndim - 1))) # reduce batch/seq -> [dim]
return (d_x.cast(call.src[1].dtype).uop, d_weight.cast(call.src[2].dtype).uop)
def rmsnorm_mul(x_in:Tensor, weight:Tensor, eps:float) -> tuple[Tensor, Tensor]:
fxn = _rmsnorm_mul_fwd_fxn(x_in.as_param(0).uop, weight.as_param(1).uop, eps, x_in.device)
call = UOp.maketuple(fxn[0].uop, fxn[1].uop).call(x_in.uop, weight.uop, grad_fxn=_rmsnorm_mul_bwd)
return Tensor(call.gettuple(0)), Tensor(call.gettuple(1))
tinygrad_repo/extra/hcq/hcq_smi.py
+14 -5
View File
@@ -16,9 +16,12 @@ def _do_reset_device(pci_bus): os.system(f"sudo sh -c 'echo 1 > /sys/bus/pci/dev
def _is_module_loaded(name: str) -> bool: return os.path.isdir(f"/sys/module/{name}")
def cmd_remove_module(args):
modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia", "ast"] if args.backend == "nv" else ["amdgpu"]
modules = ["nvidia_drm", "nvidia_modeset", "nvidia_uvm", "nvidia"] if args.backend == "nv" else ["amdgpu"]
to_unload = [m for m in modules if _is_module_loaded(m)]
if not to_unload: print("Kernel modules are not loaded")
elif getattr(args, "expect", False):
print(f"Kernel modules are loaded: {to_unload}")
sys.exit(1)
else:
print("Removing kernel modules:", ", ".join(to_unload))
try: subprocess.run(["sudo", "modprobe", "-r", *to_unload], check=True)
@@ -60,17 +63,19 @@ def cmd_show_pids(args):
def cmd_kill_pids(args):
devs = scan_devs_based_on_lock(prefix:={"amd":"am", "nv":"nv"}[args.backend], args)
use_sudo = not getattr(args, "sudoless", False)
for dev in devs:
for i in range(128):
if i > 0: time.sleep(0.2)
try:
try: pid = subprocess.check_output(['sudo', 'lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
try: pid = subprocess.check_output((['sudo'] if use_sudo else []) +
['lsof', temp(f'{prefix}_{dev}.lock')]).decode('utf-8').strip().split('\n')[1].split()[1]
except subprocess.CalledProcessError: break
print(f"Killing process {pid} (which uses {dev})")
subprocess.run(['sudo', 'kill', '-9', pid], check=True)
subprocess.run((['sudo'] if use_sudo else []) + ['kill', '-9', pid], check=True)
except subprocess.CalledProcessError as e:
print(f"Failed to kill process for device {dev}: {e}", file=sys.stderr)
@@ -79,6 +84,7 @@ def add_common_commands(parent_subparsers):
p_insmod.set_defaults(func=cmd_insert_module)
p_rmmod = parent_subparsers.add_parser("rmmod", help="Remove a kernel module")
p_rmmod.add_argument("--expect", action="store_true", help="Just assert that module is already unloaded")
p_rmmod.set_defaults(func=cmd_remove_module)
p_reset = parent_subparsers.add_parser("reset", help="Reset a device")
@@ -91,17 +97,20 @@ def add_common_commands(parent_subparsers):
p_reset = parent_subparsers.add_parser("kill_pids", help="Kill pids of processes using the device")
p_reset.add_argument("--pci_bus", default="", help="PCI bus ID of the device")
p_reset.add_argument("--sudoless", action="store_true", help="Do not use sudo when detecting or killing pids")
p_reset.set_defaults(func=cmd_kill_pids)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
backend_subparsers = parser.add_subparsers(dest="backend", required=True, metavar="{nv,amd}", help="Hardware backend to target")
nv_parser = backend_subparsers.add_parser("nv", help="NVIDIA GPUs")
nv_parser = backend_subparsers.add_parser("nv", aliases=["NV"], help="NVIDIA GPUs")
nv_parser.set_defaults(backend="nv")
nv_commands = nv_parser.add_subparsers(dest="command", required=True)
add_common_commands(nv_commands)
amd_parser = backend_subparsers.add_parser("amd", help="AMD GPUs")
amd_parser = backend_subparsers.add_parser("amd", aliases=["AMD"], help="AMD GPUs")
amd_parser.set_defaults(backend="amd")
amd_commands = amd_parser.add_subparsers(dest="command", required=True)
add_common_commands(amd_commands)
tinygrad_repo/extra/hcq2/ops_amd2.py
+174 -161
View File
@@ -1,33 +1,32 @@
from __future__ import annotations
from typing import cast, Any, Callable
import os, ctypes, struct, hashlib, functools, importlib, mmap, errno, array, contextlib, sys, weakref, itertools, collections, atexit
from typing import cast
import os, ctypes, struct, functools, importlib, mmap, errno, contextlib, sys, itertools, atexit
assert sys.platform != 'win32'
from dataclasses import dataclass
from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, HCQ2Buffer, encode_kernargs_clike, make_cmdbuf
from tinygrad.runtime.support.hcq2 import make_binary_patch
from tinygrad.runtime.support.hcq2 import HCQ2Compiled, HCQAllocator, make_buf, hcq_size_var
from tinygrad.uop.ops import sint, UOp
from tinygrad.device import Compiled, BufferSpec, Buffer, Device
from tinygrad.device import BufferSpec, Buffer, Device
from tinygrad.dtype import dtypes
from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, ProfileEvent, lo32, hi32, colored, prod, ContextVar, TracingKey
from tinygrad.helpers import VIZ, ceildiv, unwrap, pluralize, to_tuple
from tinygrad.helpers import getenv, round_up, data64_le, DEBUG, PROFILE, lo32, hi32
from tinygrad.helpers import ceildiv, unwrap, pluralize, to_tuple
from tinygrad.renderer.cstyle import HIPRenderer, HIPCCRenderer
from tinygrad.renderer.llvmir import AMDLLVMRenderer
from tinygrad.runtime.autogen import kfd, hsa, sqtt, amdgpu_kd, amdgpu_drm
from tinygrad.runtime.autogen import kfd, hsa, amdgpu_kd, amdgpu_drm
from tinygrad.runtime.autogen.am import am
from tinygrad.runtime.support.elf import elf_loader
from tinygrad.runtime.support.hcq import FileIOInterface, HCQBuffer, MMIOInterface, hcq_filter_visible_devices
from tinygrad.runtime.support.am.amdev import AMDev, AMMemoryManager
from tinygrad.runtime.support.amd import AMDReg, AMDIP, import_module, import_soc, import_pmc
from tinygrad.runtime.support.system import PCIIfaceBase, PCIAllocationMeta, USBPCIDevice, MAP_FIXED, MAP_NORESERVE
from tinygrad.runtime.support.usb import USB3
from tinygrad.runtime.support.usb import USB3, pm_usb_bufferize
from tinygrad.runtime.support.memory import AddrSpace, BumpAllocator
from tinygrad.runtime.ops_amd import SQTT, SQTT_ITRACE_SE_MASK, SQTT_LIMIT_SE, SQTT_SIMD_SEL, SQTT_TOKEN_EXCLUDE, PMC
from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_EQ, WAIT_REG_MEM_FUNCTION_NEQ, WAIT_REG_MEM_FUNCTION_GEQ
from tinygrad.runtime.ops_amd import SQTT, PMC
from tinygrad.runtime.ops_amd import EVENT_INDEX_PARTIAL_FLUSH, WAIT_REG_MEM_FUNCTION_GEQ
if getenv("IOCTL"): import extra.hip_gpu_driver.hip_ioctl # noqa: F401 # pylint: disable=unused-import
from tinygrad.engine.realize import get_runtime, pm_flatten_linear
from tinygrad.engine.realize import get_call_arg_uops, get_call_var_uops
from tinygrad.uop import FastEnum, auto
from tinygrad.uop.ops import Ops, UPat, PatternMatcher, graph_rewrite
from tinygrad.uop.ops import Ops, UPat, PatternMatcher
# *****************
# PM4
@@ -36,9 +35,11 @@ class PM4Ops(FastEnum):
SET_SH_REG = auto(); SET_UCONFIG_REG = auto(); WAIT_REG_MEM = auto(); ACQUIRE_MEM = auto() # noqa: E702
RELEASE_MEM = auto(); DISPATCH_DIRECT = auto(); EVENT_WRITE = auto() # noqa: E702
def _dw(vals) -> int: return sum(2 if isinstance(x, UOp) and x.dtype.itemsize == 8 else 1 for x in vals)
def pkt3(ctx, op:PM4Ops, *vals):
return UOp(Ops.INS, arg=op, src=tuple(UOp.const(x, dtypes.uint32)
for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), len(vals) - 1), *vals)))
return UOp(Ops.LINEAR, src=tuple(x if isinstance(x, UOp) else UOp.const(x, dtypes.uint32)
for x in (ctx.pm4.PACKET3(getattr(ctx.pm4, f"PACKET3_{op.name}"), _dw(vals) - 1), *vals)))
def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
if bool(args) == bool(kwargs): raise RuntimeError('One (and only one) of *args or **kwargs must be specified')
@@ -52,7 +53,7 @@ def wreg(ctx, reg:AMDReg, *args:sint, **kwargs:int):
def wait_reg_mem(ctx, value, mask=0xffffffff, mem=None, reg=None, reg_done=0, op=WAIT_REG_MEM_FUNCTION_GEQ):
wrm_info_dw = ctx.pm4.WAIT_REG_MEM_MEM_SPACE(int(mem is not None)) | ctx.pm4.WAIT_REG_MEM_OPERATION(int(mem is None and reg_done > 0)) \
| ctx.pm4.WAIT_REG_MEM_FUNCTION(op) | ctx.pm4.WAIT_REG_MEM_ENGINE(0)
return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *(data64_le(mem) if mem is not None else (reg, reg_done)), value, mask, 4)
return pkt3(ctx, PM4Ops.WAIT_REG_MEM, wrm_info_dw, *((mem,) if mem is not None else (reg, reg_done)), value, mask, 4)
def acquire_mem(ctx, addr=0x0, sz=(1 << 64)-1, gli=1, glm=1, glk=1, glv=1, gl1=1, gl2=1):
if ctx.target[0] != 9:
@@ -83,16 +84,18 @@ def release_mem(ctx, address=0x0, value=0, data_sel=0, int_sel=2, ctxid=0, cache
event_dw = ctx.pm4.EVENT_TYPE(ctx.pm4.CACHE_FLUSH_AND_INV_TS_EVENT) | ctx.pm4.EVENT_INDEX(ctx.pm4.event_index__mec_release_mem__end_of_pipe)
memsel_dw = ctx.pm4.DATA_SEL(data_sel) | ctx.pm4.INT_SEL(int_sel)
ctxid = 0
return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, *data64_le(address), *data64_le(value), ctxid)
addr_w = address if isinstance(address, UOp) else UOp.const(address, dtypes.uint64)
val_w = value.cast(dtypes.uint64) if isinstance(value, UOp) else UOp.const(value, dtypes.uint64)
return pkt3(ctx, PM4Ops.RELEASE_MEM, event_dw | cache_flags_dw, memsel_dw, addr_w, val_w, ctxid)
def memory_barrier(ctx):
pf = '' if ctx.nbio.version[0] == 2 else '0' if ctx.nbio.version[:2] != (7, 11) else '1'
return UOp(Ops.LINEAR, dtypes.void, (
return UOp(Ops.LINEAR, src=(
wait_reg_mem(ctx, reg=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_REQ').addr[0],
reg_done=getattr(ctx.nbio, f'regBIF_BX_PF{pf}_GPU_HDP_FLUSH_DONE').addr[0], value=0xffffffff),
acquire_mem(ctx)))
def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val, mem=dst.getaddr(ctx.devs))
def pm4_wait(ctx, dst, val): return wait_reg_mem(ctx, val.cast(dtypes.uint32), mem=dst.getaddr(ctx.devs))
def pm4_barrier(ctx): return memory_barrier(ctx)
@@ -106,155 +109,130 @@ def pm4_timestamp(ctx, dst):
ctx.pm4.int_sel__mec_release_mem__none)
def pm4_program(ctx, call, prg):
data, info = prg.arg
lib_gpu = prg.src[0]
args = encode_kernargs_clike(call, prg, ctx.devs)
prog_addr = lib_gpu.getaddr(ctx.devs) + data.entry_point_offset
scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
args_addr = args.getaddr(ctx.devs)
data, lib = amd_build_program(ctx.dev, prg)
info = prg.arg
user_regs = []
if data.enable_private_segment_sgpr:
scratch_hilo = data64_le(scratch_addr)
user_regs = [scratch_hilo[0], scratch_hilo[1] | 1 << 31, 0xffffffff, 0x20c14000]
if data.enable_dispatch_ptr: user_regs += [*data64_le(args_addr + data.kernargs_segment_size)]
user_regs += [*data64_le(args_addr)]
# kernargs: a nested blob linear inside a getaddr, input addresses and variable values are filled per call through the input table
ka_words = [get_call_arg_uops(call)[gi].getaddr(ctx.devs) for gi in info.globals] + list(get_call_var_uops(call, prg))
pad = data.kernargs_alloc_size - sum(w.dtype.itemsize for w in ka_words)
assert pad >= 0 and pad % 4 == 0, f"bad kernargs padding {pad}"
ka = UOp(Ops.LINEAR, src=tuple(ka_words) + (UOp.const(0, dtypes.uint32),) * (pad // 4)).rtag("kernargs")
prog_addr = lib.getaddr(ctx.devs) + data.entry_point_offset
scratch_addr = UOp.placeholder((data.private_segment_size,), dtypes.uint8, 0, device=ctx.devs).rtag("scratch").getaddr(ctx.devs)
args_addr = ka.getaddr(ctx.devs)
user_regs:list = []
if data.enable_private_segment_sgpr: user_regs = [scratch_addr | (1 << 63), 0xffffffff, 0x20c14000]
if data.enable_dispatch_ptr: user_regs += [args_addr + data.kernargs_segment_size]
user_regs += [args_addr]
dispatch_init = ctx.gc.regCOMPUTE_DISPATCH_INITIATOR.encode(
**({'cs_w32_en': int(data.wave32)} if ctx.target[0] != 9 else {}), force_start_at_000=1, compute_shader_en=1)
ins = [acquire_mem(ctx, gli=0, gl2=0),
wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, *data64_le(prog_addr >> 8)),
wreg(ctx, ctx.gc.regCOMPUTE_PGM_LO, prog_addr >> 8),
wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC1, data.rsrc1, data.rsrc2),
wreg(ctx, ctx.gc.regCOMPUTE_PGM_RSRC3, data.rsrc3),
wreg(ctx, ctx.gc.regCOMPUTE_TMPRING_SIZE, ctx.tmpring_size(data.private_segment_size))]
ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, *data64_le((scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8))
ins += [wreg(ctx, ctx.gc.regCOMPUTE_DISPATCH_SCRATCH_BASE_LO, (scratch_addr + data.private_segment_size // ctx.xccs * xcc_id) >> 8)
for xcc_id in range(ctx.xccs)]
ins += [wreg(ctx, ctx.gc.regCOMPUTE_RESTART_X, 0, 0, 0),
wreg(ctx, ctx.gc.regCOMPUTE_USER_DATA_0, *user_regs),
wreg(ctx, ctx.gc.regCOMPUTE_RESOURCE_LIMITS, ctx.gc.regCOMPUTE_RESOURCE_LIMITS.encode(waves_per_sh=getenv("WAVES_PER_SH"))),
wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *(info.local_size or (1, 1, 1)), 0, 0),
wreg(ctx, ctx.gc.regCOMPUTE_START_X, 0, 0, 0, *info.local_size, 0, 0),
pkt3(ctx, PM4Ops.DISPATCH_DIRECT, *info.global_size, dispatch_init),
pkt3(ctx, PM4Ops.EVENT_WRITE, ctx.pm4.EVENT_TYPE(ctx.soc.CS_PARTIAL_FLUSH) | ctx.pm4.EVENT_INDEX(EVENT_INDEX_PARTIAL_FLUSH))]
return UOp(Ops.LINEAR, dtypes.void, tuple(ins))
return UOp(Ops.LINEAR, src=tuple(ins))
pm_pm4_opsel = PatternMatcher([
def pm4_ib(ctx, submit:UOp, lin:UOp) -> UOp|None:
# the ring only carries a packet pointing at the ib: the host fence at the start of the batch guarantees the ib is free to reuse
if lin.tag is not None or any(w.op in {Ops.CALL, Ops.INS, Ops.LINEAR, Ops.NOOP} for w in lin.src): return None # wait for the flat word linear
assert (size_dw:=sum(w.dtype.itemsize for w in lin.src) // 4) < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
pkt = (UOp.const(ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), dtypes.uint32), lin.rtag("indirect").getaddr(ctx.devs),
UOp.const(size_dw | ctx.pm4.INDIRECT_BUFFER_VALID, dtypes.uint32))
return submit.replace(src=(UOp(Ops.LINEAR, src=pkt, arg=lin.arg).rtag(("cmdbuf", ctx.queue)),))
pm_pm4_encode = PatternMatcher([
(UPat(Ops.CALL, src=(UPat(Ops.PROGRAM, name="prg"),), name="call", allow_any_len=True), pm4_program),
(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="lin"),), name="submit"), pm4_ib),
(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
(UPat(Ops.INS, arg="barrier"), pm4_barrier),
(UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),)), pm4_timestamp),
(UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
(UPat(Ops.INS, arg=("barrier", dtypes.void)), pm4_barrier),
(UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), pm4_wait),
(UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), pm4_timestamp),
(UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), pm4_store),
])
def pm4_submit(ctx, lin):
# ensure compute queues are allocated
for d in (devs:=ctx.devs): q = Device[d].compute_queue
ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COMPUTE:0_{name}")
for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
# the host fence at the start of the batch guarantees the ib is free to reuse
size_dw = sum(len(ins.src) for ins in lin.src)
assert size_dw < (1 << 20), f"indirect buffer of {size_dw} dwords doesn't fit one packet"
ib = UOp.placeholder((size_dw,), dtypes.uint32, next(UOp.unique_num), device=devs, volatile=True).rtag("cmdbuf")
cmdbuf = make_cmdbuf(lin, devs, buf=ib)
# the ring itself only carries a packet pointing at the ib, wrapping the ring
put = put_ptr.index(zero:=UOp.const(0, dtypes.int))
pkt = (ctx.pm4.PACKET3(ctx.pm4.PACKET3_INDIRECT_BUFFER, 2), *data64_le(cmdbuf.getaddr(devs)), size_dw | ctx.pm4.INDIRECT_BUFFER_VALID)
write_pkt = UOp.barrier(*[ring.index(((put + off) % q.ring.size).cast(dtypes.int)).store(UOp.const(x, dtypes.uint32)) for off,x in enumerate(pkt)])
# advance the put/write pointers past the packet
bump_put_ptr = put_ptr.index(zero).store(put + len(pkt))
bump_wptr = wptr.index(zero).store(put + len(pkt))
flush = UOp.barrier(write_pkt, bump_put_ptr, bump_wptr)
return doorbell.after(flush).index(zero).store(put + len(pkt))
pm_pm4_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"), pm4_submit)])
# *****************
# SDMA
class SDMAOps(FastEnum): COPY = auto(); POLL_REGMEM = auto(); FENCE = auto(); TRAP = auto(); TIMESTAMP = auto() # noqa: E702
def sdma_copy(ctx, call):
sz = call.src[2].max_numel() * call.src[2].dtype.itemsize
src_addr, dst_addr = call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs)
return call.ins(SDMAOps.COPY, src=tuple(UOp.const(x, dtypes.uint32) for off in range(0, sz, ctx.max_copy_size) for x in (
ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR),
ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0,
*data64_le(src_addr+UOp.const(off, dtypes.uint64)), *data64_le(dst_addr+UOp.const(off, dtypes.uint64)))))
hdr = ctx.sdma.SDMA_OP_COPY | ctx.sdma.SDMA_PKT_COPY_LINEAR_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_COPY_LINEAR)
return UOp(Ops.LINEAR, src=tuple(x for off in range(0, sz, ctx.max_copy_size) for x in (
*(UOp.const(v, dtypes.uint32) for v in (hdr, ctx.sdma.SDMA_PKT_COPY_LINEAR_COUNT_COUNT(min(sz-off, ctx.max_copy_size)-1), 0)),
*(a + UOp.const(off, dtypes.uint64) if off else a for a in (call.src[2].getaddr(ctx.devs), call.src[1].getaddr(ctx.devs))))))
def sdma_wait(ctx, ins, dst, val):
def sdma_wait(ctx, dst, val):
op = ctx.sdma.SDMA_OP_POLL_REGMEM | ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_FUNC(WAIT_REG_MEM_FUNCTION_GEQ) \
| ctx.sdma.SDMA_PKT_POLL_REGMEM_HEADER_MEM_POLL(1)
return ins.ins(SDMAOps.POLL_REGMEM, src=tuple(UOp.const(x, dtypes.uint32) for x in (
op, *data64_le(dst.getaddr(ctx.devs)), val, 0xffffffff,
ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff))))
def sdma_store(ctx, ins, dst, val):
op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
return UOp(Ops.LINEAR, src=(
ins.ins(SDMAOps.FENCE, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)), val))),
ins.ins(SDMAOps.TRAP, src=tuple(UOp.const(x, dtypes.uint32) for x in (ctx.sdma.SDMA_OP_TRAP, 0)))))
UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs), val.cast(dtypes.uint32), UOp.const(0xffffffff, dtypes.uint32),
UOp.const(ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_INTERVAL(0x04) | ctx.sdma.SDMA_PKT_POLL_REGMEM_DW5_RETRY_COUNT(0xfff), dtypes.uint32)))
def sdma_timestamp(ctx, ins, dst):
def sdma_store(ctx, dst, val): # a fence packet then a trap
op = ctx.sdma.SDMA_OP_FENCE | (ctx.sdma.SDMA_PKT_FENCE_HEADER_MTYPE(3) if ctx.target[0] != 9 else 0)
return UOp(Ops.LINEAR, src=(UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs), val.cast(dtypes.uint32),
UOp.const(ctx.sdma.SDMA_OP_TRAP, dtypes.uint32), UOp.const(0, dtypes.uint32)))
def sdma_timestamp(ctx, dst):
op = ctx.sdma.SDMA_OP_TIMESTAMP | ctx.sdma.SDMA_PKT_TIMESTAMP_GET_HEADER_SUB_OP(ctx.sdma.SDMA_SUBOP_TIMESTAMP_GET_GLOBAL)
return ins.ins(SDMAOps.TIMESTAMP, src=tuple(UOp.const(x, dtypes.uint32) for x in (op, *data64_le(dst.getaddr(ctx.devs)))))
return UOp(Ops.LINEAR, src=(UOp.const(op, dtypes.uint32), dst.getaddr(ctx.devs)))
pm_sdma_opsel = PatternMatcher([
pm_sdma_encode = PatternMatcher([
(UPat(Ops.CALL, src=(UPat(Ops.COPY),), name="call", allow_any_len=True), sdma_copy),
(UPat(Ops.INS, arg="barrier"), lambda: UOp(Ops.NOOP, dtypes.void, ())),
(UPat(Ops.INS, arg="wait", src=(UPat(name="dst"), UPat(name="val")), name="ins"), sdma_wait),
(UPat(Ops.INS, arg="timestamp", src=(UPat(name="dst"),), name="ins"), sdma_timestamp),
(UPat(Ops.INS, arg="store", src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val")), name="ins"), sdma_store),
(UPat(Ops.INS, arg=("barrier", dtypes.void)), lambda: UOp(Ops.LINEAR)),
(UPat(Ops.INS, arg=("wait", dtypes.void), src=(UPat(name="dst"), UPat(name="val"))), sdma_wait),
(UPat(Ops.INS, arg=("timestamp", dtypes.void), src=(UPat(name="dst"),)), sdma_timestamp),
(UPat(Ops.INS, arg=("store", dtypes.void), src=(UPat((Ops.BUFFER, Ops.PARAM), name="dst"), UPat(name="val"))), sdma_store),
])
def sdma_submit(cmdbuf, devs):
# the cmdbuf to submit + the patch writes that fill it
size_dw, zero = cmdbuf.nbytes() // dtypes.uint32.itemsize, UOp.const(0, dtypes.int)
# *****************
# queue submit
# the sdma queue's ring and its host-side ring/write/put pointers
for d in devs: q = Device[d].sdma_queue(0)
ring, wptr, doorbell, put_ptr = (UOp.placeholder((b.size,), b.dtype, 0, device=devs).rtag(f"COPY:0_{name}")
for name, b in (("ring", q.ring), ("write_ptr", q.write_ptr), ("doorbell", q.doorbell), ("put_value", q.put_value)))
def _queue_bufs(ctx, q:AMDQueueDesc) -> tuple[UOp, UOp, UOp, UOp]:
ring = UOp.placeholder((q.ring.size,), q.ring.dtype, 0, device=ctx.devs, volatile=True).rtag(f"{ctx.queue}_ring")
return (ring, *(make_buf(ctx.devs, tag=f"{ctx.queue}_{n}") for n in ("write_ptr", "doorbell", "put_value")))
# sdma needs the cmdbuf contiguous: if it won't fit before the ring end, restart at 0 and zero the tail
put_b = put_ptr.index(zero)
tail_off_dw = ((put_b % (q.ring.size * 4)) // 4).cast(dtypes.int)
fits = (size_dw <= q.ring.size - tail_off_dw).cast(dtypes.int)
start_dw = fits * tail_off_dw
zero_amt_dw = (1 - fits) * (q.ring.size - tail_off_dw)
def pm4_submit(ctx, cmdbuf:UOp) -> UOp:
for d in ctx.devs: q = Device[d].compute_queue
ring, wptr, doorbell, put = _queue_bufs(ctx, q)
p, size_dw = put.after(cmdbuf).index(0).load(), hcq_size_var(cmdbuf) // 4
i = UOp.range(size_dw, 10, dtype=dtypes.int, src=(cmdbuf, ring))
copy = ring.index(((p + i.cast(p.dtype)) % q.ring.size).cast(dtypes.int)).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
next_put = p + size_dw.cast(p.dtype)
flush = UOp.barrier(copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
return doorbell.after(flush).index(0).store(next_put)
# zero the wrapped tail, then copy the cmdbuf into the ring
zi = UOp.range(zero_amt_dw, 0, dtype=dtypes.int, src=(cmdbuf,))
zero_tail = ring.index(tail_off_dw + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
i = UOp.range(UOp.const(size_dw, dtypes.int), 0, dtype=dtypes.int, src=(cmdbuf,))
copy_to_ring = ring.index(start_dw + i).store(cmdbuf.index(i).load()).end(i)
def sdma_submit(ctx, cmdbuf:UOp) -> UOp:
# sdma needs the cmdbuf contiguous in the ring: if it won't fit before the ring end, restart at 0 and zero the tail
for d in ctx.devs: q = unwrap(Device[d].sdma_queue(int(ctx.queue.split(":")[1])))
(ring, wptr, doorbell, put), rs = _queue_bufs(ctx, q), q.ring.size
size_dw = hcq_size_var(cmdbuf) // 4
put_b = put.after(cmdbuf).index(0).load()
tail = ((put_b % (rs * 4)) // 4).cast(dtypes.int)
fits = (size_dw <= rs - tail).cast(dtypes.int)
start_dw, zero_amt = fits * tail, (1 - fits) * (rs - tail)
zi = UOp.range(zero_amt, 10, dtype=dtypes.int, src=(ring,))
zero_tail = ring.index(tail + zi).store(UOp.const(0, dtypes.uint32)).end(zi)
i = UOp.range(size_dw, 11, dtype=dtypes.int, src=(cmdbuf, ring))
copy = ring.index(start_dw + i).store(cmdbuf.bitcast(dtypes.uint32).index(i).load()).end(i)
next_put = put_b + ((zero_amt + size_dw) * 4).cast(put_b.dtype)
flush = UOp.barrier(zero_tail, copy, put.index(0).store(next_put), wptr.index(0).store(next_put))
return doorbell.after(flush).index(0).store(next_put)
# advance the put/write pointers past the zeroed tail and the cmdbuf
next_put_b = put_b + ((zero_amt_dw + size_dw) * 4).cast(put_b.dtype)
bump_put_ptr = put_ptr.index(zero).store(next_put_b)
bump_wptr = wptr.index(zero).store(next_put_b)
# ring the doorbell once the writes have landed
flush = UOp.barrier(zero_tail, copy_to_ring, bump_put_ptr, bump_wptr)
return doorbell.after(flush).index(zero).store(next_put_b)
pm_sdma_submit = PatternMatcher([(UPat(Ops.LINEAR, name="lin"),
lambda ctx, lin: sdma_submit(make_cmdbuf(lin, ctx.devs), ctx.devs))])
@dataclass(frozen=True)
class AMDEncodeCtx: # encode-time constants for one queue: devs (every cmdbuf address resolves into these) + gfx version + packet/ip modules
devs: tuple[str, ...]; target: tuple[int, ...]; pm4: Any; sdma: Any; soc: Any # noqa: E702
gc: AMDIP; nbio: AMDIP; xccs: int; max_copy_size: int; tmpring_size: Callable # noqa: E702
def encode_queue(q:UOp) -> UOp|None:
d = Device[(devs:=to_tuple(q.arg[0]))[0]]
ctx = AMDEncodeCtx(devs, d.target, d.pm4, d.sdma, d.soc, d.gc, d.nbio, d.xccs, d.max_copy_size, d.tmpring_size)
opsel, submit = (pm_pm4_opsel, pm_pm4_submit) if q.arg[1].startswith("COMPUTE") else (pm_sdma_opsel, pm_sdma_submit)
return submit.rewrite(graph_rewrite(q, opsel + pm_flatten_linear, walk=True, ctx=ctx, name=f"{q.arg[1]} opsel"), ctx)
pm_pm4_submit = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(name="cmdbuf"),)), pm4_submit)])
pm_sdma_submit = PatternMatcher([(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(name="cmdbuf"),)), sdma_submit)])
@dataclass(frozen=True)
class AMDProgramData:
@@ -262,10 +240,10 @@ class AMDProgramData:
private_segment_size:int; kernargs_segment_size:int; kernargs_alloc_size:int
enable_dispatch_ptr:int; enable_private_segment_sgpr:int
_amd_program_cache:dict[tuple[bytes,str], tuple[AMDProgramData,bytes]] = {}
def amd_build_program(prg:UOp) -> UOp:
dev = Device[to_tuple(prg.device)[0]] # TODO: rm this
if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, dev.device))) is None:
_amd_program_cache:dict[tuple[bytes, tuple[str, ...]], tuple[AMDProgramData, UOp]] = {}
def amd_build_program(dev, prg:UOp) -> tuple[AMDProgramData, UOp]:
# key on the full device tuple: the same lib can be built for different device sets, each needs its own program buffer
if (cached:=_amd_program_cache.get(key:=(lib:=prg.src[3].arg, to_tuple(prg.device)))) is None:
image, sections, relocs = elf_loader(lib)
rodata = next(sh.header.sh_addr for sh in sections if sh.name == ".rodata")
for off, sym, typ, addent in relocs:
@@ -282,20 +260,23 @@ def amd_build_program(prg:UOp) -> UOp:
wave32=bool(desc.kernel_code_properties & 0x400), private_segment_size=desc.private_segment_fixed_size, kernargs_segment_size=desc.kernarg_size,
kernargs_alloc_size=desc.kernarg_size + (ctypes.sizeof(hsa.hsa_kernel_dispatch_packet_t) if edp else 0), enable_dispatch_ptr=edp,
enable_private_segment_sgpr=desc.kernel_code_properties & hsa.AMD_KERNEL_CODE_PROPERTIES_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER)
image = bytes(image).ljust(round_up(len(image), 4), b"\x00") # the program is uploaded as whole dwords
buf = UOp.placeholder((len(image),), dtypes.uint8, next(UOp.unique_num), device=prg.device).rtag("program")
cached = _amd_program_cache[key] = prg.replace(src=(buf.after(make_binary_patch(buf, bytes(image))),), arg=(data, prg.arg))
cached = _amd_program_cache[key] = (data, buf.after(buf.store(UOp(Ops.BINARY, src=(), arg=image).bitcast(buf.dtype))))
return cached
class AMDAllocator(HCQAllocator['AMDDevice']):
def __init__(self, dev:AMDDevice):
super().__init__(dev, supports_copy_from_disk=dev.has_sdma_queue, supports_transfer=dev.has_sdma_queue and not dev.is_usb())
super().__init__(dev, supports_copy_from_disk=dev.has_copy_queue, supports_transfer=dev.has_copy_queue and not dev.is_usb)
def _alloc(self, size:int, options:BufferSpec) -> HCQ2Buffer:
return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_sdma_queue)
def _alloc(self, size:int, options:BufferSpec) -> HCQBuffer:
return self.dev.iface.alloc(size, host=options.host, uncached=options.uncached, cpu_access=options.cpu_access or not self.dev.has_copy_queue)
def _do_free(self, opaque, options:BufferSpec): self.dev.iface.free(opaque)
def _do_map(self, buf:HCQ2Buffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
def _do_map(self, buf:HCQBuffer): return self.dev.iface.map(buf._base if buf._base is not None else buf)
def _do_unmap(self, buf:HCQBuffer): self.dev.iface.unmap(buf)
@dataclass
class AMDQueueDesc:
@@ -388,15 +369,24 @@ class KFDIface:
return hcqbuf
def free(self, mem):
self._unmap(mem)
if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
def unmap(self, mem):
self._unmap(mem)
if getattr(mem, '_owns_kfd_handle', False): kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
def _unmap(self, mem):
gpus = (ctypes.c_int32 * 1)(self.gpu_id)
stm = kfd.AMDKFD_IOC_UNMAP_MEMORY_FROM_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(gpus), n_devices=1)
assert stm.n_success == 1
if mem.owner == self.dev:
if mem.va_addr: FileIOInterface.munmap(mem.va_addr, mem.size)
kfd.AMDKFD_IOC_FREE_MEMORY_OF_GPU(self.kfd, handle=mem.meta.handle)
def map(self, mem):
if mem.owner is not None and mem.owner._is_cpu(): return self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
if mem.owner is not None and mem.owner._is_cpu():
mapped = self.alloc(mem.size, host=True, cpu_addr=mem.va_addr)
mapped._owns_kfd_handle = True
return mapped
c_gpus = (ctypes.c_int32 * 1)(self.gpu_id)
stm = kfd.AMDKFD_IOC_MAP_MEMORY_TO_GPU(self.kfd, handle=mem.meta.handle, device_ids_array_ptr=ctypes.addressof(c_gpus), n_devices=1)
@@ -468,6 +458,7 @@ class PCIIface(PCIIfaceBase):
def require_profile_mode(self): return True
def is_wgp_active(self, xcc, se, sa, wgp) -> bool: return True # TODO: account for WGP disablement on some asics.
def unmap(self, mem): self.free(mem)
def _compute_props(self):
self.ip_versions = self.dev_impl.ip_ver
@@ -512,8 +503,7 @@ class PCIIface(PCIIfaceBase):
cq = d.compute_queue
for b in (cq.put_value, cq.read_ptr, cq.write_ptr): b._buf.view.view(fmt='Q')[0] = 0
d.iface.dev_impl.gfx.setup_ring(*cq.params)
d.signal('timeline')._buf.cpu_view().mv.cast('Q')[0] = \
d.signal('value', 1).as_memoryview(force_zero_copy=True, no_sync=True).cast('Q')[0] - 1
d.signal('timeline')._buf.cpu_view().view(fmt='Q')[0] = d.signal('value', 1, device="CPU")._buf.cpu_view().view(fmt='Q')[0] - 1
def sleep(self, timeout):
if hasattr(self.pci_dev, 'irq_poller') and self.pci_dev.irq_poller is not None and (events_cnt:=len(self.pci_dev.irq_poller.poll(timeout))):
@@ -527,31 +517,50 @@ class PCIIface(PCIIfaceBase):
def device_fini(self): self.dev_impl.fini()
class USBIface(PCIIface):
def __init__(self, dev, dev_id): # pylint: disable=super-init-not-called
if dev_id >= len(visible:=hcq_filter_visible_devices(USB3.list_devices(0xADD1, 0x0001) + USB3.list_devices(0x3801, 0x0001), "AMD")):
raise RuntimeError(f"AMD:{dev_id} does not exist ({pluralize('device', len(visible))} available)")
self.dev, self.pci_dev, self.vram_bar, self.count = dev, USBPCIDevice("AM", *visible[dev_id]), 0, len(visible)
self.dev_impl = AMDev(self.pci_dev)
self._compute_props()
self.sram = self._dma_region(ctrl_addr=0xf000, sys_addr=0x200000, size=0x80000)
self.cq_buf = self._dma_region(ctrl_addr=0xb800, sys_addr=0x822000, size=0x1000) # +12 is the dword that releases an armed read
self.usb_handle = unwrap(ctypes.cast(self.pci_dev.usb.usb.handle, ctypes.c_void_p).value)
def _dma_region(self, ctrl_addr, sys_addr, size):
region = self.dev_impl.mm.map_range(vaddr:=self.dev_impl.mm.alloc_vaddr(size=size), size, [(sys_addr, size)], aspace=AddrSpace.SYS, uncached=True)
return HCQBuffer(vaddr, size, meta=PCIAllocationMeta(region, has_cpu_mapping=False), view=self.pci_dev.dma_view(ctrl_addr, size), owner=self.dev)
def alloc(self, size:int, host=False, uncached=False, cpu_access=False, contiguous=False, force_devmem=False, **kwargs) -> HCQBuffer:
# everything, even host-style signals, lives in vram: gpu writes into the bridge's own memory collide with an armed 0xF2 read stream
return super().alloc(size, host=False, uncached=uncached, cpu_access=cpu_access or host, contiguous=contiguous, force_devmem=True, **kwargs)
def sleep(self, timeout): pass
# we don't own the sram region, so the buffer never frees it
@functools.cached_property
def usb_sram(self) -> Buffer:
return Buffer(self.dev.device, (b:=self.sram).size, dtypes.uint8, options=BufferSpec(external_ptr=b.va_addr, nolru=True)).allocate(opaque=b)
def _mock(iface, name=None): return type(name or f"MOCK{iface.__name__}", (iface,), {})
class AMDDevice(HCQ2Compiled):
pm_lower = PatternMatcher([
# prep program
(UPat(Ops.PROGRAM, src=(UPat(), UPat(), UPat(), UPat(Ops.BINARY)), name="prg"), amd_build_program),
# encoding of cmdbuf
(UPat(Ops.CUSTOM_FUNCTION, arg="submit_cmdbuf", src=(UPat(Ops.LINEAR, name="q"),)), encode_queue),
])
timestamp_divider = 100.0 # AMD GPU clock: ticks/us
max_scratch_psize = 0
pm_encode = {"COMPUTE": pm_pm4_encode, "COPY": pm_sdma_encode}
pm_lower = {"COMPUTE": pm_pm4_submit, "COPY": pm_sdma_submit}
ifaces = [KFDIface, PCIIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface)]
ifaces = [KFDIface, PCIIface, USBIface, _mock(KFDIface, "MOCKIface"), _mock(KFDIface), _mock(PCIIface), _mock(USBIface)]
def device_props(self): return self.iface.props
def is_am(self) -> bool: return isinstance(self.iface, (PCIIface,))
def is_usb(self) -> bool: return False
def __init__(self, device:str=""):
self.device_id = int(device.split(":")[1]) if ":" in device else 0
self.iface = self._select_iface()
self.iface = self._select_iface(device)
self.is_usb = isinstance(self.iface, USBIface)
if self.is_usb: self.rt_nbytes = 4 << 20
self.target:tuple[int, ...] = ((trgt:=self.iface.props['gfx_target_version']) // 10000, (trgt // 100) % 100, trgt % 100)
self.arch = "gfx%d%x%x" % self.target
@@ -576,12 +585,12 @@ class AMDDevice(HCQ2Compiled):
self.is_aql = getenv("AMD_AQL", int(self.xccs > 1))
if self.is_aql:
self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb() else (16 << 20), uncached=True, cpu_access=True)
self.pm4_ibs = self.iface.alloc(0x2000 if self.is_usb else (16 << 20), uncached=True, cpu_access=True)
self.pm4_ib_alloc = BumpAllocator(self.pm4_ibs.size, wrap=True)
self.max_copy_size = 0x40000000 if self.iface.ip_versions[am.SDMA0_HWIP][0] >= 5 else 0x400000
self.sdma_queues:dict = {}
self.has_sdma_queue = True # self.sdma_queue(0) is not None, TODO: think of this
self.has_copy_queue = not getenv("AMD_DISABLE_SDMA")
super().__init__(device, AMDAllocator(self), [HIPRenderer, AMDLLVMRenderer, HIPCCRenderer], None, can_recover=self.is_am(), arch=self.arch)
@@ -589,6 +598,10 @@ class AMDDevice(HCQ2Compiled):
self.max_private_segment_size = 0
self.pm_bufferize = PatternMatcher([(UPat(Ops.PARAM, tag="scratch", name="b"), lambda ctx, b: ctx[0].scratch_buffer(b.max_numel()))]) + self.pm_bufferize
if self.is_usb:
self.pm_bufferize = pm_usb_bufferize + self.pm_bufferize
raise NotImplementedError("usb amd is not migrated to sealed submits yet") # a usb pm_lower can override the whole submit graph
self.pmc_enabled:bool = PROFILE > 0 and PMC > 0
if self.pmc_enabled:
self.iface.require_profile_mode()
@@ -649,7 +662,7 @@ class AMDDevice(HCQ2Compiled):
wg_data_size = round_up((vgpr_size_per_cu + sgrp_size_per_cu + lds_size_per_cu + hwreg_size_per_cu) * self.cu_cnt, mmap.PAGESIZE)
ctl_stack_size = round_up((12 if self.target[0] != 9 else 8) * self.wave_cnt + 8 + 40, mmap.PAGESIZE)
return self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_COMPUTE_AQL if self.is_aql else kfd.KFD_IOC_QUEUE_TYPE_COMPUTE,
0x2000 if self.is_usb() else (16 << 20), eop_buffer_size=0x1000,
0x2000 if self.is_usb else (16 << 20), eop_buffer_size=0x1000,
ctx_save_restore_size=0 if self.is_am() else wg_data_size + ctl_stack_size, ctl_stack_size=ctl_stack_size,
debug_memory_size=round_up(self.wave_cnt * 32, 64))
@@ -657,7 +670,7 @@ class AMDDevice(HCQ2Compiled):
if getenv("AMD_DISABLE_SDMA"): return None
if idx in self.sdma_queues: return self.sdma_queues[idx]
with contextlib.suppress(OSError):
self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x200 if self.is_usb() else (16 << 20), idx=idx)
self.sdma_queues[idx] = self.create_queue(kfd.KFD_IOC_QUEUE_TYPE_SDMA, 0x2000 if self.is_usb else (16 << 20), idx=idx)
return self.sdma_queues.get(idx, None)
def tmpring_size(self, private_segment_size):
tinygrad_repo/extra/llama_kernels/__init__.py
+2 -4
View File
@@ -1,8 +1,6 @@
from __future__ import annotations
import functools, pathlib
from dataclasses import replace
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import shape_to_shape_arg
from tinygrad import Tensor
from tinygrad.runtime.support.compiler_amd import HIPCCCompiler
FP8_MAX = 448.0
@@ -12,7 +10,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
@functools.cache
def _local_abs_max_fxn(x_p, device):
x = Tensor(x_p, device=device)
inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x # the per-shard view of the flat param
return (inner.abs().max(),)
def local_abs_max(x:Tensor) -> Tensor:
tinygrad_repo/extra/llama_kernels/quantize_fp8_delayed/__init__.py
+1 -1
View File
@@ -50,7 +50,7 @@ def _custom_quantize_fp8_with_amax(fp8_out:UOp, amax_out:UOp, x:UOp, amax_state:
else: raise NotImplementedError(f"no atomic max for device {device}")
amax_idx = amax_out.reshape((1,)).index(UOp.const(0))
max_val = lds[0].load()
atomic = UOp(Ops.CUSTOM, dtypes.void, (amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=atomic_arg)
atomic = UOp(Ops.CUSTOM, src=(amax_idx, max_val.bitcast(dtypes.int32), max_val, amax_idx.load()), arg=(atomic_arg, dtypes.void))
return atomic.end(tid, wg).sink(arg=KernelInfo(f"quantize_fp8_with_amax_{n_elems}", opts_to_apply=()))
@functools.cache
tinygrad_repo/extra/llama_kernels/quantize_mxfp4/__init__.py
+1 -1
View File
@@ -12,7 +12,7 @@ def _custom_quantize_mxfp4(row_fp4:UOp, row_scale:UOp, col_fp4:UOp, col_scale:UO
mem = M*N*2 + M*N + M*N//16 # read bf16, write row+col fp4 + e8m0
outputs = (row_fp4, row_scale, col_fp4, col_scale)
sink = UOp.sink(*(o.base for o in outputs), x.base,
*(UOp(Ops.CUSTOM, dtypes.void, (o.base.index(0),), arg="") for o in outputs),
*(UOp(Ops.CUSTOM, src=(o.base.index(0),), arg=("", dtypes.void)) for o in outputs),
UOp.special(256, "lidx0"), UOp.special(M//128, "gidx0"), UOp.special(N//64, "gidx1"),
arg=KernelInfo(name, estimates=Estimates(ops=12*M*N, mem=mem)))
src = (pathlib.Path(__file__).parent/"quantize_mxfp4.cpp").read_text()
tinygrad_repo/extra/mmapeak/mmapeak.py
+2 -2
View File
@@ -3,7 +3,7 @@ import os
# TODO: there is a timing bug without this
os.environ["AMD_AQL"] = "1"
from tinygrad import Tensor, Device, GlobalCounters, Context
from tinygrad import Tensor, Device, GlobalCounters, Context, dtypes
from tinygrad.helpers import getenv, DEV
from tinygrad.uop.ops import UOp, Ops, KernelInfo
from tinygrad.renderer import Estimates
@@ -37,7 +37,7 @@ def launchBenchmark(instruction, vgprIndices, dense=True, accum=False, **kwargs)
gidx = UOp.special(NUM_WORKGROUPS, "gidx0")
FLOPs = FLOPS_PER_MATMUL * NUM_WAVES * NUM_WORKGROUPS * INTERNAL_LOOP * INSTRUCTIONS_PER_LOOP
sink = UOp.sink(A.base, threads, gidx, arg=KernelInfo(inst.op.name.lower(), estimates=Estimates(ops=FLOPs, mem=0)))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=x) for x in insts]))))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple([UOp(Ops.INS, arg=(x, dtypes.void)) for x in insts]))))
dummy = Tensor.zeros(1).contiguous().realize()
out = Tensor.custom_kernel(dummy, fxn=fxn)[0]
linear = out.schedule_linear()
tinygrad_repo/extra/models/llama.py
+3 -3
View File
@@ -5,9 +5,9 @@ from tinygrad.helpers import getenv, DEBUG
# https://github.com/facebookresearch/llama/blob/1076b9c51c77ad06e9d7ba8a4c6df775741732bd/llama/model.py#L47
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0) -> Tensor:
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2)[:(dim // 2)] / dim))
freqs = Tensor.arange(end).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).reshape(1, end, 1, dim//2, 2)
freqs = 1.0 / (theta ** (Tensor.arange(0, dim, 2, dtype=dtypes.float32)[:(dim // 2)] / dim))
freqs = Tensor.arange(end, dtype=dtypes.float32).unsqueeze(dim=1) * freqs.unsqueeze(dim=0)
return Tensor.stack(freqs.cos(), freqs.sin(), dim=-1).cast(dtypes.default_float).reshape(1, end, 1, dim//2, 2)
# matches meta, non hugging face weights
# (a+i*b) * (c+i*d) = (ac-bd) + i*(ad+bc)
tinygrad_repo/extra/onnx_helpers.py
+2
View File
@@ -1,10 +1,12 @@
from tinygrad import Tensor
import os
from tinygrad.tensor import _to_np_dtype
from tinygrad.nn.onnx import OnnxRunner, OnnxValue
import numpy as np
import onnxruntime as ort
ort_options = ort.SessionOptions()
ort_options.log_severity_level = 3
ort_options.intra_op_num_threads = os.cpu_count() or 1
def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
"""
tinygrad_repo/extra/optimization/test_beam_search.py
+4 -3
View File
@@ -89,7 +89,8 @@ class TestBeamSearch(unittest.TestCase):
s.apply_opt(Opt(OptOps.TC, 0, (-1, 0, 1)))
up = prod([x for x, t in zip(s.full_shape, s.axis_types) if t in (AxisType.UPCAST, AxisType.UNROLL)])
actions = get_kernel_actions(s, include_0=False, max_up=int(up))
upcasted = [s for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]
upcasted = [s for s in actions.values() if any(o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)
for o in s.applied_opts)]
assert len(upcasted) > 0, f"expected upcast/unroll actions after TC with max_up={up}, but got none"
def test_max_up(self):
@@ -98,8 +99,8 @@ class TestBeamSearch(unittest.TestCase):
s = Scheduler(ast, Device[Device.DEFAULT].renderer)
for max_up in (2, 4):
actions = get_kernel_actions(s, include_0=False, max_up=max_up)
for up_opts in [s.applied_opts for s in actions.values() if any(opt.op in (OptOps.UPCAST, OptOps.UNROLL) for opt in s.applied_opts)]:
assert len([opt for opt in up_opts if opt.arg > max_up]) == 0 and len([op for op in up_opts if op.arg <= max_up]) > 0
up_opts = [o for s in actions.values() for o in s.applied_opts if o.op is OptOps.SPLIT and o.arg[1] in (AxisType.UPCAST, AxisType.UNROLL)]
assert len([opt for opt in up_opts if opt.arg[0] > max_up]) == 0 and len([op for op in up_opts if op.arg[0] <= max_up]) > 0
if __name__ == '__main__':
unittest.main()
tinygrad_repo/extra/runbook_digitalocean_mi350x.md
+197
View File
@@ -0,0 +1,197 @@
# Runbook: Llama 3 8B Training on DigitalOcean MI350X
## Machine Specs
- 8x MI350X GPUs (gfx950, device ID 75b0), 288GB VRAM each
- 2TB RAM, 192 CPUs, 2TB disk
- ROCm 7.14 at `/opt/rocm` (NOT `/opt/rocm-7.1.1` like the submission scripts assume)
- Python 3.12
## Phase 1: System Setup
### 1.1 Install packages
```bash
apt-get update
apt-get install -y python3-pip python3-venv git tmux rclone clang
```
### 1.2 Install Python deps
```bash
python3 -m pip install --break-system-packages --ignore-installed typing-extensions numpy tqdm wandb tiktoken sentencepiece
```
Note: `--ignore-installed typing-extensions` is needed because the base image ships typing-extensions 4.10.0 without a RECORD file, so pip cannot uninstall it.
### 1.3 Install ROCm dev headers
The base image has ROCm runtime but NOT the HIP dev headers. Need:
```bash
apt-get install -y amdrocm-core-dev
```
This installs `hip/hip_runtime.h` at `/opt/rocm/core-7.14/include/hip/hip_runtime.h`.
The symlink `/opt/rocm/include` → `/opt/rocm/core-7.14/include` makes it available at `/opt/rocm/include/hip/hip_runtime.h`.
### 1.4 Configure ROCm comgr
ROCm 7.14 ships comgr 3.3 at `/opt/rocm/lib/libamd_comgr.so`. tinygrad's DLL loader needs explicit env vars to find it (it searches for `libcomgr.so*` by default, not `libamd_comgr.so*`). Set these in the run command:
```bash
export COMGR_PATH=/opt/rocm/lib/libamd_comgr.so
export COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so
```
Also add ROCm libs to ldconfig so comgr's shared library dependencies resolve:
```bash
cat > /etc/ld.so.conf.d/rocm.conf << 'EOF'
/opt/rocm/lib
/opt/rocm/lib/llvm/lib
/opt/rocm/lib/rocm_sysdeps/lib
EOF
ldconfig
```
### 1.5 Install geohot tmux config
```bash
curl -sL https://raw.githubusercontent.com/geohot/configuration/master/.tmux.conf -o ~/.tmux.conf
```
### 1.6 Verify GPU PCI access
The AM userspace driver accesses the GPUs directly over PCI. Do not load `amdgpu`. `/dev/kfd` is not required.
```bash
rmmod amdgpu
lspci -nnk -d 1002:
```
The MI350X devices should not show a `Kernel driver in use: amdgpu`.
## Phase 2: Clone tinygrad
```bash
cd /root
git clone https://github.com/tinygrad/tinygrad.git
cd tinygrad
python3 -m pip install --break-system-packages -e .
```
## Phase 3: Download C4 Dataset
The C4 data is on the MLCommons Cloudflare R2 bucket in Megatron-LM indexed format.
```bash
rclone config create mlc-training s3 provider=Cloudflare \
access_key_id=76ea42eadb867e854061a1806220ee1e \
secret_access_key=a53625c4d45e3ca8ac0df8a353ea3a41ffc3292aa25259addd8b7dc5a6ce2936 \
endpoint=c2686074cb2caf5cbaf6d134bdba8b47.r2.cloudflarestorage.com
mkdir -p /raid/datasets/c4-8b
(rclone copy mlc-training:mlcommons-training-wg-public/llama3_1/datasets/c4/llama3_1_8b/ /raid/datasets/c4-8b/ -P && \
PYTHONPATH=. python3 examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/buid_dataset_cache.py) \
> /root/dataset_cache.log 2>&1 &
```
Leave this running and proceed to the beam step while the dataset downloads and its cache builds.
### 3.1 Smoke test (beam search, 2 layers, fake data)
Always run beam first to validate the pipeline:
```bash
tmux new-session -d -s beam 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_beam.sh 2>&1 | tee /root/beam.log'
```
The beam test runs 10 training steps with 2 layers. Expected results:
- ~0.29s per step after warmup
- ~700K GFLOPS, ~7% MFU (low because only 2 layers)
- ~380 GB VRAM used
- Loss stable at ~12.55 with random init
Files downloaded (~85GB total, ~6 minutes):
- `c4-train.en_6_text_document.bin` (79 GB)
- `c4-train.en_6_text_document.idx` (870 MB)
- `c4-validation-91205-samples.en_text_document.bin` (159 MB)
- `c4-validation-91205-samples.en_text_document.idx` (1.8 MB)
- `LICENSE.txt`, `NOTICE.txt`
**Wait for rclone to fully complete before starting training.** Starting training while the dataset is still downloading will read a truncated .bin file, causing `ValueError: all input arrays must have the same shape` in the dataloader. The stale `.index_cache` and `.blend_cache` files must also be deleted if this happens:
```bash
rm -f /raid/datasets/c4-8b/*.index_cache /raid/datasets/c4-8b/*.blend_cache
```
## Phase 4: wandb Login
```bash
wandb login
```
Enter API key from https://wandb.ai/authorize
Alternatively, pass the key directly:
```bash
wandb login <API_KEY>
```
## Phase 5: Run Training
Run training in tmux so it survives SSH disconnects:
```bash
tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
```
Attach with `tmux attach -t train`.
### 5.1 Full training run
```bash
tmux new-session -d -s train 'cd /root/tinygrad && COMGR_PATH=/opt/rocm/lib/libamd_comgr.so COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so CC=/opt/rocm/core-7.14/lib/llvm/bin/clang DEV=PCI+AMD:HIP ROCM_PATH=/opt/rocm WANDB=1 bash examples/mlperf/training_submission_v6.0/tinycorp/benchmarks/llama31_8b/implementations/tinybox_8xMI350X/dev_run.sh 2>&1 | tee /root/train.log'
```
## Environment Variable Reference
| Variable | Value | Why |
|---|---|---|
| `COMGR_PATH` | `/opt/rocm/lib/libamd_comgr.so` | tinygrad's DLL loader needs explicit path to find comgr 3.3 |
| `COMGR_3_PATH` | `/opt/rocm/lib/libamd_comgr.so` | comgr 3.x uses a separate `comgr_3` module with its own path var |
| `CC` | `/opt/rocm/core-7.14/lib/llvm/bin/clang` | System clang doesn't know gfx950; must use ROCm's bundled clang |
| `DEV` | `PCI+AMD:HIP` | Force HIPRenderer (comgr-based) over HIPCCRenderer (hipcc subprocess) |
| `ROCM_PATH` | `/opt/rocm` | Script defaults to `/opt/rocm-7.1.1` which doesn't exist |
| `WANDB` | `1` | Enable wandb logging (off by default) |
## Architecture
| Component | Source file |
|---|---|
| Model | `examples/mlperf/models/flat_llama.py` — FlatTransformer, FP8 MXFP4 weights, fused QKV, flash attention |
| Trainer | `examples/mlperf/model_train.py` → `train_llama3()` |
| Optimizer | `examples/mlperf/optim.py` — GradAccClipAdamW, master weights, FP8 re-quant |
| LR schedule | `examples/mlperf/lr_schedulers.py` — CosineAnnealingLRWithWarmup |
| Dataloader | `examples/mlperf/dataloader.py` — Megatron-LM indexed bin format |
| ASM GEMM | `extra/gemm/cdna_asm_gemm.py` — gfx950 MFMA assembly, MXFP4 |
| Flash attention | `extra/thunder/amd/fa.py` |
| Fused kernels | `extra/llama_kernels/` — rmsnorm, silu, quantize, fused_ce |
| GPU driver | `tinygrad/runtime/ops_amd.py` — HCQ, using the AM userspace PCI interface |
| Renderer | `tinygrad/renderer/cstyle.py` — HIPRenderer for gfx950 |
| comgr compiler | `tinygrad/runtime/support/compiler_amd.py` — HIPCompiler using comgr 3.3 |
## Troubleshooting
### `'hip/hip_runtime.h' file not found`
Install `amdrocm-core-dev`:
```bash
apt-get install -y amdrocm-core-dev
```
### `'gfx950' is not a recognized processor` + LLVM crash
System clang doesn't know gfx950. Set `CC=/opt/rocm/core-7.14/lib/llvm/bin/clang`.
### `comgr not available: try setting COMGR_PATH?`
Add ROCm libs to ldconfig and set `COMGR_PATH` and `COMGR_3_PATH`:
```bash
# /etc/ld.so.conf.d/rocm.conf should contain /opt/rocm/lib paths
ldconfig
```
### `comgr not available: try setting COMGR_3_PATH?`
comgr 3.x uses a separate module. Set `COMGR_3_PATH=/opt/rocm/lib/libamd_comgr.so` too.
### `No such file or directory: 'clang'`
Install clang: `apt-get install -y clang` (for CPU compilation).
For gfx950 HIP compilation, comgr (not clang) is used — ensure the ROCm 7.14 comgr 3.3 is properly loaded via `COMGR_PATH` and `COMGR_3_PATH`.
## Appendix: KVM Virtualization Observations
### Virtualization detection
```
$ systemd-detect-virt
kvm
$ lspci -nn | grep AMD
83:00.0 ... Device [1002:75b0]
```
CPU flags include `hypervisor`. `dmesg` shows `Hypervisor detected: KVM`.
### No fan control
No `fan*` or `pwm*` hwmon entries exist. Only `temp*`, `power*`, `freq*` are exposed. GPU temps read 56-63°C, power ~265W per GPU.
tinygrad_repo/extra/sqtt/examples/generate_examples.py
+4 -1
View File
@@ -1,4 +1,4 @@
import os, subprocess, sys, shlex
import os, subprocess, sys, shlex, pickle
from pathlib import Path
from tinygrad.helpers import temp, getenv
@@ -23,5 +23,8 @@ if __name__ == "__main__":
# AM_RESET=1 gets a clear trace, does not work on mi300 machines
subprocess.run([sys.executable, *shlex.split(test)], cwd=EXAMPLES_DIR.parent.parent.parent,
env={**os.environ, "DEV":"AMD", "AM_RESET":"1" if not arch.startswith("gfx9") else "0", "VIZ":"-2", "PYTHONPATH":"."})
with open(PROFILE_PATH, "rb") as f: events = pickle.load(f)
with open(PROFILE_PATH, "wb") as f:
pickle.dump([e for e in events if type(e).__name__ in {"ProfilePMCEvent", "ProfileSQTTEvent", "ProfileProgramEvent"}], f)
PROFILE_PATH.rename(dest:=EXAMPLES_DIR/arch/f"profile_{name}_run_{i}.pkl")
print(f"saved SQTT trace to {dest}")
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_empty_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_empty_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_gemm_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_gemm_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_plus_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_plus_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_sync_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1100/profile_sync_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_empty_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_empty_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_gemm_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_gemm_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_handwritten_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_handwritten_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_plus_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_plus_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_sync_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx1200/profile_sync_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_empty_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_empty_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_gemm_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_gemm_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_plus_run_0.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/examples/gfx950/profile_plus_run_1.pkl
BIN
View File
Binary file not shown.
tinygrad_repo/extra/sqtt/roc.py
+35 -84
View File
@@ -1,13 +1,11 @@
#!/usr/bin/env python3
import ctypes, pathlib, argparse, pickle, dataclasses, threading, itertools
from decimal import Decimal
from typing import Generator
from typing import Any, Generator
from tinygrad.helpers import temp, unwrap, DEBUG
from tinygrad.runtime.ops_amd import ProfileSQTTEvent
from tinygrad.runtime.autogen import rocprof
from tinygrad.renderer.amd.dsl import Inst
from tinygrad.helpers import ProfileEvent, ProfileRangeEvent, ProfilePointEvent
from tinygrad.device import ProfileProgramEvent
from tinygrad.device import ProfileDeviceEvent, ProfileProgramEvent
from test.amd.disasm import disasm
@dataclasses.dataclass(frozen=True)
@@ -39,17 +37,17 @@ class WaveExec(WaveSlot):
insts_array = (struct*(len(self.insts)//sz)).from_buffer(self.insts)
for inst in insts_array:
inst_typ = rocprof.enum_rocprofiler_thread_trace_decoder_inst_category_t.get(inst.category)
yield InstExec(inst_typ, inst.pc.address, inst.stall, inst.duration, inst.time)
yield InstExec(inst_typ or "UNKNOWN", inst.pc.address, inst.stall, inst.duration, inst.time)
@dataclasses.dataclass(frozen=True)
class OccEvent(WaveSlot):
time:int
start:int
RunKey = tuple[str, int]
RunKey = tuple[int, int]
class _ROCParseCtx:
def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]):
def __init__(self, sqtt_evs:list[ProfileSQTTEvent], disasms:dict[int, dict[int, Inst]]):
self.sqtt_evs, self.disasms = iter(sqtt_evs), {k:{k2:(disasm(v2), v2.size()) for k2,v2 in v.items()} for k,v in disasms.items()}
self.inst_execs:dict[RunKey, list[WaveExec]] = {}
self.occ_events:dict[RunKey, list[OccEvent]] = {}
@@ -76,7 +74,7 @@ class _ROCParseCtx:
self.inst_execs.setdefault(unwrap(self.active_run), []).append(WaveExec(ev.wave_id, ev.cu, ev.simd, unwrap(self.active_se), ev.begin_time,
ev.end_time, insts_blob))
def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]]) -> _ROCParseCtx:
def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[int, dict[int, Inst]]) -> _ROCParseCtx:
ROCParseCtx = _ROCParseCtx(sqtt_evs, disasms)
@rocprof.rocprof_trace_decoder_se_data_callback_t
@@ -129,44 +127,7 @@ def decode(sqtt_evs:list[ProfileSQTTEvent], disasms:dict[str, dict[int, Inst]])
raise exc
return ROCParseCtx
def unpack_occ(viz_data, i:int, j:int, key:tuple[str, int], data:list, p:ProfileProgramEvent, target:str) -> dict:
from tinygrad.viz.serve import amd_decode, create_step, row_tuple
steps = viz_data.ctxs[i]["steps"]
if len(steps[j+1:]) > 0: return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
base = unwrap(p.base)
disasm:dict[int, Inst] = {addr+base:inst for addr,inst in amd_decode(unwrap(p.lib), target).items()}
rctx = decode(data, {p.tag:disasm})
cu_events:dict[str, list[ProfileEvent]] = {}
# ** inst traces
wave_insts:dict[str, dict[str, dict]] = {}
inst_units:dict[str, itertools.count] = {}
for w in rctx.inst_execs.get(key, []):
if (u:=w.wave_loc) not in inst_units: inst_units[u] = itertools.count(0)
n = next(inst_units[u])
if (events:=cu_events.get(w.cu_loc)) is None: cu_events[w.cu_loc] = events = []
events.append(ProfileRangeEvent(f"SIMD:{w.simd}", loc:=f"INST WAVE:{w.wave_id} N:{n}", Decimal(w.begin_time), Decimal(w.end_time)))
wave_insts.setdefault(w.cu_loc, {})[f"{u} N:{n}"] = {"wave":w, "disasm":disasm, "prg":p, "run_number":n, "loc":loc}
# ** occ traces (only WAVESTART/WAVEEND)
units:dict[str, itertools.count] = {}
wave_start:dict[str, int] = {}
for occ in rctx.occ_events.get(key, []):
if (u:=occ.wave_loc) not in units: units[u] = itertools.count(0)
if u in inst_units: continue
if occ.start: wave_start[u] = occ.time
else:
if (events:=cu_events.get(occ.cu_loc)) is None: cu_events[occ.cu_loc] = events = []
events.append(ProfileRangeEvent(f"SIMD:{occ.simd}", f"OCC WAVE:{occ.wave_id} N:{next(units[u])}", Decimal(wave_start.pop(u)),Decimal(occ.time)))
# ** split graph by CU
for cu in sorted(cu_events, key=row_tuple):
steps.append(create_step(f"{cu} {len(cu_events[cu])}", ("/cu-sqtt", i, len(steps)), depth=1,
data=[ProfilePointEvent(unit, "start", unit, ts=Decimal(0)) for unit in units]+cu_events[cu]))
for k in sorted(wave_insts.get(cu, []), key=row_tuple):
wd = wave_insts[cu][k]
steps.append(create_step(k.replace(cu, ""), ("/amd-sqtt-insts", i, len(steps)), loc=wd["loc"], depth=2,
data={"fxn":unpack_insts, "args":(wd,)}))
return {"steps":[{k:v for k,v in s.items() if k != "data"} for s in steps[j+1:]]}
def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
def unpack_insts(w:WaveExec, pc_to_inst:dict[int, Inst]) -> dict:
columns = ["PC", "Instruction", "Hits", "Cycles", "Stall", "Type"]
inst_columns = ["N", "Clk", "Idle", "Dur", "Stall"]
# Idle: The total time gap between the completion of previous instruction and the beginning of the current instruction.
@@ -176,36 +137,24 @@ def unpack_insts(viz_data, i:int, j:int, data:dict) -> dict:
# * Instruction cache miss
# Stall: The total number of cycles the hardware pipe couldn't issue an instruction.
# Duration: Total latency in cycles, defined as "Stall time + Issue time" for gfx9 or "Stall time + Execute time" for gfx10+.
prev_instr = (w:=data["wave"]).begin_time
pc_to_inst = data["disasm"]
prev_instr = w.begin_time
start_pc = None
rows:dict[int, dict] = {}
rows:dict[int, dict[str, Any]] = {}
for pc, inst in pc_to_inst.items():
if start_pc is None: start_pc = pc
rows[pc] = {"pc":pc-start_pc, "inst":str(inst), "hit_count":0, "dur":0, "stall":0, "type":"", "hits":{"cols":inst_columns, "rows":[]}}
for e in w.unpack_insts():
if not (inst:=rows[e.pc]).get("type"): inst["type"] = str(e.typ).split("_")[-1]
inst["hit_count"] += 1
inst["dur"] += e.dur
inst["stall"] += e.stall
inst["hits"]["rows"].append((inst["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
if not (row:=rows[e.pc]).get("type"): row["type"] = str(e.typ).split("_")[-1]
row["hit_count"] += 1
row["dur"] += e.dur
row["stall"] += e.stall
row["hits"]["rows"].append((row["hit_count"]-1, e.time, max(0, e.time-prev_instr), e.dur, e.stall))
prev_instr = max(prev_instr, e.time + e.dur)
summary = [{"label":"Total Cycles", "value":w.end_time-w.begin_time}, {"label":"SE", "value":w.se}, {"label":"CU", "value":w.cu},
{"label":"SIMD", "value":w.simd}, {"label":"Wave ID", "value":w.wave_id}, {"label":"Run number", "value":data["run_number"]}]
return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns, "metadata":[summary], "ref":viz_data.ref_map.get(data["prg"].name)}
def print_data(data:dict) -> None:
from tabulate import tabulate
# plaintext
if "src" in data: print(data["src"])
# table format
elif "cols" in data:
print(tabulate([r[:len(data["cols"])] for r in data["rows"]], headers=data["cols"], tablefmt="github"))
return {"rows":[tuple(v.values()) for v in rows.values()], "cols":columns}
def main() -> None:
import tinygrad.viz.serve as viz
from tinygrad.uop.ops import RewriteTrace
data = viz.VizData()
from tabulate import tabulate
from tinygrad.viz.serve import amd_decode
parser = argparse.ArgumentParser()
parser.add_argument('--profile', type=pathlib.Path, metavar="PATH", help='Path to profile (optional file, default: latest profile)',
@@ -216,26 +165,28 @@ def main() -> None:
with args.profile.open("rb") as f: profile = pickle.load(f)
viz.get_profile(profile, data=data)
# List all kernels
if args.kernel is None:
for c in data.ctxs:
print(c["name"])
for s in c["steps"]: print(" "+s["name"])
for p in profile:
if isinstance(p, ProfileProgramEvent) and p.device.startswith("AMD"): print(p.name)
return None
# Find kernel trace
trace = next((c for c in data.ctxs if c["name"] == f"SQTT {args.kernel}"), None)
if not trace: raise RuntimeError(f"no matching trace for {args.kernel}")
n = 0
for s in trace["steps"]:
if "PKTS" in s["name"]: continue
print(s["name"])
ret = viz.get_render(data, s["query"])
print_data(ret)
n += 1
if n > args.n: break
prg = next((p for p in profile if isinstance(p, ProfileProgramEvent) and p.name == args.kernel), None)
dev = next((p for p in profile if isinstance(p, ProfileDeviceEvent) and p.device == prg.device), None)
assert prg is not None and dev is not None, "must have program binary and device props"
target = f"gfx{dev.props['gfx_target_version']//1000}"
sqtt = [p for p in profile if isinstance(p, ProfileSQTTEvent) and p.kern == prg.tag]
pc_to_inst = {addr+prg.base:inst for addr,inst in amd_decode(prg.lib, target).items()}
rctx = decode(sqtt, {prg.tag:pc_to_inst})
waves = sorted(itertools.chain.from_iterable(rctx.inst_execs.values()), key=lambda w:(w.se, w.cu, w.simd, w.wave_id, w.begin_time))
if not waves: raise RuntimeError(f"no instruction traces for {args.kernel}")
run_numbers:dict[str, itertools.count] = {}
for w in itertools.islice(waves, args.n):
if w.wave_loc not in run_numbers: run_numbers[w.wave_loc] = itertools.count()
print(f"{w.wave_loc} N:{next(run_numbers[w.wave_loc])} Total Cycles:{w.end_time-w.begin_time}")
table = unpack_insts(w, pc_to_inst)
print(tabulate([r[:len(table["cols"])] for r in table["rows"]], headers=table["cols"], tablefmt="github"))
if __name__ == "__main__":
main()
tinygrad_repo/extra/thunder/amd/fa.py
+36 -20
View File
@@ -19,16 +19,33 @@ def _sharded_empty(shape:Tensor, ref:Tensor, axis:int|None, dtype:DTypeLike|None
@functools.cache
def custom_fused_qkv_rope_forward(q:UOp, k:UOp, v:UOp, xqkv:UOp, freqs_cis:UOp,
device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
code = (pathlib.Path(__file__).parent / "fused_qkv_rope.cpp").read_text()
threads = 256
thread_idx = UOp.special(threads, "lidx0")
block_idx_x, block_idx_y = UOp.special(B, "gidx0"), UOp.special(N, "gidx1")
sink = UOp.sink(q.base, k.base, v.base, xqkv.base, freqs_cis.base, thread_idx, block_idx_x, block_idx_y,
arg=KernelInfo(name="fused_qkv_rope_forward"))
compile_args = ["-std=c++20", "-ffast-math", f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}",
f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DTHREADS_PER_BLOCK={threads}"]
lib = HIPCCCompiler(arch, compile_args).compile_cached(code)
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
group_size = H // H_KV
q, k, v = q.reshape(B, N, H, D), k.reshape(B, N, H_KV, D), v.reshape(B, N, H_KV, D)
xqkv = xqkv.reshape(B, N, H_KV, group_size + 2, D)
b, n = UOp.range(B, 0), UOp.range(N, 1)
pair = UOp.range(D // 2, 2)
even = pair * 2
c = freqs_cis[0, n, 0, pair, 0].cast(dtypes.float)
s = freqs_cis[0, n, 0, pair, 1].cast(dtypes.float)
ordered:UOp|None = None
for kvh in range(H_KV):
q_out, k_out, v_out = (x.after(ordered) if ordered is not None else x for x in (q, k, v))
x_in = xqkv.after(ordered) if ordered is not None else xqkv
stores:list[UOp] = []
for rep in range(group_size):
a = x_in[b, n, kvh, rep, even].cast(dtypes.float)
bb = x_in[b, n, kvh, rep, even + 1].cast(dtypes.float)
h = kvh * group_size + rep
stores += [q_out[b, n, h, even].store((a * c - bb * s).cast(q.dtype)), q_out[b, n, h, even + 1].store((a * s + bb * c).cast(q.dtype))]
a = x_in[b, n, kvh, group_size, even].cast(dtypes.float)
bb = x_in[b, n, kvh, group_size, even + 1].cast(dtypes.float)
stores += [k_out[b, n, kvh, even].store((a * c - bb * s).cast(k.dtype)),
k_out[b, n, kvh, even + 1].store((a * s + bb * c).cast(k.dtype)),
v_out[b, n, kvh, even].store(x_in[b, n, kvh, group_size + 1, even]),
v_out[b, n, kvh, even + 1].store(x_in[b, n, kvh, group_size + 1, even + 1])]
ordered = UOp.group(*stores)
assert ordered is not None
return ordered.end(pair, n, b).sink(arg=KernelInfo(name="fused_qkv_rope_forward"))
@functools.cache
def custom_fused_qkv_rope_backward(dxqkv:UOp, dq:UOp, dk:UOp, dv:UOp, freqs_cis:UOp,
@@ -109,8 +126,7 @@ def fused_qkv_rope(xqkv:Tensor, freqs_cis:Tensor, n_heads:int, n_kv_heads:int, h
def _sharded_empty_like(ref:Tensor, axis:int|None=None) -> Tensor:
return _sharded_empty(ref.shape, ref, axis)
@functools.cache
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink):
def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=0):
def grad(dou:UOp, ker:UOp) -> tuple:
do = Tensor(dou, device=dou.device)
attn = Tensor(ker.src[1].after(ker), device=ker.src[1].device)
@@ -129,7 +145,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
delta_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
delta_vec, dq = Tensor.custom_kernel(delta_vec, dq, attn, do, fxn=functools.partial(custom_fa_backward_pre, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:2]
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D))[:3]
dq, dk_partial, dv_partial = Tensor.custom_kernel(dq, dk_partial, dv_partial, do, xq, xk, xv, l_vec, delta_vec, fxn=functools.partial(custom_fa_backward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, window=window))[:3]
if D == 64:
dq = dq.reshape(B, H, N//16, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2).permute(0, 1, 2, 8, 9, 10, 11, 3, 4, 6, 7, 5, 12).reshape(B, H, N, D).transpose(1, 2)
@@ -149,7 +165,7 @@ def _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, sha
return grad
# TODO: remove write_flat once scheduler can remove reshapes between custom_kernel. TestCustomKernel.test_simple_reshape
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None):
def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False, write_flat:bool=False, sinks:Tensor|None=None, window:int=0):
assert attn_mask is None, "attn_mask not supported"
assert is_causal, "only causal attention supported"
@@ -176,18 +192,18 @@ def flash_attention(xq, xk, xv, attn_mask:Tensor|None=None, is_causal:bool=False
attn = _sharded_empty((B, N, H * D), xq, axis=shard_axis) if write_flat else _sharded_empty_like(xq, axis=shard_axis)
l_vec = _sharded_empty((B, H, 1, N), xq, dtype=dtypes.float32, axis=shard_axis_t)
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink)
grad = _fa_grad_fxn(B, H, N, D, H_local, H_KV_local, H_KV, B_local, shard_axis, shard_axis_t, single_device, arch, has_sink, window=window)
fwd_inputs = (attn, l_vec, xq, xk, xv) + ((sinks,) if has_sink else ())
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink), grad_fxn=grad)[:2]
attn, l_vec = Tensor.custom_kernel(*fwd_inputs, fxn=functools.partial(custom_fa_forward, device=single_device, arch=arch, B=B_local, N=N, H=H_local, H_KV=H_KV_local, D=D, has_sink=has_sink, window=window), grad_fxn=grad)[:2]
return attn, attn, l_vec
@functools.cache
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True):
def custom_fa_forward(o:UOp, l_vec:UOp, q:UOp, k:UOp, v:UOp, sinks:UOp|None=None, *, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, has_sink:bool=True, window:int=0):
code = (pathlib.Path(__file__).parent / "fa_fwd_causal.cpp").read_text()
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}"]
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DATTN_SINK={int(has_sink)}", f"-DWINDOW={window}"]
Q_BLOCK_SIZE = 32
NUM_WARPS = 8
@@ -247,10 +263,10 @@ def custom_fa_backward_pre(delta_vec:UOp, dq:UOp, o:UOp, do:UOp, device:str, arc
src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)), UOp(Ops.SOURCE, arg=code), UOp(Ops.BINARY, arg=lib)))
@functools.cache
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int):
def custom_fa_backward(dq:UOp, dk:UOp, dv:UOp, do:UOp, q:UOp, k:UOp, v:UOp, l_vec:UOp, delta_vec:UOp, device:str, arch:str, B:int, N:int, H:int, H_KV:int, D:int, window:int=0):
code = (pathlib.Path(__file__).parent / "fa_bwd_causal.cpp").read_text()
compile_args = [f"-I{(pathlib.Path(__file__).parent / 'include').as_posix()}", "-std=c++20", "-DKITTENS_CDNA4", "-DHIP_ENABLE_WARP_SYNC_BUILTINS", "-ffast-math",
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}"]
f"-DATTN_B={B}", f"-DATTN_N={N}", f"-DATTN_H={H}", f"-DATTN_H_KV={H_KV}", f"-DATTN_D={D}", f"-DWINDOW={window}"]
BLOCK_SIZE_KV = 256
GROUP_SIZE = H // H_KV
tinygrad_repo/extra/thunder/amd/fa_fwd_causal.cpp
+20
View File
@@ -269,7 +269,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
qo_tile<D, float> q_reg_fl;
load<1, qo_tile<D, float>, _gl_QKVO>(q_reg_fl, g.Qg, {batch_idx, tile_idx, head_idx, 0});
#if !WINDOW
mul(q_reg_fl, q_reg_fl, TEMPERATURE_SCALE); // Use sqrtf for clarity
#endif
copy(q_reg, q_reg_fl);
transpose(q_reg_transposed, q_reg);
@@ -288,6 +290,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
__builtin_amdgcn_sched_barrier(0);
if constexpr (causal) {
const int kv_end_pos = (min_tile + 1) * KV_BLOCK_SIZE;
@@ -337,6 +342,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
#if WINDOW
// window masks interior tiles that causal skips
mask_kv_tile(att_block[1], tile_idx, j - 2, neg_inf_v, lane);
@@ -401,6 +409,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK1
exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -469,6 +480,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK2
exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -535,6 +549,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[0]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[0], k_reg_transposed, q_reg_transposed, att_block[0]);
#if WINDOW
mul(att_block[0], att_block[0], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK3
exp2(att_block[1].tiles[1][0], att_block[1].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
@@ -597,6 +614,9 @@ __global__ void attend_ker(bf16 *O_ptr, float *L_vec_ptr, bf16 *Q_ptr, bf16 *K_p
zero(att_block[1]);
transpose(k_reg_transposed, k_reg);
mma_AtB(att_block[1], k_reg_transposed, q_reg_transposed, att_block[1]);
#if WINDOW
mul(att_block[1], att_block[1], TEMPERATURE_SCALE);
#endif
// Finish softmax for QK4
exp2(att_block[0].tiles[1][0], att_block[0].tiles[1][0]);
mul(norm_vec, norm_vec, scale_vec);
tinygrad_repo/extra/thunder/amd/fused_qkv_rope.cpp
-69
View File
@@ -1,69 +0,0 @@
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#ifndef ATTN_B
#define ATTN_B 2
#endif
#ifndef ATTN_N
#define ATTN_N 8192
#endif
#ifndef ATTN_H
#define ATTN_H 32
#endif
#ifndef ATTN_H_KV
#define ATTN_H_KV 8
#endif
#ifndef ATTN_D
#define ATTN_D 128
#endif
#ifndef THREADS_PER_BLOCK
#define THREADS_PER_BLOCK 256
#endif
constexpr int GROUP_SIZE = ATTN_H / ATTN_H_KV;
constexpr int HALF_D = ATTN_D / 2;
constexpr int PACKED_D = (GROUP_SIZE + 2) * ATTN_D;
extern "C" __global__ __launch_bounds__(THREADS_PER_BLOCK) void
fused_qkv_rope_forward(
__hip_bfloat16* __restrict__ q,
__hip_bfloat16* __restrict__ k,
__hip_bfloat16* __restrict__ v,
const __hip_bfloat16* __restrict__ xqkv,
const __hip_bfloat16* __restrict__ freqs_cis) {
const int b = blockIdx.x;
const int n = blockIdx.y;
const int bn = b * ATTN_N + n;
const int packed_bn = bn * ATTN_H_KV * PACKED_D;
const int q_bn = bn * ATTN_H * ATTN_D;
const int kv_bn = bn * ATTN_H_KV * ATTN_D;
if (threadIdx.x < HALF_D) {
const int pair = threadIdx.x;
const int even = pair << 1;
const float c = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 0]);
const float s = static_cast<float>(freqs_cis[((n * HALF_D + pair) * 2) + 1]);
for (int kvh = 0; kvh < ATTN_H_KV; kvh++) {
const int base = packed_bn + kvh * PACKED_D;
for (int rep = 0; rep < GROUP_SIZE; rep++) {
const int qbase = base + rep * ATTN_D;
const int h = kvh * GROUP_SIZE + rep;
const float a = static_cast<float>(xqkv[qbase + even]);
const float bb = static_cast<float>(xqkv[qbase + even + 1]);
const int out = q_bn + h * ATTN_D + even;
q[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
q[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
}
const float a = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even]);
const float bb = static_cast<float>(xqkv[base + GROUP_SIZE * ATTN_D + even + 1]);
const int out = kv_bn + kvh * ATTN_D + even;
k[out] = static_cast<__hip_bfloat16>(a * c - bb * s);
k[out + 1] = static_cast<__hip_bfloat16>(a * s + bb * c);
v[out] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even];
v[out + 1] = xqkv[base + (GROUP_SIZE + 1) * ATTN_D + even + 1];
}
}
}
tinygrad_repo/extra/thunder/tiny/tk/tiles.py
+1 -1
View File
@@ -209,7 +209,7 @@ class ST:
return cls(uop, rows, cols, layout, base_shape, ker)
def swizzle(self, row, col):
swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype.scalar())
swizzled_offset = self.base_shape.swizzle(row, col, self._uop.dtype)
row = swizzled_offset // self.base_shape.cols
col = swizzled_offset % self.base_shape.cols
tinygrad_repo/extra/tinyfs/fetch_file.py
-38
View File
@@ -1,38 +0,0 @@
from tinygrad.tensor import Tensor
from tinygrad.helpers import CHUNK_SIZE
from tinygrad.nn.state import fs_load
import argparse, math, hashlib
def _python_hash_1mb(data:bytes|bytearray):
chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
def hash_file(data: bytes|bytearray):
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
base_chunks = math.ceil(len(data) / CHUNK_SIZE)
tree_depth = math.ceil(math.log(base_chunks, CHUNK_SIZE // 16))
for _ in range(tree_depth + 1):
data_chunks = [data[i:i+CHUNK_SIZE] for i in range(0, len(data), CHUNK_SIZE)]
data_chunk_hashes = [_python_hash_1mb(chunk) for chunk in data_chunks]
data = b''.join(data_chunk_hashes)
if len(data) % CHUNK_SIZE != 0: data += bytes(CHUNK_SIZE - len(data) % CHUNK_SIZE)
return data[:16]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--hash", type=str, required=True, help="file hash to fetch")
parser.add_argument("--len", type=int, required=True, help="file length to fetch")
parser.add_argument("--dest", type=str, required=True, help="destination path to save the file")
parser.add_argument("--check", action="store_true", help="verify the file hash after fetching")
args = parser.parse_args()
fs_load(Tensor(bytes.fromhex(args.hash), device="CPU"), args.len).to(f"disk:{args.dest}").realize()
if args.check:
with open(args.dest, "rb") as f:
data = f.read()
assert hash_file(data) == bytes.fromhex(args.hash), "Hash mismatch after fetching file"
print("File hash verified successfully!")
tinygrad_repo/extra/tinyfs/fetch_raid.py
-42
View File
@@ -1,42 +0,0 @@
import json, multiprocessing, functools
from pathlib import Path
from tinygrad.tensor import Tensor
from tinygrad.helpers import tqdm, getenv
from tinygrad.nn.state import fs_load
raid_root = Path(getenv("RAID_ROOT", "/raid"))
def fetch_file(item):
path, info = item
h, size = info["hash"], info["size"]
path = raid_root / Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
try:
pt = fs_load(Tensor(bytes.fromhex(h), device="CPU"), size).to(f"disk:{path.as_posix()}").realize()
except Exception as e:
print(f"error fetching {path}, {h}, {size}: {e}")
raise
pt.uop.buffer.deallocate()
def fetch_mapping(h, l):
mapping_tensor = fs_load(Tensor(bytes.fromhex(h)), l).realize()
mapping = mapping_tensor.data().tobytes().decode()
mapping = json.loads(mapping)
mapped_files = mapping.items()
return list(mapped_files)
if __name__ == "__main__":
h, l = getenv("HASH", "d734f5e3be9f1e9d863bfaa4fc6c1ef2"), getenv("LENGTH", 175866113)
with multiprocessing.Pool(processes=1) as pool:
mapped_files = pool.apply(functools.partial(fetch_mapping, h, l))
print(f"fetched mapping for {len(mapped_files)} files")
with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
for _ in tqdm(pool.imap_unordered(fetch_file, mapped_files), total=len(mapped_files)):
pass
tinygrad_repo/extra/tinyfs/upload_raid.py
-32
View File
@@ -1,32 +0,0 @@
from pathlib import Path
import multiprocessing, json
from tinygrad.tensor import Tensor
from tinygrad.helpers import tqdm
from tinygrad.nn.state import fs_store
raid_root = Path("/raid")
def upload_file(path: Path):
pt = Tensor(path).realize()
h = fs_store(pt).realize()
pt.uop.realized.deallocate()
return h.data().hex(), path, pt.nbytes()
if __name__ == "__main__":
raid_files = sorted([p for p in raid_root.rglob("*") if p.is_file()])
print(f"found {len(raid_files)} files in /raid")
mapping = {}
with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
for h, p, s in tqdm(pool.imap_unordered(upload_file, raid_files), total=len(raid_files)):
mapping[p.relative_to(raid_root).as_posix()] = {"hash": h, "size": s}
# sort the mapping by key
mapping = dict(sorted(mapping.items()))
mapping = json.dumps(mapping).encode()
mapping_tensor = Tensor(mapping, device="CPU")
h = fs_store(mapping_tensor).realize()
print(f"final hash: {h.data().hex()}, size: {len(mapping)}")
tinygrad_repo/extra/torch_backend/backend.py
+87 -125
View File
@@ -4,7 +4,7 @@
# A006 Lambda argument `input` is shadowing a Python builtin
from tinygrad import Tensor, dtypes, Device
from tinygrad.uop.ops import Ops, GroupOp
from tinygrad.helpers import getenv, prod, strides_for_shape, argfix
from tinygrad.helpers import getenv, prod, strides_for_shape
import torch.lib
TORCH_DEBUG = getenv("TORCH_DEBUG")
import torch, pathlib, operator, functools, weakref
@@ -73,6 +73,12 @@ def wrap_view_op(fn):
return wrap(ret)
return _wrap
# NOTE: list assignment raises IndexError on an out of range dim, and the index must be a tuple: a list of all ints is one advanced index
def _index_dim(self, dim, idx):
idxs = [slice(None)] * self.ndim
idxs[dim] = idx
return self[tuple(idxs)]
view_ops = {
"aten.view": Tensor.reshape,
"aten._unsafe_view": Tensor.reshape, # when are views unsafe, and do we care?
@@ -82,15 +88,13 @@ view_ops = {
"aten.transpose.int": Tensor.transpose,
"aten.squeeze.dim": Tensor.squeeze,
"aten.unsqueeze": Tensor.unsqueeze,
"aten.select.int": lambda self, dim, idx: self[(slice(None),) * (dim%self.ndim) + (idx,)],
"aten.select.int": _index_dim,
"aten.permute": Tensor.permute,
"aten.alias": lambda self: self,
"aten.diagonal": Tensor.diagonal,
"aten.slice.Tensor": lambda self, dim=0, start=None, end=None, step=1: _index_dim(self, dim, slice(start, end, step)),
}
# torch 2.10 handles this natively
if tuple(map(int, torch.__version__.split('.')[:2])) < (2, 10): view_ops.update({"aten.detach": Tensor.detach})
for k,v in view_ops.items(): torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(wrap_view_op(v))
def _get_view_ops(view): return getattr(view, "_view_ops", [])
@@ -99,46 +103,21 @@ def _apply_view_ops(target, ops):
for fn, args, kwargs in ops: target = fn(target, *args, **kwargs)
return target
# similar to https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/InferSize.h
def _reshape_target_shape(shape:tuple[int, ...], args) -> tuple[int, ...]|None:
if not (req := argfix(*args)): return None
new_shape, infer_idx = [], -1
for i, s in enumerate(req):
if s is None: s = shape[i] if i < len(shape) else None
if not isinstance(s, int): return None
if s == -1:
if infer_idx != -1: return None
infer_idx = len(new_shape)
new_shape.append(s)
total = prod(shape)
if infer_idx != -1:
known = prod(x for x in new_shape if x != -1)
if known == 0:
if total != 0: return None
new_shape[infer_idx] = 0
else: new_shape[infer_idx] = total // known
return tuple(new_shape) if prod(new_shape) == total else None
# TODO: can we get rid of this? only for test_flatten_reshape_add
# a chain of reshapes is undone by reshaping the value back to the base
def _try_simple_reshape_view_write(base: Tensor, view: Tensor, val: Tensor) -> bool:
if not (ops := _get_view_ops(view)): return False
shapes = [base.shape]
for fn, args, _ in ops:
if fn is Tensor.reshape:
if not (next_shape := _reshape_target_shape(shapes[-1], args)): return False
shapes.append(next_shape)
if shapes[-1] != view.shape: return False
for s in reversed(shapes[:-1]): val = val.reshape(s)
base.assign(val)
if any(fn is not Tensor.reshape for fn, _, _ in ops): return False
base.assign(val.reshape(base.shape))
return True
def _view_write(base: Tensor, view: Tensor, value: Tensor) -> None:
val = value if value.dtype == base.dtype else value.cast(base.dtype)
if view.shape == base.shape: return base.assign(val)
if _try_simple_reshape_view_write(base, view, val): return
idx_base = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape)
idx_view = _apply_view_ops(idx_base, _get_view_ops(view)).reshape(-1)
flat_base = base.reshape(base.numel()).contiguous()
# clone, not contiguous: contiguous() on a base that already owns its buffer returns the base itself, and scattering
# into that is an in-place write to a buffer other tensors still hold, which setitem refuses
flat_base = base.reshape(base.numel()).clone()
flat_base[idx_view] = val.reshape(-1)
base.assign(flat_base.reshape(base.shape))
@@ -166,11 +145,6 @@ def _index_put_impl_(self, indices, values, accumulate=False, unsafe=False):
def index_put(self, indices, values, accumulate=False):
return aten.index_put(self.cpu(), [z.cpu() if isinstance(z, torch.Tensor) else None for z in indices], values.clone().cpu(), accumulate).tiny()
@torch.library.impl("aten::isin.Tensor_Tensor_out", "privateuseone")
def isin_tensor_tensor_out(x, y, *, assume_unique=False, invert=False, out=None):
result = (unwrap(x).unsqueeze(-1) == unwrap(y).flatten()).any(-1)
return out.copy_(wrap(~result if invert else result))
@torch.library.impl("aten::randperm.generator_out", "privateuseone")
def randperm_generator(n, generator=None, out=None):
if generator is not None: raise NotImplementedError("tinygrad torch backend does not support torch.Generator for randperm")
@@ -231,49 +205,6 @@ def as_strided(tensor:torch.Tensor, size, stride, storage_offset=None):
def _reshape_alias(tensor:torch.Tensor, size, stride):
return _as_strided(tensor, size, stride)
@torch.library.impl("aten::empty_strided", "privateuseone")
def empty_strided(size, stride, dtype=None, layout=None, device=None, pin_memory=False):
if TORCH_DEBUG: print(f"empty_strided {size=} {stride=} {dtype=} {layout=} {device=} {pin_memory=}")
ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
# TODO: should return with requested strides
return wrap(ret)
@torch.library.impl("aten::empty.memory_format", "privateuseone")
def empty_memory_format(size, dtype=None, layout=None, device=None, pin_memory=False, memory_format=None):
if TORCH_DEBUG: print(f"empty.memory_format {size=} {dtype=} {layout=} {device=} {pin_memory=} {memory_format=}")
ret = Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
return wrap(ret)
@torch.library.impl("aten::max_pool2d_with_indices", "privateuseone")
def max_pool2d_with_indices(self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False):
# TODO: supprt stride [] in tinygrad?
if stride is not None and len(stride) == 0: stride = None
ret, idx = unwrap(self).max_pool2d(kernel_size, stride, dilation, padding, ceil_mode, return_indices=True)
return (wrap(ret), wrap(idx.cast(dtypes.int64)))
@torch.library.impl("aten::max_pool2d_with_indices_backward", "privateuseone")
def max_pool2d_with_indices_backward(grad_out:torch.Tensor, self:torch.Tensor, kernel_size:tuple[int, ...], stride=None, padding=0, dilation=1, ceil_mode=False, indices=None):
return wrap(Tensor.max_unpool2d(unwrap(grad_out), unwrap(indices), output_size=unwrap(self).shape))
@torch.library.impl("aten::max_unpool2d", "privateuseone")
def max_unpool2d(self:torch.Tensor, indices:torch.Tensor, output_size):
return wrap(unwrap(self).max_unpool2d(unwrap(indices), output_size=output_size))
@torch.library.impl("aten::arange", "privateuseone")
def arange(end, dtype=None, device=None, pin_memory=None):
has_float = isinstance(end, float)
return wrap(Tensor.arange(0, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::arange.start", "privateuseone")
def arange_start(start, end, dtype=None, device=None, pin_memory=None):
has_float = any(isinstance(x, float) for x in (start, end))
return wrap(Tensor.arange(start, end, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::arange.start_step", "privateuseone")
def arange_start_step(start, end, step, dtype=None, device=None, pin_memory=None):
has_float = any(isinstance(x, float) for x in (start, end, step))
return wrap(Tensor.arange(start, end, step, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if has_float else torch.int64))))
@torch.library.impl("aten::convolution_overrideable", "privateuseone")
def convolution_overrideable(input, weight, bias, stride, padding, dilation, transposed, output_padding, groups):
if TORCH_DEBUG >= 1:
@@ -294,12 +225,27 @@ def convolution_backward_overrideable(grad_out, input, weight, stride, padding,
grads = out.gradient(*[t for t,m in zip([input, weight, bias], output_mask) if m], gradient=grad_out)
return tuple([wrap(grads.pop(0)) if m else None for m in output_mask])
@torch.library.impl("aten::slice.Tensor", "privateuseone")
@wrap_view_op
def slice_tensor(self, dim=0, start=None, end=None, step=1):
slices = [slice(None)] * self.ndim
slices[dim] = slice(start, end, step)
return self[slices]
# the functional scatters. without an impl aten falls back to a path that assumes a real storage: "self.has_storage() INTERNAL ASSERT FAILED"
def _scatter_into(self, src, dim, index):
out = unwrap(self).clone()
slices = [slice(None)] * out.ndim
slices[dim] = index
out[slices] = unwrap(src).cast(out.dtype) # torch casts src to self's dtype, tinygrad setitem demands they already match
return wrap(out)
@torch.library.impl("aten::slice_scatter", "privateuseone")
def slice_scatter(self, src, dim=0, start=None, end=None, step=1): return _scatter_into(self, src, dim, slice(start, end, step))
@torch.library.impl("aten::select_scatter", "privateuseone")
def select_scatter(self, src, dim, index): return _scatter_into(self, src, dim, index)
@torch.library.impl("aten::diagonal_scatter", "privateuseone")
def diagonal_scatter(self, src, offset=0, dim1=0, dim2=1):
# a diagonal is not one axis, so scatter through the flat indices it picks out
base, out = unwrap(self), unwrap(self).clone().reshape(-1)
idx = Tensor.arange(base.numel(), dtype=dtypes.int32).reshape(base.shape).diagonal(offset, dim1, dim2).reshape(-1)
out[idx] = unwrap(src).cast(base.dtype).reshape(-1)
return wrap(out.reshape(base.shape))
@torch.library.impl("aten::slice_backward", "privateuseone")
def slice_backward(grad_out, input_sizes, dim, start, end, step):
@@ -341,19 +287,14 @@ for dim in [1, 2, 3]:
torch.library.impl(f"aten::{pad_type}_pad{dim}d", "privateuseone")(functools.partial(pad_forward, mode=mode))
torch.library.impl(f"aten::{pad_type}_pad{dim}d_backward", "privateuseone")(functools.partial(pad_backward, mode=mode))
def upsample(self, size, align_corners=False, mode=None): return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=align_corners))
# the schemas are all positional: (self, output_size, align_corners, *scales) for linear, (self, output_size, *scales) for nearest.
def upsample(self, size, *args, mode=None):
return wrap(Tensor.interpolate(unwrap(self), size, mode=mode, align_corners=args[0] if mode == "linear" else False))
for i,pre in enumerate(["", "bi", "tri"]):
torch.library.impl(f"aten::upsample_{pre}linear{i+1}d", "privateuseone")(functools.partial(upsample, mode="linear"))
torch.library.impl(f"aten::upsample_nearest{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest"))
torch.library.impl(f"aten::_upsample_nearest_exact{i+1}d", "privateuseone")(functools.partial(upsample, mode="nearest-exact"))
@torch.library.impl("aten::scatter_add.out", "privateuseone")
def scatter_add(self, dim, index, src, out):
self, index, src, out_unwrapped = unwrap(self), unwrap(index), unwrap(src), unwrap(out)
if self.shape == (): _apply_inplace(out_unwrapped, src)
else: _apply_inplace(out_unwrapped, Tensor.scatter_reduce(self, dim, index, src, reduce='sum'))
return out
def _copy_between_devices(src, dest, cast_dtype, to_device, non_blocking=False):
if src.is_tiny and dest.is_tiny:
src_t, dest_t = unwrap(src), unwrap(dest)
@@ -404,15 +345,11 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
_apply_inplace(unwrap(indices), out_indices.cast(dtypes.int64))
return values, indices
@torch.library.impl("aten::_linalg_svd", "privateuseone")
def _linalg_svd(self, full_matrices=False):
U, S, Vh = unwrap(self).svd(full_matrices)
return wrap(U), wrap(S), wrap(Vh)
# register some decompositions
from torch._decomp import get_decompositions
decomps = [
aten.native_layer_norm_backward,
aten.native_group_norm_backward,
aten.linalg_cross,
aten.addmm,
aten.addcmul,
@@ -447,12 +384,20 @@ decomps = [
aten._softmax_backward_data, aten.embedding_dense_backward,
aten.linalg_vector_norm,
aten.binary_cross_entropy, aten.binary_cross_entropy_backward,
# the C++ mse/smooth_l1 kernels resize their out tensor, and a tiny tensor has no storage to resize
aten.mse_loss, aten.mse_loss_backward,
aten.smooth_l1_loss, aten.smooth_l1_loss_backward,
aten.upsample_nearest2d.out,
# NOTE: only the "out" overload, the "vec" one is CompositeImplicitAutograd and overriding it loses the autograd kernel
aten.upsample_bicubic2d.out,
aten._adaptive_avg_pool2d,
# activations
aten.hardswish, aten.hardswish_backward,
aten.hardtanh, aten.hardtanh_backward,
aten.gelu, aten.gelu_backward,
aten.logical_and,
# NOTE: no aten.logical_or here, its decomposition reaches aten.bitwise_or through a path that checks aliasing by
# reading storage, which a tiny tensor has none of. it gets a direct impl below instead
aten.logical_and, aten.logical_xor,
aten.randint,
aten.eye,
aten.hardsigmoid_backward,
@@ -495,7 +440,7 @@ simple_tensor_methods = [
# reduce
"all", "any", "argmax", "argmin", "cumsum", "cumprod",
# complex
"avg_pool2d", "linspace"]
"linspace"]
tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_methods}, **{
"aten.add.out": lambda input,other,alpha=1: input+alpha*other,
@@ -540,6 +485,8 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
"aten.where.self_out": Tensor.where,
"aten.prod.int_out": Tensor.prod,
"aten.scatter.src_out": Tensor.scatter,
"aten.scatter_add.out": lambda self,dim,index,src: src if self.shape == () else Tensor.scatter_reduce(self, dim, index, src, reduce="sum"),
"aten.isin.Tensor_Tensor_out": lambda x,y,assume_unique=False,invert=False: (x.unsqueeze(-1)==y.flatten()).any(-1) != invert,
# NOTE: axis=[] in torch means all, change tinygrad?
"aten.sum.IntList_out": lambda self,axis,keepdim=False,dtype=None:
self.sum(axis if axis is None or len(axis) else None, keepdim,
@@ -555,10 +502,9 @@ def wrap_out(f):
assert out.shape == assigned.shape, f"shape mismatch: {assigned.shape} -> {out.shape}"
assert out.device == assigned.device or out.device is None or assigned.device is None, f"device mismatch: {assigned.device} -> {out.device}"
assert out.dtype == assigned.dtype, f"dtype mismatch: {assigned.dtype} -> {out.dtype}"
# an out= that is a view has to be written through its base, and _apply_inplace gives a deviceless base its buffer first
if canonical_base(out) is not out: return _apply_inplace(out, assigned) or out
if out.device is None and assigned.device is not None: out.replace(out.empty_like(device=assigned.device))
return out.assign(assigned)
# writing out= is an in-place write like any other: through the base if it is a view, refreshing any derived views
_apply_inplace(out, assigned)
return out
return _wrap_out
def _inplace_op(t, new_value):
@@ -566,7 +512,14 @@ def _inplace_op(t, new_value):
else: _apply_inplace(t, new_value)
return t
tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
# the three arange overloads are one function at different arity, and dtype/layout/device/pin_memory are keyword only in all of them
def _arange(*args, dtype=None, **_):
return Tensor.arange(*args, dtype=_from_torch_dtype(dtype or (torch.get_default_dtype() if any(isinstance(x, float) for x in args) else torch.int64)))
def _empty(size, dtype=None, device=None, **_):
return Tensor.empty(*size, dtype=_from_torch_dtype(dtype or torch.get_default_dtype()), device=_from_torch_device(device))
tiny_backend = {**tiny_backend_out, **{
"aten.remainder.Scalar_Tensor": lambda x,y: x%y,
"aten.floor_divide": lambda x,y: x//y,
"aten.floor_divide_.Tensor": lambda x,y: x//y,
@@ -579,8 +532,8 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
# inplace ops using replace for fusion
"aten.zero_": lambda x: x.const_like(0),
"aten.fill_.Scalar": lambda x, y: x.const_like(y),
"aten.add_.Tensor": lambda self, other, alpha=1.0: self + other * alpha,
"aten.add_.Scalar": lambda self, other, alpha=1.0: self + other * alpha,
"aten.add_.Tensor": lambda self, other, alpha=1: self + other * alpha,
"aten.add_.Scalar": lambda self, other, alpha=1: self + other * alpha,
"aten.mul_.Tensor": lambda self, other: self * other,
"aten.mul_.Scalar": lambda self, other: self * other,
# relu doesn't have an out form?
@@ -613,7 +566,9 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
# these don't work in out form, they have size 0
"aten.abs": Tensor.abs,
"aten.logical_not": Tensor.logical_not,
"aten.logical_or_": lambda x, y: x | y,
# compare against zero first: logical_* is bool-valued for any input dtype, while | is bitwise
"aten.logical_or": lambda x, y: (x != 0) | (y != 0),
"aten.logical_or_": lambda x, y: (x != 0) | (y != 0),
"aten.multinomial": Tensor.multinomial,
"aten.masked_fill_.Scalar": lambda self, mask, value: self.masked_fill(mask, value),
"aten.masked_fill_.Tensor": lambda self, mask, value: self.masked_fill(mask, value),
@@ -622,14 +577,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
"aten.masked_select": Tensor.masked_select,
"aten.all": Tensor.all,
"aten.sgn": Tensor.sign,
"aten.acos": Tensor.acos,
"aten.any": Tensor.any,
"aten.bitwise_not": Tensor.bitwise_not,
"aten.argmax": Tensor.argmax,
"aten.argmin": Tensor.argmin,
"aten.asinh": Tensor.asinh,
"aten.mul": Tensor.mul,
"aten.atanh": Tensor.atanh,
"aten.fill_.Tensor": lambda self, value: self.const_like(value.reshape(()).item()),
"aten.flip": Tensor.flip,
"aten.scatter_reduce.two": Tensor.scatter_reduce,
@@ -640,10 +588,22 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
"aten.add.Tensor": lambda input,other,alpha=1: input+alpha*other,
"aten.linspace": lambda start, stop, steps, dtype=None, **kwargs:
Tensor.linspace(start, stop, steps, **({"dtype": _from_torch_dtype(dtype)} if dtype is not None else {})),
# the functional copy_. without an impl the fallback segfaults on a tensor with no storage
"aten.copy": lambda self,src,non_blocking=False: src.cast(self.dtype).to(self.device).expand(self.shape),
"aten.arange": lambda end, **kwargs: _arange(0, end, **kwargs),
"aten.arange.start": _arange,
"aten.arange.start_step": _arange,
# empty_strided takes the strides and drops them: we always allocate contiguous
"aten.empty_strided": lambda size, stride, **kwargs: _empty(size, **kwargs),
"aten.empty.memory_format": _empty,
# TODO: supprt stride [] in tinygrad?
"aten.max_pool2d_with_indices": lambda self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False: ((r:=Tensor.max_pool2d(self, kernel_size, stride or None, dilation, padding, ceil_mode, return_indices=True))[0], r[1].cast(dtypes.int64)),
"aten.max_pool2d_with_indices_backward": lambda grad_out,self,kernel_size,stride=None,padding=0,dilation=1,ceil_mode=False,indices=None: Tensor.max_unpool2d(grad_out, indices, output_size=self.shape),
"aten.max_unpool2d": lambda self,indices,output_size: Tensor.max_unpool2d(self, indices, output_size=output_size),
"aten._linalg_svd": lambda self,full_matrices=False: Tensor.svd(self, full_matrices),
"aten.topk": Tensor.topk,
"aten.constant_pad_nd": lambda self, padding, value=0.0: self.pad(padding, mode="constant", value=value).contiguous(),
# TODO: input contiguous is needed to prevent CFGContext circular dependency assertion for shapes >512 (see test_cumsum_arange_large)
"aten.cumsum": lambda self, dim: self.contiguous().cumsum(dim),
"aten.cumsum": lambda self, dim: self.cumsum(dim),
"aten.logsumexp": lambda self, axis, keepdim=False: self.logsumexp(axis[0], keepdim=keepdim),
"aten.roll": Tensor.roll,
"aten.logcumsumexp": Tensor.logcumsumexp,
@@ -652,6 +612,7 @@ tiny_backend = {**{k:wrap_out(v) for k,v in tiny_backend_out.items()}, **{
self.ones_like(**{k: v for k, v in {"dtype": _from_torch_dtype(dtype) if dtype else None,
"device": _from_torch_device(device) if device else None}.items() if v is not None}),
"aten.max.dim": lambda self, dim, keepdim=False: (self.max(dim, keepdim), self.argmax(dim, keepdim).cast(dtype=dtypes.int64)),
"aten.min.dim": lambda self, dim, keepdim=False: (self.min(dim, keepdim), self.argmin(dim, keepdim).cast(dtype=dtypes.int64)),
"aten.cummax": lambda self, dim: ((r := self.cummax(dim))[0], r[1].cast(dtypes.int64)),
"aten.cummin": lambda self, dim: ((r := self.cummin(dim))[0], r[1].cast(dtypes.int64)),
"aten.nonzero": Tensor.nonzero,
@@ -713,15 +674,16 @@ def wrap_inplace_view_op(f):
return nf
# the aten schema says how an op is called: an inplace view retargets the view, a writable first arg is inplace,
# and a writable out arg must have come from tiny_backend_out so that wrap_out was applied
# and a writable out arg gets wrap_out's dtype cast, shape assert, and view write-through
for k,v in tiny_backend.items():
name, _, overload = k.removeprefix("aten.").partition(".")
op = getattr(getattr(aten, name), overload or "default")
writes = [a.name for a in op._schema.arguments if a.alias_info is not None and a.alias_info.is_write]
if torch.Tag.inplace_view in op.tags: fxn = wrap_inplace_view_op(v)
elif writes == [op._schema.arguments[0].name] and op._schema.returns: fxn = wrap_inplace(v)
elif not writes or (writes == ["out"] and k in tiny_backend_out): fxn = wrap_fxn(k, v)
else: raise RuntimeError(f"{k} writes {writes}: expected an inplace first arg, or an out arg with {k} in tiny_backend_out")
elif not writes: fxn = wrap_fxn(k, v)
elif writes == ["out"]: fxn = wrap_fxn(k, wrap_out(v))
else: raise RuntimeError(f"{k} writes {writes}: unhandled writable arg in schema")
torch.library.impl(k.replace("aten.", "aten::"), "privateuseone")(fxn)
@torch.library.impl("aten::equal", "privateuseone")
tinygrad_repo/extra/torch_backend/test.py
+120
View File
@@ -83,6 +83,12 @@ class TestTorchBackend(unittest.TestCase):
torch.add(torch.ones(5, device=device), torch.ones(5, device=device), out=a)
self.assertEqual(a.detach().storage_offset(), 3)
def test_out_refreshes_views_of_base(self):
a = torch.zeros(4, device=device)
v = a[2:]
torch.add(torch.ones(4, device=device), torch.ones(4, device=device), out=a)
np.testing.assert_equal(v.cpu().numpy(), [2., 2.])
@unittest.expectedFailure # TODO: storage offset assumes a contiguous source, use UOp.contiguous_view_offset
def test_storage_offset_non_contiguous_source(self):
a = torch.arange(12., device=device).reshape(3,4)
@@ -166,6 +172,15 @@ class TestTorchBackend(unittest.TestCase):
expected = np.array([[1.5, 5.2, 9.0], [13.2, 17.1, 18.4]], dtype=np.float32)
np.testing.assert_equal(y3.cpu().numpy(), expected)
def test_argmax_argmin(self):
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
c = a.cpu()
for got, want in [(a.argmax(), c.argmax()), (a.argmin(0), c.argmin(0)), (a.argmax(1, keepdim=True), c.argmax(1, keepdim=True)),
(torch.min(a, 1).indices, torch.min(c, 1).indices), (torch.max(a, 1).indices, torch.max(c, 1).indices),
(torch.min(a, 1).values, torch.min(c, 1).values), (torch.min(a, 1, keepdim=True).indices, torch.min(c, 1, keepdim=True).indices)]:
self.assertEqual(got.dtype, want.dtype) # torch's arg reduces are int64, tinygrad's are int32
np.testing.assert_equal(got.cpu().numpy(), want.numpy())
def test_isfinite(self):
a = torch.ones(4, device=device)
np.testing.assert_equal(torch.isfinite(a).cpu().numpy(), [True, True, True, True])
@@ -373,6 +388,22 @@ class TestTorchBackend(unittest.TestCase):
for bwd_eps in [1e-5, 0.3]:
for got, want in zip(run(device, bwd_eps), run("cpu", bwd_eps)): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
def test_groupnorm_backward(self):
def run(dev):
x = torch.arange(24., device=dev).reshape(2, 4, 3).requires_grad_()
w = torch.linspace(0.5, 2.0, 4).to(dev).requires_grad_()
torch.nn.functional.group_norm(x, 2, w, torch.zeros(4, device=dev)).square().sum().backward()
return x.grad.cpu().numpy(), w.grad.cpu().numpy()
for got, want in zip(run(device), run("cpu")): np.testing.assert_allclose(got, want, atol=1e-4, rtol=1e-3)
def test_mse_smooth_l1_loss_backward(self):
def run(dev, loss):
x = torch.arange(4., device=dev).requires_grad_()
loss(x, torch.ones(4, device=dev)).backward()
return x.grad.cpu().numpy()
for loss in [torch.nn.functional.mse_loss, torch.nn.functional.smooth_l1_loss]:
np.testing.assert_allclose(run(device, loss), run("cpu", loss), atol=1e-6)
def test_batchnorm_unsqueeze(self):
bn = torch.nn.BatchNorm2d(4).to(device)
x = torch.randn(8, 4, 3, 3, device=device)
@@ -516,6 +547,15 @@ class TestTorchBackend(unittest.TestCase):
cpu_res = torch.arange(20, dtype=torch.float32)[::2][1:4].numpy()
np.testing.assert_equal(torch_res, cpu_res)
def test_select_out_of_range_dim(self):
a = torch.arange(12, dtype=torch.int32, device=device).reshape(3, 4)
with self.assertRaises(IndexError): a.select(5, 0)
def test_select_collapses_the_only_dim(self):
a = torch.arange(3, dtype=torch.int32, device=device)
self.assertEqual(a.select(0, 1).shape, ())
np.testing.assert_equal(a.select(0, 1).cpu().numpy(), 1)
def test_slice_negative_dim(self):
a = torch.arange(13, dtype=torch.int32, device=device).repeat(8, 1)
torch_chunks = a.chunk(3, -1)
@@ -796,6 +836,86 @@ class TestTorchBackend(unittest.TestCase):
np.testing.assert_allclose(w_tiny.grad.cpu().numpy(), w_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
np.testing.assert_allclose(b_tiny.grad.cpu().numpy(), b_cpu.grad.numpy(), atol=1e-4, rtol=1e-3)
def test_write_through_detach_of_unrealized(self):
a = torch.empty(4, device=device)
a.detach().fill_(3)
np.testing.assert_equal(a.cpu().numpy(), [3, 3, 3, 3])
def test_square_transpose_inplace(self):
# a same-shape transpose is not a reshape: writing the transposed values straight back would scramble the base
a = torch.tensor([[0., 1., 2.], [3., 4., 5.], [6., 7., 8.]], device=device)
a.transpose(0, 1).add_(100)
np.testing.assert_equal(a.cpu().numpy(), [[100., 101., 102.], [103., 104., 105.], [106., 107., 108.]])
def test_interpolate(self):
a = torch.arange(4, dtype=torch.float32, device=device).reshape(1, 1, 2, 2)
nearest = torch.nn.functional.interpolate(a, scale_factor=2.0)
np.testing.assert_equal(nearest.cpu().numpy()[0, 0], [[0, 0, 1, 1], [0, 0, 1, 1], [2, 2, 3, 3], [2, 2, 3, 3]])
linear = torch.nn.functional.interpolate(a, size=(4, 4), mode="bilinear", align_corners=False)
ref = torch.nn.functional.interpolate(a.cpu(), size=(4, 4), mode="bilinear", align_corners=False)
np.testing.assert_allclose(linear.cpu().numpy(), ref.numpy(), rtol=1e-5)
def test_interpolate_bicubic_area(self):
a = torch.arange(32, dtype=torch.float32, device=device).reshape(1, 2, 4, 4)
for mode, scale in [("bicubic", 2.0), ("area", 0.5)]:
ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=scale, mode=mode)
np.testing.assert_allclose(torch.nn.functional.interpolate(a, scale_factor=scale, mode=mode).cpu().numpy(), ref.numpy(), atol=1e-4)
@unittest.expectedFailure
def test_interpolate_bicubic_backward(self):
# the forward comes from a decomposition, but aten::upsample_bicubic2d_backward has none (nor does
# aten::_adaptive_avg_pool2d_backward, for area), so training through these modes needs a real kernel
x = torch.arange(32., dtype=torch.float32, device=device).reshape(1, 2, 4, 4).requires_grad_()
torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bicubic").sum().backward()
@unittest.expectedFailure
def test_interpolate_inexact_scale(self):
# torch forwards the raw scale_factor, Tensor.interpolate recomputes it from output_size, and they disagree here
a = torch.arange(6, dtype=torch.float32, device=device).reshape(1, 1, 2, 3)
tiny = torch.nn.functional.interpolate(a, scale_factor=2.5, mode="bilinear")
ref = torch.nn.functional.interpolate(a.cpu(), scale_factor=2.5, mode="bilinear")
np.testing.assert_allclose(tiny.cpu().numpy(), ref.numpy(), rtol=1e-5)
def test_logical_or_xor(self):
a = torch.tensor([True, True, False, False], device=device)
b = torch.tensor([True, False, True, False], device=device)
np.testing.assert_equal(torch.logical_or(a, b).cpu().numpy(), [True, True, True, False])
np.testing.assert_equal(torch.logical_xor(a, b).cpu().numpy(), [False, True, True, False])
# bool-valued whatever the input dtype, so this is not | and ^
i, j = torch.tensor([2, 0, 5, 0], device=device), torch.tensor([0, 0, 1, 1], device=device)
np.testing.assert_equal(torch.logical_or(i, j).cpu().numpy(), [True, False, True, True])
np.testing.assert_equal(torch.logical_xor(i, j).cpu().numpy(), [True, False, False, True])
def test_slice_scatter(self):
# the scatters are functional: they return a new tensor and must leave the one they were given alone
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
out = torch.slice_scatter(a, torch.ones(1, 4, device=device), 0, 0, 1)
np.testing.assert_equal(out.cpu().numpy(), [[1, 1, 1, 1], [4, 5, 6, 7], [8, 9, 10, 11]])
np.testing.assert_equal(a.cpu().numpy(), np.arange(12, dtype=np.float32).reshape(3, 4))
def test_slice_scatter_casts_src(self):
a = torch.zeros(3, 4, device=device)
out = torch.slice_scatter(a, torch.ones(1, 4, dtype=torch.int32, device=device), 0, 0, 1)
self.assertEqual(out.dtype, torch.float32)
np.testing.assert_equal(out.cpu().numpy()[0], np.ones(4, dtype=np.float32))
def test_select_scatter(self):
a = torch.arange(12, dtype=torch.float32, device=device).reshape(3, 4)
out = torch.select_scatter(a, torch.ones(4, device=device), 0, 1)
np.testing.assert_equal(out.cpu().numpy(), [[0, 1, 2, 3], [1, 1, 1, 1], [8, 9, 10, 11]])
def test_diagonal_scatter(self):
a = torch.zeros(3, 3, device=device)
out = torch.diagonal_scatter(a, torch.arange(3, dtype=torch.float32, device=device))
np.testing.assert_equal(out.cpu().numpy(), np.diag([0., 1., 2.]))
np.testing.assert_equal(a.cpu().numpy(), np.zeros((3, 3), dtype=np.float32))
def test_copy_functional(self):
# without an impl this segfaults rather than fails: a regression here takes the whole run down
a = torch.arange(4, dtype=torch.float32, device=device)
out = torch.ops.aten.copy(a, torch.zeros(4, device=device))
np.testing.assert_equal(out.cpu().numpy(), [0., 0., 0., 0.])
np.testing.assert_equal(a.cpu().numpy(), [0., 1., 2., 3.])
from tinygrad import Tensor
class TestBackendHelpers(unittest.TestCase):
tinygrad_repo/pyproject.toml
+6 -2
View File
@@ -1,6 +1,6 @@
[project]
name = "tinygrad"
version = "0.13.0"
version = "0.14.0"
description = "You like pytorch? You like micrograd? You love tinygrad! <3"
authors = [{ name = "George Hotz" }]
@@ -84,7 +84,7 @@ testing = [
"pillow",
"onnx==1.19.0",
"onnx2torch",
"onnxruntime",
"onnxruntime==1.24.1",
"opencv-python",
"transformers",
"sentencepiece",
@@ -111,6 +111,10 @@ docs = [
"numpy",
]
mesa = ["tinymesa==25.2.7.2"]
autogen = [
"pyyaml",
"mako",
]
[tool.mutmut]
tinygrad_repo/spec/tinyspec.pdf
BIN
View File
Binary file not shown.
tinygrad_repo/spec/tinyspec.tex
+33 -33
View File
@@ -23,7 +23,6 @@
\definecolor{axblue}{HTML}{1565C0} % GLOBAL
\definecolor{axcyan}{HTML}{00838F} % LOCAL
\definecolor{axbrcyan}{HTML}{00ACC1} % WARP
\definecolor{axbrblue}{HTML}{42A5F5} % THREAD
\definecolor{axwhite}{HTML}{616161} % LOOP (gray on white paper)
\definecolor{axred}{HTML}{C62828} % REDUCE
\definecolor{axbrred}{HTML}{E53935} % GROUP_REDUCE
@@ -50,10 +49,12 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
\toprule
\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
\midrule
\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
Placeholder with shape $\mathbf{s}$. Substituted in \op{Function}. \\[4pt]
\op{Buffer} & $(\mathbf{s})$ & slot, dtype, device, addrspace &
Concrete buffer slot with shape $\mathbf{s}$. If device is a tuple, it creates the fully sized buffer across multiple devices. \\
\op{Param} & () & \texttt{ParamArg} &
Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
\op{Returned} & () & \texttt{ParamArg} &
Call output buffer placeholder; slot is its call argument index. \\[4pt]
\op{Buffer} & () & \texttt{ParamArg} &
Concrete buffer slot with flat storage of $\mathrm{size}$ elements. \\
\op{Const} & () & value, dtype &
A scalar constant with shape $(\ )$. \\
& & & Form vector consts with \op{Stack} \\
@@ -62,7 +63,21 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
\end{tabular}
\smallskip
\textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
\texttt{ParamArg} contains slot, dtype, concrete size (or \textsc{null} for a scalar), value bounds, alignment, name, addrspace, device, volatility, and optional image shape. \textbf{addrspace} is \texttt{GLOBAL}, \texttt{LOCAL}, or \texttt{REG}.
%% ============================================================
\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
\begin{tabular}{@{}l l l l@{}}
\toprule
\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
\midrule
\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in body with $a_k$. \\
\bottomrule
\end{tabular}
\smallskip
A value \op{Call} is void: its \op{Sink} body stores to \op{Param}s bound positionally to \op{Returned} arguments; output $a_k$ is \op{After}$(a_k, \op{Call})$.
%% ============================================================
\subsection*{{\color{movgreen}Movement Ops} \normalfont\small--- no arithmetic; view, indexing, and reinterpretation only}
@@ -95,20 +110,6 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
\bottomrule
\end{tabular}
%% ============================================================
\subsection*{{\color{callblue}Call Ops} \normalfont\small--- function abstraction, like the lambda calculus}
\begin{tabular}{@{}l l l l@{}}
\toprule
\textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
\midrule
\op{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
\op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
\bottomrule
\end{tabular}
%% ============================================================
\subsection*{{\color{loadred}Load Ops} \normalfont\small--- can change device or addrspace}
@@ -256,9 +257,10 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
\toprule
\textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
\midrule
\op{Buffer} & from arg & from $\mathrm{src}[0]$ & from arg & dtype range \\
\op{Buffer} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\
\op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
\op{Param} & from arg & from arg ($\mathrm{size}$) & from arg & from src or dtype range \\
\op{Returned} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\[3pt]
Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
\op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
\op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
@@ -272,7 +274,7 @@ ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$
Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
\op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
\op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
\op{Call} & void & --- & first non-null src device & --- \\
\op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
\op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
\op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
@@ -304,7 +306,6 @@ Each kernel's iteration space is a set of \op{Range} axes. Every range has an \t
{\color{axblue}\texttt{GLOBAL}} & \texttt{g} & --- & --- & GPU global workgroup dimension. \\
{\color{axcyan}\texttt{LOCAL}} & \texttt{l} & g, L & inner & Workgroup local dimension (shared memory). \\
{\color{axbrcyan}\texttt{WARP}} & \texttt{w} & \multicolumn{2}{l}{(created by \op{TC})} & Warp-level lanes for tensor cores. \\
{\color{axbrblue}\texttt{THREAD}} & \texttt{t} & g & outer & CPU thread parallelism. \\
{\color{axwhite}\texttt{LOOP}} & \texttt{L} & --- & --- & Generic sequential loop (initial state). \\
{\color{axred}\texttt{REDUCE}} & \texttt{R} & --- & --- & Reduction axis. \\
{\color{axbrred}\texttt{GROUP\_REDUCE}} & \texttt{G} & R & inner/outer & Shared-memory group reduction. \\
@@ -327,8 +328,6 @@ An optimization is a triple $(\mathrm{op},\;\mathrm{axis},\;\mathrm{arg})$:
Pad axis to next multiple of $m$ with validity masks. \\[4pt]
\op{Swap} & axis$_i$ & axis$_j$ &
Swap two axes $i \leftrightarrow j$. \\
\op{Nolocals} & --- & --- &
Disable local memory; no workgroup dims emitted. \\
\op{TC} & reduce idx & (tc, opt, mode) &
Apply tensor core \op{Wmma}: split reduce/output axes into \texttt{WARP}, \texttt{UPCAST}, and \texttt{UNROLL} dims. \\
\bottomrule
@@ -421,7 +420,7 @@ def allreduce(T):
%% ============================================================
\subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}
The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Call} node.
\begin{lstlisting}
@function
@@ -433,14 +432,15 @@ When \texttt{f(x, y)} is called, the decorator:
\begin{enumerate}[leftmargin=1.5em, itemsep=2pt]
\item \textbf{Extracts inputs}: walks all arguments to find every Tensor, deduplicates by identity.
\item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
\item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
\item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
\op{Function}(\op{Tuple}(body), $x$, $y$).
\item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
\item \textbf{Runs the function} lazily (no device execution), building a UOp graph from each returned value.
\item \textbf{Parameterizes inputs}: replaces each input UOp with a positional \op{Param}$(k)$ placeholder.
\item \textbf{Parameterizes outputs}: for each returned value $v_i$, creates an output \op{Param}$(m+i)$ and a matching \op{Returned}$(m+i)$, where $m$ is the number of inputs.
\item \textbf{Builds the call}: stores every $v_i$ into its output parameter and creates\\
\op{Call}(\op{Sink}(\op{Store}(\op{Param}$(m)$, $v_0$), \ldots), $x$, $y$, \op{Returned}$(m)$, \ldots).
\item \textbf{Returns values}: exposes each result as \op{After}(\op{Returned}$(m+i)$, \op{Call}).
\end{enumerate}
The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers, and a single call can return any number of values. At schedule time, an ordinary value-producing \op{Call} is inlined by positional \op{Param} substitution and each output \op{After} resolves to the value stored in the body. A precompiled call instead materializes real output buffers in the \op{Returned} slots and lowers the body to an opaque call that writes them.
%% ============================================================
\subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
tinygrad_repo/test/amd/hw/test_sop.py
+33
View File
@@ -1002,6 +1002,39 @@ class TestBarrier(unittest.TestCase):
for tid in range(64):
self.assertEqual(st.vgpr[tid][0], tid + 100 + 1000, f"tid={tid}")
class TestSMaxMinSCCRegressions(unittest.TestCase):
"""Regression test: S_MAX sets SCC only on strict inequality (equal operands -> SCC=0)."""
def test_s_max_i32_equal_scc(self):
st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
self.assertEqual(st.scc, 0)
self.assertEqual(st.sgpr[6], 64)
st = run_program([s_mov_b32(s[4], 65), s_mov_b32(s[5], 64), s_max_i32(s[6], s[4], s[5])], n_lanes=1)
self.assertEqual(st.scc, 1) # still set when strictly greater
def test_s_max_u32_equal_scc(self):
st = run_program([s_mov_b32(s[4], 64), s_mov_b32(s[5], 64), s_max_u32(s[6], s[4], s[5])], n_lanes=1)
self.assertEqual(st.scc, 0)
class TestAbsdiffOverflowRegressions(unittest.TestCase):
"""Regression test: S_ABSDIFF_I32 computes abs on the WRAPPED 32-bit difference (found by random difftest vs hardware)."""
def test_s_absdiff_wrapped(self):
# |45 - (-2147483647)| overflows int32; hardware takes abs of the wrapped 32-bit difference
instructions = [s_mov_b32(s[4], 45), s_mov_b32(s[5], 0x80000001), s_absdiff_i32(s[6], s[4], s[5])]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0x7FFFFFD4)
self.assertEqual(st.scc, 1)
# INT_MIN - 1 wraps to +2147483647, already positive
instructions = [s_mov_b32(s[4], 0x80000000), s_mov_b32(s[5], 1), s_absdiff_i32(s[6], s[4], s[5])]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0x7FFFFFFF)
# equality -> 0 and SCC=0
instructions = [s_mov_b32(s[4], 7), s_mov_b32(s[5], 7), s_absdiff_i32(s[6], s[4], s[5])]
st = run_program(instructions, n_lanes=1)
self.assertEqual(st.sgpr[6], 0)
self.assertEqual(st.scc, 0)
if __name__ == '__main__':
unittest.main()
tinygrad_repo/test/amd/hw/test_vop1.py
+61
View File
@@ -1629,5 +1629,66 @@ class TestSwap(unittest.TestCase):
self.assertEqual(st.vgpr[0][1], 0x55555555)
class TestCvtFrexpRegressions(unittest.TestCase):
"""Regression tests for float<->int conversion and FREXP corner cases (found by random difftest vs hardware)."""
def test_cvt_i32_f32_nan_is_zero(self):
"""v_cvt_i32_f32 of NaN is 0, not INT_MIN (x86 cvttss2si returns INT_MIN)."""
for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
def test_cvt_i32_f32_positive_overflow(self):
"""v_cvt_i32_f32 saturates positive overflow/inf to INT_MAX, not INT_MIN."""
for bits in (0x7F800000, 0x4F000000, 0x4F800000): # +inf, 2^31, ~2^32
st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x7FFFFFFF, f"bits=0x{bits:08x}")
def test_cvt_i32_f32_negative_overflow(self):
"""v_cvt_i32_f32 saturates negative overflow/-inf to INT_MIN."""
for bits in (0xFF800000, 0xCF000001): # -inf, below -2^31
st = run_program([v_mov_b32_e32(v[0], bits), v_cvt_i32_f32_e32(v[1], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0x80000000, f"bits=0x{bits:08x}")
def test_cvt_u32_f32_nan_is_zero(self):
"""v_cvt_u32_f32 of NaN is 0, not UINT_MAX."""
for nan in (0x7FC00000, 0xFFC00000, 0x7F800001):
st = run_program([v_mov_b32_e32(v[0], nan), v_cvt_u32_f32_e32(v[1], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1], 0, f"nan=0x{nan:08x}")
def test_cvt_i32_f64_nan_and_overflow(self):
"""v_cvt_i32_f64: NaN -> 0, positive overflow/+inf -> INT_MAX."""
st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x7FF80000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0)
st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0x41F00000), v_cvt_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x7FFFFFFF) # 2^32 -> INT_MAX
def test_frexp_f32_denormal(self):
"""v_frexp_exp/mant_f32 of denormal/zero inputs is (0, signed zero) on hardware."""
for bits in (0x00000001, 0x007FFFFF, 0x00000000):
st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0]), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"exp bits=0x{bits:08x}")
self.assertEqual(st.vgpr[0][2], bits & 0x80000000, f"mant bits=0x{bits:08x}")
# negative denormal: mant is -0.0
st = run_program([v_mov_b32_e32(v[0], 0x80000001), v_frexp_mant_f32_e32(v[2], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0x80000000)
def test_frexp_f64_denormal(self):
"""v_frexp_exp_f64 of a denormal returns the normalized exponent (-1073 for min-denormal); zero -> 0."""
st = run_program([v_mov_b32_e32(v[0], 1), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0xFFFFFBCF) # -1073
st = run_program([v_mov_b32_e32(v[0], 0), v_mov_b32_e32(v[1], 0), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
self.assertEqual(st.vgpr[0][2], 0)
def test_frexp_exp_inf_nan(self):
"""v_frexp_exp of +/-inf and NaN is 0 on hardware (host frexp gives 129/1024), for both f32 and f64."""
for bits in (0x7F800000, 0xFF800000, 0x7FC00000):
st = run_program([v_mov_b32_e32(v[0], bits), v_frexp_exp_i32_f32_e32(v[1], v[0])], n_lanes=1)
self.assertEqual(st.vgpr[0][1] & 0xFFFFFFFF, 0, f"f32 bits=0x{bits:08x}")
for lo, hi in ((0, 0x7FF00000), (0, 0xFFF00000), (0, 0x7FF80000), (1, 0x7FF00000)):
st = run_program([v_mov_b32_e32(v[0], lo), v_mov_b32_e32(v[1], hi), v_frexp_exp_i32_f64_e32(v[2], v[0:1])], n_lanes=1)
self.assertEqual(st.vgpr[0][2] & 0xFFFFFFFF, 0, f"f64 bits=0x{hi:08x}{lo:08x}")
if __name__ == '__main__':
unittest.main()

Some files were not shown because too many files have changed in this diff Show More